{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_HADDOCK hide #-}
module Data.BitSet where
import Data.Bits
import Prelude hiding ( foldl, foldr )
import qualified Data.List as List
import GHC.Exts ( IsList, build )
import qualified GHC.Exts as Exts
newtype BitSet c a = BitSet { getBits :: c }
deriving Eq
instance (Enum a, Show a, Bits c, Num c) => Show (BitSet c a) where
showsPrec p bs
= showParen (p > 10)
$ showString "fromList " . shows (toList bs)
instance (Enum a, Bits c) => Semigroup (BitSet c a) where
(<>) = union
instance (Enum a, Bits c, Num c) => Monoid (BitSet c a) where
mempty = empty
instance (Enum a, Bits c, Num c) => IsList (BitSet c a) where
type Item (BitSet c a) = a
fromList = fromList
toList = toList
{-# INLINE fromList #-}
{-# INLINE toList #-}
{-# INLINE null #-}
null :: (Eq c, Num c) => BitSet c a -> Bool
null (BitSet bits) = bits == 0
{-# INLINE size #-}
size :: Bits c => BitSet c a -> Int
size (BitSet bits) = popCount bits
{-# INLINE member #-}
member :: (Enum a , Bits c) => a -> BitSet c a -> Bool
member x (BitSet bits) = bits `testBit` fromEnum x
{-# INLINE empty #-}
empty :: (Enum a, Bits c, Num c) => BitSet c a
empty = BitSet 0
{-# INLINE singleton #-}
singleton :: (Enum a, Bits c, Num c) => a -> BitSet c a
singleton x = BitSet $! bit (fromEnum x)
{-# INLINE insert #-}
insert :: (Enum a, Bits c) => a -> BitSet c a -> BitSet c a
insert x (BitSet bits) = BitSet $! bits `setBit` fromEnum x
{-# INLINE delete #-}
delete :: (Enum a, Bits c) => a -> BitSet c a -> BitSet c a
delete x (BitSet bits ) = BitSet $! bits `clearBit` fromEnum x
{-# INLINE union #-}
union :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
union (BitSet bits1) (BitSet bits2) = BitSet $! bits1 .|. bits2
{-# INLINE difference #-}
difference :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
difference (BitSet bits1) (BitSet bits2) = BitSet $! bits1 .&. complement bits2
infix 5 \\
{-# INLINE (\\) #-}
(\\) :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
(\\) = difference
{-# INLINE intersection #-}
intersection :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
intersection (BitSet bits1) (BitSet bits2) = BitSet $! bits1 .&. bits2
{-# INLINE map #-}
map :: (Enum a, Enum b, Bits c, Num c) => (a -> b) -> BitSet c a -> BitSet c b
map f = foldl' (\bs a -> f a `insert` bs) empty
{-# INLINE foldl' #-}
foldl' :: (Enum a, Bits c) => (b -> a -> b) -> b -> BitSet c a -> b
foldl' f z (BitSet bits) = go z (popCount bits) 0
where
go !acc 0 !_ = acc
go !acc !n !b = if bits `testBit` b
then go (f acc $ toEnum b) (pred n) (succ b)
else go acc n (succ b)
{-# INLINE foldr #-}
foldr :: (Enum a, Bits c) => (a -> b -> b) -> b -> BitSet c a -> b
foldr f z (BitSet bits) = go (popCount bits) 0
where
go 0 !_ = z
go !n !b = if bits `testBit` b
then toEnum b `f` go (pred n) (succ b)
else go n (succ b)
{-# INLINE [0] toList #-}
toList :: (Enum a, Bits c, Num c) => BitSet c a -> [a]
toList bs = build (\k z -> foldr k z bs)
{-# INLINE [0] fromList #-}
fromList :: (Enum a, Bits c, Num c) => [a] -> BitSet c a
fromList xs = BitSet $! List.foldl' (\i x -> i `setBit` fromEnum x) 0 xs
{-# RULES
"fromList/toList" forall bs. fromList (toList bs) = bs
#-}