Optimize IntMap.Bin (haskell#995)

* IntMap: Merge prefix and mask in Bin

* Update lookupGE benchmarks

* Add note on the relation to the Okasaki-Gill paper

* Adopt BurningWitness's lower/upper helpers

* Add property tests for is{Proper}SubmapOfBy

* Generate large keys in Arbitrary IntMap for better coverage

* Add GHC flag to prevent benchmarks from changing due to unpredictable alignment changes

containers-tests/benchmarks/LookupGE/LookupGE_IntMap.hs

@@ -12,18 +12,18 @@ lookupGE1 k m =
 
 lookupGE2 :: Key -> IntMap a -> Maybe (Key,a)
 lookupGE2 k t = case t of
-    Bin _ m l r | m < 0 -> if k >= 0
+    Bin p l r | signBranch p -> if k >= 0
       then go l
       else case go r of
         Nothing -> Just $ findMin l
         justx -> justx
     _ -> go t
   where
-    go (Bin p m l r)
-      | nomatch k p m = if k < p
+    go (Bin p l r)
+      | nomatch k p = if k < unPrefix p
         then Just $ findMin l
         else Nothing
-      | zero k m = case go l of
+      | left k p = case go l of
         Nothing -> Just $ findMin r
         justx -> justx
       | otherwise = go r
@@ -34,14 +34,14 @@ lookupGE2 k t = case t of
 
 lookupGE3 :: Key -> IntMap a -> Maybe (Key,a)
 lookupGE3 k t = k `seq` case t of
-    Bin _ m l r | m < 0 -> if k >= 0
+    Bin p l r | signBranch p -> if k >= 0
       then go Nothing l
       else go (Just (findMin l)) r
     _ -> go Nothing t
   where
-    go def (Bin p m l r)
-      | nomatch k p m = if k < p then Just $ findMin l else def
-      | zero k m  = go (Just $ findMin r) l
+    go def (Bin p l r)
+      | nomatch k p = if k < unPrefix p then Just $ findMin l else def
+      | left k p  = go (Just $ findMin r) l
       | otherwise = go def r
     go def (Tip ky y)
       | k > ky    = def
@@ -50,13 +50,13 @@ lookupGE3 k t = k `seq` case t of
 
 lookupGE4 :: Key -> IntMap a -> Maybe (Key,a)
 lookupGE4 k t = k `seq` case t of
-    Bin _ m l r | m < 0 -> if k >= 0 then go Nil l
-                                     else go l r
+    Bin p l r | signBranch p -> if k >= 0 then go Nil l
+                                          else go l r
     _ -> go Nil t
   where
-    go def (Bin p m l r)
-      | nomatch k p m = if k < p then fMin l else fMin def
-      | zero k m  = go r l
+    go def (Bin p l r)
+      | nomatch k p = if k < unPrefix p then fMin l else fMin def
+      | left k p  = go r l
       | otherwise = go def r
     go def (Tip ky y)
       | k > ky    = fMin def
@@ -66,7 +66,7 @@ lookupGE4 k t = k `seq` case t of
     fMin :: IntMap a -> Maybe (Key, a)
     fMin Nil = Nothing
     fMin (Tip ky y) = Just (ky, y)
-    fMin (Bin _ _ l _) = fMin l
+    fMin (Bin _ l _) = fMin l
 
 -------------------------------------------------------------------------------
 -- Utilities

containers-tests/containers-tests.cabal

@@ -57,6 +57,10 @@ common benchmark-deps
     , deepseq           >=1.1.0.0 && <1.6
     , tasty-bench       >=0.3.1   && <0.4
 
+  -- Flags recommended by tasty-bench
+  if impl(ghc >= 8.6)
+    ghc-options:    -fproc-alignment=64
+
 -- Copy of containers library,
 library
   import: deps
@@ -70,7 +74,11 @@ library
 
   include-dirs:     include
   hs-source-dirs:   src, tests
+
   ghc-options:      -O2 -Wall
+  if impl(ghc >= 8.6)
+    ghc-options:    -fproc-alignment=64
+
   other-extensions:
     BangPatterns
     CPP

containers-tests/tests/IntMapValidity.hs

@@ -1,9 +1,14 @@
-module IntMapValidity (valid) where
+module IntMapValidity
+  ( valid
+  , hasPrefix
+  , hasPrefixSimple
+  ) where
 
-import Data.Bits (xor, (.&.))
+import Data.Bits (finiteBitSize, testBit, xor, (.&.))
+import Data.List (intercalate, elemIndex)
 import Data.IntMap.Internal
+import Numeric (showHex)
 import Test.Tasty.QuickCheck (Property, counterexample, property, (.&&.))
-import Utils.Containers.Internal.BitUtil (bitcount)
 
 {--------------------------------------------------------------------
   Assertions
@@ -12,54 +17,61 @@ import Utils.Containers.Internal.BitUtil (bitcount)
 valid :: IntMap a -> Property
 valid t =
   counterexample "nilNeverChildOfBin" (nilNeverChildOfBin t) .&&.
-  counterexample "commonPrefix" (commonPrefix t) .&&.
-  counterexample "maskRespected" (maskRespected t)
+  counterexample "prefixOk" (prefixOk t)
 
 -- Invariant: Nil is never found as a child of Bin.
 nilNeverChildOfBin :: IntMap a  -> Bool
 nilNeverChildOfBin t =
   case t of
     Nil -> True
     Tip _ _ -> True
-    Bin _ _ l r -> noNilInSet l && noNilInSet r
+    Bin _ l r -> noNilInSet l && noNilInSet r
   where
     noNilInSet t' =
       case t' of
         Nil -> False
         Tip _ _ -> True
-        Bin _ _ l' r' -> noNilInSet l' && noNilInSet r'
+        Bin _ l' r' -> noNilInSet l' && noNilInSet r'
 
--- Invariant: The Mask is a power of 2. It is the largest bit position at which
---            two keys of the map differ.
-maskPowerOfTwo :: IntMap a -> Bool
-maskPowerOfTwo t =
+-- Invariants:
+-- * All keys in a Bin start with the Bin's shared prefix.
+-- * All keys in the Bin's left child have the Prefix's mask bit unset.
+-- * All keys in the Bin's right child have the Prefix's mask bit set.
+prefixOk :: IntMap a -> Property
+prefixOk t =
   case t of
-    Nil -> True
-    Tip _ _ -> True
-    Bin _ m l r ->
-      bitcount 0 (fromIntegral m) == 1 && maskPowerOfTwo l && maskPowerOfTwo r
+    Nil -> property ()
+    Tip _ _ -> property ()
+    Bin p l r ->
+      let px = unPrefix p
+          m = px .&. (-px)
+          keysl = keys l
+          keysr = keys r
+          debugStr = concat
+            [ "px=" ++ showIntHex px
+            , ", keysl=[" ++ intercalate "," (fmap showIntHex keysl) ++ "]"
+            , ", keysr=[" ++ intercalate "," (fmap showIntHex keysr) ++ "]"
+            ]
+      in counterexample debugStr $
+           counterexample "mask bit absent" (px /= 0) .&&.
+           counterexample "prefix not shared" (all (`hasPrefix` p) (keysl ++ keysr)) .&&.
+           counterexample "left child, mask found set" (all (\x -> x .&. m == 0) keysl) .&&.
+           counterexample "right child, mask found unset" (all (\x -> x .&. m /= 0) keysr)
 
--- Invariant: Prefix is the common high-order bits that all elements share to
---            the left of the Mask bit.
-commonPrefix :: IntMap a -> Bool
-commonPrefix t =
-  case t of
-    Nil -> True
-    Tip _ _ -> True
-    b@(Bin p _ l r) -> all (sharedPrefix p) (keys b) && commonPrefix l && commonPrefix r
+hasPrefix :: Int -> Prefix -> Bool
+hasPrefix i p = not (nomatch i p)
+
+-- We test that hasPrefix behaves the same as hasPrefixSimple.
+hasPrefixSimple :: Int -> Prefix -> Bool
+hasPrefixSimple k p = case elemIndex True pbits of
+  Nothing -> error "no mask bit" -- should not happen
+  Just i -> drop (i+1) kbits == drop (i+1) pbits
   where
-    sharedPrefix :: Prefix -> Int -> Bool
-    sharedPrefix p a = p == p .&. a
+    kbits = toBits k
+    pbits = toBits (unPrefix p)
 
--- Invariant: In Bin prefix mask left right, left consists of the elements that
---            don't have the mask bit set; right is all the elements that do.
-maskRespected :: IntMap a -> Bool
-maskRespected t =
-  case t of
-    Nil -> True
-    Tip _ _ -> True
-    Bin _ binMask l r ->
-      all (\x -> zero x binMask) (keys l) &&
-      all (\x -> not (zero x binMask)) (keys r) &&
-      maskRespected l &&
-      maskRespected r
+    -- Bits from lowest to highest.
+    toBits x = fmap (testBit x) [0 .. finiteBitSize (0 :: Int) - 1]
+
+showIntHex :: Int -> String
+showIntHex x = "0x" ++ showHex (fromIntegral x :: Word) ""

containers-tests/tests/intmap-properties.hs

@@ -10,7 +10,8 @@ import Data.IntMap.Internal (traverseMaybeWithKey)
 import Data.IntMap.Merge.Lazy
 #endif
 import Data.IntMap.Internal.Debug (showTree)
-import IntMapValidity (valid)
+import Data.IntMap.Internal (Prefix(..))
+import IntMapValidity (hasPrefix, hasPrefixSimple, valid)
 
 import Control.Applicative (Applicative(..))
 import Control.Monad ((<=<))
@@ -134,6 +135,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testCase "minimum" test_minimum
              , testCase "maximum" test_maximum
              , testProperty "valid"                prop_valid
+             , testProperty "hasPrefix"            prop_hasPrefix
              , testProperty "empty valid"          prop_emptyValid
              , testProperty "insert to singleton"  prop_singleton
              , testProperty "insert then lookup"   prop_insertLookup
@@ -209,6 +211,8 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "traverseMaybeWithKey identity"         prop_traverseMaybeWithKey_identity
              , testProperty "traverseMaybeWithKey->mapMaybeWithKey" prop_traverseMaybeWithKey_degrade_to_mapMaybeWithKey
              , testProperty "traverseMaybeWithKey->traverseWithKey" prop_traverseMaybeWithKey_degrade_to_traverseWithKey
+             , testProperty "isProperSubmapOfBy"   prop_isProperSubmapOfBy
+             , testProperty "isSubmapOfBy"         prop_isSubmapOfBy
              ]
 
 apply2 :: Fun (a, b) c -> a -> b -> c
@@ -223,12 +227,23 @@ apply3 f a b c = apply f (a, b, c)
 --------------------------------------------------------------------}
 
 instance Arbitrary a => Arbitrary (IntMap a) where
-  arbitrary = fmap fromList arbitrary
+  arbitrary = oneof [go arbitrary, go (getLarge <$> arbitrary)]
+    where
+      go kgen = fromList <$> listOf ((,) <$> kgen <*> arbitrary)
+  shrink = fmap fromList . shrink . toAscList
 
 newtype NonEmptyIntMap a = NonEmptyIntMap {getNonEmptyIntMap :: IntMap a} deriving (Eq, Show)
 
 instance Arbitrary a => Arbitrary (NonEmptyIntMap a) where
-  arbitrary = fmap (NonEmptyIntMap . fromList . getNonEmpty) arbitrary
+  arbitrary = oneof [go arbitrary, go (getLarge <$> arbitrary)]
+    where
+      go kgen = NonEmptyIntMap . fromList <$> listOf1 ((,) <$> kgen <*> arbitrary)
+  shrink =
+    fmap (NonEmptyIntMap . fromList) .
+    List.filter (not . List.null) .
+    shrink .
+    toAscList .
+    getNonEmptyIntMap
 
 
 ------------------------------------------------------------------------
@@ -1150,6 +1165,10 @@ forValidUnitTree f = forValid f
 prop_valid :: Property
 prop_valid = forValidUnitTree $ \t -> valid t
 
+prop_hasPrefix :: Int -> NonZero Int -> Property
+prop_hasPrefix i (NonZero p) =
+  hasPrefix i (Prefix p) === hasPrefixSimple i (Prefix p)
+
 ----------------------------------------------------------------
 -- QuickCheck
 ----------------------------------------------------------------
@@ -1641,3 +1660,21 @@ prop_traverseMaybeWithKey_degrade_to_traverseWithKey fun mp =
         -- so this also checks the order of traversing is the same.
   where f k v = (show k, applyFun2 fun k v)
         g k v = fmap Just $ f k v
+
+prop_isProperSubmapOfBy :: Fun (A, A) Bool -> IntMap A -> IntMap A -> Property
+prop_isProperSubmapOfBy f m1 m2 =
+  isProperSubmapOfBy (applyFun2 f) m1 m2 ===
+  (length xs == size m1 && size m1 < size m2)
+  where
+    xs = List.intersectBy
+           (\(k1,x1) (k2,x2) -> k1 == k2 && applyFun2 f x1 x2)
+           (assocs m1) (assocs m2)
+
+prop_isSubmapOfBy :: Fun (A, A) Bool -> IntMap A -> IntMap A -> Property
+prop_isSubmapOfBy f m1 m2 =
+  isSubmapOfBy (applyFun2 f) m1 m2 ===
+  (length xs == size m1)
+  where
+    xs = List.intersectBy
+           (\(k1,x1) (k2,x2) -> k1 == k2 && applyFun2 f x1 x2)
+           (assocs m1) (assocs m2)