module Agda.Utils.Update
  ( ChangeT
  , runChangeT, mapChangeT
  , UpdaterT
  , runUpdaterT
  , Change
  , MonadChange(..)
  , runChange
  , Updater
  , sharing
  , runUpdater
  , dirty
  , ifDirty
  , Updater1(..)
  , Updater2(..)
  ) where

-- Control.Monad.Fail import is redundant since GHC 8.8.1
import Control.Monad.Fail (MonadFail)

import Control.Monad.Identity
import Control.Monad.Trans
import Control.Monad.Trans.Control
-- NB: Control.Monad.Trans.Identity is already exported by Control.Monad.Identity
-- since version mtl 2.2.2, but this needs at least ghc 8.2.2
import Control.Monad.Trans.Identity
import Control.Monad.Writer.Strict ( MonadWriter(..), Writer, WriterT, mapWriterT, runWriterT )

import Data.Monoid ( Any(..) )

import Agda.Utils.Tuple

-- * Change monad.

-- | The class of change monads.
class Monad m => MonadChange m where
  tellDirty   :: m () -- ^ Mark computation as having changed something.
  listenDirty :: m a -> m (a, Bool)

-- | The @ChangeT@ monad transformer.
newtype ChangeT m a = ChangeT { forall (m :: * -> *) a. ChangeT m a -> WriterT Any m a
fromChangeT :: WriterT Any m a }
  deriving (forall a b. a -> ChangeT m b -> ChangeT m a
forall a b. (a -> b) -> ChangeT m a -> ChangeT m b
forall (m :: * -> *) a b.
Functor m =>
a -> ChangeT m b -> ChangeT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ChangeT m a -> ChangeT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> ChangeT m b -> ChangeT m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> ChangeT m b -> ChangeT m a
fmap :: forall a b. (a -> b) -> ChangeT m a -> ChangeT m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ChangeT m a -> ChangeT m b
Functor, forall a. a -> ChangeT m a
forall a b. ChangeT m a -> ChangeT m b -> ChangeT m a
forall a b. ChangeT m a -> ChangeT m b -> ChangeT m b
forall a b. ChangeT m (a -> b) -> ChangeT m a -> ChangeT m b
forall a b c.
(a -> b -> c) -> ChangeT m a -> ChangeT m b -> ChangeT m c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall {m :: * -> *}. Applicative m => Functor (ChangeT m)
forall (m :: * -> *) a. Applicative m => a -> ChangeT m a
forall (m :: * -> *) a b.
Applicative m =>
ChangeT m a -> ChangeT m b -> ChangeT m a
forall (m :: * -> *) a b.
Applicative m =>
ChangeT m a -> ChangeT m b -> ChangeT m b
forall (m :: * -> *) a b.
Applicative m =>
ChangeT m (a -> b) -> ChangeT m a -> ChangeT m b
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> ChangeT m a -> ChangeT m b -> ChangeT m c
<* :: forall a b. ChangeT m a -> ChangeT m b -> ChangeT m a
$c<* :: forall (m :: * -> *) a b.
Applicative m =>
ChangeT m a -> ChangeT m b -> ChangeT m a
*> :: forall a b. ChangeT m a -> ChangeT m b -> ChangeT m b
$c*> :: forall (m :: * -> *) a b.
Applicative m =>
ChangeT m a -> ChangeT m b -> ChangeT m b
liftA2 :: forall a b c.
(a -> b -> c) -> ChangeT m a -> ChangeT m b -> ChangeT m c
$cliftA2 :: forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> ChangeT m a -> ChangeT m b -> ChangeT m c
<*> :: forall a b. ChangeT m (a -> b) -> ChangeT m a -> ChangeT m b
$c<*> :: forall (m :: * -> *) a b.
Applicative m =>
ChangeT m (a -> b) -> ChangeT m a -> ChangeT m b
pure :: forall a. a -> ChangeT m a
$cpure :: forall (m :: * -> *) a. Applicative m => a -> ChangeT m a
Applicative, forall a. a -> ChangeT m a
forall a b. ChangeT m a -> ChangeT m b -> ChangeT m b
forall a b. ChangeT m a -> (a -> ChangeT m b) -> ChangeT m b
forall {m :: * -> *}. Monad m => Applicative (ChangeT m)
forall (m :: * -> *) a. Monad m => a -> ChangeT m a
forall (m :: * -> *) a b.
Monad m =>
ChangeT m a -> ChangeT m b -> ChangeT m b
forall (m :: * -> *) a b.
Monad m =>
ChangeT m a -> (a -> ChangeT m b) -> ChangeT m b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> ChangeT m a
$creturn :: forall (m :: * -> *) a. Monad m => a -> ChangeT m a
>> :: forall a b. ChangeT m a -> ChangeT m b -> ChangeT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
ChangeT m a -> ChangeT m b -> ChangeT m b
>>= :: forall a b. ChangeT m a -> (a -> ChangeT m b) -> ChangeT m b
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
ChangeT m a -> (a -> ChangeT m b) -> ChangeT m b
Monad, forall (m :: * -> *) a. Monad m => m a -> ChangeT m a
forall (t :: (* -> *) -> * -> *).
(forall (m :: * -> *) a. Monad m => m a -> t m a) -> MonadTrans t
lift :: forall (m :: * -> *) a. Monad m => m a -> ChangeT m a
$clift :: forall (m :: * -> *) a. Monad m => m a -> ChangeT m a
MonadTrans, forall a. String -> ChangeT m a
forall (m :: * -> *).
Monad m -> (forall a. String -> m a) -> MonadFail m
forall {m :: * -> *}. MonadFail m => Monad (ChangeT m)
forall (m :: * -> *) a. MonadFail m => String -> ChangeT m a
fail :: forall a. String -> ChangeT m a
$cfail :: forall (m :: * -> *) a. MonadFail m => String -> ChangeT m a
MonadFail, forall a. IO a -> ChangeT m a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
forall {m :: * -> *}. MonadIO m => Monad (ChangeT m)
forall (m :: * -> *) a. MonadIO m => IO a -> ChangeT m a
liftIO :: forall a. IO a -> ChangeT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> ChangeT m a
MonadIO)

-- This instance cannot be derived in older ghcs like 8.0
-- because of the associated type synonym.
-- 8.4 can derive it, but needs UndecidableInstances.
instance MonadTransControl ChangeT where
  type StT ChangeT a = (a, Any) -- StT (WriterT Any) a  would require UndecidableInstances
  liftWith :: forall (m :: * -> *) a.
Monad m =>
(Run ChangeT -> m a) -> ChangeT m a
liftWith Run ChangeT -> m a
f = forall (m :: * -> *) a. WriterT Any m a -> ChangeT m a
ChangeT forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTransControl t, Monad m) =>
(Run t -> m a) -> t m a
liftWith forall a b. (a -> b) -> a -> b
$ \ Run (WriterT Any)
runWriterT -> Run ChangeT -> m a
f forall a b. (a -> b) -> a -> b
$ Run (WriterT Any)
runWriterT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. ChangeT m a -> WriterT Any m a
fromChangeT
  -- Andreas, 2020-04-17: these point-free variants do not seem to type check:
  -- liftWith f = ChangeT $ liftWith $ f . (. fromChangeT)
  -- liftWith = ChangeT . liftWith . (. (. fromChangeT))
  restoreT :: forall (m :: * -> *) a. Monad m => m (StT ChangeT a) -> ChangeT m a
restoreT = forall (m :: * -> *) a. WriterT Any m a -> ChangeT m a
ChangeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTransControl t, Monad m) =>
m (StT t a) -> t m a
restoreT

instance Monad m => MonadChange (ChangeT m) where
  tellDirty :: ChangeT m ()
tellDirty     = forall (m :: * -> *) a. WriterT Any m a -> ChangeT m a
ChangeT forall a b. (a -> b) -> a -> b
$ forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell forall a b. (a -> b) -> a -> b
$ Bool -> Any
Any Bool
True
  listenDirty :: forall a. ChangeT m a -> ChangeT m (a, Bool)
listenDirty ChangeT m a
m = forall (m :: * -> *) a. WriterT Any m a -> ChangeT m a
ChangeT forall a b. (a -> b) -> a -> b
$ do
    (a
a, Any Bool
dirty) <- forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (forall (m :: * -> *) a. ChangeT m a -> WriterT Any m a
fromChangeT ChangeT m a
m)
    forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Bool
dirty)

-- | Run a 'ChangeT' computation, returning result plus change flag.
runChangeT :: Functor m => ChangeT m a -> m (a, Bool)
runChangeT :: forall (m :: * -> *) a. Functor m => ChangeT m a -> m (a, Bool)
runChangeT = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall b d a. (b -> d) -> (a, b) -> (a, d)
mapSnd Any -> Bool
getAny) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. ChangeT m a -> WriterT Any m a
fromChangeT

-- | Run a 'ChangeT' computation, but ignore change flag.
execChangeT :: Functor m => ChangeT m a -> m a -- A library function, so keep
execChangeT :: forall (m :: * -> *) a. Functor m => ChangeT m a -> m a
execChangeT = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Functor m => ChangeT m a -> m (a, Bool)
runChangeT

-- | Map a 'ChangeT' computation (monad transformer action).
mapChangeT :: (m (a, Any) -> n (b, Any)) -> ChangeT m a -> ChangeT n b
mapChangeT :: forall (m :: * -> *) a (n :: * -> *) b.
(m (a, Any) -> n (b, Any)) -> ChangeT m a -> ChangeT n b
mapChangeT m (a, Any) -> n (b, Any)
f (ChangeT WriterT Any m a
m) = forall (m :: * -> *) a. WriterT Any m a -> ChangeT m a
ChangeT (forall (m :: * -> *) a w (n :: * -> *) b w'.
(m (a, w) -> n (b, w')) -> WriterT w m a -> WriterT w' n b
mapWriterT m (a, Any) -> n (b, Any)
f WriterT Any m a
m)

-- Don't actually track changes with the identity monad:

-- | A mock change monad.  Always assume change has happened.
instance MonadChange Identity where
  tellDirty :: Identity ()
tellDirty   = forall (m :: * -> *) a. Monad m => a -> m a
return ()
  listenDirty :: forall a. Identity a -> Identity (a, Bool)
listenDirty = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,Bool
True)

instance Monad m => MonadChange (IdentityT m) where
  tellDirty :: IdentityT m ()
tellDirty   = forall {k} (f :: k -> *) (a :: k). f a -> IdentityT f a
IdentityT    forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return ()
  listenDirty :: forall a. IdentityT m a -> IdentityT m (a, Bool)
listenDirty = forall {k1} {k2} (m :: k1 -> *) (a :: k1) (n :: k2 -> *) (b :: k2).
(m a -> n b) -> IdentityT m a -> IdentityT n b
mapIdentityT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,Bool
True)

-- * Pure endo function and updater

type UpdaterT m a = a -> ChangeT m a

-- | Blindly run an updater.
runUpdaterT :: Functor m => UpdaterT m a -> a -> m (a, Bool)
runUpdaterT :: forall (m :: * -> *) a.
Functor m =>
UpdaterT m a -> a -> m (a, Bool)
runUpdaterT UpdaterT m a
f a
a = forall (m :: * -> *) a. Functor m => ChangeT m a -> m (a, Bool)
runChangeT forall a b. (a -> b) -> a -> b
$ UpdaterT m a
f a
a

type EndoFun a = a -> a
type Change  a = ChangeT Identity a
type Updater a = UpdaterT Identity a

-- NB:: Defined but not used
fromChange :: Change a -> Writer Any a
fromChange :: forall a. Change a -> Writer Any a
fromChange = forall (m :: * -> *) a. ChangeT m a -> WriterT Any m a
fromChangeT

-- | Run a 'Change' computation, returning result plus change flag.
{-# INLINE runChange #-}
runChange :: Change a -> (a, Bool)
runChange :: forall a. Change a -> (a, Bool)
runChange = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Functor m => ChangeT m a -> m (a, Bool)
runChangeT

-- | Blindly run an updater.
{-# INLINE runUpdater #-}
runUpdater :: Updater a -> a -> (a, Bool)
runUpdater :: forall a. Updater a -> a -> (a, Bool)
runUpdater Updater a
f a
a = forall a. Change a -> (a, Bool)
runChange forall a b. (a -> b) -> a -> b
$ Updater a
f a
a

-- | Mark a computation as dirty.
dirty :: Monad m => UpdaterT m a
dirty :: forall (m :: * -> *) a. Monad m => UpdaterT m a
dirty a
a = do
  forall (m :: * -> *). MonadChange m => m ()
tellDirty
  forall (m :: * -> *) a. Monad m => a -> m a
return a
a

{-# SPECIALIZE ifDirty :: Change a -> (a -> Change b) -> (a -> Change b) -> Change b #-}
{-# SPECIALIZE ifDirty :: Identity a -> (a -> Identity b) -> (a -> Identity b) -> Identity b #-}
ifDirty :: (Monad m, MonadChange m) => m a -> (a -> m b) -> (a -> m b) -> m b
ifDirty :: forall (m :: * -> *) a b.
(Monad m, MonadChange m) =>
m a -> (a -> m b) -> (a -> m b) -> m b
ifDirty m a
m a -> m b
f a -> m b
g = do
  (a
a, Bool
dirty) <- forall (m :: * -> *) a. MonadChange m => m a -> m (a, Bool)
listenDirty m a
m
  if Bool
dirty then a -> m b
f a
a else a -> m b
g a
a

-- * Proper updater (Q-combinators)

-- | Replace result of updating with original input if nothing has changed.
sharing :: Monad m => UpdaterT m a -> UpdaterT m a
sharing :: forall (m :: * -> *) a. Monad m => UpdaterT m a -> UpdaterT m a
sharing UpdaterT m a
f a
a = do
  (a
a', Bool
changed) <- forall (m :: * -> *) a. MonadChange m => m a -> m (a, Bool)
listenDirty forall a b. (a -> b) -> a -> b
$ UpdaterT m a
f a
a
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if Bool
changed then a
a' else a
a

-- | Eval an updater (using 'sharing').
evalUpdater :: Updater a -> EndoFun a
evalUpdater :: forall a. Updater a -> EndoFun a
evalUpdater Updater a
f a
a = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. Change a -> (a, Bool)
runChange forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => UpdaterT m a -> UpdaterT m a
sharing Updater a
f a
a

-- END REAL STUFF

-- * Updater transformer classes

-- ** Unary (functors)

-- | Like 'Functor', but preserving sharing.
class Traversable f => Updater1 f where
  updater1 :: Updater a -> Updater (f a)
  updates1 :: Updater a -> Updater (f a) -- ^ @= sharing . updater1@
  update1  :: Updater a -> EndoFun (f a)

  updater1   = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse
  updates1 Updater a
f = forall (m :: * -> *) a. Monad m => UpdaterT m a -> UpdaterT m a
sharing forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Updater1 f => Updater a -> Updater (f a)
updater1 Updater a
f
  update1  Updater a
f = forall a. Updater a -> EndoFun a
evalUpdater forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Updater1 f => Updater a -> Updater (f a)
updater1 Updater a
f

instance Updater1 Maybe where

instance Updater1 [] where
  updater1 :: forall a. Updater a -> Updater [a]
updater1 Updater a
f []       = forall (m :: * -> *) a. Monad m => a -> m a
return []
  updater1 Updater a
f (a
x : [a]
xs) = (:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Updater a
f a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Updater1 f => Updater a -> Updater (f a)
updates1 Updater a
f [a]
xs

-- ** Binary (bifunctors)

-- | Like 'Bifunctor', but preserving sharing.
class Updater2 f where
  updater2 :: Updater a -> Updater b -> Updater (f a b)
  updates2 :: Updater a -> Updater b -> Updater (f a b)
  update2  :: Updater a -> Updater b -> EndoFun (f a b)

  updates2 Updater a
f1 Updater b
f2 = forall (m :: * -> *) a. Monad m => UpdaterT m a -> UpdaterT m a
sharing forall a b. (a -> b) -> a -> b
$ forall (f :: * -> * -> *) a b.
Updater2 f =>
Updater a -> Updater b -> Updater (f a b)
updater2 Updater a
f1 Updater b
f2
  update2  Updater a
f1 Updater b
f2 = forall a. Updater a -> EndoFun a
evalUpdater forall a b. (a -> b) -> a -> b
$ forall (f :: * -> * -> *) a b.
Updater2 f =>
Updater a -> Updater b -> Updater (f a b)
updater2 Updater a
f1 Updater b
f2

instance Updater2 (,) where
  updater2 :: forall a b. Updater a -> Updater b -> Updater (a, b)
updater2 Updater a
f1 Updater b
f2 (a
a,b
b) = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. Monad m => UpdaterT m a -> UpdaterT m a
sharing Updater a
f1 a
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) a. Monad m => UpdaterT m a -> UpdaterT m a
sharing Updater b
f2 b
b

instance Updater2 Either where
  updater2 :: forall a b. Updater a -> Updater b -> Updater (Either a b)
updater2 Updater a
f1 Updater b
f2 (Left a
a)  = forall a b. a -> Either a b
Left forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Updater a
f1 a
a
  updater2 Updater a
f1 Updater b
f2 (Right b
b) = forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Updater b
f2 b
b


{-- BEGIN MOCK

-- * Mock updater

type Change = Identity

-- | Replace result of updating with original input if nothing has changed.
{-# INLINE sharing #-}
sharing :: Updater a -> Updater a
sharing f a = f a

-- | Run an updater.
{-# INLINE evalUpdater #-}
evalUpdater :: Updater a -> EndoFun a
evalUpdater f a = runIdentity (f a)

-- | Mark a computation as dirty.
{-# INLINE dirty #-}
dirty :: Updater a
dirty = Identity

{-# INLINE ifDirty #-}
ifDirty :: Identity a -> (a -> Identity b) -> (a -> Identity b) -> Identity b
ifDirty m f g = m >>= f

-- END MOCK -}