Skip to content
95 changes: 2 additions & 93 deletions Batteries/Data/Stream.lean
Original file line number Diff line number Diff line change
@@ -1,93 +1,2 @@
/-
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.Finite
160 changes: 160 additions & 0 deletions Batteries/Data/Stream/Basic.lean
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
/-
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`. -/
private 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

private 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]

private theorem fst_takeTR [Stream σ α] (s : σ) (n : Nat) : (takeTR s n).fst = (take s n).fst :=
fst_takeTR_loop ..

private 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]

private theorem snd_takeTR [Stream σ α] (s : σ) (n : Nat) :
(takeTR s n).snd = drop s n := snd_takeTR_loop ..

@[csimp] private 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]

/-## Stream Iterators

The standard library provides iterators that behave much like streams but include many additional
features to support monadic effects, among other things. This added functionality makes the user
interface more complicated. By comparison, the stream interface is quite simple and easy to use.
The standard library provides a stream interface for productive pure iterators.

The following provide the reverse translation. The function `Stream.iter` converts a `Stream σ α`
into a productive pure iterator of type `StreamIterator σ α`. Finiteness properties for streams are
provided in `Batteries.Stream.Finite`.

In addition to providing a back and forth translation between stream types and productive pure
iterators. This functionality is also useful for existing stream-based libraries to incrementally
convert to iterators.
-/

/-- The underlying type of a stream iterator. -/
structure StreamIterator (σ α) [Stream σ α] where
/-- Underlying stream of a stream iterator. -/
stream : σ

/--
Returns a productive pure iterator for the given stream.

**Termination properties:**

* `Finite` instance: maybe available
* `Productive` instance: always
-/
@[always_inline, inline]
def Stream.iter [Stream σ α] (stream : σ) : Std.Iterators.Iter (α := StreamIterator σ α) α :=
Std.Iterators.toIterM { stream } Id α |>.toIter

@[always_inline, inline]
instance (σ α) [Stream σ α] : Std.Iterators.Iterator (StreamIterator σ α) Id α where
IsPlausibleStep it
| .yield it' _ =>
∃ x, Stream.next? it.internalState.stream = some (x, it'.internalState.stream)
| .skip _ => False
| .done => Stream.next? it.internalState.stream = none
step it :=
match Stream.next? it.internalState.stream with
| some (out, stream) => .yield ⟨⟨stream⟩⟩ out ⟨out, rfl⟩
| none => .done rfl

private def StreamIterator.instProductivenessRelation [Stream σ α] :
Std.Iterators.ProductivenessRelation (StreamIterator σ α) Id where
rel := emptyWf.rel
wf := emptyWf.wf
subrelation h := by cases h

instance StreamIterator.instProductive [Stream σ α] :
Std.Iterators.Productive (StreamIterator σ α) Id :=
Std.Iterators.Productive.of_productivenessRelation StreamIterator.instProductivenessRelation

instance StreamIterator.instIteratorLoop [Stream σ α] [Monad n] :
Std.Iterators.IteratorLoop (StreamIterator σ α) Id n := .defaultImplementation

instance StreamIterator.instIteratorLoopPartial [Stream σ α] [Monad n] :
Std.Iterators.IteratorLoopPartial (StreamIterator σ α) Id n := .defaultImplementation

instance StreamIterator.instIteratorCollect [Stream σ α] [Monad n] :
Std.Iterators.IteratorCollect (StreamIterator σ α) Id n := .defaultImplementation

instance StreamIterator.instIteratorCollectPartial [Stream σ α] [Monad n] :
Std.Iterators.IteratorCollectPartial (StreamIterator σ α) Id n := .defaultImplementation
Loading