{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.BitSet
-- Copyright   : [2019..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

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


-- | A space-efficient implementation of a set data structure for
-- enumerated data types.
--
newtype BitSet c a = BitSet { BitSet c a -> c
getBits :: c }
  deriving BitSet c a -> BitSet c a -> Bool
(BitSet c a -> BitSet c a -> Bool)
-> (BitSet c a -> BitSet c a -> Bool) -> Eq (BitSet c a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall c a. Eq c => BitSet c a -> BitSet c a -> Bool
/= :: BitSet c a -> BitSet c a -> Bool
$c/= :: forall c a. Eq c => BitSet c a -> BitSet c a -> Bool
== :: BitSet c a -> BitSet c a -> Bool
$c== :: forall c a. Eq c => BitSet c a -> BitSet c a -> Bool
Eq

instance (Enum a, Show a, Bits c, Num c) => Show (BitSet c a) where
  showsPrec :: Int -> BitSet c a -> ShowS
showsPrec Int
p BitSet c a
bs
    = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10)
    (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String -> ShowS
showString String
"fromList " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> ShowS
forall a. Show a => a -> ShowS
shows (BitSet c a -> [a]
forall a c. (Enum a, Bits c, Num c) => BitSet c a -> [a]
toList BitSet c a
bs)

instance (Enum a, Bits c) => Semigroup (BitSet c a) where
  <> :: BitSet c a -> BitSet c a -> BitSet c a
(<>) = BitSet c a -> BitSet c a -> BitSet c a
forall c a. Bits c => BitSet c a -> BitSet c a -> BitSet c a
union

instance (Enum a, Bits c, Num c) => Monoid (BitSet c a) where
  mempty :: BitSet c a
mempty = BitSet c a
forall a c. (Enum a, Bits c, Num c) => BitSet c a
empty

instance (Enum a, Bits c, Num c) => IsList (BitSet c a) where
  type Item (BitSet c a) = a
  fromList :: [Item (BitSet c a)] -> BitSet c a
fromList = [Item (BitSet c a)] -> BitSet c a
forall a c. (Enum a, Bits c, Num c) => [a] -> BitSet c a
fromList
  toList :: BitSet c a -> [Item (BitSet c a)]
toList   = BitSet c a -> [Item (BitSet c a)]
forall a c. (Enum a, Bits c, Num c) => BitSet c a -> [a]
toList
  {-# INLINE fromList #-}
  {-# INLINE toList   #-}

-- | Is the bit set empty?
--
{-# INLINE null #-}
null :: (Eq c, Num c) => BitSet c a -> Bool
null :: BitSet c a -> Bool
null (BitSet c
bits) = c
bits c -> c -> Bool
forall a. Eq a => a -> a -> Bool
== c
0

-- | The number of elements in the bit set.
--
{-# INLINE size #-}
size :: Bits c => BitSet c a -> Int
size :: BitSet c a -> Int
size (BitSet c
bits) = c -> Int
forall a. Bits a => a -> Int
popCount c
bits

-- | Ask whether the item is in the bit set.
--
{-# INLINE member #-}
member :: (Enum a , Bits c) => a -> BitSet c a -> Bool
member :: a -> BitSet c a -> Bool
member a
x (BitSet c
bits) = c
bits c -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` a -> Int
forall a. Enum a => a -> Int
fromEnum a
x

-- | The empty bit set.
--
{-# INLINE empty #-}
empty :: (Enum a, Bits c, Num c) => BitSet c a
empty :: BitSet c a
empty = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet c
0

-- | Create a singleton set.
--
{-# INLINE singleton #-}
singleton :: (Enum a, Bits c, Num c) => a -> BitSet c a
singleton :: a -> BitSet c a
singleton a
x = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! Int -> c
forall a. Bits a => Int -> a
bit (a -> Int
forall a. Enum a => a -> Int
fromEnum a
x)

-- | Insert an item into the bit set.
--
{-# INLINE insert #-}
insert :: (Enum a, Bits c) => a -> BitSet c a -> BitSet c a
insert :: a -> BitSet c a -> BitSet c a
insert a
x (BitSet c
bits) = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! c
bits c -> Int -> c
forall a. Bits a => a -> Int -> a
`setBit` a -> Int
forall a. Enum a => a -> Int
fromEnum a
x

-- | Delete an item from the bit set.
{-# INLINE delete #-}
delete :: (Enum a, Bits c) => a -> BitSet c a -> BitSet c a
delete :: a -> BitSet c a -> BitSet c a
delete a
x (BitSet c
bits ) = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! c
bits c -> Int -> c
forall a. Bits a => a -> Int -> a
`clearBit` a -> Int
forall a. Enum a => a -> Int
fromEnum a
x

-- | The union of two bit sets.
--
{-# INLINE union #-}
union :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
union :: BitSet c a -> BitSet c a -> BitSet c a
union (BitSet c
bits1) (BitSet c
bits2) = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! c
bits1 c -> c -> c
forall a. Bits a => a -> a -> a
.|. c
bits2

-- | Difference of two bit sets.
--
{-# INLINE difference #-}
difference :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
difference :: BitSet c a -> BitSet c a -> BitSet c a
difference (BitSet c
bits1) (BitSet c
bits2) = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! c
bits1 c -> c -> c
forall a. Bits a => a -> a -> a
.&. c -> c
forall a. Bits a => a -> a
complement c
bits2

-- | See 'difference'.
--
infix 5 \\ -- comment to fool cpp: https://downloads.haskell.org/~ghc/latest/docs/html/users_guide/phases.html#cpp-and-string-gaps
{-# INLINE (\\) #-}
(\\) :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
\\ :: BitSet c a -> BitSet c a -> BitSet c a
(\\) = BitSet c a -> BitSet c a -> BitSet c a
forall c a. Bits c => BitSet c a -> BitSet c a -> BitSet c a
difference

-- | The intersection of two bit sets.
--
{-# INLINE intersection #-}
intersection :: Bits c => BitSet c a -> BitSet c a -> BitSet c a
intersection :: BitSet c a -> BitSet c a -> BitSet c a
intersection (BitSet c
bits1) (BitSet c
bits2) = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! c
bits1 c -> c -> c
forall a. Bits a => a -> a -> a
.&. c
bits2

-- | Transform this bit set by applying a function to every value.
-- Resulting bit set may be smaller then the original.
--
{-# INLINE map #-}
map :: (Enum a, Enum b, Bits c, Num c) => (a -> b) -> BitSet c a -> BitSet c b
map :: (a -> b) -> BitSet c a -> BitSet c b
map a -> b
f = (BitSet c b -> a -> BitSet c b)
-> BitSet c b -> BitSet c a -> BitSet c b
forall a c b.
(Enum a, Bits c) =>
(b -> a -> b) -> b -> BitSet c a -> b
foldl' (\BitSet c b
bs a
a -> a -> b
f a
a b -> BitSet c b -> BitSet c b
forall a c. (Enum a, Bits c) => a -> BitSet c a -> BitSet c a
`insert` BitSet c b
bs) BitSet c b
forall a c. (Enum a, Bits c, Num c) => BitSet c a
empty

-- | Reduce this bit set by applying a binary function to all elements,
-- using the given starting value. Each application of the operator is
-- evaluated before before using the result in the next application. This
-- function is strict in the starting value.
--
{-# INLINE foldl' #-}
foldl' :: (Enum a, Bits c) => (b -> a -> b) -> b -> BitSet c a -> b
foldl' :: (b -> a -> b) -> b -> BitSet c a -> b
foldl' b -> a -> b
f b
z (BitSet c
bits) = b -> Int -> Int -> b
go b
z (c -> Int
forall a. Bits a => a -> Int
popCount c
bits) Int
0
  where
    go :: b -> Int -> Int -> b
go !b
acc Int
0  !Int
_ = b
acc
    go !b
acc !Int
n !Int
b = if c
bits c -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` Int
b
                      then b -> Int -> Int -> b
go (b -> a -> b
f b
acc (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ Int -> a
forall a. Enum a => Int -> a
toEnum Int
b) (Int -> Int
forall a. Enum a => a -> a
pred Int
n) (Int -> Int
forall a. Enum a => a -> a
succ Int
b)
                      else b -> Int -> Int -> b
go b
acc Int
n (Int -> Int
forall a. Enum a => a -> a
succ Int
b)

-- | Reduce this bit set by applying a binary function to all elements,
-- using the given starting value.
--
{-# INLINE foldr #-}
foldr :: (Enum a, Bits c) => (a -> b -> b) -> b -> BitSet c a -> b
foldr :: (a -> b -> b) -> b -> BitSet c a -> b
foldr a -> b -> b
f b
z (BitSet c
bits) = Int -> Int -> b
go (c -> Int
forall a. Bits a => a -> Int
popCount c
bits) Int
0
  where
    go :: Int -> Int -> b
go Int
0  !Int
_ = b
z
    go !Int
n !Int
b = if c
bits c -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` Int
b
                 then Int -> a
forall a. Enum a => Int -> a
toEnum Int
b a -> b -> b
`f` Int -> Int -> b
go (Int -> Int
forall a. Enum a => a -> a
pred Int
n) (Int -> Int
forall a. Enum a => a -> a
succ Int
b)
                 else Int -> Int -> b
go Int
n (Int -> Int
forall a. Enum a => a -> a
succ Int
b)

-- | Convert this bit set set to a list of elements.
--
{-# INLINE [0] toList #-}
toList :: (Enum a, Bits c, Num c) => BitSet c a -> [a]
toList :: BitSet c a -> [a]
toList BitSet c a
bs = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
k b
z -> (a -> b -> b) -> b -> BitSet c a -> b
forall a c b.
(Enum a, Bits c) =>
(a -> b -> b) -> b -> BitSet c a -> b
foldr a -> b -> b
k b
z BitSet c a
bs)

-- | Make a bit set from a list of elements.
--
{-# INLINE [0] fromList #-}
fromList :: (Enum a, Bits c, Num c) => [a] -> BitSet c a
fromList :: [a] -> BitSet c a
fromList [a]
xs = c -> BitSet c a
forall c a. c -> BitSet c a
BitSet (c -> BitSet c a) -> c -> BitSet c a
forall a b. (a -> b) -> a -> b
$! (c -> a -> c) -> c -> [a] -> c
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\c
i a
x -> c
i c -> Int -> c
forall a. Bits a => a -> Int -> a
`setBit` a -> Int
forall a. Enum a => a -> Int
fromEnum a
x) c
0 [a]
xs

{-# RULES
"fromList/toList" forall bs. fromList (toList bs) = bs
  #-}