{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS -Wno-orphans #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_HADDOCK hide #-}

module Control.Functor.Linear.Internal.Reader
  ( --  ReaderT monad transformer
    Reader,
    reader,
    runReader,
    mapReader,
    withReader,
    ReaderT (..),
    runReaderT,
    mapReaderT,
    withReaderT,
    ask,
    local,
    asks,
  )
where

import Control.Functor.Linear.Internal.Class
import Control.Functor.Linear.Internal.Instances ()
import Control.Functor.Linear.Internal.MonadTrans
import qualified Control.Monad as NonLinear ()
import qualified Control.Monad.Trans.Reader as NonLinear
import Data.Functor.Identity
import qualified Data.Functor.Linear.Internal.Applicative as Data
import qualified Data.Functor.Linear.Internal.Functor as Data
import Data.Kind (FUN)
import Data.Unrestricted.Linear.Internal.Consumable
import Data.Unrestricted.Linear.Internal.Dupable
import GHC.Exts (Multiplicity (..))
import Prelude.Linear.Internal (runIdentity', ($), (.))

-- # Linear ReaderT
-------------------------------------------------------------------------------

-- | A linear reader monad transformer.
-- This reader monad requires that use of the read-only state is explict.
--
-- The monad instance requires that @r@ be 'Dupable'.  This means that you
-- should use the linear reader monad just like the non-linear monad, except
-- that the type system ensures that you explicity use or discard the
-- read-only state (with the 'Consumable' instance).
newtype ReaderT r m a = ReaderT (r %1 -> m a)

-- XXX: Replace with a newtype deconstructor once it can be inferred as linear.

-- | Provide an intial read-only state and run the monadic computation in
-- a reader monad transformer
runReaderT :: ReaderT r m a %1 -> r %1 -> m a
runReaderT (ReaderT f) = f

instance (Data.Functor m) => Data.Functor (ReaderT r m) where
  fmap f = mapReaderT (Data.fmap f)

instance (Functor m) => Functor (ReaderT r m) where
  fmap f = mapReaderT (fmap f)

instance (Data.Applicative m, Dupable r) => Data.Applicative (ReaderT r m) where
  pure x = ReaderT $ \r -> lseq r (Data.pure x)
  ReaderT f <*> ReaderT x = ReaderT ((\(r1, r2) -> f r1 Data.<*> x r2) . dup)

instance (Applicative m, Dupable r) => Applicative (ReaderT r m) where
  pure x = ReaderT $ \r -> lseq r (pure x)
  ReaderT f <*> ReaderT x = ReaderT ((\(r1, r2) -> f r1 Data.<*> x r2) . dup)

instance (Monad m, Dupable r) => Monad (ReaderT r m) where
  ReaderT x >>= f = ReaderT ((\(r1, r2) -> x r1 >>= (\a -> runReaderT (f a) r2)) . dup)

type Reader r = ReaderT r Identity

ask :: (Applicative m) => ReaderT r m r
ask = ReaderT pure

withReaderT :: (r' %1 -> r) %1 -> ReaderT r m a %1 -> ReaderT r' m a
withReaderT f m = ReaderT $ runReaderT m . f

local :: (r %1 -> r) %1 -> ReaderT r m a %1 -> ReaderT r m a
local = withReaderT

reader :: (Monad m) => (r %1 -> a) %1 -> ReaderT r m a
reader f = ReaderT (return . f)

runReader :: Reader r a %1 -> r %1 -> a
runReader m = runIdentity' . runReaderT m

mapReader :: (a %1 -> b) %1 -> Reader r a %1 -> Reader r b
mapReader f = mapReaderT (Identity . f . runIdentity')

mapReaderT :: (m a %1 -> n b) %1 -> ReaderT r m a %1 -> ReaderT r n b
mapReaderT f m = ReaderT (f . runReaderT m)

withReader :: (r' %1 -> r) %1 -> Reader r a %1 -> Reader r' a
withReader = withReaderT

asks :: (Monad m) => (r %1 -> a) %1 -> ReaderT r m a
asks f = ReaderT (return . f)

instance (Dupable r) => MonadTrans (ReaderT r) where
  lift x = ReaderT (`lseq` x)

-- # Instances for nonlinear ReaderT
-------------------------------------------------------------------------------

instance (Functor m) => Functor (NonLinear.ReaderT r m) where
  fmap f (NonLinear.ReaderT g) = NonLinear.ReaderT $ \r -> fmap f (g r)

instance (Applicative m) => Applicative (NonLinear.ReaderT r m) where
  pure x = NonLinear.ReaderT $ \_ -> pure x
  NonLinear.ReaderT f <*> NonLinear.ReaderT x = NonLinear.ReaderT $ \r -> f r <*> x r

instance (Monad m) => Monad (NonLinear.ReaderT r m) where
  NonLinear.ReaderT x >>= f = NonLinear.ReaderT $ \r -> x r >>= (\a -> runReaderT' (f a) r)

-- XXX: Temporary, until newtype record projections are linear.
runReaderT' :: NonLinear.ReaderT r m a %1 -> r -> m a
runReaderT' (NonLinear.ReaderT f) = f

instance MonadTrans (NonLinear.ReaderT r) where
  lift x = NonLinear.ReaderT (\_ -> x)

deriving via Reader r instance (Dupable r) => Data.Applicative (FUN 'One r)

deriving via Reader r instance (Dupable r) => Applicative (FUN 'One r)

deriving via Reader r instance (Dupable r) => Monad (FUN 'One r)
