Extending do notation with `+=`

Stream: lean4

[Tomas Skrivan]

I would like to extend do notation with x += 1, x[0] := 1, x[0] += 1 and alike.

I have figured out how to extend the parser, but I do not know how to correctly follow up with elaboration.

import Lean.Elab.Term
import Lean.Parser.Term
import Lean.Parser.Do

open Lean Parser Term Elab

@[doElemParser] def Lean.Term.Parser.doFoo := leading_parser "foo"

#eval Id.run do
  foo -- Error: unexpected do-element of kind Lean.Term.Parser.doFoo: foo

def letAddDeclNoBinders := node `Lean.Parser.Term.letIdDecl $ atomic (Lean.Parser.Term.ident >> pushNone >> optType >> " += ") >> termParser
@[doElemParser] def Lean.Term.Parser.doAddReassign := leading_parser notFollowedByRedefinedTermToken >> letAddDeclNoBinders

-- This should be called in Lean.Elab.Term.Do.ToCodeBlock.doSeqToCode
def getDoAddReassignVars (doReassign : Syntax) : TermElabM (Array Name) := sorry

#eval Id.run do
  let mut x := 0
  x += 1 -- Error: unexpected do-element of kind Lean.Term.Parser.doAddReassign: x += 1
  x

The code is failing inside of doSeqToCode . There does not seem to be a path for a custom do-element :( How can I fix it?

[Mario Carneiro]

you could try handling it as a macro instead of an elab

[Mario Carneiro]

syntax atomic(ident Term.optType) " += " term : doElem
macro_rules
| `(doElem| $x:ident $[: $ty]? += $e) => `(doElem| $x:ident $[: $ty]? := $x:ident + $e)

#eval Id.run do
  let mut x := 0
  x += 1
  pure x

[Tomas Skrivan]

Thanks a lot! Definitely a better approach. I was trying to do lots of different things with syntax and macro_rules, but I have always failed to get anywhere.

[Mario Carneiro]

it is still limited in that you have to produce a doSeq syntax at the end of the day, there is no equivalent of elab for such things

[Mario Carneiro]

but for += that's probably fine

[Tomas Skrivan]

This is nice, x[0] := 42 is working too!

syntax atomic(Lean.Parser.Term.ident) noWs "[" term "] := " term : doElem

macro_rules
| `(doElem| $x:ident[ $i:term ] := $xi) => `(doElem| $x:ident := ($x:ident).set! $i $xi)

#eval Id.run do
  let mut x : Array Nat := #[0,1,2,3]
  x[0] := 42
  x

There is one problem that there has to be exactly one space between x[0] and :=. Replacing "] := " with "]" " := " in the syntax definition complains about multiple interpretations.

[Mario Carneiro]

yes, because lean already defines this syntax

[Mario Carneiro]

I think it is a TODO right now though

[Mario Carneiro]

you can make your syntax higher priority

[Tomas Skrivan]

Mario Carneiro said:

yes, because lean already defines this syntax

Ohh, is the notation x[0] := 42 new? I though that it is not possible.

[Tomas Skrivan]

And how can I find the definition of the built in x[0] := 42 syntax? So I can make sure my types are supporting it.

[Mario Carneiro]

I take it back, there is no special support in this particular position. However the syntax term := term is accepted because the left term can be a pattern, and x[0] := 42 falls into this

[Tomas Skrivan]

As you have mentioned, changing priority solves the problem:

import Lean.Parser.Term

open Lean Parser Term

syntax (priority := high) atomic(Term.ident) noWs "[" term "]" " := " term : doElem

macro_rules
| `(doElem| $x:ident[ $i:term ] := $xi) => `(doElem| $x:ident := ($x:ident).set! $i $xi)

#eval Id.run do
  let mut x : Array Nat := #[0,1,2,3]
  x[0] := 42
  x[1]   := 13
  x   --- prints: #[42, 13, 2, 3]

[James Gallicchio]

Was this ever committed somewhere? Seems a good addition to Std, particularly if we extend it to x[i][j] esque stuff

[Gabriel Ebner]

See also the later discussion at https://github.com/leanprover/lean4/issues/905

[James Gallicchio]

I'm not gonna lie, most of those syntax examples are pretty meaningless to me at first glance :joy: imperative programming is hard...

[James Gallicchio]

I'll stick to the functional stuff for now.