-- A concrete use case for the type which is to '(->)' as 'Day' is to '(,)'.
-- I call it "FunDay", but I don't know what its proper name is. I've been
-- trying to find a use for 'FunDay', and I think I've found a pretty neat one.

{-# LANGUAGE FlexibleContexts, FlexibleInstances, PolyKinds, RankNTypes, TypeSynonymInstances #-}
module Main where
import Test.DocTest

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer


-- Suppose you have a computation which uses many different effects.
--
-- Here I only use three, for illustration purposes, but obviously this list can
-- easily get much longer.
myRWST :: ( MonadReader String m
          , MonadWriter String m
          , MonadState String m
          )
       => m String
myRWST = do
  r <- ask
  tell "w"
  modify (++ "'")
  pure (r ++ "esult")

-- The normal way to discharge those effects is to call a bunch of 'run'
-- functions one after the other. As the number of constraints becomes smaller
-- and smaller, the return type becomes larger and larger. Once you're done
-- discharging everything, the result probably doesn't have the shape you want
-- yet, so you need to pattern-match on it and rearrange it.
--
-- Can you figure out which of 's1', 's2', and 's3' correspond to the new state,
-- the value accumulated by 'Writer', and the result returned by 'myRWST'?
runMyRWST :: Monad m
          => m String
runMyRWST = do
  ((s1, s2), s3) <- flip runStateT "s" . runWriterT . flip runReaderT "r" $ myRWST
  pure (s1 ++ ", " ++ s2 ++ ", " ++ s3)

-- |
-- The answer is below:
--
-- ...
--
--
--
--
--
--
--
--
--
--
--
--
--
--
--
--
--
-- >>> runMyRWST
-- "result, w, s'"
--
-- So 's1' is the result, 's2' is the value accumulated by 'Writer', and s3 is
-- the new state. Did you get it right? How long did it take you to figure it
-- out?
--
-- I think there is room for improvement here. How about Applicative syntax?
-- Conceptually, each component of our expression produces a different piece of
-- the tuple: @flip runStateT "s"@ produces the updated state, 'runWriterT'
-- produces the accumulated value, and 'myRWST' produces the result. So I would
-- like to write something like this:
--
-- > (\() w s result -> result ++ ", " ++ w ++ ", " ++ s)
-- >   <$> flip runReaderT "r"
-- >   <*> runWriterT
-- >   <*> flip runStateT "s"
-- >   <*> myRWST
--
-- I did manage to get something like that syntax, with one small differences:
-- I am using an indexed Applicative instead of the ordinary Applicative. It
-- looks like this:
--
-- >>> :{
-- runFunDay $ (\() w s result -> result ++ ", " ++ w ++ ", " ++ s)
--        <$$> funReaderT "r"
--        <**> funWriterT
--        <**> funStateT "s"
--        <**> funday1 myRWST
-- :}
-- "result, w, s'"
--
-- And the magic behind it is our new friend, 'FunDay'!

newtype FunDay f g a = FunDay
  { unFunDay :: forall x y. (a -> x -> y) -> f x -> g y }

-- or equivalently

newtype HomDay f g a = HomDay
  { unHomDay :: forall r. f (a -> r) -> g r }

-- My intuition for @FunDay f g@ is that it is a way to transform an 'f'-based
-- computation into a 'g'-based computation. Since 'runWriterT' and friends are
-- converting a @WriterT w m@ computation into an 'm' computation, this means
-- 'FunDay' is a good match for our task!
--
-- But what about the 'a'? To understand that, let's take a step back and look
-- at a version of 'HomDay' which is only indexed by a single 'f' rather than by
-- both an 'f' and a 'g':

type FunDay1 f = FunDay f f
type HomDay1 f = HomDay f f

-- 'HomDay1' is very similar to 'DList', 'Yoneda' and 'Codensity':

data DList a = DList
  { unDList :: [a] -> [a] }

data Yoneda f a = Yoneda
  { unYoneda :: forall r. (a -> r) -> f r }

data Codensity f a = Codensity
  { unCodensity :: forall r. (a -> f r) -> f r }

-- It's the same trick we have seen over and over: a difference list is holding
-- a hidden @[a]@, and in order to extract that list we have to give it a
-- suffix, typically '[]'. And 'Yoneda' and 'Codensity' are both holding a
-- hidden @f a@ computation, and in order to extract that computation, we need
-- to give them a post-computation, typically 'id' or 'return'. 'Yoneda' is
-- based on 'fmap', so that post-computation is a pure function @a -> r@, while
-- 'Codensity' is based on '(>>=)', so that post-computation is an @a -> f r@.
-- 'HomDay1' fills the gap in between: it is based on '(<*>)', so its
-- post-computation is an @f (a -> r)@, or equivalently, both an @f x@ and some
-- pure function to combine that 'x' with the 'a' which the hidden @f a@
-- computation has computed so far.


-------------------
-- DETOUR BEGINS --
-------------------

-- As a slight detour, this means that 'HomDay1' has the same performance
-- benefit we typically get from that trick.
--
-- Suppose we have a binary method like '(<>)', '(<*>)', or '(>>=)') whose cost
-- is proportional to the size of its left argument, because it needs to
-- traverse it in order to reach the leaves. If have a left-associative chain of
-- method calls, in which the output of each call is used as the left argument
-- to the next call, we'll get accidentally-quadratic performance. Thankfully,
-- those methods obey laws which allow us to rewrite such a chain in a
-- right-associative way, in which the output of each call is instead used as
-- the right argument to the next call.
--
-- > (((([] <> [1]) <> [2]) <> [3]) <> [4]) <> [5]
-- > [1] <> ([2] <> ([3] <> ([4] <> ([5] <> []))))
--
-- > (((\x y z -> x + y + z) <$> liftF [1]) <*> liftF [2]) <*> liftF [3]
-- > (\(x,(y,z)) -> x + y + z) <$> (liftF [(1,)] <*> (liftF [(2,)] <*> liftF [3]))
--
-- > ((liftF [()] >> liftF [()]) >> liftF [()]) >> liftF [()]
-- > liftF [()] >> (liftF [()] >> (liftF [()] >> liftF [()]))
--
-- We can get a performance boost by rewriting a left-associative chain of calls
-- into a right-associative chain of calls, but having to remember to do so is a
-- burden on the programmer, which makes these methods a poor API. Difference
-- lists, 'Yoneda' and 'Codensity' provide a better API because they relieve the
-- programmer from having to think about such low-level implementation details:
-- their API is just as fast in the left- and right-associative styles, because
-- they internally rewrite the chain in the right-associative style:

dlist :: [a] -> DList a
dlist xs = DList (\nil -> xs <> nil)

-- > (dlist [x] <> dlist [y]) <> dlist [z]
--
-- becomes
--
-- > DList (\nil -> [x] <> ([y] <> ([z] <> nil)))
instance Semigroup (DList a) where
  DList f <> DList g = DList (\nil -> f (g nil))


funday1 :: Applicative m
        => m a -> FunDay1 m a
funday1 ma = FunDay (\l mx -> l <$> ma <*> mx)

-- > ((ff <$> funday1 ma) <*> funday1 mb) <*> funday1 mc
--
-- becomes
--
-- > FunDay $ \l mx
-- >       -> (\(a,(b,(c,x))) -> l (f a b c) x)
-- >      <$> f (g (h mx))
-- >   where
-- >     f mx = (,) <$> ma <*> mx
-- >     g mx = (,) <$> mb <*> mx
-- >     h mx = (,) <$> mc <*> mx
instance Functor m => Applicative (FunDay1 m) where
  pure a = FunDay $ \l mx -> l a <$> mx
  fundayF <*> fundayA = FunDay $ \l mx
                     -> (\(a2b,(a,x)) -> l (a2b a) x) <$> f (g mx)
    where f = unFunDay fundayF (,)
          g = unFunDay fundayA (,)


codensity :: Monad m
          => m a -> Codensity m a
codensity ma = Codensity $ \cc -> do
  a <- ma
  cc a

-- > (codensity ma >> codensity mb) >> codensity mc
--
-- becomes
--
-- > Codensity $ \cc -> do
-- >   ma >> (mb >> (mc >>= cc))
instance Monad (Codensity m) where
  codensityA >>= f = Codensity $ \cc -> do
    unCodensity codensityA $ \a -> do
      unCodensity (f a) $ \b -> do
        cc b

-----------------
-- DETOUR ENDS --
-----------------


-- Okay, so 'FunDay1' is holding a hidden @f a@ computation, and that
-- computation is waiting for a post-computation @f x@ and a pure function
-- @a -> x -> r@ combining the two results. One last bit of complexity is that
-- unlike 'FunDay1', 'FunDay' is indexed by both 'f' and 'g'. What does that
-- change?
--
-- Well, 'FunDay' is holding a hidden @g a@ computation, and that computation is
-- waiting for a post-computation @f x@ and a pure function @a -> x -> r@
-- combining the two results. That's how @FunDay f g@ is transforming an
-- 'f'-based computation into a 'g'-based computation: it already has a
-- 'g'-based computation, but that computation is waiting for an 'f'-based
-- post-computation, so it "converts" the 'f'-based computation into a 'g'-based
-- computation by running the 'f'-based computation at the end of the hidden
-- 'g'-based computation.
--
-- How is it possible for a 'g'-based computation to be waiting for an 'f'-based
-- computation rather than another 'g'-based computation? Simple: the hidden
-- computation is not simply planning to run one computation after the other, it
-- is planning to interpret the 'f'-based computation into 'g'! This is exactly
-- what 'runWriterT' and friends are doing, interpreting a @WriterT w m@
-- computation into 'm'.

-- @funReaderT r@ has a hidden @m ()@ computation which is waiting for a
-- @ReaderT r m x@ post-computation. So it transforms a @ReaderT r m@
-- computation into an @m@ computation and returns a '()'.
funReaderT :: Monad m
           => r -> FunDay (ReaderT r m) m ()
funReaderT r = FunDay $ \l ccX -> do
  x <- runReaderT ccX r
  pure (l () x)

-- 'funWriterT' has a hidden @m w@ computation which is waiting for a
-- @WriterT w m x@ post-computation. So it transforms a @WriterT w m@
-- computation into an @m@ computation and returns the accumulated 'w'.
funWriterT :: Monad m
           => FunDay (WriterT w m) m w
funWriterT = FunDay $ \l ccX -> do
  (x, w) <- runWriterT ccX
  pure (l w x)

-- @funStateT s@ has a hidden @m s@ computation which is waiting for a
-- @StateT s m x@ post-computation. So it transforms a @StateT s m@
-- computation into an @m@ computation and return the new 's' state.
funStateT :: Monad m
          => s -> FunDay (StateT s m) m s
funStateT s = FunDay $ \l ccX -> do
  (x, s') <- runStateT ccX s
  pure (l s' x)

-- Note that I was careful to pick monad transformers which always return a
-- value. If I try this with 'ExceptT', for example, that doesn't work because
-- if 'runExceptT' returns a 'Left', we won't have an 'x' to give to 'l'.

funExceptT :: Monad m
           => FunDay (ExceptT e m) m ()
funExceptT = FunDay $ \l ccX -> do
  r <- runExceptT ccX
  case r of
    Right x -> pure (l () x)
    Left  _ -> error "what now?"

-- Another thing which doesn't work is to use 'Codensity' instead of 'FunDay'.
-- We really do need 'FunDay' for this example!

-- There is no such thing as an 'IxYoneda'
data IxCodensity f g a = IxCodensity
  { unIxCodensity :: forall r. (a -> f r) -> g r }

codenStateT :: Monad m
            => s -> IxCodensity (StateT s m) m s
codenStateT s = IxCodensity $ \cc -> do
  (r, s') <- runStateT (cc s) s
  -- oops! we gave @s@ to @cc@, but we wanted to give @s'@ instead.
  -- and now that @s'@ is in scope, it's too late because we no longer have
  -- access to the 'StateT' effects which 'cc' needs!
  undefined r s'

-- All right, now that we have a bunch of 'FunDay's, how do we combine them? I
-- want to use something like Applicative syntax, but the types are not quite
-- right, as '(<*>)' expects the same 'f' on both sides and our 'FunDay's all
-- have different indices. So we need an indexed Applicative instead:

infixl 4 <$$>
infixl 4 <**>
class IxApplicative f where
  -- should probably be moved to an 'IxFunctor'
  (<$$>) :: (a -> b) -> f m n a -> f m n b

  -- the indices match those of function composition:
  --
  -- > (.)    :: (m -> n)  -> (l -> m)  -> (l -> n)
  -- > (<**>) :: f m n ... -> f l m ... -> f l n ...
  (<**>) :: f m n (a -> b) -> f l m a -> f l n b

instance IxApplicative FunDay where
  a2b <$$> fundayA = FunDay $ \l ccX
                  -> unFunDay fundayA (go l)
                   $ ccX
    where go l a x = l (a2b a) x
  fundayF <**> fundayA = FunDay $ \l ccX
                      -> unFunDay fundayF (go l)
                       $ unFunDay fundayA (,)
                       $ ccX
    where go l a2b (a,x) = l (a2b a) x

-- |
-- We now have everything we need to run our original example:
--
-- >>> :{
-- runFunDay $ (\() w s result -> result ++ ", " ++ w ++ ", " ++ s)
--        <$$> funReaderT "r"
--        <**> funWriterT
--        <**> funStateT "s"
--        <**> funday1 myRWST
-- :}
-- "result, w, s'"
--
-- @funday1 ma@ has a hidden @m1 a@ computation which is waiting for an @m1 x@
-- post-computation. So it keeps the index the same, and returns an 'a'. In this
-- case, 'm1' is @StateT String (WriterT String (ReaderT String m))@.
--
-- @funStateT "s"@ transforms 'm1' into @WriterT String (ReaderT String m)@,
-- which 'funWriterT' then transforms into @ReaderT String m@, which
-- @funReaderT "r"@ finally transforms into @m@. Overall, we transform an 'm1'
-- computation into an 'm' computation, so we have a @FunDay m1 m@. Since 'm2'
-- has an Applicative instance, we can use 'runFunDay' to extract the final
-- @m String@ computation by providing the empty post-computation @pure ()@:

runFunDay :: Applicative f
          => FunDay f g a -> g a
runFunDay funday = unFunDay funday const
                 $ pure ()

-- Finally,the REPL instantiates the @m String@ to @IO String@ and prints the
-- resulting String.
--
-- The effects occur in the opposite order: first @funReaderT@'s, then
-- @funWriterT@'s, then @funStateT@'s, then @funday1@'s, and finally, the
-- @pure ()@ post-computation given by 'runFunDay' to extract the hidden
-- computation. Then @funday1@ may clean up, then 'funStateT', then
-- 'funWriterT', then 'funReaderT'.
--
-- Throughout the computation, the 4 results computed by the 4 steps are
-- threaded by the various 'l's, and finally the '(<$$>)' receives those 4
-- results and gives them to one final pure function, which combines them into
-- the final result.



instance Functor (FunDay1 m) where
  fmap a2b fundayA = FunDay $ \l mx
                  -> unFunDay fundayA (\a x -> l (a2b a) x) mx


instance Functor (Codensity m) where
  fmap a2b codensityA = Codensity $ \cc -> do
    unCodensity codensityA $ \a -> do
      cc (a2b a)

instance Applicative (Codensity m) where
  pure a = Codensity $ \cc -> do
    cc a
  (<*>) = ap


main :: IO ()
main = doctest ["src/Main.hs"]