wip: experiments with eliminating MProd

src/Init/Data/MProd.lean

@@ -0,0 +1,121 @@
+/-
+Copyright (c) 2025 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author: Kim Morrison
+-/
+
+
+/-!
+Recall that `MProd` is a universe monomorphic version of `Prod`,
+used by `do` notation to avoid universe unification issues.
+
+However, this means that code generated by `do` notation used `MProd`
+internally, and we have essentially no lemmas about `MProd`.
+
+This file develops basic `simp` lemmas for eliminating appearances of `MProd` in favour of `Prod`.
+All the simp lemmas work by replacing constructors of `MProd` with the constructor of `Prod`
+followed by `toMProd`, and then pulling `toMProd` outwards until it hits a `toProd` or an
+`MProd.fst` or `MProd.snd`.
+-/
+
+/-- Convert an `MProd` to a `Prod`. -/
+def MProd.toProd {α β} : MProd α β → α × β
+| ⟨a, b⟩ => (a, b)
+
+/-- Convert a `Prod` to an `MProd`. -/
+def Prod.toMProd {α β} : α × β → MProd α β
+| (a, b) => ⟨a, b⟩
+
+namespace Prod
+
+@[simp] theorem toProd_toMProd {α β} (p : α × β) : p.toMProd.toProd = p := by
+  cases p
+  rfl
+
+end Prod
+
+open Prod
+
+namespace MProd
+
+@[simp] theorem toMProd_toProd {α β} (p : MProd α β) : p.toProd.toMProd = p := by
+  cases p
+  rfl
+
+@[simp]
+theorem mk_eq_toMProd_mk {α β : Type u} (a : α) (b : β) : MProd.mk a b = (a, b).toMProd := by
+  simp [toMProd]
+
+@[simp] theorem fst_eq {α β} (p : MProd α β) : p.fst = p.toProd.fst := by
+  cases p
+  rfl
+
+@[simp] theorem snd_eq {α β} (p : MProd α β) : p.snd = p.toProd.snd := by
+  cases p
+  rfl
+
+end MProd
+
+namespace List
+
+-- WIP: this is insufficient! It works fine for a pairwise product,
+-- but in general we will have an iterated `MProd`, with one factor for each mutable variable.
+-- We need a simproc, indexed on `List.foldl _ (Prod.toMProd _)`, which will descend into the
+-- iterated product.
+
+/-- Pull `toMProd` outwards through `List.foldl`. -/
+@[simp] theorem foldl_toMProd {α β γ} (f : MProd α β → γ → α × β) (init : α × β) (l : List γ) :
+    List.foldl (fun p g => (f p g).toMProd) init.toMProd l =
+      (List.foldl (fun p g => f p.toMProd g) init l).toMProd := by
+  induction l generalizing init <;> simp_all
+
+/-- Pull `toMProd` outwards through `List.foldr`. -/
+@[simp] theorem foldr_toMProd {α β γ} (f : γ → MProd α β → α × β) (init : α × β) (l : List γ) :
+    List.foldr (fun g p => (f g p).toMProd) init.toMProd l =
+      (List.foldr (fun g p => f g p.toMProd) init l).toMProd := by
+  induction l <;> simp_all
+
+end List
+
+def fib_spec : Nat → Nat
+| 0 => 0
+| 1 => 1
+| n+2 => fib_spec n + fib_spec (n+1)
+
+def fib_impl (n : Nat) := Id.run do
+  if n = 0 then return 0
+  let mut a := 0
+  let mut b := 0
+  b := b + 1
+  for _ in [1:n] do
+    let a' := a
+    a := b
+    b := a' + b
+  return b
+
+theorem fib_correct {n} : fib_impl n = fib_spec n := by
+  simp [fib_impl, Id.run]
+  match n with
+  | 0 => simp [fib_spec]
+  | n+1 =>
+    suffices ((List.range' 1 n).foldl (fun b a ↦ (b.snd, b.fst + b.snd)) (0, 1)) =
+        ⟨fib_spec n, fib_spec (n + 1)⟩ by simp_all
+    induction n with
+    | zero => rfl
+    | succ n ih => simp [fib_spec, List.range'_1_concat, ih]
+
+
+def triple (n : Nat) := Id.run do
+  let mut a := 0
+  let mut b := 0
+  let mut c := 0
+  for _ in [1:n] do
+    let a' := a
+    a := b
+    b := c
+    c := a'
+  return a + b + c
+
+theorem triple_correct {n} : triple n = 0 := by
+  simp [triple, Id.run]
+  sorry

tests/lean/run/fib_correct.lean

@@ -0,0 +1,46 @@
+/-!
+This is an example for monadic reasoning.
+
+The eventual goal is to provide a nice user experience for proving `fib_impl n = fib_spec n`
+and related goals.
+
+Currently, this file just contains a proof that uses `simp` lemmas to convert the `do` notation
+and for loop into a `List.foldl`, and then gives a "functional" proof.
+(This is *not* the nice user experience we are aiming for!)
+
+Even in this setup, there is an awkward problem that `do` blocks handle multiple mutable variables
+via the universe monomorphic `MProd` type, to avoid universe unification issues arising when using
+`Prod`. We have to jump through some additional hoops to handle that.
+We could provide simp lemmas, simprocs, and possibly custom tactics to eliminator `MProd` from the
+terms produced by `do` notation.
+-/
+
+def fib_spec : Nat → Nat
+| 0 => 0
+| 1 => 1
+| n+2 => fib_spec n + fib_spec (n+1)
+
+def fib_impl (n : Nat) := Id.run do
+  if n = 0 then return 0
+  let mut a := 0
+  let mut b := 0
+  b := b + 1
+  for _ in [1:n] do
+    let a' := a
+    a := b
+    b := a' + b
+  return b
+
+theorem fib_correct {n} : fib_impl n = fib_spec n := by
+  -- The default simp set eliminates the binds generated by `do` notation,
+  -- and converts the `for` loop into a `List.foldl` over `List.range'`.
+  simp [fib_impl, Id.run]
+  match n with
+  | 0 => simp [fib_spec]
+  | n+1 =>
+    -- Note here that we have to use `⟨x, y⟩ : MProd _ _`, because these are not `Prod` products.
+    suffices ((List.range' 1 n).foldl (fun b a ↦ ⟨b.snd, b.fst + b.snd⟩) (⟨0, 1⟩ : MProd _ _)) =
+        ⟨fib_spec n, fib_spec (n + 1)⟩ by simp_all
+    induction n with
+    | zero => rfl
+    | succ n ih => simp [fib_spec, List.range'_1_concat, ih]