Hunter S Thompson once typed out the entirety of "The Great Gatsby" so that he could feel what it was like to write the great American novel. In this spirit I am reimplementing core concepts in functional programming and trying to work out the harder puzzles myself.

gatsby.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE AllowAmbiguousTypes #-}

-- The following functions are re-implemented
import Prelude hiding ( filter
                      , sequence
                      , sequence_
                      , mapM_
                      , maybe
                      , foldr
                      , foldl
                      , map
                      , concat
                      , (++)
                      , unzip
                      , iterate
                      , length
                      , empty
                      , Monoid
                      , succ
                      , replicate
                      , take
                      , repeat
                      , print
                      , Either
                      , either
                      , Left
                      , Right
                      , join
                      , zip
                      , zipWith
                      , reverse
                      , unlines
                      )

-- Pragmatic utilities for displaying values in the console and making some
-- slightly more interesting examples
import qualified Prelude as P (concat, (++), length, foldr, filter)
import System.IO (isEOF)
import Data.Char (toUpper)

-- Machinery Haskell expects me to use (for example, any monad instances I
-- define must also be Applicative instances; can't have one without the
-- other).
import Control.Applicative hiding ( some
                                  , getConst
                                  , Const
                                  , empty
                                  , Alternative(..)
                                  )
import Control.Monad.Trans.Class

{-
 - Prologue: typeclass trickery and building blocks
 -
 - The Functor and Monad typeclasses are two of the most pervasive and useful
 - typeclasses in Haskell. Functor looks like this:
 -
 -     class Functor f where
 -         fmap :: (a -> b) -> f a -> f b
 -         (<$) :: a -> f b -> f a
 -         (<$) = fmap . const
 -
 - Monad looks like this:
 -
 -     class Monad m where
 -         return :: a -> m a
 -         (>>=)  :: forall a b. m a -> (a -> m b) -> m b
 -         (>>)   :: forall a b. m a -> m b -> m b
 -         mv >> f = mv >>= \_ -> f
 -         fail :: String -> m a
 -         fail s = error s
 -
 - Were I truly implementing the language from scratch these wouldn't be
 - commented out. However re-implementing them would do nothing for me
 - educationally and would make working with the rest of Haskell really
 - annoying.
 -
 - For completion, here is the definition of Applicative, which is a functor
 - which may be used to perform a number of actions in a sequence and collect
 - the results. Since I'm importing Haskell's monad class, I need to also use
 - its Applicative class.
 -
 -     class Functor f => Applicative f where
 -         pure :: a -> f a
 -         (<*>) :: f (a -> b) -> f a -> f b
 -
 -}

-- | The identity functor
newtype Box a = Box { unbox :: a } deriving Show
instance Functor Box where
    fmap f i = Box $ f $ unbox i

instance Foldable Box where
    foldMap f (Box x) = f x

instance Applicative Box where
    pure a = Box a
    Box f <*> Box x = Box (f x)

instance Monad Box where
    return a = Box a
    m >>= k = k (unbox m)

instance Comonad Box where
    duplicate = Box
    {-# INLINE duplicate #-}
    extract = unbox
    {-# INLINE extract #-}

newtype Const a b = Const { getConst :: a }
instance Functor (Const a) where
    fmap _ (Const x) = Const x

newtype Pair a b = Pair { unpair :: forall r. (a -> b -> r) -> r }
instance Functor (Pair a) where
    fmap f p = unpair p (\x y -> pair x (f y))

pair :: a -> b -> Pair a b
pair x y = Pair $ \f -> f x y

fst' :: Pair r b -> r
fst' p = unpair p (\f _ -> f)

snd' :: Pair a r -> r
snd' p = unpair p (\_ s -> s)

{-
 - Bifunctors
 -
 - A functor which is parameterized over two variables
 -}

class Bifunctor p where
    bimap :: (a -> b) -> (c -> d) -> p a c -> p b d
    bimap f g = bifirst f . bisecond g

    bifirst :: (a -> b) -> p a c -> p b c
    bifirst f = bimap f id

    bisecond :: (b -> c) -> p a b -> p a c
    bisecond = bimap id

instance Bifunctor (,) where
    bimap f g ~(a, b) = (f a, g b)

instance Bifunctor Const where
    bimap f _ (Const a) = Const (f a)
    {-# INLINE bimap #-}

{-
 - Monoids
 -
 - A monoidal type is a type with:
 -
 - 1. An associative binary operation; and
 - 2. An identity for this operation
 -
 - A list can be thought of as the free monoid: it can wrap any type - monoid
 - or not - and produce a type which is a monoid.
 -}

class Monoid a where
    empty  :: a
    append :: a -> a -> a
    mconcat :: List a -> a
    mconcat xs = listFoldr xs append empty

(<>) :: Monoid a => a -> a -> a
(<>) = append

instance Monoid b => Monoid (a -> b) where
    empty _ = empty
    append f g x = f x `append` g x

instance Monoid String where
    empty = ""
    append x y = x P.++ y

-- | The monoid of endomorphisms under composition
newtype Endo a = Endo { appEndo :: a -> a }

instance Monoid (Endo a) where
    empty = Endo id
    Endo f `append` Endo g = Endo (f . g)

-- | The dual of a monoid, by swapping the arguments of append
newtype MonoidDual a = MonoidDual { getMonoidDual :: a }

instance Monoid a => Monoid (MonoidDual a) where
    empty = MonoidDual empty
    MonoidDual x `append` MonoidDual y = MonoidDual (y `append` x)

{- |
 - Alternative
 -
 - An Alternative is a sub-type of Applicative which are also monoids.
 -}

class Applicative f => Alternative f where
    zilch :: f a
    (<|>) :: f a -> f a -> f a

class Functor f => Alt f where
    (<!>) :: f a -> f a -> f a

    some :: Applicative f => f a -> f (List a)
    some v = some_v where
        many_v = some_v <!> pure nil
        some_v = cons <$> v <*> many_v

    many :: Applicative f => f a -> f (List a)
    many v = many_v where
        many_v = some_v <!> pure nil
        some_v = cons <$> v <*> many_v

{- |
 - Foldable
 -
 - A foldable type is one which can be folded (surprise).
 - A minimal complete definition: `foldMap` or `foldr`.
 -}

class Foldable t where
    -- | Combine the elements of a structure using a monoid
    fold :: Monoid m => t m -> m
    fold = foldMap id

    -- | Map each element of the structure to a monoid, and combine the results
    foldMap :: Monoid m => (a -> m) -> t a -> m
    foldMap f = foldr (append . f) empty

    -- | Right-associative fold of a structure
    foldr :: (a -> b -> b) -> b -> t a -> b
    foldr f z t = appEndo (foldMap (Endo . f) t) z

    -- | Left-associative fold of a structure
    foldl :: (b -> a -> b) -> b -> t a -> b
    foldl f z t = appEndo (getMonoidDual (foldMap (MonoidDual . Endo . flip f) t)) z

    -- | Left-associative fold of a structure with strict application
    foldl' :: (b -> a -> b) -> b -> t a -> b
    foldl' f z0 xs = foldr f' id xs z0 where
        f' x k z = k $! f z x

    -- | If you can fold it, you can filter it and map it
    filter :: (Monoid (t a), Applicative t) => (a -> Bool) -> t a -> t a
    filter p = foldMap (\a -> if p a then pure a else empty)

    -- | Monadic left-associative fold over the elements of a structure
    foldlM :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m b
    foldlM f z0 xs = foldr f' return xs z0
        where f' x k z = f z x >>= k


{- |
 - Bifoldable
 -
 - A foldable structure with two varieties of elements.
 - Basically, sums and products.
 -}

class Bifoldable p where
    bifold :: Monoid m => p m m -> m
    bifold = bifoldMap id id
    {-# INLINE bifold #-}

    bifoldMap :: Monoid m => (a -> m) -> (b -> m) -> p a b -> m
    bifoldMap f g = bifoldr (append . f) (append . g) empty

    bifoldr :: (a -> c -> c) -> (b -> c -> c) -> c -> p a b -> c
    bifoldr f g z t = appEndo (bifoldMap (Endo . f) (Endo . g) t) z

instance Bifoldable (,) where
    bifoldMap f g ~(a, b) = f a `append` g b
    {-# INLINE bifoldMap #-}

instance Bifoldable Const where
    bifoldMap f _ (Const a) = f a
    {-# INLINE bifoldMap #-}

{- |
 - Optional types
 - Signifies uncertainty: a function which returns an Opt value may return the
 - expected result, or none at all.
 -}

newtype Opt a = Opt { maybe :: forall r. (a -> r) -> r -> r }

just :: a -> Opt a
just x = Opt $ \s n -> s x

none :: Opt a
none = Opt $ \s n -> n

instance Functor Opt where
    fmap f o = maybe o (\x -> just (f x)) none

instance Applicative Opt where
    pure = return
    (<*>) = ap

instance Monad Opt where
    return = just
    o >>= f = maybe o (\x -> f x) none
    fail _ = none

instance Alternative Opt where
    zilch = none
    o <|> r = maybe o (\j -> just j) r

instance Foldable Opt where
    foldr f z o = maybe o (\x -> f x z) z

instance Monoid a => Monoid (Opt a) where
    empty = none
    append x y = maybe x (\x' ->
                    maybe y (\y' -> just (x' `append` y'))
                            x
                 ) y

instance Show a => Show (Opt a) where
    show o = maybe o (\x -> (P.++) "Value: " (show x)) "(none)"

{- |
 - Either!
 -
 - A coproduct of two different types. Similar to `Opt` except both cases store
 - values, which may be of different types.
 -}

newtype Either a b = Either { either :: forall r. (a -> r) -> (b -> r) -> r }


left :: a -> Either a b
left x = Either $ \l r -> l x
{-# INLINE left #-}

right :: b -> Either a b
right y = Either $ \l r -> r y
{-# INLINE right #-}

instance Functor (Either a) where
    fmap f e = either e (\x -> left x) (\y -> right (f y))

instance Monad (Either a) where
    return = right
    e >>= k = either e left (\y -> k y)

instance Applicative (Either a) where
    pure = return
    (<*>) = ap

-- | This class basically exists to overload function names
class Sequence s where
    repeat :: a -> s a
    zipWith :: (a -> b -> c) -> s a -> s b -> s c
    unfoldr :: (b -> (a, b)) -> b -> s a

{- |
 - Fusable lists
 -
 - A Mendler encoding of a list. foldr is the catamorphism, accepting as
 - arguments two possible continuations: one for receiving the next value in a
 - list, and one for halting.
 -
 - Expresses the idea of recursion and iteration.
 -}

newtype List a = List { listFoldr :: forall r. (a -> r -> r) -> r -> r }

-- | Maps a function over every element in a list.
map :: (a -> b) -> List a -> List b
map f xs = listFoldr xs (\y ys -> cons (f y) ys) nil

instance Functor List where
    fmap = map

instance Applicative List where
    pure = return
    (<*>) = ap

instance Monad List where
    return x = cons x nil
    xs >>= f = concat (map f xs)

instance Alternative List where
    zilch = nil
    (<|>) = (++)

instance Monoid (List a) where
    empty = nil
    append x y = x ++ y

instance Foldable List where
    foldr c n xs = listFoldr xs c n
    filter = listFilter -- probably more efficient this way

instance Foldable [] where
    foldr = P.foldr
    filter = P.filter

-- | Construct a new list by prepending a value to an existing list
cons :: a -> List a -> List a
cons x xs = List $ \c e -> c x $ listFoldr xs c e

-- | Construct a new, empty list
nil :: List a
nil = List $ \c e -> e

-- | Computes the fixed point of a function. See `diverge`
fix :: (a -> a) -> a
fix f = f (fix f)

-- | Computes the fixed point of the identity function, resulting in a nonsense
-- value which can represent any type. See `car`.
diverge :: a
diverge = fix id

-- | Splits a list into its head and tail. Handles the empty case.
split :: List a -> (Opt a, List a)
split xs = listFoldr xs f (none, nil) where
    f y ys = (just y, List (\c e ->
        maybe (fst ys) (\x -> c x (listFoldr (snd ys) c e))
                e))

-- | Name is a lisp term. `diverge` is used because `maybe` requires a third
-- argument, but in a non-empty list we will never actually need the value.
car :: List a -> a
car xs = maybe (fst $ split xs) id diverge

-- | Name is a lisp term. Grabs the head of the list, or none.
car' :: List a -> Opt a
car' xs = fst . split $ xs

-- | Name is a lisp term. Grabs the tail of a list.
cdr :: List a -> List a
cdr = snd . split

-- | Constructs a new list from an old one, omitting elements which do not
-- satisfy a predicate. List-specific version
listFilter :: (a -> Bool) -> List a -> List a
listFilter pred xs = listFoldr xs f nil where
    f y ys = List $ \c e -> if (pred y)
                            then c y (listFoldr ys c e)
                            else listFoldr ys c e

-- | Appends two lists
(++) :: List a -> List a -> List a
xs ++ ys = listFoldr xs cons ys

-- | Zip the elements of two lists into a list of pairs
zip :: List a -> List b -> List (a, b)
zip as bs =
    let (ha, ta) = split as
        (hb, tb) = split bs
    in  maybe ha (\a ->
                    maybe hb (\b -> cons (a, b) (zip ta tb)) nil)
                 nil

-- | Generalizes 'zip' by zipping with the function given
listZipWith :: (a -> b -> c) -> List a -> List b -> List c
listZipWith f as bs =
    let (mha, ta) = split as
        (mhb, tb) = split bs
    in  maybe mha (\ha ->
        maybe mhb (\hb -> cons (f ha hb)
                          (listZipWith f ta tb)) diverge
                          ) diverge

-- | Flattens a list of lists.
concat :: List (List a) -> List a
concat xs = listFoldr xs (++) nil

-- | Splits a list of tuples into a tuple of lists.
unzip :: List (a, b) -> ((List a), (List b))
unzip xs = listFoldr xs (\(a,b) (as,bs) -> (cons a as, cons b bs)) (nil, nil)

-- | Given a successor function and a seed value, computes an infinite list.
iterate :: (a -> a) -> a -> List a
iterate f x = cons x $ iterate f (f x)

-- | Length of a list
length :: List a -> Int
length xs = listFoldr xs (\_ -> (+1)) 0

-- | Take the first n elements of a list
take :: Int -> List a -> List a
take = fix take' where
    take' t n xs = listFoldr xs (\y ys ->
        if n == 0 then nil else cons y (t (n-1) ys)) nil

-- | Repeat a value in an infinite list
listRepeat :: a -> List a
listRepeat x = xs where xs = cons x xs

-- | Replicate a value just finite number of times in a list
replicate :: Int -> a -> List a
replicate n x = take n (repeat x)

-- | Construct a singleton list
single :: a -> List a
single = flip cons nil

-- | Reverse a list
reverse :: Foldable t => t a -> List a
reverse = foldl (flip cons) nil

-- | Join lines, after appending a terminating newline to each
unlines :: List String -> String
unlines = foldMap (P.++ "\n")

-- | Because we are doing this in Haskell, to represent things we must use
-- `String`, which implicitly means conversion to `[]`, to show our list.
toPrimList :: List a -> [a]
toPrimList xs = listFoldr xs (\y ys -> y : ys) []

instance Show a => Show (List a) where
    show = show . toPrimList where

instance Sequence List where
    repeat = listRepeat
    zipWith = listZipWith
    unfoldr f c =
        let (x, d)  = f c
        in  cons x (unfoldr f d)

{-
 - The free monad and free monad transformer
 -
 - For any Functor type f, Mu gives a monad instance for f "for free." A monad
 - may be thought of either as a terminal (or "unit" value) or a wrapped
 - continuing computation.
 -
 - A monad transformer is a monad which can be layered on top of another
 - existing monad, resulting in a new custom composition of the two. Here the
 - canonical free monad, Mu, is defined in terms of the free monad transformer
 - and Box, the identity functor.
 -}

-- | All free monad implementations will become instances of this class. See
-- `improve`.
class (Functor f, Monad m) => FreeMonad f m | m -> f where
    wrap :: f (m a) -> m a

liftF :: (FreeMonad f m) => f a -> m a
liftF = wrap . fmap return

-- | The base functor of a free monad
data MuF f a b = Pure a | Wrap (f b)
    deriving (Eq, Ord, Show, Read)

instance Functor f => Functor (MuF f a) where
    fmap _ (Pure a)  = Pure a
    fmap f (Wrap as) = Wrap (fmap f as)
    {-# INLINE fmap #-}

instance (Functor f, Monad (MuF f a)) => Applicative (MuF f a) where
    pure = return
    (<*>) = ap

instance Foldable f => Foldable (MuF f a) where
    foldMap f (Wrap as) = foldMap f as
    foldMap _ _         = empty
    {-# INLINE foldMap #-}

-- | The free monad transformer
newtype MuT f m a = MuT { runMuT :: m (MuF f a (MuT f m a)) }

-- | Finally, our first definition of a full free monad.
type Mu f = MuT f Box

runMu :: Mu f a -> MuF f a (Mu f a)
runMu = unbox . runMuT

free :: MuF f a (Mu f a) -> Mu f a
free = MuT . Box
{-# INLINE free #-}

instance (Functor f, Monad m) => Functor (MuT f m) where
    fmap f (MuT m) = MuT (liftM f' m) where
        f' (Pure a)  = Pure (f a)
        f' (Wrap as) = Wrap (fmap (fmap f) as)

instance (Functor f, Monad m) => Applicative (MuT f m) where
    pure a = MuT (return (Pure a))
    {-# INLINE pure #-}
    (<*>) = ap

instance (Functor f, Monad m) => Monad (MuT f m) where
    return a = MuT (return (Pure a))
    {-# INLINE return #-}
    MuT m >>= f = MuT $ m >>= \v -> case v of
        Pure a  -> runMuT (f a)
        Wrap w  -> return (Wrap (fmap (>>= f) w))

{-
 - Pedagogical aside: the monad bind operation
 -
 - The Monad typeclass calls this function `>>=` for reasons which will be
 - illuminated shortly.
 -
 - Its first argument is a monadic value of just type `a`, and its second
 - argument is a function `a -> m b`, where `m b` is a monadic value of type
 - `b`. If you consider the second argument to be an expression, then `>>=`
 - uses anonymous functions to bind the first argument in the expression.
 -
 - Without any syntax sugar, monadic computations look like this:
 -
 -     mf = getX              >>= \x ->
 -          getY              >>= \y ->
 -          foo x y           >>= \z ->
 -          dojustthingWith z >>
 -          return z
 -
 - Formatting this way is helpfully suggestive, I hope. It's not a stretch to
 - see how this is transformed from do-notation:
 -
 -     mf = do
 -       x <- getX
 -       y <- getY
 -       z <- foo x y
 -       dojustthingWith z
 -       return z
 -
 - There is another function, called `>>` or *next* which evaluates a monadic
 - function and disregards the result. Given a working implementation of `>>=`
 - it is written for you.
 -}

-- | Our monad transformer is, in fact, a monad transformer. `liftM` is defined
-- below along with other monad utilities.
instance MonadTrans (MuT f) where
    lift = MuT . liftM Pure
    {-# INLINE lift #-}

instance (Functor f, Monad m) => FreeMonad f (MuT f m) where
    wrap = MuT . return . Wrap
    {-# INLINE wrap #-}

{- |
 - Mendler-encoded monad transformer
 -
 - The above recursive definition of a monad is very useful and easy to work
 - with. However, it requires building up a structure which will simply be torn
 - down when evaluated. Most times, this is totally fine. Sometimes, though,
 - this leads to quadratic space usage.
 -
 - Instead of building up a structure, an alternate way to think of a monad is
 - a function which accepts two continuation functions, one for each
 - constructor above. Depending on the value of the monad, it selects either
 - continuation.
 -}

newtype MT f m a = MT { runMT :: forall r. (a -> m r) -> (f (m r) -> m r) -> m r }

instance Functor (MT f m) where
    fmap f (MT k) = MT $ \a fr -> k (a . f) fr

instance Applicative (MT f m) where
    pure a = MT $ \k _ -> k a
    MT fk <*> MT ak = MT $ \b fr -> ak (\d -> fk (\e -> b (e d)) fr) fr

instance Monad (MT f m) where
    return = pure
    MT fk >>= f = MT $ \b fr -> fk (\d -> runMT (f d) b fr) fr

instance Functor f => FreeMonad f (MT f m) where
    wrap f = MT $ \u w -> w (fmap (\(MT m) -> m u w) f)

type M f = MT f Box

runM :: Functor f
     => M f a
     -> (forall r. (a -> r) -> (f r -> r) -> r)
runM (MT m) = \u w -> unbox $ m (return . u) (return . w . fmap unbox)

-- | Generate a Mendler-encoded free monad from a `Mu`
toMT :: (Monad m, Functor f) => MuT f m a -> MT f m a
toMT (MuT f) = MT $ \ka kfr -> do
    muf <- f
    case muf of
        Pure a  -> ka a
        Wrap fb -> kfr $ fmap (($ kfr) . ($ ka) . runMT . toMT) fb

toM :: (Functor f) => Mu f a -> M f a
toM = toMT
{-# INLINE toM #-}

-- | fromM can convert from M to any other FreeMonad instance.
fromM :: (Functor f, FreeMonad f m) => M f a -> m a
fromM m = runM m return wrap
{-# INLINE fromM #-}

{- |
 - improve uses fromM to constrain the argument type as M, but then casts the
 - resulting value as a Mu. Thus, generic FreeMonad values may be wrapped in
 - improve and automatically see asymptotic improvements in performance.
 -}
improve :: Functor f => (forall m. FreeMonad f m => m a) -> Mu f a
improve m = fromM m
{-# INLINE improve #-}

{-
 - Simple monad utilities
 -
 - These are generic to any monad constructed using either free monad
 - machinery.
 -}

when :: (FreeMonad f m) => Bool -> m () -> m ()
when p s = if p then s else return ()

-- | TODO: Implement Traversable so this isn't List-specific
sequence :: (Foldable t, Monad m) => t (m a) -> m (List a)
sequence ms = foldr k (return nil) ms where
    k m m' = do { x <- m; xs <- m'; return (cons x xs) }
{-# INLINE sequence #-}

-- | TODO: Implement Traversable so this isn't List-specific
mapM :: Monad m => (a -> m b) -> List a -> m (List b)
mapM f as = sequence (fmap f as)
{-# INLINE mapM #-}

sequence_ :: (Foldable t, Monad m) => t (m a) -> m ()
sequence_ ms = foldr (>>) (return ()) ms
{-# INLINE sequence_ #-}

mapM_ :: (Foldable f, Functor f, Monad m) => (a -> m b) -> f a -> m ()
mapM_ f as = sequence_ (fmap f as)
{-# INLINE mapM_ #-}

forever :: Monad m => m a -> m b
forever x = let x' = (>>) x x' in x'
{-# INLINE forever #-}

replicateM_ :: Monad m => Int -> m a -> m ()
replicateM_ n x = sequence_ (replicate n x)

mapAndUnzipM f xs = sequence (map f xs) >>= return . unzip

liftM :: (Monad m) => (a1 -> r) -> m a1 -> m r
liftM f m1 = do { x1 <- m1; return (f x1) }

liftM2 :: (Monad m) => (a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 f m1 m2 = do { x1 <- m1; x2 <- m2; return (f x1 x2) }

ap :: (Monad m) => m (a -> b) -> m a -> m b
ap = liftM2 id

unless :: (Monad m) => Bool -> m () -> m ()
unless p s = if p then return () else s

(>=>) :: Monad m => (t -> m a) -> (a -> m b) -> t -> m b
f >=> g = \x -> f x >>= g

join :: (Monad m) => m (m a) -> m a
join x = x >>= id

{-
 - Utility monads
 -
 - There are a few handy monad transformers one might want to stack together to
 - create others. While psilo will, in all likelihood, come with some
 - combination of the following pre-fab, it's important to see how they work
 - independently for verification purposes.
 -}

-- | The StateT monad transformer and the State monad derived from it
newtype StateT s m a = StateT { runStateT :: s -> m (a, s) } deriving Functor
type State s = StateT s Box

runState :: State s a
         -> s
         -> (a, s)
runState m = unbox . runStateT m

evalState = flip runState
evalStateT = flip runStateT

instance (Functor m, Monad m) => Applicative (StateT s m) where
    pure = return
    (<*>) = ap

instance (Monad m) => Monad (StateT s m) where
    return a = StateT $ \s -> return (a, s)
    m >>= k = StateT $ \s -> do
        (a, s') <- runStateT m s
        runStateT (k a) s'
    fail str = StateT $ \_ -> fail str

instance MonadTrans (StateT s) where
    lift m = StateT $ \s -> do
        a <- m
        return (a, s)

-- | The requirements to be a State monad
class Monad m => MonadState s m | m -> s where
    get :: m s
    put :: s -> m ()

-- | State now satisfies its own requirements
instance Monad m => MonadState s (StateT s m) where
    get   = StateT $ \s -> return (s, s)
    put s = StateT $ \_ -> return ((), s)

instance (Functor f, FreeMonad f m) => FreeMonad f (StateT s m) where
    wrap fm = StateT $ \s -> wrap $ flip runStateT s <$> fm

{- |
 - The Continuation Monad & Transformer
 -
 - Provides the ability to suspend a computation to be continued elsewhere or
 - at a later time. Continuations are sufficient to express any kind of control
 - flow mechanism.
 -}

newtype ContT r m a = ContT { runContT :: (a -> m r) -> m r } deriving (Functor)

evalContT :: (Monad m) => ContT r m r -> m r
evalContT m = runContT m return

withContT :: ((b -> m r) -> (a -> m r)) -> ContT r m a -> ContT r m b
withContT f m = ContT $ runContT m . f

instance Applicative (ContT r m) where
    pure x = ContT ($ x)
    f <*> v = ContT $ \c -> runContT f $ \g -> runContT v (c . g)

instance Monad (ContT r m) where
    return x = ContT ($ x)
    m >>= k = ContT $ \c -> runContT m (\x -> runContT (k x) c)

instance MonadTrans (ContT r) where
    lift m = ContT (m >>=)

callCC' :: ((a -> ContT r m b) -> ContT r m a) -> ContT r m a
callCC' f = ContT $ \c -> runContT (f (\x -> ContT $ \_ -> c x)) c

resetT :: (Monad m) => ContT r m r -> ContT r' m r
resetT = lift . evalContT

shiftT :: (Monad m) => ((a -> m r) -> ContT r m r) -> ContT r m a
shiftT f = ContT (evalContT . f)

type Cont r = ContT r Box

cont :: ((a -> r) -> r) -> Cont r a
cont f = ContT (\c -> Box (f (unbox . c)))

runCont :: Cont r a -> (a -> r) -> r
runCont m k = unbox (runContT m (Box . k))

evalCont :: Cont r r -> r
evalCont m = unbox (evalContT m)

withCont :: ((b -> r) -> (a -> r)) -> Cont r a -> Cont r b
withCont f = withContT ((Box .) . f . (unbox .))

reset :: Cont r r -> Cont r' r
reset = resetT

shift :: ((a -> r) -> Cont r r) -> Cont r a
shift f = shiftT (f . (unbox .))

class Monad m => MonadCont m where
    callCC :: ((a -> m b) -> m a) -> m a

instance MonadCont (ContT r m) where
    callCC = callCC'

-- | The ReaderT monad transformer models a stack-like binding environment
newtype ReaderT r m a = ReaderT { runReaderT :: r -> m a } deriving Functor
type Reader r = ReaderT r Box

reader :: (Monad m) => (r -> a) -> ReaderT r m a
reader f = ReaderT (return . f)

runReader :: Reader r a
          -> r
          -> a
runReader m = unbox . runReaderT m

instance (Applicative m) => Applicative (ReaderT r m) where
    pure = liftReaderT . pure
    f <*> v = ReaderT $ \ r -> runReaderT f r <*> runReaderT v r

instance (Monad m) => Monad (ReaderT r m) where
    return = lift . return
    m >>= k = ReaderT $ \r -> do
        a <- runReaderT m r
        runReaderT (k a) r

instance MonadTrans (ReaderT r) where
    lift = liftReaderT

liftReaderT :: m a -> ReaderT r m a
liftReaderT m = ReaderT (const m)

class (Monad m) => MonadReader r m | m -> r where
    ask :: m r
    local :: (r -> r)   -- ^ The function to modify the environment
          -> m a        -- ^ Reader to run in the modified environment
          -> m a

withReaderT :: (r' -> r) -> ReaderT r m a -> ReaderT r' m a
withReaderT f m = ReaderT $ runReaderT m . f

instance (Monad m) => MonadReader r (ReaderT r m) where
    ask = ReaderT return
    local = withReaderT

instance (Functor f, FreeMonad f m) => FreeMonad f (ReaderT r m) where
    wrap fm = ReaderT $ \e -> wrap $ flip runReaderT e <$> fm

class (Monad m) => MonadIO m where
    liftIO :: IO a -> m a

instance MonadIO IO where
    liftIO = id

instance (MonadIO m) => MonadIO (StateT s m) where
    liftIO = lift . liftIO

instance (MonadIO m) => MonadIO (ContT r m) where
    liftIO = lift . liftIO

instance (MonadIO m) => MonadIO (ReaderT s m) where
    liftIO = lift . liftIO

{- | WriterT Monad Transformer -}

newtype WriterT w m a = WriterT { runWriterT :: m (a, w) }

mapWriterT :: (m (a, w) -> n (b, w')) -> WriterT w m a -> WriterT w' n b
mapWriterT f m = WriterT $ f (runWriterT m)

instance (Functor m) => Functor (WriterT w m) where
    fmap f = mapWriterT $ fmap $ \ (a, w) -> (f a, w)

instance (Foldable f) => Foldable (WriterT w f) where
    foldMap f = foldMap (f . fst) . runWriterT

instance (Monoid w, Applicative m) => Applicative (WriterT w m) where
    pure a  = WriterT $ pure (a, empty)
    f <*> v = WriterT $ liftA2 k (runWriterT f) (runWriterT v)
      where k (a, w) (b, w') = (a b, w `append` w')

instance (Monoid w, Alternative m) => Alternative (WriterT w m) where
    zilch = WriterT zilch
    m <|> n = WriterT $ runWriterT m <|> runWriterT n

instance (Monoid w, Monad m) => Monad (WriterT w m) where
    return a = writer (a, empty)
    m >>= k  = WriterT $ do
        (a, w)  <- runWriterT m
        (b, w') <- runWriterT (k a)
        return (b, w `append` w')
    fail msg = WriterT $ fail msg

instance (Monoid w, MonadPlus m) => MonadPlus (WriterT w m) where
    mzero       = WriterT mzero
    m `mplus` n = WriterT $ runWriterT m `mplus` runWriterT n

instance (Monoid w) => MonadTrans (WriterT w) where
    lift m = WriterT $ do
        a <- m
        return (a, empty)

instance (Monoid w, MonadIO m) => MonadIO (WriterT w m) where
    liftIO = lift . liftIO

type Writer w = WriterT w Box

-- | Construct a writer computation from a (result, output) pair.
-- (The inverse of 'runWriter'.)
writer :: (Monad m) => (a, w) -> WriterT w m a
writer = WriterT . return

-- | Unwrap a writer computation as a (result, output) pair.
-- (The inverse of 'writer'.)
runWriter :: Writer w a -> (a, w)
runWriter = unbox . runWriterT

-- | Extract the output from a writer computation.
--
-- * @'execWriter' m = 'snd' ('runWriter' m)@
execWriter :: Writer w a -> w
execWriter m = snd (runWriter m)

-- | Map both the return value and output of a computation using
-- the given function.
--
-- * @'runWriter' ('mapWriter' f m) = f ('runWriter' m)@
mapWriter :: ((a, w) -> (b, w')) -> Writer w a -> Writer w' b
mapWriter f = mapWriterT (Box . f . unbox)

{- Already provided by GHC base libraries
instance Monad ((->) r) where
    return = const
    f >>= k = \ r -> k (f r) r
-}
instance (Monoid w) => Monad ((,) w) where
    return x = (empty, x)
    (w, x) >>= f = let (w', y) = f x in (w <> w', y)

{- |
 - MonadPlus
 -
 - Monads that support choice and failure
 -}

class Monad m => MonadPlus m where
    mzero :: m a
    mplus :: m a -> m a -> m a

instance MonadPlus List where
    mzero = nil
    mplus = (++)

instance MonadPlus Opt where
    mzero = none
    o `mplus` ys = maybe o (\xs -> just xs) ys

{- |
 - Comonads
 -
 - Represents a computation in context. The dual concept to monads.
 -
 - Where monads have the following functions defined
 -
 -   return :: a -> m a
 -   (>>=)  :: m a -> (a -> m b) -> m b
 -
 - a comonad has
 -
 -   extract :: w a -> a
 -   extend  :: (w a -> b) -> w a -> w b
 -
 - A comonad is a computation wherein an operation is performed simultaneously
 - in all possible future states.
 -
 - Applications include stream processing, reactive programming, and grid
 - computing.
 -
 -}

class Functor w => Comonad w where
    extract :: w a -> a

    duplicate :: w a -> w (w a)
    duplicate = extend id

    extend :: (w a -> b) -> w a -> w b
    extend f = fmap f . duplicate

class (Functor f, Comonad w) => CofreeComonad f w | w -> f where
    unwrap :: w a -> f (w a)

instance Comonad ((,) e) where
    extract = snd
    duplicate p = (fst p, p)

instance Comonad (Pair e) where
    extract = snd'
    duplicate p = pair (p .! _fst) p

instance (Monoid m) => Comonad ((->) m) where
    extract f = f empty
    extend f wa = \x -> f $ \y -> (wa (append x y))
--    duplicate wa x = wa . append x

-- | The underlying functor of a comonad, here called NuF.
data NuF f a b = a :< f b deriving (Eq, Ord, Show, Read)

headF :: NuF f a b -> a
headF (a :< _) = a

tailF :: NuF f a b -> f b
tailF (_ :< as) = as

instance Functor f => Functor (NuF f a) where
    fmap f (a :< as) = a :< fmap f as

instance Foldable f => Foldable (NuF f a) where
    foldMap f (_ :< as) = foldMap f as

instance Functor f => Bifunctor (NuF f) where
    bimap f g (a :< as) = f a :< fmap g as

instance Foldable f => Bifoldable (NuF f) where
    bifoldMap f g (a :< as) = f a `append` foldMap g as

-- | Comonad transformer: analogous to monad transformer
newtype NuT f w a = NuT {
    runNuT :: w (NuF f a (NuT f w a)) }

-- | The simplified cofree comonad, Nu
type Nu f = NuT f Box

cofree :: NuF f a (Nu f a) -> Nu f a
cofree = NuT . Box
{-# INLINE cofree #-}

runNu :: Nu f a -> NuF f a (Nu f a)
runNu = unbox . runNuT
{-# INLINE runNu #-}

instance (Functor f, Functor w) => Functor (NuT f w) where
    fmap f = NuT . fmap (bimap f (fmap f)) . runNuT

instance (Foldable f, Foldable w) => Foldable (NuT f w) where
    foldMap f = foldMap (bifoldMap f (foldMap f)) . runNuT

instance (Functor f, Comonad w) => Comonad (NuT f w) where
    extract = headF . extract . runNuT
    extend f = NuT . extend (\w -> f (NuT w) :< (extend f
        <$> tailF (extract w))) . runNuT

instance (Functor f, Comonad w) => CofreeComonad f (NuT f w) where
    unwrap = tailF . extract . runNuT

-- | A comonad based on a monoidal functor - such as Alternative - is a monad!
instance (Alternative f, Monad w) => Monad (NuT f w) where
    return = NuT . return . (:< zilch)
    {-# INLINE return #-}
    NuT cx >>= f = NuT $ do
        a :< m <- cx
        b :< n <- runNuT $ f a
        return $ b :< (n <|> fmap (>>= f) m)

-- | Applicative definition, for completeness
instance (Alternative f, Applicative w) => Applicative (NuT f w) where
    pure = NuT . pure . (:< zilch)
    {-# INLINE pure #-}
    wf <*> wa = NuT $ go <$> runNuT wf <*> runNuT wa where
        go (f :< t) a = case bimap f (fmap f) a of
            b :< n -> b :< (n <|> fmap (<*> wa) t)
    {-# INLINE (<*>) #-}

instance Alternative f => MonadTrans (NuT f) where
    lift = NuT . liftM (:< zilch)

-- | Comonad transformer class
class ComonadTrans t where
    lower :: Comonad w => t w a -> w a

instance Functor f => ComonadTrans (NuT f) where
    lower = fmap headF . runNuT

class Comonad w => ComonadApply w where
  (<@>) :: w (a -> b) -> w a -> w b

  (@>) :: w a -> w b -> w b

  (<@) :: w a -> w b -> w a

instance ComonadApply Box where
    (<@>) = (<*>)
    (<@)  = (<*)
    (@>)  = (*>)

{-
 - Comonad utilities
 -}

coiter :: Functor f => (a -> f a) -> a -> Nu f a
coiter psi a = cofree $ a :< (coiter psi <$> psi a)

unfold :: Functor f => (b -> (a, f b)) -> b -> Nu f a
unfold f c = cofree $ case f c of
    (x, d) -> x :< fmap (unfold f) d

coiterT :: (Functor f, Comonad w) => (w a -> f (w a)) -> w a -> NuT f w a
coiterT psi = NuT . extend (\w -> extract w :< fmap (coiterT psi) (psi w))

(=>>) :: Comonad w => w a -> (w a -> b) -> w b
a =>> cb = extend cb a

wfix :: Comonad w => w (w a -> a) -> a
wfix w = extract w (extend wfix w)

cfix :: Comonad w => (w a -> a) -> w a
cfix f = fix (extend f)
{-# INLINE cfix #-}

(=>=) :: Comonad w => (w a -> b) -> (w b -> c) -> w a -> c
f =>= g = g . extend f
{-# INLINE (=>=) #-}

{-
 - Useful comonads
 -}

{- |
 - Store is the dual of State
 -}

data StoreT s w a = StoreT (w (s -> a)) s deriving (Functor)
type Store s = StoreT s Box

store :: (s -> a) -> s -> Store s a
store f s = StoreT (Box f) s

runStore :: Store s a -> (s -> a, s)
runStore (StoreT (Box f) s) = (f, s)

instance Comonad w => Comonad (StoreT s w) where
    duplicate (StoreT wf s) = StoreT (extend StoreT wf) s
    extend f (StoreT wf s) = StoreT (extend (\wf' s' -> f (StoreT wf' s')) wf) s
    extract (StoreT wf s) = extract wf s

instance ComonadTrans (StoreT s) where
    lower (StoreT f s) = fmap ($ s) f

class Comonad w => ComonadStore s w | w -> s where
    pos :: w a -> s
    peek :: s -> w a -> a
    peeks :: (s -> s) -> w a -> a
    seek :: s -> w a -> w a
    seeks :: (s -> s) -> w a -> w a

instance Comonad w => ComonadStore s (StoreT s w) where
    pos (StoreT _ s) = s
    peek s (StoreT g _) = extract g s
    peeks f (StoreT g s) = extract g (f s)
    seek s (StoreT f _) = StoreT f s
    seeks f (StoreT g s) = StoreT g (f s)

experiment :: (Comonad w, Functor f) => (s -> f s) -> StoreT s w a -> f a
experiment f (StoreT wf s) = extract wf <$> f s

{- |
 - Env is the dual of Reader
 -}
data EnvT e w a = EnvT e (w a)
type Env e = EnvT e Box

-- | Create an Env using an environment and a value
env :: e -> a -> Env e a
env e a = EnvT e (Box a)

runEnv :: Env e a -> (e, a)
runEnv (EnvT e (Box a)) = (e, a)

runEnvT :: EnvT e w a -> (e, w a)
runEnvT (EnvT e wa) = (e, wa)

instance Functor w => Functor (EnvT e w) where
  fmap g (EnvT e wa) = EnvT e (fmap g wa)

instance Comonad w => Comonad (EnvT e w) where
  duplicate (EnvT e wa) = EnvT e (extend (EnvT e) wa)
  extract (EnvT _ wa) = extract wa

instance ComonadTrans (EnvT e) where
  lower (EnvT _ wa) = wa

lowerEnvT :: EnvT e w a -> w a
lowerEnvT (EnvT _ wa) = wa

instance Foldable w => Foldable (EnvT e w) where
    foldMap f (EnvT _ w) = foldMap f w

class Comonad w => ComonadEnv e w | w -> e where
    query :: w a -> e

instance Comonad w => ComonadEnv e (EnvT e w) where
    query (EnvT e _) = e

queries :: (e -> f) -> EnvT e w a -> f
queries f (EnvT e _) = f e

modify :: (e -> e') -> EnvT e w a -> EnvT e' w a
modify f (EnvT e wa) = EnvT (f e) wa

{-
 - Example Env usage: searching a binary tree
 -
 - source: https://gist.github.com/ruicc/5435acba4be89aed7d6a
 -}
data Bin a = Node a (Bin a) (Bin a) | Leaf a

bintree_ex :: Bin Int
bintree_ex = Node 5 (Leaf 3) (Node 8 (Leaf 7) (Leaf 10))

searchTree :: Int -> Bin Int-> Opt Int
searchTree n t = let w = env t () in
    extract $
    w =>>
    query =>> \wbt -> case extract wbt of
        Leaf a
            | a == n -> just a
            | otherwise -> none
        Node a l r
            | n == a -> just a
            | n > a -> searchTree n r
            | otherwise -> searchTree n l

test_search_binary_tree :: Opt Int
test_search_binary_tree = searchTree 8 bintree_ex

-- | Traced
newtype TracedT m w a = TracedT { runTracedT :: w (m -> a) }
type Traced m = TracedT m Box

traced :: (m -> a) -> Traced m a
traced f = TracedT (Box f)

runTraced :: Traced m a -> m -> a
runTraced (TracedT (Box f)) = f

instance Functor w => Functor (TracedT m w) where
    fmap g = TracedT . fmap (g .) . runTracedT

instance (Comonad w, Monoid m) => Comonad (TracedT m w) where
    extend f = TracedT . extend
        (\wf m -> f (TracedT (fmap (. append m) wf))) . runTracedT
    extract (TracedT wf) = extract wf empty

instance (Monoid m) => ComonadTrans (TracedT m) where
    lower = fmap ($ empty) . runTracedT

class Comonad w => ComonadTraced m w | w -> m where
    trace :: m -> w a -> a

instance (Comonad w, Monoid m) => ComonadTraced m (TracedT m w) where
    trace m (TracedT wf) = extract wf m

listen :: Functor w => TracedT m w a -> TracedT m w (a, m)
listen = TracedT . fmap (\f m -> (f m, m)) . runTracedT

listens :: Functor w => (m -> b) -> TracedT m w a -> TracedT m w (a, b)
listens g = TracedT . fmap (\f m -> (f m, g m)) . runTracedT

censor :: Functor w => (m -> m) -> TracedT m w a -> TracedT m w a
censor g = TracedT . fmap (. g) . runTracedT

{- |
 - Lenses
 -
 - A lens allows you to *focus* on an element nested inside just data structure
 - for O(1) accesses and updates. The concept is very flexible and powerful,
 - such that lenses may be mechanically generated for most any type.
 -
 - Lenses also give the ability to chain modifications to an object one after
 - the other.
 -
 - One interesting application is in stream processors: receive an object from
 - upstream and yield a modified version in one expression.
 -}
type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t

over :: Lens s t a b -> (a -> b) -> s -> t
over l f s = unbox (l (Box . f) s)

view :: Lens s t a b -> s -> a
view l s = getConst (l Const s)

(.!) :: s -> Lens s t a b -> a
s .! l = view l s
infixr 8 .!

(#) :: a -> (a -> b) -> b
x # y = y $ x
infixl 1 #

(.=) :: Lens s t a b -> b -> s -> t
l .= f = unbox . l (Box . (\_ -> f))
infixr 4 .=

(.$) :: Lens s t a b -> (a -> b) -> s -> t
l .$ f = unbox . l (Box . f)
infixr 4 .$

-- | Instances for types we've seen so far
hd :: Lens [a] [a] a a
hd f (a:as) = fmap go (f a) where
    go a' = a': as

hd' :: Lens (List a) (List a) a a
hd' f xs = let (my, ys) = split xs
               go a'    = cons a' ys
            in maybe my (\x -> fmap go (f x)) diverge

_1 :: Lens (a, b) (a', b) a a'
_1 f (a,b) = fmap (\x -> (x, b)) (f a)

_2 :: Lens (a, b) (a, b') b b'
_2 f (a,b) = fmap (\x -> (a, x)) (f b)

_fst :: Lens (Pair a b) (Pair a' b) a a'
_fst f p = unpair p (\a b -> fmap (\x -> (pair x b)) (f a))

_snd :: Lens (Pair a b) (Pair a b') b b'
_snd f p = unpair p (\a b -> fmap (\x -> (pair a x)) (f b))

stored f aStore =
    let (ev, c) = runStore aStore
    in fmap (\c' -> store ev c') (f c)

-- | Opens up a box
deref :: Lens (Box a) (Box b) a b
deref f (Box v) = fmap Box (f v)


{- |
 - Task
 -
 - Task is the combination of two different ideas: sources and sinks. The free
 - monad transformer of `((,) t)` for any value of type `t` allows a monadic
 - computation to suspend itself and yield some intermediate value.
 -
 -     type Source t = MuT ((,) t)
 -
 - Similarly, the free monad transformer of `((->) t)` for any value of type
 - `t` allows a monadic computation to suspend and wait for more input.
 -
 -     type Sink t = MuT ((->) t)
 -
 - A `yield` operator like in Python can then be defined as follows:
 -
 -     yield :: Monad m => t -> Source t m ()
 -     yield x = liftF (x, ())
 -
 - And an `await` operator can be defined like so:
 -
 -     await :: Monad m => Sink t m t
 -     await = liftF id
 -
 - Thus, a computation which can suspend execution to yield or demand
 - intermediate results is a task.
 -
 - Combined into one type, they become a tool for composing sophisticated
 - stream processing pipelines.
 -}

data TaskF a b k
    = Await (a -> k)
    | Yield b k
    deriving (Functor)

type Task a b     = MuT (TaskF a b)
type Source b m r = Task ()  b m r
type Sink   a m r = Task a  () m r
type Result   m r = Task () () m r

-- Task building blocks

yield :: FreeMonad (TaskF a b) m => b -> m ()
yield x = liftF $ Yield x ()

await :: FreeMonad (TaskF a b) m => m a
await   = liftF $ Await id

liftT = lift . runMuT

cat :: FreeMonad (TaskF a a) m => m b
cat = forever $ await >>= yield

run = runMuT

-- Task composition functions

-- | Connect a task to a continuation yielding another task
(>-) :: Monad m
     => Task a b m r
     -> (b -> Task b c m r)
     -> Task a c m r
p >- f = liftT p >>= go where
    go (Pure x) = return x
    go (Wrap (Await f')) = wrap $ Await (\a -> (f' a) >- f)
    go (Wrap (Yield v k)) = k >< f v

-- | Compose two tasks in a pull-based stream
(><) :: Monad m
     => Task a b m r
     -> Task b c m r
     -> Task a c m r
a >< b = liftT b >>= go where
    go (Pure x) = return x
    go (Wrap (Yield v k)) = wrap $ Yield v $ liftT k >>= go
    go (Wrap (Await f)) = a >- f

infixl 3 ><

instance (Monad m, Monoid r) => Monoid (Task a b m r) where
    empty = return empty
    append p1 p2 = liftT p1 >>= go where
        go (Wrap (Await f))   = wrap $ Await (\a -> liftT (f a) >>= go)
        go (Wrap (Yield v k)) = wrap $ Yield v $ liftT k >>= go
        go (Pure r)           = fmap (append r) p2

for :: Monad m
    => Task a b m r
    -> (b -> Task a c m s)
    -> Task a c m r
for src body = liftT src >>= go where
    go (Wrap (Await f)) = wrap $ Await (\x -> liftT (f x) >>= go)
    go (Wrap (Yield v k)) = do
        body v
        liftT k >>= go
    go (Pure x) = return x

each :: (Monad m, Functor t, Foldable t) => t b -> Task a b m ()
each = mapM_ yield

each' xs = (each xs >< taskmap just) >> yield none

next :: Monad m => Source a m r -> m (Either r (a, Source a m r))
next src = runMuT src >>= go where
    go (Wrap (Yield v k)) = return (right (v, k))
    go (Pure r)           = return (left r)

tasktake :: Monad m => Int -> Task a a m ()
tasktake n = replicateM_ n $ do
    x <- await
    yield x

-- | Strict left fold of a Source
reduce :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Source a m () -> m b
reduce step begin done p0 = runMuT p0 >>= \p' -> loop p' begin where
    loop p x = case p of
        Wrap (Yield v k) -> runMuT k >>= \k' -> loop k' $! step x v
        Pure _           -> return (done x)

taskmap :: (Monad m) => (a -> b) -> Task a b m r
taskmap f = for cat $ \x -> yield (f x)

-- | Convert a Foldable (such as a Stream) into a Source
toSource :: (Foldable t, Functor t, Monad m)
         => t c
         -> Task a c m ()
toSource strm = for (each strm) yield

-- Examples; every iteratee / pipes library does this already

insult :: FreeMonad (TaskF String String) m => m b
insult = forever $ do
    it <- await
    yield $ it P.++ " sucks"

backpedal :: FreeMonad (TaskF String String) m => m b
backpedal = forever $ do
    it <- await
    yield $ it P.++ " (but not really)"

print = forever $ do
    thing <- await
    liftIO . putStrLn . show $ thing

prompt :: Source String IO ()
prompt = forever $ do
    liftIO $ putStr "> "
    str <- liftIO getLine
    yield str

pipeline_2 = runMuT $ (forever prompt) >< insult >< backpedal >< print

instance (MonadIO m) => MonadIO (Task a b m) where
    liftIO = lift . liftIO

instance (MonadState s m) => MonadState s (Task a b m) where
    get = lift get
    put = lift . put

print' :: (Show a) => Task a a IO ()
print' = forever $ do
    it <- await
    liftIO . putStrLn . show $ it
    yield it

withState :: Monad m => b -> MuT f (StateT b m) a -> m b
withState s t = (evalStateT s (runMuT t)) >>= return . snd

(+>) = withState
infixl 4 +>

pipeline_3 = (each (cons 1 (cons 2 (cons 3 nil)))) >< print' # reduce (+) 0 id

instance Monad m => Show (Task a b m r) where
    show _ = "<Task>"

-- Concurrency primitives for tasks

broadcast :: (Functor f, Monad m )
          =>     Task a b m r
          -> f ( Task b c m r )
          -> f ( Task a c m r )
broadcast src tsks = fmap (\t -> src >< t) tsks

(*<) :: (Functor f, Monad m)
     =>     Task a b m r
     -> f ( Task b c m r )
     -> f ( Task a c m r )
(*<) = broadcast
infixr 4 *<

merge :: ( Functor t, Foldable t, Monad m )
      => t (Task a c m s )
      -> Task a c m ()
merge tasks = for (each tasks) $ \t -> for t yield

(>*) :: ( Functor t, Foldable t, Monad m )
     => t ( Task a b m s )
     ->     Task b c m ()
     ->     Task a c m ()
tasks >* k = (merge tasks) >< k
infixr 0 >*

{-|
  Arithmetic parsing example
 -}

data Arithmetic
    = Value Int
    | Add Arithmetic Arithmetic
    | Mul Arithmetic Arithmetic

evalArith :: Arithmetic -> Int
evalArith (Value v) = v
evalArith (Add l r) = (evalArith l) + (evalArith r)
evalArith (Mul l r) = (evalArith l) * (evalArith r)

type Token = Opt String

tokenize :: Monad m => Task (Opt Char) Token m ()
tokenize = loop "" where
    loop acc = do
        token <- await
        maybe token
            (\t -> case t of
                ' ' -> (yield (just acc)) >> loop ""
                _   -> loop $ acc P.++ [t])
            (yield (just acc) >> yield none)

parseArith :: Monad m => Source Token m () -> m Int
parseArith = reduce step nil (evalArith . car) where
    step stack token = maybe token
        (\t -> case t of
            "+" -> performOp (Add) stack
            "*" -> performOp (Mul) stack
            ""  -> stack
            _   -> cons (Value (read t)) stack)
        stack

    performOp op stack =
        if length stack < 2
            then stack
            else
                let (m1, r1) = split stack
                    (m2, rst ) = split r1
                in maybe m1
                     (\r -> maybe m2
                     (\l -> cons (op l r) rst)
                     diverge) diverge

compute :: Monad m => String -> m Int
compute x = parseArith $ each' x >< tokenize

calculator :: IO ()
calculator = forever $ do
    putStr $ "> "
    expr <- getLine
    result <- compute expr
    putStrLn . show $ result

-- | arbitrary data type
data Person = Person
    { _name :: String
    , _age  :: Int
    } deriving (Show)

-- lenses
name f (Person n a) = fmap (\n' -> Person n' a) (f n)
age  f (Person n a) = fmap (\a' -> Person n a') (f a)

g = Person "gatlin" 26
w = Person "washington" 283

people = cons g $
         cons w nil

birthday :: Monad m => Task Person Person m ()
birthday = taskmap (# age .$ (+1))

people_example = run $ each people
                    >< birthday
                    >< taskmap show
                    >< print

{- |
 - List Zippers and Plane Zippers
 -}

-- | Simple case: Löb's theorem applied to functors
loeb :: (Functor f) => f (f a -> a) -> f a
loeb x = fmap (\a -> a (loeb x)) x

test_loeb_1 = cons length (cons car nil)
-- loeb test_loeb_1 == [ 2 , 2 ]

test_loeb_2 = cons (const 0) $
              cons (\xs -> (car xs) + 5) $
              cons (const 7) $
              cons (\xs -> (car (cdr xs)) * 2) $
              cons (\xs -> (car xs) - 3)
              nil

-- | A bidirectional stream with a focus
data Cursor a = Cursor
    { _viewL :: List a
    , _focus :: a
    , _viewR :: List a
    } deriving (Functor, Show)

focus g (Cursor l f r) = fmap (\f' -> Cursor l f' r) (g f)

seed :: (c -> (a, c))
     -> (c -> a)
     -> (c -> (a, c))
     -> c
     -> Cursor a
seed prev center next =
    Cursor <$> unfoldr prev <*> center <*> unfoldr next

iter :: (a -> a)
     -> (a -> a)
     -> a
     -> Cursor a
iter prev next =
    seed (dup . prev) id (dup . next)
    where dup a = (a, a)

moveL :: Cursor a -> Cursor a
moveL (Cursor lxs c rxs) =
    let lh = car lxs
        lt = cdr lxs
    in  Cursor lt lh (cons c rxs)

moveR :: Cursor a -> Cursor a
moveR (Cursor lxs c rxs) =
    let rh = car rxs
        rt = cdr rxs
    in  Cursor (cons c lxs) rh rt

instance Comonad Cursor where
    extract (Cursor _ c _) = c
    duplicate = iter moveL moveR

-- | Löb's theorem but for comonads
evaluate :: (Comonad w) => w (w a -> a) -> w a
evaluate = extend wfix -- a very happy accident after playing type tetris

cursor_1 :: Cursor (Cursor Int -> Int)
cursor_1 = let n = const 0 in Cursor (repeat n) n (repeat n)

cursor_2 = cursor_1 #
           moveL #
           insert (\t -> 2 + (t # moveR # extract)) #
           moveR

cursor_2_ev = evaluate cursor_2

slice :: Int -> Cursor a -> List a
slice n (Cursor ls x rs) = (single x) ++ take n rs

-- cursor_2 # moveL # extract => 2

data Plane a = Plane (Cursor (Cursor a))

up :: Plane a -> Plane a
up (Plane p) = Plane (moveL p)

down :: Plane a -> Plane a
down (Plane p) = Plane (moveR p)

moveLeft :: Plane a -> Plane a
moveLeft (Plane p) = Plane (fmap moveL p)

moveRight :: Plane a -> Plane a
moveRight (Plane p) = Plane (fmap moveR p)

class Insertable i where
    insert :: a -> i a -> i a

instance Insertable Cursor where
    insert x (Cursor l _ r) = Cursor l x r

instance Insertable Plane where
    insert x (Plane p) =
        Plane $ insert newLine p where
            newLine = insert x oldLine
            oldLine = extract p

instance Functor Plane where
    fmap f (Plane p) = Plane (fmap (fmap f) p)

horizontal :: Plane a -> Cursor (Plane a)
horizontal = iter moveLeft moveRight

vertical :: Plane a -> Cursor (Plane a)
vertical = iter up down

instance Comonad Plane where
    extract (Plane p) = extract $ extract p
    duplicate z =
        Plane $ fmap horizontal $ vertical z

makePlane :: a -> List (List a) -> Plane a
makePlane def grid = Plane $ Cursor (repeat fz) fz rs where
    rs = (map line grid) ++ repeat fz
    dl = repeat def
    fz = Cursor dl def dl
    line l = Cursor dl def (l ++ dl)

sheet1 :: Plane (Plane Int -> Int)
sheet1 = makePlane (const 0)
    (cons (
        cons (\c -> 15 + (c # moveLeft # extract)) (
        cons (\c -> 10 + (c # moveLeft # extract)) nil ) )
    (cons (
        cons (\c -> 2 * (c # up # extract)) nil )
    nil )

    )

neighborhood :: Int -> Plane a -> List (List a)
neighborhood n (Plane cs) = slice n $ fmap (slice n) cs

viewPlane pln = run $ each (pln # evaluate # neighborhood 5)
                   >< taskmap show
                   >< print

--

type StreamT = NuT Opt
type Stream = StreamT Box

nums :: Stream Int
nums = coiter (\x -> just (x+1)) 0

doubleIt :: Int -> Stream Int
doubleIt n = return $ n * 2

squareIt :: Int -> Stream Int
squareIt n = return $ n * n

toList :: (Foldable t, Functor t, Monad m)
       => Int
       -> t a
       -> m (List a)
toList n ss = reduce (++) nil id $ each ss >< tasktake n >< taskmap single

-- | nice, informatively-named alias
stream :: Monad m => Stream a -> Source a m ()
stream = toSource

stream_1 = nums >>= doubleIt >=> squareIt # toList 10

-- Conway game!

-- | Extract the neighbors of a Plane's focus (a sub-plane)
neighbors :: List (Plane a -> Plane a)
neighbors =
    horiz ++ vert ++ liftM2 (.) horiz vert where
        horiz   = cons moveLeft (cons moveRight nil)
        vert    = cons up (cons down nil)

-- | Count how many neighbors are alive
aliveNeighbors :: Plane Bool -> Int
aliveNeighbors z =
    card $ map (\ dir -> extract $ dir z) neighbors

-- | Cardinality, ie number of True values in a list of booleans
card :: List Bool -> Int
card = length . filter (== True)

-- | A particular Conway rule
rule :: Plane Bool -> Bool
rule z =
    case aliveNeighbors z of
        2   -> extract z
        3   -> True
        _   -> False

evolve :: Plane Bool -> Plane Bool
evolve = extend rule

-- | Display helper for Cursors
dispLine :: Cursor Bool -> String
dispLine z =
    toPrimList $ fmap dispC $ slice 6 z where
        dispC   True    = '*'
        dispC   False   = ' '

-- | Display helper for Planes
disp :: Plane Bool -> String
disp (Plane z) =
    unlines $ fmap dispLine $ slice 6 z

-- | Initial conditions to create a "glider"
glider :: Plane Bool
glider = makePlane f $
    cons ((cons f (cons t (cons f nil)))) $
    cons ((cons f (cons f (cons t nil)))) $
    cons ((cons t (cons t (cons t nil)))) nil
    where f = False
          t = True

-- | Now we will stream the iterations of the glider, potentially infinitely
glider_stream :: Stream (Plane Bool)
glider_stream = unfold (\g -> (g, just (evolve g))) glider

print_sink f n = loop (n :: Int) where
    loop 0 = return ()
    loop n = do
        it <- await
        liftIO . putStrLn $ f it
        loop (n - 1)

{- |
 - Using this technique I have constructed a corecursive computation which I
 - can evaluate an arbitrary number of times, receiving a new result each time.
 -}
glider_task n = run $ stream glider_stream >< print_sink disp n

{- |
 - Similarly I can take a cursor computation and loop over it endlessly,
 - shifting the result value back to the input slot before each iteration.
 -}
eval cursor = stream $ unfold ( \c ->
    let r = evaluate c
        c' = insert (const (extract (moveL r))) cursor
    in  (r, just c') ) cursor

cursor_loop n = run $ (eval cursor_2) >< print_sink (show . extract) n