forStream.lean

import Std.Data.RBMap

-- Summary: Currently `forIn` provides a way to specify how
-- to iterate over values of a particular type in a `for`
-- loop in a way that conveniently guarantees termination
-- and allows arbitrary computation within the loop.

-- This file presents a slightly more general version called
-- `forStream` which aims to allow additional values to be
-- iterated over in-parallel with the first sequence. I.e.,
-- it is essentially identical to `forIn` except it accepts
-- an additional `stream` parameter which represents the
-- stream of other value(s) being iterated over in the for
-- loop (if any exist). The `stream` parameter will be of
-- type `σ` with constraint `Stream σ σ'` which essentially
-- says `σ` can be treated as a stream of values of type
-- `σ'`.

-- So the idea is instead of 
--
-- ```
-- for x in xs do 
--   loopBody
--```
--
-- expanding to _roughly_
-- ```
-- xs.forIn init (λ x => loopBody)
-- ```
--
-- we would have
-- ```
-- for x in xs[, y in ys]* do 
--   loopBody
--```
--
-- expanding to _roughly_
-- ```
-- xs.forStream init ([ys,]*) (λ x ([y,]*) => loopBody)
-- ```
--
--
-- (The details would be slightly more complicated. E.g.,
--  the `loopBody` function is actually passed a single
--  "stream state" representing the position in _all_ of the
--  streams (ys, ...), and before the code in `loopBody` is
--  actually executed, the stream state is probed to extract
--  the next stream value(s) and subsequent stream state (if
--  any exists). The "manually expanded examples at the
--  bottom might make this more clear. I tried to write in a
--  comment the rough "for-loop" syntax along with a _rough_
--  manual expansion of the code.)
--
--
-- This file defines the `Stream` class for `Range`, `List`,
-- and `RBMap`, and at the bottom are several hand-written
-- examples attempting to show roughly what the `for` macro
-- should expand into.

universes u₁ u₂ u₃ u₄


open Std


-- A (possibly infinite) stream of values.
-- `collection` is the type which we are "streaming"
-- `value` is the type streamed from the `collection`
class Stream (collection : Type u₁) (value : outParam (Type u₂)) : Type (max (u₁+1) u₂) where
  pos   : Type u₁
  -- ^ Position in the stream.
  init  : collection → pos
  -- ^ How to create a stream from a `collection`.
  next? : pos →  Option (value × pos)
  -- ^ How to access the next element of a stream (if one exists).


-- Unit is treated as an infinite stream of units
instance : Stream Unit Unit where
  pos  := Unit
  init  := λ _ => ()
  next? := λ _ => some ((),())



namespace Stream

variables {α : Type u₁} {α' : Type u₂} {β : Type u₃} {β' : Type u₄}

-- `prodStream` allows us to combine arbitrarily many
-- streams into a single stream which accesses the original
-- streams 'in-parallel'.
instance prodStream [Stream α α'] [Stream β β'] : Stream (α × β) (α' × β') where
  pos := (Stream.pos α × Stream.pos β)
  init  := λ (a, b) => (Stream.init a, Stream.init b)
  next? := λ (pa, pb) => 
             match Stream.next? pa, Stream.next? pb with
             | some (a, pa), some (b, pb) => some ((a, b), (pa, pb))
             | _,_ => none


end Stream

inductive ForStreamStep (α : Type u₁) (β : Type u₂) [Stream α β] (δ : Type u₃) where
  | done  : δ → ForStreamStep α β δ
  -- ^ One of the streams has terminated or the user broke out of the loop.
  | yield : Stream.pos α → δ → ForStreamStep α β δ
  -- ^ The updated stream position and the yielded value from the loop body.


namespace Std
namespace Range

variables {σ σ' δ m}

instance : Stream Range Nat where
  pos := Nat × Range
  init  := λ r => (r.start, r)
  next? := λ (n, r) => let next := n + r.step
                  if next < r.stop
                  then some (next, (next, r))
                  else none


@[inline] def forStream [Monad m] [Stream σ σ'] (range : Range) (stream : σ) (init : δ) (f : Nat → Stream.pos σ → δ → m (ForStreamStep σ σ' δ)) : m δ :=
  let rec @[specialize] loop (i : Nat) (j : Nat) (p : Stream.pos σ) (d : δ) : m δ := do
    if j ≥ range.stop then
       pure d
    else match i with
     | 0   => pure d
     | i+1 => match ← f j p d with
        | ForStreamStep.done b  => pure b
        | ForStreamStep.yield p b => loop i (j + range.step) p b
  loop range.stop range.start (Stream.init stream) init

end Range
end Std


namespace List

variables {α σ σ' δ m}

instance : Stream (List α) α where
  pos   := List α
  init  := λ l => l
  next? := λ l => match l with
                  | [] => none
                  | a::as => some (a, as)


@[inline] def forStream [Monad m] [Stream σ σ'] (as : List α) (stream : σ) (init : δ) (f : α → Stream.pos σ → δ → m (ForStreamStep σ σ' δ)) : m δ :=
  let rec @[specialize] loop
    | [], p, d    => pure d
    | a::as, p, d => do
      match (← f a p d) with
      | ForStreamStep.done d  => pure d
      | ForStreamStep.yield p d => loop as p d
  loop as (Stream.init stream) init

end List




namespace Std
namespace RBNode

inductive StreamPos (α : Type u₁) (β : α → Type u₂) where
  | done : StreamPos α β
  | here : (k : α) → β k → RBNode α β → StreamPos α β → StreamPos α β



variables {α : Type u₁} {β : α → Type u₂} {σ : Type u₃} {σ' : Type u₄} {δ : Type u₃}

private def toStreamPos : RBNode α β → StreamPos α β → StreamPos α β
  | leaf, prev => prev
  | node _ l k v r, prev => toStreamPos l (StreamPos.here k v r prev)


@[inline] def forStream [Monad m] [Stream σ σ'] (rb : RBNode α β) (stream : σ) (init : δ) (f : (k : α) → β k → Stream.pos σ → δ → m (ForStreamStep σ σ' δ)) : m δ := do
  let rec @[specialize] visit : RBNode α β → Stream.pos σ → δ → m (ForStreamStep σ σ' δ)
    | leaf, p, d           => return ForStreamStep.yield p d
    | node _ l k v r, p, d => do
      match (← visit l p d) with
      | r@(ForStreamStep.done _) => return r
      | ForStreamStep.yield p d =>
        match (← f k v p d) with
        | r@(ForStreamStep.done _) => return r
        | ForStreamStep.yield p d    => visit r p d
  match ← visit rb (Stream.init stream) init with
  | ForStreamStep.done d    => pure d
  | ForStreamStep.yield _ d => pure d

end RBNode

namespace RBMap

variables {α β σ σ' δ m lt}

@[inline] def forStream [Monad m] [Stream σ σ'] (t : RBMap α β lt) (s : σ) (init : δ) (f : (α × β) → Stream.pos σ → δ → m (ForStreamStep σ σ' δ)) : m δ :=
  t.val.forStream s init (fun a b p acc => f (a, b) p acc)

instance : Stream (RBMap α β lt) (α × β) where
  pos   := RBNode.StreamPos α (λ _ => β)
  init  := λ ⟨t,_⟩ => RBNode.toStreamPos t RBNode.StreamPos.done
  next? := λ sp => match sp with
                   | RBNode.StreamPos.done => none
                   | RBNode.StreamPos.here k v r sp => some ((k,v), RBNode.toStreamPos r sp)


end RBMap

end Std





-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- single Range iteration test helpers 
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- 
-- Roughly what I would expect a `for` loop using
-- `forStream` to expand into when iterating over a single
-- `Range` with no additional sequences being iterated over,
-- i.e. (sans the reps part for timing): 
-- ```
-- let mut sum := 0 
-- for n in rng do
--   sum := sum + n
-- pure sum 
-- ```
def sumRangeForStream {m} [Monad m] (rng : Range) (reps : Nat) : m Nat := do
  let loopBody : Nat → Stream.pos Unit → Nat →  m (ForStreamStep Unit Unit Nat) :=
    λ n p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some (tt, p) => ForStreamStep.yield p (n + sum)
  let mut sum := 0
  for _ in [0:reps] do
    let n ← Std.Range.forStream rng () 0 loopBody
    sum := sum + n
  pure sum

def timeSumRangeForStream (low high reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing range ["++(repr low)++":"++(repr high)++"] "++(repr reps)++" times..."
  let rng := [low:high]
  let n ← timeit msg $ sumRangeForStream rng reps
  pure n


def sumRangeForIn {m} [Monad m] (rng : Range) (reps : Nat) : m Nat := do
  let loopBody : Nat → Nat →  m (ForInStep Nat) :=
    λ n acc => ForInStep.yield (n + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← Std.Range.forIn rng 0 loopBody
    sum := sum + n
  pure sum

def timeSumRangeForIn (low high reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing range ["++(repr low)++":"++(repr high)++"] "++(repr reps)++" times..."
  let rng := [low:high]
  let n ← timeit msg $ sumRangeForIn rng reps
  pure n



-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- single List iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Like the above example except instead of a `Range` we iterate over a `List Nat`.
def sumListForStream {m} [Monad m] (rng : List Nat) (reps : Nat) : m Nat := do
  let loopBody : Nat → Stream.pos Unit → Nat →  m (ForStreamStep Unit Unit Nat) :=
    λ n p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some (tt, p) => ForStreamStep.yield p (n + sum)
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forStream rng () 0 loopBody
    sum := sum + n
  pure sum

def timeSumListForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing list (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng := List.range n
  let n ← timeit msg $ sumListForStream rng reps
  pure n


def sumListForIn {m} [Monad m] (rng : List Nat) (reps : Nat) : m Nat := do
  let loopBody : Nat → Nat →  m (ForInStep Nat) :=
    λ n acc => ForInStep.yield (n + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forIn rng 0 loopBody
    sum := sum + n
  pure sum

def timeSumListForIn (n reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing list (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng := List.range n
  let n ← timeit msg $ sumListForIn rng reps
  pure n


-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- single RBMap iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Like the previous examples but an RBMap this time (sans the `reps` part for timing): 
-- ```
-- let mut sum := 0 
-- for (n, m) in rng do
--   sum := sum + n + m
-- pure sum 
-- ```

def sumRBMapForStream {m} [Monad m] (rng : Std.RBMap Nat Nat (λ x y => x < y)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat) → Stream.pos Unit → Nat →  m (ForStreamStep Unit Unit Nat) :=
    λ (n, m) p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some (_, p) => ForStreamStep.yield p (n + m + sum)
  let mut sum := 0
  for _ in [0:reps] do
    let n ← Std.RBMap.forStream rng () 0 loopBody
    sum := sum + n
  pure sum

def timeSumRBMapForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing RBMap nodes and keys"++(repr reps)++" times..."
  let rng := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sumRBMapForStream rng reps
  pure n


def sumRBMapForIn {m} [Monad m] (rng : Std.RBMap Nat Nat (λ x y => x < y)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat) → Nat →  m (ForInStep Nat) :=
    λ (n, m) acc => ForInStep.yield (n + m + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← RBMap.forIn rng 0 loopBody
    sum := sum + n
  pure sum

def timeSumRBMapForIn (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing RBMap nodes and keys"++(repr reps)++" times..."
  let rng := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sumRBMapForIn rng reps
  pure n


-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- 2 list iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Roughly the expansion of (sans the reps part for timing):
-- ```
-- let mut sum := 0
-- for x in xs, y in ys do
--   sum := sum + x + y
-- pure sum 
-- ```
def sum2ListsForStream {m} [Monad m] (xs ys : List Nat) (reps : Nat) : m Nat := do
  let loopBody : Nat → Stream.pos (List Nat) → Nat →  m (ForStreamStep (List Nat) Nat Nat) :=
    λ x p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some (y, p) => ForStreamStep.yield p (sum + x + y)
  let mut sum := 0
  for _ in [0:reps] do
    sum ← List.forStream xs ys sum loopBody
  pure sum

def timeSum2ListsForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing 2 (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng1 := List.range n
  let rng2 := List.range n
  let n ← timeit msg $ sum2ListsForStream rng1 rng2 reps
  pure n


def sum2ListsAsOneForIn {m} [Monad m] (rng : List (Nat × Nat)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat) → Nat →  m (ForInStep Nat) :=
    λ (n, m) acc => ForInStep.yield (n + m + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forIn rng 0 loopBody
    sum := sum + n
  pure sum

def timeSum2ListsAsOneForIn (n reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing 2 (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng := List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sum2ListsAsOneForIn rng reps
  pure n

-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- 3 list iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Roughly the expansion of (sans the reps part for timing):
-- ```
-- let mut sum := 0
-- for x in xs, y in ys, z in zs do
--   sum := sum + x + y + z
-- pure sum 
-- ```
def sum3ListsForStream {m} [Monad m] (xs ys zs : List Nat) (reps : Nat) : m Nat := do
  let loopBody : Nat → Stream.pos (List Nat × List Nat) → Nat →  m (ForStreamStep (List Nat × List Nat) (Nat × Nat) Nat) :=
    λ x p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some ((y, z), p) => ForStreamStep.yield p (sum + x + y + z)
  let mut sum := 0
  for _ in [0:reps] do
    sum ← List.forStream xs (ys, zs) sum loopBody
  pure sum

def timeSum3ListsForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing 3 (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng1 := List.range n
  let rng2 := List.range n
  let rng3 := List.range n
  let n ← timeit msg $ sum3ListsForStream rng1 rng2 rng3 reps
  pure n


def sum3ListsAsOneForIn {m} [Monad m] (rng : List (Nat × Nat × Nat)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat × Nat) → Nat →  m (ForInStep Nat) :=
    λ (n, m₁, m₂) acc => ForInStep.yield (n + m₁ + m₂ + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forIn rng 0 loopBody
    sum := sum + n
  pure sum

def timeSum3ListsAsOneForIn (n reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing 3 (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng := List.zip (List.range n) (List.zip (List.range n) (List.range n))
  let n ← timeit msg $ sum3ListsAsOneForIn rng reps
  pure n


-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- 4 list iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Roughly the expansion of the following (sans the reps part for timing):
-- ```
-- let mut sum := 0
-- for x in xs, y in ys, z in zs, q in qs do
--   sum := sum + x + y + z
-- pure sum 
-- ```
def sum4ListsForStream {m} [Monad m] (xs ys zs qs : List Nat) (reps : Nat) : m Nat := do
  let loopBody : Nat → Stream.pos (List Nat × List Nat × List Nat) → Nat →  m (ForStreamStep (List Nat × List Nat × List Nat) (Nat × Nat × Nat) Nat) :=
    λ x p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some ((y, z, q), p) => ForStreamStep.yield p (sum + x + y + z + q)
  let mut sum := 0
  for _ in [0:reps] do
    sum ← List.forStream xs (ys, zs, qs) sum loopBody
  pure sum

def timeSum4ListsForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing 4 (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng1 := List.range n
  let rng2 := List.range n
  let rng3 := List.range n
  let rng4 := List.range n
  let n ← timeit msg $ sum4ListsForStream rng1 rng2 rng3 rng4 reps
  pure n


def sum4ListsAsOneForIn {m} [Monad m] (rng : List (Nat × Nat × Nat × Nat)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat × Nat × Nat) → Nat →  m (ForInStep Nat) :=
    λ (n, m₁, m₂, m₃) acc => ForInStep.yield (n + m₁ + m₂ + m₃ + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forIn rng 0 loopBody
    sum := sum + n
  pure sum

def timeSum4ListsAsOneForIn (n reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing 4 (List.range "++(repr n)++") "++(repr reps)++" times..."
  let rng := List.zip (List.range n) (List.zip (List.range n) (List.zip (List.range n) (List.range n)))
  let n ← timeit msg $ sum4ListsAsOneForIn rng reps
  pure n


-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- 2 RBMap iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Roughly the expansion of (sans the reps part for timing):
-- ```
-- let mut sum := 0
-- for (x₁, x₂) in xs, (y₁, y₂) in ys do
--   sum := sum + x₁ + x₂ + y₁ + y₂
-- pure sum 
-- ```
def sum2RBMapsForStream {m} [Monad m] (xs ys : RBMap Nat Nat (λ x y => x < y)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat) → Stream.pos (RBMap Nat Nat (λ x y => x < y)) → Nat →  m (ForStreamStep (RBMap Nat Nat (λ x y => x < y)) (Nat × Nat) Nat) :=
    λ (x₁, x₂) p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some ((y₁, y₂), p) => ForStreamStep.yield p (sum + x₁ + x₂ + y₁ + y₂)
  let mut sum := 0
  for _ in [0:reps] do
    let n ← Std.RBMap.forStream xs ys 0 loopBody
    sum := sum + n
  pure sum

def timeSum2RBMapsForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing 2 (RBMap Nat Nat)s with "++(repr n)++" elements "++(repr reps)++" times..."
  let rng1 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let rng2 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sum2RBMapsForStream rng1 rng2 reps
  pure n


def sum2RBMapsForIn {m} [Monad m] (rng other : RBMap Nat Nat (λ x y => x < y)) (reps : Nat) : m Nat := do
  let loopBody : ((Nat × Nat) × (Nat × Nat)) → Nat →  m (ForInStep Nat) :=
    λ ((n₁, n₂), (m₁, m₂)) acc => ForInStep.yield (n₁ + n₂ + m₁ + m₂ + acc) 
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forIn (List.zip rng.toList other.toList) 0 loopBody
    sum := sum + n
  pure sum

def timeSum2RBMapsForIn (n reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing 2 (RBMap Nat Nat)s with "++(repr n)++" elements "++(repr reps)++" times..."
  let rng1 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let rng2 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sum2RBMapsForIn rng1 rng2 reps
  pure n


-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- 3 RBMap iteration test helpers
-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- --
-- Roughly the expansion of (sans the reps part for timing):
-- ```
-- let mut sum := 0
-- for (x₁, x₂) in xs, (y₁, y₂) in ys, (z₁, z₂) in zs do
--   sum := sum + x₁ + x₂ + y₁ + y₂ + z₁ + z₂
-- pure sum 
-- ```
def sum3RBMapsForStream {m} [Monad m] (xs ys zs : RBMap Nat Nat (λ x y => x < y)) (reps : Nat) : m Nat := do
  let loopBody : (Nat × Nat)
                 → Stream.pos (RBMap Nat Nat (λ x y => x < y) × RBMap Nat Nat (λ x y => x < y)) 
                 → Nat
                 →  m (ForStreamStep ((RBMap Nat Nat (λ x y => x < y)) × (RBMap Nat Nat (λ x y => x < y))) ((Nat × Nat) × (Nat × Nat)) Nat) :=
    λ (x₁, x₂) p sum => match Stream.next? p with
                 | none => ForStreamStep.done sum
                 | some (((y₁, y₂), (z₃, z₄)), p) => ForStreamStep.yield p (sum + x₁ + x₂ + y₁ + y₂ + z₃ + z₄)
  let mut sum := 0
  for _ in [0:reps] do
    let n ← Std.RBMap.forStream xs (ys, zs) 0 loopBody
    sum := sum + n
  pure sum

def timeSum3RBMapsForStream (n reps : Nat) : IO Nat := do
  let msg := "[forStream] Summing 3 (RBMap Nat Nat)s with "++(repr n)++" elements "++(repr reps)++" times..."
  let rng1 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let rng2 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let rng3 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sum3RBMapsForStream rng1 rng2 rng3 reps
  pure n


def sum3RBMapsForIn {m} [Monad m] (rng other1 other2 : RBMap Nat Nat (λ x y => x < y)) (reps : Nat) : m Nat := do
  let loopBody : ((Nat × Nat) × (Nat × Nat) × (Nat × Nat)) → Nat →  m (ForInStep Nat) :=
    λ ((n₁, n₂), (m₁, m₂), (m₃, m₄)) acc => ForInStep.yield (n₁ + n₂ + m₁ + m₂ + m₃ + m₄ + acc)
  let mut sum := 0
  for _ in [0:reps] do
    let n ← List.forIn (List.zip rng.toList (List.zip other1.toList other2.toList)) 0 loopBody
    sum := sum + n
  pure sum

def timeSum3RBMapsForIn (n reps : Nat) : IO Nat := do
  let msg := "[forIn    ] Summing 3 (RBMap Nat Nat)s with "++(repr n)++" elements "++(repr reps)++" times..."
  let rng1 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let rng2 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let rng3 := Std.RBMap.fromList (lt := λ x y => x < y) $ List.zip (List.range n) (List.range n)
  let n ← timeit msg $ sum3RBMapsForIn rng1 rng2 rng3 reps
  pure n


def main (xs : List String) : IO Unit := do
  -- single range iteration timing tests for comparison
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing Range forIn vs forStream iteration"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSumRangeForIn     0 100 100000
  _ ← timeSumRangeForStream 0 100 100000
  _ ← timeSumRangeForIn     0 1000 10000
  _ ← timeSumRangeForStream 0 1000 10000
  _ ← timeSumRangeForIn     0 10000 1000
  _ ← timeSumRangeForStream 0 10000 1000
  -- single list iteration timing tests for comparison
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing (List Nat) forIn vs forStream iteration"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSumListForIn     100 100000
  _ ← timeSumListForStream 100 100000
  _ ← timeSumListForIn     1000 10000
  _ ← timeSumListForStream 1000 10000
  _ ← timeSumListForIn     10000 1000
  _ ← timeSumListForStream 10000 1000
  -- single list RBMap timing tests for comparison
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing (RBMap Nat Nat) forIn vs forStream iteration"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSumListForIn     100 100000
  _ ← timeSumListForStream 100 100000
  _ ← timeSumListForIn     1000 10000
  _ ← timeSumListForStream 1000 10000
  _ ← timeSumListForIn     10000 1000
  _ ← timeSumListForStream 10000 1000
  -- double range iteration timing tests for comparison
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing 2x(List Nat) forIn vs forStream iteration (sort of, forIn zips _before timed computation_ to simulate fusion)"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSum2ListsAsOneForIn  100 100000
  _ ← timeSum2ListsForStream   100 100000
  _ ← timeSum2ListsAsOneForIn  1000 10000
  _ ← timeSum2ListsForStream   1000 10000
  _ ← timeSum2ListsAsOneForIn  10000 1000
  _ ← timeSum2ListsForStream   10000 1000
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing 3x(List Nat) forIn vs forStream iteration (sort of, forIn zips _before timed computation_ to simulate fusion)"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSum3ListsAsOneForIn  100 100000
  _ ← timeSum3ListsForStream   100 100000
  _ ← timeSum3ListsAsOneForIn  1000 10000
  _ ← timeSum3ListsForStream   1000 10000
  _ ← timeSum3ListsAsOneForIn  10000 1000
  _ ← timeSum3ListsForStream   10000 1000
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing 4x(List Nat) forIn vs forStream iteration (sort of, forIn zips _before timed computation_ to simulate fusion)"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSum4ListsAsOneForIn  100 100000
  _ ← timeSum4ListsForStream   100 100000
  _ ← timeSum4ListsAsOneForIn  1000 10000
  _ ← timeSum4ListsForStream   1000 10000
  _ ← timeSum4ListsAsOneForIn  10000 1000
  _ ← timeSum4ListsForStream   10000 1000
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing 2x(RBMap Nat Nat) forIn vs forStream iteration (sort of, forIn toList and zips them _during_ timed computation)"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSum2RBMapsForIn       100 100000
  _ ← timeSum2RBMapsForStream   100 100000
  _ ← timeSum2RBMapsForIn       1000 10000
  _ ← timeSum2RBMapsForStream   1000 10000
  _ ← timeSum2RBMapsForIn       10000 1000
  _ ← timeSum2RBMapsForStream   10000 1000
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  IO.println "Comparing 3x(RBMap Nat Nat) forIn vs forStream iteration (sort of, forIn toList and zips them _during_ timed computation)"
  IO.println "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -"
  _ ← timeSum3RBMapsForIn       100 100000
  _ ← timeSum3RBMapsForStream   100 100000
  _ ← timeSum3RBMapsForIn       1000 10000
  _ ← timeSum3RBMapsForStream   1000 10000
  _ ← timeSum3RBMapsForIn       10000 1000
  _ ← timeSum3RBMapsForStream   10000 1000

  pure ()


/-
- - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing Range forIn vs forStream iteration
- - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing range [0:100] 100000 times... 48.5ms
[forStream] Summing range [0:100] 100000 times... 37.1ms
[forIn    ] Summing range [0:1000] 10000 times... 36.8ms
[forStream] Summing range [0:1000] 10000 times... 34.5ms
[forIn    ] Summing range [0:10000] 1000 times... 37.5ms
[forStream] Summing range [0:10000] 1000 times... 37.5ms
- - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing (List Nat) forIn vs forStream iteration
- - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing list (List.range 100) 100000 times... 24.6ms
[forStream] Summing list (List.range 100) 100000 times... 25.8ms
[forIn    ] Summing list (List.range 1000) 10000 times... 22.3ms
[forStream] Summing list (List.range 1000) 10000 times... 24.4ms
[forIn    ] Summing list (List.range 10000) 1000 times... 23.7ms
[forStream] Summing list (List.range 10000) 1000 times... 23.7ms
- - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing (RBMap Nat Nat) forIn vs forStream iteration
- - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing list (List.range 100) 100000 times... 23.6ms
[forStream] Summing list (List.range 100) 100000 times... 25.5ms
[forIn    ] Summing list (List.range 1000) 10000 times... 21.4ms
[forStream] Summing list (List.range 1000) 10000 times... 23.7ms
[forIn    ] Summing list (List.range 10000) 1000 times... 21.9ms
[forStream] Summing list (List.range 10000) 1000 times... 23.7ms
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing 2x(List Nat) forIn vs forStream iteration (sort of, forIn zips _before timed computation_ to simulate fusion)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing 2 (List.range 100) 100000 times... 35.2ms
[forStream] Summing 2 (List.range 100) 100000 times... 36.9ms
[forIn    ] Summing 2 (List.range 1000) 10000 times... 32.9ms
[forStream] Summing 2 (List.range 1000) 10000 times... 35.3ms
[forIn    ] Summing 2 (List.range 10000) 1000 times... 33.2ms
[forStream] Summing 2 (List.range 10000) 1000 times... 35.6ms
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing 3x(List Nat) forIn vs forStream iteration (sort of, forIn zips _before timed computation_ to simulate fusion)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing 3 (List.range 100) 100000 times... 42.1ms
[forStream] Summing 3 (List.range 100) 100000 times... 96.7ms
[forIn    ] Summing 3 (List.range 1000) 10000 times... 40.6ms
[forStream] Summing 3 (List.range 1000) 10000 times... 96.6ms
[forIn    ] Summing 3 (List.range 10000) 1000 times... 41.7ms
[forStream] Summing 3 (List.range 10000) 1000 times... 93.6ms
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing 4x(List Nat) forIn vs forStream iteration (sort of, forIn zips _before timed computation_ to simulate fusion)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing 4 (List.range 100) 100000 times... 48.6ms
[forStream] Summing 4 (List.range 100) 100000 times... 157ms
[forIn    ] Summing 4 (List.range 1000) 10000 times... 47.8ms
[forStream] Summing 4 (List.range 1000) 10000 times... 151ms
[forIn    ] Summing 4 (List.range 10000) 1000 times... 53.5ms
[forStream] Summing 4 (List.range 10000) 1000 times... 152ms
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing 2x(RBMap Nat Nat) forIn vs forStream iteration (sort of, forIn toList and zips them _during_ timed computation)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing 2 (RBMap Nat Nat)s with 100 elements 100000 times... 742ms
[forStream] Summing 2 (RBMap Nat Nat)s with 100 elements 100000 times... 549ms
[forIn    ] Summing 2 (RBMap Nat Nat)s with 1000 elements 10000 times... 758ms
[forStream] Summing 2 (RBMap Nat Nat)s with 1000 elements 10000 times... 562ms
[forIn    ] Summing 2 (RBMap Nat Nat)s with 10000 elements 1000 times... 825ms
[forStream] Summing 2 (RBMap Nat Nat)s with 10000 elements 1000 times... 587ms
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Comparing 3x(RBMap Nat Nat) forIn vs forStream iteration (sort of, forIn toList and zips them _during_ timed computation)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
[forIn    ] Summing 3 (RBMap Nat Nat)s with 100 elements 100000 times... 1.19s
[forStream] Summing 3 (RBMap Nat Nat)s with 100 elements 100000 times... 916ms
[forIn    ] Summing 3 (RBMap Nat Nat)s with 1000 elements 10000 times... 1.26s
[forStream] Summing 3 (RBMap Nat Nat)s with 1000 elements 10000 times... 935ms
[forIn    ] Summing 3 (RBMap Nat Nat)s with 10000 elements 1000 times... 1.36s
[forStream] Summing 3 (RBMap Nat Nat)s with 10000 elements 1000 times... 948ms

-/