perf: annotate our allocators as such

This PR implements the function `getMaxFVar?` for implementing `SymM`
primitives.

src/Lean/Meta/Sym.lean

@@ -8,6 +8,7 @@ prelude
 public import Lean.Meta.Sym.SymM
 public import Lean.Meta.Sym.Main
 public import Lean.Meta.Sym.Util
+public import Lean.Meta.Sym.MaxFVar
 
 /-!
 # Symbolic simulation support.

src/Lean/Meta/Sym/Main.lean

@@ -8,11 +8,12 @@ prelude
 public import Lean.Meta.Sym.SymM
 public import Lean.Meta.Sym.Util
 public import Lean.Meta.Tactic.Grind.Main
+public section
 namespace Lean.Meta.Sym
 open Grind (Params)
 
 def SymM.run (x : SymM α) (params : Params) : MetaM α := do
-  x.run' {} |>.run params
+  x params |>.run' {} |>.run params
 
 def SymM.run' (x : SymM α) (config : Grind.Config := {}) : MetaM α := do
   let params ← Grind.mkDefaultParams config

src/Lean/Meta/Sym/MaxFVar.lean

@@ -0,0 +1,51 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Sym.SymM
+public section
+namespace Lean.Meta.Sym
+
+private def max (fvarId1? : Option FVarId) (fvarId2? : Option FVarId) : MetaM (Option FVarId) := do
+  let some fvarId1 := fvarId1? | return fvarId2?
+  let some fvarId2 := fvarId2? | return fvarId1?
+  if fvarId1 == fvarId2 then return fvarId1?
+  let localDecl1 ← fvarId1.getDecl
+  let localDecl2 ← fvarId2.getDecl
+  if localDecl1.index > localDecl2.index then
+    return fvarId1?
+  else
+    return fvarId2?
+
+private abbrev check (e : Expr) (k : SymM (Option FVarId)) : SymM (Option FVarId) := do
+  if let some fvarId? := (← get).maxFVar.find? { expr := e } then
+    return fvarId?
+  else
+    let fvarId? ← k
+    modify fun s => { s with maxFVar := s.maxFVar.insert { expr := e } fvarId? }
+    return fvarId?
+
+/--
+Returns the maximal free variable occurring in `e`.
+This function assumes the current local context contains all free variables used in `e`.
+-/
+def getMaxFVar? (e : Expr) : SymM (Option FVarId) := do
+  match e with
+  | .fvar fvarId => return some fvarId
+  | .const .. | .bvar .. | .sort .. | .lit .. => return none
+  | .app f a => check e do max (← getMaxFVar? f) (← getMaxFVar? a)
+  | .proj _ _ s => check e do getMaxFVar? s
+  | .lam _ d b _ | .forallE _ d b _ => check e do max (← getMaxFVar? d) (← getMaxFVar? b)
+  | .letE _ t v b _ => check e do max (← max (← getMaxFVar? t) (← getMaxFVar? v)) (← getMaxFVar? b)
+  | .mdata _ e => check e do getMaxFVar? e
+  | .mvar mvarId => check e do
+    let mvarDecl ← mvarId.getDecl
+    let lctx := mvarDecl.lctx
+    if lctx.decls.isEmpty then return none
+    let some last := lctx.lastDecl | unreachable! -- SymM local contexts do not have holes
+    return some last.fvarId
+
+end Lean.Meta.Sym

src/Lean/Meta/Sym/SymM.lean

@@ -6,12 +6,14 @@ Authors: Leonardo de Moura
 module
 prelude
 public import Lean.Meta.Tactic.Grind.Types
+public import Lean.Meta.Tactic.Grind.Main
+public section
 namespace Lean.Meta.Sym
 export Grind (ExprPtr Goal)
 
 structure State where
-  maxFVar : PHashMap ExprPtr Expr := {}
+  maxFVar : PHashMap ExprPtr (Option FVarId) := {}
 
-public abbrev SymM := StateRefT State Grind.GrindM
+abbrev SymM := ReaderT Grind.Params StateRefT State Grind.GrindM
 
 end Lean.Meta.Sym

src/Lean/Meta/Tactic/Grind.lean

@@ -48,6 +48,8 @@ public import Lean.Meta.Tactic.Grind.Filter
 public import Lean.Meta.Tactic.Grind.CollectParams
 public import Lean.Meta.Tactic.Grind.Finish
 public import Lean.Meta.Tactic.Grind.RegisterCommand
+public import Lean.Meta.Tactic.Grind.AlphaShareCommon
+public import Lean.Meta.Tactic.Grind.AlphaShareBuilder
 public section
 namespace Lean
 

src/Lean/Meta/Tactic/Grind/AlphaShareBuilder.lean

@@ -0,0 +1,85 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Types
+public section
+namespace Lean.Meta.Grind
+/-!
+Helper functions for constructing maximally shared expressions from maximally shared expressions.
+-/
+private def share (e : Expr) : GrindM Expr := do
+  let key : AlphaKey := ⟨e⟩
+  if let some r := (← get).scState.set.find? key then
+    return r.expr
+  else
+    modify fun s => { s with scState.set := s.scState.set.insert key }
+    return e
+
+private def assertShared (e : Expr) : GrindM Unit := do
+  assert! (← get).scState.set.contains ⟨e⟩
+
+def mkLitS (l : Literal) : GrindM Expr :=
+  share <| .lit l
+
+def mkConstS (declName : Name) (us : List Level := []) : GrindM Expr :=
+  share <| .const declName us
+
+def mkBVarS (idx : Nat) : GrindM Expr :=
+  share <| .bvar idx
+
+def mkSortS (u : Level) : GrindM Expr :=
+  share <| .sort u
+
+def mkFVarS (fvarId : FVarId) : GrindM Expr :=
+  share <| .fvar fvarId
+
+def mkMVarS (mvarId : MVarId) : GrindM Expr := do
+  share <| .mvar mvarId
+
+def mkMDataS (m : MData) (e : Expr) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared e
+  share <| .mdata m e
+
+def mkProjS (structName : Name) (idx : Nat) (struct : Expr) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared struct
+  share <| .proj structName idx struct
+
+def mkAppS (f a : Expr) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared f
+    assertShared a
+  share <| .app f a
+
+def mkLambdaS (x : Name) (bi : BinderInfo) (t : Expr) (b : Expr) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared t
+    assertShared b
+  share <| .lam x t b bi
+
+def mkForallS (x : Name) (bi : BinderInfo) (t : Expr) (b : Expr) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared t
+    assertShared b
+  share <| .forallE x t b bi
+
+def mkLetS (x : Name) (t : Expr) (v : Expr) (b : Expr) (nondep : Bool := false) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared t
+    assertShared v
+    assertShared b
+  share <| .letE x t v b nondep
+
+def mkHaveS (x : Name) (t : Expr) (v : Expr) (b : Expr) : GrindM Expr := do
+  if (← isDebugEnabled) then
+    assertShared t
+    assertShared v
+    assertShared b
+  share <| .letE x t v b true
+
+end Lean.Meta.Grind

src/include/lean/lean.h

@@ -335,6 +335,7 @@ static inline unsigned lean_get_slot_idx(unsigned sz) {
     return sz / LEAN_OBJECT_SIZE_DELTA - 1;
 }
 
+__attribute__((malloc)) __attribute__((alloc_size(1))) __attribute__((returns_nonnull))
 LEAN_EXPORT void * lean_alloc_small(unsigned sz, unsigned slot_idx);
 LEAN_EXPORT void lean_free_small(void * p);
 LEAN_EXPORT unsigned lean_small_mem_size(void * p);
@@ -344,6 +345,7 @@ LEAN_EXPORT void lean_inc_heartbeat(void);
 void * malloc(size_t);  // avoid including big `stdlib.h`
 #endif
 
+__attribute__((malloc)) __attribute__((alloc_size(1))) __attribute__((returns_nonnull))
 static inline lean_object * lean_alloc_small_object(unsigned sz) {
 #ifdef LEAN_SMALL_ALLOCATOR
     sz = lean_align(sz, LEAN_OBJECT_SIZE_DELTA);
@@ -438,6 +440,7 @@ static inline void lean_free_small_object(lean_object * o) {
 #endif
 }
 
+__attribute__((malloc)) __attribute__((alloc_size(1))) __attribute__((returns_nonnull))
 LEAN_EXPORT lean_object * lean_alloc_object(size_t sz);
 LEAN_EXPORT void lean_free_object(lean_object * o);
 

tests/lean/run/getMaxFVar.lean

@@ -0,0 +1,34 @@
+import Lean.Meta.Sym
+import Lean.Meta.Tactic.Grind.AlphaShareBuilder
+open Lean Meta Grind Sym
+
+def checkMaxFVar (e : Expr) (x : Expr) : SymM Unit := do
+  let some fvarId ← getMaxFVar? e | unreachable!
+  assert! x.fvarId! == fvarId
+
+def test1 : MetaM Unit := do
+  let nat := mkConst ``Nat
+  withLocalDeclD `x nat fun x => do
+  let m₁ ← mkFreshExprMVar nat
+  withLocalDeclD `y nat fun y => do
+  let m₂ ← mkFreshExprMVar nat
+  withLocalDeclD `z nat fun z => do
+  SymM.run' do
+  let x ← shareCommon x
+  let y ← shareCommon y
+  let z ← shareCommon z
+  let m₁ ← shareCommon m₁
+  let m₂ ← shareCommon m₂
+  let f ← mkConstS `f
+  let e₁ ← mkAppS f x
+  checkMaxFVar e₁ x
+  let e₂ ← mkAppS e₁ m₁
+  checkMaxFVar e₂ x
+  let e₂ ← mkAppS e₁ m₂
+  checkMaxFVar e₂ y
+  let e₃ ← mkAppS e₂ (← mkProjS ``Prod 0 (← mkAppS f z))
+  checkMaxFVar e₃ z
+  let e₄ ← mkAppS f (← shareCommon (mkNatLit 3))
+  assert! (← getMaxFVar? e₄).isNone
+
+#eval test1

tests/lean/run/grind_alphaShare_builder.lean

@@ -0,0 +1,24 @@
+import Lean.Meta.Tactic.Grind
+import Lean.Meta.Sym.Main
+open Lean Meta Grind Sym
+set_option grind.debug true
+
+def test : GrindM Unit := do
+  let f  ← mkConstS `f
+  let f₁ := mkConst `f
+  let f₂ ← mkConstS `f
+  assert! isSameExpr f f₂
+  assert! !isSameExpr f f₁
+  let x₁ ← mkBVarS 0
+  let x₂ ← mkBVarS 0
+  assert! isSameExpr
+    (← mkLambdaS `x .default (← mkConstS ``Nat) (← mkMDataS {} (← mkProjS ``Prod 0 (← mkAppS f x₁))))
+    (← mkLambdaS `y .default (← mkConstS ``Nat) (← mkMDataS {} (← mkProjS ``Prod 0 (← mkAppS f₂ x₂))))
+  assert! !isSameExpr (← mkAppS f x₁) (mkApp f x₁)
+  assert!
+    mkLambda `x .default (mkConst ``Nat) (mkMData {} (mkProj ``Prod 0 (mkApp f x₁)))
+    ==
+    (← mkLambdaS `y .default (← mkConstS ``Nat) (← mkMDataS {} (← mkProjS ``Prod 0 (← mkAppS f₂ x₂))))
+
+#eval SymM.run' do
+  test