-- Code for the article on generators

module GenCode where

import Control.Monad
import Control.Monad.Trans
import Data.List (findIndices, isPrefixOf, tails)

import Control.Monad.Writer		-- for the Writer non-example

import qualified Data.Traversable as T	-- for `Generators from any Traversable'
import qualified Data.Map as M		-- example

import LogicT
import SFKT				-- 2CPS implementation of LogicT

-- ------------------------------------------------------------------------
-- Simple examples of Icon's generators
-- described in the section `Simple generators and lazy lists'

sentence = "Store it in the neighboring harbor"

-- Lazily obtain the list of all positions at which pat occurs in str
findIL :: String -> String -> [Int]
findIL pat str = findIndices (isPrefixOf pat) . tails $ str

-- Implementing the example from the Icon overview
--      sentence := "Store it in the neighboring harbor"
--      if (i := find("or", sentence)) > 5 then write(i)
-- (in Icon, indices are 1-based)

find_ex1 = do
  case filter (>4) $ findIL "or" sentence of
   (i:_) -> print i
   []    -> return ()
-- 22


-- Implementing the example from the Icon overview
--     every i := find("or", sentence)
--        do write(i)
-- which can also be written as
--      every write(find("or", sentence))

print_all = mapM_ print $ findIL "or" sentence

{-
*GenCode> print_all
2
22
32
-}

-- ------------------------------------------------------------------------
-- Examples from the section `Generators and LogicT'

-- Implementing the Icon example
--       (i | j | k) = (0 | 1)

-- First we define the equivalent of Icon's comparison,
-- which fails if the equality does not hold
equal :: MonadPlus m => m Int -> m Int -> m Int
equal i j = do
	    iv <- i
	    jv <- j
	    if iv == jv then return iv else mzero

tcomp i j k = (i `mplus` j `mplus` k) `equal`
	      (return 0 `mplus` return 1)

tcomp_ex1 :: (LogicT t, Monad m, MonadPlus (t m)) => t m String
tcomp_ex1 = ifte (tcomp (return 2) (return 1) (return 3))
	         (\i -> return $ "Yes: " ++ show i)
	         (return "No")

tcomp_ex1r :: IO ()
tcomp_ex1r = observe tcomp_ex1 >>= putStrLn
-- Yes: 1

tcomp2_run :: IO ()
tcomp2_run = observe $
	     ifte (tcomp (return 4) (return 2) (return 3))
	          (\i -> liftIO . putStrLn $ "Yes: " ++ show i)
	          (liftIO $ putStrLn "No")
-- No


-- Implementing the Icon example
--     every write(find("or", sentence1 | sentence2))|]]

-- 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 :: (Monad m, LogicT t, MonadPlus (t m)) => Maybe Int -> t m () -> t m ()
-- iter n m = bagofN n m >> return ()

-- the following is an optimized implementation of the above
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


-- The MonadPlus analogue of lazy list cons
mcons :: MonadPlus m => a -> m a -> m a
mcons h t = return h `mplus` t

-- First we `lift' findIL to use MonadPlus
-- The operation foldr below replaces nil with mzero and cons with mcons
findIM :: MonadPlus m => String -> String -> m Int
findIM pat str = foldr mcons mzero $ findIL pat str

sentence1, sentence2 :: (MonadIO (t IO), LogicT t) => t IO String
sentence1 = liftIO (putStrLn "sentence1") >> return sentence
sentence2 = liftIO (putStrLn "sentence2") >> return "Sort of"

twosen = observe $ iter Nothing $
	 (liftIO . print =<< findIM "or" =<< sentence1 `mplus` sentence2)
	 

{-
*GenCode> twosen
sentence1
2
22
32
sentence2
1
-} 


-- ------------------------------------------------------------------------
-- The code from the section ``Suspended computations as generators''

-- Implementing Icon's example
--     procedure findodd(s1, s2)
--        every i := find(s1, s2) do
--           if i % 2 = 1 then suspend i
--     end

findodd :: MonadPlus m => String -> String -> m Int
findodd s1 s2 = do
  i <- findIM s1 s2
  if i `mod` 2 == 1 then return i else mzero

findodd_r :: IO ()
findodd_r = observe $ iter Nothing $
	    findodd "a" "abracadabra" >>= liftIO . print
{-
*GenCode> findodd_r
3
5
7
-}

-- Implementing Python's example of in-order tree traversal

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

make_full_tree :: Int -> Tree
make_full_tree depth = loop 1 depth
 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

{-
*GenCode> tree1
Node 1 (Node 2 (Node 4 Leaf Leaf) (Node 5 Leaf Leaf)) 
       (Node 3 (Node 6 Leaf Leaf) (Node 7 Leaf Leaf))
-}

-- First attempt: using alternation (mplus)
in_order1 Leaf = mzero
in_order1 (Node label left right) =
    in_order1 left `mplus` 
    return label `mplus`
    in_order1 right

-- Accumulating the labels in a list, and then print
in_order1_r = observe (bagofN Nothing $ in_order1 tree1) >>= print
-- [4,2,5,1,6,3,7]

-- Printing as we go
in_order1_r' :: IO ()
in_order1_r' = observe $ iter Nothing $ 
	      in_order1 tree1 >>= liftIO . print

{-
*GenCode> in_order1_r'
4
2
5
1
6
3
7
-}

-- Using the Writer monad: yield = tell . (:[])
-- Alas, it is too strict (non-incremental)
-- The result below shows that although we needed only three first
-- labels, we still have traversed the whole tree

-- in_orderW :: (MonadWriter [Label] m, MonadIO m) => Tree -> m ()
in_orderW Leaf = return ()
in_orderW (Node label left right) = do
    in_orderW left
    liftIO . putStrLn $ "traversing: " ++ show label
    tell [label]			-- yielding, imperfectly
    in_orderW right

in_orderW_r = do
  (_,labels) <- runWriterT $ in_orderW tree1 
  let some_labels = take 3 labels
  print some_labels

{-
*GenCode> in_orderW_r
traversing: 4
traversing: 2
traversing: 5
traversing: 1
traversing: 6
traversing: 3
traversing: 7
[4,2,5]
-}


-- ------------------------------------------------------------------------
-- Implementing the generator examples using the derived yield

-- The derived yield
yield :: MonadPlus m => e -> ErrorT e m ()
yield x = raise x `mplus` return ()

-- We start with the in-order traversal example

-- This time, we implement Python code idiomatically
in_order2 :: (MonadIO m, MonadPlus m) => Tree -> ErrorT 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

-- The trace shows that we indeed yield as we visit
{-
*GenCode> in_order2_r
traversing: 4
Generated: 4
traversing: 2
Generated: 2
traversing: 5
Generated: 5
traversing: 1
Generated: 1
traversing: 6
Generated: 6
traversing: 3
Generated: 3
traversing: 7
Generated: 7
-}

-- 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


{-
*GenCode> in_order2_r'
traversing: 4
Generated: 4
traversing: 2
Generated: 2
-}


-- 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 -> ErrorT 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

{-
*GenCode> post_order_r
4
5
11
6
7
16
28
28
-}


-- Re-implementing Icon's example, using yield this time
--     procedure findodd(s1, s2)
--        every i := find(s1, s2) do
--           if i % 2 = 1 then suspend i
--     end


findodd2 :: (Monad m, LogicT t, MonadPlus (t m)) =>
     String -> String -> ErrorT Int (t m) ()
findodd2 s1 s2 = iterE Nothing $ do
  i <- findIM s1 s2
  if i `mod` 2 == 1 then yield i else return ()


findodd2_r :: IO ()
findodd2_r = observe $ iter Nothing $
	     catchError' (findodd2 "a" "abracadabra") >>= liftIO . print
{-
*GenCode> findodd2_r
3
5
7
-}

-- ------------------------------------------------------------------------
-- Generator from any Traversable

-- Demonstrating that a Generator can be derived from any Traversable
-- (compare with `Zipper from any traversible', Haskell/ZipperTraversable.hs)

-- We use Data.Map as a sample Traversable data structure to turn
-- into a Generator. That is a bit silly since Data.Map has a very
-- rich interface; it does not need Generators.
-- However, Data.Map is the only non-trivial instance of Traversable
-- defined in the standard library.

traversable_gen :: (MonadPlus m, T.Traversable t) => t a -> ErrorT a m ()
traversable_gen t = T.mapM yield t >> return ()

-- sample collections

tmap = M.fromList [ (v,product [1..v]) | v <- [1..10] ]

-- extract a few sample elements from the collection
trav :: (Monad m, T.Traversable t) => t a -> m [a]
trav t = observe $ bagofN (Just 3) $
	 catchError' $ traversable_gen t 

travm = trav tmap >>= print
-- [1,2,6]


-- ------------------------------------------------------------------------
-- Running generators side-by-side

-- Implementing the famous same-fringe problem:
-- check to see if two binary trees have the same fringe (that is,
-- the sequence of labels when traversed in a particular order).
-- We must stop the traversal as soon as the mismatch is found,
-- returning it. To demonstrate this property, we print labels
-- as they are reached by the traversal. Because of the IO action,
-- we cannot use regular Haskell lazy lists in this example.


-- We solve the problem generally, defining a zipWith function
-- for two LogicT expressions.
-- To handle the case when two streams have different lengths,
-- the zipping function receives (Maybe a) values (with Nothing
-- signifying the end of the stream)
zipWithL :: (Monad m, LogicT t, MonadPlus (t m)) =>
	    (Maybe a -> Maybe b -> c) -> t m a -> t m b -> t m c
zipWithL f m1 m2 = do
  r1 <- msplit m1
  r2 <- msplit m2
  case (r1,r2) of
   (Nothing, Nothing)    -> return $ f Nothing Nothing
   (Just (v1,_),Nothing) -> return $ f (Just v1) Nothing
   (Nothing,Just (v2,_)) -> return $ f Nothing (Just v2)
   (Just (v1,t1),Just (v2,t2)) -> return (f (Just v1) (Just v2))
				  `mplus` zipWithL f t1 t2

-- The function same_fringe takes two generator expressions and
-- runs them side-by-side, returning the stream of mismatches.
-- The one-element stream of (Nothing,Nothing) signifies the complete match.
same_fringe :: (Monad m, LogicT t, MonadPlus (t m), Eq a) =>
	       ErrorT a (t m) () -> ErrorT a (t m) () -> t m (Maybe a, Maybe a)
same_fringe m1 m2 = zipWithL (,) (catchError' m1) (catchError' m2) >>= check
 where
 check r@(Nothing,_) = return r
 check r@(_,Nothing) = return r
 check r@(Just x,Just y) | x /= y = return r
 check _             = mzero


-- Traversing two trees, using the in_order2 generator defined earlier.

sfringe_test :: MonadIO m => Tree -> Tree -> m (Maybe Label, Maybe Label)
sfringe_test t1 t2 = observe $ same_fringe (in_order2 t1) (in_order2 t2)


-- First test: comparing the identical trees
-- We have to traverse them completely. The trace shows the traversal
-- done side-by-side
sfringe_test1_r = sfringe_test tree1 tree1 >>= print

{-
*GenCode> sfringe_test1_r
traversing: 4
traversing: 4
traversing: 2
traversing: 2
traversing: 5
traversing: 5
traversing: 1
traversing: 1
traversing: 6
traversing: 6
traversing: 3
traversing: 3
traversing: 7
traversing: 7
(Nothing,Nothing)
-}

-- Comparing two different complete trees in-order
-- The trees differ already in the first label
-- The trace demonstrates the traversal is finished immediately.
sfringe_test2_r = sfringe_test tree1 (make_full_tree 4) >>= print
{-
*GenCode> sfringe_test2_r
traversing: 4
traversing: 8
(Just 4,Just 8)
-}

-- Comparing the same trees pre-order
-- This time, we proceed farther before encountering the mis-match

-- pre-order differs from in_order2 in the transposition of two lines
pre_order :: (MonadIO m, MonadPlus m) => Tree -> ErrorT Label m ()
pre_order Leaf = return ()
pre_order (Node label left right) = do
    liftIO . putStrLn $ "traversing: " ++ show label
    yield label
    pre_order left
    pre_order right

sfringe_test3_r = 
 observe (same_fringe (pre_order tree1) (pre_order (make_full_tree 4))) >>=
 print

{-
*GenCode> sfringe_test3_r
traversing: 1
traversing: 1
traversing: 2
traversing: 2
traversing: 4
traversing: 4
traversing: 5
traversing: 8
(Just 5,Just 8)
-}

-- A small modification lets us see the first 3 mismatches
sfringe_test4_r :: IO ()
sfringe_test4_r = 
 observe $ iter (Just 3) $
  (same_fringe (pre_order tree1) (pre_order (make_full_tree 4))) >>=
  liftIO. print

{-
*GenCode> sfringe_test4_r
traversing: 1
traversing: 1
traversing: 2
traversing: 2
traversing: 4
traversing: 4
traversing: 5
traversing: 8
(Just 5,Just 8)
traversing: 3
traversing: 9
(Just 3,Just 9)
traversing: 6
traversing: 5
(Just 6,Just 5)
-}

-- ------------------------------------------------------------------------
-- The following is the standard Error monad transformer,
-- rewritten without the stupid and annoying restriction on
-- the type of the exception. I should be able to throw exceptions
-- of any type, printable or not!

data ErrorT e m a = ErrorT{unErrorT :: m (Either e a)}

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

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

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

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

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

-- Functions to `run' the ErrorT monad
catchError :: Monad m => ErrorT e m e -> m e
catchError (ErrorT m) = m >>= check
 where
 check (Left x)  = return x
 check (Right x) = return x

-- 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 => ErrorT e m () -> m e
catchError' (ErrorT m) = m >>= check
 where
 check (Left x)  = return x
 check (Right x) = mzero


-- Lifting iter to the ErrorT-transformed LogicT
-- We propagate the exceptions
iterE :: (Monad m, LogicT t, MonadPlus (t m)) => 
	 Maybe Int -> ErrorT e (t m) () -> ErrorT e (t m) ()
iterE n (ErrorT m) = ErrorT $ 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 = unErrorT $ iterE (liftM pred n) (ErrorT t)