{-# LANGUAGE CPP #-}

module Data.List.Extra where

import Control.Applicative (liftA2)

-- | Monadic version of 'Data.List.partition'
partitionM :: (Monad m) => (a -> m Bool) -> [a] -> m ([a], [a])
partitionM :: (a -> m Bool) -> [a] -> m ([a], [a])
partitionM a -> m Bool
_ []     = ([a], [a]) -> m ([a], [a])
forall (m :: Type -> Type) a. Monad m => a -> m a
return ([], [])
partitionM a -> m Bool
p (a
x:[a]
xs) = do
  Bool
test      <- a -> m Bool
p a
x
  ([a]
ys, [a]
ys') <- (a -> m Bool) -> [a] -> m ([a], [a])
forall (m :: Type -> Type) a.
Monad m =>
(a -> m Bool) -> [a] -> m ([a], [a])
partitionM a -> m Bool
p [a]
xs
  ([a], [a]) -> m ([a], [a])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (([a], [a]) -> m ([a], [a])) -> ([a], [a]) -> m ([a], [a])
forall a b. (a -> b) -> a -> b
$ if Bool
test then (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys, [a]
ys') else ([a]
ys, a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys')

-- | Monadic version of 'Data.List.mapAccumL'
mapAccumLM
  :: (Monad m)
  => (acc -> x -> m (acc,y))
  -> acc
  -> [x]
  -> m (acc,[y])
mapAccumLM :: (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
_ acc
acc [] = (acc, [y]) -> m (acc, [y])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (acc
acc,[])
mapAccumLM acc -> x -> m (acc, y)
f acc
acc (x
x:[x]
xs) = do
  (acc
acc',y
y) <- acc -> x -> m (acc, y)
f acc
acc x
x
  (acc
acc'',[y]
ys) <- (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
forall (m :: Type -> Type) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
f acc
acc' [x]
xs
  (acc, [y]) -> m (acc, [y])
forall (m :: Type -> Type) a. Monad m => a -> m a
return (acc
acc'',y
yy -> [y] -> [y]
forall a. a -> [a] -> [a]
:[y]
ys)

-- | Monadic version of 'iterate'. A carbon copy ('iterateM') would not
-- terminate, hence the first argument.
iterateNM
  :: Monad m
  => Word
  -- ^ Only iterate /n/ times. Note that /n/ is the length of the resulting
  -- list, _not_ the number of times the iteration function has been invoked
  -> (a -> m a)
  -- ^ Iteration function
  -> a
  -- ^ Start value
  -> m [a]
iterateNM :: Word -> (a -> m a) -> a -> m [a]
iterateNM Word
0 a -> m a
_f a
_a = [a] -> m [a]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure []
iterateNM Word
limit a -> m a
f a
a = ([a] -> [a]) -> m [a] -> m [a]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (Word -> a -> m [a]
forall t. (Eq t, Num t) => t -> a -> m [a]
go (Word
limit Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) a
a)
 where
  go :: t -> a -> m [a]
go t
0 a
_a0 = [a] -> m [a]
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure []
  go t
n a
a0 = do
    a
a1 <- a -> m a
f a
a0
    ([a] -> [a]) -> m [a] -> m [a]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (a
a1a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (t -> a -> m [a]
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1) a
a1)

infixr 5 <:>
-- | Applicative version of 'GHC.Types.(:)'
(<:>) :: Applicative f => f a -> f [a] -> f [a]
<:> :: f a -> f [a] -> f [a]
(<:>) = (a -> [a] -> [a]) -> f a -> f [a] -> f [a]
forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (:)

-- | Safe indexing, returns a 'Nothing' if the index does not exist
indexMaybe :: [a] -> Int -> Maybe a
indexMaybe :: [a] -> Int -> Maybe a
indexMaybe [] Int
_     = Maybe a
forall a. Maybe a
Nothing
indexMaybe (a
x:[a]
_)  Int
0 = a -> Maybe a
forall a. a -> Maybe a
Just a
x
indexMaybe (a
_:[a]
xs) Int
n = [a] -> Int -> Maybe a
forall a. [a] -> Int -> Maybe a
indexMaybe [a]
xs (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

splitAtList :: [b] -> [a] -> ([a], [a])
splitAtList :: [b] -> [a] -> ([a], [a])
splitAtList [] [a]
xs         = ([], [a]
xs)
splitAtList [b]
_ xs :: [a]
xs@[]       = ([a]
xs, [a]
xs)
splitAtList (b
_:[b]
xs) (a
y:[a]
ys) = (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys', [a]
ys'')
    where
      ([a]
ys', [a]
ys'') = [b] -> [a] -> ([a], [a])
forall b a. [b] -> [a] -> ([a], [a])
splitAtList [b]
xs [a]
ys

equalLength :: [a] -> [b] -> Bool
equalLength :: [a] -> [b] -> Bool
equalLength [] [] = Bool
True
equalLength (a
_:[a]
as) (b
_:[b]
bs) = [a] -> [b] -> Bool
forall a b. [a] -> [b] -> Bool
equalLength [a]
as [b]
bs
equalLength [a]
_ [b]
_ = Bool
False

-- | Return number of occurrences of an item in a list
countEq
  :: Eq a
  => a
  --  ^ Needle
  -> [a]
  -- ^ Haystack
  -> Int
  -- ^ Times needle was found in haystack
countEq :: a -> [a] -> Int
countEq a
a [a]
as = [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a) [a]
as)

-- | Zip two lists of equal length
--
-- NB Errors out for a DEBUG compiler when the two lists are not of equal length
zipEqual
  :: [a] -> [b] -> [(a,b)]
#if !defined(DEBUG)
zipEqual = zip
#else
zipEqual :: [a] -> [b] -> [(a, b)]
zipEqual [] [] = []
zipEqual (a
a:[a]
as) (b
b:[b]
bs) = (a
a,b
b) (a, b) -> [(a, b)] -> [(a, b)]
forall a. a -> [a] -> [a]
: [a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
zipEqual [a]
as [b]
bs
zipEqual [a]
_ [b]
_ = [Char] -> [(a, b)]
forall a. HasCallStack => [Char] -> a
error [Char]
"zipEqual"
#endif

-- | Short-circuiting monadic version of 'any'
anyM
  :: (Monad m)
  => (a -> m Bool)
  -> [a]
  -> m Bool
anyM :: (a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
_ []     = Bool -> m Bool
forall (m :: Type -> Type) a. Monad m => a -> m a
return Bool
False
anyM a -> m Bool
p (a
x:[a]
xs) = do
  Bool
q <- a -> m Bool
p a
x
  if Bool
q then
    Bool -> m Bool
forall (m :: Type -> Type) a. Monad m => a -> m a
return Bool
True
  else
    (a -> m Bool) -> [a] -> m Bool
forall (m :: Type -> Type) a.
Monad m =>
(a -> m Bool) -> [a] -> m Bool
anyM a -> m Bool
p [a]
xs

allM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
allM :: (a -> m Bool) -> [a] -> m Bool
allM a -> m Bool
_ [] = Bool -> m Bool
forall (m :: Type -> Type) a. Monad m => a -> m a
return Bool
True
allM a -> m Bool
p (a
x:[a]
xs) = do
  Bool
q <- a -> m Bool
p a
x
  if Bool
q then
    (a -> m Bool) -> [a] -> m Bool
forall (m :: Type -> Type) a.
Monad m =>
(a -> m Bool) -> [a] -> m Bool
allM a -> m Bool
p [a]
xs
  else
    Bool -> m Bool
forall (m :: Type -> Type) a. Monad m => a -> m a
return Bool
False

-- | short-circuiting monadic version of 'or'
orM
  :: (Monad m)
  => [m Bool]
  -> m Bool
orM :: [m Bool] -> m Bool
orM [] = Bool -> m Bool
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Bool
False
orM (m Bool
x:[m Bool]
xs) = do
  Bool
p <- m Bool
x
  if Bool
p then
    Bool -> m Bool
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Bool
True
  else
    [m Bool] -> m Bool
forall (m :: Type -> Type). Monad m => [m Bool] -> m Bool
orM [m Bool]
xs