feat: improve macroscope encoding

src/Init/Lean/Elab/BuiltinNotation.lean

@@ -86,7 +86,7 @@ adaptExpander $ fun stx => match_syntax stx with
   some declName ← getDeclName?
     | throwError stx "invalid `parser!` macro, it must be used in definitions";
   match extractMacroScopes declName with
-  | (Name.str _ s _, scps) => do
+  | { name := Name.str _ s _, scopes :=  scps, .. } => do
     let kind := quote declName;
     let s    := quote s;
     p ← `(Lean.Parser.leadingNode $kind $e);

src/Init/Lean/Elab/Command.lean

@@ -117,9 +117,8 @@ when (checkTraceOption opts cls) $ logTrace cls ref (msg ())
 def throwUnsupportedSyntax {α} : CommandElabM α :=
 throw Elab.Exception.unsupportedSyntax
 
-protected def getCurrMacroScope : CommandElabM Nat := do
-ctx ← read;
-pure ctx.currMacroScope
+protected def getCurrMacroScope : CommandElabM Nat  := do ctx ← read; pure ctx.currMacroScope
+protected def getMainModule     : CommandElabM Name := do env ← getEnv; pure env.mainModule
 
 @[inline] protected def withFreshMacroScope {α} (x : CommandElabM α) : CommandElabM α := do
 fresh ← modifyGet (fun st => (st.nextMacroScope, { st with nextMacroScope := st.nextMacroScope + 1 }));
@@ -127,6 +126,7 @@ adaptReader (fun (ctx : Context) => { ctx with currMacroScope := fresh }) x
 
 instance CommandElabM.MonadQuotation : MonadQuotation CommandElabM := {
   getCurrMacroScope   := Command.getCurrMacroScope,
+  getMainModule       := Command.getMainModule,
   withFreshMacroScope := @Command.withFreshMacroScope
 }
 
@@ -580,12 +580,11 @@ levelNames      ←
   };
 let ref := declId;
 -- extract (optional) namespace part of id, after decoding macro scopes that would interfere with the check
-let (id, scps) := extractMacroScopes id;
-match id with
-| Name.str pre s _ =>
+match extractMacroScopes id with
+| { name := Name.str pre s _, scopes := scps, mainModule := mainModule } =>
   /- Add back macro scopes. We assume a declaration like `def a.b[1,2] ...` with macro scopes `[1,2]`
      is always meant to mean `namespace a def b[1,2] ...`. -/
-  let id := addMacroScopes (mkNameSimple s) scps;
+  let id := addMacroScopes mainModule (mkNameSimple s) scps;
   withNamespace ref pre $ do
     modifyScope $ fun scope => { levelNames := levelNames, .. scope };
     finally (f id) (modifyScope $ fun scope => { levelNames := savedLevelNames, .. scope })

src/Init/Lean/Elab/ResolveName.lean

@@ -66,10 +66,11 @@ private def resolveOpenDecls (env : Environment) (id : Name) : List OpenDecl →
   let resolvedIds := if id == openedId then resolvedId :: resolvedIds else resolvedIds;
   resolveOpenDecls openDecls resolvedIds
 
-private def resolveGlobalNameAux (env : Environment) (ns : Name) (openDecls : List OpenDecl) (scps : List MacroScope) : Name → List String → List (Name × List String)
+private def resolveGlobalNameAux (env : Environment) (ns : Name) (openDecls : List OpenDecl)
+    (extractionResult : ExtractMacroScopesResult) : Name → List String → List (Name × List String)
 | id@(Name.str p s _), projs =>
   -- NOTE: we assume that macro scopes always belong to the projected constant, not the projections
-  let id := addMacroScopes id scps;
+  let id := addMacroScopes extractionResult.mainModule id extractionResult.scopes;
   match resolveUsingNamespace env id ns with
   | resolvedIds@(_ :: _) => resolvedIds.eraseDups.map $ fun id => (id, projs)
   | [] =>
@@ -86,8 +87,8 @@ private def resolveGlobalNameAux (env : Environment) (ns : Name) (openDecls : Li
 
 def resolveGlobalName (env : Environment) (ns : Name) (openDecls : List OpenDecl) (id : Name) : List (Name × List String) :=
 -- decode macro scopes from name before recursion
-let (id, scps) := extractMacroScopes id;
-resolveGlobalNameAux env ns openDecls scps id []
+let extractionResult := extractMacroScopes id;
+resolveGlobalNameAux env ns openDecls extractionResult extractionResult.name []
 
 /- Namespace resolution -/
 

src/Init/Lean/Elab/Tactic/Basic.lean

@@ -99,6 +99,7 @@ when (ctx.currRecDepth == ctx.maxRecDepth) $ throwError ref maxRecDepthErrorMess
 adaptReader (fun (ctx : Context) => { currRecDepth := ctx.currRecDepth + 1, .. ctx }) x
 
 protected def getCurrMacroScope : TacticM MacroScope := do ctx ← read; pure ctx.currMacroScope
+protected def getMainModule     : TacticM Name       := do env ← getEnv; pure env.mainModule
 
 @[inline] protected def withFreshMacroScope {α} (x : TacticM α) : TacticM α := do
 fresh ← modifyGet (fun st => (st.nextMacroScope, { st with nextMacroScope := st.nextMacroScope + 1 }));
@@ -106,6 +107,7 @@ adaptReader (fun (ctx : Context) => { ctx with currMacroScope := fresh }) x
 
 instance monadQuotation : MonadQuotation TacticM := {
   getCurrMacroScope   := Tactic.getCurrMacroScope,
+  getMainModule       := Tactic.getMainModule,
   withFreshMacroScope := @Tactic.withFreshMacroScope
 }
 

src/Init/Lean/Elab/Term.lean

@@ -139,9 +139,8 @@ ctx ← read;
 when (ctx.currRecDepth == ctx.maxRecDepth) $ throwError ref maxRecDepthErrorMessage;
 adaptReader (fun (ctx : Context) => { currRecDepth := ctx.currRecDepth + 1, .. ctx }) x
 
-protected def getCurrMacroScope : TermElabM MacroScope := do
-ctx ← read;
-pure ctx.currMacroScope
+protected def getCurrMacroScope : TermElabM MacroScope := do ctx ← read; pure ctx.currMacroScope
+protected def getMainModule     : TermElabM Name := do env ← getEnv; pure env.mainModule
 
 @[inline] protected def withFreshMacroScope {α} (x : TermElabM α) : TermElabM α := do
 fresh ← modifyGet (fun st => (st.nextMacroScope, { st with nextMacroScope := st.nextMacroScope + 1 }));
@@ -149,6 +148,7 @@ adaptReader (fun (ctx : Context) => { ctx with currMacroScope := fresh }) x
 
 instance monadQuotation : MonadQuotation TermElabM := {
   getCurrMacroScope   := Term.getCurrMacroScope,
+  getMainModule       := Term.getMainModule,
   withFreshMacroScope := @Term.withFreshMacroScope
 }
 

src/Init/Lean/Hygiene.lean

@@ -27,7 +27,8 @@ namespace Lean
 abbrev Unhygienic := ReaderT MacroScope $ StateM MacroScope
 namespace Unhygienic
 instance MonadQuotation : MonadQuotation Unhygienic := {
-  getCurrMacroScope := read,
+  getCurrMacroScope   := read,
+  getMainModule       := pure `UnhygienicMain,
   withFreshMacroScope := fun α x => do
     fresh ← modifyGet (fun n => (n, n + 1));
     adaptReader (fun _ => fresh) x
@@ -36,7 +37,8 @@ protected def run {α : Type} (x : Unhygienic α) : α := run x 0 1
 end Unhygienic
 
 instance monadQuotationTrans {m n : Type → Type} [MonadQuotation m] [HasMonadLift m n] [MonadFunctorT m m n n] : MonadQuotation n :=
-{ getCurrMacroScope   := liftM (getCurrMacroScope : m Nat),
+{ getCurrMacroScope   := liftM (getCurrMacroScope : m MacroScope),
+  getMainModule       := liftM (getMainModule : m Name),
   withFreshMacroScope := fun α => monadMap (fun α => (withFreshMacroScope : m α → m α)) }
 
 end Lean

src/Init/Lean/LocalContext.lean

@@ -158,7 +158,7 @@ private partial def getUnusedNameAux (lctx : LocalContext) (suggestion : Name) :
 
 @[export lean_local_ctx_get_unused_name]
 def getUnusedName (lctx : LocalContext) (suggestion : Name) : Name :=
-let (suggestion, _) := extractMacroScopes suggestion;
+let suggestion := suggestion.eraseMacroScopes;
 if lctx.usesUserName suggestion then (getUnusedNameAux lctx suggestion 1).1
 else suggestion
 

src/Init/LeanInit.lean

@@ -225,6 +225,7 @@ abbrev MacroScope := Nat
 class MonadQuotation (m : Type → Type) :=
 -- Get the fresh scope of the current macro invocation
 (getCurrMacroScope {} : m MacroScope)
+(getMainModule {}     : m Name)
 /- Execute action in a new macro invocation context. This transformer should be
    used at all places that morally qualify as the beginning of a "macro call",
    e.g. `elabCommand` and `elabTerm` in the case of the elaborator. However, it
@@ -239,26 +240,83 @@ class MonadQuotation (m : Type → Type) :=
 (withFreshMacroScope {α : Type} : m α → m α)
 export MonadQuotation
 
--- TODO: try harder to avoid clashes with other autogenerated names
+/-
+We represent a name with macro scopes as
+```
+<actual name>._@.<main_module>._hyg.<scopes>
+```
+Example: suppose the main module name is `Init.Data.List.Basic`, and name is `foo.bla`, and macroscopes [2, 5]
+```
+foo.bla._@.Init.Data.List.Basic._hyg.2.5
+```
+
+We use two delimiters for improving `hasMacroScopes` performance.
+-/
+
+private def mkMacroScopeName (mainModule : Name) (n : Name) : Name :=
+mkNameStr (mkNameStr n "_@" ++ mainModule) "_hyg"
+
+def addMacroScopeExt (mainModule : Name) (n : Name) (scp : MacroScope) : Name :=
+mkNameNum (mkMacroScopeName mainModule n) scp
+
 def addMacroScope (n : Name) (scp : MacroScope) : Name :=
-mkNameNum n scp
+addMacroScopeExt Name.anonymous n scp
 
-def addMacroScopes (n : Name) (scps : List MacroScope) : Name :=
-scps.foldl addMacroScope n
+def addMacroScopes (mainModule : Name) (n : Name) (scps : List MacroScope) : Name :=
+match scps with
+| [] => n
+| _  => scps.foldl mkNameNum (mkMacroScopeName mainModule n)
 
-private def extractMacroScopesAux : Name → List MacroScope → Name × List MacroScope
-| Name.num n scp _, acc => extractMacroScopesAux n (scp::acc)
-| n               , acc => (n, acc.reverse)
+def Name.hasMacroScopes : Name → Bool
+| Name.str _ str _ => str == "_hyg"
+| Name.num p _   _ => Name.hasMacroScopes p
+| _                => false
+
+structure ExtractMacroScopesResult :=
+(name       : Name)
+(mainModule : Name)
+(scopes     : List MacroScope)
+
+instance ExtractMacroScopesResult.inhabited : Inhabited ExtractMacroScopesResult :=
+⟨⟨arbitrary _, arbitrary _, arbitrary _⟩⟩
+
+private def assembleExtractedName : List Name → Name → Name
+| [],                     acc => acc
+| (Name.str _ s _) :: ps, acc => assembleExtractedName ps (mkNameStr acc s)
+| (Name.num _ n _) :: ps, acc => assembleExtractedName ps (mkNameNum acc n)
+| _,                      _   => unreachable!
+
+private def extractMainModule (scps : List MacroScope) : Name → List Name → ExtractMacroScopesResult
+| n@(Name.str p str _), acc =>
+  if str == "_@" then
+    { name := p, mainModule := assembleExtractedName acc Name.anonymous, scopes := scps }
+  else
+    extractMainModule p (n :: acc)
+| n@(Name.num p num _), acc => extractMainModule p (n :: acc)
+| _,                    _   => unreachable!
+
+private def extractMacroScopesAux : Name → List MacroScope → ExtractMacroScopesResult
+| Name.num p scp _, acc => extractMacroScopesAux p (scp::acc)
+| Name.str p str _, acc => extractMainModule acc.reverse p [] -- str must be "_hyg"
+| _,                _   => unreachable!
 
 /--
-  Revert all `addMacroScope` calls. `(n', scps) = extractMacroScopes n → n = addMacroScopes n' scps`.
+  Revert all `addMacroScope` calls. `⟨mainModule, n', scps⟩ = extractMacroScopes n → n = addMacroScopes mainModule n' scps`.
   This operation is useful for analyzing/transforming the original identifiers, then adding back
   the scopes (via `addMacroScopes`). -/
-def extractMacroScopes (n : Name) : Name × List MacroScope :=
-extractMacroScopesAux n []
+def extractMacroScopes (n : Name) : ExtractMacroScopesResult :=
+if n.hasMacroScopes then
+  extractMacroScopesAux n []
+else
+  { name := n, scopes := [], mainModule := Name.anonymous }
+
+private def eraseMacroScopesAux : Name → Name
+| Name.str p str _ => if str == "_@" then p else eraseMacroScopesAux p
+| Name.num p _ _   => eraseMacroScopesAux p
+| Name.anonymous   => unreachable!
 
 def Name.eraseMacroScopes (n : Name) : Name :=
-(extractMacroScopes n).1
+if n.hasMacroScopes then eraseMacroScopesAux n else n
 
 namespace Macro
 
@@ -277,8 +335,16 @@ end Macro
 
 abbrev MacroM := ReaderT Macro.Context (EStateM Macro.Exception Macro.State)
 
+def Macro.mkFreshName (namePrefix : Name) : MacroM Name := do
+ctx ← read;
+s   ← get;
+let id := namePrefix ++ ctx.mainModule ++ s.ngen.curr;
+modify $ fun s => { ngen := s.ngen.next, .. s };
+pure id
+
 instance MacroM.monadQuotation : MonadQuotation MacroM :=
 { getCurrMacroScope   := fun ctx => pure ctx.currMacroScope,
+  getMainModule       := fun ctx => pure ctx.mainModule,
   withFreshMacroScope := fun _ x => x }
 
 abbrev Macro := Syntax → MacroM Syntax