feat: sanitize Syntax in messages

Fixes #182

src/Lean/Hygiene.lean

@@ -33,4 +33,77 @@ instance MonadQuotation : MonadQuotation Unhygienic := {
 protected def run {α : Type} (x : Unhygienic α) : α := run x firstFrontendMacroScope (firstFrontendMacroScope+1)
 end Unhygienic
 
+private def mkInaccessibleUserNameAux (unicode : Bool) (name : Name) (idx : Nat) : Name :=
+if unicode then
+  if idx == 0 then
+    name.appendAfter "✝"
+  else
+    name.appendAfter ("✝" ++ idx.toSuperscriptString)
+else
+  name ++ mkNameNum "_inaccessible" idx
+
+private def mkInaccessibleUserName (unicode : Bool) : Name → Name
+| Name.num p@(Name.str _ _ _) idx _ =>
+  mkInaccessibleUserNameAux unicode p idx
+| Name.num Name.anonymous idx _     =>
+  mkInaccessibleUserNameAux unicode Name.anonymous idx
+| Name.num p idx _ =>
+  if unicode then
+    (mkInaccessibleUserName p).appendAfter ("⁻" ++ idx.toSuperscriptString)
+  else
+    mkNameNum (mkInaccessibleUserName p) idx
+| n => n
+
+@[export lean_is_inaccessible_user_name]
+def isInaccessibleUserName : Name → Bool
+| Name.str _ s _   => s.contains '✝' || s == "_inaccessible"
+| Name.num p idx _ => isInaccessibleUserName p
+| _                => false
+
+def sanitizeNamesDefault := true
+@[init] def sanitizeNamesOption : IO Unit :=
+registerOption `pp.sanitizeNames { defValue := sanitizeNamesDefault, group := "pp", descr := "add suffix '_{<idx>}' to shadowed/inaccessible variables when pretty printing" }
+def getSanitizeNames (o : Options) : Bool:= o.get `pp.sanitizeNames sanitizeNamesDefault
+
+structure NameSanitizerState :=
+(options            : Options)
+-- `x` ~> 2 if we're already using `x✝`, `x✝¹`
+(nameStem2Idx       : NameMap Nat := {})
+-- `x._hyg...` ~> `x✝`
+(userName2Sanitized : NameMap Name := {})
+
+private partial def mkFreshInaccessibleUserName (userName : Name) : Nat → StateM NameSanitizerState Name
+| idx => do
+  st ← get;
+  let userNameNew := mkInaccessibleUserName (Format.getUnicode st.options) (mkNameNum userName idx);
+  if st.nameStem2Idx.contains userNameNew then
+    mkFreshInaccessibleUserName (idx+1)
+  else do
+    modify fun st => { st with nameStem2Idx := st.nameStem2Idx.insert userName (idx+1) };
+    pure userNameNew
+
+def sanitizeName (userName : Name) : StateM NameSanitizerState Name := do
+st ← get;
+let stem := userName.eraseMacroScopes;
+let idx  := (st.nameStem2Idx.find? stem).getD 0;
+san ← mkFreshInaccessibleUserName stem idx;
+modify fun st => { st with userName2Sanitized := st.userName2Sanitized.insert userName san };
+pure san
+
+private partial def sanitizeSyntaxAux : Syntax → StateM NameSanitizerState Syntax
+| Syntax.ident _ _ n _ => do
+  st ← get;
+  mkIdent <$> match st.userName2Sanitized.find? n with
+  | some n' => pure n'
+  | none    => if n.hasMacroScopes then sanitizeName n else pure n
+| Syntax.node k args => Syntax.node k <$> args.mapM sanitizeSyntaxAux
+| stx => pure stx
+
+def sanitizeSyntax (stx : Syntax) : StateM NameSanitizerState Syntax := do
+st ← get;
+if getSanitizeNames st.options then
+  sanitizeSyntaxAux stx
+else
+  pure stx
+
 end Lean

src/Lean/LocalContext.lean

@@ -371,66 +371,23 @@ Id.run $ lctx.anyM p
 @[inline] def all (lctx : LocalContext) (p : LocalDecl → Bool) : Bool :=
 Id.run $ lctx.allM p
 
-private def mkInaccessibleUserNameAux (unicode : Bool) (name : Name) (idx : Nat) : Name :=
-if unicode then
-  if idx == 0 then
-    name.appendAfter "✝"
-  else
-    name.appendAfter ("✝" ++ idx.toSuperscriptString)
-else
-  name ++ mkNameNum "_inaccessible" idx
-
-private def mkInaccessibleUserName (unicode : Bool) : Name → Name
-| Name.num p@(Name.str _ _ _) idx _ =>
-  mkInaccessibleUserNameAux unicode p idx
-| Name.num Name.anonymous idx _     =>
-  mkInaccessibleUserNameAux unicode Name.anonymous idx
-| Name.num p idx _ =>
-  if unicode then
-    (mkInaccessibleUserName p).appendAfter ("⁻" ++ idx.toSuperscriptString)
-  else
-    mkNameNum (mkInaccessibleUserName p) idx
-| n => n
-
-@[export lean_is_inaccessible_user_name]
-def isInaccessibleUserName : Name → Bool
-| Name.str _ s _   => s.contains '✝' || s == "_inaccessible"
-| Name.num p idx _ => isInaccessibleUserName p
-| _                => false
-
-private partial def mkFreshInaccessibleUserName (unicode : Bool) (usedName2Idx : NameMap Nat) (userName : Name) : Nat → Name × NameMap Nat
-| idx =>
-  let userNameNew := mkInaccessibleUserName unicode (mkNameNum userName idx);
-  if usedName2Idx.contains userNameNew then
-    mkFreshInaccessibleUserName (idx+1)
-  else
-   (userNameNew, usedName2Idx.insert userName (idx+1))
-
-def sanitizeNamesDefault := true
-@[init] def sanitizeNamesOption : IO Unit :=
-registerOption `pp.sanitizeNames { defValue := sanitizeNamesDefault, group := "pp", descr := "add suffix '_{<idx>}' to shadowed variables when pretty printing" }
-def getSanitizeNames (o : Options) : Bool:= o.get `pp.sanitizeNames sanitizeNamesDefault
-
-def sanitizeNames (lctx : LocalContext) (options : Options) : LocalContext :=
-if !getSanitizeNames options then lctx else
-let unicode   := Format.getUnicode options;
-let (lctx, _, _) := lctx.decls.size.foldRev
-  (fun i (acc : LocalContext × NameSet × NameMap Nat) =>
-    let (lctx, usedNames, usedName2Idx) := acc;
-    match lctx.decls.get! i with
-    | none      => acc
-    | some decl =>
-      let userName := decl.userName;
-      if userName.hasMacroScopes || usedNames.contains userName then
-        let userName := userName.eraseMacroScopes;
-        let idx      := (usedName2Idx.find? userName).getD 0;
-        let (userNameNew, usedName2Idx) := mkFreshInaccessibleUserName unicode usedName2Idx userName idx;
-        let lctx := lctx.setUserName decl.fvarId userNameNew;
-        (lctx, usedNames, usedName2Idx)
-      else
-        (lctx, usedNames.insert userName, usedName2Idx))
-  (lctx, {}, {});
-lctx
+def sanitizeNames (lctx : LocalContext) : StateM NameSanitizerState LocalContext := do
+st ← get;
+if !getSanitizeNames st.options then pure lctx else
+  flip StateT.run' ({} : NameSet) $
+    lctx.decls.size.foldRevM
+      (fun i lctx =>
+        match lctx.decls.get! i with
+        | none      => pure lctx
+        | some decl => do
+          usedNames ← get;
+          set $ usedNames.insert decl.userName;
+          if decl.userName.hasMacroScopes || usedNames.contains decl.userName then do
+            userNameNew ← liftM $ sanitizeName decl.userName;
+            pure $ lctx.setUserName decl.fvarId userNameNew
+          else
+            pure lctx)
+      lctx
 
 end LocalContext
 

src/Lean/Message.lean

@@ -52,9 +52,6 @@ namespace MessageData
 
 instance : Inhabited MessageData := ⟨MessageData.ofFormat (arbitrary _)⟩
 
-private def sanitizeNames (ctx : MessageDataContext) : MessageDataContext :=
-{ ctx with lctx := ctx.lctx.sanitizeNames ctx.opts }
-
 def mkPPContext (nCtx : NamingContext) (ctx : MessageDataContext) : PPContext :=
 { env := ctx.env, mctx := ctx.mctx, lctx := ctx.lctx, opts := ctx.opts,
   currNamespace := nCtx.currNamespace, openDecls := nCtx.openDecls }
@@ -69,7 +66,7 @@ partial def formatAux : NamingContext → Option MessageDataContext → MessageD
 | nCtx, some ctx,  ofExpr e                 => ppExpr (mkPPContext nCtx ctx) e
 | _,    none,      ofGoal mvarId            => pure $ "goal " ++ format (mkMVar mvarId)
 | nCtx, some ctx,  ofGoal mvarId            => ppGoal (mkPPContext nCtx ctx) mvarId
-| nCtx, _,         withContext ctx d        => formatAux nCtx (some $ sanitizeNames ctx) d
+| nCtx, _,         withContext ctx d        => formatAux nCtx ctx d
 | _,    ctx,       withNamingContext nCtx d => formatAux nCtx ctx d
 | nCtx, ctx,       tagged cls d             => do d ← formatAux nCtx ctx d; pure $ Format.sbracket (format cls) ++ " " ++ d
 | nCtx, ctx,       nest n d                 => Format.nest n <$> formatAux nCtx ctx d

src/Lean/PrettyPrinter.lean

@@ -18,12 +18,18 @@ ppCtx.runCoreM $ x.run' { lctx := ppCtx.lctx } { mctx := ppCtx.mctx }
 
 namespace PrettyPrinter
 
-def ppTerm (stx : Syntax) : CoreM Format :=
+def ppTerm (stx : Syntax) : CoreM Format := do
+opts ← getOptions;
+let stx := (sanitizeSyntax stx).run' { options := opts };
 parenthesizeTerm stx >>= formatTerm
 
 def ppExpr (currNamespace : Name) (openDecls : List OpenDecl) (e : Expr) : MetaM Format := do
-stx ← delab currNamespace openDecls e;
-liftM $ ppTerm stx
+lctx ← getLCtx;
+opts ← getOptions;
+let lctx := lctx.sanitizeNames.run' { options := opts };
+Meta.withLCtx lctx #[] $ do
+  stx ← delab currNamespace openDecls e;
+  liftM $ ppTerm stx
 
 def ppCommand (stx : Syntax) : CoreM Format :=
 parenthesizeCommand stx >>= formatCommand

src/Lean/PrettyPrinter/Formatter.lean

@@ -289,7 +289,7 @@ let id := id.simpMacroScopes;
 let s := id.toString;
 if id.isAnonymous then
   pushToken "[anonymous]"
-else if LocalContext.isInaccessibleUserName id || id.components.any Name.isNum ||
+else if isInaccessibleUserName id || id.components.any Name.isNum ||
   -- loose bvar
   "#".isPrefixOf s then
   -- not parsable anyway, output as-is

src/Lean/Util/PPGoal.lean

@@ -22,7 +22,7 @@ match mctx.findDecl? mvarId with
   let indent       := 2; -- Use option
   let ppAuxDecls   := getAuxDeclsOption opts;
   let lctx         := mvarDecl.lctx;
-  let lctx         := lctx.sanitizeNames opts;
+  let lctx         := lctx.sanitizeNames.run' { options := opts };
   let ppCtx        := { ppCtx with lctx := lctx };
   let pp (e : Expr) : IO Format := ppExpr ppCtx e;
   let instMVars (e : Expr) : Expr := (mctx.instantiateMVars e).1;

tests/lean/sanitizeMacroScopes.lean

@@ -0,0 +1,6 @@
+new_frontend
+
+macro "m" x:term : term => `(fun x => $x)
+
+set_option trace.Elab.step true in
+#check fun x => m x

tests/lean/sanitizeMacroScopes.lean.expected.out

@@ -0,0 +1,9 @@
+[Elab.step] #check fun  x => m x
+[Elab.step] none fun  x => m x
+[Elab.step] Sort _ _
+[Elab.step] none m x
+[Elab.step] none fun  x✝ => x
+[Elab.step] Sort _ _
+[Elab.step] none x
+fun (x : ?m) (x_1 : ?m x) => x : (x : ?m) → ?m x → ?m
+[Elab.step] end