feat: add instantiateS and variants

src/Lean/Meta/Sym.lean

@@ -11,6 +11,7 @@ public import Lean.Meta.Sym.Util
 public import Lean.Meta.Sym.MaxFVar
 public import Lean.Meta.Sym.ReplaceS
 public import Lean.Meta.Sym.LooseBVarsS
+public import Lean.Meta.Sym.InstantiateS
 
 /-!
 # Symbolic simulation support.

src/Lean/Meta/Sym/InstantiateS.lean

@@ -0,0 +1,89 @@
+/-
+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
+import Lean.Meta.Sym.LooseBVarsS
+import Init.Grind
+public section
+namespace Lean.Meta.Sym
+
+/--
+Similar to `Lean.Expr.instantiateRevRange`.
+It assumes the input is maximally shared, and ensures the output is too.
+It assumes `beginIdx ≤ endIdx` and `endIdx ≤ subst.size`
+-/
+def instantiateRevRangeS (e : Expr) (beginIdx endIdx : Nat) (subst : Array Expr) : SymM Expr :=
+  if _ : beginIdx > endIdx then unreachable! else
+  if _ : endIdx > subst.size then unreachable! else
+  let s := beginIdx
+  let n := endIdx - beginIdx
+  replaceS e fun e offset => do
+    let s₁ := s + offset
+    match e with
+    | .bvar idx =>
+      if _h : idx >= s₁ then
+        if _h : idx < offset + n then
+          let v := subst[n - (idx - offset) - 1]
+          liftLooseBVarsS' v 0 offset
+        else
+          Grind.mkBVarS (idx - n)
+      else
+        return some e
+    | .lit _ | .mvar _ | .fvar _ | .const _ _ | .sort _ =>
+      return some e
+    | _ =>
+      if s₁ >= e.looseBVarRange then
+        return some e
+      else
+        return none
+
+/--
+Similar to `Lean.Expr.instantiateRev`.
+It assumes the input is maximally shared, and ensures the output is too.
+-/
+@[inline] def instantiateRevS (e : Expr) (subst : Array Expr) : SymM Expr :=
+  instantiateRevRangeS e 0 subst.size subst
+
+/--
+Similar to `Lean.Expr.instantiateRange`.
+It assumes the input is maximally shared, and ensures the output is too.
+It assumes `beginIdx ≤ endIdx` and `endIdx ≤ subst.size`
+-/
+def instantiateRangeS (e : Expr) (beginIdx endIdx : Nat) (subst : Array Expr) : SymM Expr :=
+  if _ : beginIdx > endIdx then unreachable! else
+  if _ : endIdx > subst.size then unreachable! else
+  let s := beginIdx
+  let n := endIdx - beginIdx
+  replaceS e fun e offset => do
+    let s₁ := s + offset
+    match e with
+    | .bvar idx =>
+      if _h : idx >= s₁ then
+        if _h : idx < s₁ + n then
+          let v := subst[idx - s₁]
+          liftLooseBVarsS' v 0 offset
+        else
+          Grind.mkBVarS (idx - n)
+      else
+        return some e
+    | .lit _ | .mvar _ | .fvar _ | .const _ _ | .sort _ =>
+      return some e
+    | _ =>
+      if s₁ >= e.looseBVarRange then
+        return some e
+      else
+        return none
+
+/--
+Similar to `Lean.Expr.instantiate`.
+It assumes the input is maximally shared, and ensures the output is too.
+-/
+def instantiateS (e : Expr) (subst : Array Expr) : SymM Expr :=
+  instantiateRangeS e 0 subst.size subst
+
+end Lean.Meta.Sym

src/Lean/Meta/Sym/LooseBVarsS.lean

@@ -6,7 +6,7 @@ Authors: Leonardo de Moura
 module
 prelude
 public import Lean.Meta.Sym.SymM
-import Lean.Meta.Sym.ReplaceS
+public import Lean.Meta.Sym.ReplaceS
 public section
 namespace Lean.Meta.Sym
 
@@ -16,8 +16,8 @@ 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
+def lowerLooseBVarsS' (e : Expr) (s d : Nat) : AlphaShareBuilderM 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₁
@@ -29,12 +29,16 @@ def lowerLooseBVarsS (e : Expr) (s d : Nat) : SymM Expr := do
         return none
     | _ => return none
 
+@[inherit_doc lowerLooseBVarsS']
+def lowerLooseBVarsS (e : Expr) (s d : Nat) : SymM Expr :=
+  liftBuilderM <| lowerLooseBVarsS' e s d
+
 /--
 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
+def liftLooseBVarsS' (e : Expr) (s d : Nat) : AlphaShareBuilderM 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₁
@@ -46,4 +50,8 @@ def liftLooseBVarsS (e : Expr) (s d : Nat) : SymM Expr := do
         return none
     | _ => return none
 
+@[inherit_doc liftLooseBVarsS']
+def liftLooseBVarsS (e : Expr) (s d : Nat) : SymM Expr :=
+  liftBuilderM <| liftLooseBVarsS' e s d
+
 end Lean.Meta.Sym

src/Lean/Meta/Sym/ReplaceS.lean

@@ -14,6 +14,8 @@ A version of `replace_fn.h` that ensures the resulting expression is maximally s
 open Grind
 abbrev M := StateT (Std.HashMap (ExprPtr × Nat) Expr) AlphaShareBuilderM
 
+export Grind (AlphaShareBuilderM liftBuilderM)
+
 def save (key : ExprPtr × Nat) (r : Expr) : M Expr := do
   modify fun cache => cache.insert key r
   return r
@@ -44,11 +46,14 @@ 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
+@[inline] public def replaceS' (e : Expr) (f : Expr → Nat → AlphaShareBuilderM (Option Expr)) : AlphaShareBuilderM Expr := 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' {}
 
+@[inline] public def replaceS (e : Expr) (f : Expr → Nat → AlphaShareBuilderM (Option Expr)) : SymM Expr := do
+  liftBuilderM <| replaceS' e f
+
 end Lean.Meta.Sym

tests/lean/run/sym_instantiate.lean

@@ -0,0 +1,35 @@
+import Lean.Meta.Sym
+
+open Lean Meta Grind Sym
+
+def tst1 : SymM Unit := do
+  let f ← mkConstS `f
+  let e ← mkAppS (← mkAppS (← mkAppS f (← mkBVarS 0)) (← mkBVarS 1)) (← shareCommon (mkNatLit 1))
+  let a ← mkConstS `a
+  let b ← mkConstS `b
+  let r ← instantiateRevS e #[a, b]
+  assert! r == e.instantiateRev #[a, b]
+  logInfo r
+  let r ← instantiateS e #[a, b]
+  assert! r == e.instantiate #[a, b]
+  logInfo r
+  let e ← mkLambdaS `x .default (← mkConstS ``Nat) e
+  let r ← instantiateRevS e #[a, b]
+  assert! r == e.instantiateRev #[a, b]
+  logInfo r
+  let r ← instantiateS e #[a, b]
+  assert! r == e.instantiate #[a, b]
+  logInfo r
+
+
+/--
+info: f b a 1
+---
+info: f a b 1
+---
+info: fun x => f x b 1
+---
+info: fun x => f x a 1
+-/
+#guard_msgs in
+#eval SymM.run' tst1