1 files changed, 81 insertions, 0 deletions
diff --git a/AvlVerification/Tree.lean b/AvlVerification/Tree.lean
new file mode 100644
index 0000000..3316e9e
--- /dev/null
+++ b/AvlVerification/Tree.lean
@@ -0,0 +1,81 @@
+import «AvlVerification»
+
+namespace Tree
+
+variable {T: Type}
+
+open avl_verification
+
+-- Otherwise, Lean cannot prove termination by itself.
+@[reducible]
+def AVLTree (U: Type) := Option (AVLNode U)
+def AVLTree.nil: AVLTree T := none
+
+def AVLTree.val (t: AVLTree T): Option T := match t with
+| none => none
+| some (AVLNode.mk x _ _) => some x
+
+def AVLTree.left (t: AVLTree T): AVLTree T := match t with
+| none => none
+| some (AVLNode.mk _ left _) => left
+
+def AVLTree.right (t: AVLTree T): AVLTree T := match t with
+| none => none
+| some (AVLNode.mk _ _ right) => right
+
+def AVLNode.left (t: AVLNode T): AVLTree T := match t with
+| AVLNode.mk _ left _ => left
+
+def AVLNode.right (t: AVLNode T): AVLTree T := match t with
+| AVLNode.mk _ _ right => right
+
+def AVLNode.val (t: AVLNode T): T := match t with
+| AVLNode.mk x _ _ => x
+
+mutual
+def AVLTree.height_node (tree: AVLNode T): Nat := match tree with
+| AVLNode.mk y left right => 1 + AVLTree.height left + AVLTree.height right
+
+def AVLTree.height (tree: AVLTree T): Nat := match tree with 
+| none => 0
+| some n => 1 + AVLTree.height_node n
+end
+
+def AVLTreeSet.nil: AVLTreeSet T := { root := AVLTree.nil }
+
+-- Automatic synthesization of this seems possible at the Lean level?
+instance: Inhabited (AVLTree T) where
+  default := AVLTree.nil
+
+instance: Inhabited (AVLTreeSet T) where
+  default := AVLTreeSet.nil
+
+instance: Coe (Option (AVLNode T)) (AVLTree T) where
+  coe x := x
+
+-- TODO: ideally, it would be nice if we could generalize
+-- this to any `BinaryTree` typeclass.
+
+def AVLTree.mem (tree: AVLTree T) (x: T) :=
+  match tree with
+  | none => False
+  | some (AVLNode.mk y left right) => x = y ∨ AVLTree.mem left x ∨ AVLTree.mem right x
+
+@[simp]
+def AVLTree.mem_none: AVLTree.mem none = ({}: Set T) := rfl
+
+@[simp]
+def AVLTree.mem_some {x: T} {left right: AVLTree T}: AVLTree.mem (some (AVLNode.mk x left right)) = (({x}: Set T) ∪ AVLTree.mem left ∪ AVLTree.mem right) := by
+  ext y
+  rw [AVLTree.mem]
+  simp [Set.insert_union]
+  simp [Set.insert_def, Set.setOf_set, Set.mem_def, Set.union_def]
+
+-- TODO(reinforcement): ∪ is actually disjoint if we prove this is a binary search tree.
+
+def AVLTree.set (t: AVLTree T): Set T := _root_.setOf (AVLTree.mem t)
+
+def AVLTree.setOf_left_incl (t: AVLTree T): t.left.set ⊂ t.set := sorry
+def AVLTree.setOf_right_incl (t: AVLTree T): t.right.set ⊂ t.set := sorry
+
+end Tree