feat: to_additive calls addDeclarationRanges (#461)

Co-authored-by: Scott Morrison <scott.morrison@gmail.com>
Co-authored-by: Mario Carneiro <di.gama@gmail.com>

Author: kim-em

Mathlib/Tactic/ToAdditive.lean

@@ -10,6 +10,7 @@ import Mathlib.Lean.Expr.ReplaceRec
 import Mathlib.Lean.Expr
 import Lean
 import Lean.Data
+import Lean.Elab.Term
 import Std.Lean.NameMapAttribute
 
 /-!
@@ -166,7 +167,7 @@ def applyReplacementFun : Expr → MetaM Expr :=
   Lean.Expr.replaceRecMeta fun r e => do
     trace[to_additive_detail] "applyReplacementFun: replace at {e}"
     match e with
-    | .lit (.natVal 1) => pure <| mkNatLit 0
+    | .lit (.natVal 1) => pure <| mkRawNatLit 0
     | .const n₀ ls => do
       let n₁ := Name.mapPrefix (findTranslation? <|← getEnv) n₀
       trace[to_additive_detail] "applyReplacementFun: {n₀} → {n₁}"
@@ -258,15 +259,14 @@ using the transforms dictionary.
 `pre` is the declaration that got the `@[to_additive]` attribute and `tgt_pre` is the target of this
 declaration. -/
 partial def transformDeclAux
-  (pre tgt_pre : Name) : Name → CoreM Unit := fun src => do
+  (ref : Option Syntax) (pre tgt_pre : Name) : Name → CoreM Unit := fun src => do
   -- if this declaration is not `pre` or an internal declaration, we do nothing.
   if not (src == pre || isInternal' src) then
     if (findTranslation? (← getEnv) src).isSome then
-      return ()
-    else throwError
-       ("The declaration {pre} depends on the declaration {src} which is in the namespace {pre}," ++
-        " but does not have the `@[to_additive]` attribute. This is not supported. " ++
-        "Workaround: move {src} to a different namespace.")
+      return
+    throwError "The declaration {pre} depends on the declaration {src} which is in the namespace {
+      pre}, but does not have the `@[to_additive]` attribute. This is not supported.\n{""
+      }Workaround: move {src} to a different namespace."
   let env ← getEnv
   -- we find the additive name of `src`
   let tgt := src.mapPrefix (fun n => if n == pre then some tgt_pre else none)
@@ -276,10 +276,10 @@ partial def transformDeclAux
   let srcDecl ← getConstInfo src
   -- we first transform all the declarations of the form `pre._proof_i`
   for n in srcDecl.type.listNamesWithPrefix pre do
-    transformDeclAux pre tgt_pre n
+    transformDeclAux none pre tgt_pre n
   if let some value := srcDecl.value? then
     for n in value.listNamesWithPrefix pre do
-      transformDeclAux pre tgt_pre n
+      transformDeclAux none pre tgt_pre n
   -- now transform the source declaration
   let trgDecl : ConstantInfo ← MetaM.run' $ updateDecl tgt srcDecl
   if ¬ trgDecl.hasValue then
@@ -296,6 +296,18 @@ partial def transformDeclAux
       Failed to add declaration\n{trgDecl.name}:\n{msg}"
     | _ => panic! "unreachable"
   addAndCompile trgDecl.toDeclaration!
+  -- now add declaration ranges so jump-to-definition works
+  addDeclarationRanges tgt {
+    range := ← getDeclarationRange (← getRef)
+    selectionRange := ← getDeclarationRange (ref.getD (← getRef))
+  }
+  if let some ref := ref then
+    -- TODO: make a function for this
+    pushInfoLeaf <| .ofTermInfo {
+      elaborator := .anonymous, lctx := {}, expectedType? := none
+      stx := ref, isBinder := true
+      expr := ← mkConstWithLevelParams trgDecl.name
+    }
   if isProtected (← getEnv) src then
     setEnv $ addProtected (← getEnv) tgt
 
@@ -308,34 +320,33 @@ allow us to iterate the attributes applied to a given decalaration.
 -/
 def copyAttributes (src tgt : Name) : CoreM Unit := do
   -- [todo] other simp theorems
-  let some ext ← getSimpExtension? `simp | return ()
+  let some ext ← getSimpExtension? `simp | return
   let thms ← ext.getTheorems
   if (¬ thms.isLemma (.decl src)) || thms.isLemma (.decl tgt) then
-    return ()
+    return
   -- [todo] how to get prio data from SimpTheorems?
-  MetaM.run' $ Lean.Meta.addSimpTheorem ext tgt
+  Lean.Meta.addSimpTheorem ext tgt
     (post := true)
     (inv := false)
     (attrKind := AttributeKind.global)
-    (prio := 1000)
-  return ()
+    (prio := 1000) |>.run'
 
 /--
 Make a new copy of a declaration, replacing fragments of the names of identifiers in the type and
 the body using the `translations` dictionary.
 This is used to implement `@[to_additive]`.
 -/
-def transformDecl (src tgt : Name) : CoreM Unit := do
-  transformDeclAux src tgt src
+def transformDecl (ref : Option Syntax) (src tgt : Name) : CoreM Unit := do
+  transformDeclAux ref src tgt src
   let eqns? ← MetaM.run' (getEqnsFor? src true)
   -- now transform all of the equational lemmas
   if let some eqns := eqns? then
     for src_eqn in eqns do
-      transformDeclAux src tgt src_eqn
+      transformDeclAux none src tgt src_eqn
       -- [todo] copy attributes for equations
       -- [todo] add equation lemmas to tgt_eqn
   copyAttributes src tgt
-  return ()
+
 /--
 Find the first argument of `nm` that has a multiplicative type-class on it.
 Returns 1 if there are no types with a multiplicative class as arguments.
@@ -360,7 +371,6 @@ def firstMultiplicativeArg (nm : Name) : MetaM (Option Nat) := do
     | [] => return none
     | (head :: tail) => return some <| tail.foldr Nat.min head
 
-
 /-- `ValueType` is the type of the arguments that can be provided to `to_additive`. -/
 structure ValueType : Type where
   /-- Replace all multiplicative declarations, do not use the heuristic. -/
@@ -375,6 +385,10 @@ structure ValueType : Type where
   /-- If `allow_auto_name` is `false` (default) then
   `@[to_additive]` will check whether the given name can be auto-generated. -/
   allowAutoName : Bool := false
+  /-- The `Syntax` element corresponding to the original multiplicative declaration
+  (or the `to_additive` attribute if it is added later),
+  which we need for adding definition ranges. -/
+  ref : Syntax
   deriving Repr
 
 /-- `add_comm_prefix x s` returns `"comm_" ++ s` if `x = tt` and `s` otherwise. -/
@@ -464,18 +478,19 @@ def proceedFields (src tgt : Name) : CoreM Unit := do
   -- [todo] run to_additive on the constructors of n:
   -- aux (fun n => (env.constructorsOf n).mmap $ ...
 
-private def elabToAdditiveAux
-  (replaceAll trace : Bool) (tgt : Option Syntax) (doc : Option Syntax) : ValueType :=
+private def elabToAdditiveAux (ref : Syntax) (replaceAll trace : Bool) (tgt : Option Syntax)
+    (doc : Option Syntax) : ValueType :=
   { replaceAll := replaceAll
     trace := trace
     tgt := match tgt with | some tgt => tgt.getId | none => Name.anonymous
     doc := doc.bind (·.isStrLit?)
     allowAutoName := false
+    ref
   }
 
 private def elabToAdditive : Syntax → CoreM ValueType
-  | `(attr| to_additive $[!%$replaceAll]? $[?%$trace]? $[$tgt]? $[$doc]?) =>
-    return elabToAdditiveAux replaceAll.isSome trace.isSome tgt doc
+  | `(attr| to_additive%$tk $[!%$replaceAll]? $[?%$trace]? $[$tgt]? $[$doc]?) =>
+    return elabToAdditiveAux ((tgt.map (·.raw)).getD tk) replaceAll.isSome trace.isSome tgt doc
   | _ => throwUnsupportedSyntax
 
 /-!
@@ -701,10 +716,12 @@ initialize registerBuiltinAttribute {
         withOptions
           (fun o => o |>.setBool `to_additive.replaceAll val.replaceAll
                       |>.setBool `trace.to_additive shouldTrace)
-          (transformDecl src tgt)
+          (transformDecl val.ref src tgt)
       if let some doc := val.doc then
         addDocString tgt doc
-      return ()
+    -- Because `@[simp]` runs after compilation,
+    -- we have to as well to be able to copy attributes correctly.
+    applicationTime := .afterCompilation
   }
 
 

test/toAdditive.lean

@@ -95,7 +95,7 @@ if some_def.in_namespace then x * x else x
 
 -- cannot apply `@[to_additive]` to `some_def` if `some_def.in_namespace` doesn't have the attribute
 run_cmd do
-  Elab.Command.liftCoreM <| successIfFail (ToAdditive.transformDecl `Test.some_def `Test.add_some_def)
+  Elab.Command.liftCoreM <| successIfFail (ToAdditive.transformDecl (← getRef) `Test.some_def `Test.add_some_def)
 
 
 attribute [to_additive some_other_name] some_def.in_namespace