    {-# LANGUAGE GeneralizedNewtypeDeriving #-}

    module Agda.Utils.Update
      ( Change
      , MonadChange(..)
      , runChange
      , Updater
      , sharing
      , runUpdater
      , dirty
      , ifDirty
      , Updater1(..)
      , Updater2(..)
      ) where

    import Control.Monad.Identity
    import Control.Monad.Trans
    import Control.Monad.Writer.Strict

    import Data.Traversable (Traversable(..), traverse)

    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 { fromChangeT :: WriterT Any m a }
      deriving (Functor, Applicative, Monad, MonadTrans)

    instance Monad m => MonadChange (ChangeT m) where
      tellDirty     = ChangeT $ tell $ Any True
      listenDirty m = ChangeT $ do
        (a, Any dirty) <- listen (fromChangeT m)
        return (a, dirty)

    -- | A mock change monad.
    instance MonadChange Identity where
      tellDirty                = Identity ()
      listenDirty (Identity a) = Identity (a, True)

    -- * Pure endo function and updater

    type EndoFun a = a -> a
    type Updater a = a -> Change a

    -- BEGIN REAL STUFF

    -- | The @Change@ monad.
    newtype Change a = Change { fromChange :: Writer Any a }
      deriving (Functor, Applicative, Monad)

    instance MonadChange Change where
      tellDirty = Change $ tell $ Any True
      listenDirty m = Change $ do
        (a, Any dirty) <- listen (fromChange m)
        return (a, dirty)

    -- | Run a 'Change' computation, returning result plus change flag.
    runChange :: Change a -> (a, Bool)
    runChange = mapSnd getAny . runWriter . fromChange

    -- | Blindly run an updater.
    runUpdater :: Updater a -> a -> (a, Bool)
    runUpdater f a = runChange $ f a

    -- | Mark a computation as dirty.
    dirty :: Updater a
    dirty a = do
      tellDirty
      return 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 :: MonadChange m => m a -> (a -> m b) -> (a -> m b) -> m b
    ifDirty m f g = do
      (a, dirty) <- listenDirty m
      if dirty then f a else g a

    -- * Proper updater (Q-combinators)

    -- | Replace result of updating with original input if nothing has changed.
    sharing :: Updater a -> Updater a
    sharing f a = do
      (a', changed) <- listenDirty $ f a
      return $ if changed then a' else a

    -- | Eval an updater (using 'sharing').
    evalUpdater :: Updater a -> EndoFun a
    evalUpdater f a = fst $ runWriter $ fromChange $ sharing f 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   = traverse
      updates1 f = sharing $ updater1 f
      update1  f = evalUpdater $ updater1 f

    instance Updater1 Maybe where

    instance Updater1 [] where
      updater1 f []       = return []
      updater1 f (x : xs) = (:) <$> f x <*> updates1 f 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 f1 f2 = sharing $ updater2 f1 f2
      update2  f1 f2 = evalUpdater $ updater2 f1 f2

    instance Updater2 (,) where
      updater2 f1 f2 (a,b) = (,) <$> sharing f1 a <*> sharing f2 b

    instance Updater2 Either where
      updater2 f1 f2 (Left a)  = Left <$> f1 a
      updater2 f1 f2 (Right b) = Right <$> f2 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 -}