test: and fix

src/Lean/Meta/Sym.lean

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

src/Lean/Meta/Sym/LooseBVarsS.lean

@@ -0,0 +1,49 @@
+/-
+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
+import Lean.Meta.Sym.ReplaceS
+public section
+namespace Lean.Meta.Sym
+
+/--
+Lowers the loose bound variables `>= s` in `e` by `d`.
+That is, a loose bound variable `bvar i` with `i >= s` is mapped to `bvar (i-d)`.
+It assumes the input is maximally shared, and ensures the output is too.
+Remark: if `s < d`, then the result is `e`.
+-/
+def lowerLooseBVarsS (e : Expr) (s d : Nat) : SymM Expr := do
+  replaceS e fun e offset => do
+    let s₁ := s + offset
+    if s₁ >= e.looseBVarRange then
+      return some e -- `e` does not contain bound variables with idx >= s₁
+    else match e with
+    | .bvar idx =>
+      if idx >= s₁ then
+        return some (← Grind.mkBVarS (idx - d))
+      else
+        return none
+    | _ => return none
+
+/--
+Lifts loose bound variables `>= s` in `e` by `d`.
+It assumes the input is maximally shared, and ensures the output is too.
+-/
+def liftLooseBVarsS (e : Expr) (s d : Nat) : SymM Expr := do
+  replaceS e fun e offset => do
+    let s₁ := s + offset
+    if s₁ >= e.looseBVarRange then
+      return some e -- `e` does not contain bound variables with idx >= s₁
+    else match e with
+    | .bvar idx =>
+      if idx >= s₁ then
+        return some (← Grind.mkBVarS (idx + d))
+      else
+        return none
+    | _ => return none
+
+end Lean.Meta.Sym

src/Lean/Meta/Sym/ReplaceS.lean

@@ -0,0 +1,54 @@
+/-
+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 import Lean.Meta.Tactic.Grind.AlphaShareBuilder
+namespace Lean.Meta.Sym
+/-!
+A version of `replace_fn.h` that ensures the resulting expression is maximally shared.
+-/
+open Grind
+abbrev M := StateT (Std.HashMap (ExprPtr × Nat) Expr) AlphaShareBuilderM
+
+def save (key : ExprPtr × Nat) (r : Expr) : M Expr := do
+  modify fun cache => cache.insert key r
+  return r
+
+mutual
+@[specialize] def visitChild (e : Expr) (offset : Nat) (f : Expr → Nat → AlphaShareBuilderM (Option Expr)) : M Expr := do
+  let key := (⟨e⟩, offset)
+  if let some r := (← get)[key]? then
+    return r
+  else if let some r ← f e offset then
+    save key r
+  else match e with
+    | .lit _ | .mvar _ | .bvar _ | .fvar _ | .const _ _ | .sort _ => save key e
+    | e => save key (← visit e offset f)
+
+@[specialize] def visit (e : Expr) (offset : Nat) (fn : Expr → Nat → AlphaShareBuilderM (Option Expr)) : M Expr := do
+  match e with
+  | .lit _ | .mvar _ | .bvar _ | .fvar _ | .const _ _ | .sort _ => unreachable!
+  | .app f a => mkAppS (← visitChild f offset fn) (← visitChild a offset fn)
+  | .mdata m a => mkMDataS m (← visitChild a offset fn)
+  | .proj s i a => mkProjS s i (← visitChild a offset fn)
+  | .forallE n d b bi => mkForallS n bi (← visitChild d offset fn) (← visitChild b (offset+1) fn)
+  | .lam n d b bi => mkLambdaS n bi (← visitChild d offset fn) (← visitChild b (offset+1) fn)
+  | .letE n t v b d => mkLetS n (← visitChild t offset fn) (← visitChild v offset fn) (← visitChild b (offset+1) fn) (nondep := d)
+end
+
+/--
+Similar to `replace_fn` in the kernel, but assumes input is maximally shared, and ensures
+output is also maximally shared.
+-/
+@[inline] public def replaceS (e : Expr) (f : Expr → Nat → AlphaShareBuilderM (Option Expr)) : SymM Expr := liftBuilderM do
+  if let some r ← f e 0 then
+    return r
+  match e with
+  | .lit _ | .mvar _ | .bvar _ | .fvar _ | .const _ _ | .sort _ => return e
+  | _ => visit e 0 f |>.run' {}
+
+end Lean.Meta.Sym

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

@@ -11,7 +11,7 @@ namespace Lean.Meta.Grind
 /-!
 Helper functions for constructing maximally shared expressions from maximally shared expressions.
 -/
-private def share (e : Expr) : GrindM Expr := do
+protected def Grind.share (e : Expr) : GrindM Expr := do
   let key : AlphaKey := ⟨e⟩
   if let some r := (← get).scState.set.find? key then
     return r.expr
@@ -19,64 +19,117 @@ private def share (e : Expr) : GrindM Expr := do
     modify fun s => { s with scState.set := s.scState.set.insert key }
     return e
 
-private def assertShared (e : Expr) : GrindM Unit := do
+protected def Grind.assertShared (e : Expr) : GrindM Unit := do
   assert! (← get).scState.set.contains ⟨e⟩
 
-def mkLitS (l : Literal) : GrindM Expr :=
+class MonadShareCommon (m : Type → Type) where
+  share : Expr → m Expr
+  assertShared : Expr → m Unit
+  isDebugEnabled : m Bool
+
+@[always_inline]
+instance (m n) [MonadLift m n] [MonadShareCommon m] : MonadShareCommon n where
+  share e := liftM (MonadShareCommon.share e : m Expr)
+  assertShared e := liftM (MonadShareCommon.assertShared e : m Unit)
+  isDebugEnabled := liftM (MonadShareCommon.isDebugEnabled : m Bool)
+
+instance : MonadShareCommon GrindM where
+  share := Grind.share
+  assertShared := Grind.assertShared
+  isDebugEnabled := isDebugEnabled
+
+/--
+Helper monad for constructing maximally shared terms.
+The `Bool` flag indicates whether it is debug-mode or not.
+-/
+abbrev AlphaShareBuilderM := ReaderT Bool AlphaShareCommonM
+
+/--
+Helper function for lifting a `AlphaShareBuilderM` action to `GrindM`
+-/
+abbrev liftBuilderM (k : AlphaShareBuilderM α) : GrindM α := do
+  let scState ← modifyGet fun s => (s.scState, { s with scState := {} })
+  let (a, scState) := k (← isDebugEnabled) scState
+  modify fun s => { s with scState }
+  return a
+
+protected def Builder.share (e : Expr) : AlphaShareBuilderM Expr := do
+  let key : AlphaKey := ⟨e⟩
+  if let some r := (← get).set.find? key then
+    return r.expr
+  else
+    modify fun s => { s with set := s.set.insert key }
+    return e
+
+protected def Builder.assertShared (e : Expr) : AlphaShareBuilderM Unit := do
+  assert! (← get).set.contains ⟨e⟩
+
+instance : MonadShareCommon AlphaShareBuilderM where
+  share := Builder.share
+  assertShared := Builder.assertShared
+  isDebugEnabled := read
+
+open MonadShareCommon (share assertShared)
+
+variable [MonadShareCommon m]
+
+def mkLitS (l : Literal) : m Expr :=
   share <| .lit l
 
-def mkConstS (declName : Name) (us : List Level := []) : GrindM Expr :=
+def mkConstS (declName : Name) (us : List Level := []) : m Expr :=
   share <| .const declName us
 
-def mkBVarS (idx : Nat) : GrindM Expr :=
+def mkBVarS (idx : Nat) : m Expr :=
   share <| .bvar idx
 
-def mkSortS (u : Level) : GrindM Expr :=
+def mkSortS (u : Level) : m Expr :=
   share <| .sort u
 
-def mkFVarS (fvarId : FVarId) : GrindM Expr :=
+def mkFVarS (fvarId : FVarId) : m Expr :=
   share <| .fvar fvarId
 
-def mkMVarS (mvarId : MVarId) : GrindM Expr := do
+def mkMVarS (mvarId : MVarId) : m Expr :=
   share <| .mvar mvarId
 
-def mkMDataS (m : MData) (e : Expr) : GrindM Expr := do
-  if (← isDebugEnabled) then
-    assertShared e
-  share <| .mdata m e
+variable [Monad m]
 
-def mkProjS (structName : Name) (idx : Nat) (struct : Expr) : GrindM Expr := do
-  if (← isDebugEnabled) then
+def mkMDataS (d : MData) (e : Expr) : m Expr := do
+  if (← MonadShareCommon.isDebugEnabled) then
+    assertShared e
+  share <| .mdata d e
+
+def mkProjS (structName : Name) (idx : Nat) (struct : Expr) : m Expr := do
+  if (← MonadShareCommon.isDebugEnabled) then
     assertShared struct
   share <| .proj structName idx struct
 
-def mkAppS (f a : Expr) : GrindM Expr := do
-  if (← isDebugEnabled) then
+def mkAppS (f a : Expr) : m Expr := do
+  if (← MonadShareCommon.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
+def mkLambdaS (x : Name) (bi : BinderInfo) (t : Expr) (b : Expr) : m Expr := do
+  if (← MonadShareCommon.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
+def mkForallS (x : Name) (bi : BinderInfo) (t : Expr) (b : Expr) : m Expr := do
+  if (← MonadShareCommon.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
+def mkLetS (x : Name) (t : Expr) (v : Expr) (b : Expr) (nondep : Bool := false) : m Expr := do
+  if (← MonadShareCommon.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
+def mkHaveS (x : Name) (t : Expr) (v : Expr) (b : Expr) : m Expr := do
+  if (← MonadShareCommon.isDebugEnabled) then
     assertShared t
     assertShared v
     assertShared b

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

@@ -123,7 +123,7 @@ private def go (e : Expr) : M Expr := do
     visit e (return mkProj n i (← go b))
 
 /-- Similar to `shareCommon`, but handles alpha-equivalence. -/
-@[inline] def shareCommonAlpha (e : Expr) (s : AlphaShareCommon.State) : Expr × AlphaShareCommon.State :=
+@[inline] def shareCommonAlpha (e : Expr) : AlphaShareCommonM Expr := fun s =>
   if let some r := s.set.find? { expr := e } then
     (r.expr, s)
   else

tests/lean/run/sym_liftLower.lean

@@ -0,0 +1,23 @@
+import Lean.Meta.Sym
+
+open Lean Meta Grind Sym
+
+def tst1 : SymM Unit := do
+  let x1 ← mkBVarS 0
+  let x2 ← mkBVarS 1
+  let t1 ← mkAppS (← mkAppS (← mkConstS `f) x1) x2
+  let t2 ← mkForallS `x BinderInfo.default (← mkConstS `Nat) t1
+  IO.println (← liftLooseBVarsS t1 0 1)
+  IO.println (← liftLooseBVarsS t2 0 1)
+  let t3 ← lowerLooseBVarsS (← liftLooseBVarsS t2 0 1) 1 1
+  IO.println t3
+  assert! t2 == t3
+  assert! isSameExpr t2 t3
+
+/--
+info: f #1 #2
+forall (x : Nat), f x #1
+forall (x : Nat), f x #0
+-/
+#guard_msgs in
+#eval SymM.run' tst1