-- | Utility functions for lists.

module Agda.Utils.List where

import Control.Arrow (first, second)

import Data.Array (Array, array, listArray)
import qualified Data.Array as Array
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 qualified Agda.Utils.Bag as Bag
import Agda.Utils.Function (applyWhen)
import Agda.Utils.Tuple

---------------------------------------------------------------------------
-- * Variants of list case, cons, head, tail, init, last
---------------------------------------------------------------------------

-- | Append a single element at the end.
--   Time: O(length); use only on small lists.
snoc :: [a] -> a -> [a]
snoc :: [a] -> a -> [a]
snoc xs :: [a]
xs x :: a
x = [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
x]

-- | Case distinction for lists, with list first.
--   O(1).
--
--   Cf. 'Agda.Utils.Null.ifNull'.
caseList :: [a] -> b -> (a -> [a] -> b) -> b
caseList :: [a] -> b -> (a -> [a] -> b) -> b
caseList xs :: [a]
xs n :: b
n c :: a -> [a] -> b
c = b -> (a -> [a] -> b) -> [a] -> b
forall b a. b -> (a -> [a] -> b) -> [a] -> b
listCase b
n a -> [a] -> b
c [a]
xs

-- | Case distinction for lists, with list first.
--   O(1).
--
--   Cf. 'Agda.Utils.Null.ifNull'.
caseListM :: Monad m => m [a] -> m b -> (a -> [a] -> m b) -> m b
caseListM :: m [a] -> m b -> (a -> [a] -> m b) -> m b
caseListM mxs :: m [a]
mxs n :: m b
n c :: a -> [a] -> m b
c = m b -> (a -> [a] -> m b) -> [a] -> m b
forall b a. b -> (a -> [a] -> b) -> [a] -> b
listCase m b
n a -> [a] -> m b
c ([a] -> m b) -> m [a] -> m b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m [a]
mxs

-- | Case distinction for lists, with list last.
--   O(1).
--
listCase :: b -> (a -> [a] -> b) -> [a] -> b
listCase :: b -> (a -> [a] -> b) -> [a] -> b
listCase n :: b
n c :: a -> [a] -> b
c []     = b
n
listCase n :: b
n c :: a -> [a] -> b
c (x :: a
x:xs :: [a]
xs) = a -> [a] -> b
c a
x [a]
xs

-- | Head function (safe). Returns a default value on empty lists.
--   O(1).
--
-- > headWithDefault 42 []      = 42
-- > headWithDefault 42 [1,2,3] = 1
headWithDefault :: a -> [a] -> a
headWithDefault :: a -> [a] -> a
headWithDefault def :: a
def = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
def (Maybe a -> a) -> ([a] -> Maybe a) -> [a] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe

-- | Tail function (safe).
--   O(1).
tailMaybe :: [a] -> Maybe [a]
tailMaybe :: [a] -> Maybe [a]
tailMaybe = ((a, [a]) -> [a]) -> Maybe (a, [a]) -> Maybe [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, [a]) -> [a]
forall a b. (a, b) -> b
snd (Maybe (a, [a]) -> Maybe [a])
-> ([a] -> Maybe (a, [a])) -> [a] -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Maybe (a, [a])
forall a. [a] -> Maybe (a, [a])
uncons

-- | Tail function (safe).  Returns a default list on empty lists.
--   O(1).
tailWithDefault :: [a] -> [a] -> [a]
tailWithDefault :: [a] -> [a] -> [a]
tailWithDefault def :: [a]
def = [a] -> Maybe [a] -> [a]
forall a. a -> Maybe a -> a
fromMaybe [a]
def (Maybe [a] -> [a]) -> ([a] -> Maybe [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Maybe [a]
forall a. [a] -> Maybe [a]
tailMaybe

-- | Last element (safe).
--   O(n).
lastMaybe :: [a] -> Maybe a
lastMaybe :: [a] -> Maybe a
lastMaybe [] = Maybe a
forall a. Maybe a
Nothing
lastMaybe xs :: [a]
xs = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ [a] -> a
forall a. [a] -> a
last [a]
xs

-- | Last two elements (safe).
--   O(n).
last2 :: [a] -> Maybe (a, a)
last2 :: [a] -> Maybe (a, a)
last2 (x :: a
x : y :: a
y : xs :: [a]
xs) = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just ((a, a) -> Maybe (a, a)) -> (a, a) -> Maybe (a, a)
forall a b. (a -> b) -> a -> b
$ a -> a -> [a] -> (a, a)
forall t. t -> t -> [t] -> (t, t)
loop a
x a
y [a]
xs
  where
  loop :: t -> t -> [t] -> (t, t)
loop x :: t
x y :: t
y []     = (t
x, t
y)
  loop x :: t
x y :: t
y (z :: t
z:xs :: [t]
xs) = t -> t -> [t] -> (t, t)
loop t
y t
z [t]
xs
last2 _ = Maybe (a, a)
forall a. Maybe a
Nothing

-- | Opposite of cons @(:)@, safe.
--   O(1).
uncons :: [a] -> Maybe (a, [a])
uncons :: [a] -> Maybe (a, [a])
uncons []     = Maybe (a, [a])
forall a. Maybe a
Nothing
uncons (x :: a
x:xs :: [a]
xs) = (a, [a]) -> Maybe (a, [a])
forall a. a -> Maybe a
Just (a
x,[a]
xs)

-- | Maybe cons.
--   O(1).
--   @mcons ma as = maybeToList ma ++ as@
mcons :: Maybe a -> [a] -> [a]
mcons :: Maybe a -> [a] -> [a]
mcons ma :: Maybe a
ma as :: [a]
as = [a] -> (a -> [a]) -> Maybe a -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [a]
as (a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as) Maybe a
ma

-- | 'init' and 'last' in one go, safe.
--   O(n).
initLast :: [a] -> Maybe ([a],a)
initLast :: [a] -> Maybe ([a], a)
initLast []     = Maybe ([a], a)
forall a. Maybe a
Nothing
initLast (a :: a
a:as :: [a]
as) = ([a], a) -> Maybe ([a], a)
forall a. a -> Maybe a
Just (([a], a) -> Maybe ([a], a)) -> ([a], a) -> Maybe ([a], a)
forall a b. (a -> b) -> a -> b
$ a -> [a] -> ([a], a)
forall t. t -> [t] -> ([t], t)
loop a
a [a]
as where
  loop :: t -> [t] -> ([t], t)
loop a :: t
a []      = ([], t
a)
  loop a :: t
a (b :: t
b : bs :: [t]
bs) = ([t] -> [t]) -> ([t], t) -> ([t], t)
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (t
at -> [t] -> [t]
forall a. a -> [a] -> [a]
:) (([t], t) -> ([t], t)) -> ([t], t) -> ([t], t)
forall a b. (a -> b) -> a -> b
$ t -> [t] -> ([t], t)
loop t
b [t]
bs

-- | @init@, safe.
--   O(n).
initMaybe :: [a] -> Maybe [a]
initMaybe :: [a] -> Maybe [a]
initMaybe = (([a], a) -> [a]) -> Maybe ([a], a) -> Maybe [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([a], a) -> [a]
forall a b. (a, b) -> a
fst (Maybe ([a], a) -> Maybe [a])
-> ([a] -> Maybe ([a], a)) -> [a] -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Maybe ([a], a)
forall a. [a] -> Maybe ([a], a)
initLast

---------------------------------------------------------------------------
-- * Lookup and indexing
---------------------------------------------------------------------------

-- | Lookup function (partially safe).
--   O(min n index).
(!!!) :: [a] -> Int -> Maybe a
[]       !!! :: [a] -> Int -> Maybe a
!!! _         = Maybe a
forall a. Maybe a
Nothing
(x :: a
x : _)  !!! 0         = a -> Maybe a
forall a. a -> Maybe a
Just a
x
(_ : xs :: [a]
xs) !!! n :: Int
n         = [a]
xs [a] -> Int -> Maybe a
forall a. [a] -> Int -> Maybe a
!!! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)

-- | Lookup function with default value for index out of range.
--   O(min n index).
--
--   The name is chosen akin to 'Data.List.genericIndex'.
indexWithDefault :: a -> [a] -> Int -> a
indexWithDefault :: a -> [a] -> Int -> a
indexWithDefault a :: a
a []       _ = a
a
indexWithDefault a :: a
a (x :: a
x : _)  0 = a
x
indexWithDefault a :: a
a (_ : xs :: [a]
xs) n :: Int
n = a -> [a] -> Int -> a
forall a. a -> [a] -> Int -> a
indexWithDefault a
a [a]
xs (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)

-- | Find an element satisfying a predicate and return it with its index.
--   O(n) in the worst case, e.g. @findWithIndex f xs = Nothing@.
--
--   TODO: more efficient implementation!?
findWithIndex :: (a -> Bool) -> [a] -> Maybe (a, Int)
findWithIndex :: (a -> Bool) -> [a] -> Maybe (a, Int)
findWithIndex p :: a -> Bool
p as :: [a]
as = [(a, Int)] -> Maybe (a, Int)
forall a. [a] -> Maybe a
listToMaybe ([(a, Int)] -> Maybe (a, Int)) -> [(a, Int)] -> Maybe (a, Int)
forall a b. (a -> b) -> a -> b
$ ((a, Int) -> Bool) -> [(a, Int)] -> [(a, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> Bool
p (a -> Bool) -> ((a, Int) -> a) -> (a, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Int) -> a
forall a b. (a, b) -> a
fst) ([(a, Int)] -> [(a, Int)]) -> [(a, Int)] -> [(a, Int)]
forall a b. (a -> b) -> a -> b
$ [a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
as [0..]

-- | A generalised variant of 'elemIndex'.
-- O(n).
genericElemIndex :: (Eq a, Integral i) => a -> [a] -> Maybe i
genericElemIndex :: a -> [a] -> Maybe i
genericElemIndex x :: a
x xs :: [a]
xs =
  [i] -> Maybe i
forall a. [a] -> Maybe a
listToMaybe ([i] -> Maybe i) -> [i] -> Maybe i
forall a b. (a -> b) -> a -> b
$
  ((i, Bool) -> i) -> [(i, Bool)] -> [i]
forall a b. (a -> b) -> [a] -> [b]
map (i, Bool) -> i
forall a b. (a, b) -> a
fst ([(i, Bool)] -> [i]) -> [(i, Bool)] -> [i]
forall a b. (a -> b) -> a -> b
$
  ((i, Bool) -> Bool) -> [(i, Bool)] -> [(i, Bool)]
forall a. (a -> Bool) -> [a] -> [a]
filter (i, Bool) -> Bool
forall a b. (a, b) -> b
snd ([(i, Bool)] -> [(i, Bool)]) -> [(i, Bool)] -> [(i, Bool)]
forall a b. (a -> b) -> a -> b
$
  [i] -> [Bool] -> [(i, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip [0..] ([Bool] -> [(i, Bool)]) -> [Bool] -> [(i, Bool)]
forall a b. (a -> b) -> a -> b
$
  (a -> Bool) -> [a] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x) [a]
xs

-- | @downFrom n = [n-1,..1,0]@.
--   O(n).
downFrom :: Integral a => a -> [a]
downFrom :: a -> [a]
downFrom n :: a
n | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= 0     = []
           | Bool
otherwise = let n' :: a
n' = a
na -> a -> a
forall a. Num a => a -> a -> a
-1 in a
n' a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a]
forall a. Integral a => a -> [a]
downFrom a
n'

---------------------------------------------------------------------------
-- * Update
---------------------------------------------------------------------------

-- | Update the first element of a list, if it exists.
--   O(1).
updateHead :: (a -> a) -> [a] -> [a]
updateHead :: (a -> a) -> [a] -> [a]
updateHead _ []       = []
updateHead f :: a -> a
f (a :: a
a : as :: [a]
as) = a -> a
f a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as

-- | Update the last element of a list, if it exists.
--   O(n).
updateLast :: (a -> a) -> [a] -> [a]
updateLast :: (a -> a) -> [a] -> [a]
updateLast _ [] = []
updateLast f :: a -> a
f (a :: a
a : as :: [a]
as) = a -> [a] -> [a]
loop a
a [a]
as
  -- Using a helper function to minimize the pattern matching.
  where
  loop :: a -> [a] -> [a]
loop a :: a
a []       = [a -> a
f a
a]
  loop a :: a
a (b :: a
b : bs :: [a]
bs) = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a] -> [a]
loop a
b [a]
bs

-- | Update nth element of a list, if it exists.
--   @O(min index n)@.
--
--   Precondition: the index is >= 0.
updateAt :: Int -> (a -> a) -> [a] -> [a]
updateAt :: Int -> (a -> a) -> [a] -> [a]
updateAt _ _ [] = []
updateAt 0 f :: a -> a
f (a :: a
a : as :: [a]
as) = a -> a
f a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as
updateAt n :: Int
n f :: a -> a
f (a :: a
a : as :: [a]
as) = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> (a -> a) -> [a] -> [a]
forall a. Int -> (a -> a) -> [a] -> [a]
updateAt (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) a -> a
f [a]
as

---------------------------------------------------------------------------
-- * Sublist extraction and partitioning
---------------------------------------------------------------------------

type Prefix a = [a]  -- ^ The list before the split point.
type Suffix a = [a]  -- ^ The list after the split point.

-- | @splitExactlyAt n xs = Just (ys, zs)@ iff @xs = ys ++ zs@
--   and @genericLength ys = n@.
splitExactlyAt :: Integral n => n -> [a] -> Maybe (Prefix a, Suffix a)
splitExactlyAt :: n -> [a] -> Maybe ([a], [a])
splitExactlyAt 0 xs :: [a]
xs       = ([a], [a]) -> Maybe ([a], [a])
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [a]
xs)
splitExactlyAt n :: n
n []       = Maybe ([a], [a])
forall a. Maybe a
Nothing
splitExactlyAt n :: n
n (x :: a
x : xs :: [a]
xs) = ([a] -> [a]) -> ([a], [a]) -> ([a], [a])
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (([a], [a]) -> ([a], [a])) -> Maybe ([a], [a]) -> Maybe ([a], [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> n -> [a] -> Maybe ([a], [a])
forall n a. Integral n => n -> [a] -> Maybe ([a], [a])
splitExactlyAt (n
nn -> n -> n
forall a. Num a => a -> a -> a
-1) [a]
xs

-- | Drop from the end of a list.
--   O(length).
--
--   @dropEnd n = reverse . drop n . reverse@
--
--   Forces the whole list even for @n==0@.
dropEnd :: forall a. Int -> [a] -> Prefix a
dropEnd :: Int -> [a] -> [a]
dropEnd n :: Int
n = (Int, [a]) -> [a]
forall a b. (a, b) -> b
snd ((Int, [a]) -> [a]) -> ([a] -> (Int, [a])) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (Int, [a]) -> (Int, [a])) -> (Int, [a]) -> [a] -> (Int, [a])
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (Int, [a]) -> (Int, [a])
f (Int
n, [])
  where
  f :: a -> (Int, [a]) -> (Int, [a])
  f :: a -> (Int, [a]) -> (Int, [a])
f x :: a
x (n :: Int
n, xs :: [a]
xs) = (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1, Bool -> ([a] -> [a]) -> [a] -> [a]
forall a. Bool -> (a -> a) -> a -> a
applyWhen (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= 0) (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) [a]
xs)

-- | Split off the largest suffix whose elements satisfy a predicate.
--   O(n).
--
--   @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] -> (Prefix a, Suffix a)
spanEnd :: (a -> Bool) -> [a] -> ([a], [a])
spanEnd p :: a -> Bool
p = (Bool, ([a], [a])) -> ([a], [a])
forall a b. (a, b) -> b
snd ((Bool, ([a], [a])) -> ([a], [a]))
-> ([a] -> (Bool, ([a], [a]))) -> [a] -> ([a], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (Bool, ([a], [a])) -> (Bool, ([a], [a])))
-> (Bool, ([a], [a])) -> [a] -> (Bool, ([a], [a]))
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (Bool, ([a], [a])) -> (Bool, ([a], [a]))
f (Bool
True, ([], []))
  where
  f :: a -> (Bool, ([a], [a])) -> (Bool, ([a], [a]))
  f :: a -> (Bool, ([a], [a])) -> (Bool, ([a], [a]))
f x :: a
x (b' :: Bool
b', (xs :: [a]
xs, ys :: [a]
ys)) = (Bool
b, if Bool
b then ([a]
xs, a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys) else (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs, [a]
ys))
    where b :: Bool
b = Bool
b' Bool -> Bool -> Bool
&& a -> Bool
p a
x

-- | A generalized version of @takeWhile@.
--   (Cf. @mapMaybe@ vs. @filter@).
--   @O(length . takeWhileJust f).
--
--   @takeWhileJust f = fst . spanJust f@.
takeWhileJust :: (a -> Maybe b) -> [a] -> Prefix b
takeWhileJust :: (a -> Maybe b) -> [a] -> Prefix b
takeWhileJust p :: a -> Maybe b
p = [a] -> Prefix b
loop
  where
    loop :: [a] -> Prefix b
loop (a :: a
a : as :: [a]
as) | Just b :: b
b <- a -> Maybe b
p a
a = b
b b -> Prefix b -> Prefix b
forall a. a -> [a] -> [a]
: [a] -> Prefix b
loop [a]
as
    loop _ = []

-- | A generalized version of @span@.
--   @O(length . fst . spanJust f)@.
spanJust :: (a -> Maybe b) -> [a] -> (Prefix b, Suffix a)
spanJust :: (a -> Maybe b) -> [a] -> (Prefix b, [a])
spanJust p :: a -> Maybe b
p = [a] -> (Prefix b, [a])
loop
  where
    loop :: [a] -> (Prefix b, [a])
loop (a :: a
a : as :: [a]
as) | Just b :: b
b <- a -> Maybe b
p a
a = (Prefix b -> Prefix b) -> (Prefix b, [a]) -> (Prefix b, [a])
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (b
b b -> Prefix b -> Prefix b
forall a. a -> [a] -> [a]
:) ((Prefix b, [a]) -> (Prefix b, [a]))
-> (Prefix b, [a]) -> (Prefix b, [a])
forall a b. (a -> b) -> a -> b
$ [a] -> (Prefix b, [a])
loop [a]
as
    loop as :: [a]
as                       = ([], [a]
as)

-- | Partition a list into 'Nothing's and 'Just's.
--   O(n).
--
--   @partitionMaybe f = partitionEithers . map (\ a -> maybe (Left a) Right (f a))@
--
--   Note: @'mapMaybe' f = snd . partitionMaybe f@.
partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe :: (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe f :: a -> Maybe b
f = [a] -> ([a], [b])
loop
  where
    loop :: [a] -> ([a], [b])
loop []       = ([], [])
    loop (a :: a
a : as :: [a]
as) = case a -> Maybe b
f a
a of
      Nothing -> ([a] -> [a]) -> ([a], [b]) -> ([a], [b])
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (([a], [b]) -> ([a], [b])) -> ([a], [b]) -> ([a], [b])
forall a b. (a -> b) -> a -> b
$ [a] -> ([a], [b])
loop [a]
as
      Just b :: b
b  -> ([b] -> [b]) -> ([a], [b]) -> ([a], [b])
forall b d a. (b -> d) -> (a, b) -> (a, d)
mapSnd (b
b b -> [b] -> [b]
forall a. a -> [a] -> [a]
:) (([a], [b]) -> ([a], [b])) -> ([a], [b]) -> ([a], [b])
forall a b. (a -> b) -> a -> b
$ [a] -> ([a], [b])
loop [a]
as

-- | Like 'filter', but additionally return the last partition
--   of the list where the predicate is @False@ everywhere.
--   O(n).
filterAndRest :: (a -> Bool) -> [a] -> ([a], Suffix a)
filterAndRest :: (a -> Bool) -> [a] -> ([a], [a])
filterAndRest p :: a -> Bool
p = (a -> Maybe a) -> [a] -> ([a], [a])
forall a b. (a -> Maybe b) -> [a] -> ([b], [a])
mapMaybeAndRest ((a -> Maybe a) -> [a] -> ([a], [a]))
-> (a -> Maybe a) -> [a] -> ([a], [a])
forall a b. (a -> b) -> a -> b
$ \ a :: a
a -> if a -> Bool
p a
a then a -> Maybe a
forall a. a -> Maybe a
Just a
a else Maybe a
forall a. Maybe a
Nothing

-- | Like 'mapMaybe', but additionally return the last partition
--   of the list where the function always returns @Nothing@.
--   O(n).
mapMaybeAndRest :: (a -> Maybe b) -> [a] -> ([b], Suffix a)
mapMaybeAndRest :: (a -> Maybe b) -> [a] -> ([b], [a])
mapMaybeAndRest f :: a -> Maybe b
f = [a] -> [a] -> ([b], [a])
loop [] where
  loop :: [a] -> [a] -> ([b], [a])
loop acc :: [a]
acc = \case
    []                   -> ([], [a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc)
    x :: a
x:xs :: [a]
xs | Just y :: b
y <- a -> Maybe b
f a
x -> ([b] -> [b]) -> ([b], [a]) -> ([b], [a])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (b
yb -> [b] -> [b]
forall a. a -> [a] -> [a]
:) (([b], [a]) -> ([b], [a])) -> ([b], [a]) -> ([b], [a])
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> ([b], [a])
loop [] [a]
xs
         | Bool
otherwise     -> [a] -> [a] -> ([b], [a])
loop (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
acc) [a]
xs

-- | Sublist relation.
isSublistOf :: Eq a => [a] -> [a] -> Bool
isSublistOf :: [a] -> [a] -> Bool
isSublistOf = [a] -> [a] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
List.isSubsequenceOf

-- | All ways of removing one element from a list.
--   O(n²).
holes :: [a] -> [(a, [a])]
holes :: [a] -> [(a, [a])]
holes []     = []
holes (x :: a
x:xs :: [a]
xs) = (a
x, [a]
xs) (a, [a]) -> [(a, [a])] -> [(a, [a])]
forall a. a -> [a] -> [a]
: ((a, [a]) -> (a, [a])) -> [(a, [a])] -> [(a, [a])]
forall a b. (a -> b) -> [a] -> [b]
map (([a] -> [a]) -> (a, [a]) -> (a, [a])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:)) ([a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
holes [a]
xs)

---------------------------------------------------------------------------
-- * Prefix and suffix
---------------------------------------------------------------------------

-- ** Prefix

-- | Compute the common prefix of two lists.
--   O(min n m).
commonPrefix :: Eq a => [a] -> [a] -> Prefix a
commonPrefix :: [a] -> [a] -> [a]
commonPrefix [] _ = []
commonPrefix _ [] = []
commonPrefix (x :: a
x:xs :: [a]
xs) (y :: a
y:ys :: [a]
ys)
  | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y    = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. Eq a => [a] -> [a] -> [a]
commonPrefix [a]
xs [a]
ys
  | Bool
otherwise = []

-- | Drops from both lists simultaneously until one list is empty.
--   O(min n m).
dropCommon :: [a] -> [b] -> (Suffix a, Suffix b)
dropCommon :: [a] -> [b] -> ([a], [b])
dropCommon (x :: a
x : xs :: [a]
xs) (y :: b
y : ys :: [b]
ys) = [a] -> [b] -> ([a], [b])
forall a b. [a] -> [b] -> ([a], [b])
dropCommon [a]
xs [b]
ys
dropCommon xs :: [a]
xs ys :: [b]
ys = ([a]
xs, [b]
ys)

-- | Check if a list has a given prefix. If so, return the list
--   minus the prefix.
--   O(length prefix).
stripPrefixBy :: (a -> a -> Bool) -> Prefix a -> [a] -> Maybe (Suffix a)
stripPrefixBy :: (a -> a -> Bool) -> Prefix a -> Prefix a -> Maybe (Prefix a)
stripPrefixBy eq :: a -> a -> Bool
eq = Prefix a -> Prefix a -> Maybe (Prefix a)
loop
  where
  loop :: Prefix a -> Prefix a -> Maybe (Prefix a)
loop []    rest :: Prefix a
rest = Prefix a -> Maybe (Prefix a)
forall a. a -> Maybe a
Just Prefix a
rest
  loop (_:_) []   = Maybe (Prefix a)
forall a. Maybe a
Nothing
  loop (p :: a
p:pat :: Prefix a
pat) (r :: a
r:rest :: Prefix a
rest)
    | a -> a -> Bool
eq a
p a
r    = Prefix a -> Prefix a -> Maybe (Prefix a)
loop Prefix a
pat Prefix a
rest
    | Bool
otherwise = Maybe (Prefix a)
forall a. Maybe a
Nothing

-- ** Suffix

-- | Compute the common suffix of two lists.
--   O(n + m).
commonSuffix :: Eq a => [a] -> [a] -> Suffix a
commonSuffix :: [a] -> [a] -> [a]
commonSuffix xs :: [a]
xs ys :: [a]
ys = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ ([a] -> [a] -> [a]
forall a. Eq a => [a] -> [a] -> [a]
commonPrefix ([a] -> [a] -> [a]) -> ([a] -> [a]) -> [a] -> [a] -> [a]
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` [a] -> [a]
forall a. [a] -> [a]
reverse) [a]
xs [a]
ys

-- | @stripSuffix suf xs = Just pre@ iff @xs = pre ++ suf@.
-- O(n).
stripSuffix :: Eq a => Suffix a -> [a] -> Maybe (Prefix a)
stripSuffix :: Suffix a -> Suffix a -> Maybe (Suffix a)
stripSuffix [] = Suffix a -> Maybe (Suffix a)
forall a. a -> Maybe a
Just
stripSuffix s :: Suffix a
s  = Suffix a -> Suffix a -> Maybe (Suffix a)
forall a.
Eq a =>
ReversedSuffix a -> ReversedSuffix a -> Maybe (ReversedSuffix a)
stripReversedSuffix (Suffix a -> Suffix a
forall a. [a] -> [a]
reverse Suffix a
s)

type ReversedSuffix a = [a]

-- | @stripReversedSuffix rsuf xs = Just pre@ iff @xs = pre ++ reverse suf@.
--   O(n).
stripReversedSuffix :: forall a. Eq a => ReversedSuffix a -> [a] -> Maybe (Prefix a)
stripReversedSuffix :: ReversedSuffix a -> ReversedSuffix a -> Maybe (ReversedSuffix a)
stripReversedSuffix rs :: ReversedSuffix a
rs = StrSufSt a -> Maybe (ReversedSuffix a)
final (StrSufSt a -> Maybe (ReversedSuffix a))
-> (ReversedSuffix a -> StrSufSt a)
-> ReversedSuffix a
-> Maybe (ReversedSuffix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> StrSufSt a -> StrSufSt a)
-> StrSufSt a -> ReversedSuffix a -> StrSufSt a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> StrSufSt a -> StrSufSt a
step (ReversedSuffix a -> StrSufSt a
forall a. ReversedSuffix a -> StrSufSt a
SSSStrip ReversedSuffix a
rs)
  where
  -- Step of the automaton (reading input from right to left).
  step :: a -> StrSufSt a -> StrSufSt a
  step :: a -> StrSufSt a -> StrSufSt a
step x :: a
x = \case
    SSSMismatch   -> StrSufSt a
forall a. StrSufSt a
SSSMismatch
    SSSResult xs :: ReversedSuffix a
xs  -> ReversedSuffix a -> StrSufSt a
forall a. ReversedSuffix a -> StrSufSt a
SSSResult (a
xa -> ReversedSuffix a -> ReversedSuffix a
forall a. a -> [a] -> [a]
:ReversedSuffix a
xs)
    SSSStrip []   -> ReversedSuffix a -> StrSufSt a
forall a. ReversedSuffix a -> StrSufSt a
SSSResult [a
x]
    SSSStrip (y :: a
y:ys :: ReversedSuffix a
ys)
      | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y    -> ReversedSuffix a -> StrSufSt a
forall a. ReversedSuffix a -> StrSufSt a
SSSStrip ReversedSuffix a
ys
      | Bool
otherwise -> StrSufSt a
forall a. StrSufSt a
SSSMismatch

  -- Output of the automaton.
  final :: StrSufSt a -> Maybe (Prefix a)
  final :: StrSufSt a -> Maybe (ReversedSuffix a)
final = \case
    SSSResult xs :: ReversedSuffix a
xs -> ReversedSuffix a -> Maybe (ReversedSuffix a)
forall a. a -> Maybe a
Just ReversedSuffix a
xs
    SSSStrip []  -> ReversedSuffix a -> Maybe (ReversedSuffix a)
forall a. a -> Maybe a
Just []
    _            -> Maybe (ReversedSuffix a)
forall a. Maybe a
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.

---------------------------------------------------------------------------
-- * Groups and chunks
---------------------------------------------------------------------------

-- | @'groupOn' f = 'groupBy' (('==') \`on\` f) '.' 'List.sortBy' ('compare' \`on\` f)@.
-- O(n log n).
groupOn :: Ord b => (a -> b) -> [a] -> [[a]]
groupOn :: (a -> b) -> [a] -> [[a]]
groupOn f :: a -> b
f = (a -> a -> Bool) -> [a] -> [[a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
List.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
`on` a -> b
f) ([a] -> [[a]]) -> ([a] -> [a]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (b -> b -> Ordering) -> (a -> b) -> a -> a -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` a -> b
f)

-- | A variant of 'List.groupBy' which applies the predicate to consecutive
-- pairs.
-- O(n).
groupBy' :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy' :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy' _ []           = []
groupBy' p :: a -> a -> Bool
p xxs :: [a]
xxs@(x :: a
x : xs :: [a]
xs) = a -> [(Bool, a)] -> [[a]]
forall a. a -> [(Bool, a)] -> [[a]]
grp a
x ([(Bool, a)] -> [[a]]) -> [(Bool, a)] -> [[a]]
forall a b. (a -> b) -> a -> b
$ (a -> a -> (Bool, a)) -> [a] -> [a] -> [(Bool, a)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\x :: a
x y :: a
y -> (a -> a -> Bool
p a
x a
y, a
y)) [a]
xxs [a]
xs
  where
  grp :: a -> [(Bool, a)] -> [[a]]
grp x :: a
x ys :: [(Bool, a)]
ys = (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: ((Bool, a) -> a) -> [(Bool, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Bool, a) -> a
forall a b. (a, b) -> b
snd [(Bool, a)]
xs) [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [[a]]
tail
    where (xs :: [(Bool, a)]
xs, rest :: [(Bool, a)]
rest) = ((Bool, a) -> Bool) -> [(Bool, a)] -> ([(Bool, a)], [(Bool, a)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Bool, a) -> Bool
forall a b. (a, b) -> a
fst [(Bool, a)]
ys
          tail :: [[a]]
tail = case [(Bool, a)]
rest of
                   []            -> []
                   ((_, z :: a
z) : zs :: [(Bool, a)]
zs) -> a -> [(Bool, a)] -> [[a]]
grp a
z [(Bool, a)]
zs

-- | Split a list into sublists. Generalisation of the prelude function
--   @words@.
--   O(n).
--
--   > words xs == wordsBy isSpace xs
wordsBy :: (a -> Bool) -> [a] -> [[a]]
wordsBy :: (a -> Bool) -> [a] -> [[a]]
wordsBy p :: a -> Bool
p xs :: [a]
xs = [a] -> [[a]]
yesP [a]
xs
    where
        yesP :: [a] -> [[a]]
yesP xs :: [a]
xs = [a] -> [[a]]
noP ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile a -> Bool
p [a]
xs)

        noP :: [a] -> [[a]]
noP []  = []
        noP xs :: [a]
xs  = [a]
ys [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
yesP [a]
zs
            where
                (ys :: [a]
ys,zs :: [a]
zs) = (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break a -> Bool
p [a]
xs

-- | Chop up a list in chunks of a given length.
-- O(n).
chop :: Int -> [a] -> [[a]]
chop :: Int -> [a] -> [[a]]
chop _ [] = []
chop n :: Int
n xs :: [a]
xs = [a]
ys [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chop Int
n [a]
zs
    where (ys :: [a]
ys,zs :: [a]
zs) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
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.
--   O(n).
--
--    > intercalate [x] (chopWhen (== x) xs) == xs
chopWhen :: (a -> Bool) -> [a] -> [[a]]
chopWhen :: (a -> Bool) -> [a] -> [[a]]
chopWhen p :: a -> Bool
p [] = []
chopWhen p :: a -> Bool
p xs :: [a]
xs = [a] -> [[a]]
loop [a]
xs
  where
  -- Local function to avoid unnecessary pattern matching.
  loop :: [a] -> [[a]]
loop xs :: [a]
xs = case (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break a -> Bool
p [a]
xs of
    (w :: [a]
w, [])     -> [[a]
w]
    (w :: [a]
w, [_])    -> [[a]
w, []]
    (w :: [a]
w, _ : ys :: [a]
ys) -> [a]
w [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
loop [a]
ys  -- here we already know that ys /= []

---------------------------------------------------------------------------
-- * List as sets
---------------------------------------------------------------------------

-- | Check membership for the same list often.
--   Use partially applied to create membership predicate
--   @hasElem xs :: a -> Bool@.
--
--   * First time: @O(n log n)@ in the worst case.
--   * Subsequently: @O(log n)@.
--
--   Specification: @hasElem xs == (`elem` xs)@.
hasElem :: Ord a => [a] -> a -> Bool
hasElem :: [a] -> a -> Bool
hasElem xs :: [a]
xs = (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList [a]
xs)

-- | Check whether a list is sorted.
-- O(n).
--
-- Assumes that the 'Ord' instance implements a partial order.

sorted :: Ord a => [a] -> Bool
sorted :: [a] -> Bool
sorted [] = Bool
True
sorted xs :: [a]
xs = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (a -> a -> Bool) -> [a] -> [a] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=) [a]
xs ([a] -> [a]
forall a. [a] -> [a]
tail [a]
xs)

-- | Check whether all elements in a list are distinct from each other.
--   Assumes that the 'Eq' instance stands for an equivalence relation.
--
--   O(n²) in the worst case @distinct xs == True@.
distinct :: Eq a => [a] -> Bool
distinct :: [a] -> Bool
distinct []     = Bool
True
distinct (x :: a
x:xs :: [a]
xs) = a
x a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [a]
xs Bool -> Bool -> Bool
&& [a] -> Bool
forall a. Eq a => [a] -> Bool
distinct [a]
xs

-- | An optimised version of 'distinct'.
--   O(n log n).
--
--   Precondition: The list's length must fit in an 'Int'.

fastDistinct :: Ord a => [a] -> Bool
fastDistinct :: [a] -> Bool
fastDistinct xs :: [a]
xs = Set a -> Int
forall a. Set a -> Int
Set.size ([a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList [a]
xs) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs

-- | Returns an (arbitrary) representative for each list element
--   that occurs more than once.
--   O(n log n).
duplicates :: Ord a => [a] -> [a]
duplicates :: [a] -> [a]
duplicates = ([a] -> Maybe a) -> [[a]] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [a] -> Maybe a
forall a. [a] -> Maybe a
dup ([[a]] -> [a]) -> ([a] -> [[a]]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bag a -> [[a]]
forall a. Bag a -> [[a]]
Bag.groups (Bag a -> [[a]]) -> ([a] -> Bag a) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bag a
forall a. Ord a => [a] -> Bag a
Bag.fromList
  where
    dup :: [a] -> Maybe a
dup (a :: a
a : _ : _) = a -> Maybe a
forall a. a -> Maybe a
Just a
a
    dup _           = Maybe a
forall a. Maybe a
Nothing

-- | Remove the first representative for each list element.
--   Thus, returns all duplicate copies.
--   O(n log n).
--
--   @allDuplicates xs == sort $ xs \\ nub xs@.
allDuplicates :: Ord a => [a] -> [a]
allDuplicates :: [a] -> [a]
allDuplicates = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> ([a] -> [[a]]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop 1 ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse) ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bag a -> [[a]]
forall a. Bag a -> [[a]]
Bag.groups (Bag a -> [[a]]) -> ([a] -> Bag a) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bag a
forall a. Ord a => [a] -> Bag a
Bag.fromList
  -- The reverse is necessary to actually remove the *first* occurrence
  -- of each element.

-- | Efficient variant of 'nubBy' for lists, using a set to store already seen elements.
-- O(n log n)
--
-- Specification:
--
-- > nubOn f xs == 'nubBy' ((==) `'on'` f) xs.

nubOn :: Ord b => (a -> b) -> [a] -> [a]
nubOn :: (a -> b) -> [a] -> [a]
nubOn f :: a -> b
f = Set b -> [a] -> [a]
loop Set b
forall a. Set a
Set.empty
  where
  loop :: Set b -> [a] -> [a]
loop s :: Set b
s [] = []
  loop s :: Set b
s (a :: a
a:as :: [a]
as)
    | b
b b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set b
s = Set b -> [a] -> [a]
loop Set b
s [a]
as
    | Bool
otherwise        = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Set b -> [a] -> [a]
loop (b -> Set b -> Set b
forall a. Ord a => a -> Set a -> Set a
Set.insert b
b Set b
s) [a]
as
    where b :: b
b = a -> b
f a
a

-- Andreas, 2019-11-16
-- The implementation of nubOn using Set can be more than 1000-times
-- faster than the following old one using List.sort.
-- (Tested using criterion and -O on some lists of length 100.000.)

-- -- | Efficient variant of 'nubBy' for finite lists (using sorting).
-- -- O(n log n)
-- --
-- -- Specification:
-- --
-- -- > nubOn2 f xs == 'nubBy' ((==) `'on'` f) xs.
--
-- nubOn2 :: Ord b => (a -> b) -> [a] -> [a]
-- nubOn2 tag =
--     -- Throw away numbering
--   map snd
--     -- Restore original order
--   . List.sortBy (compare `on` fst)
--     -- Retain first entry of each @tag@ group
--   . map (snd . head)
--   . List.groupBy ((==) `on` fst)
--     -- Sort by tag (stable)
--   . List.sortBy (compare `on` fst)
--     -- Tag with @tag@ and sequential numbering
--   . map (\p@(_, x) -> (tag x, p))
--   . zip [1..]

-- | Efficient variant of 'nubBy' for finite lists.
-- O(n log n).
--
-- 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 id == Set.toAscList . Set.fromList
--
uniqOn :: Ord b => (a -> b) -> [a] -> [a]
uniqOn :: (a -> b) -> [a] -> [a]
uniqOn key :: a -> b
key = Map b a -> [a]
forall k a. Map k a -> [a]
Map.elems (Map b a -> [a]) -> ([a] -> Map b a) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(b, a)] -> Map b a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(b, a)] -> Map b a) -> ([a] -> [(b, a)]) -> [a] -> Map 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 :: a
a -> (a -> b
key a
a, a
a))

-- | Checks if all the elements in the list are equal. Assumes that
--   the 'Eq' instance stands for an equivalence relation.
--   O(n).
allEqual :: Eq a => [a] -> Bool
allEqual :: [a] -> Bool
allEqual []       = Bool
True
allEqual (x :: a
x : xs :: [a]
xs) = (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x) [a]
xs

---------------------------------------------------------------------------
-- * Zipping
---------------------------------------------------------------------------

-- | Requires both lists to have the same length.
--   O(n).
--
--   Otherwise, @Nothing@ is returned.

zipWith' :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
zipWith' :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
zipWith' f :: a -> b -> c
f = [a] -> [b] -> Maybe [c]
loop
  where
  loop :: [a] -> [b] -> Maybe [c]
loop []        []      = [c] -> Maybe [c]
forall a. a -> Maybe a
Just []
  loop (x :: a
x : xs :: [a]
xs) (y :: b
y : ys :: [b]
ys) = (a -> b -> c
f a
x b
y c -> [c] -> [c]
forall a. a -> [a] -> [a]
:) ([c] -> [c]) -> Maybe [c] -> Maybe [c]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [b] -> Maybe [c]
loop [a]
xs [b]
ys
  loop []       (_ : _)  = Maybe [c]
forall a. Maybe a
Nothing
  loop (_ : _)  []       = Maybe [c]
forall a. Maybe a
Nothing

-- | Like 'zipWith' but keep the rest of the second list as-is
--   (in case the second list is longer).
-- O(n).
--
-- @
--   zipWithKeepRest f as bs == zipWith f as bs ++ drop (length as) bs
-- @
zipWithKeepRest :: (a -> b -> b) -> [a] -> [b] -> [b]
zipWithKeepRest :: (a -> b -> b) -> [a] -> [b] -> [b]
zipWithKeepRest f :: a -> b -> b
f = [a] -> [b] -> [b]
loop
  where
  loop :: [a] -> [b] -> [b]
loop []       bs :: [b]
bs       = [b]
bs
  loop as :: [a]
as       []       = []
  loop (a :: a
a : as :: [a]
as) (b :: b
b : bs :: [b]
bs) = a -> b -> b
f a
a b
b b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [a] -> [b] -> [b]
loop [a]
as [b]
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


---------------------------------------------------------------------------
-- * Unzipping
---------------------------------------------------------------------------

unzipWith :: (a -> (b, c)) -> [a] -> ([b], [c])
unzipWith :: (a -> (b, c)) -> [a] -> ([b], [c])
unzipWith f :: a -> (b, c)
f = [(b, c)] -> ([b], [c])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(b, c)] -> ([b], [c])) -> ([a] -> [(b, c)]) -> [a] -> ([b], [c])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (b, c)) -> [a] -> [(b, c)]
forall a b. (a -> b) -> [a] -> [b]
map a -> (b, c)
f

---------------------------------------------------------------------------
-- * Edit distance
---------------------------------------------------------------------------

-- | Implemented using tree recursion, don't run me at home!
--   O(3^(min n m)).
editDistanceSpec :: Eq a => [a] -> [a] -> Int
editDistanceSpec :: [a] -> [a] -> Int
editDistanceSpec [] ys :: [a]
ys = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys
editDistanceSpec xs :: [a]
xs [] = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
editDistanceSpec (x :: a
x : xs :: [a]
xs) (y :: a
y : ys :: [a]
ys)
  | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y    = [a] -> [a] -> Int
forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec [a]
xs [a]
ys
  | Bool
otherwise = 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ [a] -> [a] -> Int
forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs) [a]
ys
                            , [a] -> [a] -> Int
forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec [a]
xs (a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys)
                            , [a] -> [a] -> Int
forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec [a]
xs [a]
ys ]

-- | Implemented using dynamic programming and @Data.Array@.
--   O(n*m).
editDistance :: forall a. Eq a => [a] -> [a] -> Int
editDistance :: [a] -> [a] -> Int
editDistance xs :: [a]
xs ys :: [a]
ys = Int -> Int -> Int
editD 0 0
  where
  editD :: Int -> Int -> Int
editD i :: Int
i j :: Int
j = Array (Int, Int) Int
tbl Array (Int, Int) Int -> (Int, Int) -> Int
forall i e. Ix i => Array i e -> i -> e
Array.! (Int
i, Int
j)
  -- Tabulate editD' in immutable boxed array (content computed lazily).
  tbl :: Array (Int,Int) Int
  tbl :: Array (Int, Int) Int
tbl = ((Int, Int), (Int, Int))
-> [((Int, Int), Int)] -> Array (Int, Int) Int
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((0,0), (Int
n,Int
m)) [ ((Int
i, Int
j), Int -> Int -> Int
editD' Int
i Int
j) | Int
i <- [0..Int
n], Int
j <- [0..Int
m] ]
  editD' :: Int -> Int -> Int
editD' i :: Int
i j :: Int
j =
    case (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
n, Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
j Int
m) of
      -- Interior
      (LT, LT)
        | Array Int a
xsA Array Int a -> Int -> a
forall i e. Ix i => Array i e -> i -> e
Array.! Int
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Array Int a
ysA Array Int a -> Int -> a
forall i e. Ix i => Array i e -> i -> e
Array.! Int
j
                    -> Int -> Int -> Int
editD Int
i' Int
j'
        | Bool
otherwise -> 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ Int -> Int -> Int
editD Int
i' Int
j, Int -> Int -> Int
editD Int
i Int
j', Int -> Int -> Int
editD Int
i' Int
j' ]
      -- Border: one list is empty
      (EQ, LT)      ->  Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j
      (LT, EQ)      ->  Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i
      -- Corner (EQ, EQ): both lists are empty
      _             -> 0
      -- GT cases are impossible.
    where (i' :: Int
i',j' :: Int
j') = (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1, Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
  n :: Int
n   = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
  m :: Int
m   = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys
  xsA, ysA :: Array Int a
  xsA :: Array Int a
xsA = (Int, Int) -> [a] -> Array Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (0,Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) [a]
xs
  ysA :: Array Int a
ysA = (Int, Int) -> [a] -> Array Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (0,Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) [a]
ys