val W:ℕ; // maximum weight
val N:ℕ; // maximum node number

type Weight = ℕ[W];
type Node = ℕ[N];
type Edge = Set[Node] with |value| = 2;
type Graph = Map[Edge,Weight];
type Seq = Array[N+2,Node];
type Index = ℕ[N+2];

// a "walk" s of length l ≥ 1 from node n1 to node n2 in edge set E
// is a sequence of nodes that starts with n1 and ends with n2 and
// in which every successive pair of nodes is connected by an edge in E
pred walk(s:Seq,l:Index,n1:Node,n2:Node,E:Set[Edge]) ⇔
  ⊤
;
  
// a "path" s is a walk in which all nodes are distinct
pred path(s:Seq,l:Index,n1:Node,n2:Node,E:Set[Edge]) ⇔
  ⊤
;

// a "trail" s is a walk in which all edges 
// connecting two successive nodes are distinct
pred trail(s:Seq,l:Index,n1:Node,n2:Node,E:Set[Edge]) ⇔
  ⊤
;

// edge set E is "connected" if
// for every pair of nodes n1,n2 there exists a path from n1 to n2
pred connected(E:Set[Edge]) ⇔
  ⊤
;
    
// edge set E is "acyclic" if
// there does not exist a trail of length 2 or more from a node to itself
pred acyclic(E:Set[Edge]) ⇔
  ⊤
;

// edge set E is "consistent" with graph G if 
// E contains only edges to which G assigns non-zero weight
pred consistent(E:Set[Edge],G:Graph) ⇔ ∀e∈E. G[e] ≠ 0;

// edge set E is a "spanning tree" for G if
// E is consistent with G, connected, and acyclic
pred stree(E:Set[Edge],G:Graph) ⇔ consistent(E,G) ∧ connected(E) ∧ acyclic(E);
 
// the "weigth" of an edge set E for graph G is the sum of the weights 
// that G assigns to the edges of E
fun weight(E:Set[Edge],G:Graph):ℕ[(N+1)⋅W] = ∑e∈E. G[e];

// edge set E is a "minimum spanning tree" for G if
// E is a spanning tree for G of minimal weight
pred minstree(E:Set[Edge],G:Graph) ⇔ 
  ⊤
;

// the "edge set" of a graph G is the set of edges to which G 
// assigns non-zero weight
fun edges(G:Graph):Set[Edge] = { e | e:Edge with G[e] ≠ 0 };
   
// Kruskal's algorithm: given a connected graph G, it computes a
// minimum spanning tree E for G (even more, if G is not connected,
// then E is a "minimum spanning forest" for G)
proc kruskal(G:Graph): Set[Edge]
  requires connected(edges(G));
  ensures minstree(result,G);
{
  var S: Set[Edge] = edges(G);
  var T: Set[Set[Node]] = { { v } | v: Node };
  var E: Set[Edge] = ∅[Edge];
  choose e∈S with ∀e0∈S. G[e] ≤ G[e0] do
  {
    S ≔ S\{e}; 
    choose n1∈e, n2∈e with n1 ≠ n2;
    choose t1∈T, t2∈T with n1 ∈ t1 ∧ n2 ∈ t2;
    if t1 ≠ t2 then 
    {
      E ≔ E ∪ {e};
      T ≔ (T \ { t1, t2 }) ∪ { t1 ∪ t2 };
    }
  }
  return E;
}