IntSet vs Int8Set

Test.hs

import Control.Monad
import Control.Exception (evaluate)
import qualified Data.IntSet as IS
import qualified Data.Map as M
import qualified Data.ByteString as B
import System.IO
import Data.Int
import Data.Word
import Data.Bits
import Criterion.Main
import System.Random.MWC

---------- Int8Set implementation

data Int8Set = Int8Set {
    word0 :: {-# UNPACK #-} ! Word64,
    word1 :: {-# UNPACK #-} ! Word64,
    word2 :: {-# UNPACK #-} ! Word64,
    word3 :: {-# UNPACK #-} ! Word64
  }
  deriving (Show)

-- data Int8Set = Int8Set {
--     word0 :: ! Word64,
--     word1 :: ! Word64,
--     word2 :: ! Word64,
--     word3 :: ! Word64
--   }
--   deriving (Show)

bitNumber :: Word8 -> (Int, Int)
bitNumber x =
  let x' = fromIntegral x
  in  (x' `div` 64, x' `mod` 64)

empty8 :: Int8Set
empty8 = Int8Set 0 0 0 0

insert8 :: Word8 -> Int8Set -> Int8Set
insert8 x set =
  let (word, bit) = bitNumber x
  in  case word of
        0 -> set {word0 = setBit (word0 set) bit}
        1 -> set {word1 = setBit (word1 set) bit}
        2 -> set {word2 = setBit (word2 set) bit}
        3 -> set {word3 = setBit (word3 set) bit}

delete8 :: Word8 -> Int8Set -> Int8Set
delete8 x set =
  let (word, bit) = bitNumber x
  in  case word of
        0 -> set {word0 = clearBit (word0 set) bit}
        1 -> set {word1 = clearBit (word1 set) bit}
        2 -> set {word2 = clearBit (word2 set) bit}
        3 -> set {word3 = clearBit (word3 set) bit}

member8 :: Word8 -> Int8Set -> Bool
member8 x set =
  let (word, bit) = bitNumber x
  in  case word of
        0 -> testBit (word0 set) bit
        1 -> testBit (word1 set) bit
        2 -> testBit (word2 set) bit
        3 -> testBit (word3 set) bit

fromList8 :: [Word8] -> Int8Set
fromList8 list = foldr insert8 empty8 list

---- Test IntSet

initialIntSet :: IS.IntSet
initialIntSet = IS.fromList [50..100]

testInsertIS :: [Int] -> IS.IntSet -> IS.IntSet
testInsertIS numbers x = foldr IS.insert x numbers

testDeleteIS :: [Int] -> IS.IntSet -> IS.IntSet
testDeleteIS numbers x = foldr IS.delete x numbers

testMoveIS :: [Int] -> IS.IntSet -> IS.IntSet
testMoveIS numbers x = foldr (\i set -> IS.delete i $ IS.insert i set) x numbers

testCheckIS :: [Int] -> IS.IntSet -> Bool
testCheckIS numbers x = and [IS.member i x | i <- numbers]

----- Test Int8Set

initialInt8Set :: Int8Set
initialInt8Set = fromList8 [50..100]

testInsert8 :: [Word8] -> Int8Set -> Int8Set
testInsert8 numbers x = foldr insert8 x numbers

testDelete8 :: [Word8] -> Int8Set -> Int8Set
testDelete8 numbers x = foldr delete8 x numbers

testMove8 :: [Word8] -> Int8Set -> Int8Set
testMove8 numbers x = foldr (\i set -> delete8 i $ insert8 i set) x numbers

testCheck8 :: [Word8] -> Int8Set -> Bool
testCheck8 numbers x = and [member8 i x | i <- numbers]

-----------------------------------

generate :: (Variate a, Integral a) => Int -> IO [a]
generate n =
  withSystemRandom . asGenIO $ \gen ->
    replicateM n $ uniformR (1, 254) gen

setupEnv :: (Variate a, Integral a) => Int -> IO ([a], [a])
setupEnv n = do
  xs <- generate n
  ys <- generate n
  return (xs, ys)

------------------------------

n = 1000*1000

main = defaultMain [
    env (setupEnv n) $ \ ~(initial, list) -> bgroup "IntSet" [
      bench "insert" $ whnfIO $ evaluate $ (testInsertIS list) (IS.fromList initial),
      bench "delete" $ whnfIO $ evaluate $ (testDeleteIS list) (IS.fromList initial),
      bench "move" $ whnfIO $ evaluate $ (testMoveIS list) (IS.fromList initial),
      bench "check" $ whnfIO $ evaluate $ (testCheckIS list) (IS.fromList initial)
    ],
    env (setupEnv n) $ \ ~(initial, list) -> bgroup "Int8Set" [
      bench "insert" $ whnfIO $ evaluate $ (testInsert8 list) (fromList8 initial),
      bench "delete" $ whnfIO $ evaluate $ (testDelete8 list) (fromList8 initial),
      bench "move" $ whnfIO $ evaluate $ (testMove8 list) (fromList8 initial),
      bench "check" $ whnfIO $ evaluate $ (testCheck8 list) (fromList8 initial)
    ]
  ]