|
-- ghc -O2 -hide-package containers IntMapFAL.hs && ./IntMapFAL --small | tee IntMapFAL.out |
|
|
|
{-# LANGUAGE BangPatterns #-} |
|
{-# LANGUAGE RankNTypes #-} |
|
|
|
import Control.DeepSeq (rnf) |
|
import Control.Exception (evaluate) |
|
import Gauge (bench, bgroup, env, defaultMain, whnf) |
|
import Data.List (foldl') |
|
import qualified Data.IntMap as M |
|
import qualified Data.IntMap.Strict as MS |
|
import Data.Maybe (fromMaybe) |
|
import Prelude hiding (lookup) |
|
|
|
import GHC.Exts (inline) |
|
|
|
import Data.IntMap.Internal (IntMap (..), Prefix, Mask, Key, size, link, branchMask, mask, shorter, nomatch, zero) |
|
import qualified Data.IntMap.Internal as I |
|
|
|
main = do |
|
defaultMain $ [test (s*sz) sk | |
|
sz <- [10^i | i <- [0..5]], |
|
(s, sk) <- if sz == 1 then [(1,1)] else [(1,1), (-1,51791)]] |
|
|
|
test sz sk = |
|
env (let m = M.fromAscList elems :: M.IntMap Int in evaluate $ rnf [m]) $ \m -> bgroup n |
|
[ bench "fromList" $ whnf M.fromList elems |
|
, bench "fromList1" $ whnf fromList1 elems |
|
, bench "fromAscList" $ whnf M.fromAscList elems |
|
, bench "fromAscList1" $ whnf fromAscList1 elems |
|
, bench "fromAscList1a" $ whnf fromAscList1a elems |
|
, bench "fromAscList1b" $ whnf fromAscList1b elems |
|
, bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems |
|
, bench "fromDistinctAscList1" $ whnf fromDistinctAscList1 elems |
|
, bench "fromDistinctAscList1a" $ whnf fromDistinctAscList1a elems |
|
, bench "fromDistinctAscList1b" $ whnf fromDistinctAscList1b elems |
|
, bench "fromDistinctAscList1c" $ whnf fromDistinctAscList1c elems |
|
, bench "fromDistinctAscList1d" $ whnf fromDistinctAscList1d elems |
|
, bench "fromDistinctAscList1e" $ whnf fromDistinctAscList1e elems |
|
] |
|
where |
|
n = "[" ++ show sz ++ "," ++ show sk ++ "]" |
|
elems = zip keys values |
|
keys = map (sk*) (if sz < 0 then [2*sz `div` 3.. -sz `div` 3] else [0..sz]) |
|
values = [1..] |
|
|
|
data Inserted a = Inserted !(IntMap a) ![(Key,a)] |
|
|
|
fromDistinctAscList1 :: [(Key,a)] -> IntMap a |
|
fromDistinctAscList1 [] = Nil |
|
fromDistinctAscList1 ((kx,vx) : zs0) = addAll kx (Tip kx vx) zs0 |
|
where |
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `tx`. |
|
addAll !kx !tx [] |
|
= tx |
|
addAll !kx !tx ((ky,vy) : zs) |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- addMany m ky (Tip ky vy) zs |
|
= addAll kx (link ky ty kx tx) zs' |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany (branchMask kx ky) ky (Tip ky vy) zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromDistinctAscList1a :: [(Key,a)] -> IntMap a |
|
fromDistinctAscList1a [] = Nil |
|
fromDistinctAscList1a ((kx,vx) : zs0) = addAll kx (Tip kx vx) zs0 |
|
where |
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `tx`. |
|
addAll !kx tx [] |
|
= tx |
|
addAll !kx tx ((ky,vy) : zs) |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- addMany m ky (Tip ky vy) zs |
|
= addAll kx (link ky ty kx tx) zs' |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| m' <- branchMask kx ky |
|
, Inserted ty zs' <- addMany m' ky (Tip ky vy) zs |
|
= addMany m kx (Bin (mask kx m') m' tx ty) zs' |
|
|
|
fromDistinctAscList1b :: [(Key,a)] -> IntMap a |
|
fromDistinctAscList1b [] = Nil |
|
fromDistinctAscList1b ((kx,vx) : zs1) = addAll' kx vx zs1 |
|
where |
|
addAll' !kx vx [] = inline addAll kx (Tip kx vx) [] |
|
addAll' !kx vx ((ky,vy) : zs) = inline addAll kx (Tip kx vx) ((ky,vy) : zs) |
|
|
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `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 (link ky ty kx tx) zs' |
|
|
|
addMany' !m !kx vx [] = inline addMany m kx (Tip kx vx) [] |
|
addMany' !m !kx vx ((ky,vy) : zs) = inline addMany m kx (Tip kx vx) ((ky,vy) : zs) |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromDistinctAscList1c :: [(Key,a)] -> IntMap a |
|
fromDistinctAscList1c [] = Nil |
|
fromDistinctAscList1c ((kx,vx) : zs1) = inline addAll kx (Tip kx vx) zs1 |
|
where |
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `tx`. |
|
addAll !kx !tx [] |
|
= tx |
|
addAll !kx !tx ((ky,vy) : zs) |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- addMany m ky (Tip ky vy) zs |
|
= addAll kx (link ky ty kx tx) zs' |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany (branchMask kx ky) ky (Tip ky vy) zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromDistinctAscList1d :: [(Key,a)] -> IntMap a |
|
fromDistinctAscList1d [] = Nil |
|
fromDistinctAscList1d ((kx,vx) : zs1) = inline addAll kx (Tip kx vx) zs1 |
|
where |
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `tx`. |
|
addAll !kx !tx [] |
|
= tx |
|
addAll !kx !tx ((ky,vy) : zs) |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- inline addMany m ky (Tip ky vy) zs |
|
= addAll kx (link ky ty kx tx) zs' |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany (branchMask kx ky) ky (Tip ky vy) zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromDistinctAscList1e :: [(Key,a)] -> IntMap a |
|
fromDistinctAscList1e [] = Nil |
|
fromDistinctAscList1e ((kx,vx) : zs1) = addAll' kx vx zs1 |
|
where |
|
addAll' !kx vx [] = Tip kx vx |
|
addAll' !kx vx ((ky,vy) : zs) |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- addMany' m ky vy zs |
|
= addAll kx (link ky ty kx (Tip kx vx)) zs' |
|
|
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `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 (link ky ty kx tx) zs' |
|
|
|
addMany' !m !kx vx [] = Inserted (Tip kx vx) [] |
|
addMany' !m !kx vx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted (Tip kx vx) zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx (Tip kx vx)) zs' |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromAscList1 :: [(Key,a)] -> IntMap a |
|
fromAscList1 [] = Nil |
|
fromAscList1 ((kx,vx) : zs0) = addAll' kx vx zs0 |
|
where |
|
addAll' !kx vx [] = Tip kx vx |
|
addAll' !kx vx ((ky,vy) : zs) |
|
| kx == ky = addAll' ky vy zs |
|
| otherwise = addAll kx (Tip kx vx) ((ky,vy) : zs) |
|
|
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `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 (link ky ty kx tx) zs' |
|
|
|
addMany' !m !kx vx [] = Inserted (Tip kx vx) [] |
|
addMany' !m !kx vx ((ky,vy) : zs) |
|
| kx == ky = addMany' m ky vy zs |
|
| otherwise = addMany m kx (Tip kx vx) ((ky,vy) : zs) |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromAscList1a :: [(Key,a)] -> IntMap a |
|
fromAscList1a [] = Nil |
|
fromAscList1a ((kx,vx) : zs1) = addAll' kx vx zs1 |
|
where |
|
-- `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) |
|
| kx == ky |
|
= addAll' ky vy zs |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- addMany' m ky vy zs |
|
= addAll kx (link 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 (link 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) |
|
| kx == ky |
|
= addMany' m ky vy zs |
|
| mask kx m /= mask ky m |
|
= Inserted (Tip kx vx) zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link 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 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromAscList1b :: [(Key,a)] -> IntMap a |
|
fromAscList1b [] = Nil |
|
fromAscList1b ((kx,vx) : zs0) = addAll' kx vx zs0 |
|
where |
|
addAll' !kx vx [] = Tip kx vx |
|
addAll' !kx vx ((ky,vy) : zs) |
|
| kx == ky = addAll' ky vy zs |
|
| otherwise = inline addAll kx (Tip kx vx) ((ky,vy) : zs) |
|
|
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `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 (link ky ty kx tx) zs' |
|
|
|
addMany' !m !kx vx [] = Inserted (Tip kx vx) [] |
|
addMany' !m !kx vx ((ky,vy) : zs) |
|
| kx == ky = addMany' m ky vy zs |
|
| otherwise = inline addMany m kx (Tip kx vx) ((ky,vy) : zs) |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
fromAscList1c :: [(Key,a)] -> IntMap a |
|
fromAscList1c [] = Nil |
|
fromAscList1c ((kx,vx) : zs1) = addAll' kx vx zs1 |
|
where |
|
addAll' !kx vx [] = Tip kx vx |
|
addAll' !kx vx ((ky,vy) : zs) |
|
| kx == ky |
|
= addAll' ky vy zs |
|
| m <- branchMask kx ky |
|
, Inserted ty zs' <- addMany' m ky vy zs |
|
= addAll kx (link ky ty kx (Tip kx vx)) zs' |
|
|
|
-- invariant for `addAll` and `addMany`: `kx` is a key in the map `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 (link ky ty kx tx) zs' |
|
|
|
addMany' !m !kx vx [] = Inserted (Tip kx vx) [] |
|
addMany' !m !kx vx zs0@((ky,vy) : zs) |
|
| kx == ky |
|
= addMany' m ky vy zs |
|
| mask kx m /= mask ky m |
|
= Inserted (Tip kx vx) zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx (Tip kx vx)) zs' |
|
|
|
-- addMany adds all elements that have the same prefix as `kx` w.r.t. |
|
-- the branch mask `m` to `tx`. |
|
addMany !m !kx tx [] |
|
= Inserted tx [] |
|
addMany !m !kx tx zs0@((ky,vy) : zs) |
|
| mask kx m /= mask ky m |
|
= Inserted tx zs0 |
|
| Inserted ty zs' <- addMany' (branchMask kx ky) ky vy zs |
|
= addMany m kx (link ky ty kx tx) zs' |
|
|
|
------------------------------------------------------------------------------ |
|
-- fromList implementation from #653 |
|
------------------------------------------------------------------------------ |
|
|
|
fromList1 :: [(Key,a)] -> IntMap a |
|
fromList1 = insertAll Nil |
|
{-# NOINLINE fromList1 #-} |
|
|
|
-- [Note: fromList] |
|
-- |
|
-- The obvious way to build a map from a list is just to fold over the list |
|
-- inserting each entry into the accumulator map. The problem is that this |
|
-- rebuilds the path from the root *every single time*. To avoid this, we |
|
-- insert as many elements as we can into the current subtree, backing out |
|
-- one level at a time when necessary. |
|
|
|
insertAll :: IntMap a -> [(Key, a)] -> IntMap a |
|
insertAll m [] = m |
|
insertAll m ((k,x) : kxs) |
|
| Inserted m' r <- insertSome m k x kxs |
|
= insertAll m' r |
|
|
|
-- | Insert at least one entry into an 'IntMap' or subtree. If |
|
-- others fit in the same resulting subtree, insert them too. |
|
-- Return the new map and remaining values. |
|
insertSome :: IntMap a -> Key -> a -> [(Key, a)] -> Inserted a |
|
insertSome t@(Bin p m l r) !k x kxs |
|
| nomatch k p m |
|
= insertMany (link k (Tip k x) p t) kxs |
|
|
|
| zero k m |
|
, Inserted l' kxs' <- insertSome l k x kxs |
|
= insertMany (Bin p m l' r) kxs' |
|
|
|
| Inserted r' kxs' <- insertSome r k x kxs |
|
= insertMany (Bin p m l r') kxs' |
|
|
|
insertSome t@(Tip ky _) k x kxs |
|
| k == ky |
|
= insertMany (Tip k x) kxs |
|
| otherwise |
|
= insertMany (link k (Tip k x) ky t) kxs |
|
|
|
insertSome Nil k x kxs = insertMany (Tip k x) kxs |
|
|
|
-- | Try to insert some entries into a subtree of an 'IntMap'. If |
|
-- they belong in some other subtree, just don't insert them. |
|
insertMany :: IntMap a -> [(Key, a)] -> Inserted a |
|
insertMany t [] = Inserted t [] |
|
insertMany t@(Bin p m _ _) kxs@((k, x) : kxs') |
|
| nomatch k p m |
|
= Inserted t kxs |
|
| otherwise |
|
= insertSome t k x kxs' |
|
insertMany t@(Tip ky _) kxs@((k, x) : kxs') |
|
| k==ky = insertSome t k x kxs' |
|
| otherwise = Inserted t kxs |
|
insertMany Nil kxs = Inserted Nil kxs -- Unused case |