diff --git a/dependent-map.cabal b/dependent-map.cabal index 29f15fb..455c0dd 100644 --- a/dependent-map.cabal +++ b/dependent-map.cabal @@ -1,5 +1,5 @@ name: dependent-map -version: 0.2.5.1 +version: 0.2.6.0 stability: provisional cabal-version: >= 1.6 @@ -34,7 +34,10 @@ Library ghc-options: -fwarn-unused-imports -fwarn-unused-binds exposed-modules: Data.Dependent.Map, Data.Dependent.Map.Lens, - Data.Dependent.Map.Internal + Data.Dependent.Map.NonEmpty, + Data.Dependent.Map.NonEmpty.Lens, + Data.Dependent.Map.Internal, + Data.Dependent.Map.Internal2 other-modules: Data.Dependent.Map.PtrEquality if impl(ghc < 7.8) other-modules: Data.Dependent.Map.Typeable @@ -42,6 +45,7 @@ Library build-depends: semigroups build-depends: base >= 3 && < 5, containers, - dependent-sum >= 0.3.2 + dependent-sum >= 0.3.2, + constraints-extras >= 0.2.3.0 && < 0.3 if impl(ghc >= 7.2) && impl(ghc < 7.8) ghc-options: -trust base -trust dependent-sum diff --git a/src/Data/Dependent/Map.hs b/src/Data/Dependent/Map.hs index 5160ae1..d16949e 100644 --- a/src/Data/Dependent/Map.hs +++ b/src/Data/Dependent/Map.hs @@ -8,12 +8,14 @@ #endif #if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708 {-# LANGUAGE PolyKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} #endif module Data.Dependent.Map ( DMap , DSum(..), Some(..) , GCompare(..), GOrdering(..) - + -- * Operators , (!), (\\) @@ -24,7 +26,7 @@ module Data.Dependent.Map , notMember , lookup , findWithDefault - + -- * Construction , empty , singleton @@ -37,10 +39,11 @@ module Data.Dependent.Map , insertWithKey' , insertLookupWithKey , insertLookupWithKey' - + -- ** Delete\/Update , delete , adjust + , adjustF , adjustWithKey , adjustWithKey' , update @@ -52,7 +55,7 @@ module Data.Dependent.Map -- * Combine -- ** Union - , union + , union , unionWithKey , unions , unionsWithKey @@ -60,9 +63,9 @@ module Data.Dependent.Map -- ** Difference , difference , differenceWithKey - + -- ** Intersection - , intersection + , intersection , intersectionWithKey -- * Traversal @@ -84,7 +87,7 @@ module Data.Dependent.Map -- * Conversion , keys , assocs - + -- ** Lists , toList , fromList @@ -97,7 +100,7 @@ module Data.Dependent.Map , fromAscListWithKey , fromDistinctAscList - -- * Filter + -- * Filter , filter , filterWithKey , partitionWithKey @@ -106,14 +109,14 @@ module Data.Dependent.Map , mapMaybeWithKey , mapEitherWithKey - , split - , splitLookup + , split + , splitLookup -- * Submap , isSubmapOf, isSubmapOfBy , isProperSubmapOf, isProperSubmapOfBy - -- * Indexed + -- * Indexed , lookupIndex , findIndex , elemAt @@ -133,7 +136,7 @@ module Data.Dependent.Map , updateMaxWithKey , minViewWithKey , maxViewWithKey - + -- * Debugging , showTree , showTreeWith @@ -146,12 +149,15 @@ import qualified Prelude import Control.Applicative (Applicative(..), (<$>)) #endif import Data.Dependent.Map.Internal +import Data.Dependent.Map.Internal2 #if !MIN_VERSION_base(4,7,0) import Data.Dependent.Map.Typeable ({- instance Typeable ... -}) #endif import Data.Dependent.Sum +import Data.Constraint.Extras import Data.GADT.Compare +import Data.GADT.Show import Data.Maybe (isJust) #if !MIN_VERSION_base(4,8,0) import Data.Monoid @@ -159,7 +165,6 @@ import Data.Monoid import Data.Semigroup import Data.Some import Text.Read -import Data.Dependent.Map.PtrEquality instance (GCompare k) => Monoid (DMap k f) where mempty = empty @@ -188,27 +193,42 @@ infixl 9 !,\\ -- m1 \\ m2 = difference m1 m2 -- #if __GLASGOW_HASKELL__ --- +-- -- {-------------------------------------------------------------------- --- A Data instance +-- A Data instance -- --------------------------------------------------------------------} --- +-- -- -- This instance preserves data abstraction at the cost of inefficiency. -- -- We omit reflection services for the sake of data abstraction. --- +-- -- instance (Data k, Data a, GCompare k) => Data (DMap k) where -- gfoldl f z m = z fromList `f` toList m -- toConstr _ = error "toConstr" -- gunfold _ _ = error "gunfold" -- dataTypeOf _ = mkNoRepType "Data.Map.Map" -- dataCast2 f = gcast2 f --- +-- -- #endif +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 +singleton :: k v -> f v -> DMap k f +singleton = singletonE + {-------------------------------------------------------------------- Query --------------------------------------------------------------------} +-- | /O(1)/. The number of elements in the map. +size :: DMap k f -> Int +size = sizeE + -- | /O(log n)/. Is the key a member of the map? See also 'notMember'. member :: GCompare k => k a -> DMap k f -> Bool member k = isJust . lookup k @@ -217,6 +237,13 @@ member k = isJust . lookup k notMember :: GCompare k => k v -> DMap k f -> Bool notMember k m = not (member k m) +-- | /O(log n)/. 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 :: forall k f v. GCompare k => k v -> DMap k f -> Maybe (f v) +lookup = lookupE + -- | /O(log n)/. Find the value at a key. -- Calls 'error' when the element can not be found. -- Consider using 'lookup' when elements may not be present. @@ -242,44 +269,16 @@ findWithDefault def k m = case lookup k m of -- replaced with the supplied value. 'insert' is equivalent to -- @'insertWith' 'const'@. insert :: forall k f v. GCompare k => k v -> f v -> DMap k f -> DMap k f -insert kx x = kx `seq` go - where - go :: DMap k f -> DMap k f - go Tip = singleton kx x - go t@(Bin sz ky y l r) = case gcompare kx ky of - GLT -> let !l' = go l - in if l' `ptrEq` l - then t - else balance ky y l' r - GGT -> let !r' = go r - in if r' `ptrEq` r - then t - else balance ky y l r' - GEQ - | kx `ptrEq` ky && x `ptrEq` y -> t - | otherwise -> Bin sz kx x l r +insert kx x = Bin' . makeInsert kx x -- | /O(log n)/. Insert a new key and value in the map if the key -- is not already present. If the key is already present, @insertR@ -- does nothing. insertR :: forall k f v. GCompare k => k v -> f v -> DMap k f -> DMap k f -insertR kx x = kx `seq` go - where - go :: DMap k f -> DMap k f - go Tip = singleton kx x - go t@(Bin sz ky y l r) = case gcompare kx ky of - GLT -> let !l' = go l - in if l' `ptrEq` l - then t - else balance ky y l' r - GGT -> let !r' = go r - in if r' `ptrEq` r - then t - else balance ky y l r' - GEQ -> t +insertR kx x = Bin' . makeInsertR kx x -- | /O(log n)/. Insert with a function, combining new value and old value. --- @'insertWith' f key value mp@ +-- @'insertWith' f key value mp@ -- will insert the entry @key :=> value@ into @mp@ if key does -- not exist in the map. If the key does exist, the function will -- insert the entry @key :=> f new_value old_value@. @@ -292,33 +291,17 @@ insertWith' :: GCompare k => (f v -> f v -> f v) -> k v -> f v -> DMap k f -> DM insertWith' f = insertWithKey' (\_ x' y' -> f x' y') -- | /O(log n)/. Insert with a function, combining key, new value and old value. --- @'insertWithKey' f key value mp@ +-- @'insertWithKey' f key value mp@ -- will insert the entry @key :=> value@ into @mp@ if key does -- not exist in the map. If the key does exist, the function will -- insert the entry @key :=> f key new_value old_value@. -- Note that the key passed to f is the same key passed to 'insertWithKey'. insertWithKey :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> DMap k f -> DMap k f -insertWithKey f kx x = kx `seq` go - where - go :: DMap k f -> DMap k f - go Tip = singleton kx x - go (Bin sy ky y l r) = - case gcompare kx ky of - GLT -> balance ky y (go l) r - GGT -> balance ky y l (go r) - GEQ -> Bin sy kx (f kx x y) l r +insertWithKey f kx x = Bin' . makeInsertWithKey f kx x -- | Same as 'insertWithKey', but the combining function is applied strictly. insertWithKey' :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> DMap k f -> DMap k f -insertWithKey' f kx x = kx `seq` go - where - go :: DMap k f -> DMap k f - go Tip = singleton kx $! x - go (Bin sy ky y l r) = - case gcompare kx ky of - GLT -> balance ky y (go l) r - GGT -> balance ky y l (go r) - GEQ -> let x' = f kx x y in seq x' (Bin sy kx x' l r) +insertWithKey' f kx x = Bin' . makeInsertWithKey' f kx x -- | /O(log n)/. Combines insert operation with old value retrieval. -- The expression (@'insertLookupWithKey' f k x map@) @@ -326,32 +309,12 @@ insertWithKey' f kx x = kx `seq` go -- and the second element equal to (@'insertWithKey' f k x map@). insertLookupWithKey :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> DMap k f -> (Maybe (f v), DMap k f) -insertLookupWithKey f kx x = kx `seq` go - where - go :: DMap k f -> (Maybe (f v), DMap k f) - go Tip = (Nothing, singleton kx x) - go (Bin sy ky y l r) = - case gcompare kx ky of - GLT -> let (found, l') = go l - in (found, balance ky y l' r) - GGT -> let (found, r') = go r - in (found, balance ky y l r') - GEQ -> (Just y, Bin sy kx (f kx x y) l r) +insertLookupWithKey f kx x = fmap Bin' . makeInsertLookupWithKey f kx x -- | /O(log n)/. A strict version of 'insertLookupWithKey'. insertLookupWithKey' :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> DMap k f -> (Maybe (f v), DMap k f) -insertLookupWithKey' f kx x = kx `seq` go - where - go :: DMap k f -> (Maybe (f v), DMap k f) - go Tip = x `seq` (Nothing, singleton kx x) - go (Bin sy ky y l r) = - case gcompare kx ky of - GLT -> let (found, l') = go l - in (found, balance ky y l' r) - GGT -> let (found, r') = go r - in (found, balance ky y l r') - GEQ -> let x' = f kx x y in x' `seq` (Just y, Bin sy kx x' l r) +insertLookupWithKey' f kx x = fmap Bin' . makeInsertLookupWithKey' f kx x {-------------------------------------------------------------------- Deletion @@ -361,15 +324,7 @@ insertLookupWithKey' f kx x = kx `seq` go -- | /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 :: forall k f v. GCompare k => k v -> DMap k f -> DMap k f -delete k = k `seq` go - where - go :: DMap k f -> DMap k f - go Tip = Tip - go (Bin _ kx x l r) = - case gcompare k kx of - GLT -> balance kx x (go l) r - GGT -> balance kx x l (go r) - GEQ -> glue l r +delete k = fst $ makeDelete k -- | /O(log n)/. Update a value at a specific key with the result of the provided function. -- When the key is not @@ -377,30 +332,23 @@ delete k = k `seq` go adjust :: GCompare k => (f v -> f v) -> k v -> DMap k f -> DMap k f adjust f = adjustWithKey (\_ x -> f x) +-- | Works the same as 'adjust' except the new value is return in some 'Applicative' @f@. +adjustF + :: forall k f v g + . (GCompare k, Applicative f) + => k v + -> (g v -> f (g v)) + -> DMap k g -> f (DMap k g) +adjustF k f = fst $ makeAdjustF f k + -- | /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 :: GCompare k => (k v -> f v -> f v) -> k v -> DMap k f -> DMap k f -adjustWithKey f0 !k0 = go f0 k0 - where - go :: GCompare k => (k v -> f v -> f v) -> k v -> DMap k f -> DMap k f - go _f _k Tip = Tip - go f k (Bin sx kx x l r) = - case gcompare k kx of - GLT -> Bin sx kx x (go f k l) r - GGT -> Bin sx kx x l (go f k r) - GEQ -> Bin sx kx (f kx x) l r +adjustWithKey f k = fst $ makeAdjustWithKey f k -- | /O(log n)/. A strict version of 'adjustWithKey'. adjustWithKey' :: GCompare k => (k v -> f v -> f v) -> k v -> DMap k f -> DMap k f -adjustWithKey' f0 !k0 = go f0 k0 - where - go :: GCompare k => (k v -> f v -> f v) -> k v -> DMap k f -> DMap k f - go _f _k Tip = Tip - go f k (Bin sx kx x l r) = - case gcompare k kx of - GLT -> Bin sx kx x (go f k l) r - GGT -> Bin sx kx x l (go f k r) - GEQ -> let !x' = f kx x in Bin sx kx x' l r +adjustWithKey' f k = fst $ makeAdjustWithKey' f k -- | /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 @@ -413,64 +361,24 @@ update f = updateWithKey (\_ x -> f x) -- the element is deleted. If it is (@'Just' y@), the key @k@ is bound -- to the new value @y@. updateWithKey :: forall k f v. GCompare k => (k v -> f v -> Maybe (f v)) -> k v -> DMap k f -> DMap k f -updateWithKey f k = k `seq` go - where - go :: DMap k f -> DMap k f - go Tip = Tip - go (Bin sx kx x l r) = - case gcompare k kx of - GLT -> balance kx x (go l) r - GGT -> balance kx x l (go r) - GEQ -> case f kx x of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r +updateWithKey f k = fst $ makeUpdateWithKey f k -- | /O(log n)/. Lookup and update. See also 'updateWithKey'. -- The function returns changed value, if it is updated. --- Returns the original key value if the map entry is deleted. +-- Returns the original key value if the map entry is deleted. updateLookupWithKey :: forall k f v. GCompare k => (k v -> f v -> Maybe (f v)) -> k v -> DMap k f -> (Maybe (f v), DMap k f) -updateLookupWithKey f k = k `seq` go - where - go :: DMap k f -> (Maybe (f v), DMap k f) - go Tip = (Nothing,Tip) - go (Bin sx kx x l r) = - case gcompare k kx of - GLT -> let (found,l') = go l in (found,balance kx x l' r) - GGT -> let (found,r') = go r in (found,balance kx x l r') - GEQ -> case f kx x of - Just x' -> (Just x',Bin sx kx x' l r) - Nothing -> (Just x,glue l r) +updateLookupWithKey f k = fst $ makeUpdateLookupWithKey f k -- | /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 :: forall k f v. GCompare k => (Maybe (f v) -> Maybe (f v)) -> k v -> DMap k f -> DMap k f -alter f k = k `seq` go - where - go :: DMap k f -> DMap k f - go Tip = case f Nothing of - Nothing -> Tip - Just x -> singleton k x - - go (Bin sx kx x l r) = case gcompare k kx of - GLT -> balance kx x (go l) r - GGT -> balance kx x l (go r) - GEQ -> case f (Just x) of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r +alter k f = fst $ makeAlter k f -- | Works the same as 'alter' except the new value is return in some 'Functor' @f@. -- In short : @(\v' -> alter (const v') k dm) <$> f (lookup k dm)@ alterF :: forall k f v g. (GCompare k, Functor f) => k v -> (Maybe (g v) -> f (Maybe (g v))) -> DMap k g -> f (DMap k g) -alterF k f = go - where - go :: DMap k g -> f (DMap k g) - go Tip = maybe Tip (singleton k) <$> f Nothing - - go (Bin sx kx x l r) = case gcompare k kx of - GLT -> (\l' -> balance kx x l' r) <$> go l - GGT -> (\r' -> balance kx x l r') <$> go r - GEQ -> maybe (glue l r) (\x' -> Bin sx kx x' l r) <$> f (Just x) +alterF k f = fst $ makeAlterF k f {-------------------------------------------------------------------- Indexing @@ -488,42 +396,17 @@ findIndex k t -- | /O(log n)/. Lookup the /index/ of a key. The index is a number from -- /0/ up to, but not including, the 'size' of the map. lookupIndex :: forall k f v. GCompare k => k v -> DMap k f -> Maybe Int -lookupIndex k = k `seq` go 0 - where - go :: Int -> DMap k f -> Maybe Int - go !idx Tip = idx `seq` Nothing - go !idx (Bin _ kx _ l r) - = case gcompare k kx of - GLT -> go idx l - GGT -> go (idx + size l + 1) r - GEQ -> Just (idx + size l) +lookupIndex k = fst $ makeLookupIndex k -- | /O(log n)/. Retrieve an element by /index/. Calls 'error' when an -- invalid index is used. elemAt :: Int -> DMap k f -> DSum k f -elemAt _ Tip = error "Map.elemAt: index out of range" -elemAt i (Bin _ kx x l r) - = case compare i sizeL of - LT -> elemAt i l - GT -> elemAt (i-sizeL-1) r - EQ -> kx :=> x - where - sizeL = size l +elemAt = fst makElemAt -- | /O(log n)/. Update the element at /index/. Does nothing when an -- invalid index is used. updateAt :: (forall v. k v -> f v -> Maybe (f v)) -> Int -> DMap k f -> DMap k f -updateAt f i0 t = i0 `seq` go i0 t - where - go _ Tip = Tip - go i (Bin sx kx x l r) = case compare i sizeL of - LT -> balance kx x (go i l) r - GT -> balance kx x l (go (i-sizeL-1) r) - EQ -> case f kx x of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r - where - sizeL = size l +updateAt f i0 = fst $ makeUpdateAt f i0 -- | /O(log n)/. Delete the element at /index/. -- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@). @@ -543,13 +426,7 @@ findMin m = case lookupMin m of Nothing -> error "Map.findMin: empty map has no minimal element" lookupMin :: DMap k f -> Maybe (DSum k f) -lookupMin m = case m of - Tip -> Nothing - Bin _ kx x l _ -> Just $! go kx x l - where - go :: k v -> f v -> DMap k f -> DSum k f - go kx x Tip = kx :=> x - go _ _ (Bin _ kx x l _) = go kx x l +lookupMin = fst makeLookupMin -- | /O(log n)/. The maximal key of the map. Calls 'error' is the map is empty. findMax :: DMap k f -> DSum k f @@ -558,48 +435,54 @@ findMax m = case lookupMax m of Nothing -> error "Map.findMax: empty map has no maximal element" lookupMax :: DMap k f -> Maybe (DSum k f) -lookupMax m = case m of - Tip -> Nothing - Bin _ kx x _ r -> Just $! go kx x r - where - go :: k v -> f v -> DMap k f -> DSum k f - go kx x Tip = kx :=> x - go _ _ (Bin _ kx x _ r) = go kx x r +lookupMax = fst makeLookupMax -- | /O(log n)/. Delete the minimal key. Returns an empty map if the map is empty. deleteMin :: DMap k f -> DMap k f -deleteMin (Bin _ _ _ Tip r) = r -deleteMin (Bin _ kx x l r) = balance kx x (deleteMin l) r -deleteMin Tip = Tip +deleteMin = fst makeDeleteMin -- | /O(log n)/. Delete the maximal key. Returns an empty map if the map is empty. deleteMax :: DMap k f -> DMap k f -deleteMax (Bin _ _ _ l Tip) = l -deleteMax (Bin _ kx x l r) = balance kx x l (deleteMax r) -deleteMax Tip = Tip +deleteMax = fst makeDeleteMax + +-- | /O(log n)/. Delete and find the minimal element. +-- +-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) +-- > deleteFindMin Error: can not return the minimal element of an empty map +deleteFindMin :: DMap k f -> (DSum k f, DMap k f) +deleteFindMin t = case minViewWithKey t of + Nothing -> error "Error: can not return the minimal element of an empty map" + Just x -> x + +-- | /O(log n)/. Delete and find the maximal element. +-- +-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) +-- > deleteFindMax empty Error: can not return the maximal element of an empty map +deleteFindMax :: DMap k f -> (DSum k f, DMap k f) +deleteFindMax t = case maxViewWithKey t of + Nothing -> error "Error: can not return the maximal element of an empty map" + Just x -> x -- | /O(log n)/. Update the value at the minimal key. updateMinWithKey :: (forall v. k v -> f v -> Maybe (f v)) -> DMap k f -> DMap k f -updateMinWithKey f = go - where - go (Bin sx kx x Tip r) = case f kx x of - Nothing -> r - Just x' -> Bin sx kx x' Tip r - go (Bin _ kx x l r) = balance kx x (go l) r - go Tip = Tip +updateMinWithKey f = fst $ makeUpdateMinWithKey f -- | /O(log n)/. Update the value at the maximal key. updateMaxWithKey :: (forall v. k v -> f v -> Maybe (f v)) -> DMap k f -> DMap k f -updateMaxWithKey f = go - where - go (Bin sx kx x l Tip) = case f kx x of - Nothing -> l - Just x' -> Bin sx kx x' l Tip - go (Bin _ kx x l r) = balance kx x l (go r) - go Tip = Tip +updateMaxWithKey f = fst $ makeUpdateMaxWithKey f + +-- | /O(log n)/. Retrieves the minimal (key :=> value) entry of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +minViewWithKey :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f) +minViewWithKey = minViewWithKeyE + +-- | /O(log n)/. Retrieves the maximal (key :=> value) entry of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +maxViewWithKey :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f) +maxViewWithKey = maxViewWithKeyE {-------------------------------------------------------------------- - Union. + Union. --------------------------------------------------------------------} -- | The union of a list of maps: @@ -615,20 +498,11 @@ unionsWithKey f ts = foldlStrict (unionWithKey f) empty ts -- | /O(m*log(n\/m + 1)), m <= n/. --- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@. +-- 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'@). union :: GCompare k => DMap k f -> DMap k f -> DMap k f -union t1 Tip = t1 -union t1 (Bin _ kx x Tip Tip) = insertR kx x t1 -union Tip t2 = t2 -union (Bin _ kx x Tip Tip) t2 = insert kx x t2 -union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of - (l2, r2) - | l1 `ptrEq` l1l2 && r1 `ptrEq` r1r2 -> t1 - | otherwise -> combine k1 x1 l1l2 r1r2 - where !l1l2 = l1 `union` l2 - !r1r2 = r1 `union` r2 +union t1 t2 = fst makeUnion t1 t2 {-------------------------------------------------------------------- Union with a combining function @@ -637,47 +511,23 @@ union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of -- | /O(n+m)/. -- Union with a combining function. unionWithKey :: GCompare k => (forall v. k v -> f v -> f v -> f v) -> DMap k f -> DMap k f -> DMap k f -unionWithKey _ t1 Tip = t1 -unionWithKey _ Tip t2 = t2 -unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of - (l2, mx2, r2) -> case mx2 of - Nothing -> combine k1 x1 l1l2 r1r2 - Just x2 -> combine k1 (f k1 x1 x2) l1l2 r1r2 - where !l1l2 = unionWithKey f l1 l2 - !r1r2 = unionWithKey f r1 r2 +unionWithKey f = fst $ makeUnionWithKey f {-------------------------------------------------------------------- Difference --------------------------------------------------------------------} --- | /O(m * log (n\/m + 1)), m <= n/. Difference of two maps. +-- | /O(m * log (n\/m + 1)), m <= n/. Difference of two maps. -- Return elements of the first map not existing in the second map. difference :: GCompare k => DMap k f -> DMap k g -> DMap k f -difference Tip _ = Tip -difference t1 Tip = t1 -difference t1 (Bin _ k2 _x2 l2 r2) = case split k2 t1 of - (l1, r1) - | size t1 == size l1l2 + size r1r2 -> t1 - | otherwise -> merge l1l2 r1r2 - where - !l1l2 = l1 `difference` l2 - !r1r2 = r1 `difference` r2 +difference = fst makeDifference -- | /O(n+m)/. Difference with a combining function. When two equal keys are -- encountered, the combining function is applied to the key and both values. -- If it returns 'Nothing', the element is discarded (proper set difference). If --- it returns (@'Just' y@), the element is updated with a new value @y@. +-- it returns (@'Just' y@), the element is updated with a new value @y@. differenceWithKey :: GCompare k => (forall v. k v -> f v -> g v -> Maybe (f v)) -> DMap k f -> DMap k g -> DMap k f -differenceWithKey _ Tip _ = Tip -differenceWithKey _ t1 Tip = t1 -differenceWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of - (l2, mx2, r2) -> case mx2 of - Nothing -> combine k1 x1 l1l2 r1r2 - Just x2 -> case f k1 x1 x2 of - Nothing -> merge l1l2 r1r2 - Just x1x2 -> combine k1 x1x2 l1l2 r1r2 - where !l1l2 = differenceWithKey f l1 l2 - !r1r2 = differenceWithKey f r1 r2 +differenceWithKey f = fst $ makeDifferenceWithKey f {-------------------------------------------------------------------- Intersection @@ -687,29 +537,11 @@ differenceWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of -- Return data in the first map for the keys existing in both maps. -- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). intersection :: GCompare k => DMap k f -> DMap k f -> DMap k f -intersection Tip _ = Tip -intersection _ Tip = Tip -intersection t1@(Bin s1 k1 x1 l1 r1) t2 = - let !(l2, found, r2) = splitMember k1 t2 - !l1l2 = intersection l1 l2 - !r1r2 = intersection r1 r2 - in if found - then if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 - then t1 - else combine k1 x1 l1l2 r1r2 - else merge l1l2 r1r2 +intersection = fst makeIntersection -- | /O(m * log (n\/m + 1), m <= n/. Intersection with a combining function. intersectionWithKey :: GCompare k => (forall v. k v -> f v -> g v -> h v) -> DMap k f -> DMap k g -> DMap k h -intersectionWithKey _ Tip _ = Tip -intersectionWithKey _ _ Tip = Tip -intersectionWithKey f (Bin s1 k1 x1 l1 r1) t2 = - let !(l2, found, r2) = splitLookup k1 t2 - !l1l2 = intersectionWithKey f l1 l2 - !r1r2 = intersectionWithKey f r1 r2 - in case found of - Nothing -> merge l1l2 r1r2 - Just x2 -> combine k1 (f k1 x1 x2) l1l2 r1r2 +intersectionWithKey f = fst $ makeIntersectionWithKey f {-------------------------------------------------------------------- Submap @@ -717,8 +549,11 @@ intersectionWithKey f (Bin s1 k1 x1 l1 r1) t2 = -- | /O(n+m)/. -- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' 'eqTagged')@). -- -isSubmapOf :: (GCompare k, EqTag k f) => DMap k f -> DMap k f -> Bool -isSubmapOf m1 m2 = isSubmapOfBy eqTagged m1 m2 +isSubmapOf + :: forall k f + . (GCompare k, Has' Eq k f) + => DMap k f -> DMap k f -> Bool +isSubmapOf m1 m2 = isSubmapOfBy (\k _ x0 x1 -> has' @Eq @f k (x0 == x1)) m1 m2 {- | /O(n+m)/. The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if @@ -730,26 +565,22 @@ isSubmapOfBy f t1 t2 = (size t1 <= size t2) && (submap' f t1 t2) submap' :: GCompare k => (forall v. k v -> k v -> f v -> g v -> Bool) -> DMap k f -> DMap k g -> Bool -submap' _ Tip _ = True -submap' _ _ Tip = False -submap' f (Bin _ kx x l r) t - = case found of - Nothing -> False - Just (ky, y) -> f kx ky x y && submap' f l lt && submap' f r gt - where - (lt,found,gt) = splitLookupWithKey kx t +submap' f t1 t2 = fst (makeSubmap' f) t1 t2 --- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). +-- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' 'eqTagged'@). -isProperSubmapOf :: (GCompare k, EqTag k f) => DMap k f -> DMap k f -> Bool +isProperSubmapOf + :: forall k f + . (GCompare k, Has' Eq k f) + => DMap k f -> DMap k f -> Bool isProperSubmapOf m1 m2 - = isProperSubmapOfBy eqTagged m1 m2 + = isProperSubmapOfBy (\k _ x0 x1 -> has' @Eq @f k (x0 == x1)) m1 m2 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when @m1@ and @m2@ are not equal, all keys in @m1@ are in @m2@, and when @f@ returns 'True' when - applied to their respective keys and values. + applied to their respective keys and values. -} isProperSubmapOfBy :: GCompare k => (forall v. k v -> k v -> f v -> g v -> Bool) -> DMap k f -> DMap k g -> Bool isProperSubmapOfBy f t1 t2 @@ -761,31 +592,13 @@ isProperSubmapOfBy f t1 t2 -- | /O(n)/. Filter all keys\/values that satisfy the predicate. filterWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> DMap k f -filterWithKey p = go - where - go Tip = Tip - go t@(Bin _ kx x l r) - | p kx x = if l' `ptrEq` l && r' `ptrEq` r - then t - else combine kx x l' r' - | otherwise = merge l' r' - where !l' = go l - !r' = go r +filterWithKey = makeFilterWithKey -- | /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. See also 'split'. partitionWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> (DMap k f, DMap k f) -partitionWithKey p0 m0 = toPair (go p0 m0) - where - go :: GCompare k => (forall v. k v -> f v -> Bool) -> DMap k f -> (DMap k f :*: DMap k f) - go _ Tip = (Tip :*: Tip) - go p (Bin _ kx x l r) - | p kx x = (combine kx x l1 r1 :*: merge l2 r2) - | otherwise = (merge l1 r1 :*: combine kx x l2 r2) - where - (l1 :*: l2) = go p l - (r1 :*: r2) = go p r +partitionWithKey = makePartitionWithKey -- | /O(n)/. Map values and collect the 'Just' results. mapMaybe :: GCompare k => (forall v. f v -> Maybe (g v)) -> DMap k f -> DMap k g @@ -793,28 +606,12 @@ mapMaybe f = mapMaybeWithKey (const f) -- | /O(n)/. Map keys\/values and collect the 'Just' results. mapMaybeWithKey :: GCompare k => (forall v. k v -> f v -> Maybe (g v)) -> DMap k f -> DMap k g -mapMaybeWithKey f = go - where - go Tip = Tip - go (Bin _ kx x l r) = case f kx x of - Just y -> combine kx y (go l) (go r) - Nothing -> merge (go l) (go r) +mapMaybeWithKey = makeMapMaybeWithKey -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. mapEitherWithKey :: GCompare k => (forall v. k v -> f v -> Either (g v) (h v)) -> DMap k f -> (DMap k g, DMap k h) -mapEitherWithKey f0 = toPair . go f0 - where - go :: GCompare k - => (forall v. k v -> f v -> Either (g v) (h v)) - -> DMap k f -> (DMap k g :*: DMap k h) - go _ Tip = (Tip :*: Tip) - go f (Bin _ kx x l r) = case f kx x of - Left y -> (combine kx y l1 r1 :*: merge l2 r2) - Right z -> (merge l1 r1 :*: combine kx z l2 r2) - where - (l1,l2) = mapEitherWithKey f l - (r1,r2) = mapEitherWithKey f r +mapEitherWithKey = makeMapEitherWithKey {-------------------------------------------------------------------- Mapping @@ -822,56 +619,32 @@ mapEitherWithKey f0 = toPair . go f0 -- | /O(n)/. Map a function over all values in the map. map :: (forall v. f v -> g v) -> DMap k f -> DMap k g -map f = go - where - go Tip = Tip - go (Bin sx kx x l r) = Bin sx kx (f x) (go l) (go r) +map f = fst $ makeMap f -- | /O(n)/. Map a function over all values in the map. mapWithKey :: (forall v. k v -> f v -> g v) -> DMap k f -> DMap k g -mapWithKey f = go - where - go Tip = Tip - go (Bin sx kx x l r) = Bin sx kx (f kx x) (go l) (go r) +mapWithKey f = fst $ makeMapWithKey f -- | /O(n)/. -- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ -- That is, behaves exactly like a regular 'traverse' except that the traversing -- function also has access to the key associated with a value. traverseWithKey :: Applicative t => (forall v. k v -> f v -> t (g v)) -> DMap k f -> t (DMap k g) -traverseWithKey f = go - where - go Tip = pure Tip - go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v - go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r +traverseWithKey f = fst $ makeTraverseWithKey f -- | /O(n)/. The function 'mapAccumLWithKey' threads an accumulating -- argument throught the map in ascending order of keys. mapAccumLWithKey :: (forall v. a -> k v -> f v -> (a, g v)) -> a -> DMap k f -> (a, DMap k g) -mapAccumLWithKey f = go - where - go a Tip = (a,Tip) - go a (Bin sx kx x l r) = - let (a1,l') = go a l - (a2,x') = f a1 kx x - (a3,r') = go a2 r - in (a3,Bin sx kx x' l' r') +mapAccumLWithKey f = fst $ makeMapAccumLWithKey f -- | /O(n)/. The function 'mapAccumRWithKey' threads an accumulating -- argument through the map in descending order of keys. mapAccumRWithKey :: (forall v. a -> k v -> f v -> (a, g v)) -> a -> DMap k f -> (a, DMap k g) -mapAccumRWithKey f = go - where - go a Tip = (a,Tip) - go a (Bin sx kx x l r) = - let (a1,r') = go a r - (a2,x') = f a1 kx x - (a3,l') = go a2 l - in (a3,Bin sx kx x' l' r') +mapAccumRWithKey f = fst $ makeMapAccumRWithKey f -- | /O(n*log n)/. -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. --- +-- -- The size of the result may be smaller if @f@ maps two or more distinct -- keys to the same new key. In this case the associated values will be -- combined using @c@. @@ -886,20 +659,18 @@ mapKeysWith c f = fromListWithKey c . Prelude.map fFirst . toList -- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. -- /The precondition is not checked./ -- Semi-formally, we have: --- --- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] -- > ==> mapKeysMonotonic f s == mapKeys f s -- > where ls = keys s -- -- This means that @f@ maps distinct original keys to distinct resulting keys. -- This function has better performance than 'mapKeys'. mapKeysMonotonic :: (forall v. k1 v -> k2 v) -> DMap k1 f -> DMap k2 f -mapKeysMonotonic _ Tip = Tip -mapKeysMonotonic f (Bin sz k x l r) = - Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r) +mapKeysMonotonic f = fst $ makeMapKeysMonotonic f {-------------------------------------------------------------------- - Folds + Folds --------------------------------------------------------------------} -- | /O(n)/. Fold the keys and values in the map, such that @@ -914,18 +685,12 @@ foldWithKey = foldrWithKey -- | /O(n)/. Post-order fold. The function will be applied from the lowest -- value to the highest. foldrWithKey :: (forall v. k v -> f v -> b -> b) -> b -> DMap k f -> b -foldrWithKey f = go - where - go z Tip = z - go z (Bin _ kx x l r) = go (f kx x (go z r)) l +foldrWithKey = makeFoldrWithKey -- | /O(n)/. Pre-order fold. The function will be applied from the highest -- value to the lowest. foldlWithKey :: (forall v. b -> k v -> f v -> b) -> b -> DMap k f -> b -foldlWithKey f = go - where - go z Tip = z - go z (Bin _ kx x l r) = go (f (go z l) kx x) r +foldlWithKey = makeFoldlWithKey {- -- | /O(n)/. A strict version of 'foldlWithKey'. @@ -937,7 +702,7 @@ foldlWithKey' f = go -} {-------------------------------------------------------------------- - List variations + List variations --------------------------------------------------------------------} -- | /O(n)/. Return all keys of the map in ascending order. @@ -955,7 +720,7 @@ assocs m = toList m {-------------------------------------------------------------------- - Lists + Lists use [foldlStrict] to reduce demand on the control-stack --------------------------------------------------------------------} @@ -963,7 +728,7 @@ assocs m -- If the list contains more than one value for the same key, the last value -- for the key is retained. fromList :: GCompare k => [DSum k f] -> DMap k f -fromList xs +fromList xs = foldlStrict ins empty xs where ins :: GCompare k => DMap k f -> DSum k f -> DMap k f @@ -971,7 +736,7 @@ fromList xs -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. fromListWithKey :: GCompare k => (forall v. k v -> f v -> f v -> f v) -> [DSum k f] -> DMap k f -fromListWithKey f xs +fromListWithKey f xs = foldlStrict (ins f) empty xs where ins :: GCompare k => (forall v. k v -> f v -> f v -> f v) -> DMap k f -> DSum k f -> DMap k f @@ -991,8 +756,8 @@ toDescList t = foldlWithKey (\xs k x -> (k :=> x):xs) [] t {-------------------------------------------------------------------- Building trees from ascending/descending lists can be done in linear time. - - Note that if [xs] is ascending that: + + Note that if [xs] is ascending that: fromAscList xs == fromList xs fromAscListWith f xs == fromListWith f xs --------------------------------------------------------------------} @@ -1006,7 +771,7 @@ fromAscList xs -- | /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 :: GEq k => (forall v. k v -> f v -> f v -> f v) -> [DSum k f] -> DMap k f +fromAscListWithKey :: GEq k => (forall v. k v -> f v -> f v -> f v) -> [DSum k f] -> DMap k f fromAscListWithKey f xs = fromDistinctAscList (combineEq f xs) where @@ -1028,30 +793,8 @@ fromAscListWithKey f xs -- | /O(n)/. Build a map from an ascending list of distinct elements in linear time. -- /The precondition is not checked./ fromDistinctAscList :: [DSum k f] -> DMap k f -fromDistinctAscList xs - = build const (length xs) xs - where - -- 1) use continutations so that we use heap space instead of stack space. - -- 2) special case for n==5 to build bushier trees. - - build :: (DMap k f -> [DSum k f] -> b) -> Int -> [DSum k f] -> b - build c 0 xs' = c Tip xs' - build c 5 xs' = case xs' of - ((k1:=>x1):(k2:=>x2):(k3:=>x3):(k4:=>x4):(k5:=>x5):xx) - -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx - _ -> error "fromDistinctAscList build" - build c n xs' = seq nr $ build (buildR nr c) nl xs' - where - nl = n `div` 2 - nr = n - nl - 1 - - buildR :: Int -> (DMap k f -> [DSum k f] -> b) -> DMap k f -> [DSum k f] -> b - buildR n c l ((k:=>x):ys) = build (buildB l k x c) n ys - buildR _ _ _ [] = error "fromDistinctAscList buildR []" - - buildB :: DMap k f -> k v -> f v -> (DMap k f -> a -> b) -> DMap k f -> a -> b - buildB l k x c r zs = c (bin k x l r) zs - +fromDistinctAscList = fst makeFromDistinctAscList + {-------------------------------------------------------------------- Split --------------------------------------------------------------------} @@ -1060,71 +803,43 @@ fromDistinctAscList xs -- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. -- Any key equal to @k@ is found in neither @map1@ nor @map2@. split :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, DMap k f) -split k = toPair . go - where - go :: DMap k f -> (DMap k f :*: DMap k f) - go Tip = (Tip :*: Tip) - go (Bin _ kx x l r) = case gcompare k kx of - GLT -> let !(lt :*: gt) = go l in (lt :*: combine kx x gt r) - GGT -> let !(lt :*: gt) = go r in (combine kx x l lt :*: gt) - GEQ -> (l :*: r) +split = fst . makeSplit {-# INLINABLE split #-} -- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just -- like 'split' but also returns @'lookup' k map@. splitLookup :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Maybe (f v), DMap k f) -splitLookup k = toTriple . go - where - go :: DMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f) - go Tip = Triple' Tip Nothing Tip - go (Bin _ kx x l r) = case gcompare k kx of - GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r) - GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt - GEQ -> Triple' l (Just x) r +splitLookup = fst . makeSplitLookup -- | /O(log n)/. The expression (@'splitMember' k map@) splits a map just -- like 'split' but also returns @'member' k map@. splitMember :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Bool, DMap k f) -splitMember k = toTriple . go - where - go :: DMap k f -> Triple' (DMap k f) Bool (DMap k f) - go Tip = Triple' Tip False Tip - go (Bin _ kx x l r) = case gcompare k kx of - GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r) - GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt - GEQ -> Triple' l True r +splitMember = fst . makeSplitMember -- | /O(log n)/. splitLookupWithKey :: forall k f v. GCompare k => k v -> DMap k f -> (DMap k f, Maybe (k v, f v), DMap k f) -splitLookupWithKey k = toTriple . go - where - go :: DMap k f -> Triple' (DMap k f) (Maybe (k v, f v)) (DMap k f) - go Tip = Triple' Tip Nothing Tip - go (Bin _ kx x l r) = case gcompare k kx of - GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (combine kx x gt r) - GGT -> let !(Triple' lt z gt) = go r in Triple' (combine kx x l lt) z gt - GEQ -> Triple' l (Just (kx, x)) r +splitLookupWithKey = fst . makeSplitLookupWithKey {-------------------------------------------------------------------- - Eq converts the tree to a list. In a lazy setting, this - actually seems one of the faster methods to compare two trees + Eq converts the tree to a list. In a lazy setting, this + actually seems one of the faster methods to compare two trees and it is certainly the simplest :-) --------------------------------------------------------------------} -instance EqTag k f => Eq (DMap k f) where +instance (GEq k, Has' Eq k f) => Eq (DMap k f) where t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2) {-------------------------------------------------------------------- - Ord + Ord --------------------------------------------------------------------} -instance OrdTag k f => Ord (DMap k f) where - compare m1 m2 = compare (toAscList m1) (toAscList m2) +instance (GCompare k, Has' Eq k f, Has' Ord k f) => Ord (DMap k f) where + compare m1 m2 = compare (toAscList m1) (toAscList m2) {-------------------------------------------------------------------- Read --------------------------------------------------------------------} -instance (GCompare k, ReadTag k f) => Read (DMap k f) where +instance (GCompare k, GRead k, Has' Read k f) => Read (DMap k f) where readPrec = parens $ prec 10 $ do Ident "fromList" <- lexP xs <- readPrec @@ -1135,7 +850,7 @@ instance (GCompare k, ReadTag k f) => Read (DMap k f) where {-------------------------------------------------------------------- Show --------------------------------------------------------------------} -instance ShowTag k f => Show (DMap k f) where +instance (GShow k, Has' Show k f) => Show (DMap k f) where showsPrec p m = showParen (p>10) ( showString "fromList " . showsPrec 11 (toList m) @@ -1143,11 +858,11 @@ instance ShowTag k f => Show (DMap k f) where -- | /O(n)/. Show the tree that implements the map. The tree is shown -- in a compressed, hanging format. See 'showTreeWith'. -showTree :: ShowTag k f => DMap k f -> String +showTree :: (GShow k, Has' Show k f) => DMap k f -> String showTree m = showTreeWith showElem True False m where - showElem :: ShowTag k f => k v -> f v -> String + showElem :: (GShow k, Has' Show k f) => k v -> f v -> String showElem k x = show (k :=> x) @@ -1162,48 +877,10 @@ showTreeWith showelem hang wide t | otherwise = (showsTree showelem wide [] [] t) "" showsTree :: (forall v. k v -> f v -> String) -> Bool -> [String] -> [String] -> DMap k f -> ShowS -showsTree showelem wide lbars rbars t - = case t of - Tip -> showsBars lbars . showString "|\n" - Bin _ kx x Tip Tip - -> showsBars lbars . showString (showelem kx x) . showString "\n" - Bin _ kx x l r - -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . showString (showelem kx x) . showString "\n" . - showWide wide lbars . - showsTree showelem wide (withEmpty lbars) (withBar lbars) l +showsTree showelem wide = fst $ makeShowsTree showelem wide showsTreeHang :: (forall v. k v -> f v -> String) -> Bool -> [String] -> DMap k f -> ShowS -showsTreeHang showelem wide bars t - = case t of - Tip -> showsBars bars . showString "|\n" - Bin _ kx x Tip Tip - -> showsBars bars . showString (showelem kx x) . showString "\n" - Bin _ kx x l r - -> showsBars bars . showString (showelem kx x) . showString "\n" . - showWide wide bars . - showsTreeHang showelem wide (withBar bars) l . - showWide wide bars . - showsTreeHang showelem wide (withEmpty bars) r - -showWide :: Bool -> [String] -> String -> String -showWide wide bars - | wide = showString (concat (reverse bars)) . showString "|\n" - | otherwise = id - -showsBars :: [String] -> ShowS -showsBars bars - = case bars of - [] -> id - _ -> showString (concat (reverse (tail bars))) . showString node - -node :: String -node = "+--" - -withBar, withEmpty :: [String] -> [String] -withBar bars = "| ":bars -withEmpty bars = " ":bars +showsTreeHang showelem wide = fst $ makeShowsTreeHang showelem wide {-------------------------------------------------------------------- Assertions @@ -1215,39 +892,11 @@ valid t = balanced t && ordered t && validsize t ordered :: GCompare k => DMap k f -> Bool -ordered t - = bounded (const True) (const True) t - where - bounded :: GCompare k => (Some k -> Bool) -> (Some k -> Bool) -> DMap k f -> Bool - bounded lo hi t' - = case t' of - Tip -> True - Bin _ kx _ l r -> (lo (This kx)) && (hi (This kx)) && bounded lo (< This kx) l && bounded (> This kx) hi r +ordered = fst makeOrdered -- | Exported only for "Debug.QuickCheck" balanced :: DMap k f -> Bool -balanced t - = case t of - Tip -> True - Bin _ _ _ l r -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && - balanced l && balanced r +balanced = fst makeBalanced validsize :: DMap k f -> Bool -validsize t - = (realsize t == Just (size t)) - where - realsize t' - = case t' of - Tip -> Just 0 - Bin sz _ _ l r -> case (realsize l,realsize r) of - (Just n,Just m) | n+m+1 == sz -> Just sz - _ -> Nothing -{-------------------------------------------------------------------- - Utilities ---------------------------------------------------------------------} -foldlStrict :: (a -> b -> a) -> a -> [b] -> a -foldlStrict f = go - where - go z [] = z - go z (x:xs) = z `seq` go (f z x) xs - +validsize = fst makeValidsize diff --git a/src/Data/Dependent/Map/Internal.hs b/src/Data/Dependent/Map/Internal.hs index 5596144..d14dd48 100644 --- a/src/Data/Dependent/Map/Internal.hs +++ b/src/Data/Dependent/Map/Internal.hs @@ -10,8 +10,12 @@ #if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708 {-# LANGUAGE PolyKinds #-} #endif +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} module Data.Dependent.Map.Internal where +import Prelude hiding (lookup) + import Data.Dependent.Sum import Data.GADT.Compare import Data.Some @@ -19,10 +23,10 @@ import Data.Some import Data.Typeable (Typeable) #endif --- |Dependent maps: 'k' is a GADT-like thing with a facility for +-- |Dependent maps: 'k' is a GADT-like thing with a facility for -- rediscovering its type parameter, elements of which function as identifiers -- tagged with the type of the thing they identify. Real GADTs are one --- useful instantiation of @k@, as are 'Tag's from "Data.Unique.Tag" in the +-- useful instantiation of @k@, as are 'Tag's from "Data.Unique.Tag" in the -- 'prim-uniq' package. -- -- Semantically, @'DMap' k f@ is equivalent to a set of @'DSum' k f@ where no two @@ -31,18 +35,37 @@ import Data.Typeable (Typeable) -- More informally, 'DMap' is to dependent products as 'M.Map' is to @(->)@. -- Thus it could also be thought of as a partial (in the sense of \"partial -- function\") dependent product. -data DMap k f where - Tip :: DMap k f - Bin :: {- sz -} !Int - -> {- key -} !(k v) - -> {- value -} f v - -> {- left -} !(DMap k f) - -> {- right -} !(DMap k f) - -> DMap k f +data DMap k f + = Tip + | Bin' {-# UNPACK #-} !(NonEmptyDMap k f) #if MIN_VERSION_base(4,7,0) deriving Typeable #endif +pattern Bin s k v l r = Bin' (NonEmptyDMap s k v l r) + +data NonEmptyDMap k f where + NonEmptyDMap + :: {- sz -} !Int + -> {- key -} !(k v) + -> {- value -} f v + -> {- left -} !(DMap k f) + -> {- right -} !(DMap k f) + -> NonEmptyDMap k f +#if MIN_VERSION_base(4,7,0) + deriving Typeable +#endif + +nonEmpty :: DMap k f -> Maybe (NonEmptyDMap k f) +nonEmpty = \case + Tip -> Nothing + Bin' ne -> Just ne + +fromNonEmpty :: Maybe (NonEmptyDMap k f) -> DMap k f +fromNonEmpty = \case + Nothing -> Tip + Just ne -> Bin' ne + {-------------------------------------------------------------------- Construction --------------------------------------------------------------------} @@ -58,8 +81,15 @@ empty = Tip -- -- > singleton 1 'a' == fromList [(1, 'a')] -- > size (singleton 1 'a') == 1 -singleton :: k v -> f v -> DMap k f -singleton k x = Bin 1 k x Tip Tip +singletonNE :: k v -> f v -> NonEmptyDMap k f +singletonNE k x = NonEmptyDMap 1 k x Tip Tip + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singletonE 1 'a') == 1 +singletonE :: k v -> f v -> DMap k f +singletonE k x = Bin' $! singletonNE k x {-------------------------------------------------------------------- Query @@ -68,52 +98,65 @@ singleton k x = Bin 1 k x Tip Tip -- | /O(1)/. Is the map empty? null :: DMap k f -> Bool null Tip = True -null Bin{} = False +null Bin'{} = False + +-- | /O(1)/. The number of elements in the map. +sizeNE :: NonEmptyDMap k f -> Int +sizeNE (NonEmptyDMap n _ _ _ _) = n -- | /O(1)/. The number of elements in the map. -size :: DMap k f -> Int -size Tip = 0 -size (Bin n _ _ _ _) = n +sizeE :: DMap k f -> Int +sizeE Tip = 0 +sizeE (Bin' ne) = sizeNE ne -- | /O(log n)/. 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 :: forall k f v. GCompare k => k v -> DMap k f -> Maybe (f v) -lookup k = k `seq` go - where - go :: DMap k f -> Maybe (f v) - go Tip = Nothing - go (Bin _ kx x l r) = - case gcompare k kx of - GLT -> go l - GGT -> go r - GEQ -> Just x - -lookupAssoc :: forall k f v. GCompare k => Some k -> DMap k f -> Maybe (DSum k f) -lookupAssoc (This k) = k `seq` go - where - go :: DMap k f -> Maybe (DSum k f) - go Tip = Nothing - go (Bin _ kx x l r) = - case gcompare k kx of - GLT -> go l - GGT -> go r - GEQ -> Just (kx :=> x) +lookupNE :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> Maybe (f v) +lookupNE k ne = k `seq` lookupNE' k ne + +-- | Skips the seq +{-# INLINE lookupNE' #-} +lookupNE' :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> Maybe (f v) +lookupNE' k (NonEmptyDMap _ kx x l r) = case gcompare k kx of + GLT -> lookupE' k l + GGT -> lookupE' k r + GEQ -> Just x + +-- | /O(log n)/. 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. +lookupE :: forall k f v. GCompare k => k v -> DMap k f -> Maybe (f v) +lookupE k m = k `seq` lookupE' k m + +-- | Skips the seq +{-# INLINE lookupE' #-} +lookupE' :: forall k f v. GCompare k => k v -> DMap k f -> Maybe (f v) +lookupE' k = \case + Tip -> Nothing + (Bin' ne) -> lookupNE' k ne + +lookupAssocNE :: forall k f v. GCompare k => Some k -> NonEmptyDMap k f -> Maybe (DSum k f) +lookupAssocNE (This k) m = (k :=>) <$> lookupNE k m + +lookupAssocE :: forall k f v. GCompare k => Some k -> DMap k f -> Maybe (DSum k f) +lookupAssocE (This k) m = (k :=>) <$> lookupE k m {-------------------------------------------------------------------- Utility functions that maintain the balance properties of the tree. All constructors assume that all values in [l] < [k] and all values in [r] > [k], and that [l] and [r] are valid trees. - + In order of sophistication: - [Bin sz k x l r] The type constructor. - [bin k x l r] Maintains the correct size, assumes that both [l] + [Bin' sz k x l r] The type constructor. + [bin k x l r] Maintains the correct sizeE, assumes that both [l] and [r] are balanced with respect to each other. - [balance k x l r] Restores the balance and size. + [balance k x l r] Restores the balance and sizeE. Assumes that the original tree was balanced and that [l] or [r] has changed by at most one element. - [combine k x l r] Restores balance and size. + [combine k x l r] Restores balance and sizeE. Furthermore, we can construct a new tree from two trees. Both operations assume that all values in [l] < all values in [r] and that [l] and [r] @@ -123,70 +166,87 @@ lookupAssoc (This k) = k `seq` go [merge l r] Merges two trees and restores balance. Note: in contrast to Adam's paper, we use (<=) comparisons instead - of (<) comparisons in [combine], [merge] and [balance]. - Quickcheck (on [difference]) showed that this was necessary in order - to maintain the invariants. It is quite unsatisfactory that I haven't - been able to find out why this is actually the case! Fortunately, it + of (<) comparisons in [combine], [merge] and [balance]. + Quickcheck (on [difference]) showed that this was necessary in order + to maintain the invariants. It is quite unsatisfactory that I haven't + been able to find out why this is actually the case! Fortunately, it doesn't hurt to be a bit more conservative. --------------------------------------------------------------------} {-------------------------------------------------------------------- Combine --------------------------------------------------------------------} -combine :: GCompare k => k v -> f v -> DMap k f -> DMap k f -> DMap k f -combine kx x Tip r = insertMin kx x r -combine kx x l Tip = insertMax kx x l -combine kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) - | delta*sizeL <= sizeR = balance kz z (combine kx x l lz) rz - | delta*sizeR <= sizeL = balance ky y ly (combine kx x ry r) - | otherwise = bin kx x l r +combineNE :: GCompare k => k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +combineNE kx x l@(NonEmptyDMap sizeEL ky y ly ry) r@(NonEmptyDMap sizeER kz z lz rz) + | delta*sizeEL <= sizeER = balanceNEL kz z (combineNEL kx x l lz) rz + | delta*sizeER <= sizeEL = balanceNER ky y ly (combineNER kx x ry r) + | otherwise = binNE kx x (Bin' l) (Bin' r) + +combineNEL :: GCompare k => k v -> f v -> NonEmptyDMap k f -> DMap k f -> NonEmptyDMap k f +combineNEL kx x l Tip = insertMaxNE kx x l +combineNEL kx x l (Bin' r) = combineNE kx x l r + +combineNER :: GCompare k => k v -> f v -> DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +combineNER kx x Tip r = insertMinNE kx x r +combineNER kx x (Bin' l) r = combineNE kx x l r +combineE :: GCompare k => k v -> f v -> DMap k f -> DMap k f -> NonEmptyDMap k f +combineE kx x Tip r = insertMin kx x r +combineE kx x l Tip = insertMax kx x l +combineE kx x (Bin' l) (Bin' r) = combineNE kx x l r -- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: k v -> f v -> DMap k f -> DMap k f -insertMax kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balance ky y l (insertMax kx x r) - -insertMin kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balance ky y (insertMin kx x l) r - + +insertMaxNE, insertMinNE :: k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insertMaxNE kx x (NonEmptyDMap _ ky y l r) = balanceNER ky y l (insertMax kx x r) +insertMinNE kx x (NonEmptyDMap _ ky y l r) = balanceNEL ky y (insertMin kx x l) r + +insertMax, insertMin :: k v -> f v -> DMap k f -> NonEmptyDMap k f +insertMax kx x t = case t of + Tip -> singletonNE kx x + Bin' ne -> insertMaxNE kx x ne + +insertMin kx x t = case t of + Tip -> singletonNE kx x + Bin' ne -> insertMinNE kx x ne + {-------------------------------------------------------------------- [merge l r]: merges two trees. --------------------------------------------------------------------} -merge :: DMap k f -> DMap k f -> DMap k f -merge Tip r = r -merge l Tip = l -merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) - | delta*sizeL <= sizeR = balance ky y (merge l ly) ry - | delta*sizeR <= sizeL = balance kx x lx (merge rx r) - | otherwise = glue l r +mergeNE :: NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +mergeNE l@(NonEmptyDMap sizeEL kx x lx rx) r@(NonEmptyDMap sizeER ky y ly ry) + | delta*sizeEL <= sizeER = balanceNEL ky y (mergeNEL l ly) ry + | delta*sizeER <= sizeEL = balanceNER kx x lx (mergeNER rx r) + | otherwise = glueNE l r + +mergeNEL :: NonEmptyDMap k f -> DMap k f -> NonEmptyDMap k f +mergeNEL l Tip = l +mergeNEL l (Bin' r) = mergeNE l r + +mergeNER :: DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +mergeNER Tip r = r +mergeNER (Bin' l) r = mergeNE l r + +mergeE :: DMap k f -> DMap k f -> DMap k f +mergeE Tip r = r +mergeE l Tip = l +mergeE (Bin' l) (Bin' r) = Bin' $! mergeNE l r {-------------------------------------------------------------------- [glue l r]: glues two trees together. Assumes that [l] and [r] are already balanced with respect to each other. --------------------------------------------------------------------} -glue :: DMap k f -> DMap k f -> DMap k f -glue Tip r = r -glue l Tip = l -glue l r - | size l > size r = case deleteFindMax l of (km :=> m,l') -> balance km m l' r - | otherwise = case deleteFindMin r of (km :=> m,r') -> balance km m l r' - --- | /O(log n)/. Delete and find the minimal element. --- --- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) --- > deleteFindMin Error: can not return the minimal element of an empty map - -deleteFindMin :: DMap k f -> (DSum k f, DMap k f) -deleteFindMin t = case minViewWithKey t of - Nothing -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) - Just p -> p +glueNE :: NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +glueNE l r + | sizeNE l > sizeNE r = case maxViewWithKeyNE l of + (km :=> m, l') -> balanceNER km m l' r + | otherwise = case minViewWithKeyNE r of + (km :=> m, r') -> balanceNEL km m l r' + +glueE :: DMap k f -> DMap k f -> DMap k f +glueE Tip r = r +glueE l Tip = l +glueE (Bin' l) (Bin' r) = Bin' $! glueNE l r -- | A strict pair. data (:*:) a b = !a :*: !b @@ -205,48 +265,47 @@ toTriple (Triple' a b c) = (a, b, c) {-# INLINE toTriple #-} -- | /O(log n)/. Retrieves the minimal (key :=> value) entry of the map, and --- the map stripped of that element, or 'Nothing' if passed an empty map. -minViewWithKey :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f) -minViewWithKey Tip = Nothing -minViewWithKey (Bin _ k0 x0 l0 r0) = Just $! toPair $ go k0 x0 l0 r0 +-- the map stripped of that element +minViewWithKeyNE :: forall k f . NonEmptyDMap k f -> (DSum k f, DMap k f) +minViewWithKeyNE (NonEmptyDMap _ k0 x0 l0 r0) = toPair $ go k0 x0 l0 r0 where go :: k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f go k x Tip r = (k :=> x) :*: r go k x (Bin _ kl xl ll lr) r = let !(km :*: l') = go kl xl ll lr - in (km :*: balance k x l' r) + in (km :*: (Bin' $! balanceE k x l' r)) + +-- | /O(log n)/. Retrieves the minimal (key :=> value) entry of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +minViewWithKeyE :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f) +minViewWithKeyE Tip = Nothing +minViewWithKeyE (Bin' ne) = Just $! minViewWithKeyNE ne -- | /O(log n)/. Retrieves the maximal (key :=> value) entry of the map, and -- the map stripped of that element, or 'Nothing' if passed an empty map. -maxViewWithKey :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f) -maxViewWithKey Tip = Nothing -maxViewWithKey (Bin _ k0 x0 l0 r0) = Just $! toPair $ go k0 x0 l0 r0 +maxViewWithKeyNE :: forall k f . NonEmptyDMap k f -> (DSum k f, DMap k f) +maxViewWithKeyNE (NonEmptyDMap _ k0 x0 l0 r0) = toPair $ go k0 x0 l0 r0 where go :: k v -> f v -> DMap k f -> DMap k f -> DSum k f :*: DMap k f go k x l Tip = (k :=> x) :*: l go k x l (Bin _ kr xr rl rr) = let !(km :*: r') = go kr xr rl rr - in (km :*: balance k x l r') + in (km :*: (Bin' $! balanceE k x l r')) --- | /O(log n)/. Delete and find the maximal element. --- --- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) --- > deleteFindMax empty Error: can not return the maximal element of an empty map +-- | /O(log n)/. Retrieves the maximal (key :=> value) entry of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +maxViewWithKeyE :: forall k f . DMap k f -> Maybe (DSum k f, DMap k f) +maxViewWithKeyE Tip = Nothing +maxViewWithKeyE (Bin' ne) = Just $! maxViewWithKeyNE ne -deleteFindMax :: DMap k f -> (DSum k f, DMap k f) -deleteFindMax t - = case t of - Bin _ k x l Tip -> (k :=> x,l) - Bin _ k x l r -> let (km,r') = deleteFindMax r in (km,balance k x l r') - Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) {-------------------------------------------------------------------- [balance l x r] balances two trees with value x. - The sizes of the trees should balance after decreasing the - size of one of them. (a rotation). + The sizeEs of the trees should balance after decreasing the + sizeE of one of them. (a rotation). - [delta] is the maximal relative difference between the sizes of + [delta] is the maximal relative difference between the sizeEs of two trees, it corresponds with the [w] in Adams' paper. [ratio] is the ratio between an outer and inner sibling of the heavier subtree in an unbalanced setting. It determines @@ -257,7 +316,7 @@ deleteFindMax t Note that: - [delta] should be larger than 4.646 with a [ratio] of 2. - [delta] should be larger than 3.745 with a [ratio] of 1.534. - + - A lower [delta] leads to a more 'perfectly' balanced tree. - A higher [delta] performs less rebalancing. @@ -275,49 +334,103 @@ delta,ratio :: Int delta = 4 ratio = 2 -balance :: k v -> f v -> DMap k f -> DMap k f -> DMap k f -balance k x l r - | sizeL + sizeR <= 1 = Bin sizeX k x l r - | sizeR >= delta*sizeL = rotateL k x l r - | sizeL >= delta*sizeR = rotateR k x l r - | otherwise = Bin sizeX k x l r +-- TODO delete these + +balanceNE :: k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +balanceNE k x l r + | sizeEL + sizeER <= 1 = NonEmptyDMap sizeEX k x (Bin' l) (Bin' r) + | sizeER >= delta*sizeEL = rotateL k x (Bin' l) r + | sizeEL >= delta*sizeER = rotateR k x l (Bin' r) + | otherwise = NonEmptyDMap sizeEX k x (Bin' l) (Bin' r) + where + sizeEL = sizeNE l + sizeER = sizeNE r + sizeEX = sizeEL + sizeER + 1 + +balanceNEL :: k v -> f v -> NonEmptyDMap k f -> DMap k f -> NonEmptyDMap k f +balanceNEL k x l r + | sizeEL + sizeER <= 1 = NonEmptyDMap sizeEX k x (Bin' l) r + | sizeER >= delta*sizeEL = let + r' = case r of + Tip -> error "balanceNEL: Declared sizeE of 'r' is greater than 1, shouldn't be empty" + Bin' ne -> ne + in rotateL k x (Bin' l) r' + | sizeEL >= delta*sizeER = rotateR k x l r + | otherwise = NonEmptyDMap sizeEX k x (Bin' l) r + where + sizeEL = sizeNE l + sizeER = sizeE r + sizeEX = sizeEL + sizeER + 1 + +balanceNER :: k v -> f v -> DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +balanceNER k x l r + | sizeEL + sizeER <= 1 = NonEmptyDMap sizeEX k x l (Bin' r) + | sizeER >= delta*sizeEL = rotateL k x l r + | sizeEL >= delta*sizeER = let + l' = case l of + Tip -> error "balanceNER: Declared sizeE of 'l' is greater than 1, shouldn't be empty" + Bin' ne -> ne + in rotateR k x l' (Bin' r) + | otherwise = NonEmptyDMap sizeEX k x l (Bin' r) where - sizeL = size l - sizeR = size r - sizeX = sizeL + sizeR + 1 + sizeEL = sizeE l + sizeER = sizeNE r + sizeEX = sizeEL + sizeER + 1 + + +balanceE :: k v -> f v -> DMap k f -> DMap k f -> NonEmptyDMap k f +balanceE k x l r + | sizeEL + sizeER <= 1 = NonEmptyDMap sizeEX k x l r + | sizeER >= delta*sizeEL = let + r' = case r of + Tip -> error "balance: Declared sizeE of 'r' is greater than 1, shouldn't be empty" + Bin' ne -> ne + in rotateL k x l r' + | sizeEL >= delta*sizeER = let + l' = case l of + Tip -> error "balance: Declared sizeE of 'l' is greater than 1, shouldn't be empty" + Bin' ne -> ne + in rotateR k x l' r + | otherwise = NonEmptyDMap sizeEX k x l r + where + sizeEL = sizeE l + sizeER = sizeE r + sizeEX = sizeEL + sizeER + 1 -- rotate -rotateL :: k v -> f v -> DMap k f -> DMap k f -> DMap k f -rotateL k x l r@(Bin _ _ _ ly ry) - | size ly < ratio*size ry = singleL k x l r +rotateL :: k v -> f v -> DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +rotateL k x l r@(NonEmptyDMap _ _ _ ly ry) + | sizeE ly < ratio*sizeE ry = singleL k x l r | otherwise = doubleL k x l r -rotateL _ _ _ Tip = error "rotateL Tip" -rotateR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f -rotateR k x l@(Bin _ _ _ ly ry) r - | size ry < ratio*size ly = singleR k x l r +rotateR :: k v -> f v -> NonEmptyDMap k f -> DMap k f -> NonEmptyDMap k f +rotateR k x l@(NonEmptyDMap _ _ _ ly ry) r + | sizeE ry < ratio*sizeE ly = singleR k x l r | otherwise = doubleR k x l r -rotateR _ _ Tip _ = error "rotateR Tip" -- basic rotations -singleL, singleR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f -singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 -singleL _ _ _ Tip = error "singleL Tip" -singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) -singleR _ _ Tip _ = error "singleR Tip" - -doubleL, doubleR :: k v -> f v -> DMap k f -> DMap k f -> DMap k f -doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4) +singleL :: k v -> f v -> DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +singleL k1 x1 t1 (NonEmptyDMap _ k2 x2 t2 t3) = binNE k2 x2 (binE k1 x1 t1 t2) t3 + +singleR :: k v -> f v -> NonEmptyDMap k f -> DMap k f -> NonEmptyDMap k f +singleR k1 x1 (NonEmptyDMap _ k2 x2 t1 t2) t3 = binNE k2 x2 t1 (binE k1 x1 t2 t3) + +doubleL :: k v -> f v -> DMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +doubleL k1 x1 t1 (NonEmptyDMap _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = binNE k3 x3 (binE k1 x1 t1 t2) (binE k2 x2 t3 t4) doubleL _ _ _ _ = error "doubleL" -doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4) + +doubleR :: k v -> f v -> NonEmptyDMap k f -> DMap k f -> NonEmptyDMap k f +doubleR k1 x1 (NonEmptyDMap _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = binNE k3 x3 (binE k2 x2 t1 t2) (binE k1 x1 t3 t4) doubleR _ _ _ _ = error "doubleR" {-------------------------------------------------------------------- - The bin constructor maintains the size of the tree + The bin constructor maintains the sizeE of the tree --------------------------------------------------------------------} -bin :: k v -> f v -> DMap k f -> DMap k f -> DMap k f -bin k x l r - = Bin (size l + size r + 1) k x l r +binNE :: k v -> f v -> DMap k f -> DMap k f -> NonEmptyDMap k f +binNE k x l r = NonEmptyDMap (sizeE l + sizeE r + 1) k x l r + +binE :: k v -> f v -> DMap k f -> DMap k f -> DMap k f +binE k x l r = Bin' $! binNE k x l r {-------------------------------------------------------------------- Utility functions that return sub-ranges of the original @@ -332,7 +445,7 @@ bin k x l r [split k t] Returns two trees [l] and [r] where all keys in [l] are <[k] and all keys in [r] are >[k]. - [splitLookup k t] Just like [split] but also returns whether [k] + [splitLookupE k t] Just like [split] but also returns whether [k] was found in the tree. --------------------------------------------------------------------} @@ -349,15 +462,15 @@ trim cmplo cmphi t@(Bin _ kx _ l r) GT -> t _ -> trim cmplo cmphi l _ -> trim cmplo cmphi r - -trimLookupLo :: GCompare k => Some k -> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f) -trimLookupLo _ _ Tip = (Nothing,Tip) -trimLookupLo lo cmphi t@(Bin _ kx x l r) + +trimLookupELo :: GCompare k => Some k -> (Some k -> Ordering) -> DMap k f -> (Maybe (DSum k f), DMap k f) +trimLookupELo _ _ Tip = (Nothing,Tip) +trimLookupELo lo cmphi t@(Bin _ kx x l r) = case compare lo (This kx) of LT -> case cmphi (This kx) of - GT -> (lookupAssoc lo t, t) - _ -> trimLookupLo lo cmphi l - GT -> trimLookupLo lo cmphi r + GT -> (lookupAssocE lo t, t) + _ -> trimLookupELo lo cmphi l + GT -> trimLookupELo lo cmphi r EQ -> (Just (kx :=> x),trim (compare lo) cmphi r) @@ -370,7 +483,7 @@ filterGt cmp = go where go Tip = Tip go (Bin _ kx x l r) = case cmp (This kx) of - LT -> combine kx x (go l) r + LT -> Bin' $! combineE kx x (go l) r GT -> go r EQ -> r @@ -380,5 +493,5 @@ filterLt cmp = go go Tip = Tip go (Bin _ kx x l r) = case cmp (This kx) of LT -> go l - GT -> combine kx x l (go r) + GT -> Bin' $! combineE kx x l (go r) EQ -> l diff --git a/src/Data/Dependent/Map/Internal2.hs b/src/Data/Dependent/Map/Internal2.hs new file mode 100644 index 0000000..8f8fd7e --- /dev/null +++ b/src/Data/Dependent/Map/Internal2.hs @@ -0,0 +1,1190 @@ +{-# LANGUAGE GADTs, RankNTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE CPP #-} +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702 +{-# LANGUAGE Trustworthy #-} +#endif +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} +#endif +module Data.Dependent.Map.Internal2 where + +import Prelude hiding (null, lookup, map) +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative(..), (<$>)) +#endif +import Data.Dependent.Map.Internal +#if !MIN_VERSION_base(4,7,0) +import Data.Dependent.Map.Typeable ({- instance Typeable ... -}) +#endif +import Data.Dependent.Map.PtrEquality + +import Data.Dependent.Sum +import Data.GADT.Compare +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.List.NonEmpty as NEL +import Data.Some + +{-------------------------------------------------------------------- + Insertion Worker Builders +--------------------------------------------------------------------} + +makeInsert + :: forall k f v + . GCompare k + => k v + -> f v + -> DMap k f + -> NonEmptyDMap k f +makeInsert kx x = kx `seq` go + where + go :: DMap k f -> NonEmptyDMap k f + go Tip = singletonNE kx x + go (Bin' t@(NonEmptyDMap sz ky y l r)) = case gcompare kx ky of + GLT -> let !l' = go l + in case l of + Bin' nel | l' `ptrEq` nel -> t + _ -> balanceNEL ky y l' r + GGT -> let !r' = go r + in case r of + Bin' ner | r' `ptrEq` ner -> t + _ -> balanceNER ky y l r' + GEQ + | kx `ptrEq` ky && x `ptrEq` y -> t + | otherwise -> NonEmptyDMap sz kx x l r + +makeInsertR + :: forall k f v + . GCompare k + => k v + -> f v + -> DMap k f + -> NonEmptyDMap k f +makeInsertR kx x = kx `seq` go + where + go :: DMap k f -> NonEmptyDMap k f + go Tip = singletonNE kx x + go (Bin' t@(NonEmptyDMap sz ky y l r)) = case gcompare kx ky of + GLT -> let !l' = go l + in case l of + Bin' nel | l' `ptrEq` nel -> t + _ -> balanceNEL ky y l' r + GGT -> let !r' = go r + in case r of + Bin' ner | r' `ptrEq` ner -> t + _ -> balanceNER ky y l r' + GEQ -> t + +makeInsertWithKey + :: forall k f v + . GCompare k + => (k v -> f v -> f v -> f v) + -> k v + -> f v + -> DMap k f + -> NonEmptyDMap k f +makeInsertWithKey f kx x = kx `seq` go + where + go :: DMap k f -> NonEmptyDMap k f + go Tip = singletonNE kx x + go (Bin sy ky y l r) = + case gcompare kx ky of + GLT -> balanceNEL ky y (go l) r + GGT -> balanceNER ky y l (go r) + GEQ -> NonEmptyDMap sy kx (f kx x y) l r + +makeInsertWithKey' + :: forall k f v + . GCompare k + => (k v -> f v -> f v -> f v) + -> k v + -> f v + -> DMap k f + -> NonEmptyDMap k f +makeInsertWithKey' f kx x = kx `seq` go + where + go :: DMap k f -> NonEmptyDMap k f + go Tip = singletonNE kx $! x + go (Bin sy ky y l r) = + case gcompare kx ky of + GLT -> balanceNEL ky y (go l) r + GGT -> balanceNER ky y l (go r) + GEQ -> let x' = f kx x y in seq x' (NonEmptyDMap sy kx x' l r) + +makeInsertLookupWithKey + :: forall k f v + . GCompare k + => (k v -> f v -> f v -> f v) + -> k v + -> f v + -> DMap k f + -> (Maybe (f v), NonEmptyDMap k f) +makeInsertLookupWithKey f kx x = kx `seq` go + where + go :: DMap k f -> (Maybe (f v), NonEmptyDMap k f) + go Tip = (Nothing, singletonNE kx x) + go (Bin sy ky y l r) = + case gcompare kx ky of + GLT -> let (found, l') = go l + in (found, balanceNEL ky y l' r) + GGT -> let (found, r') = go r + in (found, balanceNER ky y l r') + GEQ -> (Just y, NonEmptyDMap sy kx (f kx x y) l r) + +makeInsertLookupWithKey' + :: forall k f v + . GCompare k + => (k v -> f v -> f v -> f v) + -> k v + -> f v + -> DMap k f + -> (Maybe (f v), NonEmptyDMap k f) +makeInsertLookupWithKey' f kx x = kx `seq` go + where + go :: DMap k f -> (Maybe (f v), NonEmptyDMap k f) + go Tip = x `seq` (Nothing, singletonNE kx x) + go (Bin sy ky y l r) = + case gcompare kx ky of + GLT -> let (found, l') = go l + in (found, balanceNEL ky y l' r) + GGT -> let (found, r') = go r + in (found, balanceNER ky y l r') + GEQ -> let x' = f kx x y in x' `seq` (Just y, NonEmptyDMap sy kx x' l r) + +{-------------------------------------------------------------------- + Deletion Worker Builders +--------------------------------------------------------------------} + +makeDelete + :: forall k f v + . GCompare k + => k v + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) + ) +makeDelete k = (k `seq` go, k `seq` nonEmpty . go') + where + go :: DMap k f -> DMap k f + go Tip = Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> DMap k f + go' (NonEmptyDMap _ kx x l r) = + case gcompare k kx of + GLT -> Bin' $! balanceE kx x (go l) r + GGT -> Bin' $! balanceE kx x l (go r) + GEQ -> glueE l r + +makeAdjustF + :: forall f k v g + . (GCompare k, Applicative f) + => (g v -> f (g v)) + -> k v + -> ( DMap k g -> f (DMap k g) + , NonEmptyDMap k g -> f (NonEmptyDMap k g) + ) +makeAdjustF f k = (k `seq` go, k `seq` go') + where + go :: DMap k g -> f (DMap k g) + go Tip = pure Tip + go (Bin' ne) = Bin' <$> go' ne + + go' :: NonEmptyDMap k g -> f (NonEmptyDMap k g) + go' (NonEmptyDMap sx kx x l r) = + case gcompare k kx of + GLT -> NonEmptyDMap sx kx x <$> go l <*> pure r + GGT -> NonEmptyDMap sx kx x l <$> go r + GEQ -> NonEmptyDMap sx kx <$> f x <*> pure l <*> pure r + +makeAdjustWithKey + :: forall k f v + . GCompare k + => (k v -> f v -> f v) + -> k v + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k f + ) +makeAdjustWithKey f k = (k `seq` go, k `seq` go') + where + go :: DMap k f -> DMap k f + go Tip = Tip + go (Bin' ne) = Bin' $! go' ne + + go' :: NonEmptyDMap k f -> NonEmptyDMap k f + go' (NonEmptyDMap sx kx x l r) = + case gcompare k kx of + GLT -> NonEmptyDMap sx kx x (go l) r + GGT -> NonEmptyDMap sx kx x l (go r) + GEQ -> NonEmptyDMap sx kx (f kx x) l r + +makeAdjustWithKey' + :: forall k f v + . GCompare k + => (k v -> f v -> f v) + -> k v + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k f + ) +makeAdjustWithKey' f k = (k `seq` go, k `seq` go') + where + go :: DMap k f -> DMap k f + go Tip = Tip + go (Bin' ne) = Bin' $! go' ne + + go' :: NonEmptyDMap k f -> NonEmptyDMap k f + go' (NonEmptyDMap sx kx x l r) = + case gcompare k kx of + GLT -> NonEmptyDMap sx kx x (go l) r + GGT -> NonEmptyDMap sx kx x l (go r) + GEQ -> let !x' = f kx x in NonEmptyDMap sx kx x' l r + +makeUpdateWithKey + :: forall k f v + . GCompare k + => (k v -> f v -> Maybe (f v)) + -> k v + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) + ) +makeUpdateWithKey f k = (k `seq` go, k `seq` nonEmpty . go') + where + go :: DMap k f -> DMap k f + go Tip = Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> DMap k f + go' (NonEmptyDMap sx kx x l r) = + case gcompare k kx of + GLT -> Bin' $! balanceE kx x (go l) r + GGT -> Bin' $! balanceE kx x l (go r) + GEQ -> case f kx x of + Just x' -> Bin sx kx x' l r + Nothing -> glueE l r + +makeUpdateLookupWithKey + :: forall k f v + . GCompare k + => (k v -> f v -> Maybe (f v)) + -> k v + -> ( DMap k f -> (Maybe (f v), DMap k f) + , NonEmptyDMap k f -> (Maybe (f v), Maybe (NonEmptyDMap k f)) + ) +makeUpdateLookupWithKey f k = (k `seq` go, k `seq` fmap nonEmpty . go') + where + go :: DMap k f -> (Maybe (f v), DMap k f) + go Tip = (Nothing, Tip) + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> (Maybe (f v), DMap k f) + go' (NonEmptyDMap sx kx x l r) = + case gcompare k kx of + GLT -> let (found,l') = go l in (found,Bin' $! balanceE kx x l' r) + GGT -> let (found,r') = go r in (found,Bin' $! balanceE kx x l r') + GEQ -> case f kx x of + Just x' -> (Just x', Bin sx kx x' l r) + Nothing -> (Just x, glueE l r) + +makeAlter + :: forall k f v + . GCompare k + => (Maybe (f v) -> Maybe (f v)) + -> k v + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) + ) +makeAlter f k = (k `seq` go, k `seq` nonEmpty . go') + where + go :: DMap k f -> DMap k f + go Tip = case f Nothing of + Nothing -> Tip + Just x -> singletonE k x + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> DMap k f + go' (NonEmptyDMap sx kx x l r) = case gcompare k kx of + GLT -> Bin' $! balanceE kx x (go l) r + GGT -> Bin' $! balanceE kx x l (go r) + GEQ -> case f (Just x) of + Just x' -> Bin sx kx x' l r + Nothing -> glueE l r + +{-------------------------------------------------------------------- + Indexing Worker Builders +--------------------------------------------------------------------} + +makeAlterF + :: forall k f v g + . (GCompare k, Functor f) + => k v + -> (Maybe (g v) -> f (Maybe (g v))) + -> ( DMap k g -> f (DMap k g) + , NonEmptyDMap k g -> f (Maybe (NonEmptyDMap k g)) + ) +makeAlterF k f = (go, fmap nonEmpty . go') + where + go :: DMap k g -> f (DMap k g) + go Tip = maybe Tip (singletonE k) <$> f Nothing + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k g -> f (DMap k g) + go' (NonEmptyDMap sx kx x l r) = case gcompare k kx of + GLT -> (\l' -> Bin' $! balanceE kx x l' r) <$> go l + GGT -> (\r' -> Bin' $! balanceE kx x l r') <$> go r + GEQ -> maybe (glueE l r) (\x' -> Bin sx kx x' l r) <$> f (Just x) + +makeLookupIndex + :: forall k f v + . GCompare k + => k v + -> ( DMap k f -> Maybe Int + , NonEmptyDMap k f -> Maybe Int + ) +makeLookupIndex k = (k `seq` go 0, k `seq` go' 0) + where + go :: Int -> DMap k f -> Maybe Int + go !idx Tip = idx `seq` Nothing + go !idx (Bin' ne) = go' idx ne + + go' :: Int -> NonEmptyDMap k f -> Maybe Int + go' !idx (NonEmptyDMap _ kx _ l r) + = case gcompare k kx of + GLT -> go idx l + GGT -> go (idx + sizeE l + 1) r + GEQ -> Just (idx + sizeE l) + +makElemAt + :: ( Int -> DMap k f -> DSum k f + , Int -> NonEmptyDMap k f -> DSum k f + ) +makElemAt = (go, go') + where + go _ Tip = error "Map.elemAt: index out of range" + go i (Bin' ne) = go' i ne + + go' :: Int -> NonEmptyDMap k f -> DSum k f + go' i (NonEmptyDMap _ kx x l r) + = case compare i sizeL of + LT -> go i l + GT -> go (i-sizeL-1) r + EQ -> kx :=> x + where + sizeL = sizeE l + +makeUpdateAt + :: forall k f v + . (forall v. k v -> f v -> Maybe (f v)) + -> Int + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) + ) +makeUpdateAt f i0 = (i0 `seq` go i0, i0 `seq` nonEmpty . go' i0) + where + go :: Int -> DMap k f -> DMap k f + go _ Tip = Tip + go i (Bin' ne) = go' i ne + + go' :: Int -> NonEmptyDMap k f -> DMap k f + go' i (NonEmptyDMap sx kx x l r) = case compare i sizeL of + LT -> Bin' $! balanceE kx x (go i l) r + GT -> Bin' $! balanceE kx x l (go (i-sizeL-1) r) + EQ -> case f kx x of + Just x' -> Bin sx kx x' l r + Nothing -> glueE l r + where + sizeL = sizeE l + +{-------------------------------------------------------------------- + Minimal, Maximal Worker Builders +--------------------------------------------------------------------} + +makeLookupMin + :: ( DMap k f -> Maybe (DSum k f) + , NonEmptyDMap k f -> DSum k f + ) +makeLookupMin = (go, go') + where + go m = case m of + Tip -> Nothing + Bin _ kx x l _ -> Just $! goInner kx x l + + go' (NonEmptyDMap _ kx x l _) = goInner kx x l + + goInner :: k v -> f v -> DMap k f -> DSum k f + goInner kx x Tip = kx :=> x + goInner kx x (Bin' ne) = goInner' kx x ne + + goInner' :: k v -> f v -> NonEmptyDMap k f -> DSum k f + goInner' _ _ (NonEmptyDMap _ kx x l _) = goInner kx x l + +makeLookupMax + :: ( DMap k f -> Maybe (DSum k f) + , NonEmptyDMap k f -> DSum k f + ) +makeLookupMax = (go, go') + where + go m = case m of + Tip -> Nothing + Bin _ kx x _ r -> Just $! goInner kx x r + + go' (NonEmptyDMap _ kx x _ r) = goInner kx x r + + goInner :: k v -> f v -> DMap k f -> DSum k f + goInner kx x Tip = kx :=> x + goInner kx x (Bin' ne) = goInner' kx x ne + + goInner' :: k v -> f v -> NonEmptyDMap k f -> DSum k f + goInner' _ _ (NonEmptyDMap _ kx x _ r) = goInner kx x r + +makeDeleteMin + :: forall k f + . ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) + ) +makeDeleteMin = (go, nonEmpty . go') + where + go :: DMap k f -> DMap k f + go Tip = Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> DMap k f + go' (NonEmptyDMap _ _ _ Tip r) = r + go' (NonEmptyDMap _ kx x l r) = Bin' $! balanceE kx x (go l) r + +makeDeleteMax + :: forall k f + . ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) + ) +makeDeleteMax = (go, nonEmpty . go') + where + go :: DMap k f -> DMap k f + go Tip = Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> DMap k f + go' (NonEmptyDMap _ _ _ l Tip) = l + go' (NonEmptyDMap _ kx x l r) = Bin' $! balanceE kx x l (go r) + + +-- | /O(log n)/. Update the value at the minimal key. +makeUpdateMinWithKey + :: (forall v. k v -> f v -> Maybe (f v)) + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f)) +makeUpdateMinWithKey f = (go, nonEmpty . go') + where + go Tip = Tip + go (Bin' ne) = go' ne + + go' (NonEmptyDMap sx kx x Tip r) = case f kx x of + Nothing -> r + Just x' -> Bin sx kx x' Tip r + go' (NonEmptyDMap _ kx x l r) = Bin' $! balanceE kx x (go l) r + +makeUpdateMaxWithKey + :: (forall v. k v -> f v -> Maybe (f v)) + -> ( DMap k f -> DMap k f + , NonEmptyDMap k f -> Maybe (NonEmptyDMap k f)) +makeUpdateMaxWithKey f = (go, nonEmpty . go') + where + go Tip = Tip + go (Bin' ne) = go' ne + + go' (NonEmptyDMap sx kx x l Tip) = case f kx x of + Nothing -> l + Just x' -> Bin sx kx x' l Tip + go' (NonEmptyDMap _ kx x l r) = Bin' $! balanceE kx x l (go r) + +{-------------------------------------------------------------------- + Union Worker Builders +--------------------------------------------------------------------} + +makeUnion + :: GCompare k + => ( DMap k f -> DMap k f -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f) +makeUnion = (go, go') + where + go t1 Tip = t1 + go t1 (Bin _ kx x Tip Tip) = Bin' $! makeInsertR kx x t1 + go Tip t2 = t2 + go (Bin _ kx x Tip Tip) t2 = Bin' $! makeInsert kx x t2 + go t1@(Bin _ k1 x1 l1 r1) t2 = case fst (makeSplit k1) t2 of + (l2, r2) + | l1 `ptrEq` l1l2 && r1 `ptrEq` r1r2 -> t1 + | otherwise -> Bin' $! combineE k1 x1 l1l2 r1r2 + where !l1l2 = l1 `go` l2 + !r1r2 = r1 `go` r2 + + go' t1 (NonEmptyDMap _ kx x Tip Tip) = makeInsertR kx x (Bin' t1) + go' (NonEmptyDMap _ kx x Tip Tip) t2 = makeInsert kx x (Bin' t2) + go' t1@(NonEmptyDMap _ k1 x1 l1 r1) t2 = case snd (makeSplit k1) t2 of + (l2, r2) + | l1 `ptrEq` l1l2 && r1 `ptrEq` r1r2 -> t1 + | otherwise -> combineE k1 x1 l1l2 r1r2 + where !l1l2 = l1 `go` l2 + !r1r2 = r1 `go` r2 + +makeUnionWithKey + :: GCompare k + => (forall v. k v -> f v -> f v -> f v) + -> ( DMap k f -> DMap k f -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f) +makeUnionWithKey f = (go, go') + where + go t1 Tip = t1 + go Tip t2 = t2 + go (Bin _ k1 x1 l1 r1) t2 = case fst (makeSplitLookup k1) t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> Bin' $! combineE k1 x1 l1l2 r1r2 + Just x2 -> Bin' $! combineE k1 (f k1 x1 x2) l1l2 r1r2 + where !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + + go' (NonEmptyDMap _ k1 x1 l1 r1) t2 = case snd (makeSplitLookup k1) t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> combineE k1 x1 l1l2 r1r2 + Just x2 -> combineE k1 (f k1 x1 x2) l1l2 r1r2 + where !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + +{-------------------------------------------------------------------- + Difference Worker Builders +--------------------------------------------------------------------} + +makeDifference + :: GCompare k + => ( DMap k f -> DMap k g -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k g -> DMap k f + ) +makeDifference = (go, go') + where + go Tip _ = Tip + go t1 Tip = t1 + go t1 (Bin _ k2 _x2 l2 r2) = case fst (makeSplit k2) t1 of + (l1, r1) + | sizeE t1 == sizeE l1l2 + sizeE r1r2 -> t1 + | otherwise -> mergeE l1l2 r1r2 + where + !l1l2 = l1 `go` l2 + !r1r2 = r1 `go` r2 + + go' t1 (NonEmptyDMap _ k2 _x2 l2 r2) = case snd (makeSplit k2) t1 of + (l1, r1) + | sizeNE t1 == sizeE l1l2 + sizeE r1r2 -> Bin' t1 + | otherwise -> mergeE l1l2 r1r2 + where + !l1l2 = l1 `go` l2 + !r1r2 = r1 `go` r2 + +makeDifferenceWithKey + :: GCompare k + => (forall v. k v -> f v -> g v -> Maybe (f v)) + -> ( DMap k f -> DMap k g -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k g -> DMap k f + ) +makeDifferenceWithKey f = (go, go') + where + go Tip _ = Tip + go t1 Tip = t1 + go (Bin _ k1 x1 l1 r1) t2 = case fst (makeSplitLookup k1) t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> Bin' $! combineE k1 x1 l1l2 r1r2 + Just x2 -> case f k1 x1 x2 of + Nothing -> mergeE l1l2 r1r2 + Just x1x2 -> Bin' $! combineE k1 x1x2 l1l2 r1r2 + where !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + + go' (NonEmptyDMap _ k1 x1 l1 r1) t2 = case snd (makeSplitLookup k1) t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> Bin' $! combineE k1 x1 l1l2 r1r2 + Just x2 -> case f k1 x1 x2 of + Nothing -> mergeE l1l2 r1r2 + Just x1x2 -> Bin' $! combineE k1 x1x2 l1l2 r1r2 + where !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + +{-------------------------------------------------------------------- + Intersection Worker Builders +--------------------------------------------------------------------} + +makeIntersection + :: GCompare k + => ( DMap k f -> DMap k f -> DMap k f + , NonEmptyDMap k f -> NonEmptyDMap k f -> DMap k f + ) +makeIntersection = (go, go') + where + go Tip _ = Tip + go _ Tip = Tip + go t1@(Bin s1 k1 x1 l1 r1) t2 = + let !(l2, found, r2) = fst (makeSplitMember k1) t2 + !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + in if found + then if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 + then t1 + else Bin' $! combineE k1 x1 l1l2 r1r2 + else mergeE l1l2 r1r2 + + go' t1@(NonEmptyDMap s1 k1 x1 l1 r1) t2 = + let !(l2, found, r2) = snd (makeSplitMember k1) t2 + !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + in if found + then if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 + then Bin' t1 + else Bin' $! combineE k1 x1 l1l2 r1r2 + else mergeE l1l2 r1r2 + +makeIntersectionWithKey + :: forall k f g h + . GCompare k + => (forall v. k v -> f v -> g v -> h v) + -> ( DMap k f -> DMap k g -> DMap k h + , NonEmptyDMap k f -> NonEmptyDMap k g -> DMap k h + ) +makeIntersectionWithKey f = (go, go') + where + go :: DMap k f -> DMap k g -> DMap k h + go Tip _ = Tip + go _ Tip = Tip + go (Bin s1 k1 x1 l1 r1) t2 = + let !(l2, found, r2) = fst (makeSplitLookup k1) t2 + !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + in case found of + Nothing -> mergeE l1l2 r1r2 + Just x2 -> Bin' $! combineE k1 (f k1 x1 x2) l1l2 r1r2 + + go' :: NonEmptyDMap k f -> NonEmptyDMap k g -> DMap k h + go' (NonEmptyDMap s1 k1 x1 l1 r1) t2 = + let !(l2, found, r2) = snd (makeSplitLookup k1) t2 + !l1l2 = go l1 l2 + !r1r2 = go r1 r2 + in case found of + Nothing -> mergeE l1l2 r1r2 + Just x2 -> Bin' $! combineE k1 (f k1 x1 x2) l1l2 r1r2 + +{-------------------------------------------------------------------- + Submap Worker Builders +--------------------------------------------------------------------} + +makeSubmap' + :: GCompare k + => (forall v. k v -> k v -> f v -> g v -> Bool) + -> ( DMap k f -> DMap k g -> Bool + , NonEmptyDMap k f -> NonEmptyDMap k g -> Bool + ) +makeSubmap' f = (go, go') + where + go Tip _ = True + go _ Tip = False + go (Bin _ kx x l r) t + = case found of + Nothing -> False + Just (ky, y) -> f kx ky x y && go l lt && go r gt + where + (lt, found, gt) = fst (makeSplitLookupWithKey kx) t + + go' (NonEmptyDMap _ kx x l r) t + = case found of + Nothing -> False + Just (ky, y) -> f kx ky x y && go l lt && go r gt + where + (lt, found, gt) = snd (makeSplitLookupWithKey kx) t + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} + +makeFilterWithKey + :: GCompare k + => (forall v. k v -> f v -> Bool) + -> DMap k f -> DMap k f +makeFilterWithKey p = go + where + go Tip = Tip + go t@(Bin _ kx x l r) + | p kx x = if l' `ptrEq` l && r' `ptrEq` r + then t + else Bin' $! combineE kx x l' r' + | otherwise = mergeE l' r' + where !l' = go l + !r' = go r + +makePartitionWithKey + :: forall k f g h + . GCompare k + => (forall v. k v -> f v -> Bool) + -> DMap k f + -> (DMap k f, DMap k f) +makePartitionWithKey p m0 = toPair (go m0) + where + go :: DMap k f -> (DMap k f :*: DMap k f) + go Tip = (Tip :*: Tip) + go (Bin _ kx x l r) + | p kx x = ((Bin' $! combineE kx x l1 r1) :*: mergeE l2 r2) + | otherwise = (mergeE l1 r1 :*: (Bin' $! combineE kx x l2 r2)) + where + (l1 :*: l2) = go l + (r1 :*: r2) = go r + +makeMapMaybeWithKey + :: GCompare k + => (forall v. k v -> f v -> Maybe (g v)) + -> DMap k f + -> DMap k g +makeMapMaybeWithKey f = go + where + go Tip = Tip + go (Bin _ kx x l r) = case f kx x of + Just y -> Bin' $! combineE kx y (go l) (go r) + Nothing -> mergeE (go l) (go r) + +makeMapEitherWithKey + :: forall k f g h + . GCompare k + => (forall v. k v -> f v -> Either (g v) (h v)) + -> DMap k f + -> (DMap k g, DMap k h) +makeMapEitherWithKey f = toPair . go + where + go :: GCompare k + => DMap k f -> (DMap k g :*: DMap k h) + go Tip = (Tip :*: Tip) + go (Bin _ kx x l r) = case f kx x of + Left y -> ((Bin' $! combineE kx y l1 r1) :*: mergeE l2 r2) + Right z -> (mergeE l1 r1 :*: (Bin' $! combineE kx z l2 r2)) + where + (l1 :*: l2) = go l + (r1 :*: r2) = go r + +{-------------------------------------------------------------------- + Mapping Worker Builders +--------------------------------------------------------------------} + +makeMap + :: (forall v. f v -> g v) + -> ( DMap k f -> DMap k g + , NonEmptyDMap k f -> NonEmptyDMap k g + ) +makeMap f = (go, go') + where + go Tip = Tip + go (Bin' ne) = Bin' $! go' ne + + go' (NonEmptyDMap sx kx x l r) = NonEmptyDMap sx kx (f x) (go l) (go r) + +makeMapWithKey + :: (forall v. k v -> f v -> g v) + -> ( DMap k f -> DMap k g + , NonEmptyDMap k f -> NonEmptyDMap k g + ) +makeMapWithKey f = (go, go') + where + go Tip = Tip + go (Bin' ne) = Bin' $! go' ne + + go' (NonEmptyDMap sx kx x l r) = NonEmptyDMap sx kx (f kx x) (go l) (go r) + +makeTraverseWithKey + :: Applicative t + => (forall v. k v -> f v -> t (g v)) + -> ( DMap k f -> t (DMap k g) + , NonEmptyDMap k f -> t (NonEmptyDMap k g) + ) +makeTraverseWithKey f = (go, go') + where + go Tip = pure Tip + go (Bin' ne) = Bin' <$> go' ne + + go' (NonEmptyDMap 1 k v _ _) = (\v' -> NonEmptyDMap 1 k v' Tip Tip) <$> f k v + go' (NonEmptyDMap s k v l r) = flip (NonEmptyDMap s k) <$> go l <*> f k v <*> go r + +makeMapAccumLWithKey + :: (forall v. a -> k v -> f v -> (a, g v)) + -> ( a -> DMap k f -> (a, DMap k g) + , a -> NonEmptyDMap k f -> (a, NonEmptyDMap k g) + ) +makeMapAccumLWithKey f = (go, go') + where + go a Tip = (a,Tip) + go a (Bin' ne) = Bin' <$> go' a ne + + go' a (NonEmptyDMap sx kx x l r) = + let (a1,l') = go a l + (a2,x') = f a1 kx x + (a3,r') = go a2 r + in (a3, NonEmptyDMap sx kx x' l' r') + +makeMapAccumRWithKey + :: (forall v. a -> k v -> f v -> (a, g v)) + -> ( a -> DMap k f -> (a, DMap k g) + , a -> NonEmptyDMap k f -> (a, NonEmptyDMap k g) + ) +makeMapAccumRWithKey f = (go, go') + where + go a Tip = (a,Tip) + go a (Bin' ne) = Bin' <$> go' a ne + + go' a (NonEmptyDMap sx kx x l r) = + let (a1,r') = go a r + (a2,x') = f a1 kx x + (a3,l') = go a2 l + in (a3, NonEmptyDMap sx kx x' l' r') + +makeMapKeysMonotonic + :: (forall v. k1 v -> k2 v) + -> ( DMap k1 f -> DMap k2 f + , NonEmptyDMap k1 f -> NonEmptyDMap k2 f + ) +makeMapKeysMonotonic f = (go, go') + where + go Tip = Tip + go (Bin' ne) = Bin' $! go' ne + + go' (NonEmptyDMap sz k x l r) = NonEmptyDMap sz (f k) x (go l) (go r) + +{-------------------------------------------------------------------- + Folds +--------------------------------------------------------------------} + +makeFoldrWithKey + :: (forall v. k v -> f v -> b -> b) + -> b + -> DMap k f + -> b +makeFoldrWithKey f = go + where + go z Tip = z + go z (Bin _ kx x l r) = go (f kx x (go z r)) l + +makeFoldr1WithKey + :: (forall v. k v -> f v -> b -> b) + -> (forall v. k v -> f v -> b) + -> NonEmptyDMap k f -> b +makeFoldr1WithKey f g = go + where + go (NonEmptyDMap _ kx x l Tip) = makeFoldrWithKey f (g kx x) l + go (NonEmptyDMap _ kx x l (Bin' r)) = makeFoldrWithKey f (f kx x (go r)) l + +makeFoldlWithKey + :: (forall v. b -> k v -> f v -> b) + -> b + -> DMap k f + -> b +makeFoldlWithKey f = go + where + go z Tip = z + go z (Bin _ kx x l r) = go (f (go z l) kx x) r + +makeFoldl1WithKey + :: (forall v. b -> k v -> f v -> b) + -> (forall v. k v -> f v -> b) + -> NonEmptyDMap k f -> b +makeFoldl1WithKey f g = go + where + go (NonEmptyDMap _ kx x Tip r) = makeFoldlWithKey f (g kx x) r + go (NonEmptyDMap _ kx x (Bin' l) r) = makeFoldlWithKey f (f (go l) kx x) r + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists worker builders +--------------------------------------------------------------------} + +makeFromDistinctAscList + :: ( [DSum k f] -> DMap k f + , NonEmpty (DSum k f) -> NonEmptyDMap k f + ) +makeFromDistinctAscList + = ( \xs -> build const (length xs) xs + , \xs -> build' const (length xs) xs + ) + where + -- 1) use continutations so that we use heap space instead of stack space. + -- 2) special case for n==5 to build bushier trees. + + build + :: (DMap k f -> [DSum k f] -> b) + -> Int + -> [DSum k f] + -> b + build c 0 xs' = c Tip xs' + build c 5 xs' = case xs' of + ((k1:=>x1):(k2:=>x2):(k3:=>x3):(k4:=>x4):(k5:=>x5):xx) + -> c (binE k4 x4 (binE k2 x2 (singletonE k1 x1) (singletonE k3 x3)) (singletonE k5 x5)) xx + _ -> error "fromDistinctAscList build" + build c n xs' = seq nr $ build (buildR nr c) nl xs' + where + nl = n `div` 2 + nr = n - nl - 1 + + build' + :: (NonEmptyDMap k f -> [DSum k f] -> b) + -> Int + -> NonEmpty (DSum k f) + -> b + build' c 5 xs' = case xs' of + (k1:=>x1) :| ((k2:=>x2):(k3:=>x3):(k4:=>x4):(k5:=>x5):xx) + -> c (binNE k4 x4 (binE k2 x2 (singletonE k1 x1) (singletonE k3 x3)) (singletonE k5 x5)) xx + _ -> error "fromDistinctAscList build'" + build' c n xs' = seq nr $ build' (build'R nr c) nl xs' + where + nl = n `div` 2 + nr = n - nl - 1 + + -- TODO use 'NonEmpty' to hoist partiality to caller + buildR + :: Int + -> (DMap k f -> [DSum k f] -> b) + -> DMap k f -> [DSum k f] -> b + buildR n c l ((k:=>x):ys) = build (buildB l k x c) n ys + buildR _ _ _ [] = error "fromDistinctAscList buildR []" + + build'R + :: Int + -> (NonEmptyDMap k f -> [DSum k f] -> b) + -> NonEmptyDMap k f -> [DSum k f] -> b + build'R n c l ((k:=>x):ys) = build (build'B l k x c) n ys + build'R _ _ _ [] = error "fromDistinctAscList build'R []" + + buildB + :: DMap k f -> k v -> f v + -> (DMap k f -> a -> b) + -> DMap k f -> a -> b + buildB l k x c r zs = c (binE k x l r) zs + + build'B + :: NonEmptyDMap k f -> k v -> f v + -> (NonEmptyDMap k f -> a -> b) + -> DMap k f -> a -> b + build'B l k x c r zs = c (binNE k x (Bin' l) r) zs + +{-------------------------------------------------------------------- + Split Worker Builders +--------------------------------------------------------------------} + +makeSplit + :: forall k f v. GCompare k + => k v + -> ( DMap k f -> (DMap k f, DMap k f) + , NonEmptyDMap k f -> (DMap k f, DMap k f) + ) +makeSplit k = (toPair . go, toPair . go') + where + go :: DMap k f -> (DMap k f :*: DMap k f) + go Tip = (Tip :*: Tip) + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> (DMap k f :*: DMap k f) + go' (NonEmptyDMap _ kx x l r) = case gcompare k kx of + GLT -> let !(lt :*: gt) = go l in (lt :*: (Bin' $! combineE kx x gt r)) + GGT -> let !(lt :*: gt) = go r in ((Bin' $! combineE kx x l lt) :*: gt) + GEQ -> (l :*: r) +{-# INLINABLE makeSplit #-} + +makeSplitLookup + :: forall k f v + . GCompare k + => k v + -> ( DMap k f -> (DMap k f, Maybe (f v), DMap k f) + , NonEmptyDMap k f -> (DMap k f, Maybe (f v), DMap k f) + ) +makeSplitLookup k = (toTriple . go, toTriple . go') + where + go :: DMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f) + go Tip = Triple' Tip Nothing Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> Triple' (DMap k f) (Maybe (f v)) (DMap k f) + go' (NonEmptyDMap _ kx x l r) = case gcompare k kx of + GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (Bin' $! combineE kx x gt r) + GGT -> let !(Triple' lt z gt) = go r in Triple' (Bin' $! combineE kx x l lt) z gt + GEQ -> Triple' l (Just x) r + +makeSplitMember + :: forall k f v + . GCompare k + => k v + -> ( DMap k f -> (DMap k f, Bool, DMap k f) + , NonEmptyDMap k f -> (DMap k f, Bool, DMap k f) + ) +makeSplitMember k = (toTriple . go, toTriple . go') + where + go :: DMap k f -> Triple' (DMap k f) Bool (DMap k f) + go Tip = Triple' Tip False Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> Triple' (DMap k f) Bool (DMap k f) + go' (NonEmptyDMap _ kx x l r) = case gcompare k kx of + GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (Bin' $! combineE kx x gt r) + GGT -> let !(Triple' lt z gt) = go r in Triple' (Bin' $! combineE kx x l lt) z gt + GEQ -> Triple' l True r + +makeSplitLookupWithKey + :: forall k f v + . GCompare k + => k v + -> ( DMap k f -> (DMap k f, Maybe (k v, f v), DMap k f) + , NonEmptyDMap k f -> (DMap k f, Maybe (k v, f v), DMap k f) + ) +makeSplitLookupWithKey k = (toTriple . go, toTriple . go') + where + go :: DMap k f -> Triple' (DMap k f) (Maybe (k v, f v)) (DMap k f) + go Tip = Triple' Tip Nothing Tip + go (Bin' ne) = go' ne + + go' :: NonEmptyDMap k f -> Triple' (DMap k f) (Maybe (k v, f v)) (DMap k f) + go' (NonEmptyDMap _ kx x l r) = case gcompare k kx of + GLT -> let !(Triple' lt z gt) = go l in Triple' lt z (Bin' $! combineE kx x gt r) + GGT -> let !(Triple' lt z gt) = go r in Triple' (Bin' $! combineE kx x l lt) z gt + GEQ -> Triple' l (Just (kx, x)) r + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} + +makeShowsTree + :: (forall v. k v -> f v -> String) + -> Bool + -> ( [String] -> [String] -> DMap k f -> ShowS + , [String] -> [String] -> NonEmptyDMap k f -> ShowS + ) +makeShowsTree showelem wide = (go, go') + where + go lbars rbars t = case t of + Tip -> showsBars lbars . showString "|\n" + Bin' ne -> go' lbars rbars ne + + go' lbars rbars t = case t of + NonEmptyDMap _ kx x Tip Tip + -> showsBars lbars . showString (showelem kx x) . showString "\n" + NonEmptyDMap _ kx x l r + -> go (withBar rbars) (withEmpty rbars) r . + showWide wide rbars . + showsBars lbars . showString (showelem kx x) . showString "\n" . + showWide wide lbars . + go (withEmpty lbars) (withBar lbars) l + +makeShowsTreeHang + :: (forall v. k v -> f v -> String) + -> Bool + -> ( [String] -> DMap k f -> ShowS + , [String] -> NonEmptyDMap k f -> ShowS + ) +makeShowsTreeHang showelem wide = (go, go') + where + go bars t = case t of + Tip -> showsBars bars . showString "|\n" + Bin' ne -> go' bars ne + + go' bars t = case t of + NonEmptyDMap _ kx x Tip Tip + -> showsBars bars . showString (showelem kx x) . showString "\n" + NonEmptyDMap _ kx x l r + -> showsBars bars . showString (showelem kx x) . showString "\n" . + showWide wide bars . + go (withBar bars) l . + showWide wide bars . + go (withEmpty bars) r + +showWide :: Bool -> [String] -> String -> String +showWide wide bars + | wide = showString (concat (reverse bars)) . showString "|\n" + | otherwise = id + +showsBars :: [String] -> ShowS +showsBars bars + = case bars of + [] -> id + _ -> showString (concat (reverse (tail bars))) . showString node + +node :: String +node = "+--" + +withBar, withEmpty :: [String] -> [String] +withBar bars = "| ":bars +withEmpty bars = " ":bars + +{-------------------------------------------------------------------- + Assertions +--------------------------------------------------------------------} + +makeOrdered + :: forall k f + . GCompare k + => ( DMap k f -> Bool + , NonEmptyDMap k f -> Bool + ) +makeOrdered + = ( bounded (const True) (const True) + , bounded' (const True) (const True) + ) + where + bounded + :: (Some k -> Bool) + -> (Some k -> Bool) + -> DMap k f-> Bool + bounded lo hi t' = case t' of + Tip -> True + Bin' ne -> bounded' lo hi ne + + bounded' + :: (Some k -> Bool) + -> (Some k -> Bool) + -> NonEmptyDMap k f-> Bool + bounded' lo hi t' + = case t' of + NonEmptyDMap _ kx _ l r -> (lo (This kx)) && (hi (This kx)) && bounded lo (< This kx) l && bounded (> This kx) hi r + +-- | Exported only for "Debug.QuickCheck" +makeBalanced + :: ( DMap k f -> Bool + , NonEmptyDMap k f -> Bool + ) +makeBalanced = (go, go') + where + go t = case t of + Tip -> True + Bin' ne -> go' ne + + go' t = case t of + NonEmptyDMap _ _ _ l r -> (sizeE l + sizeE r <= 1 || (sizeE l <= delta*sizeE r && sizeE r <= delta*sizeE l)) && + go l && go r + +makeValidsize + :: ( DMap k f -> Bool + , NonEmptyDMap k f -> Bool + ) +makeValidsize = + ( \t -> realsize t == Just (sizeE t) + , \t -> realsize' t == Just (sizeNE t) + ) + where + realsize t = case t of + Tip -> Just 0 + Bin' ne -> realsize' ne + realsize' t = case t of + NonEmptyDMap sz _ _ l r -> case (realsize l,realsize r) of + (Just n,Just m) | n+m+1 == sz -> Just sz + _ -> Nothing + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} + +foldlStrict :: (a -> b -> a) -> a -> [b] -> a +foldlStrict f = go + where + go z [] = z + go z (x:xs) = z `seq` go (f z x) xs + +foldl1Strict + :: (a -> a -> a) + -> NonEmpty a -> a +foldl1Strict f (x NEL.:| xs) = foldlStrict f x xs diff --git a/src/Data/Dependent/Map/Lens.hs b/src/Data/Dependent/Map/Lens.hs index 0847528..51f9699 100644 --- a/src/Data/Dependent/Map/Lens.hs +++ b/src/Data/Dependent/Map/Lens.hs @@ -15,7 +15,7 @@ import Prelude hiding (lookup) import Control.Applicative (Applicative (pure)) #endif -import Data.Dependent.Map (DMap, alterF, insert, lookup) +import Data.Dependent.Map (DMap, alterF, adjustF) import Data.GADT.Compare (GCompare) @@ -83,5 +83,5 @@ dmat k f = alterF k f -- >>> DMap.fromList [AString :=> Identity "Shoe", AFloat :=> Identity 3.5] ^? dmix AInt -- Nothing dmix :: (GCompare k, Applicative f) => k v -> (g v -> f (g v)) -> DMap k g -> f (DMap k g) -dmix k f dmap = maybe (pure dmap) (fmap (flip (insert k) dmap) . f) $ lookup k dmap +dmix = adjustF {-# INLINE dmix #-} diff --git a/src/Data/Dependent/Map/NonEmpty.hs b/src/Data/Dependent/Map/NonEmpty.hs new file mode 100644 index 0000000..c2e793d --- /dev/null +++ b/src/Data/Dependent/Map/NonEmpty.hs @@ -0,0 +1,920 @@ +{-# LANGUAGE GADTs, RankNTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE CPP #-} +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702 +{-# LANGUAGE Trustworthy #-} +#endif +#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UndecidableInstances #-} +#endif +module Data.Dependent.Map.NonEmpty + ( NonEmptyDMap + , DSum(..), Some(..) + , GCompare(..), GOrdering(..) + + -- * non-empty specific + , nonEmpty + , toDMap + + -- * Operators + , (!), (\\) + + -- * Query + , size + , member + , notMember + , lookup + , findWithDefault + + -- * Construction + , singleton + + -- ** Insertion + , insert + , insertWith + , insertWith' + , insertWithKey + , insertWithKey' + , insertLookupWithKey + , insertLookupWithKey' + + -- ** Delete\/Update + , delete + , adjust + , adjustF + , adjustWithKey + , adjustWithKey' + , update + , updateWithKey + , updateLookupWithKey + , alter + , alterF + + -- * Combine + + -- ** Union + , union + , unionWithKey + , unions + , unionsWithKey + + -- ** Difference + , difference + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWithKey + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , mapAccumLWithKey + , mapAccumRWithKey + , mapKeysWith + , mapKeysMonotonic + + -- ** Fold + , foldWithKey + , foldrWithKey + , foldlWithKey + -- , foldlWithKey' + + -- * Conversion + , keys + , assocs + + -- ** Lists + , toList + , fromList + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , filter + , filterWithKey + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEitherWithKey + + , split + , splitLookup + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + + -- * Min\/Max + , findMin + , findMax + , lookupMin + , lookupMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMinWithKey + , updateMaxWithKey + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + ) where + +import Prelude hiding (null, lookup, map) +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative(..), (<$>)) +#endif +import Data.Dependent.Map.Internal +import Data.Dependent.Map.Internal2 +#if !MIN_VERSION_base(4,7,0) +import Data.Dependent.Map.Typeable ({- instance Typeable ... -}) +#endif + +import Data.Bifunctor (bimap) +import Data.Dependent.Sum +import Data.Constraint.Extras +import Data.GADT.Compare +import Data.GADT.Show +import Data.Maybe (isJust) +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.List.NonEmpty as NEL +#if !MIN_VERSION_base(4,8,0) +import Data.Monoid +#endif +import Data.Semigroup +import Data.Some +import Text.Read + +instance (GCompare k) => Semigroup (NonEmptyDMap k f) where + (<>) = union + +{-------------------------------------------------------------------- + NonEmpty* Specific +--------------------------------------------------------------------} + +toDMap :: NonEmptyDMap k f -> DMap k f +toDMap = Bin' + +{-------------------------------------------------------------------- + Operators +--------------------------------------------------------------------} +infixl 9 !,\\ -- + +-- | /O(log n)/. Find the value at a key. +-- Calls 'error' when the element can not be found. +-- +-- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map +-- > fromList [(5,'a'), (3,'b')] ! 5 == 'a' + +(!) :: GCompare k => NonEmptyDMap k f -> k v -> f v +(!) m k = find k m + +-- | Same as 'difference'. +(\\) :: GCompare k => NonEmptyDMap k f -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +m1 \\ m2 = difference m1 m2 + +-- #if __GLASGOW_HASKELL__ +-- +-- {-------------------------------------------------------------------- +-- A Data instance +-- --------------------------------------------------------------------} +-- +-- -- This instance preserves data abstraction at the cost of inefficiency. +-- -- We omit reflection services for the sake of data abstraction. +-- +-- instance (Data k, Data a, GCompare k) => Data (NonEmptyDMap k) where +-- gfoldl f z m = z fromList `f` toList m +-- toConstr _ = error "toConstr" +-- gunfold _ _ = error "gunfold" +-- dataTypeOf _ = mkNoRepType "Data.Map.Map" +-- dataCast2 f = gcast2 f +-- +-- #endif + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 +singleton :: k v -> f v -> NonEmptyDMap k f +singleton = singletonNE + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} + +-- | /O(1)/. The number of elements in the map. +size :: NonEmptyDMap k f -> Int +size = sizeNE + +-- | /O(log n)/. Is the key a member of the map? See also 'notMember'. +member :: GCompare k => k a -> NonEmptyDMap k f -> Bool +member k = isJust . lookupNE k + +-- | /O(log n)/. Is the key not a member of the map? See also 'member'. +notMember :: GCompare k => k v -> NonEmptyDMap k f -> Bool +notMember k m = not (member k m) + +-- | /O(log n)/. 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 :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> Maybe (f v) +lookup = lookupNE + +-- | /O(log n)/. Find the value at a key. +-- Calls 'error' when the element can not be found. +-- Consider using 'lookup' when elements may not be present. +find :: GCompare k => k v -> NonEmptyDMap k f -> f v +find k m = case lookupNE k m of + Nothing -> error "NonEmptyDMap.find: element not in the map" + Just v -> v + +-- | /O(log n)/. 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 :: GCompare k => f v -> k v -> NonEmptyDMap k f -> f v +findWithDefault def k m = case lookupNE k m of + Nothing -> def + Just v -> v + +{-------------------------------------------------------------------- + Insertion +--------------------------------------------------------------------} + +-- | /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. 'insert' is equivalent to +-- @'insertWith' 'const'@. +insert :: forall k f v. GCompare k => k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insert kx x = makeInsert kx x . Bin' + +-- | /O(log n)/. Insert a new key and value in the map if the key +-- is not already present. If the key is already present, @insertR@ +-- does nothing. +insertR :: forall k f v. GCompare k => k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insertR kx x = makeInsertR kx x . Bin' + +-- | /O(log n)/. Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ +-- will insert the entry @key :=> value@ into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the entry @key :=> f new_value old_value@. +insertWith :: GCompare k => (f v -> f v -> f v) -> k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insertWith f = insertWithKey (\_ x' y' -> f x' y') + +-- | Same as 'insertWith', but the combining function is applied strictly. +-- This is often the most desirable behavior. +insertWith' :: GCompare k => (f v -> f v -> f v) -> k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insertWith' f = insertWithKey' (\_ x' y' -> f x' y') + +-- | /O(log n)/. Insert with a function, combining key, new value and old value. +-- @'insertWithKey' f key value mp@ +-- will insert the entry @key :=> value@ into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the entry @key :=> f key new_value old_value@. +-- Note that the key passed to f is the same key passed to 'insertWithKey'. +insertWithKey :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insertWithKey f kx x = makeInsertWithKey f kx x . Bin' + +-- | Same as 'insertWithKey', but the combining function is applied strictly. +insertWithKey' :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> NonEmptyDMap k f -> NonEmptyDMap k f +insertWithKey' f kx x = makeInsertWithKey' f kx x . Bin' + +-- | /O(log n)/. Combines insert operation with old value retrieval. +-- 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@). +insertLookupWithKey :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> NonEmptyDMap k f + -> (Maybe (f v), NonEmptyDMap k f) +insertLookupWithKey f kx x = makeInsertLookupWithKey f kx x . Bin' + +-- | /O(log n)/. A strict version of 'insertLookupWithKey'. +insertLookupWithKey' :: forall k f v. GCompare k => (k v -> f v -> f v -> f v) -> k v -> f v -> NonEmptyDMap k f + -> (Maybe (f v), NonEmptyDMap k f) +insertLookupWithKey' f kx x = makeInsertLookupWithKey' f kx x . Bin' + +{-------------------------------------------------------------------- + Deletion + [delete] is the inlined version of [deleteWith (\k x -> Nothing)] +--------------------------------------------------------------------} + +-- | /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 :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +delete k = snd $ makeDelete k + +-- | /O(log n)/. Update a value at a specific key with the result of the provided function. +-- When the key is not +-- a member of the map, the original map is returned. +adjust :: GCompare k => (f v -> f v) -> k v -> NonEmptyDMap k f -> NonEmptyDMap k f +adjust f = adjustWithKey (\_ x -> f x) + +-- | Works the same as 'adjust' except the new value is return in some 'Applicative' @f@. +adjustF + :: forall k f v g + . (GCompare k, Applicative f) + => k v + -> (g v -> f (g v)) + -> NonEmptyDMap k g -> f (NonEmptyDMap k g) +adjustF k f = snd $ makeAdjustF f k + +-- | /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 :: GCompare k => (k v -> f v -> f v) -> k v -> NonEmptyDMap k f -> NonEmptyDMap k f +adjustWithKey f k = snd $ makeAdjustWithKey f k + +-- | /O(log n)/. A strict version of 'adjustWithKey'. +adjustWithKey' :: GCompare k => (k v -> f v -> f v) -> k v -> NonEmptyDMap k f -> NonEmptyDMap k f +adjustWithKey' f k = snd $ makeAdjustWithKey' f k + +-- | /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 :: GCompare k => (f v -> Maybe (f v)) -> k v -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +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 :: forall k f v. GCompare k => (k v -> f v -> Maybe (f v)) -> k v -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +updateWithKey f k = snd $ makeUpdateWithKey f k + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +updateLookupWithKey :: forall k f v. GCompare k => (k v -> f v -> Maybe (f v)) -> k v -> NonEmptyDMap k f -> (Maybe (f v), Maybe (NonEmptyDMap k f)) +updateLookupWithKey f k = snd $ makeUpdateLookupWithKey f k + +-- | /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 :: forall k f v. GCompare k => (Maybe (f v) -> Maybe (f v)) -> k v -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +alter k f = snd $ makeAlter k f + +-- | Works the same as 'alter' except the new value is return in some 'Functor' @f@. +-- In short : @(\v' -> alter (const v') k dm) <$> f (lookup k dm)@ +alterF :: forall k f v g. (GCompare k, Functor f) => k v -> (Maybe (g v) -> f (Maybe (g v))) -> NonEmptyDMap k g -> f (Maybe (NonEmptyDMap k g)) +alterF k f = snd $ makeAlterF k f + +{-------------------------------------------------------------------- + Indexing +--------------------------------------------------------------------} + +-- | /O(log n)/. Return the /index/ of a key. The index is a number from +-- /0/ up to, but not including, the 'size' of the map. Calls 'error' when +-- the key is not a 'member' of the map. +findIndex :: GCompare k => k v -> NonEmptyDMap k f -> Int +findIndex k t + = case lookupIndex k t of + Nothing -> error "Map.findIndex: element is not in the map" + Just idx -> idx + +-- | /O(log n)/. Lookup the /index/ of a key. The index is a number from +-- /0/ up to, but not including, the 'size' of the map. +lookupIndex :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> Maybe Int +lookupIndex k = snd $ makeLookupIndex k + +-- | /O(log n)/. Retrieve an element by /index/. Calls 'error' when an +-- invalid index is used. +elemAt :: Int -> NonEmptyDMap k f -> DSum k f +elemAt = snd makElemAt + +-- | /O(log n)/. Update the element at /index/. Does nothing when an +-- invalid index is used. +updateAt :: (forall v. k v -> f v -> Maybe (f v)) -> Int -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +updateAt f i0 = snd $ makeUpdateAt f i0 + +-- | /O(log n)/. Delete the element at /index/. +-- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@). +deleteAt :: Int -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +deleteAt i m + = updateAt (\_ _ -> Nothing) i m + + +{-------------------------------------------------------------------- + Minimal, Maximal +--------------------------------------------------------------------} + +-- | /O(log n)/. The minimal key of the map. Calls 'error' is the map is empty. +findMin :: NonEmptyDMap k f -> DSum k f +findMin m = lookupMin m + +lookupMin :: NonEmptyDMap k f -> DSum k f +lookupMin = snd makeLookupMin + +-- | /O(log n)/. The maximal key of the map. Calls 'error' is the map is empty. +findMax :: NonEmptyDMap k f -> DSum k f +findMax m = lookupMax m + +lookupMax :: NonEmptyDMap k f -> DSum k f +lookupMax = snd makeLookupMax + +-- | /O(log n)/. Delete the minimal key. Returns an empty map if the map is empty. +deleteMin :: NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +deleteMin = snd makeDeleteMin + +-- | /O(log n)/. Delete the maximal key. Returns an empty map if the map is empty. +deleteMax :: NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +deleteMax = snd makeDeleteMax + +-- | /O(log n)/. Delete and find the minimal element. +-- +-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) +-- > deleteFindMin Error: can not return the minimal element of an empty map +deleteFindMin :: NonEmptyDMap k f -> (DSum k f, Maybe (NonEmptyDMap k f)) +deleteFindMin = minViewWithKey + +-- | /O(log n)/. Delete and find the maximal element. +-- +-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) +-- > deleteFindMax empty Error: can not return the maximal element of an empty map +deleteFindMax :: NonEmptyDMap k f -> (DSum k f, Maybe (NonEmptyDMap k f)) +deleteFindMax = maxViewWithKey + +-- | /O(log n)/. Update the value at the minimal key. +updateMinWithKey :: (forall v. k v -> f v -> Maybe (f v)) -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +updateMinWithKey f = snd $ makeUpdateMinWithKey f + +-- | /O(log n)/. Update the value at the maximal key. +updateMaxWithKey :: (forall v. k v -> f v -> Maybe (f v)) -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +updateMaxWithKey f = snd $ makeUpdateMaxWithKey f + +-- | /O(log n)/. Retrieves the minimal (key :=> value) entry of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +minViewWithKey :: forall k f . NonEmptyDMap k f -> (DSum k f, Maybe (NonEmptyDMap k f)) +minViewWithKey = fmap nonEmpty . minViewWithKeyNE + +-- | /O(log n)/. Retrieves the maximal (key :=> value) entry of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +maxViewWithKey :: forall k f . NonEmptyDMap k f -> (DSum k f, Maybe (NonEmptyDMap k f)) +maxViewWithKey = fmap nonEmpty . maxViewWithKeyNE + +{-------------------------------------------------------------------- + Union. +--------------------------------------------------------------------} + +-- | The union of a non-empty list of maps: +-- (@'unions' == 'Prelude.foldl1' 'union'@). +unions :: GCompare k => NonEmpty (NonEmptyDMap k f) -> NonEmptyDMap k f +unions ts + = foldl1Strict union ts + +-- | The union of a list of maps, with a combining operation: +-- (@'unionsWithKey' f == 'Prelude.foldl' ('unionWithKey' f) 'empty'@). +unionsWithKey :: GCompare k => (forall v. k v -> f v -> f v -> f v) -> NonEmpty (NonEmptyDMap k f) -> NonEmptyDMap k f +unionsWithKey f ts + = foldl1Strict (unionWithKey f) ts + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- 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'@). +union :: GCompare k => NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +union t1 t2 = snd makeUnion t1 t2 + +{-------------------------------------------------------------------- + Union with a combining function +--------------------------------------------------------------------} + +-- | /O(n+m)/. +-- Union with a combining function. +unionWithKey :: GCompare k => (forall v. k v -> f v -> f v -> f v) -> NonEmptyDMap k f -> NonEmptyDMap k f -> NonEmptyDMap k f +unionWithKey f t1 t2 = snd (makeUnionWithKey f) t1 t2 + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} + +-- | /O(m * log (n\/m + 1)), m <= n/. Difference of two maps. +-- Return elements of the first map not existing in the second map. +difference :: GCompare k => NonEmptyDMap k f -> NonEmptyDMap k g -> Maybe (NonEmptyDMap k f) +difference t1 t2 = nonEmpty $ snd makeDifference t1 t2 + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +differenceWithKey :: GCompare k => (forall v. k v -> f v -> g v -> Maybe (f v)) -> NonEmptyDMap k f -> NonEmptyDMap k g -> Maybe (NonEmptyDMap k f) +differenceWithKey f t1 t2 = nonEmpty $ snd (makeDifferenceWithKey f) t1 t2 + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} + +-- | /O(m * log (n\/m + 1), m <= n/. Intersection of two maps. +-- Return data in the first map for the keys existing in both maps. +-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). +intersection :: GCompare k => NonEmptyDMap k f -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +intersection t1 t2 = nonEmpty $ snd makeIntersection t1 t2 + +-- | /O(m * log (n\/m + 1), m <= n/. Intersection with a combining function. +intersectionWithKey :: GCompare k => (forall v. k v -> f v -> g v -> h v) -> NonEmptyDMap k f -> NonEmptyDMap k g -> Maybe (NonEmptyDMap k h) +intersectionWithKey f t1 t2 = nonEmpty $ snd (makeIntersectionWithKey f) t1 t2 + +{-------------------------------------------------------------------- + Submap +--------------------------------------------------------------------} +-- | /O(n+m)/. +-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' 'eqTagged')@). +-- +isSubmapOf + :: forall k f + . (GCompare k, Has' Eq k f) + => NonEmptyDMap k f -> NonEmptyDMap k f -> Bool +isSubmapOf m1 m2 = isSubmapOfBy (\k _ x0 x1 -> has' @Eq @f k (x0 == x1)) m1 m2 + +{- | /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 keys and values. +-} +isSubmapOfBy :: GCompare k => (forall v. k v -> k v -> f v -> g v -> Bool) -> NonEmptyDMap k f -> NonEmptyDMap k g -> Bool +isSubmapOfBy f t1 t2 + = (sizeNE t1 <= sizeNE t2) && (submap' f t1 t2) + +submap' :: GCompare k => (forall v. k v -> k v -> f v -> g v -> Bool) -> NonEmptyDMap k f -> NonEmptyDMap k g -> Bool +submap' f t1 t2 = snd (makeSubmap' f) t1 t2 + +-- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). +-- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' 'eqTagged'@). +isProperSubmapOf + :: forall k f + . (GCompare k, Has' Eq k f) + => NonEmptyDMap k f -> NonEmptyDMap k f -> Bool +isProperSubmapOf m1 m2 + = isProperSubmapOfBy (\k _ x0 x1 -> has' @Eq @f k (x0 == x1)) m1 m2 + +{- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). + The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when + @m1@ and @m2@ are not equal, + all keys in @m1@ are in @m2@, and when @f@ returns 'True' when + applied to their respective keys and values. +-} +isProperSubmapOfBy :: GCompare k => (forall v. k v -> k v -> f v -> g v -> Bool) -> NonEmptyDMap k f -> NonEmptyDMap k g -> Bool +isProperSubmapOfBy f t1 t2 + = (sizeNE t1 < sizeNE t2) && (submap' f t1 t2) + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} + +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +filterWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k f) +filterWithKey p = nonEmpty . makeFilterWithKey p . Bin' + +-- | /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. See also 'split'. +partitionWithKey :: GCompare k => (forall v. k v -> f v -> Bool) -> NonEmptyDMap k f -> (Maybe (NonEmptyDMap k f), Maybe (NonEmptyDMap k f)) +partitionWithKey p t = (nonEmpty t1, nonEmpty t2) + where (t1, t2) = makePartitionWithKey p $ Bin' t + +-- | /O(n)/. Map values and collect the 'Just' results. +mapMaybe :: GCompare k => (forall v. f v -> Maybe (g v)) -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k g) +mapMaybe f = mapMaybeWithKey (const f) + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +mapMaybeWithKey :: GCompare k => (forall v. k v -> f v -> Maybe (g v)) -> NonEmptyDMap k f -> Maybe (NonEmptyDMap k g) +mapMaybeWithKey f = nonEmpty . makeMapMaybeWithKey f . Bin' + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +mapEitherWithKey :: GCompare k => + (forall v. k v -> f v -> Either (g v) (h v)) -> NonEmptyDMap k f -> (Maybe (NonEmptyDMap k g), Maybe (NonEmptyDMap k h)) +mapEitherWithKey f t = (nonEmpty t1, nonEmpty t2) + where (t1, t2) = makeMapEitherWithKey f $ Bin' t + +{-------------------------------------------------------------------- + Mapping +--------------------------------------------------------------------} + +-- | /O(n)/. Map a function over all values in the map. +map :: (forall v. f v -> g v) -> NonEmptyDMap k f -> NonEmptyDMap k g +map f = snd $ makeMap f + +-- | /O(n)/. Map a function over all values in the map. +mapWithKey :: (forall v. k v -> f v -> g v) -> NonEmptyDMap k f -> NonEmptyDMap k g +mapWithKey f = snd $ makeMapWithKey f + +-- | /O(n)/. +-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- That is, behaves exactly like a regular 'traverse' except that the traversing +-- function also has access to the key associated with a value. +traverseWithKey :: Applicative t => (forall v. k v -> f v -> t (g v)) -> NonEmptyDMap k f -> t (NonEmptyDMap k g) +traverseWithKey f = snd $ makeTraverseWithKey f + +-- | /O(n)/. The function 'mapAccumLWithKey' threads an accumulating +-- argument throught the map in ascending order of keys. +mapAccumLWithKey :: (forall v. a -> k v -> f v -> (a, g v)) -> a -> NonEmptyDMap k f -> (a, NonEmptyDMap k g) +mapAccumLWithKey f = snd $ makeMapAccumLWithKey f + +-- | /O(n)/. The function 'mapAccumRWithKey' threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey :: (forall v. a -> k v -> f v -> (a, g v)) -> a -> NonEmptyDMap k f -> (a, NonEmptyDMap k g) +mapAccumRWithKey f = snd $ makeMapAccumRWithKey f + +-- | /O(n*log n)/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. +mapKeysWith :: GCompare k2 => (forall v. k2 v -> f v -> f v -> f v) -> (forall v. k1 v -> k2 v) -> NonEmptyDMap k1 f -> NonEmptyDMap k2 f +mapKeysWith c f = fromListWithKey c . NEL.map fFirst . toList + where fFirst (x :=> y) = (f x :=> y) + + +-- | /O(n)/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has better performance than 'mapKeys'. +mapKeysMonotonic :: (forall v. k1 v -> k2 v) -> NonEmptyDMap k1 f -> NonEmptyDMap k2 f +mapKeysMonotonic f = snd $ makeMapKeysMonotonic f + +{-------------------------------------------------------------------- + Folds +--------------------------------------------------------------------} + +-- | /O(n)/. Fold the keys and values in the map, such that +-- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. +-- +-- This is identical to 'foldrWithKey', and you should use that one instead of +-- this one. This name is kept for backward compatibility. +foldWithKey :: (forall v. k v -> f v -> b -> b) -> b -> NonEmptyDMap k f -> b +foldWithKey = foldrWithKey +{-# DEPRECATED foldWithKey "Use foldrWithKey instead" #-} + +-- | /O(n)/. Post-order fold. The function will be applied from the lowest +-- value to the highest. +foldrWithKey :: (forall v. k v -> f v -> b -> b) -> b -> NonEmptyDMap k f -> b +foldrWithKey f b = makeFoldrWithKey f b . Bin' + +-- | /O(n)/. Pre-order fold. The function will be applied from the highest +-- value to the lowest. +foldlWithKey :: (forall v. b -> k v -> f v -> b) -> b -> NonEmptyDMap k f -> b +foldlWithKey f b = makeFoldlWithKey f b . Bin' + +foldr1WithKey + :: (forall v. k v -> f v -> b -> b) + -> (forall v. k v -> f v -> b) + -> NonEmptyDMap k f -> b +foldr1WithKey = makeFoldr1WithKey + +foldl1WithKey + :: (forall v. b -> k v -> f v -> b) + -> (forall v. k v -> f v -> b) + -> NonEmptyDMap k f -> b +foldl1WithKey = makeFoldl1WithKey + +{- +-- | /O(n)/. A strict version of 'foldlWithKey'. +foldlWithKey' :: (b -> k -> a -> b) -> b -> NonEmptyDMap k -> b +foldlWithKey' f = go + where + go z Tip = z + go z (NonEmptyDMap _ kx x l r) = z `seq` go (f (go z l) kx x) r +-} + +{-------------------------------------------------------------------- + List variations +--------------------------------------------------------------------} + +-- | /O(n)/. Return all keys of the map in ascending order. +-- +-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] +-- > keys empty == [] + +keys :: NonEmptyDMap k f -> NonEmpty (Some k) +keys m = (\(k :=> _) -> This k) <$> assocs m + +-- | /O(n)/. Return all key\/value pairs in the map in ascending key order. +assocs :: NonEmptyDMap k f -> NonEmpty (DSum k f) +assocs m = toList m + +{-------------------------------------------------------------------- + Lists + use [foldlStrict] to reduce demand on the control-stack +--------------------------------------------------------------------} + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'. +-- If the list contains more than one value for the same key, the last value +-- for the key is retained. +fromList :: GCompare k => NonEmpty (DSum k f) -> NonEmptyDMap k f +fromList xs + = unions $ ins <$> xs + where + ins :: GCompare k => DSum k f -> NonEmptyDMap k f + ins (k :=> x) = singletonNE k x + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +fromListWithKey :: GCompare k => (forall v. k v -> f v -> f v -> f v) -> NonEmpty (DSum k f) -> NonEmptyDMap k f +fromListWithKey f xs + = unionsWithKey f $ ins <$> xs + where + ins :: GCompare k => DSum k f -> NonEmptyDMap k f + ins (k :=> x) = singletonNE k x + +-- | /O(n)/. Convert to a list of key\/value pairs. +toList :: NonEmptyDMap k f -> NonEmpty (DSum k f) +toList t = toAscList t + +-- | /O(n)/. Convert to an ascending list. +toAscList :: NonEmptyDMap k f -> NonEmpty (DSum k f) +toAscList t = foldr1WithKey + (\k x xs -> NEL.cons (k :=> x) xs) + (\k x -> pure $ k :=> x) + t + +-- | /O(n)/. Convert to a descending list. +toDescList :: NonEmptyDMap k f -> NonEmpty (DSum k f) +toDescList t = foldl1WithKey + (\xs k x1 -> NEL.cons (k :=> x1) xs) + (\k x -> pure $ k :=> x) + t + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists can be done in linear time. + + Note that if [xs] is ascending that: + fromAscList xs == fromList xs + fromAscListWith f xs == fromListWith f xs +--------------------------------------------------------------------} + +-- | /O(n)/. Build a map from an ascending list in linear time. +-- /The precondition (input list is ascending) is not checked./ +fromAscList :: GEq k => NonEmpty (DSum k f) -> NonEmptyDMap k f +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs + +-- | /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 + :: forall k f + . GEq k + => (forall v. k v -> f v -> f v -> f v) + -> NonEmpty (DSum k f) + -> NonEmptyDMap k f +fromAscListWithKey f xs + = fromDistinctAscList $ combineEq xs + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq :: NonEmpty (DSum k f) -> NonEmpty (DSum k f) + combineEq (z :| []) = z :| [] + combineEq (z@(kz :=> zz) :| (x@(kx :=> xx):xs')) = + case geq kx kz of + Just Refl -> let yy = f kx xx zz in combineEq $ (kx :=> yy) :| xs' + Nothing -> NEL.cons z $ combineEq $ x :| xs' + + +-- | /O(n)/. Build a map from an ascending list of distinct elements in linear time. +-- /The precondition is not checked./ +fromDistinctAscList :: NonEmpty (DSum k f) -> NonEmptyDMap k f +fromDistinctAscList = snd makeFromDistinctAscList + +{-------------------------------------------------------------------- + Split +--------------------------------------------------------------------} + +-- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ where +-- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. +-- Any key equal to @k@ is found in neither @map1@ nor @map2@. +split :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> (Maybe (NonEmptyDMap k f), Maybe (NonEmptyDMap k f)) +split = fmap (bimap nonEmpty nonEmpty) . snd . makeSplit +{-# INLINABLE split #-} + +-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just +-- like 'split' but also returns @'lookup' k map@. +splitLookup :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> (Maybe (NonEmptyDMap k f), Maybe (f v), Maybe (NonEmptyDMap k f)) +splitLookup = fmap f . snd . makeSplitLookup + where f (a, b, c) = (nonEmpty a, b, nonEmpty c) + +-- | /O(log n)/. The expression (@'splitMember' k map@) splits a map just +-- like 'split' but also returns @'member' k map@. +splitMember :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> (Maybe (NonEmptyDMap k f), Bool, Maybe (NonEmptyDMap k f)) +splitMember = fmap f . snd . makeSplitMember + where f (a, b, c) = (nonEmpty a, b, nonEmpty c) + +-- | /O(log n)/. +splitLookupWithKey :: forall k f v. GCompare k => k v -> NonEmptyDMap k f -> (Maybe (NonEmptyDMap k f), Maybe (k v, f v), Maybe (NonEmptyDMap k f)) +splitLookupWithKey = fmap f . snd . makeSplitLookupWithKey + where f (a, b, c) = (nonEmpty a, b, nonEmpty c) + +{-------------------------------------------------------------------- + Eq converts the tree to a list. In a lazy setting, this + actually seems one of the faster methods to compare two trees + and it is certainly the simplest :-) +--------------------------------------------------------------------} +instance (GEq k, Has' Eq k f) => Eq (NonEmptyDMap k f) where + t1 == t2 = (sizeNE t1 == sizeNE t2) && (toAscList t1 == toAscList t2) + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance (GCompare k, Has' Eq k f, Has' Ord k f) => Ord (NonEmptyDMap k f) where + compare m1 m2 = compare (toAscList m1) (toAscList m2) + +{-------------------------------------------------------------------- + Read +--------------------------------------------------------------------} + +instance (GCompare k, GRead k, Has' Read k f) => Read (NonEmptyDMap k f) where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} +instance (GShow k, Has' Show k f) => Show (NonEmptyDMap k f) where + showsPrec p m = showParen (p>10) + ( showString "fromList " + . showsPrec 11 (toList m) + ) + +-- | /O(n)/. Show the tree that implements the map. The tree is shown +-- in a compressed, hanging format. See 'showTreeWith'. +showTree :: (GShow k, Has' Show k f) => NonEmptyDMap k f -> String +showTree m + = showTreeWith showElem True False m + where + showElem :: (GShow k, Has' Show k f) => k v -> f v -> String + showElem k x = show (k :=> x) + + +{- | /O(n)/. The expression (@'showTreeWith' showelem hang wide map@) shows + the tree that implements the map. Elements are shown using the @showElem@ function. If @hang@ is + 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If + @wide@ is 'True', an extra wide version is shown. +-} +showTreeWith :: (forall v. k v -> f v -> String) -> Bool -> Bool -> NonEmptyDMap k f -> String +showTreeWith showelem hang wide t + | hang = (showsTreeHang showelem wide [] t) "" + | otherwise = (showsTree showelem wide [] [] t) "" + +showsTree :: (forall v. k v -> f v -> String) -> Bool -> [String] -> [String] -> NonEmptyDMap k f -> ShowS +showsTree showelem wide = snd $ makeShowsTree showelem wide + +showsTreeHang :: (forall v. k v -> f v -> String) -> Bool -> [String] -> NonEmptyDMap k f -> ShowS +showsTreeHang showelem wide = snd $ makeShowsTreeHang showelem wide + +{-------------------------------------------------------------------- + Assertions +--------------------------------------------------------------------} + +-- | /O(n)/. Test if the internal map structure is valid. +valid :: GCompare k => NonEmptyDMap k f -> Bool +valid t + = balanced t && ordered t && validsize t + +ordered :: GCompare k => NonEmptyDMap k f -> Bool +ordered = snd makeOrdered + +-- | Exported only for "Debug.QuickCheck" +balanced :: NonEmptyDMap k f -> Bool +balanced = snd makeBalanced + +validsize :: NonEmptyDMap k f -> Bool +validsize = snd makeValidsize diff --git a/src/Data/Dependent/Map/NonEmpty/Lens.hs b/src/Data/Dependent/Map/NonEmpty/Lens.hs new file mode 100644 index 0000000..1503140 --- /dev/null +++ b/src/Data/Dependent/Map/NonEmpty/Lens.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE CPP #-} +-- | +-- Some functions for using lenses with 'NonEmptyDMap'. +module Data.Dependent.Map.NonEmpty.Lens + ( -- * At + dmat + -- * Ix + , dmix + ) + where + +import Prelude hiding (lookup) + +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative (pure)) +#endif + +import Data.Dependent.Map.NonEmpty (NonEmptyDMap, adjustF, alterF) + +import Data.GADT.Compare (GCompare) + +-- | +-- These functions have been specialised for use with 'NonEmptyDMap' but without any of the +-- specific 'lens' types used so that we have compatilibity without needing the +-- dependency just for these functions. +-- + +-- | +-- This is equivalent to the from : +-- +-- @ +-- type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t +-- +-- at :: Index m -> Lens' m (Maybe (IxValue m)) +-- @ +-- +-- So the type of 'dmat' is equivalent to: +-- +-- @ +-- dmat :: GCompare k => Lens' (NonEmptyDMap k f) (Maybe (f v)) +-- @ +-- +-- >>> NonEmptyDMap.fromList [AInt :=> Identity 33, AFloat :=> Identity 3.5] & dmat AString ?~ "Hat" +-- NonEmptyDMap.fromList [AString :=> Identity "Hat", AInt :=> Identity 33, AFloat :=> Identity 3.5] +-- +-- >>> NonEmptyDMap.fromList [AString :=> Identity "Shoe", AInt :=> Identity 33, AFloat :=> Identity 3.5] ^? dmat AFloat +-- Just (AFloat :=> 3.5) +-- +dmat :: (GCompare k, Functor f) => k v -> (Maybe (g v) -> f (Maybe (g v))) -> NonEmptyDMap k g -> f (Maybe (NonEmptyDMap k g)) +dmat k f = alterF k f +{-# INLINE dmat #-} + +-- | +-- This is equivalent to the from : +-- +-- @ +-- type Traversal s t a b = forall f. Applicative f => (a -> f b) -> s -> f t +-- +-- ix :: Index m -> Traversal' m (IxValue m) +-- @ +-- +-- So the type of 'dmix' is equivalent to: +-- +-- @ +-- dmix :: GCompare k => k v -> Traversal' (NonEmptyDMap k f) (f v) +-- @ +-- +-- /NB:/ Setting the value of this +-- +-- will only set the value in 'dmix' if it is already present. +-- +-- If you want to be able to insert /missing/ values, you want 'dmat'. +-- +-- >>> NonEmptyDMap.fromList [AString :=> Identity "Shoe", AInt :=> Identity 33, AFloat :=> Identity 3.5] & dmix AInt %~ f +-- NonEmptyDMap.fromList [AString :=> Identity "Shoe", AInt :=> Identity (f 33), AFloat :=> Identity 3.5] +-- +-- >>> NonEmptyDMap.fromList [AString :=> Identity "Shoe", AInt :=> Identity 33, AFloat :=> Identity 3.5] & dmix AString .~ "Hat" +-- NonEmptyDMap.fromList [AString :=> Identity "Hat", AInt :=> Identity 33, AFloat :=> Identity 3.5] +-- +-- >>> NonEmptyDMap.fromList [AString :=> Identity "Shoe", AInt :=> Identity 33, AFloat :=> Identity 3.5] ^? dmix AFloat +-- Just (AFloat :=> 3.5) +-- +-- >>> NonEmptyDMap.fromList [AString :=> Identity "Shoe", AFloat :=> Identity 3.5] ^? dmix AInt +-- Nothing +dmix :: (GCompare k, Applicative f) => k v -> (g v -> f (g v)) -> NonEmptyDMap k g -> f (NonEmptyDMap k g) +dmix = adjustF +{-# INLINE dmix #-}