Setting elements of arrays with GetElem

Stream: std4

[Frédéric Dupuis]

Do we have something like this?

def natSet (as : Array α) (i : Nat) (v : α) (hi : i < as.size := by get_elem_tactic) : Array α :=
  Array.set as ⟨i, hi⟩ v

As far as I can tell, we only have Array.set which takes in a Fin index, Array.setD which does nothing if the index is out of bound, and Array.set! which panics.

[Joachim Breitner]

While we are at it, would it be good to have notation like

do
  a[i] := x

that desugars to

do
  a := a.natSet i x

?

[Mac Malone]

@Joachim Breitner See lean4#905 for an old discussion of better do notation assignment syntax. :grinning_face_with_smiling_eyes:

[Kyle Miller]

Patrick and I experimented earlier this year with a[i] := x notation. We did set a[i] := x since that made it be unambiguous (the normal x := v parser parses any term x for the LHS and interprets it as a pattern -- maybe we could have carved out indexing notation, but I don't remember if we tried).

In action it looks like

def matrix.swapRowsp (A : matrix m n) (row row' : ℕ)
    (hrow : row < m := by set_elem_tac) (hrow' : row' < m := by set_elem_tac) :
    matrix m n := Id.run do
  if row == row' then return A
  let mut B := A
  for h : j in [0:n] do
    let tmp := B[row, j]
    set B[row, j] := B[row', j]
    set B[row', j] := tmp
  return B

(set_elem_tac is like get_elem_tac, and I don't remember if it's any different. These were extended to deal with tuple indexing though.)

code

/-! # SetElem class
-/

class SetElem (cont : Type u) (idx : Type v) (elem : outParam (Type w)) (dom : outParam (cont → idx → Prop)) where
  /--
  Like `GetElem.getElem` but for setting.

  Given a `mut` variable `arr` in a `do` block:
  * `set arr[i] := v`: proves the proof side goal by `get_elem_tactic`
  * `set arr[i]! := v`: panics `i` is out of bounds
  * `set arr[i]? := v`: does nothing if `i` is out of bounds
  * `set arr[i]'h := v`: uses `h` to prove the side goal
  -/
  setElem (xs : cont) (i : idx) (v : elem) (h : dom xs i) : cont

export SetElem (setElem)

@[never_extract]
private def outOfBounds [Inhabited α] : α :=
  panic! "index out of bounds"

/--
Set `xs[i] := v` or panic if `i` is out of bounds.
-/
@[inline] def setElem! [SetElem cont idx elem dom] (xs : cont) (i : idx) (v : elem) [Decidable (dom xs i)] : cont :=
  if h : _ then
    setElem xs i v h
  else
    have : Inhabited cont := ⟨xs⟩
    outOfBounds

/--
Set `xs[i] := v` if `i` is within bounds, otherwise do nothing.
-/
@[inline] def setElem? [SetElem cont idx elem dom] (xs : cont) (i : idx) (v : elem) [Decidable (dom xs i)] : cont :=
  if h : _ then setElem xs i v h else xs

declare_syntax_cat setExpr

syntax ident : setExpr
@[inherit_doc setElem] syntax ident noWs "[" term,+ "]" : setExpr
@[inherit_doc setElem] syntax ident noWs "[" term,+ "]'" term:max : setExpr
@[inherit_doc setElem] syntax ident noWs "[" term,+ "]" noWs "!" : setExpr
@[inherit_doc setElem] syntax ident noWs "[" term,+ "]" noWs "?" : setExpr

/-- Notation for updating a value. -/
syntax "set " setExpr " := " term : doElem

macro_rules
  | `(doElem| set $x:ident := $val) => `(doElem| $x:term := $val)
  | `(doElem| set $x:ident[$i] := $val) => `(doElem| $x:term := setElem $x $i $val (by get_elem_tactic))
  | `(doElem| set $x:ident[$i]'$h := $val) => `(doElem| $x:term := setElem $x $i $val $h)
  | `(doElem| set $x:ident[$i]! := $val) => `(doElem| $x:term := setElem! $x $i $val)
  | `(doElem| set $x:ident[$i]? := $val) => `(doElem| $x:term := setElem? $x $i $val)

-- Also support `set x[i1,i2,i3] := v` for `set x[(i1,i2,i3)] := v`
macro_rules
  | `(doElem| set $x:ident[$i,$is,*] := $val) => `(doElem| set $x:ident[($i,$is,*)] := $val)
  | `(doElem| set $x:ident[$i,$is,*]'$h := $val) => `(doElem| set $x:ident[($i,$is,*)]'$h := $val)
  | `(doElem| set $x:ident[$i,$is,*]! := $val) => `(doElem| set $x:ident[($i,$is,*)]! := $val)
  | `(doElem| set $x:ident[$i,$is,*]? := $val) => `(doElem| set $x:ident[($i,$is,*)]? := $val)

[Kyle Miller]

It's nice having array indexing have a symmetry before and after the := for code like this.

[Kyle Miller]

More generally, it's nice when a language has a concept of lvalues, which are expressions that come with assignment semantics. Common Lisp has a cool system of composable "setf macros", where in (setf p v) you explain for a given form p how to do the assignment v -> p.

Using Lean notation, that could mean macros that turn set s.x[1] := v into s := {s with x := s.x.set 1 v}. The macros could be written to be more careful about affine usage of the variables than this.

[Kyle Miller]

(In fact, setf macros were what justified prefixing these assignments with set for me. It also serves as a warning that there might possibly be non-affine uses of variables.)

[Sebastian Ullrich]

Lean already has lvalues, they just haven't made it to do yet

#check fun (p : Prod (Prod (Array Nat) Unit) Unit) => { p with fst.fst[0] := 1 }

[Kyle Miller]

If we had setf-style lvalues, then we could do something more general than what's available to structure notation, or at least I think -- I don't know if what you can do with structures is extensible.

The idea of setf is to make accessing and storing look the same, no matter the accessor function. For example, if you have a hash map m, you can do m.find! k to access, so you should be able to define a set macro so that set m.find! k := v means m := m.insert k v. (Justification for using find!: after this, m.find! k evaluates to v.)

I believe the lenses were invented for this sort of transformation, but we're in Lean and we have metaprogramming :-)

[Kyle Miller]

Common Lisp Hyperspec: setf and kinds of places