1{-# LANGUAGE TupleSections, ConstraintKinds #-}
2
3-- | This module extends "Data.List" with extra functions of a similar nature.
4--   The package also exports the existing "Data.List" functions.
5--   Some of the names and semantics were inspired by the
6--   <https://hackage.haskell.org/package/text text> package.
7module Data.List.Extra(
8    module Data.List,
9    -- * String operations
10    lower, upper, trim, trimStart, trimEnd, word1, line1,
11    escapeHTML, escapeJSON,
12    unescapeHTML, unescapeJSON,
13    -- * Splitting
14    dropEnd, takeEnd, splitAtEnd, breakEnd, spanEnd,
15    dropWhileEnd', takeWhileEnd,
16    stripSuffix, stripInfix, stripInfixEnd,
17    dropPrefix, dropSuffix,
18    wordsBy, linesBy,
19    breakOn, breakOnEnd, splitOn, split, chunksOf,
20    -- * Basics
21    headDef, lastDef, notNull, list, unsnoc, cons, snoc,
22    drop1, dropEnd1, mconcatMap, compareLength, comparingLength,
23    -- * Enum operations
24    enumerate,
25    -- * List operations
26    groupSort, groupSortOn, groupSortBy,
27    nubOrd, nubOrdBy, nubOrdOn,
28    nubOn, groupOn,
29    nubSort, nubSortBy, nubSortOn,
30    maximumOn, minimumOn,
31    sum', product',
32    sumOn', productOn',
33    disjoint, disjointOrd, disjointOrdBy, allSame, anySame,
34    repeatedly, firstJust,
35    concatUnzip, concatUnzip3,
36    zipFrom, zipWithFrom, zipWithLongest,
37    replace, merge, mergeBy,
38    ) where
39
40import Partial
41import Data.List
42import Data.Maybe
43import Data.Function
44import Data.Char
45import Data.Tuple.Extra
46import Data.Monoid
47import Numeric
48import Data.Functor
49import Data.Foldable
50import Prelude
51
52
53-- | Apply some operation repeatedly, producing an element of output
54--   and the remainder of the list.
55--
56-- > \xs -> repeatedly (splitAt 3) xs  == chunksOf 3 xs
57-- > \xs -> repeatedly word1 (trim xs) == words xs
58-- > \xs -> repeatedly line1 xs == lines xs
59repeatedly :: ([a] -> (b, [a])) -> [a] -> [b]
60repeatedly f [] = []
61repeatedly f as = b : repeatedly f as'
62    where (b, as') = f as
63
64
65-- | Are two lists disjoint, with no elements in common.
66--
67-- > disjoint [1,2,3] [4,5] == True
68-- > disjoint [1,2,3] [4,1] == False
69disjoint :: Eq a => [a] -> [a] -> Bool
70disjoint xs = null . intersect xs
71
72-- | /O((m+n) log m), m <= n/. Are two lists disjoint, with no elements in common.
73--
74-- @disjointOrd@ is more strict than `disjoint`. For example, @disjointOrd@ cannot
75-- terminate if both lists are inifite, while `disjoint` can.
76--
77-- > disjointOrd [1,2,3] [4,5] == True
78-- > disjointOrd [1,2,3] [4,1] == False
79disjointOrd :: Ord a => [a] -> [a] -> Bool
80disjointOrd = disjointOrdBy compare
81
82-- | A version of 'disjointOrd' with a custom predicate.
83--
84-- > disjointOrdBy (compare `on` (`mod` 7)) [1,2,3] [4,5] == True
85-- > disjointOrdBy (compare `on` (`mod` 7)) [1,2,3] [4,8] == False
86disjointOrdBy :: (a -> a -> Ordering) -> [a] -> [a] -> Bool
87disjointOrdBy cmp xs ys
88    | shorter xs ys = go xs ys
89    | otherwise = go ys xs
90  where
91    shorter _ [] = False
92    shorter [] _ = True
93    shorter (_:xs) (_:ys) = shorter xs ys
94
95    go xs = not . any (\a -> memberRB cmp a tree)
96      where
97        tree = foldl' (flip (insertRB cmp)) E xs
98
99-- | Is there any element which occurs more than once.
100--
101-- > anySame [1,1,2] == True
102-- > anySame [1,2,3] == False
103-- > anySame (1:2:1:undefined) == True
104-- > anySame [] == False
105-- > \xs -> anySame xs == (length (nub xs) < length xs)
106anySame :: Eq a => [a] -> Bool
107anySame = f []
108    where
109        f seen (x:xs) = x `elem` seen || f (x:seen) xs
110        f seen [] = False
111
112-- | Are all elements the same.
113--
114-- > allSame [1,1,2] == False
115-- > allSame [1,1,1] == True
116-- > allSame [1]     == True
117-- > allSame []      == True
118-- > allSame (1:1:2:undefined) == False
119-- > \xs -> allSame xs == (length (nub xs) <= 1)
120allSame :: Eq a => [a] -> Bool
121allSame [] = True
122allSame (x:xs) = all (x ==) xs
123
124
125-- | A total 'head' with a default value.
126--
127-- > headDef 1 []      == 1
128-- > headDef 1 [2,3,4] == 2
129-- > \x xs -> headDef x xs == fromMaybe x (listToMaybe xs)
130headDef :: a -> [a] -> a
131headDef d [] = d
132headDef _ (x:_) = x
133
134
135-- | A total 'last' with a default value.
136--
137-- > lastDef 1 []      == 1
138-- > lastDef 1 [2,3,4] == 4
139-- > \x xs -> lastDef x xs == last (x:xs)
140lastDef :: a -> [a] -> a
141lastDef d xs = foldl (\_ x -> x) d xs -- I know this looks weird, but apparently this is the fastest way to do this: https://hackage.haskell.org/package/base-4.12.0.0/docs/src/GHC.List.html#last
142{-# INLINE lastDef #-}
143
144
145-- | A composition of 'not' and 'null'.
146--
147-- > notNull []  == False
148-- > notNull [1] == True
149-- > \xs -> notNull xs == not (null xs)
150notNull :: [a] -> Bool
151notNull = not . null
152
153-- | Non-recursive transform over a list, like 'maybe'.
154--
155-- > list 1 (\v _ -> v - 2) [5,6,7] == 3
156-- > list 1 (\v _ -> v - 2) []      == 1
157-- > \nil cons xs -> maybe nil (uncurry cons) (uncons xs) == list nil cons xs
158list :: b -> (a -> [a] -> b) -> [a] -> b
159list nil cons [] = nil
160list nil cons (x:xs) = cons x xs
161
162-- | If the list is empty returns 'Nothing', otherwise returns the 'init' and the 'last'.
163--
164-- > unsnoc "test" == Just ("tes",'t')
165-- > unsnoc ""     == Nothing
166-- > \xs -> unsnoc xs == if null xs then Nothing else Just (init xs, last xs)
167unsnoc :: [a] -> Maybe ([a], a)
168unsnoc [] = Nothing
169unsnoc [x] = Just ([], x)
170unsnoc (x:xs) = Just (x:a, b)
171    where Just (a,b) = unsnoc xs
172
173-- | Append an element to the start of a list, an alias for '(:)'.
174--
175-- > cons 't' "est" == "test"
176-- > \x xs -> uncons (cons x xs) == Just (x,xs)
177cons :: a -> [a] -> [a]
178cons = (:)
179
180-- | Append an element to the end of a list, takes /O(n)/ time.
181--
182-- > snoc "tes" 't' == "test"
183-- > \xs x -> unsnoc (snoc xs x) == Just (xs,x)
184snoc :: [a] -> a -> [a]
185snoc xs x = xs ++ [x]
186
187
188-- | Enumerate all the values of an 'Enum', from 'minBound' to 'maxBound'.
189--
190-- > enumerate == [False, True]
191enumerate :: (Enum a, Bounded a) => [a]
192enumerate = [minBound..maxBound]
193
194-- | Take a number of elements from the end of the list.
195--
196-- > takeEnd 3 "hello"  == "llo"
197-- > takeEnd 5 "bye"    == "bye"
198-- > takeEnd (-1) "bye" == ""
199-- > \i xs -> takeEnd i xs `isSuffixOf` xs
200-- > \i xs -> length (takeEnd i xs) == min (max 0 i) (length xs)
201takeEnd :: Int -> [a] -> [a]
202takeEnd i xs
203    | i <= 0 = []
204    | otherwise = f xs (drop i xs)
205    where f (x:xs) (y:ys) = f xs ys
206          f xs _ = xs
207
208-- | Drop a number of elements from the end of the list.
209--
210-- > dropEnd 3 "hello"  == "he"
211-- > dropEnd 5 "bye"    == ""
212-- > dropEnd (-1) "bye" == "bye"
213-- > \i xs -> dropEnd i xs `isPrefixOf` xs
214-- > \i xs -> length (dropEnd i xs) == max 0 (length xs - max 0 i)
215-- > \i -> take 3 (dropEnd 5 [i..]) == take 3 [i..]
216dropEnd :: Int -> [a] -> [a]
217dropEnd i xs
218    | i <= 0 = xs
219    | otherwise = f xs (drop i xs)
220    where f (x:xs) (y:ys) = x : f xs ys
221          f _ _ = []
222
223
224-- | @'splitAtEnd' n xs@ returns a split where the second element tries to
225--   contain @n@ elements.
226--
227-- > splitAtEnd 3 "hello" == ("he","llo")
228-- > splitAtEnd 3 "he"    == ("", "he")
229-- > \i xs -> uncurry (++) (splitAt i xs) == xs
230-- > \i xs -> splitAtEnd i xs == (dropEnd i xs, takeEnd i xs)
231splitAtEnd :: Int -> [a] -> ([a], [a])
232splitAtEnd i xs
233    | i <= 0 = (xs, [])
234    | otherwise = f xs (drop i xs)
235    where f (x:xs) (y:ys) = first (x:) $ f xs ys
236          f xs _ = ([], xs)
237
238
239-- | 'zip' against an enumeration.
240--   Never truncates the output - raises an error if the enumeration runs out.
241--
242-- > \i xs -> zip [i..] xs == zipFrom i xs
243-- > zipFrom False [1..3] == [(False,1),(True, 2)]
244zipFrom :: Enum a => a -> [b] -> [(a, b)]
245zipFrom = zipWithFrom (,)
246
247-- | 'zipFrom' generalised to any combining operation.
248--   Never truncates the output - raises an error if the enumeration runs out.
249--
250-- > \i xs -> zipWithFrom (,) i xs == zipFrom i xs
251zipWithFrom :: Enum a => (a -> b -> c) -> a -> [b] -> [c]
252-- would love to deforest the intermediate [a..] list
253-- but would require Bounded and Eq as well, so better go for simplicit
254zipWithFrom f a = zipWith f [a..]
255
256
257-- | A merging of 'unzip' and 'concat'.
258--
259-- > concatUnzip [("a","AB"),("bc","C")] == ("abc","ABC")
260concatUnzip :: [([a], [b])] -> ([a], [b])
261concatUnzip = (concat *** concat) . unzip
262
263-- | A merging of 'unzip3' and 'concat'.
264--
265-- > concatUnzip3 [("a","AB",""),("bc","C","123")] == ("abc","ABC","123")
266concatUnzip3 :: [([a],[b],[c])] -> ([a],[b],[c])
267concatUnzip3 xs = (concat a, concat b, concat c)
268    where (a,b,c) = unzip3 xs
269
270
271-- | A version of 'takeWhile' operating from the end.
272--
273-- > takeWhileEnd even [2,3,4,6] == [4,6]
274takeWhileEnd :: (a -> Bool) -> [a] -> [a]
275takeWhileEnd f = reverse . takeWhile f . reverse
276
277
278-- | Remove spaces from the start of a string, see 'trim'.
279trimStart :: String -> String
280trimStart = dropWhile isSpace
281
282-- | Remove spaces from the end of a string, see 'trim'.
283trimEnd :: String -> String
284trimEnd = dropWhileEnd isSpace
285
286-- | Remove spaces from either side of a string. A combination of 'trimEnd' and 'trimStart'.
287--
288-- > trim      "  hello   " == "hello"
289-- > trimStart "  hello   " == "hello   "
290-- > trimEnd   "  hello   " == "  hello"
291-- > \s -> trim s == trimEnd (trimStart s)
292trim :: String -> String
293trim = trimEnd . trimStart
294
295-- | Convert a string to lower case.
296--
297-- > lower "This is A TEST" == "this is a test"
298-- > lower "" == ""
299lower :: String -> String
300lower = map toLower
301
302-- | Convert a string to upper case.
303--
304-- > upper "This is A TEST" == "THIS IS A TEST"
305-- > upper "" == ""
306upper :: String -> String
307upper = map toUpper
308
309
310-- | Split the first word off a string. Useful for when starting to parse the beginning
311--   of a string, but you want to accurately preserve whitespace in the rest of the string.
312--
313-- > word1 "" == ("", "")
314-- > word1 "keyword rest of string" == ("keyword","rest of string")
315-- > word1 "  keyword\n  rest of string" == ("keyword","rest of string")
316-- > \s -> fst (word1 s) == concat (take 1 $ words s)
317-- > \s -> words (snd $ word1 s) == drop 1 (words s)
318word1 :: String -> (String, String)
319word1 = second trimStart . break isSpace . trimStart
320
321-- | Split the first line off a string.
322--
323-- > line1 "" == ("", "")
324-- > line1 "test" == ("test","")
325-- > line1 "test\n" == ("test","")
326-- > line1 "test\nrest" == ("test","rest")
327-- > line1 "test\nrest\nmore" == ("test","rest\nmore")
328line1 :: String -> (String, String)
329line1 = second drop1 . break (== '\n')
330
331-- | Escape a string such that it can be inserted into an HTML document or @\"@ attribute
332--   without any special interpretation. This requires escaping the @<@, @>@, @&@ and @\"@ characters.
333--   Note that it will escape @\"@ and @\'@ even though that is not required in an HTML body (but is not harmful).
334--
335-- > escapeHTML "this is a test" == "this is a test"
336-- > escapeHTML "<b>\"g&t\"</n>" == "&lt;b&gt;&quot;g&amp;t&quot;&lt;/n&gt;"
337-- > escapeHTML "t'was another test" == "t&#39;was another test"
338escapeHTML :: String -> String
339escapeHTML = concatMap f
340    where
341        f '>' = "&gt;"
342        f '<' = "&lt;"
343        f '&' = "&amp;"
344        f '\"' = "&quot;"
345        f '\'' = "&#39;"
346        f x = [x]
347
348-- | Invert of 'escapeHTML' (does not do general HTML unescaping)
349--
350-- > \xs -> unescapeHTML (escapeHTML xs) == xs
351unescapeHTML :: String -> String
352unescapeHTML ('&':xs)
353    | Just xs <- stripPrefix "lt;" xs = '<' : unescapeHTML xs
354    | Just xs <- stripPrefix "gt;" xs = '>' : unescapeHTML xs
355    | Just xs <- stripPrefix "amp;" xs = '&' : unescapeHTML xs
356    | Just xs <- stripPrefix "quot;" xs = '\"' : unescapeHTML xs
357    | Just xs <- stripPrefix "#39;" xs = '\'' : unescapeHTML xs
358unescapeHTML (x:xs) = x : unescapeHTML xs
359unescapeHTML [] = []
360
361
362-- | Escape a string so it can form part of a JSON literal.
363--   This requires escaping the special whitespace and control characters. Additionally,
364--   Note that it does /not/ add quote characters around the string.
365--
366-- > escapeJSON "this is a test" == "this is a test"
367-- > escapeJSON "\ttab\nnewline\\" == "\\ttab\\nnewline\\\\"
368-- > escapeJSON "\ESC[0mHello" == "\\u001b[0mHello"
369escapeJSON :: String -> String
370escapeJSON x = concatMap f x
371    where f '\"' = "\\\""
372          f '\\' = "\\\\"
373          -- the spaces are technically optional, but we include them so the JSON is readable
374          f '\b' = "\\b"
375          f '\f' = "\\f"
376          f '\n' = "\\n"
377          f '\r' = "\\r"
378          f '\t' = "\\t"
379          f x | isControl x = "\\u" ++ takeEnd 4 ("0000" ++ showHex (ord x) "")
380          f x = [x]
381
382-- | General JSON unescaping, inversion of 'escapeJSON' and all other JSON escapes.
383--
384-- > \xs -> unescapeJSON (escapeJSON xs) == xs
385unescapeJSON :: String -> String
386unescapeJSON ('\\':x:xs)
387    | x == '\"' = '\"' : unescapeJSON xs
388    | x == '\\' = '\\' : unescapeJSON xs
389    | x == '/' = '/' : unescapeJSON xs
390    | x == 'b' = '\b' : unescapeJSON xs
391    | x == 'f' = '\f' : unescapeJSON xs
392    | x == 'n' = '\n' : unescapeJSON xs
393    | x == 'r' = '\r' : unescapeJSON xs
394    | x == 't' = '\t' : unescapeJSON xs
395    | x == 'u', let (a,b) = splitAt 4 xs, length a == 4, [(i, "")] <- readHex a = chr i : unescapeJSON b
396unescapeJSON (x:xs) = x : unescapeJSON xs
397unescapeJSON [] = []
398
399
400-- | A version of 'group' where the equality is done on some extracted value.
401groupOn :: Eq b => (a -> b) -> [a] -> [[a]]
402groupOn f = groupBy ((==) `on2` f)
403    -- redefine on so we avoid duplicate computation for most values.
404    where (.*.) `on2` f = \x -> let fx = f x in \y -> fx .*. f y
405
406
407-- | /DEPRECATED/ Use 'nubOrdOn', since this function is _O(n^2)_.
408--
409--   A version of 'nub' where the equality is done on some extracted value.
410--   @nubOn f@ is equivalent to @nubBy ((==) `on` f)@, but has the
411--   performance advantage of only evaluating @f@ once for each element in the
412--   input list.
413{-# DEPRECATED nubOn "Use nubOrdOn, since this function is O(n^2)" #-}
414nubOn :: Eq b => (a -> b) -> [a] -> [a]
415nubOn f = map snd . nubBy ((==) `on` fst) . map (\x -> let y = f x in y `seq` (y, x))
416
417-- | A version of 'maximum' where the comparison is done on some extracted value.
418--   Raises an error if the list is empty. Only calls the function once per element.
419--
420-- > maximumOn id [] == undefined
421-- > maximumOn length ["test","extra","a"] == "extra"
422maximumOn :: (Partial, Ord b) => (a -> b) -> [a] -> a
423maximumOn f [] = error "Data.List.Extra.maximumOn: empty list"
424maximumOn f (x:xs) = g x (f x) xs
425    where
426        g v mv [] = v
427        g v mv (x:xs) | mx > mv = g x mx xs
428                      | otherwise = g v mv xs
429            where mx = f x
430
431
432-- | A version of 'minimum' where the comparison is done on some extracted value.
433--   Raises an error if the list is empty. Only calls the function once per element.
434--
435-- > minimumOn id [] == undefined
436-- > minimumOn length ["test","extra","a"] == "a"
437minimumOn :: (Partial, Ord b) => (a -> b) -> [a] -> a
438minimumOn f [] = error "Data.List.Extra.minimumOn: empty list"
439minimumOn f (x:xs) = g x (f x) xs
440    where
441        g v mv [] = v
442        g v mv (x:xs) | mx < mv = g x mx xs
443                      | otherwise = g v mv xs
444            where mx = f x
445
446-- | A combination of 'group' and 'sort'.
447--
448-- > groupSort [(1,'t'),(3,'t'),(2,'e'),(2,'s')] == [(1,"t"),(2,"es"),(3,"t")]
449-- > \xs -> map fst (groupSort xs) == sort (nub (map fst xs))
450-- > \xs -> concatMap snd (groupSort xs) == map snd (sortOn fst xs)
451groupSort :: Ord k => [(k, v)] -> [(k, [v])]
452groupSort = map (\x -> (fst $ head x, map snd x)) . groupOn fst . sortOn fst
453
454-- | A combination of 'group' and 'sort', using a part of the value to compare on.
455--
456-- > groupSortOn length ["test","of","sized","item"] == [["of"],["test","item"],["sized"]]
457groupSortOn :: Ord b => (a -> b) -> [a] -> [[a]]
458groupSortOn f = map (map snd) . groupBy ((==) `on` fst) . sortBy (compare `on` fst) . map (f &&& id)
459
460-- | A combination of 'group' and 'sort', using a predicate to compare on.
461--
462-- > groupSortBy (compare `on` length) ["test","of","sized","item"] == [["of"],["test","item"],["sized"]]
463groupSortBy :: (a -> a -> Ordering) -> [a] -> [[a]]
464groupSortBy f = groupBy (\a b -> f a b == EQ) . sortBy f
465
466
467-- | A strict version of 'sum'.
468--   Unlike 'sum' this function is always strict in the `Num` argument,
469--   whereas the standard version is only strict if the optimiser kicks in.
470--
471-- > sum' [1, 2, 3] == 6
472sum' :: (Num a) => [a] -> a
473sum' = foldl' (+) 0
474
475-- | A strict version of 'sum', using a custom valuation function.
476--
477-- > sumOn' read ["1", "2", "3"] == 6
478sumOn' :: (Num b) => (a -> b) -> [a] -> b
479sumOn' f = foldl' (\acc x -> acc + f x) 0
480
481-- | A strict version of 'product'.
482--
483-- > product' [1, 2, 4] == 8
484product' :: (Num a) => [a] -> a
485product' = foldl' (*) 1
486
487-- | A strict version of 'product', using a custom valuation function.
488--
489-- > productOn' read ["1", "2", "4"] == 8
490productOn' :: (Num b) => (a -> b) -> [a] -> b
491productOn' f = foldl' (\acc x -> acc * f x) 1
492
493-- | Merge two lists which are assumed to be ordered.
494--
495-- > merge "ace" "bd" == "abcde"
496-- > \xs ys -> merge (sort xs) (sort ys) == sort (xs ++ ys)
497merge :: Ord a => [a] -> [a] -> [a]
498merge = mergeBy compare
499
500
501-- | Like 'merge', but with a custom ordering function.
502mergeBy :: (a -> a -> Ordering) -> [a] -> [a] -> [a]
503mergeBy f xs [] = xs
504mergeBy f [] ys = ys
505mergeBy f (x:xs) (y:ys)
506    | f x y /= GT = x : mergeBy f xs (y:ys)
507    | otherwise = y : mergeBy f (x:xs) ys
508
509
510-- | Replace a subsequence everywhere it occurs. The first argument must
511--   not be the empty list.
512--
513-- > replace "el" "_" "Hello Bella" == "H_lo B_la"
514-- > replace "el" "e" "Hello"       == "Helo"
515-- > replace "" "e" "Hello"         == undefined
516-- > \xs ys -> not (null xs) ==> replace xs xs ys == ys
517replace :: (Partial, Eq a) => [a] -> [a] -> [a] -> [a]
518replace [] _ _ = error "Extra.replace, first argument cannot be empty"
519replace from to xs | Just xs <- stripPrefix from xs = to ++ replace from to xs
520replace from to (x:xs) = x : replace from to xs
521replace from to [] = []
522
523
524-- | Break, but from the end.
525--
526-- > breakEnd isLower "youRE" == ("you","RE")
527-- > breakEnd isLower "youre" == ("youre","")
528-- > breakEnd isLower "YOURE" == ("","YOURE")
529-- > \f xs -> breakEnd (not . f) xs == spanEnd f  xs
530breakEnd :: (a -> Bool) -> [a] -> ([a], [a])
531breakEnd f = swap . both reverse . break f . reverse
532
533-- | Span, but from the end.
534--
535-- > spanEnd isUpper "youRE" == ("you","RE")
536-- > spanEnd (not . isSpace) "x y z" == ("x y ","z")
537-- > \f xs -> uncurry (++) (spanEnd f xs) == xs
538-- > \f xs -> spanEnd f xs == swap (both reverse (span f (reverse xs)))
539spanEnd :: (a -> Bool) -> [a] -> ([a], [a])
540spanEnd f = breakEnd (not . f)
541
542
543-- | A variant of 'words' with a custom test. In particular,
544--   adjacent separators are discarded, as are leading or trailing separators.
545--
546-- > wordsBy (== ':') "::xyz:abc::123::" == ["xyz","abc","123"]
547-- > \s -> wordsBy isSpace s == words s
548wordsBy :: (a -> Bool) -> [a] -> [[a]]
549wordsBy f s = case dropWhile f s of
550    [] -> []
551    x:xs -> (x:w) : wordsBy f (drop1 z)
552        where (w,z) = break f xs
553
554-- | A variant of 'lines' with a custom test. In particular,
555--   if there is a trailing separator it will be discarded.
556--
557-- > linesBy (== ':') "::xyz:abc::123::" == ["","","xyz","abc","","123",""]
558-- > \s -> linesBy (== '\n') s == lines s
559-- > linesBy (== ';') "my;list;here;" == ["my","list","here"]
560linesBy :: (a -> Bool) -> [a] -> [[a]]
561linesBy f [] = []
562linesBy f s = cons $ case break f s of
563    (l, s) -> (l,) $ case s of
564        [] -> []
565        _:s -> linesBy f s
566  where
567    cons ~(h, t) = h : t -- to fix a space leak, see the GHC defn of lines
568
569-- | Find the first element of a list for which the operation returns 'Just', along
570--   with the result of the operation. Like 'find' but useful where the function also
571--   computes some expensive information that can be reused. Particular useful
572--   when the function is monadic, see 'firstJustM'.
573--
574-- > firstJust id [Nothing,Just 3]  == Just 3
575-- > firstJust id [Nothing,Nothing] == Nothing
576firstJust :: (a -> Maybe b) -> [a] -> Maybe b
577firstJust f = listToMaybe . mapMaybe f
578
579
580-- | Equivalent to @drop 1@, but likely to be faster and a single lexeme.
581--
582-- > drop1 ""         == ""
583-- > drop1 "test"     == "est"
584-- > \xs -> drop 1 xs == drop1 xs
585drop1 :: [a] -> [a]
586drop1 [] = []
587drop1 (x:xs) = xs
588
589
590-- | Equivalent to @dropEnd 1@, but likely to be faster and a single lexeme.
591--
592-- > dropEnd1 ""         == ""
593-- > dropEnd1 "test"     == "tes"
594-- > \xs -> dropEnd 1 xs == dropEnd1 xs
595dropEnd1 :: [a] -> [a]
596dropEnd1 [] = []
597dropEnd1 (x:xs) = foldr (\z f y -> y : f z) (const []) xs x
598
599
600-- | Version on `concatMap` generalised to a `Monoid` rather than just a list.
601--
602-- > mconcatMap Sum [1,2,3] == Sum 6
603-- > \f xs -> mconcatMap f xs == concatMap f xs
604mconcatMap :: Monoid b => (a -> b) -> [a] -> b
605mconcatMap f = mconcat . map f
606
607
608-- | Find the first instance of @needle@ in @haystack@.
609-- The first element of the returned tuple
610-- is the prefix of @haystack@ before @needle@ is matched.  The second
611-- is the remainder of @haystack@, starting with the match.
612-- If you want the remainder /without/ the match, use 'stripInfix'.
613--
614-- > breakOn "::" "a::b::c" == ("a", "::b::c")
615-- > breakOn "/" "foobar"   == ("foobar", "")
616-- > \needle haystack -> let (prefix,match) = breakOn needle haystack in prefix ++ match == haystack
617breakOn :: Eq a => [a] -> [a] -> ([a], [a])
618breakOn needle haystack | needle `isPrefixOf` haystack = ([], haystack)
619breakOn needle [] = ([], [])
620breakOn needle (x:xs) = first (x:) $ breakOn needle xs
621
622-- | Similar to 'breakOn', but searches from the end of the
623-- string.
624--
625-- The first element of the returned tuple is the prefix of @haystack@
626-- up to and including the last match of @needle@.  The second is the
627-- remainder of @haystack@, following the match.
628--
629-- > breakOnEnd "::" "a::b::c" == ("a::b::", "c")
630breakOnEnd :: Eq a => [a] -> [a] -> ([a], [a])
631breakOnEnd needle haystack = both reverse $ swap $ breakOn (reverse needle) (reverse haystack)
632
633
634-- | Break a list into pieces separated by the first
635-- list argument, consuming the delimiter. An empty delimiter is
636-- invalid, and will cause an error to be raised.
637--
638-- > splitOn "\r\n" "a\r\nb\r\nd\r\ne" == ["a","b","d","e"]
639-- > splitOn "aaa"  "aaaXaaaXaaaXaaa"  == ["","X","X","X",""]
640-- > splitOn "x"    "x"                == ["",""]
641-- > splitOn "x"    ""                 == [""]
642-- > \s x -> s /= "" ==> intercalate s (splitOn s x) == x
643-- > \c x -> splitOn [c] x                           == split (==c) x
644splitOn :: (Partial, Eq a) => [a] -> [a] -> [[a]]
645splitOn [] _ = error "splitOn, needle may not be empty"
646splitOn _ [] = [[]]
647splitOn needle haystack = a : if null b then [] else splitOn needle $ drop (length needle) b
648    where (a,b) = breakOn needle haystack
649
650
651-- | Splits a list into components delimited by separators,
652-- where the predicate returns True for a separator element.  The
653-- resulting components do not contain the separators.  Two adjacent
654-- separators result in an empty component in the output.
655--
656-- > split (== 'a') "aabbaca" == ["","","bb","c",""]
657-- > split (== 'a') ""        == [""]
658-- > split (== ':') "::xyz:abc::123::" == ["","","xyz","abc","","123","",""]
659-- > split (== ',') "my,list,here" == ["my","list","here"]
660split :: (a -> Bool) -> [a] -> [[a]]
661split f [] = [[]]
662split f (x:xs) | f x = [] : split f xs
663split f (x:xs) | y:ys <- split f xs = (x:y) : ys
664
665
666-- | A version of 'dropWhileEnd' but with different strictness properties.
667--   The function 'dropWhileEnd' can be used on an infinite list and tests the property
668--   on each character. In contrast, 'dropWhileEnd'' is strict in the spine of the list
669--   but only tests the trailing suffix.
670--   This version usually outperforms 'dropWhileEnd' if the list is short or the test is expensive.
671--   Note the tests below cover both the prime and non-prime variants.
672--
673-- > dropWhileEnd  isSpace "ab cde  " == "ab cde"
674-- > dropWhileEnd' isSpace "ab cde  " == "ab cde"
675-- > last (dropWhileEnd  even [undefined,3]) == undefined
676-- > last (dropWhileEnd' even [undefined,3]) == 3
677-- > head (dropWhileEnd  even (3:undefined)) == 3
678-- > head (dropWhileEnd' even (3:undefined)) == undefined
679dropWhileEnd' :: (a -> Bool) -> [a] -> [a]
680dropWhileEnd' p = foldr (\x xs -> if null xs && p x then [] else x : xs) []
681
682
683-- | Drops the given prefix from a list.
684--   It returns the original sequence if the sequence doesn't start with the given prefix.
685--
686-- > dropPrefix "Mr. " "Mr. Men" == "Men"
687-- > dropPrefix "Mr. " "Dr. Men" == "Dr. Men"
688dropPrefix :: Eq a => [a] -> [a] -> [a]
689dropPrefix a b = fromMaybe b $ stripPrefix a b
690
691
692-- | Drops the given suffix from a list.
693--   It returns the original sequence if the sequence doesn't end with the given suffix.
694--
695-- > dropSuffix "!" "Hello World!"  == "Hello World"
696-- > dropSuffix "!" "Hello World!!" == "Hello World!"
697-- > dropSuffix "!" "Hello World."  == "Hello World."
698dropSuffix :: Eq a => [a] -> [a] -> [a]
699dropSuffix a b = fromMaybe b $ stripSuffix a b
700
701-- | Return the prefix of the second list if its suffix
702--   matches the entire first list.
703--
704-- Examples:
705--
706-- > stripSuffix "bar" "foobar" == Just "foo"
707-- > stripSuffix ""    "baz"    == Just "baz"
708-- > stripSuffix "foo" "quux"   == Nothing
709stripSuffix :: Eq a => [a] -> [a] -> Maybe [a]
710stripSuffix a b = reverse <$> stripPrefix (reverse a) (reverse b)
711
712
713-- | Return the the string before and after the search string,
714--   or 'Nothing' if the search string is not present.
715--
716-- Examples:
717--
718-- > stripInfix "::" "a::b::c" == Just ("a", "b::c")
719-- > stripInfix "/" "foobar"   == Nothing
720stripInfix :: Eq a => [a] -> [a] -> Maybe ([a], [a])
721stripInfix needle haystack | Just rest <- stripPrefix needle haystack = Just ([], rest)
722stripInfix needle [] = Nothing
723stripInfix needle (x:xs) = first (x:) <$> stripInfix needle xs
724
725
726-- | Similar to 'stripInfix', but searches from the end of the
727-- string.
728--
729-- > stripInfixEnd "::" "a::b::c" == Just ("a::b", "c")
730stripInfixEnd :: Eq a => [a] -> [a] -> Maybe ([a], [a])
731stripInfixEnd needle haystack = both reverse . swap <$> stripInfix (reverse needle) (reverse haystack)
732
733
734-- | Split a list into chunks of a given size. The last chunk may contain
735--   fewer than n elements. The chunk size must be positive.
736--
737-- > chunksOf 3 "my test" == ["my ","tes","t"]
738-- > chunksOf 3 "mytest"  == ["myt","est"]
739-- > chunksOf 8 ""        == []
740-- > chunksOf 0 "test"    == undefined
741chunksOf :: Partial => Int -> [a] -> [[a]]
742chunksOf i xs | i <= 0 = error $ "chunksOf, number must be positive, got " ++ show i
743chunksOf i xs = repeatedly (splitAt i) xs
744
745
746-- | /O(n log n)/. The 'nubSort' function sorts and removes duplicate elements from a list.
747-- In particular, it keeps only the first occurrence of each element.
748--
749-- > nubSort "this is a test" == " aehist"
750-- > \xs -> nubSort xs == nub (sort xs)
751nubSort :: Ord a => [a] -> [a]
752nubSort = nubSortBy compare
753
754-- | A version of 'nubSort' which operates on a portion of the value.
755--
756-- > nubSortOn length ["a","test","of","this"] == ["a","of","test"]
757nubSortOn :: Ord b => (a -> b) -> [a] -> [a]
758nubSortOn f = nubSortBy (compare `on` f)
759
760-- | A version of 'nubSort' with a custom predicate.
761--
762-- > nubSortBy (compare `on` length) ["a","test","of","this"] == ["a","of","test"]
763nubSortBy :: (a -> a -> Ordering) -> [a] -> [a]
764nubSortBy cmp = f . sortBy cmp
765    where f (x1:x2:xs) | cmp x1 x2 == EQ = f (x1:xs)
766          f (x:xs) = x : f xs
767          f [] = []
768
769-- | /O(n log n)/. The 'nubOrd' function removes duplicate elements from a list.
770-- In particular, it keeps only the first occurrence of each element.
771-- Unlike the standard 'nub' operator, this version requires an 'Ord' instance
772-- and consequently runs asymptotically faster.
773--
774-- > nubOrd "this is a test" == "this ae"
775-- > nubOrd (take 4 ("this" ++ undefined)) == "this"
776-- > \xs -> nubOrd xs == nub xs
777nubOrd :: Ord a => [a] -> [a]
778nubOrd = nubOrdBy compare
779
780-- | A version of 'nubOrd' which operates on a portion of the value.
781--
782-- > nubOrdOn length ["a","test","of","this"] == ["a","test","of"]
783nubOrdOn :: Ord b => (a -> b) -> [a] -> [a]
784nubOrdOn f = map snd . nubOrdBy (compare `on` fst) . map (f &&& id)
785
786-- | A version of 'nubOrd' with a custom predicate.
787--
788-- > nubOrdBy (compare `on` length) ["a","test","of","this"] == ["a","test","of"]
789nubOrdBy :: (a -> a -> Ordering) -> [a] -> [a]
790nubOrdBy cmp xs = f E xs
791    where f seen [] = []
792          f seen (x:xs) | memberRB cmp x seen = f seen xs
793                        | otherwise = x : f (insertRB cmp x seen) xs
794
795---------------------------------------------------------------------
796-- OKASAKI RED BLACK TREE
797-- Taken from https://www.cs.kent.ac.uk/people/staff/smk/redblack/Untyped.hs
798-- But with the Color = R|B fused into the tree
799
800data RB a = E | T_R (RB a) a (RB a) | T_B (RB a) a (RB a) deriving Show
801
802{- Insertion and membership test as by Okasaki -}
803insertRB :: (a -> a -> Ordering) -> a -> RB a -> RB a
804insertRB cmp x s = case ins s of
805    T_R a z b -> T_B a z b
806    x -> x
807    where
808    ins E = T_R E x E
809    ins s@(T_B a y b) = case cmp x y of
810        LT -> lbalance (ins a) y b
811        GT -> rbalance a y (ins b)
812        EQ -> s
813    ins s@(T_R a y b) = case cmp x y of
814        LT -> T_R (ins a) y b
815        GT -> T_R a y (ins b)
816        EQ -> s
817
818memberRB :: (a -> a -> Ordering) -> a -> RB a -> Bool
819memberRB cmp x E = False
820memberRB cmp x (T_R a y b) = case cmp x y of
821    LT -> memberRB cmp x a
822    GT -> memberRB cmp x b
823    EQ -> True
824memberRB cmp x (T_B a y b) = case cmp x y of
825    LT -> memberRB cmp x a
826    GT -> memberRB cmp x b
827    EQ -> True
828
829{- balance: first equation is new,
830   to make it work with a weaker invariant -}
831lbalance, rbalance :: RB a -> a -> RB a -> RB a
832lbalance (T_R a x b) y (T_R c z d) = T_R (T_B a x b) y (T_B c z d)
833lbalance (T_R (T_R a x b) y c) z d = T_R (T_B a x b) y (T_B c z d)
834lbalance (T_R a x (T_R b y c)) z d = T_R (T_B a x b) y (T_B c z d)
835lbalance a x b = T_B a x b
836rbalance (T_R a x b) y (T_R c z d) = T_R (T_B a x b) y (T_B c z d)
837rbalance a x (T_R b y (T_R c z d)) = T_R (T_B a x b) y (T_B c z d)
838rbalance a x (T_R (T_R b y c) z d) = T_R (T_B a x b) y (T_B c z d)
839rbalance a x b = T_B a x b
840
841
842-- | Like 'zipWith', but keep going to the longest value. The function
843--   argument will always be given at least one 'Just', and while both
844--   lists have items, two 'Just' values.
845--
846-- > zipWithLongest (,) "a" "xyz" == [(Just 'a', Just 'x'), (Nothing, Just 'y'), (Nothing, Just 'z')]
847-- > zipWithLongest (,) "a" "x" == [(Just 'a', Just 'x')]
848-- > zipWithLongest (,) "" "x" == [(Nothing, Just 'x')]
849zipWithLongest :: (Maybe a -> Maybe b -> c) -> [a] -> [b] -> [c]
850zipWithLongest f [] [] = []
851zipWithLongest f (x:xs) (y:ys) = f (Just x) (Just y) : zipWithLongest f xs ys
852zipWithLongest f [] ys = map (f Nothing . Just) ys
853zipWithLongest f xs [] = map ((`f` Nothing) . Just) xs
854
855-- | Lazily compare the length of a 'Foldable' with a number.
856--
857-- > compareLength [1,2,3] 1 == GT
858-- > compareLength [1,2] 2 == EQ
859-- > \(xs :: [Int]) n -> compareLength xs n == compare (length xs) n
860-- > compareLength (1:2:3:undefined) 2 == GT
861compareLength :: (Ord b, Num b, Foldable f) => f a -> b -> Ordering
862compareLength = foldr (\_ acc n -> if n > 0 then acc (n - 1) else GT) (compare 0)
863
864-- | Lazily compare the length of two 'Foldable's.
865-- > comparingLength [1,2,3] [False] == GT
866-- > comparingLength [1,2] "ab" == EQ
867-- > \(xs :: [Int]) (ys :: [Int]) -> comparingLength xs ys == Data.Ord.comparing length xs ys
868-- > comparingLength [1,2] (1:2:3:undefined) == LT
869-- > comparingLength (1:2:3:undefined) [1,2] == GT
870comparingLength :: (Foldable f1, Foldable f2) => f1 a -> f2 b -> Ordering
871comparingLength x y = go (toList x) (toList y)
872  where
873    go [] [] = EQ
874    go [] (_:_) = LT
875    go (_:_) [] = GT
876    go (_:xs) (_:ys) = go xs ys
877