Modifying the `reassoc` attribute for monadic computations

Stream: metaprogramming / tactics

[Brendan Seamas Murphy]

I would like to have a variant of reassoc which generates associated lemmas for theorems like ∀ (s₁ s₂ : σ), (do set s₁; set s₂ : m PUnit) = set s₂ which states a universally quantified equality between stuff of the form m α for some monad m (possible being quantified over). Here is an excerpt of my intended applications

class LawfulMonadStateOf (σ : Type u) (m : Type u → Type v)
    [MonadStateOf σ m] [Monad m] extends LawfulMonad m where
  set_set : ∀ (s₁ s₂ : σ), (do set s₁; set s₂ : m PUnit) = set s₂
  set_get : ∀ (s : σ), (do set s; get : m σ) = (do set s; pure s)
  get_set : (do let s ← (get : m σ); set s) = pure PUnit.unit
  get_get_pair : (do let s₁ ← (get : m σ); let s₂ ← get; pure (s₁, s₂))
               = (do let s ← get; pure (s, s))
  modifyGet_def : ∀ {α} (f : σ → α × σ),
    (modifyGet f : m α) = (do let p := f (← get); set p.snd; pure p.fst)

attribute [reassocM (attr := simp)] set_set
-- attribute [reassocM (attr := simp)] set_get get_set modifyGet_def

And here is my attempt at modifying the Mathlib.Tactic.CategoryTheory.Assoc file:

import Mathlib.Util.AddRelatedDecl
import Mathlib.Lean.Meta.Simp
import Mathlib.Lean.Meta

open Lean Meta Elab Tactic
open Mathlib.Tactic

namespace LawfulMonad

universe u v

theorem eq_bind {m : Type u → Type v} [Monad m] {α} {x y : m α} (w : x = y)
  {β} (f : α → m β) : x >>= f = y >>= f := by rw [w]

def monadSimp (e : Expr) : MetaM Simp.Result :=
  getAllSimpDecls `monad_norm >>=
  (simpOnlyNames . e (config := { decide := false }))

def reassocMExpr (e : Expr) : MetaM Expr := do
  mapForallTelescope (fun e => do simpType monadSimp (← mkAppM ``eq_bind #[e])) e

syntax (name := reassocM) "reassocM" (" (" &"attr" ":=" Parser.Term.attrInstance,* ")")? : attr

initialize registerBuiltinAttribute {
  name := `reassocM
  descr := ""
  applicationTime := .afterCompilation
  add := fun src ref kind => match ref with
  | `(attr| reassocM $[(attr := $stx?,*)]?) => MetaM.run' do
    if (kind != AttributeKind.global) then
      throwError "`reassocM` can only be used as a global attribute"
    addRelatedDecl src "_assoc" ref stx? fun type value levels => do
      pure (← reassocMExpr (← mkExpectedTypeHint value type), levels)
  | _ => throwUnsupportedSyntax }

open Term in
elab "reassocM_of% " t:term : term => do
  reassocMExpr (← elabTerm t none)

end LawfulMonad

I'm currently getting an error that looks like

... has type
  (do set s₁; set s₂) = set s₂ : Prop
but is expected to have type
  @Eq.{?u.1614 + 1} (?m ?α) ?x ?y  : Prop

so my assumption is that lean is having trouble inferring the m and α arguments to eq_bind in reassocMExpr. Does anyone know how I can make this work?

[Kyle Miller]

You need to set up the additional binders for m and α so that they come out as universally quantified

[Brendan Seamas Murphy]

I thought that's what mapForallTelescope was doing, but I guess I misunderstood?

[Kyle Miller]

You could look at the elementwise attribute instead, which might be more similar to what you want to do

[Kyle Miller]

I might be misunderstanding what you need to do exactly -- maybe take a look at mkAppOptM as well, if you know the m, etc., that you need to pass in.

[Brendan Seamas Murphy]

It's supposed to work for theorems involving universally quantified monads or even multiple monads. I guess what I had in mind is that it would figure out the monad by matching the type of the eqation (after all universal quantifiers) with m a. So maybe I'll try to do that match and then use mkAppOptMwith the discovered m passed in?

[Kyle Miller]

(Btw, mapForallTelescope is able to take the binders from the type of e itself, but if you need to introduce new variables you'll have to do that inside the body of the fun using withLocalDecl to create the local binding and then use mkForallFVars or mkLambdaFVars (depending on what you want) before leaving withLocalDecl)

[Kyle Miller]

One way you could do this is by using forallMetaTelescope on the type of eq_bind, unifying the correct argument with e, and then passing the remaining arguments to mkAppN ``eq_bind metavars, and then use mkLambdaFVars on this with the unsolved-for metavariables. It turns out mkLambdaFVars is able to take metavariables too, despite the name.

[Kyle Miller]

That's for the fully general multi-monad lemma, if I'm understanding this right.

[Brendan Seamas Murphy]

also re: what I want this to do, an example would be to generate

@[simp]
theorem set_set_assoc {σ : Type u} {m : Type u → Type v} [MonadStateOf σ m]
    [Monad m] [LawfulMonad m] [LawfulMonadStateOf σ m]
  : ∀ (s₁ s₂ : σ) {β} (f : PUnit → m β), (do set s₁; let u ← set s₂; f u) = (do let u ← set s₂; f u)

from the annotation attribute [reassocM (attr := simp)] set_set, where set_set : ∀ (s₁ s₂ : σ), (do set s₁; set s₂ : m PUnit) = set s₂ is a field of the typeclass LawfulMonadStateOf

[Brendan Seamas Murphy]

as a less simple example, I think this should be able to generate

@[simp]
theorem get_set_pure_eq_get_assoc {α} (f : σ → m α)
    : (do let s ← get; set s; f s) = get >>= f

from get_set_pure_eq_get : (do let s ← get; set s; pure s) = (get : m σ)

[Brendan Seamas Murphy]

The multi monad stuff would be if I wanted to use this for a hypothetical LawfulMonadLift class, where axioms or theorem statements quantify over multiple monads

[Kyle Miller]

@Brendan Seamas Murphy I get it now -- it was a failure to do higher-order unification I think. Here's a replacement function:

def reassocMExpr (e : Expr) : MetaM Expr := do
  mapForallTelescope (fun e => do
    let ety ← instantiateMVars (← inferType e)
    let some (ty, _, _) := ety.eq? | throwError "Expecting an equality, not{indentD ety}"
    let .app m _ := ty | throwError "Expecting an equality for a monad, not{indentD ty}"
    simpType monadSimp (← mkAppOptM ``eq_bind #[some m, none, none, none, none, some e])) e

[Kyle Miller]

Does this look right?

attribute [reassocM (attr := simp)] LawfulMonadStateOf.set_set

#print LawfulMonadStateOf.set_set_assoc
/-
⊢ ∀ {σ : Type u} {m : Type u → Type v} [inst : MonadStateOf σ m] [inst_1 : Monad m] [self : LawfulMonadStateOf σ m]
  (s₁ s₂ : σ) {β : Type u} (f : PUnit → m β),
  (do
        set s₁
        set s₂) >>=
      f =
    set s₂ >>= f
-/

[Brendan Seamas Murphy]

Hm, it should be applying the simplifier with bind_assoc. Maybe the monad_norm tag needs to be in scope when I define reassocM?

[Brendan Seamas Murphy]

I don't understand metaprogramming in lean well :sweat_smile:

[Kyle Miller]

I just have the code you gave me -- I don't have any monad_norm lemmas

[Brendan Seamas Murphy]

ahh, of course

[Brendan Seamas Murphy]

That looks awesome then, thanks!

[Kyle Miller]

If you make the m argument explicit in eq_bind, then you can switch the mkAppOptM to mkAppM ``eq_bind #[m, e]

[Brendan Seamas Murphy]

Is there a reason that's better? Also, it simplifies correctly when in a context with monad_norm lemmas so hooray!

[Kyle Miller]

I feel like it makes the meta code cleaner, rather than needing to count implicit arguments carefully. I think it's good making the lemmas that support meta code be tailored to it.

[Brendan Seamas Murphy]

Makes sense, I'll change that then. Thanks for the help!

[Kyle Miller]

You don't need to, but it's something you could do if you wanted.