indexed-transformers-0.2.0.0: src/Control/Monad/Trans/Indexed/Free.hs

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

The free indexed monad transformer.

#example#

== Example

Free and freer indexed monad transformers can be used as
a domain specific language generated by primitive commands like this
[Conor McBride example]
(https://stackoverflow.com/questions/28690448/what-is-indexed-monad).

>>> :set -XGADTs -XDataKinds
>>> import Data.Kind
>>> type DVD = String
>>> :{
data DVDCommand
  :: Bool -- ^ drive is full before command
  -> Bool -- ^ drive is full after command
  -> Type -- ^ return type
  -> Type where
  Insert :: DVD -> DVDCommand 'False 'True ()
  Eject :: DVDCommand 'True 'False DVD
:}

>>> :{
insert :: Monad m => DVD -> FreerIx FreeIx DVDCommand 'False 'True m ()
insert dvd = liftFreerIx (Insert dvd)
:}

>>> :{
eject :: Monad m => FreerIx FreeIx DVDCommand 'True 'False m DVD
eject = liftFreerIx Eject
:}

>>> :set -XQualifiedDo
>>> import qualified Control.Monad.Trans.Indexed.Do as Indexed
>>> :{
swap :: Monad m => DVD -> FreerIx FreeIx DVDCommand 'True 'True m DVD
swap dvd = Indexed.do
  dvd' <- eject
  insert dvd
  return dvd'
:}

>>> import Control.Monad.Trans
>>> :{
printDVD :: FreerIx FreeIx DVDCommand 'True 'True IO ()
printDVD = Indexed.do
  dvd <- eject
  insert dvd
  lift $ putStrLn dvd
:}
-}

module Control.Monad.Trans.Indexed.Free
  ( IxMonadTransFree (..)
  , FreeIx (..)
  , improveIx, coerceFreeIx
  , FreerIx, liftFreerIx, hoistFreerIx, foldFreerIx
  , CoyonedaIx (..)
  , WrapFreeIx (..)
  , FoldFreeIx (..)
  , ImproveFreeIx (..)
  ) where

import Control.Monad
import Control.Monad.Free
import Control.Monad.Trans
import Control.Monad.Trans.Indexed
import Control.Monad.Trans.Indexed.Codensity

{- |
The free `IxMonadTrans` generated by an indexed `Functor`
is characterized by the `IxMonadTransFree` class
up to the isomorphism `coerceFreeIx`.
-}
class
  ( forall f. (forall k l. Functor (f k l)) => IxMonadTrans (freeIx f)
  , forall f m i j. (forall k l. Functor (f k l), Monad m, i ~ j)
    => MonadFree (f i j) (freeIx f i j m)
  ) => IxMonadTransFree freeIx where
  -- | lift a computation 
  liftFreeIx
    :: (forall k l. Functor (f k l), Monad m)
    => f i j x -- ^ computation
    -> freeIx f i j m x
  -- | hoist a transformation
  hoistFreeIx
    :: (forall k l. Functor (f k l), forall k l. Functor (g k l), Monad m)
    => (forall k l a. f k l a -> g k l a) -- ^ transformation
    -> freeIx f i j m x -> freeIx g i j m x
  -- | fold with a monadic transformation 
  foldFreeIx
    :: (forall k l. Functor (f k l), IxMonadTrans t, Monad m)
    => (forall k l a. f k l a -> t k l m a) -- ^ monadic transformation
    -> freeIx f i j m x -> t i j m x

{- |
> prop> coerceFreeIx = foldFreeIx liftFreeIx
> prop> id = coerceFreeIx . coerceFreeIx
-}
coerceFreeIx
  :: ( IxMonadTransFree freeIx0
     , IxMonadTransFree freeIx1
     , forall k l. Functor (f k l)
     , Monad m
     )
  => freeIx0 f i j m x -- ^ from free
  -> freeIx1 f i j m x -- ^ to free
coerceFreeIx = foldFreeIx liftFreeIx

{- | Right associate all binds in a computation that generates an indexed free monad transformer.

This can improve the asymptotic efficiency of the result, while preserving semantics.

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

<https://www.janis-voigtlaender.eu/papers/AsymptoticImprovementOfComputationsOverFreeMonads.pdf>
-}
improveIx
  :: (forall k l. Functor (f k l), Monad m)
  => (forall freeIx. IxMonadTransFree freeIx => freeIx f i j m a) -- ^ improve this
  -> FreeIx f i j m a
improveIx m = lowerCodensityIx (runImproveFreeIx m)

{- |
Combining `IxMonadTransFree` with `CoyonedaIx` gives the "freer" `IxMonadTrans`,
modeled on this [Oleg Kiselyov explanation]
(https://okmij.org/ftp/Computation/free-monad.html#freer).
See the [example]("Control.Monad.Trans.Indexed.Free#example").
-}
type FreerIx freeIx f = freeIx (CoyonedaIx f)

{- | `CoyonedaIx` is the free indexed `Functor`. -}
data CoyonedaIx f i j x where
  CoyonedaIx :: (x -> y) -> f i j x -> CoyonedaIx f i j y
instance Functor (CoyonedaIx f i j) where
  fmap g (CoyonedaIx f x) = CoyonedaIx (g . f) x

-- | Lift a computation to `FreerIx`.
liftFreerIx
  :: (IxMonadTransFree freeIx, Monad m)
  => f i j x -- ^ computation
  -> FreerIx freeIx f i j m x
liftFreerIx x = liftFreeIx (CoyonedaIx id x)

-- | Hoist a transformation to `FreerIx`
hoistFreerIx
  :: (IxMonadTransFree freeIx, Monad m)
  => (forall k l a. f k l a -> g k l a) -- ^ transformation
  -> FreerIx freeIx f i j m x -> FreerIx freeIx g i j m x
hoistFreerIx f = hoistFreeIx (\(CoyonedaIx g x) -> CoyonedaIx g (f x))

-- | Fold with a monadic transformation over `FreerIx`.
foldFreerIx
  :: (IxMonadTransFree freeIx, IxMonadTrans t, Monad m)
  => (forall k l a. f k l a -> t k l m a) -- ^ monadic transformation
  -> FreerIx freeIx f i j m x -> t i j m x
foldFreerIx f x = foldFreeIx (\(CoyonedaIx g y) -> g <$> f y) x

-- | A helper type for `FreeIx`.
data WrapFreeIx f i j m x where
  Unwrap :: x -> WrapFreeIx f i i m x
  Wrap :: f i j (FreeIx f j k m x) -> WrapFreeIx f i k m x
instance (forall k l. Functor (f k l), Monad m)
  => Functor (WrapFreeIx f i j m) where
    fmap f = \case
      Unwrap x -> Unwrap $ f x
      Wrap fm -> Wrap $ fmap (fmap f) fm

-- | The free indexed monad transformer.
newtype FreeIx f i j m x = FreeIx {runFreeIx :: m (WrapFreeIx f i j m x)}
instance (forall k l. Functor (f k l), Monad m)
  => Functor (FreeIx f i j m) where
    fmap f (FreeIx m) = FreeIx $ fmap (fmap f) m
instance (forall k l. Functor (f k l), i ~ j, Monad m)
  => Applicative (FreeIx f i j m) where
    pure = FreeIx . pure . Unwrap
    (<*>) = apIx
instance (forall k l. Functor (f k l), i ~ j, Monad m)
  => Monad (FreeIx f i j m) where
    return = pure
    (>>=) = flip bindIx
instance (forall k l. Functor (f k l), i ~ j)
  => MonadTrans (FreeIx f i j) where
    lift = FreeIx . fmap Unwrap
instance (forall k l. Functor (f k l)) => IxMonadTrans (FreeIx f) where
  joinIx (FreeIx mm) = FreeIx $ mm >>= \case
    Unwrap (FreeIx m) -> m
    Wrap fm -> return $ Wrap $ fmap joinIx fm
instance
  ( forall k l. Functor (f k l)
  , Monad m
  , i ~ j
  ) => MonadFree (f i j) (FreeIx f i j m) where
    wrap = FreeIx . return . Wrap
instance IxMonadTransFree FreeIx where
  liftFreeIx = FreeIx . return . Wrap . fmap return
  hoistFreeIx f (FreeIx m) = FreeIx (fmap hoist_f m)
    where
      hoist_f = \case
        Unwrap x -> Unwrap x
        Wrap y -> Wrap (f (fmap (hoistFreeIx f) y))
  foldFreeIx f (FreeIx m) = bindIx foldMap_f (lift m)
    where
      foldMap_f = \case
        Unwrap x -> return x
        Wrap y -> bindIx (foldFreeIx f) (f y)

-- | The free indexed monad transformer, encoded as its `foldFreeIx` function.
newtype FoldFreeIx g i j m x = FoldFreeIx
  {runFoldFreeIx :: forall t. (IxMonadTrans t, Monad m)
    => (forall k l a. g k l a -> t k l m a) -> t i j m x}
instance (forall k l. Functor (f k l), Monad m) => Functor (FoldFreeIx f i j m) where
  fmap f (FoldFreeIx k) = FoldFreeIx $ \step -> fmap f (k step)
instance (forall k l. Functor (f k l), i ~ j, Monad m)
  => Applicative (FoldFreeIx f i j m) where
    pure x = FoldFreeIx $ const $ pure x
    (<*>) = apIx
instance (forall k l. Functor (f k l), i ~ j, Monad m)
  => Monad (FoldFreeIx f i j m) where
    return = pure
    (>>=) = flip bindIx
instance (forall k l. Functor (f k l), i ~ j)
  => MonadTrans (FoldFreeIx f i j) where
    lift m = FoldFreeIx $ const $ lift m
instance (forall k l. Functor (f k l))
  => IxMonadTrans (FoldFreeIx f) where
    joinIx (FoldFreeIx g) = FoldFreeIx $ \k -> bindIx (\(FoldFreeIx f) -> f k) (g k)
instance
  ( forall k l. Functor (f k l)
  , Monad m
  , i ~ j
  ) => MonadFree (f i j) (FoldFreeIx f i j m) where
    wrap = join . liftFreeIx
instance IxMonadTransFree FoldFreeIx where
  liftFreeIx m = FoldFreeIx $ \k -> k m
  hoistFreeIx f (FoldFreeIx k) = FoldFreeIx $ \g -> k (g . f)
  foldFreeIx f (FoldFreeIx k) = k f

-- | A helper type for `improveIx`.
newtype ImproveFreeIx f i j m a = ImproveFreeIx
  { runImproveFreeIx :: CodensityIx (FreeIx f) i j m a }
deriving newtype instance Functor (ImproveFreeIx f i j m)
deriving newtype instance i ~ j => Applicative (ImproveFreeIx f i j m)
deriving newtype instance i ~ j => Monad (ImproveFreeIx f i j m)
deriving newtype instance (forall k l. Functor (f k l), i ~ j) => MonadTrans (ImproveFreeIx f i j)
deriving newtype instance (forall k l. Functor (f k l)) => IxMonadTrans (ImproveFreeIx f)
instance
  ( forall k l. Functor (f k l)
  , Monad m
  , i ~ j
  ) => MonadFree (f i j) (ImproveFreeIx f i j m) where
    wrap t = ImproveFreeIx $ CodensityIx $ \h ->
      joinIx (liftFreeIx (fmap (\p -> runCodensityIx (runImproveFreeIx p) h) t))
instance IxMonadTransFree ImproveFreeIx where
    liftFreeIx = ImproveFreeIx . liftCodensityIx . liftFreeIx
    hoistFreeIx f
      = ImproveFreeIx . liftCodensityIx
      . hoistFreeIx f
      . lowerCodensityIx . runImproveFreeIx
    foldFreeIx f = foldFreeIx f . lowerCodensityIx . runImproveFreeIx