120 lines
4.3 KiB
Text
120 lines
4.3 KiB
Text
/-
|
||
Copyright (c) 2020 Wojciech Nawrocki. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
|
||
Authors: Wojciech Nawrocki
|
||
-/
|
||
import Init.System.IO
|
||
|
||
namespace IO
|
||
|
||
universe u v
|
||
|
||
/-- An async IO list is like a lazy list but instead of being *unevaluated* `Thunk`s,
|
||
lazy tails are `Task`s *being evaluated asynchronously*. A tail can signal the end
|
||
of computation (successful or due to a failure) with a terminating value of type `ε`. -/
|
||
inductive AsyncList (ε : Type u) (α : Type v) where
|
||
| cons (hd : α) (tl : AsyncList ε α)
|
||
| asyncTail (tl : Task $ Except ε $ AsyncList ε α)
|
||
| nil
|
||
|
||
namespace AsyncList
|
||
|
||
instance : Inhabited (AsyncList ε α) := ⟨nil⟩
|
||
|
||
-- TODO(WN): tail-recursion without forcing sync?
|
||
partial def append : AsyncList ε α → AsyncList ε α → AsyncList ε α
|
||
| cons hd tl, s => cons hd (append tl s)
|
||
| asyncTail ttl, s => asyncTail (ttl.map $ Except.map (append · s))
|
||
| nil, s => s
|
||
|
||
instance : Append (AsyncList ε α) := ⟨append⟩
|
||
|
||
def ofList : List α → AsyncList ε α :=
|
||
List.foldr AsyncList.cons AsyncList.nil
|
||
|
||
instance : Coe (List α) (AsyncList ε α) := ⟨ofList⟩
|
||
|
||
/-- A stateful step computation `f` is applied iteratively, forming an async
|
||
stream. The stream ends once `f` returns `none` for the first time.
|
||
|
||
For cooperatively cancelling an ongoing computation, we recommend referencing
|
||
a cancellation token in `f` and checking it when appropriate. -/
|
||
partial def unfoldAsync (f : StateT σ (EIO ε) $ Option α) (init : σ)
|
||
: BaseIO (AsyncList ε α) := do
|
||
let rec step (s : σ) : EIO ε (AsyncList ε α) := do
|
||
let (aNext, sNext) ← f s
|
||
match aNext with
|
||
| none => return nil
|
||
| some aNext => do
|
||
let tNext ← EIO.asTask (step sNext)
|
||
return cons aNext $ asyncTail tNext
|
||
|
||
let tInit ← EIO.asTask (step init)
|
||
return asyncTail tInit
|
||
|
||
/-- The computed, synchronous list. If an async tail was present, returns also
|
||
its terminating value. -/
|
||
partial def getAll : AsyncList ε α → List α × Option ε
|
||
| cons hd tl =>
|
||
let ⟨l, e?⟩ := tl.getAll
|
||
⟨hd :: l, e?⟩
|
||
| nil => ⟨[], none⟩
|
||
| asyncTail tl =>
|
||
match tl.get with
|
||
| Except.ok tl => tl.getAll
|
||
| Except.error e => ⟨[], some e⟩
|
||
|
||
/-- Spawns a `Task` waiting on the prefix of elements for which `p` is true. -/
|
||
partial def waitAll (p : α → Bool := fun _ => true) : AsyncList ε α → BaseIO (Task (List α × Option ε))
|
||
| cons hd tl => do
|
||
if p hd then
|
||
let t ← tl.waitAll p
|
||
return t.map fun ⟨l, e?⟩ => ⟨hd :: l, e?⟩
|
||
else
|
||
return Task.pure ⟨[hd], none⟩
|
||
| nil => return Task.pure ⟨[], none⟩
|
||
| asyncTail tl => do
|
||
BaseIO.bindTask tl fun
|
||
| Except.ok tl => tl.waitAll p
|
||
| Except.error e => return Task.pure ⟨[], some e⟩
|
||
|
||
/-- Spawns a `Task` acting like `List.find?` but which will wait for tail evalution
|
||
when necessary to traverse the list. If the tail terminates before a matching element
|
||
is found, the task throws the terminating value. -/
|
||
partial def waitFind? (p : α → Bool) : AsyncList ε α → BaseIO (Task $ Except ε $ Option α)
|
||
| nil => return Task.pure <| Except.ok none
|
||
| cons hd tl => do
|
||
if p hd then return Task.pure <| Except.ok <| some hd
|
||
else tl.waitFind? p
|
||
| asyncTail tl => do
|
||
BaseIO.bindTask tl fun
|
||
| Except.ok tl => tl.waitFind? p
|
||
| Except.error e => return Task.pure <| Except.error e
|
||
|
||
/-- Extends the `finishedPrefix` as far as possible. If computation was ongoing
|
||
and has finished, also returns the terminating value. -/
|
||
partial def updateFinishedPrefix : AsyncList ε α → BaseIO (AsyncList ε α × Option ε)
|
||
| cons hd tl => do
|
||
let ⟨tl, e?⟩ ← tl.updateFinishedPrefix
|
||
pure ⟨cons hd tl, e?⟩
|
||
| nil => pure ⟨nil, none⟩
|
||
| l@(asyncTail tl) => do
|
||
if (← hasFinished tl) then
|
||
match tl.get with
|
||
| Except.ok tl => tl.updateFinishedPrefix
|
||
| Except.error e => pure ⟨nil, some e⟩
|
||
else pure ⟨l, none⟩
|
||
|
||
private partial def finishedPrefixAux : List α → AsyncList ε α → List α
|
||
| acc, cons hd tl => finishedPrefixAux (hd :: acc) tl
|
||
| acc, nil => acc
|
||
| acc, asyncTail _ => acc
|
||
|
||
/-- The longest already-computed prefix of the list. -/
|
||
def finishedPrefix : AsyncList ε α → List α :=
|
||
List.reverse ∘ (finishedPrefixAux [])
|
||
|
||
end AsyncList
|
||
|
||
end IO
|