-- |
-- This is an internal module. You probably don't need to import this. Use
-- "Data.Seqn.Seq", "Data.Seqn.MSeq", or "Data.Seqn.PQueue" instead.
--
-- The only reason to use this module is to use the constructs defined here with
-- other internal modules.
--
module Data.Seqn.Internal.Util
  ( Biapplicative(..)
  , S2(..)
  , S3(..)
  , SMaybe(..)
  , Tagged(..)
  , SStateT(..)
  , evalSStateT
  , SState
  , sState
  , evalSState
  , (#.)
  ) where

import qualified Control.Applicative -- for before liftA2 in Prelude
import Data.Bifunctor (Bifunctor(..))
import Data.Coerce (Coercible, coerce)
import Data.Functor.Const (Const(..))
import Data.Functor.Identity (Identity(..))

class Bifunctor p => Biapplicative p where
  bipure :: a -> b -> p a b
  biliftA2 :: (a -> b -> c) -> (d -> e -> f) -> p a d -> p b e -> p c f

instance Biapplicative Const where
  bipure :: forall a b. a -> b -> Const a b
bipure a
x b
_ = a -> Const a b
forall a b. Coercible a b => a -> b
coerce a
x
  biliftA2 :: forall a b c d e f.
(a -> b -> c)
-> (d -> e -> f) -> Const a d -> Const b e -> Const c f
biliftA2 a -> b -> c
f d -> e -> f
_ = (a -> b -> c) -> Const a d -> Const b e -> Const c f
forall a b. Coercible a b => a -> b
coerce a -> b -> c
f

data S2 a b = S2 !a !b

instance Functor (S2 a) where
  fmap :: forall a b. (a -> b) -> S2 a a -> S2 a b
fmap a -> b
f (S2 a
x a
y) = a -> b -> S2 a b
forall a b. a -> b -> S2 a b
S2 a
x (a -> b
f a
y)

instance Bifunctor S2 where
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> S2 a c -> S2 b d
bimap a -> b
f c -> d
g (S2 a
x c
y) = b -> d -> S2 b d
forall a b. a -> b -> S2 a b
S2 (a -> b
f a
x) (c -> d
g c
y)

instance Biapplicative S2 where
  bipure :: forall a b. a -> b -> S2 a b
bipure = a -> b -> S2 a b
forall a b. a -> b -> S2 a b
S2
  biliftA2 :: forall a b c d e f.
(a -> b -> c) -> (d -> e -> f) -> S2 a d -> S2 b e -> S2 c f
biliftA2 a -> b -> c
f d -> e -> f
g (S2 a
x1 d
y1) (S2 b
x2 e
y2) = c -> f -> S2 c f
forall a b. a -> b -> S2 a b
S2 (a -> b -> c
f a
x1 b
x2) (d -> e -> f
g d
y1 e
y2)

data S3 a b c = S3 !a !b !c

data SMaybe a
  = SNothing
  | SJust !a

newtype Tagged a b = Tagged { forall a b. Tagged a b -> b
unTagged :: b }

instance Functor (Tagged a) where
  fmap :: forall a b. (a -> b) -> Tagged a a -> Tagged a b
fmap = (a -> b) -> Tagged a a -> Tagged a b
forall a b. Coercible a b => a -> b
coerce

instance Bifunctor Tagged where
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> Tagged a c -> Tagged b d
bimap a -> b
_ = (c -> d) -> Tagged a c -> Tagged b d
forall a b. Coercible a b => a -> b
coerce

instance Biapplicative Tagged where
  bipure :: forall a b. a -> b -> Tagged a b
bipure a
_ = b -> Tagged a b
forall a b. Coercible a b => a -> b
coerce
  biliftA2 :: forall a b c d e f.
(a -> b -> c)
-> (d -> e -> f) -> Tagged a d -> Tagged b e -> Tagged c f
biliftA2 a -> b -> c
_ = (d -> e -> f) -> Tagged a d -> Tagged b e -> Tagged c f
forall a b. Coercible a b => a -> b
coerce

-- Strict in the state, value, and bind.
newtype SStateT s m a = SStateT { forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
runSStateT :: s -> m (S2 s a) }

evalSStateT :: Functor m => SStateT s m a -> s -> m a
evalSStateT :: forall (m :: * -> *) s a. Functor m => SStateT s m a -> s -> m a
evalSStateT SStateT s m a
m s
s = (S2 s a -> a) -> m (S2 s a) -> m a
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(S2 s
_ a
x) -> a
x) (SStateT s m a -> s -> m (S2 s a)
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
runSStateT SStateT s m a
m s
s)
{-# INLINE evalSStateT #-}

type SState s = SStateT s Identity

sState :: (s -> S2 s a) -> SState s a
sState :: forall s a. (s -> S2 s a) -> SState s a
sState = (s -> S2 s a) -> SState s a
forall a b. Coercible a b => a -> b
coerce

evalSState :: SState s a -> s -> a
evalSState :: forall s a. SState s a -> s -> a
evalSState SState s a
m s
s = case SState s a -> s -> Identity (S2 s a)
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
runSStateT SState s a
m s
s of Identity (S2 s
_ a
x) -> a
x
{-# INLINE evalSState #-}

instance Functor m => Functor (SStateT s m) where
  fmap :: forall a b. (a -> b) -> SStateT s m a -> SStateT s m b
fmap a -> b
f SStateT s m a
m = (s -> m (S2 s b)) -> SStateT s m b
forall s (m :: * -> *) a. (s -> m (S2 s a)) -> SStateT s m a
SStateT ((s -> m (S2 s b)) -> SStateT s m b)
-> (s -> m (S2 s b)) -> SStateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s -> ((S2 s a -> S2 s b) -> m (S2 s a) -> m (S2 s b)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((S2 s a -> S2 s b) -> m (S2 s a) -> m (S2 s b))
-> ((a -> b) -> S2 s a -> S2 s b)
-> (a -> b)
-> m (S2 s a)
-> m (S2 s b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> S2 s a -> S2 s b
forall a b. (a -> b) -> S2 s a -> S2 s b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) a -> b
f (SStateT s m a -> s -> m (S2 s a)
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
runSStateT SStateT s m a
m s
s)
  {-# INLINE fmap #-}

instance Monad m => Applicative (SStateT s m) where
  pure :: forall a. a -> SStateT s m a
pure a
x = (s -> m (S2 s a)) -> SStateT s m a
forall s (m :: * -> *) a. (s -> m (S2 s a)) -> SStateT s m a
SStateT ((s -> m (S2 s a)) -> SStateT s m a)
-> (s -> m (S2 s a)) -> SStateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> S2 s a -> m (S2 s a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (S2 s a -> m (S2 s a)) -> S2 s a -> m (S2 s a)
forall a b. (a -> b) -> a -> b
$ s -> a -> S2 s a
forall a b. a -> b -> S2 a b
S2 s
s a
x
  {-# INLINE pure #-}

  liftA2 :: forall a b c.
(a -> b -> c) -> SStateT s m a -> SStateT s m b -> SStateT s m c
liftA2 a -> b -> c
f SStateT s m a
m1 SStateT s m b
m2 = (s -> m (S2 s c)) -> SStateT s m c
forall s (m :: * -> *) a. (s -> m (S2 s a)) -> SStateT s m a
SStateT ((s -> m (S2 s c)) -> SStateT s m c)
-> (s -> m (S2 s c)) -> SStateT s m c
forall a b. (a -> b) -> a -> b
$ \s
s -> do
    S2 s
s1 a
x <- SStateT s m a -> s -> m (S2 s a)
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
runSStateT SStateT s m a
m1 s
s
    S2 s
s2 b
y <- SStateT s m b -> s -> m (S2 s b)
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
runSStateT SStateT s m b
m2 s
s1
    S2 s c -> m (S2 s c)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (S2 s c -> m (S2 s c)) -> S2 s c -> m (S2 s c)
forall a b. (a -> b) -> a -> b
$ s -> c -> S2 s c
forall a b. a -> b -> S2 a b
S2 s
s2 (a -> b -> c
f a
x b
y)
  {-# INLINE liftA2 #-}

-- Borrow a trick from base in case GHC has trouble optimizing certain
-- function coercions.
--
-- See Note [Function coercion] in
-- https://gitlab.haskell.org/ghc/ghc/-/blob/8d67f247c3e4ca3810712654e1becbf927405f6b/libraries/ghc-internal/src/GHC/Internal/Data/Functor/Utils.hs#L134-160
(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c)
#. :: forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
(#.) b -> c
_ = (a -> b) -> a -> c
forall a b. Coercible a b => a -> b
coerce
{-# INLINE (#.) #-}

infixr 9 #.