content: {-# LANGUAGE CPP, LambdaCase, TupleSections, DeriveTraversable #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving, DerivingVia #-}
  {-# OPTIONS_GHC -fno-warn-unused-imports #-}
  -----------------------------------------------------------------------------
  -- |
  -- Module     : Control.Selective
  -- Copyright  : (c) Andrey Mokhov 2018-2024
  -- License    : MIT (see the file LICENSE)
  -- Maintainer : andrey.mokhov@gmail.com
  -- Stability  : experimental
  --
  -- This is a library for /selective applicative functors/, or just
  -- /selective functors/ for short, an abstraction between applicative functors
  -- and monads, introduced in this paper: https://dl.acm.org/doi/10.1145/3341694.
  --
  -----------------------------------------------------------------------------
  module Control.Selective (
      -- * Type class
      Selective (..), (<*?), branch, selectA, selectT, apS, selectM,
  
      -- * Conditional combinators
      ifS, whenS, fromMaybeS, orElse, andAlso, untilRight, whileS, (<||>), (<&&>),
      foldS, anyS, allS, bindS, Cases, casesEnum, cases, matchS, matchM,
  
      -- * Selective functors
      SelectA (..), SelectM (..), Over (..), Under (..), Validation (..),
  
      -- * Miscellaneous
      swapEither, ComposeEither (..), ComposeTraversable (..)
      ) where
  
  import Control.Applicative
  import Control.Applicative.Lift
  import Control.Arrow
  import Control.Monad.ST
  import Control.Monad.Trans.Cont
  import Control.Monad.Trans.Identity
  import Control.Monad.Trans.Maybe
  import Control.Monad.Trans.Reader
  import Control.Monad.Trans.RWS
  import Control.Monad.Trans.State
  import Control.Monad.Trans.Writer
  import Data.Bool
  import Data.Function
  import Data.Functor.Compose
  import Data.Functor.Identity
  import Data.Functor.Product
  import Data.List.NonEmpty
  import Data.Proxy
  import Data.Semigroup (Semigroup (..))
  import GHC.Conc (STM)
  
  import qualified Control.Monad.Trans.RWS.Strict    as S
  import qualified Control.Monad.Trans.State.Strict  as S
  import qualified Control.Monad.Trans.Writer.Strict as S
  
  -- | Selective applicative functors. You can think of 'select' as a selective
  -- function application: when given a value of type 'Left' @a@, you __must apply__
  -- the given function, but when given a 'Right' @b@, you __may skip__ the
  -- function and associated effects, and simply return the @b@.
  --
  -- Note that it is not a requirement for selective functors to skip unnecessary
  -- effects. It may be counterintuitive, but this makes them more useful. Why?
  -- Typically, when executing a selective computation, you would want to skip the
  -- effects (saving work); but on the other hand, if your goal is to statically
  -- analyse a given selective computation and extract the set of all possible
  -- effects (without actually executing them), then you do not want to skip any
  -- effects, because that defeats the purpose of static analysis.
  --
  -- The type signature of 'select' is reminiscent of both '<*>' and '>>=', and
  -- indeed a selective functor is in some sense a composition of an applicative
  -- functor and the 'Either' monad.
  --
  -- Laws:
  --
  -- * Identity:
  --
  -- @
  -- x \<*? pure id = either id id \<$\> x
  -- @
  --
  -- * Distributivity; note that @y@ and @z@ have the same type @f (a -> b)@:
  --
  -- @
  -- pure x \<*? (y *\> z) = (pure x \<*? y) *\> (pure x \<*? z)
  -- @
  --
  -- * Associativity:
  --
  -- @
  -- x \<*? (y \<*? z) = (f \<$\> x) \<*? (g \<$\> y) \<*? (h \<$\> z)
  --   where
  --     f x = Right \<$\> x
  --     g y = \a -\> bimap (,a) ($a) y
  --     h z = uncurry z
  -- @
  --
  -- * Monadic 'select' (for selective functors that are also monads):
  --
  -- @
  -- select = selectM
  -- @
  --
  -- There are also a few useful theorems:
  --
  -- * Apply a pure function to the result:
  --
  -- @
  -- f \<$\> select x y = select (fmap f \<$\> x) (fmap f \<$\> y)
  -- @
  --
  -- * Apply a pure function to the @Left@ case of the first argument:
  --
  -- @
  -- select (first f \<$\> x) y = select x ((. f) \<$\> y)
  -- @
  --
  -- * Apply a pure function to the second argument:
  --
  -- @
  -- select x (f \<$\> y) = select (first (flip f) \<$\> x) ((&) \<$\> y)
  -- @
  --
  -- * Generalised identity:
  --
  -- @
  -- x \<*? pure y = either y id \<$\> x
  -- @
  --
  -- * A selective functor is /rigid/ if it satisfies '<*>' @=@ 'apS'. The
  -- following /interchange/ law holds for rigid selective functors:
  --
  -- @
  -- x *\> (y \<*? z) = (x *\> y) \<*? z
  -- @
  --
  -- If f is also a 'Monad', we require that 'select' = 'selectM', from which one
  -- can prove '<*>' @=@ 'apS'.
  class Applicative f => Selective f where
      select :: f (Either a b) -> f (a -> b) -> f b
  
  {- Why do we have skew associativity, where we can reassociate effects to the
     left but not to the right?
  
     The following two tables, which list all possible combinations of effect
     execution and skipping, might give you some intuition on why this happens.
  
       ---------------
       (x <*? y) <*? z
       ---------------
        1     0      0
        1     1      0
        1     0      1
        1     1      1
  
       ---------------
       x <*? (y <*? z)
       ---------------
       1      0     0
       1      1     0
       1      1     1
  
     A key observation is that when effects are associated to the right, we can't
     skip the effect y and execute the effect z: combination 101 is impossible.
  -}
  
  -- | An operator alias for 'select', which is sometimes convenient. It tries to
  -- follow the notational convention for 'Applicative' operators. The angle
  -- bracket pointing to the left means we always use the corresponding value.
  -- The value on the right, however, may be skipped, hence the question mark.
  (<*?) :: Selective f => f (Either a b) -> f (a -> b) -> f b
  (<*?) = select
  
  infixl 4 <*?
  
  -- | The 'branch' function is a natural generalisation of 'select': instead of
  -- skipping an unnecessary effect, it chooses which of the two given effectful
  -- functions to apply to a given argument; the other effect is unnecessary. It
  -- is possible to implement 'branch' in terms of 'select', which is a good
  -- puzzle (give it a try!).
  --
  -- We can also implement 'select' via 'branch':
  --
  -- @
  -- selectB :: Selective f => f (Either a b) -> f (a -> b) -> f b
  -- selectB x y = branch x y (pure id)
  -- @
  --
  branch :: Selective f => f (Either a b) -> f (a -> c) -> f (b -> c) -> f c
  branch x l r = fmap (fmap Left) x <*? fmap (fmap Right) l <*? r
  
  -- | We can write a function with the type signature of 'select' using the
  -- 'Applicative' type class, but it will always execute the effects associated
  -- with the second argument, hence being potentially less efficient.
  selectA :: Applicative f => f (Either a b) -> f (a -> b) -> f b
  selectA x y = (\e f -> either f id e) <$> x <*> y
  
  -- | For traversable functors, we can implement 'select' in another interesting
  -- way: the effects associated with the second argument can be skipped as long
  -- as the first argument contains only 'Right' values.
  selectT :: Traversable f => f (Either a b) -> f (a -> b) -> f b
  selectT x y = case sequenceA x of
      Left  a  -> ($a) <$> y
      Right fb -> fb
  
  {-| Recover the application operator '<*>' from 'select'. /Rigid/ selective
  functors satisfy the law '<*>' @=@ 'apS' and furthermore, the resulting
  applicative functor satisfies all laws of 'Applicative':
  
  * Identity:
  
      > pure id <*> v = v
  
  * Homomorphism:
  
      > pure f <*> pure x = pure (f x)
  
  * Interchange:
  
      > u <*> pure y = pure ($y) <*> u
  
  * Composition:
  
      > (.) <$> u <*> v <*> w = u <*> (v <*> w)
  -}
  apS :: Selective f => f (a -> b) -> f a -> f b
  apS f x = select (Left <$> f) ((&) <$> x)
  
  -- | One can easily implement a monadic 'selectM' that satisfies the laws,
  -- hence any 'Monad' is 'Selective'.
  selectM :: Monad f => f (Either a b) -> f (a -> b) -> f b
  selectM x y = x >>= \case Left  a -> ($a) <$> y -- execute y
                            Right b -> pure b     -- skip y
  
  -- Many useful 'Monad' combinators can be implemented with 'Selective'
  
  -- | Branch on a Boolean value, skipping unnecessary effects.
  ifS :: Selective f => f Bool -> f a -> f a -> f a
  ifS x t e = branch (bool (Right ()) (Left ()) <$> x) (const <$> t) (const <$> e)
  
  -- Implementation used in the paper:
  -- ifS x t e = branch selector (fmap const t) (fmap const e)
  --   where
  --     selector = bool (Right ()) (Left ()) <$> x -- NB: convert True to Left ()
  
  -- | Eliminate a specified value @a@ from @f (Either a b)@ by replacing it
  -- with a given @f b@.
  eliminate :: (Eq a, Selective f) => a -> f b -> f (Either a b) -> f (Either a b)
  eliminate x fb fa = select (match x <$> fa) (const . Right <$> fb)
    where
      match _ (Right y) = Right (Right y)
      match x (Left  y) = if x == y then Left () else Right (Left y)
  
  -- | A list of values, equipped with a fast membership test.
  data Cases a = Cases [a] (a -> Bool)
  
  -- | The list of all possible values of an enumerable data type.
  casesEnum :: (Bounded a, Enum a) => Cases a
  casesEnum = Cases [minBound..maxBound] (const True)
  
  -- | Embed a list of values into 'Cases' using the trivial but slow membership
  -- test based on 'elem'.
  cases :: Eq a => [a] -> Cases a
  cases as = Cases as (`elem` as)
  
  -- | Eliminate all specified values @a@ from @f (Either a b)@ by replacing each
  -- of them with a given @f a@.
  matchS :: (Eq a, Selective f) => Cases a -> f a -> (a -> f b) -> f (Either a b)
  matchS (Cases cs _) x f = foldr (\c -> eliminate c (f c)) (Left <$> x) cs
  
  -- | Eliminate all specified values @a@ from @f (Either a b)@ by replacing each
  -- of them with a given @f a@.
  matchM :: Monad m => Cases a -> m a -> (a -> m b) -> m (Either a b)
  matchM (Cases _ p) mx f = do
      x <- mx
      if p x then Right <$> f x else pure (Left x)
  
  -- TODO: Add a type-safe version based on @KnownNat@.
  -- | A restricted version of monadic bind. Fails with an error if the 'Bounded'
  -- and 'Enum' instances for @a@ do not cover all values of @a@.
  bindS :: (Bounded a, Enum a, Eq a, Selective f) => f a -> (a -> f b) -> f b
  bindS x f = fromRight <$> matchS casesEnum x f
    where
      fromRight (Right b) = b
      fromRight _ = error "Selective.bindS: incorrect Bounded and/or Enum instance"
  
  -- | Conditionally perform an effect.
  whenS :: Selective f => f Bool -> f () -> f ()
  whenS x y = select (bool (Right ()) (Left ()) <$> x) (const <$> y)
  
  -- Implementation used in the paper:
  -- whenS x y = selector <*? effect
  --   where
  --     selector = bool (Right ()) (Left ()) <$> x -- NB: maps True to Left ()
  --     effect   = const                     <$> y
  
  -- | A lifted version of 'Data.Maybe.fromMaybe'.
  fromMaybeS :: Selective f => f a -> f (Maybe a) -> f a
  fromMaybeS fNothing fMaybe =
      select (toEither <$> fMaybe) (toUnitFunction <$> fNothing)
    where
      toEither :: Maybe a -> Either () a
      toEither Nothing  = Left ()
      toEither (Just a) = Right a
      toUnitFunction :: a -> () -> a
      toUnitFunction x () = x
  
  -- | Return the first @Right@ value. If both are @Left@'s, accumulate errors.
  orElse :: (Selective f, Semigroup e) => f (Either e a) -> f (Either e a) -> f (Either e a)
  orElse x y = select (prepare <$> x) (combine <$> y)
    where
      prepare :: Either e a -> Either e (Either e a)
      prepare = fmap Right
  
      combine :: Semigroup e => Either e a -> e -> Either e a
      combine (Left ey) ex = Left (ex <> ey)
      combine (Right a) _  = Right a
  
  -- | Accumulate the @Right@ values, or return the first @Left@.
  andAlso :: (Selective f, Semigroup a) => f (Either e a) -> f (Either e a) -> f (Either e a)
  andAlso x y = swapEither <$> orElse (swapEither <$> x) (swapEither <$> y)
  
  -- | Swap @Left@ and @Right@.
  swapEither :: Either a b -> Either b a
  swapEither = either Right Left
  
  -- | Keep checking an effectful condition while it holds.
  whileS :: Selective f => f Bool -> f ()
  whileS act = whenS act (whileS act)
  
  -- | Keep running an effectful computation until it returns a @Right@ value,
  -- collecting the @Left@'s using a supplied @Monoid@ instance.
  untilRight :: (Monoid a, Selective f) => f (Either a b) -> f (a, b)
  untilRight x = select y h
    where
      -- y :: f (Either a (a, b))
      y = fmap (mempty,) <$> x
      -- h :: f (a -> (a, b))
      h = (\(as, b) a -> (mappend a as, b)) <$> untilRight x
  
  -- | A lifted version of lazy Boolean OR.
  (<||>) :: Selective f => f Bool -> f Bool -> f Bool
  a <||> b = ifS a (pure True) b
  
  -- | A lifted version of lazy Boolean AND.
  (<&&>) :: Selective f => f Bool -> f Bool -> f Bool
  a <&&> b = ifS a b (pure False)
  
  -- | A lifted version of 'any'. Retains the short-circuiting behaviour.
  anyS :: Selective f => (a -> f Bool) -> [a] -> f Bool
  anyS p = foldr ((<||>) . p) (pure False)
  
  -- | A lifted version of 'all'. Retains the short-circuiting behaviour.
  allS :: Selective f => (a -> f Bool) -> [a] -> f Bool
  allS p = foldr ((<&&>) . p) (pure True)
  
  -- | Generalised folding with the short-circuiting behaviour.
  foldS :: (Selective f, Foldable t, Monoid a) => t (f (Either e a)) -> f (Either e a)
  foldS = foldr andAlso (pure (Right mempty))
  
  -- Instances
  
  -- | Any applicative functor can be given a 'Selective' instance by defining
  -- 'select' @=@ 'selectA'. This data type captures this pattern, so you can use
  -- it in combination with the @DerivingVia@ extension as follows:
  --
  -- @
  -- newtype Over m a = Over m
  --     deriving (Functor, Applicative, Selective) via SelectA (Const m)
  -- @
  --
  newtype SelectA f a = SelectA { getSelectA :: f a }
      deriving (Functor, Applicative)
  
  instance Applicative f => Selective (SelectA f) where
      select = selectA
  
  -- Note: Validation e a ~ Lift (Under e) a
  instance Selective f => Selective (Lift f) where
      select (Pure (Right x)) _         = Pure x -- Lazy in the second argument
      select              x   (Pure  y) = either y id <$> x
      select (Pure (Left  x)) (Other y) = Other $ ($x) <$> y
      select (Other       x ) (Other y) = Other $   x  <*? y
  
  -- | Any monad can be given a 'Selective' instance by defining
  -- 'select' @=@ 'selectM'. This data type captures this pattern, so you can use
  -- it in combination with the @DerivingVia@ extension as follows:
  --
  -- @
  -- newtype V1 a = V1 a
  --     deriving (Functor, Applicative, Selective, Monad) via SelectM Identity
  -- @
  --
  newtype SelectM f a = SelectM { getSelectM :: f a }
      deriving (Functor, Applicative, Monad)
  
  instance Monad f => Selective (SelectM f) where
      select = selectM
  
  -- | Static analysis of selective functors with over-approximation.
  newtype Over m a = Over { getOver :: m }
      deriving (Eq, Functor, Ord, Show)
      deriving Applicative via (Const m)
  
  -- select = selectA
  instance Monoid m => Selective (Over m) where
      select (Over x) (Over y) = Over (mappend x y)
  
  -- | Static analysis of selective functors with under-approximation.
  newtype Under m a = Under { getUnder :: m }
      deriving (Eq, Functor, Ord, Show, Foldable, Traversable)
      deriving Applicative via (Const m)
  
  -- select = selectT
  instance Monoid m => Selective (Under m) where
      select (Under m) _ = Under m
  
  -- The 'Selective' 'ZipList' instance corresponds to the SIMT execution model.
  -- Quoting https://en.wikipedia.org/wiki/Single_instruction,_multiple_threads:
  --
  -- "...to handle an IF-ELSE block where various threads of a processor execute
  -- different paths, all threads must actually process both paths (as all threads
  -- of a processor always execute in lock-step), but masking is used to disable
  -- and enable the various threads as appropriate..."
  instance Selective ZipList where select = selectA
  
  -- | Selective instance for the standard applicative functor Validation. This is
  -- a good example of a non-trivial selective functor which is not a monad.
  data Validation e a = Failure e | Success a deriving (Eq, Functor, Ord, Show)
  
  instance Semigroup e => Applicative (Validation e) where
      pure = Success
      Failure e1 <*> Failure e2 = Failure (e1 <> e2)
      Failure e1 <*> Success _  = Failure e1
      Success _  <*> Failure e2 = Failure e2
      Success f  <*> Success a  = Success (f a)
  
  instance Semigroup e => Selective (Validation e) where
      select (Success (Left  a)) f = ($a) <$> f
      select (Success (Right b)) _ = Success b
      select (Failure e        ) _ = Failure e
  
  instance (Selective f, Selective g) => Selective (Product f g) where
      select (Pair fx gx) (Pair fy gy) = Pair (select fx fy) (select gx gy)
  
  instance Selective f => Selective (IdentityT f) where
      select (IdentityT x) (IdentityT y) = IdentityT (select x y)
  
  instance Selective f => Selective (ReaderT env f) where
      select (ReaderT x) (ReaderT y) = ReaderT $ \env -> select (x env) (y env)
  
  distributeEither :: (Either a b, w) -> Either (a, w) (b, w)
  distributeEither (Left  a, w) = Left  (a, w)
  distributeEither (Right b, w) = Right (b, w)
  
  distributeFunction :: Monoid w => (a -> b, w) -> (a, w) -> (b, w)
  distributeFunction (f, wf) (x, wx) = (f x, mappend wx wf)
  
  instance (Monoid w, Selective f) => Selective (WriterT w f) where
      select (WriterT x) (WriterT f) =
          WriterT $ select (distributeEither <$> x) (distributeFunction <$> f)
  
  instance (Monoid w, Selective f) => Selective (S.WriterT w f) where
      select (S.WriterT x) (S.WriterT f) =
          S.WriterT $ select (distributeEither <$> x) (distributeFunction <$> f)
  
  -- TODO: Is this a useful instance? Note that composition of 'Alternative'
  -- requires @f@ to be 'Alternative', and @g@ to be 'Applicative', which is
  -- opposite to what we have here:
  --
  -- instance (Alternative f, Applicative g) => Alternative (Compose f g) ...
  --
  instance (Applicative f, Selective g) => Selective (Compose f g) where
      select (Compose x) (Compose y) = Compose (select <$> x <*> y)
  
  {- Here is why composing selective functors is tricky.
  
  Consider @Compose Maybe IO@. The only sensible implementation is:
  
  > select :: Maybe (IO (Either a b)) -> Maybe (IO (a -> b)) -> Maybe (IO b)
  > select Nothing  _        = Nothing
  > select (Just x) (Just y) = Just (select x y)
  > select (Just x) Nothing  = Nothing -- Can't use Just: we don't have @a -> b@!
  
  In other words, we have to be 'Applicative' on the outside functor 'Maybe',
  because there is no way to peek inside 'IO', which forces us to statically
  choose between 'Just', which doesn't work since we have no function @a -> b@,
  and 'Nothing' which corresponds to the behaviour of 'selectA'.
  
  -}
  
  -- Monad instances
  
  -- As a quick experiment, try: ifS (pure True) (print 1) (print 2)
  instance Selective IO where select = selectM
  
  -- And... we need to write a lot more instances
  instance             Selective []         where select = selectM
  instance Monoid a => Selective ((,) a)    where select = selectM
  instance             Selective ((->) a)   where select = selectM
  instance             Selective (Either e) where select = selectM
  instance             Selective Identity   where select = selectM
  instance             Selective Maybe      where select = selectM
  instance             Selective NonEmpty   where select = selectM
  instance             Selective Proxy      where select = selectM
  instance             Selective (ST s)     where select = selectM
  instance             Selective STM        where select = selectM
  
  instance                        Selective (ContT      r m) where select = selectM
  instance            Monad m  => Selective (MaybeT       m) where select = selectM
  instance (Monoid w, Monad m) => Selective (RWST   r w s m) where select = selectM
  instance (Monoid w, Monad m) => Selective (S.RWST r w s m) where select = selectM
  instance            Monad m  => Selective (StateT     s m) where select = selectM
  instance            Monad m  => Selective (S.StateT   s m) where select = selectM
  
  ------------------------------------ Arrows ------------------------------------
  -- See the following standard definitions in "Control.Arrow".
  -- newtype ArrowMonad a o = ArrowMonad (a () o)
  -- instance Arrow a => Functor (ArrowMonad a)
  -- instance Arrow a => Applicative (ArrowMonad a)
  
  instance ArrowChoice a => Selective (ArrowMonad a) where
      select (ArrowMonad x) y = ArrowMonad $ x >>> (toArrow y ||| returnA)
  
  toArrow :: Arrow a => ArrowMonad a (i -> o) -> a i o
  toArrow (ArrowMonad f) = arr ((),) >>> first f >>> arr (uncurry ($))
  
  ------------------------------ ComposeTraversable ------------------------------
  -- | Composition of a selective functor @f@ and an applicative traversable
  -- functor @g@.
  newtype ComposeTraversable f g a = ComposeTraversable (f (g a))
      deriving (Functor, Applicative) via Compose f g
  
  instance (Selective f, Applicative g, Traversable g) => Selective (ComposeTraversable f g) where
      select (ComposeTraversable x) (ComposeTraversable f) = ComposeTraversable $
          select (prepare <$> x) (combine <$> f)
        where
          prepare :: Traversable g => g (Either a b) -> Either a (g b)
          prepare = sequenceA
  
          combine :: Traversable g => g (a -> b) -> a -> g b
          combine = sequenceA
  
  --------------------------------- ComposeEither --------------------------------
  -- | Composition of a selective functor @f@ with the 'Either' monad.
  newtype ComposeEither f e a = ComposeEither (f (Either e a))
      deriving Functor via Compose f (Either e)
      deriving Selective via ComposeTraversable f (Either e)
  
  instance Selective f => Applicative (ComposeEither f e) where
      pure = ComposeEither . pure . Right
  
      ComposeEither f <*> ComposeEither a = ComposeEither $
          select (prepare <$> f) (combine <$> a)
        where
          prepare :: Either e (a -> b) -> Either (a -> b) (Either e b)
          prepare = either (Right . Left) Left
  
          combine :: Either e a -> (a -> b) -> Either e b
          combine = flip fmap
  
  ---------------------------------- Alternative ---------------------------------
  instance (Selective f, Monoid e) => Alternative (ComposeEither f e) where
      empty                               = ComposeEither (pure $ Left mempty)
      ComposeEither x <|> ComposeEither y = ComposeEither (x `orElse` y)
  
  {- One could also try implementing 'select' via 'Alternative' as follows:
  
  selectAlt :: Alternative f => f (Either a b) -> f (a -> b) -> f b
  selectAlt x y = failIfLeft x <|> selectA x y
    where
      failIfLeft :: Alternative f => f (Either a b) -> f b
      failIfLeft = undefined
  
  This has two issues:
  
  1) A generic 'failIfLeft' if not possible, although many actual instances should
     be able to implement it.
  
  2) More importantly, this requires duplication of work: if we failed because we
     happened to parse a 'Left' value in the first parser, then we need to rerun
     it, obtain a 'Left' once again, and then execute the second parser. Again, a
     specific instance may be able to cache the result and reuse it without
     duplicating any work, but this does not seem to be possible to achieve
     generically for any Alternative.
  
  -}
  