module Combinators where

import Control.Applicative
import Control.Monad
import Control.Monad.State.Strict
import Data.Foldable
import Data.Function
import Data.Int
import Data.Traversable
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

-- * Traversal

-- |
-- Indexed version of 'forM'.
{-# INLINE iforM #-}
iforM :: (Monad m, Traversable f) => f a -> (Int -> a -> m b) -> m (f b)
iforM :: forall (m :: * -> *) (f :: * -> *) a b.
(Monad m, Traversable f) =>
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