-- Frap/RedBlack.lean

inductive Color where
  | red
  | black

inductive Tree (α : Type u) where
  | empty  : Tree α
  | tree (c: Color) (l : Tree α) (k : Nat) (v : α) (r : Tree α) : Tree α

namespace Tree

open Color

def ex_tree : Tree String :=
  tree black (tree red empty 2 "two" empty) 4 "four" (tree red empty 5 "five" empty)

def balance {α : Type u} (c : Color) (l : Tree α) (k : Nat) (vk : α) (r : Tree α) : Tree α :=
  match c with
  | red => tree red l k vk r
  | black =>
    match (l, k, vk, r) with
    | (tree red (tree red a x vx b) y vy c, z, vz, d)
    | (tree red a x vx (tree red b y vy c), z, vz, d)
    | (a, x, vx, tree red (tree red b y vy c) z vz d)
    | (a, x, vx, tree red b y vy (tree red c z vz d))
        => tree red (tree black a x vx b) y vy (tree black c z vz d)
    | _ => tree black l k vk r

inductive ForallTree {α : Type u} (p : Nat → α → Prop) : Tree α → Prop where
  | empty : ForallTree p empty
  | tree : ∀ c l k v r,
      p k v → ForallTree p l → ForallTree p r
      → ForallTree p (tree c l k v r)

inductive BST {α : Type u} : Tree α → Prop where
  | empty : BST empty
  | tree : ∀ c l k v r,
      ForallTree (fun x _ => x < k) l
      → ForallTree (fun x _ => x > k) r
      → BST l
      → BST r
      → BST (tree c l k v r)

 

example : BST ex_tree := by admit /- proof elided -/

-- Repository definitions used:
content: ...

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: inductive Color where
  | red
  | black

content: inductive Tree (α : Type u) where
  | empty  : Tree α
  | tree (c: Color) (l : Tree α) (k : Nat) (v : α) (r : Tree α) : Tree α

content: def ex_tree : Tree String :=
  tree black (tree red empty 2 "two" empty) 4 "four" (tree red empty 5 "five" empty)

content: def balance {α : Type u} (c : Color) (l : Tree α) (k : Nat) (vk : α) (r : Tree α) : Tree α :=
  match c with
  | red => tree red l k vk r
  | black =>
    match (l, k, vk, r) with
    | (tree red (tree red a x vx b) y vy c, z, vz, d)
    | (tree red a x vx (tree red b y vy c), z, vz, d)
    | (a, x, vx, tree red (tree red b y vy c) z vz d)
    | (a, x, vx, tree red b y vy (tree red c z vz d))
        => tree red (tree black a x vx b) y vy (tree black c z vz d)
    | _ => tree black l k vk r

content: inductive ForallTree {α : Type u} (p : Nat → α → Prop) : Tree α → Prop where
  | empty : ForallTree p empty
  | tree : ∀ c l k v r,
      p k v → ForallTree p l → ForallTree p r
      → ForallTree p (tree c l k v r)

content: inductive BST {α : Type u} : Tree α → Prop where
  | empty : BST empty
  | tree : ∀ c l k v r,
      ForallTree (fun x _ => x < k) l
      → ForallTree (fun x _ => x > k) r
      → BST l
      → BST r
      → BST (tree c l k v r)

 

example : BST ex_tree := by admit /- proof elided -/

-- Local lemmas used:
content: theorem forallTree_imp {α : Type u} (P Q : Nat → α → Prop) t
    : ForallTree P t → (∀ k v, P k v → Q k v) → ForallTree Q t

content: theorem forallTree_lt {α : Type u} (t : Tree α) k k'
    : ForallTree (fun x _ => x < k) t → k < k'
      → ForallTree (fun x _ => x < k') t

content: theorem forallTree_gt {α : Type u} (t : Tree α) k k'
    : ForallTree (fun x _ => x > k) t → k > k'
      → ForallTree (fun x _ => x > k') t

-- Target theorem and proof:
theorem balance_BST {α : Type u} c (l : Tree α) k vk (r : Tree α)
    : ForallTree (fun x _ => x < k) l
      -> ForallTree (fun x _ => x > k) r
      -> BST l
      -> BST r
      -> BST (balance c l k vk r)
:= by
  intro hlk hkr hbl hbr; simp
  split
  . constructor <;> assumption
  . split
    . -- we are in the case where `x` is left child of `y`
      cases l <;> repeat simp [*] at *
      rcases hbl with ⟨⟩ | ⟨_, _, _, _, _, hxy, hyc, hbx, hbc⟩
      rcases hlk with ⟨⟩ | ⟨_, _, _, _, _, hyz, hxz, hzc⟩
      rcases hbx with ⟨⟩ | ⟨_, _, _, _, _, hax, hxb, hba, hbb⟩
      rcases hxy with ⟨⟩ | ⟨_, _, _, _, _, hxy, hay, hby⟩
      rcases hxz with ⟨⟩ | ⟨_, _, _, _, _, haz, hxz, hbz⟩
      /-
      Here, we have destructed all the assumptions to the atomic level.
      The goal remains to show that the resulting tree is a BST.
      We apply the constructor tactic repeatedly.
      Most of the resulting subgoals are trivial by assumption.
      -/
      (repeat' constructor) <;> try assumption
      /-
      The only interesting subgoal is to show that each node in the right subtree of the original root is greater than the left child of the original root.
      This should follow from assumptions, although not trivially.
      To facilitate reasoning at this point, we prove lemmas `forallTree_imp`, `forallTree_lt`, and `forallTree_gt` above.
      -/
      apply forallTree_gt <;> assumption
    . -- we are in the case where `y` is right child of `x`
      cases l <;> repeat simp [*] at *
      rcases hbl with ⟨⟩ | ⟨_, _, _, _, _, hax, hxy, hba, hby⟩
      rcases hlk with ⟨⟩ | ⟨_, _, _, _, _, hxz, haz, hyz⟩
      rcases hby with ⟨⟩ | ⟨_, _, _, _, _, hby, hyc, hbb, hbc⟩
      rcases hxy with ⟨⟩ | ⟨_, _, _, _, _, hxy, hxb, hxc⟩
      rcases hyz with ⟨⟩ | ⟨_, _, _, _, _, hyz, hbz, hcz⟩
      (repeat' constructor) <;> try assumption
      . apply forallTree_lt
        . exact hax
        . assumption
      . apply forallTree_gt <;> assumption
    . -- we are in the case where `y` is left child of `z`
      cases r <;> repeat simp [*] at *
      rcases hbr with ⟨⟩ | ⟨_, _, _, _, _, hyz, hzd, hby, hbd⟩
      rcases hkr with ⟨⟩ | ⟨_, _, _, _, _, hxz, hxy, hxd⟩
      rcases hby with ⟨⟩ | ⟨_, _, _, _, _, hby, hyc, hbb, hbc⟩
      rcases hxy with ⟨⟩ | ⟨_, _, _, _, _, hxy, hxb, hxc⟩
      rcases hyz with ⟨⟩ | ⟨_, _, _, _, _, hyz, hbz, hcz⟩
      (repeat' constructor) <;> try assumption
      . apply forallTree_lt <;> assumption
      . apply forallTree_gt
        . exact hzd
        . assumption
    . -- we are in the case where `z` is right child of `y`
      cases r <;> repeat simp [*] at *
      rcases hbr with ⟨⟩ | ⟨_, _, _, _, _, hby, hyz, hbb, hbz⟩
      rcases hkr with ⟨⟩ | ⟨_, _, _, _, _, hxy, hxb, hxz⟩
      rcases hbz with ⟨⟩ | ⟨_, _, _, _, _, hcz, hzd, hbc, hbd⟩
      rcases hxz with ⟨⟩ | ⟨_, _, _, _, _, hxz, hxc, hxd⟩
      rcases hyz with ⟨⟩ | ⟨_, _, _, _, _, hyz, hyc, hyd⟩
      (repeat' constructor) <;> try assumption
      apply forallTree_lt <;> assumption
    . constructor <;> assumption