From 1af02dcaca3ca0665a37254aaa30ea196ec77c90 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 17 Mar 2021 19:11:16 -0700
Subject: [PATCH] feat: allow users to mark definitions with `[simp]`

cc @JasonGross @Kha
---
 src/Lean/Elab/Tactic/Simp.lean       |  5 ++---
 src/Lean/Meta/Tactic/Simp/Main.lean  |  4 ++--
 src/Lean/Meta/Tactic/Simp/Types.lean |  1 -
 tests/lean/run/simpDefToUnfold.lean  | 20 ++++++++++++++++++++
 4 files changed, 24 insertions(+), 6 deletions(-)
 create mode 100644 tests/lean/run/simpDefToUnfold.lean

diff --git a/src/Lean/Elab/Tactic/Simp.lean b/src/Lean/Elab/Tactic/Simp.lean
index 7896625f0c..d20702a4d1 100644
--- a/src/Lean/Elab/Tactic/Simp.lean
+++ b/src/Lean/Elab/Tactic/Simp.lean
@@ -93,7 +93,6 @@ private def elabSimpLemmas (stx : Syntax) (ctx : Simp.Context) : TacticM Simp.Co
     -/
     withMainContext do
       let mut lemmas := ctx.simpLemmas
-      let mut toUnfold : NameSet := {}
       for arg in stx[1].getSepArgs do
         if arg.getKind == ``Lean.Parser.Tactic.simpErase then
           let declName ← resolveGlobalConstNoOverload arg[1].getId
@@ -112,13 +111,13 @@ private def elabSimpLemmas (stx : Syntax) (ctx : Simp.Context) : TacticM Simp.Co
               if (← isProp info.type) then
                 lemmas ← lemmas.addConst declName post
               else
-                toUnfold := toUnfold.insert declName
+                lemmas := lemmas.addDeclToUnfold declName
             else
               lemmas ← lemmas.add e post
           | _ =>
             let arg ← elabTerm arg[1] none (mayPostpone := false)
             lemmas ← lemmas.add arg post
-      return { ctx with simpLemmas := lemmas, toUnfold := toUnfold }
+      return { ctx with simpLemmas := lemmas }
 where
   resolveSimpIdLemma? (simpArgTerm : Syntax) : TacticM (Option Expr) := do
     if simpArgTerm.isIdent then
diff --git a/src/Lean/Meta/Tactic/Simp/Main.lean b/src/Lean/Meta/Tactic/Simp/Main.lean
index 4063689d58..e621e3a471 100644
--- a/src/Lean/Meta/Tactic/Simp/Main.lean
+++ b/src/Lean/Meta/Tactic/Simp/Main.lean
@@ -58,7 +58,7 @@ private def reduceProjFn? (e : Expr) : SimpM (Option Expr) := do
     | none => return none
     | some projInfo =>
       if projInfo.fromClass then
-        if (← read).toUnfold.contains cinfo.name then
+        if (← read).simpLemmas.isDeclToUnfold cinfo.name then
           -- We only unfold class projections when the user explicitly requested them to be unfolded.
           -- Recall that `unfoldDefinition?` has support for unfolding this kind of projection.
           withReducibleAndInstances <| unfoldDefinition? e
@@ -88,7 +88,7 @@ private def unfold? (e : Expr) : SimpM (Option Expr) := do
   let fName := f.constName!
   if (← isProjectionFn fName) then
     return none -- should be reduced by `reduceProjFn?`
-  if (← read).toUnfold.contains e.getAppFn.constName! then
+  if (← read).simpLemmas.isDeclToUnfold e.getAppFn.constName! then
     withDefault <| unfoldDefinition? e
   else
     return none
diff --git a/src/Lean/Meta/Tactic/Simp/Types.lean b/src/Lean/Meta/Tactic/Simp/Types.lean
index ea9b294670..e2b4bb4678 100644
--- a/src/Lean/Meta/Tactic/Simp/Types.lean
+++ b/src/Lean/Meta/Tactic/Simp/Types.lean
@@ -21,7 +21,6 @@ structure Context where
   config      : Config
   simpLemmas  : SimpLemmas
   congrLemmas : CongrLemmas
-  toUnfold    : NameSet := {}
   parent?     : Option Expr := none
 
 structure State where
diff --git a/tests/lean/run/simpDefToUnfold.lean b/tests/lean/run/simpDefToUnfold.lean
new file mode 100644
index 0000000000..62da0f7912
--- /dev/null
+++ b/tests/lean/run/simpDefToUnfold.lean
@@ -0,0 +1,20 @@
+@[simp] def f (x : Nat) := x + 1
+
+@[simp] def g (x : Nat) := id x
+
+def h (x : Nat) := g (g x)
+
+namespace Extra
+  attribute [scoped simp] h
+end Extra
+
+theorem ex1 : f (g (h x)) = x + 1 := by
+  simp -- did not unfold h
+  simp [h]
+
+theorem ex2 : f (g (h x)) = x + 1 := by
+  open Extra in simp -- unfold f,g,h
+
+theorem ex3 : f (g x) = x + 1 := by
+  simp [-f] -- did not unfold f
+  simp