More general command to modify variables

Stream: mathlib4

[Raghuram]

This is about a proposed more powerful version of the variable command.

Motivation

Often when developing some theory one first wants to prove something in a general case and then that stronger things hold in a "special" case. E.g., developing some module theory over a ring then proving extra results over a commutative ring, or proving some results about continuous functions on compact topological spaces, then "specialising" to compact uniform spaces to prove results about uniform continuity or something. (To be slightly more precise, B is a "special case" of A if hypothesis or structure B induces hypothesis or structure A in a canonical way that we have in mind.) I will illustrate with the example of Ring and CommRing for the rest of this message.

Lean allows us to do this in the following form:

variable {R M : Type*}

section ForRings
variable [Ring R] [AddCommGroup M] [Module R M]

-- Results true for general rings...

end ForRings

section ForCommRings
variable [CommRing R] [AddCommGroup M] [Module R M]

-- Results true for commutative rings...
-- We can use the results from the previous section, since there is an instance [CommRing R] -> [Ring R]

end ForCommRings

However, this seems non-ideal in a few ways. Even though R will in fact be a ring throughout the whole file, we have to separate the variable declarations of R and [Ring R] and introduce a lot of sections and end commands. If we were also introducing an R-module M, we would have to duplicate at least the [Module R M] variable declaration as well (and more, if one particularly about the order of the variables). Moreover, I don't think this is how mathematicians naturally think - as in the language I used above, one thinks of "specialising" from Ring to CommRing temporarily, and certainly doesn't forget that M is an R-module and remember it after specialising.

I think this is because the variable command can only add to the (end of the) context, forcing these workarounds when one wants to change the context in other ways. (omit doesn't remove variables from the context, only prevents them from being added to declarations - for example, I couldn't use omit [Ring R]; variable [CommRing R] because there's no way to tell Lean to replace the Ring R parameter to Module R M from the [Ring R] variable to the Ring R instance inferred from the CommRing R variable. omit/include is probably sufficient and works great if all structures and hypotheses are completely unbundled, but that's not the case in Mathlib.)

Proposal

I propose a new command to accomplish this (which could be called variable and shadow the default one if that's possible, or could have a new name like modify_variables, replace_variables, etc.) The syntax of the command would roughly be: specify sets of variables A to replace, B to introduce to replace A with, C to remove; and a way to produce terms of the types of A in terms of the remaining existing context plus the introduced variables of type B. Then the command would modify the current telescope of variables by removing A and C, introducing B, and replacing all references to the variables in A in subsequent variable types with the expressions for them in terms of B.

For example, the above example would become (with a rough syntax for the new command, where you provide the "conversion" terms in tactic mode starting with one goal for each variable in A):

variable {R M : Type*} [Ring R] [AddCommGroup M] [Module R M]

-- Theorems for general rings...

section ForCommRings
replace_variables [Ring R] with [CommRing R] by
  . infer_instance
-- The [Ring R] in the type of the variable [Module R M] would now
-- have been automatically replaced with CommRing.toRing [CommRing R].

-- Theorems true for commutative rings...

end ForCommRings

-- Old context restored - back to non-commutative rings
-- It should also be possible to use `replace_variables` with `in` for specialising hypotheses for a single conclusion.

1. Is something resembling this capability already in Lean 4 or Mathlib?
2. If not, would it be desirable?

P.S.

Some notes about the proposed behaviour:

• I am not sure whether in practice the functionality of removing variables will be necessary, but on the other hand it seems like it would be the easiest to implement anyway. I included it because I feel like there might be a case where some subset C of the old context can't be given in terms of B, but we don't need it and thus need it to be gone to do the replacement. I haven't thought of any examples of this yet, though.
• One subtlety in the implementation would be deciding where in the telescope of variables to insert B. I am not sure whether it would be desirable to allow the user to specify this or compute the "best" place to insert them. I think it should be straightforward, if not easy, to compute constraints on where each variable in B can be inserted, but it's not clear to me whether they should be inserted as late as possible or as early as possible, etc. Some user control might also be helpful for scenarios where the exact order of variables is important to give definitions made using them a nice interface.

[Kyle Miller]

Currently variable telescopes are stored completely syntactically (they are re-elaborated for every declaration).

I think it wouldn't be hard for someone to make an elaborator for the syntax

variable [CommRing R] → [Ring R]

that would replace the given binder with another. It would be a hack that looks at internals, but it's doable.

[Raghuram]

I see. I might try to implement a basic proof-of-concept then.

[Kyle Miller]

Another thing that's possible, which I haven't seen people do yet, is make macros for telescopes. E.g.

macro "variable_module " R:ident M:ident : command => `(variable [AddCommGroup $M] [Module $R $M])

and then you can do

variable [Ring R]
variable_module R M

and later in another section do

variable [CommRing R]
variable_module R M

It's clunky, but it's a way to use telescopes as near-first-class objects.

[Kyle Miller]

In fact, it wouldn't be hard to make a telescope command that defines such parameterized macros.

[Raghuram]

In fact, it was recently learning about variable?, which achieves a similar effect (in a narrower context than arbitrary telescope macros), that got me thinking about this.

[Kyle Miller]

I almost mentioned variable?, but it didn't seem like what you were looking for.

[Kyle Miller]

Here's a telescope-construction creator. Example:

telescope "mytele " α => [DecidableEq α] [Repr α]

This makes it so that if you use the mytele X command, you get [DecidableEq X] [Repr X].

import Lean

open Lean

syntax "telescope" str (ppSpace ident)* " =>" (ppSpace colGt bracketedBinder)* : command

/-- Wrap all occurrences of the given `ident` nodes in antiquotations -/
private partial def antiquote (vars : Array Syntax) (stx : Syntax) : Syntax :=
  match stx with
  | `($id:ident) =>
    if vars.any (fun var => var.getId == id.getId) then
      Syntax.mkAntiquotNode id (kind := `term) (isPseudoKind := true)
    else
      stx
  | _ =>
    match stx with
    | Syntax.node i k args => Syntax.node i k (args.map (antiquote vars))
    | stx => stx

macro_rules
  | `(command| telescope $name:str $vars:ident* => $binders*) => do
    let binders' : TSyntaxArray ``Parser.Term.bracketedBinder :=
      binders.map fun binder => ⟨antiquote vars binder⟩
    let args : TSyntaxArray ``Parser.Command.macroArg ←
      vars.mapM fun var => `(Parser.Command.macroArg| $var:ident:ident)
    `(macro $name:str $[$args]* : command => `(command| variable $binders'*))

telescope "mytele " α => [DecidableEq α] [Repr α]

variable (X : Type)
mytele X

#where
-- variable (X : Type) [DecidableEq✝ X] [Repr✝ X]

def test (a b : X) : String := if a = b then "equal" else toString (repr a)

[Kyle Miller]

(This is a reiteration of something I tried in Lean 3 a few years ago, but it's really nice this time being able to work with Syntax instead of strings!)

[Raghuram]

To confirm I'm reading the implementation correctly,

telescope "my_tele " a => [DecidableEq a] [Repr a]

expands to

macro "my_tele " a:ident : command => `(command| variable [DecidableEq $a] [Repr $a])

?

[Raghuram]

Kyle Miller said:

Currently variable telescopes are stored completely syntactically (they are re-elaborated for every declaration).

In that case, why does variable {R M : Type*} [Module R M] fail? Shouldn't it wait until a declaration before it elaborates the Module R M and fails to synthesise Semiring R etc.?

[Raghuram]

Is it because of this application of Term.elabBinders in docs#Lean.Elab.Command.elabVariable?

    runTermElabM fun _ => Term.withSynthesize <| Term.withAutoBoundImplicit <|
      Term.elabBinders binders fun xs => do
        ...

[Kyle Miller]

Yes, variable also elaborates the binders to check that they elaborate, but then it throws the result away.

[Raghuram]

ReplaceVariables.lean
Here is a first pass at replace_variables. It's still rather buggy, but I would appreciate any feedback from using it or reading the implementation.

The code inline:

import Lean

/-! # Util -/
section Util

def Array.findFinIdxM? {α : Type _} {m : Type _ → Type _} [Monad m]
    (as : Array α) (p : (i : Nat ) → i < as.size → α → m Bool) : m (Option (Fin as.size)) :=
  let rec @[specialize] find (i : Nat) : m (Option (Fin as.size)) :=
    if h:i < as.size then do
      if ← p i h as[i] then
        pure (some ⟨i, h⟩)
      else
        have := Nat.sub_succ_lt_self as.size i h
        find (i+1)
    else
      pure none
  termination_by as.size - i
  find 0

-- Strangely, `open Lean.Parser.Term in` and `open Lean.Parser.Term (bracketedBinderF)` (within a section) work,
-- but `open Lean.Parser.Term (bracketedBinderF) in` doesn't.
open Lean.Parser.Term (bracketedBinder bracketedBinderF)

/-- `Lean.Elab.Command.getBracketedBinderIds`, but for a more general kind of monad. -/
def Lean.Elab.Term.getBracketedBinderIds {m} [Monad m] [MonadError m] : Syntax → m (Array Name)
  | `(bracketedBinderF|($ids* $[: $ty?]? $(_annot?)?)) => return ids.map Syntax.getId
  | `(bracketedBinderF|{$ids* $[: $ty?]?})             => return ids.map Syntax.getId
  | `(bracketedBinderF|⦃$ids* : $_⦄)                   => return ids.map Syntax.getId
  | `(bracketedBinderF|[$id : $_])                     => return #[id.getId]
  | `(bracketedBinderF|[$_])                           => return #[Name.anonymous]
  | _                                                  => throwUnsupportedSyntax

/-- Splits a `bracketedBinder` into a series of binders, one per variable introduced. -/
def Lean.Elab.Term.splitBinder {m} [Monad m] [MonadQuotation m] [MonadError m] :
    TSyntax ``bracketedBinder → m (Array (TSyntax ``bracketedBinder))
| `(bracketedBinderF| ($ids* $[: $type?]?)) =>
  ids.mapM fun id ↦ `(bracketedBinderF| ($id $[: $type?]?))
| `(bracketedBinderF| {$ids* $[: $type?]?}) =>
  ids.mapM fun id ↦ `(bracketedBinderF| {$id $[: $type?]?})
| `(bracketedBinderF| ⦃$ids* $[: $type?]?⦄) =>
  ids.mapM fun id ↦ `(bracketedBinderF| ⦃$id $[: $type?]?⦄)
| `(bracketedBinderF| [$id : $type]) =>
  do pure #[← `(bracketedBinderF| [$id : $type])]
| `(bracketedBinderF| [$type]) =>
  do pure #[← `(bracketedBinderF| [$type])]
| _ => throwUnsupportedSyntax

end Util

/-! # Syntax -/
namespace Lean.Parser.Command

open Lean.Parser.Term (instBinder bracketedBinder) in
/-- A single clause for the `replace_variables` command. -/
def replaceVariablesItem := leading_parser
  (ident <|> instBinder) >> unicodeSymbol " → " " -> " >> checkColGt >> many Term.bracketedBinder

/-- `replace_variables id₁ → b₁, …` replaces variable `id₁` with new ones `b₁`, and so on.

Each `id` specifies a currently declared variable, either by name using the syntax `name` or by type using the syntax `[type]`.
Each `b` is a list of new variable binders to replace the variable specified by `id` with. -/
syntax (name := replaceVariables) "replace_variables " replaceVariablesItem,* : command

end Lean.Parser.Command

/-! # Elab -/
namespace Lean.Elab.Command

open Lean.Parser.Command (replaceVariablesItem replaceVariables)
open Lean.Parser.Term (bracketedBinder)

/-- Parse a `replace_variables` clause into pattern and new binders. -/
-- Based on code in the implementation of `Lean.Elab.Command.elabOmit`.
private def parseReplaceVariablesItem : TSyntax ``replaceVariablesItem →
    TermElabM (Syntax × (Name ⊕ Expr) × Array (TSyntax ``bracketedBinder))
-- TODO how to get `%$tk` on the correct node?
| `(replaceVariablesItem| $var:ident →%$tk $binders*)
| `(replaceVariablesItem| [$var : $_] →%$tk $binders*) =>
    pure (tk, .inl var.getId, binders)
| `(replaceVariablesItem| [$type] →%$tk $binders*) => do
    pure (tk, .inr (← Term.withoutErrToSorry <| Term.elabTermAndSynthesize type none),
          binders)
| _ => throwUnsupportedSyntax

open Lean.Meta (isDefEq) in
/-- Matches the variable `id` against `pattern`.

If `pattern` is a `Name`, it names a variable, so `id` must have the same name.
If `pattern` is an `Expr`, it specifies the type, so the type of `id` must unify with it. -/
-- Based on code in the implementation of `Lean.Elab.Command.elabOmit`.
private def variableMatch? (id : FVarId) : Sum Name Expr → MetaM Bool
| .inl name => do return (← id.getDecl).userName == name
| .inr type => do
    let decl ← id.getDecl
    -- Might as well allow `[type]` as syntax for matching any variable with of type `type`, rather than only instance-implicit ones.
    /- if !decl.binderInfo.isInstImplicit then pure false else -/
    let mctx ← getMCtx; isDefEq type decl.type <* setMCtx mctx

-- Heavily based on `elabOmit`, `replaceBinderAnnotation` and `elabVariable`.
@[command_elab replaceVariables] def elabReplaceVariables : CommandElab
| `(replace_variables%$tk $replacements:replaceVariablesItem,*) => do
  -- • Get some stuff from the scope that we'll need inside the `TermElabM`.
  -- `binderData` is an array of the same length as the total number of variables,
  -- consisting of the binders introducing each variable
  -- (thus, each binder is repeated as many times as it introduces variables).
  -- `oldVarUIds` is simply the current value of `scope.varUIds`.
  let (binderData, oldVarUIds) ← do
    let scope ← getScope
    pure (← scope.varDecls.flatMapM fun decl ↦ liftTermElabM (Term.splitBinder decl),
          scope.varUIds)
  -- • Elaborate current variables, as in `elabOmit`.
  let (newVarDecls, newVarUIds) ← runTermElabM fun oldVars ↦ do
    Term.synthesizeSyntheticMVarsNoPostponing
    Core.resetMessageLog
    /- Lean.logInfo m!"{oldVars.zip (binderData.zip oldVarUIds)}" -/

    -- • Parse the arguments to command (using the elaborated context).
    let mut replacements ← replacements.getElems.mapM parseReplaceVariablesItem

    -- • Build new list of variable declarations (and `varUId`s for the sake of `Scope`).
    let mut newVarDecls : Array (TSyntax ``bracketedBinder) := #[]
    let mut newVarUIds : Array Name := #[]
    for var in oldVars.reverse, decl in binderData.reverse, varUId in oldVarUIds.reverse do
      let fvar := var.fvarId!
      if let some idx ← replacements.findFinIdxM? fun _ _ (_, pattern, _) ↦
                          variableMatch? fvar pattern
      then
        -- Replace this variable with a list of binders `newBinders`.
        let (_, _, newBinders) := replacements[idx]
        replacements := replacements.eraseIdx idx
        newVarDecls := newVarDecls ++ newBinders
        newVarUIds := newVarUIds ++
          -- This is from `Lean.Elab.Command.replaceBinderAnnotations`.
          (← (← newBinders.flatMapM Term.getBracketedBinderIds)
                        |>.mapM (withFreshMacroScope ∘ MonadQuotation.addMacroScope))
      else
        -- Preserve this variable.
        newVarDecls := newVarDecls.push decl
        newVarUIds := newVarUIds.push varUId
    /- Lean.logInfo m!"Leftover replacements: {replacements.map (·.2.1)}" -/
    for (tk, pattern, _) in replacements do
      throwErrorAt tk match pattern with
                      | .inl name => m!"no variable found with name {name}"
                      | .inr type => m!"no variable found with type unifying with {type}"
    pure (newVarDecls, newVarUIds)

  -- • Elaborate the new variable telescope to check that it is valid, as in `elabVariable`.
  liftTermElabM <| Term.withSynthesize <| Term.withAutoBoundImplicit <|
  Term.elabBinders newVarDecls fun xs => do
    let _ ← Term.addAutoBoundImplicits xs (tk.getTailPos? (canonicalOnly := true))

  -- • Finally, modify the scope.
  modifyScope fun scope ↦ { scope with
    varDecls := newVarDecls.reverse,
    varUIds := newVarUIds.reverse,
    includedVars := scope.includedVars.filter newVarUIds.contains,
    omittedVars := scope.omittedVars.filter newVarUIds.contains
  }
| _ => throwUnsupportedSyntax

end Lean.Elab.Command

/-!
Known issues:
• If re-elaborating after modifying results in errors, they are shown on the original variable declaration.
• If `autoImplicit` is false while executing `replace_variables`, there are errors even for variables which were declared.
• Errors for variables not found are on the arrow.
• Unused variable linting is applied to the newly introduced binders for some reason.
Notes:
• It is possible to create accepted variable telescopes that are rejected when used in a declaration, if `autoImplicit` is true for `replace_variables` but not for the declaration.
  But this is true for `variable` as well.
• The elaboration check for `replace_variables` does not seem to consider previous variables in typeclass inference, the way that `variable` does.
-/

[Raghuram]

Also the rudimentary tests I've written so far:

section Test1

variable {R M : Type _} [Add R] [Mul R] [Add M] [SMul R M]

-- variable {R M : Type _} [Add R] [Mul R] [Add M] [SMul R M]
#where

-- TODO stop it from linting that α is unused
replace_variables R → (α : Type _)
-- variable (α : Type _) {M : Type _} [Add R] [Mul R] [Add M] [SMul R M]
#where

-- no variable found with type unifying with `Pow R`
replace_variables [Pow R] → [Mul N]

-- With autoImplicit false, would say "unknown identifier 'N'".
-- However, this causes errors on the first `variable` line.
-- set_option autoImplicit false in
replace_variables α → (_ : Mul N)
#where

end Test1

section TestInstances

class Ring (R : Type _)

class CommRing (R : Type _) extends Ring R

-- `inst.toRing`
variable {R} [inst : CommRing R] in #synth Ring R
-- failed to synthesize `CommRing R`
variable {R} [inst : Ring R] in #synth CommRing R

class Module (R M : Type _) [Ring R]

class CommModule (R M : Type _) [CommRing R]

section
variable {R} [Ring R] {M} [Module R M]
-- TODO fix: causes error at `Module R M`
replace_variables [Ring R] → [Ring R]
end

end TestInstances