-----------------------------------------------------------------------------
-- |
-- Module : Data.Map.AVL
-- Copyright : (c) Adrian Hey 2005,2006
-- License : BSD3
--
-- Maintainer : http://homepages.nildram.co.uk/~ahey/em.png
-- Stability : provisional
-- Portability : portable
--
-- This module provides an AVL tree based clone of the base package Data.Map.
--
-- There are some differences though..
--
-- * 'size' is O(n), not O(1). Consequently, indexed access is disabled.
--
-- * The showTree and showTreeWith functions are not implemented.
--
-- * Some other functions are not yet implemented.
--
-----------------------------------------------------------------------------
module Data.Map.AVL (
-- * Map type
Map
-- * Operators
, (!)
, (\\)
-- * Query
, null
, size
, member
, lookup
, findWithDefault
-- * Construction
, empty
, singleton
-- ** Insertion
, insert
, insertWith, insertWithKey, insertLookupWithKey
-- ** Delete\/Update
, delete
, adjust
, alter
, adjustWithKey
, update
, updateWithKey
, updateLookupWithKey
-- * Combine
-- ** Union
, union
, unionWith
, unionWithKey
, unions
, unionsWith
-- ** Difference
, difference
, differenceWith
, differenceWithKey
-- ** Intersection
, intersection
, intersectionWith
, intersectionWithKey
-- * Traversal
-- ** Map
, map
, mapWithKey
, mapAccum
-- , mapAccumWithKey
-- , mapKeys
-- , mapKeysWith
-- , mapKeysMonotonic
-- ** Fold
, fold
, foldWithKey
-- * Conversion
, elems
, keys
, keysSet
, liftKeysSet
, assocs
, unsafeFromTree
, toTree
, toList
, fromList
, fromListWith
, fromListWithKey
-- ** Ordered lists
, toAscList
, fromAscList
, fromAscListWith
, fromAscListWithKey
, fromDistinctAscList
-- * Filter
, filter
, filterWithKey
, partition
, partitionWithKey
, split
, splitLookup
-- * Submap
, isSubmapOf
, isSubmapOfBy
-- , isProperSubmapOf, isProperSubmapOfBy
-- * Indexed
-- , lookupIndex
-- , findIndex
-- , elemAt
-- , updateAt
-- , deleteAt
-- * Min\/Max
, findMin
, findMax
, deleteMin
, deleteMax
, deleteFindMin
, deleteFindMax
-- , updateMin
-- , updateMax
-- , updateMinWithKey
-- , updateMaxWithKey
-- * Debugging
-- , showTree
-- , showTreeWith
-- , valid
) where
import Prelude hiding (lookup,map,filter,foldr,foldl,null)
import qualified Data.List as List
import Data.Monoid
-- import qualified Data.Maybe as Maybe
import qualified Data.Set.AVL as Set
-- import Data.Monoid
import Data.Foldable hiding (toList, find, fold)
import qualified Data.COrdering as COrdering
import qualified Data.Tree.AVL as AVL
-- import qualified Data.Tree.AVL.Test.Utils as AVL
import Data.Typeable
#include "Typeable.h"
INSTANCE_TYPEABLE2(Map,mapTc,"Data.Map.AVL")
------------------------------------------------------
---- local combining comparison utilities ----------
------------------------------------------------------
readValCC :: Ord k => k -> (k, a) -> COrdering.COrdering a
readValCC k (k', a) = case compare k k' of
LT -> COrdering.Lt
EQ -> COrdering.Eq a
GT -> COrdering.Gt
mcmp :: Ord a => (a, b) -> (a, c) -> COrdering.COrdering (a, b)
mcmp (k, a) (k', _) = case compare k k' of
LT -> COrdering.Lt
EQ -> COrdering.Eq (k, a)
GT -> COrdering.Gt
mfcmp :: Ord k => (k -> a -> b -> c) -> (k, a) -> (k, b)
-> COrdering.COrdering (k, c)
mfcmp f (k, a) (k', b) = case compare k k' of
LT -> COrdering.Lt
EQ -> COrdering.Eq (k, f k a b)
GT -> COrdering.Gt
mmfcmp :: (Functor f, Ord k) => (k -> a -> b -> f c) -> (k, a)
-> (k, b) -> COrdering.COrdering (f (k, c))
mmfcmp f (k, a) (k', b) = case compare k k' of
LT -> COrdering.Lt
EQ -> COrdering.Eq $ fmap (\c -> (k, c)) $ f k a b
GT -> COrdering.Gt
infixl 9 !, \\ --
toOrdering :: COrdering.COrdering a -> Ordering
toOrdering c = case c of
COrdering.Lt -> LT
COrdering.Eq _ -> EQ
COrdering.Gt -> GT
toOrd :: (a -> b -> COrdering.COrdering c) -> a -> b -> Ordering
toOrd f a = toOrdering . f a
-- | A Map from keys @k@ to values @a@.
newtype Map k a = Map (AVL.AVL (k, a))
-- deriving (Eq, Ord, Show)
instance (Eq k, Eq a) => Eq (Map k a) where
m1 == m2 = toList m1 == toList m2
instance (Ord k, Ord a) => Ord (Map k a) where
compare m1 m2 = compare (toList m1) (toList m2)
showSet :: (Show a) => [a] -> ShowS
showSet []
= showString "{}"
showSet (x:xs)
= showChar '{' . shows x . showTail xs
where
showTail [] = showChar '}'
showTail (x':xs') = showString ", " . shows x' . showTail xs'
instance (Show k, Show a) => Show (Map k a) where
showsPrec _ (Map t) = showSet (AVL.asListL t)
-- | /O(1)/. The empty map.
empty :: Map k a
empty = Map (AVL.empty)
-- | /O(1)/. A map with a single element.
singleton :: k -> a -> Map k a
singleton k a = k `seq` Map (AVL.singleton (k, a))
-- | /O(1)/. Is the map empty?
null :: Map k a -> Bool
null (Map t) = AVL.isEmpty t
-- | /O(n)/. The number of elements in the map.
size :: Map k a -> Int
size (Map t) = AVL.size t
-- | /O(log n)/. Is the key a member of the map?
member :: Ord k => k -> Map k a -> Bool
member k (Map t) = k `seq` AVL.genContains t (compare k . fst)
-- | /O(log n)/. Find the value at a key.
-- Calls 'error' when the element can not be found.
(!) :: Ord k => Map k a -> k -> a
(!) m k = find k m
-- | /O(log n)/. Find the value at a key.
-- Calls 'error' when the element can not be found.
find :: Ord k => k -> Map k a -> a
find = findWithDefault (error "Map.find: element not in the map")
-- | /O(log n)/. Lookup the value at a key in the map.
lookup :: (Monad m,Ord k) => k -> Map k a -> m a
lookup k (Map t) = k `seq` maybe (fail "AvlMap.lookup: Key not found")
return (AVL.genTryRead t (readValCC k))
-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
-- the value at key @k@ or returns @def@ when the key is not in the map.
findWithDefault :: Ord k => a -> k -> Map k a -> a
findWithDefault def k (Map t) = k `seq` AVL.genDefaultRead def t (readValCC k)
-- | /O(log n)/. Insert a new key and value in the map.
-- If the key is already present in the map, the associated value is
-- replaced with the supplied value, i.e. 'insert' is equivalent to
-- @'insertWith' 'const'@.
insert :: Ord k => k -> a -> Map k a -> Map k a
insert k a (Map t) = k `seq` Map (AVL.genPush (mcmp (k, a)) (k, a) t)
-- | /O(log n)/. Insert with a combining function.
insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith f = insertWithKey (\_ z y -> f z y)
-- | /O(log n)/. Insert with a combining function.
insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey f k a (Map t) =
k `seq` Map (AVL.genPush (mfcmp f (k, a)) (k, a) t)
-- | /O(log n)/. The expression (@'insertLookupWithKey' f k x map@)
-- is a pair where the first element is equal to (@'lookup' k map@)
-- and the second element equal to (@'insertWithKey' f k x map@).
--
-- TODO: only one traversal. This requires fiddling with AVL.Push.
insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a,Map k a)
insertLookupWithKey f k a m = (lookup k m, insertWithKey f k a m)
-- | /O(log n)/. Delete a key and its value from the map. When the key is not
-- a member of the map, the original map is returned.
delete :: Ord k => k -> Map k a -> Map k a
delete k (Map t) = k `seq` Map (AVL.genDel (compare k . fst) t)
-- | /O(n)/. Map a function over all values in the map.
map :: (a -> b) -> Map k a -> Map k b
map f = mapWithKey (\_ x -> f x)
-- | /O(n)/. Map a function over all values in the map.
mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
mapWithKey f (Map t) = Map (AVL.mapAVL mf t)
where mf (k', a') = (k', f k' a')
-- | /O(n)/. The function 'mapAccum' threads an accumulating
-- argument through the map in ascending order of keys.
mapAccum :: Ord k => (a -> b -> (a, c)) -> a -> Map k b -> (a, Map k c)
mapAccum f a = foldWithKey ( \ k b (s, m) ->
let (r, c) = f s b in (r, insert k c m)) (a, empty)
-- | /O(n)/. Filter all values that satisfy the predicate.
filter :: Ord k => (a -> Bool) -> Map k a -> Map k a
filter p (Map t) = Map (AVL.filterViaList (p . snd) t)
-- | /O(n)/. Filter all keys\/values that satisfy the predicate.
filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
filterWithKey p (Map t) = Map (AVL.filterViaList (mp p) t)
mp :: (k -> a -> Bool) -> (k, a) -> Bool
mp p (k, a) = p k a
-- | /O(n)/. partition the map according to a predicate. The first
-- map contains all elements that satisfy the predicate, the second all
-- elements that fail the predicate.
partition :: Ord k => (a -> Bool) -> Map k a -> (Map k a,Map k a)
partition p = partitionWithKey (\_ x -> p x)
-- | /O(n)/. partition the map according to a predicate. The first
-- map contains all elements that satisfy the predicate, the second all
-- elements that fail the predicate.
partitionWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> (Map k a,Map k a)
partitionWithKey p (Map t) = let (t1, t2) = AVL.partitionAVL (mp p) t in
(Map t1, Map t2)
-- | /O(log n)/. The expression (@'split' x set@) is a pair @(set1,set2)@
-- where all elements in @set1@ are lower than @x@ and all elements in
-- @set2@ larger than @x@. @x@ is not found in neither @set1@ nor @set2@.
split :: Ord k => k -> Map k a -> (Map k a,Map k a)
split k (Map t) = (Map lessT, Map greaterT)
where (lessT, _, greaterT) = AVL.genFork (readValCC k) t
-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
-- like 'split' but also returns @'lookup' k map@.
splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a)
splitLookup k (Map t) = (Map lessT, a, Map greaterT)
where (lessT, a, greaterT) = AVL.genFork (readValCC k) t
-- | /O(log n)/. The minimal key of the map.
findMin :: Map k a -> (k,a)
findMin (Map t) = AVL.assertReadL t
-- | /O(log n)/. Delete the minimal key.
deleteMin :: Map k a -> Map k a
deleteMin (Map t) = Map $ maybe (error "Set.deleteMin") id $ AVL.tryDelL t
-- | /O(log n)/. Delete and find the minimal element.
deleteFindMin :: Map k a -> ((k,a),Map k a)
deleteFindMin (Map t) = let ((m, v), s) = AVL.assertPopL t in ((m, v), Map s)
-- | /O(log n)/. Delete and find the maximal element.
deleteFindMax :: Map k a -> ((k,a),Map k a)
deleteFindMax (Map t) = let (s, (m, v)) = AVL.assertPopR t in ((m, v), Map s)
-- | /O(log n)/. The minimal key of the map.
findMax :: Map k a -> (k,a)
findMax (Map t) = AVL.assertReadR t
-- | /O(log n)/. Delete the minimal key.
deleteMax :: Map k a -> Map k a
deleteMax (Map t) = Map $ maybe (error "Set.deleteMax") id $ AVL.tryDelR t
-- | /O(n+m)/. Intersection of two maps. The values in the first
-- map are returned, i.e.
-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@).
intersection :: Ord k => Map k a -> Map k b -> Map k a
intersection (Map t1) (Map t2) = Map (AVL.genIntersection mcmp t1 t2)
-- | /O(n+m)/. Intersection with a combining function.
intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c
intersectionWith f = intersectionWithKey (\_ x y -> f x y)
-- | /O(n+m)/. Intersection with a combining function.
-- Intersection is more efficient on (bigset `intersection` smallset)
intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b
-> Map k c
intersectionWithKey f (Map t1) (Map t2) =
Map (AVL.genIntersection (mfcmp f) t1 t2)
-- | /O(n)/. Convert to a list of key\/value pairs.
toList :: Map k a -> [(k,a)]
toList (Map t) = AVL.asListL t
-- | /O(n)/. Convert to a list of key\/value pairs.
toAscList :: Map k a -> [(k,a)]
toAscList = toList
-- | /O(n)/. Convert to a list of key\/value pairs.
assocs :: Map k a -> [(k,a)]
assocs = toList
-- | /O(n)/. Convert to a list of keys.
keys :: Map k a -> [k]
keys = List.map fst . toList
-- | /O(n)/. The set of all keys of the map.
keysSet :: Map k a -> Set.Set k
keysSet = Set.unsafeFromTree . fmap fst . toTree
-- | /O(n)/. Apply a function to each element of a set and return the resulting map.
liftKeysSet :: (k -> b) -> Set.Set k -> Map k b
liftKeysSet f = unsafeFromTree . fmap (\k -> (k,f k)) . Set.toTree
-- | /O(n)/. Convert to a list of values.
elems :: Map k a -> [a]
elems (Map t) = List.map snd (AVL.asListL t)
-- | /O(n)/. Fold the values in the map, such that
-- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
-- For example,
--
-- > elems map = fold (:) [] map
--
fold :: (a -> b -> b) -> b -> Map k a -> b
fold f = foldWithKey (\_ x c -> f x c)
foldWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
foldWithKey f z (Map t) = AVL.foldlAVL' (\ c (k, a) -> f k a c) z t
-- | /O(n+m)/. See 'difference'.
(\\) :: Ord k => Map k a -> Map k b -> Map k a
m1 \\ m2 = difference m1 m2
-- | /O(n+m)/. Difference of two maps.
difference :: Ord k => Map k a -> Map k b -> Map k a
difference (Map t1) (Map t2) = Map (AVL.genDifference (toOrd mcmp) t1 t2)
-- | /O(n+m)/. Difference with a combining function.
differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
differenceWith f = differenceWithKey (\_ x y -> f x y)
differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b
-> Map k a
differenceWithKey f (Map t1) (Map t2) =
Map (AVL.genDifferenceMaybe (mmfcmp f) t1 t2)
-- | /O(n)/. Build a map from an ascending list of distinct elements in linear time.
-- /The precondition is not checked./
fromDistinctAscList :: [(k,a)] -> Map k a
fromDistinctAscList = Map . AVL.asTreeL
-- | /O(n)/. Build a map from an ascending list in linear time.
-- /The precondition (input list is ascending) is not checked./
fromAscList :: Eq k => [(k,a)] -> Map k a
fromAscList = fromAscListWithKey (\_ x _ -> x)
-- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys.
-- /The precondition (input list is ascending) is not checked./
fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
fromAscListWith f = fromAscListWithKey (\_ x y -> f x y)
-- | /O(n)/. Build a map from an ascending list in linear time with a
-- combining function for equal keys.
-- /The precondition (input list is ascending) is not checked./
fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
fromAscListWithKey f = fromDistinctAscList . combineEq
where
-- [combineEq xs] combines equal elements with function [f] in an ordered list [xs]
combineEq xs
= case xs of
[] -> []
[x] -> [x]
(x:xx) -> combineEq' x xx
combineEq' z [] = [z]
combineEq' z@(kz,zz) (x@(kx,xx):xs)
| kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs
| otherwise = z:combineEq' x xs
fromList :: Ord k => [(k,a)] -> Map k a
fromList l = Map (AVL.genAsTree mcmp l)
-- | The union of a list of maps:
-- (@'unions' == 'Prelude.foldl' 'union' 'empty'@).
unions :: Ord k => [Map k a] -> Map k a
unions ts
= foldlStrict union empty ts
-- | The union of a list of maps, with a combining operation:
-- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a
unionsWith f ts
= foldlStrict (unionWith f) empty ts
-- | /O(n+m)/.
-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@.
-- It prefers @t1@ when duplicate keys are encountered,
-- i.e. (@'union' == 'unionWith' 'const'@).
-- The implementation uses the efficient /hedge-union/ algorithm.
-- Hedge-union is more efficient on (bigset `union` smallset)?
union :: Ord k => Map k a -> Map k a -> Map k a
union = unionWith const
-- | /O(n+m)/. Union with a combining function.
unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWith f = unionWithKey (\_ x y -> f x y)
-- | /O(n+m)/.
-- Union with a combining function.
unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
unionWithKey f (Map t1) (Map t2) = Map (AVL.genUnion (mfcmp f) t1 t2)
-- | /O(n+m)/.
-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
isSubmapOf = isSubmapOfBy (==)
{- | /O(n+m)/.
The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if
all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when
applied to their respective values. For example, the following
expressions are all 'True':
> isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
> isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
> isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)])
But the following are all 'False':
> isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)])
> isSubmapOfBy (<) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
> isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)])
-}
isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool
isSubmapOfBy f (Map s) (Map t) = AVL.genIsSubsetOf
(\ (k, a) (k', b) -> case compare k k' of
LT -> LT
GT -> GT
EQ -> if f a b then EQ else LT) s t
-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
-- 'alter' can be used to insert, delete, or update a value in a 'Map'.
-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
alter f k m = case f (lookup k m) of
Just a -> insert k a m
Nothing -> delete k m
-- TODO: add support for this in Data.Tree.AVL
-- | /O(log n)/. Adjust a value at a specific key. When the key is not
-- a member of the map, the original map is returned.
adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
adjust f = adjustWithKey (\_ x -> f x)
-- | /O(log n)/. Adjust a value at a specific key. When the key is not
-- a member of the map, the original map is returned.
adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
adjustWithKey f = updateWithKey (\k x -> Just (f k x))
-- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@
-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
update f = updateWithKey (\_ x -> f x)
-- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the
-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing',
-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound
-- to the new value @y@.
updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
updateWithKey f k (Map t) = let
cc (k', a) = case compare k k' of
LT -> COrdering.Lt
EQ -> COrdering.Eq $ fmap ( \ c -> (k', c)) $ f k' a
GT -> COrdering.Gt
in Map (AVL.genDelMaybe cc t)
-- | /O(log n)/. Lookup and update.
--
-- TODO: only one traversal. This requires fiddling with AVL.Push.
updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a)
updateLookupWithKey f k m = (lookup k m, updateWithKey f k m)
-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
fromListWith f xs
= fromListWithKey (\_k x y -> f x y) xs
-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
fromListWithKey f xs
= foldlStrict ins empty xs
where
ins t (k,x) = insertWithKey f k x t
------------------------------
-- Conversion from/to raw tree.
-- | /O(1)/. Convert a /sorted/ AVL tree to an AVL tree based Set (as provided by this module).
-- This function does not check the input AVL tree is sorted.
{-# INLINE unsafeFromTree #-}
unsafeFromTree :: AVL.AVL (k,a) -> Map k a
unsafeFromTree = Map
-- | /O(1)/. Convert an AVL tree based Set (as provided by this module) to a sorted AVL tree.
{-# INLINE toTree #-}
toTree :: Map k a -> AVL.AVL (k,a)
toTree (Map t) = t
-----------------------------
-- Instances
instance Foldable (Map k) where
foldMap f (Map t) = foldMap (f . snd) t
instance Ord k => Monoid (Map k a) where
mempty = empty
mappend = union
instance Functor (Map k) where
fmap f (Map t) = Map (fmap f' t)
where f' (k,a) = (k,f a)
-------------------------------------------------
-- Utilities
foldlStrict :: (a -> b -> a) -> a -> [b] -> a
foldlStrict f z xs
= case xs of
[] -> z
(x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)