module Data.ByteSet (
ByteSet
, Word8
, null
, size
, member
, notMember
, empty
, singleton
, insert
, delete
, union
, unions
, difference
, intersection
, filter
, map
, foldr
, elems
, toList
, fromList
) where
import Prelude
( Eq (..), Ord (..)
, (+), (), ($), (*)
, fromIntegral
, Show (..)
)
import Data.Word (Word8, Word64)
import Data.Int (Int)
import Data.Bits
import Data.Bool
import Control.Category
import Control.Applicative (Applicative (..))
import qualified Data.Foldable as F
import qualified Data.List as L
import Data.Functor
import GHC.Generics (Generic)
import Data.Binary (Binary (..))
import Data.Binary.Put (putWord64le)
import Data.Binary.Get (getWord64le)
data ByteSet = ByteSet
!Word64
!Word64
!Word64
!Word64
deriving (Eq, Ord, Generic)
generalGet :: (Word64 -> Int -> a) -> Word8 -> ByteSet -> a
generalGet f w (ByteSet s1 s2 s3 s4)
| w < 64 = f s1 i
| w < 128 = f s2 (i 64)
| w < 192 = f s3 (i 128)
| otherwise = f s4 (i 192)
where
i = fromIntegral w
generalSet :: (Word64 -> Int -> Word64) -> Word8 -> ByteSet -> ByteSet
generalSet f w (ByteSet s1 s2 s3 s4)
| w < 64 = ByteSet (f s1 i ) s2 s3 s4
| w < 128 = ByteSet s1 (f s2 $ i 64) s3 s4
| w < 192 = ByteSet s1 s2 (f s3 $ i 128) s4
| otherwise = ByteSet s1 s2 s3 (f s4 $ i 192)
where
i = fromIntegral w
generalOp :: (Word64 -> Word64 -> Word64 )
-> ByteSet -> ByteSet -> ByteSet
generalOp f (ByteSet s1 s2 s3 s4)
(ByteSet t1 t2 t3 t4) =
ByteSet (f s1 t1) (f s2 t2)
(f s3 t3) (f s4 t4)
generalFun :: (Int -> Word64 -> Word64) -> ByteSet -> ByteSet
generalFun f (ByteSet s1 s2 s3 s4) =
ByteSet (f 0 s1) (f 1 s2) (f 2 s3) (f 3 s4)
size :: ByteSet -> Int
size (ByteSet s1 s2 s3 s4) =
popCount s1 + popCount s2
+ popCount s3 + popCount s4
null :: ByteSet -> Bool
null = (==0) . size
member :: Word8 -> ByteSet -> Bool
member = generalGet testBit
notMember :: Word8 -> ByteSet -> Bool
notMember w = not . member w
empty :: ByteSet
empty = ByteSet 0 0 0 0
insert :: Word8 -> ByteSet -> ByteSet
insert = generalSet setBit
singleton :: Word8 -> ByteSet
singleton w = insert w empty
delete :: Word8 -> ByteSet -> ByteSet
delete = generalSet clearBit
union :: ByteSet -> ByteSet -> ByteSet
union = generalOp (.|.)
unions :: [ByteSet] -> ByteSet
unions = F.foldl' union empty
difference :: ByteSet -> ByteSet -> ByteSet
difference = generalOp $ \w1 w2 -> w1 .&. complement w2
intersection :: ByteSet -> ByteSet -> ByteSet
intersection = generalOp (.&.)
filter :: (Word8 -> Bool) -> ByteSet -> ByteSet
filter f = generalFun $ \i w ->
let b0 = i * 64
go acc (1) = acc
go !acc n = if testBit w n && f (fromIntegral $ b0 + n)
then go (setBit acc n) $ n 1
else go acc $ n 1
in go 0 63
bits :: Word64 -> [Int]
bits w = L.filter (testBit w) [0..63]
foldr :: (Word8 -> a -> a) -> a -> ByteSet -> a
foldr f r0 (ByteSet s1 s2 s3 s4) =
let g = f . fromIntegral
r1 = F.foldr g r0 $ fmap (+192) $ bits s4
r2 = F.foldr g r1 $ fmap (+128) $ bits s3
r3 = F.foldr g r2 $ fmap (+ 64) $ bits s2
in F.foldr g r3 $ bits s1
map :: (Word8 -> Word8) -> ByteSet -> ByteSet
map f = foldr (insert . f) empty
elems :: ByteSet -> [Word8]
elems = foldr (:) []
toList :: ByteSet -> [Word8]
toList = elems
fromList :: [Word8] -> ByteSet
fromList = F.foldr insert empty
instance Show ByteSet where
show = show . elems
instance Binary ByteSet where
put (ByteSet s1 s2 s3 s4) =
putWord64le s1 *> putWord64le s2
*> putWord64le s3 *> putWord64le s4
get = ByteSet <$> getWord64le <*> getWord64le
<*> getWord64le <*> getWord64le