Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve fromAscList and friends. #658

Merged
merged 3 commits into from
Jul 15, 2019
Merged

Improve fromAscList and friends. #658

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e2e018b
improve Data.IntSet.from*AscList (#654)
int-e Jul 6, 2019
207e24d
improve Data.IntMap.from*AscList* (#654)
int-e Jul 6, 2019
48c1ca0
improve Data.IntMap.Strict.from*AscList* (#654)
int-e Jul 6, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 77 additions & 42 deletions containers/src/Data/IntMap/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,7 @@ module Data.IntMap.Internal (
, natFromInt
, intFromNat
, link
, linkWithMask
, bin
, binCheckLeft
, binCheckRight
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -3132,50 +3133,80 @@ 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.
-- /The precondition (input list is strictly ascending) is not checked./
--
-- > 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
int-e marked this conversation as resolved.
Show resolved Hide resolved
, 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
Expand Down Expand Up @@ -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
int-e marked this conversation as resolved.
Show resolved Hide resolved
m = branchMask p1 p2
p = mask p1 m
{-# INLINE link #-}
{-# INLINE linkWithMask #-}

{--------------------------------------------------------------------
@bin@ assures that we never have empty trees within a tree.
Expand Down
102 changes: 71 additions & 31 deletions containers/src/Data/IntMap/Strict/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,7 @@ import Data.IntMap.Internal
, binCheckLeft
, binCheckRight
, link
, linkWithMask

, (\\)
, (!)
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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
Loading