-- | Utility functions for lists.

module Agda.Utils.List where

import Control.Monad (filterM)

import Data.Array (Array, array, listArray)
import qualified Data.Array as Array
import Data.Bifunctor
import Data.Function
import Data.Hashable
import qualified Data.List as List
import qualified Data.List.NonEmpty as List1
import Data.List.NonEmpty (pattern (:|), (<|))
import Data.Maybe
import qualified Data.Map as Map
import qualified Data.HashMap.Strict as HMap
import qualified Data.Set as Set

import qualified Agda.Utils.Bag as Bag
import Agda.Utils.CallStack.Base
import Agda.Utils.Function (applyWhen)
import Agda.Utils.Functor  ((<.>))
import Agda.Utils.Tuple

import {-# SOURCE #-} Agda.Utils.List1 (List1)

import Agda.Utils.Impossible

---------------------------------------------------------------------------
-- * 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 :: forall a. [a] -> a -> [a]
snoc [a]
xs a
x = [a]
xs 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 :: forall a b. [a] -> b -> (a -> [a] -> b) -> b
caseList [a]
xs b
n a -> [a] -> b
c = 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 :: forall (m :: * -> *) a b.
Monad m =>
m [a] -> m b -> (a -> [a] -> m b) -> m b
caseListM m [a]
mxs m b
n a -> [a] -> m b
c = forall b a. b -> (a -> [a] -> b) -> [a] -> b
listCase m b
n a -> [a] -> m b
c 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 :: forall b a. b -> (a -> [a] -> b) -> [a] -> b
listCase b
n a -> [a] -> b
c []     = b
n
listCase b
n a -> [a] -> b
c (a
x:[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 :: forall a. a -> [a] -> a
headWithDefault a
def = forall a. a -> Maybe a -> a
fromMaybe a
def forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> Maybe a
listToMaybe

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

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

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

-- | Last element (safe).  Returns a default list on empty lists.
--   O(n).
lastWithDefault :: a -> [a] -> a
lastWithDefault :: forall a. a -> [a] -> a
lastWithDefault = forall a. a -> [a] -> a
last1

-- | Last element of non-empty list (safe).
--   O(n).
--   @last1 a as = last (a : as)@
last1 :: a -> [a] -> a
last1 :: forall a. a -> [a] -> a
last1 a
a = \case
  [] -> a
a
  a
b:[a]
bs -> forall a. a -> [a] -> a
last1 a
b [a]
bs

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

-- | Opposite of cons @(:)@, safe.
--   O(1).
uncons :: [a] -> Maybe (a, [a])
uncons :: forall a. [a] -> Maybe (a, [a])
uncons []     = forall a. Maybe a
Nothing
uncons (a
x:[a]
xs) = 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 :: forall a. Maybe a -> [a] -> [a]
mcons Maybe a
ma [a]
as = forall b a. b -> (a -> b) -> Maybe a -> b
maybe [a]
as (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 :: forall a. [a] -> Maybe ([a], a)
initLast []     = forall a. Maybe a
Nothing
initLast (a
a:[a]
as) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. a -> [a] -> ([a], a)
initLast1 a
a [a]
as

-- | 'init' and 'last' of non-empty list, safe.
--   O(n).
--   @initLast1 a as = (init (a:as), last (a:as)@
initLast1 :: a -> [a] -> ([a], a)
initLast1 :: forall a. a -> [a] -> ([a], a)
initLast1 a
a = \case
  []   -> ([], a
a)
  a
b:[a]
bs -> forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a
aforall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ forall a. a -> [a] -> ([a], a)
initLast1 a
b [a]
bs

-- | 'init' of non-empty list, safe.
--   O(n).
--   @init1 a as = init (a:as)@
init1 :: a -> [a] -> [a]
init1 :: forall a. a -> [a] -> [a]
init1 a
a = \case
  []   -> []
  a
b:[a]
bs -> a
a forall a. a -> [a] -> [a]
: forall a. a -> [a] -> [a]
init1 a
b [a]
bs

-- | @init@, safe.
--   O(n).
initMaybe :: [a] -> Maybe [a]
initMaybe :: forall a. [a] -> Maybe [a]
initMaybe = \case
  []   -> forall a. Maybe a
Nothing
  a
a:[a]
as -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. a -> [a] -> [a]
init1 a
a [a]
as

-- | @init@, safe.
--   O(n).
initWithDefault :: [a] -> [a] -> [a]
initWithDefault :: forall a. [a] -> [a] -> [a]
initWithDefault [a]
as []     = [a]
as
initWithDefault [a]
_  (a
a:[a]
as) = forall a. a -> [a] -> [a]
init1 a
a [a]
as

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

-- | Lookup function (safe).
--   O(min n index).
(!!!) :: [a] -> Int -> Maybe a
[a]
xs !!! :: forall a. [a] -> Int -> Maybe a
!!! (!Int
i)
  | Int
i forall a. Ord a => a -> a -> Bool
< Int
0     = forall a. Maybe a
Nothing
  | Bool
otherwise = forall {t} {a}. (Eq t, Num t) => [a] -> t -> Maybe a
index [a]
xs Int
i
  where
  index :: [a] -> t -> Maybe a
index []       !t
i = forall a. Maybe a
Nothing
  index (a
x : [a]
xs) t
0  = forall a. a -> Maybe a
Just a
x
  index (a
x : [a]
xs) t
i  = [a] -> t -> Maybe a
index [a]
xs (t
i forall a. Num a => a -> a -> a
- t
1)

-- | A variant of 'Prelude.!!' that might provide more informative
-- error messages if the index is out of bounds.
--
-- Precondition: The index should not be out of bounds.

(!!) :: HasCallStack => [a] -> Int -> a
[a]
xs !! :: forall a. HasCallStack => [a] -> Int -> a
!! Int
i = case [a]
xs forall a. [a] -> Int -> Maybe a
!!! Int
i of
  Just a
x  -> a
x
  Maybe a
Nothing -> forall a. HasCallStack => a
__IMPOSSIBLE__

-- | 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 :: forall a. a -> [a] -> Int -> a
indexWithDefault a
a []       Int
_ = a
a
indexWithDefault a
a (a
x : [a]
_)  Int
0 = a
x
indexWithDefault a
a (a
_ : [a]
xs) Int
n = forall a. a -> [a] -> Int -> a
indexWithDefault a
a [a]
xs (Int
n forall a. Num a => a -> a -> a
- Int
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 :: forall a. (a -> Bool) -> [a] -> Maybe (a, Int)
findWithIndex a -> Bool
p [a]
as = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
List.find (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) (forall a b. [a] -> [b] -> [(a, b)]
zip [a]
as [Int
0..])

-- | A generalised variant of 'elemIndex'.
-- O(n).
genericElemIndex :: (Eq a, Integral i) => a -> [a] -> Maybe i
genericElemIndex :: forall a i. (Eq a, Integral i) => a -> [a] -> Maybe i
genericElemIndex a
x [a]
xs =
  forall a. [a] -> Maybe a
listToMaybe forall a b. (a -> b) -> a -> b
$
  forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$
  forall a. (a -> Bool) -> [a] -> [a]
filter forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$
  forall a b. [a] -> [b] -> [(a, b)]
zip [i
0..] forall a b. (a -> b) -> a -> b
$
  forall a b. (a -> b) -> [a] -> [b]
map (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 :: forall a. Integral a => a -> [a]
downFrom a
n | a
n forall a. Ord a => a -> a -> Bool
<= a
0     = []
           | Bool
otherwise = let n' :: a
n' = a
nforall a. Num a => a -> a -> a
-a
1 in a
n' forall 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 :: forall a. (a -> a) -> [a] -> [a]
updateHead a -> a
_ []       = []
updateHead a -> a
f (a
a : [a]
as) = a -> a
f 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 :: forall a. (a -> a) -> [a] -> [a]
updateLast a -> a
_ [] = []
updateLast a -> a
f (a
a : [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
f a
a]
  loop a
a (a
b : [a]
bs) = 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 :: forall a. Int -> (a -> a) -> [a] -> [a]
updateAt Int
_ a -> a
_ [] = []
updateAt Int
0 a -> a
f (a
a : [a]
as) = a -> a
f a
a forall a. a -> [a] -> [a]
: [a]
as
updateAt Int
n a -> a
f (a
a : [a]
as) = a
a forall a. a -> [a] -> [a]
: forall a. Int -> (a -> a) -> [a] -> [a]
updateAt (Int
nforall a. Num a => a -> a -> a
-Int
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 :: forall n a. Integral n => n -> [a] -> Maybe ([a], [a])
splitExactlyAt n
0 [a]
xs       = forall (m :: * -> *) a. Monad m => a -> m a
return ([], [a]
xs)
splitExactlyAt n
n []       = forall a. Maybe a
Nothing
splitExactlyAt n
n (a
x : [a]
xs) = forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (a
x forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall n a. Integral n => n -> [a] -> Maybe ([a], [a])
splitExactlyAt (n
nforall a. Num a => a -> a -> a
-n
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 :: forall a. Int -> [a] -> [a]
dropEnd Int
n = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 a
x (Int
n, [a]
xs) = (Int
nforall a. Num a => a -> a -> a
-Int
1, forall a. Bool -> (a -> a) -> a -> a
applyWhen (Int
n forall a. Ord a => a -> a -> Bool
<= Int
0) (a
xforall 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 :: forall a. (a -> Bool) -> [a] -> ([a], [a])
spanEnd a -> Bool
p = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 a
x (Bool
b', ([a]
xs, [a]
ys)) = (Bool
b, if Bool
b then ([a]
xs, a
xforall a. a -> [a] -> [a]
:[a]
ys) else (a
xforall a. a -> [a] -> [a]
:[a]
xs, [a]
ys))
    where b :: Bool
b = Bool
b' Bool -> Bool -> Bool
&& a -> Bool
p a
x

-- | Breaks a list just /after/ an element satisfying the predicate is
--   found.
--
--   >>> breakAfter1 even 1 [3,5,2,4,7,8]
--   ([1,3,5,2],[4,7,8])

breakAfter1 :: (a -> Bool) -> a -> [a] -> (List1 a, [a])
breakAfter1 :: forall a. (a -> Bool) -> a -> [a] -> (List1 a, [a])
breakAfter1 a -> Bool
p = a -> [a] -> (NonEmpty a, [a])
loop
  where
  loop :: a -> [a] -> (NonEmpty a, [a])
loop a
x = \case
    xs :: [a]
xs@[]         -> (a
x forall a. a -> [a] -> NonEmpty a
:| [], [a]
xs)
    xs :: [a]
xs@(a
y : [a]
ys)
      | a -> Bool
p a
x       -> (a
x forall a. a -> [a] -> NonEmpty a
:| [], [a]
xs)
      | Bool
otherwise -> let (NonEmpty a
vs, [a]
ws) = a -> [a] -> (NonEmpty a, [a])
loop a
y [a]
ys in (a
x forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty a
vs, [a]
ws)

-- | Breaks a list just /after/ an element satisfying the predicate is
--   found.
--
--   >>> breakAfter even [1,3,5,2,4,7,8]
--   ([1,3,5,2],[4,7,8])

breakAfter :: (a -> Bool) -> [a] -> ([a], [a])
breakAfter :: forall a. (a -> Bool) -> [a] -> ([a], [a])
breakAfter a -> Bool
p = \case
  []   -> ([], [])
  a
x:[a]
xs -> forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a. NonEmpty a -> [a]
List1.toList forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> a -> [a] -> (List1 a, [a])
breakAfter1 a -> Bool
p a
x [a]
xs

-- | 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 :: forall a b. (a -> Maybe b) -> [a] -> Prefix b
takeWhileJust a -> Maybe b
p = [a] -> [b]
loop
  where
    loop :: [a] -> [b]
loop (a
a : [a]
as) | Just b
b <- a -> Maybe b
p a
a = b
b forall a. a -> [a] -> [a]
: [a] -> [b]
loop [a]
as
    loop [a]
_ = []

-- | A generalized version of @span@.
--   @O(length . fst . spanJust f)@.
spanJust :: (a -> Maybe b) -> [a] -> (Prefix b, Suffix a)
spanJust :: forall a b. (a -> Maybe b) -> [a] -> (Prefix b, [a])
spanJust a -> Maybe b
p = [a] -> ([b], [a])
loop
  where
    loop :: [a] -> ([b], [a])
loop (a
a : [a]
as) | Just b
b <- a -> Maybe b
p a
a = forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (b
b forall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ [a] -> ([b], [a])
loop [a]
as
    loop [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 :: forall a b. (a -> Maybe b) -> [a] -> ([a], [b])
partitionMaybe a -> Maybe b
f = [a] -> ([a], [b])
loop
  where
    loop :: [a] -> ([a], [b])
loop []       = ([], [])
    loop (a
a : [a]
as) = case a -> Maybe b
f a
a of
      Maybe b
Nothing -> forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (a
a forall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ [a] -> ([a], [b])
loop [a]
as
      Just b
b  -> forall b d a. (b -> d) -> (a, b) -> (a, d)
mapSnd (b
b forall a. a -> [a] -> [a]
:) 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 :: forall a. (a -> Bool) -> [a] -> ([a], [a])
filterAndRest a -> Bool
p = forall a b. (a -> Maybe b) -> [a] -> (Prefix b, [a])
mapMaybeAndRest forall a b. (a -> b) -> a -> b
$ \ a
a -> if a -> Bool
p a
a then forall a. a -> Maybe a
Just a
a else 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 :: forall a b. (a -> Maybe b) -> [a] -> (Prefix b, [a])
mapMaybeAndRest a -> Maybe b
f = [a] -> [a] -> ([b], [a])
loop [] where
  loop :: [a] -> [a] -> ([b], [a])
loop [a]
acc = \case
    []                   -> ([], forall a. [a] -> [a]
reverse [a]
acc)
    a
x:[a]
xs | Just b
y <- a -> Maybe b
f a
x -> forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (b
yforall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> ([b], [a])
loop [] [a]
xs
         | Bool
otherwise     -> [a] -> [a] -> ([b], [a])
loop (a
xforall a. a -> [a] -> [a]
:[a]
acc) [a]
xs

-- | Sublist relation.
isSublistOf :: Eq a => [a] -> [a] -> Bool
isSublistOf :: forall a. Eq a => [a] -> [a] -> Bool
isSublistOf = 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 :: forall a. [a] -> [(a, [a])]
holes []     = []
holes (a
x:[a]
xs) = (a
x, [a]
xs) forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (a
xforall 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 :: forall a. Eq a => [a] -> [a] -> [a]
commonPrefix [] [a]
_ = []
commonPrefix [a]
_ [] = []
commonPrefix (a
x:[a]
xs) (a
y:[a]
ys)
  | a
x forall a. Eq a => a -> a -> Bool
== a
y    = a
x forall 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 :: forall a b. [a] -> [b] -> ([a], [b])
dropCommon (a
x : [a]
xs) (b
y : [b]
ys) = forall a b. [a] -> [b] -> ([a], [b])
dropCommon [a]
xs [b]
ys
dropCommon [a]
xs [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 :: forall a.
(a -> a -> Bool) -> Prefix a -> Prefix a -> Maybe (Prefix a)
stripPrefixBy a -> a -> Bool
eq = [a] -> [a] -> Maybe [a]
loop
  where
  loop :: [a] -> [a] -> Maybe [a]
loop []    [a]
rest = forall a. a -> Maybe a
Just [a]
rest
  loop (a
_:[a]
_) []   = forall a. Maybe a
Nothing
  loop (a
p:[a]
pat) (a
r:[a]
rest)
    | a -> a -> Bool
eq a
p a
r    = [a] -> [a] -> Maybe [a]
loop [a]
pat [a]
rest
    | Bool
otherwise = forall a. Maybe a
Nothing

-- ** Suffix

-- | Compute the common suffix of two lists.
--   O(n + m).
commonSuffix :: Eq a => [a] -> [a] -> Suffix a
commonSuffix :: forall a. Eq a => [a] -> [a] -> [a]
commonSuffix [a]
xs [a]
ys = forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ (forall a. Eq a => [a] -> [a] -> [a]
commonPrefix forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` 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 :: forall a. Eq a => Suffix a -> Suffix a -> Maybe (Suffix a)
stripSuffix [] = forall a. a -> Maybe a
Just
stripSuffix [a]
s  = forall a. Eq a => Suffix a -> Suffix a -> Maybe (Suffix a)
stripReversedSuffix (forall a. [a] -> [a]
reverse [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 :: forall a. Eq a => Suffix a -> Suffix a -> Maybe (Suffix a)
stripReversedSuffix ReversedSuffix a
rs = StrSufSt a -> Maybe (ReversedSuffix a)
final forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> StrSufSt a -> StrSufSt a
step (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 a
x = \case
    StrSufSt a
SSSMismatch   -> forall a. StrSufSt a
SSSMismatch
    SSSResult ReversedSuffix a
xs  -> forall a. ReversedSuffix a -> StrSufSt a
SSSResult (a
xforall a. a -> [a] -> [a]
:ReversedSuffix a
xs)
    SSSStrip []   -> forall a. ReversedSuffix a -> StrSufSt a
SSSResult [a
x]
    SSSStrip (a
y:ReversedSuffix a
ys)
      | a
x forall a. Eq a => a -> a -> Bool
== a
y    -> forall a. ReversedSuffix a -> StrSufSt a
SSSStrip ReversedSuffix a
ys
      | Bool
otherwise -> forall a. StrSufSt a
SSSMismatch

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

-- | Returns a list with one boolean for each non-empty suffix of the
-- list, starting with the longest suffix (the entire list). Each
-- boolean is 'True' exactly when every element in the corresponding
-- suffix satisfies the predicate.
--
-- An example:
-- @
--  'suffixesSatisfying' 'Data.Char.isLower' "AbCde" =
--  [False, False, False, True, True]
-- @
--
-- For total predicates @p@ and finite and total lists @xs@ the
-- following holds:
-- @
--  'suffixesSatisfying' p xs = 'map' ('all' p) ('List.init' ('List.tails' xs))
-- @
suffixesSatisfying :: (a -> Bool) -> [a] -> [Bool]
suffixesSatisfying :: forall a. (a -> Bool) -> [a] -> [Bool]
suffixesSatisfying a -> Bool
p =
  forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x (Bool
b, [Bool]
bs) -> let !b' :: Bool
b' = a -> Bool
p a
x Bool -> Bool -> Bool
&& Bool
b in (Bool
b', Bool
b' forall a. a -> [a] -> [a]
: [Bool]
bs))
        (Bool
True, [])

-- ** Finding overlap

-- | Find the longest suffix of the first string @xs@
--   that is a prefix of the second string @ys@.
--   So, basically, find the overlap where the strings can be glued together.
--   Returns the index where the overlap starts and the length of the overlap.
--   The length of the overlap plus the index is the length of the first string.
--   Note that in the worst case, the empty overlap @(length xs,0)@ is returned.
--
--   Worst-case time complexity is quadratic: @O(min(n,m)²)@
--   where @n = length xs@ and @m = length ys@.
--
--   There might be asymptotically better implementations following
--   Knuth-Morris-Pratt (KMP), but for rather short lists this is good enough.
--
findOverlap :: forall a. Eq a => [a] -> [a] -> (Int, Int)
findOverlap :: forall a. Eq a => [a] -> [a] -> (Int, Int)
findOverlap [a]
xs [a]
ys =
  forall a. a -> [a] -> a
headWithDefault forall a. HasCallStack => a
__IMPOSSIBLE__ forall a b. (a -> b) -> a -> b
$ forall a b. (a -> Maybe b) -> [a] -> Prefix b
mapMaybe (Int, [a]) -> Maybe (Int, Int)
maybePrefix forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (forall a. [a] -> [[a]]
List.tails [a]
xs)
  where
  maybePrefix :: (Int, [a]) -> Maybe (Int, Int)
  maybePrefix :: (Int, [a]) -> Maybe (Int, Int)
maybePrefix (Int
k, [a]
xs')
    | [a]
xs' forall a. Eq a => [a] -> [a] -> Bool
`List.isPrefixOf` [a]
ys = forall a. a -> Maybe a
Just (Int
k, forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs')
    | Bool
otherwise                = forall a. Maybe a
Nothing

---------------------------------------------------------------------------
-- * 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 :: forall b a. Ord b => (a -> b) -> [a] -> [[a]]
groupOn a -> b
f = forall a. (a -> a -> Bool) -> [a] -> [[a]]
List.groupBy (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` a -> b
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (forall a. Ord a => a -> a -> Ordering
compare 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).
-- DEPRECATED in favor of 'Agda.Utils.List1.groupBy''.
groupBy' :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy' :: forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy' a -> a -> Bool
_ []           = []
groupBy' a -> a -> Bool
p xxs :: [a]
xxs@(a
x : [a]
xs) = forall {a}. a -> [(Bool, a)] -> [[a]]
grp a
x forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
x a
y -> (a -> a -> Bool
p a
x a
y, a
y)) [a]
xxs [a]
xs
  where
  grp :: a -> [(Bool, a)] -> [[a]]
grp a
x [(Bool, a)]
ys = (a
x forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Bool, a)]
xs) forall a. a -> [a] -> [a]
: [[a]]
tail
    where ([(Bool, a)]
xs, [(Bool, a)]
rest) = forall a. (a -> Bool) -> [a] -> ([a], [a])
span forall a b. (a, b) -> a
fst [(Bool, a)]
ys
          tail :: [[a]]
tail = case [(Bool, a)]
rest of
                   []            -> []
                   ((Bool
_, a
z) : [(Bool, a)]
zs) -> a -> [(Bool, a)] -> [[a]]
grp a
z [(Bool, a)]
zs

-- | Chop up a list in chunks of a given length.
-- O(n).
chop :: Int -> [a] -> [[a]]
chop :: forall a. Int -> [a] -> [[a]]
chop Int
_ [] = []
chop Int
n [a]
xs = [a]
ys forall a. a -> [a] -> [a]
: forall a. Int -> [a] -> [[a]]
chop Int
n [a]
zs
    where ([a]
ys,[a]
zs) = 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 :: forall a. (a -> Bool) -> [a] -> [[a]]
chopWhen :: forall a. (a -> Bool) -> [a] -> [[a]]
chopWhen a -> Bool
p []     = []
chopWhen a -> Bool
p (a
x:[a]
xs) = List1 a -> [[a]]
loop (a
x forall a. a -> [a] -> NonEmpty a
:| [a]
xs)
  where
  -- Local function to avoid unnecessary pattern matching.
  loop :: List1 a -> [[a]]
  loop :: List1 a -> [[a]]
loop List1 a
xs = case forall a. (a -> Bool) -> NonEmpty a -> ([a], [a])
List1.break a -> Bool
p List1 a
xs of
    ([a]
w, []        ) -> [[a]
w]
    ([a]
w, a
_ : []    ) -> [[a]
w, []]
    ([a]
w, a
_ : a
y : [a]
ys) -> [a]
w forall a. a -> [a] -> [a]
: List1 a -> [[a]]
loop (a
y forall a. a -> [a] -> NonEmpty a
:| [a]
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 :: forall a. Ord a => [a] -> a -> Bool
hasElem [a]
xs = (forall a. Ord a => a -> Set a -> Bool
`Set.member` 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 :: forall a. Ord a => [a] -> Bool
sorted [] = Bool
True
sorted [a]
xs = forall (t :: * -> *). Foldable t => t Bool -> Bool
and forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Ord a => a -> a -> Bool
(<=) [a]
xs (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 :: forall a. Eq a => [a] -> Bool
distinct []     = Bool
True
distinct (a
x:[a]
xs) = a
x forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [a]
xs Bool -> Bool -> 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 :: forall a. Ord a => [a] -> Bool
fastDistinct [a]
xs = forall a. Set a -> Int
Set.size (forall a. Ord a => [a] -> Set a
Set.fromList [a]
xs) forall a. Eq a => a -> a -> Bool
== 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 :: forall a. Ord a => [a] -> [a]
duplicates = forall a b. (a -> Maybe b) -> [a] -> Prefix b
mapMaybe forall a. [a] -> Maybe a
dup forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> [[a]]
Bag.groups forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> Bag a
Bag.fromList
  where
    dup :: [a] -> Maybe a
dup (a
a : a
_ : [a]
_) = forall a. a -> Maybe a
Just a
a
    dup [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 :: forall a. Ord a => [a] -> [a]
allDuplicates = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. Int -> [a] -> [a]
drop Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> [[a]]
Bag.groups forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> Bag a
Bag.fromList
  -- The reverse is necessary to actually remove the *first* occurrence
  -- of each element.

-- | Partition a list into first and later occurrences of elements
--   (modulo some quotient given by a representation function).
--
--  Time: O(n log n).
--
--  Specification:
--
--  > nubAndDuplicatesOn f xs = (ys, xs List.\\ ys)
--  >   where ys = nubOn f xs

nubAndDuplicatesOn :: Ord b => (a -> b) -> [a] -> ([a], [a])
nubAndDuplicatesOn :: forall b a. Ord b => (a -> b) -> [a] -> ([a], [a])
nubAndDuplicatesOn a -> b
f = Set b -> [a] -> ([a], [a])
loop forall a. Set a
Set.empty
  where
  loop :: Set b -> [a] -> ([a], [a])
loop Set b
s [] = ([], [])
  loop Set b
s (a
a:[a]
as)
    | b
b forall a. Ord a => a -> Set a -> Bool
`Set.member` Set b
s = forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (a
aforall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ Set b -> [a] -> ([a], [a])
loop Set b
s [a]
as
    | Bool
otherwise        = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first  (a
aforall a. a -> [a] -> [a]
:) forall a b. (a -> b) -> a -> b
$ Set b -> [a] -> ([a], [a])
loop (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

-- | 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 :: forall b a. Ord b => (a -> b) -> [a] -> [a]
nubOn a -> b
f = Set b -> [a] -> [a]
loop forall a. Set a
Set.empty
  where
  loop :: Set b -> [a] -> [a]
loop Set b
s [] = []
  loop Set b
s (a
a:[a]
as)
    | b
b 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 forall a. a -> [a] -> [a]
: Set b -> [a] -> [a]
loop (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

-- | A variant of 'nubOn' that is parametrised by a function that is
-- used to select which element from a group of equal elements that is
-- returned. The returned elements keep the order that they had in the
-- input list.
--
-- Precondition: The length of the input list must be at most
-- @'maxBound' :: 'Int'@.

nubFavouriteOn
  :: forall a b c. (Ord b, Eq c, Hashable c)
  => (a -> b)
     -- ^ The values returned by this function are used to determine
     -- which element from a group of equal elements that is returned:
     -- the smallest one is chosen (and if two elements are equally
     -- small, then the first one is chosen).
  -> (a -> c)
     -- ^ Two elements are treated as equal if this function returns
     -- the same value for both elements.
  -> [a] -> [a]
nubFavouriteOn :: forall a b c.
(Ord b, Eq c, Hashable c) =>
(a -> b) -> (a -> c) -> [a] -> [a]
nubFavouriteOn a -> b
fav a -> c
f = Int -> HashMap c ((b, Int), a) -> [a] -> [a]
go Int
0 forall k v. HashMap k v
HMap.empty
  where
  go :: Int -> HMap.HashMap c ((b, Int), a) -> [a] -> [a]
  go :: Int -> HashMap c ((b, Int), a) -> [a] -> [a]
go !Int
pos !HashMap c ((b, Int), a)
acc (a
x : [a]
xs) =
    Int -> HashMap c ((b, Int), a) -> [a] -> [a]
go (Int
1 forall a. Num a => a -> a -> a
+ Int
pos)
       (forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HMap.insertWith
          (\((b, Int), a)
new ((b, Int), a)
old -> if forall a b. (a, b) -> a
fst ((b, Int), a)
new forall a. Ord a => a -> a -> Bool
< forall a b. (a, b) -> a
fst ((b, Int), a)
old then ((b, Int), a)
new else ((b, Int), a)
old)
          (a -> c
f a
x) ((a -> b
fav a
x, Int
pos), a
x) HashMap c ((b, Int), a)
acc)
       [a]
xs
  go Int
_ HashMap c ((b, Int), a)
acc [] =
    forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (forall a. Ord a => a -> a -> Ordering
compare forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$
    forall k v. HashMap k v -> [v]
HMap.elems HashMap c ((b, Int), a)
acc

-- -- | 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).
--
-- > uniqOn f == 'List.sortBy' (compare `'on'` f) . 'nubBy' ((==) `'on'` f)
--
-- If there are several elements with the same @f@-representative,
-- the first of these is kept.
--
uniqOn :: Ord b => (a -> b) -> [a] -> [a]
uniqOn :: forall b a. Ord b => (a -> b) -> [a] -> [a]
uniqOn a -> b
key = forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith (\ a
_ -> forall a. a -> a
id) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\ 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 :: forall a. Eq a => [a] -> Bool
allEqual []       = Bool
True
allEqual (a
x : [a]
xs) = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== a
x) [a]
xs

-- | Non-efficient, monadic 'nub'.
-- O(n²).
nubM :: Monad m => (a -> a -> m Bool) -> [a] -> m [a]
nubM :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> m Bool) -> [a] -> m [a]
nubM a -> a -> m Bool
eq = [a] -> m [a]
loop where
  loop :: [a] -> m [a]
loop []     = forall (m :: * -> *) a. Monad m => a -> m a
return []
  loop (a
a:[a]
as) = (a
a forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do [a] -> m [a]
loop forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM (Bool -> Bool
not forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> (a -> m b) -> a -> m c
<.> a -> a -> m Bool
eq a
a) [a]
as

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

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

zipWith' :: (a -> b -> c) -> [a] -> [b] -> Maybe [c]
zipWith' :: forall a b c. (a -> b -> c) -> [a] -> [b] -> Maybe [c]
zipWith' a -> b -> c
f = [a] -> [b] -> Maybe [c]
loop
  where
  loop :: [a] -> [b] -> Maybe [c]
loop []        []      = forall a. a -> Maybe a
Just []
  loop (a
x : [a]
xs) (b
y : [b]
ys) = (a -> b -> c
f a
x b
y forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [b] -> Maybe [c]
loop [a]
xs [b]
ys
  loop []       (b
_ : [b]
_)  = forall a. Maybe a
Nothing
  loop (a
_ : [a]
_)  []       = 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 :: forall a b. (a -> b -> b) -> [a] -> [b] -> [b]
zipWithKeepRest a -> b -> b
f = [a] -> [b] -> [b]
loop
  where
  loop :: [a] -> [b] -> [b]
loop []       [b]
bs       = [b]
bs
  loop [a]
as       []       = []
  loop (a
a : [a]
as) (b
b : [b]
bs) = a -> b -> b
f a
a 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 :: forall a b c. (a -> (b, c)) -> [a] -> ([b], [c])
unzipWith a -> (b, c)
f = forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> 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 :: forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec [] [a]
ys = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys
editDistanceSpec [a]
xs [] = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
editDistanceSpec (a
x : [a]
xs) (a
y : [a]
ys)
  | a
x forall a. Eq a => a -> a -> Bool
== a
y    = forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec [a]
xs [a]
ys
  | Bool
otherwise = Int
1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec (a
x forall a. a -> [a] -> [a]
: [a]
xs) [a]
ys
                            , forall a. Eq a => [a] -> [a] -> Int
editDistanceSpec [a]
xs (a
y forall a. a -> [a] -> [a]
: [a]
ys)
                            , 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 :: forall a. Eq a => [a] -> [a] -> Int
editDistance [a]
xs [a]
ys = Int -> Int -> Int
editD Int
0 Int
0
  where
  editD :: Int -> Int -> Int
editD Int
i Int
j = Array (Int, Int) Int
tbl 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 = forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ((Int
0,Int
0), (Int
n,Int
m)) [ ((Int
i, Int
j), Int -> Int -> Int
editD' Int
i Int
j) | Int
i <- [Int
0..Int
n], Int
j <- [Int
0..Int
m] ]
  editD' :: Int -> Int -> Int
editD' Int
i Int
j =
    case (forall a. Ord a => a -> a -> Ordering
compare Int
i Int
n, forall a. Ord a => a -> a -> Ordering
compare Int
j Int
m) of
      -- Interior
      (Ordering
LT, Ordering
LT)
        | Array Int a
xsA forall i e. Ix i => Array i e -> i -> e
Array.! Int
i forall a. Eq a => a -> a -> Bool
== Array Int a
ysA forall i e. Ix i => Array i e -> i -> e
Array.! Int
j
                    -> Int -> Int -> Int
editD Int
i' Int
j'
        | Bool
otherwise -> Int
1 forall a. Num a => a -> a -> a
+ 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
      (Ordering
EQ, Ordering
LT)      ->  Int
m forall a. Num a => a -> a -> a
- Int
j
      (Ordering
LT, Ordering
EQ)      ->  Int
n forall a. Num a => a -> a -> a
- Int
i
      -- Corner (EQ, EQ): both lists are empty
      (Ordering, Ordering)
_             -> Int
0
      -- GT cases are impossible.
    where (Int
i',Int
j') = (Int
iforall a. Num a => a -> a -> a
+Int
1, Int
jforall a. Num a => a -> a -> a
+Int
1)
  n :: Int
n   = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
  m :: Int
m   = forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ys
  xsA, ysA :: Array Int a
  xsA :: Array Int a
xsA = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0,Int
nforall a. Num a => a -> a -> a
-Int
1) [a]
xs
  ysA :: Array Int a
ysA = forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0,Int
mforall a. Num a => a -> a -> a
-Int
1) [a]
ys


mergeStrictlyOrderedBy :: (a -> a -> Bool) -> [a] -> [a] -> Maybe [a]
mergeStrictlyOrderedBy :: forall a.
(a -> a -> Bool) -> Prefix a -> Prefix a -> Maybe (Prefix a)
mergeStrictlyOrderedBy a -> a -> Bool
(<) = [a] -> [a] -> Maybe [a]
loop where
  loop :: [a] -> [a] -> Maybe [a]
loop [] [a]
ys = forall a. a -> Maybe a
Just [a]
ys
  loop [a]
xs [] = forall a. a -> Maybe a
Just [a]
xs
  loop (a
x:[a]
xs) (a
y:[a]
ys)
    | a
x a -> a -> Bool
< a
y = (a
xforall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [a] -> Maybe [a]
loop [a]
xs (a
yforall a. a -> [a] -> [a]
:[a]
ys)
    | a
y a -> a -> Bool
< a
x = (a
yforall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [a] -> Maybe [a]
loop (a
xforall a. a -> [a] -> [a]
:[a]
xs) [a]
ys
    | Bool
otherwise = forall a. Maybe a
Nothing