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

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

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

#if !MIN_VERSION_base(4,18,0)
import Control.Applicative (liftA2)
#endif

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

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

-- | Like 'all', but the predicate operates over two lists. Asserts to 'False'
-- when the lists are of unequal length
all2 :: (a -> b -> Bool) -> [a] -> [b] -> Bool
all2 :: (a -> b -> Bool) -> [a] -> [b] -> Bool
all2 a -> b -> Bool
_ [] [] = Bool
True
all2 a -> b -> Bool
p (a
a:[a]
as) (b
b:[b]
bs) = a -> b -> Bool
p a
a b
b Bool -> Bool -> Bool
&& (a -> b -> Bool) -> [a] -> [b] -> Bool
forall a b. (a -> b -> Bool) -> [a] -> [b] -> Bool
all2 a -> b -> Bool
p [a]
as [b]
bs
all2 a -> b -> Bool
_ [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