streamstream.v

Require Import List.
Import ListNotations.
Open Scope list_scope.

CoInductive stream (a : Type) :=
| scons : a -> stream a -> stream a.

Arguments scons {a}.

Definition repeat {a} (x : a) : stream a :=
  cofix self := scons x self.

Definition append_stream {a} (xx : list a) (s : stream a) :=
  (cofix go (xx : list a) :=
     match xx with
     | [] => s
     | x :: xx' => scons x (go xx')
     end) xx.  

CoFixpoint zip_with {a b c} (f : a -> b -> c)
           (sa : stream a) (sb : stream b) : stream c :=
  match sa, sb with
  | scons x sa', scons y sb' => scons (f x y) (zip_with f sa' sb')
  end.

CoFixpoint flatten' {b}
           (acc : stream (list b))
           (s : stream (stream b)) :=
  match acc, s with
  | scons yy acc', scons (scons x xx) s' =>
    scons x (append_stream yy (flatten' (zip_with cons xx acc') s'))
  end.

Definition flatten {b} (s : stream (stream b)) : stream b :=
  flatten' (repeat []) s.

Fixpoint take {a} (n : nat) (s : stream a) : list a :=
  match n, s with
  | O, _ => []
  | S n', scons x s' => x :: take n' s'
  end.

(* Example *)

Definition grid : stream (stream (nat * nat)) :=
  (cofix grid_x (n : nat) :=
     scons
       ((cofix grid_y (m : nat) :=
           scons (n, m) (grid_y (S m))
        ) 0)
       (grid_x (S n))
  ) 0.

Definition example : stream (nat * nat) := flatten grid.

Compute (take 20 example).
Compute (List.map (fun '(x, y) => x + y) (take 20 example)).