{-|
Extra utility functions to manipulate 'Test.SmallCheck.Series'.
-}
module Test.SmallCheck.Series.Utils
  (
  -- * Zipping
    zipLogic
  , zipLogic3
  ) where

import Control.Applicative (empty, (<|>))
import Control.Monad ((<=<))
import Control.Monad.Logic (MonadLogic(msplit))
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT(..), runMaybeT)

-- $setup
-- >>> import Data.Char
-- >>> import Data.Functor.Identity
-- >>> import Data.Text (Text)
-- >>> import Test.SmallCheck.Series
-- >>> import Test.SmallCheck.Series.Instances

-- | /One-to-One/ zipping of 2 'MonadLogic' instances. You can use for
--   'Test.SmallCheck.Series' like this:
--
-- >>> list 2 $ (series :: Series Identity Char) `zipLogic` (series :: Series Identity Int)
-- [('a',0),('b',1),('c',-1)]
--
-- Notice the difference with 'Test.SmallCheck.Series.><':
--
-- >>> list 2 $ (series :: Series Identity Char) >< (series :: Series Identity Int)
-- [('a',0),('b',0),('a',1),('c',0),('a',-1),...,('b',-2),('c',-2)]

-- Thanks to Roman Cheplyaka: https://groups.google.com/d/msg/haskell-tasty/k0dXCx9EBsc/XYkCTjYKqswJ
zipLogic :: MonadLogic m => m a -> m b -> m (a, b)
zipLogic :: m a -> m b -> m (a, b)
zipLogic m a
gx m b
gy =
  m (a, b) -> ((a, b) -> m (a, b)) -> Maybe (a, b) -> m (a, b)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m (a, b)
forall (f :: * -> *) a. Alternative f => f a
empty (a, b) -> m (a, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, b) -> m (a, b))
-> (MaybeT m (a, b) -> m (Maybe (a, b)))
-> MaybeT m (a, b)
-> m (a, b)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< MaybeT m (a, b) -> m (Maybe (a, b))
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT m (a, b) -> m (a, b)) -> MaybeT m (a, b) -> m (a, b)
forall a b. (a -> b) -> a -> b
$ do
    (a
x, m a
rx) <- m (Maybe (a, m a)) -> MaybeT m (a, m a)
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m a
gx)
    (b
y, m b
ry) <- m (Maybe (b, m b)) -> MaybeT m (b, m b)
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m b -> m (Maybe (b, m b))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m b
gy)
    m (a, b) -> MaybeT m (a, b)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (a, b) -> MaybeT m (a, b)) -> m (a, b) -> MaybeT m (a, b)
forall a b. (a -> b) -> a -> b
$ (a, b) -> m (a, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, b
y) m (a, b) -> m (a, b) -> m (a, b)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> m a -> m b -> m (a, b)
forall (m :: * -> *) a b. MonadLogic m => m a -> m b -> m (a, b)
zipLogic m a
rx m b
ry

-- | /One-to-One/ zipping of 3 'MonadLogic' instances. You can use for
--   'Test.SmallCheck.Series' like this:
--
-- >>> list 3 $ zipLogic3 (series :: Series Identity Char) (series :: Series Identity Int) (series :: Series Identity Text)
-- [('a',0,""),('b',1,"a"),('c',-1,"b"),('d',2,"aa")]

-- Thanks to Roman Cheplyaka: https://groups.google.com/d/msg/haskell-tasty/k0dXCx9EBsc/XYkCTjYKqswJ
zipLogic3 :: MonadLogic m => m a -> m b -> m c -> m (a, b, c)
zipLogic3 :: m a -> m b -> m c -> m (a, b, c)
zipLogic3 m a
gx m b
gy m c
gz =
  m (a, b, c)
-> ((a, b, c) -> m (a, b, c)) -> Maybe (a, b, c) -> m (a, b, c)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m (a, b, c)
forall (f :: * -> *) a. Alternative f => f a
empty (a, b, c) -> m (a, b, c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (a, b, c) -> m (a, b, c))
-> (MaybeT m (a, b, c) -> m (Maybe (a, b, c)))
-> MaybeT m (a, b, c)
-> m (a, b, c)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< MaybeT m (a, b, c) -> m (Maybe (a, b, c))
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT m (a, b, c) -> m (a, b, c))
-> MaybeT m (a, b, c) -> m (a, b, c)
forall a b. (a -> b) -> a -> b
$ do
    (a
x, m a
rx) <- m (Maybe (a, m a)) -> MaybeT m (a, m a)
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m a -> m (Maybe (a, m a))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m a
gx)
    (b
y, m b
ry) <- m (Maybe (b, m b)) -> MaybeT m (b, m b)
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m b -> m (Maybe (b, m b))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m b
gy)
    (c
z, m c
rz) <- m (Maybe (c, m c)) -> MaybeT m (c, m c)
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m c -> m (Maybe (c, m c))
forall (m :: * -> *) a. MonadLogic m => m a -> m (Maybe (a, m a))
msplit m c
gz)
    m (a, b, c) -> MaybeT m (a, b, c)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (a, b, c) -> MaybeT m (a, b, c))
-> m (a, b, c) -> MaybeT m (a, b, c)
forall a b. (a -> b) -> a -> b
$ (a, b, c) -> m (a, b, c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, b
y, c
z) m (a, b, c) -> m (a, b, c) -> m (a, b, c)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> m a -> m b -> m c -> m (a, b, c)
forall (m :: * -> *) a b c.
MonadLogic m =>
m a -> m b -> m c -> m (a, b, c)
zipLogic3 m a
rx m b
ry m c
rz