{-# LANGUAGE BangPatterns, DeriveGeneric #-}

-- | Inspired in the @Data.IntSet@ API, a similar API where the elements
--   of the set are bytes (values of type 'Word8').
module Data.ByteSet (
    -- * Types
    ByteSet
  , Word8
    -- * Query
  , null
  , size
  , member
  , notMember
    -- * Construction
  , empty
  , singleton
  , insert
  , delete
    -- * Combine
  , union
  , unions
  , difference
  , intersection
    -- * Filter
  , filter
    -- * Map
  , map
    -- * Folds
  , foldr
    -- * List conversion
  , elems
  , toList
  , fromList
  ) where

import Prelude
  ( Eq (..), Ord (..)
  , (+), (-), ($), (*)
  , fromIntegral
  , Show (..)
    )
import Data.Word (Word8, Word64)
import Data.Int (Int)
import Data.Bits
import Data.Bool
import Control.Category
import Control.Applicative (Applicative (..))
import qualified Data.Foldable as F
import qualified Data.List as L
import Data.Functor
import GHC.Generics (Generic)
import Data.Binary (Binary (..))
import Data.Binary.Put (putWord64le)
import Data.Binary.Get (getWord64le)

-- | Set of bytes ('Word8'). Note that NF and WHNF are equivalent
--   for values of type 'ByteSet'.
data ByteSet = ByteSet
  {-# UNPACK #-} !Word64
  {-# UNPACK #-} !Word64
  {-# UNPACK #-} !Word64
  {-# UNPACK #-} !Word64
  deriving (ByteSet -> ByteSet -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ByteSet -> ByteSet -> Bool
$c/= :: ByteSet -> ByteSet -> Bool
== :: ByteSet -> ByteSet -> Bool
$c== :: ByteSet -> ByteSet -> Bool
Eq, Eq ByteSet
ByteSet -> ByteSet -> Bool
ByteSet -> ByteSet -> Ordering
ByteSet -> ByteSet -> ByteSet
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
min :: ByteSet -> ByteSet -> ByteSet
$cmin :: ByteSet -> ByteSet -> ByteSet
max :: ByteSet -> ByteSet -> ByteSet
$cmax :: ByteSet -> ByteSet -> ByteSet
>= :: ByteSet -> ByteSet -> Bool
$c>= :: ByteSet -> ByteSet -> Bool
> :: ByteSet -> ByteSet -> Bool
$c> :: ByteSet -> ByteSet -> Bool
<= :: ByteSet -> ByteSet -> Bool
$c<= :: ByteSet -> ByteSet -> Bool
< :: ByteSet -> ByteSet -> Bool
$c< :: ByteSet -> ByteSet -> Bool
compare :: ByteSet -> ByteSet -> Ordering
$ccompare :: ByteSet -> ByteSet -> Ordering
Ord, forall x. Rep ByteSet x -> ByteSet
forall x. ByteSet -> Rep ByteSet x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep ByteSet x -> ByteSet
$cfrom :: forall x. ByteSet -> Rep ByteSet x
Generic)

generalGet :: (Word64 -> Int -> a) -> Word8 -> ByteSet -> a
{-# INLINE generalGet #-}
generalGet :: forall a. (Word64 -> Int -> a) -> Word8 -> ByteSet -> a
generalGet Word64 -> Int -> a
f Word8
w (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4)
  | Word8
w forall a. Ord a => a -> a -> Bool
<  Word8
64   = Word64 -> Int -> a
f Word64
s1  Int
i
  | Word8
w forall a. Ord a => a -> a -> Bool
< Word8
128   = Word64 -> Int -> a
f Word64
s2 (Int
i forall a. Num a => a -> a -> a
-  Int
64)
  | Word8
w forall a. Ord a => a -> a -> Bool
< Word8
192   = Word64 -> Int -> a
f Word64
s3 (Int
i forall a. Num a => a -> a -> a
- Int
128)
  | Bool
otherwise = Word64 -> Int -> a
f Word64
s4 (Int
i forall a. Num a => a -> a -> a
- Int
192)
  where
    i :: Int
i = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w

generalSet :: (Word64 -> Int -> Word64) -> Word8 -> ByteSet -> ByteSet
{-# INLINE generalSet #-}
generalSet :: (Word64 -> Int -> Word64) -> Word8 -> ByteSet -> ByteSet
generalSet Word64 -> Int -> Word64
f Word8
w (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4)
  | Word8
w forall a. Ord a => a -> a -> Bool
<  Word8
64   = Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet          (Word64 -> Int -> Word64
f Word64
s1   Int
i      ) Word64
s2 Word64
s3 Word64
s4
  | Word8
w forall a. Ord a => a -> a -> Bool
< Word8
128   = Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet Word64
s1       (Word64 -> Int -> Word64
f Word64
s2 forall a b. (a -> b) -> a -> b
$ Int
i forall a. Num a => a -> a -> a
-  Int
64) Word64
s3 Word64
s4
  | Word8
w forall a. Ord a => a -> a -> Bool
< Word8
192   = Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet Word64
s1 Word64
s2    (Word64 -> Int -> Word64
f Word64
s3 forall a b. (a -> b) -> a -> b
$ Int
i forall a. Num a => a -> a -> a
- Int
128) Word64
s4
  | Bool
otherwise = Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet Word64
s1 Word64
s2 Word64
s3 (Word64 -> Int -> Word64
f Word64
s4 forall a b. (a -> b) -> a -> b
$ Int
i forall a. Num a => a -> a -> a
- Int
192)
  where
    i :: Int
i = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w

generalOp :: (Word64  -> Word64  -> Word64 )
          ->  ByteSet -> ByteSet -> ByteSet
{-# INLINE generalOp #-}
generalOp :: (Word64 -> Word64 -> Word64) -> ByteSet -> ByteSet -> ByteSet
generalOp Word64 -> Word64 -> Word64
f (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4)
            (ByteSet Word64
t1 Word64
t2 Word64
t3 Word64
t4) =
  Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet (Word64 -> Word64 -> Word64
f Word64
s1 Word64
t1) (Word64 -> Word64 -> Word64
f Word64
s2 Word64
t2)
          (Word64 -> Word64 -> Word64
f Word64
s3 Word64
t3) (Word64 -> Word64 -> Word64
f Word64
s4 Word64
t4)

generalFun :: (Int -> Word64 -> Word64) -> ByteSet -> ByteSet -- tons of fun!
{-# INLINE generalFun #-}
generalFun :: (Int -> Word64 -> Word64) -> ByteSet -> ByteSet
generalFun Int -> Word64 -> Word64
f (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4) =
  Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet (Int -> Word64 -> Word64
f Int
0 Word64
s1) (Int -> Word64 -> Word64
f Int
1 Word64
s2) (Int -> Word64 -> Word64
f Int
2 Word64
s3) (Int -> Word64 -> Word64
f Int
3 Word64
s4)

----------------------------------------------------------------------
----------------------------------------------------------------------
-- API

-- | /O(1)/. Cardinality of the byteset.
size :: ByteSet -> Int
size :: ByteSet -> Int
size (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4) =
    forall a. Bits a => a -> Int
popCount Word64
s1 forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount Word64
s2
  forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount Word64
s3 forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount Word64
s4

-- | /O(1)/. Is the byteset empty?
null :: ByteSet -> Bool
null :: ByteSet -> Bool
null = (forall a. Eq a => a -> a -> Bool
==Int
0) forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ByteSet -> Int
size

-- | /O(1)/. Is the value a member of the byteset?
member :: Word8 -> ByteSet -> Bool
member :: Word8 -> ByteSet -> Bool
member = forall a. (Word64 -> Int -> a) -> Word8 -> ByteSet -> a
generalGet forall a. Bits a => a -> Int -> Bool
testBit

-- | /O(1)/. Is the element not in the set?
notMember :: Word8 -> ByteSet -> Bool
notMember :: Word8 -> ByteSet -> Bool
notMember Word8
w = Bool -> Bool
not forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word8 -> ByteSet -> Bool
member Word8
w

-- | /O(1)/. The empty byteset.
empty :: ByteSet
empty :: ByteSet
empty = Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet Word64
0 Word64
0 Word64
0 Word64
0

-- | /O(1)/. Add a value to the byteset.
insert :: Word8 -> ByteSet -> ByteSet
insert :: Word8 -> ByteSet -> ByteSet
insert = (Word64 -> Int -> Word64) -> Word8 -> ByteSet -> ByteSet
generalSet forall a. Bits a => a -> Int -> a
setBit

-- | /O(1)/. A byteset of one element.
singleton :: Word8 -> ByteSet
singleton :: Word8 -> ByteSet
singleton Word8
w = Word8 -> ByteSet -> ByteSet
insert Word8
w ByteSet
empty

-- | /O(1)/. Delete a byte in the byteset. Returns the original byteset when the byte was not present.
delete :: Word8 -> ByteSet -> ByteSet
delete :: Word8 -> ByteSet -> ByteSet
delete = (Word64 -> Int -> Word64) -> Word8 -> ByteSet -> ByteSet
generalSet forall a. Bits a => a -> Int -> a
clearBit

-- | /O(1)/. The union of two bytesets.
union :: ByteSet -> ByteSet -> ByteSet
union :: ByteSet -> ByteSet -> ByteSet
union = (Word64 -> Word64 -> Word64) -> ByteSet -> ByteSet -> ByteSet
generalOp forall a. Bits a => a -> a -> a
(.|.)

-- | The union of a list of bytesets. Just a fold over the list using 'union'.
unions :: [ByteSet] -> ByteSet
unions :: [ByteSet] -> ByteSet
unions = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' ByteSet -> ByteSet -> ByteSet
union ByteSet
empty

-- | /O(1)/. Difference between two bytesets.
difference :: ByteSet -> ByteSet -> ByteSet
difference :: ByteSet -> ByteSet -> ByteSet
difference = (Word64 -> Word64 -> Word64) -> ByteSet -> ByteSet -> ByteSet
generalOp forall a b. (a -> b) -> a -> b
$ \Word64
w1 Word64
w2 -> Word64
w1 forall a. Bits a => a -> a -> a
.&. forall a. Bits a => a -> a
complement Word64
w2

-- | /O(1)/. The intersection of two bytesets.
intersection :: ByteSet -> ByteSet -> ByteSet
intersection :: ByteSet -> ByteSet -> ByteSet
intersection = (Word64 -> Word64 -> Word64) -> ByteSet -> ByteSet -> ByteSet
generalOp forall a. Bits a => a -> a -> a
(.&.)

-- | /O(n)/. Filter all elements that satisfy some predicate.
filter :: (Word8 -> Bool) -> ByteSet -> ByteSet
filter :: (Word8 -> Bool) -> ByteSet -> ByteSet
filter Word8 -> Bool
f = (Int -> Word64 -> Word64) -> ByteSet -> ByteSet
generalFun forall a b. (a -> b) -> a -> b
$ \Int
i Word64
w ->
  let b0 :: Int
b0 = Int
i forall a. Num a => a -> a -> a
* Int
64
      go :: t -> Int -> t
go t
acc (-1) = t
acc
      go !t
acc Int
n = if forall a. Bits a => a -> Int -> Bool
testBit Word64
w Int
n Bool -> Bool -> Bool
&& Word8 -> Bool
f (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Int
b0 forall a. Num a => a -> a -> a
+ Int
n)
                     then t -> Int -> t
go (forall a. Bits a => a -> Int -> a
setBit t
acc Int
n) forall a b. (a -> b) -> a -> b
$ Int
n forall a. Num a => a -> a -> a
- Int
1
                     else t -> Int -> t
go         t
acc    forall a b. (a -> b) -> a -> b
$ Int
n forall a. Num a => a -> a -> a
- Int
1
  in  forall a. Bits a => a -> Int -> a
go Word64
0 Int
63

bits :: Word64 -> [Int]
bits :: Word64 -> [Int]
bits Word64
w = forall a. (a -> Bool) -> [a] -> [a]
L.filter (forall a. Bits a => a -> Int -> Bool
testBit Word64
w) [Int
0..Int
63]

-- | /O(n)/. Fold the elements in the byteset using the given right-associative binary operator.
foldr :: (Word8 -> a -> a) -> a -> ByteSet -> a
foldr :: forall a. (Word8 -> a -> a) -> a -> ByteSet -> a
foldr Word8 -> a -> a
f a
r0 (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4) =
  let g :: Int -> a -> a
g  = Word8 -> a -> a
f forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral
      r1 :: a
r1 = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr Int -> a -> a
g a
r0 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Num a => a -> a -> a
+Int
192) forall a b. (a -> b) -> a -> b
$ Word64 -> [Int]
bits Word64
s4
      r2 :: a
r2 = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr Int -> a -> a
g a
r1 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Num a => a -> a -> a
+Int
128) forall a b. (a -> b) -> a -> b
$ Word64 -> [Int]
bits Word64
s3
      r3 :: a
r3 = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr Int -> a -> a
g a
r2 forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Num a => a -> a -> a
+ Int
64) forall a b. (a -> b) -> a -> b
$ Word64 -> [Int]
bits Word64
s2
  in       forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr Int -> a -> a
g a
r3 forall a b. (a -> b) -> a -> b
$               Word64 -> [Int]
bits Word64
s1

-- | /O(n)/. Map a function over a byteset.
map :: (Word8 -> Word8) -> ByteSet -> ByteSet
map :: (Word8 -> Word8) -> ByteSet -> ByteSet
map Word8 -> Word8
f = forall a. (Word8 -> a -> a) -> a -> ByteSet -> a
foldr (Word8 -> ByteSet -> ByteSet
insert forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word8 -> Word8
f) ByteSet
empty

-- | /O(n)/. The elements of a byteset in ascending order.
elems :: ByteSet -> [Word8]
elems :: ByteSet -> [Word8]
elems = forall a. (Word8 -> a -> a) -> a -> ByteSet -> a
foldr (:) []

-- | /O(n)/. An alias of 'elems'.
toList :: ByteSet -> [Word8]
{-# INLINE toList #-}
toList :: ByteSet -> [Word8]
toList = ByteSet -> [Word8]
elems

-- | /O(n)/. Create a byteset from a list of bytes.
fromList :: [Word8] -> ByteSet
fromList :: [Word8] -> ByteSet
fromList = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr Word8 -> ByteSet -> ByteSet
insert ByteSet
empty

instance Show ByteSet where
  show :: ByteSet -> String
show = forall a. Show a => a -> String
show forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ByteSet -> [Word8]
elems

instance Binary ByteSet where
  put :: ByteSet -> Put
put (ByteSet Word64
s1 Word64
s2 Word64
s3 Word64
s4) =
       Word64 -> Put
putWord64le Word64
s1 forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Word64 -> Put
putWord64le Word64
s2
    forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Word64 -> Put
putWord64le Word64
s3 forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Word64 -> Put
putWord64le Word64
s4
  get :: Get ByteSet
get = Word64 -> Word64 -> Word64 -> Word64 -> ByteSet
ByteSet forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word64
getWord64le forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word64
getWord64le
                forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word64
getWord64le forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word64
getWord64le