feat: close `find` / `insert` proofs

After a complete 180 with the Order theory, we close the goals of find
and insert and we give an example of U32 order that we will upstream to
Aeneas directly.

Signed-off-by: Raito Bezarius <masterancpp@gmail.com>

6 files changed, 577 insertions, 0 deletions

diff --git a/Verification/BinarySearchTree.lean b/Verification/BinarySearchTree.lean
new file mode 100644
index 0000000..a49be5e
--- /dev/null
+++ b/Verification/BinarySearchTree.lean
@@ -0,0 +1,114 @@
+import Verification.Tree
+import AvlVerification
+
+namespace BST
+
+open Primitives (Result)
+open avl_verification (AVLNode Ordering)
+open Tree (AVLTree AVLNode.left AVLNode.right AVLNode.val)
+
+inductive ForallNode (p: T -> Prop): AVLTree T -> Prop
+| none : ForallNode p none
+| some (a: T) (left: AVLTree T) (right: AVLTree T) : ForallNode p left -> p a -> ForallNode p right -> ForallNode p (some (AVLNode.mk a left right))
+
+theorem ForallNode.left {p: T -> Prop} {t: AVLTree T}: ForallNode p t -> ForallNode p t.left := by
+  intro Hpt
+  cases Hpt with 
+  | none => simp [AVLTree.left, ForallNode.none]
+  | some a left right f_pleft f_pa f_pright => simp [AVLTree.left, f_pleft]
+
+theorem ForallNode.right {p: T -> Prop} {t: AVLTree T}: ForallNode p t -> ForallNode p t.right := by
+  intro Hpt
+  cases Hpt with 
+  | none => simp [AVLTree.right, ForallNode.none]
+  | some a left right f_pleft f_pa f_pright => simp [AVLTree.right, f_pright]
+
+theorem ForallNode.label {a: T} {p: T -> Prop} {left right: AVLTree T}: ForallNode p (AVLNode.mk a left right) -> p a := by
+  intro Hpt
+  cases Hpt with 
+  | some a left right f_pleft f_pa f_pright => exact f_pa
+
+theorem ForallNode.not_mem {a: T} (p: T -> Prop) (t: Option (AVLNode T)): ¬ p a -> ForallNode p t -> a ∉ AVLTree.set t := fun Hnpa Hpt => by
+  cases t with 
+  | none => simp [AVLTree.set]; tauto
+  | some t =>
+    cases Hpt with 
+    | some b left right f_pbleft f_pb f_pbright =>
+      simp [AVLTree.set_some]
+      push_neg
+      split_conjs
+      . by_contra hab; rw [hab] at Hnpa; exact Hnpa f_pb
+      . exact ForallNode.not_mem p left Hnpa f_pbleft
+      . exact ForallNode.not_mem p right Hnpa f_pbright
+
+theorem ForallNode.not_mem' {a: T} (p: T -> Prop) (t: Option (AVLNode T)): p a -> ForallNode (fun x =>  ¬p x) t -> a ∉ AVLTree.set t := fun Hpa Hnpt => by 
+  refine' ForallNode.not_mem (fun x => ¬ p x) t _ _
+  simp [Hpa]
+  exact Hnpt
+
+theorem ForallNode.imp {p q: T -> Prop} {t: AVLTree T}: (∀ x, p x -> q x) -> ForallNode p t -> ForallNode q t := fun Himp Hpt => by 
+  induction Hpt
+  . simp [ForallNode.none]
+  . constructor
+    . assumption
+    . apply Himp; assumption
+    . assumption
+
+-- This is the binary search invariant.
+variable [LinearOrder T]
+inductive Invariant: AVLTree T -> Prop
+| none : Invariant none
+| some (a: T) (left: AVLTree T) (right: AVLTree T) : 
+  ForallNode (fun v => v < a) left -> ForallNode (fun v => a < v) right
+  -> Invariant left -> Invariant right -> Invariant (some (AVLNode.mk a left right))
+
+@[simp]
+theorem singleton_bst {a: T}: Invariant (some (AVLNode.mk a none none)) := by
+  apply Invariant.some
+  all_goals simp [ForallNode.none, Invariant.none] 
+
+theorem left {t: AVLTree T}: Invariant t -> Invariant t.left := by
+  intro H
+  induction H with 
+  | none => exact Invariant.none
+  | some _ _ _ _ _ _ _ _ _ => simp [AVLTree.left]; assumption
+
+theorem right {t: AVLTree T}: Invariant t -> Invariant t.right := by
+  intro H
+  induction H with 
+  | none => exact Invariant.none
+  | some _ _ _ _ _ _ _ _ _ => simp [AVLTree.right]; assumption
+
+-- TODO: ask at most for LT + Irreflexive (lt_irrefl) + Trichotomy (le_of_not_lt)?
+theorem left_pos {left right: Option (AVLNode T)} {a x: T}: BST.Invariant (some (AVLNode.mk a left right)) -> x ∈ AVLTree.set (AVLNode.mk a left right) -> x < a -> x ∈ AVLTree.set left := fun Hbst Hmem Hxa => by
+   simp [AVLTree.set_some] at Hmem
+   rcases Hmem with (Heq | Hleft) | Hright
+   . rewrite [Heq] at Hxa; exact absurd Hxa (lt_irrefl _)
+   . assumption
+   . exfalso
+ 
+     -- Hbst -> x ∈ right -> ForallNode (fun v => ¬ v < a)
+     refine' ForallNode.not_mem' (fun v => v < a) right Hxa _ _
+     simp [le_of_not_lt]
+     cases Hbst with 
+     | some _ _ _ _ Hforall _ =>
+       refine' ForallNode.imp _ Hforall
+       exact fun x => le_of_lt
+     assumption
+ 
+theorem right_pos {left right: Option (AVLNode T)} {a x: T}: BST.Invariant (some (AVLNode.mk a left right)) -> x ∈ AVLTree.set (AVLNode.mk a left right) -> a < x -> x ∈ AVLTree.set right := fun Hbst Hmem Hax => by
+   simp [AVLTree.set_some] at Hmem
+   rcases Hmem with (Heq | Hleft) | Hright
+   . rewrite [Heq] at Hax; exact absurd Hax (lt_irrefl _)
+   . exfalso
+     refine' ForallNode.not_mem' (fun v => a < v) left Hax _ _
+     simp [le_of_not_lt]
+     cases Hbst with 
+     | some _ _ _ Hforall _ _ =>
+       refine' ForallNode.imp _ Hforall
+       exact fun x => le_of_lt
+     assumption  
+   . assumption
+ 
+
+end BST
diff --git a/Verification/Find.lean b/Verification/Find.lean
new file mode 100644
index 0000000..764a685
--- /dev/null
+++ b/Verification/Find.lean
@@ -0,0 +1,40 @@
+import Verification.Tree
+import Verification.BinarySearchTree
+import Verification.Specifications
+import AvlVerification
+
+namespace Implementation
+
+open Primitives
+open avl_verification
+open Tree (AVLTree AVLTree.set)
+open Specifications (OrdSpecLinearOrderEq infallible ltOfRustOrder gtOfRustOrder)
+
+variable (T: Type) (H: avl_verification.Ord T) [DecidableEq T] [LinearOrder T] (Ospec: OrdSpecLinearOrderEq H)
+
+@[pspec]
+def AVLTreeSet.find_loop_spec
+  (a: T) (t: Option (AVLNode T)):
+  BST.Invariant t -> a ∈ AVLTree.set t -> AVLTreeSet.find_loop _ H a t = Result.ok true := fun Hbst Hmem => by
+  match t with 
+  | none => trivial
+  | some (AVLNode.mk b left right) =>
+    rw [AVLTreeSet.find_loop]
+    dsimp only
+    have : ∀ a b, ∃ o, H.cmp a b = .ok o := infallible H
+    progress keep Hordering as ⟨ ordering ⟩
+    cases ordering
+    all_goals dsimp only
+    . refine' AVLTreeSet.find_loop_spec a right (BST.right Hbst) (BST.right_pos Hbst Hmem _)
+      exact ltOfRustOrder _ _ _ Hordering
+    . refine' AVLTreeSet.find_loop_spec a left (BST.left Hbst) (BST.left_pos Hbst Hmem _)
+      exact gtOfRustOrder _ _ _ Hordering
+
+def AVLTreeSet.find_spec
+  (a: T) (t: AVLTreeSet T):
+  BST.Invariant t.root -> a ∈ AVLTree.set t.root ->
+    t.find _ H a = Result.ok true := fun Hbst Hmem => by
+    rw [AVLTreeSet.find]; progress
+
+end Implementation
+
diff --git a/Verification/Insert.lean b/Verification/Insert.lean
new file mode 100644
index 0000000..260eaa1
--- /dev/null
+++ b/Verification/Insert.lean
@@ -0,0 +1,134 @@
+import Verification.Tree
+import Verification.BinarySearchTree
+import Verification.Specifications
+
+namespace Implementation
+
+open Primitives
+open avl_verification
+open Tree (AVLTree AVLTree.set)
+open Specifications (OrdSpecLinearOrderEq infallible ltOfRustOrder gtOfRustOrder)
+
+variable (T: Type) (H: avl_verification.Ord T) [LinearOrder T] (Ospec: OrdSpecLinearOrderEq H)
+
+@[pspec]
+theorem AVLTreeSet.insert_loop_spec_local (p: T -> Prop)
+  (Hcmp_spec: ∀ a b, ∃ o, H.cmp a b = Result.ok o)
+  (a: T) (t: Option (AVLNode T)):
+  ∃ added t_new, AVLTreeSet.insert_loop T H a t = Result.ok ⟨ added, t_new ⟩
+--  ∧ AVLTree.set t'.2 = insert a (AVLTree.set t)
+  ∧ (BST.ForallNode p t -> p a -> BST.ForallNode p t_new)
+  := by match t with
+  | none =>
+    simp [AVLTreeSet.insert_loop, AVLTree.set, setOf]
+    intros _ Hpa
+    constructor; all_goals try simp [BST.ForallNode.none]
+    exact Hpa
+  | some (AVLNode.mk b left right) =>
+    rw [AVLTreeSet.insert_loop]
+    simp only []
+    progress keep Hordering as ⟨ ordering ⟩
+    cases ordering
+    all_goals simp only []
+    {
+      progress keep Htree as ⟨ added₁, right₁, Hnode ⟩
+      refine' ⟨ added₁, ⟨ some (AVLNode.mk b left right₁), _ ⟩ ⟩
+      simp only [true_and]
+      intros Hptree Hpa
+      constructor
+      exact Hptree.left
+      exact Hptree.label
+      exact Hnode Hptree.right Hpa
+    }
+    {
+      simp; tauto
+    }
+    {
+      -- TODO: investigate wlog.
+      -- Symmetric case of left.
+      progress keep Htree as ⟨ added₁, left₁, Hnode ⟩
+      refine' ⟨ added₁, ⟨ some (AVLNode.mk b left₁ right), _ ⟩ ⟩
+      simp only [true_and]
+      intros Hptree Hpa
+      constructor
+      exact Hnode Hptree.left Hpa
+      exact Hptree.label
+      exact Hptree.right
+    }
+
+@[pspec]
+lemma AVLTreeSet.insert_loop_spec_global
+  (a: T) (t: Option (AVLNode T))
+  :
+  BST.Invariant t -> ∃ added t_new, AVLTreeSet.insert_loop T H a t = Result.ok ⟨ added, t_new ⟩
+  ∧ BST.Invariant t_new ∧ AVLTree.set t_new = {a} ∪ AVLTree.set t := by 
+  intro Hbst
+  match t with
+  | none => simp [AVLTreeSet.insert_loop, AVLTree.set, setOf]
+  | some (AVLNode.mk b left right) =>
+    rw [AVLTreeSet.insert_loop]
+    simp only []
+    have : ∀ a b, ∃ o, H.cmp a b = .ok o := infallible H
+    progress keep Hordering as ⟨ ordering ⟩
+    cases ordering
+    all_goals simp only []
+    {
+      have ⟨ added₂, right₂, ⟨ H_result, ⟨ H_bst, H_set ⟩ ⟩ ⟩ := AVLTreeSet.insert_loop_spec_global a right (BST.right Hbst)
+      progress keep Htree with AVLTreeSet.insert_loop_spec_local as ⟨ added₁, right₁, Hnode ⟩
+      exact (fun x => b < x)
+      rewrite [Htree] at H_result; simp at H_result
+      refine' ⟨ added₁, ⟨ some (AVLNode.mk b left right₁), _ ⟩ ⟩
+      simp only [true_and]
+      split_conjs
+      cases' Hbst with _ _ _ H_forall_left H_forall_right H_bst_left H_bst_right
+      constructor
+      exact H_forall_left
+      apply Hnode; exact H_forall_right
+      exact (ltOfRustOrder H b a Hordering)
+      exact H_bst_left
+      convert H_bst
+      exact H_result.2
+      simp [AVLTree.set_some]
+      rewrite [H_result.2, H_set]
+      simp
+    }
+    {
+      simp; split_conjs
+      . tauto
+      . simp [Ospec.equivalence _ _ Hordering]
+    }
+    {
+      have ⟨ added₂, left₂, ⟨ H_result, ⟨ H_bst, H_set ⟩ ⟩ ⟩ := AVLTreeSet.insert_loop_spec_global a left (BST.left Hbst)
+      progress keep Htree with AVLTreeSet.insert_loop_spec_local as ⟨ added₁, left₁, Hnode ⟩
+      exact (fun x => x < b)
+      rewrite [Htree] at H_result; simp at H_result
+      refine' ⟨ added₁, ⟨ some (AVLNode.mk b left₁ right), _ ⟩ ⟩
+      simp only [true_and]
+      split_conjs
+      cases' Hbst with _ _ _ H_forall_left H_forall_right H_bst_left H_bst_right
+      constructor
+      apply Hnode; exact H_forall_left
+      exact (gtOfRustOrder H b a Hordering)
+      exact H_forall_right
+      convert H_bst
+      exact H_result.2
+      exact H_bst_right
+      simp [AVLTree.set_some]
+      rewrite [H_result.2, H_set]
+      simp [Set.singleton_union, Set.insert_comm, Set.insert_union]
+    }
+
+@[pspec]
+def AVLTreeSet.insert_spec 
+  (a: T) (t: AVLTreeSet T):
+  BST.Invariant t.root -> (∃ t' added,t.insert _ H a = Result.ok (added, t')
+  -- it's still a binary search tree.
+  ∧ BST.Invariant t'.root
+  ∧ AVLTree.set t'.root = {a} ∪ AVLTree.set t.root)
+  := by
+  rw [AVLTreeSet.insert]; intro Hbst
+  progress keep h as ⟨ t', Hset ⟩; 
+  simp; assumption
+
+end Implementation
+
diff --git a/Verification/Order.lean b/Verification/Order.lean
new file mode 100644
index 0000000..396a524
--- /dev/null
+++ b/Verification/Order.lean
@@ -0,0 +1,57 @@
+import Verification.Specifications
+
+namespace Implementation
+
+open Primitives
+open avl_verification
+open Specifications (OrdSpecLinearOrderEq ltOfRustOrder gtOfRustOrder)
+
+instance ScalarU32DecidableLE : DecidableRel (· ≤ · : U32 -> U32 -> Prop) := by
+  simp [instLEScalar]
+  -- Lift this to the decidability of the Int version.
+  infer_instance
+
+instance : LinearOrder (Scalar .U32) where
+  le_antisymm := fun a b Hab Hba => by
+    apply (Scalar.eq_equiv a b).2; exact (Int.le_antisymm ((Scalar.le_equiv _ _).1 Hab) ((Scalar.le_equiv _ _).1 Hba))
+  le_total := fun a b => by
+    rcases (Int.le_total a b) with H | H 
+    left; exact (Scalar.le_equiv _ _).2 H 
+    right; exact (Scalar.le_equiv _ _).2 H
+  decidableLE := ScalarU32DecidableLE
+
+instance : OrdSpecLinearOrderEq OrdU32 where
+  infallible := fun a b => by
+    unfold Ord.cmp
+    unfold OrdU32
+    unfold OrdU32.cmp
+    rw [LinearOrder.compare_eq_compareOfLessAndEq, compareOfLessAndEq]
+    if hlt : a < b then 
+      use .Less
+      simp [hlt]
+    else
+      if heq: a = b
+      then
+      use .Equal
+      simp [hlt]
+      rw [heq]
+      -- TODO: simp [hlt, heq] breaks everything???
+      else
+        use .Greater
+        simp [hlt, heq]
+  symmetry := fun a b => by
+    rw [Ordering.toDualOrdering, LinearOrder.compare_eq_compareOfLessAndEq, compareOfLessAndEq]
+    rw [compare, Ord.opposite]
+    simp [LinearOrder.compare_eq_compareOfLessAndEq, compareOfLessAndEq]
+    split_ifs with hab hba hba' hab' hba'' _ hba₃ _ <;> tauto
+    exact lt_irrefl _ (lt_trans hab hba)
+    rw [hba'] at hab; exact lt_irrefl _ hab
+    rw [hab'] at hba''; exact lt_irrefl _ hba''
+    -- The order is total, therefore, we have at least one case where we are comparing something.
+    cases (lt_trichotomy a b) <;> tauto
+  equivalence := fun a b => by
+    unfold Ord.cmp
+    unfold OrdU32
+    unfold OrdU32.cmp
+    simp only []
+    split_ifs <;> simp only [Result.ok.injEq, not_false_eq_true, neq_imp, IsEmpty.forall_iff]; tauto; try assumption
diff --git a/Verification/Specifications.lean b/Verification/Specifications.lean
new file mode 100644
index 0000000..392c438
--- /dev/null
+++ b/Verification/Specifications.lean
@@ -0,0 +1,150 @@
+import «AvlVerification»
+
+namespace Primitives
+
+namespace Result
+
+def map {A B: Type} (x: Result A) (f: A -> B): Result B := match x with
+| .ok y => .ok (f y)
+| .fail e => .fail e
+| .div => .div
+
+@[inline]
+def isok {A: Type} : Result A -> Bool
+| .ok _ => true
+| .fail _ => false
+| .div => false
+
+@[inline]
+def get? {A: Type}: (r: Result A) -> isok r -> A
+| .ok x, _ => x
+
+end Result
+
+end Primitives
+
+namespace avl_verification
+
+@[simp]
+def Ordering.toLeanOrdering (o: avl_verification.Ordering): _root_.Ordering := match o with 
+| .Less => .lt
+| .Equal => .eq
+| .Greater => .gt
+
+def Ordering.ofLeanOrdering (o: _root_.Ordering): avl_verification.Ordering := match o with
+| .lt => .Less
+| .eq => .Equal
+| .gt => .Greater
+
+@[simp]
+def Ordering.toDualOrdering (o: avl_verification.Ordering): avl_verification.Ordering := match o with 
+| .Less => .Greater
+| .Equal => .Equal
+| .Greater => .Less
+
+@[simp]
+theorem Ordering.toLeanOrdering.injEq (x y: avl_verification.Ordering): (x.toLeanOrdering = y.toLeanOrdering) = (x = y) := by
+  apply propext
+  cases x <;> cases y <;> simp
+
+@[simp]
+theorem ite_eq_lt_distrib (c : Prop) [Decidable c] (a b : Ordering) :
+    ((if c then a else b) = .Less) = if c then a = .Less else b = .Less := by
+  by_cases c <;> simp [*]
+
+@[simp]
+theorem ite_eq_eq_distrib (c : Prop) [Decidable c] (a b : Ordering) :
+    ((if c then a else b) = .Equal) = if c then a = .Equal else b = .Equal := by
+  by_cases c <;> simp [*]
+
+@[simp]
+theorem ite_eq_gt_distrib (c : Prop) [Decidable c] (a b : Ordering) :
+    ((if c then a else b) = .Greater) = if c then a = .Greater else b = .Greater := by
+  by_cases c <;> simp [*]
+
+end avl_verification
+
+namespace Specifications
+
+open Primitives
+open Result
+
+variable {T: Type} (H: outParam (avl_verification.Ord T))
+
+@[simp]
+def _root_.Ordering.toDualOrdering (o: _root_.Ordering): _root_.Ordering := match o with 
+| .lt => .gt
+| .eq => .eq
+| .gt => .lt
+
+
+@[simp]
+theorem toDualOrderingOfToLeanOrdering (o: avl_verification.Ordering): o.toDualOrdering.toLeanOrdering = o.toLeanOrdering.toDualOrdering := by
+  cases o <;> simp
+
+@[simp]
+theorem toDualOrderingIdempotency (o: _root_.Ordering): o.toDualOrdering.toDualOrdering = o := by 
+  cases o <;> simp
+
+-- TODO: reason about raw bundling vs. refined bundling.
+-- raw bundling: hypothesis with Rust extracted objects.
+-- refined bundling: lifted hypothesis with Lean native objects.
+class OrdSpec [Ord T] where
+  infallible: ∀ a b, ∃ (o: avl_verification.Ordering), H.cmp a b = .ok o ∧ compare a b = o.toLeanOrdering
+  
+class OrdSpecSymmetry [O: Ord T] extends OrdSpec H where
+  symmetry: ∀ a b, O.compare a b = (O.opposite.compare a b).toDualOrdering
+
+-- Must be R decidableRel and an equivalence relationship?
+class OrdSpecRel [O: Ord T] (R: outParam (T -> T -> Prop)) extends OrdSpec H where
+  equivalence: ∀ a b, H.cmp a b = .ok .Equal -> R a b
+
+class OrdSpecLinearOrderEq [O: Ord T] extends OrdSpecSymmetry H, OrdSpecRel H Eq
+
+theorem infallible [Ord T] [OrdSpec H]: ∀ a b, ∃ o, H.cmp a b = .ok o := fun a b => by 
+  obtain ⟨ o, ⟨ H, _ ⟩ ⟩ := OrdSpec.infallible a b
+  exact ⟨ o, H ⟩
+
+instance: Coe (avl_verification.Ordering) (_root_.Ordering) where
+  coe a := a.toLeanOrdering
+
+theorem rustCmpEq [Ord T] [O: OrdSpec H]: H.cmp a b = .ok o <-> compare a b = o.toLeanOrdering := by
+  apply Iff.intro
+  . intro Hcmp
+    obtain ⟨ o', ⟨ Hcmp', Hcompare ⟩ ⟩ := O.infallible a b
+    rw [Hcmp', ok.injEq] at Hcmp
+    simp [Hcompare, Hcmp', Hcmp]
+  . intro Hcompare
+    obtain ⟨ o', ⟨ Hcmp', Hcompare' ⟩ ⟩ := O.infallible a b
+    rw [Hcompare', avl_verification.Ordering.toLeanOrdering.injEq] at Hcompare
+    simp [Hcompare.symm, Hcmp']
+
+
+theorem oppositeOfOpposite {x y: _root_.Ordering}: x.toDualOrdering = y ↔ x = y.toDualOrdering := by
+  cases x <;> cases y <;> simp
+theorem oppositeRustOrder [Ord T] [Spec: OrdSpecSymmetry H] {a b}: H.cmp b a = .ok o ↔ H.cmp a b = .ok o.toDualOrdering := by
+  rw [rustCmpEq, Spec.symmetry, compare, Ord.opposite, oppositeOfOpposite, rustCmpEq, toDualOrderingOfToLeanOrdering]
+
+theorem ltOfRustOrder 
+  [LO: LinearOrder T]
+  [Spec: OrdSpec H]: 
+  ∀ a b, H.cmp a b = .ok .Less -> a < b := by
+  intros a b
+  intro Hcmp
+  -- why the typeclass search doesn't work here?
+  refine' (@compare_lt_iff_lt T LO).1 _
+  obtain ⟨ o, ⟨ Hcmp', Hcompare ⟩ ⟩ := Spec.infallible a b 
+  simp only [Hcmp', ok.injEq] at Hcmp
+  simp [Hcompare, Hcmp, avl_verification.Ordering.toLeanOrdering]
+
+theorem gtOfRustOrder 
+  [LinearOrder T]
+  [Spec: OrdSpecSymmetry H]: 
+  ∀ a b, H.cmp a b = .ok .Greater -> b < a := by
+  intros a b
+  intro Hcmp
+  refine' @ltOfRustOrder _ H _ Spec.toOrdSpec _ _ _
+  rewrite [oppositeRustOrder]
+  simp [Hcmp]
+
+end Specifications
diff --git a/Verification/Tree.lean b/Verification/Tree.lean
new file mode 100644
index 0000000..d6a4f80
--- /dev/null
+++ b/Verification/Tree.lean
@@ -0,0 +1,82 @@
+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)
+
+@[simp]
+def AVLTree.set_some {x: T} {left right: AVLTree T}: AVLTree.set (some (AVLNode.mk x left right)) = {x} ∪ AVLTree.set left ∪ AVLTree.set right := by 
+  simp [set, setOf]
+
+end Tree