modifying local context in MetaM

Stream: metaprogramming / tactics

[Tomas Skrivan]

I would like to have a variant of MetaM monad that can modify the local context. I find passing local context around or working with continuation(e.g. lambdaTelescope) a bit cumbersome sometimes.

Effectively instead of working with

abbrev MetaM  := ReaderT Context $ StateRefT State CoreM

I would like to work with

abbrev MetaLCtxM  := StateT Context $ StateRefT State CoreM

Is this a good/bad idea? Is this done in core Lean in some capacity?

---

Using this monad I would like to write code like

def lambdaIntro (e : Expr) : MetaLCtxM (Expr × Array Expr) := do
  Meta.lambdaTelescope e fun xs e' => do
    let ctx ← getThe ContextCtx
    fun _ _ => return ((e',xs), ctx)

open Lean Meta Qq in
#eval show MetaLCtxM Unit from do
  let e := q(fun a b => a + b + 42)

  let (b, xs) ← lambdaIntro e

  IO.println s!"lambda body: {← ppExpr b}"
  IO.println s!"lambda vars: {← liftM <| xs.mapM ppExpr}"
  IO.println s!"lambda: {← ppExpr (← mkLambdaFVars xs b)}"

My application is a tactic dealing with a big expression with lots of let bindings. I want to just keep on collecting (let) free variables and not worry about passing around the local context.

---

Here is my current attempt at setting up this monad. My main concern is the instance MonadControl MetaM MetaLCtxM, not sure if I got it right and I just hacked the code together until it typechecks.

code

import Lean
import Qq

open Lean Meta

structure _root_.Lean.Meta.ContextCtx where
  lctx              : LocalContext         := {}
  localInstances    : LocalInstances       := #[]

structure _root_.Lean.Meta.ContextCfg where
  config            : Meta.Config               := {}
  defEqCtx?         : Option DefEqContext  := none
  synthPendingDepth : Nat                  := 0
  canUnfold?        : Option (Meta.Config → ConstantInfo → CoreM Bool) := none

def _root_.Lean.Meta.Context.mkCtxCfg (ctx : ContextCtx) (cfg : ContextCfg) : Meta.Context :=
  { config := cfg.config
    lctx := ctx.lctx
    localInstances := ctx.localInstances
    defEqCtx? := cfg.defEqCtx?
    synthPendingDepth := cfg.synthPendingDepth
    canUnfold? := cfg.canUnfold? }

abbrev MetaLCtxM  := ReaderT Meta.ContextCfg $ StateT Meta.ContextCtx $ StateRefT Meta.State CoreM

@[always_inline]
instance : Monad MetaLCtxM := let i := inferInstanceAs (Monad MetaLCtxM); { pure := i.pure, bind := i.bind }

instance {α} [Inhabited α] : Inhabited (MetaLCtxM α) where
  default := fun _ ctx => do pure (default,ctx)

instance : MonadLCtx MetaLCtxM where
  getLCtx := return (← get).lctx

instance : MonadMCtx MetaLCtxM where
  getMCtx    := return (← getThe Meta.State).mctx
  modifyMCtx f := modifyThe Meta.State fun s => { s with mctx := f s.mctx }

instance : MonadEnv MetaLCtxM where
  getEnv      := return (← getThe Core.State).env
  modifyEnv f := do modifyThe Core.State fun s => { s with env := f s.env, cache := {} }; modifyThe Meta.State fun s => { s with cache := {} }

instance : AddMessageContext MetaLCtxM where
  addMessageContext := addMessageContextFull

instance : MonadLift MetaM MetaLCtxM where
  monadLift x := fun cfg ctx => do let r ← x (.mkCtxCfg ctx cfg); pure (r, ctx)

protected def MetaLCtxM.saveState : MetaLCtxM (SavedState×ContextCtx) :=
  return ({ core := (← getThe Core.State), meta := (← get) }, ⟨← getLCtx, ← getLocalInstances⟩)

def MetaLCtxM.restore (b : SavedState) (ctx : ContextCtx) : MetaLCtxM Unit := do
  Core.restore b.core
  modify fun s => { s with mctx := b.meta.mctx, zetaDeltaFVarIds := b.meta.zetaDeltaFVarIds, postponed := b.meta.postponed }
  modifyThe ContextCtx fun _ => ctx

instance : MonadBacktrack (SavedState×ContextCtx) MetaLCtxM where
  saveState      := MetaLCtxM.saveState
  restoreState s := MetaLCtxM.restore s.1 s.2

@[inline] def MetaLCtxM.run (x : MetaLCtxM α) (cfg : ContextCfg := {}) (ctx : ContextCtx := {}) (s : Meta.State := {}) :
    CoreM (α × Meta.State) := do
  let ((r,_),s) ← x cfg ctx |>.run s
  return (r,s)

@[inline] def MetaLCtxM.run' (x : MetaLCtxM α) (cfg : ContextCfg := {}) (ctx : ContextCtx := {}) (s : Meta.State := {}) : CoreM α := do
  let ((r,_),_) ← x cfg ctx |>.run s
  return r


instance [MetaEval α] : MetaEval (MetaLCtxM α) :=
  ⟨fun env opts x _ => MetaEval.eval env opts x.run' true⟩


instance : MonadControl MetaM MetaLCtxM where
  stM      := fun α => α × ContextCtx
  liftWith := fun f => do
    let cfg ← readThe ContextCfg
    let f' := (f (fun x c s => do
                      let (x',ctx') ← x cfg ⟨c.lctx,c.localInstances⟩ s
                      return (x', ctx')))
    f'
  restoreM := fun x => do let (a, s) ← liftM x; set s; pure a


def lambdaIntro (e : Expr) : MetaLCtxM (Expr × Array Expr) := do
  Meta.lambdaTelescope e fun xs e' => do
    let ctx ← getThe ContextCtx
    fun _ _ => return ((e',xs), ctx)

open Lean Meta Qq in
#eval show MetaLCtxM Unit from do
  let e := q(fun a b => a + b + 42)

  let (b, xs) ← lambdaIntro e

  IO.println s!"lambda body: {← ppExpr b}"
  IO.println s!"lambda vars: {← liftM <| xs.mapM ppExpr}"
  IO.println s!"lambda: {← ppExpr (← mkLambdaFVars xs b)}"

[Eric Wieser]

In #5863 (here) my solution to this (thanks I think to @Kyle Miller) was to work in ContT to transform the continuations back into a monad

[Eric Wieser]

But that ran into the rather thorny #5897 (discussed here), which found that you cant really make ContT support throw without running into issues with catch.

[Tomas Skrivan]

Using ContT is a great idea! I will play around with it.