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

// (incomplete version) Two children instead of a sequence of an array
class Node {
  var left: Node;
  var right: Node;
  var marked: bool;
  var l: bool;
  var r: bool; 
}

class Main {
  method RecursiveMark(root: Node, S: seq<Node>)
    requires acc(S[*].marked, 50) && acc(S[*].*, 50);
    requires root != null && root in S;
    // S is closed under 'left' and 'right':
    requires forall n in S :: n != null && 
      ((n.left != null ==> n.left in S) && 
       (n.right != null ==> n.right in S));
    requires forall n in S :: ! n.marked;
    ensures acc(S[*].marked, 50) && acc(S[*].*, 50);
    ensures root.marked;
    // nodes reachable from 'root' are marked:
    ensures forall n in S :: n.marked ==> 
      ((n.left != null ==> n.left in S && n.left.marked) &&
       (n.right != null ==> n.right in S && n.right.marked));
  {
    var stack: seq<Node> := [];
    call RecursiveMarkWorker(root, S, stack);
  }

  method RecursiveMarkWorker(root: Node, S: seq<Node>, stack: seq<Node>) 
    requires acc(S[*].marked, 50) && acc(S[*].*, 50);
    requires root != null && root in S;
    requires forall n in S :: n != null &&
      ((n.left != null ==> n.left in S) &&
       (n.right != null ==> n.right in S))
    requires forall n in S :: n.marked ==> 
      (n in stack || 
       ((n.left != null ==> n.left.marked) &&
        (n.right != null ==> n.right.marked)));
    requires forall n in stack :: n != null && n in S && n.marked;
    ensures acc(S[*].marked, 50) && acc(S[*].*, 50);
    ensures forall n in S :: n.left == old(n.left) && n.right == old(n.right);
    ensures forall n in S :: n.marked ==> 
      (n in stack ||
       ((n.left != null ==> n.left.marked) &&
        (n.right != null ==> n.right.marked)));    
    ensures forall n in S :: old(n.marked) ==> n.marked;
    ensures root.marked;
  {
    if (! root.marked) {
      root.marked := true;
      var next:seq<Node> := [root] ++ stack;
      assert next[0] == root;
      if (root.left != null) {
        call RecursiveMarkWorker(root.left, S, next);
      }

      if (root.right != null) {
        call RecursiveMarkWorker(root.right, S, next);
      }
    } 
  }

  method IterativeMark(root: Node, S: seq<Node>)
    requires acc(S[*].*);
    requires root in S;
    requires forall n in S :: n != null && 
      (n.left != null ==> n.left in S) && 
      (n.right != null ==> n.right in S);
    requires forall n in S :: ! n.marked;
    requires forall n in S :: ! n.l && ! n.r;
    ensures acc(S[*].*);
    ensures forall n in S :: n.left == old(n.left) && n.right == old(n.right);
    ensures root.marked;
    ensures forall n in S :: n.marked ==> 
      (n.left != null ==> n.left.marked) && 
      (n.right != null ==> n.right.marked);
    ensures forall n in S :: ! n.l && ! n.r;
  {
    var t:Node := root;
    t.marked := true;
    var stack: seq<Node> := [];

    var stop := false;
    while(!stop)
      invariant acc(S[*].*);
      invariant root.marked && t in S && t !in stack;
      invariant forall n in stack :: n in S;
      invariant forall i in [0..|stack|] :: forall j in [i+1..|stack|] :: stack[i] != stack[j];    
      invariant forall n in S :: (n in stack || n == t) ==>
        n.marked &&
        (n.r ==> n.l) &&
        (n.l && n.left != null ==> n.left in S && n.left.marked) &&
        (n.r && n.right != null ==> n.right in S && n.right.marked)
      // stack is linked
      invariant forall i in [1..|stack|] :: stack[i-1] == (stack[i].l ? stack[i].right : stack[i].left);
      invariant 0 < |stack| ==> t == (stack[0].l ? stack[0].right : stack[0].left);
      // goal
      invariant forall n in S :: n.marked && n !in stack && n != t ==> 
        (n.left != null ==> n.left in S && n.left.marked) && 
        (n.right != null ==> n.right in S && n.right.marked);
      // preservation
      invariant forall n in S :: n !in stack && n != t ==> ! n.l && ! n.r;      
      invariant forall n in S :: n.left == old(n.left) && n.right == old(n.right);     
      invariant stop ==> |stack| == 0 && ! t.l && ! t.r && 
        (t.left != null ==> t.left.marked) && 
        (t.right != null ==> t.right.marked);
    {  
      if (! t.l && (t.left == null || t.left.marked)) {
        // advance
        t.l := true;         
      } else if (t.l && ! t.r && (t.right == null || t.right.marked)) {
        // advance
        t.r := true;
      } else if (t.r) {
        // pop
        t.l := false;
        t.r := false;
        if (|stack| == 0) {       
          stop := true;
        } else {
          t := stack[0];
          stack := stack[1..];
          if (t.l) {t.r := true} else {t.l := true}
        }
      } else if (!t.l) {
        // push
        stack := [t] ++ stack;
        assert stack[0] == t;
        t := t.left;
        t.marked := true;        
      } else if (!t.r) {
        // push
        assert t.l;
        stack := [t] ++ stack;
        assert stack[0] == t;
        t := t.right;
        t.marked := true;               
      }
    }
  }
}