TrieMap-1.0.0: Data/TrieMap.hs

{-# LANGUAGE TypeFamilies, FlexibleContexts, UnboxedTuples #-}

module Data.TrieMap (
	-- * Map type
	TKey,
	TMap,
	-- * Operators
	(!),
	(\\),
	-- * Query
	null,
	size,
	member,
	notMember,
	lookup,
	findWithDefault,
	-- * Construction
	empty,
-- 	showMap,
	singleton,
	-- ** Insertion
	insert,
	insertWith,
	insertWithKey,
	-- ** Delete/Update
	delete,
	adjust,
	adjustWithKey,
	update,
	updateWithKey,
	alter,
	-- * Combine
	-- ** Union
	union,
	unionWith,
	unionWithKey,
	unionMaybeWith,
	unionMaybeWithKey,
	symmetricDifference,
	-- ** Difference
	difference,
	differenceWith,
	differenceWithKey,
	-- ** Intersection
	intersection,
	intersectionWith,
	intersectionWithKey,
	intersectionMaybeWith,
	intersectionMaybeWithKey,
	-- * Traversal
	-- ** Map
	map,
	mapWithKey,
	mapKeys,
	mapKeysWith,
	mapKeysMonotonic,
	-- ** Traverse
	traverseWithKey,
	-- ** Fold
	fold,
	foldWithKey,
	foldrWithKey,
	foldlWithKey,
	-- * Conversion
	elems,
	keys,
	keysSet,
	assocs,
	-- ** Lists
	fromList,
	fromListWith,
	fromListWithKey,
	-- ** Ordered lists
	fromAscList,
	fromAscListWith,
	fromAscListWithKey,
	fromDistinctAscList,
	-- * Filter
	filter,
	filterWithKey,
	partition,
	partitionWithKey,
	mapMaybe,
	mapMaybeWithKey,
	mapEither,
	mapEitherWithKey,
	split,
	splitLookup,
	-- * Submap
	isSubmapOf,
	isSubmapOfBy,
	-- * Min/Max
	findMin,
	findMax,
	deleteMin,
	deleteMax,
	deleteFindMin,
	deleteFindMax,
	updateMin,
	updateMax,
	updateMinWithKey,
	updateMaxWithKey,
	minView,
	maxView,
	minViewWithKey,
	maxViewWithKey
	) where

import Data.TrieMap.Class
import Data.TrieMap.Class.Instances()
import Data.TrieMap.TrieKey
import Data.TrieMap.Applicative
import Data.TrieMap.Rep
import Data.TrieMap.Rep.Instances ()
import Data.TrieMap.Sized

import Control.Applicative hiding (empty)
import Control.Arrow
import Control.Monad
import Data.Maybe hiding (mapMaybe)
import Data.Monoid(Monoid(..), First(..), Last(..))

import GHC.Exts (build)

import Prelude hiding (lookup, foldr, null, map, filter, reverse)

instance (Show k, Show a, TKey k) => Show (TMap k a) where
	show m = "fromList " ++ show (assocs m)

instance (Eq k, TKey k, Eq a) => Eq (TMap k a) where
	m1 == m2 = assocs m1 == assocs m2

instance (Ord k, TKey k, Ord a) => Ord (TMap k a) where
	m1 `compare` m2 = assocs m1 `compare` assocs m2

instance TKey k => Monoid (TMap k a) where
	mempty = empty
	mappend = union

-- | The empty map.
empty :: TKey k => TMap k a
empty = TMap emptyM

-- | A map with a single element.
singleton :: TKey k => k -> a -> TMap k a
singleton k a = insert k a empty

-- | Is the map empty?
null :: TKey k => TMap k a -> Bool
null (TMap m) = nullM m

-- | Lookup the value at a key in the map.
-- 
-- The function will return the corresponding value as @('Just' value)@, or 'Nothing' if the key isn't in the map.
lookup :: TKey k => k -> TMap k a -> Maybe a
lookup k (TMap m) = getElem <$> lookupM (toRep k) m

-- | The expression @('findWithDefault' def k map)@ returns the value at key @k@ or returns default value @def@
-- when the key is not in the map.
findWithDefault :: TKey k => a -> k -> TMap k a -> a
findWithDefault a = fromMaybe a .: lookup

-- | Find the value at a key. Calls 'error' when the element can not be found.
(!) :: TKey k => TMap k a -> k -> a
m ! k = fromMaybe (error "Element not found") (lookup k m)

-- | 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 'TMap'. In short:
-- @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
alter :: TKey k => (Maybe a -> Maybe a) -> k -> TMap k a -> TMap k a
alter f k (TMap m) = TMap (alterM elemSize (fmap Elem . f . fmap getElem) (toRep k) m)

extract :: (TKey k, MonadPlus m) => (k -> a -> m (x, Maybe a)) -> TMap k a -> m (x, TMap k a)
extract f m = unwrapMonad (extractA (WrapMonad .: f) m)

-- | Projects information out of, and modifies or deletes, an individual association pair, 
-- alternating over all associations in the map.
-- 
-- If @assocs m == [(k1, a1), ..., (kn, an)]@, then
-- 
-- > extract f m = let upd k (x, maybeA) = (x, alter (const maybeA) k m) in
-- >   (upd k1 <$> f kn an) <|> ... <|> (upd kn <$> f kn an)
-- 
-- This generalizes a large number of operations, including
-- 
-- > minViewWithKey == getFirst (extract (\ k a -> return ((k, a), Nothing)))
-- > updateMaxWithKey f m == maybe m snd (getLast (extract (\ k a -> return ((), f k a)) m))
-- 
-- In addition,
-- 
-- > getFirst (extract (\ k a -> if p k a then return ((k, a), Nothing) else mzero) m)
-- 
-- finds and removes the first association pair satisfying the predicate |p|.
extractA :: (TKey k, Alternative f) => (k -> a -> f (x, Maybe a)) -> TMap k a -> f (x, TMap k a)
extractA f (TMap m) = fmap TMap <$> extractM elemSize (\ k (Elem a) -> fmap (fmap (fmap Elem)) (f (fromRep k) a)) m

insert :: TKey k => k -> a -> TMap k a -> TMap k a
insert = insertWith const

insertWith :: TKey k => (a -> a -> a) -> k -> a -> TMap k a -> TMap k a
insertWith = insertWithKey . const

insertWithKey :: TKey k => (k -> a -> a -> a) -> k -> a -> TMap k a -> TMap k a
insertWithKey f k a = alter f' k where
	f' = Just . maybe a (f k a)

delete :: TKey k => k -> TMap k a -> TMap k a
delete = alter (const Nothing)

adjust :: TKey k => (a -> a) -> k -> TMap k a -> TMap k a
adjust = adjustWithKey . const

adjustWithKey :: TKey k => (k -> a -> a) -> k -> TMap k a -> TMap k a
adjustWithKey f = updateWithKey (Just .: f)

update :: TKey k => (a -> Maybe a) -> k -> TMap k a -> TMap k a
update f = alter (>>= f)

updateWithKey :: TKey k => (k -> a -> Maybe a) -> k -> TMap k a -> TMap k a
updateWithKey f k = update (f k) k

fold :: TKey k => (a -> b -> b) -> b -> TMap k a -> b
fold = foldWithKey . const

foldWithKey, foldrWithKey :: TKey k => (k -> a -> b -> b) -> b -> TMap k a -> b
foldWithKey f z (TMap m) = foldWithKeyM (\ k (Elem a) -> f (fromRep k) a) m z
foldrWithKey = foldWithKey

foldlWithKey :: TKey k => (b -> k -> a -> b) -> b -> TMap k a -> b
foldlWithKey f z (TMap m) = foldlWithKeyM (\ k z (Elem a) -> f z (fromRep k) a) m z

traverseWithKey :: (TKey k, Applicative f) => (k -> a -> f b) -> TMap k a -> f (TMap k b)
traverseWithKey f (TMap m) = TMap <$> traverseWithKeyM elemSize (\ k (Elem a) -> Elem <$> f (fromRep k) a) m

map :: TKey k => (a -> b) -> TMap k a -> TMap k b
map = fmap

mapWithKey :: TKey k => (k -> a -> b) -> TMap k a -> TMap k b
mapWithKey f (TMap m) = TMap (mapWithKeyM elemSize (\ k (Elem a) -> Elem (f (fromRep k) a)) m)

mapKeys :: (TKey k, TKey k') => (k -> k') -> TMap k a -> TMap k' a
mapKeys f m = fromList [(f k, a) | (k, a) <- assocs m]

mapKeysWith :: (TKey k, TKey k') => (a -> a -> a) -> (k -> k') -> TMap k a -> TMap k' a
mapKeysWith g f m = fromListWith g [(f k, a) | (k, a) <- assocs m]

mapKeysMonotonic :: (TKey k, TKey k') => (k -> k') -> TMap k a -> TMap k' a
mapKeysMonotonic f m = fromDistinctAscList [(f k, a) | (k, a) <- assocs m]

union :: TKey k => TMap k a -> TMap k a -> TMap k a
union = unionWith const

unionWith :: TKey k => (a -> a -> a) -> TMap k a -> TMap k a -> TMap k a
unionWith = unionWithKey . const

unionWithKey :: TKey k => (k -> a -> a -> a) -> TMap k a -> TMap k a -> TMap k a
unionWithKey f = unionMaybeWithKey (\ k a b -> Just (f k a b))

unionMaybeWith :: TKey k => (a -> a -> Maybe a) -> TMap k a -> TMap k a -> TMap k a
unionMaybeWith = unionMaybeWithKey . const

unionMaybeWithKey :: TKey k => (k -> a -> a -> Maybe a) -> TMap k a -> TMap k a -> TMap k a
unionMaybeWithKey f (TMap m1) (TMap m2) = TMap (unionM elemSize f' m1 m2) where
	f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b

symmetricDifference :: TKey k => TMap k a -> TMap k a -> TMap k a
symmetricDifference = unionMaybeWith (\ _ _ -> Nothing)

intersection :: TKey k => TMap k a -> TMap k b -> TMap k a
intersection = intersectionWith const

intersectionWith :: TKey k => (a -> b -> c) -> TMap k a -> TMap k b -> TMap k c
intersectionWith = intersectionWithKey . const

intersectionWithKey :: TKey k => (k -> a -> b -> c) -> TMap k a -> TMap k b -> TMap k c
intersectionWithKey f = intersectionMaybeWithKey (\ k a b -> Just (f k a b))

intersectionMaybeWith :: TKey k => (a -> b -> Maybe c) -> TMap k a -> TMap k b -> TMap k c
intersectionMaybeWith = intersectionMaybeWithKey . const

intersectionMaybeWithKey :: TKey k => (k -> a -> b -> Maybe c) -> TMap k a -> TMap k b -> TMap k c
intersectionMaybeWithKey f (TMap m1) (TMap m2) = TMap (isectM elemSize f' m1 m2) where
	f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b

difference, (\\) :: TKey k => TMap k a -> TMap k b -> TMap k a
difference = differenceWith (\ _ _ -> Nothing)

(\\) = difference

differenceWith :: TKey k => (a -> b -> Maybe a) -> TMap k a -> TMap k b -> TMap k a
differenceWith = differenceWithKey . const

differenceWithKey :: TKey k => (k -> a -> b -> Maybe a) -> TMap k a -> TMap k b -> TMap k a
differenceWithKey f (TMap m1) (TMap m2) = TMap (diffM elemSize f' m1 m2) where
	f' k (Elem a) (Elem b) = Elem <$> f (fromRep k) a b

minView, maxView :: TKey k => TMap k a -> Maybe (a, TMap k a)
minView m = first snd <$> minViewWithKey m
maxView m = first snd <$> maxViewWithKey m

findMin, findMax :: TKey k => TMap k a -> (k, a)
findMin = maybe (error "empty map has no minimal element") fst . minViewWithKey
findMax = maybe (error "empty map has no maximal element") fst . maxViewWithKey

deleteMin, deleteMax :: TKey k => TMap k a -> TMap k a
deleteMin m = maybe m snd (minViewWithKey m)
deleteMax m = maybe m snd (maxViewWithKey m)

updateMin, updateMax :: TKey k => (a -> Maybe a) -> TMap k a -> TMap k a
updateMin = updateMinWithKey . const
updateMax = updateMaxWithKey . const

updateMinWithKey, updateMaxWithKey :: TKey k => (k -> a -> Maybe a) -> TMap k a -> TMap k a
updateMinWithKey f m = maybe m snd (getFirst (extract (\ k a -> return ((), f k a)) m))
updateMaxWithKey f m = maybe m snd (getLast (extract (\ k a -> return ((), f k a)) m))

deleteFindMin, deleteFindMax :: TKey k => TMap k a -> ((k, a), TMap k a)
deleteFindMin m = fromMaybe (error "Cannot return the minimal element of an empty map") (minViewWithKey m)
deleteFindMax m = fromMaybe (error "Cannot return the maximal element of an empty map") (maxViewWithKey m)

minViewWithKey, maxViewWithKey :: TKey k => TMap k a -> Maybe ((k, a), TMap k a)
minViewWithKey = getFirst . extract (\ k a -> return ((k, a), Nothing))
maxViewWithKey = getLast . extract (\ k a -> return ((k, a), Nothing))

elems :: TKey k => TMap k a -> [a]
elems = fmap snd . assocs

keys :: TKey k => TMap k a -> [k]
keys = fmap fst . assocs

assocs :: TKey k => TMap k a -> [(k, a)]
assocs m = build (\ c n -> foldWithKey (curry c) n m)

mapEither :: TKey k => (a -> Either b c) -> TMap k a -> (TMap k b, TMap k c)
mapEither = mapEitherWithKey . const

mapEitherWithKey :: TKey k => (k -> a -> Either b c) -> TMap k a -> (TMap k b, TMap k c)
mapEitherWithKey f (TMap m) = case mapEitherM elemSize elemSize f' m of
	(# mL, mR #) -> (TMap mL, TMap mR) 
	where	f' k (Elem a) = case f (fromRep k) a of
			Left b	-> (# Just (Elem b), Nothing #)
			Right c	-> (# Nothing, Just (Elem c) #)

mapMaybe :: TKey k => (a -> Maybe b) -> TMap k a -> TMap k b
mapMaybe = mapMaybeWithKey . const

mapMaybeWithKey :: TKey k => (k -> a -> Maybe b) -> TMap k a -> TMap k b
mapMaybeWithKey f (TMap m) = TMap (mapMaybeM elemSize (\ k (Elem a) -> Elem <$> f (fromRep k) a) m)

partition :: TKey k => (a -> Bool) -> TMap k a -> (TMap k a, TMap k a)
partition = partitionWithKey . const

partitionWithKey :: TKey k => (k -> a -> Bool) -> TMap k a -> (TMap k a, TMap k a)
partitionWithKey p = mapEitherWithKey (\ k a -> (if p k a then Left else Right) a)

filter :: TKey k => (a -> Bool) -> TMap k a -> TMap k a
filter = filterWithKey . const

filterWithKey :: TKey k => (k -> a -> Bool) -> TMap k a -> TMap k a
filterWithKey p = mapMaybeWithKey (\ k a -> if p k a then Just a else Nothing)

split :: TKey k => k -> TMap k a -> (TMap k a, TMap k a)
split k m = case splitLookup k m of
	(mL, _, mR) -> (mL, mR)

splitLookup :: TKey k => k -> TMap k a -> (TMap k a, Maybe a, TMap k a)
splitLookup k (TMap m) = case splitLookupM elemSize f (toRep k) m of
	(# mL, x, mR #) -> (TMap mL, x, TMap mR) 
	where	f (Elem x) = (# Nothing, Just x, Nothing #)

isSubmapOf :: (TKey k, Eq a) => TMap k a -> TMap k a -> Bool
isSubmapOf = isSubmapOfBy (==)

isSubmapOfBy :: TKey k => (a -> b -> Bool) -> TMap k a -> TMap k b -> Bool
isSubmapOfBy (<=) (TMap m1) (TMap m2) = isSubmapM (<<=) m1 m2 where
	Elem a <<= Elem b = a <= b

fromList, fromAscList :: TKey k => [(k, a)] -> TMap k a
fromList = fromListWith const
fromAscList = fromAscListWith const

fromListWith, fromAscListWith :: TKey k => (a -> a -> a) -> [(k, a)] -> TMap k a
fromListWith = fromListWithKey . const
fromAscListWith = fromAscListWithKey . const

fromListWithKey, fromAscListWithKey :: TKey k => (k -> a -> a -> a) -> [(k, a)] -> TMap k a
fromListWithKey f xs = TMap (fromListM elemSize (\ k (Elem a) (Elem b) -> Elem (f (fromRep k) a b)) [(toRep k, Elem a) | (k, a) <- xs])
fromAscListWithKey f xs = TMap (fromAscListM elemSize (\ k (Elem a) (Elem b) -> Elem (f (fromRep k) a b)) [(toRep k, Elem a) | (k, a) <- xs])

fromDistinctAscList :: TKey k => [(k, a)] -> TMap k a
fromDistinctAscList xs = TMap (fromDistAscListM elemSize [(toRep k, Elem a) | (k, a) <- xs])

size :: TKey k => TMap k a -> Int
size (TMap m) = sizeM elemSize m

member :: TKey k => k -> TMap k a -> Bool
member = isJust .: lookup

notMember :: TKey k => k -> TMap k a -> Bool
notMember = not .: member

keysSet :: TKey k => TMap k a -> TSet k
keysSet m = TSet (() <$ m)