feat: add well founded streams

Batteries/Data/Stream.lean

@@ -1,93 +1,3 @@
-/-
-Copyright (c) 2024 François G. Dorais. All rights reserved.
-Released under Apache 2. license as described in the file LICENSE.
-Authors: François G. Dorais
--/
-
-namespace Stream
-
-/-- Drop up to `n` values from the stream `s`. -/
-def drop [Stream σ α] (s : σ) : Nat → σ
-  | 0 => s
-  | n+1 =>
-    match next? s with
-    | none => s
-    | some (_, s) => drop s n
-
-/-- Read up to `n` values from the stream `s` as a list from first to last. -/
-def take [Stream σ α] (s : σ) : Nat → List α × σ
-  | 0 => ([], s)
-  | n+1 =>
-    match next? s with
-    | none => ([], s)
-    | some (a, s) =>
-      match take s n with
-      | (as, s) => (a :: as, s)
-
-@[simp] theorem fst_take_zero [Stream σ α] (s : σ) :
-    (take s 0).fst = [] := rfl
-
-theorem fst_take_succ [Stream σ α] (s : σ) :
-    (take s (n+1)).fst = match next? s with
-      | none => []
-      | some (a, s) => a :: (take s n).fst := by
-  simp only [take]; split <;> rfl
-
-@[simp] theorem snd_take_eq_drop [Stream σ α] (s : σ) (n : Nat) :
-    (take s n).snd = drop s n := by
-  induction n generalizing s with
-  | zero => rfl
-  | succ n ih =>
-    simp only [take, drop]
-    split <;> simp [ih]
-
-/-- Tail recursive version of `Stream.take`. -/
-def takeTR [Stream σ α] (s : σ) (n : Nat) : List α × σ :=
-  loop s [] n
-where
-  /-- Inner loop for `Stream.takeTR`. -/
-  loop (s : σ) (acc : List α)
-    | 0 => (acc.reverse, s)
-    | n+1 => match next? s with
-      | none => (acc.reverse, s)
-      | some (a, s) => loop s (a :: acc) n
-
-theorem fst_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
-    (takeTR.loop s acc n).fst = acc.reverseAux (take s n).fst := by
-  induction n generalizing acc s with
-  | zero => rfl
-  | succ n ih => simp only [take, takeTR.loop]; split; rfl; simp [ih]
-
-theorem fst_takeTR [Stream σ α] (s : σ) (n : Nat) : (takeTR s n).fst = (take s n).fst :=
-  fst_takeTR_loop ..
-
-theorem snd_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
-    (takeTR.loop s acc n).snd = drop s n := by
-  induction n generalizing acc s with
-  | zero => rfl
-  | succ n ih => simp only [takeTR.loop, drop]; split; rfl; simp [ih]
-
-theorem snd_takeTR [Stream σ α] (s : σ) (n : Nat) :
-    (takeTR s n).snd = drop s n := snd_takeTR_loop ..
-
-@[csimp] theorem take_eq_takeTR : @take = @takeTR := by
-  funext; ext : 1; rw [fst_takeTR]; rw [snd_takeTR, snd_take_eq_drop]
-
-end Stream
-
-@[simp] theorem List.stream_drop_eq_drop (l : List α) : Stream.drop l n = l.drop n := by
-  induction n generalizing l with
-  | zero => rfl
-  | succ n ih =>
-    match l with
-    | [] => rfl
-    | _::_ => simp [Stream.drop, List.drop_succ_cons, ih]
-
-@[simp] theorem List.stream_take_eq_take_drop (l : List α) :
-    Stream.take l n = (l.take n, l.drop n) := by
-  induction n generalizing l with
-  | zero => rfl
-  | succ n ih =>
-    match l with
-    | [] => rfl
-    | _::_ => simp [Stream.take, ih]
+import Batteries.Data.Stream.Basic
+import Batteries.Data.Stream.Fold
+import Batteries.Data.Stream.WF

Batteries/Data/Stream/Basic.lean

@@ -0,0 +1,93 @@
+/-
+Copyright (c) 2024 François G. Dorais. All rights reserved.
+Released under Apache 2. license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+
+namespace Stream
+
+/-- Drop up to `n` values from the stream `s`. -/
+def drop [Stream σ α] (s : σ) : Nat → σ
+  | 0 => s
+  | n+1 =>
+    match next? s with
+    | none => s
+    | some (_, s) => drop s n
+
+/-- Read up to `n` values from the stream `s` as a list from first to last. -/
+def take [Stream σ α] (s : σ) : Nat → List α × σ
+  | 0 => ([], s)
+  | n+1 =>
+    match next? s with
+    | none => ([], s)
+    | some (a, s) =>
+      match take s n with
+      | (as, s) => (a :: as, s)
+
+@[simp] theorem fst_take_zero [Stream σ α] (s : σ) :
+    (take s 0).fst = [] := rfl
+
+theorem fst_take_succ [Stream σ α] (s : σ) :
+    (take s (n+1)).fst = match next? s with
+      | none => []
+      | some (a, s) => a :: (take s n).fst := by
+  simp only [take]; split <;> rfl
+
+@[simp] theorem snd_take_eq_drop [Stream σ α] (s : σ) (n : Nat) :
+    (take s n).snd = drop s n := by
+  induction n generalizing s with
+  | zero => rfl
+  | succ n ih =>
+    simp only [take, drop]
+    split <;> simp [ih]
+
+/-- Tail recursive version of `Stream.take`. -/
+def takeTR [Stream σ α] (s : σ) (n : Nat) : List α × σ :=
+  loop s [] n
+where
+  /-- Inner loop for `Stream.takeTR`. -/
+  loop (s : σ) (acc : List α)
+    | 0 => (acc.reverse, s)
+    | n+1 => match next? s with
+      | none => (acc.reverse, s)
+      | some (a, s) => loop s (a :: acc) n
+
+theorem fst_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
+    (takeTR.loop s acc n).fst = acc.reverseAux (take s n).fst := by
+  induction n generalizing acc s with
+  | zero => rfl
+  | succ n ih => simp only [take, takeTR.loop]; split; rfl; simp [ih]
+
+theorem fst_takeTR [Stream σ α] (s : σ) (n : Nat) : (takeTR s n).fst = (take s n).fst :=
+  fst_takeTR_loop ..
+
+theorem snd_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
+    (takeTR.loop s acc n).snd = drop s n := by
+  induction n generalizing acc s with
+  | zero => rfl
+  | succ n ih => simp only [takeTR.loop, drop]; split; rfl; simp [ih]
+
+theorem snd_takeTR [Stream σ α] (s : σ) (n : Nat) :
+    (takeTR s n).snd = drop s n := snd_takeTR_loop ..
+
+@[csimp] theorem take_eq_takeTR : @take = @takeTR := by
+  funext; ext : 1; rw [fst_takeTR]; rw [snd_takeTR, snd_take_eq_drop]
+
+end Stream
+
+@[simp] theorem List.stream_drop_eq_drop (l : List α) : Stream.drop l n = l.drop n := by
+  induction n generalizing l with
+  | zero => rfl
+  | succ n ih =>
+    match l with
+    | [] => rfl
+    | _::_ => simp [Stream.drop, List.drop_succ_cons, ih]
+
+@[simp] theorem List.stream_take_eq_take_drop (l : List α) :
+    Stream.take l n = (l.take n, l.drop n) := by
+  induction n generalizing l with
+  | zero => rfl
+  | succ n ih =>
+    match l with
+    | [] => rfl
+    | _::_ => simp [Stream.take, ih]

Batteries/Data/Stream/Fold.lean

@@ -0,0 +1,112 @@
+/-
+Copyright (c) 2024 François G. Dorais. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: François G. Dorais
+-/
+import Batteries.Data.Stream.WF
+
+namespace Stream
+
+/-! ### foldlM -/
+
+/-- Folds a monadic function over a well founded stream from left to right. (Tail recursive.) -/
+@[specialize] def foldlM [Monad m] [Stream.WF σ α] (f : β → α → m β) (init : β) (s : σ) : m β :=
+  match _hint : next? s with
+  | none => pure init
+  | some (x, t) => f init x >>= (foldlM f · t)
+termination_by s
+
+theorem foldlM_init [Monad m] [Stream.WF σ α] {f : β → α → m β} {init : β} {s : σ}
+    (h : next? s = none) : foldlM f init s = pure init := by rw [foldlM, h]
+
+theorem foldlM_next [Monad m] [Stream.WF σ α] {f : β → α → m β} {init : β} {s t : σ} {x}
+    (h : next? s = some (x, t)) : foldlM f init s = f init x >>= (foldlM f · t) := by rw [foldlM, h]
+
+theorem foldlM_eq_foldlM_toList [Monad m] [Stream.WF σ α] (f : β → α → m β) (init : β) (s : σ) :
+    foldlM f init s = (toList s).foldlM f init := by
+  induction s using Stream.recWF generalizing init with
+  | init h => simp only [foldlM_init h, toList_init h, List.foldlM_nil]
+  | next h ih => simp only [foldlM_next h, toList_next h, List.foldlM_cons, ih]
+
+/-! ### foldl -/
+
+/-- Folds a function over a well founded stream from left to right. (Tail recursive.) -/
+@[inline] def foldl [Stream.WF σ α] (f : β → α → β) (init : β) (s : σ) : β :=
+  foldlM (m := Id) f init s
+
+theorem foldl_init [Stream.WF σ α] {f : β → α → β} {init : β} {s : σ}
+    (h : next? s = none) : foldl f init s = init := foldlM_init h
+
+theorem foldl_next [Stream.WF σ α] {f : β → α → β} {init : β} {s t : σ}
+    (h : next? s = some (x, t)) : foldl f init s = foldl f (f init x) t := foldlM_next h
+
+theorem foldl_eq_foldl_toList [Stream.WF σ α] (f : β → α → β) (init : β) (s : σ) :
+    foldl f init s = (toList s).foldl f init := by
+  induction s using Stream.recWF generalizing init with
+  | init h => simp only [foldl_init h, toList_init h, List.foldl_nil]
+  | next h ih => simp only [foldl_next h, toList_next h, List.foldl_cons, ih]
+
+/-! ### foldrM -/
+
+/-- Folds a monadic function over a well founded stream from left to right. -/
+@[specialize] def foldrM [Monad m] [Stream.WF σ α] (f : α → β → m β) (init : β) (s : σ) : m β :=
+  match _hint : next? s with
+  | none => pure init
+  | some (x, t) => foldrM f init t >>= f x
+termination_by s
+
+theorem foldrM_init [Monad m] [Stream.WF σ α] {f : α → β → m β} {init : β} {s : σ}
+    (h : next? s = none) : foldrM f init s = pure init := by rw [foldrM, h]
+
+theorem foldrM_next [Monad m] [Stream.WF σ α] {f : α → β → m β} {init : β} {s t : σ}
+    (h : next? s = some (x, t)) : foldrM f init s = foldrM f init t >>= f x := by rw [foldrM, h]
+
+theorem foldrM_eq_foldrM_toList [Monad m] [LawfulMonad m] [Stream.WF σ α]
+    (f : α → β → m β) (init : β) (s : σ) : foldrM f init s = (toList s).foldrM f init := by
+  induction s using Stream.recWF with
+  | init h => simp only [foldrM_init h, toList_init h, List.foldrM_nil]
+  | next h ih => simp only [foldrM_next h, toList_next h, List.foldrM_cons, ih]
+
+/-! ### foldr -/
+
+/-- Folds a function over a well founded stream from left to right. -/
+@[inline] def foldr [Stream.WF σ α] (f : α → β → β) (init : β) (s : σ) : β :=
+  foldrM (m := Id) f init s
+
+theorem foldr_init [Stream.WF σ α] {f : α → β → β} {init : β} {s : σ}
+    (h : next? s = none) : foldr f init s = init := foldrM_init h
+
+theorem foldr_next [Stream.WF σ α] {f : α → β → β} {init : β} {s t : σ}
+    (h : next? s = some (x, t)) : foldr f init s = f x (foldr f init t) := foldrM_next h
+
+theorem foldr_eq_foldr_toList [Stream.WF σ α] (f : α → β → β) (init : β) (s : σ) :
+    foldr f init s = (toList s).foldr f init := by
+  induction s using Stream.recWF with
+  | init h => rw [foldr_init h, toList_init h, List.foldr_nil]
+  | next h ih => rw [foldr_next h, toList_next h, List.foldr_cons, ih]
+
+end Stream
+
+theorem List.stream_foldlM_eq_foldlM [Monad m] (f : β → α → m β) (init : β) (l : List α) :
+    Stream.foldlM f init l = l.foldlM f init := by
+  induction l generalizing init with
+  | nil => rw [Stream.foldlM_init, foldlM_nil]; rfl
+  | cons x l ih => rw [Stream.foldlM_next, foldlM_cons]; congr; funext; apply ih; rfl
+
+theorem List.stream_foldl_eq_foldl (f : β → α → β) (init : β) (l : List α) :
+    Stream.foldl f init l = l.foldl f init := by
+  induction l generalizing init with
+  | nil => rw [Stream.foldl_init, foldl_nil]; rfl
+  | cons x l ih => rw [Stream.foldl_next, foldl_cons]; congr; funext; apply ih; rfl
+
+theorem List.stream_foldrM_eq_foldrM [Monad m] [LawfulMonad m] (f : α → β → m β) (init : β)
+    (l : List α) : Stream.foldrM f init l = l.foldrM f init := by
+  induction l generalizing init with
+  | nil => rw [Stream.foldrM_init, foldrM_nil]; rfl
+  | cons x l ih => rw [Stream.foldrM_next, foldrM_cons]; congr; funext; apply ih; rfl
+
+theorem List.stream_foldr_eq_foldr (f : α → β → β) (init : β) (l : List α) :
+    Stream.foldr f init l = l.foldr f init := by
+  induction l generalizing init with
+  | nil => rw [Stream.foldr_init, foldr_nil]; rfl
+  | cons x l ih => rw [Stream.foldr_next, foldr_cons]; congr; funext; apply ih; rfl