-------------------------------------------------------------------------
-- |
-- Module      : Control.Monad.Logic.Class
-- Copyright   : (c) 2007-2014 Dan Doel,
--               (c) 2011-2013 Edward Kmett,
--               (c) 2014      Roman Cheplyaka,
--               (c) 2020-2021 Andrew Lelechenko,
--               (c) 2020-2021 Kevin Quick
-- License     : BSD3
-- Maintainer  : Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Adapted from the paper
-- <http://okmij.org/ftp/papers/LogicT.pdf Backtracking, Interleaving, and Terminating Monad Transformers>
-- by Oleg Kiselyov, Chung-chieh Shan, Daniel P. Friedman, Amr Sabry.
-- Note that the paper uses 'MonadPlus' vocabulary
-- ('mzero' and 'mplus'),
-- while examples below prefer 'empty' and '<|>'
-- from 'Alternative'.
-------------------------------------------------------------------------

{-# LANGUAGE CPP #-}

#if __GLASGOW_HASKELL__ >= 704
{-# LANGUAGE Safe #-}
#endif

module Control.Monad.Logic.Class (MonadLogic(..), reflect) where

import Control.Applicative
import Control.Monad
import Control.Monad.Reader (ReaderT(..))
import Control.Monad.Trans (MonadTrans(..))
import qualified Control.Monad.State.Lazy as LazyST
import qualified Control.Monad.State.Strict as StrictST

-- | A backtracking, logic programming monad.
class (Monad m, Alternative m) => MonadLogic m where
    -- | Attempts to __split__ the computation, giving access to the first
    --   result. Satisfies the following laws:
    --
    --   > msplit empty          == pure Nothing
    --   > msplit (pure a <|> m) == pure (Just (a, m))
    msplit     :: m a -> m (Maybe (a, m a))

    -- | __Fair disjunction.__ It is possible for a logical computation
    --   to have an infinite number of potential results, for instance:
    --
    --   > odds = pure 1 <|> fmap (+ 2) odds
    --
    --   Such computations can cause problems in some circumstances. Consider:
    --
    --   > two = do x <- odds <|> pure 2
    --   >          if even x then pure x else empty
    --
    --   >>> observe two
    --   ...never completes...
    --
    --   Such a computation may never consider 'pure' @2@, and
    --   therefore even 'Control.Monad.Logic.observe' @two@ will
    --   never return any results. By
    --   contrast, using 'interleave' in place of
    --   'Control.Applicative.<|>' ensures fair consideration of both
    --   branches of a disjunction.
    --
    --   > fairTwo = do x <- odds `interleave` pure 2
    --   >              if even x then pure x else empty
    --
    --   >>> observe fairTwo
    --   2
    --
    --   Note that even with 'interleave' this computation will never
    --   terminate after returning 2: only the first value can be
    --   safely observed, after which each odd value becomes 'Control.Applicative.empty'
    --   (equivalent to
    --   <http://lpn.swi-prolog.org/lpnpage.php?pagetype=html&pageid=lpn-htmlse45 Prolog's fail>)
    --   which does not stop the evaluation but indicates there is no
    --   value to return yet.
    --
    --   Unlike '<|>', 'interleave' is not associative:
    --
    --   >>> let x = [1,2,3]; y = [4,5,6]; z = [7,8,9] :: [Int]
    --   >>> x `interleave` y
    --   [1,4,2,5,3,6]
    --   >>> (x `interleave` y) `interleave` z
    --   [1,7,4,8,2,9,5,3,6]
    --   >>> y `interleave` z
    --   [4,7,5,8,6,9]
    --   >>> x `interleave` (y `interleave` z)
    --   [1,4,2,7,3,5,8,6,9]
    --
    interleave :: m a -> m a -> m a

    -- | __Fair conjunction.__ Similarly to the previous function, consider
    --   the distributivity law, naturally expected from 'MonadPlus':
    --
    --   > (a <|> b) >>= k = (a >>= k) <|> (b >>= k)
    --
    --   If @a@ '>>=' @k@ can backtrack arbitrarily many times, @b@ '>>=' @k@
    --   may never be considered. In logic statements,
    --   "backtracking" is the process of discarding the current
    --   possible solution value and returning to a previous decision
    --   point where a new value can be obtained and tried.  For
    --   example:
    --
    --   >>> do { x <- pure 0 <|> pure 1 <|> pure 2; if even x then pure x else empty } :: [Int]
    --   [0,2]
    --
    --   Here, the @x@ value can be produced three times, where
    --   'Control.Applicative.<|>' represents the decision points of that
    --   production.  The subsequent @if@ statement specifies
    --   'Control.Applicative.empty' (fail)
    --   if @x@ is odd, causing it to be discarded and a return
    --   to an 'Control.Applicative.<|>' decision point to get the next @x@.
    --
    --   The statement "@a@ '>>=' @k@ can backtrack arbitrarily many
    --   times" means that the computation is resulting in 'Control.Applicative.empty' and
    --   that @a@ has an infinite number of 'Control.Applicative.<|>' applications to
    --   return to.  This is called a conjunctive computation because
    --   the logic for @a@ /and/ @k@ must both succeed (i.e. 'pure'
    --   a value instead of 'Control.Applicative.empty').
    --
    --   Similar to the way 'interleave' allows both branches of a
    --   disjunctive computation, the '>>-' operator takes care to
    --   consider both branches of a conjunctive computation.
    --
    --   Consider the operation:
    --
    --   > odds = pure 1 <|> fmap (2 +) odds
    --   >
    --   > oddsPlus n = odds >>= \a -> pure (a + n)
    --   >
    --   > g = do x <- (pure 0 <|> pure 1) >>= oddsPlus
    --   >        if even x then pure x else empty
    --
    --   >>> observeMany 3 g
    --   ...never completes...
    --
    --   This will never produce any value because all values produced
    --   by the @do@ program come from the 'pure' @1@ driven operation
    --   (adding one to the sequence of odd values, resulting in the
    --   even values that are allowed by the test in the second line),
    --   but the 'pure' @0@ input to @oddsPlus@ generates an infinite
    --   number of 'Control.Applicative.empty' failures so the even values generated by
    --   the 'pure' @1@ alternative are never seen.  Using
    --   'interleave' here instead of 'Control.Applicative.<|>' does not help due
    --   to the aforementioned distributivity law.
    --
    --   Also note that the @do@ notation desugars to '>>=' bind
    --   operations, so the following would also fail:
    --
    --   > do a <- pure 0 <|> pure 1
    --   >    x <- oddsPlus a
    --   >    if even x then pure x else empty
    --
    --   The solution is to use the '>>-' in place of the normal
    --   monadic bind operation '>>=' when fairness between
    --   alternative productions is needed in a conjunction of
    --   statements (rules):
    --
    --   > h = do x <- (pure 0 <|> pure 1) >>- oddsPlus
    --   >        if even x then pure x else empty
    --
    --   >>> observeMany 3 h
    --   [2,4,6]
    --
    --   However, a bit of care is needed when using '>>-' because,
    --   unlike '>>=', it is not associative.  For example:
    --
    --   >>> let m = [2,7] :: [Int]
    --   >>> let k x = [x, x + 1]
    --   >>> let h x = [x, x * 2]
    --   >>> m >>= (\x -> k x >>= h)
    --   [2,4,3,6,7,14,8,16]
    --   >>> (m >>= k) >>= h -- same as above
    --   [2,4,3,6,7,14,8,16]
    --   >>> m >>- (\x -> k x >>- h)
    --   [2,7,3,8,4,14,6,16]
    --   >>> (m >>- k) >>- h -- central elements are different
    --   [2,7,4,3,14,8,6,16]
    --
    --   This means that the following will be productive:
    --
    --   > (pure 0 <|> pure 1) >>-
    --   >   oddsPlus >>-
    --   >     \x -> if even x then pure x else empty
    --
    --   Which is equivalent to
    --
    --   > ((pure 0 <|> pure 1) >>- oddsPlus) >>-
    --   >   (\x -> if even x then pure x else empty)
    --
    --   But the following will /not/ be productive:
    --
    --   > (pure 0 <|> pure 1) >>-
    --   >   (\a -> (oddsPlus a >>- \x -> if even x then pure x else empty))
    --
    --   Since do notation desugaring results in the latter, the
    --   @RebindableSyntax@ language pragma cannot easily be used
    --   either.  Instead, it is recommended to carefully use explicit
    --   '>>-' only when needed.
    --
    (>>-)      :: m a -> (a -> m b) -> m b
    infixl 1 >>-

    -- | __Pruning.__ Selects one result out of many. Useful for when multiple
    --   results of a computation will be equivalent, or should be treated as
    --   such.
    --
    --   As an example, here's a way to determine if a number is
    --   <https://wikipedia.org/wiki/Composite_number composite>
    --   (has non-trivial integer divisors, i.e. not a
    --   prime number):
    --
    --   > choose = foldr ((<|>) . pure) empty
    --   >
    --   > divisors n = do a <- choose [2..n-1]
    --   >                 b <- choose [2..n-1]
    --   >                 guard (a * b == n)
    --   >                 pure (a, b)
    --   >
    --   > composite_ v = do _ <- divisors v
    --   >                   pure "Composite"
    --
    --   While this works as intended, it actually does too much work:
    --
    --   >>> observeAll (composite_ 20)
    --   ["Composite", "Composite", "Composite", "Composite"]
    --
    --   Because there are multiple divisors of 20, and they can also
    --   occur in either order:
    --
    --   >>> observeAll (divisors 20)
    --   [(2,10), (4,5), (5,4), (10,2)]
    --
    --   Clearly one could just use 'Control.Monad.Logic.observe' here to get the first
    --   non-prime result, but if the call to @composite@ is in the
    --   middle of other logic code then use 'once' instead.
    --
    --   > composite v = do _ <- once (divisors v)
    --   >                  pure "Composite"
    --
    --   >>> observeAll (composite 20)
    --   ["Composite"]
    --
    once       :: m a -> m a

    -- | __Inverts__ a logic computation. If @m@ succeeds with at least one value,
    --   'lnot' @m@ fails. If @m@ fails, then 'lnot' @m@ succeeds with the value @()@.
    --
    --   For example, evaluating if a number is prime can be based on
    --   the failure to find divisors of a number:
    --
    --   > choose = foldr ((<|>) . pure) empty
    --   >
    --   > divisors n = do d <- choose [2..n-1]
    --   >                 guard (n `rem` d == 0)
    --   >                 pure d
    --   >
    --   > prime v = do _ <- lnot (divisors v)
    --   >              pure True
    --
    --   >>> observeAll (prime 20)
    --   []
    --   >>> observeAll (prime 19)
    --   [True]
    --
    --   Here if @divisors@ never succeeds, then the 'lnot' will
    --   succeed and the number will be declared as prime.
    lnot :: m a -> m ()

    -- | Logical __conditional.__ The equivalent of
    --   <http://lpn.swi-prolog.org/lpnpage.php?pagetype=html&pageid=lpn-htmlse44 Prolog's soft-cut>.
    --   If its first argument succeeds at all,
    --   then the results will be fed into the success
    --   branch. Otherwise, the failure branch is taken.  The failure
    --   branch is never considered if the first argument has any
    --   successes.  The 'ifte' function satisfies the following laws:
    --
    --   > ifte (pure a) th el       == th a
    --   > ifte empty th el          == el
    --   > ifte (pure a <|> m) th el == th a <|> (m >>= th)
    --
    --   For example, the previous @prime@ function returned nothing
    --   if the number was not prime, but if it should return 'False'
    --   instead, the following can be used:
    --
    --   > choose = foldr ((<|>) . pure) empty
    --   >
    --   > divisors n = do d <- choose [2..n-1]
    --   >                 guard (n `rem` d == 0)
    --   >                 pure d
    --   >
    --   > prime v = once (ifte (divisors v)
    --   >                   (const (pure True))
    --   >                   (pure False))
    --
    --   >>> observeAll (prime 20)
    --   [False]
    --   >>> observeAll (prime 19)
    --   [True]
    --
    --   Notice that this cannot be done with a simple @if-then-else@
    --   because @divisors@ either generates values or it does not, so
    --   there's no "false" condition to check with a simple @if@
    --   statement.
    ifte       :: m a -> (a -> m b) -> m b -> m b

    -- All the class functions besides msplit can be derived from msplit, if
    -- desired
    interleave m a
m1 m a
m2 = m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m a
m1 m (Maybe (a, m a)) -> (Maybe (a, m a) -> m a) -> m a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
                        m a -> ((a, m a) -> m a) -> Maybe (a, m a) -> m a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m a
m2 (\(a
a, m a
m1') -> a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a m a -> m a -> m a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> m a -> m a -> m a
forall (m :: * -> *) a. MonadLogic m => m a -> m a -> m a
interleave m a
m2 m a
m1')

    m a
m >>- a -> m b
f = do (a
a, m a
m') <- m (a, m a)
-> ((a, m a) -> m (a, m a)) -> Maybe (a, m a) -> m (a, m a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m (a, m a)
forall (f :: * -> *) a. Alternative f => f a
empty (a, m a) -> m (a, m a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, m a) -> m (a, m a)) -> m (Maybe (a, m a)) -> m (a, m a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m a
m
                 m b -> m b -> m b
forall (m :: * -> *) a. MonadLogic m => m a -> m a -> m a
interleave (a -> m b
f a
a) (m a
m' m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. MonadLogic m => m a -> (a -> m b) -> m b
>>- a -> m b
f)

    ifte m a
t a -> m b
th m b
el = m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m a
t m (Maybe (a, m a)) -> (Maybe (a, m a) -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m b -> ((a, m a) -> m b) -> Maybe (a, m a) -> m b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m b
el (\(a
a,m a
m) -> a -> m b
th a
a m b -> m b -> m b
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (m a
m m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m b
th))

    once m a
m = do (a
a, m a
_) <- m (a, m a)
-> ((a, m a) -> m (a, m a)) -> Maybe (a, m a) -> m (a, m a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m (a, m a)
forall (f :: * -> *) a. Alternative f => f a
empty (a, m a) -> m (a, m a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, m a) -> m (a, m a)) -> m (Maybe (a, m a)) -> m (a, m a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m a
m
                a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

    lnot m a
m = m a -> (a -> m ()) -> m () -> m ()
forall (m :: * -> *) a b.
MonadLogic m =>
m a -> (a -> m b) -> m b -> m b
ifte (m a -> m a
forall (m :: * -> *) a. MonadLogic m => m a -> m a
once m a
m) (m () -> a -> m ()
forall a b. a -> b -> a
const m ()
forall (f :: * -> *) a. Alternative f => f a
empty) (() -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())


-------------------------------------------------------------------------------
-- | The inverse of 'msplit'. Satisfies the following law:
--
-- > msplit m >>= reflect == m
reflect :: Alternative m => Maybe (a, m a) -> m a
reflect :: forall (m :: * -> *) a. Alternative m => Maybe (a, m a) -> m a
reflect Maybe (a, m a)
Nothing = m a
forall (f :: * -> *) a. Alternative f => f a
empty
reflect (Just (a
a, m a
m)) = a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a m a -> m a -> m a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> m a
m

-- An instance of MonadLogic for lists
instance MonadLogic [] where
    msplit :: forall a. [a] -> [Maybe (a, [a])]
msplit []     = Maybe (a, [a]) -> [Maybe (a, [a])]
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (a, [a])
forall a. Maybe a
Nothing
    msplit (a
x:[a]
xs) = Maybe (a, [a]) -> [Maybe (a, [a])]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, [a]) -> [Maybe (a, [a])])
-> Maybe (a, [a]) -> [Maybe (a, [a])]
forall a b. (a -> b) -> a -> b
$ (a, [a]) -> Maybe (a, [a])
forall a. a -> Maybe a
Just (a
x, [a]
xs)

-- | Note that splitting a transformer does
-- not allow you to provide different input
-- to the monadic object returned.
-- For instance, in:
--
-- > let Just (_, rm') = runReaderT (msplit rm) r in runReaderT rm' r'
--
-- @r'@ will be ignored, because @r@ was already threaded through the
-- computation.
instance MonadLogic m => MonadLogic (ReaderT e m) where
    msplit :: forall a. ReaderT e m a -> ReaderT e m (Maybe (a, ReaderT e m a))
msplit ReaderT e m a
rm = (e -> m (Maybe (a, ReaderT e m a)))
-> ReaderT e m (Maybe (a, ReaderT e m a))
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((e -> m (Maybe (a, ReaderT e m a)))
 -> ReaderT e m (Maybe (a, ReaderT e m a)))
-> (e -> m (Maybe (a, ReaderT e m a)))
-> ReaderT e m (Maybe (a, ReaderT e m a))
forall a b. (a -> b) -> a -> b
$ \e
e -> do Maybe (a, m a)
r <- m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit (m a -> m (Maybe (a, m a))) -> m a -> m (Maybe (a, m a))
forall a b. (a -> b) -> a -> b
$ ReaderT e m a -> e -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT e m a
rm e
e
                                   case Maybe (a, m a)
r of
                                     Maybe (a, m a)
Nothing -> Maybe (a, ReaderT e m a) -> m (Maybe (a, ReaderT e m a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (a, ReaderT e m a)
forall a. Maybe a
Nothing
                                     Just (a
a, m a
m) -> Maybe (a, ReaderT e m a) -> m (Maybe (a, ReaderT e m a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, ReaderT e m a) -> Maybe (a, ReaderT e m a)
forall a. a -> Maybe a
Just (a
a, m a -> ReaderT e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m a
m))

-- | See note on splitting above.
instance (MonadLogic m, MonadPlus m) => MonadLogic (StrictST.StateT s m) where
    msplit :: forall a. StateT s m a -> StateT s m (Maybe (a, StateT s m a))
msplit StateT s m a
sm = (s -> m (Maybe (a, StateT s m a), s))
-> StateT s m (Maybe (a, StateT s m a))
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StrictST.StateT ((s -> m (Maybe (a, StateT s m a), s))
 -> StateT s m (Maybe (a, StateT s m a)))
-> (s -> m (Maybe (a, StateT s m a), s))
-> StateT s m (Maybe (a, StateT s m a))
forall a b. (a -> b) -> a -> b
$ \s
s ->
                    do Maybe ((a, s), m (a, s))
r <- m (a, s) -> m (Maybe ((a, s), m (a, s)))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit (StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m a
sm s
s)
                       case Maybe ((a, s), m (a, s))
r of
                            Maybe ((a, s), m (a, s))
Nothing          -> (Maybe (a, StateT s m a), s) -> m (Maybe (a, StateT s m a), s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, StateT s m a)
forall a. Maybe a
Nothing, s
s)
                            Just ((a
a,s
s'), m (a, s)
m) ->
                                (Maybe (a, StateT s m a), s) -> m (Maybe (a, StateT s m a), s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, StateT s m a) -> Maybe (a, StateT s m a)
forall a. a -> Maybe a
Just (a
a, (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StrictST.StateT (m (a, s) -> s -> m (a, s)
forall a b. a -> b -> a
const m (a, s)
m)), s
s')

    interleave :: forall a. StateT s m a -> StateT s m a -> StateT s m a
interleave StateT s m a
ma StateT s m a
mb = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StrictST.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s ->
                        StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m a
ma s
s m (a, s) -> m (a, s) -> m (a, s)
forall (m :: * -> *) a. MonadLogic m => m a -> m a -> m a
`interleave` StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m a
mb s
s

    StateT s m a
ma >>- :: forall a b. StateT s m a -> (a -> StateT s m b) -> StateT s m b
>>- a -> StateT s m b
f = (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StrictST.StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s ->
                StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m a
ma s
s m (a, s) -> ((a, s) -> m (b, s)) -> m (b, s)
forall (m :: * -> *) a b. MonadLogic m => m a -> (a -> m b) -> m b
>>- \(a
a,s
s') -> StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT (a -> StateT s m b
f a
a) s
s'

    ifte :: forall a b.
StateT s m a -> (a -> StateT s m b) -> StateT s m b -> StateT s m b
ifte StateT s m a
t a -> StateT s m b
th StateT s m b
el = (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StrictST.StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s -> m (a, s) -> ((a, s) -> m (b, s)) -> m (b, s) -> m (b, s)
forall (m :: * -> *) a b.
MonadLogic m =>
m a -> (a -> m b) -> m b -> m b
ifte (StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m a
t s
s)
                                                (\(a
a,s
s') -> StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT (a -> StateT s m b
th a
a) s
s')
                                                (StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m b
el s
s)

    once :: forall a. StateT s m a -> StateT s m a
once StateT s m a
ma = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StrictST.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> m (a, s) -> m (a, s)
forall (m :: * -> *) a. MonadLogic m => m a -> m a
once (StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
StrictST.runStateT StateT s m a
ma s
s)

-- | See note on splitting above.
instance (MonadLogic m, MonadPlus m) => MonadLogic (LazyST.StateT s m) where
    msplit :: forall a. StateT s m a -> StateT s m (Maybe (a, StateT s m a))
msplit StateT s m a
sm = (s -> m (Maybe (a, StateT s m a), s))
-> StateT s m (Maybe (a, StateT s m a))
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
LazyST.StateT ((s -> m (Maybe (a, StateT s m a), s))
 -> StateT s m (Maybe (a, StateT s m a)))
-> (s -> m (Maybe (a, StateT s m a), s))
-> StateT s m (Maybe (a, StateT s m a))
forall a b. (a -> b) -> a -> b
$ \s
s ->
                    do Maybe ((a, s), m (a, s))
r <- m (a, s) -> m (Maybe ((a, s), m (a, s)))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit (StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m a
sm s
s)
                       case Maybe ((a, s), m (a, s))
r of
                            Maybe ((a, s), m (a, s))
Nothing -> (Maybe (a, StateT s m a), s) -> m (Maybe (a, StateT s m a), s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, StateT s m a)
forall a. Maybe a
Nothing, s
s)
                            Just ((a
a,s
s'), m (a, s)
m) ->
                                (Maybe (a, StateT s m a), s) -> m (Maybe (a, StateT s m a), s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, StateT s m a) -> Maybe (a, StateT s m a)
forall a. a -> Maybe a
Just (a
a, (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
LazyST.StateT (m (a, s) -> s -> m (a, s)
forall a b. a -> b -> a
const m (a, s)
m)), s
s')

    interleave :: forall a. StateT s m a -> StateT s m a -> StateT s m a
interleave StateT s m a
ma StateT s m a
mb = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
LazyST.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s ->
                        StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m a
ma s
s m (a, s) -> m (a, s) -> m (a, s)
forall (m :: * -> *) a. MonadLogic m => m a -> m a -> m a
`interleave` StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m a
mb s
s

    StateT s m a
ma >>- :: forall a b. StateT s m a -> (a -> StateT s m b) -> StateT s m b
>>- a -> StateT s m b
f = (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
LazyST.StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s ->
                StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m a
ma s
s m (a, s) -> ((a, s) -> m (b, s)) -> m (b, s)
forall (m :: * -> *) a b. MonadLogic m => m a -> (a -> m b) -> m b
>>- \(a
a,s
s') -> StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT (a -> StateT s m b
f a
a) s
s'

    ifte :: forall a b.
StateT s m a -> (a -> StateT s m b) -> StateT s m b -> StateT s m b
ifte StateT s m a
t a -> StateT s m b
th StateT s m b
el = (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
LazyST.StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ \s
s -> m (a, s) -> ((a, s) -> m (b, s)) -> m (b, s) -> m (b, s)
forall (m :: * -> *) a b.
MonadLogic m =>
m a -> (a -> m b) -> m b -> m b
ifte (StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m a
t s
s)
                                              (\(a
a,s
s') -> StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT (a -> StateT s m b
th a
a) s
s')
                                              (StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m b
el s
s)

    once :: forall a. StateT s m a -> StateT s m a
once StateT s m a
ma = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
LazyST.StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> m (a, s) -> m (a, s)
forall (m :: * -> *) a. MonadLogic m => m a -> m a
once (StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
LazyST.runStateT StateT s m a
ma s
s)