{-# LANGUAGE CPP #-}

{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

-}



-- | Set-like operations on lists
--
-- Avoid using them as much as possible
module GHC.Data.List.SetOps (
        unionLists, unionListsOrd, minusList,

        -- Association lists
        Assoc, assoc, assocMaybe, assocUsing, assocDefault, assocDefaultUsing,

        -- Duplicate handling
        hasNoDups, removeDups, removeDupsOn, nubOrdBy, findDupsEq,
        equivClasses,

        -- Indexing
        getNth,

        -- Membership
        isIn, isn'tIn,
   ) where

import GHC.Prelude

import GHC.Utils.Outputable
import GHC.Utils.Panic
import GHC.Utils.Misc

import qualified Data.List as L
import qualified Data.List.NonEmpty as NE
import Data.List.NonEmpty (NonEmpty(..))
import Data.Ord (comparing)
import qualified Data.Set as S

getNth :: Outputable a => [a] -> Int -> a
getNth :: forall a. Outputable a => [a] -> Int -> a
getNth [a]
xs Int
n = Bool -> SDoc -> a -> a
forall a. HasCallStack => Bool -> SDoc -> a -> a
assertPpr ([a]
xs [a] -> Int -> Bool
forall a. [a] -> Int -> Bool
`lengthExceeds` Int
n) (Int -> SDoc
forall a. Outputable a => a -> SDoc
ppr Int
n SDoc -> SDoc -> SDoc
forall doc. IsDoc doc => doc -> doc -> doc
$$ [a] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [a]
xs) (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$
             [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
n

{-
************************************************************************
*                                                                      *
        Treating lists as sets
        Assumes the lists contain no duplicates, but are unordered
*                                                                      *
************************************************************************
-}



-- | Combines the two lists while keeping their order, placing the first argument
-- first in the result.
--
-- Uses a set internally to record duplicates. This makes it slightly slower for
-- very small lists but avoids quadratic behaviour for large lists.
unionListsOrd :: (HasDebugCallStack, Outputable a, Ord a) => [a] -> [a] -> [a]
unionListsOrd :: forall a.
(HasDebugCallStack, Outputable a, Ord a) =>
[a] -> [a] -> [a]
unionListsOrd [a]
xs [a]
ys
  -- Since both arguments don't have internal duplicates we can just take all of xs
  -- and every element of ys that's not already in xs.
  = let set_ys :: Set a
set_ys = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
ys
    in ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (\a
e -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member a
e Set a
set_ys) [a]
xs) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
ys

-- | Assumes that the arguments contain no duplicates
unionLists :: (HasDebugCallStack, Outputable a, Eq a) => [a] -> [a] -> [a]
-- We special case some reasonable common patterns.
unionLists :: forall a.
(HasDebugCallStack, Outputable a, Eq a) =>
[a] -> [a] -> [a]
unionLists [a]
xs [] = [a]
xs
unionLists [] [a]
ys = [a]
ys
unionLists [a
x] [a]
ys
  | String -> a -> [a] -> Bool
forall a. Eq a => String -> a -> [a] -> Bool
isIn String
"unionLists" a
x [a]
ys = [a]
ys
  | Bool
otherwise = a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys
unionLists [a]
xs [a
y]
  | String -> a -> [a] -> Bool
forall a. Eq a => String -> a -> [a] -> Bool
isIn String
"unionLists" a
y [a]
xs = [a]
xs
  | Bool
otherwise = a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs
unionLists [a]
xs [a]
ys
  = Bool -> String -> SDoc -> [a] -> [a]
forall a. HasCallStack => Bool -> String -> SDoc -> a -> a
warnPprTrace ([a] -> Int -> Bool
forall a. [a] -> Int -> Bool
lengthExceeds [a]
xs Int
100 Bool -> Bool -> Bool
|| [a] -> Int -> Bool
forall a. [a] -> Int -> Bool
lengthExceeds [a]
ys Int
100) String
"unionLists" ([a] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [a]
xs SDoc -> SDoc -> SDoc
forall doc. IsDoc doc => doc -> doc -> doc
$$ [a] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [a]
ys) ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$
    [a
x | a
x <- [a]
xs, String -> a -> [a] -> Bool
forall a. Eq a => String -> a -> [a] -> Bool
isn'tIn String
"unionLists" a
x [a]
ys] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
ys

-- | Calculate the set difference of two lists. This is
-- /O((m + n) log n)/, where we subtract a list of /n/ elements
-- from a list of /m/ elements.
--
-- Extremely short cases are handled specially:
-- When /m/ or /n/ is 0, this takes /O(1)/ time. When /m/ is 1,
-- it takes /O(n)/ time.
minusList :: Ord a => [a] -> [a] -> [a]
-- There's no point building a set to perform just one lookup, so we handle
-- extremely short lists specially. It might actually be better to use
-- an O(m*n) algorithm when m is a little longer (perhaps up to 4 or even 5).
-- The tipping point will be somewhere in the area of where /m/ and /log n/
-- become comparable, but we probably don't want to work too hard on this.
minusList :: forall a. Ord a => [a] -> [a] -> [a]
minusList [] [a]
_ = []
minusList xs :: [a]
xs@[a
x] [a]
ys
  | a
x a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a]
ys = []
  | Bool
otherwise = [a]
xs
-- Using an empty set or a singleton would also be silly, so let's not.
minusList [a]
xs [] = [a]
xs
minusList [a]
xs [a
y] = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
y) [a]
xs
-- When each list has at least two elements, we build a set from the
-- second argument, allowing us to filter the first argument fairly
-- efficiently.
minusList [a]
xs [a]
ys = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set a
yss) [a]
xs
  where
    yss :: Set a
yss = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
ys

{-
************************************************************************
*                                                                      *
\subsection[Utils-assoc]{Association lists}
*                                                                      *
************************************************************************

Inefficient finite maps based on association lists and equality.
-}

-- | A finite mapping based on equality and association lists.
type Assoc a b = [(a,b)]

assoc             :: (Eq a) => String -> Assoc a b -> a -> b
assocDefault      :: (Eq a) => b -> Assoc a b -> a -> b
assocUsing        :: (a -> a -> Bool) -> String -> Assoc a b -> a -> b
-- | Lookup key, fail gracefully using Nothing if not found.
assocMaybe        :: (Eq a) => Assoc a b -> a -> Maybe b
assocDefaultUsing :: (a -> a -> Bool) -> b -> Assoc a b -> a -> b

assocDefaultUsing :: forall a b. (a -> a -> Bool) -> b -> Assoc a b -> a -> b
assocDefaultUsing a -> a -> Bool
_  b
deflt []             a
_   = b
deflt
assocDefaultUsing a -> a -> Bool
eq b
deflt ((a
k,b
v) : [(a, b)]
rest) a
key
  | a
k a -> a -> Bool
`eq` a
key = b
v
  | Bool
otherwise  = (a -> a -> Bool) -> b -> [(a, b)] -> a -> b
forall a b. (a -> a -> Bool) -> b -> Assoc a b -> a -> b
assocDefaultUsing a -> a -> Bool
eq b
deflt [(a, b)]
rest a
key

assoc :: forall a b. Eq a => String -> Assoc a b -> a -> b
assoc String
crash_msg         Assoc a b
list a
key = (a -> a -> Bool) -> b -> Assoc a b -> a -> b
forall a b. (a -> a -> Bool) -> b -> Assoc a b -> a -> b
assocDefaultUsing a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (String -> b
forall a. HasCallStack => String -> a
panic (String
"Failed in assoc: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
crash_msg)) Assoc a b
list a
key
assocDefault :: forall a b. Eq a => b -> Assoc a b -> a -> b
assocDefault b
deflt      Assoc a b
list a
key = (a -> a -> Bool) -> b -> Assoc a b -> a -> b
forall a b. (a -> a -> Bool) -> b -> Assoc a b -> a -> b
assocDefaultUsing a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) b
deflt Assoc a b
list a
key
assocUsing :: forall a b. (a -> a -> Bool) -> String -> Assoc a b -> a -> b
assocUsing a -> a -> Bool
eq String
crash_msg Assoc a b
list a
key = (a -> a -> Bool) -> b -> Assoc a b -> a -> b
forall a b. (a -> a -> Bool) -> b -> Assoc a b -> a -> b
assocDefaultUsing a -> a -> Bool
eq (String -> b
forall a. HasCallStack => String -> a
panic (String
"Failed in assoc: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
crash_msg)) Assoc a b
list a
key

assocMaybe :: forall a b. Eq a => Assoc a b -> a -> Maybe b
assocMaybe Assoc a b
alist a
key
  = Assoc a b -> Maybe b
lookup Assoc a b
alist
  where
    lookup :: Assoc a b -> Maybe b
lookup []             = Maybe b
forall a. Maybe a
Nothing
    lookup ((a
tv,b
ty):Assoc a b
rest) = if a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
tv then b -> Maybe b
forall a. a -> Maybe a
Just b
ty else Assoc a b -> Maybe b
lookup Assoc a b
rest

{-
************************************************************************
*                                                                      *
\subsection[Utils-dups]{Duplicate-handling}
*                                                                      *
************************************************************************
-}

hasNoDups :: (Eq a) => [a] -> Bool

hasNoDups :: forall a. Eq a => [a] -> Bool
hasNoDups [a]
xs = [a] -> [a] -> Bool
f [] [a]
xs
  where
    f :: [a] -> [a] -> Bool
f [a]
_           []     = Bool
True
    f [a]
seen_so_far (a
x:[a]
xs) = if a
x a -> [a] -> Bool
`is_elem` [a]
seen_so_far
                           then Bool
False
                           else [a] -> [a] -> Bool
f (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
seen_so_far) [a]
xs

    is_elem :: a -> [a] -> Bool
is_elem = String -> a -> [a] -> Bool
forall a. Eq a => String -> a -> [a] -> Bool
isIn String
"hasNoDups"

equivClasses :: (a -> a -> Ordering) -- Comparison
             -> [a]
             -> [NonEmpty a]

equivClasses :: forall a. (a -> a -> Ordering) -> [a] -> [NonEmpty a]
equivClasses a -> a -> Ordering
_   []      = []
equivClasses a -> a -> Ordering
_   [a
stuff] = [a
stuff a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| []]
equivClasses a -> a -> Ordering
cmp [a]
items   = (a -> a -> Bool) -> [a] -> [NonEmpty a]
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> Bool) -> f a -> [NonEmpty a]
NE.groupBy a -> a -> Bool
eq ((a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy a -> a -> Ordering
cmp [a]
items)
  where
    eq :: a -> a -> Bool
eq a
a a
b = case a -> a -> Ordering
cmp a
a a
b of { Ordering
EQ -> Bool
True; Ordering
_ -> Bool
False }

-- | Remove the duplicates from a list using the provided
-- comparison function. Might change the order of elements.
--
-- Returns the list without duplicates, and accumulates
-- all the duplicates in the second component of its result.
removeDups :: (a -> a -> Ordering) -- Comparison function
           -> [a]
           -> ([a],          -- List with no duplicates
               [NonEmpty a]) -- List of duplicate groups.  One representative
                             -- from each group appears in the first result

removeDups :: forall a. (a -> a -> Ordering) -> [a] -> ([a], [NonEmpty a])
removeDups a -> a -> Ordering
_   []  = ([], [])
removeDups a -> a -> Ordering
_   [a
x] = ([a
x],[])
removeDups a -> a -> Ordering
cmp [a]
xs
  = case ([NonEmpty a] -> NonEmpty a -> ([NonEmpty a], a))
-> [NonEmpty a] -> [NonEmpty a] -> ([NonEmpty a], [a])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumR [NonEmpty a] -> NonEmpty a -> ([NonEmpty a], a)
forall a. [NonEmpty a] -> NonEmpty a -> ([NonEmpty a], a)
collect_dups [] ((a -> a -> Ordering) -> [a] -> [NonEmpty a]
forall a. (a -> a -> Ordering) -> [a] -> [NonEmpty a]
equivClasses a -> a -> Ordering
cmp [a]
xs) of { ([NonEmpty a]
dups, [a]
xs') ->
    ([a]
xs', [NonEmpty a]
dups) }
  where
    collect_dups :: [NonEmpty a] -> NonEmpty a -> ([NonEmpty a], a)
    collect_dups :: forall a. [NonEmpty a] -> NonEmpty a -> ([NonEmpty a], a)
collect_dups [NonEmpty a]
dups_so_far (a
x :| [])     = ([NonEmpty a]
dups_so_far,      a
x)
    collect_dups [NonEmpty a]
dups_so_far dups :: NonEmpty a
dups@(a
x :| [a]
_) = (NonEmpty a
dupsNonEmpty a -> [NonEmpty a] -> [NonEmpty a]
forall a. a -> [a] -> [a]
:[NonEmpty a]
dups_so_far, a
x)

removeDupsOn :: Ord b => (a -> b) -> [a] -> ([a], [NonEmpty a])
removeDupsOn :: forall b a. Ord b => (a -> b) -> [a] -> ([a], [NonEmpty a])
removeDupsOn a -> b
f [a]
x = (a -> a -> Ordering) -> [a] -> ([a], [NonEmpty a])
forall a. (a -> a -> Ordering) -> [a] -> ([a], [NonEmpty a])
removeDups ((a -> b) -> a -> a -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing a -> b
f) [a]
x

-- | Remove the duplicates from a list using the provided
-- comparison function.
nubOrdBy :: (a -> a -> Ordering) -> [a] -> [a]
nubOrdBy :: forall a. (a -> a -> Ordering) -> [a] -> [a]
nubOrdBy a -> a -> Ordering
cmp [a]
xs = ([a], [NonEmpty a]) -> [a]
forall a b. (a, b) -> a
fst ((a -> a -> Ordering) -> [a] -> ([a], [NonEmpty a])
forall a. (a -> a -> Ordering) -> [a] -> ([a], [NonEmpty a])
removeDups a -> a -> Ordering
cmp [a]
xs)

findDupsEq :: (a->a->Bool) -> [a] -> [NonEmpty a]
findDupsEq :: forall a. (a -> a -> Bool) -> [a] -> [NonEmpty a]
findDupsEq a -> a -> Bool
_  [] = []
findDupsEq a -> a -> Bool
eq (a
x:[a]
xs) | [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null [a]
eq_xs  = (a -> a -> Bool) -> [a] -> [NonEmpty a]
forall a. (a -> a -> Bool) -> [a] -> [NonEmpty a]
findDupsEq a -> a -> Bool
eq [a]
xs
                     | Bool
otherwise     = (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
eq_xs) NonEmpty a -> [NonEmpty a] -> [NonEmpty a]
forall a. a -> [a] -> [a]
: (a -> a -> Bool) -> [a] -> [NonEmpty a]
forall a. (a -> a -> Bool) -> [a] -> [NonEmpty a]
findDupsEq a -> a -> Bool
eq [a]
neq_xs
    where ([a]
eq_xs, [a]
neq_xs) = (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (a -> a -> Bool
eq a
x) [a]
xs

-- Debugging/specialising versions of \tr{elem} and \tr{notElem}

# if !defined(DEBUG)
isIn, isn'tIn :: Eq a => String -> a -> [a] -> Bool
isIn :: forall a. Eq a => String -> a -> [a] -> Bool
isIn    String
_msg a
x [a]
ys = a
x a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a]
ys
isn'tIn :: forall a. Eq a => String -> a -> [a] -> Bool
isn'tIn String
_msg a
x [a]
ys = a
x a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [a]
ys

# else /* DEBUG */
isIn, isn'tIn :: (HasDebugCallStack, Eq a) => String -> a -> [a] -> Bool
isIn msg x ys
  = elem100 0 x ys
  where
    elem100 :: Eq a => Int -> a -> [a] -> Bool
    elem100 _ _ [] = False
    elem100 i x (y:ys)
      | i > 100 = warnPprTrace True ("Over-long elem in " ++ msg) empty (x `elem` (y:ys))
      | otherwise = x == y || elem100 (i + 1) x ys

isn'tIn msg x ys
  = notElem100 0 x ys
  where
    notElem100 :: Eq a => Int -> a -> [a] -> Bool
    notElem100 _ _ [] =  True
    notElem100 i x (y:ys)
      | i > 100 = warnPprTrace True ("Over-long notElem in " ++ msg) empty (x `notElem` (y:ys))
      | otherwise = x /= y && notElem100 (i + 1) x ys
# endif /* DEBUG */