diff options
Diffstat (limited to 'compiler/GHC/Data/EnumSet.hs')
| -rw-r--r-- | compiler/GHC/Data/EnumSet.hs | 31 |
1 files changed, 31 insertions, 0 deletions
diff --git a/compiler/GHC/Data/EnumSet.hs b/compiler/GHC/Data/EnumSet.hs index a7949c7e71..d364dc05d7 100644 --- a/compiler/GHC/Data/EnumSet.hs +++ b/compiler/GHC/Data/EnumSet.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} -- | A tiny wrapper around 'IntSet.IntSet' for representing sets of 'Enum' -- things. module GHC.Data.EnumSet @@ -12,10 +14,12 @@ module GHC.Data.EnumSet ) where import GHC.Prelude +import GHC.Utils.Binary import qualified Data.IntSet as IntSet newtype EnumSet a = EnumSet IntSet.IntSet + deriving (Semigroup, Monoid) member :: Enum a => a -> EnumSet a -> Bool member x (EnumSet s) = IntSet.member (fromEnum x) s @@ -37,3 +41,30 @@ empty = EnumSet IntSet.empty difference :: EnumSet a -> EnumSet a -> EnumSet a difference (EnumSet a) (EnumSet b) = EnumSet (IntSet.difference a 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_ bh = put_ bh . enumSetToBitArray + get bh = bitArrayToEnumSet <$> get 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 (EnumSet int_set) = + IntSet.foldl' setBit 0 int_set + +bitArrayToEnumSet :: BitArray -> EnumSet a +bitArrayToEnumSet ba = EnumSet (go (popCount ba) 0 IntSet.empty) + where + go 0 _ !int_set = int_set + go n i !int_set = + if ba `testBit` i + then go (pred n) (succ i) (IntSet.insert i int_set) + else go n (succ i) int_set |