// 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 {
  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(root: Node, S: seq<Node>)
    requires acc(S[*].*);
    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(S[*].*);
    // graph structure is the same
    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;
    var stack: seq<Node> := [];
    t.path := stack;

    // no termination check
    var stop := false;
    while(!stop)
      invariant acc(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);
    {
      var n: Node;

      // stupid syntactic trick to make transformers go around call
      if (true) {
        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 {
  transforms IterativeMark(root: Node, S: seq<Node>) 
  {
    assume forall n in S :: n.parent == null;
    _
    while 
      invariant forall n in S :: n !in stack ==> n.parent == null;
      invariant forall i in [1..|stack|] :: stack[i-1].parent == stack[i];
    {
      _
      if {*}
      if {
        // push
        if (|stack| > 0) {
          t.parent := stack[0];
        }
        *
      } else {
        // pop
        if {*} else {
          _
          t.parent := null;
        }
      }
    }
  }

  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..]);
      }
    }
  }
}