202 lines
7.5 KiB
Text
202 lines
7.5 KiB
Text
universes u v w r s
|
||
|
||
inductive coroutineResultCore (coroutine : Type (max u v w)) (α : Type u) (δ : Type v) (β : Type w) : Type (max u v w)
|
||
| done : β → coroutineResultCore
|
||
| yielded : δ → coroutine → coroutineResultCore
|
||
|
||
/--
|
||
Asymmetric coroutines `coroutine α δ β` takes inputs of Type `α`, yields elements of Type `δ`,
|
||
and produces an element of Type `β`.
|
||
|
||
Asymmetric coroutines are so called because they involve two types of control transfer operations:
|
||
one for resuming/invoking a coroutine and one for suspending it, the latter returning
|
||
control to the coroutine invoker. An asymmetric coroutine can be regarded as subordinate
|
||
to its caller, the relationship between them being similar to that between a called and
|
||
a calling routine.
|
||
-/
|
||
inductive coroutine (α : Type u) (δ : Type v) (β : Type w) : Type (max u v w)
|
||
| mk : (α → coroutineResultCore coroutine α δ β) → coroutine
|
||
|
||
abbrev coroutineResult (α : Type u) (δ : Type v) (β : Type w) : Type (max u v w) :=
|
||
coroutineResultCore (coroutine α δ β) α δ β
|
||
|
||
namespace coroutine
|
||
variables {α : Type u} {δ : Type v} {β γ : Type w}
|
||
|
||
export coroutineResultCore (done yielded)
|
||
|
||
/-- `resume c a` resumes/invokes the coroutine `c` with input `a`. -/
|
||
@[inline] def resume : coroutine α δ β → α → coroutineResult α δ β
|
||
| mk k, a => k a
|
||
|
||
@[inline] protected def pure (b : β) : coroutine α δ β :=
|
||
mk $ fun _ => done b
|
||
|
||
/-- Read the input argument passed to the coroutine.
|
||
Remark: should we use a different Name? I added an instance [MonadReader] later. -/
|
||
@[inline] protected def read : coroutine α δ α :=
|
||
mk $ fun a => done a
|
||
|
||
/-- Run nested coroutine with transformed input argument. Like `ReaderT.adapt`, but
|
||
cannot change the input Type. -/
|
||
@[inline] protected def adapt (f : α → α) (c : coroutine α δ β) : coroutine α δ β :=
|
||
mk $ fun a => c.resume (f a)
|
||
|
||
/-- Return the control to the invoker with Result `d` -/
|
||
@[inline] protected def yield (d : δ) : coroutine α δ PUnit :=
|
||
mk $ fun a => yielded d (coroutine.pure ⟨⟩)
|
||
|
||
/-
|
||
TODO(Leo): following relations have been commented because Lean4 is currently
|
||
accepting non-terminating programs.
|
||
|
||
/-- Auxiliary relation for showing that bind/pipe terminate -/
|
||
inductive directSubcoroutine : coroutine α δ β → coroutine α δ β → Prop
|
||
| mk : ∀ (k : α → coroutineResult α δ β) (a : α) (d : δ) (c : coroutine α δ β), k a = yielded d c → directSubcoroutine c (mk k)
|
||
|
||
theorem directSubcoroutineWf : WellFounded (@directSubcoroutine α δ β) :=
|
||
begin
|
||
Constructor, intro c,
|
||
apply @coroutine.ind _ _ _
|
||
(fun c => Acc directSubcoroutine c)
|
||
(fun r => ∀ (d : δ) (c : coroutine α δ β), r = yielded d c → Acc directSubcoroutine c),
|
||
{ intros k ih, dsimp at ih, Constructor, intros c' h, cases h, apply ih hA hD, assumption },
|
||
{ intros, contradiction },
|
||
{ intros d c ih d₁ c₁ Heq, injection Heq, subst c, assumption }
|
||
end
|
||
|
||
/-- Transitive closure of directSubcoroutine. It is not used here, but may be useful when defining
|
||
more complex procedures. -/
|
||
def subcoroutine : coroutine α δ β → coroutine α δ β → Prop :=
|
||
Tc directSubcoroutine
|
||
|
||
theorem subcoroutineWf : WellFounded (@subcoroutine α δ β) :=
|
||
Tc.wf directSubcoroutineWf
|
||
|
||
-- Local instances for proving termination by well founded relation
|
||
|
||
def bindWfInst : HasWellFounded (Σ' a : coroutine α δ β, (β → coroutine α δ γ)) :=
|
||
{ r := Psigma.Lex directSubcoroutine (fun _ => emptyRelation),
|
||
wf := Psigma.lexWf directSubcoroutineWf (fun _ => emptyWf) }
|
||
|
||
def pipeWfInst : HasWellFounded (Σ' a : coroutine α δ β, coroutine δ γ β) :=
|
||
{ r := Psigma.Lex directSubcoroutine (fun _ => emptyRelation),
|
||
wf := Psigma.lexWf directSubcoroutineWf (fun _ => emptyWf) }
|
||
|
||
local attribute [instance] wfInst₁ wfInst₂
|
||
|
||
open wellFoundedTactics
|
||
|
||
-/
|
||
|
||
/- TODO: remove `unsafe` keyword after we restore well-founded recursion -/
|
||
|
||
@[inlineIfReduce] protected unsafe def bind : coroutine α δ β → (β → coroutine α δ γ) → coroutine α δ γ
|
||
| mk k, f => mk $ fun a =>
|
||
match k a, rfl : ∀ (n : _), n = k a → _ with
|
||
| done b, _ => coroutine.resume (f b) a
|
||
| yielded d c, h =>
|
||
-- have directSubcoroutine c (mk k), { apply directSubcoroutine.mk k a d, rw h },
|
||
yielded d (bind c f)
|
||
-- usingWellFounded { decTac := unfoldWfRel >> processLex (tactic.assumption) }
|
||
|
||
unsafe def pipe : coroutine α δ β → coroutine δ γ β → coroutine α γ β
|
||
| mk k₁, mk k₂ => mk $ fun a =>
|
||
match k₁ a, rfl : ∀ (n : _), n = k₁ a → _ with
|
||
| done b, h => done b
|
||
| yielded d k₁', h =>
|
||
match k₂ d with
|
||
| done b => done b
|
||
| yielded r k₂' =>
|
||
-- have directSubcoroutine k₁' (mk k₁), { apply directSubcoroutine.mk k₁ a d, rw h },
|
||
yielded r (pipe k₁' k₂')
|
||
-- usingWellFounded { decTac := unfoldWfRel >> processLex (tactic.assumption) }
|
||
|
||
private unsafe def finishAux (f : δ → α) : coroutine α δ β → α → List δ → List δ × β
|
||
| mk k, a, ds =>
|
||
match k a with
|
||
| done b => (ds.reverse, b)
|
||
| yielded d k' => finishAux k' (f d) (d::ds)
|
||
|
||
/-- Run a coroutine to completion, feeding back yielded items after transforming them with `f`. -/
|
||
unsafe def finish (f : δ → α) : coroutine α δ β → α → List δ × β :=
|
||
fun k a => finishAux f k a []
|
||
|
||
unsafe instance : Monad (coroutine α δ) :=
|
||
{ pure := @coroutine.pure _ _,
|
||
bind := @coroutine.bind _ _ }
|
||
|
||
unsafe instance : MonadReaderOf α (coroutine α δ) :=
|
||
{ read := @coroutine.read _ _ }
|
||
|
||
end coroutine
|
||
|
||
/-- Auxiliary class for lifiting `yield` -/
|
||
class monadCoroutine (α : outParam (Type u)) (δ : outParam (Type v)) (m : Type w → Type r) :=
|
||
(yield : δ → m PUnit)
|
||
|
||
instance (α : Type u) (δ : Type v) : monadCoroutine α δ (coroutine α δ) :=
|
||
{ yield := coroutine.yield }
|
||
|
||
instance monadCoroutineTrans (α : Type u) (δ : Type v) (m : Type w → Type r) (n : Type w → Type s)
|
||
[monadCoroutine α δ m] [MonadLift m n] : monadCoroutine α δ n :=
|
||
{ yield := fun d => monadLift (monadCoroutine.yield d : m _) }
|
||
|
||
export monadCoroutine (yield)
|
||
|
||
open coroutine
|
||
|
||
namespace ex1
|
||
|
||
inductive tree (α : Type u)
|
||
| leaf : tree
|
||
| Node : tree → α → tree → tree
|
||
|
||
/-- Coroutine as generators/iterators -/
|
||
unsafe def visit {α : Type v} : tree α → coroutine Unit α Unit
|
||
| tree.leaf => pure ()
|
||
| tree.Node l a r => do
|
||
visit l;
|
||
yield a;
|
||
visit r
|
||
|
||
unsafe def tst {α : Type} [HasToString α] (t : tree α) : IO Unit :=
|
||
do c ← pure $ visit t;
|
||
(yielded v₁ c) ← pure (resume c ()) | throw $ IO.userError "failed";
|
||
(yielded v₂ c) ← pure (resume c ()) | throw $ IO.userError "failed";
|
||
IO.println $ toString v₁;
|
||
IO.println $ toString v₂;
|
||
pure ()
|
||
|
||
-- #eval tst (tree.Node (tree.Node (tree.Node tree.leaf 5 tree.leaf) 10 (tree.Node tree.leaf 20 tree.leaf)) 30 tree.leaf)
|
||
|
||
end ex1
|
||
|
||
namespace ex2
|
||
|
||
unsafe def ex : StateT Nat (coroutine Nat String) Unit :=
|
||
do
|
||
x ← read;
|
||
y ← get;
|
||
set (y+5);
|
||
yield ("1) val: " ++ toString (x+y));
|
||
x ← read;
|
||
y ← get;
|
||
yield ("2) val: " ++ toString (x+y));
|
||
pure ()
|
||
|
||
unsafe def tst2 : IO Unit :=
|
||
do let c := StateT.run ex 5;
|
||
(yielded r c₁) ← pure $ resume c 10 | throw $ IO.userError "failed";
|
||
IO.println r;
|
||
(yielded r c₂) ← pure $ resume c₁ 20 | throw $ IO.userError "failed";
|
||
IO.println r;
|
||
(done _) ← pure $ resume c₂ 30 | throw $ IO.userError "failed";
|
||
(yielded r c₃) ← pure $ resume c₁ 100 | throw $ IO.userError "failed";
|
||
IO.println r;
|
||
IO.println "done";
|
||
pure ()
|
||
|
||
-- #eval tst2
|
||
|
||
end ex2
|