{-# LANGUAGE FlexibleContexts, ScopedTypeVariables, ViewPatterns #-}

-- | Functions on association lists that make use of an 'Eq' constraint
-- on the type of the keys.

module Data.AssocList.ListLike.Eq
    (

    -- * Related modules
    -- $relatedModules

    -- * Lookup
      lookupFirst
    , lookupAll

    -- * Removal
    , removeFirst
    , removeAll

    -- * Mapping
    -- $mapping
    , mapFirst
    , mapAll

    -- * Alteration
    -- $alteration
    , alterFirst
    , alterAll

    -- * Grouping
    , partition
    , break
    , breakPartition

    -- * Operators
    , (!)
    , (!?)

    ) where

import Data.AssocList.ListLike.Concept

-- base
import Control.Exception (throw)
import Prelude (Eq (..), Maybe (..), maybe, error, otherwise, (<$>))

-- ListLike
import Data.ListLike (cons, uncons)
import qualified Data.ListLike as LL

-- $setup
-- >>> import Prelude (fmap, map, negate, take)

-- $relatedModules
-- Some other modules that are a lot like this one:
--
-- * "Data.AssocList.ListLike.Equivalence" - Functions on association
--   lists that involve 'Equivalence's on the keys
-- * "Data.AssocList.ListLike.Predicate" - Functions on association
--   lists that involve 'Predicate's on the keys

-- | Obtain the first value associated with a particular key.
--
-- >>> [('A',1), ('B',2), ('B',3), ('C',4)] ! 'B'
-- 2
--
-- This function is to be used only when the key must be known to
-- be present in the mapping. If @x@ is not mapped by any entry in
-- @AssocList@ @l@, then @l '!' x@ throws 'MissingAssocListKey'.
-- The exclamation mark is intended as a reminder of this danger.
--
-- >>> [('A', 1), ('B', 2), ('B', 3), ('C', 4)] ! 'D'
-- *** Exception: MissingAssocListKey
--
-- There is a related operator called '!?' which maps the
-- missing-key condition to 'Nothing' instead.

(!) :: forall l a b. (AssocList l a b, Eq a)
    => l -> a -> b
(l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing) ! :: l -> a -> b
! a
key               = MissingAssocListKey -> b
forall a e. Exception e => e -> a
throw MissingAssocListKey
MissingAssocListKey
(l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Just ((a
x, b
y), l
xys)) ! a
key
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x                      = b
y
        | Bool
otherwise                     = l
xys l -> a -> b
forall l a b. (AssocList l a b, Eq a) => l -> a -> b
! a
key

-- | Obtain the first value associated with a particular key, if such
-- a mapping is present.
--
-- >>> [('A',1), ('B',2), ('B',3), ('C',4)] !? 'B'
-- Just 2
--
-- The result is 'Nothing' if the key is not mapped by any entry in
-- the list.
--
-- >>> [('A',1), ('B',2), ('B',3), ('C',4)] !? 'D'
-- Nothing
--
-- This function is the same as 'lookupFirst' but for the order of
-- its arguments.

(!?) :: forall l a b. (AssocList l a b, Eq a)
    => l -> a -> Maybe b
l
l !? :: l -> a -> Maybe b
!? a
key = a -> l -> Maybe b
forall l a b. (AssocList l a b, Eq a) => a -> l -> Maybe b
lookupFirst a
key l
l

-- | Obtain the first value associated with a particular key, if such
-- a mapping is present.
--
-- >>> lookupFirst 'B' [('A',1), ('B',2), ('B',3), ('C',4)]
-- Just 2
--
-- The result is 'Nothing' if the key is not mapped by any entry in
-- the list.
--
-- >>> lookupFirst 'D' [('A',1), ('B',2), ('B',3), ('C',4)]
-- Nothing
--
-- This function is the same as '!?' but for the order of its
-- arguments.

lookupFirst :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> Maybe b
lookupFirst :: a -> l -> Maybe b
lookupFirst a
_key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing)    =  Maybe b
forall a. Maybe a
Nothing
lookupFirst a
key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Just ((a
x, b
y), l
xys))
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x                      =  b -> Maybe b
forall a. a -> Maybe a
Just b
y
        | Bool
otherwise                     =  a -> l -> Maybe b
forall l a b. (AssocList l a b, Eq a) => a -> l -> Maybe b
lookupFirst a
key l
xys

-- | Obtain all values associated with a particular key, in the
-- order in which the mappings appear in the list.
--
-- >>> lookupAll 'B' [('A',1), ('B',2), ('B',3), ('C',4), ('B',3)]
-- [2,3,3]

lookupAll :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> [b]
lookupAll :: a -> l -> [b]
lookupAll a
_key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing)      =  []
lookupAll a
key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Just ((a
x, b
y), l
xys))
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x                      =  b
y b -> [b] -> [b]
forall a. a -> [a] -> [a]
: a -> l -> [b]
forall l a b. (AssocList l a b, Eq a) => a -> l -> [b]
lookupAll a
key l
xys
        | Bool
otherwise                     =      a -> l -> [b]
forall l a b. (AssocList l a b, Eq a) => a -> l -> [b]
lookupAll a
key l
xys

-- | Produce a modified version of the association list in which the
-- first occurrence of a particular key has been removed.
--
-- >>> removeFirst 'B' [('A',1), ('B',2), ('B',3), ('C',4)]
-- [('A',1),('B',3),('C',4)]
--
-- If the key is not present in the mapping, then the original list
-- is returned.
--
-- >>> removeFirst 'C' [('A',1), ('B',2), ('B',3)]
-- [('A',1),('B',2),('B',3)]

removeFirst :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> l
removeFirst :: a -> l -> l
removeFirst a
_key l :: l
l@(l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing)  =  l
l
removeFirst a
key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Just (xy :: (a, b)
xy@(a
x, b
y), l
xys))
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x                      =  l
xys
        | Bool
otherwise                     =  (a, b) -> l -> l
forall full item. ListLike full item => item -> full -> full
cons (a, b)
xy (a -> l -> l
forall l a b. (AssocList l a b, Eq a) => a -> l -> l
removeFirst a
key l
xys)

-- | Produce a modified version of the association list in which all
-- occurrences of a particular key have been removed.
--
-- >>> removeAll 'B' [('A',1), ('B',2), ('B',3), ('C',4)]
-- [('A',1),('C',4)]
--
-- If the key is not present in the mapping, then the original list
-- is returned.
--
-- >>> removeAll 'C' [('A',1), ('B',2), ('B',3)]
-- [('A',1),('B',2),('B',3)]

removeAll :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> l
removeAll :: a -> l -> l
removeAll a
_key l :: l
l@(l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing)    =  l
l
removeAll a
key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Just (xy :: (a, b)
xy@(a
x, b
y), l
xys))
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x                      =       a -> l -> l
forall l a b. (AssocList l a b, Eq a) => a -> l -> l
removeAll a
key l
xys
        | Bool
otherwise                     =  (a, b) -> l -> l
forall full item. ListLike full item => item -> full -> full
cons (a, b)
xy (a -> l -> l
forall l a b. (AssocList l a b, Eq a) => a -> l -> l
removeAll a
key l
xys)

-- | Produces a tuple of two results:
--
-- 1. All values associated with a particular key
-- 2. All of the other key-value pairs
--
-- @'partition' x l = ('lookupAll' x l, 'removeAll' x l)@
--
-- >>> partition 'B' [('A',1), ('B',2), ('B',3), ('C',4), ('B',3)]
-- ([2,3,3],[('A',1),('C',4)])

partition :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> ([b], l)
partition :: a -> l -> ([b], l)
partition a
_key l :: l
l@(l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing)    = ([], l
l)
partition a
key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Just (xy :: (a, b)
xy@(a
x, b
y), l
xys))
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x                      = (b
y b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
yes ,         l
no)
        | Bool
otherwise                     = (    [b]
yes , (a, b) -> l -> l
forall full item. ListLike full item => item -> full -> full
cons (a, b)
xy l
no)
    where
        ([b]
yes, l
no) = a -> l -> ([b], l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> ([b], l)
partition a
key l
xys

-- | Produces a tuple of two results:
--
-- 1. The longest prefix of the association list that does /not/ contain
--    a particular key
-- 2. The remainder of the list
--
-- >>> break 'B' [('A',1), ('B',2), ('B',3), ('C',4)]
-- ([('A',1)],[('B',2),('B',3),('C',4)])
--
-- If the first mapping in the list contains the given key, then the first
-- part of the resulting tuple is empty, and the second part of the result
-- is the entire list.
--
-- >>> break 'A' [('A',1), ('B',2), ('B',3), ('C',4)]
-- ([],[('A',1),('B',2),('B',3),('C',4)])
--
-- If the key is not present in the list, then the first part of the
-- resulting tuple is the entire list, and the second part of the result
-- is empty.
--
-- >>> break 'D' [('A',1), ('B',2), ('B',3), ('C',4)]
-- ([('A',1),('B',2),('B',3),('C',4)],[])

break :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> (l, l)
break :: a -> l -> (l, l)
break a
key = ((a, b) -> Bool) -> l -> (l, l)
forall full item.
ListLike full item =>
(item -> Bool) -> full -> (full, full)
LL.break (\(a
x, b
y) -> a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x)

-- | 'break' on a key, then 'partition' the remainder.
--
-- @'breakPartition' key l@ separates @l@ into three parts:
--
-- 1. The key-value pairs for which the key is /not/ @key@ that
--    occur in the list /before/ the first occurrence of @key@
--    (@fst ('break' key l)@)
-- 2. All values associated with @key@ (@'lookupAll' key l@)
-- 3. The key-value pairs for which the key is /not/ @key@ that
--    occur in the list /after/ the first occurrence of @key@
--    (@'removeAll' key (snd ('break' key l))@)
--
-- >>> breakPartition 'B' [('A',1),('B',2),('C',3),('B',4)]
-- ([('A',1)],[2,4],[('C',3)])
--
-- If the key is not present in the list, then the first part of the
-- result is the entire list, and the other parts are empty.
--
-- >>> breakPartition 'D' [('A',1),('B',2),('C',3),('B',4)]
-- ([('A',1),('B',2),('C',3),('B',4)],[],[])

breakPartition :: forall l a b. (AssocList l a b, Eq a)
    => a -> l -> (l, [b], l)
breakPartition :: a -> l -> (l, [b], l)
breakPartition a
key l
l =
    let
        (l
before, l
l') = a -> l -> (l, l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> (l, l)
break     a
key l
l
        ([b]
xs, l
after)  = a -> l -> ([b], l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> ([b], l)
partition a
key l
l'
    in
        (l
before, [b]
xs, l
after)

-- $mapping
-- The "map" functions modify values while preserving the structure of
-- the assocative list. The resulting list has the same size and order
-- as the original.

-- | At the position where a particular key first appears in the list,
-- apply a function to the corresponding value.
--
-- >>> mapFirst 'B' negate [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',-4),('C',2),('B',6)]
--
-- If the key does not appear in the list, then the original list is
-- returned without modification.
--
-- >>> mapFirst 'D' negate [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',4),('C',2),('B',6)]

mapFirst :: forall l a b. (AssocList l a b, Eq a)
    => a -> (b -> b) -> l -> l
mapFirst :: a -> (b -> b) -> l -> l
mapFirst a
key b -> b
f l
l =
    let
        (l
before, l
l') = a -> l -> (l, l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> (l, l)
break a
key l
l
    in
        l
before l -> l -> l
forall full item. ListLike full item => full -> full -> full
`LL.append`
        case (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons l
l') of
            Maybe ((a, b), l)
Nothing               ->  l
l'
            Just ((a
x, b
y), l
after)  ->  (a, b) -> l -> l
forall full item. ListLike full item => item -> full -> full
cons (a
x, b -> b
f b
y) l
after

-- | At each position where a particular key appears in the list,
-- apply a function to the corresponding value.
--
-- >>> mapAll 'B' negate [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',-4),('C',2),('B',-6)]
--
-- If the key does not appear in the list, then the original list is
-- returned without modification.
--
-- >>> mapAll 'D' negate [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',4),('C',2),('B',6)]

mapAll :: forall l a b. (AssocList l a b, Eq a)
    => a -> (b -> b) -> l -> l
mapAll :: a -> (b -> b) -> l -> l
mapAll a
key b -> b
f =
    ((a, b) -> (a, b)) -> l -> l
forall full item full' item'.
(ListLike full item, ListLike full' item') =>
(item -> item') -> full -> full'
LL.map (a, b) -> (a, b)
g
  where
    g :: (a, b) -> (a, b)
g xy :: (a, b)
xy@(a
x, b
y)
        | a
key a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x   =  (a
x, b -> b
f b
y)
        | Bool
otherwise  =  (a, b)
xy

-- $alteration
-- The "alter" functions provide an all-in-one way to do insertion,
-- modification, and removal.

-- | Insert, modify, or delete a single value corresponding to
-- the first place where a particular key appears in the list.
--
-- __Modification__ - If the key first appears in the list with a
-- corresponding value of @x@, and @f x = 'Just' x'@, then that value
-- @x@ will be replaced with @x'@ in the resulting list.
--
-- >>> alterFirst 'B' (fmap negate) [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',-4),('C',2),('B',6)]
--
-- __Removal__ - If the key first appears in the list with a corresponding
-- value of @x@, and @f x = 'Nothing'@, then that mapping will be removed
-- in the resulting list.
--
-- >>> alterFirst 'B' (\_ -> Nothing) [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('C',2),('B',6)]
--
-- __Insertion__ - If the key does not appear in the list and
-- @f 'Nothing' = 'Just' x@, then @x@ be appended to the /end/ of the list.
--
-- >>> alterFirst 'D' (\_ -> Just 0) [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',4),('C',2),('B',6),('D',0)]

alterFirst :: forall l a b. (AssocList l a b, Eq a)
    => a -> (Maybe b -> Maybe b) -- ^ @f@
    -> l -> l
alterFirst :: a -> (Maybe b -> Maybe b) -> l -> l
alterFirst a
key Maybe b -> Maybe b
f l
l =
    let (l
before, l
l') = a -> l -> (l, l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> (l, l)
break a
key l
l
    in  l
before l -> l -> l
forall full item. ListLike full item => full -> full -> full
`LL.append`
        case l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
LL.uncons l
l' of
            Maybe ((a, b), l)
Nothing               ->  l -> ((a, b) -> l) -> Maybe (a, b) -> l
forall b a. b -> (a -> b) -> Maybe a -> b
maybe l
forall full item. ListLike full item => full
LL.empty (a, b) -> l
forall full item. ListLike full item => item -> full
LL.singleton ((,) a
key (b -> (a, b)) -> Maybe b -> Maybe (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe b -> Maybe b
f Maybe b
forall a. Maybe a
Nothing)
            Just ((a
x, b
y), l
after)  ->  l -> ((a, b) -> l) -> Maybe (a, b) -> l
forall b a. b -> (a -> b) -> Maybe a -> b
maybe l
forall full item. ListLike full item => full
LL.empty (a, b) -> l
forall full item. ListLike full item => item -> full
LL.singleton ((,) a
x   (b -> (a, b)) -> Maybe b -> Maybe (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe b -> Maybe b
f (b -> Maybe b
forall a. a -> Maybe a
Just b
y))
                                      l -> l -> l
forall full item. ListLike full item => full -> full -> full
`LL.append` l
after

-- | Modify the list of values that correspond to a particular key.
--
-- __Mapping__ - For example, to negate all values of @'B'@:
--
-- >>> alterAll 'B' (map negate) [('A', 1), ('B', 4), ('B', 5), ('C', 2)]
-- [('A',1),('B',-4),('B',-5),('C',2)]
--
-- __Length alteration__ - For example, to limit the number of occurrences
-- of 'B' to at most two:
--
-- >>> alterAll 'B' (take 2) [('A', 1), ('B', 4), ('B', 5), ('B', 6), ('C', 2)]
-- [('A',1),('B',4),('B',5),('C',2)]
--
-- __Removal__ - If @f@ returns an empty list, then the key will be removed
-- from the list entirely.
--
-- >>> alterAll 'B' (\_ -> []) [('A', 1), ('B', 4), ('B', 5), ('C', 2)]
-- [('A',1),('C',2)]
--
-- __Reordering__ - The key may appear in multiple noncontiguous positions
-- in the input list, but all of the new mappings for the key in the output
-- will be in one contiguous sequence starting at the position where the
-- key /first/ appears in the input list.
--
-- >>> alterAll 'B' (map negate) [('A', 1), ('B', 4), ('C', 2), ('B', 5), ('D', 3), ('B', 6)]
-- [('A',1),('B',-4),('B',-5),('B',-6),('C',2),('D',3)]
--
-- __Insertion__ - If the key does not appear in the list, then any result
-- from @f@ will be appended to the /end/ of the list.
--
-- >>> alterAll 'D' (\_ -> [7, 8]) [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',4),('C',2),('B',6),('D',7),('D',8)]

alterAll :: forall l a b. (AssocList l a b, Eq a)
    => a -> ([b] -> [b]) -- ^ @f@
    -> l -> l
alterAll :: a -> ([b] -> [b]) -> l -> l
alterAll a
key [b] -> [b]
f l
l =
    let (l
before, l
l') = a -> l -> (l, l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> (l, l)
break a
key l
l
    in  l
before l -> l -> l
forall full item. ListLike full item => full -> full -> full
`LL.append`
        case (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons l
l') of
            Maybe ((a, b), l)
Nothing  ->  [Item l] -> l
forall l. IsList l => [Item l] -> l
LL.fromList ((,) a
key (b -> (a, b)) -> [b] -> [(a, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [b] -> [b]
f [])
            Maybe ((a, b), l)
_        ->  let ([b]
ys, l
after) = a -> l -> ([b], l)
forall l a b. (AssocList l a b, Eq a) => a -> l -> ([b], l)
partition a
key l
l'
                         in  [Item l] -> l
forall l. IsList l => [Item l] -> l
LL.fromList ((,) a
key (b -> (a, b)) -> [b] -> [(a, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [b] -> [b]
f [b]
ys) l -> l -> l
forall full item. ListLike full item => full -> full -> full
`LL.append` l
after