Skip to content

Commit a268e1c

Browse files
fgdoraisfgdorais
committed
feat: add well founded streams
1 parent f007bfe commit a268e1c

File tree

4 files changed

+424
-93
lines changed

4 files changed

+424
-93
lines changed

‎Batteries/Data/Stream.lean‎

Lines changed: 3 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -1,93 +1,3 @@
1-
/-
2-
Copyright (c) 2024 François G. Dorais. All rights reserved.
3-
Released under Apache 2. license as described in the file LICENSE.
4-
Authors: François G. Dorais
5-
-/
6-
7-
namespace Stream
8-
9-
/-- Drop up to `n` values from the stream `s`. -/
10-
def drop [Stream σ α] (s : σ) : Nat → σ
11-
| 0 => s
12-
| n+1 =>
13-
match next? s with
14-
| none => s
15-
| some (_, s) => drop s n
16-
17-
/-- Read up to `n` values from the stream `s` as a list from first to last. -/
18-
def take [Stream σ α] (s : σ) : Nat → List α × σ
19-
| 0 => ([], s)
20-
| n+1 =>
21-
match next? s with
22-
| none => ([], s)
23-
| some (a, s) =>
24-
match take s n with
25-
| (as, s) => (a :: as, s)
26-
27-
@[simp] theorem fst_take_zero [Stream σ α] (s : σ) :
28-
(take s 0).fst = [] := rfl
29-
30-
theorem fst_take_succ [Stream σ α] (s : σ) :
31-
(take s (n+1)).fst = match next? s with
32-
| none => []
33-
| some (a, s) => a :: (take s n).fst := by
34-
simp only [take]; split <;> rfl
35-
36-
@[simp] theorem snd_take_eq_drop [Stream σ α] (s : σ) (n : Nat) :
37-
(take s n).snd = drop s n := by
38-
induction n generalizing s with
39-
| zero => rfl
40-
| succ n ih =>
41-
simp only [take, drop]
42-
split <;> simp [ih]
43-
44-
/-- Tail recursive version of `Stream.take`. -/
45-
def takeTR [Stream σ α] (s : σ) (n : Nat) : List α × σ :=
46-
loop s [] n
47-
where
48-
/-- Inner loop for `Stream.takeTR`. -/
49-
loop (s : σ) (acc : List α)
50-
| 0 => (acc.reverse, s)
51-
| n+1 => match next? s with
52-
| none => (acc.reverse, s)
53-
| some (a, s) => loop s (a :: acc) n
54-
55-
theorem fst_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
56-
(takeTR.loop s acc n).fst = acc.reverseAux (take s n).fst := by
57-
induction n generalizing acc s with
58-
| zero => rfl
59-
| succ n ih => simp only [take, takeTR.loop]; split; rfl; simp [ih]
60-
61-
theorem fst_takeTR [Stream σ α] (s : σ) (n : Nat) : (takeTR s n).fst = (take s n).fst :=
62-
fst_takeTR_loop ..
63-
64-
theorem snd_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
65-
(takeTR.loop s acc n).snd = drop s n := by
66-
induction n generalizing acc s with
67-
| zero => rfl
68-
| succ n ih => simp only [takeTR.loop, drop]; split; rfl; simp [ih]
69-
70-
theorem snd_takeTR [Stream σ α] (s : σ) (n : Nat) :
71-
(takeTR s n).snd = drop s n := snd_takeTR_loop ..
72-
73-
@[csimp] theorem take_eq_takeTR : @take = @takeTR := by
74-
funext; ext : 1; rw [fst_takeTR]; rw [snd_takeTR, snd_take_eq_drop]
75-
76-
end Stream
77-
78-
@[simp] theorem List.stream_drop_eq_drop (l : List α) : Stream.drop l n = l.drop n := by
79-
induction n generalizing l with
80-
| zero => rfl
81-
| succ n ih =>
82-
match l with
83-
| [] => rfl
84-
| _::_ => simp [Stream.drop, List.drop_succ_cons, ih]
85-
86-
@[simp] theorem List.stream_take_eq_take_drop (l : List α) :
87-
Stream.take l n = (l.take n, l.drop n) := by
88-
induction n generalizing l with
89-
| zero => rfl
90-
| succ n ih =>
91-
match l with
92-
| [] => rfl
93-
| _::_ => simp [Stream.take, ih]
1+
import Batteries.Data.Stream.Basic
2+
import Batteries.Data.Stream.Fold
3+
import Batteries.Data.Stream.WF

‎Batteries/Data/Stream/Basic.lean‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/-
2+
Copyright (c) 2024 François G. Dorais. All rights reserved.
3+
Released under Apache 2. license as described in the file LICENSE.
4+
Authors: François G. Dorais
5+
-/
6+
7+
namespace Stream
8+
9+
/-- Drop up to `n` values from the stream `s`. -/
10+
def drop [Stream σ α] (s : σ) : Nat → σ
11+
| 0 => s
12+
| n+1 =>
13+
match next? s with
14+
| none => s
15+
| some (_, s) => drop s n
16+
17+
/-- Read up to `n` values from the stream `s` as a list from first to last. -/
18+
def take [Stream σ α] (s : σ) : Nat → List α × σ
19+
| 0 => ([], s)
20+
| n+1 =>
21+
match next? s with
22+
| none => ([], s)
23+
| some (a, s) =>
24+
match take s n with
25+
| (as, s) => (a :: as, s)
26+
27+
@[simp] theorem fst_take_zero [Stream σ α] (s : σ) :
28+
(take s 0).fst = [] := rfl
29+
30+
theorem fst_take_succ [Stream σ α] (s : σ) :
31+
(take s (n+1)).fst = match next? s with
32+
| none => []
33+
| some (a, s) => a :: (take s n).fst := by
34+
simp only [take]; split <;> rfl
35+
36+
@[simp] theorem snd_take_eq_drop [Stream σ α] (s : σ) (n : Nat) :
37+
(take s n).snd = drop s n := by
38+
induction n generalizing s with
39+
| zero => rfl
40+
| succ n ih =>
41+
simp only [take, drop]
42+
split <;> simp [ih]
43+
44+
/-- Tail recursive version of `Stream.take`. -/
45+
def takeTR [Stream σ α] (s : σ) (n : Nat) : List α × σ :=
46+
loop s [] n
47+
where
48+
/-- Inner loop for `Stream.takeTR`. -/
49+
loop (s : σ) (acc : List α)
50+
| 0 => (acc.reverse, s)
51+
| n+1 => match next? s with
52+
| none => (acc.reverse, s)
53+
| some (a, s) => loop s (a :: acc) n
54+
55+
theorem fst_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
56+
(takeTR.loop s acc n).fst = acc.reverseAux (take s n).fst := by
57+
induction n generalizing acc s with
58+
| zero => rfl
59+
| succ n ih => simp only [take, takeTR.loop]; split; rfl; simp [ih]
60+
61+
theorem fst_takeTR [Stream σ α] (s : σ) (n : Nat) : (takeTR s n).fst = (take s n).fst :=
62+
fst_takeTR_loop ..
63+
64+
theorem snd_takeTR_loop [Stream σ α] (s : σ) (acc : List α) (n : Nat) :
65+
(takeTR.loop s acc n).snd = drop s n := by
66+
induction n generalizing acc s with
67+
| zero => rfl
68+
| succ n ih => simp only [takeTR.loop, drop]; split; rfl; simp [ih]
69+
70+
theorem snd_takeTR [Stream σ α] (s : σ) (n : Nat) :
71+
(takeTR s n).snd = drop s n := snd_takeTR_loop ..
72+
73+
@[csimp] theorem take_eq_takeTR : @take = @takeTR := by
74+
funext; ext : 1; rw [fst_takeTR]; rw [snd_takeTR, snd_take_eq_drop]
75+
76+
end Stream
77+
78+
@[simp] theorem List.stream_drop_eq_drop (l : List α) : Stream.drop l n = l.drop n := by
79+
induction n generalizing l with
80+
| zero => rfl
81+
| succ n ih =>
82+
match l with
83+
| [] => rfl
84+
| _::_ => simp [Stream.drop, List.drop_succ_cons, ih]
85+
86+
@[simp] theorem List.stream_take_eq_take_drop (l : List α) :
87+
Stream.take l n = (l.take n, l.drop n) := by
88+
induction n generalizing l with
89+
| zero => rfl
90+
| succ n ih =>
91+
match l with
92+
| [] => rfl
93+
| _::_ => simp [Stream.take, ih]

‎Batteries/Data/Stream/Fold.lean‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/-
2+
Copyright (c) 2024 François G. Dorais. All rights reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: François G. Dorais
5+
-/
6+
import Batteries.Data.Stream.WF
7+
8+
namespace Stream
9+
10+
/-! ### foldlM -/
11+
12+
/-- Folds a monadic function over a well founded stream from left to right. (Tail recursive.) -/
13+
@[specialize] def foldlM [Monad m] [Stream.WF σ α] (f : β → α → m β) (init : β) (s : σ) : m β :=
14+
match _hint : next? s with
15+
| none => pure init
16+
| some (x, t) => f init x >>= (foldlM f · t)
17+
termination_by s
18+
19+
theorem foldlM_init [Monad m] [Stream.WF σ α] {f : β → α → m β} {init : β} {s : σ}
20+
(h : next? s = none) : foldlM f init s = pure init := by rw [foldlM, h]
21+
22+
theorem foldlM_next [Monad m] [Stream.WF σ α] {f : β → α → m β} {init : β} {s t : σ} {x}
23+
(h : next? s = some (x, t)) : foldlM f init s = f init x >>= (foldlM f · t) := by rw [foldlM, h]
24+
25+
theorem foldlM_eq_foldlM_toList [Monad m] [Stream.WF σ α] (f : β → α → m β) (init : β) (s : σ) :
26+
foldlM f init s = (toList s).foldlM f init := by
27+
induction s using Stream.recWF generalizing init with
28+
| init h => simp only [foldlM_init h, toList_init h, List.foldlM_nil]
29+
| next h ih => simp only [foldlM_next h, toList_next h, List.foldlM_cons, ih]
30+
31+
/-! ### foldl -/
32+
33+
/-- Folds a function over a well founded stream from left to right. (Tail recursive.) -/
34+
@[inline] def foldl [Stream.WF σ α] (f : β → α → β) (init : β) (s : σ) : β :=
35+
foldlM (m := Id) f init s
36+
37+
theorem foldl_init [Stream.WF σ α] {f : β → α → β} {init : β} {s : σ}
38+
(h : next? s = none) : foldl f init s = init := foldlM_init h
39+
40+
theorem foldl_next [Stream.WF σ α] {f : β → α → β} {init : β} {s t : σ}
41+
(h : next? s = some (x, t)) : foldl f init s = foldl f (f init x) t := foldlM_next h
42+
43+
theorem foldl_eq_foldl_toList [Stream.WF σ α] (f : β → α → β) (init : β) (s : σ) :
44+
foldl f init s = (toList s).foldl f init := by
45+
induction s using Stream.recWF generalizing init with
46+
| init h => simp only [foldl_init h, toList_init h, List.foldl_nil]
47+
| next h ih => simp only [foldl_next h, toList_next h, List.foldl_cons, ih]
48+
49+
/-! ### foldrM -/
50+
51+
/-- Folds a monadic function over a well founded stream from left to right. -/
52+
@[specialize] def foldrM [Monad m] [Stream.WF σ α] (f : α → β → m β) (init : β) (s : σ) : m β :=
53+
match _hint : next? s with
54+
| none => pure init
55+
| some (x, t) => foldrM f init t >>= f x
56+
termination_by s
57+
58+
theorem foldrM_init [Monad m] [Stream.WF σ α] {f : α → β → m β} {init : β} {s : σ}
59+
(h : next? s = none) : foldrM f init s = pure init := by rw [foldrM, h]
60+
61+
theorem foldrM_next [Monad m] [Stream.WF σ α] {f : α → β → m β} {init : β} {s t : σ}
62+
(h : next? s = some (x, t)) : foldrM f init s = foldrM f init t >>= f x := by rw [foldrM, h]
63+
64+
theorem foldrM_eq_foldrM_toList [Monad m] [LawfulMonad m] [Stream.WF σ α]
65+
(f : α → β → m β) (init : β) (s : σ) : foldrM f init s = (toList s).foldrM f init := by
66+
induction s using Stream.recWF with
67+
| init h => simp only [foldrM_init h, toList_init h, List.foldrM_nil]
68+
| next h ih => simp only [foldrM_next h, toList_next h, List.foldrM_cons, ih]
69+
70+
/-! ### foldr -/
71+
72+
/-- Folds a function over a well founded stream from left to right. -/
73+
@[inline] def foldr [Stream.WF σ α] (f : α → β → β) (init : β) (s : σ) : β :=
74+
foldrM (m := Id) f init s
75+
76+
theorem foldr_init [Stream.WF σ α] {f : α → β → β} {init : β} {s : σ}
77+
(h : next? s = none) : foldr f init s = init := foldrM_init h
78+
79+
theorem foldr_next [Stream.WF σ α] {f : α → β → β} {init : β} {s t : σ}
80+
(h : next? s = some (x, t)) : foldr f init s = f x (foldr f init t) := foldrM_next h
81+
82+
theorem foldr_eq_foldr_toList [Stream.WF σ α] (f : α → β → β) (init : β) (s : σ) :
83+
foldr f init s = (toList s).foldr f init := by
84+
induction s using Stream.recWF with
85+
| init h => rw [foldr_init h, toList_init h, List.foldr_nil]
86+
| next h ih => rw [foldr_next h, toList_next h, List.foldr_cons, ih]
87+
88+
end Stream
89+
90+
theorem List.stream_foldlM_eq_foldlM [Monad m] (f : β → α → m β) (init : β) (l : List α) :
91+
Stream.foldlM f init l = l.foldlM f init := by
92+
induction l generalizing init with
93+
| nil => rw [Stream.foldlM_init, foldlM_nil]; rfl
94+
| cons x l ih => rw [Stream.foldlM_next, foldlM_cons]; congr; funext; apply ih; rfl
95+
96+
theorem List.stream_foldl_eq_foldl (f : β → α → β) (init : β) (l : List α) :
97+
Stream.foldl f init l = l.foldl f init := by
98+
induction l generalizing init with
99+
| nil => rw [Stream.foldl_init, foldl_nil]; rfl
100+
| cons x l ih => rw [Stream.foldl_next, foldl_cons]; congr; funext; apply ih; rfl
101+
102+
theorem List.stream_foldrM_eq_foldrM [Monad m] [LawfulMonad m] (f : α → β → m β) (init : β)
103+
(l : List α) : Stream.foldrM f init l = l.foldrM f init := by
104+
induction l generalizing init with
105+
| nil => rw [Stream.foldrM_init, foldrM_nil]; rfl
106+
| cons x l ih => rw [Stream.foldrM_next, foldrM_cons]; congr; funext; apply ih; rfl
107+
108+
theorem List.stream_foldr_eq_foldr (f : α → β → β) (init : β) (l : List α) :
109+
Stream.foldr f init l = l.foldr f init := by
110+
induction l generalizing init with
111+
| nil => rw [Stream.foldr_init, foldr_nil]; rfl
112+
| cons x l ih => rw [Stream.foldr_next, foldr_cons]; congr; funext; apply ih; rfl

0 commit comments

Comments
 (0)