simp and vector notation

Stream: mathlib4

[Michael Stoll]

would it be possible to make (say)

import Mathlib

example {F} [Field F] : (![(0 : Int), 1, 0, -1, 0, -1, 0, 1] 6 : F) = 0 := by
  simp

work?

(Replacing 6 by 4, it does work because of docs#Matrix.cons_val_four, but there is docs#Matrix.cons_val_five etc. Without the coercion, it's rfl.)

Is this something a simproc could do?

[Eric Wieser]

Yes, this is a good area for a simproc

[Eric Wieser]

Things like docs#FinVec.map_eq are also waiting for a simproc

[Eric Wieser]

Things like docs#FinVec.map_eq are also waiting for a simproc

[Michael Stoll]

(This may be a regression from lean3.)

[Eric Wieser]

In lean3 we had a clever trick that exploited the binary representation of numbers

[Eric Wieser]

But a simproc makes more sense here

[Michael Stoll]

Anybody wants to write it? :smile:

[Kim Morrison]

(deleted)

[Michael Stoll]

(Vector notation is defined in Mathlib, I think, that's why I was asking in this channel. Not sure if the question is relevant outside Mathlib.)

[Eric Wieser]

I suspect Kim was thinking of the analogous array / list versions?

[Eric Wieser]

Here's a prototype:

import Mathlib.Data.Fin.VecNotation
import Mathlib.Util.Qq


open Lean Qq

partial def matchVecApp {u : Level} {α : Q(Type u)} (x : Q($α)) :
  MetaM <| Option (Σ' (n : Q(ℕ)) (f : Q(Fin $n → $α)) (i : Q(Fin $n)), $x =Q $f $i) := do
    Meta.withNewMCtxDepth do
      let n : Q(ℕ) ← mkFreshExprMVarQ q(ℕ)
      let f : Q(Fin $n → $α) ← mkFreshExprMVarQ q(Fin $n → $α)
      let i : Q(Fin $n) ← mkFreshExprMVarQ q(Fin $n)
      let .defEq _h ← isDefEqQ q($x) q($f $i) | pure none
      return some ⟨← instantiateMVars n, ← instantiateMVars f, ← instantiateMVars i, ⟨⟩⟩

partial def matchVecCons {u : Level} {α : Q(Type u)}
    {n : Q(ℕ)} (f : Q(Fin ($n + 1) → $α)) :
  MetaM <| Option (Σ' (a : Q($α)) (f' : Q(Fin $n → $α)), $f =Q Matrix.vecCons (α := $α) (n := $n) $a $f') := do
    Meta.withNewMCtxDepth do
      let a : Q($α) ← mkFreshExprMVarQ q($α)
      let f' : Q(Fin $n → $α) ← mkFreshExprMVarQ q(Fin $n → $α)
      let .defEq _h ← isDefEqQ (u := u.succ) (α := q(Fin ($n + 1) → $α)) q($f) q(Matrix.vecCons (α := $α) (n := $n) $a $f') | pure none
      return some ⟨← instantiateMVars a, ← instantiateMVars f', ⟨⟩⟩


partial def reduce {u : Level} {α : Q(Type u)}
    {n : Q(ℕ)} (f : Q(Fin $n → $α)) (i : Q(Fin $n)) :
    MetaM (Option (Σ a : Q($α), Q($f $i = $a))) := do
  let .some iv := i.int? | pure none
  let .some nv := n.nat? | pure none
  let .succ nv' := nv | pure none
  have n' : Q(ℕ) := mkNatLit nv'
  have _hn : $n =Q ($n' + 1) := ⟨⟩
  let .some ⟨a, f', _h⟩ ← matchVecCons (n := n') (α := α) f | pure none
  match (iv % (nv' + 1)).toNat with
  | 0 =>
    have _hi : $i =Q 0 := ⟨⟩
    return some ⟨a, q(Matrix.cons_val_zero _ _)⟩
  | iv' + 1 =>
    have ilit : Q(ℕ) := mkRawNatLit iv'
    have : Q(NeZero $n') := q(sorry) -- TODO
    have i' : Q(Fin $n') := q(OfNat.ofNat $ilit)
    have _hi : $i =Q (Fin.succ $i') := ⟨⟩
    let .some ⟨x, pf⟩ ← reduce (u := u) (α := α) (n := n') q($f') i' | pure none
    return some ⟨x, q((Matrix.cons_val_succ $x $f' $i').trans $pf)⟩

simproc Matrix.cons_val (Matrix.vecCons _ _ _) := fun e => do
  let ⟨_u, _α, x⟩ ← inferTypeQ' e
  let .some ⟨_n, f, i, _h⟩ ← matchVecApp x | failure
  let .some ⟨e, pf⟩ ← reduce f i | failure
  return .continue (some { expr := e, proof? := pf})

lemma foo {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] 33 = b := by simp
lemma bar {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] (-1) = d := by simp

[Eric Wieser]

I now realize the above is more complicated than it needs to be: the proof is always rfl, so I can avoid all the effort to track the proof

[Eric Wieser]

I've made a more polished attempt in #13578

[Eric Wieser]

What I can't work out is how to obtain the length of the vector; seemingly .nat? fails

[Eric Wieser]

I suppose repeatedly unifying with Nat.succ is probably better

[Heather Macbeth]

I've wanted things like this several times, and indeed I've tried to arrange things so that the matrix indices which are produced in the earlier part of the work are all iterated .succs. This works:

example {F} [Field F] :
    (![(0 : Int), 1, 0, -1, 0, -1, 0, 1] (0 : Fin 2).succ.succ.succ.succ.succ.succ : F) = 0 := by
  simp

[Kyle Miller]

I'm not sure about all the isDefEq in your simproc @Eric Wieser. It seems like this should be more syntactically driven, only working for things that are obviously vector literals.

Here's a short implementation using core style. The expr matches handle mdata and annotation cleanup.

import Mathlib.Data.Fin.VecNotation

open Lean

partial def matchVecLit (e : Expr) : Option (List Expr) :=
  match_expr e with
  | Matrix.vecEmpty _ => some []
  | Matrix.vecCons _ _ x xs => (x :: ·) <$> matchVecLit xs
  | _ => none

dsimproc Matrix.cons_val (Matrix.vecCons _ _ _) := fun e => do
  let_expr Matrix.vecCons _ _ x xs' n := e | return .continue
  let some n := n.int? | return .continue -- The docstring claims only nat or int, but works fine for Fin
  let some xs' := matchVecLit xs' | return .continue
  let xs := x :: xs'
  let n' := (n % xs.length).toNat
  return .continue (xs.get! n')

lemma foo {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] 33 = b := by dsimp
lemma bar {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] (-1) = d := by simp

[Eric Wieser]

One difference is that I think my version works for vecCons a (vecCons b a_three_tuple) 6 = b

[Eric Wieser]

I'm not sure I agree re defeq; shouldn't _everything_ see through reducible defeq? (Probably I forgot to make it reducible)

[Kyle Miller]

If you imagine that the simproc is supposed to be looking for a dynamic discrimination tree key, simulating an infinite collection of simp lemmas, then maybe it shouldn't see through reducible defeq when matching?

It's not clear to me that this simproc should handle "improper" vecCons expressions (those that don't terminate in vecNil) since these aren't vec literals (they don't use the ![...] notation), but in any case it's not hard to add support, at the cost of putting the length of the array into WHNF at default transparency.

import Mathlib.Data.Fin.VecNotation

open Lean

partial def matchVecLitPrefix (e : Expr) : List Expr :=
  match_expr e with
  | Matrix.vecCons _ _ x xs => x :: matchVecLitPrefix xs
  | _ => []

dsimproc Matrix.cons_val (Matrix.vecCons _ _ _) := fun e => do
  let_expr Matrix.vecCons _ _ x xs' n := e | return .continue
  let some n' := n.int? | return .continue -- The docstring claims only nat or int, but works fine for Fin
  let_expr Fin length := (← instantiateMVars (← Meta.inferType n)) | return .continue
  let Expr.lit (.natVal length) ← Meta.whnfD length | return .continue
  let xs := x :: matchVecLitPrefix xs'
  let n' := (n' % length).toNat
  guard <| n' < xs.length
  return .continue (xs.get! n')

lemma foo {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] 33 = b := by dsimp
lemma bar {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] (-1) = d := by simp
example {a b c d e : Int} : ![a, b, c, d, e] 123456789 = e := by simp

[Eric Wieser]

I was under the impression that simp saw through abbrev. Is that the case, and is that also true of ler_expr?

[Eric Wieser]

I just checked; it sounds like my concern is real, but can be fixed by inserting ← Meta.whnfRs before the match_expr / let_expr

[Eric Wieser]

Here's a revised version; I think I agree that using Qq is not really appropriate here:

import Mathlib.Data.Fin.VecNotation

open Lean Qq

/-- Extracts a chain of vecCons calls -/
partial def matchVecConsPrefix (e : Expr) : MetaM <| List Expr × Expr := do
  match_expr ← Meta.whnfR e with
  | Matrix.vecCons _ _ x xs => do
    let (elems, tail) ← matchVecConsPrefix xs
    return (x :: elems, tail)
  | _ => return ([], e)

dsimproc Matrix.cons_val (Matrix.vecCons _ _ _) := fun e => do
  let_expr Matrix.vecCons _ _ x xs' n := ← Meta.whnfR e | return .continue
  let some n' := n.int? | return .continue -- The docstring claims only nat or int, but works fine for Fin
  let_expr Fin length := (← instantiateMVars (← Meta.inferType n)) | return .continue
  let Expr.lit (.natVal length) ← Meta.whnfD length | return .continue
  let (xs, tail) ← matchVecConsPrefix xs'
  let xs := x :: xs
  let n' := (n' % length).toNat
  if n' < xs.length then
    return .continue (xs.get! n')
  else if 0 < xs.length then
    let newn := unsafe toExpr (⟨n' - xs.length, lcProof⟩ : Fin (length - xs.length))
    return .continue (.some <| .app tail <| newn)
  else
    return .continue

abbrev id' (x : α) := x

open Matrix

attribute [-simp] Matrix.cons_val_four

example {a b c d : ℕ} : id' ![a, b, c, d, a, b, c, d, a, b, c, d] 8 = a := by dsimp
example {a b c : ℕ} {f : Fin 3 → ℕ} : vecCons a (id' <| vecCons b (vecCons c f)) 7 = b := by dsimp
example {a b c : ℕ} {f : Fin 3 → ℕ} : vecCons a (vecCons b (vecCons c f)) 4 = f 1 := by dsimp
example {a b c d : ℕ} : ![a, b, c, d, a, b, c, d, a, b, c, d] (-1) = d := by dsimp
example {a b c d e : Int} : ![a, b, c, d, e] 123456789 = e := by simp

[Eric Wieser]

Separately, I think having ![a, b, c, ..f] as notation for that third case might be reasonable; do you think core would accept anything similar for List?

[Yaël Dillies]

... and ![f.., a, b c] for the other direction?

[Yaël Dillies]

Note that #11142 is trying to add notation for vecCons

[Kyle Miller]

@Eric Wieser ![a, b, c, ..f] seems reasonable to me.

For List, an immediate question would be whether it would pretty print, and whether .. is accepted elsewhere. I think there's a path forward here, but I think we'll have to wait for some time before it will be acceptable in core.

[Kyle Miller]

Why does this line need to be unsafe?

let newn := unsafe toExpr (⟨n' - xs.length, lcProof⟩ : Fin (length - xs.length))

Is there another way to construct that?

[Eric Wieser]

It's unsafe because of the use of lcProof, and because there's not much point providing the proof when it's thrown away by toExpr anyway

[Eric Wieser]

Of course we could inline the implement of Fin.toExpr there and the unsafe would go away

[Eric Wieser]

I supposed I could also just use elabTerm to build the fin literal, which I now do on the other branch anyway

[Eric Wieser]

I've updated #13578 to avoid both lcProof and elabTerm