From 93bfc64812449784dfcfad3bf195ceaf09170d51 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 25 Sep 2024 10:41:57 +0200
Subject: [PATCH] fix: call hierarchy into `initialize`

---
 src/Lean/DeclarationRange.lean |  5 +++--
 src/Lean/Elab/Declaration.lean | 23 ++++++++++++++++-------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/Lean/DeclarationRange.lean b/src/Lean/DeclarationRange.lean
index 551614398bf6..9a0e14465200 100644
--- a/src/Lean/DeclarationRange.lean
+++ b/src/Lean/DeclarationRange.lean
@@ -42,8 +42,9 @@ builtin_initialize declRangeExt : MapDeclarationExtension DeclarationRanges ←
 def addBuiltinDeclarationRanges (declName : Name) (declRanges : DeclarationRanges) : IO Unit :=
   builtinDeclRanges.modify (·.insert declName declRanges)
 
-def addDeclarationRanges [MonadEnv m] (declName : Name) (declRanges : DeclarationRanges) : m Unit :=
-  modifyEnv fun env => declRangeExt.insert env declName declRanges
+def addDeclarationRanges [Monad m] [MonadEnv m] (declName : Name) (declRanges : DeclarationRanges) : m Unit := do
+  unless declRangeExt.contains (← getEnv) declName do
+    modifyEnv fun env => declRangeExt.insert env declName declRanges
 
 def findDeclarationRangesCore? [Monad m] [MonadEnv m] (declName : Name) : m (Option DeclarationRanges) :=
   return declRangeExt.find? (← getEnv) declName
diff --git a/src/Lean/Elab/Declaration.lean b/src/Lean/Elab/Declaration.lean
index 720c1c899d2e..5df132532e6c 100644
--- a/src/Lean/Elab/Declaration.lean
+++ b/src/Lean/Elab/Declaration.lean
@@ -382,19 +382,28 @@ def elabMutual : CommandElab := fun stx => do
     for attrName in toErase do
       Attribute.erase declName attrName
 
-@[builtin_macro Lean.Parser.Command.«initialize»] def expandInitialize : Macro
+@[builtin_command_elab Lean.Parser.Command.«initialize»] def elabInitialize : CommandElab
   | stx@`($declModifiers:declModifiers $kw:initializeKeyword $[$id? : $type? ←]? $doSeq) => do
     let attrId := mkIdentFrom stx <| if kw.raw[0].isToken "initialize" then `init else `builtin_init
     if let (some id, some type) := (id?, type?) then
       let `(Parser.Command.declModifiersT| $[$doc?:docComment]? $[@[$attrs?,*]]? $(vis?)? $[unsafe%$unsafe?]?) := stx[0]
-        | Macro.throwErrorAt declModifiers "invalid initialization command, unexpected modifiers"
-      `($[unsafe%$unsafe?]? def initFn : IO $type := with_decl_name% ?$id do $doSeq
-        $[$doc?:docComment]? @[$attrId:ident initFn, $(attrs?.getD ∅),*] $(vis?)? opaque $id : $type)
+        | throwErrorAt declModifiers "invalid initialization command, unexpected modifiers"
+      let defStx ← `($[$doc?:docComment]? @[$attrId:ident initFn, $(attrs?.getD ∅),*] $(vis?)? opaque $id : $type)
+      let mut fullId := (← getCurrNamespace) ++ id.getId
+      if vis?.any (·.raw.isOfKind ``Parser.Command.private) then
+        fullId := mkPrivateName (← getEnv) fullId
+      -- We need to add `id`'s ranges *before* elaborating `initFn` (and then `id` itself) as
+      -- otherwise the info context created by `with_decl_name` will be incomplete and break the
+      -- call hierarchy
+      addDeclarationRanges fullId defStx
+      elabCommand (← `(
+        $[unsafe%$unsafe?]? def initFn : IO $type := with_decl_name% $(mkIdent fullId) do $doSeq
+        $defStx:command))
     else
       let `(Parser.Command.declModifiersT| $[$doc?:docComment]? ) := declModifiers
-        | Macro.throwErrorAt declModifiers "invalid initialization command, unexpected modifiers"
-      `($[$doc?:docComment]? @[$attrId:ident] def initFn : IO Unit := do $doSeq)
-  | _ => Macro.throwUnsupported
+        | throwErrorAt declModifiers "invalid initialization command, unexpected modifiers"
+      elabCommand (← `($[$doc?:docComment]? @[$attrId:ident] def initFn : IO Unit := do $doSeq))
+  | _ => throwUnsupportedSyntax
 
 builtin_initialize
   registerTraceClass `Elab.axiom