content: {-# LANGUAGE CPP                        #-}
  
  {-# LANGUAGE BangPatterns               #-}
  {-# LANGUAGE DeriveDataTypeable         #-}
  {-# LANGUAGE DeriveGeneric              #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE LambdaCase                 #-}
  {-# LANGUAGE MagicHash                  #-}
  {-# LANGUAGE MultiParamTypeClasses      #-}
  {-# LANGUAGE TypeFamilies               #-}
  {-# LANGUAGE UnboxedTuples              #-}
  {-# LANGUAGE ViewPatterns               #-}
  
  #ifndef BITVEC_THREADSAFE
  module Data.Bit.Internal
  #else
  module Data.Bit.InternalTS
  #endif
    ( Bit(..)
    , U.Vector(BitVec)
    , U.MVector(BitMVec)
    , indexWord
    , readWord
    , writeWord
    , unsafeFlipBit
    , flipBit
    , modifyByteArray
    ) where
  
  #if MIN_VERSION_vector(0,13,0)
  import Data.Vector.Internal.Check (checkIndex, Checks(..))
  #else
  #include "vector.h"
  #endif
  
  import Control.DeepSeq
  import Control.Exception
  import Control.Monad.Primitive
  import Control.Monad.ST
  import Data.Bits
  import Data.Bit.Utils
  import Data.Primitive.ByteArray
  import Data.Ratio
  import qualified Data.Vector.Generic as V
  import qualified Data.Vector.Generic.Mutable as MV
  import qualified Data.Vector.Unboxed as U
  import GHC.Generics
  
  #ifdef BITVEC_THREADSAFE
  import GHC.Exts
  #endif
  
  #ifndef BITVEC_THREADSAFE
  -- | A newtype wrapper with a custom instance
  -- for "Data.Vector.Unboxed", which packs booleans
  -- as efficient as possible (8 values per byte).
  -- Unboxed vectors of t'Bit' use 8x less memory
  -- than unboxed vectors of 'Bool' (which store one value per byte),
  -- but random writes are slightly slower.
  --
  -- @since 0.1
  newtype Bit = Bit {
    unBit :: Bool -- ^ @since 0.2.0.0
    } deriving
    (Bounded, Enum, Eq, Ord
    , FiniteBits -- ^ @since 0.2.0.0
    , Bits
    , Generic    -- ^ @since 1.0.1.0
    , NFData     -- ^ @since 1.0.1.0
    )
  #else
  -- | A newtype wrapper with a custom instance
  -- for "Data.Vector.Unboxed", which packs booleans
  -- as efficient as possible (8 values per byte).
  -- Unboxed vectors of t'Bit' use 8x less memory
  -- than unboxed vectors of 'Bool' (which store one value per byte),
  -- but random writes are slightly slower.
  --
  -- @since 1.0.0.0
  newtype Bit = Bit {
    unBit :: Bool -- ^ @since 0.2.0.0
    } deriving
    (Bounded, Enum, Eq, Ord
    , FiniteBits -- ^ @since 0.2.0.0
    , Bits
    , Generic    -- ^ @since 1.0.1.0
    , NFData     -- ^ @since 1.0.1.0
    )
  #endif
  
  -- | There is only one lawful 'Num' instance possible
  -- with '+' = 'xor' and
  -- 'fromInteger' = v'Bit' . 'odd'.
  --
  -- @since 1.0.1.0
  instance Num Bit where
    Bit a * Bit b = Bit (a && b)
    Bit a + Bit b = Bit (a /= b)
    Bit a - Bit b = Bit (a /= b)
    negate = id
    abs    = id
    signum = id
    fromInteger = Bit . odd
  
  -- | @since 1.0.1.0
  instance Real Bit where
    toRational = fromIntegral
  
  -- | @since 1.0.1.0
  instance Integral Bit where
    quotRem _ (Bit False) = throw DivideByZero
    quotRem x (Bit True)  = (x, Bit False)
    toInteger (Bit False) = 0
    toInteger (Bit True)  = 1
  
  -- | @since 1.0.1.0
  instance Fractional Bit where
    fromRational x = fromInteger (numerator x) / fromInteger (denominator x)
    (/) = quot
  
  instance Show Bit where
    showsPrec _ (Bit False) = showString "0"
    showsPrec _ (Bit True ) = showString "1"
  
  instance Read Bit where
    readsPrec p (' ' : rest) = readsPrec p rest
    readsPrec _ ('0' : rest) = [(Bit False, rest)]
    readsPrec _ ('1' : rest) = [(Bit True, rest)]
    readsPrec _ _            = []
  
  instance U.Unbox Bit
  
  -- Ints are offset and length in bits
  data instance U.MVector s Bit = BitMVec !Int !Int !(MutableByteArray s)
  data instance U.Vector    Bit = BitVec  !Int !Int !ByteArray
  
  readBit :: Int -> Word -> Bit
  readBit i w = Bit (w .&. (1 `unsafeShiftL` i) /= 0)
  {-# INLINE readBit #-}
  
  extendToWord :: Bit -> Word
  extendToWord (Bit False) = 0
  extendToWord (Bit True ) = complement 0
  
  -- | Read a word at the given bit offset in little-endian order (i.e., the LSB will correspond to the bit at the given address, the 2's bit will correspond to the address + 1, etc.).  If the offset is such that the word extends past the end of the vector, the result is padded with memory garbage.
  indexWord :: U.Vector Bit -> Int -> Word
  indexWord (BitVec _ 0 _) _ = 0
  indexWord (BitVec off len' arr) !i' = word
   where
    len    = off + len'
    i      = off + i'
    nMod   = modWordSize i
    loIx   = divWordSize i
    loWord = indexByteArray arr loIx
    hiWord = indexByteArray arr (loIx + 1)
  
    word
      | nMod == 0
      = loWord
      | loIx == divWordSize (len - 1)
      = loWord `unsafeShiftR` nMod
      | otherwise
      = (loWord `unsafeShiftR` nMod) .|. (hiWord `unsafeShiftL` (wordSize - nMod))
  {-# INLINE indexWord #-}
  
  -- | Read a word at the given bit offset in little-endian order (i.e., the LSB will correspond to the bit at the given address, the 2's bit will correspond to the address + 1, etc.).  If the offset is such that the word extends past the end of the vector, the result is padded with memory garbage.
  readWord :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m Word
  readWord (BitMVec _ 0 _) _ = pure 0
  readWord (BitMVec off len' arr) !i' = do
    let len  = off + len'
        i    = off + i'
        nMod = modWordSize i
        loIx = divWordSize i
    loWord <- readByteArray arr loIx
  
    if nMod == 0
      then pure loWord
      else if loIx == divWordSize (len - 1)
        then pure (loWord `unsafeShiftR` nMod)
        else do
          hiWord <- readByteArray arr (loIx + 1)
          pure
            $   (loWord `unsafeShiftR` nMod)
            .|. (hiWord `unsafeShiftL` (wordSize - nMod))
  {-# SPECIALIZE readWord :: U.MVector s Bit -> Int -> ST s Word #-}
  {-# INLINE readWord #-}
  
  modifyByteArray
    :: PrimMonad m
    => MutableByteArray (PrimState m)
    -> Int
    -> Word
    -> Word
    -> m ()
  #ifndef BITVEC_THREADSAFE
  modifyByteArray arr ix msk new = do
    old <- readByteArray arr ix
    writeByteArray arr ix (old .&. msk .|. new)
  {-# INLINE modifyByteArray #-}
  #else
  modifyByteArray (MutableByteArray mba) (I# ix) (W# msk) (W# new) = do
    primitive $ \state ->
      let !(# state',  _ #) = fetchAndIntArray# mba ix (word2Int# msk) state  in
      let !(# state'', _ #) = fetchOrIntArray#  mba ix (word2Int# new) state' in
      (# state'', () #)
  
  -- https://gitlab.haskell.org/ghc/ghc/issues/17334
  #if __GLASGOW_HASKELL__ == 808 && __GLASGOW_HASKELL_PATCHLEVEL1__ == 1
  {-# NOINLINE modifyByteArray #-}
  #else
  {-# INLINE modifyByteArray #-}
  #endif
  
  #endif
  
  -- | Write a word at the given bit offset in little-endian order (i.e., the LSB will correspond to the bit at the given address, the 2's bit will correspond to the address + 1, etc.).  If the offset is such that the word extends past the end of the vector, the word is truncated and as many low-order bits as possible are written.
  writeWord :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> Word -> m ()
  writeWord (BitMVec _ 0 _) !_ !_ = pure ()
  writeWord (BitMVec off len' arr) !i' !x
    | iMod == 0
    = if len >= i + wordSize
      then writeByteArray arr iDiv x
      else modifyByteArray arr iDiv (hiMask lenMod) (x .&. loMask lenMod)
    | iDiv == divWordSize (len - 1)
    = if lenMod == 0
      then modifyByteArray arr iDiv (loMask iMod) (x `unsafeShiftL` iMod)
      else modifyByteArray arr iDiv (loMask iMod .|. hiMask lenMod) ((x `unsafeShiftL` iMod) .&. loMask lenMod)
    | iDiv + 1 == divWordSize (len - 1)
    = do
      modifyByteArray arr iDiv (loMask iMod) (x `unsafeShiftL` iMod)
      if lenMod == 0
      then modifyByteArray arr (iDiv + 1) (hiMask iMod) (x `unsafeShiftR` (wordSize - iMod))
      else modifyByteArray arr (iDiv + 1) (hiMask iMod .|. hiMask lenMod) (x `unsafeShiftR` (wordSize - iMod) .&. loMask lenMod)
    | otherwise
    = do
      modifyByteArray arr iDiv (loMask iMod) (x `unsafeShiftL` iMod)
      modifyByteArray arr (iDiv + 1) (hiMask iMod) (x `unsafeShiftR` (wordSize - iMod))
    where
      len    = off + len'
      lenMod = modWordSize len
      i      = off + i'
      iMod   = modWordSize i
      iDiv   = divWordSize i
  
  {-# SPECIALIZE writeWord :: U.MVector s Bit -> Int -> Word -> ST s () #-}
  {-# INLINE writeWord #-}
  
  instance MV.MVector U.MVector Bit where
    {-# INLINE basicInitialize #-}
    basicInitialize vec = MV.basicSet vec (Bit False)
  
    {-# INLINE basicUnsafeNew #-}
    basicUnsafeNew n
      | n < 0 = error $ "Data.Bit.basicUnsafeNew: negative length: " ++ show n
      | otherwise = do
        arr <- newByteArray (wordsToBytes $ nWords n)
        pure $ BitMVec 0 n arr
  
    {-# INLINE basicUnsafeReplicate #-}
    basicUnsafeReplicate n x
      | n < 0 =  error $  "Data.Bit.basicUnsafeReplicate: negative length: " ++ show n
      | otherwise = do
        arr <- newByteArray (wordsToBytes $ nWords n)
        setByteArray arr 0 (nWords n) (extendToWord x :: Word)
        pure $ BitMVec 0 n arr
  
    {-# INLINE basicOverlaps #-}
    basicOverlaps (BitMVec i' m' arr1) (BitMVec j' n' arr2) =
      sameMutableByteArray arr1 arr2
        && (between i j (j + n) || between j i (i + m))
     where
      i = divWordSize i'
      m = nWords (i' + m') - i
      j = divWordSize j'
      n = nWords (j' + n') - j
      between x y z = x >= y && x < z
  
    {-# INLINE basicLength #-}
    basicLength (BitMVec _ n _) = n
  
    {-# INLINE basicUnsafeRead #-}
    basicUnsafeRead (BitMVec off _ arr) !i' = do
      let i = off + i'
      word <- readByteArray arr (divWordSize i)
      pure $ readBit (modWordSize i) word
  
    {-# INLINE basicUnsafeWrite #-}
  #ifndef BITVEC_THREADSAFE
    basicUnsafeWrite (BitMVec off _ arr) !i' !x = do
      let i  = off + i'
          j  = divWordSize i
          k  = modWordSize i
          kk = 1 `unsafeShiftL` k :: Word
      word <- readByteArray arr j
      writeByteArray arr j (if unBit x then word .|. kk else word .&. complement kk)
  #else
    basicUnsafeWrite (BitMVec off _ (MutableByteArray mba)) !i' (Bit b) = do
      let i       = off + i'
          !(I# j) = divWordSize i
          !(I# k) = 1 `unsafeShiftL` modWordSize i
      primitive $ \state ->
        let !(# state', _ #) =
                (if b
                  then fetchOrIntArray# mba j k state
                  else fetchAndIntArray# mba j (notI# k) state
                )
        in  (# state', () #)
  #endif
  
    {-# INLINE basicSet #-}
    basicSet (BitMVec off len arr) (extendToWord -> x) | offBits == 0 =
      case modWordSize len of
        0    -> setByteArray arr offWords lWords (x :: Word)
        nMod -> do
          setByteArray arr offWords (lWords - 1) (x :: Word)
          modifyByteArray arr (offWords + lWords - 1) (hiMask nMod) (x .&. loMask nMod)
     where
      offBits  = modWordSize off
      offWords = divWordSize off
      lWords   = nWords (offBits + len)
    basicSet (BitMVec off len arr) (extendToWord -> x) =
      case modWordSize (off + len) of
        0 -> do
          modifyByteArray arr offWords (loMask offBits) (x .&. hiMask offBits)
          setByteArray arr (offWords + 1) (lWords - 1) (x :: Word)
        nMod -> if lWords == 1
          then do
            let lohiMask = loMask offBits .|. hiMask nMod
            modifyByteArray arr offWords lohiMask (x .&. complement lohiMask)
          else do
            modifyByteArray arr offWords (loMask offBits) (x .&. hiMask offBits)
            setByteArray arr (offWords + 1) (lWords - 2) (x :: Word)
            modifyByteArray arr (offWords + lWords - 1) (hiMask nMod) (x .&. loMask nMod)
     where
      offBits  = modWordSize off
      offWords = divWordSize off
      lWords   = nWords (offBits + len)
  
    {-# INLINE basicUnsafeCopy #-}
    basicUnsafeCopy (BitMVec offDst lenDst dst) (BitMVec offSrc _ src)
      | offDstBits == 0, offSrcBits == 0 = case modWordSize lenDst of
        0 -> copyMutableByteArray dst
                                  (wordsToBytes offDstWords)
                                  src
                                  (wordsToBytes offSrcWords)
                                  (wordsToBytes lDstWords)
        nMod -> do
          copyMutableByteArray dst
                               (wordsToBytes offDstWords)
                               src
                               (wordsToBytes offSrcWords)
                               (wordsToBytes $ lDstWords - 1)
  
          lastWordSrc <- readByteArray src (offSrcWords + lDstWords - 1)
          modifyByteArray dst (offDstWords + lDstWords - 1) (hiMask nMod) (lastWordSrc .&. loMask nMod)
     where
      offDstBits  = modWordSize offDst
      offDstWords = divWordSize offDst
      lDstWords   = nWords (offDstBits + lenDst)
      offSrcBits  = modWordSize offSrc
      offSrcWords = divWordSize offSrc
    basicUnsafeCopy (BitMVec offDst lenDst dst) (BitMVec offSrc _ src)
      | offDstBits == offSrcBits = case modWordSize (offSrc + lenDst) of
        0 -> do
          firstWordSrc <- readByteArray src offSrcWords
          modifyByteArray dst offDstWords (loMask offSrcBits) (firstWordSrc .&. hiMask offSrcBits)
          copyMutableByteArray dst
                               (wordsToBytes $ offDstWords + 1)
                               src
                               (wordsToBytes $ offSrcWords + 1)
                               (wordsToBytes $ lDstWords - 1)
        nMod -> if lDstWords == 1
          then do
            let lohiMask = loMask offSrcBits .|. hiMask nMod
            theOnlyWordSrc <- readByteArray src offSrcWords
            modifyByteArray dst offDstWords lohiMask (theOnlyWordSrc .&. complement lohiMask)
          else do
            firstWordSrc <- readByteArray src offSrcWords
            modifyByteArray dst offDstWords (loMask offSrcBits) (firstWordSrc .&. hiMask offSrcBits)
            copyMutableByteArray dst
                                 (wordsToBytes $ offDstWords + 1)
                                 src
                                 (wordsToBytes $ offSrcWords + 1)
                                 (wordsToBytes $ lDstWords - 2)
            lastWordSrc <- readByteArray src (offSrcWords + lDstWords - 1)
            modifyByteArray dst (offDstWords + lDstWords - 1) (hiMask nMod) (lastWordSrc .&. loMask nMod)
     where
      offDstBits  = modWordSize offDst
      offDstWords = divWordSize offDst
      lDstWords   = nWords (offDstBits + lenDst)
      offSrcBits  = modWordSize offSrc
      offSrcWords = divWordSize offSrc
  
    basicUnsafeCopy dst@(BitMVec _ len _) src = do_copy 0
     where
      n = alignUp len
  
      do_copy i
        | i < n = do
          x <- readWord src i
          writeWord dst i x
          do_copy (i + wordSize)
        | otherwise = pure ()
  
    {-# INLINE basicUnsafeMove #-}
    basicUnsafeMove !dst src@(BitMVec srcShift srcLen _)
      | MV.basicOverlaps dst src = do
            -- Align shifts of src and srcCopy to speed up basicUnsafeCopy srcCopy src
        srcCopy <- MV.drop (modWordSize srcShift)
          <$> MV.basicUnsafeNew (modWordSize srcShift + srcLen)
        MV.basicUnsafeCopy srcCopy src
        MV.basicUnsafeCopy dst srcCopy
      | otherwise = MV.basicUnsafeCopy dst src
  
    {-# INLINE basicUnsafeSlice #-}
    basicUnsafeSlice offset n (BitMVec off _ arr) = BitMVec (off + offset) n arr
  
    {-# INLINE basicUnsafeGrow #-}
    basicUnsafeGrow (BitMVec off len src) byBits
      | byWords == 0 = pure $ BitMVec off (len + byBits) src
      | otherwise = do
        dst <- newByteArray (wordsToBytes newWords)
        copyMutableByteArray dst 0 src 0 (wordsToBytes oldWords)
        pure $ BitMVec off (len + byBits) dst
     where
      oldWords = nWords (off + len)
      newWords = nWords (off + len + byBits)
      byWords  = newWords - oldWords
  
  #ifndef BITVEC_THREADSAFE
  
  -- | Flip the bit at the given position.
  -- No bound checks are performed.
  -- Equivalent to 'flip' 'Data.Vector.Unboxed.Mutable.unsafeModify' 'Data.Bits.complement',
  -- but up to 2x faster.
  --
  -- In general there is no reason to 'Data.Vector.Unboxed.Mutable.unsafeModify' bit vectors:
  -- either you modify it with 'id' (which is 'id' altogether)
  -- or with 'Data.Bits.complement' (which is 'unsafeFlipBit').
  --
  -- >>> :set -XOverloadedLists
  -- >>> Data.Vector.Unboxed.modify (`unsafeFlipBit` 2) [1,1,1,1]
  -- [1,1,0,1]
  --
  -- @since 1.0.0.0
  unsafeFlipBit :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m ()
  unsafeFlipBit v i =
  #if MIN_VERSION_vector(0,13,0)
    checkIndex Unsafe
  #else
    UNSAFE_CHECK(checkIndex) "flipBit"
  #endif
      i (MV.length v) $ basicFlipBit v i
  {-# INLINE unsafeFlipBit #-}
  
  basicFlipBit :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m ()
  basicFlipBit (BitMVec off _ arr) !i' = do
    let i  = off + i'
        j  = divWordSize i
        k  = modWordSize i
        kk = 1 `unsafeShiftL` k :: Word
    word <- readByteArray arr j
    writeByteArray arr j (word `xor` kk)
  {-# INLINE basicFlipBit #-}
  
  -- | Flip the bit at the given position.
  -- Equivalent to 'flip' 'Data.Vector.Unboxed.Mutable.modify' 'Data.Bits.complement',
  -- but up to 2x faster.
  --
  -- In general there is no reason to 'Data.Vector.Unboxed.Mutable.modify' bit vectors:
  -- either you modify it with 'id' (which is 'id' altogether)
  -- or with 'Data.Bits.complement' (which is 'flipBit').
  --
  -- >>> :set -XOverloadedLists
  -- >>> Data.Vector.Unboxed.modify (`flipBit` 2) [1,1,1,1]
  -- [1,1,0,1]
  --
  -- @since 1.0.0.0
  flipBit :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m ()
  flipBit v i =
  #if MIN_VERSION_vector(0,13,0)
    checkIndex Bounds
  #else
    BOUNDS_CHECK(checkIndex) "flipBit"
  #endif
      i (MV.length v) $
        unsafeFlipBit v i
  {-# INLINE flipBit #-}
  
  #else
  
  -- | Flip the bit at the given position.
  -- No bound checks are performed.
  -- Equivalent to 'flip' 'Data.Vector.Unboxed.Mutable.unsafeModify' 'Data.Bits.complement',
  -- but up to 33% faster and atomic.
  --
  -- In general there is no reason to 'Data.Vector.Unboxed.Mutable.unsafeModify' bit vectors:
  -- either you modify it with 'id' (which is 'id' altogether)
  -- or with 'Data.Bits.complement' (which is 'unsafeFlipBit').
  --
  -- >>> Data.Vector.Unboxed.modify (\v -> unsafeFlipBit v 1) (read "[1,1,1]")
  -- [1,0,1]
  unsafeFlipBit :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m ()
  unsafeFlipBit v i =
  #if MIN_VERSION_vector(0,13,0)
    checkIndex Unsafe
  #else
    UNSAFE_CHECK(checkIndex) "flipBit"
  #endif
      i (MV.length v) $ basicFlipBit v i
  {-# INLINE unsafeFlipBit #-}
  
  basicFlipBit :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m ()
  basicFlipBit (BitMVec off _ (MutableByteArray mba)) !i' = do
    let i       = off + i'
        !(I# j) = divWordSize i
        !(I# k) = 1 `unsafeShiftL` modWordSize i
    primitive $ \state ->
      let !(# state', _ #) = fetchXorIntArray# mba j k state in (# state', () #)
  {-# INLINE basicFlipBit #-}
  
  -- | Flip the bit at the given position.
  -- Equivalent to 'flip' 'Data.Vector.Unboxed.Mutable.modify' 'Data.Bits.complement',
  -- but up to 33% faster and atomic.
  --
  -- In general there is no reason to 'Data.Vector.Unboxed.Mutable.modify' bit vectors:
  -- either you modify it with 'id' (which is 'id' altogether)
  -- or with 'Data.Bits.complement' (which is 'flipBit').
  --
  -- >>> Data.Vector.Unboxed.modify (\v -> flipBit v 1) (read "[1,1,1]")
  -- [1,0,1]
  flipBit :: PrimMonad m => U.MVector (PrimState m) Bit -> Int -> m ()
  flipBit v i =
  #if MIN_VERSION_vector(0,13,0)
    checkIndex Bounds
  #else
    BOUNDS_CHECK(checkIndex) "flipBit"
  #endif
      i (MV.length v) $ basicFlipBit v i
  {-# INLINE flipBit #-}
  
  #endif
  
  instance V.Vector U.Vector Bit where
    basicUnsafeFreeze (BitMVec s n v) = BitVec s n <$> unsafeFreezeByteArray v
    basicUnsafeThaw (BitVec s n v) = BitMVec s n <$> unsafeThawByteArray v
    basicLength (BitVec _ n _) = n
  
    basicUnsafeIndexM (BitVec off _ arr) !i' = do
      let i = off + i'
      pure $! readBit (modWordSize i) (indexByteArray arr (divWordSize i))
  
    basicUnsafeCopy dst src = do
      src1 <- V.basicUnsafeThaw src
      MV.basicUnsafeCopy dst src1
  
    {-# INLINE basicUnsafeSlice #-}
    basicUnsafeSlice offset n (BitVec off _ arr) = BitVec (off + offset) n arr
  