Haskell implementation of python's yield

module Lib
    ( someFunc
    ) where

import Control.Monad
import Control.Monad.Trans
import Control.Monad.Logic
import Control.Monad.Logic.Class

someFunc :: IO ()
someFunc = putStrLn "someFunc"

newtype EitherT e m a = EitherT{runEitherT :: m (Either e a)}

instance Monad m => Monad (EitherT e m) where
    return  = EitherT . return . Right
    EitherT m >>= f = EitherT $ m >>= check
      where
      check (Right a) = runEitherT $ f a
      check (Left  e) = return $ Left e

instance MonadPlus m => MonadPlus (EitherT e m) where
    mzero = EitherT mzero
    mplus (EitherT m1) (EitherT m2) = EitherT $ m1 `mplus` m2

instance MonadTrans (EitherT e) where
    lift m = EitherT $ m >>= return . Right

instance MonadIO m => MonadIO (EitherT e m) where
    liftIO = lift . liftIO

raise :: Monad m => e -> EitherT e m a
raise = EitherT . return . Left

yield :: MonadPlus m => e -> EitherT e m ()
yield x = raise x `mplus` return ()

-- We start with the in-order traversal example

-- A variant of catchError when we don't care about the
-- return type, and the normal return of an expression is mapped
-- to mzero. This is common for the normal return from a generator
catchError' :: MonadPlus m => EitherT e m () -> m e
catchError' (EitherT m) = m >>= check
  where
  check (Left x)  = return x
  check (Right x) = mzero

-- Lifting iter to the EitherT-transformed LogicT
-- We propagate the exceptions
iterE :: (Monad m, MonadLogic (t m), MonadPlus (t m)) =>
  Maybe Int -> EitherT e (t m) () -> EitherT e (t m) ()
iterE n (EitherT m) = EitherT $ msplit m >>= check n
  where
  check _ Nothing = return (Right ())
  check (Just n) _ | n <= 1  = return (Right ())
  check n (Just (Right _,t)) = next n t
  check n (Just (Left e,t))  = return (Left e) `mplus` next n t
  next n t = runEitherT $ iterE (liftM pred n) (EitherT t)

-- A version of bagofN that doesn't care about the result of
-- the computation (which is unit). No need to accumulate it in a list
-- iter n m = bagofN n m >> return ()
-- the following is an optimized implementation of the above
iter :: (Monad m, MonadLogic (t m), MonadPlus (t m)) => Maybe Int -> t m () -> t m ()
iter n m = msplit m >>= check n
  where
  check _ Nothing = return ()
  check (Just n) _ | n <= 1 = return ()
  check n (Just (_,t)) = iter (liftM pred n) t


type Label = Int
data Tree = Leaf | Node Label Tree Tree deriving Show

make_full_tree :: Int -> Tree
make_full_tree = loop 1
 where
 loop label 0 = Leaf
 loop label n = Node label (loop (2*label) (pred n)) (loop (2*label+1) (pred n))

tree1 = make_full_tree 3

-- This time, we implement Python code idiomatically
in_order2 :: (MonadIO m, MonadPlus m) => Tree -> EitherT Label m ()
in_order2 Leaf = return ()
in_order2 (Node label left right) = do
    in_order2 left
    liftIO . putStrLn $ "traversing: " ++ show label
    yield label
    in_order2 right

in_order2_r :: IO ()
in_order2_r = observe $ iter Nothing $ do
  i <- catchError' (in_order2 tree1)
  liftIO . putStrLn $ "Generated: " ++ show i

-- Stopping the generator earlier: request only two generated values
-- The trace shows that we stop the traversal after consuming
-- the needed two values.
-- We indeed traverse on-demand.
in_order2_r' :: IO ()
in_order2_r' = observe $ iter (Just 2) $ do
  i <- catchError' (in_order2 tree1)
  liftIO . putStrLn $ "Generated: " ++ show i


-- The post-order traversal example:
-- traverse a tree post-order and print out the sum of the current
-- label and the labels in the left and the right branches.
-- Now the generator has to return a useful value.

post_order :: MonadPlus m => Tree -> EitherT Label m Label
post_order Leaf = return 0
post_order (Node label left right) = do
  sum_left  <- post_order left
  sum_right <- post_order right
  let sum = sum_left + sum_right + label
  yield sum
  return sum

post_order_r :: IO ()
post_order_r = observe $ iter Nothing $
           catchError (post_order tree1) >>= liftIO . print