diff --git a/containers/src/Data/IntMap/Internal.hs b/containers/src/Data/IntMap/Internal.hs index 3b8fbb446..9cf26968a 100644 --- a/containers/src/Data/IntMap/Internal.hs +++ b/containers/src/Data/IntMap/Internal.hs @@ -266,6 +266,7 @@ module Data.IntMap.Internal ( , natFromInt , intFromNat , link + , linkWithMask , bin , binCheckLeft , binCheckRight @@ -3111,8 +3112,8 @@ fromListWithKey f xs -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] fromAscList :: [(Key,a)] -> IntMap a -fromAscList xs - = fromAscListWithKey (\_ x _ -> x) xs +fromAscList = fromMonoListWithKey Nondistinct (\_ x _ -> x) +{-# NOINLINE fromAscList #-} -- | /O(n)/. Build a map from a list of key\/value pairs where -- the keys are in ascending order, with a combining function on equal keys. @@ -3121,8 +3122,8 @@ fromAscList xs -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f xs - = fromAscListWithKey (\_ x y -> f x y) xs +fromAscListWith f = fromMonoListWithKey Nondistinct (\_ x y -> f x y) +{-# NOINLINE fromAscListWith #-} -- | /O(n)/. Build a map from a list of key\/value pairs where -- the keys are in ascending order, with a combining function on equal keys. @@ -3132,14 +3133,8 @@ fromAscListWith f xs -- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")] fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey _ [] = Nil -fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) - where - -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] - combineEq z [] = [z] - combineEq z@(kz,zz) (x@(kx,xx):xs) - | kx==kz = let yy = f kx xx zz in combineEq (kx,yy) xs - | otherwise = z:combineEq x xs +fromAscListWithKey f = fromMonoListWithKey Nondistinct f +{-# NOINLINE fromAscListWithKey #-} -- | /O(n)/. Build a map from a list of key\/value pairs where -- the keys are in ascending order and all distinct. @@ -3147,35 +3142,71 @@ fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) -- -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] -#if __GLASGOW_HASKELL__ -fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a -#else -fromDistinctAscList :: [(Key,a)] -> IntMap a -#endif -fromDistinctAscList [] = Nil -fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada - where - work (kx,vx) [] stk = finish kx (Tip kx vx) stk - work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk - -#if __GLASGOW_HASKELL__ - reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a -#endif - reduce z zs _ px tx Nada = work z zs (Push px tx Nada) - reduce z zs m px tx stk@(Push py ty stk') = - let mxy = branchMask px py - pxy = mask px mxy - in if shorter m mxy - then reduce z zs m pxy (Bin pxy mxy ty tx) stk' - else work z zs (Push px tx stk) - - finish _ t Nada = t - finish px tx (Push py ty stk) = finish p (link py ty px tx) stk - where m = branchMask px py - p = mask px m +fromDistinctAscList :: [(Key,a)] -> IntMap a +fromDistinctAscList = fromMonoListWithKey Distinct (\_ x _ -> x) +{-# NOINLINE fromDistinctAscList #-} -data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada +-- | /O(n)/. Build a map from a list of key\/value pairs with monotonic keys +-- and a combining function. +-- +-- The precise conditions under which this function works are subtle: +-- For any branch mask, keys with the same prefix w.r.t. the branch +-- mask must occur consecutively in the list. +fromMonoListWithKey :: Distinct -> (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromMonoListWithKey distinct f = go + where + go [] = Nil + go ((kx,vx) : zs1) = addAll' kx vx zs1 + + -- `addAll'` collects all keys equal to `kx` into a single value, + -- and then proceeds with `addAll`. + addAll' !kx vx [] + = Tip kx vx + addAll' !kx vx ((ky,vy) : zs) + | Nondistinct <- distinct, kx == ky + = let v = f kx vy vx in addAll' ky v zs + -- inlined: | otherwise = addAll kx (Tip kx vx) (ky : zs) + | m <- branchMask kx ky + , Inserted ty zs' <- addMany' m ky vy zs + = addAll kx (linkWithMask m ky ty {-kx-} (Tip kx vx)) zs' + + -- for `addAll` and `addMany`, kx is /a/ key inside the tree `tx` + -- `addAll` consumes the rest of the list, adding to the tree `tx` + addAll !kx !tx [] + = tx + addAll !kx !tx ((ky,vy) : zs) + | m <- branchMask kx ky + , Inserted ty zs' <- addMany' m ky vy zs + = addAll kx (linkWithMask m ky ty {-kx-} tx) zs' + + -- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`. + addMany' !m !kx vx [] + = Inserted (Tip kx vx) [] + addMany' !m !kx vx zs0@((ky,vy) : zs) + | Nondistinct <- distinct, kx == ky + = let v = f kx vy vx in addMany' m ky v zs + -- inlined: | otherwise = addMany m kx (Tip kx vx) (ky : zs) + | mask kx m /= mask ky m + = Inserted (Tip kx vx) zs0 + | mxy <- branchMask kx ky + , Inserted ty zs' <- addMany' mxy ky vy zs + = addMany m kx (linkWithMask mxy ky ty {-kx-} (Tip kx vx)) zs' + + -- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `kx`. + addMany !m !kx tx [] + = Inserted tx [] + addMany !m !kx tx zs0@((ky,vy) : zs) + | mask kx m /= mask ky m + = Inserted tx zs0 + | mxy <- branchMask kx ky + , Inserted ty zs' <- addMany' mxy ky vy zs + = addMany m kx (linkWithMask mxy ky ty {-kx-} tx) zs' +{-# INLINE fromMonoListWithKey #-} + +data Inserted a = Inserted !(IntMap a) ![(Key,a)] + +data Distinct = Distinct | Nondistinct {-------------------------------------------------------------------- Eq @@ -3297,13 +3328,17 @@ INSTANCE_TYPEABLE1(IntMap) Link --------------------------------------------------------------------} link :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a -link p1 t1 p2 t2 +link p1 t1 p2 t2 = linkWithMask (branchMask p1 p2) p1 t1 {-p2-} t2 +{-# INLINE link #-} + +-- `linkWithMask` is useful when the `branchMask` has already been computed +linkWithMask :: Mask -> Prefix -> IntMap a -> IntMap a -> IntMap a +linkWithMask m p1 t1 {-p2-} t2 | zero p1 m = Bin p m t1 t2 | otherwise = Bin p m t2 t1 where - m = branchMask p1 p2 p = mask p1 m -{-# INLINE link #-} +{-# INLINE linkWithMask #-} {-------------------------------------------------------------------- @bin@ assures that we never have empty trees within a tree. diff --git a/containers/src/Data/IntMap/Strict/Internal.hs b/containers/src/Data/IntMap/Strict/Internal.hs index f4097c438..e7345418f 100644 --- a/containers/src/Data/IntMap/Strict/Internal.hs +++ b/containers/src/Data/IntMap/Strict/Internal.hs @@ -270,6 +270,7 @@ import Data.IntMap.Internal , binCheckLeft , binCheckRight , link + , linkWithMask , (\\) , (!) @@ -1098,8 +1099,8 @@ fromListWithKey f xs -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] fromAscList :: [(Key,a)] -> IntMap a -fromAscList xs - = fromAscListWithKey (\_ x _ -> x) xs +fromAscList = fromMonoListWithKey Nondistinct (\_ x _ -> x) +{-# NOINLINE fromAscList #-} -- | /O(n)/. Build a map from a list of key\/value pairs where -- the keys are in ascending order, with a combining function on equal keys. @@ -1108,8 +1109,8 @@ fromAscList xs -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f xs - = fromAscListWithKey (\_ x y -> f x y) xs +fromAscListWith f = fromMonoListWithKey Nondistinct (\_ x y -> f x y) +{-# NOINLINE fromAscListWith #-} -- | /O(n)/. Build a map from a list of key\/value pairs where -- the keys are in ascending order, with a combining function on equal keys. @@ -1118,14 +1119,8 @@ fromAscListWith f xs -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey _ [] = Nil -fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) - where - -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] - combineEq z [] = [z] - combineEq z@(kz,zz) (x@(kx,xx):xs) - | kx==kz = let !yy = f kx xx zz in combineEq (kx,yy) xs - | otherwise = z:combineEq x xs +fromAscListWithKey f = fromMonoListWithKey Nondistinct f +{-# NOINLINE fromAscListWithKey #-} -- | /O(n)/. Build a map from a list of key\/value pairs where -- the keys are in ascending order and all distinct. @@ -1134,24 +1129,69 @@ fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] fromDistinctAscList :: [(Key,a)] -> IntMap a -fromDistinctAscList [] = Nil -fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada - where - work (kx,!vx) [] stk = finish kx (Tip kx vx) stk - work (kx,!vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk - - reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a - reduce z zs _ px tx Nada = work z zs (Push px tx Nada) - reduce z zs m px tx stk@(Push py ty stk') = - let mxy = branchMask px py - pxy = mask px mxy - in if shorter m mxy - then reduce z zs m pxy (Bin pxy mxy ty tx) stk' - else work z zs (Push px tx stk) - - finish _ t Nada = t - finish px tx (Push py ty stk) = finish p (link py ty px tx) stk - where m = branchMask px py - p = mask px m +fromDistinctAscList = fromMonoListWithKey Distinct (\_ x _ -> x) +{-# NOINLINE fromDistinctAscList #-} data Stack a = Push {-# UNPACK #-} !Prefix !(IntMap a) !(Stack a) | Nada + +-- | /O(n)/. Build a map from a list of key\/value pairs with monotonic keys +-- and a combining function. +-- +-- The precise conditions under which this function works are subtle: +-- For any branch mask, keys with the same prefix w.r.t. the branch +-- mask must occur consecutively in the list. + +fromMonoListWithKey :: Distinct -> (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromMonoListWithKey distinct f = go + where + go [] = Nil + go ((kx,vx) : zs1) = addAll' kx vx zs1 + + -- `addAll'` collects all keys equal to `kx` into a single value, + -- and then proceeds with `addAll`. + addAll' !kx vx [] + = Tip kx $! vx + addAll' !kx vx ((ky,vy) : zs) + | Nondistinct <- distinct, kx == ky + = let !v = f kx vy vx in addAll' ky v zs + -- inlined: | otherwise = addAll kx (Tip kx $! vx) (ky : zs) + | m <- branchMask kx ky + , Inserted ty zs' <- addMany' m ky vy zs + = addAll kx (linkWithMask m ky ty {-kx-} (Tip kx $! vx)) zs' + + -- for `addAll` and `addMany`, kx is /a/ key inside the tree `tx` + -- `addAll` consumes the rest of the list, adding to the tree `tx` + addAll !kx !tx [] + = tx + addAll !kx !tx ((ky,vy) : zs) + | m <- branchMask kx ky + , Inserted ty zs' <- addMany' m ky vy zs + = addAll kx (linkWithMask m ky ty {-kx-} tx) zs' + + -- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`. + addMany' !m !kx vx [] + = Inserted (Tip kx $! vx) [] + addMany' !m !kx vx zs0@((ky,vy) : zs) + | Nondistinct <- distinct, kx == ky + = let !v = f kx vy vx in addMany' m ky v zs + -- inlined: | otherwise = addMany m kx (Tip kx $! vx) (ky : zs) + | mask kx m /= mask ky m + = Inserted (Tip kx $! vx) zs0 + | mxy <- branchMask kx ky + , Inserted ty zs' <- addMany' mxy ky vy zs + = addMany m kx (linkWithMask mxy ky ty {-kx-} (Tip kx $! vx)) zs' + + -- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `kx`. + addMany !m !kx tx [] + = Inserted tx [] + addMany !m !kx tx zs0@((ky,vy) : zs) + | mask kx m /= mask ky m + = Inserted tx zs0 + | mxy <- branchMask kx ky + , Inserted ty zs' <- addMany' mxy ky vy zs + = addMany m kx (linkWithMask mxy ky ty {-kx-} tx) zs' +{-# INLINE fromMonoListWithKey #-} + +data Inserted a = Inserted !(IntMap a) ![(Key,a)] + +data Distinct = Distinct | Nondistinct diff --git a/containers/src/Data/IntSet/Internal.hs b/containers/src/Data/IntSet/Internal.hs index 7782ed48f..df7c04afd 100644 --- a/containers/src/Data/IntSet/Internal.hs +++ b/containers/src/Data/IntSet/Internal.hs @@ -1060,41 +1060,74 @@ fromList xs -- | /O(n)/. Build a set from an ascending list of elements. -- /The precondition (input list is ascending) is not checked./ fromAscList :: [Key] -> IntSet -fromAscList [] = Nil -fromAscList (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0) - where - combineEq x' [] = [x'] - combineEq x' (x:xs) - | x==x' = combineEq x' xs - | otherwise = x' : combineEq x xs +fromAscList = fromMonoList +{-# NOINLINE fromAscList #-} -- | /O(n)/. Build a set from an ascending list of distinct elements. -- /The precondition (input list is strictly ascending) is not checked./ fromDistinctAscList :: [Key] -> IntSet -fromDistinctAscList [] = Nil -fromDistinctAscList (z0 : zs0) = work (prefixOf z0) (bitmapOf z0) zs0 Nada - where - -- 'work' accumulates all values that go into one tip, before passing this Tip - -- to 'reduce' - work kx bm [] stk = finish kx (Tip kx bm) stk - work kx bm (z:zs) stk | kx == prefixOf z = work kx (bm .|. bitmapOf z) zs stk - work kx bm (z:zs) stk = reduce z zs (branchMask z kx) kx (Tip kx bm) stk - - reduce z zs _ px tx Nada = work (prefixOf z) (bitmapOf z) zs (Push px tx Nada) - reduce z zs m px tx stk@(Push py ty stk') = - let mxy = branchMask px py - pxy = mask px mxy - in if shorter m mxy - then reduce z zs m pxy (Bin pxy mxy ty tx) stk' - else work (prefixOf z) (bitmapOf z) zs (Push px tx stk) - - finish _ t Nada = t - finish px tx (Push py ty stk) = finish p (link py ty px tx) stk - where m = branchMask px py - p = mask px m - -data Stack = Push {-# UNPACK #-} !Prefix !IntSet !Stack | Nada +fromDistinctAscList = fromAscList +{-# INLINE fromDistinctAscList #-} +-- | /O(n)/. Build a set from a monotonic list of elements. +-- +-- The precise conditions under which this function works are subtle: +-- For any branch mask, keys with the same prefix w.r.t. the branch +-- mask must occur consecutively in the list. +fromMonoList :: [Key] -> IntSet +fromMonoList [] = Nil +fromMonoList (kx : zs1) = addAll' (prefixOf kx) (bitmapOf kx) zs1 + where + -- `addAll'` collects all keys with the prefix `px` into a single + -- bitmap, and then proceeds with `addAll`. + addAll' !px !bm [] + = Tip px bm + addAll' !px !bm (ky : zs) + | px == prefixOf ky + = addAll' px (bm .|. bitmapOf ky) zs + -- inlined: | otherwise = addAll px (Tip px bm) (ky : zs) + | py <- prefixOf ky + , m <- branchMask px py + , Inserted ty zs' <- addMany' m py (bitmapOf ky) zs + = addAll px (linkWithMask m py ty {-px-} (Tip px bm)) zs' + + -- for `addAll` and `addMany`, px is /a/ prefix inside the tree `tx` + -- `addAll` consumes the rest of the list, adding to the tree `tx` + addAll !px !tx [] + = tx + addAll !px !tx (ky : zs) + | py <- prefixOf ky + , m <- branchMask px py + , Inserted ty zs' <- addMany' m py (bitmapOf ky) zs + = addAll px (linkWithMask m py ty {-px-} tx) zs' + + -- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`. + addMany' !m !px !bm [] + = Inserted (Tip px bm) [] + addMany' !m !px !bm zs0@(ky : zs) + | px == prefixOf ky + = addMany' m px (bm .|. bitmapOf ky) zs + -- inlined: | otherwise = addMany m px (Tip px bm) (ky : zs) + | mask px m /= mask ky m + = Inserted (Tip (prefixOf px) bm) zs0 + | py <- prefixOf ky + , mxy <- branchMask px py + , Inserted ty zs' <- addMany' mxy py (bitmapOf ky) zs + = addMany m px (linkWithMask mxy py ty {-px-} (Tip px bm)) zs' + + -- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `px`. + addMany !m !px tx [] + = Inserted tx [] + addMany !m !px tx zs0@(ky : zs) + | mask px m /= mask ky m + = Inserted tx zs0 + | py <- prefixOf ky + , mxy <- branchMask px py + , Inserted ty zs' <- addMany' mxy py (bitmapOf ky) zs + = addMany m px (linkWithMask mxy py ty {-px-} tx) zs' +{-# INLINE fromMonoList #-} + +data Inserted = Inserted !IntSet ![Key] {-------------------------------------------------------------------- Eq @@ -1249,13 +1282,17 @@ withEmpty bars = " ":bars Link --------------------------------------------------------------------} link :: Prefix -> IntSet -> Prefix -> IntSet -> IntSet -link p1 t1 p2 t2 +link p1 t1 p2 t2 = linkWithMask (branchMask p1 p2) p1 t1 {-p2-} t2 +{-# INLINE link #-} + +-- `linkWithMask` is useful when the `branchMask` has already been computed +linkWithMask :: Mask -> Prefix -> IntSet -> IntSet -> IntSet +linkWithMask m p1 t1 {-p2-} t2 | zero p1 m = Bin p m t1 t2 | otherwise = Bin p m t2 t1 where - m = branchMask p1 p2 p = mask p1 m -{-# INLINE link #-} +{-# INLINE linkWithMask #-} {-------------------------------------------------------------------- @bin@ assures that we never have empty trees within a tree.