-- | 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 :: forall k v. Ord k => AssocList k v -> k -> Maybe v
apply AssocList k v
m = (forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` 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 :: forall k v. AssocList k v -> [k]
keys = forall a b. (a -> b) -> [a] -> [b]
map 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 :: forall k v. k -> v -> AssocList k v -> AssocList k v
insert k
k v
v = ((k
k,v
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 :: forall k v. Eq k => k -> v -> AssocList k v -> AssocList k v
update k
k v
v = forall k v. Eq k => k -> (v -> v) -> AssocList k v -> AssocList k v
updateAt k
k forall a b. (a -> b) -> a -> b
$ 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 :: forall k v. Eq k => k -> AssocList k v -> AssocList k v
delete k
k = forall a. (a -> a -> Bool) -> a -> [a] -> [a]
List.deleteBy (forall a. Eq a => a -> a -> Bool
(==) forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) (k
k, 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 :: forall k v. Eq k => k -> (v -> v) -> AssocList k v -> AssocList k v
updateAt k
k v -> v
f = [(k, v)] -> [(k, v)]
loop where
  loop :: [(k, v)] -> [(k, v)]
loop []       = forall a. HasCallStack => a
__IMPOSSIBLE__
  loop (p :: (k, v)
p@(k
k',v
v) : [(k, v)]
ps)
    | k
k forall a. Eq a => a -> a -> Bool
== k
k'   = (k
k, v -> v
f v
v) forall a. a -> [a] -> [a]
: [(k, v)]
ps
    | Bool
otherwise = (k, v)
p forall a. a -> [a] -> [a]
: [(k, v)] -> [(k, v)]
loop [(k, v)]
ps

-- | O(n).
--   Map over an association list, preserving the order.
mapWithKey :: (k -> v -> v) -> AssocList k v -> AssocList k v
mapWithKey :: forall k v. (k -> v -> v) -> AssocList k v -> AssocList k v
mapWithKey k -> v -> v
f = forall a b. (a -> b) -> [a] -> [b]
map 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 :: forall (m :: * -> *) k v.
Applicative m =>
(k -> v -> m v) -> AssocList k v -> m (AssocList k v)
mapWithKeyM k -> v -> m v
f = forall {f :: * -> *} {t} {a}.
Applicative f =>
(t -> f a) -> [t] -> f [a]
mapM forall a b. (a -> b) -> a -> b
$ \ (k
k,v
v) -> (k
k,) 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 [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    mapM t -> f a
g (t
x : [t]
xs) = (:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> f a
g t
x 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 :: forall k k' v. (k -> k') -> AssocList k v -> AssocList k' v
mapKeysMonotonic k -> k'
f = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a -> b) -> a -> b
$ forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst k -> k'
f