feat: explicit defeq attribute

src/Init/Data/Fin/Fold.lean

@@ -183,10 +183,7 @@ theorem foldrM_loop [Monad m] [LawfulMonad m] (f : Fin (n+1) → α → m α) (x
   | zero =>
     rw [foldrM_loop_zero, foldrM_loop_succ, pure_bind]
     conv => rhs; rw [←bind_pure (f 0 x)]
-    congr
-    try  -- TODO: block can be deleted after bootstrapping
-      funext
-      simp [foldrM_loop_zero]
+    rfl
   | succ i ih =>
     rw [foldrM_loop_succ, foldrM_loop_succ, bind_assoc]
     congr; funext; exact ih ..

src/Lean.lean

@@ -41,3 +41,4 @@ import Lean.PrivateName
 import Lean.PremiseSelection
 import Lean.Namespace
 import Lean.EnvExtension
+import Lean.DefEqAttrib

src/Lean/AddDecl.lean

@@ -64,13 +64,6 @@ Checks whether the declaration was originally declared as a theorem; see also
 def wasOriginallyTheorem (env : Environment) (declName : Name) : Bool :=
   getOriginalConstKind? env declName |>.map (· matches .thm) |>.getD false
 
--- HACK: remove together with MutualDef HACK when `[dsimp]` is introduced
-private def isSimpleRflProof (proof : Expr) : Bool :=
-  if let .lam _ _ proof _ := proof then
-    isSimpleRflProof proof
-  else
-    proof.isAppOfArity ``rfl 2
-
 private def looksLikeRelevantTheoremProofType (type : Expr) : Bool :=
   if let .forallE _ _ type _ := type then
     looksLikeRelevantTheoremProofType type
@@ -90,9 +83,6 @@ def addDecl (decl : Declaration) : CoreM Unit := do
   let (name, info, kind) ← match decl with
     | .thmDecl thm =>
       let exportProof := !(← getEnv).header.isModule ||
-        -- We should preserve rfl theorems but also we should not override a decision to hide by the
-        -- MutualDef elaborator via `withoutExporting`
-        (← getEnv).isExporting && isSimpleRflProof thm.value ||
         -- TODO: this is horrible...
         looksLikeRelevantTheoremProofType thm.type
       if !exportProof then

src/Lean/Attributes.lean

@@ -135,7 +135,9 @@ structure TagAttribute where
   deriving Inhabited
 
 def registerTagAttribute (name : Name) (descr : String)
-    (validate : Name → AttrM Unit := fun _ => pure ()) (ref : Name := by exact decl_name%) (applicationTime := AttributeApplicationTime.afterTypeChecking) : IO TagAttribute := do
+    (validate : Name → AttrM Unit := fun _ => pure ()) (ref : Name := by exact decl_name%)
+    (applicationTime := AttributeApplicationTime.afterTypeChecking)
+    (asyncMode : EnvExtension.AsyncMode := .mainOnly) : IO TagAttribute := do
   let ext : PersistentEnvExtension Name Name NameSet ← registerPersistentEnvExtension {
     name            := ref
     mkInitial       := pure {}
@@ -145,6 +147,12 @@ def registerTagAttribute (name : Name) (descr : String)
       let r : Array Name := es.fold (fun a e => a.push e) #[]
       r.qsort Name.quickLt
     statsFn         := fun s => "tag attribute" ++ Format.line ++ "number of local entries: " ++ format s.size
+    asyncMode       := asyncMode
+    replay?         := some fun _ newState newConsts s =>
+      newConsts.foldl (init := s) fun s c =>
+        if newState.contains c then
+          s.insert c
+        else s
   }
   let attrImpl : AttributeImpl := {
     ref, name, descr, applicationTime
@@ -153,7 +161,9 @@ def registerTagAttribute (name : Name) (descr : String)
       unless kind == AttributeKind.global do throwError "invalid attribute '{name}', must be global"
       let env ← getEnv
       unless (env.getModuleIdxFor? decl).isNone do
-        throwError "invalid attribute '{name}', declaration is in an imported module"
+        throwError "invalid attribute '{name}', declaration {.ofConstName decl} is in an imported module"
+      unless env.asyncMayContain decl do
+        throwError "invalid attribute '{name}', declaration {.ofConstName decl} is not from the present async context {env.asyncPrefix?}"
       validate decl
       modifyEnv fun env => ext.addEntry env decl
   }
@@ -162,10 +172,25 @@ def registerTagAttribute (name : Name) (descr : String)
 
 namespace TagAttribute
 
+/-- Sets the attribute (without running `validate`) -/
+def setTag  [Monad m] [MonadError m] [MonadEnv m] (attr : TagAttribute) (decl : Name) : m Unit := do
+  let env ← getEnv
+  unless (env.getModuleIdxFor? decl).isNone do
+    throwError "invalid attribute '{attr.attr.name}', declaration {.ofConstName decl} is in an imported module"
+  unless env.asyncMayContain decl do
+    throwError "invalid attribute '{attr.attr.name}', declaration {.ofConstName decl} is not from the present async context {env.asyncPrefix?}"
+  modifyEnv fun env => attr.ext.addEntry env decl
+
 def hasTag (attr : TagAttribute) (env : Environment) (decl : Name) : Bool :=
   match env.getModuleIdxFor? decl with
   | some modIdx => (attr.ext.getModuleEntries env modIdx).binSearchContains decl Name.quickLt
-  | none        => (attr.ext.getState env).contains decl
+  | none        =>
+    if attr.ext.toEnvExtension.asyncMode matches .async then
+      -- TODO: This is obviously not correct. To discuss with Sebastian.
+      (attr.ext.findStateAsync env decl).contains decl ||
+      (attr.ext.getState env (asyncMode := .local)).contains decl
+    else
+      (attr.ext.getState env).contains decl
 
 end TagAttribute
 

src/Lean/Declaration.lean

@@ -482,7 +482,7 @@ def value! (info : ConstantInfo) (allowOpaque := false) : Expr :=
   | .defnInfo {value, ..}   => value
   | .thmInfo  {value, ..}   => value
   | .opaqueInfo {value, ..} => if allowOpaque then value else panic! "declaration with value expected"
-  | _                       => panic! "declaration with value expected"
+  | _                       => panic! s!"declaration with value expected, but {info.name} has none"
 
 def hints : ConstantInfo → ReducibilityHints
   | .defnInfo {hints, ..} => hints

src/Lean/DefEqAttrib.lean

@@ -0,0 +1,108 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+prelude
+
+import Lean.PrettyPrinter
+
+namespace Lean
+open Meta
+
+/--
+There are defeq theorems that only hold at transparency `.all`, but also others that hold
+(from the kernel's point of view) but where the defeq checker here will run out of cycles.
+
+So we try the more careful first.
+-/
+private def isDefEqCareful (e1 e2 : Expr) : MetaM Bool := do
+  withOptions (smartUnfolding.set · false) <| do
+    withDefault (isDefEq e1 e2) <||> withTransparency .all (isDefEq e1 e2)
+
+def validateDefEqAttr (declName : Name) : AttrM Unit := do
+  let info ← getConstVal declName
+  MetaM.run' do
+    withTransparency .all do -- we want to look through defs in `info.type` all the way to `Eq`
+      forallTelescopeReducing info.type fun _ type => do
+        let type ← whnf type
+        -- NB: The warning wording should work both for explicit uses of `@[defeq]` as well as the implicit `:= rfl`.
+        let some (_, lhs, rhs) := type.eq? |
+          throwError m!"Not a definitional equality: the conclusion should be an equality, but is{inlineExpr type}"
+        let ok ← isDefEqCareful lhs rhs
+        unless ok do
+          let explanation := MessageData.ofLazyM (es := #[lhs, rhs]) do
+            let (lhs, rhs) ← addPPExplicitToExposeDiff lhs rhs
+            let mut msg := m!"Not a definitional equality: the left-hand side{indentExpr lhs}\nis \
+              not definitionally equal to the right-hand side{indentExpr rhs}"
+            if (← getEnv).isExporting then
+              let okPrivately ← withoutExporting <| isDefEqCareful lhs rhs
+              if okPrivately then
+                msg := msg ++ .note m!"This theorem is exported from the current module. \
+                  This requires that all definitions that need to be unfolded to prove this \
+                  theorem must be exposed."
+            pure msg
+          throwError explanation
+
+/--
+Marks the theorem as a definitional equality.
+
+The theorem must be an equality that holds by `rfl`. This allows `dsimp` to use this theorem
+when rewriting.
+
+A theorem with a body is written exactly as `:= rfl` is implicitly marked `@[defeq]`. To avoid
+this behavior, write `:= (rfl)` instead.
+
+The attribute should be given before a `@[simp]` attribute to have effect.
+
+When using the module system, an exported theorem can only be `@[defeq]` if all definitions that
+need to be unfolded to prove the theorem are also exported.
+-/
+@[builtin_doc]
+builtin_initialize defeqAttr : TagAttribute ←
+  registerTagAttribute `defeq "mark theorem as a definitional equality, to be used by `dsimp`"
+    (validate := validateDefEqAttr) (applicationTime := .afterTypeChecking)
+    (asyncMode := .async)
+
+private partial def isRflProofCore (type : Expr) (proof : Expr) : CoreM Bool := do
+  match type with
+  | .forallE _ _ type _ =>
+    if let .lam _ _ proof _ := proof then
+      isRflProofCore type proof
+    else
+      return false
+  | _ =>
+    if type.isAppOfArity ``Eq 3 then
+      if proof.isAppOfArity ``Eq.refl 2 || proof.isAppOfArity ``rfl 2 then
+        return true
+      else if proof.isAppOfArity ``Eq.symm 4 then
+        -- `Eq.symm` of rfl proof is a rfl proof
+        isRflProofCore type proof.appArg! -- small hack: we don't need to set the exact type
+      else if proof.getAppFn.isConst then
+        -- The application of a `defeq` theorem is a `rfl` proof
+        return defeqAttr.hasTag (← getEnv) proof.getAppFn.constName!
+      else
+        return false
+    else
+      return false
+
+/--
+For automatically generated theorems (equational theorems etc.), we want to set the `defeq` attribute
+if the proof is `rfl`, essentially reproducing the behavior before the introduction of the `defeq`
+attribute. This function infers the `defeq` attribute based on the declaration value.
+-/
+def inferDefEqAttr (declName : Name) : MetaM Unit := do
+  withExporting (isExporting := !isPrivateName declName) do
+    let info ← getConstInfo declName
+    let isRfl ←
+      if let some value := info.value? then
+        isRflProofCore info.type value
+      else
+        pure false
+    if isRfl then
+      try
+        validateDefEqAttr declName -- sanity-check: would we have accepted `@[defeq]` on this?
+      catch e =>
+        logError m!"Theorem {declName} has a `rfl`-proof and was thus inferred to be `@[defeq]`, \
+          but validating that attribute failed:{indentD e.toMessageData}"
+      defeqAttr.setTag declName

src/Lean/Elab/DeclModifiers.lean

@@ -84,10 +84,14 @@ def Modifiers.isNonrec : Modifiers → Bool
   | { recKind := .nonrec, .. } => true
   | _                          => false
 
-/-- Adds attribute `attr` in `modifiers` -/
+/-- Adds attribute `attr` in `modifiers`, at the end -/
 def Modifiers.addAttr (modifiers : Modifiers) (attr : Attribute) : Modifiers :=
   { modifiers with attrs := modifiers.attrs.push attr }
 
+/-- Adds attribute `attr` in `modifiers`, at the beginning -/
+def Modifiers.addFirstAttr (modifiers : Modifiers) (attr : Attribute) : Modifiers :=
+  { modifiers with attrs := #[attr] ++ modifiers.attrs }
+
 /-- Filters attributes using `p` -/
 def Modifiers.filterAttrs (modifiers : Modifiers) (p : Attribute → Bool) : Modifiers :=
   { modifiers with attrs := modifiers.attrs.filter p }

src/Lean/Elab/DefView.lean

@@ -127,6 +127,11 @@ structure DefView where
 def DefView.isInstance (view : DefView) : Bool :=
   view.modifiers.attrs.any fun attr => attr.name == `instance
 
+/-- Prepends the `defeq` attribute, removing existing ones if there are any -/
+def DefView.markDefEq (view : DefView) : DefView :=
+  { view with modifiers :=
+      view.modifiers.filterAttrs (·.name != `defeq) |>.addFirstAttr { name := `defeq } }
+
 namespace Command
 open Meta
 

src/Lean/Elab/MutualDef.lean

@@ -1072,34 +1072,20 @@ where
     withExporting (isExporting := !expandedDeclIds.all (isPrivateName ·.declName)) do
     let headers ← elabHeaders views expandedDeclIds bodyPromises tacPromises
     let headers ← levelMVarToParamHeaders views headers
-    -- If the decl looks like a `rfl` theorem, we elaborate is synchronously as we need to wait for
-    -- the type before we can decide whether the theorem body should be exported and then waiting
-    -- for the body as well should not add any significant overhead.
-    let isRflLike := headers.all (·.value matches `(declVal| := rfl))
-    -- elaborate body in parallel when all stars align
     if let (#[view], #[declId]) := (views, expandedDeclIds) then
-      if Elab.async.get (← getOptions) && view.kind.isTheorem && !isRflLike &&
+      if Elab.async.get (← getOptions) && view.kind.isTheorem &&
           !deprecated.oldSectionVars.get (← getOptions) &&
           -- holes in theorem types is not a fatal error, but it does make parallelism impossible
           !headers[0]!.type.hasMVar then
         elabAsync headers[0]! view declId
-      else elabSync headers isRflLike
-    else elabSync headers isRflLike
+      else elabSync headers
+    else elabSync headers
     for view in views, declId in expandedDeclIds do
       -- NOTE: this should be the full `ref`, and thus needs to be done after any snapshotting
       -- that depends only on a part of the ref
       addDeclarationRangesForBuiltin declId.declName view.modifiers.stx view.ref
-  elabSync headers isRflLike := do
-    -- If the reflexivity holds publicly as well (we're still inside `withExporting` here), export
-    -- the body even if it is a theorem so that it is recognized as a rfl theorem even without
-    -- `import all`.
-    let rflPublic ← pure isRflLike <&&> pure (← getEnv).header.isModule <&&>
-      forallTelescopeReducing headers[0]!.type fun _ type => do
-        let some (_, lhs, rhs) := type.eq? | pure false
-        try
-          isDefEq lhs rhs
-        catch _ => pure false
-    finishElab (isExporting := rflPublic) headers
+  elabSync headers := do
+    finishElab headers
     processDeriving headers
   elabAsync header view declId := do
     let env ← getEnv
@@ -1147,7 +1133,7 @@ where
         (cancelTk? := cancelTk) fun _ => do profileitM Exception "elaboration" (← getOptions) do
       setEnv async.asyncEnv
       try
-        finishElab (isExporting := false) #[header]
+        finishElab #[header]
       finally
         reportDiag
         -- must introduce node to fill `infoHole` with multiple info trees
@@ -1279,6 +1265,8 @@ def elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
     if ds.size > 1 && modifiers.isNonrec then
       throwErrorAt d "invalid use of 'nonrec' modifier in 'mutual' block"
     let mut view ← mkDefView modifiers d[1]
+    if view.kind != .example && view.value matches `(declVal| := rfl) then
+      view := view.markDefEq
     let fullHeaderRef := mkNullNode #[d[0], view.headerRef]
     if let some snap := snap? then
       view := { view with headerSnap? := some {

src/Lean/Elab/PreDefinition/EqUnfold.lean

@@ -9,6 +9,7 @@ import Lean.Meta.Tactic.Util
 import Lean.Meta.Tactic.Rfl
 import Lean.Meta.Tactic.Intro
 import Lean.Meta.Tactic.Apply
+import Lean.DefEqAttrib
 
 namespace Lean.Meta
 
@@ -53,6 +54,7 @@ def getConstUnfoldEqnFor? (declName : Name) : MetaM (Option Name) := do
       name, type, value
       levelParams := info.levelParams
     }
+    inferDefEqAttr name
   return some name
 
 

src/Lean/Elab/PreDefinition/Eqns.lean

@@ -12,6 +12,7 @@ import Lean.Meta.Tactic.Split
 import Lean.Meta.Tactic.Apply
 import Lean.Meta.Tactic.Refl
 import Lean.Meta.Match.MatchEqs
+import Lean.DefEqAttrib
 
 namespace Lean.Elab.Eqns
 open Meta
@@ -429,6 +430,7 @@ where
       name, type, value
       levelParams := info.levelParams
     }
+    inferDefEqAttr name
 
 /--
   Auxiliary method for `mkUnfoldEq`. The structure is based on `mkEqnTypes`.

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

@@ -36,6 +36,7 @@ where
         name, type, value
         levelParams := info.levelParams
       }
+      inferDefEqAttr name
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
   if (← isRecursiveDefinition declName) then