feat: fine-grained equational lemmas for non-recursive functions

src/Init/Data/Int/Order.lean

@@ -1006,7 +1006,7 @@ theorem natAbs_mul_self : ∀ {a : Int}, ↑(natAbs a * natAbs a) = a * a
 theorem eq_nat_or_neg (a : Int) : ∃ n : Nat, a = n ∨ a = -↑n := ⟨_, natAbs_eq a⟩
 
 theorem natAbs_mul_natAbs_eq {a b : Int} {c : Nat}
-    (h : a * b = (c : Int)) : a.natAbs * b.natAbs = c := by rw [← natAbs_mul, h, natAbs]
+    (h : a * b = (c : Int)) : a.natAbs * b.natAbs = c := by rw [← natAbs_mul, h, natAbs.eq_def]
 
 @[simp] theorem natAbs_mul_self' (a : Int) : (natAbs a * natAbs a : Int) = a * a := by
   rw [← Int.ofNat_mul, natAbs_mul_self]

src/Init/Omega/Int.lean

@@ -116,7 +116,7 @@ theorem ofNat_max (a b : Nat) : ((max a b : Nat) : Int) = max (a : Int) (b : Int
   split <;> rfl
 
 theorem ofNat_natAbs (a : Int) : (a.natAbs : Int) = if 0 ≤ a then a else -a := by
-  rw [Int.natAbs]
+  rw [Int.natAbs.eq_def]
   split <;> rename_i n
   · simp only [Int.ofNat_eq_coe]
     rw [if_pos (Int.ofNat_nonneg n)]

src/Init/Tactics.lean

@@ -552,9 +552,9 @@ The `simp` tactic uses lemmas and hypotheses to simplify the main goal target or
 non-dependent hypotheses. It has many variants:
 - `simp` simplifies the main goal target using lemmas tagged with the attribute `[simp]`.
 - `simp [h₁, h₂, ..., hₙ]` simplifies the main goal target using the lemmas tagged
-  with the attribute `[simp]` and the given `hᵢ`'s, where the `hᵢ`'s are expressions.
-  If an `hᵢ` is a defined constant `f`, then the equational lemmas associated with
-  `f` are used. This provides a convenient way to unfold `f`.
+  with the attribute `[simp]` and the given `hᵢ`'s, where the `hᵢ`'s are expressions.-
+- If an `hᵢ` is a defined constant `f`, then `f` is unfolded. If `f` has equational lemmas associated
+  with it (and is not a projection or a `reducible` definition), these are used to rewrite with `f`.
 - `simp [*]` simplifies the main goal target using the lemmas tagged with the
   attribute `[simp]` and all hypotheses.
 - `simp only [h₁, h₂, ..., hₙ]` is like `simp [h₁, h₂, ..., hₙ]` but does not use `[simp]` lemmas.

src/Lean/Elab/PreDefinition.lean

@@ -10,3 +10,4 @@ import Lean.Elab.PreDefinition.Main
 import Lean.Elab.PreDefinition.MkInhabitant
 import Lean.Elab.PreDefinition.WF
 import Lean.Elab.PreDefinition.Eqns
+import Lean.Elab.PreDefinition.Nonrec.Eqns

src/Lean/Elab/PreDefinition/Basic.lean

@@ -12,6 +12,7 @@ import Lean.Util.NumApps
 import Lean.PrettyPrinter
 import Lean.Meta.AbstractNestedProofs
 import Lean.Meta.ForEachExpr
+import Lean.Meta.Eqns
 import Lean.Elab.RecAppSyntax
 import Lean.Elab.DefView
 import Lean.Elab.PreDefinition.TerminationHint
@@ -153,6 +154,7 @@ private def addNonRecAux (preDef : PreDefinition) (compile : Bool) (all : List N
     if compile && shouldGenCodeFor preDef then
       discard <| compileDecl decl
     if applyAttrAfterCompilation then
+      generateEagerEqns preDef.declName
       applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
 
 def addAndCompileNonRec (preDef : PreDefinition) (all : List Name := [preDef.declName]) : TermElabM Unit := do

src/Lean/Elab/PreDefinition/Eqns.lean

@@ -75,7 +75,12 @@ private def findMatchToSplit? (env : Environment) (e : Expr) (declNames : Array
           break
       unless hasFVarDiscr do
         return Expr.FindStep.visit
-      -- At least one alternative must contain a `declNames` application with loose bound variables.
+      -- For non-recursive functions (`declNames` empty), we split here
+      if declNames.isEmpty then
+          return Expr.FindStep.found
+      -- For recursive functions we only split when at least one alternatives contains a `declNames`
+      -- application with loose bound variables.
+      -- (We plan to disable this by default and treat recursive and non-recursie functions the same)
       for i in [info.getFirstAltPos : info.getFirstAltPos + info.numAlts] do
         let alt := args[i]!
         if Option.isSome <| alt.find? fun e => declNames.any e.isAppOf && e.hasLooseBVars then

src/Lean/Elab/PreDefinition/Nonrec/Eqns.lean

@@ -0,0 +1,108 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura, Joachim Breitner
+-/
+prelude
+import Lean.Meta.Tactic.Rewrite
+import Lean.Meta.Tactic.Split
+import Lean.Elab.PreDefinition.Basic
+import Lean.Elab.PreDefinition.Eqns
+import Lean.Meta.ArgsPacker.Basic
+import Init.Data.Array.Basic
+
+namespace Lean.Elab.Nonrec
+open Meta
+open Eqns
+
+/--
+Simple, coarse-grained equation theorem for nonrecursive definitions.
+-/
+private def mkSimpleEqThm (declName : Name) (suffix := Name.mkSimple unfoldThmSuffix) : MetaM (Option Name) := do
+  if let some (.defnInfo info) := (← getEnv).find? declName then
+    lambdaTelescope (cleanupAnnotations := true) info.value fun xs body => do
+      let lhs := mkAppN (mkConst info.name <| info.levelParams.map mkLevelParam) xs
+      let type  ← mkForallFVars xs (← mkEq lhs body)
+      let value ← mkLambdaFVars xs (← mkEqRefl lhs)
+      let name := declName ++ suffix
+      addDecl <| Declaration.thmDecl {
+        name, type, value
+        levelParams := info.levelParams
+      }
+      return some name
+  else
+    return none
+
+private partial def mkProof (declName : Name) (type : Expr) : MetaM Expr := do
+  trace[Elab.definition.eqns] "proving: {type}"
+  withNewMCtxDepth do
+    let main ← mkFreshExprSyntheticOpaqueMVar type
+    let (_, mvarId) ← main.mvarId!.intros
+    let rec go (mvarId : MVarId) : MetaM Unit := do
+      trace[Elab.definition.eqns] "step\n{MessageData.ofGoal mvarId}"
+      if ← withAtLeastTransparency .all (tryURefl mvarId) then
+        return ()
+      else if (← tryContradiction mvarId) then
+        return ()
+      else if let some mvarId ← simpMatch? mvarId then
+        go mvarId
+      else if let some mvarId ← simpIf? mvarId then
+        go mvarId
+      else if let some mvarId ← whnfReducibleLHS? mvarId then
+        go mvarId
+      else match (← simpTargetStar mvarId { config.dsimp := false } (simprocs := {})).1 with
+        | TacticResultCNM.closed => return ()
+        | TacticResultCNM.modified mvarId => go mvarId
+        | TacticResultCNM.noChange =>
+          if let some mvarIds ← casesOnStuckLHS? mvarId then
+            mvarIds.forM go
+          else if let some mvarIds ← splitTarget? mvarId then
+            mvarIds.forM go
+          else
+            throwError "failed to generate equational theorem for '{declName}'\n{MessageData.ofGoal mvarId}"
+
+    -- Try rfl before deltaLHS to avoid `id` checkpoints in the proof, which would make
+    -- the lemma ineligible for dsimp
+    unless ← withAtLeastTransparency .all (tryURefl mvarId) do
+      go (← deltaLHS mvarId)
+    instantiateMVars main
+
+def mkEqns (declName : Name) (info : DefinitionVal) : MetaM (Array Name) :=
+  withOptions (tactic.hygienic.set · false) do
+  let baseName := declName
+  let eqnTypes ← withNewMCtxDepth <| lambdaTelescope (cleanupAnnotations := true) info.value fun xs body => do
+    let us := info.levelParams.map mkLevelParam
+    let target ← mkEq (mkAppN (Lean.mkConst declName us) xs) body
+    let goal ← mkFreshExprSyntheticOpaqueMVar target
+    withReducible do
+      mkEqnTypes #[] goal.mvarId!
+  let mut thmNames := #[]
+  for i in [: eqnTypes.size] do
+    let type := eqnTypes[i]!
+    trace[Elab.definition.eqns] "eqnType[{i}]: {eqnTypes[i]!}"
+    let name := (Name.str baseName eqnThmSuffixBase).appendIndexAfter (i+1)
+    thmNames := thmNames.push name
+    let value ← mkProof declName type
+    let (type, value) ← removeUnusedEqnHypotheses type value
+    addDecl <| Declaration.thmDecl {
+      name, type, value
+      levelParams := info.levelParams
+    }
+  return thmNames
+
+def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
+  if (← isRecursiveDefinition declName) then
+    return none
+  if let some (.defnInfo info) := (← getEnv).find? declName then
+    if eqns.nonrecursive.get (← getOptions) then
+      mkEqns declName info
+    else
+      let o ← mkSimpleEqThm declName
+      return o.map (#[·])
+  else
+    return none
+
+builtin_initialize
+  registerGetEqnsFn getEqnsFor?
+
+end Lean.Elab.Nonrec

src/Lean/Elab/PreDefinition/Structural/Main.lean

@@ -193,6 +193,7 @@ def structuralRecursion (preDefs : Array PreDefinition) (termArg?s : Array (Opti
         registerEqnsInfo preDef (preDefs.map (·.declName)) recArgPos numFixed
     addSmartUnfoldingDef preDef recArgPos
     markAsRecursive preDef.declName
+    generateEagerEqns preDef.declName
   applyAttributesOf preDefsNonRec AttributeApplicationTime.afterCompilation
 
 

src/Lean/Elab/PreDefinition/WF/Main.lean

@@ -145,6 +145,7 @@ def wfRecursion (preDefs : Array PreDefinition) (termArg?s : Array (Option Termi
   registerEqnsInfo preDefs preDefNonRec.declName fixedPrefixSize argsPacker
   for preDef in preDefs do
     markAsRecursive preDef.declName
+    generateEagerEqns preDef.declName
     applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
     -- Unless the user asks for something else, mark the definition as irreducible
     unless preDef.modifiers.attrs.any fun a =>

src/Lean/Elab/Print.lean

@@ -134,7 +134,7 @@ private def printAxiomsOf (constName : Name) : CommandElabM Unit := do
   | _ => throwUnsupportedSyntax
 
 private def printEqnsOf (constName : Name) : CommandElabM Unit := do
-  let some eqns ← liftTermElabM <| Meta.getEqnsFor? constName (nonRec := true) |
+  let some eqns ← liftTermElabM <| Meta.getEqnsFor? constName |
     logInfo m!"'{constName}' does not have equations"
   let mut m := m!"equations:"
   for eq in eqns do

src/Lean/Elab/Tactic/Rewrite.lean

@@ -52,9 +52,11 @@ def withRWRulesSeq (token : Syntax) (rwRulesSeqStx : Syntax) (x : (symm : Bool)
           else
             -- Try to get equation theorems for `id`.
             let declName ← try realizeGlobalConstNoOverload id catch _ => return (← x symm term)
-            let some eqThms ← getEqnsFor? declName (nonRec := true) | x symm term
+            let some eqThms ← getEqnsFor? declName | x symm term
+            let hint := if eqThms.size = 1 then m!"" else
+              m!" Try rewriting with '{Name.str declName unfoldThmSuffix}'."
             let rec go : List Name →  TacticM Unit
-              | [] => throwError "failed to rewrite using equation theorems for '{declName}'"
+              | [] => throwError "failed to rewrite using equation theorems for '{declName}'.{hint}"
               -- Remark: we prefix `eqThm` with `_root_` to ensure it is resolved correctly.
               -- See test: `rwPrioritizesLCtxOverEnv.lean`
               | eqThm::eqThms => (x symm (mkIdentFrom id (`_root_ ++ eqThm))) <|> go eqThms

src/Lean/Meta/Eqns.lean

@@ -11,6 +11,23 @@ import Lean.Meta.AppBuilder
 import Lean.Meta.Match.MatcherInfo
 
 namespace Lean.Meta
+
+register_builtin_option eqns.nonrecursive : Bool := {
+    defValue := true
+    descr    := "Create fine-grained equational lemmas even for non-recursive definitions."
+  }
+
+
+/--
+These options affect the generation of equational theorems in a significant way. For these, their
+value at definition time, not realization time, should matter.
+
+This is implemented by
+ * eagerly realizing the equations when they are set to a non-default vaule
+ * when realizing them lazily, reset the options to their default
+-/
+def eqnAffectingOptions : Array (Lean.Option Bool) := #[eqns.nonrecursive]
+
 /--
 Environment extension for storing which declarations are recursive.
 This information is populated by the `PreDefinition` module, but the simplifier
@@ -105,7 +122,8 @@ def registerGetEqnsFn (f : GetEqnsFn) : IO Unit := do
 /-- Returns `true` iff `declName` is a definition and its type is not a proposition. -/
 private def shouldGenerateEqnThms (declName : Name) : MetaM Bool := do
   if let some (.defnInfo info) := (← getEnv).find? declName then
-    return !(← isProp info.type)
+    if (← isProp info.type) then return false
+    return true
   else
     return false
 
@@ -170,12 +188,7 @@ private partial def alreadyGenerated? (declName : Name) : MetaM (Option (Array N
   else
     return none
 
-/--
-Returns equation theorems for the given declaration.
-By default, we do not create equation theorems for nonrecursive definitions.
-You can use `nonRec := true` to override this behavior, a dummy `rfl` proof is created on the fly.
--/
-def getEqnsFor? (declName : Name) (nonRec := false) : MetaM (Option (Array Name)) := withLCtx {} {} do
+private def getEqnsFor?Core (declName : Name) : MetaM (Option (Array Name)) := withLCtx {} {} do
   if let some eqs := eqnsExt.getState (← getEnv) |>.map.find? declName then
     return some eqs
   else if let some eqs ← alreadyGenerated? declName then
@@ -185,13 +198,25 @@ def getEqnsFor? (declName : Name) (nonRec := false) : MetaM (Option (Array Name)
       if let some r ← f declName then
         registerEqnThms declName r
         return some r
-    if nonRec then
-      let some eqThm ← mkSimpleEqThm declName (suffix := Name.mkSimple eqn1ThmSuffix) | return none
-      let r := #[eqThm]
-      registerEqnThms declName r
-      return some r
   return none
 
+/--
+Returns equation theorems for the given declaration.
+-/
+def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := withLCtx {} {} do
+  -- This is the entry point for lazy equaion generation. Ignore the current value
+  -- of the options, and revert to the default.
+  withOptions (eqnAffectingOptions.foldl fun os o => o.set os o.defValue) do
+    getEqnsFor?Core declName
+
+/--
+If any equation theorem affecting option is not the default value, create the equations now.
+-/
+def generateEagerEqns (declName : Name) : MetaM Unit := do
+  let opts ← getOptions
+  if eqnAffectingOptions.any fun o => o.get opts != o.defValue then
+    let _ ← getEqnsFor?Core declName
+
 def GetUnfoldEqnFn := Name → MetaM (Option Name)
 
 private builtin_initialize getUnfoldEqnFnsRef : IO.Ref (List GetUnfoldEqnFn) ← IO.mkRef []
@@ -251,7 +276,7 @@ builtin_initialize
     let .str p s := name | return false
     unless (← getEnv).isSafeDefinition p do return false
     if isEqnReservedNameSuffix s then
-      return (← MetaM.run' <| getEqnsFor? p (nonRec := true)).isSome
+      return (← MetaM.run' <| getEqnsFor? p).isSome
     if isUnfoldReservedNameSuffix s then
       return (← MetaM.run' <| getUnfoldEqnFor? p (nonRec := true)).isSome
     return false

src/Lean/Meta/Tactic/Simp/Attr.lean

@@ -30,11 +30,13 @@ def mkSimpAttr (attrName : Name) (attrDescr : String) (ext : SimpExtension)
           if (← isProp info.type) then
             addSimpTheorem ext declName post (inv := false) attrKind prio
           else if info.hasValue then
-            if let some eqns ← getEqnsFor? declName then
+            if (← SimpTheorems.ignoreEquations declName) then
+              ext.add (SimpEntry.toUnfold declName) attrKind
+            else if let some eqns ← getEqnsFor? declName then
               for eqn in eqns do
                 addSimpTheorem ext eqn post (inv := false) attrKind prio
               ext.add (SimpEntry.toUnfoldThms declName eqns) attrKind
-              if hasSmartUnfoldingDecl (← getEnv) declName then
+              if (← SimpTheorems.unfoldEvenWithEqns declName) then
                 ext.add (SimpEntry.toUnfold declName) attrKind
             else
               ext.add (SimpEntry.toUnfold declName) attrKind

src/Lean/Meta/Tactic/Simp/SimpTheorems.lean

@@ -463,8 +463,38 @@ def SimpTheorems.add (s : SimpTheorems) (id : Origin) (levelParams : Array Name)
     let simpThms ← mkSimpTheorems id levelParams proof post inv prio
     return simpThms.foldl addSimpTheoremEntry s
 
+/--
+Reducible functions and projection functions should always be put in `toUnfold`, instead
+of trying to use equational theorems.
+
+The simplifiers has special support for structure and class projections, and gets
+confused when they suddenly rewrite, so ignore equations for them
+-/
+def SimpTheorems.ignoreEquations (declName : Name) : CoreM Bool := do
+  return (← isProjectionFn declName) || (← isReducible declName)
+
+/--
+Even if a function has equation theorems,
+we also store it in the `toUnfold` set in the following two cases:
+1- It was defined by structural recursion and has a smart-unfolding associated declaration.
+2- It is non-recursive.
+
+Reason: `unfoldPartialApp := true` or conditional equations may not apply.
+
+Remark: In the future, we are planning to disable this
+behavior unless `unfoldPartialApp := true`.
+Moreover, users will have to use `f.eq_def` if they want to force the definition to be
+unfolded.
+-/
+def SimpTheorems.unfoldEvenWithEqns (declName : Name) : CoreM Bool := do
+  if hasSmartUnfoldingDecl (← getEnv) declName then return true
+  unless (← isRecursiveDefinition declName) do return true
+  return false
+
 def SimpTheorems.addDeclToUnfold (d : SimpTheorems) (declName : Name) : MetaM SimpTheorems := do
-  if let some eqns ← getEqnsFor? declName then
+  if (← ignoreEquations declName) then
+    return d.addDeclToUnfoldCore declName
+  else if let some eqns ← getEqnsFor? declName then
     let mut d := d
     for h : i in [:eqns.size] do
       let eqn := eqns[i]
@@ -484,20 +514,7 @@ def SimpTheorems.addDeclToUnfold (d : SimpTheorems) (declName : Name) : MetaM Si
       else
         100 - i
       d ← SimpTheorems.addConst d eqn (prio := prio)
-    /-
-    Even if a function has equation theorems,
-    we also store it in the `toUnfold` set in the following two cases:
-    1- It was defined by structural recursion and has a smart-unfolding associated declaration.
-    2- It is non-recursive.
-
-    Reason: `unfoldPartialApp := true` or conditional equations may not apply.
-
-    Remark: In the future, we are planning to disable this
-    behavior unless `unfoldPartialApp := true`.
-    Moreover, users will have to use `f.eq_def` if they want to force the definition to be
-    unfolded.
-    -/
-    if hasSmartUnfoldingDecl (← getEnv) declName || !(← isRecursiveDefinition declName) then
+    if (← unfoldEvenWithEqns declName) then
       d := d.addDeclToUnfoldCore declName
     return d
   else