coroutine.lean

import system.io 

universes u v w r 

/-- 
   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. 
 -/
mutual inductive coroutine, coroutine_result (α : Type u) (δ : Type v) (β : Type w)
with coroutine : Type (max u v w)
| mk    {} : (α → coroutine_result) → coroutine 
with coroutine_result : Type (max u v w)
| done     {} : β → coroutine_result
| yielded {}  : δ → coroutine → coroutine_result 

/- TODO(Leo): delete. This is the recursor thas was generated by Lean 2, 
   and will be generated by Lean 4. We use it here to show that the direct_subcoroutine relation is
   well founded -/
axiom coroutine.ind {α : Type u} {δ : Type v} {β : Type w} 
                    {C₁ : coroutine α δ β → Sort r}
                    {C₂ : coroutine_result α δ β → Sort r}
                    (m₁ : Π (k : α → coroutine_result α δ β), (Π a, C₂ (k a)) → C₁ (coroutine.mk k))
                    (m₂ : Π (b : β), C₂ (coroutine_result.done b))
                    (m₃ : Π (d : δ) (c : coroutine α δ β), C₁ c → C₂ (coroutine_result.yielded d c))
                    : Π (c : coroutine α δ β), C₁ c 

namespace coroutine
variables {α : Type u} {δ : Type v} {β γ : Type w}

export coroutine_result (done yielded)

@[inline] protected def pure (b : β) : coroutine α δ β :=
mk $ λ _, done b

/-- Read the input argument passed to the coroutine. 

    Remark: should we use a different name? I added an instance [monad_reader] later. -/
@[inline] protected def read : coroutine α δ α :=
mk $ λ a, done a 

/-- `resume c a` resumes/invokes the coroutine `c` with input `a`. -/
@[inline] def resume : coroutine α δ β → α → coroutine_result α δ β  
| (mk k) a := k a

/-- Return the control to the invoker with result `d` -/
@[inline] def yield (d : δ) : coroutine α δ punit :=
mk $ λ a : α, yielded d (coroutine.pure ⟨⟩)

/-- Auxiliary relation for showing that bind/pipe terminate -/
inductive direct_subcoroutine : coroutine α δ β → coroutine α δ β → Prop
| mk : ∀ (k : α → coroutine_result α δ β) (a : α) (d : δ) (c : coroutine α δ β), k a = yielded d c → direct_subcoroutine c (mk k)

theorem direct_subcoroutine_wf : well_founded (@direct_subcoroutine α δ β) :=
begin
  constructor, intro c, 
  apply @coroutine.ind _ _ _ 
          (λ c, acc direct_subcoroutine c) 
          (λ r, ∀ (d : δ) (c : coroutine α δ β), r = yielded d c → acc direct_subcoroutine c),
  { intros k ih, dsimp at ih, constructor, intros c' h, cases h, apply ih h_a h_d, assumption },
  { intros, contradiction },
  { intros d c ih d₁ c₁ heq, injection heq, subst c, assumption } 
end 

/-- Transitive closure of direct_subcoroutine. It is not used here, but may be useful when defining
    more complex procedures. -/
def subcoroutine : coroutine α δ β → coroutine α δ β → Prop :=
tc direct_subcoroutine 

theorem subcoroutine_wf : well_founded (@subcoroutine α δ β) :=
tc.wf direct_subcoroutine_wf

-- Local instances for proving termination by well founded relation

def wf_inst₁ : has_well_founded (Σ' a : coroutine α δ β, (β → coroutine α δ γ)) :=
{ r  := psigma.lex direct_subcoroutine (λ _, empty_relation),
  wf := psigma.lex_wf direct_subcoroutine_wf (λ _, empty_wf) }

def wf_inst₂ : has_well_founded (Σ' a : coroutine α δ β, coroutine δ γ β) :=
{ r  := psigma.lex direct_subcoroutine (λ _, empty_relation),
  wf := psigma.lex_wf direct_subcoroutine_wf (λ _, empty_wf) }

local attribute [instance] wf_inst₁ wf_inst₂ 

open well_founded_tactics

protected def bind : coroutine α δ β → (β → coroutine α δ γ) → coroutine α δ γ 
| (mk k) f := mk $ λ a, 
    match k a, rfl : ∀ (n : _), n = k a → _ with 
    | done b, _      := resume (f b) a
    | yielded d c, h := 
      have direct_subcoroutine c (mk k), from 
      begin
        apply direct_subcoroutine.mk k a d,  
        rw h  
      end, 
      yielded d (bind c f) 
    end 
 using_well_founded { dec_tac := unfold_wf_rel >> process_lex (tactic.assumption) }

def pipe : coroutine α δ β → coroutine δ γ β → coroutine α γ β 
| (mk k₁) (mk k₂) := mk $ λ 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 direct_subcoroutine k₁' (mk k₁), from 
      begin
        apply direct_subcoroutine.mk k₁ a d,  
        rw h  
      end ,
      yielded r (pipe k₁' k₂')
    end      
  end   
  using_well_founded { dec_tac := unfold_wf_rel >> process_lex (tactic.assumption) }

/- 
-- We can also defined `bind` using `coroutine.ind`, but the generated code will not be efficient.
protected def bind (x : coroutine α δ β) : (β → coroutine α δ γ) → coroutine α δ γ :=
@coroutine.ind α δ β 
   (λ _, (β → coroutine α δ γ) → coroutine α δ γ) 
   (λ _, α → (β → coroutine α δ γ) → coroutine_result α δ γ)
   (λ k ih f, coroutine.mk $ λ a, ih a a f)
   (λ b a f, resume (f b) a)
   (λ d c ih a f, coroutine_result.yielded d (ih f))
   x 
-/

instance : monad (coroutine α δ) :=
{ pure := @coroutine.pure _ _,
  bind := @coroutine.bind _ _ }

instance : monad_reader α (coroutine α δ) :=
{ read := @coroutine.read _ _ }

end coroutine

/- Examples -/
open coroutine 

inductive tree (α : Type u)
| leaf {} : tree
| node    : tree → α → tree → tree

/-- Coroutine as generators/iterators -/
def visit {α : Type v} : tree α → coroutine unit α unit  
| tree.leaf         := pure () 
| (tree.node l a r) := do
  visit l,
  yield a,
  visit r  

def tst {α : Type} [has_to_string α] (t : tree α) : io unit :=
do c  ← pure $ visit t,
   (yielded v₁ c) ← pure $ resume c (),
   (yielded v₂ c) ← pure $ resume c (),
   io.print_ln $ to_string v₁,
   io.print_ln $ to_string v₂,
   return ()

#eval tst (tree.node (tree.node (tree.node tree.leaf 5 tree.leaf) 10 (tree.node tree.leaf 20 tree.leaf)) 30 tree.leaf)

def ex : state_t nat (coroutine nat string) unit :=
do 
  x ← read,   
  y ← get, 
  put (y+5),
  monad_lift $ (yield ("1) val: " ++ to_string (x+y)) : coroutine nat string unit),
  x ← read,
  y ← get, 
  monad_lift $ (yield ("2) val: " ++ to_string (x+y)) : coroutine nat string unit),
  return ()

def tst2 : io unit :=
do let c := state_t.run ex 5,
   (yielded r c₁) ← pure $ resume c 10,
   io.print_ln r,
   (yielded r c₂) ← pure $ resume c₁ 20,
   io.print_ln r,
   (done _) ← pure $ resume c₂ 30,
   (yielded r c₃) ← pure $ resume c₁ 100,
   io.print_ln r,
   io.print_ln "done",
   return ()

#eval tst2