{- |
Similar to Data.Edison.Coll.EnumSet
but it allows to choose the underlying type for bit storage.
This is really a low-level module for type-safe foreign function interfaces.
-}
module Data.EnumSet (
   T(Cons, decons),
   fromEnum,
   fromEnums,
   toEnums,
   intToEnums,
   empty,
   (.&.),
   (.|.),
   xor,
   unions,
   get,
   put,
   set,
   clear,
   flip,
   fromBool,
   ) where

import qualified Data.Bits as B
import Data.Bits (Bits, )

import qualified Foreign.Storable.Newtype as Store
import Foreign.Storable (Storable(..), )

import qualified Prelude as P
import Prelude hiding (fromEnum, toEnum, flip, )


newtype T word enum = Cons {decons :: word}
   deriving (Eq)

instance (Storable word, Enum enum) => Storable (T word enum) where
   sizeOf = Store.sizeOf decons
   alignment = Store.alignment decons
   peek = Store.peek Cons
   poke = Store.poke decons


fromEnum :: (Enum a, Bits w) => a -> T w a
fromEnum = Cons . B.bit . P.fromEnum

fromEnums :: (Enum a, Bits w) => [a] -> T w a
fromEnums = Cons . foldl B.setBit 0 . map P.fromEnum

toEnums :: (Enum a, Bits w) => T w a -> [a]
toEnums =
   map fst . filter (P.flip B.testBit 0 . snd) .
   zip [P.toEnum 0 ..] . iterate (P.flip B.shiftR 1) .
   decons

intToEnums :: (Enum a, Integral w) => T w a -> [a]
intToEnums =
   map fst . filter (odd . snd) .
   zip [P.toEnum 0 ..] . iterate (P.flip div 2) .
   decons


empty :: (Enum a, Bits w) => T w a
empty = Cons 0


{-# INLINE lift2 #-}
lift2 :: (w -> w -> w) -> (T w a -> T w a -> T w a)
lift2 f (Cons x) (Cons y) = Cons (f x y)

-- fixities like in Data.Bits
infixl 7 .&.
infixl 5 .|.

(.&.), (.|.), xor :: (Enum a, Bits w) => T w a -> T w a -> T w a
(.&.) = lift2 (B..&.)
(.|.) = lift2 (B..|.)
xor   = lift2 B.xor

unions :: (Enum a, Bits w) => [T w a] -> T w a
unions = foldl (.|.) empty


get :: (Enum a, Bits w) => a -> T w a -> Bool
get n = P.flip B.testBit (P.fromEnum n) . decons

put :: (Enum a, Bits w) => a -> Bool -> T w a -> T w a
put n b s =
   fromBool n b .|. clear n s


{-# INLINE lift1 #-}
lift1 ::
   (Enum a, Bits w) =>
   (w -> Int -> w) -> (a -> T w a -> T w a)
lift1 f n (Cons vec) = Cons (f vec (P.fromEnum n))

set :: (Enum a, Bits w) => a -> T w a -> T w a
set = lift1 B.setBit

clear :: (Enum a, Bits w) => a -> T w a -> T w a
clear = lift1 B.clearBit

flip :: (Enum a, Bits w) => a -> T w a -> T w a
flip = lift1 B.complementBit

fromBool :: (Enum a, Bits w) => a -> Bool -> T w a
fromBool n b =
   Cons (B.shiftL (fromIntegral $ P.fromEnum b) (P.fromEnum n))