----------------------------------------------------------------------------- -- | -- 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)