content: {-# LANGUAGE BangPatterns               #-}
  {-# LANGUAGE CPP                        #-}
  {-# LANGUAGE MagicHash                  #-}
  
  module Data.Bit.Utils
    ( lgWordSize
    , modWordSize
    , divWordSize
    , mulWordSize
    , wordSize
    , wordsToBytes
    , nWords
    , aligned
    , alignUp
    , selectWord
    , reverseWord
    , reversePartialWord
    , masked
    , meld
    , ffs
    , loMask
    , hiMask
    , sparseBits
    , fromPrimVector
    , toPrimVector
    ) where
  
  import Data.Bits
  import qualified Data.Vector.Primitive as P
  import qualified Data.Vector.Unboxed as U
  import qualified Data.Vector.Unboxed.Base as UB
  #if __GLASGOW_HASKELL__ >= 810
  import GHC.Exts
  #endif
  
  import Data.Bit.PdepPext
  
  -- | The number of bits in a 'Word'.  A handy constant to have around when defining 'Word'-based bulk operations on bit vectors.
  wordSize :: Int
  wordSize = finiteBitSize (0 :: Word)
  
  -- | The base 2 logarithm of 'wordSize'.
  lgWordSize :: Int
  lgWordSize = case wordSize of
    32 -> 5
    64 -> 6
    _  -> error "lgWordSize: unknown architecture"
  
  wordSizeMask :: Int
  wordSizeMask = wordSize - 1
  
  wordSizeMaskC :: Int
  wordSizeMaskC = complement wordSizeMask
  
  divWordSize :: Bits a => a -> a
  divWordSize x = unsafeShiftR x lgWordSize
  {-# INLINE divWordSize #-}
  
  modWordSize :: Int -> Int
  modWordSize x = x .&. (wordSize - 1)
  {-# INLINE modWordSize #-}
  
  mulWordSize :: Bits a => a -> a
  mulWordSize x = unsafeShiftL x lgWordSize
  {-# INLINE mulWordSize #-}
  
  -- number of words needed to store n bits
  nWords :: Int -> Int
  nWords ns = divWordSize (ns + wordSize - 1)
  
  wordsToBytes :: Int -> Int
  wordsToBytes ns = case wordSize of
    32 -> ns `unsafeShiftL` 2
    64 -> ns `unsafeShiftL` 3
    _  -> error "wordsToBytes: unknown architecture"
  
  aligned :: Int -> Bool
  aligned x = x .&. wordSizeMask == 0
  
  -- round a number of bits up to the nearest multiple of word size
  alignUp :: Int -> Int
  alignUp x | x == x'   = x'
            | otherwise = x' + wordSize
    where x' = alignDown x
  
  -- round a number of bits down to the nearest multiple of word size
  alignDown :: Int -> Int
  alignDown x = x .&. wordSizeMaskC
  
  -- create a mask consisting of the lower n bits
  mask :: Int -> Word
  mask b
    | b >= wordSize = complement 0
    | b < 0         = 0
    | otherwise     = bit b - 1
  
  masked :: Int -> Word -> Word
  masked b x = x .&. mask b
  
  -- meld 2 words by taking the low 'b' bits from 'lo' and the rest from 'hi'
  meld :: Int -> Word -> Word -> Word
  meld b lo hi = (lo .&. m) .|. (hi .&. complement m) where m = mask b
  {-# INLINE meld #-}
  
  #if __GLASGOW_HASKELL__ >= 810
  
  reverseWord :: Word -> Word
  reverseWord (W# w#) = W# (bitReverse# w#)
  
  #else
  
  reverseWord :: Word -> Word
  reverseWord = case wordSize of
    32 -> reverseWord32
    64 -> reverseWord64
    _  -> error "reverseWord: unknown architecture"
  
  reverseWord64 :: Word -> Word
  reverseWord64 x0 = x6
   where
    x1 = ((x0 .&. 0x5555555555555555) `shiftL`  1) .|. ((x0 .&. 0xAAAAAAAAAAAAAAAA) `shiftR`  1)
    x2 = ((x1 .&. 0x3333333333333333) `shiftL`  2) .|. ((x1 .&. 0xCCCCCCCCCCCCCCCC) `shiftR`  2)
    x3 = ((x2 .&. 0x0F0F0F0F0F0F0F0F) `shiftL`  4) .|. ((x2 .&. 0xF0F0F0F0F0F0F0F0) `shiftR`  4)
    x4 = ((x3 .&. 0x00FF00FF00FF00FF) `shiftL`  8) .|. ((x3 .&. 0xFF00FF00FF00FF00) `shiftR`  8)
    x5 = ((x4 .&. 0x0000FFFF0000FFFF) `shiftL` 16) .|. ((x4 .&. 0xFFFF0000FFFF0000) `shiftR` 16)
    x6 = ((x5 .&. 0x00000000FFFFFFFF) `shiftL` 32) .|. ((x5 .&. 0xFFFFFFFF00000000) `shiftR` 32)
  
  reverseWord32 :: Word -> Word
  reverseWord32 x0 = x5
   where
    x1 = ((x0 .&. 0x55555555) `shiftL`  1) .|. ((x0 .&. 0xAAAAAAAA) `shiftR`  1)
    x2 = ((x1 .&. 0x33333333) `shiftL`  2) .|. ((x1 .&. 0xCCCCCCCC) `shiftR`  2)
    x3 = ((x2 .&. 0x0F0F0F0F) `shiftL`  4) .|. ((x2 .&. 0xF0F0F0F0) `shiftR`  4)
    x4 = ((x3 .&. 0x00FF00FF) `shiftL`  8) .|. ((x3 .&. 0xFF00FF00) `shiftR`  8)
    x5 = ((x4 .&. 0x0000FFFF) `shiftL` 16) .|. ((x4 .&. 0xFFFF0000) `shiftR` 16)
  
  #endif
  
  reversePartialWord :: Int -> Word -> Word
  reversePartialWord n w
    | n >= wordSize = reverseWord w
    | otherwise     = reverseWord w `shiftR` (wordSize - n)
  
  ffs :: Word -> Maybe Int
  ffs 0 = Nothing
  ffs x = Just $! (popCount (x `xor` complement (-x)) - 1)
  {-# INLINE ffs #-}
  
  selectWord :: Word -> Word -> (Int, Word)
  selectWord msk src = (popCount msk, pext src msk)
  {-# INLINE selectWord #-}
  
  -- | Insert 0 between each consecutive bits of an input.
  -- xyzw --> (x0y0, z0w0)
  sparseBits :: Word -> (Word, Word)
  sparseBits = case wordSize of
    32 -> sparseBits32
    64 -> sparseBits64
    _  -> error "sparseBits: unknown architecture"
  
  sparseBits64 :: Word -> (Word, Word)
  sparseBits64 w = (x, y)
    where
      x = sparseBitsInternal64 (w .&. loMask 32)
      y = sparseBitsInternal64 (w `shiftR` 32)
  
  sparseBitsInternal64 :: Word -> Word
  sparseBitsInternal64 x = x4
    where
      t  = (x  `xor` (x  `shiftR` 16)) .&. 0x00000000ffff0000
      x0 = x  `xor` (t  `xor` (t  `shiftL` 16));
  
      t0 = (x0 `xor` (x0 `shiftR` 8)) .&. 0x0000ff000000ff00;
      x1 = x0 `xor` (t0 `xor` (t0 `shiftL` 8));
      t1 = (x1 `xor` (x1 `shiftR` 4)) .&. 0x00f000f000f000f0;
      x2 = x1 `xor` (t1 `xor` (t1 `shiftL` 4));
      t2 = (x2 `xor` (x2 `shiftR` 2)) .&. 0x0c0c0c0c0c0c0c0c;
      x3 = x2 `xor` (t2 `xor` (t2 `shiftL` 2));
      t3 = (x3 `xor` (x3 `shiftR` 1)) .&. 0x2222222222222222;
      x4 = x3 `xor` (t3 `xor` (t3 `shiftL` 1));
  
  sparseBits32 :: Word -> (Word, Word)
  sparseBits32 w = (x, y)
    where
      x = sparseBitsInternal32 (w .&. loMask 16)
      y = sparseBitsInternal32 (w `shiftR` 16)
  
  sparseBitsInternal32 :: Word -> Word
  sparseBitsInternal32 x0 = x4
    where
      t0 = (x0 `xor` (x0 `shiftR` 8)) .&. 0x0000ff00;
      x1 = x0 `xor` (t0 `xor` (t0 `shiftL` 8));
      t1 = (x1 `xor` (x1 `shiftR` 4)) .&. 0x00f000f0;
      x2 = x1 `xor` (t1 `xor` (t1 `shiftL` 4));
      t2 = (x2 `xor` (x2 `shiftR` 2)) .&. 0x0c0c0c0c;
      x3 = x2 `xor` (t2 `xor` (t2 `shiftL` 2));
      t3 = (x3 `xor` (x3 `shiftR` 1)) .&. 0x22222222;
      x4 = x3 `xor` (t3 `xor` (t3 `shiftL` 1));
  
  loMask :: Int -> Word
  loMask n = 1 `unsafeShiftL` n - 1
  {-# INLINE loMask #-}
  
  hiMask :: Int -> Word
  hiMask n = complement (1 `unsafeShiftL` n - 1)
  {-# INLINE hiMask #-}
  
  fromPrimVector :: P.Vector Word -> U.Vector Word
  fromPrimVector = UB.V_Word
  {-# INLINE fromPrimVector #-}
  
  toPrimVector :: U.Vector Word -> P.Vector Word
  toPrimVector (UB.V_Word ws) = ws
  {-# INLINE toPrimVector #-}
  