{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}

-- | List utils
module Haskus.Utils.List
   ( at
   , unsafeAt
   , checkLength
   , (++)
   , replicate
   , drop
   , length
   , take
   , chunksOf
   , pick1
   , enumList
   , zipLeftWith
   , zipRightWith
   , L.partition
   , L.nub
   , L.sort
   , L.intersperse
   , L.foldl'
   , L.head
   , L.tail
   , L.zipWith
   , L.repeat
   , nubOn
   , L.nubBy
   , L.sortOn
   , L.sortBy
   , groupOn
   , L.groupBy
   , L.transpose
   , (L.\\)
   , L.intersect
   , L.find
   , L.zip3
   , L.zip4
   , L.zip5
   , L.zip6
   , L.zip7
   , L.stripPrefix
   , L.isPrefixOf
   , L.deleteBy
   , L.isSuffixOf
   , L.elem
   , L.notElem
   -- * Split
   , splitAt
   , split
   , splitOn
   , breakOn
   )
where

import Prelude hiding (replicate, length, drop, take,splitAt)

import Data.Bifunctor
import Data.Function (on)
import qualified Data.List as L

-- | Safely index into a list
--
-- >>> [0,1,2,3] `at` 10
-- Nothing
--
-- >>> [0,1,2,3] `at` 2
-- Just 2
at :: [a] -> Word -> Maybe a
{-# INLINABLE at #-}
at :: [a] -> Word -> Maybe a
at = [a] -> Word -> Maybe a
forall t a. (Eq t, Num t) => [a] -> t -> Maybe a
go
   where
      go :: [a] -> t -> Maybe a
go []       t
_ = Maybe a
forall a. Maybe a
Nothing
      go (a
x:[a]
_xs)  t
0 = a -> Maybe a
forall a. a -> Maybe a
Just a
x
      go (a
_x:[a]
xs) !t
n = [a] -> t -> Maybe a
go [a]
xs (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1)

-- | Unsafe `a`
--
-- >>> [0,1,2,3] `unsafeAt` 2
-- 2
unsafeAt :: [a] -> Word -> a
{-# INLINABLE unsafeAt #-}
unsafeAt :: [a] -> Word -> a
unsafeAt [a]
vs Word
k = [a] -> Word -> a
forall t p. (Eq t, Num t) => [p] -> t -> p
go [a]
vs Word
k
   where
      go :: [p] -> t -> p
go []       t
_ = [Char] -> p
forall a. HasCallStack => [Char] -> a
error ([Char]
"Unsafe list index too large: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Word -> [Char]
forall a. Show a => a -> [Char]
show Word
k)
      go (p
x:[p]
_xs)  t
0 = p
x
      go (p
_x:[p]
xs) !t
n = [p] -> t -> p
go [p]
xs (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1)

-- | Check that a list has the given length (support infinite lists)
checkLength :: Word -> [a] -> Bool
checkLength :: Word -> [a] -> Bool
checkLength Word
0 []     = Bool
True
checkLength Word
0 [a]
_      = Bool
False
checkLength Word
_ []     = Bool
False
checkLength Word
i (a
_:[a]
xs) = Word -> [a] -> Bool
forall a. Word -> [a] -> Bool
checkLength (Word
iWord -> Word -> Word
forall a. Num a => a -> a -> a
-Word
1) [a]
xs

-- | Replicate
replicate :: Word -> a -> [a]
replicate :: Word -> a -> [a]
replicate Word
n a
a = Int -> a -> [a]
forall a. Int -> a -> [a]
L.replicate (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n) a
a

-- | Take
take :: Word -> [a] -> [a]
take :: Word -> [a] -> [a]
take Word
n = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
L.take (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)

-- | Length
length :: Foldable t => t a -> Word
length :: t a -> Word
length = Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word) -> (t a -> Int) -> t a -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
L.length

-- | Drop
drop :: Word -> [a] -> [a]
drop :: Word -> [a] -> [a]
drop Word
n = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
L.drop (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n)

-- | Apply some operation repeatedly, producing an element of output
--   and the remainder of the list.
repeatedly :: ([a] -> (b, [a])) -> [a] -> [b]
repeatedly :: ([a] -> (b, [a])) -> [a] -> [b]
repeatedly [a] -> (b, [a])
_ [] = []
repeatedly [a] -> (b, [a])
f [a]
as = b
b b -> [b] -> [b]
forall a. a -> [a] -> [a]
: ([a] -> (b, [a])) -> [a] -> [b]
forall a b. ([a] -> (b, [a])) -> [a] -> [b]
repeatedly [a] -> (b, [a])
f [a]
as'
    where (b
b, [a]
as') = [a] -> (b, [a])
f [a]
as

-- | Split a list into chunks of a given size. The last chunk may contain fewer
-- than n elements.
--
-- >>> chunksOf 3 "my test"
-- ["my ","tes","t"]
--
-- >>> chunksOf 3 "mytest"
-- ["myt","est"]
--
-- >>> chunksOf 8 ""
-- []
--
-- >> chunksOf 0 "test"
-- undefined
chunksOf :: Word -> [a] -> [[a]]
chunksOf :: Word -> [a] -> [[a]]
chunksOf Word
i [a]
xs = ([a] -> ([a], [a])) -> [a] -> [[a]]
forall a b. ([a] -> (b, [a])) -> [a] -> [b]
repeatedly (Word -> [a] -> ([a], [a])
forall n a. Integral n => n -> [a] -> ([a], [a])
splitAt Word
i) [a]
xs

-- | Pick each element and return the element and the rest of the list
--
-- >>> pick1 [1,2,3,4]
-- [(1,[2,3,4]),(2,[1,3,4]),(3,[1,2,4]),(4,[1,2,3])]
pick1 :: [a] -> [(a,[a])]
pick1 :: [a] -> [(a, [a])]
pick1 = [a] -> [a] -> [(a, [a])]
forall a. [a] -> [a] -> [(a, [a])]
go []
   where
      go :: [a] -> [a] -> [(a, [a])]
go [a]
_  []     = []
      go [a]
ys (a
x:[a]
xs) = (a
x,[a] -> [a]
forall a. [a] -> [a]
reverse [a]
ys[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
xs) (a, [a]) -> [(a, [a])] -> [(a, [a])]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [(a, [a])]
go (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys) [a]
xs

-- | Get members of a bounded enum in a list
--
-- >>> :set -XTypeApplications
-- >>> data Letters = A | B | C | D deriving (Bounded,Enum,Show)
-- >>> enumList @Letters
-- [A,B,C,D]
enumList :: forall a. (Bounded a,Enum a) => [a]
enumList :: [a]
enumList = a -> [a]
forall a. Enum a => a -> [a]
enumFrom a
forall a. Bounded a => a
minBound

-- | Zip left with something extracted from each value
--
-- >>> zipLeftWith odd [0..5]
-- [(False,0),(True,1),(False,2),(True,3),(False,4),(True,5)]
zipLeftWith :: (a -> b) -> [a] -> [(b,a)]
zipLeftWith :: (a -> b) -> [a] -> [(b, a)]
zipLeftWith a -> b
f [a]
xs = (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f [a]
xs [b] -> [a] -> [(b, a)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [a]
xs

-- | Zip right with something extracted from each value
--
-- >>> zipRightWith odd [0..5]
-- [(0,False),(1,True),(2,False),(3,True),(4,False),(5,True)]
zipRightWith :: (a -> b) -> [a] -> [(a,b)]
zipRightWith :: (a -> b) -> [a] -> [(a, b)]
zipRightWith a -> b
f [a]
xs = [a]
xs [a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f [a]
xs


-- | A version of 'nub' where the equality is done on some extracted value.
--   @nubOn f@ is equivalent to @nubBy ((==) `on` f)@, but has the
--   performance advantage of only evaluating @f@ once for each element in the
--   input list.
nubOn :: Eq b => (a -> b) -> [a] -> [a]
nubOn :: (a -> b) -> [a] -> [a]
nubOn a -> b
f = ((b, a) -> a) -> [(b, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (b, a) -> a
forall a b. (a, b) -> b
snd ([(b, a)] -> [a]) -> ([a] -> [(b, a)]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, a) -> (b, a) -> Bool) -> [(b, a)] -> [(b, a)]
forall a. (a -> a -> Bool) -> [a] -> [a]
L.nubBy (b -> b -> Bool
forall a. Eq a => a -> a -> Bool
(==) (b -> b -> Bool) -> ((b, a) -> b) -> (b, a) -> (b, a) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (b, a) -> b
forall a b. (a, b) -> a
fst) ([(b, a)] -> [(b, a)]) -> ([a] -> [(b, a)]) -> [a] -> [(b, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (b, a)) -> [a] -> [(b, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> let y :: b
y = a -> b
f a
x in b
y b -> (b, a) -> (b, a)
`seq` (b
y, a
x))


-- | A version of 'group' where the equality is done on some extracted value.
groupOn :: Eq b => (a -> b) -> [a] -> [[a]]
groupOn :: (a -> b) -> [a] -> [[a]]
groupOn a -> b
f = (a -> a -> Bool) -> [a] -> [[a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
L.groupBy (b -> b -> Bool
forall a. Eq a => a -> a -> Bool
(==) (b -> b -> Bool) -> (a -> b) -> a -> a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on2` a -> b
f)
    -- redefine on so we avoid duplicate computation for most values.
    where t -> t -> t
(.*.) on2 :: (t -> t -> t) -> (p -> t) -> p -> p -> t
`on2` p -> t
g = \p
x -> let fx :: t
fx = p -> t
g p
x in \p
y -> t
fx t -> t -> t
.*. p -> t
g p
y

---------------------------------------
-- Split
---------------------------------------

-- | 'splitAt' @n xs@ returns a tuple where first element is @xs@ prefix of
-- length @n@ and second element is the remainder of the list:
--
-- > splitAt 6 "Hello World!" == ("Hello ","World!")
-- > splitAt 3 [1,2,3,4,5] == ([1,2,3],[4,5])
-- > splitAt 1 [1,2,3] == ([1],[2,3])
-- > splitAt 3 [1,2,3] == ([1,2,3],[])
-- > splitAt 4 [1,2,3] == ([1,2,3],[])
-- > splitAt 0 [1,2,3] == ([],[1,2,3])
-- > splitAt (-1) [1,2,3] == ([],[1,2,3])
--
-- It is equivalent to @('take' n xs, 'drop' n xs)@ when @n@ is not @_|_@
-- (@splitAt _|_ xs = _|_@).
splitAt :: Integral n => n -> [a] -> ([a],[a])
splitAt :: n -> [a] -> ([a], [a])
splitAt n
n [a]
xs = n -> [a] -> ([a], [a])
forall n a. Integral n => n -> [a] -> ([a], [a])
L.genericSplitAt n
n [a]
xs

-- | Break a list into pieces separated by the first
-- list argument, consuming the delimiter. An empty delimiter is
-- invalid, and will cause an error to be raised.
--
-- > splitOn "\r\n" "a\r\nb\r\nd\r\ne" == ["a","b","d","e"]
-- > splitOn "aaa"  "aaaXaaaXaaaXaaa"  == ["","X","X","X",""]
-- > splitOn "x"    "x"                == ["",""]
-- > splitOn "x"    ""                 == [""]
-- > \s x -> s /= "" ==> intercalate s (splitOn s x) == x
-- > \c x -> splitOn [c] x                           == split (==c) x
splitOn :: (Eq a) => [a] -> [a] -> [[a]]
splitOn :: [a] -> [a] -> [[a]]
splitOn [] [a]
_ = [Char] -> [[a]]
forall a. HasCallStack => [Char] -> a
error [Char]
"splitOn, needle may not be empty"
splitOn [a]
_ [] = [[]]
splitOn [a]
needle [a]
haystack = [a]
a [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: if [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
b then [] else [a] -> [a] -> [[a]]
forall a. Eq a => [a] -> [a] -> [[a]]
splitOn [a]
needle ([a] -> [[a]]) -> [a] -> [[a]]
forall a b. (a -> b) -> a -> b
$ Word -> [a] -> [a]
forall a. Word -> [a] -> [a]
drop ([a] -> Word
forall (t :: * -> *) a. Foldable t => t a -> Word
length [a]
needle) [a]
b
    where ([a]
a,[a]
b) = [a] -> [a] -> ([a], [a])
forall a. Eq a => [a] -> [a] -> ([a], [a])
breakOn [a]
needle [a]
haystack


-- | Splits a list into components delimited by separators,
-- where the predicate returns True for a separator element.  The
-- resulting components do not contain the separators.  Two adjacent
-- separators result in an empty component in the output.
--
-- > split (== 'a') "aabbaca" == ["","","bb","c",""]
-- > split (== 'a') ""        == [""]
-- > split (== ':') "::xyz:abc::123::" == ["","","xyz","abc","","123","",""]
-- > split (== ',') "my,list,here" == ["my","list","here"]
split :: (a -> Bool) -> [a] -> [[a]]
split :: (a -> Bool) -> [a] -> [[a]]
split a -> Bool
_ [] = [[]]
split a -> Bool
f (a
x:[a]
xs) | a -> Bool
f a
x                   = [] [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: (a -> Bool) -> [a] -> [[a]]
forall a. (a -> Bool) -> [a] -> [[a]]
split a -> Bool
f [a]
xs
split a -> Bool
f (a
x:[a]
xs) | ~([a]
y:[[a]]
ys) <- (a -> Bool) -> [a] -> [[a]]
forall a. (a -> Bool) -> [a] -> [[a]]
split a -> Bool
f [a]
xs = (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
y) [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [[a]]
ys

-- | Find the first instance of @needle@ in @haystack@.
-- The first element of the returned tuple
-- is the prefix of @haystack@ before @needle@ is matched.  The second
-- is the remainder of @haystack@, starting with the match.
-- If you want the remainder /without/ the match, use 'stripInfix'.
--
-- > breakOn "::" "a::b::c" == ("a", "::b::c")
-- > breakOn "/" "foobar"   == ("foobar", "")
-- > \needle haystack -> let (prefix,match) = breakOn needle haystack in prefix ++ match == haystack
breakOn :: Eq a => [a] -> [a] -> ([a], [a])
breakOn :: [a] -> [a] -> ([a], [a])
breakOn [a]
needle [a]
haystack
   | [a]
needle [a] -> [a] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`L.isPrefixOf` [a]
haystack = ([], [a]
haystack)
breakOn [a]
_      []                   = ([], [])
breakOn [a]
needle (a
x:[a]
xs)               = ([a] -> [a]) -> ([a], [a]) -> ([a], [a])
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (([a], [a]) -> ([a], [a])) -> ([a], [a]) -> ([a], [a])
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> ([a], [a])
forall a. Eq a => [a] -> [a] -> ([a], [a])
breakOn [a]
needle [a]
xs