{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- | A tiny wrapper around 'IntSet.IntSet' for representing sets of 'Enum'
-- things.
module GHC.Data.EnumSet
    ( EnumSet
    , member
    , insert
    , delete
    , toList
    , fromList
    , empty
    , difference
    ) where

import GHC.Prelude
import GHC.Utils.Binary

import qualified Data.IntSet as IntSet

newtype EnumSet a = EnumSet IntSet.IntSet
  deriving (NonEmpty (EnumSet a) -> EnumSet a
EnumSet a -> EnumSet a -> EnumSet a
(EnumSet a -> EnumSet a -> EnumSet a)
-> (NonEmpty (EnumSet a) -> EnumSet a)
-> (forall b. Integral b => b -> EnumSet a -> EnumSet a)
-> Semigroup (EnumSet a)
forall b. Integral b => b -> EnumSet a -> EnumSet a
forall a. NonEmpty (EnumSet a) -> EnumSet a
forall a. EnumSet a -> EnumSet a -> EnumSet a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
forall a b. Integral b => b -> EnumSet a -> EnumSet a
$c<> :: forall a. EnumSet a -> EnumSet a -> EnumSet a
<> :: EnumSet a -> EnumSet a -> EnumSet a
$csconcat :: forall a. NonEmpty (EnumSet a) -> EnumSet a
sconcat :: NonEmpty (EnumSet a) -> EnumSet a
$cstimes :: forall a b. Integral b => b -> EnumSet a -> EnumSet a
stimes :: forall b. Integral b => b -> EnumSet a -> EnumSet a
Semigroup, Semigroup (EnumSet a)
EnumSet a
Semigroup (EnumSet a)
-> EnumSet a
-> (EnumSet a -> EnumSet a -> EnumSet a)
-> ([EnumSet a] -> EnumSet a)
-> Monoid (EnumSet a)
[EnumSet a] -> EnumSet a
EnumSet a -> EnumSet a -> EnumSet a
forall a. Semigroup (EnumSet a)
forall a. EnumSet a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall a. [EnumSet a] -> EnumSet a
forall a. EnumSet a -> EnumSet a -> EnumSet a
$cmempty :: forall a. EnumSet a
mempty :: EnumSet a
$cmappend :: forall a. EnumSet a -> EnumSet a -> EnumSet a
mappend :: EnumSet a -> EnumSet a -> EnumSet a
$cmconcat :: forall a. [EnumSet a] -> EnumSet a
mconcat :: [EnumSet a] -> EnumSet a
Monoid)

member :: Enum a => a -> EnumSet a -> Bool
member :: forall a. Enum a => a -> EnumSet a -> Bool
member a
x (EnumSet IntSet
s) = Key -> IntSet -> Bool
IntSet.member (a -> Key
forall a. Enum a => a -> Key
fromEnum a
x) IntSet
s

insert :: Enum a => a -> EnumSet a -> EnumSet a
insert :: forall a. Enum a => a -> EnumSet a -> EnumSet a
insert a
x (EnumSet IntSet
s) = IntSet -> EnumSet a
forall a. IntSet -> EnumSet a
EnumSet (IntSet -> EnumSet a) -> IntSet -> EnumSet a
forall a b. (a -> b) -> a -> b
$ Key -> IntSet -> IntSet
IntSet.insert (a -> Key
forall a. Enum a => a -> Key
fromEnum a
x) IntSet
s

delete :: Enum a => a -> EnumSet a -> EnumSet a
delete :: forall a. Enum a => a -> EnumSet a -> EnumSet a
delete a
x (EnumSet IntSet
s) = IntSet -> EnumSet a
forall a. IntSet -> EnumSet a
EnumSet (IntSet -> EnumSet a) -> IntSet -> EnumSet a
forall a b. (a -> b) -> a -> b
$ Key -> IntSet -> IntSet
IntSet.delete (a -> Key
forall a. Enum a => a -> Key
fromEnum a
x) IntSet
s

toList :: Enum a => EnumSet a -> [a]
toList :: forall a. Enum a => EnumSet a -> [a]
toList (EnumSet IntSet
s) = (Key -> a) -> [Key] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Key -> a
forall a. Enum a => Key -> a
toEnum ([Key] -> [a]) -> [Key] -> [a]
forall a b. (a -> b) -> a -> b
$ IntSet -> [Key]
IntSet.toList IntSet
s

fromList :: Enum a => [a] -> EnumSet a
fromList :: forall a. Enum a => [a] -> EnumSet a
fromList = IntSet -> EnumSet a
forall a. IntSet -> EnumSet a
EnumSet (IntSet -> EnumSet a) -> ([a] -> IntSet) -> [a] -> EnumSet a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Key] -> IntSet
IntSet.fromList ([Key] -> IntSet) -> ([a] -> [Key]) -> [a] -> IntSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Key) -> [a] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
map a -> Key
forall a. Enum a => a -> Key
fromEnum

empty :: EnumSet a
empty :: forall a. EnumSet a
empty = IntSet -> EnumSet a
forall a. IntSet -> EnumSet a
EnumSet IntSet
IntSet.empty

difference :: EnumSet a -> EnumSet a -> EnumSet a
difference :: forall a. EnumSet a -> EnumSet a -> EnumSet a
difference (EnumSet IntSet
a) (EnumSet IntSet
b) = IntSet -> EnumSet a
forall a. IntSet -> EnumSet a
EnumSet (IntSet -> IntSet -> IntSet
IntSet.difference IntSet
a IntSet
b)

-- | Represents the 'EnumSet' as a bit set.
--
-- Assumes that all elements are non-negative.
--
-- This is only efficient for values that are sufficiently small,
-- for example in the lower hundreds.
instance Binary (EnumSet a) where
  put_ :: BinHandle -> EnumSet a -> IO ()
put_ BinHandle
bh = BinHandle -> BitArray -> IO ()
forall a. Binary a => BinHandle -> a -> IO ()
put_ BinHandle
bh (BitArray -> IO ())
-> (EnumSet a -> BitArray) -> EnumSet a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnumSet a -> BitArray
forall a. EnumSet a -> BitArray
enumSetToBitArray
  get :: BinHandle -> IO (EnumSet a)
get BinHandle
bh = BitArray -> EnumSet a
forall a. BitArray -> EnumSet a
bitArrayToEnumSet (BitArray -> EnumSet a) -> IO BitArray -> IO (EnumSet a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BinHandle -> IO BitArray
forall a. Binary a => BinHandle -> IO a
get BinHandle
bh

-- TODO: Using 'Natural' instead of 'Integer' should be slightly more efficient
-- but we don't currently have a 'Binary' instance for 'Natural'.
type BitArray = Integer

enumSetToBitArray :: EnumSet a -> BitArray
enumSetToBitArray :: forall a. EnumSet a -> BitArray
enumSetToBitArray (EnumSet IntSet
int_set) =
    (BitArray -> Key -> BitArray) -> BitArray -> IntSet -> BitArray
forall a. (a -> Key -> a) -> a -> IntSet -> a
IntSet.foldl' BitArray -> Key -> BitArray
forall a. Bits a => a -> Key -> a
setBit BitArray
0 IntSet
int_set

bitArrayToEnumSet :: BitArray -> EnumSet a
bitArrayToEnumSet :: forall a. BitArray -> EnumSet a
bitArrayToEnumSet BitArray
ba = IntSet -> EnumSet a
forall a. IntSet -> EnumSet a
EnumSet (Key -> Key -> IntSet -> IntSet
go (BitArray -> Key
forall a. Bits a => a -> Key
popCount BitArray
ba) Key
0 IntSet
IntSet.empty)
  where
    go :: Key -> Key -> IntSet -> IntSet
go Key
0 Key
_ !IntSet
int_set = IntSet
int_set
    go Key
n Key
i !IntSet
int_set =
      if BitArray
ba BitArray -> Key -> Bool
forall a. Bits a => a -> Key -> Bool
`testBit` Key
i
        then Key -> Key -> IntSet -> IntSet
go (Key -> Key
forall a. Enum a => a -> a
pred Key
n) (Key -> Key
forall a. Enum a => a -> a
succ Key
i) (Key -> IntSet -> IntSet
IntSet.insert Key
i IntSet
int_set)
        else Key -> Key -> IntSet -> IntSet
go Key
n        (Key -> Key
forall a. Enum a => a -> a
succ Key
i) IntSet
int_set