// Derived from an example in the Boogie distribution

class IntSet {
  ghost var Contents: set<int>;
  ghost var Repr: set<object>;

  var root: TreeNode;

  predicate Valid
    reads this, Repr;
  {
    this in Repr &&
    (root == null ==> Contents == {}) &&
    (root != null ==>
       root in Repr && root.Repr <= Repr && this !in root.Repr &&
       root.Valid &&
       Contents == root.Contents)
  }

  constructor Init()
    modifies this;
    ensures Valid && fresh(Repr - {this});
    ensures Contents == {};
  {
    root := null;
    Repr := {this};
    Contents := {};
  }

  method Find(x: int) returns (present: bool)
    requires Valid;
    ensures present <==> x in Contents;
  {
    if (root == null) {
      present := false;
    } else {
      present := root.Find(x);
    }
  }

  method Insert(x: int)
    requires Valid;
    modifies Repr;
    ensures Valid && fresh(Repr - old(Repr));
    ensures Contents == old(Contents) + {x};
  { 
    if (root == null) {
        root := new TreeNode.Init(x);
    } else {
        root.Insert(x);
    }
    Contents := root.Contents;
    Repr := root.Repr + {this};
  }

  method Remove(x: int)
    requires Valid;
    modifies Repr;
    ensures Valid && fresh(Repr - old(Repr));
    ensures Contents == old(Contents) - {x};
  {
    if (root != null) {
      var newRoot := root.Remove(x);
      root := newRoot;
      if (root == null) {
        Contents := {};
        Repr := {this};
      } else {
        Contents := root.Contents;
        Repr := root.Repr + {this};
      }
    }
  }
}

class TreeNode {
  ghost var Contents: set<int>;
  ghost var Repr: set<object>;

  var data: int;
  var left: TreeNode;
  var right: TreeNode;

  predicate Valid
    reads this, Repr;
  {
    this in Repr &&
    null !in Repr &&
    (left != null ==>
      left in Repr &&
      left.Repr <= Repr && this !in left.Repr &&
      left.Valid &&
      (forall y :: y in left.Contents ==> y < data)) &&
    (right != null ==>
      right in Repr &&
      right.Repr <= Repr && this !in right.Repr &&
      right.Valid &&
      (forall y :: y in right.Contents ==> data < y)) &&
    (left != null && right != null ==> left.Repr !! right.Repr) &&
    Contents == (if left == null then {} else left.Contents) +
                (if right == null then {} else right.Contents) +
                {data}
  }

  constructor Init(x: int)
    modifies this;
    ensures Valid && fresh(Repr - {this});
    ensures Contents == {x};
  {
    data := x;
    left := null;
    right := null;
    Contents := {x};
    Repr := {this};
  }

  method Insert(x: int)
    requires Valid;
    modifies Repr;
    ensures Valid;
    ensures Contents == old(Contents) + {x};
    ensures fresh(Repr - old(Repr));
    decreases Repr;
  {
    if (x == data) {
        return;
    } 
    if (x < data) {
        if (left == null) {
            left := new TreeNode.Init(x);
        } else {
            left.Insert(x);
        }
        Repr := Repr + left.Repr;
    } else {
        if (right == null) {
            right := new TreeNode.Init(x);
        } else {
            right.Insert(x);
        }
        Repr := Repr + right.Repr;
    }
    Contents := Contents + {x};
  }

  method Find(x: int) returns (present: bool)
    requires Valid;
    ensures present <==> x in Contents;
    decreases Repr;
  {
    if (x == data) {
      present := true;
    } else if (left != null && x < data) {
      present := left.Find(x);
    } else if (right != null && data < x) {
      present := right.Find(x);
    } else {
      present := false;
    }
  }

  method Remove(x: int) returns (TreeNode: TreeNode)
    requires Valid;
    modifies Repr;
    ensures fresh(Repr - old(Repr));
    ensures TreeNode != null ==> TreeNode.Valid;
    ensures TreeNode == null ==> old(Contents) <= {x};
    ensures TreeNode != null ==> TreeNode.Repr <= Repr && TreeNode.Contents == old(Contents) - {x};
    decreases Repr;
  {
    TreeNode := this;
    if (left != null && x < data) {
      var t := left.Remove(x);
      left := t;
      Contents := Contents - {x};
      if (left != null) { Repr := Repr + left.Repr; }
    } else if (right != null && data < x) {
      var t := right.Remove(x);
      right := t;
      Contents := Contents - {x};
      if (right != null) { Repr := Repr + right.Repr; }
    } else if (x == data) {
      if (left == null && right == null) {
        TreeNode := null;
      } else if (left == null) {
        TreeNode := right;
      } else if (right == null) {
        TreeNode := left;
      } else {
        // rotate
        var min, r := right.RemoveMin();
        data := min;  right := r;
        Contents := Contents - {x};
        if (right != null) { Repr := Repr + right.Repr; }
      }
    }
  }

  method RemoveMin() returns (min: int, TreeNode: TreeNode)
    requires Valid;
    modifies Repr;
    ensures fresh(Repr - old(Repr));
    ensures TreeNode != null ==> TreeNode.Valid;
    ensures TreeNode == null ==> old(Contents) == {min};
    ensures TreeNode != null ==> TreeNode.Repr <= Repr && TreeNode.Contents == old(Contents) - {min};
    ensures min in old(Contents) && (forall x :: x in old(Contents) ==> min <= x);
    decreases Repr;
  {
    if (left == null) {
      min := data;
      TreeNode := right;
    } else {
      var t;
      min, t := left.RemoveMin();
      left := t;
      TreeNode := this;
      Contents := Contents - {min};
      if (left != null) { Repr := Repr + left.Repr; }
    }
  }
}

class Main {
  method Client0(x: int)
  {
    var s := new IntSet.Init();

    s.Insert(12);
    s.Insert(24);
    var present := s.Find(x);
    assert present <==> x == 12 || x == 24;
  }

  method Client1(s: IntSet, x: int)
    requires s != null && s.Valid;
    modifies s.Repr;
  {
    s.Insert(x);
    s.Insert(24);
    assert old(s.Contents) - {x,24} == s.Contents - {x,24};
  }
}