-- | Additional functions for association lists.

module Agda.Utils.AssocList
  ( module Agda.Utils.AssocList
  , lookup
  ) where

import Prelude hiding (lookup)

import Data.Function
import Data.List (lookup)
import qualified Data.List as List
import qualified Data.Map  as Map

import Agda.Utils.Tuple

import Agda.Utils.Impossible

-- | A finite map, represented as a set of pairs.
--
--   Invariant: at most one value per key.
type AssocList k v = [(k,v)]

-- Lookup, reexported from Data.List.
-- O(n).
-- lookup :: Eq k => k -> AssocList k v -> Maybe v

-- | Lookup keys in the same association list often.
--   Use partially applied to create partial function
--   @apply m :: k -> Maybe v@.
--
--   * First time: @O(n log n)@ in the worst case.
--   * Subsequently: @O(log n)@.
--
--   Specification: @apply m == (`lookup` m)@.
apply :: Ord k => AssocList k v -> k -> Maybe v
apply :: AssocList k v -> k -> Maybe v
apply AssocList k v
m = (k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` (v -> v -> v) -> AssocList k v -> Map k v
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith (\ v
_new v
old -> v
old) AssocList k v
m)

-- | O(n).
--   Get the domain (list of keys) of the finite map.
keys :: AssocList k v -> [k]
keys :: AssocList k v -> [k]
keys = ((k, v) -> k) -> AssocList k v -> [k]
forall a b. (a -> b) -> [a] -> [b]
map (k, v) -> k
forall a b. (a, b) -> a
fst

-- | O(1).
--   Add a new binding.
--   Assumes the binding is not yet in the list.
insert :: k -> v -> AssocList k v -> AssocList k v
insert :: k -> v -> AssocList k v -> AssocList k v
insert k
k v
v = ((k
k,v
v) (k, v) -> AssocList k v -> AssocList k v
forall a. a -> [a] -> [a]
:)

-- | O(n).
--   Update the value at a key.
--   The key must be in the domain of the finite map.
--   Otherwise, an internal error is raised.
update :: Eq k => k -> v -> AssocList k v -> AssocList k v
update :: k -> v -> AssocList k v -> AssocList k v
update k
k v
v = k -> (v -> v) -> AssocList k v -> AssocList k v
forall k v. Eq k => k -> (v -> v) -> AssocList k v -> AssocList k v
updateAt k
k ((v -> v) -> AssocList k v -> AssocList k v)
-> (v -> v) -> AssocList k v -> AssocList k v
forall a b. (a -> b) -> a -> b
$ v -> v -> v
forall a b. a -> b -> a
const v
v

-- | O(n).
--   Delete a binding.
--   The key must be in the domain of the finite map.
--   Otherwise, an internal error is raised.
delete :: Eq k => k -> AssocList k v -> AssocList k v
delete :: k -> AssocList k v -> AssocList k v
delete k
k = ((k, v) -> (k, v) -> Bool)
-> (k, v) -> AssocList k v -> AssocList k v
forall a. (a -> a -> Bool) -> a -> [a] -> [a]
List.deleteBy (k -> k -> Bool
forall a. Eq a => a -> a -> Bool
(==) (k -> k -> Bool) -> ((k, v) -> k) -> (k, v) -> (k, v) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (k, v) -> k
forall a b. (a, b) -> a
fst) (k
k, v
forall a. HasCallStack => a
__IMPOSSIBLE__)

-- | O(n).
--   Update the value at a key with a certain function.
--   The key must be in the domain of the finite map.
--   Otherwise, an internal error is raised.
updateAt :: Eq k => k -> (v -> v) -> AssocList k v -> AssocList k v
updateAt :: k -> (v -> v) -> AssocList k v -> AssocList k v
updateAt k
k v -> v
f = AssocList k v -> AssocList k v
loop where
  loop :: AssocList k v -> AssocList k v
loop []       = AssocList k v
forall a. HasCallStack => a
__IMPOSSIBLE__
  loop (p :: (k, v)
p@(k
k',v
v) : AssocList k v
ps)
    | k
k k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k'   = (k
k, v -> v
f v
v) (k, v) -> AssocList k v -> AssocList k v
forall a. a -> [a] -> [a]
: AssocList k v
ps
    | Bool
otherwise = (k, v)
p (k, v) -> AssocList k v -> AssocList k v
forall a. a -> [a] -> [a]
: AssocList k v -> AssocList k v
loop AssocList k v
ps

-- | O(n).
--   Map over an association list, preserving the order.
mapWithKey :: (k -> v -> v) -> AssocList k v -> AssocList k v
mapWithKey :: (k -> v -> v) -> AssocList k v -> AssocList k v
mapWithKey k -> v -> v
f = ((k, v) -> (k, v)) -> AssocList k v -> AssocList k v
forall a b. (a -> b) -> [a] -> [b]
map (((k, v) -> (k, v)) -> AssocList k v -> AssocList k v)
-> ((k, v) -> (k, v)) -> AssocList k v -> AssocList k v
forall a b. (a -> b) -> a -> b
$ \ (k
k,v
v) -> (k
k, k -> v -> v
f k
k v
v)

-- | O(n).
--   If called with a effect-producing function, violation of the invariant
--   could matter here (duplicating effects).
mapWithKeyM :: Applicative m => (k -> v -> m v) -> AssocList k v -> m (AssocList k v)
mapWithKeyM :: (k -> v -> m v) -> AssocList k v -> m (AssocList k v)
mapWithKeyM k -> v -> m v
f = ((k, v) -> m (k, v)) -> AssocList k v -> m (AssocList k v)
forall (f :: * -> *) t a.
Applicative f =>
(t -> f a) -> [t] -> f [a]
mapM (((k, v) -> m (k, v)) -> AssocList k v -> m (AssocList k v))
-> ((k, v) -> m (k, v)) -> AssocList k v -> m (AssocList k v)
forall a b. (a -> b) -> a -> b
$ \ (k
k,v
v) -> (k
k,) (v -> (k, v)) -> m v -> m (k, v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> v -> m v
f k
k v
v
  where
    -- mapM is applicative!
    mapM :: (t -> f a) -> [t] -> f [a]
mapM t -> f a
g [] = [a] -> f [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    mapM t -> f a
g (t
x : [t]
xs) = (:) (a -> [a] -> [a]) -> f a -> f ([a] -> [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> f a
g t
x f ([a] -> [a]) -> f [a] -> f [a]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (t -> f a) -> [t] -> f [a]
mapM t -> f a
g [t]
xs

-- | O(n).
--   Named in analogy to 'Data.Map.mapKeysMonotonic'.
--   To preserve the invariant, it is sufficient that the key
--   transformation is injective (rather than monotonic).
mapKeysMonotonic :: (k -> k') -> AssocList k v -> AssocList k' v
mapKeysMonotonic :: (k -> k') -> AssocList k v -> AssocList k' v
mapKeysMonotonic k -> k'
f = ((k, v) -> (k', v)) -> AssocList k v -> AssocList k' v
forall a b. (a -> b) -> [a] -> [b]
map (((k, v) -> (k', v)) -> AssocList k v -> AssocList k' v)
-> ((k, v) -> (k', v)) -> AssocList k v -> AssocList k' v
forall a b. (a -> b) -> a -> b
$ (k -> k') -> (k, v) -> (k', v)
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst k -> k'
f