From a06cd40e29a385c6939a4e120a3a7985a40ff042 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 28 Mar 2022 17:15:13 -0700
Subject: [PATCH] feat: improve match expression support at `simp`

---
 src/Lean/Meta/Match/MatchEqs.lean      | 24 +++--------------
 src/Lean/Meta/Match/MatchEqsExt.lean   | 36 ++++++++++++++++++++++++++
 src/Lean/Meta/Tactic/Simp/Rewrite.lean | 17 ++++++++++++
 src/Lean/Meta/Tactic/Split.lean        |  9 +++----
 tests/lean/run/def10.lean              |  2 +-
 tests/lean/run/simpMatch.lean          | 10 +++++++
 tests/lean/run/split2.lean             |  2 +-
 7 files changed, 72 insertions(+), 28 deletions(-)
 create mode 100644 src/Lean/Meta/Match/MatchEqsExt.lean
 create mode 100644 tests/lean/run/simpMatch.lean

diff --git a/src/Lean/Meta/Match/MatchEqs.lean b/src/Lean/Meta/Match/MatchEqs.lean
index c641902783..e5b91506f7 100644
--- a/src/Lean/Meta/Match/MatchEqs.lean
+++ b/src/Lean/Meta/Match/MatchEqs.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Meta.Match.Match
+import Lean.Meta.Match.MatchEqsExt
 import Lean.Meta.Tactic.Apply
 import Lean.Meta.Tactic.Delta
 import Lean.Meta.Tactic.SplitIf
@@ -52,25 +53,6 @@ def casesOnStuckLHS? (mvarId : MVarId) : MetaM (Option (Array MVarId)) := do
 
 namespace Match
 
-structure MatchEqns where
-  eqnNames             : Array Name
-  splitterName         : Name
-  splitterAltNumParams : Array Nat
-  deriving Inhabited, Repr
-
-def MatchEqns.size (e : MatchEqns) : Nat :=
-  e.eqnNames.size
-
-structure MatchEqnsExtState where
-  map : Std.PHashMap Name MatchEqns := {}
-  deriving Inhabited
-
-/- We generate the equations and splitter on demand, and do not save them on .olean files. -/
-builtin_initialize matchEqnsExt : EnvExtension MatchEqnsExtState ←
-  registerEnvExtension (pure {})
-
-private def registerMatchEqns (matchDeclName : Name) (matchEqns : MatchEqns) : CoreM Unit :=
-  modifyEnv fun env => matchEqnsExt.modifyState env fun s => { s with map := s.map.insert matchDeclName matchEqns }
 
 def unfoldNamedPattern (e : Expr) : MetaM Expr := do
   let visit (e : Expr) : MetaM TransformStep := do
@@ -491,7 +473,9 @@ private partial def mkEquationsFor (matchDeclName : Name) :  MetaM MatchEqns :=
       registerMatchEqns matchDeclName result
       return result
 
-def getEquationsFor (matchDeclName : Name) : MetaM MatchEqns := do
+/- See header at `MatchEqsExt.lean` -/
+@[export lean_get_match_equations_for]
+def getEquationsForImpl (matchDeclName : Name) : MetaM MatchEqns := do
   match matchEqnsExt.getState (← getEnv) |>.map.find? matchDeclName with
   | some matchEqns => return matchEqns
   | none => mkEquationsFor matchDeclName
diff --git a/src/Lean/Meta/Match/MatchEqsExt.lean b/src/Lean/Meta/Match/MatchEqsExt.lean
new file mode 100644
index 0000000000..36b3ef6901
--- /dev/null
+++ b/src/Lean/Meta/Match/MatchEqsExt.lean
@@ -0,0 +1,36 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Meta.Basic
+
+namespace Lean.Meta.Match
+
+structure MatchEqns where
+  eqnNames             : Array Name
+  splitterName         : Name
+  splitterAltNumParams : Array Nat
+  deriving Inhabited, Repr
+
+def MatchEqns.size (e : MatchEqns) : Nat :=
+  e.eqnNames.size
+
+structure MatchEqnsExtState where
+  map : Std.PHashMap Name MatchEqns := {}
+  deriving Inhabited
+
+/- We generate the equations and splitter on demand, and do not save them on .olean files. -/
+builtin_initialize matchEqnsExt : EnvExtension MatchEqnsExtState ←
+  registerEnvExtension (pure {})
+
+def registerMatchEqns (matchDeclName : Name) (matchEqns : MatchEqns) : CoreM Unit :=
+  modifyEnv fun env => matchEqnsExt.modifyState env fun s => { s with map := s.map.insert matchDeclName matchEqns }
+
+/-
+  Forward definition. We want to use `getEquationsFor` in the simplifier,
+ `getEquationsFor` depends on `mkEquationsfor` which uses the simplifier. -/
+@[extern "lean_get_match_equations_for"]
+constant getEquationsFor (matchDeclName : Name) : MetaM MatchEqns
+
+end Lean.Meta.Match
diff --git a/src/Lean/Meta/Tactic/Simp/Rewrite.lean b/src/Lean/Meta/Tactic/Simp/Rewrite.lean
index 352ce3882d..89a3a7576b 100644
--- a/src/Lean/Meta/Tactic/Simp/Rewrite.lean
+++ b/src/Lean/Meta/Tactic/Simp/Rewrite.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Meta.ACLt
+import Lean.Meta.Match.MatchEqsExt
 import Lean.Meta.AppBuilder
 import Lean.Meta.SynthInstance
 import Lean.Meta.Tactic.Simp.Types
@@ -202,6 +203,21 @@ def simpArith? (e : Expr) : SimpM (Option Step) := do
   let some (e', h) ← Linear.simp? e (← read).parent? | return none
   return Step.visit { expr := e', proof? := h }
 
+def simpMatchCore? (app : MatcherApp) (e : Expr) (discharge? : Expr → SimpM (Option Expr)) : SimpM (Option Step) := do
+  for matchEq in (← Match.getEquationsFor app.matcherName).eqnNames do
+    -- Try lemma
+    match (← withReducible <| Simp.tryTheorem? e { proof := mkConst matchEq, name? := some matchEq } discharge?) with
+    | none   => pure ()
+    | some r => return some (Simp.Step.done r)
+  return none
+
+def simpMatch? (discharge? : Expr → SimpM (Option Expr)) (e : Expr) : SimpM (Option Step) := do
+  if (← read).config.iota then
+    let some app ← matchMatcherApp? e | return none
+    simpMatchCore? app e discharge?
+  else
+    return none
+
 def rewritePre (e : Expr) (discharge? : Expr → SimpM (Option Expr)) : SimpM Step := do
   for thms in (← read).simpTheorems do
     let r ← rewrite e thms.pre thms.erased discharge? (tag := "pre")
@@ -222,6 +238,7 @@ def preDefault (e : Expr) (discharge? : Expr → SimpM (Option Expr)) : SimpM St
 
 def postDefault (e : Expr) (discharge? : Expr → SimpM (Option Expr)) : SimpM Step := do
   let s ← rewritePost e discharge?
+  let s ← andThen s (simpMatch? discharge?)
   let s ← andThen s simpArith?
   let s ← andThen s tryRewriteUsingDecide?
   andThen s tryRewriteCtorEq?
diff --git a/src/Lean/Meta/Tactic/Split.lean b/src/Lean/Meta/Tactic/Split.lean
index 82fdee682a..1350300da5 100644
--- a/src/Lean/Meta/Tactic/Split.lean
+++ b/src/Lean/Meta/Tactic/Split.lean
@@ -25,12 +25,9 @@ where
     match (← reduceRecMatcher? e) with
     | some e' => return Simp.Step.done { expr := e' }
     | none    =>
-      for matchEq in (← Match.getEquationsFor app.matcherName).eqnNames do
-        -- Try lemma
-        match (← withReducible <| Simp.tryTheorem? e { proof := mkConst matchEq, name? := some matchEq } SplitIf.discharge?) with
-        | none   => pure ()
-        | some r => return Simp.Step.done r
-      return Simp.Step.visit { expr := e }
+      match (← Simp.simpMatchCore? app e SplitIf.discharge?) with
+      | some r => return r
+      | none => return Simp.Step.visit { expr := e }
 
 def simpMatchTarget (mvarId : MVarId) : MetaM MVarId := withMVarContext mvarId do
   let target ← instantiateMVars (← getMVarType mvarId)
diff --git a/tests/lean/run/def10.lean b/tests/lean/run/def10.lean
index a91dff8191..874119b304 100644
--- a/tests/lean/run/def10.lean
+++ b/tests/lean/run/def10.lean
@@ -17,7 +17,7 @@ theorem ex4 : g false false false = 2 := rfl
 theorem ex5 : g false true true = 3 := rfl
 
 theorem f_eq (h : a = true → b = true → False) : f a b = 3 := by
-  simp [f]
+  simp (config := { iota := false }) [f]
   split
   · contradiction
   · rfl
diff --git a/tests/lean/run/simpMatch.lean b/tests/lean/run/simpMatch.lean
new file mode 100644
index 0000000000..68637ee6ee
--- /dev/null
+++ b/tests/lean/run/simpMatch.lean
@@ -0,0 +1,10 @@
+def f (x y : Nat) : Nat :=
+  match x, y with
+  | 0, 0 => 1
+  | _, _ => 2
+
+example (h : f x y = 1) : f x y ≠ 2 := by
+  simp [f] at *
+  split
+  next => decide
+  next x' y' hnp => simp [hnp] at h
diff --git a/tests/lean/run/split2.lean b/tests/lean/run/split2.lean
index 434a655f8f..178fcb61aa 100644
--- a/tests/lean/run/split2.lean
+++ b/tests/lean/run/split2.lean
@@ -7,7 +7,7 @@ def f (x y z : Nat) : Nat :=
 
 example (x y z : Nat) : x ≠ 5 → y ≠ 5 → z ≠ 5 → f x y z = 1 := by
   intros
-  simp [f]
+  simp (config := { iota := false }) [f]
   split
   · contradiction
   · contradiction