From 506cf01d9494075c3da025233cd36e1d1128f17b Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 12 Sep 2022 16:02:19 -0700
Subject: [PATCH] fix: bug at `simpCasesOnCtor?`

---
 src/Lean/Compiler/LCNF/Simp.lean       | 17 ++++++++++++++---
 tests/lean/run/simpCasesOnCtorBug.lean | 24 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)
 create mode 100644 tests/lean/run/simpCasesOnCtorBug.lean

diff --git a/src/Lean/Compiler/LCNF/Simp.lean b/src/Lean/Compiler/LCNF/Simp.lean
index 891d444f72..a4e013bf2f 100644
--- a/src/Lean/Compiler/LCNF/Simp.lean
+++ b/src/Lean/Compiler/LCNF/Simp.lean
@@ -842,7 +842,7 @@ Try to simplify `cases` of `constructor`
 partial def simpCasesOnCtor? (cases : Cases) : SimpM (Option Code) := do
   let discr ← normFVar cases.discr
   let discrExpr ← findCtor (.fvar discr)
-  let some (ctorVal, ctorArgs) := discrExpr.constructorApp? (← getEnv) | return none
+  let some (ctorVal, ctorArgs) := discrExpr.constructorApp? (← getEnv) (useRaw := true) | return none
   let (alt, cases) := cases.extractAlt! ctorVal.name
   eraseFVarsAt (.cases cases)
   markSimplified
@@ -850,11 +850,22 @@ partial def simpCasesOnCtor? (cases : Cases) : SimpM (Option Code) := do
   | .default k => simp k
   | .alt _ params k =>
     let fields := ctorArgs[ctorVal.numParams:]
+    let mut auxDecls := #[]
     for param in params, field in fields do
-      addSubst param.fvarId field
+      /-
+      `field` may not be a free variable. Recall that `constructorApp?` has special support for numerals,
+      and `ctorArgs` contains a numeral if `discrExpr` is a numeral. We may have other cases in the future.
+      To make the code robust, we add auxiliary declarations whenever the `field` is not a free variable.
+      -/
+      if field.isFVar then
+        addSubst param.fvarId field
+      else
+        let auxDecl ← mkAuxLetDecl field
+        auxDecls := auxDecls.push (CodeDecl.let auxDecl)
+        addSubst param.fvarId (.fvar auxDecl.fvarId)
     let k ← simp k
     eraseParams params
-    return k
+    attachCodeDecls auxDecls k
 
 /--
 Simplify `code`
diff --git a/tests/lean/run/simpCasesOnCtorBug.lean b/tests/lean/run/simpCasesOnCtorBug.lean
new file mode 100644
index 0000000000..8e7c90d3a1
--- /dev/null
+++ b/tests/lean/run/simpCasesOnCtorBug.lean
@@ -0,0 +1,24 @@
+import Lean
+
+def test1 (y : Nat) : Nat :=
+  let x := 3
+  match x with
+  | 0 => y+1
+  | .succ x => y + x
+
+#eval Lean.Compiler.compile #[``test1]
+
+def test2 (y : Nat) : Nat :=
+  let x := 3
+  match x with
+  | 0 => y+1
+  | .succ x =>
+    match x with
+    | 0 => y+2
+    | .succ x => y + x
+
+#eval Lean.Compiler.compile #[``test2]
+
+set_option trace.Compiler.result true
+#eval Lean.Compiler.compile #[``test1]
+#eval Lean.Compiler.compile #[``test2]