{-# LANGUAGE ScopedTypeVariables #-}
-- | Utility functions for lists.
module Agda.Utils.List where
import Control.Arrow (first)
import Data.Functor ((<$>))
import Data.Function
import qualified Data.List as List
import Data.Maybe
import qualified Data.Map as Map
import qualified Data.Set as Set
import Text.Show.Functions ()
import qualified Agda.Utils.Bag as Bag
import Agda.Utils.Tuple
-- | Append a single element at the end.
-- Time: O(length); use only on small lists.
snoc :: [a] -> a -> [a]
snoc xs x = xs ++ [x]
-- | Case distinction for lists, with list first.
-- Cf. 'Agda.Utils.Null.ifNull'.
caseList :: [a] -> b -> (a -> [a] -> b) -> b
caseList xs n c = listCase n c xs
-- | Case distinction for lists, with list first.
-- Cf. 'Agda.Utils.Null.ifNull'.
caseListM :: Monad m => m [a] -> m b -> (a -> [a] -> m b) -> m b
caseListM mxs n c = listCase n c =<< mxs
-- | Case distinction for lists, with list last.
listCase :: b -> (a -> [a] -> b) -> [a] -> b
listCase n c [] = n
listCase n c (x:xs) = c x xs
-- | Head function (safe).
headMaybe :: [a] -> Maybe a
headMaybe = listToMaybe
-- | Head function (safe). Returns a default value on empty lists.
--
-- > headWithDefault 42 [] = 42
-- > headWithDefault 42 [1,2,3] = 1
headWithDefault :: a -> [a] -> a
headWithDefault def = fromMaybe def . headMaybe
-- | Tail function (safe).
tailMaybe :: [a] -> Maybe [a]
tailMaybe = fmap snd . uncons
-- | Tail function (safe). Returns a default list on empty lists.
tailWithDefault :: [a] -> [a] -> [a]
tailWithDefault def = fromMaybe def . tailMaybe
-- | Last element (safe).
lastMaybe :: [a] -> Maybe a
lastMaybe [] = Nothing
lastMaybe xs = Just $ last xs
-- | Last two elements (safe).
last2 :: [a] -> Maybe (a, a)
last2 (x : y : xs) = Just $ loop x y xs
where
loop x y [] = (x, y)
loop x y (z:xs) = loop y z xs
last2 _ = Nothing
-- | Drop from the end of a list.
-- @dropEnd n = reverse . drop n . reverse@
-- (Forces the whole list even for @n==0@.)
dropEnd :: forall a. Int -> [a] -> [a]
dropEnd n = snd . foldr f (n, [])
where
f :: a -> (Int, [a]) -> (Int, [a])
f x (n, xs)
| n <= 0 = (0, x:xs)
| otherwise = (n-1, xs)
-- | Opposite of cons @(:)@, safe.
uncons :: [a] -> Maybe (a, [a])
uncons [] = Nothing
uncons (x:xs) = Just (x,xs)
-- | Maybe cons. @mcons ma as = maybeToList ma ++ as@
mcons :: Maybe a -> [a] -> [a]
mcons ma as = maybe as (:as) ma
-- | 'init' and 'last' in one go, safe.
initLast :: [a] -> Maybe ([a],a)
initLast [] = Nothing
initLast (a:as) = Just $ loop a as where
loop a [] = ([], a)
loop a (b : bs) = mapFst (a:) $ loop b bs
-- | init, safe.
initMaybe :: [a] -> Maybe [a]
initMaybe = fmap fst . initLast
-- | Lookup function (partially safe).
(!!!) :: [a] -> Int -> Maybe a
[] !!! _ = Nothing
(x : _) !!! 0 = Just x
(_ : xs) !!! n = xs !!! (n - 1)
-- | Lookup function with default value for index out of range.
-- The name is chosen akin to 'Data.List.genericIndex'.
indexWithDefault :: a -> [a] -> Int -> a
indexWithDefault a [] _ = a
indexWithDefault a (x : _) 0 = x
indexWithDefault a (_ : xs) n = indexWithDefault a xs (n - 1)
-- | Find an element satisfying a predicate and return it with its index.
-- TODO: more efficient implementation!?
findWithIndex :: (a -> Bool) -> [a] -> Maybe (a, Int)
findWithIndex p as = headMaybe $ filter (p . fst) $ zip as [0..]
-- | downFrom n = [n-1,..1,0]
downFrom :: Integral a => a -> [a]
downFrom n | n <= 0 = []
| otherwise = let n' = n-1 in n' : downFrom n'
-- | Update the first element of a list, if it exists.
updateHead :: (a -> a) -> [a] -> [a]
updateHead f [] = []
updateHead f (a : as) = f a : as
-- | Update the last element of a list, if it exists.
updateLast :: (a -> a) -> [a] -> [a]
updateLast f [] = []
updateLast f [a] = [f a]
updateLast f (a : as@(_ : _)) = a : updateLast f as
-- | Update nth element of a list, if it exists.
-- Precondition: the index is >= 0.
updateAt :: Int -> (a -> a) -> [a] -> [a]
updateAt _ f [] = []
updateAt 0 f (a : as) = f a : as
updateAt n f (a : as) = a : updateAt (n-1) f as
-- | A generalized version of @partition@.
-- (Cf. @mapMaybe@ vs. @filter@).
mapEither :: (a -> Either b c) -> [a] -> ([b],[c])
{-# INLINE mapEither #-}
mapEither f xs = foldr (deal f) ([],[]) xs
deal :: (a -> Either b c) -> a -> ([b],[c]) -> ([b],[c])
deal f a ~(bs,cs) = case f a of
Left b -> (b:bs, cs)
Right c -> (bs, c:cs)
-- | Split off the largest suffix whose elements satisfy a predicate.
--
-- @spanEnd p xs = (ys, zs)@
-- where @xs = ys ++ zs@
-- and @all p zs@
-- and @maybe True (not . p) (lastMaybe yz)@.
spanEnd :: forall a. (a -> Bool) -> [a] -> ([a], [a])
spanEnd p = snd . foldr f (True, ([], []))
where
f :: a -> (Bool, ([a], [a])) -> (Bool, ([a], [a]))
f x (b', (xs, ys)) = (b, if b then (xs, x:ys) else (x:xs, ys))
where b = b' && p x
-- | A generalized version of @takeWhile@.
-- (Cf. @mapMaybe@ vs. @filter@).
takeWhileJust :: (a -> Maybe b) -> [a] -> [b]
takeWhileJust p = loop
where
loop (a : as) | Just b <- p a = b : loop as
loop _ = []
-- | A generalized version of @span@.
spanJust :: (a -> Maybe b) -> [a] -> ([b], [a])
spanJust p = loop
where
loop (a : as) | Just b <- p a = mapFst (b :) $ loop as
loop as = ([], as)
-- | Partition a list into 'Nothing's and 'Just's.
-- @'mapMaybe' f = snd . partitionMaybe f@.
partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe f = loop
where
loop [] = ([], [])
loop (a : as) = case f a of
Nothing -> mapFst (a :) $ loop as
Just b -> mapSnd (b :) $ loop as
-- | Like 'filter', but additionally return the last partition
-- of the list where the predicate is @False@ everywhere.
filterAndRest :: (a -> Bool) -> [a] -> ([a],[a])
filterAndRest p = mapMaybeAndRest $ \ a -> if p a then Just a else Nothing
-- | Like 'mapMaybe', but additionally return the last partition
-- of the list where the function always returns @Nothing@.
mapMaybeAndRest :: (a -> Maybe b) -> [a] -> ([b],[a])
mapMaybeAndRest f = loop [] where
loop acc = \case
[] -> ([], reverse acc)
x:xs | Just y <- f x -> first (y:) $ loop [] xs
| otherwise -> loop (x:acc) xs
-- | Drops from both lists simultaneously until one list is empty.
dropCommon :: [a] -> [b] -> ([a],[b])
dropCommon (x : xs) (y : ys) = dropCommon xs ys
dropCommon xs ys = (xs, ys)
-- | Sublist relation.
isSublistOf :: Eq a => [a] -> [a] -> Bool
isSublistOf [] ys = True
isSublistOf (x : xs) ys =
case dropWhile (x /=) ys of
[] -> False
(_:ys) -> isSublistOf xs ys
type Prefix a = [a]
type Suffix a = [a]
-- | Check if a list has a given prefix. If so, return the list
-- minus the prefix.
stripPrefixBy :: (a -> a -> Bool) -> Prefix a -> [a] -> Maybe (Suffix a)
stripPrefixBy eq = loop
where
loop [] rest = Just rest
loop (_:_) [] = Nothing
loop (p:pat) (r:rest)
| eq p r = loop pat rest
| otherwise = Nothing
-- | @stripSuffix suf xs = Just pre@ iff @xs = pre ++ suf@.
stripSuffix :: Eq a => Suffix a -> [a] -> Maybe (Prefix a)
stripSuffix [] = Just
stripSuffix s = stripReversedSuffix (reverse s)
type ReversedSuffix a = [a]
-- | @stripReversedSuffix rsuf xs = Just pre@ iff @xs = pre ++ reverse suf@.
stripReversedSuffix :: forall a. Eq a => ReversedSuffix a -> [a] -> Maybe (Prefix a)
stripReversedSuffix rs = final . foldr step (SSSStrip rs)
where
-- Step of the automaton (reading input from right to left).
step :: a -> StrSufSt a -> StrSufSt a
step x = \case
SSSMismatch -> SSSMismatch
SSSResult xs -> SSSResult (x:xs)
SSSStrip [] -> SSSResult [x]
SSSStrip (y:ys)
| x == y -> SSSStrip ys
| otherwise -> SSSMismatch
-- Output of the automaton.
final :: StrSufSt a -> Maybe (Prefix a)
final = \case
SSSResult xs -> Just xs
SSSStrip [] -> Just []
_ -> Nothing -- We have not stripped the whole suffix or encountered a mismatch.
-- | Internal state for stripping suffix.
data StrSufSt a
= SSSMismatch -- ^ Error.
| SSSStrip (ReversedSuffix a) -- ^ "Negative string" to remove from end. List may be empty.
| SSSResult [a] -- ^ "Positive string" (result). Non-empty list.
-- | Split a list into sublists. Generalisation of the prelude function
-- @words@.
--
-- > words xs == wordsBy isSpace xs
wordsBy :: (a -> Bool) -> [a] -> [[a]]
wordsBy p xs = yesP xs
where
yesP xs = noP (dropWhile p xs)
noP [] = []
noP xs = ys : yesP zs
where
(ys,zs) = break p xs
-- | Chop up a list in chunks of a given length.
chop :: Int -> [a] -> [[a]]
chop _ [] = []
chop n xs = ys : chop n zs
where (ys,zs) = splitAt n xs
-- | Chop a list at the positions when the predicate holds. Contrary to
-- 'wordsBy', consecutive separator elements will result in an empty segment
-- in the result.
-- > intercalate [x] (chopWhen (== x) xs) == xs
chopWhen :: (a -> Bool) -> [a] -> [[a]]
chopWhen p [] = []
chopWhen p xs =
case break p xs of
(w, []) -> [w]
(w, [_]) -> [w, []]
(w, _ : ys) -> w : chopWhen p ys
-- | All ways of removing one element from a list.
holes :: [a] -> [(a, [a])]
holes [] = []
holes (x:xs) = (x, xs) : map (id -*- (x:)) (holes xs)
-- | Check whether a list is sorted.
--
-- Assumes that the 'Ord' instance implements a partial order.
sorted :: Ord a => [a] -> Bool
sorted [] = True
sorted xs = and $ zipWith (<=) xs (tail xs)
-- | Check whether all elements in a list are distinct from each
-- other. Assumes that the 'Eq' instance stands for an equivalence
-- relation.
distinct :: Eq a => [a] -> Bool
distinct [] = True
distinct (x:xs) = x `notElem` xs && distinct xs
-- | An optimised version of 'distinct'.
--
-- Precondition: The list's length must fit in an 'Int'.
fastDistinct :: Ord a => [a] -> Bool
fastDistinct xs = Set.size (Set.fromList xs) == length xs
-- | Checks if all the elements in the list are equal. Assumes that
-- the 'Eq' instance stands for an equivalence relation.
allEqual :: Eq a => [a] -> Bool
allEqual [] = True
allEqual (x : xs) = all (== x) xs
-- | Returns an (arbitrary) representative for each list element
-- that occurs more than once.
duplicates :: Ord a => [a] -> [a]
duplicates = mapMaybe dup . Bag.groups . Bag.fromList
where
dup (a : _ : _) = Just a
dup _ = Nothing
-- | A variant of 'List.groupBy' which applies the predicate to consecutive
-- pairs.
groupBy' :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy' _ [] = []
groupBy' p xxs@(x : xs) = grp x $ zipWith (\x y -> (p x y, y)) xxs xs
where
grp x ys = (x : map snd xs) : tail
where (xs, rest) = span fst ys
tail = case rest of
[] -> []
((_, z) : zs) -> grp z zs
-- | @'groupOn' f = 'groupBy' (('==') \`on\` f) '.' 'List.sortBy' ('compare' \`on\` f)@.
groupOn :: Ord b => (a -> b) -> [a] -> [[a]]
groupOn f = List.groupBy ((==) `on` f) . List.sortBy (compare `on` f)
-- | @splitExactlyAt n xs = Just (ys, zs)@ iff @xs = ys ++ zs@
-- and @genericLength ys = n@.
splitExactlyAt :: Integral n => n -> [a] -> Maybe ([a], [a])
splitExactlyAt 0 xs = return ([], xs)
splitExactlyAt n [] = Nothing
splitExactlyAt n (x : xs) = mapFst (x :) <$> splitExactlyAt (n-1) xs
-- | A generalised variant of 'elemIndex'.
genericElemIndex :: (Eq a, Integral i) => a -> [a] -> Maybe i
genericElemIndex x xs =
listToMaybe $
map fst $
filter snd $
zip [0..] $
map (== x) xs
-- | Requires both lists to have the same length.
--
-- Otherwise, @Nothing@ is returned.
zipWith' :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
zipWith' f = loop
where
loop [] [] = Just []
loop (x : xs) (y : ys) = (f x y :) <$> loop xs ys
loop [] (_ : _) = Nothing
loop (_ : _) [] = Nothing
-- | Like 'zipWith' but keep the rest of the second list as-is
-- (in case the second list is longer).
--
-- @
-- zipWithKeepRest f as bs == zipWith f as bs ++ drop (length as) bs
-- @
zipWithKeepRest :: (a -> b -> b) -> [a] -> [b] -> [b]
zipWithKeepRest f = loop
where
loop [] bs = bs
loop as [] = []
loop (a : as) (b : bs) = f a b : loop as bs
-- -- UNUSED; a better type would be
-- -- zipWithTails :: (a -> b -> c) -> [a] -> [b] -> ([c], Either [a] [b])
-- -- | Like zipWith, but returns the leftover elements of the input lists.
-- zipWithTails :: (a -> b -> c) -> [a] -> [b] -> ([c], [a] , [b])
-- zipWithTails f xs [] = ([], xs, [])
-- zipWithTails f [] ys = ([], [] , ys)
-- zipWithTails f (x : xs) (y : ys) = (f x y : zs , as , bs)
-- where (zs , as , bs) = zipWithTails f xs ys
-- | Efficient variant of 'nubBy' for finite lists.
--
-- Specification:
--
-- > nubOn f xs == 'nubBy' ((==) `'on'` f) xs.
nubOn :: Ord b => (a -> b) -> [a] -> [a]
nubOn tag =
map snd
. List.sortBy (compare `on` fst)
. map (snd . head)
. List.groupBy ((==) `on` fst)
. List.sortBy (compare `on` fst)
. map (\p@(_, x) -> (tag x, p))
. zip [1..]
-- | Efficient variant of 'nubBy' for finite lists.
--
-- Specification: For each list @xs@ there is a list @ys@ which is a
-- permutation of @xs@ such that
--
-- > uniqOn f xs == 'nubBy' ((==) `'on'` f) ys.
--
-- Furthermore
--
-- > List.sortBy (compare `on` f) (uniqOn f xs) == uniqOn f xs.
uniqOn :: Ord b => (a -> b) -> [a] -> [a]
uniqOn key = Map.elems . Map.fromList . map (\ a -> (key a, a))
-- | Compute the common suffix of two lists.
commonSuffix :: Eq a => [a] -> [a] -> [a]
commonSuffix xs ys = reverse $ (commonPrefix `on` reverse) xs ys
-- | Compute the common prefix of two lists.
commonPrefix :: Eq a => [a] -> [a] -> [a]
commonPrefix [] _ = []
commonPrefix _ [] = []
commonPrefix (x:xs) (y:ys)
| x == y = x : commonPrefix xs ys
| otherwise = []
editDistanceSpec :: Eq a => [a] -> [a] -> Int
editDistanceSpec [] ys = length ys
editDistanceSpec xs [] = length xs
editDistanceSpec (x : xs) (y : ys)
| x == y = editDistanceSpec xs ys
| otherwise = 1 + minimum [ editDistanceSpec (x : xs) ys
, editDistanceSpec xs (y : ys)
, editDistanceSpec xs ys ]
editDistance :: Eq a => [a] -> [a] -> Int
editDistance xs ys = editD 0 0
where xss = List.tails xs
yss = List.tails ys
tbl = Map.fromList [ ((i, j), editD' i j) | i <- [0..length xss - 1], j <- [0..length yss - 1] ]
editD i j = tbl Map.! (i, j)
editD' i j =
case (xss !! i, yss !! j) of
([], ys) -> length ys
(xs, []) -> length xs
(x : xs, y : ys)
| x == y -> editD (i + 1) (j + 1)
| otherwise -> 1 + minimum [ editD (i + 1) j, editD i (j + 1), editD (i + 1) (j + 1) ]