diff --git a/containers-tests/benchmarks/LookupGE/LookupGE_Map.hs b/containers-tests/benchmarks/LookupGE/LookupGE_Map.hs index 56cabf999..5721f464e 100644 --- a/containers-tests/benchmarks/LookupGE/LookupGE_Map.hs +++ b/containers-tests/benchmarks/LookupGE/LookupGE_Map.hs @@ -14,11 +14,11 @@ lookupGE2 :: Ord k => k -> Map k a -> Maybe (k,a) lookupGE2 = go where go !_ Tip = Nothing - go !k (Bin _ kx x l r) = + go !k (NE (Bin' _ kx x l r)) = case compare k kx of LT -> case go k l of Nothing -> Just (kx,x) - ret -> ret + ret -> ret GT -> go k r EQ -> Just (kx,x) {-# INLINABLE lookupGE2 #-} @@ -27,7 +27,7 @@ lookupGE3 :: Ord k => k -> Map k a -> Maybe (k,a) lookupGE3 = go Nothing where go def !_ Tip = def - go def !k (Bin _ kx x l r) = + go def !k (NE (Bin' _ kx x l r)) = case compare k kx of LT -> go (Just (kx,x)) k l GT -> go def k r @@ -38,16 +38,16 @@ lookupGE4 :: Ord k => k -> Map k a -> Maybe (k,a) lookupGE4 k = k `seq` goNothing where goNothing Tip = Nothing - goNothing (Bin _ kx x l r) = case compare k kx of - LT -> goJust kx x l - EQ -> Just (kx, x) - GT -> goNothing r + goNothing (NE (Bin' _ kx x l r)) = case compare k kx of + LT -> goJust kx x l + EQ -> Just (kx, x) + GT -> goNothing r goJust ky y Tip = Just (ky, y) - goJust ky y (Bin _ kx x l r) = case compare k kx of - LT -> goJust kx x l - EQ -> Just (kx, x) - GT -> goJust ky y r + goJust ky y (NE (Bin' _ kx x l r)) = case compare k kx of + LT -> goJust kx x l + EQ -> Just (kx, x) + GT -> goJust ky y r {-# INLINABLE lookupGE4 #-} ------------------------------------------------------------------------------- @@ -55,9 +55,9 @@ lookupGE4 k = k `seq` goNothing ------------------------------------------------------------------------------- findMinMaybe :: Map k a -> Maybe (k,a) -findMinMaybe (Bin _ kx x Tip _) = Just (kx,x) -findMinMaybe (Bin _ _ _ l _) = findMinMaybe l -findMinMaybe Tip = Nothing +findMinMaybe (NE (Bin' _ kx x Tip _)) = Just (kx,x) +findMinMaybe (NE (Bin' _ _ _ l _)) = findMinMaybe l +findMinMaybe Tip = Nothing #ifdef TESTING ------------------------------------------------------------------------------- diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs index bed8d2b7e..f00a8a6f0 100644 --- a/containers-tests/tests/map-properties.hs +++ b/containers-tests/tests/map-properties.hs @@ -7,7 +7,7 @@ import Data.Map.Merge.Strict import Data.Map.Lazy as Data.Map hiding (showTree, showTreeWith) import Data.Map.Merge.Lazy #endif -import Data.Map.Internal (Map (..), link2, link, bin) +import Data.Map.Internal (Map (..), NonEmptyMap(..), link2, link, bin) import Data.Map.Internal.Debug (showTree, showTreeWith, balanced) import Control.Applicative (Const(Const, getConst), pure, (<$>), (<*>)) @@ -332,15 +332,15 @@ mkArb step n vOuter <- liftGen arbitrary vInner <- liftGen arbitrary if dir - then return (Bin 2 q vOuter (singleton p vInner) Tip) - else return (Bin 2 p vOuter Tip (singleton q vInner)) + then return (NE (Bin' 2 q vOuter (singleton p vInner) Tip)) + else return (NE (Bin' 2 p vOuter Tip (singleton q vInner))) | otherwise = do -- This assumes a balance factor of delta = 3 let upper = (3*(n - 1)) `quot` 4 let lower = (n + 2) `quot` 4 ln <- liftGen $ choose (lower, upper) let rn = n - ln - 1 - liftM4 (\lt x v rt -> Bin n x v lt rt) (mkArb step ln) step (liftGen arbitrary) (mkArb step rn) + liftM4 (\lt x v rt -> NE (Bin' n x v lt rt)) (mkArb step ln) step (liftGen arbitrary) (mkArb step rn) -- A type with a peculiar Eq instance designed to make sure keys -- come from where they're supposed to. diff --git a/containers-tests/tests/set-properties.hs b/containers-tests/tests/set-properties.hs index 08658bf2e..f47aa0c36 100644 --- a/containers-tests/tests/set-properties.hs +++ b/containers-tests/tests/set-properties.hs @@ -247,15 +247,15 @@ mkArb step n p <- step q <- step if dir - then return (Bin 2 q (singleton p) Tip) - else return (Bin 2 p Tip (singleton q)) + then return (NE $ Bin' 2 q (singleton p) Tip) + else return (NE $ Bin' 2 p Tip (singleton q)) | otherwise = do -- This assumes a balance factor of delta = 3 let upper = (3*(n - 1)) `quot` 4 let lower = (n + 2) `quot` 4 ln <- liftGen $ choose (lower, upper) let rn = n - ln - 1 - liftM3 (\lt x rt -> Bin n x lt rt) (mkArb step ln) step (mkArb step rn) + liftM3 (\lt x rt -> NE $ Bin' n x lt rt) (mkArb step ln) step (mkArb step rn) -- | Given a strictly increasing list of elements, produce an arbitrarily -- shaped set with exactly those elements. diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index 4b4d809fe..7ef4fe4ee 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -3,6 +3,7 @@ {-# LANGUAGE PatternGuards #-} #if defined(__GLASGOW_HASKELL__) {-# LANGUAGE DeriveLift #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RoleAnnotations #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE Trustworthy #-} @@ -124,12 +125,14 @@ -- Currently in GHC 7.0, when type has 2 constructors, a forward conditional -- jump is made when successfully matching second constructor. Successful match -- of first constructor results in the forward jump not taken. --- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip +-- On GHC 7.0, reordering constructors from Tip | NE to NE | Tip -- improves the benchmark by up to 10% on x86. module Data.Map.Internal ( -- * Map type Map(..) -- instance Eq,Show,Read + , pattern Bin + , NonEmptyMap(..) -- instance Eq,Show,Read , Size -- * Operators @@ -137,57 +140,59 @@ module Data.Map.Internal ( -- * Query , null - , size - , member - , notMember - , lookup - , findWithDefault - , lookupLT - , lookupGT - , lookupLE - , lookupGE + , nonEmpty + , size, sizeNE + , member, memberNE + , notMember, notMemberNE + , lookup, lookupNE + , findWithDefault, findWithDefaultNE + , lookupLT, lookupLTNE + , lookupGT, lookupGTNE + , lookupLE, lookupLENE + , lookupGE, lookupGENE -- * Construction , empty - , singleton + , singleton, singletonNE -- ** Insertion - , insert - , insertWith - , insertWithKey - , insertLookupWithKey + , insert, insertNE + , insertWith, insertWithNE + , insertWithKey, insertWithKeyNE + , insertLookupWithKey, insertLookupWithKeyNE -- ** Delete\/Update - , delete - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - , alterF + , delete, deleteNE + , adjust, adjustNE + , adjustWithKey, adjustWithKeyNE + , update, updateNE + , updateWithKey, updateWithKeyNE + , updateLookupWithKey, updateLookupWithKeyNE + , alter, alterNE + , alterF, alterFNE -- * Combine -- ** Union - , union - , unionWith - , unionWithKey + , union, unionNE + , unionWith, unionWithNE + , unionWithKey, unionWithKeyNE , unions , unionsWith -- ** Difference - , difference - , differenceWith - , differenceWithKey + , difference, differenceNE + , differenceWith, differenceWithNE + , differenceWithKey, differenceWithKeyNE -- ** Intersection - , intersection - , intersectionWith - , intersectionWithKey + , intersection, intersectionNE + , intersectionWith, intersectionWithNE + , intersectionWithKey, intersectionWithKeyNE -- ** Disjoint - , disjoint + , disjoint, disjointNE + -- ** Compose , compose @@ -197,7 +202,7 @@ module Data.Map.Internal ( , SimpleWhenMatched , runWhenMatched , runWhenMissing - , merge + , merge, mergeNE -- *** @WhenMatched@ tactics , zipWithMaybeMatched , zipWithMatched @@ -212,7 +217,7 @@ module Data.Map.Internal ( -- ** Applicative general combining function , WhenMissing (..) , WhenMatched (..) - , mergeA + , mergeA, mergeANE -- *** @WhenMatched@ tactics -- | The tactics described for 'merge' work for @@ -235,16 +240,16 @@ module Data.Map.Internal ( -- * Traversal -- ** Map - , map + , map, mapNE , mapWithKey - , traverseWithKey - , traverseMaybeWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey + , traverseWithKey, traverseWithKeyNE + , traverseMaybeWithKey, traverseMaybeWithKeyNE + , mapAccum, mapAccumNE + , mapAccumWithKey, mapAccumWithKeyNE + , mapAccumRWithKey, mapAccumRWithKeyNE , mapKeys , mapKeysWith - , mapKeysMonotonic + , mapKeysMonotonic, mapKeysMonotonicNE -- * Folds , foldr @@ -294,7 +299,7 @@ module Data.Map.Internal ( , dropWhileAntitone , spanAntitone - , restrictKeys + , restrictKeys, restrictKeysNE , withoutKeys , partition , partitionWithKey @@ -305,7 +310,7 @@ module Data.Map.Internal ( , mapEitherWithKey , split - , splitLookup + , splitLookup, splitLookupNE , splitRoot -- * Submap @@ -346,7 +351,7 @@ module Data.Map.Internal ( #ifdef __GLASGOW_HASKELL__ , atKeyPlain #endif - , bin + , bin, binNE , balance , balanceL , balanceR @@ -458,18 +463,28 @@ m1 \\ m2 = difference m1 m2 -- their union @m1 <> m2@ maps @k@ to @a1@. -- See Note: Order of constructors -data Map k a = Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a) +data Map k a = NE {-# UNPACK #-} !(NonEmptyMap k a) | Tip +data NonEmptyMap k a = Bin' {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a) + type Size = Int +#if __GLASGOW_HASKELL__ >= 802 +{-# COMPLETE Bin, Tip #-} +#endif #ifdef __GLASGOW_HASKELL__ +pattern Bin :: Size -> k -> a -> Map k a -> Map k a -> Map k a +pattern Bin s k a l r = NE (Bin' s k a l r) + type role Map nominal representational +type role NonEmptyMap nominal representational #endif #ifdef __GLASGOW_HASKELL__ -- | @since FIXME deriving instance (Lift k, Lift a) => Lift (Map k a) +deriving instance (Lift k, Lift a) => Lift (NonEmptyMap k a) #endif instance (Ord k) => Monoid (Map k v) where @@ -516,10 +531,16 @@ mapDataType = mkDataType "Data.Map.Internal.Map" [fromListConstr] -- > Data.Map.null (singleton 1 'a') == False null :: Map k a -> Bool -null Tip = True -null (Bin {}) = False +null Tip = True +null (NE (Bin' {})) = False {-# INLINE null #-} +-- | /O(1)/. Return 'Just' if the set is not empty. +nonEmpty :: Map k a -> Maybe (NonEmptyMap k a) +nonEmpty Tip = Nothing +nonEmpty (NE ne) = Just ne +{-# INLINE nonEmpty #-} + -- | \(O(1)\). The number of elements in the map. -- -- > size empty == 0 @@ -527,10 +548,12 @@ null (Bin {}) = False -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 size :: Map k a -> Int -size Tip = 0 -size (Bin sz _ _ _ _) = sz +size Tip = 0 +size (NE ne) = sizeNE ne {-# INLINE size #-} +sizeNE :: NonEmptyMap k a -> Int +sizeNE (Bin' sz _ _ _ _) = sz -- | \(O(\log n)\). Lookup the value at a key in the map. -- @@ -561,37 +584,39 @@ size (Bin sz _ _ _ _) = sz -- > John's currency: Just "Euro" -- > Pete's currency: Nothing lookup :: Ord k => k -> Map k a -> Maybe a -lookup = go - where - go !_ Tip = Nothing - go k (Bin _ kx x l r) = case compare k kx of - LT -> go k l - GT -> go k r - EQ -> Just x +lookup !_ Tip = Nothing +lookup k (NE ne) = lookupNE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE lookup #-} #else {-# INLINE lookup #-} #endif +lookupNE :: Ord k => k -> NonEmptyMap k a -> Maybe a +lookupNE k (Bin' _ kx x l r) = case compare k kx of + LT -> lookup k l + GT -> lookup k r + EQ -> Just x + -- | \(O(\log n)\). Is the key a member of the map? See also 'notMember'. -- -- > member 5 (fromList [(5,'a'), (3,'b')]) == True -- > member 1 (fromList [(5,'a'), (3,'b')]) == False member :: Ord k => k -> Map k a -> Bool -member = go - where - go !_ Tip = False - go k (Bin _ kx _ l r) = case compare k kx of - LT -> go k l - GT -> go k r - EQ -> True +member !_ Tip = False +member k (NE ne) = memberNE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE member #-} #else {-# INLINE member #-} #endif +memberNE :: Ord k => k -> NonEmptyMap k a -> Bool +memberNE k (Bin' _ kx _ l r) = case compare k kx of + LT -> member k l + GT -> member k r + EQ -> True + -- | \(O(\log n)\). Is the key not a member of the map? See also 'member'. -- -- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False @@ -605,22 +630,26 @@ notMember k m = not $ member k m {-# INLINE notMember #-} #endif +notMemberNE :: Ord k => k -> NonEmptyMap k a -> Bool +notMemberNE k m = not $ memberNE k m + -- | \(O(\log n)\). Find the value at a key. -- Calls 'error' when the element can not be found. find :: Ord k => k -> Map k a -> a -find = go - where - go !_ Tip = error "Map.!: given key is not an element in the map" - go k (Bin _ kx x l r) = case compare k kx of - LT -> go k l - GT -> go k r - EQ -> x +find !_ Tip = error "Map.!: given key is not an element in the map" +find k (NE ne) = findNE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE find #-} #else {-# INLINE find #-} #endif +findNE :: Ord k => k -> NonEmptyMap k a -> a +findNE k (Bin' _ kx x l r) = case compare k kx of + LT -> find k l + GT -> find k r + EQ -> x + -- | \(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. @@ -628,61 +657,64 @@ find = go -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' findWithDefault :: Ord k => a -> k -> Map k a -> a -findWithDefault = go - where - go def !_ Tip = def - go def k (Bin _ kx x l r) = case compare k kx of - LT -> go def k l - GT -> go def k r - EQ -> x +findWithDefault def !_ Tip = def +findWithDefault def k (NE ne) = findWithDefaultNE def k ne #if __GLASGOW_HASKELL__ {-# INLINABLE findWithDefault #-} #else {-# INLINE findWithDefault #-} #endif +findWithDefaultNE :: Ord k => a -> k -> NonEmptyMap k a -> a +findWithDefaultNE def k (Bin' _ kx x l r) = case compare k kx of + LT -> findWithDefault def k l + GT -> findWithDefault def k r + EQ -> x + -- | \(O(\log n)\). Find largest key smaller than the given one and return the -- corresponding (key, value) pair. -- -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') lookupLT :: Ord k => k -> Map k v -> Maybe (k, v) -lookupLT = goNothing - where - goNothing !_ Tip = Nothing - goNothing k (Bin _ kx x l r) | k <= kx = goNothing k l - | otherwise = goJust k kx x r - - goJust !_ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) | k <= kx = goJust k kx' x' l - | otherwise = goJust k kx x r +lookupLT !_ Tip = Nothing +lookupLT k (NE ne) = lookupLTNE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE lookupLT #-} #else {-# INLINE lookupLT #-} #endif +lookupLTNE :: Ord k => k -> NonEmptyMap k v -> Maybe (k, v) +lookupLTNE k (Bin' _ kx x l r) | k <= kx = lookupLT k l + | otherwise = go kx x r + where + go kx' x' Tip = Just (kx', x') + go kx' x' (NE (Bin' _ ky y l' r')) | k <= ky = go kx' x' l' + | otherwise = go ky y r' + -- | \(O(\log n)\). Find smallest key greater than the given one and return the -- corresponding (key, value) pair. -- -- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') -- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing lookupGT :: Ord k => k -> Map k v -> Maybe (k, v) -lookupGT = goNothing - where - goNothing !_ Tip = Nothing - goNothing k (Bin _ kx x l r) | k < kx = goJust k kx x l - | otherwise = goNothing k r - - goJust !_ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) | k < kx = goJust k kx x l - | otherwise = goJust k kx' x' r +lookupGT !_ Tip = Nothing +lookupGT k (NE ne) = lookupGTNE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE lookupGT #-} #else {-# INLINE lookupGT #-} #endif +lookupGTNE :: Ord k => k -> NonEmptyMap k v -> Maybe (k, v) +lookupGTNE k (Bin' _ kx x l r) | k < kx = go kx x l + | otherwise = lookupGT k r + where + go kx' x' Tip = Just (kx', x') + go kx' x' (NE (Bin' _ ky y l' r')) | k < ky = go ky y l' + | otherwise = go kx' x' r' + -- | \(O(\log n)\). Find largest key smaller or equal to the given one and return -- the corresponding (key, value) pair. -- @@ -690,23 +722,26 @@ lookupGT = goNothing -- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') -- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') lookupLE :: Ord k => k -> Map k v -> Maybe (k, v) -lookupLE = goNothing - where - goNothing !_ Tip = Nothing - goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goNothing k l - EQ -> Just (kx, x) - GT -> goJust k kx x r - - goJust !_ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l - EQ -> Just (kx, x) - GT -> goJust k kx x r +lookupLE !_ Tip = Nothing +lookupLE k (NE ne) = lookupLENE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE lookupLE #-} #else {-# INLINE lookupLE #-} #endif +lookupLENE :: Ord k => k -> NonEmptyMap k v -> Maybe (k, v) +lookupLENE k (Bin' _ kx x l r) = case compare k kx of + LT -> lookupLE k l + EQ -> Just (kx, x) + GT -> go kx x r + where + go kx' x' Tip = Just (kx', x') + go kx' x' (NE (Bin' _ ky y l' r')) = case compare k ky of + LT -> go kx' x' l' + EQ -> Just (ky, y) + GT -> go ky y r' + -- | \(O(\log n)\). Find smallest key greater or equal to the given one and return -- the corresponding (key, value) pair. -- @@ -714,23 +749,26 @@ lookupLE = goNothing -- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') -- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing lookupGE :: Ord k => k -> Map k v -> Maybe (k, v) -lookupGE = goNothing - where - goNothing !_ Tip = Nothing - goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l - EQ -> Just (kx, x) - GT -> goNothing k r - - goJust !_ kx' x' Tip = Just (kx', x') - goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l - EQ -> Just (kx, x) - GT -> goJust k kx' x' r +lookupGE !_ Tip = Nothing +lookupGE k (NE ne) = lookupGENE k ne #if __GLASGOW_HASKELL__ {-# INLINABLE lookupGE #-} #else {-# INLINE lookupGE #-} #endif +lookupGENE :: Ord k => k -> NonEmptyMap k v -> Maybe (k, v) +lookupGENE k (Bin' _ kx x l r) = case compare k kx of + LT -> go kx x l + EQ -> Just (kx, x) + GT -> lookupGE k r + where + go kx' x' Tip = Just (kx', x') + go kx' x' (NE (Bin' _ ky y l' r')) = case compare k ky of + LT -> go ky y l' + EQ -> Just (ky, y) + GT -> go kx' x' r' + {-------------------------------------------------------------------- Construction --------------------------------------------------------------------} @@ -749,9 +787,12 @@ empty = Tip -- > size (singleton 1 'a') == 1 singleton :: k -> a -> Map k a -singleton k x = Bin 1 k x Tip Tip +singleton k x = NE $ Bin' 1 k x Tip Tip {-# INLINE singleton #-} +singletonNE :: k -> a -> NonEmptyMap k a +singletonNE k x = Bin' 1 k x Tip Tip + {-------------------------------------------------------------------- Insertion --------------------------------------------------------------------} @@ -767,28 +808,42 @@ singleton k x = Bin 1 k x Tip Tip -- See Note: Type of local 'go' function -- See Note: Avoiding worker/wrapper insert :: Ord k => k -> a -> Map k a -> Map k a -insert kx0 = go kx0 kx0 - where - -- Unlike insertR, we only get sharing here - -- when the inserted value is at the same address - -- as the present value. We try anyway; this condition - -- seems particularly likely to occur in 'union'. - go :: Ord k => k -> k -> a -> Map k a -> Map k a - go orig !_ x Tip = singleton (lazy orig) x - go orig !kx x t@(Bin sz ky y l r) = - case compare kx ky of - LT | l' `ptrEq` l -> t - | otherwise -> balanceL ky y l' r - where !l' = go orig kx x l - GT | r' `ptrEq` r -> t - | otherwise -> balanceR ky y l r' - where !r' = go orig kx x r - EQ | x `ptrEq` y && (lazy orig `seq` (orig `ptrEq` ky)) -> t - | otherwise -> Bin sz (lazy orig) x l r +insert k x0 m0 = case insertReturningDifferent k k x0 m0 of + Nothing -> m0 + Just q -> NE q + +insertNE :: Ord k => k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertNE k x0 m0 = case insertReturningDifferentNE k k x0 m0 of + Nothing -> m0 + Just q -> q + +-- | Returns 'Nothing' if the element is already in the Map, and 'Just s' if a +-- new set had to be created to contain it. +-- +-- Unlike insertR, we only get sharing here when the inserted value is at the +-- same address as the present value. We try anyway; this condition seems +-- particularly likely to occur in 'union'. +insertReturningDifferent :: Ord k => k -> k -> a -> Map k a -> Maybe (NonEmptyMap k a) +insertReturningDifferent orig !_ x Tip = Just $ singletonNE (lazy orig) x +insertReturningDifferent orig !k x (NE ne) = insertReturningDifferentNE orig k x ne + +insertReturningDifferentNE :: Ord k => k -> k -> a -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +insertReturningDifferentNE orig !kx x (Bin' sz ky y l r) = case compare kx ky of + LT -> case insertReturningDifferent orig kx x l of + Nothing -> Nothing + Just l' -> Just $! balanceLNE ky y l' r + GT -> case insertReturningDifferent orig kx x r of + Nothing -> Nothing + Just r' -> Just $! balanceRNE ky y l r' + EQ | x `ptrEq` y && (lazy orig `seq` (orig `ptrEq` ky)) -> Nothing + | otherwise -> Just $ Bin' sz (lazy orig) x l r + #if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} +{-# INLINABLE insertNE #-} #else {-# INLINE insert #-} +{-# INLINE insertNE #-} #endif #ifndef __GLASGOW_HASKELL__ @@ -816,23 +871,35 @@ lazy a = a -- See Note: Type of local 'go' function -- See Note: Avoiding worker/wrapper insertR :: Ord k => k -> a -> Map k a -> Map k a -insertR kx0 = go kx0 kx0 - where - go :: Ord k => k -> k -> a -> Map k a -> Map k a - go orig !_ x Tip = singleton (lazy orig) x - go orig !kx x t@(Bin _ ky y l r) = - case compare kx ky of - LT | l' `ptrEq` l -> t - | otherwise -> balanceL ky y l' r - where !l' = go orig kx x l - GT | r' `ptrEq` r -> t - | otherwise -> balanceR ky y l r' - where !r' = go orig kx x r - EQ -> t +insertR k x0 m0 = case insertRReturningDifferent k k x0 m0 of + Nothing -> m0 + Just q -> NE q + +insertRNE :: Ord k => k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertRNE k x0 m0 = case insertRReturningDifferentNE k k x0 m0 of + Nothing -> m0 + Just q -> q + +insertRReturningDifferent :: Ord k => k -> k -> a -> Map k a -> Maybe (NonEmptyMap k a) +insertRReturningDifferent orig !_ x Tip = Just $ singletonNE (lazy orig) x +insertRReturningDifferent orig !k x (NE ne) = insertRReturningDifferentNE orig k x ne + +insertRReturningDifferentNE :: Ord k => k -> k -> a -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +insertRReturningDifferentNE orig !kx x (Bin' _ ky y l r) = case compare kx ky of + LT -> case insertRReturningDifferent orig kx x l of + Nothing -> Nothing + Just l' -> Just $! balanceLNE ky y l' r + GT -> case insertRReturningDifferent orig kx x r of + Nothing -> Nothing + Just r' -> Just $! balanceRNE ky y l r' + EQ -> Nothing + #if __GLASGOW_HASKELL__ {-# INLINABLE insertR #-} +{-# INLINABLE insertRNE #-} #else {-# INLINE insertR #-} +{-# INLINE insertRNE #-} #endif -- | \(O(\log n)\). Insert with a function, combining new value and old value. @@ -846,24 +913,30 @@ insertR kx0 = go kx0 kx0 -- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWith = go - where - -- We have no hope of making pointer equality tricks work - -- here, because lazy insertWith *always* changes the tree, - -- either adding a new entry or replacing an element with a - -- thunk. - go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a - go _ !kx x Tip = singleton kx x - go f !kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go f kx x l) r - GT -> balanceR ky y l (go f kx x r) - EQ -> Bin sy kx (f x y) l r +insertWith f k v m = NE $ insertWithToNE f k v m + +insertWithNE :: Ord k => (a -> a -> a) -> k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertWithNE f k v m = insertWithToNE f k v $ NE m + +-- We have no hope of making pointer equality tricks work +-- here, because lazy insertWith *always* changes the tree, +-- either adding a new entry or replacing an element with a +-- thunk. +insertWithToNE :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> NonEmptyMap k a +insertWithToNE _ !kx x Tip = singletonNE kx x +insertWithToNE f !kx x (NE (Bin' sy ky y l r)) = + case compare kx ky of + LT -> balanceLNE ky y (insertWithToNE f kx x l) r + GT -> balanceRNE ky y l (insertWithToNE f kx x r) + EQ -> Bin' sy kx (f x y) l r +{-# INLINE insertWithToNE #-} #if __GLASGOW_HASKELL__ {-# INLINABLE insertWith #-} +{-# INLINABLE insertWithNE #-} #else {-# INLINE insertWith #-} +{-# INLINE insertWithNE #-} #endif -- | A helper function for 'unionWith'. When the key is already in @@ -872,19 +945,26 @@ insertWith = go -- new value. insertWithR :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithR = go - where - go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a - go _ !kx x Tip = singleton kx x - go f !kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go f kx x l) r - GT -> balanceR ky y l (go f kx x r) - EQ -> Bin sy ky (f y x) l r +insertWithR f k v m = NE $ insertWithRToNE f k v m + +insertWithRNE :: Ord k => (a -> a -> a) -> k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertWithRNE f k v m = insertWithRToNE f k v $ NE m + +insertWithRToNE :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> NonEmptyMap k a +insertWithRToNE _ !kx x Tip = singletonNE kx x +insertWithRToNE f !kx x (NE (Bin' sy ky y l r)) = + case compare kx ky of + LT -> balanceLNE ky y (insertWithRToNE f kx x l) r + GT -> balanceRNE ky y l (insertWithRToNE f kx x r) + EQ -> Bin' sy ky (f y x) l r +{-# INLINE insertWithRToNE #-} + #if __GLASGOW_HASKELL__ {-# INLINABLE insertWithR #-} +{-# INLINABLE insertWithRNE #-} #else {-# INLINE insertWithR #-} +{-# INLINE insertWithRNE #-} #endif -- | \(O(\log n)\). Insert with a function, combining key, new value and old value. @@ -898,22 +978,31 @@ insertWithR = go -- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] -- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] -- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" - --- See Note: Type of local 'go' function insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithKey = go - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a - go _ !kx x Tip = singleton kx x - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go f kx x l) r - GT -> balanceR ky y l (go f kx x r) - EQ -> Bin sy kx (f kx x y) l r +insertWithKey f k v m = NE $ insertWithKeyToNE f k v m + +insertWithKeyNE :: Ord k => (k -> a -> a -> a) -> k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertWithKeyNE f k v m = insertWithKeyToNE f k v $ NE m + +-- We have no hope of making pointer equality tricks work +-- here, because lazy insertWithKey *always* changes the tree, +-- either adding a new entry or replacing an element with a +-- thunk. +insertWithKeyToNE :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> NonEmptyMap k a +insertWithKeyToNE _ !kx x Tip = singletonNE kx x +insertWithKeyToNE f !kx x (NE (Bin' sy ky y l r)) = + case compare kx ky of + LT -> balanceLNE ky y (insertWithKeyToNE f kx x l) r + GT -> balanceRNE ky y l (insertWithKeyToNE f kx x r) + EQ -> Bin' sy kx (f kx x y) l r +{-# INLINE insertWithKeyToNE #-} + #if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKey #-} +{-# INLINABLE insertWithKeyNE #-} #else {-# INLINE insertWithKey #-} +{-# INLINE insertWithKeyNE #-} #endif -- | A helper function for 'unionWithKey'. When the key is already in @@ -921,15 +1010,24 @@ insertWithKey = go -- function is flipped--it is applied to the old value and then the -- new value. insertWithKeyR :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a -insertWithKeyR = go - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a - go _ !kx x Tip = singleton kx x - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> balanceL ky y (go f kx x l) r - GT -> balanceR ky y l (go f kx x r) - EQ -> Bin sy ky (f ky y x) l r +insertWithKeyR f k v m = NE $ insertWithKeyRToNE f k v m + +insertWithKeyRNE :: Ord k => (k -> a -> a -> a) -> k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertWithKeyRNE f k v m = insertWithKeyRToNE f k v $ NE m + +-- We have no hope of making pointer equality tricks work +-- here, because lazy insertWithKeyR *always* changes the tree, +-- either adding a new entry or replacing an element with a +-- thunk. +insertWithKeyRToNE :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> NonEmptyMap k a +insertWithKeyRToNE _ !kx x Tip = singletonNE kx x +insertWithKeyRToNE f !kx x (NE (Bin' sy ky y l r)) = + case compare kx ky of + LT -> balanceLNE ky y (insertWithKeyRToNE f kx x l) r + GT -> balanceRNE ky y l (insertWithKeyRToNE f kx x r) + EQ -> Bin' sy ky (f ky y x) l r +{-# INLINE insertWithKeyRToNE #-} + #if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKeyR #-} #else @@ -955,23 +1053,30 @@ insertWithKeyR = go -- See Note: Type of local 'go' function insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a) -insertLookupWithKey f0 k0 x0 = toPair . go f0 k0 x0 - where - go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a) - go _ !kx x Tip = (Nothing :*: singleton kx x) - go f kx x (Bin sy ky y l r) = - case compare kx ky of - LT -> let !(found :*: l') = go f kx x l - !t' = balanceL ky y l' r - in (found :*: t') - GT -> let !(found :*: r') = go f kx x r - !t' = balanceR ky y l r' - in (found :*: t') - EQ -> (Just y :*: Bin sy kx (f kx x y) l r) +insertLookupWithKey f0 k0 x0 = fmap NE . toPair . insertLookupWithKeyToNE f0 k0 x0 + +insertLookupWithKeyNE :: Ord k => (k -> a -> a -> a) -> k -> a -> NonEmptyMap k a + -> (Maybe a, NonEmptyMap k a) +insertLookupWithKeyNE f0 k0 x0 = toPair . insertLookupWithKeyToNE f0 k0 x0 . NE + +insertLookupWithKeyToNE :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (NonEmptyMap k a) +insertLookupWithKeyToNE _ !kx x Tip = (Nothing :*: singletonNE kx x) +insertLookupWithKeyToNE f kx x (NE (Bin' sy ky y l r)) = + case compare kx ky of + LT -> let !(found :*: l') = insertLookupWithKeyToNE f kx x l + !t' = balanceLNE ky y l' r + in (found :*: t') + GT -> let !(found :*: r') = insertLookupWithKeyToNE f kx x r + !t' = balanceRNE ky y l r' + in (found :*: t') + EQ -> (Just y :*: Bin' sy kx (f kx x y) l r) + #if __GLASGOW_HASKELL__ {-# INLINABLE insertLookupWithKey #-} +{-# INLINABLE insertLookupWithKeyNE #-} #else {-# INLINE insertLookupWithKey #-} +{-# INLINE insertLookupWithKeyNE #-} #endif {-------------------------------------------------------------------- @@ -984,25 +1089,37 @@ insertLookupWithKey f0 k0 x0 = toPair . go f0 k0 x0 -- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] -- > delete 5 empty == empty --- See Note: Type of local 'go' function delete :: Ord k => k -> Map k a -> Map k a -delete = go - where - go :: Ord k => k -> Map k a -> Map k a - go !_ Tip = Tip - go k t@(Bin _ kx x l r) = - case compare k kx of - LT | l' `ptrEq` l -> t - | otherwise -> balanceR kx x l' r - where !l' = go k l - GT | r' `ptrEq` r -> t - | otherwise -> balanceL kx x l r' - where !r' = go k r - EQ -> glue l r +delete !_ Tip = Tip +delete k s0 = case deleteReturningDifferent k s0 of + Nothing -> s0 + Just s -> s + +deleteNE :: Ord k => k -> NonEmptyMap k a -> Map k a +deleteNE k s0 = case deleteReturningDifferentNE k s0 of + Nothing -> NE s0 + Just s -> s + +deleteReturningDifferent :: Ord k => k -> Map k a -> Maybe (Map k a) +deleteReturningDifferent !_ Tip = Nothing +deleteReturningDifferent k (NE ne) = deleteReturningDifferentNE k ne + +deleteReturningDifferentNE :: Ord k => k -> NonEmptyMap k a -> Maybe (Map k a) +deleteReturningDifferentNE !k (Bin' _ kx x l r) = case compare k kx of + LT -> case deleteReturningDifferent k l of + Nothing -> Nothing + Just l' -> Just $ balanceR kx x l' r + GT -> case deleteReturningDifferent k r of + Nothing -> Nothing + Just r' -> Just $ balanceL kx x l r' + EQ -> Just $ glue l r + #if __GLASGOW_HASKELL__ {-# INLINABLE delete #-} +{-# INLINABLE deleteNE #-} #else {-# INLINE delete #-} +{-# INLINE deleteNE #-} #endif -- | \(O(\log n)\). Update a value at a specific key with the result of the provided function. @@ -1015,10 +1132,16 @@ delete = go adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a adjust f = adjustWithKey (\_ x -> f x) + +adjustNE :: Ord k => (a -> a) -> k -> NonEmptyMap k a -> NonEmptyMap k a +adjustNE f = adjustWithKeyNE (\_ x -> f x) + #if __GLASGOW_HASKELL__ {-# INLINABLE adjust #-} +{-# INLINABLE adjustNE #-} #else {-# INLINE adjust #-} +{-# INLINE adjustNE #-} #endif -- | \(O(\log n)\). Adjust a value at a specific key. When the key is not @@ -1030,19 +1153,22 @@ adjust f = adjustWithKey (\_ x -> f x) -- > adjustWithKey f 7 empty == empty adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a -adjustWithKey = go - where - go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a - go _ !_ Tip = Tip - go f k (Bin sx kx x l r) = - case compare k kx of - LT -> Bin sx kx x (go f k l) r - GT -> Bin sx kx x l (go f k r) - EQ -> Bin sx kx (f kx x) l r +adjustWithKey _ !_ Tip = Tip +adjustWithKey f k (NE t) = NE $ adjustWithKeyNE f k t + +adjustWithKeyNE :: Ord k => (k -> a -> a) -> k -> NonEmptyMap k a -> NonEmptyMap k a +adjustWithKeyNE f k (Bin' sx kx x l r) = + case compare k kx of + LT -> Bin' sx kx x (adjustWithKey f k l) r + GT -> Bin' sx kx x l (adjustWithKey f k r) + EQ -> Bin' sx kx (f kx x) l r + #if __GLASGOW_HASKELL__ {-# INLINABLE adjustWithKey #-} +{-# INLINABLE adjustWithKeyNE #-} #else {-# INLINE adjustWithKey #-} +{-# INLINE adjustWithKeyNE #-} #endif -- | \(O(\log n)\). The expression (@'update' f k map@) updates the value @x@ @@ -1056,10 +1182,16 @@ adjustWithKey = go update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a update f = updateWithKey (\_ x -> f x) + +updateNE :: Ord k => (a -> Maybe a) -> k -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +updateNE f = updateWithKeyNE (\_ x -> f x) + #if __GLASGOW_HASKELL__ {-# INLINABLE update #-} +{-# INLINABLE updateNE #-} #else {-# INLINE update #-} +{-# INLINE updateNE #-} #endif -- | \(O(\log n)\). The expression (@'updateWithKey' f k map@) updates the @@ -1074,21 +1206,27 @@ update f = updateWithKey (\_ x -> f x) -- See Note: Type of local 'go' function updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a -updateWithKey = go - where - go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a - go _ !_ Tip = Tip - go f k(Bin sx kx x l r) = - case compare k kx of - LT -> balanceR kx x (go f k l) r - GT -> balanceL kx x l (go f k r) - EQ -> case f kx x of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r +updateWithKey _ !_ Tip = Tip +updateWithKey f k (NE t) = updateWithKeyFromNE f k t + +updateWithKeyNE :: Ord k => (k -> a -> Maybe a) -> k -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +updateWithKeyNE f k t = nonEmpty $ updateWithKeyFromNE f k t + +updateWithKeyFromNE :: Ord k => (k -> a -> Maybe a) -> k -> NonEmptyMap k a -> Map k a +updateWithKeyFromNE f k (Bin' sx kx x l r) = + case compare k kx of + LT -> balanceR kx x (updateWithKey f k l) r + GT -> balanceL kx x l (updateWithKey f k r) + EQ -> case f kx x of + Just x' -> NE $ Bin' sx kx x' l r + Nothing -> glue l r + #if __GLASGOW_HASKELL__ {-# INLINABLE updateWithKey #-} +{-# INLINABLE updateWithKeyNE #-} #else {-# INLINE updateWithKey #-} +{-# INLINE updateWithKeyNE #-} #endif -- | \(O(\log n)\). Lookup and update. See also 'updateWithKey'. @@ -1100,28 +1238,36 @@ updateWithKey = go -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") --- See Note: Type of local 'go' function -updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) -updateLookupWithKey f0 k0 = toPair . go f0 k0 - where - go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) - go _ !_ Tip = (Nothing :*: Tip) - go f k (Bin sx kx x l r) = - case compare k kx of - LT -> let !(found :*: l') = go f k l - !t' = balanceR kx x l' r - in (found :*: t') - GT -> let !(found :*: r') = go f k r - !t' = balanceL kx x l r' - in (found :*: t') - EQ -> case f kx x of - Just x' -> (Just x' :*: Bin sx kx x' l r) - Nothing -> let !glued = glue l r - in (Just x :*: glued) +updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a, Map k a) +updateLookupWithKey f k t = toPair $ updateLookupWithKey' f k t + +updateLookupWithKeyNE :: Ord k => (k -> a -> Maybe a) -> k -> NonEmptyMap k a -> (Maybe a, Maybe (NonEmptyMap k a)) +updateLookupWithKeyNE f k t = fmap nonEmpty $ toPair $ updateLookupWithKeyFromNE f k t + +updateLookupWithKey' :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) +updateLookupWithKey' _ !_ Tip = (Nothing :*: Tip) +updateLookupWithKey' f k (NE t) = updateLookupWithKeyFromNE f k t + +updateLookupWithKeyFromNE :: Ord k => (k -> a -> Maybe a) -> k -> NonEmptyMap k a -> StrictPair (Maybe a) (Map k a) +updateLookupWithKeyFromNE f k (Bin' sx kx x l r) = + case compare k kx of + LT -> let !(found :*: l') = updateLookupWithKey' f k l + !t' = balanceR kx x l' r + in (found :*: t') + GT -> let !(found :*: r') = updateLookupWithKey' f k r + !t' = balanceL kx x l r' + in (found :*: t') + EQ -> case f kx x of + Just x' -> (Just x' :*: NE (Bin' sx kx x' l r)) + Nothing -> let !glued = glue l r + in (Just x :*: glued) + #if __GLASGOW_HASKELL__ {-# INLINABLE updateLookupWithKey #-} +{-# INLINABLE updateLookupWithKeyNE #-} #else {-# INLINE updateLookupWithKey #-} +{-# INLINE updateLookupWithKeyNE #-} #endif -- | \(O(\log n)\). The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. @@ -1140,19 +1286,22 @@ updateLookupWithKey f0 k0 = toPair . go f0 k0 -- See Note: Type of local 'go' function alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a -alter = go - where - go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a - go f !k Tip = case f Nothing of - Nothing -> Tip - Just x -> singleton k x - - go f k (Bin sx kx x l r) = case compare k kx of - LT -> balance kx x (go f k l) r - GT -> balance kx x l (go f k r) - EQ -> case f (Just x) of - Just x' -> Bin sx kx x' l r - Nothing -> glue l r +alter f !k Tip = case f Nothing of + Nothing -> Tip + Just x -> singleton k x +alter f k (NE t) = alterFromNE f k t + +alterNE :: Ord k => (Maybe a -> Maybe a) -> k -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +alterNE f k t = nonEmpty $ alterFromNE f k t + +alterFromNE :: Ord k => (Maybe a -> Maybe a) -> k -> NonEmptyMap k a -> Map k a +alterFromNE f k (Bin' sx kx x l r) = case compare k kx of + LT -> balance kx x (alter f k l) r + GT -> balance kx x l (alter f k r) + EQ -> case f (Just x) of + Just x' -> NE $ Bin' sx kx x' l r + Nothing -> glue l r + #if __GLASGOW_HASKELL__ {-# INLINABLE alter #-} #else @@ -1207,15 +1356,24 @@ alterF :: (Functor f, Ord k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a) alterF f k m = atKeyImpl Lazy k f m +alterFNE + :: (Functor f, Ord k) + => (Maybe a -> f (Maybe a)) -> k -> NonEmptyMap k a -> f (Maybe (NonEmptyMap k a)) +alterFNE f k m = atKeyImplNE Lazy k f m + #ifndef __GLASGOW_HASKELL__ {-# INLINE alterF #-} +{-# INLINE alterFNE #-} #else {-# INLINABLE [2] alterF #-} +{-# INLINABLE [2] alterFNE #-} -- We can save a little time by recognizing the special case of -- `Control.Applicative.Const` and just doing a lookup. {-# RULES -"alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m +"alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m #-} +{-# RULES +"alterFNE/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterFNE f k = \m -> Const . getConst . f $ lookupNE k m #-} -- base 4.8 and above include Data.Functor.Identity, so we can @@ -1223,6 +1381,9 @@ alterF f k m = atKeyImpl Lazy k f m {-# RULES "alterF/Identity" forall k f . alterF f k = atKeyIdentity k f #-} +{-# RULES +"alterFNE/Identity" forall k f . alterFNE f k = atKeyIdentityNE k f + #-} #endif atKeyImpl :: (Functor f, Ord k) => @@ -1250,6 +1411,32 @@ atKeyImpl strict !k f m = case lookupTrace k m of {-# INLINE atKeyImpl #-} +atKeyImplNE + :: (Functor f, Ord k) + => AreWeStrict -> k -> (Maybe a -> f (Maybe a)) -> NonEmptyMap k a -> f (Maybe (NonEmptyMap k a)) +#ifdef DEFINE_ALTERF_FALLBACK +atKeyImplNE strict !k f m +-- It doesn't seem sensible to worry about overflowing the queue +-- if the word size is 61 or more. If I calculate it correctly, +-- that would take a map with nearly a quadrillion entries. + | wordSize < 61 && sizeNE m >= alterFCutoff = alterFFallbackNE strict k f m +#endif +atKeyImplNE strict !k f m = case lookupTraceNE k m of + TraceResult mv q -> (<$> f mv) $ \ fres -> + case fres of + Nothing -> case mv of + Nothing -> Just $ m + Just old -> deleteAlongNE old q m + Just new -> case strict of + Strict -> new `seq` case mv of + Nothing -> Just $ insertAlongNE q k new m + Just _ -> Just $ replaceAlongNE q new m + Lazy -> case mv of + Nothing -> Just $ insertAlongNE q k new m + Just _ -> Just $ replaceAlongNE q new m + +{-# INLINE atKeyImplNE #-} + #ifdef DEFINE_ALTERF_FALLBACK alterFCutoff :: Int #if WORD_SIZE_IN_BITS == 32 @@ -1268,30 +1455,44 @@ data TraceResult a = TraceResult (Maybe a) {-# UNPACK #-} !BitQueue -- Look up a key and return a result indicating whether it was found -- and what path was taken. lookupTrace :: Ord k => k -> Map k a -> TraceResult a -lookupTrace = go emptyQB - where - go :: Ord k => BitQueueB -> k -> Map k a -> TraceResult a - go !q !_ Tip = TraceResult Nothing (buildQ q) - go q k (Bin _ kx x l r) = case compare k kx of - LT -> (go $! q `snocQB` False) k l - GT -> (go $! q `snocQB` True) k r - EQ -> TraceResult (Just x) (buildQ q) +lookupTrace = lookupTrace' emptyQB + +lookupTrace' :: Ord k => BitQueueB -> k -> Map k a -> TraceResult a +lookupTrace' !q !_ Tip = TraceResult Nothing (buildQ q) +lookupTrace' q k (NE t) = lookupTraceNE' q k t + +lookupTraceNE :: Ord k => k -> NonEmptyMap k a -> TraceResult a +lookupTraceNE = lookupTraceNE' emptyQB + +lookupTraceNE' :: Ord k => BitQueueB -> k -> NonEmptyMap k a -> TraceResult a +lookupTraceNE' q k (Bin' _ kx x l r) = case compare k kx of + LT -> (lookupTrace' $! q `snocQB` False) k l + GT -> (lookupTrace' $! q `snocQB` True) k r + EQ -> TraceResult (Just x) (buildQ q) #ifdef __GLASGOW_HASKELL__ {-# INLINABLE lookupTrace #-} +{-# INLINABLE lookupTraceNE #-} #else {-# INLINE lookupTrace #-} +{-# INLINE lookupTraceNE #-} #endif -- Insert at a location (which will always be a leaf) -- described by the path passed in. insertAlong :: BitQueue -> k -> a -> Map k a -> Map k a -insertAlong !_ kx x Tip = singleton kx x -insertAlong q kx x (Bin sz ky y l r) = +insertAlong q kx x t = NE $ insertAlong' q kx x t + +insertAlong' :: BitQueue -> k -> a -> Map k a -> NonEmptyMap k a +insertAlong' !_ kx x Tip = singletonNE kx x +insertAlong' q kx x (NE t) = insertAlongNE q kx x t + +insertAlongNE :: BitQueue -> k -> a -> NonEmptyMap k a -> NonEmptyMap k a +insertAlongNE q kx x (Bin' sz ky y l r) = case unconsQ q of - Just (False, tl) -> balanceL ky y (insertAlong tl kx x l) r - Just (True,tl) -> balanceR ky y l (insertAlong tl kx x r) - Nothing -> Bin sz kx x l r -- Shouldn't happen + Just (False, tl) -> balanceLNE ky y (insertAlong' tl kx x l) r + Just (True,tl) -> balanceRNE ky y l (insertAlong' tl kx x r) + Nothing -> Bin' sz kx x l r -- Shouldn't happen -- Delete from a location (which will always be a node) -- described by the path passed in. @@ -1313,18 +1514,29 @@ insertAlong q kx x (Bin sz ky y l r) = -- so instead we convert the value to a magical zero-width -- proxy that's ultimately erased. deleteAlong :: any -> BitQueue -> Map k a -> Map k a -deleteAlong old !q0 !m = go (bogus old) q0 m where +deleteAlong old !q0 !m = deleteAlong' (bogus old) q0 m + +deleteAlongNE :: any -> BitQueue -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +deleteAlongNE old !q0 !m = nonEmpty $ deleteAlongFromNE' (bogus old) q0 m + #ifdef USE_MAGIC_PROXY - go :: Proxy# () -> BitQueue -> Map k a -> Map k a +deleteAlong' :: Proxy# () -> BitQueue -> Map k a -> Map k a #else - go :: any -> BitQueue -> Map k a -> Map k a +deleteAlong' :: any -> BitQueue -> Map k a -> Map k a #endif - go !_ !_ Tip = Tip - go foom q (Bin _ ky y l r) = - case unconsQ q of - Just (False, tl) -> balanceR ky y (go foom tl l) r - Just (True, tl) -> balanceL ky y l (go foom tl r) - Nothing -> glue l r +deleteAlong' !_ !_ Tip = Tip +deleteAlong' foom q (NE t) = deleteAlongFromNE' foom q t + +#ifdef USE_MAGIC_PROXY +deleteAlongFromNE' :: Proxy# () -> BitQueue -> NonEmptyMap k a -> Map k a +#else +deleteAlongFromNE' :: any -> BitQueue -> NonEmptyMap k a -> Map k a +#endif +deleteAlongFromNE' foom q (Bin' _ ky y l r) = + case unconsQ q of + Just (False, tl) -> balanceR ky y (deleteAlong' foom tl l) r + Just (True, tl) -> balanceL ky y l (deleteAlong' foom tl r) + Nothing -> glue l r #ifdef USE_MAGIC_PROXY {-# NOINLINE bogus #-} @@ -1341,50 +1553,69 @@ bogus a = a -- by the given path with a new one. replaceAlong :: BitQueue -> a -> Map k a -> Map k a replaceAlong !_ _ Tip = Tip -- Should not happen -replaceAlong q x (Bin sz ky y l r) = +replaceAlong q x (NE t) = NE $ replaceAlongNE q x t + +replaceAlongNE :: BitQueue -> a -> NonEmptyMap k a -> NonEmptyMap k a +replaceAlongNE q x (Bin' sz ky y l r) = case unconsQ q of - Just (False, tl) -> Bin sz ky y (replaceAlong tl x l) r - Just (True,tl) -> Bin sz ky y l (replaceAlong tl x r) - Nothing -> Bin sz ky x l r + Just (False, tl) -> Bin' sz ky y (replaceAlong tl x l) r + Just (True,tl) -> Bin' sz ky y l (replaceAlong tl x r) + Nothing -> Bin' sz ky x l r #ifdef __GLASGOW_HASKELL__ atKeyIdentity :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> Map k a -> Identity (Map k a) atKeyIdentity k f t = Identity $ atKeyPlain Lazy k (coerce f) t + +atKeyIdentityNE :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> NonEmptyMap k a -> Identity (Maybe (NonEmptyMap k a)) +atKeyIdentityNE k f t = Identity $ atKeyPlainNE Lazy k (coerce f) t {-# INLINABLE atKeyIdentity #-} atKeyPlain :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> Map k a -> Map k a -atKeyPlain strict k0 f0 t = case go k0 f0 t of +atKeyPlain strict k0 f0 t = case atKeyPlain' strict k0 f0 t of AltSmaller t' -> t' - AltBigger t' -> t' + AltBigger t' -> NE t' AltAdj t' -> t' AltSame -> t - where - go :: Ord k => k -> (Maybe a -> Maybe a) -> Map k a -> Altered k a - go !k f Tip = case f Nothing of - Nothing -> AltSame - Just x -> case strict of - Lazy -> AltBigger $ singleton k x - Strict -> x `seq` (AltBigger $ singleton k x) - - go k f (Bin sx kx x l r) = case compare k kx of - LT -> case go k f l of - AltSmaller l' -> AltSmaller $ balanceR kx x l' r - AltBigger l' -> AltBigger $ balanceL kx x l' r - AltAdj l' -> AltAdj $ Bin sx kx x l' r - AltSame -> AltSame - GT -> case go k f r of - AltSmaller r' -> AltSmaller $ balanceL kx x l r' - AltBigger r' -> AltBigger $ balanceR kx x l r' - AltAdj r' -> AltAdj $ Bin sx kx x l r' - AltSame -> AltSame - EQ -> case f (Just x) of - Just x' -> case strict of - Lazy -> AltAdj $ Bin sx kx x' l r - Strict -> x' `seq` (AltAdj $ Bin sx kx x' l r) - Nothing -> AltSmaller $ glue l r + +atKeyPlainNE :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +atKeyPlainNE strict k0 f0 t = case atKeyPlainNE' strict k0 f0 t of + AltSmaller t' -> nonEmpty t' + AltBigger t' -> Just t' + AltAdj t' -> Just t' + AltSame -> Just t + +atKeyPlain' :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> Map k a -> Altered Map k a +atKeyPlain' strict !k f Tip = case f Nothing of + Nothing -> AltSame + Just x -> case strict of + Lazy -> AltBigger $ singletonNE k x + Strict -> x `seq` (AltBigger $ singletonNE k x) +atKeyPlain' strict k f (NE t) = case atKeyPlainNE' strict k f t of + AltSmaller t' -> AltSmaller t' + AltBigger t' -> AltBigger t' + AltAdj t' -> AltAdj $ NE t' + AltSame -> AltSame + +atKeyPlainNE' :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> NonEmptyMap k a -> Altered NonEmptyMap k a +atKeyPlainNE' strict k f (Bin' sx kx x l r) = case compare k kx of + LT -> case atKeyPlain' strict k f l of + AltSmaller l' -> AltSmaller $ balanceR kx x l' r + AltBigger l' -> AltBigger $ balanceLNE kx x l' r + AltAdj l' -> AltAdj $ Bin' sx kx x l' r + AltSame -> AltSame + GT -> case atKeyPlain' strict k f r of + AltSmaller r' -> AltSmaller $ balanceL kx x l r' + AltBigger r' -> AltBigger $ balanceRNE kx x l r' + AltAdj r' -> AltAdj $ Bin' sx kx x l r' + AltSame -> AltSame + EQ -> case f (Just x) of + Just x' -> case strict of + Lazy -> AltAdj $ Bin' sx kx x' l r + Strict -> x' `seq` (AltAdj $ Bin' sx kx x' l r) + Nothing -> AltSmaller $ glue l r {-# INLINE atKeyPlain #-} -data Altered k a = AltSmaller !(Map k a) | AltBigger !(Map k a) | AltAdj !(Map k a) | AltSame +data Altered m k a = AltSmaller !(Map k a) | AltBigger !(NonEmptyMap k a) | AltAdj !(m k a) | AltSame #endif #ifdef DEFINE_ALTERF_FALLBACK @@ -1401,21 +1632,40 @@ alterFFallback Strict k f t = alterFYoneda k (\m q -> q . forceMaybe <$> f m) t forceMaybe may@(Just !_) = may {-# NOINLINE alterFFallback #-} -alterFYoneda :: Ord k => - k -> (Maybe a -> (Maybe a -> b) -> f b) -> Map k a -> (Map k a -> b) -> f b -alterFYoneda = go +alterFFallbackNE :: (Functor f, Ord k) + => AreWeStrict -> k -> (Maybe a -> f (Maybe a)) -> NonEmptyMap k a -> f (Maybe (NonEmptyMap k a)) +alterFFallbackNE Lazy k f t = fmap nonEmpty $ alterFYonedaNE k (\m q -> q <$> f m) t id +alterFFallbackNE Strict k f t = fmap nonEmpty $ alterFYonedaNE k (\m q -> q . forceMaybe <$> f m) t id where - go :: Ord k => - k -> (Maybe a -> (Maybe a -> b) -> f b) -> Map k a -> (Map k a -> b) -> f b - go !k f Tip g = f Nothing $ \ mx -> case mx of - Nothing -> g Tip - Just x -> g (singleton k x) - go k f (Bin sx kx x l r) g = case compare k kx of - LT -> go k f l (\m -> g (balance kx x m r)) - GT -> go k f r (\m -> g (balance kx x l m)) - EQ -> f (Just x) $ \ mx' -> case mx' of - Just x' -> g (Bin sx kx x' l r) - Nothing -> g (glue l r) + forceMaybe Nothing = Nothing + forceMaybe may@(Just !_) = may +{-# NOINLINE alterFFallbackNE #-} + +alterFYoneda + :: Ord k + => k + -> (Maybe a -> (Maybe a -> b) -> f b) + -> Map k a + -> (Map k a -> b) + -> f b +alterFYoneda !k f Tip g = f Nothing $ \ mx -> case mx of + Nothing -> g Tip + Just x -> g (singleton k x) +alterFYoneda k f (NE t) g = alterFYonedaNE k f t g + +alterFYonedaNE + :: Ord k + => k + -> (Maybe a -> (Maybe a -> b) -> f b) + -> NonEmptyMap k a + -> (Map k a -> b) + -> f b +alterFYonedaNE k f (Bin' sx kx x l r) g = case compare k kx of + LT -> alterFYoneda k f l (\m -> g (balance kx x m r)) + GT -> alterFYoneda k f r (\m -> g (balance kx x l m)) + EQ -> f (Just x) $ \ mx' -> case mx' of + Just x' -> g (NE (Bin' sx kx x' l r)) + Nothing -> g (glue l r) {-# INLINE alterFYoneda #-} #endif @@ -1438,7 +1688,7 @@ findIndex = go 0 where go :: Ord k => Int -> k -> Map k a -> Int go !_ !_ Tip = error "Map.findIndex: element is not in the map" - go idx k (Bin _ kx _ l r) = case compare k kx of + go idx k (NE (Bin' _ kx _ l r)) = case compare k kx of LT -> go idx k l GT -> go (idx + size l + 1) k r EQ -> idx + size l @@ -1461,7 +1711,7 @@ lookupIndex = go 0 where go :: Ord k => Int -> k -> Map k a -> Maybe Int go !_ !_ Tip = Nothing - go idx k (Bin _ kx _ l r) = case compare k kx of + go idx k (NE (Bin' _ kx _ l r)) = case compare k kx of LT -> go idx k l GT -> go (idx + size l + 1) k r EQ -> Just $! idx + size l @@ -1479,7 +1729,7 @@ lookupIndex = go 0 elemAt :: Int -> Map k a -> (k,a) elemAt !_ Tip = error "Map.elemAt: index out of range" -elemAt i (Bin _ kx x l r) +elemAt i (NE (Bin' _ kx x l r)) = case compare i sizeL of LT -> elemAt i l GT -> elemAt (i-sizeL-1) r @@ -1502,7 +1752,7 @@ take i0 m0 = go i0 m0 where go i !_ | i <= 0 = Tip go !_ Tip = Tip - go i (Bin _ kx x l r) = + go i (NE (Bin' _ kx x l r)) = case compare i sizeL of LT -> go i l GT -> link kx x l (go (i - sizeL - 1) r) @@ -1523,7 +1773,7 @@ drop i0 m0 = go i0 m0 where go i m | i <= 0 = m go !_ Tip = Tip - go i (Bin _ kx x l r) = + go i (NE (Bin' _ kx x l r)) = case compare i sizeL of LT -> link kx x (go i l) r GT -> go (i - sizeL - 1) r @@ -1544,7 +1794,7 @@ splitAt i0 m0 where go i m | i <= 0 = Tip :*: m go !_ Tip = Tip :*: Tip - go i (Bin _ kx x l r) + go i (NE (Bin' _ kx x l r)) = case compare i sizeL of LT -> case go i l of ll :*: lr -> ll :*: link kx x lr r @@ -1570,11 +1820,11 @@ updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a updateAt f !i t = case t of Tip -> error "Map.updateAt: index out of range" - Bin sx kx x l r -> case compare i sizeL of + NE (Bin' sx kx x l r) -> case compare i sizeL of LT -> balanceR kx x (updateAt f i l) r GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) EQ -> case f kx x of - Just x' -> Bin sx kx x' l r + Just x' -> NE $ Bin' sx kx x' l r Nothing -> glue l r where sizeL = size l @@ -1592,7 +1842,7 @@ deleteAt :: Int -> Map k a -> Map k a deleteAt !i t = case t of Tip -> error "Map.deleteAt: index out of range" - Bin _ kx x l r -> case compare i sizeL of + NE (Bin' _ kx x l r) -> case compare i sizeL of LT -> balanceR kx x (deleteAt i l) r GT -> balanceL kx x l (deleteAt (i-sizeL-1) r) EQ -> glue l r @@ -1606,7 +1856,7 @@ deleteAt !i t = lookupMinSure :: k -> a -> Map k a -> (k, a) lookupMinSure k a Tip = (k, a) -lookupMinSure _ _ (Bin _ k a l _) = lookupMinSure k a l +lookupMinSure _ _ (NE (Bin' _ k a l _)) = lookupMinSure k a l -- | \(O(\log n)\). The minimal key of the map. Returns 'Nothing' if the map is empty. -- @@ -1617,7 +1867,7 @@ lookupMinSure _ _ (Bin _ k a l _) = lookupMinSure k a l lookupMin :: Map k a -> Maybe (k,a) lookupMin Tip = Nothing -lookupMin (Bin _ k x l _) = Just $! lookupMinSure k x l +lookupMin (NE (Bin' _ k x l _)) = Just $! lookupMinSure k x l -- | \(O(\log n)\). The minimal key of the map. Calls 'error' if the map is empty. -- @@ -1636,7 +1886,7 @@ findMin t lookupMaxSure :: k -> a -> Map k a -> (k, a) lookupMaxSure k a Tip = (k, a) -lookupMaxSure _ _ (Bin _ k a _ r) = lookupMaxSure k a r +lookupMaxSure _ _ (NE (Bin' _ k a _ r)) = lookupMaxSure k a r -- | \(O(\log n)\). The maximal key of the map. Returns 'Nothing' if the map is empty. -- @@ -1647,7 +1897,7 @@ lookupMaxSure _ _ (Bin _ k a _ r) = lookupMaxSure k a r lookupMax :: Map k a -> Maybe (k, a) lookupMax Tip = Nothing -lookupMax (Bin _ k x _ r) = Just $! lookupMaxSure k x r +lookupMax (NE (Bin' _ k x _ r)) = Just $! lookupMaxSure k x r findMax :: Map k a -> (k,a) findMax t @@ -1660,8 +1910,8 @@ findMax t -- > deleteMin empty == empty deleteMin :: Map k a -> Map k a -deleteMin (Bin _ _ _ Tip r) = r -deleteMin (Bin _ kx x l r) = balanceR kx x (deleteMin l) r +deleteMin (NE (Bin' _ _ _ Tip r)) = r +deleteMin (NE (Bin' _ kx x l r)) = balanceR kx x (deleteMin l) r deleteMin Tip = Tip -- | \(O(\log n)\). Delete the maximal key. Returns an empty map if the map is empty. @@ -1670,8 +1920,8 @@ deleteMin Tip = Tip -- > deleteMax empty == empty deleteMax :: Map k a -> Map k a -deleteMax (Bin _ _ _ l Tip) = l -deleteMax (Bin _ kx x l r) = balanceL kx x l (deleteMax r) +deleteMax (NE (Bin' _ _ _ l Tip)) = l +deleteMax (NE (Bin' _ kx x l r)) = balanceL kx x l (deleteMax r) deleteMax Tip = Tip -- | \(O(\log n)\). Update the value at the minimal key. @@ -1700,10 +1950,10 @@ updateMax f m updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a updateMinWithKey _ Tip = Tip -updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of +updateMinWithKey f (NE (Bin' sx kx x Tip r)) = case f kx x of Nothing -> r - Just x' -> Bin sx kx x' Tip r -updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r + Just x' -> NE $ Bin' sx kx x' Tip r +updateMinWithKey f (NE (Bin' _ kx x l r)) = balanceR kx x (updateMinWithKey f l) r -- | \(O(\log n)\). Update the value at the maximal key. -- @@ -1712,10 +1962,10 @@ updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a updateMaxWithKey _ Tip = Tip -updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of +updateMaxWithKey f (NE (Bin' sx kx x l Tip)) = case f kx x of Nothing -> l - Just x' -> Bin sx kx x' l Tip -updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) + Just x' -> NE $ Bin' sx kx x' l Tip +updateMaxWithKey f (NE (Bin' _ kx x l r)) = balanceL kx x l (updateMaxWithKey f r) -- | \(O(\log n)\). Retrieves the minimal (key,value) pair of the map, and -- the map stripped of that element, or 'Nothing' if passed an empty map. @@ -1725,7 +1975,7 @@ updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) minViewWithKey :: Map k a -> Maybe ((k,a), Map k a) minViewWithKey Tip = Nothing -minViewWithKey (Bin _ k x l r) = Just $ +minViewWithKey (NE (Bin' _ k x l r)) = Just $ case minViewSure k x l r of MinView km xm t -> ((km, xm), t) -- We inline this to give GHC the best possible chance of getting @@ -1741,7 +1991,7 @@ minViewWithKey (Bin _ k x l r) = Just $ maxViewWithKey :: Map k a -> Maybe ((k,a), Map k a) maxViewWithKey Tip = Nothing -maxViewWithKey (Bin _ k x l r) = Just $ +maxViewWithKey (NE (Bin' _ k x l r)) = Just $ case maxViewSure k x l r of MaxView km xm t -> ((km, xm), t) -- See note on inlining at minViewWithKey @@ -1789,6 +2039,9 @@ unions ts {-# INLINABLE unions #-} #endif +-- unionsNE :: (Foldable f, Ord k) => f (NonEmptyMap k a) -> NonEmptyMap k a +-- unionsNE ts = foldl1' unionNE ts + -- | The union of a list of maps, with a combining operation: -- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). -- @@ -1811,10 +2064,10 @@ unionsWith f ts union :: Ord k => Map k a -> Map k a -> Map k a union t1 Tip = t1 -union t1 (Bin _ k x Tip Tip) = insertR k x t1 -union (Bin _ k x Tip Tip) t2 = insert k x t2 +union t1 (NE (Bin' _ k x Tip Tip)) = insertR k x t1 +union (NE (Bin' _ k x Tip Tip)) t2 = insert k x t2 union Tip t2 = t2 -union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of +union t1@(NE (Bin' _ k1 x1 l1 r1)) t2 = case split k1 t2 of (l2, r2) | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1 | otherwise -> link k1 x1 l1l2 r1r2 where !l1l2 = union l1 l2 @@ -1823,6 +2076,19 @@ union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of {-# INLINABLE union #-} #endif +unionNE :: Ord k => NonEmptyMap k a -> NonEmptyMap k a -> NonEmptyMap k a +unionNE t1 (Bin' _ k x Tip Tip) = insertRNE k x t1 +unionNE (Bin' _ k x Tip Tip) t2 = insertNE k x t2 +unionNE t1@(Bin' _ k x l1 r1) t2 = case split k (NE t2) of + (l2, r2) + | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1 + | otherwise -> linkNE k x l1l2 r1r2 + where !l1l2 = union l1 l2 + !r1r2 = union r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionNE #-} +#endif + {-------------------------------------------------------------------- Union with a combining function --------------------------------------------------------------------} @@ -1833,10 +2099,10 @@ union t1@(Bin _ k1 x1 l1 r1) t2 = case split k1 t2 of unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a -- QuickCheck says pointer equality never happens here. unionWith _f t1 Tip = t1 -unionWith f t1 (Bin _ k x Tip Tip) = insertWithR f k x t1 -unionWith f (Bin _ k x Tip Tip) t2 = insertWith f k x t2 +unionWith f t1 (NE (Bin' _ k x Tip Tip)) = insertWithR f k x t1 +unionWith f (NE (Bin' _ k x Tip Tip)) t2 = insertWith f k x t2 unionWith _f Tip t2 = t2 -unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of +unionWith f (NE (Bin' _ k1 x1 l1 r1)) t2 = case splitLookup k1 t2 of (l2, mb, r2) -> case mb of Nothing -> link k1 x1 l1l2 r1r2 Just x2 -> link k1 (f x1 x2) l1l2 r1r2 @@ -1846,6 +2112,20 @@ unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of {-# INLINABLE unionWith #-} #endif +unionWithNE :: Ord k => (a -> a -> a) -> NonEmptyMap k a -> NonEmptyMap k a -> NonEmptyMap k a +-- QuickCheck says pointer equality never happens here. +unionWithNE f t1 (Bin' _ k x Tip Tip) = insertWithRNE f k x t1 +unionWithNE f (Bin' _ k x Tip Tip) t2 = insertWithNE f k x t2 +unionWithNE f (Bin' _ k1 x1 l1 r1) t2 = case splitLookupNE k1 t2 of + (l2, mb, r2) -> case mb of + Nothing -> linkNE k1 x1 l1l2 r1r2 + Just x2 -> linkNE k1 (f x1 x2) l1l2 r1r2 + where !l1l2 = unionWith f l1 l2 + !r1r2 = unionWith f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionWithNE #-} +#endif + -- | \(O\bigl(m \log\bigl(\frac{n+1}{m+1}\bigr)\bigr), \; m \leq n\). -- Union with a combining function. -- @@ -1854,10 +2134,10 @@ unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a unionWithKey _f t1 Tip = t1 -unionWithKey f t1 (Bin _ k x Tip Tip) = insertWithKeyR f k x t1 -unionWithKey f (Bin _ k x Tip Tip) t2 = insertWithKey f k x t2 +unionWithKey f t1 (NE (Bin' _ k x Tip Tip)) = insertWithKeyR f k x t1 +unionWithKey f (NE (Bin' _ k x Tip Tip)) t2 = insertWithKey f k x t2 unionWithKey _f Tip t2 = t2 -unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of +unionWithKey f (NE (Bin' _ k1 x1 l1 r1)) t2 = case splitLookup k1 t2 of (l2, mb, r2) -> case mb of Nothing -> link k1 x1 l1l2 r1r2 Just x2 -> link k1 (f k1 x1 x2) l1l2 r1r2 @@ -1867,6 +2147,19 @@ unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of {-# INLINABLE unionWithKey #-} #endif +unionWithKeyNE :: Ord k => (k -> a -> a -> a) -> NonEmptyMap k a -> NonEmptyMap k a -> NonEmptyMap k a +unionWithKeyNE f t1 (Bin' _ k x Tip Tip) = insertWithKeyRNE f k x t1 +unionWithKeyNE f (Bin' _ k x Tip Tip) t2 = insertWithKeyNE f k x t2 +unionWithKeyNE f (Bin' _ k1 x1 l1 r1) t2 = case splitLookupNE k1 t2 of + (l2, mb, r2) -> case mb of + Nothing -> linkNE k1 x1 l1l2 r1r2 + Just x2 -> linkNE k1 (f k1 x1 x2) l1l2 r1r2 + where !l1l2 = unionWithKey f l1 l2 + !r1r2 = unionWithKey f r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionWithKeyNE #-} +#endif + {-------------------------------------------------------------------- Difference --------------------------------------------------------------------} @@ -1885,15 +2178,23 @@ unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of difference :: Ord k => Map k a -> Map k b -> Map k a difference Tip _ = Tip difference t1 Tip = t1 -difference t1 (Bin _ k _ l2 r2) = case split k t1 of +difference t1 (NE t2) = differenceNE' t1 t2 + +differenceNE :: Ord k => NonEmptyMap k a -> NonEmptyMap k a -> Maybe (NonEmptyMap k a) +differenceNE t1 t2 = nonEmpty $ differenceNE' (NE t1) t2 + +differenceNE' :: Ord k => Map k a -> NonEmptyMap k b -> Map k a +differenceNE' t1 (Bin' _ k _ l2 r2) = case split k t1 of (l1, r1) | size l1l2 + size r1r2 == size t1 -> t1 | otherwise -> link2 l1l2 r1r2 where !l1l2 = difference l1 l2 !r1r2 = difference r1 r2 + #if __GLASGOW_HASKELL__ {-# INLINABLE difference #-} +{-# INLINABLE differenceNE #-} #endif -- | \(O\bigl(m \log\bigl(\frac{n+1}{m+1}\bigr)\bigr), \; m \leq n\). Remove all keys in a 'Set' from a 'Map'. @@ -1908,7 +2209,7 @@ difference t1 (Bin _ k _ l2 r2) = case split k t1 of withoutKeys :: Ord k => Map k a -> Set k -> Map k a withoutKeys Tip _ = Tip withoutKeys m Set.Tip = m -withoutKeys m (Set.Bin _ k ls rs) = case splitMember k m of +withoutKeys m (Set.NE (Set.Bin' _ k ls rs)) = case splitMember k m of (lm, b, rm) | not b && lm' `ptrEq` lm && rm' `ptrEq` rm -> m | otherwise -> link2 lm' rm' @@ -1931,8 +2232,14 @@ withoutKeys m (Set.Bin _ k ls rs) = case splitMember k m of differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a differenceWith f = merge preserveMissing dropMissing $ zipWithMaybeMatched (\_ x y -> f x y) + +differenceWithNE :: Ord k => (a -> b -> Maybe a) -> NonEmptyMap k a -> NonEmptyMap k b -> Maybe (NonEmptyMap k a) +differenceWithNE f = mergeNE preserveMissing dropMissing $ + zipWithMaybeMatched (\_ x y -> f x y) + #if __GLASGOW_HASKELL__ {-# INLINABLE differenceWith #-} +{-# INLINABLE differenceWithNE #-} #endif -- | \(O(n+m)\). Difference with a combining function. When two equal keys are @@ -1947,8 +2254,14 @@ differenceWith f = merge preserveMissing dropMissing $ differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a differenceWithKey f = merge preserveMissing dropMissing (zipWithMaybeMatched f) + +differenceWithKeyNE :: Ord k => (k -> a -> b -> Maybe a) -> NonEmptyMap k a -> NonEmptyMap k b -> Maybe (NonEmptyMap k a) +differenceWithKeyNE f = + mergeNE preserveMissing dropMissing (zipWithMaybeMatched f) + #if __GLASGOW_HASKELL__ {-# INLINABLE differenceWithKey #-} +{-# INLINABLE differenceWithKeyNE #-} #endif @@ -1964,7 +2277,7 @@ differenceWithKey f = intersection :: Ord k => Map k a -> Map k b -> Map k a intersection Tip _ = Tip intersection _ Tip = Tip -intersection t1@(Bin _ k x l1 r1) t2 +intersection t1@(NE (Bin' _ k x l1 r1)) t2 | mb = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 then t1 else link k x l1l2 r1r2 @@ -1973,8 +2286,13 @@ intersection t1@(Bin _ k x l1 r1) t2 !(l2, mb, r2) = splitMember k t2 !l1l2 = intersection l1 l2 !r1r2 = intersection r1 r2 + +intersectionNE :: Ord k => NonEmptyMap k a -> NonEmptyMap k b -> Maybe (NonEmptyMap k a) +intersectionNE t1 t2 = nonEmpty $ intersection (NE t1) (NE t2) + #if __GLASGOW_HASKELL__ {-# INLINABLE intersection #-} +{-# INLINABLE intersectionNE #-} #endif -- | \(O\bigl(m \log\bigl(\frac{n+1}{m+1}\bigr)\bigr), \; m \leq n\). Restrict a 'Map' to only those keys @@ -1989,7 +2307,13 @@ intersection t1@(Bin _ k x l1 r1) t2 restrictKeys :: Ord k => Map k a -> Set k -> Map k a restrictKeys Tip _ = Tip restrictKeys _ Set.Tip = Tip -restrictKeys m@(Bin _ k x l1 r1) s +restrictKeys m0@(NE m) s = restrictKeysNE' m0 m s + +restrictKeysNE :: Ord k => NonEmptyMap k a -> Set k -> Maybe (NonEmptyMap k a) +restrictKeysNE m s = nonEmpty $ restrictKeysNE' (NE m) m s + +restrictKeysNE' :: Ord k => Map k a -> NonEmptyMap k a -> Set k -> Map k a +restrictKeysNE' m (Bin' _ k x l1 r1) s | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 then m else link k x l1l2 r1r2 @@ -1998,8 +2322,11 @@ restrictKeys m@(Bin _ k x l1 r1) s !(l2, b, r2) = Set.splitMember k s !l1l2 = restrictKeys l1 l2 !r1r2 = restrictKeys r1 r2 + #if __GLASGOW_HASKELL__ {-# INLINABLE restrictKeys #-} +{-# INLINABLE restrictKeysNE #-} +{-# INLINABLE restrictKeysNE' #-} #endif -- | \(O\bigl(m \log\bigl(\frac{n+1}{m+1}\bigr)\bigr), \; m \leq n\). Intersection with a combining function. @@ -2011,15 +2338,24 @@ intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c -- element in the result will be a thunk. intersectionWith _f Tip _ = Tip intersectionWith _f _ Tip = Tip -intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of +intersectionWith f (NE t1) t2 = intersectionWithNE' f t1 t2 + +intersectionWithNE :: Ord k => (a -> b -> c) -> NonEmptyMap k a -> NonEmptyMap k b -> Maybe (NonEmptyMap k c) +intersectionWithNE f t1 t2 = nonEmpty $ intersectionWithNE' f t1 (NE t2) + +intersectionWithNE' :: Ord k => (a -> b -> c) -> NonEmptyMap k a -> Map k b -> Map k c +intersectionWithNE' f (Bin' _ k x1 l1 r1) t2 = case mb of Just x2 -> link k (f x1 x2) l1l2 r1r2 Nothing -> link2 l1l2 r1r2 where !(l2, mb, r2) = splitLookup k t2 !l1l2 = intersectionWith f l1 l2 !r1r2 = intersectionWith f r1 r2 + #if __GLASGOW_HASKELL__ {-# INLINABLE intersectionWith #-} +{-# INLINABLE intersectionWithNE #-} +{-# INLINABLE intersectionWithNE' #-} #endif -- | \(O\bigl(m \log\bigl(\frac{n+1}{m+1}\bigr)\bigr), \; m \leq n\). Intersection with a combining function. @@ -2030,15 +2366,24 @@ intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c intersectionWithKey _f Tip _ = Tip intersectionWithKey _f _ Tip = Tip -intersectionWithKey f (Bin _ k x1 l1 r1) t2 = case mb of +intersectionWithKey f (NE t1) t2 = intersectionWithKeyNE' f t1 t2 + +intersectionWithKeyNE :: Ord k => (k -> a -> b -> c) -> NonEmptyMap k a -> NonEmptyMap k b -> Maybe (NonEmptyMap k c) +intersectionWithKeyNE f t1 t2 = nonEmpty $ intersectionWithKeyNE' f t1 (NE t2) + +intersectionWithKeyNE' :: Ord k => (k -> a -> b -> c) -> NonEmptyMap k a -> Map k b -> Map k c +intersectionWithKeyNE' f (Bin' _ k x1 l1 r1) t2 = case mb of Just x2 -> link k (f k x1 x2) l1l2 r1r2 Nothing -> link2 l1l2 r1r2 where !(l2, mb, r2) = splitLookup k t2 !l1l2 = intersectionWithKey f l1 l2 !r1r2 = intersectionWithKey f r1 r2 + #if __GLASGOW_HASKELL__ {-# INLINABLE intersectionWithKey #-} +{-# INLINABLE intersectionWithKeyNE #-} +{-# INLINABLE intersectionWithKeyNE' #-} #endif {-------------------------------------------------------------------- @@ -2062,8 +2407,14 @@ intersectionWithKey f (Bin _ k x1 l1 r1) t2 = case mb of disjoint :: Ord k => Map k a -> Map k b -> Bool disjoint Tip _ = True disjoint _ Tip = True -disjoint (Bin 1 k _ _ _) t = k `notMember` t -disjoint (Bin _ k _ l r) t +disjoint (NE t1) t2 = disjointNE' t1 t2 + +disjointNE :: Ord k => NonEmptyMap k a -> NonEmptyMap k b -> Bool +disjointNE t1 t2 = disjointNE' t1 (NE t2) + +disjointNE' :: Ord k => NonEmptyMap k a -> Map k b -> Bool +disjointNE' (Bin' 1 k _ _ _) t = k `notMember` t +disjointNE' (Bin' _ k _ l r) t = not found && disjoint l lt && disjoint r gt where (lt,found,gt) = splitMember k t @@ -2407,7 +2758,7 @@ preserveMissing' = WhenMissing -- Force all the values in a tree. forceTree :: Map k a -> () -forceTree (Bin _ _ v l r) = v `seq` forceTree l `seq` forceTree r `seq` () +forceTree (NE (Bin' _ _ v l r)) = v `seq` forceTree l `seq` forceTree r `seq` () forceTree Tip = () -- | Map over the entries whose keys are missing from the other map. @@ -2591,6 +2942,18 @@ merge g1 g2 f m1 m2 = runIdentity $ mergeA g1 g2 f m1 m2 {-# INLINE merge #-} +mergeNE + :: Ord k + => SimpleWhenMissing k a c -- ^ What to do with keys in @m1@ but not @m2@ + -> SimpleWhenMissing k b c -- ^ What to do with keys in @m2@ but not @m1@ + -> SimpleWhenMatched k a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> NonEmptyMap k a -- ^ Map @m1@ + -> NonEmptyMap k b -- ^ Map @m2@ + -> Maybe (NonEmptyMap k c) +mergeNE g1 g2 f m1 m2 = runIdentity $ + mergeANE g1 g2 f m1 m2 +{-# INLINE mergeNE #-} + -- | An applicative version of 'merge'. -- -- 'mergeA' takes two 'WhenMissing' tactics, a 'WhenMatched' @@ -2669,7 +3032,7 @@ mergeA where go t1 Tip = g1t t1 go Tip t2 = g2t t2 - go (Bin _ kx x1 l1 r1) t2 = case splitLookup kx t2 of + go (NE (Bin' _ kx x1 l1 r1)) t2 = case splitLookup kx t2 of (l2, mx2, r2) -> case mx2 of Nothing -> liftA3 (\l' mx' r' -> maybe link2 (link kx) mx' l' r') l1l2 (g1k kx x1) r1r2 @@ -2680,6 +3043,31 @@ mergeA !r1r2 = go r1 r2 {-# INLINE mergeA #-} +mergeANE + :: (Applicative f, Ord k) + => WhenMissing f k a c -- ^ What to do with keys in @m1@ but not @m2@ + -> WhenMissing f k b c -- ^ What to do with keys in @m2@ but not @m1@ + -> WhenMatched f k a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> NonEmptyMap k a -- ^ Map @m1@ + -> NonEmptyMap k b -- ^ Map @m2@ + -> f (Maybe (NonEmptyMap k c)) +mergeANE + w0@(WhenMissing{missingKey = g1k}) + w1 + w2@(WhenMatched f) + (Bin' _ kx x1 l1 r1) + t2 + = fmap nonEmpty $ case splitLookupNE kx t2 of + (l2, mx2, r2) -> case mx2 of + Nothing -> liftA3 (\l' mx' r' -> maybe link2 (link kx) mx' l' r') + l1l2 (g1k kx x1) r1r2 + Just x2 -> liftA3 (\l' mx' r' -> maybe link2 (link kx) mx' l' r') + l1l2 (f kx x1 x2) r1r2 + where + !l1l2 = mergeA w0 w1 w2 l1 l2 + !r1r2 = mergeA w0 w1 w2 r1 r2 +{-# INLINE mergeANE #-} + {-------------------------------------------------------------------- MergeWithKey @@ -2727,11 +3115,11 @@ mergeWithKey f g1 g2 = go where go Tip t2 = g2 t2 go t1 Tip = g1 t1 - go (Bin _ kx x l1 r1) t2 = + go (NE (Bin' _ kx x l1 r1)) t2 = case found of Nothing -> case g1 (singleton kx x) of Tip -> link2 l' r' - (Bin _ _ x' Tip Tip) -> link kx x' l' r' + (NE (Bin' _ _ x' Tip Tip)) -> link kx x' l' r' _ -> error "mergeWithKey: Given function only1 does not fulfill required conditions (see documentation)" Just x2 -> case f kx x x2 of Nothing -> link2 l' r' @@ -2787,11 +3175,11 @@ isSubmapOfBy f t1 t2 submap' :: Ord a => (b -> c -> Bool) -> Map a b -> Map a c -> Bool submap' _ Tip _ = True submap' _ _ Tip = False -submap' f (Bin 1 kx x _ _) t +submap' f (NE (Bin' 1 kx x _ _)) t = case lookup kx t of Just y -> f x y Nothing -> False -submap' f (Bin _ kx x l r) t +submap' f (NE (Bin' _ kx x l r)) t = case found of Nothing -> False Just y -> f x y @@ -2856,7 +3244,7 @@ filter p m filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a filterWithKey _ Tip = Tip -filterWithKey p t@(Bin _ kx x l r) +filterWithKey p t@(NE (Bin' _ kx x l r)) | p kx x = if pl `ptrEq` l && pr `ptrEq` r then t else link kx x pl pr @@ -2868,7 +3256,7 @@ filterWithKey p t@(Bin _ kx x l r) -- predicate. filterWithKeyA :: Applicative f => (k -> a -> f Bool) -> Map k a -> f (Map k a) filterWithKeyA _ Tip = pure Tip -filterWithKeyA p t@(Bin _ kx x l r) = +filterWithKeyA p t@(NE (Bin' _ kx x l r)) = liftA3 combine (p kx x) (filterWithKeyA p l) (filterWithKeyA p r) where combine True pl pr @@ -2889,7 +3277,7 @@ filterWithKeyA p t@(Bin _ kx x l r) = takeWhileAntitone :: (k -> Bool) -> Map k a -> Map k a takeWhileAntitone _ Tip = Tip -takeWhileAntitone p (Bin _ kx x l r) +takeWhileAntitone p (NE (Bin' _ kx x l r)) | p kx = link kx x l (takeWhileAntitone p r) | otherwise = takeWhileAntitone p l @@ -2906,7 +3294,7 @@ takeWhileAntitone p (Bin _ kx x l r) dropWhileAntitone :: (k -> Bool) -> Map k a -> Map k a dropWhileAntitone _ Tip = Tip -dropWhileAntitone p (Bin _ kx x l r) +dropWhileAntitone p (NE (Bin' _ kx x l r)) | p kx = dropWhileAntitone p r | otherwise = link kx x (dropWhileAntitone p l) r @@ -2930,7 +3318,7 @@ spanAntitone :: (k -> Bool) -> Map k a -> (Map k a, Map k a) spanAntitone p0 m = toPair (go p0 m) where go _ Tip = Tip :*: Tip - go p (Bin _ kx x l r) + go p (NE (Bin' _ kx x l r)) | p kx = let u :*: v = go p r in link kx x l u :*: v | otherwise = let u :*: v = go p l in u :*: link kx x v r @@ -2958,7 +3346,7 @@ partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a,Map k a) partitionWithKey p0 t0 = toPair $ go p0 t0 where go _ Tip = (Tip :*: Tip) - go p t@(Bin _ kx x l r) + go p t@(NE (Bin' _ kx x l r)) | p kx x = (if l1 `ptrEq` l && r1 `ptrEq` r then t else link kx x l1 r1) :*: link2 l2 r2 @@ -2985,7 +3373,7 @@ mapMaybe f = mapMaybeWithKey (\_ x -> f x) mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b mapMaybeWithKey _ Tip = Tip -mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of +mapMaybeWithKey f (NE (Bin' _ kx x l r)) = case f kx x of Just y -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r) @@ -2994,11 +3382,24 @@ mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of -- @since 0.5.8 traverseMaybeWithKey :: Applicative f => (k -> a -> f (Maybe b)) -> Map k a -> f (Map k b) -traverseMaybeWithKey = go - where - go _ Tip = pure Tip - go f (Bin _ kx x Tip Tip) = maybe Tip (\x' -> Bin 1 kx x' Tip Tip) <$> f kx x - go f (Bin _ kx x l r) = liftA3 combine (go f l) (f kx x) (go f r) +traverseMaybeWithKey f = traverseMaybeWithKey' f + +traverseMaybeWithKeyNE + :: Applicative f + => (k -> a -> f (Maybe b)) -> NonEmptyMap k a -> f (Maybe (NonEmptyMap k b)) +traverseMaybeWithKeyNE f = fmap nonEmpty . traverseMaybeWithKeyFromNE' f + +traverseMaybeWithKey' + :: Applicative f + => (k -> a -> f (Maybe b)) -> Map k a -> f (Map k b) +traverseMaybeWithKey' _ Tip = pure Tip +traverseMaybeWithKey' f (NE m) = traverseMaybeWithKeyFromNE' f m + +traverseMaybeWithKeyFromNE' + :: Applicative f + => (k -> a -> f (Maybe b)) -> NonEmptyMap k a -> f (Map k b) +traverseMaybeWithKeyFromNE' f (Bin' _ kx x Tip Tip) = maybe Tip (\x' -> NE $ Bin' 1 kx x' Tip Tip) <$> f kx x +traverseMaybeWithKeyFromNE' f (Bin' _ kx x l r) = liftA3 combine (traverseMaybeWithKey' f l) (f kx x) (traverseMaybeWithKey' f r) where combine !l' mx !r' = case mx of Nothing -> link2 l' r' @@ -3030,7 +3431,7 @@ mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) mapEitherWithKey f0 t0 = toPair $ go f0 t0 where go _ Tip = (Tip :*: Tip) - go f (Bin _ kx x l r) = case f kx x of + go f (NE (Bin' _ kx x l r)) = case f kx x of Left y -> link kx y l1 r1 :*: link2 l2 r2 Right z -> link2 l1 r1 :*: link kx z l2 r2 where @@ -3045,12 +3446,23 @@ mapEitherWithKey f0 t0 = toPair $ go f0 t0 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] map :: (a -> b) -> Map k a -> Map k b -map f = go where - go Tip = Tip - go (Bin sx kx x l r) = Bin sx kx (f x) (go l) (go r) --- We use a `go` function to allow `map` to inline. This makes +map = map' + +mapNE :: (a -> b) -> NonEmptyMap k a -> NonEmptyMap k b +mapNE = mapNE' + +map' :: (a -> b) -> Map k a -> Map k b +map' _ Tip = Tip +map' f (NE m) = NE $ mapNE' f m + +mapNE' :: (a -> b) -> NonEmptyMap k a -> NonEmptyMap k b +mapNE' f (Bin' sx kx x l r) = Bin' sx kx (f x) (map' f l) (map' f r) + +-- We use a `map'` function to allow `map` to inline. This makes -- a big difference if someone uses `map (const x) m` instead -- of `x <$ m`; it doesn't seem to do any harm. +-- +-- TODO(@Ericson2314) aren't we screwing this up with non-empty? #ifdef __GLASGOW_HASKELL__ {-# NOINLINE [1] map #-} @@ -3067,7 +3479,10 @@ map f = go where mapWithKey :: (k -> a -> b) -> Map k a -> Map k b mapWithKey _ Tip = Tip -mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r) +mapWithKey f (NE m) = NE $ mapWithKeyNE f m + +mapWithKeyNE :: (k -> a -> b) -> NonEmptyMap k a -> NonEmptyMap k b +mapWithKeyNE f (Bin' sx kx x l r) = Bin' sx kx (f kx x) (mapWithKey f l) (mapWithKey f r) #ifdef __GLASGOW_HASKELL__ {-# NOINLINE [1] mapWithKey #-} @@ -3089,12 +3504,19 @@ mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(1, 'a'), (5, 'e')]) == Just (fromList [(1, 'b'), (5, 'f')]) -- > traverseWithKey (\k v -> if odd k then Just (succ v) else Nothing) (fromList [(2, 'c')]) == Nothing traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) -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) = liftA3 (flip (Bin s k)) (go l) (f k v) (go r) +traverseWithKey _ Tip = pure Tip +traverseWithKey f (NE m) = NE <$> traverseWithKeyNE f m + +traverseWithKeyNE :: Applicative t => (k -> a -> t b) -> NonEmptyMap k a -> t (NonEmptyMap k b) +traverseWithKeyNE f (Bin' 1 k v _ _) = (\v' -> Bin' 1 k v' Tip Tip) <$> f k v +traverseWithKeyNE f (Bin' s k v l r) = liftA3 + (\l' v' -> Bin' s k v' l') + (traverseWithKey f l) + (f k v) + (traverseWithKey f r) + {-# INLINE traverseWithKey #-} +{-# INLINE traverseWithKeyNE #-} -- | \(O(n)\). The function 'mapAccum' threads an accumulating -- argument through the map in ascending order of keys. @@ -3102,39 +3524,53 @@ traverseWithKey f = go -- > let f a b = (a ++ b, b ++ "X") -- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) -mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccum :: (a -> b -> (a, c)) -> a -> Map k b -> (a, Map k c) mapAccum f a m = mapAccumWithKey (\a' _ x' -> f a' x') a m +mapAccumNE :: (a -> b -> (a, c)) -> a -> NonEmptyMap k b -> (a, NonEmptyMap k c) +mapAccumNE f a m + = mapAccumWithKeyNE (\a' _ x' -> f a' x') a m + -- | \(O(n)\). The function 'mapAccumWithKey' threads an accumulating -- argument through the map in ascending order of keys. -- -- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") -- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) -mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumWithKey :: (a -> k -> b -> (a, c)) -> a -> Map k b -> (a, Map k c) mapAccumWithKey f a t = mapAccumL f a t +mapAccumWithKeyNE :: (a -> k -> b -> (a, c)) -> a -> NonEmptyMap k b -> (a, NonEmptyMap k c) +mapAccumWithKeyNE f a t + = mapAccumLNE f a t + -- | \(O(n)\). The function 'mapAccumL' threads an accumulating -- argument through the map in ascending order of keys. -mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumL _ a Tip = (a,Tip) -mapAccumL f a (Bin sx kx x l r) = +mapAccumL :: (a -> k -> b -> (a, c)) -> a -> Map k b -> (a, Map k c) +mapAccumL _ a Tip = (a, Tip) +mapAccumL f a (NE m) = fmap NE $ mapAccumLNE f a m + +mapAccumLNE :: (a -> k -> b -> (a, c)) -> a -> NonEmptyMap k b -> (a, NonEmptyMap k c) +mapAccumLNE f a (Bin' sx kx x l r) = let (a1,l') = mapAccumL f a l (a2,x') = f a1 kx x (a3,r') = mapAccumL f a2 r - in (a3,Bin sx kx x' l' r') + in (a3, Bin' sx kx x' l' r') -- | \(O(n)\). The function 'mapAccumRWithKey' threads an accumulating -- argument through the map in descending order of keys. -mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumRWithKey :: (a -> k -> b -> (a, c)) -> a -> Map k b -> (a, Map k c) mapAccumRWithKey _ a Tip = (a,Tip) -mapAccumRWithKey f a (Bin sx kx x l r) = +mapAccumRWithKey f a (NE m) = NE <$> mapAccumRWithKeyNE f a m + +mapAccumRWithKeyNE :: (a -> k -> b -> (a, c)) -> a -> NonEmptyMap k b -> (a, NonEmptyMap k c) +mapAccumRWithKeyNE f a (Bin' sx kx x l r) = let (a1,r') = mapAccumRWithKey f a r (a2,x') = f a1 kx x (a3,l') = mapAccumRWithKey f a2 l - in (a3,Bin sx kx x' l' r') + in (a3, Bin' sx kx x' l' r') -- | \(O(n \log n)\). -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. @@ -3189,10 +3625,13 @@ mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] -- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True -- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False -mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a +mapKeysMonotonic :: (k1 -> k2) -> Map k1 a -> Map k2 a mapKeysMonotonic _ Tip = Tip -mapKeysMonotonic f (Bin sz k x l r) = - Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r) +mapKeysMonotonic f (NE m) = NE $ mapKeysMonotonicNE f m + +mapKeysMonotonicNE :: (k1 -> k2) -> NonEmptyMap k1 a -> NonEmptyMap k2 a +mapKeysMonotonicNE f (Bin' sz k x l r) = + Bin' sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r) {-------------------------------------------------------------------- Folds @@ -3211,7 +3650,7 @@ foldr :: (a -> b -> b) -> b -> Map k a -> b foldr f z = go z where go z' Tip = z' - go z' (Bin _ _ x l r) = go (f x (go z' r)) l + go z' (NE (Bin' _ _ x l r)) = go (f x (go z' r)) l {-# INLINE foldr #-} -- | \(O(n)\). A strict version of 'foldr'. Each application of the operator is @@ -3221,7 +3660,7 @@ foldr' :: (a -> b -> b) -> b -> Map k a -> b foldr' f z = go z where go !z' Tip = z' - go z' (Bin _ _ x l r) = go (f x $! go z' r) l + go z' (NE (Bin' _ _ x l r)) = go (f x $! go z' r) l {-# INLINE foldr' #-} -- | \(O(n)\). Fold the values in the map using the given left-associative @@ -3237,7 +3676,7 @@ foldl :: (a -> b -> a) -> a -> Map k b -> a foldl f z = go z where go z' Tip = z' - go z' (Bin _ _ x l r) = go (f (go z' l) x) r + go z' (NE (Bin' _ _ x l r)) = go (f (go z' l) x) r {-# INLINE foldl #-} -- | \(O(n)\). A strict version of 'foldl'. Each application of the operator is @@ -3247,11 +3686,13 @@ foldl' :: (a -> b -> a) -> a -> Map k b -> a foldl' f z = go z where go !z' Tip = z' - go z' (Bin _ _ x l r) = + go z' (NE (Bin' _ _ x l r)) = let !z'' = go z' l in go (f z'' x) r {-# INLINE foldl' #-} +-- foldl :: (b -> b -> b) -> NonEmptyMap k b -> NonEmptyMap k + -- | \(O(n)\). Fold the keys and values in the map using the given right-associative -- binary operator, such that -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. @@ -3266,7 +3707,7 @@ foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b foldrWithKey f z = go z where go z' Tip = z' - go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l + go z' (NE (Bin' _ kx x l r)) = go (f kx x (go z' r)) l {-# INLINE foldrWithKey #-} -- | \(O(n)\). A strict version of 'foldrWithKey'. Each application of the operator is @@ -3276,7 +3717,7 @@ foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b foldrWithKey' f z = go z where go !z' Tip = z' - go z' (Bin _ kx x l r) = go (f kx x $! go z' r) l + go z' (NE (Bin' _ kx x l r)) = go (f kx x $! go z' r) l {-# INLINE foldrWithKey' #-} -- | \(O(n)\). Fold the keys and values in the map using the given left-associative @@ -3293,7 +3734,7 @@ foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a foldlWithKey f z = go z where go z' Tip = z' - go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r + go z' (NE (Bin' _ kx x l r)) = go (f (go z' l) kx x) r {-# INLINE foldlWithKey #-} -- | \(O(n)\). A strict version of 'foldlWithKey'. Each application of the operator is @@ -3303,7 +3744,7 @@ foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a foldlWithKey' f z = go z where go !z' Tip = z' - go z' (Bin _ kx x l r) = + go z' (NE (Bin' _ kx x l r)) = let !z'' = go z' l in go (f z'' kx x) r {-# INLINE foldlWithKey' #-} @@ -3319,8 +3760,8 @@ foldMapWithKey :: Monoid m => (k -> a -> m) -> Map k a -> m foldMapWithKey f = go where go Tip = mempty - go (Bin 1 k v _ _) = f k v - go (Bin _ k v l r) = go l `mappend` (f k v `mappend` go r) + go (NE (Bin' 1 k v _ _)) = f k v + go (NE (Bin' _ k v l r)) = go l `mappend` (f k v `mappend` go r) {-# INLINE foldMapWithKey #-} {-------------------------------------------------------------------- @@ -3362,7 +3803,8 @@ assocs m keysSet :: Map k a -> Set.Set k keysSet Tip = Set.Tip -keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r) +keysSet (NE (Bin' sz kx _ l r)) = Set.NE $ + Set.Bin' sz kx (keysSet l) (keysSet r) -- | \(O(n)\). The set of all elements of the map contained in 'Arg's. -- @@ -3381,7 +3823,7 @@ argSet (Bin sz kx x l r) = Set.Bin sz (Arg kx x) (argSet l) (argSet r) fromSet :: (k -> a) -> Set.Set k -> Map k a fromSet _ Set.Tip = Tip -fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r) +fromSet f (Set.NE (Set.Bin' sz x l r)) = NE $ Bin' sz x (f x) (fromSet f l) (fromSet f r) -- | \(O(n)\). Build a map from a set of elements contained inside 'Arg's. -- @@ -3419,9 +3861,9 @@ instance (Ord k) => GHCExts.IsList (Map k v) where -- create, it is not inlined, so we inline it manually. fromList :: Ord k => [(k,a)] -> Map k a fromList [] = Tip -fromList [(kx, x)] = Bin 1 kx x Tip Tip -fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip) xs0 - | otherwise = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 +fromList [(kx, x)] = NE $ Bin' 1 kx x Tip Tip +fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (NE (Bin' 1 kx0 x0 Tip Tip)) xs0 + | otherwise = go (1::Int) (NE (Bin' 1 kx0 x0 Tip Tip)) xs0 where not_ordered _ [] = False not_ordered kx ((ky,_) : _) = kx >= ky @@ -3444,8 +3886,8 @@ fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip T -- ordered so far. create !_ [] = (Tip, [], []) create s xs@(xp : xss) - | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip, [], xss) - | otherwise -> (Bin 1 kx x Tip Tip, xss, []) + | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (NE $ Bin' 1 kx x Tip Tip, [], xss) + | otherwise -> (NE $ Bin' 1 kx x Tip Tip, xss, []) | otherwise = case create (s `shiftR` 1) xs of res@(_, [], _) -> res (l, [(ky, y)], zs) -> (insertMax ky y l, [], zs) @@ -3701,7 +4143,7 @@ fromDescListWithKey f xs -- create, it is not inlined, so we inline it manually. fromDistinctAscList :: [(k,a)] -> Map k a fromDistinctAscList [] = Tip -fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 +fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (NE (Bin' 1 kx0 x0 Tip Tip)) xs0 where go !_ t [] = t go s l ((kx, x) : xs) = case create s xs of @@ -3710,7 +4152,7 @@ fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 create !_ [] = (Tip :*: []) create s xs@(x' : xs') - | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip :*: xs') + | s == 1 = case x' of (kx, x) -> (NE (Bin' 1 kx x Tip Tip) :*: xs') | otherwise = case create (s `shiftR` 1) xs of res@(_ :*: []) -> res (l :*: (ky, y):ys) -> case create (s `shiftR` 1) ys of @@ -3729,7 +4171,7 @@ fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 -- create, it is not inlined, so we inline it manually. fromDistinctDescList :: [(k,a)] -> Map k a fromDistinctDescList [] = Tip -fromDistinctDescList ((kx0, x0) : xs0) = go (1 :: Int) (Bin 1 kx0 x0 Tip Tip) xs0 +fromDistinctDescList ((kx0, x0) : xs0) = go (1 :: Int) (NE (Bin' 1 kx0 x0 Tip Tip)) xs0 where go !_ t [] = t go s r ((kx, x) : xs) = case create s xs of @@ -3738,7 +4180,7 @@ fromDistinctDescList ((kx0, x0) : xs0) = go (1 :: Int) (Bin 1 kx0 x0 Tip Tip) xs create !_ [] = (Tip :*: []) create s xs@(x' : xs') - | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip :*: xs') + | s == 1 = case x' of (kx, x) -> (NE (Bin' 1 kx x Tip Tip) :*: xs') | otherwise = case create (s `shiftR` 1) xs of res@(_ :*: []) -> res (r :*: (ky, y):ys) -> case create (s `shiftR` 1) ys of @@ -3757,7 +4199,7 @@ fromDistinctDescList ((kx0, x0) : xs0) = go (1 :: Int) (Bin 1 kx0 x0 Tip Tip) xs --------------------------------------------------------------------} filterGt :: Ord k => k -> Map k v -> Map k v filterGt !_ Tip = Tip -filterGt !b (Bin _ kx x l r) = +filterGt !b (NE (Bin' _ kx x l r)) = case compare b kx of LT -> link kx x (filterGt b l) r EQ -> r GT -> filterGt b r @@ -3767,7 +4209,7 @@ filterGt !b (Bin _ kx x l r) = filterLt :: Ord k => k -> Map k v -> Map k v filterLt !_ Tip = Tip -filterLt !b (Bin _ kx x l r) = +filterLt !b (NE (Bin' _ kx x l r)) = case compare kx b of LT -> link kx x l (filterLt b r) EQ -> l GT -> filterLt b l @@ -3794,8 +4236,8 @@ split !k0 t0 = toPair $ go k0 t0 where go k t = case t of - Tip -> Tip :*: Tip - Bin _ kx x l r -> case compare k kx of + Tip -> Tip :*: Tip + NE (Bin' _ kx x l r) -> case compare k kx of LT -> let (lt :*: gt) = go k l in lt :*: link kx x gt r GT -> let (lt :*: gt) = go k r in link kx x l lt :*: gt EQ -> (l :*: r) @@ -3811,24 +4253,32 @@ split !k0 t0 = toPair $ go k0 t0 -- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a") -- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty) -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty) -splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a) -splitLookup k0 m = case go k0 m of - StrictTriple l mv r -> (l, mv, r) - where - go :: Ord k => k -> Map k a -> StrictTriple (Map k a) (Maybe a) (Map k a) - go !k t = - case t of - Tip -> StrictTriple Tip Nothing Tip - Bin _ kx x l r -> case compare k kx of - LT -> let StrictTriple lt z gt = go k l - !gt' = link kx x gt r - in StrictTriple lt z gt' - GT -> let StrictTriple lt z gt = go k r - !lt' = link kx x l lt - in StrictTriple lt' z gt - EQ -> StrictTriple l (Just x) r +splitLookup :: Ord k => k -> Map k a -> (Map k a, Maybe a, Map k a) +splitLookup k0 m = case splitLookup' k0 m of + StrictTriple l mv r -> (l, mv, r) + +splitLookupNE :: Ord k => k -> NonEmptyMap k a -> (Map k a, Maybe a, Map k a) +splitLookupNE k0 m = case splitLookupNE' k0 m of + StrictTriple l mv r -> (l, mv, r) + +splitLookup' :: Ord k => k -> Map k a -> StrictTriple (Map k a) (Maybe a) (Map k a) +splitLookup' !k t = case t of + Tip -> StrictTriple Tip Nothing Tip + NE t' -> splitLookupNE' k t' + +splitLookupNE' :: Ord k => k -> NonEmptyMap k a -> StrictTriple (Map k a) (Maybe a) (Map k a) +splitLookupNE' !k (Bin' _ kx x l r) = case compare k kx of + LT -> let StrictTriple lt z gt = splitLookup' k l + !gt' = link kx x gt r + in StrictTriple lt z gt' + GT -> let StrictTriple lt z gt = splitLookup' k r + !lt' = link kx x l lt + in StrictTriple lt' z gt + EQ -> StrictTriple l (Just x) r + #if __GLASGOW_HASKELL__ {-# INLINABLE splitLookup #-} +{-# INLINABLE splitLookupNE #-} #endif -- | A variant of 'splitLookup' that indicates only whether the @@ -3843,7 +4293,7 @@ splitMember k0 m = case go k0 m of go !k t = case t of Tip -> StrictTriple Tip False Tip - Bin _ kx x l r -> case compare k kx of + NE (Bin' _ kx x l r) -> case compare k kx of LT -> let StrictTriple lt z gt = go k l !gt' = link kx x gt r in StrictTriple lt z gt' @@ -3863,7 +4313,7 @@ data StrictTriple a b c = StrictTriple !a !b !c in [r] > [k], and that [l] and [r] are valid trees. In order of sophistication: - [Bin sz k x l r] The type constructor. + [NE sz k x l r] The type constructor. [bin k x l r] Maintains the correct size, assumes that both [l] and [r] are balanced with respect to each other. [balance k x l r] Restores the balance and size. @@ -3883,27 +4333,30 @@ data StrictTriple a b c = StrictTriple !a !b !c Link --------------------------------------------------------------------} link :: k -> a -> Map k a -> Map k a -> Map k a -link kx x Tip r = insertMin kx x r -link kx x l Tip = insertMax kx x l -link kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) - | delta*sizeL < sizeR = balanceL kz z (link kx x l lz) rz - | delta*sizeR < sizeL = balanceR ky y ly (link kx x ry r) - | otherwise = bin kx x l r +link k x l r = NE $ linkNE k x l r + +linkNE :: k -> a -> Map k a -> Map k a -> NonEmptyMap k a +linkNE kx x Tip r = insertMinNE kx x r +linkNE kx x l Tip = insertMaxNE kx x l +linkNE kx x l@(NE (Bin' sizeL ky y ly ry)) r@(NE (Bin' sizeR kz z lz rz)) + | delta*sizeL < sizeR = balanceLNE kz z (linkNE kx x l lz) rz + | delta*sizeR < sizeL = balanceRNE ky y ly (linkNE kx x ry r) + | otherwise = binNE kx x l r -- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: k -> a -> Map k a -> Map k a -insertMax kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balanceR ky y l (insertMax kx x r) - -insertMin kx x t - = case t of - Tip -> singleton kx x - Bin _ ky y l r - -> balanceL ky y (insertMin kx x l) r +insertMax, insertMin :: k -> a -> Map k a -> Map k a +insertMax kx x t = NE $ insertMaxNE kx x t +insertMin kx x t = NE $ insertMinNE kx x t + +insertMaxNE, insertMinNE :: k -> a -> Map k a -> NonEmptyMap k a +insertMaxNE kx x t = case t of + Tip -> singletonNE kx x + NE (Bin' _ ky y l r) -> balanceRNE ky y l (insertMaxNE kx x r) + +insertMinNE kx x t = case t of + Tip -> singletonNE kx x + NE (Bin' _ ky y l r) -> balanceLNE ky y (insertMinNE kx x l) r {-------------------------------------------------------------------- [link2 l r]: merges two trees. @@ -3911,7 +4364,7 @@ insertMin kx x t link2 :: Map k a -> Map k a -> Map k a link2 Tip r = r link2 l Tip = l -link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) +link2 l@(NE (Bin' sizeL kx x lx rx)) r@(NE (Bin' sizeR ky y ly ry)) | delta*sizeL < sizeR = balanceL ky y (link2 l ly) ry | delta*sizeR < sizeL = balanceR kx x lx (link2 rx r) | otherwise = glue l r @@ -3923,7 +4376,7 @@ link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) glue :: Map k a -> Map k a -> Map k a glue Tip r = r glue l Tip = l -glue l@(Bin sl kl xl ll lr) r@(Bin sr kr xr rl rr) +glue l@(NE (Bin' sl kl xl ll lr)) r@(NE (Bin' sr kr xr rl rr)) | sl > sr = let !(MaxView km m l') = maxViewSure kl xl ll lr in balanceR km m l' r | otherwise = let !(MinView km m r') = minViewSure kr xr rl rr in balanceL km m l r' @@ -3934,7 +4387,7 @@ minViewSure :: k -> a -> Map k a -> Map k a -> MinView k a minViewSure = go where go k x Tip r = MinView k x r - go k x (Bin _ kl xl ll lr) r = + go k x (NE (Bin' _ kl xl ll lr)) r = case go kl xl ll lr of MinView km xm l' -> MinView km xm (balanceR k x l' r) {-# NOINLINE minViewSure #-} @@ -3943,7 +4396,7 @@ maxViewSure :: k -> a -> Map k a -> Map k a -> MaxView k a maxViewSure = go where go k x l Tip = MaxView k x l - go k x l (Bin _ kr xr rl rr) = + go k x l (NE (Bin' _ kr xr rl rr)) = case go kr xr rl rr of MaxView km xm r' -> MaxView km xm (balanceL k x l r') {-# NOINLINE maxViewSure #-} @@ -4009,64 +4462,78 @@ ratio = 2 -- -- balance :: k -> a -> Map k a -> Map k a -> Map k a -- balance k x l r --- | sizeL + sizeR <= 1 = Bin sizeX k x l r +-- | sizeL + sizeR <= 1 = NE 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 +-- | otherwise = NE sizeX k x l r -- where -- sizeL = size l -- sizeR = size r -- sizeX = sizeL + sizeR + 1 -- -- rotateL :: a -> b -> Map a b -> Map a b -> Map a b --- rotateL k x l r@(Bin _ _ _ ly ry) | size ly < ratio*size ry = singleL k x l r +-- rotateL k x l r@(NE (Bin' _ _ _ ly ry)) | size ly < ratio*size ry = singleL k x l r -- | otherwise = doubleL k x l r -- -- rotateR :: a -> b -> Map a b -> Map a b -> Map a b --- rotateR k x l@(Bin _ _ _ ly ry) r | size ry < ratio*size ly = singleR k x l r +-- rotateR k x l@(NE (Bin' _ _ _ ly ry)) r | size ry < ratio*size ly = singleR k x l r -- | otherwise = doubleR k x l r -- -- singleL, singleR :: a -> b -> Map a b -> Map a b -> Map a b --- singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 --- singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) +-- singleL k1 x1 t1 (NE (Bin' _ k2 x2 t2 t3)) = bin k2 x2 (bin k1 x1 t1 t2) t3 +-- singleR k1 x1 (NE (Bin' _ k2 x2 t1 t2)) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) -- -- doubleL, doubleR :: a -> b -> Map a b -> Map a b -> Map a b --- 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) --- 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) +-- doubleL k1 x1 t1 (NE (Bin' _ k2 x2 (NE _ k3 x3 t2 t3)) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4) +-- doubleR k1 x1 (NE (Bin' _ k2 x2 t1 (NE _ k3 x3 t2 t3))) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4) -- -- It is only written in such a way that every node is pattern-matched only once. balance :: k -> a -> Map k a -> Map k a -> Map k a balance k x l r = case l of Tip -> case r of - Tip -> Bin 1 k x Tip Tip - (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r - (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr - (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) - (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) - | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr - | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - - (Bin ls lk lx ll lr) -> case r of + Tip -> + NE $ Bin' 1 k x Tip Tip + (NE (Bin' _ _ _ Tip Tip)) -> + NE $ Bin' 2 k x Tip r + (NE (Bin' _ rk rx Tip rr@(NE (Bin' _ _ _ _ _)))) -> + NE $ Bin' 3 rk rx (NE (Bin' 1 k x Tip Tip)) rr + (NE (Bin' _ rk rx (NE (Bin' _ rlk rlx _ _)) Tip)) -> + NE $ Bin' 3 rlk rlx (NE $ Bin' 1 k x Tip Tip) (NE $ Bin' 1 rk rx Tip Tip) + (NE (Bin' rs rk rx rl@(NE (Bin' rls rlk rlx rll rlr)) rr@(NE (Bin' rrs _ _ _ _)))) + | rls < ratio*rrs -> NE $ Bin' (1+rs) rk rx + (NE $ Bin' (1+rls) k x Tip rl) + rr + | otherwise -> NE $ Bin' (1+rs) rlk rlx + (NE $ Bin' (1+size rll) k x Tip rll) + (NE $ Bin' (1+rrs+size rlr) rk rx rlr rr) + + (NE (Bin' ls lk lx ll lr)) -> case r of Tip -> case (ll, lr) of - (Tip, Tip) -> Bin 2 k x l Tip - (Tip, (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) - ((Bin _ _ _ _ _), Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) - ((Bin lls _ _ _ _), (Bin lrs lrk lrx lrl lrr)) - | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) - | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) - (Bin rs rk rx rl rr) + (Tip, Tip) -> NE $ Bin' 2 k x l Tip + (Tip, (NE (Bin' _ lrk lrx _ _))) -> NE $ Bin' 3 lrk lrx (NE (Bin' 1 lk lx Tip Tip)) (NE (Bin' 1 k x Tip Tip)) + ((NE (Bin' _ _ _ _ _)), Tip) -> NE $ Bin' 3 lk lx ll (NE (Bin' 1 k x Tip Tip)) + ((NE (Bin' lls _ _ _ _)), (NE (Bin' lrs lrk lrx lrl lrr))) + | lrs < ratio*lls -> NE $ Bin' (1+ls) lk lx ll (NE $ Bin' (1+lrs) k x lr Tip) + | otherwise -> NE $ Bin' (1+ls) lrk lrx + (NE $ Bin' (1+lls+size lrl) lk lx ll lrl) + (NE $ Bin' (1+size lrr) k x lrr Tip) + (NE (Bin' rs rk rx rl rr)) | rs > delta*ls -> case (rl, rr) of - (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) - | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr - | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + (NE (Bin' rls rlk rlx rll rlr), NE (Bin' rrs _ _ _ _)) + | rls < ratio*rrs -> NE $ Bin' (1+ls+rs) rk rx (NE $ Bin' (1+ls+rls) k x l rl) rr + | otherwise -> NE $ Bin' (1+ls+rs) rlk rlx + (NE $ Bin' (1+ls+size rll) k x l rll) + (NE $ Bin' (1+rrs+size rlr) rk rx rlr rr) (_, _) -> error "Failure in Data.Map.balance" | ls > delta*rs -> case (ll, lr) of - (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) - | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) - | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) + (NE (Bin' lls _ _ _ _), NE (Bin' lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> NE $ Bin' (1+ls+rs) lk lx ll (NE $ Bin' (1+rs+lrs) k x lr r) + | otherwise -> NE $ Bin' (1+ls+rs) lrk lrx + (NE $ Bin' (1+lls+size lrl) lk lx ll lrl) + (NE $ Bin' (1+rs+size lrr) k x lrr r) (_, _) -> error "Failure in Data.Map.balance" - | otherwise -> Bin (1+ls+rs) k x l r + | otherwise -> NE $ Bin' (1+ls+rs) k x l r {-# NOINLINE balance #-} -- Functions balanceL and balanceR are specialised versions of balance. @@ -4078,60 +4545,135 @@ balance k x l r = case l of balanceL :: k -> a -> Map k a -> Map k a -> Map k a balanceL k x l r = case r of Tip -> case l of - Tip -> Bin 1 k x Tip Tip - (Bin _ _ _ Tip Tip) -> Bin 2 k x l Tip - (Bin _ lk lx Tip (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) - (Bin _ lk lx ll@(Bin _ _ _ _ _) Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) - (Bin ls lk lx ll@(Bin lls _ _ _ _) lr@(Bin lrs lrk lrx lrl lrr)) - | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) - | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) - - (Bin rs _ _ _ _) -> case l of - Tip -> Bin (1+rs) k x Tip r - - (Bin ls lk lx ll lr) - | ls > delta*rs -> case (ll, lr) of - (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) - | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) - | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) - (_, _) -> error "Failure in Data.Map.balanceL" - | otherwise -> Bin (1+ls+rs) k x l r + Tip -> NE $ Bin' 1 k x Tip Tip + NE nel -> NE $ balanceLNEE k x nel + + NE ner@(Bin' rs _ _ _ _) -> case l of + Tip -> NE $ Bin' (1+rs) k x Tip r + NE nel -> NE $ balanceLNENE k x nel ner {-# NOINLINE balanceL #-} +balanceLNE :: k -> a -> NonEmptyMap k a -> Map k a -> NonEmptyMap k a +balanceLNE k x nel r = case r of + Tip -> balanceLNEE k x nel + NE ner -> balanceLNENE k x nel ner +{-# NOINLINE balanceLNE #-} + +-- | Balance helper where: +-- - Left child might be too big +-- - Left child is non-empty +-- - Right child is empty +balanceLNEE :: k -> a -> NonEmptyMap k a -> NonEmptyMap k a +balanceLNEE k x nel = case nel of + Bin' _ _ _ Tip Tip -> + Bin' 2 k x (NE nel) Tip + Bin' _ lk lx Tip (NE (Bin' _ lrk lrx _ _)) -> + Bin' 3 lrk lrx (NE (Bin' 1 lk lx Tip Tip)) (NE (Bin' 1 k x Tip Tip)) + Bin' _ lk lx ll@(NE (Bin' _ _ _ _ _)) Tip -> + Bin' 3 lk lx ll (NE (Bin' 1 k x Tip Tip)) + Bin' ls lk lx ll@(NE (Bin' lls _ _ _ _)) lr@(NE (Bin' lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> + Bin' (1+ls) lk lx ll (NE $ Bin' (1+lrs) k x lr Tip) + | otherwise -> + Bin' (1+ls) lrk lrx + (NE $ Bin' (1+lls+size lrl) lk lx ll lrl) + (NE $ Bin' (1+size lrr) k x lrr Tip) +{-# INLINE balanceLNEE #-} + +-- | Balance helper where: +-- - Left child might be too big +-- - Left child is non-empty +-- - Right child is non-empty +balanceLNENE :: k -> a -> NonEmptyMap k a -> NonEmptyMap k a -> NonEmptyMap k a +balanceLNENE k x l@(Bin' ls lk lx ll lr) r@(Bin' rs _ _ _ _) + | ls > delta*rs = case (ll, lr) of + (NE (Bin' lls _ _ _ _), NE (Bin' lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> Bin' (1+ls+rs) lk lx + ll + (NE $ Bin' (1+rs+lrs) k x lr $ NE r) + | otherwise -> Bin' (1+ls+rs) lrk lrx + (NE $ Bin' (1+lls+size lrl) lk lx ll lrl) + (NE $ Bin' (1+rs+size lrr) k x lrr $ NE r) + (_, _) -> error "Failure in Data.Map.balanceL" + | otherwise = Bin' (1+ls+rs) k x (NE l) (NE r) +{-# INLINE balanceLNENE #-} + -- balanceR is called when right subtree might have been inserted to or when -- left subtree might have been deleted from. balanceR :: k -> a -> Map k a -> Map k a -> Map k a balanceR k x l r = case l of Tip -> case r of - Tip -> Bin 1 k x Tip Tip - (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r - (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr - (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) - (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) - | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr - | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - - (Bin ls _ _ _ _) -> case r of - Tip -> Bin (1+ls) k x l Tip - - (Bin rs rk rx rl rr) - | rs > delta*ls -> case (rl, rr) of - (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) - | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr - | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) - (_, _) -> error "Failure in Data.Map.balanceR" - | otherwise -> Bin (1+ls+rs) k x l r + Tip -> NE $ Bin' 1 k x Tip Tip + (NE ner) -> NE $ balanceRNEE k x ner + + (NE nel@(Bin' ls _ _ _ _)) -> case r of + Tip -> NE $ Bin' (1+ls) k x l Tip + (NE ner) -> NE $ balanceRNENE k x nel ner {-# NOINLINE balanceR #-} +-- | Balance helper where: +-- - Right child might be too big +-- - Left child is empty +-- - Right child is non-empty +balanceRNE :: k -> a -> Map k a -> NonEmptyMap k a -> NonEmptyMap k a +balanceRNE k x l ner = case l of + Tip -> balanceRNEE k x ner + NE nel -> balanceRNENE k x nel ner +{-# NOINLINE balanceRNE #-} + +-- | Balance helper where: +-- - Right child might be too big +-- - Left child is non-empty +-- - Right child is empty +balanceRNEE :: k -> a -> NonEmptyMap k a -> NonEmptyMap k a +balanceRNEE k x ner = case ner of + Bin' _ _ _ Tip Tip -> + Bin' 2 k x Tip (NE ner) + Bin' _ rk rx Tip rr@(NE (Bin' _ _ _ _ _)) -> + Bin' 3 rk rx (NE $ Bin' 1 k x Tip Tip) rr + Bin' _ rk rx (NE (Bin' _ rlk rlx _ _)) Tip -> + Bin' 3 rlk rlx + (NE (Bin' 1 k x Tip Tip)) + (NE (Bin' 1 rk rx Tip Tip)) + Bin' rs rk rx rl@(NE (Bin' rls rlk rlx rll rlr)) rr@(NE (Bin' rrs _ _ _ _)) + | rls < ratio*rrs -> Bin' (1+rs) rk rx + (NE (Bin' (1+rls) k x Tip rl)) + rr + | otherwise -> Bin' (1+rs) rlk rlx + (NE (Bin' (1+size rll) k x Tip rll)) + (NE (Bin' (1+rrs+size rlr) rk rx rlr rr)) +{-# INLINE balanceRNEE #-} + +-- | Balance helper where: +-- - Right child might be too big +-- - Left child is non-empty +-- - Right child is non-empty +balanceRNENE :: k -> a -> NonEmptyMap k a -> NonEmptyMap k a -> NonEmptyMap k a +balanceRNENE k x l@(Bin' ls _ _ _ _) r@(Bin' rs rk rx rl rr) + | rs > delta*ls = case (rl, rr) of + (NE (Bin' rls rlk rlx rll rlr), NE (Bin' rrs _ _ _ _)) + | rls < ratio*rrs -> Bin' (1+ls+rs) rk rx + (NE (Bin' (1+ls+rls) k x (NE l) rl)) + rr + | otherwise -> Bin' (1+ls+rs) rlk rlx + (NE $ Bin' (1+ls+size rll) k x (NE l) rll) + (NE $ Bin' (1+rrs+size rlr) rk rx rlr rr) + (_, _) -> error "Failure in Data.Set.balanceR" + | otherwise = Bin' (1+ls+rs) k x (NE l) (NE r) +{-# INLINE balanceRNENE #-} {-------------------------------------------------------------------- The bin constructor maintains the size of the tree --------------------------------------------------------------------} bin :: k -> a -> Map k a -> Map k a -> Map k a -bin k x l r - = Bin (size l + size r + 1) k x l r +bin k x l r = NE $ binNE k x l r {-# INLINE bin #-} +binNE :: k -> a -> Map k a -> Map k a -> NonEmptyMap k a +binNE k x l r + = Bin' (size l + size r + 1) k x l r +{-# INLINE binNE #-} + {-------------------------------------------------------------------- Eq converts the tree to a list. In a lazy setting, this @@ -4197,7 +4739,7 @@ instance Functor (Map k) where fmap f m = map f m #ifdef __GLASGOW_HASKELL__ _ <$ Tip = Tip - a <$ (Bin sx kx _ l r) = Bin sx kx a (a <$ l) (a <$ r) + a <$ (NE (Bin' sx kx _ l r)) = NE $ Bin' sx kx a (a <$ l) (a <$ r) #endif -- | Traverses in order of increasing key. @@ -4209,8 +4751,8 @@ instance Traversable (Map k) where instance Foldable.Foldable (Map k) where fold = go where go Tip = mempty - go (Bin 1 _ v _ _) = v - go (Bin _ _ v l r) = go l `mappend` (v `mappend` go r) + go (NE (Bin' 1 _ v _ _)) = v + go (NE (Bin' _ _ v l r)) = go l `mappend` (v `mappend` go r) {-# INLINABLE fold #-} foldr = foldr {-# INLINE foldr #-} @@ -4218,8 +4760,8 @@ instance Foldable.Foldable (Map k) where {-# INLINE foldl #-} foldMap f t = go t where go Tip = mempty - go (Bin 1 _ v _ _) = f v - go (Bin _ _ v l r) = go l `mappend` (f v `mappend` go r) + go (NE (Bin' 1 _ v _ _)) = f v + go (NE (Bin' _ _ v l r)) = go l `mappend` (f v `mappend` go r) {-# INLINE foldMap #-} foldl' = foldl' {-# INLINE foldl' #-} @@ -4233,21 +4775,21 @@ instance Foldable.Foldable (Map k) where {-# INLINE toList #-} elem = go where go !_ Tip = False - go x (Bin _ _ v l r) = x == v || go x l || go x r + go x (NE (Bin' _ _ v l r)) = x == v || go x l || go x r {-# INLINABLE elem #-} maximum = start where start Tip = error "Data.Foldable.maximum (for Data.Map): empty map" - start (Bin _ _ v l r) = go (go v l) r + start (NE (Bin' _ _ v l r)) = go (go v l) r go !m Tip = m - go m (Bin _ _ v l r) = go (go (max m v) l) r + go m (NE (Bin' _ _ v l r)) = go (go (max m v) l) r {-# INLINABLE maximum #-} minimum = start where start Tip = error "Data.Foldable.minimum (for Data.Map): empty map" - start (Bin _ _ v l r) = go (go v l) r + start (NE (Bin' _ _ v l r)) = go (go v l) r go !m Tip = m - go m (Bin _ _ v l r) = go (go (min m v) l) r + go m (NE (Bin' _ _ v l r)) = go (go (min m v) l) r {-# INLINABLE minimum #-} sum = foldl' (+) 0 {-# INLINABLE sum #-} @@ -4279,7 +4821,37 @@ instance Bifoldable Map where instance (NFData k, NFData a) => NFData (Map k a) where rnf Tip = () - rnf (Bin _ kx x l r) = rnf kx `seq` rnf x `seq` rnf l `seq` rnf r + rnf (NE (Bin' _ kx x l r)) = rnf kx `seq` rnf x `seq` rnf l `seq` rnf r + +instance Functor (NonEmptyMap k) where + fmap f m = mapNE f m +#ifdef __GLASGOW_HASKELL__ + a <$ (Bin' sx kx _ l r) = Bin' sx kx a (a <$ l) (a <$ r) +#endif + +instance Traversable (NonEmptyMap k) where + traverse f = traverseWithKeyNE (\_ -> f) + {-# INLINE traverse #-} + +instance Foldable.Foldable (NonEmptyMap k) where + fold = goNE + where + goNE (Bin' 1 _ v _ _) = v + goNE (Bin' _ _ v l r) = mappend (go l) (mappend v $ go r) + go Tip = mempty + go (NE xs) = goNE xs + {-# INLINABLE fold #-} + foldMap f t = goNE t + where + goNE (Bin' 1 _ v _ _) = f v + goNE (Bin' _ _ v l r) = mappend (go l) (mappend (f v) $ go r) + go Tip = mempty + go (NE xs) = goNE xs + {-# INLINABLE foldMap #-} + +instance (NFData k, NFData a) => NFData (NonEmptyMap k a) where + rnf (Bin' _ kx x l r) = rnf kx `seq` rnf x `seq` rnf l `seq` rnf r + {-------------------------------------------------------------------- Read @@ -4333,6 +4905,6 @@ instance (Show k, Show a) => Show (Map k a) where splitRoot :: Map k b -> [Map k b] splitRoot orig = case orig of - Tip -> [] - Bin _ k v l r -> [l, singleton k v, r] + Tip -> [] + NE (Bin' _ k v l r) -> [l, singleton k v, r] {-# INLINE splitRoot #-} diff --git a/containers/src/Data/Map/Internal/Debug.hs b/containers/src/Data/Map/Internal/Debug.hs index f8644007b..a23faf3ed 100644 --- a/containers/src/Data/Map/Internal/Debug.hs +++ b/containers/src/Data/Map/Internal/Debug.hs @@ -3,7 +3,7 @@ module Data.Map.Internal.Debug where -import Data.Map.Internal (Map (..), size, delta) +import Data.Map.Internal (Map (..), NonEmptyMap (..), size, delta) import Control.Monad (guard) -- | \(O(n)\). Show the tree that implements the map. The tree is shown @@ -60,9 +60,9 @@ showsTree :: (k -> a -> String) -> Bool -> [String] -> [String] -> Map k a -> Sh showsTree showelem wide lbars rbars t = case t of Tip -> showsBars lbars . showString "|\n" - Bin _ kx x Tip Tip + NE (Bin' _ kx x Tip Tip) -> showsBars lbars . showString (showelem kx x) . showString "\n" - Bin _ kx x l r + NE (Bin' _ kx x l r) -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r . showWide wide rbars . showsBars lbars . showString (showelem kx x) . showString "\n" . @@ -73,9 +73,9 @@ showsTreeHang :: (k -> a -> String) -> Bool -> [String] -> Map k a -> ShowS showsTreeHang showelem wide bars t = case t of Tip -> showsBars bars . showString "|\n" - Bin _ kx x Tip Tip + NE (Bin' _ kx x Tip Tip) -> showsBars bars . showString (showelem kx x) . showString "\n" - Bin _ kx x l r + NE (Bin' _ kx x l r) -> showsBars bars . showString (showelem kx x) . showString "\n" . showWide wide bars . showsTreeHang showelem wide (withBar bars) l . @@ -119,15 +119,18 @@ ordered t where bounded lo hi t' = case t' of - Tip -> True - Bin _ kx _ l r -> (lo kx) && (hi kx) && bounded lo (kx) hi r + Tip -> True + NE (Bin' _ kx _ l r) + -> (lo kx) && (hi kx) && bounded lo (kx) hi r -- | Test if a map obeys the balance invariants. balanced :: Map k a -> 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)) && + Tip + -> True + NE (Bin' _ _ _ l r) + -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && balanced l && balanced r -- | Test if each node of a map reports its size correctly. @@ -137,7 +140,7 @@ validsize t = case slowSize t of Just _ -> True where slowSize Tip = Just 0 - slowSize (Bin sz _ _ l r) = do + slowSize (NE (Bin' sz _ _ l r)) = do ls <- slowSize l rs <- slowSize r guard (sz == ls + rs + 1) diff --git a/containers/src/Data/Map/Internal/DeprecatedShowTree.hs b/containers/src/Data/Map/Internal/DeprecatedShowTree.hs index 0b67acfe6..d818734b9 100644 --- a/containers/src/Data/Map/Internal/DeprecatedShowTree.hs +++ b/containers/src/Data/Map/Internal/DeprecatedShowTree.hs @@ -1,5 +1,4 @@ {-# LANGUAGE CPP, FlexibleContexts, DataKinds, MonoLocalBinds #-} - #include "containers.h" -- | This module simply holds disabled copies of functions from diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs index 7b9b2b795..158eebdee 100644 --- a/containers/src/Data/Map/Strict/Internal.hs +++ b/containers/src/Data/Map/Strict/Internal.hs @@ -86,6 +86,7 @@ module Data.Map.Strict.Internal -- * Map type Map(..) -- instance Eq,Show,Read + , NonEmptyMap(..) -- instance Eq,Show,Read , L.Size -- * Operators @@ -312,6 +313,7 @@ import Prelude hiding (lookup,map,filter,foldr,foldl,null,take,drop,splitAt) import Data.Map.Internal ( Map (..) + , NonEmptyMap (..) , AreWeStrict (..) , WhenMissing (..) , WhenMatched (..) @@ -477,7 +479,7 @@ findWithDefault :: Ord k => a -> k -> Map k a -> a findWithDefault def k = k `seq` go where go Tip = def - go (Bin _ kx x l r) = case compare k kx of + go (NE (Bin' _ kx x l r)) = case compare k kx of LT -> go l GT -> go r EQ -> x @@ -497,7 +499,7 @@ findWithDefault def k = k `seq` go -- > size (singleton 1 'a') == 1 singleton :: k -> a -> Map k a -singleton k x = x `seq` Bin 1 k x Tip Tip +singleton k x = x `seq` NE (Bin' 1 k x Tip Tip) {-# INLINE singleton #-} {-------------------------------------------------------------------- @@ -518,11 +520,11 @@ insert = go where go :: Ord k => k -> a -> Map k a -> Map k a go !kx !x Tip = singleton kx x - go kx x (Bin sz ky y l r) = + go kx x (NE (Bin' sz ky y l r)) = case compare kx ky of LT -> balanceL ky y (go kx x l) r GT -> balanceR ky y l (go kx x r) - EQ -> Bin sz kx x l r + EQ -> NE $ Bin' sz kx x l r #if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} #else @@ -544,11 +546,11 @@ insertWith = go where go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a go _ !kx x Tip = singleton kx x - go f !kx x (Bin sy ky y l r) = + go f !kx x (NE (Bin' sy ky y l r)) = case compare kx ky of LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) - EQ -> let !y' = f x y in Bin sy kx y' l r + EQ -> let !y' = f x y in NE $ Bin' sy kx y' l r #if __GLASGOW_HASKELL__ {-# INLINABLE insertWith #-} #else @@ -560,11 +562,11 @@ insertWithR = go where go :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a go _ !kx x Tip = singleton kx x - go f !kx x (Bin sy ky y l r) = + go f !kx x (NE (Bin' sy ky y l r)) = case compare kx ky of LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) - EQ -> let !y' = f y x in Bin sy ky y' l r + EQ -> let !y' = f y x in NE $ Bin' sy ky y' l r #if __GLASGOW_HASKELL__ {-# INLINABLE insertWithR #-} #else @@ -591,12 +593,12 @@ insertWithKey = go -- Forcing `kx` may look redundant, but it's possible `compare` will -- be lazy. go _ !kx x Tip = singleton kx x - go f kx x (Bin sy ky y l r) = + go f kx x (NE (Bin' sy ky y l r)) = case compare kx ky of LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> let !x' = f kx x y - in Bin sy kx x' l r + in NE $ Bin' sy kx x' l r #if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKey #-} #else @@ -610,12 +612,12 @@ insertWithKeyR = go -- Forcing `kx` may look redundant, but it's possible `compare` will -- be lazy. go _ !kx x Tip = singleton kx x - go f kx x (Bin sy ky y l r) = + go f kx x (NE (Bin' sy ky y l r)) = case compare kx ky of LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> let !y' = f ky y x - in Bin sy ky y' l r + in NE $ Bin' sy ky y' l r #if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKeyR #-} #else @@ -645,14 +647,14 @@ insertLookupWithKey f0 kx0 x0 t0 = toPair $ go f0 kx0 x0 t0 where go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a) go _ !kx x Tip = Nothing :*: singleton kx x - go f kx x (Bin sy ky y l r) = + go f kx x (NE (Bin' sy ky y l r)) = case compare kx ky of LT -> let (found :*: l') = go f kx x l in found :*: balanceL ky y l' r GT -> let (found :*: r') = go f kx x r in found :*: balanceR ky y l r' EQ -> let x' = f kx x y - in x' `seq` (Just y :*: Bin sy kx x' l r) + in x' `seq` (Just y :*: NE (Bin' sy kx x' l r)) #if __GLASGOW_HASKELL__ {-# INLINABLE insertLookupWithKey #-} #else @@ -692,11 +694,11 @@ adjustWithKey = go where go :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a go _ !_ Tip = Tip - go f k (Bin sx kx x l r) = + go f k (NE (Bin' sx kx x l r)) = case compare k kx of - LT -> Bin sx kx x (go f k l) r - GT -> Bin sx kx x l (go f k r) - EQ -> Bin sx kx x' l r + LT -> NE $ Bin' sx kx x (go f k l) r + GT -> NE $ Bin' sx kx x l (go f k r) + EQ -> NE $ Bin' sx kx x' l r where !x' = f kx x #if __GLASGOW_HASKELL__ {-# INLINABLE adjustWithKey #-} @@ -737,12 +739,12 @@ updateWithKey = go where go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a go _ !_ Tip = Tip - go f k(Bin sx kx x l r) = + go f k(NE (Bin' sx kx x l r)) = case compare k kx of LT -> balanceR kx x (go f k l) r GT -> balanceL kx x l (go f k r) EQ -> case f kx x of - Just x' -> x' `seq` Bin sx kx x' l r + Just x' -> x' `seq` NE (Bin' sx kx x' l r) Nothing -> glue l r #if __GLASGOW_HASKELL__ {-# INLINABLE updateWithKey #-} @@ -765,14 +767,14 @@ updateLookupWithKey f0 k0 t0 = toPair $ go f0 k0 t0 where go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) go _ !_ Tip = (Nothing :*: Tip) - go f k (Bin sx kx x l r) = + go f k (NE (Bin' sx kx x l r)) = case compare k kx of LT -> let (found :*: l') = go f k l in found :*: balanceR kx x l' r GT -> let (found :*: r') = go f k r in found :*: balanceL kx x l r' EQ -> case f kx x of - Just x' -> x' `seq` (Just x' :*: Bin sx kx x' l r) + Just x' -> x' `seq` (Just x' :*: NE (Bin' sx kx x' l r)) Nothing -> (Just x :*: glue l r) #if __GLASGOW_HASKELL__ {-# INLINABLE updateLookupWithKey #-} @@ -803,11 +805,11 @@ alter = go Nothing -> Tip Just x -> singleton k x - go f k (Bin sx kx x l r) = case compare k kx of + go f k (NE (Bin' sx kx x l r)) = case compare k kx of LT -> balance kx x (go f k l) r GT -> balance kx x l (go f k r) EQ -> case f (Just x) of - Just x' -> x' `seq` Bin sx kx x' l r + Just x' -> x' `seq` NE (Bin' sx kx x' l r) Nothing -> glue l r #if __GLASGOW_HASKELL__ {-# INLINABLE alter #-} @@ -896,11 +898,11 @@ updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a updateAt f i t = i `seq` case t of Tip -> error "Map.updateAt: index out of range" - Bin sx kx x l r -> case compare i sizeL of + NE (Bin' sx kx x l r) -> case compare i sizeL of LT -> balanceR kx x (updateAt f i l) r GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) EQ -> case f kx x of - Just x' -> x' `seq` Bin sx kx x' l r + Just x' -> x' `seq` NE (Bin' sx kx x' l r) Nothing -> glue l r where sizeL = size l @@ -934,11 +936,11 @@ updateMax f m -- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a -updateMinWithKey _ Tip = Tip -updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of - Nothing -> r - Just x' -> x' `seq` Bin sx kx x' Tip r -updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r +updateMinWithKey _ Tip = Tip +updateMinWithKey f (NE (Bin' sx kx x Tip r)) = case f kx x of + Nothing -> r + Just x' -> x' `seq` NE (Bin' sx kx x' Tip r) +updateMinWithKey f (NE (Bin' _ kx x l r)) = balanceR kx x (updateMinWithKey f l) r -- | \(O(\log n)\). Update the value at the maximal key. -- @@ -946,11 +948,11 @@ updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r -- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a -updateMaxWithKey _ Tip = Tip -updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of - Nothing -> l - Just x' -> x' `seq` Bin sx kx x' l Tip -updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) +updateMaxWithKey _ Tip = Tip +updateMaxWithKey f (NE (Bin' sx kx x l Tip)) = case f kx x of + Nothing -> l + Just x' -> x' `seq` NE (Bin' sx kx x' l Tip) +updateMaxWithKey f (NE (Bin' _ kx x l r)) = balanceL kx x l (updateMaxWithKey f r) {-------------------------------------------------------------------- Union. @@ -978,10 +980,10 @@ unionsWith f ts unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a unionWith _f t1 Tip = t1 -unionWith f t1 (Bin _ k x Tip Tip) = insertWithR f k x t1 -unionWith f (Bin _ k x Tip Tip) t2 = insertWith f k x t2 +unionWith f t1 (NE (Bin' _ k x Tip Tip)) = insertWithR f k x t1 +unionWith f (NE (Bin' _ k x Tip Tip)) t2 = insertWith f k x t2 unionWith _f Tip t2 = t2 -unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of +unionWith f (NE (Bin' _ k1 x1 l1 r1)) t2 = case splitLookup k1 t2 of (l2, mb, r2) -> link k1 x1' (unionWith f l1 l2) (unionWith f r1 r2) where !x1' = maybe x1 (f x1) mb #if __GLASGOW_HASKELL__ @@ -996,10 +998,10 @@ unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a unionWithKey _f t1 Tip = t1 -unionWithKey f t1 (Bin _ k x Tip Tip) = insertWithKeyR f k x t1 -unionWithKey f (Bin _ k x Tip Tip) t2 = insertWithKey f k x t2 +unionWithKey f t1 (NE (Bin' _ k x Tip Tip)) = insertWithKeyR f k x t1 +unionWithKey f (NE (Bin' _ k x Tip Tip)) t2 = insertWithKey f k x t2 unionWithKey _f Tip t2 = t2 -unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of +unionWithKey f (NE (Bin' _ k1 x1 l1 r1)) t2 = case splitLookup k1 t2 of (l2, mb, r2) -> link k1 x1' (unionWithKey f l1 l2) (unionWithKey f r1 r2) where !x1' = maybe x1 (f k1 x1) mb #if __GLASGOW_HASKELL__ @@ -1053,7 +1055,7 @@ differenceWithKey f = merge preserveMissing dropMissing (zipWithMaybeMatched f) intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c intersectionWith _f Tip _ = Tip intersectionWith _f _ Tip = Tip -intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of +intersectionWith f (NE (Bin' _ k x1 l1 r1)) t2 = case mb of Just x2 -> let !x1' = f x1 x2 in link k x1' l1l2 r1r2 Nothing -> link2 l1l2 r1r2 where @@ -1072,7 +1074,7 @@ intersectionWith f (Bin _ k x1 l1 r1) t2 = case mb of intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c intersectionWithKey _f Tip _ = Tip intersectionWithKey _f _ Tip = Tip -intersectionWithKey f (Bin _ k x1 l1 r1) t2 = case mb of +intersectionWithKey f (NE (Bin' _ k x1 l1 r1)) t2 = case mb of Just x2 -> let !x1' = f k x1 x2 in link k x1' l1l2 r1r2 Nothing -> link2 l1l2 r1r2 where @@ -1245,11 +1247,11 @@ mergeWithKey f g1 g2 = go where go Tip t2 = g2 t2 go t1 Tip = g1 t1 - go (Bin _ kx x l1 r1) t2 = + go (NE (Bin' _ kx x l1 r1)) t2 = case found of Nothing -> case g1 (singleton kx x) of Tip -> link2 l' r' - (Bin _ _ x' Tip Tip) -> link kx x' l' r' + (NE (Bin' _ _ x' Tip Tip)) -> link kx x' l' r' _ -> error "mergeWithKey: Given function only1 does not fulfill required conditions (see documentation)" Just x2 -> case f kx x x2 of Nothing -> link2 l' r' @@ -1279,7 +1281,7 @@ mapMaybe f = mapMaybeWithKey (\_ x -> f x) mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b mapMaybeWithKey _ Tip = Tip -mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of +mapMaybeWithKey f (NE (Bin' _ kx x l r)) = case f kx x of Just y -> y `seq` link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r) @@ -1292,8 +1294,8 @@ traverseMaybeWithKey :: Applicative f traverseMaybeWithKey = go where go _ Tip = pure Tip - go f (Bin _ kx x Tip Tip) = maybe Tip (\ !x' -> Bin 1 kx x' Tip Tip) <$> f kx x - go f (Bin _ kx x l r) = liftA3 combine (go f l) (f kx x) (go f r) + go f (NE (Bin' _ kx x Tip Tip)) = maybe Tip (\ !x' -> NE $ Bin' 1 kx x' Tip Tip) <$> f kx x + go f (NE (Bin' _ kx x l r)) = liftA3 combine (go f l) (f kx x) (go f r) where combine !l' mx !r' = case mx of Nothing -> link2 l' r' @@ -1325,7 +1327,7 @@ mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) mapEitherWithKey f0 t0 = toPair $ go f0 t0 where go _ Tip = (Tip :*: Tip) - go f (Bin _ kx x l r) = case f kx x of + go f (NE (Bin' _ kx x l r)) = case f kx x of Left y -> y `seq` (link kx y l1 r1 :*: link2 l2 r2) Right z -> z `seq` (link2 l1 r1 :*: link kx z l2 r2) where @@ -1343,7 +1345,7 @@ map :: (a -> b) -> Map k a -> Map k b map f = go where go Tip = Tip - go (Bin sx kx x l r) = let !x' = f x in Bin sx kx x' (go l) (go r) + go (NE (Bin' sx kx x l r)) = let !x' = f x in NE $ Bin' sx kx x' (go l) (go r) -- We use `go` to let `map` inline. This is important if `f` is a constant -- function. @@ -1362,9 +1364,9 @@ map f = go mapWithKey :: (k -> a -> b) -> Map k a -> Map k b mapWithKey _ Tip = Tip -mapWithKey f (Bin sx kx x l r) = +mapWithKey f (NE (Bin' sx kx x l r)) = let x' = f kx x - in x' `seq` Bin sx kx x' (mapWithKey f l) (mapWithKey f r) + in x' `seq` NE $ Bin' sx kx x' (mapWithKey f l) (mapWithKey f r) #ifdef __GLASGOW_HASKELL__ {-# NOINLINE [1] mapWithKey #-} @@ -1396,8 +1398,8 @@ traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) 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) = liftA3 (\ l' !v' r' -> Bin s k v' l' r') (go l) (f k v) (go r) + go (NE (Bin' 1 k v _ _)) = (\ !v' -> NE $ Bin' 1 k v' Tip Tip) <$> f k v + go (NE (Bin' s k v l r)) = liftA3 (\ l' !v' r' -> NE $ Bin' s k v' l' r') (go l) (f k v) (go r) {-# INLINE traverseWithKey #-} -- | \(O(n)\). The function 'mapAccum' threads an accumulating @@ -1423,22 +1425,22 @@ mapAccumWithKey f a t -- | \(O(n)\). The function 'mapAccumL' threads an accumulating -- argument through the map in ascending order of keys. mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) -mapAccumL _ a Tip = (a,Tip) -mapAccumL f a (Bin sx kx x l r) = +mapAccumL _ a Tip = (a,Tip) +mapAccumL f a (NE (Bin' sx kx x l r)) = let (a1,l') = mapAccumL f a l (a2,x') = f a1 kx x (a3,r') = mapAccumL f a2 r - in x' `seq` (a3,Bin sx kx x' l' r') + in x' `seq` (a3, NE $ Bin' sx kx x' l' r') -- | \(O(n)\). The function 'mapAccumRWithKey' threads an accumulating -- argument through the map in descending order of keys. mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) mapAccumRWithKey _ a Tip = (a,Tip) -mapAccumRWithKey f a (Bin sx kx x l r) = +mapAccumRWithKey f a (NE (Bin' sx kx x l r)) = let (a1,r') = mapAccumRWithKey f a r (a2,x') = f a1 kx x (a3,l') = mapAccumRWithKey f a2 l - in x' `seq` (a3,Bin sx kx x' l' r') + in x' `seq` (a3, NE $ Bin' sx kx x' l' r') -- | \(O(n \log n)\). -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. @@ -1469,7 +1471,8 @@ mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] fromSet :: (k -> a) -> Set.Set k -> Map k a fromSet _ Set.Tip = Tip -fromSet f (Set.Bin sz x l r) = case f x of v -> v `seq` Bin sz x v (fromSet f l) (fromSet f r) +fromSet f (Set.NE (Set.Bin' sz x l r)) = case f x of + v -> v `seq` NE (Bin' sz x v (fromSet f l) (fromSet f r)) -- | \(O(n)\). Build a map from a set of elements contained inside 'Arg's. -- @@ -1478,7 +1481,7 @@ fromSet f (Set.Bin sz x l r) = case f x of v -> v `seq` Bin sz x v (fromSet f l) fromArgSet :: Set.Set (Arg k a) -> Map k a fromArgSet Set.Tip = Tip -fromArgSet (Set.Bin sz (Arg x v) l r) = v `seq` Bin sz x v (fromArgSet l) (fromArgSet r) +fromArgSet (Set.NE (Set.Bin' sz (Arg x v) l r)) = v `seq` NE (Bin' sz x v (fromArgSet l) (fromArgSet r)) {-------------------------------------------------------------------- Lists @@ -1498,9 +1501,9 @@ fromArgSet (Set.Bin sz (Arg x v) l r) = v `seq` Bin sz x v (fromArgSet l) (fromA -- create, it is not inlined, so we inline it manually. fromList :: Ord k => [(k,a)] -> Map k a fromList [] = Tip -fromList [(kx, x)] = x `seq` Bin 1 kx x Tip Tip -fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0 x0 Tip Tip) xs0 - | otherwise = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 +fromList [(kx, x)] = x `seq` NE (Bin' 1 kx x Tip Tip) +fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (NE $ Bin' 1 kx0 x0 Tip Tip) xs0 + | otherwise = x0 `seq` go (1::Int) (NE $ Bin' 1 kx0 x0 Tip Tip) xs0 where not_ordered _ [] = False not_ordered kx ((ky,_) : _) = kx >= ky @@ -1523,8 +1526,8 @@ fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0 -- ordered so far. create !_ [] = (Tip, [], []) create s xs@(xp : xss) - | s == 1 = case xp of (kx, x) | not_ordered kx xss -> x `seq` (Bin 1 kx x Tip Tip, [], xss) - | otherwise -> x `seq` (Bin 1 kx x Tip Tip, xss, []) + | s == 1 = case xp of (kx, x) | not_ordered kx xss -> x `seq` (NE $ Bin' 1 kx x Tip Tip, [], xss) + | otherwise -> x `seq` (NE $ Bin' 1 kx x Tip Tip, xss, []) | otherwise = case create (s `shiftR` 1) xs of res@(_, [], _) -> res (l, [(ky, y)], zs) -> y `seq` (insertMax ky y l, [], zs) @@ -1697,7 +1700,7 @@ fromDescListWithKey f xs -- create, it is not inlined, so we inline it manually. fromDistinctAscList :: [(k,a)] -> Map k a fromDistinctAscList [] = Tip -fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 +fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (NE $ Bin' 1 kx0 x0 Tip Tip) xs0 where go !_ t [] = t go s l ((kx, x) : xs) = @@ -1707,7 +1710,7 @@ fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip T create !_ [] = (Tip :*: []) create s xs@(x' : xs') - | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip :*: xs') + | s == 1 = case x' of (kx, x) -> x `seq` (NE (Bin' 1 kx x Tip Tip) :*: xs') | otherwise = case create (s `shiftR` 1) xs of res@(_ :*: []) -> res (l :*: (ky, y):ys) -> case create (s `shiftR` 1) ys of @@ -1724,7 +1727,7 @@ fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip T -- create, it is not inlined, so we inline it manually. fromDistinctDescList :: [(k,a)] -> Map k a fromDistinctDescList [] = Tip -fromDistinctDescList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 +fromDistinctDescList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (NE (Bin' 1 kx0 x0 Tip Tip)) xs0 where go !_ t [] = t go s r ((kx, x) : xs) = @@ -1734,7 +1737,7 @@ fromDistinctDescList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip create !_ [] = (Tip :*: []) create s xs@(x' : xs') - | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip :*: xs') + | s == 1 = case x' of (kx, x) -> x `seq` (NE (Bin' 1 kx x Tip Tip) :*: xs') | otherwise = case create (s `shiftR` 1) xs of res@(_ :*: []) -> res (r :*: (ky, y):ys) -> case create (s `shiftR` 1) ys of diff --git a/containers/src/Data/Set.hs b/containers/src/Data/Set.hs index 0b4bc57d4..9f5f5f5e9 100644 --- a/containers/src/Data/Set.hs +++ b/containers/src/Data/Set.hs @@ -70,6 +70,7 @@ module Data.Set ( Set -- instance Eq,Ord,Show,Read,Data #else Set(..) + , NonEmptySet(..) #endif -- * Construction diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs index 7a0f15a51..d9d606b80 100644 --- a/containers/src/Data/Set/Internal.hs +++ b/containers/src/Data/Set/Internal.hs @@ -6,9 +6,11 @@ #endif #ifdef __GLASGOW_HASKELL__ {-# LANGUAGE DeriveLift #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RoleAnnotations #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} #endif {-# OPTIONS_HADDOCK not-home #-} @@ -118,12 +120,14 @@ -- Currently in GHC 7.0, when type has 2 constructors, a forward conditional -- jump is made when successfully matching second constructor. Successful match -- of first constructor results in the forward jump not taken. --- On GHC 7.0, reordering constructors from Tip | Bin to Bin | Tip +-- On GHC 7.0, reordering constructors from Tip | NE to NE | Tip -- improves the benchmark by up to 10% on x86. module Data.Set.Internal ( -- * Set type - Set(..) -- instance Eq,Ord,Show,Read,Data + Set(..) -- instance Eq,Ord,Show,Read,Data + , pattern Bin + , NonEmptySet(..) -- instance Eq,Ord,Show,Read,Data , Size -- * Operators @@ -131,79 +135,81 @@ module Data.Set.Internal ( -- * Query , null - , size - , member - , notMember - , lookupLT - , lookupGT - , lookupLE - , lookupGE - , isSubsetOf - , isProperSubsetOf - , disjoint + , nonEmpty + , size, sizeNE + , member, memberNE + , notMember, notMemberNE + , lookupLT, lookupLTNE + , lookupGT, lookupGTNE + , lookupLE, lookupLENE + , lookupGE, lookupGENE + , isSubsetOf, isSubsetOfNE + , isProperSubsetOf, isProperSubsetOfNE + , disjoint, disjointNE, disjointNEX -- * Construction , empty - , singleton - , insert - , delete + , singleton, singletonNE + , insert, insertNE + , delete, deleteNE , alterF - , powerSet + , powerSet, powerSetNE -- * Combine - , union + , union, unionNE , unions - , difference - , intersection + , difference, differenceNE + , intersection, intersectionNE , intersections - , cartesianProduct - , disjointUnion + , cartesianProduct, cartesianProductNE + , disjointUnion, disjointUnionNE, disjointUnionNEX, disjointUnionXNE , Intersection(..) -- * Filter - , filter - , takeWhileAntitone - , dropWhileAntitone - , spanAntitone - , partition - , split - , splitMember - , splitRoot + , filter, filterNE + , takeWhileAntitone, takeWhileAntitoneNE + , dropWhileAntitone, dropWhileAntitoneNE + , spanAntitone, spanAntitoneNE + , partition, partitionNE + , split, splitNE + , splitMember, splitMemberNE + , splitRoot, splitRootNE, splitNERootNE -- * Indexed - , lookupIndex - , findIndex - , elemAt - , deleteAt - , take - , drop - , splitAt + , lookupIndex, lookupIndexNE + , findIndex, findIndexNE + , elemAt, elemAtNE + , deleteAt, deleteAtNE + , take, takeNE + , drop, dropNE + , splitAt, splitAtNE -- * Map - , map - , mapMonotonic + , map, mapNE + , mapMonotonic, mapMonotonicNE -- * Folds - , foldr - , foldl - -- ** Strict folds - , foldr' - , foldl' + , foldr, foldr1 + , foldl, foldl1 + , foldr', foldr1' + , foldr1By, foldr1By' + , foldl1By, foldl1By' + , foldl', foldl1' -- ** Legacy folds , fold -- * Min\/Max - , lookupMin - , lookupMax + , lookupMin, lookupMinNE + , lookupMax, lookupMaxNE , findMin , findMax - , deleteMin - , deleteMax + , deleteMin, deleteMinNE + , deleteMax, deleteMaxNE , deleteFindMin , deleteFindMax - , maxView - , minView + , maxView, maxViewNE + , minView, minViewNE -- * Conversion @@ -211,6 +217,9 @@ module Data.Set.Internal ( , elems , toList , fromList + , elemsNE + , toListNE + , fromListNE -- ** Ordered list , toAscList @@ -219,20 +228,24 @@ module Data.Set.Internal ( , fromDistinctAscList , fromDescList , fromDistinctDescList + , toAscListNE + , toDescListNE + , fromDistinctAscListNE + , fromDistinctDescListNE -- * Debugging - , showTree - , showTreeWith - , valid + , showTree, showTreeNE + , showTreeWith, showTreeWithNE + , valid, validNE -- Internals (for testing) - , bin - , balanced - , link - , merge + , bin, binNE + , balanced, balancedNE + , link, linkNE + , merge, mergeNE ) where -import Prelude hiding (filter,foldl,foldr,null,map,take,drop,splitAt) +import Prelude hiding (filter,foldl,foldl1,foldr,foldr1,null,map,take,drop,splitAt) import Control.Applicative (Const(..)) import qualified Data.List as List import Data.Bits (shiftL, shiftR) @@ -245,6 +258,7 @@ import Data.Semigroup (stimesIdempotentMonoid, stimesIdempotent) import Data.Functor.Classes import Data.Functor.Identity (Identity) import qualified Data.Foldable as Foldable +import qualified Data.List.NonEmpty as NEL import Control.DeepSeq (NFData(rnf)) import Utils.Containers.Internal.StrictPair @@ -278,17 +292,27 @@ m1 \\ m2 = difference m1 m2 -- | A set of values @a@. -- See Note: Order of constructors -data Set a = Bin {-# UNPACK #-} !Size !a !(Set a) !(Set a) +data Set a = NE {-# UNPACK #-} !(NonEmptySet a) | Tip +data NonEmptySet a = Bin' {-# UNPACK #-} !Size !a !(Set a) !(Set a) + type Size = Int +#if __GLASGOW_HASKELL__ >= 802 +{-# COMPLETE Bin, Tip #-} +#endif #ifdef __GLASGOW_HASKELL__ +pattern Bin :: Size -> a -> Set a -> Set a -> Set a +pattern Bin s a l r = NE (Bin' s a l r) + type role Set nominal +type role NonEmptySet nominal #endif -- | @since FIXME deriving instance Lift a => Lift (Set a) +deriving instance Lift a => Lift (NonEmptySet a) instance Ord a => Monoid (Set a) where mempty = empty @@ -304,8 +328,8 @@ instance Ord a => Semigroup (Set a) where instance Foldable.Foldable Set where fold = go where go Tip = mempty - go (Bin 1 k _ _) = k - go (Bin _ k l r) = go l `mappend` (k `mappend` go r) + go (NE (Bin' 1 k _ _)) = k + go (NE (Bin' _ k l r)) = go l `mappend` (k `mappend` go r) {-# INLINABLE fold #-} foldr = foldr {-# INLINE foldr #-} @@ -313,8 +337,8 @@ instance Foldable.Foldable Set where {-# INLINE foldl #-} foldMap f t = go t where go Tip = mempty - go (Bin 1 k _ _) = f k - go (Bin _ k l r) = go l `mappend` (f k `mappend` go r) + go (NE (Bin' 1 k _ _)) = f k + go (NE (Bin' _ k l r)) = go l `mappend` (f k `mappend` go r) {-# INLINE foldMap #-} foldl' = foldl' {-# INLINE foldl' #-} @@ -328,7 +352,7 @@ instance Foldable.Foldable Set where {-# INLINE toList #-} elem = go where go !_ Tip = False - go x (Bin _ y l r) = x == y || go x l || go x r + go x (NE (Bin' _ y l r)) = x == y || go x l || go x r {-# INLINABLE elem #-} minimum = findMin {-# INLINE minimum #-} @@ -339,6 +363,44 @@ instance Foldable.Foldable Set where product = foldl' (*) 1 {-# INLINABLE product #-} +instance Foldable.Foldable NonEmptySet where + fold = goNE + where goNE (Bin' 1 k _ _) = k + goNE (Bin' _ k l r) = go l `mappend` (k `mappend` go r) + go Tip = mempty + go (NE s) = goNE s + {-# INLINABLE fold #-} + -- foldr f z s = foldr + -- {-# INLINE foldr #-} + -- foldl = foldl + -- {-# INLINE foldl #-} + foldMap f t = goNE t + where goNE (Bin' 1 k _ _) = f k + goNE (Bin' _ k l r) = go l `mappend` (f k `mappend` go r) + go Tip = mempty + go (NE s) = goNE s + {-# INLINE foldMap #-} + -- foldl' = foldl' + -- {-# INLINE foldl' #-} + -- foldr' = foldr' + -- {-# INLINE foldr' #-} + length = sizeNE + {-# INLINE length #-} + null _ = False + {-# INLINE null #-} + toList = NEL.toList . toListNE + {-# INLINE toList #-} + elem x xs = elem x $ NE xs + {-# INLINABLE elem #-} + minimum = lookupMinNE + {-# INLINE minimum #-} + maximum = lookupMaxNE + {-# INLINE maximum #-} + -- sum = foldl' (+) 0 + -- {-# INLINABLE sum #-} + -- product = foldl' (*) 1 + -- {-# INLINABLE product #-} + #if __GLASGOW_HASKELL__ {-------------------------------------------------------------------- @@ -370,29 +432,43 @@ setDataType = mkDataType "Data.Set.Internal.Set" [fromListConstr] --------------------------------------------------------------------} -- | \(O(1)\). Is this the empty set? null :: Set a -> Bool -null Tip = True -null (Bin {}) = False +null Tip = True +null (NE _) = False {-# INLINE null #-} +-- | /O(1)/. Return 'Just' if the set is not empty. +nonEmpty :: Set a -> Maybe (NonEmptySet a) +nonEmpty Tip = Nothing +nonEmpty (NE ne) = Just ne +{-# INLINE nonEmpty #-} + -- | \(O(1)\). The number of elements in the set. size :: Set a -> Int size Tip = 0 -size (Bin sz _ _ _) = sz +size (NE ne) = sizeNE ne {-# INLINE size #-} +sizeNE :: NonEmptySet a -> Int +sizeNE (Bin' sz _ _ _) = sz +{-# INLINE sizeNE #-} + -- | \(O(\log n)\). Is the element in the set? member :: Ord a => a -> Set a -> Bool -member = go - where - go !_ Tip = False - go x (Bin _ y l r) = case compare x y of - LT -> go x l - GT -> go x r - EQ -> True +member !_ Tip = False +member x (NE t) = memberNE x t + +memberNE :: Ord a => a -> NonEmptySet a -> Bool +memberNE !a (Bin' _ x l r) = case compare a x of + EQ -> True + LT -> member a l + GT -> member a r + #if __GLASGOW_HASKELL__ {-# INLINABLE member #-} +{-# INLINABLE memberNE #-} #else {-# INLINE member #-} +{-# INLINE memberNE #-} #endif -- | \(O(\log n)\). Is the element not in the set? @@ -404,90 +480,146 @@ notMember a t = not $ member a t {-# INLINE notMember #-} #endif +notMemberNE :: Ord a => a -> NonEmptySet a -> Bool +notMemberNE a t = not $ memberNE a t +#if __GLASGOW_HASKELL__ +{-# INLINABLE notMemberNE #-} +#else +{-# INLINE notMemberNE #-} +#endif + +-------------------------------------------------------------------- + -- | \(O(\log n)\). Find largest element smaller than the given one. -- -- > lookupLT 3 (fromList [3, 5]) == Nothing -- > lookupLT 5 (fromList [3, 5]) == Just 3 lookupLT :: Ord a => a -> Set a -> Maybe a -lookupLT = goNothing - where - goNothing !_ Tip = Nothing - goNothing x (Bin _ y l r) | x <= y = goNothing x l - | otherwise = goJust x y r +lookupLT !_ Tip = Nothing +lookupLT x (NE ne) = lookupLTNE x ne + +lookupLTNE :: Ord a => a -> NonEmptySet a -> Maybe a +lookupLTNE x (Bin' _ y l r) + | x <= y = lookupLT x l + | otherwise = lookupLTWithDefault x y r + +lookupLTWithDefault :: Ord a => a -> a -> Set a -> Maybe a +lookupLTWithDefault !_ best Tip = Just best +lookupLTWithDefault x best (NE ne) = lookupLTWithDefaultNE x best ne + +lookupLTWithDefaultNE :: Ord a => a -> a -> NonEmptySet a -> Maybe a +lookupLTWithDefaultNE x best (Bin' _ y l r) + | x <= y = lookupLTWithDefault x best l + | otherwise = lookupLTWithDefault x y r - goJust !_ best Tip = Just best - goJust x best (Bin _ y l r) | x <= y = goJust x best l - | otherwise = goJust x y r #if __GLASGOW_HASKELL__ {-# INLINABLE lookupLT #-} +{-# INLINABLE lookupLTNE #-} #else {-# INLINE lookupLT #-} +{-# INLINE lookupLTNE #-} #endif +-------------------------------------------------------------------- + -- | \(O(\log n)\). Find smallest element greater than the given one. -- -- > lookupGT 4 (fromList [3, 5]) == Just 5 -- > lookupGT 5 (fromList [3, 5]) == Nothing lookupGT :: Ord a => a -> Set a -> Maybe a -lookupGT = goNothing - where - goNothing !_ Tip = Nothing - goNothing x (Bin _ y l r) | x < y = goJust x y l - | otherwise = goNothing x r +lookupGT !_ Tip = Nothing +lookupGT x (NE ne) = lookupGTNE x ne + +lookupGTNE :: Ord a => a -> NonEmptySet a -> Maybe a +lookupGTNE x (Bin' _ y l r) + | x < y = lookupGTWithDefault x y l + | otherwise = lookupGT x r + +lookupGTWithDefault :: Ord a => a -> a -> Set a -> Maybe a +lookupGTWithDefault !_ best Tip = Just best +lookupGTWithDefault x best (NE ne) = lookupGTWithDefaultNE x best ne + +lookupGTWithDefaultNE :: Ord a => a -> a -> NonEmptySet a -> Maybe a +lookupGTWithDefaultNE x best (Bin' _ y l r) + | x < y = lookupGTWithDefault x y l + | otherwise = lookupGTWithDefault x best r - goJust !_ best Tip = Just best - goJust x best (Bin _ y l r) | x < y = goJust x y l - | otherwise = goJust x best r #if __GLASGOW_HASKELL__ {-# INLINABLE lookupGT #-} +{-# INLINABLE lookupGTNE #-} #else {-# INLINE lookupGT #-} +{-# INLINE lookupGTNE #-} #endif +-------------------------------------------------------------------- + -- | \(O(\log n)\). Find largest element smaller or equal to the given one. -- -- > lookupLE 2 (fromList [3, 5]) == Nothing -- > lookupLE 4 (fromList [3, 5]) == Just 3 -- > lookupLE 5 (fromList [3, 5]) == Just 5 lookupLE :: Ord a => a -> Set a -> Maybe a -lookupLE = goNothing - where - goNothing !_ Tip = Nothing - goNothing x (Bin _ y l r) = case compare x y of LT -> goNothing x l - EQ -> Just y - GT -> goJust x y r - - goJust !_ best Tip = Just best - goJust x best (Bin _ y l r) = case compare x y of LT -> goJust x best l - EQ -> Just y - GT -> goJust x y r +lookupLE !_ Tip = Nothing +lookupLE x (NE ne) = lookupLENE x ne + +lookupLENE :: Ord a => a -> NonEmptySet a -> Maybe a +lookupLENE x (Bin' _ y l r) = case compare x y of + LT -> lookupLE x l + EQ -> Just y + GT -> lookupLEWithDefault x y r + +lookupLEWithDefault :: Ord a => a -> a -> Set a -> Maybe a +lookupLEWithDefault !_ best Tip = Just best +lookupLEWithDefault x best (NE ne) = lookupLEWithDefaultNE x best ne + +lookupLEWithDefaultNE :: Ord a => a -> a -> NonEmptySet a -> Maybe a +lookupLEWithDefaultNE x best (Bin' _ y l r) = case compare x y of + LT -> lookupLEWithDefault x best l + EQ -> Just y + GT -> lookupLEWithDefault x y r + #if __GLASGOW_HASKELL__ {-# INLINABLE lookupLE #-} +{-# INLINABLE lookupLENE #-} #else {-# INLINE lookupLE #-} +{-# INLINE lookupLENE #-} #endif +-------------------------------------------------------------------- + -- | \(O(\log n)\). Find smallest element greater or equal to the given one. -- -- > lookupGE 3 (fromList [3, 5]) == Just 3 -- > lookupGE 4 (fromList [3, 5]) == Just 5 -- > lookupGE 6 (fromList [3, 5]) == Nothing lookupGE :: Ord a => a -> Set a -> Maybe a -lookupGE = goNothing - where - goNothing !_ Tip = Nothing - goNothing x (Bin _ y l r) = case compare x y of LT -> goJust x y l - EQ -> Just y - GT -> goNothing x r - - goJust !_ best Tip = Just best - goJust x best (Bin _ y l r) = case compare x y of LT -> goJust x y l - EQ -> Just y - GT -> goJust x best r +lookupGE !_ Tip = Nothing +lookupGE x (NE ne) = lookupGENE x ne + +lookupGENE :: Ord a => a -> NonEmptySet a -> Maybe a +lookupGENE x (Bin' _ y l r) = case compare x y of + LT -> lookupGEWithDefault x y l + EQ -> Just y + GT -> lookupGE x r + +lookupGEWithDefault :: Ord a => a -> a -> Set a -> Maybe a +lookupGEWithDefault !_ best Tip = Just best +lookupGEWithDefault x best (NE ne) = lookupGEWithDefaultNE x best ne + +lookupGEWithDefaultNE :: Ord a => a -> a -> NonEmptySet a -> Maybe a +lookupGEWithDefaultNE x best (Bin' _ y l r) = case compare x y of + LT -> lookupGEWithDefault x y l + EQ -> Just y + GT -> lookupGEWithDefault x best r + #if __GLASGOW_HASKELL__ {-# INLINABLE lookupGE #-} +{-# INLINABLE lookupGENE #-} #else {-# INLINE lookupGE #-} +{-# INLINE lookupGENE #-} #endif {-------------------------------------------------------------------- @@ -500,9 +632,13 @@ empty = Tip -- | \(O(1)\). Create a singleton set. singleton :: a -> Set a -singleton x = Bin 1 x Tip Tip +singleton = NE . singletonNE {-# INLINE singleton #-} +singletonNE :: a -> NonEmptySet a +singletonNE x = Bin' 1 x Tip Tip +{-# INLINE singletonNE #-} + {-------------------------------------------------------------------- Insertion, Deletion --------------------------------------------------------------------} @@ -513,23 +649,38 @@ singleton x = Bin 1 x Tip Tip -- See Note: Type of local 'go' function -- See Note: Avoiding worker/wrapper (in Data.Map.Internal) insert :: Ord a => a -> Set a -> Set a -insert x0 = go x0 x0 - where - go :: Ord a => a -> a -> Set a -> Set a - go orig !_ Tip = singleton (lazy orig) - go orig !x t@(Bin sz y l r) = case compare x y of - LT | l' `ptrEq` l -> t - | otherwise -> balanceL y l' r - where !l' = go orig x l - GT | r' `ptrEq` r -> t - | otherwise -> balanceR y l r' - where !r' = go orig x r - EQ | lazy orig `seq` (orig `ptrEq` y) -> t - | otherwise -> Bin sz (lazy orig) l r +insert x0 s0 = case insertReturningDifferent x0 x0 s0 of + Nothing -> s0 + Just q -> NE q + +insertNE :: Ord a => a -> NonEmptySet a -> NonEmptySet a +insertNE x0 s0 = case insertReturningDifferentNE x0 x0 s0 of + Nothing -> s0 + Just q -> q + +-- | Returns 'Nothing' if the element is already in the Set, and 'Just s' if a +-- new set had to be created to contain it. +insertReturningDifferent :: Ord a => a -> a -> Set a -> Maybe (NonEmptySet a) +insertReturningDifferent orig !_ Tip = Just $ singletonNE (lazy orig) +insertReturningDifferent orig !x (NE ne) = insertReturningDifferentNE orig x ne + +insertReturningDifferentNE :: Ord a => a -> a -> NonEmptySet a -> Maybe (NonEmptySet a) +insertReturningDifferentNE orig !x (Bin' sz y l r) = case compare x y of + LT -> case insertReturningDifferent orig x l of + Nothing -> Nothing + Just l' -> Just $! balanceLNE y l' r + GT -> case insertReturningDifferent orig x r of + Nothing -> Nothing + Just r' -> Just $! balanceRNE y l r' + EQ | lazy orig `seq` (orig `ptrEq` y) -> Nothing + | otherwise -> Just $ Bin' sz (lazy orig) l r + #if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} +{-# INLINABLE insertNE #-} #else {-# INLINE insert #-} +{-# INLINE insertNE #-} #endif #ifndef __GLASGOW_HASKELL__ @@ -537,50 +688,80 @@ lazy :: a -> a lazy a = a #endif +-------------------------------------------------------------------- + -- Insert an element to the set only if it is not in the set. -- Used by `union`. -- See Note: Type of local 'go' function -- See Note: Avoiding worker/wrapper (in Data.Map.Internal) insertR :: Ord a => a -> Set a -> Set a -insertR x0 = go x0 x0 - where - go :: Ord a => a -> a -> Set a -> Set a - go orig !_ Tip = singleton (lazy orig) - go orig !x t@(Bin _ y l r) = case compare x y of - LT | l' `ptrEq` l -> t - | otherwise -> balanceL y l' r - where !l' = go orig x l - GT | r' `ptrEq` r -> t - | otherwise -> balanceR y l r' - where !r' = go orig x r - EQ -> t +insertR x0 s0 = case insertRReturningDifferent x0 x0 s0 of + Nothing -> s0 + Just s -> NE s + +insertRNE :: Ord a => a -> NonEmptySet a -> NonEmptySet a +insertRNE x0 s0 = case insertRReturningDifferentNE x0 x0 s0 of + Nothing -> s0 + Just s -> s + +insertRReturningDifferent :: Ord a => a -> a -> Set a -> Maybe (NonEmptySet a) +insertRReturningDifferent orig !_ Tip = Just $ singletonNE (lazy orig) +insertRReturningDifferent orig !x (NE ne) = insertRReturningDifferentNE orig x ne + +insertRReturningDifferentNE :: Ord a => a -> a -> NonEmptySet a -> Maybe (NonEmptySet a) +insertRReturningDifferentNE orig !x (Bin' _ y l r) = case compare x y of + LT -> case insertRReturningDifferent orig x l of + Nothing -> Nothing + Just l' -> Just $! balanceLNE y l' r + GT -> case insertRReturningDifferent orig x r of + Nothing -> Nothing + Just r' -> Just $! balanceRNE y l r' + EQ -> Nothing + #if __GLASGOW_HASKELL__ {-# INLINABLE insertR #-} +{-# INLINABLE insertRNE #-} #else {-# INLINE insertR #-} +{-# INLINE insertRNE #-} #endif +-------------------------------------------------------------------- + -- | \(O(\log n)\). Delete an element from a set. --- See Note: Type of local 'go' function delete :: Ord a => a -> Set a -> Set a -delete = go - where - go :: Ord a => a -> Set a -> Set a - go !_ Tip = Tip - go x t@(Bin _ y l r) = case compare x y of - LT | l' `ptrEq` l -> t - | otherwise -> balanceR y l' r - where !l' = go x l - GT | r' `ptrEq` r -> t - | otherwise -> balanceL y l r' - where !r' = go x r - EQ -> glue l r +delete !_ Tip = Tip +delete x s0 = case deleteReturningDifferent x s0 of + Nothing -> s0 + Just s -> s + +deleteNE :: Ord a => a -> NonEmptySet a -> Set a +deleteNE x s0 = case deleteReturningDifferentNE x s0 of + Nothing -> NE s0 + Just s -> s + +deleteReturningDifferent :: Ord a => a -> Set a -> Maybe (Set a) +deleteReturningDifferent !_ Tip = Nothing +deleteReturningDifferent x (NE ne) = deleteReturningDifferentNE x ne + +deleteReturningDifferentNE :: Ord a => a -> NonEmptySet a -> Maybe (Set a) +deleteReturningDifferentNE !x (Bin' _ y l r) = case compare x y of + LT -> case deleteReturningDifferent x l of + Nothing -> Nothing + Just l' -> Just $ balanceR y l' r + GT -> case deleteReturningDifferent x r of + Nothing -> Nothing + Just r' -> Just $ balanceL y l r' + EQ -> Just $ glue l r + #if __GLASGOW_HASKELL__ {-# INLINABLE delete #-} +{-# INLINABLE deleteNE #-} #else {-# INLINE delete #-} +{-# INLINE deleteNE #-} #endif -- | \(O(\log n)\) @('alterF' f x s)@ can delete or insert @x@ in @s@ depending on @@ -632,8 +813,8 @@ alteredSet :: Ord a => a -> Set a -> AlteredSet a alteredSet x0 s0 = go x0 s0 where go :: Ord a => a -> Set a -> AlteredSet a - go x Tip = Inserted (singleton x) - go x (Bin _ y l r) = case compare x y of + go x Tip = Inserted (singleton x) + go x (NE (Bin' _ y l r)) = case compare x y of LT -> case go x l of Deleted d -> Deleted (balanceR y d r) Inserted i -> Inserted (balanceL y i r) @@ -659,11 +840,17 @@ alteredSet x0 s0 = go x0 s0 -- @ isProperSubsetOf :: Ord a => Set a -> Set a -> Bool isProperSubsetOf s1 s2 - = size s1 < size s2 && isSubsetOfX s1 s2 + = size s1 < size s2 && isSubsetOfSkipSize s1 s2 #if __GLASGOW_HASKELL__ {-# INLINABLE isProperSubsetOf #-} #endif +isProperSubsetOfNE :: Ord a => NonEmptySet a -> NonEmptySet a -> Bool +isProperSubsetOfNE s1 s2 + = sizeNE s1 < sizeNE s2 && isSubsetOfSkipSizeNE s1 (NE s2) +#if __GLASGOW_HASKELL__ +{-# INLINABLE isProperSubsetOfNE #-} +#endif -- | \(O\bigl(m \log\bigl(\frac{n+1}{m+1}\bigr)\bigr), \; m \leq n\). -- @(s1 \`isSubsetOf\` s2)@ indicates whether @s1@ is a subset of @s2@. @@ -676,11 +863,18 @@ isProperSubsetOf s1 s2 -- @ isSubsetOf :: Ord a => Set a -> Set a -> Bool isSubsetOf t1 t2 - = size t1 <= size t2 && isSubsetOfX t1 t2 + = size t1 <= size t2 && isSubsetOfSkipSize t1 t2 #if __GLASGOW_HASKELL__ {-# INLINABLE isSubsetOf #-} #endif +isSubsetOfNE :: Ord a => NonEmptySet a -> NonEmptySet a -> Bool +isSubsetOfNE t1 t2 + = sizeNE t1 <= sizeNE t2 && isSubsetOfSkipSizeNE t1 (NE t2) +#if __GLASGOW_HASKELL__ +{-# INLINABLE isSubsetOfNE #-} +#endif + -- Test whether a set is a subset of another without the *initial* -- size test. -- @@ -689,12 +883,15 @@ isSubsetOf t1 t2 -- et al needed to account for both "split work" and "merge work", we -- only have to worry about split work here, which is the same as in -- those functions. -isSubsetOfX :: Ord a => Set a -> Set a -> Bool -isSubsetOfX Tip _ = True -isSubsetOfX _ Tip = False +isSubsetOfSkipSize :: Ord a => Set a -> Set a -> Bool +isSubsetOfSkipSize Tip _ = True +isSubsetOfSkipSize _ Tip = False +isSubsetOfSkipSize (NE ne) t = isSubsetOfSkipSizeNE ne t + -- Skip the final split when we hit a singleton. -isSubsetOfX (Bin 1 x _ _) t = member x t -isSubsetOfX (Bin _ x l r) t +isSubsetOfSkipSizeNE :: Ord a => NonEmptySet a -> Set a -> Bool +isSubsetOfSkipSizeNE (Bin' 1 x _ _) t = member x t +isSubsetOfSkipSizeNE (Bin' _ x l r) t = found && -- Cheap size checks can sometimes save expensive recursive calls when the -- result will be False. Suppose we check whether [1..10] (with root 4) is @@ -710,11 +907,12 @@ isSubsetOfX (Bin _ x l r) t -- costs without obvious benefits. It might be worth considering if we find -- a way to use it to tighten the bounds in some useful/comprehensible way. size l <= size lt && size r <= size gt && - isSubsetOfX l lt && isSubsetOfX r gt + isSubsetOfSkipSize l lt && isSubsetOfSkipSize r gt where (lt,found,gt) = splitMember x t #if __GLASGOW_HASKELL__ -{-# INLINABLE isSubsetOfX #-} +{-# INLINABLE isSubsetOfSkipSize #-} +{-# INLINABLE isSubsetOfSkipSizeNE #-} #endif {-------------------------------------------------------------------- @@ -737,9 +935,15 @@ isSubsetOfX (Bin _ x l r) t disjoint :: Ord a => Set a -> Set a -> Bool disjoint Tip _ = True disjoint _ Tip = True +disjoint (NE ne) t = disjointNEX ne t + +disjointNE :: Ord a => NonEmptySet a -> NonEmptySet a -> Bool +disjointNE ne0 ne1 = disjointNEX ne0 $ NE ne1 + -- Avoid a split for the singleton case. -disjoint (Bin 1 x _ _) t = x `notMember` t -disjoint (Bin _ x l r) t +disjointNEX :: Ord a => NonEmptySet a -> Set a -> Bool +disjointNEX (Bin' 1 x _ _) t = x `notMember` t +disjointNEX (Bin' _ x l r) t -- Analogous implementation to `subsetOfX` = not found && disjoint l lt && disjoint r gt where @@ -755,7 +959,7 @@ disjoint (Bin _ x l r) t lookupMinSure :: a -> Set a -> a lookupMinSure x Tip = x -lookupMinSure _ (Bin _ x l _) = lookupMinSure x l +lookupMinSure _ (NE ne) = lookupMinNE ne -- | \(O(\log n)\). The minimal element of a set. -- @@ -763,7 +967,10 @@ lookupMinSure _ (Bin _ x l _) = lookupMinSure x l lookupMin :: Set a -> Maybe a lookupMin Tip = Nothing -lookupMin (Bin _ x l _) = Just $! lookupMinSure x l +lookupMin (NE (Bin' _ x l _)) = Just $! lookupMinSure x l + +lookupMinNE :: NonEmptySet a -> a +lookupMinNE (Bin' _ x l _) = lookupMinSure x l -- | \(O(\log n)\). The minimal element of a set. findMin :: Set a -> a @@ -773,7 +980,7 @@ findMin t lookupMaxSure :: a -> Set a -> a lookupMaxSure x Tip = x -lookupMaxSure _ (Bin _ x _ r) = lookupMaxSure x r +lookupMaxSure _ (NE (Bin' _ x _ r)) = lookupMaxSure x r -- | \(O(\log n)\). The maximal element of a set. -- @@ -781,7 +988,10 @@ lookupMaxSure _ (Bin _ x _ r) = lookupMaxSure x r lookupMax :: Set a -> Maybe a lookupMax Tip = Nothing -lookupMax (Bin _ x _ r) = Just $! lookupMaxSure x r +lookupMax (NE (Bin' _ x _ r)) = Just $! lookupMaxSure x r + +lookupMaxNE :: NonEmptySet a -> a +lookupMaxNE (Bin' _ x l _) = lookupMaxSure x l -- | \(O(\log n)\). The maximal element of a set. findMax :: Set a -> a @@ -791,15 +1001,21 @@ findMax t -- | \(O(\log n)\). Delete the minimal element. Returns an empty set if the set is empty. deleteMin :: Set a -> Set a -deleteMin (Bin _ _ Tip r) = r -deleteMin (Bin _ x l r) = balanceR x (deleteMin l) r -deleteMin Tip = Tip +deleteMin (NE ne) = deleteMinNE ne +deleteMin Tip = Tip + +deleteMinNE :: NonEmptySet a -> Set a +deleteMinNE (Bin' _ _ Tip r) = r +deleteMinNE (Bin' _ x (NE l) r) = balanceR x (deleteMinNE l) r -- | \(O(\log n)\). Delete the maximal element. Returns an empty set if the set is empty. deleteMax :: Set a -> Set a -deleteMax (Bin _ _ l Tip) = l -deleteMax (Bin _ x l r) = balanceL x l (deleteMax r) -deleteMax Tip = Tip +deleteMax (NE ne) = deleteMaxNE ne +deleteMax Tip = Tip + +deleteMaxNE :: NonEmptySet a -> Set a +deleteMaxNE (Bin' _ _ l Tip) = l +deleteMaxNE (Bin' _ x l (NE r)) = balanceL x l (deleteMaxNE r) {-------------------------------------------------------------------- Union. @@ -815,10 +1031,10 @@ unions = Foldable.foldl' union empty -- equal elements are encountered. union :: Ord a => Set a -> Set a -> Set a union t1 Tip = t1 -union t1 (Bin 1 x _ _) = insertR x t1 -union (Bin 1 x _ _) t2 = insert x t2 +union t1 (NE (Bin' 1 x _ _)) = insertR x t1 +union (NE (Bin' 1 x _ _)) t2 = insert x t2 union Tip t2 = t2 -union t1@(Bin _ x l1 r1) t2 = case splitS x t2 of +union t1@(NE (Bin' _ x l1 r1)) t2 = case splitS x t2 of (l2 :*: r2) | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1 | otherwise -> link x l1l2 r1r2 @@ -828,6 +1044,20 @@ union t1@(Bin _ x l1 r1) t2 = case splitS x t2 of {-# INLINABLE union #-} #endif +unionNE :: Ord a => NonEmptySet a -> NonEmptySet a -> NonEmptySet a +unionNE t1 (Bin' _ x Tip Tip) = insertRNE x t1 +unionNE (Bin' _ x Tip Tip) t2 = insertNE x t2 +unionNE t1@(Bin' _ x l1 r1) t2 = case splitS x (NE t2) of + (l2 :*: r2) + | l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 -> t1 + | otherwise -> linkNE x l1l2 r1r2 + where !l1l2 = union l1 l2 + !r1r2 = union r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE unionNE #-} +#endif + + {-------------------------------------------------------------------- Difference --------------------------------------------------------------------} @@ -839,14 +1069,22 @@ union t1@(Bin _ x l1 r1) t2 = case splitS x t2 of difference :: Ord a => Set a -> Set a -> Set a difference Tip _ = Tip difference t1 Tip = t1 -difference t1 (Bin _ x l2 r2) = case split x t1 of - (l1, r1) +difference t1 (NE t2) = differenceNE' t1 t2 + +differenceNE :: Ord a => NonEmptySet a -> NonEmptySet a -> Set a +differenceNE t1 t2 = differenceNE' (NE t1) t2 + +differenceNE' :: Ord a => Set a -> NonEmptySet a -> Set a +differenceNE' t1 (Bin' _ x l2 r2) = case splitS x t1 of + (l1 :*: r1) | size l1l2 + size r1r2 == size t1 -> t1 | otherwise -> merge l1l2 r1r2 where !l1l2 = difference l1 l2 !r1r2 = difference r1 r2 + #if __GLASGOW_HASKELL__ {-# INLINABLE difference #-} +{-# INLINABLE differenceNE #-} #endif {-------------------------------------------------------------------- @@ -866,7 +1104,7 @@ difference t1 (Bin _ x l2 r2) = case split x t1 of intersection :: Ord a => Set a -> Set a -> Set a intersection Tip _ = Tip intersection _ Tip = Tip -intersection t1@(Bin _ x l1 r1) t2 +intersection t1@(NE (Bin' _ x l1 r1)) t2 | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 then t1 else link x l1l2 r1r2 @@ -879,6 +1117,20 @@ intersection t1@(Bin _ x l1 r1) t2 {-# INLINABLE intersection #-} #endif +intersectionNE :: Ord a => NonEmptySet a -> NonEmptySet a -> Set a +intersectionNE t1@(Bin' _ x l1 r1) t2 + | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 + then NE t1 + else link x l1l2 r1r2 + | otherwise = merge l1l2 r1r2 + where + !(l2, b, r2) = splitMemberNE x t2 + !l1l2 = intersection l1 l2 + !r1r2 = intersection r1 r2 +#if __GLASGOW_HASKELL__ +{-# INLINABLE intersectionNE #-} +#endif + -- | The intersection of a series of sets. Intersections are performed left-to-right. intersections :: Ord a => NonEmpty (Set a) -> Set a intersections (s0 :| ss) = List.foldr go id ss s0 @@ -901,9 +1153,12 @@ instance (Ord a) => Semigroup (Intersection a) where -- | \(O(n)\). Filter all elements that satisfy the predicate. filter :: (a -> Bool) -> Set a -> Set a filter _ Tip = Tip -filter p t@(Bin _ x l r) +filter p (NE ne) = filterNE p ne + +filterNE :: (a -> Bool) -> NonEmptySet a -> Set a +filterNE p t@(Bin' _ x l r) | p x = if l `ptrEq` l' && r `ptrEq` r' - then t + then NE t else link x l' r' | otherwise = merge l' r' where @@ -913,19 +1168,26 @@ filter p t@(Bin _ x l r) -- | \(O(n)\). Partition the set into two sets, one with all elements that satisfy -- the predicate and one with all elements that don't satisfy the predicate. -- See also 'split'. -partition :: (a -> Bool) -> Set a -> (Set a,Set a) -partition p0 t0 = toPair $ go p0 t0 - where - go _ Tip = (Tip :*: Tip) - go p t@(Bin _ x l r) = case (go p l, go p r) of - ((l1 :*: l2), (r1 :*: r2)) - | p x -> (if l1 `ptrEq` l && r1 `ptrEq` r - then t - else link x l1 r1) :*: merge l2 r2 - | otherwise -> merge l1 r1 :*: - (if l2 `ptrEq` l && r2 `ptrEq` r - then t - else link x l2 r2) +partition :: (a -> Bool) -> Set a -> (Set a, Set a) +partition p0 t0 = toPair $ partitionS p0 t0 + +partitionNE :: (a -> Bool) -> NonEmptySet a -> (Set a, Set a) +partitionNE p0 t0 = toPair $ partitionSNE p0 t0 + +partitionS :: (a -> Bool) -> Set a -> StrictPair (Set a) (Set a) +partitionS _ Tip = Tip :*: Tip +partitionS p (NE ne) = partitionSNE p ne + +partitionSNE :: (a -> Bool) -> NonEmptySet a -> StrictPair (Set a) (Set a) +partitionSNE p t@(Bin' _ x l r) = case partitionS p l :*: partitionS p r of + (l1 :*: l2) :*: (r1 :*: r2) + | p x -> (NE $ if l1 `ptrEq` l && r1 `ptrEq` r + then t + else linkNE x l1 r1) :*: merge l2 r2 + | otherwise -> merge l1 r1 :*: + (NE $ if l2 `ptrEq` l && r2 `ptrEq` r + then t + else linkNE x l2 r2) {---------------------------------------------------------------------- Map @@ -943,6 +1205,9 @@ map f = fromList . List.map f . toList {-# INLINABLE map #-} #endif +mapNE :: Ord b => (a->b) -> NonEmptySet a -> NonEmptySet b +mapNE f = fromListNE . fmap f . toListNE + -- | \(O(n)\). The -- -- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly increasing. @@ -953,9 +1218,12 @@ map f = fromList . List.map f . toList -- > ==> mapMonotonic f s == map f s -- > where ls = toList s -mapMonotonic :: (a->b) -> Set a -> Set b +mapMonotonic :: (a -> b) -> Set a -> Set b mapMonotonic _ Tip = Tip -mapMonotonic f (Bin sz x l r) = Bin sz (f x) (mapMonotonic f l) (mapMonotonic f r) +mapMonotonic p (NE ne) = NE $ mapMonotonicNE p ne + +mapMonotonicNE :: (a -> b) -> NonEmptySet a -> NonEmptySet b +mapMonotonicNE f (Bin' sz x l r) = Bin' sz (f x) (mapMonotonic f l) (mapMonotonic f r) {-------------------------------------------------------------------- Fold @@ -979,9 +1247,23 @@ foldr :: (a -> b -> b) -> b -> Set a -> b foldr f z = go z where go z' Tip = z' - go z' (Bin _ x l r) = go (f x (go z' r)) l + go z' (NE (Bin' _ x l r)) = go (f x (go z' r)) l {-# INLINE foldr #-} +foldr1 :: (b -> b -> b) -> NonEmptySet b -> b +foldr1 f = foldr1By f id +{-# INLINE foldr1 #-} + +foldr1By :: forall a b . (a -> b -> b) -> (a -> b) -> NonEmptySet a -> b +foldr1By f g = go + where + finish :: Set a -> b -> b + finish l acc = foldr f acc l + go :: NonEmptySet a -> b + go (Bin' _ v l (NE r)) = finish l (f v (go r)) + go (Bin' _ v l Tip) = finish l (g v) +{-# INLINE foldr1By #-} + -- | \(O(n)\). A strict version of 'foldr'. Each application of the operator is -- evaluated before using the result in the next application. This -- function is strict in the starting value. @@ -989,9 +1271,25 @@ foldr' :: (a -> b -> b) -> b -> Set a -> b foldr' f z = go z where go !z' Tip = z' - go z' (Bin _ x l r) = go (f x $! go z' r) l + go z' (NE (Bin' _ x l r)) = go (f x $! go z' r) l {-# INLINE foldr' #-} +foldr1' :: (b -> b -> b) -> NonEmptySet b -> b +foldr1' f = go + where + go (Bin' _ x l Tip) = foldr' f x l + go (Bin' _ x l (NE r)) = foldr' f (f x (go r)) l +{-# INLINE foldr1' #-} + +foldr1By' :: forall a b . (a -> b -> b) -> (a -> b) -> NonEmptySet a -> b +foldr1By' f g = go + where + finish :: Set a -> b -> b + finish l !acc = foldr' f acc l + go :: NonEmptySet a -> b + go (Bin' _ v l (NE r)) = finish l (f v (go r)) + go (Bin' _ v l Tip) = finish l (g v) +{-# INLINE foldr1By' #-} -- | \(O(n)\). Fold the elements in the set using the given left-associative -- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'toAscList'@. -- @@ -1002,9 +1300,26 @@ foldl :: (a -> b -> a) -> a -> Set b -> a foldl f z = go z where go z' Tip = z' - go z' (Bin _ x l r) = go (f (go z' l) x) r + go z' (NE (Bin' _ x l r)) = go (f (go z' l) x) r {-# INLINE foldl #-} +foldl1 :: (b -> b -> b) -> NonEmptySet b -> b +foldl1 f = go + where + go (Bin' _ x Tip r) = foldl f x r + go (Bin' _ x (NE l) r) = foldl f (f (go l) x) r +{-# INLINE foldl1 #-} + +foldl1By :: forall a b . (b -> a -> b) -> (a -> b) -> NonEmptySet a -> b +foldl1By f g = go + where + finish :: b -> Set a -> b + finish acc r = foldl f acc r + go :: NonEmptySet a -> b + go (Bin' _ v (NE l) r) = finish (f (go l) v) r + go (Bin' _ v Tip r) = finish (g v) r +{-# INLINE foldl1By #-} + -- | \(O(n)\). A strict version of 'foldl'. Each application of the operator is -- evaluated before using the result in the next application. This -- function is strict in the starting value. @@ -1012,11 +1327,25 @@ foldl' :: (a -> b -> a) -> a -> Set b -> a foldl' f z = go z where go !z' Tip = z' - go z' (Bin _ x l r) = + go z' (NE (Bin' _ x l r)) = let !z'' = go z' l in go (f z'' x) r {-# INLINE foldl' #-} +foldl1' :: (b -> b -> b) -> NonEmptySet b -> b +foldl1' f = foldl1By' f id +{-# INLINE foldl1' #-} + +foldl1By' :: forall a b . (b -> a -> b) -> (a -> b) -> NonEmptySet a -> b +foldl1By' f g = go + where + finish :: b -> Set a -> b + finish !acc r = foldl' f acc r + go :: NonEmptySet a -> b + go (Bin' _ v (NE l) r) = finish (f (go l) v) r + go (Bin' _ v Tip r) = finish (g v) r +{-# INLINE foldl1By' #-} + {-------------------------------------------------------------------- List variations --------------------------------------------------------------------} @@ -1025,6 +1354,9 @@ foldl' f z = go z elems :: Set a -> [a] elems = toAscList +elemsNE :: NonEmptySet a -> NEL.NonEmpty a +elemsNE = toAscListNE + {-------------------------------------------------------------------- Lists --------------------------------------------------------------------} @@ -1041,15 +1373,24 @@ instance (Ord a) => GHCExts.IsList (Set a) where toList :: Set a -> [a] toList = toAscList +toListNE :: NonEmptySet a -> NEL.NonEmpty a +toListNE = toAscListNE + -- | \(O(n)\). Convert the set to an ascending list of elements. Subject to list fusion. toAscList :: Set a -> [a] toAscList = foldr (:) [] +toAscListNE :: NonEmptySet a -> NEL.NonEmpty a +toAscListNE = foldr1 (<>) . mapMonotonicNE pure + -- | \(O(n)\). Convert the set to a descending list of elements. Subject to list -- fusion. toDescList :: Set a -> [a] toDescList = foldl (flip (:)) [] +toDescListNE :: NonEmptySet a -> NEL.NonEmpty a +toDescListNE = foldl1 (<>) . mapMonotonicNE pure + -- List fusion for the list generating functions. #if __GLASGOW_HASKELL__ -- The foldrFB and foldlFB are foldr and foldl equivalents, used for list fusion. @@ -1088,23 +1429,34 @@ foldlFB = foldl -- create, it is not inlined, so we inline it manually. fromList :: Ord a => [a] -> Set a fromList [] = Tip -fromList [x] = Bin 1 x Tip Tip -fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0 - | otherwise = go (1::Int) (Bin 1 x0 Tip Tip) xs0 +fromList (x : xs) = NE $ fromListNE' x xs +#if __GLASGOW_HASKELL__ +{-# INLINABLE fromList #-} +#endif + +fromListNE :: Ord a => NEL.NonEmpty a -> NonEmptySet a +fromListNE (x NEL.:| xs) = fromListNE' x xs +{-# INLINABLE fromListNE #-} + +fromListNE' :: Ord a => a -> [a] -> NonEmptySet a +fromListNE' x [] = Bin' 1 x Tip Tip +fromListNE' x0 xs0 + | not_ordered x0 xs0 = fromList' (Bin' 1 x0 Tip Tip) xs0 + | otherwise = go (1::Int) (Bin' 1 x0 Tip Tip) xs0 where not_ordered _ [] = False not_ordered x (y : _) = x >= y {-# INLINE not_ordered #-} fromList' t0 xs = Foldable.foldl' ins t0 xs - where ins t x = insert x t + where ins t x = insertNE x t go !_ t [] = t - go _ t [x] = insertMax x t + go _ t [x] = insertMaxNE x (NE t) go s l xs@(x : xss) | not_ordered x xss = fromList' l xs | otherwise = case create s xss of - (r, ys, []) -> go (s `shiftL` 1) (link x l r) ys - (r, _, ys) -> fromList' (link x l r) ys + (r, ys, []) -> go (s `shiftL` 1) (linkNE x (NE l) r) ys + (r, _, ys) -> fromList' (linkNE x (NE l) r) ys -- The create is returning a triple (tree, xs, ys). Both xs and ys -- represent not yet processed elements and only one of them can be nonempty. @@ -1113,8 +1465,8 @@ fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0 -- ordered so far. create !_ [] = (Tip, [], []) create s xs@(x : xss) - | s == 1 = if not_ordered x xss then (Bin 1 x Tip Tip, [], xss) - else (Bin 1 x Tip Tip, xss, []) + | s == 1 = if not_ordered x xss then (NE $ Bin' 1 x Tip Tip, [], xss) + else (NE $ Bin' 1 x Tip Tip, xss, []) | otherwise = case create (s `shiftR` 1) xs of res@(_, [], _) -> res (l, [y], zs) -> (insertMax y l, [], zs) @@ -1122,7 +1474,7 @@ fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0 | otherwise -> case create (s `shiftR` 1) yss of (r, zs, ws) -> (link y l r, zs, ws) #if __GLASGOW_HASKELL__ -{-# INLINABLE fromList #-} +{-# INLINABLE fromListNE' #-} #endif {-------------------------------------------------------------------- @@ -1170,16 +1522,22 @@ combineEq (x : xs) = combineEq' x xs -- create, it is not inlined, so we inline it manually. fromDistinctAscList :: [a] -> Set a fromDistinctAscList [] = Tip -fromDistinctAscList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0 +fromDistinctAscList (x0 : xs0) = NE $ fromDistinctAscListNE' x0 xs0 + +fromDistinctAscListNE :: NEL.NonEmpty a -> NonEmptySet a +fromDistinctAscListNE (x NEL.:| xs) = fromDistinctAscListNE' x xs + +fromDistinctAscListNE' :: a -> [a] -> NonEmptySet a +fromDistinctAscListNE' x0 xs0 = go (1::Int) (Bin' 1 x0 Tip Tip) xs0 where go !_ t [] = t go s l (x : xs) = case create s xs of - (r :*: ys) -> let !t' = link x l r + (r :*: ys) -> let !t' = linkNE x (NE l) r in go (s `shiftL` 1) t' ys create !_ [] = (Tip :*: []) create s xs@(x : xs') - | s == 1 = (Bin 1 x Tip Tip :*: xs') + | s == 1 = (NE (Bin' 1 x Tip Tip) :*: xs') | otherwise = case create (s `shiftR` 1) xs of res@(_ :*: []) -> res (l :*: (y:ys)) -> case create (s `shiftR` 1) ys of @@ -1194,16 +1552,22 @@ fromDistinctAscList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0 -- @since 0.5.8 fromDistinctDescList :: [a] -> Set a fromDistinctDescList [] = Tip -fromDistinctDescList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0 +fromDistinctDescList (x0 : xs0) = NE $ fromDistinctDescListNE' x0 xs0 + +fromDistinctDescListNE :: NEL.NonEmpty a -> NonEmptySet a +fromDistinctDescListNE (x NEL.:| xs) = fromDistinctDescListNE' x xs + +fromDistinctDescListNE' :: a -> [a] -> NonEmptySet a +fromDistinctDescListNE' x0 xs0 = go (1::Int) (Bin' 1 x0 Tip Tip) xs0 where go !_ t [] = t go s r (x : xs) = case create s xs of - (l :*: ys) -> let !t' = link x l r + (l :*: ys) -> let !t' = linkNE x l (NE r) in go (s `shiftL` 1) t' ys create !_ [] = (Tip :*: []) create s xs@(x : xs') - | s == 1 = (Bin 1 x Tip Tip :*: xs') + | s == 1 = (NE (Bin' 1 x Tip Tip) :*: xs') | otherwise = case create (s `shiftR` 1) xs of res@(_ :*: []) -> res (r :*: (y:ys)) -> case create (s `shiftR` 1) ys of @@ -1270,7 +1634,7 @@ instance (Read a, Ord a) => Read (Set a) where instance NFData a => NFData (Set a) where rnf Tip = () - rnf (Bin _ y l r) = rnf y `seq` rnf l `seq` rnf r + rnf (NE (Bin' _ y l r)) = rnf y `seq` rnf l `seq` rnf r {-------------------------------------------------------------------- Split @@ -1282,20 +1646,34 @@ split :: Ord a => a -> Set a -> (Set a,Set a) split x t = toPair $ splitS x t {-# INLINABLE split #-} +splitNE :: Ord a => a -> NonEmptySet a -> (Set a,Set a) +splitNE x t = toPair $ splitSNE x t +{-# INLINABLE splitNE #-} + splitS :: Ord a => a -> Set a -> StrictPair (Set a) (Set a) splitS _ Tip = (Tip :*: Tip) -splitS x (Bin _ y l r) +splitS x (NE ne) = splitSNE x ne +{-# INLINABLE splitS #-} + +splitSNE :: Ord a => a -> NonEmptySet a -> StrictPair (Set a) (Set a) +splitSNE x (Bin' _ y l r) = case compare x y of LT -> let (lt :*: gt) = splitS x l in (lt :*: link y gt r) GT -> let (lt :*: gt) = splitS x r in (link y l lt :*: gt) EQ -> (l :*: r) -{-# INLINABLE splitS #-} +{-# INLINABLE splitSNE #-} -- | \(O(\log n)\). Performs a 'split' but also returns whether the pivot -- element was found in the original set. splitMember :: Ord a => a -> Set a -> (Set a,Bool,Set a) splitMember _ Tip = (Tip, False, Tip) -splitMember x (Bin _ y l r) +splitMember x (NE ne) = splitMemberNE x ne +#if __GLASGOW_HASKELL__ +{-# INLINABLE splitMember #-} +#endif + +splitMemberNE :: Ord a => a -> NonEmptySet a -> (Set a, Bool, Set a) +splitMemberNE x (Bin' _ y l r) = case compare x y of LT -> let (lt, found, gt) = splitMember x l !gt' = link y gt r @@ -1305,7 +1683,7 @@ splitMember x (Bin _ y l r) in (lt', found, gt) EQ -> (l, True, r) #if __GLASGOW_HASKELL__ -{-# INLINABLE splitMember #-} +{-# INLINABLE splitMemberNE #-} #endif {-------------------------------------------------------------------- @@ -1326,18 +1704,27 @@ splitMember x (Bin _ y l r) -- See Note: Type of local 'go' function findIndex :: Ord a => a -> Set a -> Int -findIndex = go 0 - where - go :: Ord a => Int -> a -> Set a -> Int - go !_ !_ Tip = error "Set.findIndex: element is not in the set" - go idx x (Bin _ kx l r) = case compare x kx of - LT -> go idx x l - GT -> go (idx + size l + 1) x r - EQ -> idx + size l +findIndex = findIndexS 0 #if __GLASGOW_HASKELL__ {-# INLINABLE findIndex #-} #endif +findIndexNE :: Ord a => a -> NonEmptySet a -> Int +findIndexNE = findIndexSNE 0 +#if __GLASGOW_HASKELL__ +{-# INLINABLE findIndexNE #-} +#endif + +findIndexS :: Ord a => Int -> a -> Set a -> Int +findIndexS !_ !_ Tip = error "Set.findIndex: element is not in the set" +findIndexS idx x (NE ne) = findIndexSNE idx x ne + +findIndexSNE :: Ord a => Int -> a -> NonEmptySet a -> Int +findIndexSNE idx x (Bin' _ kx l r) = case compare x kx of + LT -> findIndexS idx x l + GT -> findIndexS (idx + size l + 1) x r + EQ -> idx + size l + -- | \(O(\log n)\). Lookup the /index/ of an element, which is its zero-based index in -- the sorted sequence of elements. The index is a number from /0/ up to, but not -- including, the 'size' of the set. @@ -1351,18 +1738,27 @@ findIndex = go 0 -- See Note: Type of local 'go' function lookupIndex :: Ord a => a -> Set a -> Maybe Int -lookupIndex = go 0 - where - go :: Ord a => Int -> a -> Set a -> Maybe Int - go !_ !_ Tip = Nothing - go idx x (Bin _ kx l r) = case compare x kx of - LT -> go idx x l - GT -> go (idx + size l + 1) x r - EQ -> Just $! idx + size l +lookupIndex = lookupIndexS 0 #if __GLASGOW_HASKELL__ {-# INLINABLE lookupIndex #-} #endif +lookupIndexNE :: Ord a => a -> NonEmptySet a -> Maybe Int +lookupIndexNE = lookupIndexSNE 0 +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookupIndexNE #-} +#endif + +lookupIndexS :: Ord a => Int -> a -> Set a -> Maybe Int +lookupIndexS !_ !_ Tip = Nothing +lookupIndexS idx x (NE ne) = lookupIndexSNE idx x ne + +lookupIndexSNE :: Ord a => Int -> a -> NonEmptySet a -> Maybe Int +lookupIndexSNE idx x (Bin' _ kx l r) = case compare x kx of + LT -> lookupIndexS idx x l + GT -> lookupIndexS (idx + size l + 1) x r + EQ -> Just $! idx + size l + -- | \(O(\log n)\). Retrieve an element by its /index/, i.e. by its zero-based -- index in the sorted sequence of elements. If the /index/ is out of range (less -- than zero, greater or equal to 'size' of the set), 'error' is called. @@ -1375,7 +1771,10 @@ lookupIndex = go 0 elemAt :: Int -> Set a -> a elemAt !_ Tip = error "Set.elemAt: index out of range" -elemAt i (Bin _ x l r) +elemAt i (NE ne) = elemAtNE i ne + +elemAtNE :: Int -> NonEmptySet a -> a +elemAtNE i (Bin' _ x l r) = case compare i sizeL of LT -> elemAt i l GT -> elemAt (i-sizeL-1) r @@ -1395,15 +1794,16 @@ elemAt i (Bin _ x l r) -- @since 0.5.4 deleteAt :: Int -> Set a -> Set a -deleteAt !i t = - case t of - Tip -> error "Set.deleteAt: index out of range" - Bin _ x l r -> case compare i sizeL of - LT -> balanceR x (deleteAt i l) r - GT -> balanceL x l (deleteAt (i-sizeL-1) r) - EQ -> glue l r - where - sizeL = size l +deleteAt !_ Tip = error "Set.deleteAt: index out of range" +deleteAt i (NE ne) = deleteAtNE i ne + +deleteAtNE :: Int -> NonEmptySet a -> Set a +deleteAtNE !i (Bin' _ x l r) = case compare i sizeL of + LT -> balanceR x (deleteAt i l) r + GT -> balanceL x l (deleteAt (i-sizeL-1) r) + EQ -> glue l r + where + sizeL = size l -- | Take a given number of elements in order, beginning -- with the smallest ones. @@ -1415,16 +1815,25 @@ deleteAt !i t = -- @since 0.5.8 take :: Int -> Set a -> Set a take i m | i >= size m = m -take i0 m0 = go i0 m0 - where - go i !_ | i <= 0 = Tip - go !_ Tip = Tip - go i (Bin _ x l r) = - case compare i sizeL of - LT -> go i l - GT -> link x l (go (i - sizeL - 1) r) - EQ -> l - where sizeL = size l +take i0 m0 = takeS i0 m0 + +takeNE :: Int -> NonEmptySet a -> Set a +takeNE i m | i >= sizeNE m = NE m +takeNE i !_ | i <= 0 = Tip +takeNE i0 m0 = takeSNE i0 m0 + +takeS :: Int -> Set a -> Set a +takeS i !_ | i <= 0 = Tip +takeS !_ Tip = Tip +takeS i (NE ne) = takeSNE i ne + +takeSNE :: Int -> NonEmptySet a -> Set a +takeSNE i (Bin' _ x l r) = + case compare i sizeL of + LT -> takeS i l + GT -> link x l (takeS (i - sizeL - 1) r) + EQ -> l + where sizeL = size l -- | Drop a given number of elements in order, beginning -- with the smallest ones. @@ -1436,16 +1845,25 @@ take i0 m0 = go i0 m0 -- @since 0.5.8 drop :: Int -> Set a -> Set a drop i m | i >= size m = Tip -drop i0 m0 = go i0 m0 - where - go i m | i <= 0 = m - go !_ Tip = Tip - go i (Bin _ x l r) = - case compare i sizeL of - LT -> link x (go i l) r - GT -> go (i - sizeL - 1) r - EQ -> insertMin x r - where sizeL = size l +drop i0 m0 = dropS i0 m0 + +dropNE :: Int -> NonEmptySet a -> Set a +dropNE i m | i >= sizeNE m = Tip +dropNE i m | i <= 0 = NE m +dropNE i0 m0 = dropSNE i0 m0 + +dropS :: Int -> Set a -> Set a +dropS i m | i <= 0 = m +dropS !_ Tip = Tip +dropS i (NE ne) = dropSNE i ne + +dropSNE :: Int -> NonEmptySet a -> Set a +dropSNE i (Bin' _ x l r) = + case compare i sizeL of + LT -> link x (dropS i l) r + GT -> dropS (i - sizeL - 1) r + EQ -> insertMin x r + where sizeL = size l -- | \(O(\log n)\). Split a set at a particular index. -- @@ -1455,18 +1873,28 @@ drop i0 m0 = go i0 m0 splitAt :: Int -> Set a -> (Set a, Set a) splitAt i0 m0 | i0 >= size m0 = (m0, Tip) - | otherwise = toPair $ go i0 m0 - where - go i m | i <= 0 = Tip :*: m - go !_ Tip = Tip :*: Tip - go i (Bin _ x l r) - = case compare i sizeL of - LT -> case go i l of - ll :*: lr -> ll :*: link x lr r - GT -> case go (i - sizeL - 1) r of - rl :*: rr -> link x l rl :*: rr - EQ -> l :*: insertMin x r - where sizeL = size l + | otherwise = toPair $ splitAtS i0 m0 + +splitAtNE :: Int -> NonEmptySet a -> (Set a, Set a) +splitAtNE i0 m0 + | i0 >= sizeNE m0 = (NE m0, Tip) + | i0 <= 0 = (Tip, NE m0) + | otherwise = toPair $ splitAtSNE i0 m0 + +splitAtS :: Int -> Set a -> StrictPair (Set a) (Set a) +splitAtS i m | i <= 0 = Tip :*: m +splitAtS !_ Tip = Tip :*: Tip +splitAtS i (NE ne) = splitAtSNE i ne + +splitAtSNE :: Int -> NonEmptySet a -> StrictPair (Set a) (Set a) +splitAtSNE i (Bin' _ x l r) + = case compare i sizeL of + LT -> case splitAtS i l of + ll :*: lr -> ll :*: link x lr r + GT -> case splitAtS (i - sizeL - 1) r of + rl :*: rr -> link x l rl :*: rr + EQ -> l :*: insertMin x r + where sizeL = size l -- | \(O(\log n)\). Take while a predicate on the elements holds. -- The user is responsible for ensuring that for all elements @j@ and @k@ in the set, @@ -1481,7 +1909,10 @@ splitAt i0 m0 takeWhileAntitone :: (a -> Bool) -> Set a -> Set a takeWhileAntitone _ Tip = Tip -takeWhileAntitone p (Bin _ x l r) +takeWhileAntitone p (NE ne) = takeWhileAntitoneNE p ne + +takeWhileAntitoneNE :: (a -> Bool) -> NonEmptySet a -> Set a +takeWhileAntitoneNE p (Bin' _ x l r) | p x = link x l (takeWhileAntitone p r) | otherwise = takeWhileAntitone p l @@ -1498,7 +1929,10 @@ takeWhileAntitone p (Bin _ x l r) dropWhileAntitone :: (a -> Bool) -> Set a -> Set a dropWhileAntitone _ Tip = Tip -dropWhileAntitone p (Bin _ x l r) +dropWhileAntitone p (NE ne) = dropWhileAntitoneNE p ne + +dropWhileAntitoneNE :: (a -> Bool) -> NonEmptySet a -> Set a +dropWhileAntitoneNE p (Bin' _ x l r) | p x = dropWhileAntitone p r | otherwise = link x (dropWhileAntitone p l) r @@ -1519,12 +1953,19 @@ dropWhileAntitone p (Bin _ x l r) -- @since 0.5.8 spanAntitone :: (a -> Bool) -> Set a -> (Set a, Set a) -spanAntitone p0 m = toPair (go p0 m) - where - go _ Tip = Tip :*: Tip - go p (Bin _ x l r) - | p x = let u :*: v = go p r in link x l u :*: v - | otherwise = let u :*: v = go p l in u :*: link x v r +spanAntitone p m = toPair $ spanAntitoneS p m + +spanAntitoneNE :: (a -> Bool) -> NonEmptySet a -> (Set a, Set a) +spanAntitoneNE p m = toPair $ spanAntitoneSNE p m + +spanAntitoneS :: (a -> Bool) -> Set a -> StrictPair (Set a) (Set a) +spanAntitoneS _ Tip = Tip :*: Tip +spanAntitoneS p (NE ne) = spanAntitoneSNE p ne + +spanAntitoneSNE :: (a -> Bool) -> NonEmptySet a -> StrictPair (Set a) (Set a) +spanAntitoneSNE p (Bin' _ x l r) + | p x = let u :*: v = spanAntitoneS p r in link x l u :*: v + | otherwise = let u :*: v = spanAntitoneS p l in u :*: link x v r {-------------------------------------------------------------------- @@ -1533,7 +1974,7 @@ spanAntitone p0 m = toPair (go p0 m) in [r] > [x], and that [l] and [r] are valid trees. In order of sophistication: - [Bin sz x l r] The type constructor. + [NE sz x l r] The type constructor. [bin x l r] Maintains the correct size, assumes that both [l] and [r] are balanced with respect to each other. [balance x l r] Restores the balance and size. @@ -1552,28 +1993,47 @@ spanAntitone p0 m = toPair (go p0 m) {-------------------------------------------------------------------- Link --------------------------------------------------------------------} + link :: a -> Set a -> Set a -> Set a -link x Tip r = insertMin x r -link x l Tip = insertMax x l -link x l@(Bin sizeL y ly ry) r@(Bin sizeR z lz rz) - | delta*sizeL < sizeR = balanceL z (link x l lz) rz - | delta*sizeR < sizeL = balanceR y ly (link x ry r) - | otherwise = bin x l r +link x l r = NE $ linkNE x l r + +linkNE :: a -> Set a -> Set a -> NonEmptySet a +linkNE x Tip r = insertMinNE x r +linkNE x l Tip = insertMaxNE x l +linkNE x (NE l) (NE r) = linkNENE x l r + +linkXNE :: a -> Set a -> NonEmptySet a -> NonEmptySet a +linkXNE x Tip r = insertMinNE x (NE r) +linkXNE x (NE l) r = linkNENE x l r + +linkNEX :: a -> NonEmptySet a -> Set a -> NonEmptySet a +linkNEX x l Tip = insertMaxNE x (NE l) +linkNEX x l (NE r) = linkNENE x l r + +linkNENE :: a -> NonEmptySet a -> NonEmptySet a -> NonEmptySet a +linkNENE x l@(Bin' sizeL y ly ry) r@(Bin' sizeR z lz rz) + | delta*sizeL < sizeR = balanceLNE z (linkNEX x l lz) rz + | delta*sizeR < sizeL = balanceRNE y ly (linkXNE x ry r) + | otherwise = binNE x (NE l) (NE r) -- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: a -> Set a -> Set a -insertMax x t +insertMax, insertMin :: a -> Set a -> Set a +insertMaxNE, insertMinNE :: a -> Set a -> NonEmptySet a + +insertMax x t = NE $ insertMaxNE x t +insertMaxNE x t = case t of - Tip -> singleton x - Bin _ y l r - -> balanceR y l (insertMax x r) + Tip -> singletonNE x + NE (Bin' _ y l r) + -> balanceRNE y l (insertMaxNE x r) -insertMin x t +insertMin x t = NE $ insertMinNE x t +insertMinNE x t = case t of - Tip -> singleton x - Bin _ y l r - -> balanceL y (insertMin x l) r + Tip -> singletonNE x + NE (Bin' _ y l r) + -> balanceLNE y (insertMinNE x l) r {-------------------------------------------------------------------- [merge l r]: merges two trees. @@ -1581,10 +2041,21 @@ insertMin x t merge :: Set a -> Set a -> Set a merge Tip r = r merge l Tip = l -merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry) - | delta*sizeL < sizeR = balanceL y (merge l ly) ry - | delta*sizeR < sizeL = balanceR x lx (merge rx r) - | otherwise = glue l r +merge (NE l) (NE r) = NE $ mergeNE l r + +mergeXNE :: Set a -> NonEmptySet a -> NonEmptySet a +mergeXNE Tip r = r +mergeXNE (NE l) r = mergeNE l r + +mergeNEX :: NonEmptySet a -> Set a -> NonEmptySet a +mergeNEX l Tip = l +mergeNEX l (NE r) = mergeNE l r + +mergeNE :: NonEmptySet a -> NonEmptySet a -> NonEmptySet a +mergeNE l@(Bin' sizeL x lx rx) r@(Bin' sizeR y ly ry) + | delta*sizeL < sizeR = balanceLNE y (mergeNEX l ly) ry + | delta*sizeR < sizeL = balanceRNE x lx (mergeXNE rx r) + | otherwise = glueNE l r {-------------------------------------------------------------------- [glue l r]: glues two trees together. @@ -1593,9 +2064,12 @@ merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry) glue :: Set a -> Set a -> Set a glue Tip r = r glue l Tip = l -glue l@(Bin sl xl ll lr) r@(Bin sr xr rl rr) - | sl > sr = let !(m :*: l') = maxViewSure xl ll lr in balanceR m l' r - | otherwise = let !(m :*: r') = minViewSure xr rl rr in balanceL m l r' +glue (NE l) (NE r) = NE $ glueNE l r + +glueNE :: NonEmptySet a -> NonEmptySet a -> NonEmptySet a +glueNE l@(Bin' sl xl ll lr) r@(Bin' sr xr rl rr) + | sl > sr = let !(m :*: l') = maxViewSure xl ll lr in balanceRNE m l' r + | otherwise = let !(m :*: r') = minViewSure xr rl rr in balanceLNE m l r' -- | \(O(\log n)\). Delete and find the minimal element. -- @@ -1618,7 +2092,7 @@ minViewSure :: a -> Set a -> Set a -> StrictPair a (Set a) minViewSure = go where go x Tip r = x :*: r - go x (Bin _ xl ll lr) r = + go x (NE (Bin' _ xl ll lr)) r = case go xl ll lr of xm :*: l' -> xm :*: balanceR x l' r @@ -1626,13 +2100,16 @@ minViewSure = go -- stripped of that element, or 'Nothing' if passed an empty set. minView :: Set a -> Maybe (a, Set a) minView Tip = Nothing -minView (Bin _ x l r) = Just $! toPair $ minViewSure x l r +minView (NE ne) = Just $! minViewNE ne + +minViewNE :: NonEmptySet a -> (a, Set a) +minViewNE (Bin' _ x l r) = toPair $ minViewSure x l r maxViewSure :: a -> Set a -> Set a -> StrictPair a (Set a) maxViewSure = go where go x l Tip = x :*: l - go x l (Bin _ xr rl rr) = + go x l (NE (Bin' _ xr rl rr)) = case go xr rl rr of xm :*: r' -> xm :*: balanceL x l r' @@ -1640,7 +2117,10 @@ maxViewSure = go -- stripped of that element, or 'Nothing' if passed an empty set. maxView :: Set a -> Maybe (a, Set a) maxView Tip = Nothing -maxView (Bin _ x l r) = Just $! toPair $ maxViewSure x l r +maxView (NE ne) = Just $! maxViewNE ne + +maxViewNE :: NonEmptySet a -> (a, Set a) +maxViewNE (Bin' _ x l r) = toPair $ maxViewSure x l r {-------------------------------------------------------------------- [balance x l r] balances two trees with value x. @@ -1683,29 +2163,29 @@ ratio = 2 -- -- balance :: a -> Set a -> Set a -> Set a -- balance x l r --- | sizeL + sizeR <= 1 = Bin sizeX x l r +-- | sizeL + sizeR <= 1 = NE $ Bin' sizeX x l r -- | sizeR > delta*sizeL = rotateL x l r -- | sizeL > delta*sizeR = rotateR x l r --- | otherwise = Bin sizeX x l r +-- | otherwise = NE $ Bin' sizeX x l r -- where -- sizeL = size l -- sizeR = size r -- sizeX = sizeL + sizeR + 1 -- -- rotateL :: a -> Set a -> Set a -> Set a --- rotateL x l r@(Bin _ _ ly ry) | size ly < ratio*size ry = singleL x l r +-- rotateL x l r@(NE _ _ ly ry) | size ly < ratio*size ry = singleL x l r -- | otherwise = doubleL x l r -- rotateR :: a -> Set a -> Set a -> Set a --- rotateR x l@(Bin _ _ ly ry) r | size ry < ratio*size ly = singleR x l r +-- rotateR x l@(NE _ _ ly ry) r | size ry < ratio*size ly = singleR x l r -- | otherwise = doubleR x l r -- -- singleL, singleR :: a -> Set a -> Set a -> Set a --- singleL x1 t1 (Bin _ x2 t2 t3) = bin x2 (bin x1 t1 t2) t3 --- singleR x1 (Bin _ x2 t1 t2) t3 = bin x2 t1 (bin x1 t2 t3) +-- singleL x1 t1 (NE _ x2 t2 t3) = bin x2 (bin x1 t1 t2) t3 +-- singleR x1 (NE _ x2 t1 t2) t3 = bin x2 t1 (bin x1 t2 t3) -- -- doubleL, doubleR :: a -> Set a -> Set a -> Set a --- doubleL x1 t1 (Bin _ x2 (Bin _ x3 t2 t3) t4) = bin x3 (bin x1 t1 t2) (bin x2 t3 t4) --- doubleR x1 (Bin _ x2 t1 (Bin _ x3 t2 t3)) t4 = bin x3 (bin x2 t1 t2) (bin x1 t3 t4) +-- doubleL x1 t1 (NE (Bin' _ x2 (NE _ x3 t2 t3) t4)) = bin x3 (bin x1 t1 t2) (bin x2 t3 t4) +-- doubleR x1 (NE _ x2 t1 (NE _ x3 t2 t3)) t4 = bin x3 (bin x2 t1 t2) (bin x1 t3 t4) -- -- It is only written in such a way that every node is pattern-matched only once. -- @@ -1721,59 +2201,139 @@ ratio = 2 balanceL :: a -> Set a -> Set a -> Set a balanceL x l r = case r of Tip -> case l of - Tip -> Bin 1 x Tip Tip - (Bin _ _ Tip Tip) -> Bin 2 x l Tip - (Bin _ lx Tip (Bin _ lrx _ _)) -> Bin 3 lrx (Bin 1 lx Tip Tip) (Bin 1 x Tip Tip) - (Bin _ lx ll@(Bin _ _ _ _) Tip) -> Bin 3 lx ll (Bin 1 x Tip Tip) - (Bin ls lx ll@(Bin lls _ _ _) lr@(Bin lrs lrx lrl lrr)) - | lrs < ratio*lls -> Bin (1+ls) lx ll (Bin (1+lrs) x lr Tip) - | otherwise -> Bin (1+ls) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+size lrr) x lrr Tip) - - (Bin rs _ _ _) -> case l of - Tip -> Bin (1+rs) x Tip r - - (Bin ls lx ll lr) - | ls > delta*rs -> case (ll, lr) of - (Bin lls _ _ _, Bin lrs lrx lrl lrr) - | lrs < ratio*lls -> Bin (1+ls+rs) lx ll (Bin (1+rs+lrs) x lr r) - | otherwise -> Bin (1+ls+rs) lrx (Bin (1+lls+size lrl) lx ll lrl) (Bin (1+rs+size lrr) x lrr r) - (_, _) -> error "Failure in Data.Map.balanceL" - | otherwise -> Bin (1+ls+rs) x l r + Tip -> NE $ Bin' 1 x Tip Tip + (NE nel) -> NE $ balanceLNEE x nel + + (NE ner@(Bin' rs _ _ _)) -> case l of + Tip -> NE $ Bin' (1+rs) x Tip r + (NE nel) -> NE $ balanceLNENE x nel ner {-# NOINLINE balanceL #-} +balanceLNE :: a -> NonEmptySet a -> Set a -> NonEmptySet a +balanceLNE x nel r = case r of + Tip -> balanceLNEE x nel + (NE ner) -> balanceLNENE x nel ner +{-# NOINLINE balanceLNE #-} + +-- | Balance helper where: +-- - Left child might be too big +-- - Left child is non-empty +-- - Right child is empty +balanceLNEE :: a -> NonEmptySet a -> NonEmptySet a +balanceLNEE x nel = case nel of + (Bin' _ _ Tip Tip) -> + Bin' 2 x (NE nel) Tip + (Bin' _ lx Tip (NE (Bin' _ lrx _ _))) -> + Bin' 3 lrx (NE $ Bin' 1 lx Tip Tip) (NE $ Bin' 1 x Tip Tip) + (Bin' _ lx ll@(NE (Bin' _ _ _ _)) Tip) -> + Bin' 3 lx ll (NE $ Bin' 1 x Tip Tip) + (Bin' ls lx ll@(NE (Bin' lls _ _ _)) + lr@(NE (Bin' lrs lrx lrl lrr))) + | lrs < ratio*lls -> + Bin' (1+ls) lx ll (NE $ Bin' (1+lrs) x lr Tip) + | otherwise -> + Bin' (1+ls) lrx + (NE $ Bin' (1+lls+size lrl) lx ll lrl) + (NE $ Bin' (1+size lrr) x lrr Tip) +{-# INLINE balanceLNEE #-} + +-- | Balance helper where: +-- - Left child might be too big +-- - Left child is non-empty +-- - Right child is non-empty +balanceLNENE :: a -> NonEmptySet a -> NonEmptySet a -> NonEmptySet a +balanceLNENE x l@(Bin' ls lx ll lr) r@(Bin' rs _ _ _) + | ls > delta*rs = case (ll, lr) of + (NE (Bin' lls _ _ _), NE (Bin' lrs lrx lrl lrr)) + | lrs < ratio*lls -> Bin' (1+ls+rs) lx + ll + (NE $ Bin' (1+rs+lrs) x lr $ NE r) + | otherwise -> Bin' (1+ls+rs) lrx + (NE $ Bin' (1+lls+size lrl) lx ll lrl) + (NE $ Bin' (1+rs+size lrr) x lrr $ NE r) + (_, _) -> error "Failure in Data.Set.balanceL" + | otherwise = Bin' (1+ls+rs) x (NE l) (NE r) +{-# INLINE balanceLNENE #-} + -- balanceR is called when right subtree might have been inserted to or when -- left subtree might have been deleted from. balanceR :: a -> Set a -> Set a -> Set a balanceR x l r = case l of Tip -> case r of - Tip -> Bin 1 x Tip Tip - (Bin _ _ Tip Tip) -> Bin 2 x Tip r - (Bin _ rx Tip rr@(Bin _ _ _ _)) -> Bin 3 rx (Bin 1 x Tip Tip) rr - (Bin _ rx (Bin _ rlx _ _) Tip) -> Bin 3 rlx (Bin 1 x Tip Tip) (Bin 1 rx Tip Tip) - (Bin rs rx rl@(Bin rls rlx rll rlr) rr@(Bin rrs _ _ _)) - | rls < ratio*rrs -> Bin (1+rs) rx (Bin (1+rls) x Tip rl) rr - | otherwise -> Bin (1+rs) rlx (Bin (1+size rll) x Tip rll) (Bin (1+rrs+size rlr) rx rlr rr) - - (Bin ls _ _ _) -> case r of - Tip -> Bin (1+ls) x l Tip - - (Bin rs rx rl rr) - | rs > delta*ls -> case (rl, rr) of - (Bin rls rlx rll rlr, Bin rrs _ _ _) - | rls < ratio*rrs -> Bin (1+ls+rs) rx (Bin (1+ls+rls) x l rl) rr - | otherwise -> Bin (1+ls+rs) rlx (Bin (1+ls+size rll) x l rll) (Bin (1+rrs+size rlr) rx rlr rr) - (_, _) -> error "Failure in Data.Map.balanceR" - | otherwise -> Bin (1+ls+rs) x l r + Tip -> NE $ Bin' 1 x Tip Tip + (NE ner) -> NE $ balanceRNEE x ner + + (NE nel@(Bin' ls _ _ _)) -> case r of + Tip -> NE $ Bin' (1+ls) x l Tip + (NE ner) -> NE $ balanceRNENE x nel ner {-# NOINLINE balanceR #-} +-- | Balance helper where: +-- - Right child might be too big +-- - Left child is empty +-- - Right child is non-empty +balanceRNE :: a -> Set a -> NonEmptySet a -> NonEmptySet a +balanceRNE x l ner = case l of + Tip -> balanceRNEE x ner + (NE nel) -> balanceRNENE x nel ner +{-# NOINLINE balanceRNE #-} + +-- | Balance helper where: +-- - Right child might be too big +-- - Left child is non-empty +-- - Right child is empty +balanceRNEE :: a -> NonEmptySet a -> NonEmptySet a +balanceRNEE x ner = case ner of + (Bin' _ _ Tip Tip) -> + Bin' 2 x Tip (NE ner) + (Bin' _ rx Tip rr@(NE (Bin' _ _ _ _))) -> + Bin' 3 rx (NE (Bin' 1 x Tip Tip)) rr + (Bin' _ rx (NE (Bin' _ rlx _ _)) Tip) -> + Bin' 3 rlx + (NE (Bin' 1 x Tip Tip)) + (NE (Bin' 1 rx Tip Tip)) + (Bin' rs rx + rl@(NE (Bin' rls rlx rll rlr)) + rr@(NE (Bin' rrs _ _ _))) + | rls < ratio*rrs -> Bin' (1+rs) rx + (NE (Bin' (1+rls) x Tip rl)) + rr + | otherwise -> Bin' (1+rs) rlx + (NE (Bin' (1+size rll) x Tip rll)) + (NE (Bin' (1+rrs+size rlr) rx rlr rr)) +{-# INLINE balanceRNEE #-} + +-- | Balance helper where: +-- - Right child might be too big +-- - Left child is non-empty +-- - Right child is non-empty +balanceRNENE :: a -> NonEmptySet a -> NonEmptySet a -> NonEmptySet a +balanceRNENE x l@(Bin' ls _ _ _) r@(Bin' rs rx rl rr) + | rs > delta*ls = case (rl, rr) of + (NE (Bin' rls rlx rll rlr), NE (Bin' rrs _ _ _)) + | rls < ratio*rrs -> Bin' (1+ls+rs) rx + (NE (Bin' (1+ls+rls) x (NE l) rl)) + rr + | otherwise -> Bin' (1+ls+rs) rlx + (NE $ Bin' (1+ls+size rll) x (NE l) rll) + (NE $ Bin' (1+rrs+size rlr) rx rlr rr) + (_, _) -> error "Failure in Data.Set.balanceR" + | otherwise = Bin' (1+ls+rs) x (NE l) (NE r) +{-# INLINE balanceRNENE #-} + + {-------------------------------------------------------------------- The bin constructor maintains the size of the tree --------------------------------------------------------------------} + bin :: a -> Set a -> Set a -> Set a bin x l r - = Bin (size l + size r + 1) x l r + = NE $ Bin' (size l + size r + 1) x l r {-# INLINE bin #-} +binNE :: a -> Set a -> Set a -> NonEmptySet a +binNE x l r = Bin' (size l + size r + 1) x l r +{-# INLINE binNE #-} {-------------------------------------------------------------------- Utilities @@ -1803,9 +2363,17 @@ splitRoot :: Set a -> [Set a] splitRoot orig = case orig of Tip -> [] - Bin _ v l r -> [l, singleton v, r] + NE (Bin' _ v l r) -> [l, singleton v, r] {-# INLINE splitRoot #-} +splitRootNE :: NonEmptySet a -> NEL.NonEmpty (Set a) +splitRootNE (Bin' _ v l r) = l NEL.:| [singleton v, r] + +splitNERootNE :: NonEmptySet a -> NEL.NonEmpty (NonEmptySet a) +splitNERootNE (Bin' _ v Tip Tip) = pure $ singletonNE v +splitNERootNE (Bin' _ v (NE l) Tip) = l NEL.:| [singletonNE v] +splitNERootNE (Bin' _ v Tip (NE r)) = singletonNE v NEL.:| [r] +splitNERootNE (Bin' _ v (NE l) (NE r)) = l NEL.:| [singletonNE v, r] -- | Calculate the power set of a set: the set of all its subsets. -- @@ -1825,6 +2393,15 @@ powerSet :: Set a -> Set (Set a) powerSet xs0 = insertMin empty (foldr' step Tip xs0) where step x pxs = insertMin (singleton x) (insertMin x `mapMonotonic` pxs) `glue` pxs +powerSetNE :: NonEmptySet a -> NonEmptySet (Set a) +powerSetNE xs = insertMinNE empty . NE . mapMonotonicNE NE $ nePowerSetNE xs + +nePowerSetNE :: forall a . NonEmptySet a -> NonEmptySet (NonEmptySet a) +nePowerSetNE xs = foldr1By f (singletonNE.singletonNE) xs + where + f :: a -> NonEmptySet (NonEmptySet a) -> NonEmptySet (NonEmptySet a) + f v acc = insertMinNE (singletonNE v) (NE $ mapMonotonicNE (insertMinNE v . NE) acc) `glueNE` acc + -- | \(O(mn)\) (conjectured). Calculate the Cartesian product of two sets. -- -- @ @@ -1859,10 +2436,21 @@ cartesianProduct :: Set a -> Set b -> Set (a, b) -- When the second argument has at most one element, we can be a little -- clever. cartesianProduct !_as Tip = Tip -cartesianProduct as (Bin 1 b _ _) = mapMonotonic (flip (,) b) as +cartesianProduct as (NE (Bin' 1 b _ _)) = mapMonotonic (flip (,) b) as cartesianProduct as bs = getMergeSet $ foldMap (\a -> MergeSet $ mapMonotonic ((,) a) bs) as +cartesianProductNE :: forall a b . NonEmptySet a -> NonEmptySet b -> NonEmptySet (a, b) +cartesianProductNE as (Bin' 1 b _ _) = mapMonotonicNE (flip (,) b) as +cartesianProductNE as bs = goFoldMapNE as + where + f a = mapMonotonicNE ((,) a) bs + goFoldMapNE :: NonEmptySet a -> NonEmptySet (a, b) + goFoldMapNE (Bin' 1 k _ _) = f k + goFoldMapNE (Bin' _ k l r) = goFoldMap l `mergeXNE` (f k `mergeNEX` goFoldMap r) + goFoldMap Tip = empty + goFoldMap (NE s) = NE $ goFoldMapNE s + -- A version of Set with peculiar Semigroup and Monoid instances. -- The result of xs <> ys will only be a valid set if the greatest -- element of xs is strictly less than the least element of ys. @@ -1892,6 +2480,14 @@ instance Monoid (MergeSet a) where disjointUnion :: Set a -> Set b -> Set (Either a b) disjointUnion as bs = merge (mapMonotonic Left as) (mapMonotonic Right bs) +disjointUnionNE :: NonEmptySet a -> NonEmptySet b -> NonEmptySet (Either a b) +disjointUnionNE as bs = mergeNE (mapMonotonicNE Left as) (mapMonotonicNE Right bs) + +disjointUnionNEX :: NonEmptySet a -> Set b -> NonEmptySet (Either a b) +disjointUnionNEX as bs = mergeNEX (mapMonotonicNE Left as) (mapMonotonic Right bs) + +disjointUnionXNE :: Set a -> NonEmptySet b -> NonEmptySet (Either a b) +disjointUnionXNE as bs = mergeXNE (mapMonotonic Left as) (mapMonotonicNE Right bs) {-------------------------------------------------------------------- Debugging --------------------------------------------------------------------} @@ -1901,6 +2497,10 @@ showTree :: Show a => Set a -> String showTree s = showTreeWith True False s +showTreeNE :: Show a => NonEmptySet a -> String +showTreeNE s + = showTreeWithNE True False s + {- | \(O(n)\). The expression (@showTreeWith hang wide map@) shows the tree that implements the set. If @hang@ is @@ -1942,13 +2542,23 @@ showTreeWith hang wide t | hang = (showsTreeHang wide [] t) "" | otherwise = (showsTree wide [] [] t) "" +showTreeWithNE :: Show a => Bool -> Bool -> NonEmptySet a -> String +showTreeWithNE hang wide t + | hang = (showsTreeHangNE wide [] t) "" + | otherwise = (showsTreeNE wide [] [] t) "" + showsTree :: Show a => Bool -> [String] -> [String] -> Set a -> ShowS showsTree wide lbars rbars t = case t of Tip -> showsBars lbars . showString "|\n" - Bin _ x Tip Tip + NE ne -> showsTreeNE wide lbars rbars ne + +showsTreeNE :: Show a => Bool -> [String] -> [String] -> NonEmptySet a -> ShowS +showsTreeNE wide lbars rbars t + = case t of + Bin' _ x Tip Tip -> showsBars lbars . shows x . showString "\n" - Bin _ x l r + Bin' _ x l r -> showsTree wide (withBar rbars) (withEmpty rbars) r . showWide wide rbars . showsBars lbars . shows x . showString "\n" . @@ -1959,9 +2569,14 @@ showsTreeHang :: Show a => Bool -> [String] -> Set a -> ShowS showsTreeHang wide bars t = case t of Tip -> showsBars bars . showString "|\n" - Bin _ x Tip Tip + NE ne -> showsTreeHangNE wide bars ne + +showsTreeHangNE :: Show a => Bool -> [String] -> NonEmptySet a -> ShowS +showsTreeHangNE wide bars t + = case t of + Bin' _ x Tip Tip -> showsBars bars . shows x . showString "\n" - Bin _ x l r + Bin' _ x l r -> showsBars bars . shows x . showString "\n" . showWide wide bars . showsTreeHang wide (withBar bars) l . @@ -1994,29 +2609,53 @@ valid :: Ord a => Set a -> Bool valid t = balanced t && ordered t && validsize t +validNE :: Ord a => NonEmptySet a -> Bool +validNE t + = balancedNE t && orderedNE t && validsizeNE t + +-------------------------------------------------------------------- + ordered :: Ord a => Set a -> Bool -ordered t - = bounded (const True) (const True) t - where - bounded lo hi t' - = case t' of - Tip -> True - Bin _ x l r -> (lo x) && (hi x) && bounded lo (x) hi r +ordered = bounded (const True) (const True) + +orderedNE :: Ord a => NonEmptySet a -> Bool +orderedNE = boundedNE (const True) (const True) + +bounded :: Ord a => (a -> Bool) -> (a -> Bool) -> Set a -> Bool +bounded lo hi t' = case t' of + Tip -> True + NE ne -> boundedNE lo hi ne + +boundedNE :: Ord a => (a -> Bool) -> (a -> Bool) -> NonEmptySet a -> Bool +boundedNE lo hi (Bin' _ x l r) = + (lo x) && (hi x) && bounded lo (x) hi r + +-------------------------------------------------------------------- balanced :: Set a -> 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 t = case t of + Tip -> True + NE ne -> balancedNE ne + +balancedNE :: NonEmptySet a -> Bool +balancedNE (Bin' _ _ l r) = + (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) && + balanced l && balanced r + +-------------------------------------------------------------------- validsize :: Set a -> 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 +validsize t = realsize t == Just (size t) + +validsizeNE :: NonEmptySet a -> Bool +validsizeNE t = realsizeNE t == Just (sizeNE t) + +realsize :: Set a -> Maybe Size +realsize t' = case t' of + Tip -> Just 0 + NE ne -> realsizeNE ne + +realsizeNE :: NonEmptySet a -> Maybe Size +realsizeNE (Bin' sz _ l r) = case (realsize l,realsize r) of + (Just n, Just m) | n+m+1 == sz -> Just sz + _ -> Nothing