How would you make a monad instance of this coroutine?

Question

I am rolling a Coroutine package for education purposes, here it is:

data Step a b r 
  = Stop
  | Yield b r      
  | Await (a -> r) 

instance Functor (Step a b) where
   fmap g s = case s of 
   Stop       -> Stop
   Yield b r  -> Yield b $ g r
   Await f    -> Await (g . f) 

data CoroutineT m a b = CoT { resume :: m (Step a b (CoroutineT m a b)) }

run_ :: Monad m => CoroutineT m a b -> [a] -> m [b]
run_ (CoT m) as = m >>= \step -> case step of
     Stop      -> return []
     Yield o r -> liftM (o:) $ run_ r as 
     Await k   -> case is of 
       []     -> return [] 
       (x:xs) -> run_ (k x) xs

instance Monad m => Functor (CoroutineT m a) where
  fmap g (CoT m) = CoT $ liftM ap m where 
    ap Stop        = Stop
    ap (Yield b r) = Yield (g b) (fmap g r)
    ap (Await k)   = Await $ (fmap g) . k


instance Monad m => Monad (CoroutineT m a) where
   return b      = CoT . return . Yield b $ return b    
   (CoT m) >>= g = CoT $ liftM go m where 
        go Stop        = Stop
        go (Yield b r) = undefined      -- * This line I am having trouble with
        go (Await k)   = Await $ (>>=g) . k

As you could see in the comments above, the only line I am having issue with the Yield case, I can see that

(>>=)     :: CoroutineT m a b -> (b -> CoroutineT m a c) -> CoroutineT m a c
(g b)     :: CoroutineT m a c
r         :: CoroutineT m a b
(r >>= g) :: CoroutineT m a c

But I am not sure of

1. How to put them together so it type checks
2. What the semantics of bind is in the case of Yield

Answer 1

Per Gabriel's suggestion, the alternate implementation.

data Step a b x r 
  = Done x 
  | Yield b r
  | Await (a -> r)
  | Fail

instance Functor (Step a b x) where
    fmap g s = case s of 
        Done x    -> Done x
        Yield b r -> Yield b (g r)
        Await k   -> Await $ g . k
        Fail      -> Fail


-- | note the monad type needs to parameterize over type of `x` in `Done x`
data CoroutineT a b m x = CoT { resume :: m (Step a b x (CoroutineT a b m x)) }

instance Monad m => Functor (CoroutineT a b m) where
  fmap g (CoT m) = CoT $ liftM ap m where 
    ap (Done x)    = Done $ g x
    ap (Yield b r) = Yield b (fmap g r) 
    ap (Await k)   = Await $ (fmap g) . k
    ap Fail        = Fail


instance Monad m => Monad (CoroutineT a b m) where
  return = CoT . return . Done
  (CoT m) >>= g = CoT $ m >>= \step -> case step of 
    Done x    -> resume $ g x
    Yield b r -> return . Yield b $ r >>= g
    Await k   -> return . Await $ (>>=g) . k   
    Fail      -> return Fail

Note the implementation of return makes more sense now as well.

Answer 2

The base functor should be:

data Step a b r x 
  = Stop r
  | Yield b x      
  | Await (a -> x)

... and your coroutine type should be:

data CoroutineT m a b r = CoT { resume :: m (Step a b r (CoroutineT m a b r)) }

Until you make those two fixes it won't work.