module Combinators where

import Control.Applicative
import Control.Monad
import Control.Monad.State.Strict
import Data.Bool
import Data.Foldable
import Data.Function
import Data.Int
import Data.Monoid
import Data.Traversable
import Data.Tuple
import GHC.Enum

-- * Alternation

-- |
-- Generalization over many common natural transformations, including:
--
-- - 'listToMaybe'
-- - 'maybeToList'
-- - 'toList'
-- - @'either' ('const' 'Nothing') 'Just'@
{-# INLINE alternate #-}
alternate :: (Foldable f, Alternative g) => f a -> g a
alternate :: forall (f :: * -> *) (g :: * -> *) a.
(Foldable f, Alternative g) =>
f a -> g a
alternate = (a -> g a) -> f a -> g a
forall (f :: * -> *) (g :: * -> *) a b.
(Foldable f, Alternative g) =>
(a -> g b) -> f a -> g b
alternateMapM a -> g a
forall a. a -> g a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- |
-- 'alternate' extended with ability to map the wrapped value.
{-# INLINE alternateMap #-}
alternateMap :: (Foldable f, Alternative g) => (a -> b) -> f a -> g b
alternateMap :: forall (f :: * -> *) (g :: * -> *) a b.
(Foldable f, Alternative g) =>
(a -> b) -> f a -> g b
alternateMap a -> b
mapper = (a -> g b) -> f a -> g b
forall (f :: * -> *) (g :: * -> *) a b.
(Foldable f, Alternative g) =>
(a -> g b) -> f a -> g b
alternateMapM (b -> g b
forall a. a -> g a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> g b) -> (a -> b) -> a -> g b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
mapper)

-- |
-- 'alternateMap' extended with ability to do the mapping in the 'Alternative' context.
{-# INLINE alternateMapM #-}
alternateMapM :: (Foldable f, Alternative g) => (a -> g b) -> f a -> g b
alternateMapM :: forall (f :: * -> *) (g :: * -> *) a b.
(Foldable f, Alternative g) =>
(a -> g b) -> f a -> g b
alternateMapM a -> g b
mapper = (a -> g b -> g b) -> g b -> f a -> g b
forall a b. (a -> b -> b) -> b -> f a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> g b -> g b
cons g b
forall a. g a
forall (f :: * -> *) a. Alternative f => f a
empty
  where
    cons :: a -> g b -> g b
cons a
a g b
b = a -> g b
mapper a
a g b -> g b -> g b
forall a. g a -> g a -> g a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> g b
b

-- * Folding

-- |
-- A generic version of the original list-specialized version:
--
-- > intercalate :: [a] -> [[a]] -> [a]
{-# INLINE intercalate #-}
intercalate :: (Foldable f, Monoid a) => a -> f a -> a
intercalate :: forall (f :: * -> *) a. (Foldable f, Monoid a) => a -> f a -> a
intercalate = (a -> (a -> a) -> f a -> a) -> (a -> a) -> a -> f a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> (a -> a) -> f a -> a
forall (f :: * -> *) m a.
(Foldable f, Monoid m) =>
m -> (a -> m) -> f a -> m
intercalateMap a -> a
forall a. a -> a
id

-- |
-- 'intercalate' extended with ability to map the wrapped value.
{-# INLINE intercalateMap #-}
intercalateMap :: (Foldable f, Monoid m) => m -> (a -> m) -> f a -> m
intercalateMap :: forall (f :: * -> *) m a.
(Foldable f, Monoid m) =>
m -> (a -> m) -> f a -> m
intercalateMap m
separator a -> m
proj =
  (m, Bool) -> m
forall a b. (a, b) -> a
fst
    ((m, Bool) -> m) -> (f a -> (m, Bool)) -> f a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((m, Bool) -> a -> (m, Bool)) -> (m, Bool) -> f a -> (m, Bool)
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
      ( \(m
acc, Bool
isFirst) a
element ->
          if Bool
isFirst
            then (a -> m
proj a
element, Bool
False)
            else (m
acc m -> m -> m
forall a. Semigroup a => a -> a -> a
<> m
separator m -> m -> m
forall a. Semigroup a => a -> a -> a
<> a -> m
proj a
element, Bool
False)
      )
      (m
forall a. Monoid a => a
mempty, Bool
True)

-- * Traversal

-- |
-- Indexed version of 'forM'.
{-# INLINE iforM #-}
iforM :: (Traversable f, Monad m) => f a -> (Int -> a -> m b) -> m (f b)
iforM :: forall (f :: * -> *) (m :: * -> *) a b.
(Traversable f, Monad m) =>
f a -> (Int -> a -> m b) -> m (f b)
iforM f a
collection Int -> a -> m b
f =
  f a
collection
    f a -> (f a -> StateT Int m (f b)) -> StateT Int m (f b)
forall a b. a -> (a -> b) -> b
& (a -> StateT Int m b) -> f a -> StateT Int m (f b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f a -> f (f b)
traverse
      ( \a
item -> do
          Int
i <- (Int -> (Int, Int)) -> StateT Int m Int
forall a. (Int -> (a, Int)) -> StateT Int m a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (\Int
i -> (Int
i, Int -> Int
forall a. Enum a => a -> a
succ Int
i))
          m b -> StateT Int m b
forall (m :: * -> *) a. Monad m => m a -> StateT Int m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Int -> a -> m b
f Int
i a
item)
      )
    StateT Int m (f b) -> (StateT Int m (f b) -> m (f b)) -> m (f b)
forall a b. a -> (a -> b) -> b
& (StateT Int m (f b) -> Int -> m (f b))
-> Int -> StateT Int m (f b) -> m (f b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT Int m (f b) -> Int -> m (f b)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT Int
0

-- |
-- Indexed version of 'traverse'.
{-# INLINE itraverse #-}
itraverse :: (Traversable f, Monad m) => (Int -> a -> m b) -> f a -> m (f b)
itraverse :: forall (f :: * -> *) (m :: * -> *) a b.
(Traversable f, Monad m) =>
(Int -> a -> m b) -> f a -> m (f b)
itraverse = (f a -> (Int -> a -> m b) -> m (f b))
-> (Int -> a -> m b) -> f a -> m (f b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip f a -> (Int -> a -> m b) -> m (f b)
forall (f :: * -> *) (m :: * -> *) a b.
(Traversable f, Monad m) =>
f a -> (Int -> a -> m b) -> m (f b)
iforM