Codensity.hs

{- |
Module      :  Control.Monad.Trans.Indexed.Codensity
Copyright   :  (C) 2026 Eitan Chatav
License     :  BSD 3-Clause License (see the file LICENSE)
Maintainer  :  Eitan Chatav <eitan.chatav@gmail.com>

The indexed codensity monad transformer.
-}

module Control.Monad.Trans.Indexed.Codensity
  ( CodensityIx (..)
  , lowerCodensityIx
  , liftCodensityIx
  , toCodensity
  , wrapCodensityIx
  , resetCodensityIx
  , shiftCodensityIx
  , PredensityIx (..)
  , lowerToStateIx
  , liftFromStateIx
  ) where

import Control.Applicative
import Control.Monad
import Control.Monad.Codensity
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Trans.Indexed
import Control.Monad.Trans.Indexed.State
import Data.Kind

{- |
@'CodensityIx' t@ is the indexed monad transformer generated by taking
the right Kan extension of any indexed monad transformer @t@ along itself.

This can often be more \"efficient\" to construct than @t@ itself using
repeated applications of 'bindIx',
as in `Control.Monad.Trans.Indexed.Free.ImproveFreeIx`.

See \"Asymptotic Improvement of Computations over Free Monads\" by Janis
Voigtländer for more information.

<https://www.janis-voigtlaender.eu/papers/AsymptoticImprovementOfComputationsOverFreeMonads.pdf>
-}
newtype CodensityIx t i j m a = CodensityIx
  { runCodensityIx :: forall b k. (a -> t j k m b) -> t i k m b }
  deriving Functor

{- |
This serves as the *left*-inverse (retraction) of 'liftCodensityIx'.

> prop> lowerCodensityIx . liftCodensityIx ≡ id

In general this is not a full 2-sided inverse, merely a retraction, as
@'CodensityIx' t@ is often considerably \"larger\" than @t@.
-}
lowerCodensityIx
  :: (IxMonadTrans t, Monad m)
  => CodensityIx t i j m a -> t i j m a
lowerCodensityIx (CodensityIx f) = f return

{- | Lift a computation from the argument indexed monad
to the constructed indexed monad. -}
liftCodensityIx
  :: (IxMonadTrans t, Monad m)
  => t i j m a -> CodensityIx t i j m a
liftCodensityIx m = CodensityIx $ \h -> bindIx h m

{- | Convert to `Codensity`. -}
toCodensity :: CodensityIx t i i m a -> Codensity (t i i m) a
toCodensity (CodensityIx f) = Codensity f

{- | Wrap the remainder of the 'CodensityIx' action
using the given function. -}
wrapCodensityIx
  :: (forall a k. t j k (m :: Type -> Type) a -> t i k m a) -- ^ remainder
  -> CodensityIx t i j m ()
wrapCodensityIx f = CodensityIx (\k -> f (k ()))

{- | @'resetCodensityIx' m@ delimits the continuation of any 'shiftCodensityIx' inside @m@.

> prop> resetCodensityIx (return m) = return m
-}
resetCodensityIx :: (IxMonadTrans t, Monad m) => CodensityIx t i j m a -> CodensityIx t i j m a
resetCodensityIx = liftCodensityIx . lowerCodensityIx

{- | @'shiftCodensityIx' f@ captures the continuation up to the nearest enclosing
'resetCodensityIx' and passes it to @f@:

> prop> resetCodensityIx (shiftCodensityIx f & bindIx k) = resetCodensityIx (f (lowerCodensityIx . k))
-}
shiftCodensityIx
  :: (IxMonadTrans t, Monad m)
  => (forall b k. (a -> t j k m b) -> CodensityIx t i k m b)
  -> CodensityIx t i j m a
shiftCodensityIx f = CodensityIx $ lowerCodensityIx . f

-- CodensityIx instances
instance IxMonadTrans t => IxMonadTrans (CodensityIx t) where
  joinIx (CodensityIx k) =
    CodensityIx $ \f -> k $ \(CodensityIx g) -> g f
instance i ~ j => Applicative (CodensityIx t i j m) where
  pure x = CodensityIx $ \k -> k x
  CodensityIx cf <*> CodensityIx cx =
    CodensityIx $ \ k -> cf $ \ f -> cx (k . f)
instance i ~ j => Monad (CodensityIx t i j m) where
  return = pure
  CodensityIx cx >>= k =
    CodensityIx $ \ c -> cx (\ x -> runCodensityIx (k x) c)
instance (IxMonadTrans t, i ~ j) => MonadTrans (CodensityIx t i j) where
  lift m = CodensityIx (\k -> bindIx k (lift m))
instance (i ~ j, Alternative (t i j m), IxMonadTrans t, Monad m)
  => Alternative (CodensityIx t i j m) where
    empty = liftCodensityIx empty
    x <|> y = liftCodensityIx (lowerCodensityIx x <|> lowerCodensityIx y)
instance (i ~ j, Alternative (t i j m), IxMonadTrans t, Monad m)
  => MonadPlus (CodensityIx t i j m)

{- | `PredensityIx` `ReaderT` is an efficient encoding of `StateIx`.

'lowerToStateIx' is the *left*-inverse (retraction) of 'liftFromStateIx'.

> prop> lowerToStateIx . liftFromStateIx ≡ id

In general this is not a full 2-sided inverse, merely a retraction, as
@'PredensityIx' 'ReaderT'@ is \"larger\" than `StateIx`:
it may call its continuation any number of times.
-}
newtype PredensityIx t i j m a = PredensityIx
  { runPredensityIx :: forall b. (a -> t j m b) -> t i m b }
  deriving Functor

{- | Convert to `StateIx`. -}
lowerToStateIx :: Monad m => PredensityIx ReaderT i j m a -> StateIx i j m a
lowerToStateIx (PredensityIx f) =
  StateIx . runReaderT . f $ \x -> ReaderT $ \j -> return (x, j)

{- | Convert from `StateIx`. -}
liftFromStateIx :: Monad m => StateIx i j m a -> PredensityIx ReaderT i j m a
liftFromStateIx (StateIx f) =
  PredensityIx $ \k -> ReaderT $ \i -> f i >>= \(x, j) -> runReaderT (k x) j

-- PredensityIx instances
instance (forall i. MonadTrans (t i)) => IxMonadTrans (PredensityIx t) where
  joinIx (PredensityIx k) =
    PredensityIx $ \f -> k $ \(PredensityIx g) -> g f
instance i ~ j => Applicative (PredensityIx t i j m) where
  pure x = PredensityIx $ \k -> k x
  PredensityIx cf <*> PredensityIx cx =
    PredensityIx $ \ k -> cf $ \ f -> cx (k . f)
instance i ~ j => Monad (PredensityIx t i j m) where
  return = pure
  PredensityIx cx >>= k =
    PredensityIx $ \ c -> cx (\ x -> runPredensityIx (k x) c)
instance (MonadTrans (t i), i ~ j) => MonadTrans (PredensityIx t i j) where
  lift m = PredensityIx (lift m >>=)
instance (i ~ j, Monad m) => MonadState i (PredensityIx ReaderT i j m) where
  get = getIx
  put = putIx
instance IxMonadTransState (PredensityIx ReaderT) where
  getIx = PredensityIx (ask >>=)
  putIx j = PredensityIx (\k -> withReaderT (const j) (k ()))