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

-- | Functions on association lists that involve 'Predicate's on the keys.

module Data.AssocList.ListLike.Predicate
    (

    -- * Related modules
    -- $relatedModules

    -- * Lookup
      lookupFirst
    , lookupAll

    -- * Removal
    , removeFirst
    , removeAll

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

    -- * Grouping
    , partition
    , break
    , breakPartition

    ) 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

-- contravariant
import Data.Functor.Contravariant (Predicate (..))

-- $setup
-- >>> import Prelude ((==), negate)

-- $relatedModules
-- Some other modules that are a lot like this one:
--
-- * "Data.AssocList.ListLike.Eq" - Functions on association lists that
--   make use of an 'Eq' constraint on the type of the keys
-- * "Data.AssocList.ListLike.Equivalence" - Functions on association
--   lists that involve 'Equivalence's on the keys

-- | Obtain the first value associated with a key that satisfies a
-- predicate, if such a mapping is present.
--
-- >>> lookupFirst (Predicate (== 'B')) [('A',1), ('B',2), ('B',3), ('C',4)]
-- Just 2
--
-- The result is 'Nothing' if no key in the list satisfies the predicate.
--
-- >>> lookupFirst (Predicate (== 'D')) [('A',1), ('B',2), ('B',3), ('C',4)]
-- Nothing

lookupFirst :: forall l a b. AssocList l a b
    => Predicate a -> l -> Maybe b
lookupFirst :: Predicate a -> l -> Maybe b
lookupFirst Predicate 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 Predicate 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))
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x            =  b -> Maybe b
forall a. a -> Maybe a
Just b
y
        | Bool
otherwise                     =  Predicate a -> l -> Maybe b
forall l a b. AssocList l a b => Predicate a -> l -> Maybe b
lookupFirst Predicate a
key l
xys

-- | Obtain all values associated with keys that satisfy the predicate,
-- in the order in which the mappings appear in the list.
--
-- >>> lookupAll (Predicate (== 'B')) [('A',1), ('B',2), ('B',3), ('C',4), ('B',3)]
-- [2,3,3]

lookupAll :: forall l a b. AssocList l a b
    => Predicate a -> l -> [b]
lookupAll :: Predicate a -> l -> [b]
lookupAll Predicate a
_key (l -> Maybe ((a, b), l)
forall full item. ListLike full item => full -> Maybe (item, full)
uncons -> Maybe ((a, b), l)
Nothing)      =  []
lookupAll Predicate 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))
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x            =  b
y b -> [b] -> [b]
forall a. a -> [a] -> [a]
: Predicate a -> l -> [b]
forall l a b. AssocList l a b => Predicate a -> l -> [b]
lookupAll Predicate a
key l
xys
        | Bool
otherwise                     =      Predicate a -> l -> [b]
forall l a b. AssocList l a b => Predicate a -> l -> [b]
lookupAll Predicate a
key l
xys

-- | Produce a modified version of the association list in which the
-- first occurrence of a key that satisfied the predicate has been removed.
--
-- >>> removeFirst (Predicate (== 'B')) [('A',1), ('B',2), ('B',3), ('C',4)]
-- [('A',1),('B',3),('C',4)]
--
-- If no key in the list satisfies the predicate, then the original list
-- is returned.
--
-- >>> removeFirst (Predicate (== 'C')) [('A',1), ('B',2), ('B',3)]
-- [('A',1),('B',2),('B',3)]

removeFirst :: forall l a b. AssocList l a b
    => Predicate a -> l -> l
removeFirst :: Predicate a -> l -> l
removeFirst Predicate 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 Predicate 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))
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x            =  l
xys
        | Bool
otherwise                     =  (a, b) -> l -> l
forall full item. ListLike full item => item -> full -> full
cons (a, b)
xy (Predicate a -> l -> l
forall l a b. AssocList l a b => Predicate a -> l -> l
removeFirst Predicate a
key l
xys)

-- | Produce a modified version of the association list in which all
-- occurrences of keys that satisfy the predicate have been removed.
--
-- >>> removeAll (Predicate (== '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 (Predicate (== 'C')) [('A',1), ('B',2), ('B',3)]
-- [('A',1),('B',2),('B',3)]

removeAll :: forall l a b. AssocList l a b
    => Predicate a -> l -> l
removeAll :: Predicate a -> l -> l
removeAll Predicate 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 Predicate 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))
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x            =           Predicate a -> l -> l
forall l a b. AssocList l a b => Predicate a -> l -> l
removeAll Predicate a
key l
xys
        | Bool
otherwise                     =  (a, b) -> l -> l
forall full item. ListLike full item => item -> full -> full
cons (a, b)
xy (Predicate a -> l -> l
forall l a b. AssocList l a b => Predicate a -> l -> l
removeAll Predicate a
key l
xys)

-- | Produces a tuple of two results:
--
-- 1. All values associated with keys that satify the predicate
-- 2. All of the other key-value pairs
--
-- @'partition' x l = ('lookupAll' x l, 'removeAll' x l)@
--
-- >>> partition (Predicate (== '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
    => Predicate a -> l -> ([b], l)
partition :: Predicate a -> l -> ([b], l)
partition Predicate 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 Predicate 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))
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key 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) = Predicate a -> l -> ([b], l)
forall l a b. AssocList l a b => Predicate a -> l -> ([b], l)
partition Predicate a
key l
xys

-- | Produces a tuple of two results:
--
-- 1. The longest prefix of the association list that does /not/ contain
--    a key satisfying the predict
-- 2. The remainder of the list
--
-- >>> break (Predicate (== 'B')) [('A',1), ('B',2), ('B',3), ('C',4)]
-- ([('A',1)],[('B',2),('B',3),('C',4)])
--
-- If the key of the first mapping in the list satisfies the predicate,
-- then the first part of the resulting tuple is empty, and the second
-- part of the result is the entire list.
--
-- >>> break (Predicate (== 'A')) [('A',1), ('B',2), ('B',3), ('C',4)]
-- ([],[('A',1),('B',2),('B',3),('C',4)])
--
-- If no key in the list satisfies the predicate, then the first part of
-- the resulting tuple is the entire list, and the second part of the
-- result is empty.
--
-- >>> break (Predicate (== '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
    => Predicate a -> l -> (l, l)
break :: Predicate a -> l -> (l, l)
break Predicate 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) -> Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x)

-- | 'break' on a predicate, then 'partition' the remainder.
--
-- @'breakPartition' p l@ separates @l@ into three parts:
--
-- 1. The key-value pairs for which the predicate is /not/ satisfied that
--    occur in the list /before/ the first occurrence of a key that satisfies
--    the predicate (@fst ('break' p l)@)
-- 2. All values associated with keys that satisfy the predicate
--    (@'lookupAll' p l@)
-- 3. The key-value pairs for which the predicate is /not/ satisfied that
--    occur in the list /after/ the first occurrence of a key that satisfies
--    the predicate (@'removeAll' p (snd ('break' p l))@)
--
-- >>> breakPartition (Predicate (== 'B')) [('A',1),('B',2),('C',3),('B',4)]
-- ([('A',1)],[2,4],[('C',3)])
--
-- If the predicate is not satisfied by any key in the list, then the
-- first part of the result is the entire list, and the other parts are
-- empty.
--
-- >>> breakPartition (Predicate (== '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
    => Predicate a -> l -> (l, [b], l)
breakPartition :: Predicate a -> l -> (l, [b], l)
breakPartition Predicate a
key l
l =
    let
        (l
before, l
l') = Predicate a -> l -> (l, l)
forall l a b. AssocList l a b => Predicate a -> l -> (l, l)
break     Predicate a
key l
l
        ([b]
xs, l
after)  = Predicate a -> l -> ([b], l)
forall l a b. AssocList l a b => Predicate a -> l -> ([b], l)
partition Predicate 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 key satisfying the predicate first appears
-- in the list, apply a function to the corresponding value.
--
-- >>> mapFirst (Predicate (== 'B')) negate [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',-4),('C',2),('B',6)]
--
-- If no key in the list satisfies the predicate, then the original list is
-- returned without modification.
--
-- >>> mapFirst (Predicate (== '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
    => Predicate a -> (b -> b) -> l -> l
mapFirst :: Predicate a -> (b -> b) -> l -> l
mapFirst Predicate a
key b -> b
f l
l =
    let
        (l
before, l
l') = Predicate a -> l -> (l, l)
forall l a b. AssocList l a b => Predicate a -> l -> (l, l)
break Predicate 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 in the list where the key satisfies the predicate,
-- apply a function to the corresponding value.
--
-- >>> mapAll (Predicate (== 'B')) negate [('A', 1), ('B', 4), ('C', 2), ('B', 6)]
-- [('A',1),('B',-4),('C',2),('B',-6)]
--
-- If no key in the list satisfies the predicate, then the original list is
-- returned without modification.
--
-- >>> mapAll (Predicate (== '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
    => Predicate a -> (b -> b) -> l -> l
mapAll :: Predicate a -> (b -> b) -> l -> l
mapAll Predicate 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)
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x   =  (a
x, b -> b
f b
y)
        | Bool
otherwise            =  (a, b)
xy