chore: apply suggestions from code review

Co-authored-by: Gabriel Ebner <gebner@gebner.org>

src/Init/Prelude.lean

@@ -2253,7 +2253,7 @@ export Macro (expandMacro?)
 
 namespace PrettyPrinter
 
-abbrev UnexpandM := EStateM Unit Syntax
+abbrev UnexpandM := ReaderT Syntax (EStateM Unit Unit)
 
 /--
   Function that tries to reverse macro expansions as a post-processing step of delaboration.
@@ -2262,15 +2262,10 @@ abbrev UnexpandM := EStateM Unit Syntax
 -- a `kindUnexpander` could reasonably be added later
 abbrev Unexpander := Syntax → UnexpandM Syntax
 
--- unexpanders should not need to introduce new names
 instance : MonadQuotation UnexpandM where
-  getRef              := get
-  withRef             := fun ref m => do
-    let orig ← get
-    set ref
-    let res ← m
-    set orig
-    return res
+  getRef              := read
+  withRef ref x       := withReader (fun _ => ref) x
+  -- unexpanders should not need to introduce new names
   getCurrMacroScope   := pure 0
   getMainModule       := pure `_fakeMod
   withFreshMacroScope := id

src/Lean/Elab/Notation.lean

@@ -53,16 +53,17 @@ def mkSimpleDelab (attrKind : Syntax) (vars : Array Syntax) (pat qrhs : Syntax)
     guard <| args.all (Syntax.isIdent ∘ getAntiquotTerm)
     guard <| args.allDiff
     -- replace head constant with (unused) antiquotation so we're not dependent on the exact pretty printing of the head
+    -- The reference is attached to the syntactic representation of the called function itself, not the entire function application
     `(@[$attrKind:attrKind appUnexpander $(mkIdent c):ident]
       aux_def unexpand $(mkIdent c) : Lean.PrettyPrinter.Unexpander := fun
-       | `($$tk:ident $args*) => withRef tk `($pat)
-       | _                     => throw ())
+       | `($$f:ident $args*) => withRef f `($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
-       | `($$tk:ident) => withRef tk `($pat)
-       | _              => throw ())
+       | `($$f:ident) => withRef f `($pat)
+       | _            => throw ())
   | _                  => failure
 
 private def isLocalAttrKind (attrKind : Syntax) : Bool :=

src/Lean/PrettyPrinter/Delaborator/Builtins.lean

@@ -104,13 +104,11 @@ def delabConst : Delab := do
     else
       `($(mkIdent c).{$[$(ls.toArray.map quote)],*})
 
-  let stx ← maybeAddBlockImplicit stx
-  if (←getPPOption getPPTagSymbols) then
-    let stx ← annotateCurPos stx
-    addTermInfo (←getPos) stx (←getExpr)
-    stx
-  else
-    stx
+  let mut stx ← maybeAddBlockImplicit stx
+  if (← getPPOption getPPTagAppFns) then
+    stx ← annotateCurPos stx
+    addTermInfo (← getPos) stx (← getExpr)
+  stx
 
 structure ParamKind where
   name        : Name
@@ -199,7 +197,7 @@ def unexpandRegularApp (stx : Syntax) : Delab := do
   let fs ← appUnexpanderAttribute.getValues (← getEnv) c
   let ref ← getRef
   fs.firstM fun f =>
-    match f stx |>.run ref with
+    match f stx |>.run ref |>.run () with
     | EStateM.Result.ok stx _ => pure stx
     | _ => failure
 

src/Lean/PrettyPrinter/Delaborator/Options.lean

@@ -90,10 +90,10 @@ register_builtin_option pp.safeShadowing  : Bool := {
   group    := "pp"
   descr    := "(pretty printer) allow variable shadowing if there is no collision"
 }
-register_builtin_option pp.tagSymbols : Bool := {
+register_builtin_option pp.tagAppFns : Bool := {
   defValue := false
   group    := "pp"
-  descr    := "(pretty printer) tag symbol syntax with information"
+  descr    := "(pretty printer) tag all constants that are the function in a function application"
 }
 register_builtin_option pp.proofs : Bool := {
   defValue := false
@@ -185,7 +185,7 @@ def getPPMatch (o : Options) : Bool := o.get pp.match.name (!getPPAll o)
 def getPPStructureProjections (o : Options) : Bool := o.get pp.structureProjections.name (!getPPAll o)
 def getPPStructureInstances (o : Options) : Bool := o.get pp.structureInstances.name (!getPPAll o)
 def getPPStructureInstanceType (o : Options) : Bool := o.get pp.structureInstanceTypes.name (getPPAll o)
-def getPPTagSymbols (o : Options) : Bool := o.get pp.tagSymbols.name (getPPAll o)
+def getPPTagAppFns (o : Options) : Bool := o.get pp.tagAppFns.name (getPPAll o)
 def getPPUniverses (o : Options) : Bool := o.get pp.universes.name (getPPAll o)
 def getPPFullNames (o : Options) : Bool := o.get pp.fullNames.name (getPPAll o)
 def getPPPrivateNames (o : Options) : Bool := o.get pp.privateNames.name (getPPAll o)