Add Data.IntSet.fromRange (#965)

fromRange (l,h) is faster than fromRange [l..h], primarily because we
can pack Tips with multiple elements at once.

Author: meooow25

containers-tests/benchmarks/IntSet.hs

@@ -38,6 +38,8 @@ main = do
         , bench "difference" $ whnf (IS.difference s) s_even
         , bench "intersection" $ whnf (IS.intersection s) s_even
         , bench "fromList" $ whnf IS.fromList elems
+        , bench "fromRange" $ whnf IS.fromRange (1,bound)
+        , bench "fromRange:small" $ whnf IS.fromRange (-1,0)
         , bench "fromAscList" $ whnf IS.fromAscList elems
         , bench "fromDistinctAscList" $ whnf IS.fromDistinctAscList elems
         , bench "disjoint:false" $ whnf (IS.disjoint s) s_even
@@ -56,12 +58,13 @@ main = do
         , bench "splitMember:sparse" $ whnf (IS.splitMember elem_sparse_mid) s_sparse
         ]
   where
-    elems = [1..2^12]
-    elems_even = [2,4..2^12]
-    elems_odd = [1,3..2^12]
-    elem_mid = 2^11 + 31 -- falls in the middle of a packed Tip bitmask (assuming 64-bit words)
+    bound = 2^12
+    elems = [1..bound]
+    elems_even = [2,4..bound]
+    elems_odd = [1,3..bound]
+    elem_mid = bound `div` 2 + 31 -- falls in the middle of a packed Tip bitmask (assuming 64-bit words)
     elems_sparse = map (*64) elems -- when built into a map, each Tip is a singleton
-    elem_sparse_mid = 2^11 * 64
+    elem_sparse_mid = bound `div` 2 * 64
 
 member :: [Int] -> IS.IntSet -> Int
 member xs s = foldl' (\n x -> if IS.member x s then n + 1 else n) 0 xs

containers-tests/tests/intset-properties.hs

@@ -45,6 +45,7 @@ main = defaultMain $ testGroup "intset-properties"
                    , testProperty "prop_DescList" prop_DescList
                    , testProperty "prop_AscDescList" prop_AscDescList
                    , testProperty "prop_fromList" prop_fromList
+                   , testProperty "prop_fromRange" prop_fromRange
                    , testProperty "prop_MaskPow2" prop_MaskPow2
                    , testProperty "prop_Prefix" prop_Prefix
                    , testProperty "prop_LeftRight" prop_LeftRight
@@ -277,6 +278,12 @@ prop_fromList xs
   where sort_xs = sort xs
         nub_sort_xs = List.map List.head $ List.group sort_xs
 
+prop_fromRange :: Property
+prop_fromRange = forAll (scale (*100) arbitrary) go
+  where
+    go (l,h) = valid t .&&. t === fromAscList [l..h]
+      where t = fromRange (l,h)
+
 {--------------------------------------------------------------------
   Bin invariants
 --------------------------------------------------------------------}

containers/src/Data/IntSet.hs

@@ -76,6 +76,7 @@ module Data.IntSet (
             , empty
             , singleton
             , fromList
+            , fromRange
             , fromAscList
             , fromDistinctAscList
 

containers/src/Data/IntSet/Internal.hs

@@ -123,6 +123,7 @@ module Data.IntSet.Internal (
     -- * Construction
     , empty
     , singleton
+    , fromRange
     , insert
     , delete
     , alterF
@@ -1215,6 +1216,60 @@ fromList xs
   where
     ins t x  = insert x t
 
+-- | \(O(n / W)\). Create a set from a range of integers.
+--
+-- > fromRange (low, high) == fromList [low..high]
+--
+-- @since FIXME
+fromRange :: (Key, Key) -> IntSet
+fromRange (lx,rx)
+  | lx > rx  = empty
+  | lp == rp = Tip lp (bitmapOf rx `shiftLL` 1 - bitmapOf lx)
+  | otherwise =
+      let m = branchMask lx rx
+          p = mask lx m
+      in if m < 0  -- handle negative numbers
+         then Bin 0 m (goR 0) (goL 0)
+         else Bin p m (goL (p .|. m)) (goR (p .|. m))
+  where
+    lp = prefixOf lx
+    rp = prefixOf rx
+    -- goL p0 = fromList [lx .. p0-1]
+    -- Expected: p0 is lx where one 0-bit is flipped to 1 and all bits lower than that are 0.
+    --           p0 can be 0 (pretend that bit WORD_SIZE is flipped to 1).
+    goL :: Prefix -> IntSet
+    goL !p0 = go (Tip lp (- bitmapOf lx)) (lp + lbm prefixBitMask)
+      where
+        go !l p | p == p0 = l
+        go l p =
+          let m = lbm p
+              p' = p `xor` m
+              l' = Bin p' m l (goFull p (shr1 m))
+          in go l' (p + m)
+    -- goR p0 = fromList [p0 .. rx]
+    -- Expected: p0 is a prefix of rx
+    goR :: Prefix -> IntSet
+    goR !p0 = go (Tip rp (bitmapOf rx `shiftLL` 1 - 1)) rp
+      where
+        go !r p | p == p0 = r
+        go r p =
+          let m = lbm p
+              p' = p `xor` m
+              r' = Bin p' m (goFull p' (shr1 m)) r
+          in go r' p'
+    -- goFull p m = fromList [p .. p+2*m-1]
+    -- Expected: popCount m == 1, p == mask p m
+    goFull :: Prefix -> Mask -> IntSet
+    goFull p m
+      | m < suffixBitMask = Tip p (complement 0)
+      | otherwise         = Bin p m (goFull p (shr1 m)) (goFull (p .|. m) (shr1 m))
+    lbm :: Prefix -> Prefix
+    lbm p = intFromNat (lowestBitMask (natFromInt p))
+    {-# INLINE lbm #-}
+    shr1 :: Mask -> Mask
+    shr1 m = intFromNat (natFromInt m `shiftRL` 1)
+    {-# INLINE shr1 #-}
+
 -- | \(O(n)\). Build a set from an ascending list of elements.
 -- /The precondition (input list is ascending) is not checked./
 fromAscList :: [Key] -> IntSet