Problem: bst_zigzag

Your task is to implement the Zig-Zag operation (Left-Right Rotation) in a binary search tree and verify its correctness in Lean. A Zig-Zag is performed when node X lies between parent P and grandparent G (X is right child of P, P is left child of G). Prove: (1) Set Preservation: The tree's view is unchanged. (2) BST Invariant Preservation. (3) Structural Correctness: X becomes the new root, P becomes the left child of X, and G becomes the right child of X.


import Mathlib

inductive BinarySearchTree : Type
| Empty : BinarySearchTree
| Node : Int → BinarySearchTree → BinarySearchTree → BinarySearchTree
deriving Inhabited

def view (t : BinarySearchTree) : Set Int :=
  match t with
  | BinarySearchTree.Empty => ∅
  | BinarySearchTree.Node v l r => view l ∪ view r ∪ {v}

def is_bst (t : BinarySearchTree) : Prop :=
  match t with
  | BinarySearchTree.Empty => True
  | BinarySearchTree.Node v l r =>
    (∀ x ∈ view l, x < v) ∧
    is_bst l ∧
    (∀ x ∈ view r, x > v) ∧
    is_bst r

def get_left (t : BinarySearchTree) : Option BinarySearchTree :=
  match t with
  | BinarySearchTree.Node _ l _ => some l
  | _ => none

def get_right (t : BinarySearchTree) : Option BinarySearchTree :=
  match t with
  | BinarySearchTree.Node _ _ r => some r
  | _ => none

def get_val (t : BinarySearchTree) : Option Int :=
  match t with
  | BinarySearchTree.Node v _ _ => some v
  | _ => none

def BinarySearchTree.is_node (t : BinarySearchTree) : Prop :=
  match t with
  | BinarySearchTree.Node _ _ _ => True
  | _ => False

-- Precondition definitions
@[reducible, simp]
def zigzag_precond (t : BinarySearchTree) : Prop :=
  -- !benchmark @start precond
  is_bst t ∧ 
  BinarySearchTree.is_node t ∧ 
  (match get_left t with 
   | some p => BinarySearchTree.is_node p ∧ 
               (match get_right p with 
                | some x => BinarySearchTree.is_node x 
                | none => False)
   | none => False)
  -- !benchmark @end precond

-- !benchmark @start auxcode
-- !benchmark @end auxcode

-- Main function definition
def zigzag (t : BinarySearchTree) 
    (h_precond : zigzag_precond t) : BinarySearchTree :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition definitions
@[reducible, simp]
def zigzag_postcond (t : BinarySearchTree) (result : BinarySearchTree)
    (_ : zigzag_precond t) : Prop :=
  -- !benchmark @start postcond
  is_bst result ∧
  view result = view t ∧
  -- Structural properties: X (grandchild, originally P's right child) is now root
  get_val result = get_val (get_right (get_left t).get!).get!
  -- !benchmark @end postcond

-- !benchmark @start lemma
-- !benchmark @end lemma

-- Proof content
theorem zigzag_postcond_satisfied (t : BinarySearchTree)
    (h_precond : zigzag_precond t) :
    zigzag_postcond t (zigzag t h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof