feat: Keep info in auto generated notation delaborators

hargoniX · Dec 21, 2021 · 1214d5e · 1214d5e
1 parent 42f2b40
commit 1214d5e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/src/Init/Prelude.lean b/src/Init/Prelude.lean
@@ -2251,7 +2251,7 @@ export Macro (expandMacro?)
 
 namespace PrettyPrinter
 
-abbrev UnexpandM := EStateM Unit Unit
+abbrev UnexpandM := EStateM Unit Syntax
 
 /--
   Function that tries to reverse macro expansions as a post-processing step of delaboration.
@@ -2262,8 +2262,13 @@ abbrev Unexpander := Syntax → UnexpandM Syntax
 
 -- unexpanders should not need to introduce new names
 instance : MonadQuotation UnexpandM where
-  getRef              := pure Syntax.missing
-  withRef             := fun _ => id
+  getRef              := get
+  withRef             := fun ref m => do
+    let orig ← get
+    set ref
+    let res ← m
+    set orig
+    return res
   getCurrMacroScope   := pure 0
   getMainModule       := pure `_fakeMod
   withFreshMacroScope := id

diff --git a/src/Lean/Elab/Notation.lean b/src/Lean/Elab/Notation.lean
@@ -55,13 +55,13 @@ def mkSimpleDelab (attrKind : Syntax) (vars : Array Syntax) (pat qrhs : Syntax)
     -- replace head constant with (unused) antiquotation so we're not dependent on the exact pretty printing of the head
     `(@[$attrKind:attrKind appUnexpander $(mkIdent c):ident]
       aux_def unexpand $(mkIdent c) : Lean.PrettyPrinter.Unexpander := fun
-       | `($$(_):ident $args*) => `($pat)
+       | `($$tk:ident $args*) => withRef tk `($pat)
        | _                     => throw ())
   | `($c:ident)        =>
     let [(c, [])] ← Macro.resolveGlobalName c.getId | failure
     `(@[$attrKind:attrKind appUnexpander $(mkIdent c):ident]
       aux_def unexpand $(mkIdent c) : Lean.PrettyPrinter.Unexpander := fun
-       | `($$(_):ident) => `($pat)
+       | `($$tk:ident) => withRef tk `($pat)
        | _              => throw ())
   | _                  => failure
 

diff --git a/src/Lean/PrettyPrinter/Delaborator/Builtins.lean b/src/Lean/PrettyPrinter/Delaborator/Builtins.lean
@@ -194,8 +194,9 @@ def isRegularApp : DelabM Bool := do
 def unexpandRegularApp (stx : Syntax) : Delab := do
   let Expr.const c .. ← pure (unfoldMDatas (← getExpr).getAppFn) | unreachable!
   let fs ← appUnexpanderAttribute.getValues (← getEnv) c
+  let ref ← getRef
   fs.firstM fun f =>
-    match f stx |>.run () with
+    match f stx |>.run ref with
     | EStateM.Result.ok stx _ => pure stx
     | _ => failure