class TreeNode 
{
    var data: int
    var left: TreeNode?
    var right: TreeNode?

    ghost var Repr: set<object>

    ghost var Contents: set<int>

    predicate Valid()
        reads this, Repr
    {
        this in Repr && 
        (left != null ==> 
            left in Repr &&
            left.Repr <= Repr && this !in left.Repr &&
            (forall e :: e in left.Contents ==> e < data) &&
            left.Valid()) &&
        (right != null ==> 
            right in Repr &&
            right.Repr <= Repr && this !in right.Repr &&
            (forall e :: e in right.Contents ==> data < e) &&
            right.Valid()) 
        && (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 (x: int)
        ensures fresh(Repr - {this})
        ensures Valid()
        ensures Contents == {x}
    {
        data := x;
        left := null;
        right := null;

        Repr := { this };
        Contents := { data };
    }

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

    }

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

    method Main()
    {
        var t1 := new TreeNode(1);
        var t2 := new TreeNode(2);
        t1.Insert(2);
        assert t1.Valid();
        t2.Insert(3);
        assert t1.Valid();
        var b := t1.Find(2);
        assert b;
        t1 := t1.Remove(2);
        b := t1.Find(2);
        assert !b;
    }

}