Proving correctness of do notation programs

findLeast.lean

import Lean.Elab.Tactic

def findLeast (a : Array Nat) : Nat := do
  let mut smallest := a[0]
  for i in [1:a.size] do
    if a[i] ≤ smallest then
      smallest := a[i]
  return smallest

#eval findLeast #[8, 3, 10, 4, 6]

namespace Lean.Elab.Tactic
open Meta

syntax (name := trustme) "trustme" term : tactic
@[tactic «trustme»] def evalTrustme : Tactic := fun stx =>
  match stx with
  | `(tactic| trustme $e) => closeMainGoalUsing (fun type => elabTerm e type)
  | _                   => throwUnsupportedSyntax

end Lean.Elab.Tactic

namespace Std.Range.forIn

theorem loop_eq {β : Type u} {m : Type u → Type v} [Monad m]
  (range : Std.Range) (f : Nat → β → m (ForInStep β)) (i j : Nat) (b : β) :
  loop range f i j b =
  Nat.brecOn (motive := fun _ => Nat → β → m β) i
    (fun i (F : Nat.below i) (j : Nat) (b : β) =>
      if j ≥ range.stop then pure b else
        (match i : ∀ i, Nat.below i → m β with
          | 0 => fun x => pure b
          | Nat.succ i => fun x =>
            do match ← f j b with
            | ForInStep.done b => pure b
            | ForInStep.yield b => by exact PProd.fst x.fst (j + range.step) b)
          F)
    j b := by
  trustme (rfl : loop range f i j b = _)

@[simp] theorem loop_zero {β : Type u} {m : Type u → Type v} [Monad m]
  (range : Std.Range) (f : Nat → β → m (ForInStep β)) (j : Nat) (b : β) :
  loop range f 0 j b = pure b := by
  simp [loop_eq]
  show (if range.stop ≤ j then pure b else pure b : m β) = pure b
  cases Nat.decLe range.stop j <;> rfl

@[simp] theorem loop_succ {β : Type u} {m : Type u → Type v} [Monad m]
  (range : Std.Range) (f : Nat → β → m (ForInStep β)) (i j : Nat) (b : β) :
  loop range f (Nat.succ i) j b =
  (if j ≥ range.stop then pure b else do
    match ← f j b with
    | ForInStep.done b => pure b
    | ForInStep.yield b => loop range f i (j + range.step) b) := by
  simp [loop_eq]

theorem loop_stop {β : Type u} {m : Type u → Type v} [Monad m]
  (range : Std.Range) (f : Nat → β → m (ForInStep β)) (i j : Nat) (b : β)
  (h : range.stop ≤ j) : loop range f i j b = pure b := by
  cases i <;> simp [h]

theorem loop_non_stop {β : Type u} {m : Type u → Type v} [Monad m]
  (range : Std.Range) (f : Nat → β → m (ForInStep β)) (i j : Nat) (b : β)
  (h : ¬ range.stop ≤ j) : loop range f (Nat.succ i) j b = (do
    match ← f j b with
    | ForInStep.done b => pure b
    | ForInStep.yield b => loop range f i (j + range.step) b) := by
  simp [h]

end Std.Range.forIn

theorem Nat.le_refl (a : Nat) : a ≤ a := sorry
theorem Nat.le_trans {a b c : Nat} : a ≤ b → b ≤ c → a ≤ c := sorry
theorem Nat.lt_trans {a b c : Nat} : a < b → b < c → a < c := sorry
theorem Nat.le_of_lt {a b : Nat} : a < b → a ≤ b := sorry
theorem Nat.not_le_of_lt {a b : Nat} : a < b → ¬ b ≤ a := sorry
theorem Nat.le_of_not_le {a b : Nat} : ¬ a ≤ b → b ≤ a := sorry
theorem Nat.lt_or_eq {a b : Nat} : a ≤ b → a < b ∨ a = b := sorry
theorem Nat.lt_succ_of_lt {a b : Nat} : a < b → a < b + 1 := sorry

open ForInStep in
def findLeastRec (f : Nat → Nat) (lo hi z j b : Nat) : Nat :=
Std.Range.forIn.loop (m := Id) [lo:hi]
  (fun i r => if f i ≤ r then yield (f i) else yield r) z j b

theorem findLeast_findLeastRec {a : Array Nat} :
  findLeast a = findLeastRec a.getOp 1 a.size a.size 1 a[0] := rfl

open ForInStep in
theorem findLeastRec_succ (f lo hi z j b) :
  findLeastRec f lo hi (z+1) j b = if hi ≤ j then b else
    match if f j ≤ b then yield (f j) else yield b with
    | done b => b
    | yield b => findLeastRec f lo hi z (j + 1) b := by
  simp [findLeastRec]; rfl

theorem findLeastRec_prop (f lo hi) : ∀ n j b,
  findLeastRec f lo hi n j b ≤ b ∧
  ∀ i, j ≤ i → i < j + n → i < hi → findLeastRec f lo hi n j b ≤ f i
| 0, j, b => by
  simp [findLeastRec]
  exact ⟨Nat.le_refl _, λ i h₁ h₂ => by cases Nat.not_le_of_lt h₂ h₁⟩
| n+1, j, b => by
  simp [findLeastRec_succ]
  cases Decidable.em (hi ≤ j) with simp [h]
  | inl h =>
    refine ⟨Nat.le_refl _, λ i h₁ _ h₂ => ?_⟩
    cases Nat.not_le_of_lt h₂ (Nat.le_trans h h₁)
  | inr h => cases Decidable.em (f j ≤ b) with simp [h']
    | inl h' =>
      let ⟨ih₁, ih₂⟩ := findLeastRec_prop f lo hi n (j+1) (f j)
      refine ⟨Nat.le_trans ih₁ h', λ i h₁ h₂ => ?_⟩
      cases Nat.lt_or_eq h₁ with
      | inl h => exact ih₂ _ h (by rw Nat.add_right_comm; exact h₂)
      | inr h => cases h; intro; exact ih₁
    | inr h' =>
      let ⟨ih₁, ih₂⟩ := findLeastRec_prop f lo hi n (j+1) b
      refine ⟨ih₁, λ i h₁ h₂ => ?_⟩
      cases Nat.lt_or_eq h₁ with
      | inl h => exact ih₂ _ h (by rw Nat.add_right_comm; exact h₂)
      | inr h => cases h; intro; exact Nat.le_trans ih₁ (Nat.le_of_not_le h')

theorem findLeast_min {a : Array Nat} {i} (h : i < a.size) : findLeast a ≤ a[i] := by
  rw findLeast_findLeastRec
  let ⟨ih₁, ih₂⟩ := findLeastRec_prop a.getOp 1 a.size a.size 1 a[0]
  match i, h with
  | 0, h => exact ih₁
  | i+1, h => refine ih₂ _ (by cases i <;> rfl) ?_ h; rw Nat.add_comm 1; exact Nat.lt_succ_of_lt h