[commit: packages/containers] master, revert-408-bugfix_394: defined lookupPrefix as part of optimizing restrictKeys (af7bb60)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/containers

On branches: master,revert-408-bugfix_394
Link       : http://git.haskell.org/packages/containers.git/commitdiff/af7bb60b941468ec6785eeab3dc61a837ceda4e9

>---------------------------------------------------------------

commit af7bb60b941468ec6785eeab3dc61a837ceda4e9
Author: wren romano <wren at community.haskell.org>
Date:   Wed Feb 8 22:09:06 2017 -0800

    defined lookupPrefix as part of optimizing restrictKeys


>---------------------------------------------------------------

af7bb60b941468ec6785eeab3dc61a837ceda4e9
 Data/IntMap/Internal.hs | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/Data/IntMap/Internal.hs b/Data/IntMap/Internal.hs
index 15d83a4..95e5259 100644
--- a/Data/IntMap/Internal.hs
+++ b/Data/IntMap/Internal.hs
@@ -1125,7 +1125,7 @@ restrictKeys t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
         | zero p1 m2        = restrictKeys t1 l2
         | otherwise         = restrictKeys t1 r2
 restrictKeys t1@(Bin _ _ _ _) (IntSet.Tip p2 bm2) =
-    restrictBM t1 p2 bm2 (IntSet.suffixBitMask + 1)
+    restrictBM (lookupPrefix p2 bm2 t1) p2 bm2 (IntSet.suffixBitMask + 1)
 restrictKeys (Bin _ _ _ _) IntSet.Nil = Nil
 restrictKeys t1@(Tip k1 _) t2
     | k1 `IntSet.member` t2 = t1
@@ -1136,6 +1136,21 @@ restrictKeys Nil _ = Nil
 type IntSetPrefix = Int
 type IntSetBitMap = Word
 
+-- | Find the sub-tree of @t@ which matches the prefix @kp at .
+lookupPrefix :: IntSetPrefix -> IntSetBitMap -> IntMap a -> IntMap a
+lookupPrefix !kp !bm t@(Bin p m l r)
+    | m .&. IntSet.suffixBitMask /= 0 =
+        if p .&. IntSet.prefixBitMask == kp then t else Nil
+    | nomatch kp p m = Nil
+    | zero kp m      = lookupPrefix kp bm l
+    | otherwise      = lookupPrefix kp bm r
+lookupPrefix kp bm t@(Tip kx x)
+    -- TODO(wrengr): need we manually inline 'IntSet.Member' here?
+    | kx `IntSet.member` IntSet.Tip kp bm = t
+    | otherwise = Nil
+lookupPrefix _ _ Nil = Nil
+
+
 -- TODO(wrengr): Right now this is still pretty naive. It essentially
 -- unpacks the 'IntSetBitMap' into a tree-representation, and then
 -- calls 'lookup' on each element of the set; thus, it is
@@ -1146,6 +1161,11 @@ type IntSetBitMap = Word
 -- on the recursive calls, so that the 'lookup's are cheaper. But
 -- we should be able to do even better by avoiding the call to
 -- 'lookup' entirely.
+--
+-- | The initial value passed for the final argument should be
+-- @(IntSet.suffixBitMask + 1)@. We don't set that here via
+-- worker/wrapper, because this is the worker floated to the
+-- top-level.
 restrictBM :: IntMap a -> IntSetPrefix -> IntSetBitMap -> Key -> IntMap a
 restrictBM t !prefix !_ 0 =
     case lookup prefix t of