{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
module Data.IntSet.Bounded.Imperative (
IntSet
, IOIntSet
, intSetMinBound
, intSetMaxBound
, empty
, insert
, member
, notMember
, delete
) where
import Control.DeepSeq (deepseq)
import Control.Monad.Primitive (PrimMonad(..))
import Control.Monad.ST (ST)
import Data.Bits
import Data.Primitive.ByteArray
import Data.Primitive.MutVar
import Data.Word (Word64)
data IntSet s = IntSet {
intSetMinBound# :: {-# UNPACK #-} !Word64
, intSetMaxBound# :: {-# UNPACK #-} !Word64
, intSetInBounds# :: {-# UNPACK #-} !(MutableByteArray s)
, intSetOutBounds# :: {-# UNPACK #-} !(MutVar s [Word64])
}
type IOIntSet = IntSet (PrimState IO)
intSetMinBound :: IntSet s -> Word64
intSetMinBound = intSetMinBound#
{-# INLINE intSetMinBound #-}
intSetMaxBound :: IntSet s -> Word64
intSetMaxBound = intSetMaxBound#
{-# INLINE intSetMaxBound #-}
empty :: PrimMonad m
=> Word64
-> Word64
-> m (IntSet (PrimState m))
empty !minB !maxB = do
let !numInBounds = (maxB - minB) `div` 8 + 1
set <- newByteArray (fromIntegral numInBounds)
fillByteArray set 0 (fromIntegral numInBounds) 0
outBounds <- newMutVar []
return $! IntSet {
intSetMinBound# = minB
, intSetMaxBound# = maxB
, intSetInBounds# = set
, intSetOutBounds# = outBounds
}
{-# SPECIALIZE empty :: Word64 -> Word64 -> IO (IntSet (PrimState IO)) #-}
{-# SPECIALIZE empty :: Word64 -> Word64 -> ST s (IntSet s) #-}
insert :: PrimMonad m
=> IntSet (PrimState m)
-> Word64
-> m ()
insert !set !n = do
if n >= intSetMinBound# set && n <= intSetMaxBound# set then do
let !n' = n - intSetMinBound# set
let !o = fromIntegral $ n' `shiftR` 6
let !i = fromIntegral $ n' .&. 63
let !mask = (1 :: Word64) `shiftL` i
b <- readByteArray (intSetInBounds# set) o
let !b' = b .|. mask
writeByteArray (intSetInBounds# set) o b'
else do
ns <- readMutVar (intSetOutBounds# set)
let !ns' = if n `elem` ns then ns else (n:ns)
writeMutVar (intSetOutBounds# set) ns'
{-# SPECIALIZE insert :: IntSet (PrimState IO) -> Word64 -> IO () #-}
{-# SPECIALIZE insert :: IntSet s -> Word64 -> ST s () #-}
delete :: PrimMonad m
=> IntSet (PrimState m)
-> Word64
-> m ()
delete !set !n = do
if n >= intSetMinBound# set && n <= intSetMaxBound# set then do
let !n' = n - intSetMinBound# set
let !o = fromIntegral $ n' `shiftR` 6
let !i = fromIntegral $ n' .&. 63
let !mask = (1 :: Word64) `shiftL` i
b <- readByteArray (intSetInBounds# set) o
let !b' = b .&. (complement mask)
writeByteArray (intSetInBounds# set) o b'
else do
ns <- readMutVar (intSetOutBounds# set)
let ns' = filter (/= n) ns
ns' `deepseq` writeMutVar (intSetOutBounds# set) ns'
{-# SPECIALIZE delete :: IntSet (PrimState IO) -> Word64 -> IO () #-}
{-# SPECIALIZE delete :: IntSet s -> Word64 -> ST s () #-}
member :: PrimMonad m
=> IntSet (PrimState m)
-> Word64
-> m Bool
member !set !n = do
if n >= intSetMinBound# set && n <= intSetMaxBound# set then do
let !n' = n - intSetMinBound# set
let !o = fromIntegral $ n' `shiftR` 6
let !i = fromIntegral $ n' .&. 63
let !mask = (1 :: Word64) `shiftL` i
b <- readByteArray (intSetInBounds# set) o
return $! (b .&. mask) /= 0
else do
ns <- readMutVar (intSetOutBounds# set)
return $! n `elem` ns
{-# SPECIALIZE member :: IntSet (PrimState IO) -> Word64 -> IO Bool #-}
{-# SPECIALIZE member :: IntSet s -> Word64 -> ST s Bool #-}
notMember :: PrimMonad m
=> IntSet (PrimState m)
-> Word64
-> m Bool
notMember !set !n = not <$> member set n
{-# INLINE notMember #-}
{-# SPECIALIZE notMember :: IntSet (PrimState IO) -> Word64 -> IO Bool #-}
{-# SPECIALIZE notMember :: IntSet s -> Word64 -> ST s Bool #-}