For loop over Fin range

Stream: lean4

[James Sully]

Is it possible to loop over a range of Fin? I have:

import Std
open Std

def twoSum (nums : Array Int) (target : Int) : Option (Fin nums.size × Fin nums.size) := do
  let mut seen : HashMap Int (Fin nums.size) := HashMap.empty

  for x in nums, index_ in [0:nums.size] do
    let index : Fin nums.size := ⟨index_, sorry⟩
    let complement := target - x
    match seen.find? complement with
    | some complementIndex => return (complementIndex, index)
    | none => seen := seen.insert x index

  none

#eval twoSum #[2, 7, 11, 15] 9 == some (⟨0, by simp⟩, ⟨1, by simp⟩)
#eval twoSum #[3, 2, 4] 6 == some (⟨1, by simp⟩, ⟨2, by simp⟩)
#eval twoSum #[3, 3] 6 == some (⟨0, by simp⟩, ⟨1, by simp⟩)

I'd like index_ to somehow be a Fin nums.size instead of a Nat, so that I don't have to fill in the sorry.

[Thomas Murrills]

(It's not what you've asked for on the nose, but you can use the syntax for ... h : index_ in [0:nums.size] do to include the hypothesis that index_ is in [0:nums.size] in the local context, which you can use to prove that sorry! You might be aware of this, but I wanted to make sure it was mentioned just in case.)

[James Sully]

I actually tried that, with

  for x in nums, h : index_ in [0:nums.size] do

I'm getting

▶ 7:18-7:21: error:
the proof annotation here has not been implemented yet

In the infoview. Presumably the parallel looping syntax is incomplete

[James Sully]

I'm sure there's some workaround. But yeah, ofc I would prefer not to have to prove it

[James Sully]

Ok, having searched some more, it seems like @David Thrane Christiansen has implemented exactly what I'm wanting here
https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/.E2.9C.94.20Iterating.20over.20array.20indices/near/421224993

Would be cool to see this in the standard library

[James Sully]

Unfortunately when I copy it in, I'm getting

▶ 15:1-15:43: error:
Too many extra parameters bound; the function definition only has 0 extra parameters.

On the line

termination_by _ ff action => n - ff.start

[Thomas Murrills]

Ah, gotcha, oops! I think the conventional way to do this at the moment would be for h : index in [:nums.size] do, where we simply access nums via nums[index] and construct the Fin element in the loop. Being able to iterate through Fin n with for doesn't seem like a bad idea, though—I wonder if it exists in mathlib/std anywhere yet, in some form, at least.

[James Sully]

Yeah, that's what I was doing before. In this particular case it would be slightly nicer to be iterating Fin because I'm returning the indices as Fins, so I have to get the proof of index < nums.size somehow anyway.

[James Sully]

James Sully said:

Unfortunately when I copy it in, I'm getting

I don't really understand termination_by, but after some programming-by-permutation, changing it to termination_by n - ff.start seems to have made the typechecker happy

[James Sully]

the next thing is that for parallel for loops, apparently I need to implement ToStream for the thing I'm looping over

[James Sully]

Actually maybe I can get away without it if I commute them

[James Sully]

yep, that worked. Full code, using David's FinsFrom

import Std
open Std

structure FinsFrom (n : Nat) where
  start : Nat
  ok : start ≤ n

def fins (n : Nat) : FinsFrom n := ⟨0, by simp⟩

def FinsFrom.forM [Monad m] (ff : FinsFrom n) (action : Fin n → m Unit) : m Unit := do
  if h : ff.start ≥ n then pure ()
  else
    action ⟨ff.start, by omega⟩
    forM ⟨ff.start + 1, by omega⟩ action
termination_by n - ff.start
decreasing_by simp_wf; omega

instance : ForM m (FinsFrom n) (Fin n) where
  forM := FinsFrom.forM

instance : ForIn m (FinsFrom n) (Fin n) where forIn := ForM.forIn

def twoSum (nums : Array Int) (target : Int) : Option (Fin nums.size × Fin nums.size) := do
  let mut seen : HashMap Int (Fin nums.size) := HashMap.empty

  for index in fins nums.size, x in nums do
    let complement := target - x
    match seen.find? complement with
    | some complementIndex => return (complementIndex, index)
    | none => seen := seen.insert x index

  none

#eval twoSum #[2, 7, 11, 15] 9 == some (⟨0, by simp⟩, ⟨1, by simp⟩)
#eval twoSum #[3, 2, 4] 6 == some (⟨1, by simp⟩, ⟨2, by simp⟩)
#eval twoSum #[3, 3] 6 == some (⟨0, by simp⟩, ⟨1, by simp⟩)

[James Sully]

Proposal to overload [i:j] syntax via typeclass is relevant
https://github.com/leanprover/lean4/issues/1962

[Thomas Murrills]

Honestly it would be nice to be able to simply write for x in Fin n do, but the current nature of ForIn makes using types A in that position awkward, since you would need an instance of ForIn m (Type u) A.

I wonder what the best way to solve this is: you could either have some "formalizing" map inductive F (α : Type v) where | formal (plus nice syntax) to force it to fit the current shape of the instance (this is kind of a hack), or you could overload the for ... in syntax to also try to use an instance of a new class ForInType m A, for example. (Or something else?)

[James Gallicchio]

LeanColls lets you write for x in IndexType.univ (Fin n) do for example, using a trivial wrapper structure

[James Gallicchio]

you might be able to write the similar Finset.univ (Fin n) if you have Mathlib, but that is not made for computation

[Thomas Murrills]

Nice! I wonder if IndexType.univ could somehow be a coercion from Types to IndexTypes, and if that could be picked up by for. (Or if that would mess with other things.)

[James Gallicchio]

Probably doable. I'd happily accept a PR if it works :-)

[James Sully]

Sorry to necro, posting here for my own reference
Here's a version of david's FinsFrom with Stream and ToStream instances, which gets you the parallel looping syntax. The ForIn instance comes via instForInOfStream.

structure FinRange (n : Nat) where
  curr : Nat
  ok : curr ≤ n

def fins
  (start : Nat)
  (to : Nat)
  (ok : start ≤ to := by omega)
  : FinRange to :=
  { curr := start, ok }

def allFins (n : Nat) : FinRange n := fins 0 n

instance : Stream (FinRange n) (Fin n) where
  next? (fs : FinRange n) : Option (Fin n × FinRange n) :=
    if h : fs.curr < n then
      let i : Fin n := ⟨fs.curr, h⟩
      some (i, { curr := fs.curr + 1, ok := Nat.succ_le_of_lt h })
    else
      none

instance : ToStream (FinRange n) (FinRange n) where
  toStream := id

def main : IO Unit :=
  for i in allFins 10, j in fins 10 20 do
    IO.println s!"{i.val} {j.val}"

#eval main
#synth ForIn Id (FinRange 10) (Fin 10)