{-# LANGUAGE CPP #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif

#ifndef MIN_VERSION_speculation
#define MIN_VERSION_speculation(x,y,z) 1
#endif
-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Codensity
-- Copyright   :  (C) 2008-2013 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  non-portable (rank-2 polymorphism)
--
----------------------------------------------------------------------------
module Control.Monad.Codensity
  ( Codensity(..)
  , lowerCodensity
  , codensityToAdjunction, adjunctionToCodensity
  , codensityToRan, ranToCodensity
  , codensityToComposedRep, composedRepToCodensity
  , improve
  ) where

import Control.Applicative
import Control.Concurrent.Speculation
import Control.Concurrent.Speculation.Class
import Control.Monad (ap, MonadPlus(..))
import Control.Monad.Free
import Control.Monad.IO.Class
import Control.Monad.Reader.Class
import Control.Monad.State.Class
import Control.Monad.Trans.Class
import Data.Functor.Adjunction
import Data.Functor.Apply
import Data.Functor.Kan.Ran
import Data.Functor.Plus
import Data.Functor.Representable
import Data.Key

-- |
-- @'Codensity' f@ is the Monad generated by taking the right Kan extension
-- of any 'Functor' @f@ along itself (@Ran f f@).
--
-- This can often be more \"efficient\" to construct than @f@ itself using
-- repeated applications of @(>>=)@.
--
-- See \"Asymptotic Improvement of Computations over Free Monads\" by Janis
-- Voightländer for more information about this type.
--
-- <http://www.iai.uni-bonn.de/~jv/mpc08.pdf>
newtype Codensity m a = Codensity
  { runCodensity :: forall b. (a -> m b) -> m b
  }

instance MonadSpec (Codensity m) where
  specByM f g a = Codensity $ \k -> specBy f g k a
  {-# INLINE specByM #-}
#if !(MIN_VERSION_speculation(1,5,0))
  specByM' f g a = Codensity $ \k -> specBy' f g k a
  {-# INLINE specByM' #-}
#endif

instance Functor (Codensity k) where
  fmap f (Codensity m) = Codensity (\k -> m (k . f))
  {-# INLINE fmap #-}

instance Apply (Codensity f) where
  (<.>) = ap
  {-# INLINE (<.>) #-}

instance Applicative (Codensity f) where
  pure x = Codensity (\k -> k x)
  {-# INLINE pure #-}
  (<*>) = ap
  {-# INLINE (<*>) #-}

instance Monad (Codensity f) where
  return x = Codensity (\k -> k x)
  {-# INLINE return #-}
  m >>= k = Codensity (\c -> runCodensity m (\a -> runCodensity (k a) c))
  {-# INLINE (>>=) #-}

instance MonadIO m => MonadIO (Codensity m) where
  liftIO = lift . liftIO
  {-# INLINE liftIO #-}

instance MonadTrans Codensity where
  lift m = Codensity (m >>=)
  {-# INLINE lift #-}

instance Alt v => Alt (Codensity v) where
  Codensity m <!> Codensity n = Codensity (\k -> m k <!> n k)
  {-# INLINE (<!>) #-}

instance Plus v => Plus (Codensity v) where
  zero = Codensity (const zero)
  {-# INLINE zero #-}

{-
instance Plus v => Alternative (Codensity v) where
  empty = zero
  (<|>) = (<!>)

instance Plus v => MonadPlus (Codensity v) where
  mzero = zero
  mplus = (<!>)
-}

instance Alternative v => Alternative (Codensity v) where
  empty = Codensity (\_ -> empty)
  {-# INLINE empty #-}
  Codensity m <|> Codensity n = Codensity (\k -> m k <|> n k)
  {-# INLINE (<|>) #-}

instance MonadPlus v => MonadPlus (Codensity v) where
  mzero = Codensity (\_ -> mzero)
  {-# INLINE mzero #-}
  Codensity m `mplus` Codensity n = Codensity (\k -> m k `mplus` n k)
  {-# INLINE mplus #-}

-- |
-- This serves as the *left*-inverse (retraction) of 'lift'.
--
--
-- @
-- 'lowerCodensity . lift' ≡ 'id'
-- @
--
-- In general this is not a full 2-sided inverse, merely a retraction, as
-- @'Codensity' m@ is often considerably "larger" than @m@.
--
-- e.g. @'Codensity' ((->) s)) a ~ forall r. (a -> s -> r) -> s -> r@
-- could support a full complement of @'MonadState' s@ actions, while @(->) s@
-- is limited to @'MonadReader' s@ actions.
lowerCodensity :: Monad m => Codensity m a -> m a
lowerCodensity a = runCodensity a return
{-# INLINE lowerCodensity #-}

-- | The 'Codensity' monad of a right adjoint is isomorphic to the
-- monad obtained from the 'Adjunction'.
--
-- @
-- 'codensityToAdjunction' . 'adjunctionToCodensity' ≡ 'id'
-- 'adjunctionToCodensity' . 'codensityToAdjunction' ≡ 'id'
-- @
codensityToAdjunction :: Adjunction f g => Codensity g a -> g (f a)
codensityToAdjunction r = runCodensity r unit
{-# INLINE codensityToAdjunction #-}

adjunctionToCodensity :: Adjunction f g => g (f a) -> Codensity g a
adjunctionToCodensity f = Codensity (\a -> fmap (rightAdjunct a) f)
{-# INLINE adjunctionToCodensity #-}

-- | The 'Codensity' monad of a representable 'Functor' is isomorphic to the
-- monad obtained from the 'Adjunction' for which that 'Functor' is the right
-- adjoint.
--
-- @
-- 'codensityToComposedRep' . 'composedRepToCodensity' ≡ 'id'
-- 'composedRepToCodensity' . 'codensityToComposedRep' ≡ 'id'
-- @
--
-- @
-- codensityToComposedRep = 'ranToComposedRep' . 'codensityToRan'
-- @

codensityToComposedRep :: Representable u => Codensity u a -> u (Key u, a)
codensityToComposedRep (Codensity f) = f (\a -> tabulate $ \e -> (e, a))
{-# INLINE codensityToComposedRep #-}

-- |
--
-- @
-- 'composedRepToCodensity' = 'ranToCodensity' . 'composedRepToRan'
-- @
composedRepToCodensity :: Representable u => u (Key u, a) -> Codensity u a
composedRepToCodensity hfa = Codensity $ \k -> fmap (\(e, a) -> index (k a) e) hfa
{-# INLINE composedRepToCodensity #-}

-- | The 'Codensity' 'Monad' of a 'Functor' @g@ is the right Kan extension ('Ran')
-- of @g@ along itself.
--
-- @
-- 'codensityToRan' . 'ranToCodensity' ≡ 'id'
-- 'ranToCodensity' . 'codensityToRan' ≡ 'id'
-- @
codensityToRan :: Codensity g a -> Ran g g a
codensityToRan (Codensity m) = Ran m
{-# INLINE codensityToRan #-}

ranToCodensity :: Ran g g a -> Codensity g a
ranToCodensity (Ran m) = Codensity m
{-# INLINE ranToCodensity #-}

instance (Functor f, MonadFree f m) => MonadFree f (Codensity m) where
  wrap t = Codensity (\h -> wrap (fmap (\p -> runCodensity p h) t))
  {-# INLINE wrap #-}

instance MonadReader r m => MonadState r (Codensity m) where
  get = Codensity (ask >>=)
  {-# INLINE get #-}
  put s = Codensity (\k -> local (const s) (k ()))
  {-# INLINE put #-}

-- | Right associate all binds in a computation that generates a free monad
--
-- This can improve the asymptotic efficiency of the result, while preserving
-- semantics.
--
-- See \"Asymptotic Improvement of Computations over Free Monads\" by Janis
-- Voightländer for more information about this combinator.
--
-- <http://www.iai.uni-bonn.de/~jv/mpc08.pdf>
improve :: Functor f => (forall m. MonadFree f m => m a) -> Free f a
improve m = lowerCodensity m
{-# INLINE improve #-}