{-# LANGUAGE CPP #-}
{-# LANGUAGE PackageImports #-}

module Data.List.Extra
  ( partitionM
  , mapAccumLM
  , (<:>)
  , indexMaybe
  , splitAtList
  , equalLength
  , countEq
  , zipEqual

  -- * From Control.Monad.Extra
  , anyM
  , allM
  , orM

  -- * From "extra"
  , module NeilsExtra
  ) where

import "extra" Data.List.Extra as NeilsExtra
import "extra" Control.Monad.Extra (anyM, allM, orM, partitionM)

import Control.Applicative (liftA2)

#if defined(DEBUG)
import GHC.Stack (HasCallStack)
#endif

-- | 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)

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
#if !defined(DEBUG)
zipEqual
  :: [a] -> [b] -> [(a,b)]
zipEqual = zip
{-# INLINE zipEqual #-}
#else
zipEqual ::
  HasCallStack =>
  [a] -> [b] -> [(a,b)]
zipEqual :: [a] -> [b] -> [(a, b)]
zipEqual = [a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
go
  where
    go :: [a] -> [b] -> [(a, b)]
go [] [] = []
    go (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)]
go [a]
as [b]
bs
    go (a
_:[a]
_) [] = [Char] -> [(a, b)]
forall a. HasCallStack => [Char] -> a
error [Char]
"zipEqual: left list is longer"
    go [] (b
_:[b]
_) = [Char] -> [(a, b)]
forall a. HasCallStack => [Char] -> a
error [Char]
"zipEqual: right list is longer"
#endif