// Schorr-Waite algorithm in Chalice
// (see Test/dafny1/SchorrWaite.dfy)

class Node {
  var children: seq<Node>;
  var marked: bool;
  ghost var path: seq<Node>;

  var parent: Node;
}

class DSW0 {
  var stack: seq<Node>;
  var S: seq<Node>;
  var root: Node;

  function Reachable(to: Node, p:seq<Node>, from: Node): bool
    requires acc(p[*].children);    
  {
    |p| == 0 ? to == from :
      (p[0] != null && to in p[0].children && Reachable(p[0], p[1..], from))
  }

  method IterativeMark()
    requires acc(this.*) && acc(S[*].children) && acc(S[*].marked) && acc(S[*].path);
    requires root in S;
    requires forall n in S :: n != null && (forall ch in n.children :: ch != null && ch in S); 
    requires forall n in S :: ! n.marked;
    ensures acc(this.*) && acc(S[*].children) && acc(S[*].marked) && acc(S[*].path);
    // graph structure is the same
    ensures S == old(S) && root == old(root);
    ensures forall n in S :: n.children == old(n.children);
    // all nodes reachable from root are marked
    ensures root.marked;
    ensures forall n in S :: n.marked ==> (forall ch in n.children :: ch.marked);      
    // all marked nodes are reachable from root
    ensures forall n in S :: n.marked ==> 
      (forall m in n.path :: m in S) && Reachable(n, n.path, root);
  {
    var t:Node := root;
    t.marked := true;
    stack := nil<Node>;
    t.path := stack;

    // no termination check
    var stop := false;
    while(!stop)
      invariant acc(this.*) && acc(S[*].children) && acc(S[*].marked) && acc(S[*].path);
      invariant root == old(root);
      invariant S == old(S);
      invariant root.marked;
      invariant t in S && t.marked && t !in stack;
      // stack well-formed
      invariant forall i in [0..|stack|] :: forall j in [i+1..|stack|] :: stack[i] != stack[j];    
      invariant forall n in stack :: n in S && n.marked;
      invariant forall i in [1..|stack|] :: stack[i-1] in stack[i].children;
      invariant 0 < |stack| ==> t in stack[0].children;
      // goal
      invariant forall n in S :: n.marked && n !in stack && n != t ==> 
        (forall ch in n.children :: ch in S && ch.marked);
      invariant forall n in S :: n.marked ==> 
        (forall m in n.path :: m in S) && Reachable(n, n.path, root);
      // preservation
      invariant forall n in S :: n.children == old(n.children);      
      // termination
      invariant stop ==> |stack| == 0 && (forall ch in t.children :: ch.marked);
    {
      call n := PickUnmarked(t.children); 
      if (n != null) {
        // push             
        stack := [t] ++ stack;          
        n.path := [t] ++ t.path;
        t := n;
        t.marked := true;
        assert Reachable(t.path[0], t.path[1..], root); // needed for limited function
      } else {       
        // pop
        if (|stack| == 0) {
          stop := true;
        } else {
          t := stack[0];
          stack := stack[1..];
        }
      }
    }
  }

  method PickUnmarked(p: seq<Node>) returns (x: Node)
    requires rd(p[*].marked);
    requires forall n in p :: n != null;
    ensures rd(p[*].marked);
    ensures x != null ==> x in p && ! x.marked;
    ensures x == null ==> (forall n in p :: n.marked);
  {
    var x [(exists n in p :: !n.marked) ? (x in p && !x.marked) : (x == null)]
  }
}

class DSW1 refines DSW0 {
  replaces stack by acc(S[*].parent);

  transforms IterativeMark() 
  {
    *
  }

  transforms PickUnmarked(p: seq<Node>) returns (x: Node)
  {
    replaces x by {
      if (|p| == 0) {
        x := null;
      } else if (! p[0].marked) {
        x := p[0];
      } else {
        call x := PickUnmarked(p[1..]);
      }
    }
  }
}