------------------------------------------------------------------------
-- |
-- Module           : Lang.Crucible.Utils.BitSet
-- Description      : Encode a set of enumerable elements using the bit-positions
--                    in an Integer
-- Copyright        : (c) Galois, Inc 2015-2016
-- License          : BSD3
-- Maintainer       : Joe Hendrix <jhendrix@galois.com>
-- Stability        : provisional
--
-- This module provides a simple bitset datastructure
-- built on top of GHC-native Integers.
------------------------------------------------------------------------
module Lang.Crucible.Utils.BitSet
( BitSet
, getBits
, empty
, null
, singleton
, insert
, remove
, size
, member
, isSubsetOf
, difference
, intersection
, union
, toList
, foldr
, foldl
, foldl'
) where

import Data.Bits
import Data.Word
import Data.Hashable
import qualified Data.List as List
import Prelude hiding (null, foldr, foldl)

newtype BitSet a = BitSet { forall a. BitSet a -> Integer
getBits :: Integer }
 deriving (Int -> BitSet a -> ShowS
[BitSet a] -> ShowS
BitSet a -> String
(Int -> BitSet a -> ShowS)
-> (BitSet a -> String) -> ([BitSet a] -> ShowS) -> Show (BitSet a)
forall a. Int -> BitSet a -> ShowS
forall a. [BitSet a] -> ShowS
forall a. BitSet a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Int -> BitSet a -> ShowS
showsPrec :: Int -> BitSet a -> ShowS
$cshow :: forall a. BitSet a -> String
show :: BitSet a -> String
$cshowList :: forall a. [BitSet a] -> ShowS
showList :: [BitSet a] -> ShowS
Show, BitSet a -> BitSet a -> Bool
(BitSet a -> BitSet a -> Bool)
-> (BitSet a -> BitSet a -> Bool) -> Eq (BitSet a)
forall a. BitSet a -> BitSet a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. BitSet a -> BitSet a -> Bool
== :: BitSet a -> BitSet a -> Bool
$c/= :: forall a. BitSet a -> BitSet a -> Bool
/= :: BitSet a -> BitSet a -> Bool
Eq, Eq (BitSet a)
Eq (BitSet a) =>
(BitSet a -> BitSet a -> Ordering)
-> (BitSet a -> BitSet a -> Bool)
-> (BitSet a -> BitSet a -> Bool)
-> (BitSet a -> BitSet a -> Bool)
-> (BitSet a -> BitSet a -> Bool)
-> (BitSet a -> BitSet a -> BitSet a)
-> (BitSet a -> BitSet a -> BitSet a)
-> Ord (BitSet a)
BitSet a -> BitSet a -> Bool
BitSet a -> BitSet a -> Ordering
BitSet a -> BitSet a -> BitSet a
forall a. Eq (BitSet a)
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. BitSet a -> BitSet a -> Bool
forall a. BitSet a -> BitSet a -> Ordering
forall a. BitSet a -> BitSet a -> BitSet a
$ccompare :: forall a. BitSet a -> BitSet a -> Ordering
compare :: BitSet a -> BitSet a -> Ordering
$c< :: forall a. BitSet a -> BitSet a -> Bool
< :: BitSet a -> BitSet a -> Bool
$c<= :: forall a. BitSet a -> BitSet a -> Bool
<= :: BitSet a -> BitSet a -> Bool
$c> :: forall a. BitSet a -> BitSet a -> Bool
> :: BitSet a -> BitSet a -> Bool
$c>= :: forall a. BitSet a -> BitSet a -> Bool
>= :: BitSet a -> BitSet a -> Bool
$cmax :: forall a. BitSet a -> BitSet a -> BitSet a
max :: BitSet a -> BitSet a -> BitSet a
$cmin :: forall a. BitSet a -> BitSet a -> BitSet a
min :: BitSet a -> BitSet a -> BitSet a
Ord)

instance Hashable (BitSet a) where
  hashWithSalt :: Int -> BitSet a -> Int
hashWithSalt Int
s (BitSet Integer
x) = Int -> Integer -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
s Integer
x

empty :: BitSet a
empty :: forall a. BitSet a
empty = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet Integer
forall a. Bits a => a
zeroBits

null :: BitSet a -> Bool
null :: forall a. BitSet a -> Bool
null = (Integer
0Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
==) (Integer -> Bool) -> (BitSet a -> Integer) -> BitSet a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitSet a -> Integer
forall a. BitSet a -> Integer
getBits

singleton :: Enum a => a -> BitSet a
singleton :: forall a. Enum a => a -> BitSet a
singleton a
a = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet (Int -> Integer
forall a. Bits a => Int -> a
bit (a -> Int
forall a. Enum a => a -> Int
fromEnum a
a))

insert :: Enum a => a -> BitSet a -> BitSet a
insert :: forall a. Enum a => a -> BitSet a -> BitSet a
insert a
a (BitSet Integer
x) = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
setBit Integer
x (a -> Int
forall a. Enum a => a -> Int
fromEnum a
a))

remove :: Enum a => a -> BitSet a -> BitSet a
remove :: forall a. Enum a => a -> BitSet a -> BitSet a
remove a
a (BitSet Integer
x) = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
clearBit Integer
x (a -> Int
forall a. Enum a => a -> Int
fromEnum a
a))

union :: BitSet a -> BitSet a -> BitSet a
union :: forall a. BitSet a -> BitSet a -> BitSet a
union (BitSet Integer
x) (BitSet Integer
y) = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet (Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.|. Integer
y)

intersection :: BitSet a -> BitSet a -> BitSet a
intersection :: forall a. BitSet a -> BitSet a -> BitSet a
intersection (BitSet Integer
x) (BitSet Integer
y) = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet (Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
y)

difference :: BitSet a -> BitSet a -> BitSet a
difference :: forall a. BitSet a -> BitSet a -> BitSet a
difference (BitSet Integer
x) (BitSet Integer
y) = Integer -> BitSet a
forall a. Integer -> BitSet a
BitSet (Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer -> Integer
forall a. Bits a => a -> a
complement Integer
y)

isSubsetOf :: BitSet a -> BitSet a -> Bool
isSubsetOf :: forall a. BitSet a -> BitSet a -> Bool
isSubsetOf (BitSet Integer
x) (BitSet Integer
y) = Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.|. Integer
y Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
y

member :: Enum a => a -> BitSet a -> Bool
member :: forall a. Enum a => a -> BitSet a -> Bool
member a
a (BitSet Integer
x) = Integer -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Integer
x (a -> Int
forall a. Enum a => a -> Int
fromEnum a
a)

size :: BitSet a -> Int
size :: forall a. BitSet a -> Int
size (BitSet Integer
x) = Integer -> Int
forall a. Bits a => a -> Int
popCount Integer
x

toList :: Enum a => BitSet a -> [a]
toList :: forall a. Enum a => BitSet a -> [a]
toList (BitSet Integer
bs) = Integer -> Int -> [a]
forall a. Enum a => Integer -> Int -> [a]
go Integer
bs Int
0
  where go :: Enum a => Integer -> Int -> [a]
        go :: forall a. Enum a => Integer -> Int -> [a]
go Integer
0 Int
_ = []
        go Integer
x Int
i
           | Word32
y Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0xffffffff Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0 = Integer -> Int -> [a]
forall a. Enum a => Integer -> Int -> [a]
go (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
x Int
32) (Int -> [a]) -> Int -> [a]
forall a b. (a -> b) -> a -> b
$! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
32)
           | Word32
y Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0x0000ffff Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0 = Integer -> Int -> [a]
forall a. Enum a => Integer -> Int -> [a]
go (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
x Int
16) (Int -> [a]) -> Int -> [a]
forall a b. (a -> b) -> a -> b
$! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
16)
           | Word32
y Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0x000000ff Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0 = Integer -> Int -> [a]
forall a. Enum a => Integer -> Int -> [a]
go (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
x  Int
8) (Int -> [a]) -> Int -> [a]
forall a b. (a -> b) -> a -> b
$! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8)
           | Bool
otherwise = [[a]] -> [a]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat
               [ if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
0 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
0)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
1 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
2 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
3 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
3)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
4 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
4)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
5 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
5)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
6 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
6)] else []
               , if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
y Int
7 then [Int -> a
forall a. Enum a => Int -> a
toEnum (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
7)] else []
               , Integer -> Int -> [a]
forall a. Enum a => Integer -> Int -> [a]
go (Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
shiftR Integer
x Int
8) (Int -> [a]) -> Int -> [a]
forall a b. (a -> b) -> a -> b
$! (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
8)
               ]

          where y :: Word32
                y :: Word32
y = Integer -> Word32
forall a. Num a => Integer -> a
fromInteger Integer
x

foldl' :: Enum a => (b -> a -> b) -> b -> BitSet a -> b
foldl' :: forall a b. Enum a => (b -> a -> b) -> b -> BitSet a -> b
foldl' b -> a -> b
f b
z = (b -> a -> b) -> b -> [a] -> b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' b -> a -> b
f b
z ([a] -> b) -> (BitSet a -> [a]) -> BitSet a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitSet a -> [a]
forall a. Enum a => BitSet a -> [a]
toList

foldl :: Enum a => (b -> a -> b) -> b -> BitSet a -> b
foldl :: forall a b. Enum a => (b -> a -> b) -> b -> BitSet a -> b
foldl b -> a -> b
f b
z = (b -> a -> b) -> b -> [a] -> b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl b -> a -> b
f b
z ([a] -> b) -> (BitSet a -> [a]) -> BitSet a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitSet a -> [a]
forall a. Enum a => BitSet a -> [a]
toList

foldr :: Enum a => (a -> b -> b) -> b -> BitSet a -> b
foldr :: forall a b. Enum a => (a -> b -> b) -> b -> BitSet a -> b
foldr a -> b -> b
f b
z = (a -> b -> b) -> b -> [a] -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
List.foldr a -> b -> b
f b
z ([a] -> b) -> (BitSet a -> [a]) -> BitSet a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitSet a -> [a]
forall a. Enum a => BitSet a -> [a]
toList