feat: pp.exprSizes debugging option (#5218)

src/Lean/Expr.lean

@@ -1899,6 +1899,29 @@ with initial value `a`. -/
 def foldlM {α : Type} {m} [Monad m] (f : α → Expr → m α) (init : α) (e : Expr) : m α :=
   Prod.snd <$> StateT.run (e.traverseChildren (fun e' => fun a => Prod.mk e' <$> f a e')) init
 
+/--
+Returns the size of `e` as a DAG, i.e. the number of unique `Expr` constructor objects reachable
+from `e`.
+-/
+@[extern "lean_expr_size_shared"]
+opaque sizeWithSharing (e : Expr) : Nat
+
+/--
+Returns the size of `e` as a tree, i.e. nodes reachable via multiple paths are counted multiple
+times.
+
+This is a naive implementation that visits shared subterms multiple times instead of caching their
+sizes. It is primarily meant for debugging.
+-/
+def sizeWithoutSharing : (e : Expr) → Nat
+  | .forallE _ d b _ => 1 + d.sizeWithoutSharing + b.sizeWithoutSharing
+  | .lam _ d b _     => 1 + d.sizeWithoutSharing + b.sizeWithoutSharing
+  | .mdata _ e       => 1 + e.sizeWithoutSharing
+  | .letE _ t v b _  => 1 + t.sizeWithoutSharing + v.sizeWithoutSharing + b.sizeWithoutSharing
+  | .app f a         => 1 + f.sizeWithoutSharing + a.sizeWithoutSharing
+  | .proj _ _ e      => 1 + e.sizeWithoutSharing
+  | .lit .. | .const .. | .sort .. | .mvar .. | .fvar .. | .bvar .. => 1
+
 end Expr
 
 /--

src/Lean/PrettyPrinter.lean

@@ -9,7 +9,7 @@ import Lean.PrettyPrinter.Parenthesizer
 import Lean.PrettyPrinter.Formatter
 import Lean.Parser.Module
 import Lean.ParserCompiler
-
+import Lean.Util.ShareCommon
 namespace Lean
 
 def PPContext.runCoreM {α : Type} (ppCtx : PPContext) (x : CoreM α) : IO α :=
@@ -36,8 +36,21 @@ def ppUsing (e : Expr) (delab : Expr → MetaM Term) : MetaM Format := do
   Meta.withLCtx lctx #[] do
     ppTerm (← delab e)
 
+register_builtin_option pp.exprSizes : Bool := {
+  defValue := false
+  group    := "pp"
+  descr    := "(pretty printer) prefix each embedded expression with its sizes in the format \
+    (size disregarding sharing/size with sharing/size with max sharing)"
+}
+
+private def maybePrependExprSizes (e : Expr) (f : Format) : MetaM Format :=
+  return if pp.exprSizes.get (← getOptions) then
+    f!"[size {e.sizeWithoutSharing}/{e.sizeWithSharing}/{ShareCommon.shareCommon e |>.sizeWithSharing}] {f}"
+  else
+    f
+
 def ppExpr (e : Expr) : MetaM Format := do
-  ppUsing e delab
+  ppUsing e delab >>= maybePrependExprSizes e
 
 /-- Return a `fmt` representing pretty-printed `e` together with a map from tags in `fmt`
 to `Elab.Info` nodes produced by the delaborator at various subexpressions of `e`. -/
@@ -46,7 +59,7 @@ def ppExprWithInfos (e : Expr) (optsPerPos : Delaborator.OptionsPerPos := {}) (d
   let lctx := (← getLCtx).sanitizeNames.run' { options := (← getOptions) }
   Meta.withLCtx lctx #[] do
     let (stx, infos) ← delabCore e optsPerPos delab
-    let fmt ← ppTerm stx
+    let fmt ← ppTerm stx >>= maybePrependExprSizes e
     return ⟨fmt, infos⟩
 
 @[export lean_pp_expr]

src/kernel/for_each_fn.cpp

@@ -161,6 +161,16 @@ void for_each(expr const & e, std::function<bool(expr const &, unsigned)> && f)
     return for_each_offset_fn(f)(e);
 }
 
+extern "C" LEAN_EXPORT obj_res lean_expr_size_shared(b_obj_arg e_) {
+    expr const & e = TO_REF(expr, e_);
+    size_t size = 0;
+    for_each_fn<true>([&](expr const &) {
+        size++;
+        return true;
+    })(e);
+    return usize_to_nat(size);
+}
+
 extern "C" LEAN_EXPORT obj_res lean_find_expr(b_obj_arg p, b_obj_arg e_) {
     lean_object * found = nullptr;
     expr const & e = TO_REF(expr, e_);