{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}

{-# OPTIONS_HADDOCK hide #-}

module Capability.State.Internal.Class
  ( HasState(..)
  , get
  , put
  , state
  , modify
  , modify'
  , gets
  , zoom
  ) where

import Capability.Constraints
import Capability.Derive (derive)
import Capability.Source.Internal.Class
import Capability.Sink.Internal.Class
import Data.Coerce (Coercible)
import GHC.Exts (Proxy#, proxy#)

-- | State capability
--
-- An instance should fulfill the following laws.
-- At this point these laws are not definitive,
-- see <https://github.com/haskell/mtl/issues/5>.
--
-- prop> get @t >>= \s1 -> get @t >>= \s2 -> pure (s1, s2) = get @t >>= \s -> pure (s, s)
-- prop> get @t >>= \_ -> put @t s = put @t s
-- prop> put @t s1 >> put @t s2 = put @t s2
-- prop> put @t s >> get @t = put @t s >> pure s
-- prop> state @t f = get @t >>= \s -> let (a, s') = f s in put @t s' >> pure a
class (Monad m, HasSource tag s m, HasSink tag s m)
  => HasState (tag :: k) (s :: *) (m :: * -> *) | tag m -> s
  where
    -- | For technical reasons, this method needs an extra proxy argument.
    -- You only need it if you are defining new instances of 'HasState.
    -- Otherwise, you will want to use 'state'.
    -- See 'state' for more documentation.
    state_ :: Proxy# tag -> (s -> (a, s)) -> m a

-- | @get \@tag@
-- retrieve the current state of the state capability @tag@.
get :: forall tag s m. HasState tag s m => m s
get = await @tag
{-# INLINE get #-}

-- | @put \@tag s@
-- replace the current state of the state capability @tag@ with @s@.
put :: forall tag s m. HasState tag s m => s -> m ()
put = yield @tag
{-# INLINE put #-}

-- | @state \@tag f@
-- lifts a pure state computation @f@ to a monadic action in an arbitrary
-- monad @m@ with capability @HasState@.
--
-- Given the current state @s@ of the state capability @tag@
-- and @(a, s') = f s@, update the state to @s'@ and return @a@.
state :: forall tag s m a. HasState tag s m => (s -> (a, s)) -> m a
state = state_ (proxy# @_ @tag)
{-# INLINE state #-}

-- | @modify \@tag f@
-- given the current state @s@ of the state capability @tag@
-- and @s' = f s@, updates the state of the capability @tag@ to @s'@.
modify :: forall tag s m. HasState tag s m => (s -> s) -> m ()
modify f = state @tag $ \s -> ((), f s)
{-# INLINE modify #-}

-- | Same as 'modify' but strict in the new state.
modify' :: forall tag s m. HasState tag s m => (s -> s) -> m ()
modify' f = do
  s' <- get @tag
  put @tag $! f s'
{-# INLINE modify' #-}

-- | @gets \@tag f@
-- retrieves the image, by @f@ of the current state
-- of the state capability @tag@.
--
-- prop> gets @tag f = f <$> get @tag
gets :: forall tag s m a. HasState tag s m => (s -> a) -> m a
gets f = do
  s <- get @tag
  pure (f s)
{-# INLINE gets #-}

-- | Execute the given state action on a sub-component of the current state as
-- defined by the given transformer @t@. The set of retained capabilities must
-- be passed as @cs. If no capabilities are required,
-- 'Capabilities.Constraints.None' can be used.
--
-- Examples:
--
-- > foo :: HasState "foo" Int m => m ()
-- > zoom @"foo" @(Field "foo" "foobar") @None foo
-- >   :: (HasField' "foobar" record Int, HasState "foobar" record m) => m ()
-- >
-- > zoom @"foo" @(Field "foo" "foobar") @('[MonadIO]) bar
-- >   :: ( HasField' "foobar" record Int, HasState "foobar" record m
-- >      , MonadIO m) => m ()
-- >
-- > foo :: HasState "foo" Int m => m ()
-- > bar :: (MonadIO m, HasState "foo" Int m) => m ()
--
-- Note: the 'Data.Generics.Product.Fields.HasField'' constraint comes from the
-- @generic-lens@ package.
--
-- This function is experimental and subject to change.
-- See <https://github.com/tweag/capability/issues/46>.
zoom :: forall innertag t (cs :: [Capability]) inner m a.
  ( forall x. Coercible (t m x) (m x)
  , HasState innertag inner (t m)
  , All cs m )
  => (forall m'. All (HasState innertag inner ': cs) m' => m' a) -> m a
zoom =
  derive @t @'[HasState innertag inner] @cs
{-# INLINE zoom #-}