From ff49fd6b7ea9f8b7faf6bb037e13c46c3a4eeb96 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 5 Jan 2022 15:56:48 -0800
Subject: [PATCH] fix: apply `afterCompilation` attributes after we have
 compiled *all* definitions in a mutual block

---
 src/Lean/Elab/PreDefinition/Basic.lean   | 9 +++++----
 src/Lean/Elab/PreDefinition/WF/Main.lean | 6 ++++--
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/Lean/Elab/PreDefinition/Basic.lean b/src/Lean/Elab/PreDefinition/Basic.lean
index d62ca38b42..f33fb01bdb 100644
--- a/src/Lean/Elab/PreDefinition/Basic.lean
+++ b/src/Lean/Elab/PreDefinition/Basic.lean
@@ -92,7 +92,7 @@ private def compileDecl (decl : Declaration) : TermElabM Bool := do
       throw ex
   return true
 
-private def addNonRecAux (preDef : PreDefinition) (compile : Bool) : TermElabM Unit :=
+private def addNonRecAux (preDef : PreDefinition) (compile : Bool) (applyAttrAfterCompilation := true) : TermElabM Unit :=
   withRef preDef.ref do
     let preDef ← abstractNestedProofs preDef
     let env ← getEnv
@@ -118,13 +118,14 @@ private def addNonRecAux (preDef : PreDefinition) (compile : Bool) : TermElabM U
     if compile && shouldGenCodeFor preDef then
       unless (← compileDecl decl) do
         return ()
-    applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
+    if applyAttrAfterCompilation then
+      applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
 
 def addAndCompileNonRec (preDef : PreDefinition) : TermElabM Unit := do
   addNonRecAux preDef true
 
-def addNonRec (preDef : PreDefinition) : TermElabM Unit := do
-  addNonRecAux preDef false
+def addNonRec (preDef : PreDefinition) (applyAttrAfterCompilation := true) : TermElabM Unit := do
+  addNonRecAux preDef (compile := false) (applyAttrAfterCompilation := applyAttrAfterCompilation)
 
 /--
   Eliminate recursive application annotations containing syntax. These annotations are used by the well-founded recursion module
diff --git a/src/Lean/Elab/PreDefinition/WF/Main.lean b/src/Lean/Elab/PreDefinition/WF/Main.lean
index d9c2dded70..d7d44113cb 100644
--- a/src/Lean/Elab/PreDefinition/WF/Main.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Main.lean
@@ -45,7 +45,7 @@ private partial def addNonRecPreDefs (preDefs : Array PreDefinition) (preDefNonR
       let arg ← mkSum 0 domain
       mkLambdaFVars xs (mkApp (mkConst preDefNonRec.declName us) arg)
     trace[Elab.definition.wf] "{preDef.declName} := {value}"
-    addNonRec { preDef with value }
+    addNonRec { preDef with value } (applyAttrAfterCompilation := false)
 
 def wfRecursion (preDefs : Array PreDefinition) (wfStx? : Option Syntax) (decrTactic? : Option Syntax) : TermElabM Unit := do
   let unaryPreDef ← withoutModifyingEnv do
@@ -62,8 +62,10 @@ def wfRecursion (preDefs : Array PreDefinition) (wfStx? : Option Syntax) (decrTa
   let preDefs ← preDefs.mapM fun d => eraseRecAppSyntax d
   addNonRec preDefNonRec
   addNonRecPreDefs preDefs preDefNonRec
-  addAndCompilePartialRec preDefs
   registerEqnsInfo preDefs preDefNonRec.declName
+  for preDef in preDefs do
+    applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
+  addAndCompilePartialRec preDefs
 
 builtin_initialize registerTraceClass `Elab.definition.wf