infinite_iterator.lean

structure InfRange where
  (start : Nat := 0)
  (step : Nat := 1)

@[inline] partial def InfRange.forIn {β : Type u} {m : Type u → Type v} [Monad m] (range : InfiniteRange) (init : β) (f : Nat → β → m (ForInStep β)) : m β :=
  let rec @[specialize] loop (j : Nat) (b : β) : m β := do
    match ← f j b with
      | ForInStep.done b  => pure b
      | ForInStep.yield b => loop (j + range.step) b
  loop range.start init

def main : IO Unit := do
  -- loop forever
  for i in InfRange.mk 0 1 do
    if i >= 100 then break
    println! "{i}"