{-# LANGUAGE FlexibleInstances #-}

module FreeMonad where

import Control.Monad


--
-- * Free Monad (adapted from Control.Monad.Free)
--
 
-- | The free monad for a functor f.
data Free f a
   = Pure a
   | Free (f (Free f a))

instance Functor f => Monad (Free f) where
  return = Pure
  Pure x >>= f = f x
  Free m >>= f = Free (fmap (>>= f) m)

-- | Lift a value of the functor into the monad.
liftF :: Functor f => f a -> Free f a
liftF = Free . fmap Pure


-- Question: What do we mean by "free"?
-- 
--  * Given any Functor f, Free f is a monad with no extra work on our part
--    ... it was free!
--
--  * It just builds a structure that satisfies the monad laws, but doesn't
--    "do" anything.
--
--     * The structure is a tree of alternating Frees and constructors from
--       functor type, with Pure at the leaves.
--
--     * We can use this structure to *later* recover any more specific monad.
--
--     * "Free" in the algebraic sense: https://en.wikipedia.org/wiki/Free_object
--
--  * Analogy: lists are the free monoid; given any type a, [a] is a monoid;
--    we can recover any more specific monoid on type a from [a].


--
-- * Illustrations
--

-- ** Free list monad

listMonad :: [Int]
listMonad = do
    x <- [1..3]
    y <- [10,20]
    return (x+y)

listFree :: Free [] Int
listFree = do
    x <- liftF [1..3]
    y <- liftF [10,20]
    return (x+y)

runList :: Free [] a -> [a]
runList (Pure a) = [a]
runList (Free l) = concatMap runList l


-- ** Free binary tree moand

data Tree a
   = Leaf a
   | Node (Tree a) (Tree a)
  deriving (Eq,Show)

instance Functor Tree where
  fmap f (Leaf a)   = Leaf (f a)
  fmap f (Node l r) = Node (fmap f l) (fmap f r)

instance Monad Tree where
  return         = Leaf
  Leaf a   >>= f = f a
  Node l r >>= f = Node (l >>= f) (r >>= f)

treeMonad :: Tree Int
treeMonad = do
    x <- Node (Leaf "hi") (Leaf "bye")
    y <- Node (Leaf 10) (Leaf 20)
    return (length x + y)

treeFree :: Free Tree Int
treeFree = do
    x <- liftF $ Node (Leaf "hi") (Leaf "bye")
    y <- liftF $ Node (Leaf 10) (Leaf 20)
    return (length x + y)

runTree :: Free Tree a -> Tree a
runTree (Pure a) = return a
runTree (Free t) = t >>= runTree


--
-- * A more "useful" example
--

-- | A data type that describes two primitive operations and that supports
--   a simple Functor instance.
data GetPutF s t
   = Get (s -> t)  -- ^ Get a value of type 's', use it to compute 't'.
   | Put s t       -- ^ Put a value of type 's', then compute 't'.

instance Functor (GetPutF s) where
  fmap f (Get g)   = Get (f . g)
  fmap f (Put s x) = Put s (f x)

-- | The free get-put monad.
type GetPut s = Free (GetPutF s)

-- | Generic "get" primitive.
get :: GetPut s s
get = liftF (Get id)

-- | Generic "put" primitive.
put :: s -> GetPut s ()
put s = liftF (Put s ())


-- ** Simple example programs

-- | Put the value 0.
reset :: GetPut Int ()
reset = put 0

-- | Get a value and put its successor.
inc :: GetPut Int ()
inc = do
    i <- get
    put (i+1)

-- | A simulated counter application.
counter :: GetPut Int ()
counter = do
    sequence_ (replicate 10 inc)
    reset
    sequence_ (replicate 15 inc)

-- | Put the same argument a given number of times in a row.
puts :: Int -> s -> GetPut s ()
puts n s = sequence_ (replicate n (put s))

-- | Get a value, check to see if its the terminating value,
--   if not, put it and repeat.
echo :: Eq s => s -> GetPut s ()
echo end = do
    s <- get
    if s == end
      then return ()
      else put s >> echo end


-- ** Run as more specific "monads"

-- | Run as a state monad.
runState :: GetPut s a -> s -> (a,s)
runState (Pure a)         s = (a, s)
runState (Free (Get f))   s = runState (f s) s
runState (Free (Put s x)) _ = runState x s

-- | Run as a state monad but print each use of a primitive operation.
runStateIO :: Show s => GetPut s a -> s -> IO (a,s)
runStateIO (Pure a)         s = return (a, s)
runStateIO (Free (Get f))   s = putStrLn ("get " ++ show s) >> runStateIO (f s) s
runStateIO (Free (Put s x)) _ = putStrLn ("put " ++ show s) >> runStateIO x s

-- | Run as an interactive console monad.
runConsoleIO :: GetPut String a -> IO a
runConsoleIO (Pure a)         = return a
runConsoleIO (Free (Get f))   = getLine >>= runConsoleIO . f
runConsoleIO (Free (Put s x)) = putStrLn s >> runConsoleIO x

-- | Run as a simulated console by reading/writing inputs/outputs to lists.
runConsoleList :: GetPut String a -> [String] -> (a, [String])
runConsoleList c ins = (r, reverse outs)
  where
    (r, (_, outs)) = go c (ins,[])
    go (Pure a) s                = (a,s)
    go (Free (Get f)) (i:is, os) = go (f i) (is, os)
    go (Free (Put s x)) (is, os) = go x (is, s:os)


-- ** Manipulate the monadic structure directly!

-- | Optimize a monadic computation by eliminating "redundant" puts.
optimize :: GetPut s a -> GetPut s a
optimize (Free (Put _ (Free (Put s x)))) = optimize (Free (Put s x))
optimize (Pure a) = Pure a
optimize (Free x) = Free (fmap optimize x)

-- | Interleave to monadic operations.
interleave :: Functor f => Free f a -> Free f a -> Free f a
interleave (Pure _)  c         = c
interleave c         (Pure _)  = c
interleave (Free m1) (Free m2) = do
    next1 <- liftF m1
    next2 <- liftF m2
    interleave next1 next2


-- ** Also: run as monads we couldn't implement otherwise!

-- Example: Set monad


--
-- * Boring instances
--

instance Functor f => Functor (Free f) where
  fmap = liftM

instance Functor f => Applicative (Free f) where
  pure  = return
  (<*>) = ap

instance Show a => Show (Free [] a) where
  show (Pure a) = "Pure " ++ show a
  show (Free l) = "Free " ++ show l

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

instance Show a => Show (Free Tree a) where
  show (Pure a) = "(Pure (" ++ show a ++ "))"
  show (Free t) = "(Free (" ++ show t ++ "))"