; A nasty, imperative implementation of Prim's algorithm... in egglog!
; https://en.wikipedia.org/wiki/Prim%27s_algorithm

; Weighted edge (vertex 1 * vertex 2 * weight)
(datatype edge (Edge i64 i64 i64))
(relation edge-exists (edge))

(relation mytrue ())
(mytrue)
(let infinity 99999999)  ; close enough

; ==== PROBLEM INSTANCES ====

; Graph 1
; (1)--2--(2)
;    \     |
;      1   2
;        \ |
; (3)--3--(4)
(ruleset graph1)
(rule ((mytrue))
      ((edge-exists (Edge 1 2 2))
       (edge-exists (Edge 1 4 1))
       (edge-exists (Edge 2 4 2))
       (edge-exists (Edge 3 4 3)))
       :ruleset graph1)

; Graph 2
; (1)-2-(2)  (3)
;  |\   /|   / |
;  | 3 5 |  4  |
;  5  X  2 /   5
;  | / \ |/    |
; (4)-4-(5)-7-(6)
(ruleset graph2)
(rule ((mytrue))
      ((edge-exists (Edge 1 2 1))
       (edge-exists (Edge 1 4 5))
       (edge-exists (Edge 1 5 3))
       (edge-exists (Edge 2 4 5))
       (edge-exists (Edge 2 5 2))
       (edge-exists (Edge 3 5 4))
       (edge-exists (Edge 3 6 5))
       (edge-exists (Edge 4 5 4))
       (edge-exists (Edge 5 6 7)))
       :ruleset graph2)

; ==== "INIT" RULESET ====

(ruleset init)

; Graph is undirected
(rule ((= e (Edge x y weight)))
      ((union e (Edge y x weight)))
      :ruleset init)

; Whether a vertex is included *so far* (this changes). Returns 0 or 1.
(function vertex-included (i64) i64 :merge (max old new))

; All vertices default to being not included (note vertex-included's :merge)
(rule ((edge-exists (Edge x y weight)))
      ((set (vertex-included x) 0))
      :ruleset init)

; Keep track of the current iteration
(function current-iteration () i64 :merge (max old new))

; Map iteration to best edge found so far
(function iteration-to-best-edge (i64) edge :merge new)
(function iteration-to-best-edge-weight (i64) i64 :merge new)

(rule ((mytrue))
      ((set (vertex-included 1) 1)  ; Initially just include vertex 1
       (set (current-iteration) 0)
       (set (iteration-to-best-edge-weight 0) infinity))
      :ruleset init)

; === "CHOOSE BEST EDGE" RULESET ===

(relation edge-in-mst (edge))  ; whether an edge is in our solution

(ruleset choose-best-edge)
(rule ((= i (current-iteration))
       (edge-exists (Edge x y weight))
       (= 1 (vertex-included x))
       (= 0 (vertex-included y))
       (< weight (iteration-to-best-edge-weight i)))
      ((set (iteration-to-best-edge-weight i) weight)
       (set (iteration-to-best-edge i) (Edge x y weight)))
      :ruleset choose-best-edge)

; === "FINISH ITERATION" RULESET ===

(ruleset finish-iteration)
(rule ((= i (current-iteration))
       (= (Edge x y weight) (iteration-to-best-edge i)))
      ((edge-in-mst (Edge x y weight))    ; incorporate chosen best edge
       (set (vertex-included x) 1)        ; mark its vertices as included
       (set (vertex-included y) 1)
       (set (current-iteration) (+ i 1))  ; advance iteration
       (set (iteration-to-best-edge-weight (+ i 1)) infinity))
      :ruleset finish-iteration)

; === RUN VIA SCHEDULE ===

(run-schedule
    (saturate init graph1)  ; change to graph2 to see other example
    (saturate (saturate choose-best-edge) finish-iteration)
)

; === PRINT RESULTS ===

; (print-function edge-in-mst) ; this is not very helpful

; Just copy canonical edges to solution
(relation solution (i64 i64 i64))

(ruleset finalize)
(rule ((edge-in-mst (Edge x y weight)) (< x y))
      ((solution x y weight))
      :ruleset finalize)
(run-schedule (saturate finalize))

(print-function solution 100) ; this is better