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

module Data.AssocList.List.Predicate
    (

    -- * Related modules
    -- $relatedModules

    -- * Lookup
      lookupFirst
    , lookupAll

    -- * Removal
    , removeFirst
    , removeAll

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

    -- * Grouping
    , partition
    , break
    , breakPartition

    ) where

import Data.AssocList.List.Concept

-- base
import qualified Data.List
import Data.Maybe (maybeToList)
import Prelude (Maybe (..), (++), (<$>), otherwise)

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

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

-- $relatedModules
-- Some other modules that are a lot like this one:
--
-- * "Data.AssocList.List.Eq" - Functions on 'AssocList's that make
--   use of an 'Eq' constraint on the type of the keys
-- * "Data.AssocList.List.Equivalence" - Functions on 'AssocList's
--   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 :: Predicate a -> AssocList a b -> Maybe b
lookupFirst :: Predicate a -> AssocList a b -> Maybe b
lookupFirst Predicate a
_key []                =  Maybe b
forall a. Maybe a
Nothing
lookupFirst Predicate a
key ((a
x, b
y) : AssocList a b
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 -> AssocList a b -> Maybe b
forall a b. Predicate a -> AssocList a b -> Maybe b
lookupFirst Predicate a
key AssocList a b
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 :: Predicate a -> AssocList a b -> [b]
lookupAll :: Predicate a -> AssocList a b -> [b]
lookupAll Predicate a
_key []                  =  []
lookupAll Predicate a
key ((a
x, b
y) : AssocList a b
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 -> AssocList a b -> [b]
forall a b. Predicate a -> AssocList a b -> [b]
lookupAll Predicate a
key AssocList a b
xys
        | Bool
otherwise                =      Predicate a -> AssocList a b -> [b]
forall a b. Predicate a -> AssocList a b -> [b]
lookupAll Predicate a
key AssocList a b
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 :: Predicate a -> AssocList a b -> AssocList a b
removeFirst :: Predicate a -> AssocList a b -> AssocList a b
removeFirst Predicate a
_key l :: AssocList a b
l@[]              =  AssocList a b
l
removeFirst Predicate a
key (xy :: (a, b)
xy@(a
x, b
y) : AssocList a b
xys)
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x       =  AssocList a b
xys
        | Bool
otherwise                =  (a, b)
xy (a, b) -> AssocList a b -> AssocList a b
forall a. a -> [a] -> [a]
: Predicate a -> AssocList a b -> AssocList a b
forall a b. Predicate a -> AssocList a b -> AssocList a b
removeFirst Predicate a
key AssocList a b
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 :: Predicate a -> AssocList a b -> AssocList a b
removeAll :: Predicate a -> AssocList a b -> AssocList a b
removeAll Predicate a
_key l :: AssocList a b
l@[]                =  AssocList a b
l
removeAll Predicate a
key (xy :: (a, b)
xy@(a
x, b
y) : AssocList a b
xys)
        | Predicate a -> a -> Bool
forall a. Predicate a -> a -> Bool
getPredicate Predicate a
key a
x       =       Predicate a -> AssocList a b -> AssocList a b
forall a b. Predicate a -> AssocList a b -> AssocList a b
removeAll Predicate a
key AssocList a b
xys
        | Bool
otherwise                =  (a, b)
xy (a, b) -> AssocList a b -> AssocList a b
forall a. a -> [a] -> [a]
: Predicate a -> AssocList a b -> AssocList a b
forall a b. Predicate a -> AssocList a b -> AssocList a b
removeAll Predicate a
key AssocList a b
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 :: Predicate a -> AssocList a b -> ([b], AssocList a b)
partition :: Predicate a -> AssocList a b -> ([b], AssocList a b)
partition Predicate a
_key l :: AssocList a b
l@[]                = ([], AssocList a b
l)
partition Predicate a
key (xy :: (a, b)
xy@(a
x, b
y) : AssocList a b
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 ,      AssocList a b
no)
        | Bool
otherwise                = (    [b]
yes , (a, b)
xy (a, b) -> AssocList a b -> AssocList a b
forall a. a -> [a] -> [a]
: AssocList a b
no)
    where
        ([b]
yes, AssocList a b
no) = Predicate a -> AssocList a b -> ([b], AssocList a b)
forall a b. Predicate a -> AssocList a b -> ([b], AssocList a b)
partition Predicate a
key AssocList a b
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 :: Predicate a -> AssocList a b -> (AssocList a b, AssocList a b)
break :: Predicate a -> AssocList a b -> (AssocList a b, AssocList a b)
break Predicate a
key = ((a, b) -> Bool) -> AssocList a b -> (AssocList a b, AssocList a b)
forall a. (a -> Bool) -> [a] -> ([a], [a])
Data.List.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 :: Predicate a -> AssocList a b
    -> (AssocList a b, [b], AssocList a b)
breakPartition :: Predicate a -> AssocList a b -> (AssocList a b, [b], AssocList a b)
breakPartition Predicate a
key AssocList a b
l =
    let
        (AssocList a b
before, AssocList a b
l') = Predicate a -> AssocList a b -> (AssocList a b, AssocList a b)
forall a b.
Predicate a -> AssocList a b -> (AssocList a b, AssocList a b)
break     Predicate a
key AssocList a b
l
        ([b]
xs, AssocList a b
after)  = Predicate a -> AssocList a b -> ([b], AssocList a b)
forall a b. Predicate a -> AssocList a b -> ([b], AssocList a b)
partition Predicate a
key AssocList a b
l'
    in
        (AssocList a b
before, [b]
xs, AssocList a b
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 :: Predicate a -> (b -> b) -> AssocList a b -> AssocList a b
mapFirst :: Predicate a -> (b -> b) -> AssocList a b -> AssocList a b
mapFirst Predicate a
key b -> b
f AssocList a b
l =
    let
        (AssocList a b
before, AssocList a b
l') = Predicate a -> AssocList a b -> (AssocList a b, AssocList a b)
forall a b.
Predicate a -> AssocList a b -> (AssocList a b, AssocList a b)
break Predicate a
key AssocList a b
l
    in
        AssocList a b
before AssocList a b -> AssocList a b -> AssocList a b
forall a. [a] -> [a] -> [a]
++
        case AssocList a b
l' of
            []              ->  AssocList a b
l'
            (a
x, b
y) : AssocList a b
after  ->  (a
x, b -> b
f b
y) (a, b) -> AssocList a b -> AssocList a b
forall a. a -> [a] -> [a]
: AssocList a b
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 :: Predicate a -> (b -> b) -> AssocList a b -> AssocList a b
mapAll :: Predicate a -> (b -> b) -> AssocList a b -> AssocList a b
mapAll Predicate a
key b -> b
f =
    ((a, b) -> (a, b)) -> AssocList a b -> AssocList a b
forall a b. (a -> b) -> [a] -> [b]
Data.List.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