notation `x[i]` inferring the type of `i`

Stream: lean4

[Tomas Skrivan]

The current element access notation currently does not allow you to write fun i => x[i] as the type of i does not get inferred. This is particularly unfortunate for types like Vector α n. Here is my attempt at fixing it. I would like know if there is a cleaner solution to this problem.

I'm using James' LeanColls library which defines Indexed

class Indexed (Cont : Type u) (Idx : outParam (Type v)) (Elem : outParam (Type w)) where
   ...

which determined the canonical index type Idx for the container type Cont. This is unlike GetElem class where the index type is not marked as outParam. I want to uses this class to determine the index type.

Here is my solution:

code

import Lean
import Mathlib.Logic.Function.Basic

structure ArrayN (α : Type) (n : Nat) where
  data : Array α
  h_size : n = data.size

class Indexed (Cont : Type) (Idx Elem : outParam Type) where
  get : Cont → Idx → Elem
  ofFn : (Idx → Elem) → Cont

instance : Indexed (ArrayN α n) (Fin n) α where
  get c i := c.data.get (c.h_size ▸ i)
  ofFn f := ⟨Array.ofFn f, by simp⟩

instance [Indexed Cont Idx Elem] : GetElem Cont Idx Elem (fun _ _ => True) where
  getElem c i _ := Indexed.get c i



open Lean Elab Term Meta in
@[inherit_doc getElem]
elab:max (priority:=high) x:term noWs "[" i:term "]" : term => do
  try
    let x ← elabTerm x none
    let X ← inferType x
    let Idx ← mkFreshTypeMVar
    let Elem ← mkFreshTypeMVar
    let cls := (mkAppN (← mkConstWithFreshMVarLevels ``Indexed) #[X, Idx, Elem])
    let _ ← synthInstance cls
    let Dom ← mkFreshExprMVar none
    let cls := (mkAppN (← mkConstWithFreshMVarLevels ``GetElem) #[X, (← instantiateMVars Idx), Elem, Dom])
    let instGetElem ← synthInstance cls
    let i ← elabTerm i Idx
    let dom ← mkFreshExprMVar (mkAppN Dom #[x,i])
    dom.mvarId!.assign (.const ``True.intro [])
    return ← mkAppOptM ``getElem #[X,Idx,Elem,Dom,instGetElem,x,i,dom]
  catch _ =>
    return ← elabTerm (← `(getElem $x $i (by get_elem_tactic))) none

And few test cases that I found out to be challenging when trying different approach

variable {Cont Idx Elem} [Indexed Cont Idx Elem] [Inhabited Elem]
variable {Cont'} [Indexed Cont' (Fin 10) Elem]

variable (c : Cont) (c' : Cont') (j : Idx) (b : Array Nat) (k : Fin b.size)

#check c[j]
#check c'[0]
#check fun i => c[i]
#check fun i j => (c[i],c[j])
#check ↿fun i j => (c[i],c[j])
#check b[k]

example (f : Idx → Elem) :
  Function.invFun (fun (f : Idx → Elem) => Indexed.ofFn (Cont:=Cont) f)
  =
  fun x i => x[i] := sorry

[Mac Malone]

An outParam is not used on the index type in core because some container types like Array support multiple different index types (e.g., Array supports Nat, Fin, and USize).

[Tomas Skrivan]

Yes it is completely understandable that the index type is not outParam in general. But for certain types, like Vector, it would be nice to have a canonical index type.