{-# LANGUAGE CPP                  #-}

{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE TupleSections        #-}
{-# LANGUAGE UndecidableInstances #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

#ifndef BITVEC_THREADSAFE
module Data.Bit.Immutable
#else
module Data.Bit.ImmutableTS
#endif
  ( castFromWords
  , castToWords
  , cloneToWords

  , castFromWords8
  , castToWords8
  , cloneToWords8

  , cloneFromByteString
  , cloneToByteString

  , zipBits
  , mapBits
  , invertBits
  , selectBits
  , excludeBits
  , reverseBits

  , bitIndex
  , nthBitIndex
  , countBits
  , listBits
  ) where

#include "MachDeps.h"

import Control.Monad
import Control.Monad.ST
import Data.Bits
#if UseLibGmp
import Data.Bit.Gmp
#endif
#ifndef BITVEC_THREADSAFE
import Data.Bit.Internal
import Data.Bit.Mutable
#else
import Data.Bit.InternalTS
import Data.Bit.MutableTS
#endif
import Data.Bit.PdepPext
import Data.Bit.Utils
import qualified Data.ByteString.Internal as BS
import Data.Primitive.ByteArray
import qualified Data.Vector.Primitive as P
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Base as UB
import qualified Data.Vector.Unboxed.Mutable as MU
import Data.Word

#ifdef WORDS_BIGENDIAN
import GHC.Exts
#endif

#if UseLibGmp
gmpLimbShift :: Int
gmpLimbShift = case wordSize of
  32 -> 2
  64 -> 3
  _  -> error "gmpLimbShift: unknown architecture"
#endif

instance {-# OVERLAPPING #-} Bits (Vector Bit) where
  .&. :: Vector Bit -> Vector Bit -> Vector Bit
(.&.) = (forall a. Bits a => a -> a -> a)
-> Vector Bit -> Vector Bit -> Vector Bit
zipBits forall a. Bits a => a -> a -> a
(.&.)
  .|. :: Vector Bit -> Vector Bit -> Vector Bit
(.|.) = (forall a. Bits a => a -> a -> a)
-> Vector Bit -> Vector Bit -> Vector Bit
zipBits forall a. Bits a => a -> a -> a
(.|.)
  xor :: Vector Bit -> Vector Bit -> Vector Bit
xor   = (forall a. Bits a => a -> a -> a)
-> Vector Bit -> Vector Bit -> Vector Bit
zipBits forall a. Bits a => a -> a -> a
xor
  complement :: Vector Bit -> Vector Bit
complement = Vector Bit -> Vector Bit
invertBits
  bitSize :: Vector Bit -> Int
bitSize Vector Bit
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"bitSize is undefined"
  bitSizeMaybe :: Vector Bit -> Maybe Int
bitSizeMaybe Vector Bit
_ = forall a. Maybe a
Nothing
  isSigned :: Vector Bit -> Bool
isSigned Vector Bit
_ = Bool
False
  zeroBits :: Vector Bit
zeroBits = forall a. Unbox a => Vector a
U.empty
  popCount :: Vector Bit -> Int
popCount = Vector Bit -> Int
countBits

  testBit :: Vector Bit -> Int -> Bool
testBit Vector Bit
v Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n forall a. Ord a => a -> a -> Bool
>= forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v = Bool
False
    | Bool
otherwise = Bit -> Bool
unBit (forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Bit
v Int
n)

  setBit :: Vector Bit -> Int -> Vector Bit
setBit Vector Bit
v Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n forall a. Ord a => a -> a -> Bool
>= forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v = Vector Bit
v
    | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      MVector s Bit
u <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
v
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Bit
u Int
n (Bool -> Bit
Bit Bool
True)
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
u

  clearBit :: Vector Bit -> Int -> Vector Bit
clearBit Vector Bit
v Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n forall a. Ord a => a -> a -> Bool
>= forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v = Vector Bit
v
    | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      MVector s Bit
u <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
v
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Bit
u Int
n (Bool -> Bit
Bit Bool
False)
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
u

  complementBit :: Vector Bit -> Int -> Vector Bit
complementBit Vector Bit
v Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n forall a. Ord a => a -> a -> Bool
>= forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v = Vector Bit
v
    | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      MVector s Bit
u <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
v
      forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m ()
unsafeFlipBit MVector s Bit
u Int
n
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
u

  bit :: Int -> Vector Bit
bit Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
< Int
0 = forall a. Unbox a => Vector a
U.empty
    | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      MVector s Bit
v <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MU.replicate (Int
n forall a. Num a => a -> a -> a
+ Int
1) (Bool -> Bit
Bit Bool
False)
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Bit
v Int
n (Bool -> Bit
Bit Bool
True)
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
v

  shift :: Vector Bit -> Int -> Vector Bit
shift Vector Bit
v Int
n = case Int
n forall a. Ord a => a -> a -> Ordering
`compare` Int
0 of
    -- shift right
    Ordering
LT
      | forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v forall a. Num a => a -> a -> a
+ Int
n forall a. Ord a => a -> a -> Bool
< Int
0 -> forall a. Unbox a => Vector a
U.empty
      | Bool
otherwise -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
        MVector s Bit
u <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new (forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v forall a. Num a => a -> a -> a
+ Int
n)
        forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
U.copy MVector s Bit
u (forall a. Unbox a => Int -> Vector a -> Vector a
U.drop (- Int
n) Vector Bit
v)
        forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
u
    -- do not shift
    Ordering
EQ -> Vector Bit
v
    -- shift left
    Ordering
GT -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      MVector s Bit
u <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new (forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v forall a. Num a => a -> a -> a
+ Int
n)
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MU.set (forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.take Int
n MVector s Bit
u) (Bool -> Bit
Bit Bool
False)
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
U.copy (forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.drop Int
n MVector s Bit
u) Vector Bit
v
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
u

  rotate :: Vector Bit -> Int -> Vector Bit
rotate Vector Bit
v Int
n'
    | forall a. Unbox a => Vector a -> Bool
U.null Vector Bit
v = Vector Bit
v
    | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      let l :: Int
l = forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v
          n :: Int
n = Int
n' forall a. Integral a => a -> a -> a
`mod` Int
l
      MVector s Bit
u <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
l
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
U.copy (forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.drop Int
n MVector s Bit
u) (forall a. Unbox a => Int -> Vector a -> Vector a
U.take (Int
l forall a. Num a => a -> a -> a
- Int
n) Vector Bit
v)
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
U.copy (forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.take Int
n MVector s Bit
u) (forall a. Unbox a => Int -> Vector a -> Vector a
U.drop (Int
l forall a. Num a => a -> a -> a
- Int
n) Vector Bit
v)
      forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
u

-- | Cast an unboxed vector of words
-- to an unboxed vector of bits.
-- Cf. 'Data.Bit.castFromWordsM'.
--
-- >>> :set -XOverloadedLists
-- >>> castFromWords [123]
-- [1,1,0,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
--
-- @since 1.0.0.0
castFromWords :: U.Vector Word -> U.Vector Bit
castFromWords :: Vector Word -> Vector Bit
castFromWords Vector Word
ws = Int -> Int -> ByteArray -> Vector Bit
BitVec (forall a. Bits a => a -> a
mulWordSize Int
off) (forall a. Bits a => a -> a
mulWordSize Int
len) ByteArray
arr
  where
    P.Vector Int
off Int
len ByteArray
arr = Vector Word -> Vector Word
toPrimVector Vector Word
ws

-- | Try to cast an unboxed vector of bits
-- to an unboxed vector of words.
-- It succeeds if the vector of bits is aligned.
-- Use 'cloneToWords' otherwise.
-- Cf. 'Data.Bit.castToWordsM'.
--
-- > castToWords (castFromWords v) == Just v
--
-- @since 1.0.0.0
castToWords :: U.Vector Bit -> Maybe (U.Vector Word)
castToWords :: Vector Bit -> Maybe (Vector Word)
castToWords (BitVec Int
s Int
n ByteArray
ws)
  | Int -> Bool
aligned Int
s, Int -> Bool
aligned Int
n =
    forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Vector Word -> Vector Word
fromPrimVector forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (forall a. Bits a => a -> a
divWordSize Int
s) (forall a. Bits a => a -> a
divWordSize Int
n) ByteArray
ws
  | Bool
otherwise = forall a. Maybe a
Nothing


-- | Clone an unboxed vector of bits
-- to a new unboxed vector of words.
-- If the bits don't completely fill the words,
-- the last word will be zero-padded.
-- Cf. 'Data.Bit.cloneToWordsM'.
--
-- >>> :set -XOverloadedLists
-- >>> cloneToWords [1,1,0,1,1,1,1]
-- [123]
--
-- @since 1.0.0.0
cloneToWords :: U.Vector Bit -> U.Vector Word
cloneToWords :: Vector Bit -> Vector Word
cloneToWords Vector Bit
v = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
v' <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
v
  MVector s Word
w  <- forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> m (MVector (PrimState m) Word)
cloneToWordsM MVector s Bit
v'
  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Word
w
{-# INLINABLE cloneToWords #-}

-- | Cast an unboxed vector of 'Word8'
-- to an unboxed vector of bits.
--
-- On big-endian architectures 'castFromWords8'
-- resorts to copying instead of aliasing the underlying array.
--
-- >>> :set -XOverloadedLists
-- >>> castFromWords8 [123]
-- [1,1,0,1,1,1,1,0]
--
-- @since 1.0.3.0
castFromWords8 :: U.Vector Word8 -> U.Vector Bit
castFromWords8 :: Vector Word8 -> Vector Bit
castFromWords8 Vector Word8
ws = Int -> Int -> ByteArray -> Vector Bit
BitVec (Int
off forall a. Bits a => a -> Int -> a
`shiftL` Int
3) (Int
len forall a. Bits a => a -> Int -> a
`shiftL` Int
3) ByteArray
arr
  where
#ifdef WORDS_BIGENDIAN
    UB.V_Word8 (P.Vector off' len arr') = ws
    off = 0
    arr = runST $ do
      let lenWords = nWords $ len `shiftL` 3
          len' = wordsToBytes lenWords
      marr <- newByteArray len'
      copyByteArray marr 0 arr' off' len
      fillByteArray marr len (len' - len) 0
      forM_ [0..lenWords - 1] $ \i -> do
        W# w <- readByteArray marr i
        writeByteArray marr i (W# (byteSwap# w))
      unsafeFreezeByteArray marr
#else
    UB.V_Word8 (P.Vector Int
off Int
len ByteArray
arr) = Vector Word8
ws
#endif

-- | Try to cast an unboxed vector of bits
-- to an unboxed vector of 'Word8'.
-- It succeeds if the vector of bits is aligned.
-- Use 'Data.Bit.cloneToWords8' otherwise.
--
-- > castToWords8 (castFromWords8 v) == Just v
--
-- @since 1.0.3.0
castToWords8 :: U.Vector Bit -> Maybe (U.Vector Word8)
#ifdef WORDS_BIGENDIAN
castToWords8 = const Nothing
#else
castToWords8 :: Vector Bit -> Maybe (Vector Word8)
castToWords8 (BitVec Int
s Int
n ByteArray
ws)
  | Int
s forall a. Bits a => a -> a -> a
.&. Int
7 forall a. Eq a => a -> a -> Bool
== Int
0, Int
n forall a. Bits a => a -> a -> a
.&. Int
7 forall a. Eq a => a -> a -> Bool
== Int
0
  = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Vector Word8 -> Vector Word8
UB.V_Word8 forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
s forall a. Bits a => a -> Int -> a
`shiftR` Int
3) (Int
n forall a. Bits a => a -> Int -> a
`shiftR` Int
3) ByteArray
ws
  | Bool
otherwise = forall a. Maybe a
Nothing
#endif

-- | Clone an unboxed vector of bits
-- to a new unboxed vector of 'Word8'.
-- If the bits don't completely fill the bytes,
-- the last 'Word8' will be zero-padded.
--
-- >>> :set -XOverloadedLists
-- >>> cloneToWords8 [1,1,0,1,1,1,1]
-- [123]
--
-- @since 1.0.3.0
cloneToWords8 :: U.Vector Bit -> U.Vector Word8
cloneToWords8 :: Vector Bit -> Vector Word8
cloneToWords8 Vector Bit
v = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
v' <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
v
  MVector s Word8
w  <- forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> m (MVector (PrimState m) Word8)
cloneToWords8M MVector s Bit
v'
  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Word8
w
{-# INLINABLE cloneToWords8 #-}

-- | Clone a 'BS.ByteString' to a new unboxed vector of bits.
--
-- >>> :set -XOverloadedStrings
-- >>> cloneFromByteString "abc"
-- [1,0,0,0,0,1,1,0,0,1,0,0,0,1,1,0,1,1,0,0,0,1,1,0]
--
-- @since 1.1.0.0
cloneFromByteString :: BS.ByteString -> U.Vector Bit
cloneFromByteString :: ByteString -> Vector Bit
cloneFromByteString
  = Vector Word8 -> Vector Bit
castFromWords8
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
U.convert
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c d. (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
S.unsafeFromForeignPtr
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> (ForeignPtr Word8, Int, Int)
BS.toForeignPtr

-- | Clone an unboxed vector of bits to a new 'BS.ByteString'.
-- If the bits don't completely fill the bytes,
-- the last character will be zero-padded.
--
-- >>> :set -XOverloadedLists
-- >>> cloneToByteString [1,0,0,0,0,1,1,0,0,1,0,0,0,1,1,0,1,1,0,0,0,1]
-- "ab#"
--
-- @since 1.1.0.0
cloneToByteString :: U.Vector Bit -> BS.ByteString
cloneToByteString :: Vector Bit -> ByteString
cloneToByteString
  = forall a b c d. (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 ForeignPtr Word8 -> Int -> Int -> ByteString
BS.fromForeignPtr
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Vector a -> (ForeignPtr a, Int, Int)
S.unsafeToForeignPtr
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
U.convert
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Bit -> Vector Word8
cloneToWords8

uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 :: forall a b c d. (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 a -> b -> c -> d
f (a
x, b
y, c
z) = a -> b -> c -> d
f a
x b
y c
z

-- | Zip two vectors with the given function.
-- Similar to 'Data.Vector.Unboxed.zipWith',
-- but up to 1000x (!) faster.
--
-- For sufficiently dense sets, represented as bitmaps,
-- 'zipBits' is up to 32x faster than
-- 'Data.IntSet.union', 'Data.IntSet.intersection', etc.
--
-- Users are strongly encouraged to enable the
-- @libgmp@ flag for the ultimate performance of 'zipBits'.
--
-- >>> :set -XOverloadedLists
-- >>> import Data.Bits
-- >>> zipBits (.&.) [1,1,0] [0,1,1] -- intersection
-- [0,1,0]
-- >>> zipBits (.|.) [1,1,0] [0,1,1] -- union
-- [1,1,1]
-- >>> zipBits (\x y -> x .&. complement y) [1,1,0] [0,1,1] -- difference
-- [1,0,0]
-- >>> zipBits xor [1,1,0] [0,1,1] -- symmetric difference
-- [1,0,1]
--
-- @since 1.0.0.0
zipBits
  :: (forall a . Bits a => a -> a -> a)
  -> U.Vector Bit
  -> U.Vector Bit
  -> U.Vector Bit
zipBits :: (forall a. Bits a => a -> a -> a)
-> Vector Bit -> Vector Bit -> Vector Bit
zipBits forall a. Bits a => a -> a -> a
_ (BitVec Int
_ Int
0 ByteArray
_) Vector Bit
_ = forall a. Unbox a => Vector a
U.empty
zipBits forall a. Bits a => a -> a -> a
_ Vector Bit
_ (BitVec Int
_ Int
0 ByteArray
_) = forall a. Unbox a => Vector a
U.empty
#if UseLibGmp
zipBits f (BitVec 0 l1 arg1) (BitVec 0 l2 arg2) = runST $ do
    let l = l1 `min` l2
        w = nWords l
        b = w `shiftL` gmpLimbShift
    brr <- newByteArray b
    let ff = unBit $ f (Bit False) (Bit False)
        ft = unBit $ f (Bit False) (Bit True)
        tf = unBit $ f (Bit True)  (Bit False)
        tt = unBit $ f (Bit True)  (Bit True)
    case (ff, ft, tf, tt) of
      (False, False, False, False) -> setByteArray brr 0 w (zeroBits :: Word)
      (False, False, False, True)  -> mpnAndN  brr arg1 arg2 w
      (False, False, True,  False) -> mpnAndnN brr arg1 arg2 w
      (False, False, True,  True)  -> copyByteArray brr 0 arg1 0 b
      (False, True,  False, False) -> mpnAndnN brr arg2 arg1 w
      (False, True,  False, True)  -> copyByteArray brr 0 arg2 0 b
      (False, True,  True,  False) -> mpnXorN  brr arg1 arg2 w
      (False, True,  True,  True)  -> mpnIorN  brr arg1 arg2 w
      (True,  False, False, False) -> mpnNiorN brr arg1 arg2 w
      (True,  False, False, True)  -> mpnXnorN brr arg1 arg2 w
      (True,  False, True,  False) -> mpnCom   brr arg2      w
      (True,  False, True,  True)  -> mpnIornN brr arg1 arg2 w
      (True,  True,  False, False) -> mpnCom   brr arg1      w
      (True,  True,  False, True)  -> mpnIornN brr arg2 arg1 w
      (True,  True,  True,  False) -> mpnNandN brr arg1 arg2 w
      (True,  True,  True,  True)  -> setByteArray brr 0 w (complement zeroBits :: Word)
    BitVec 0 l <$> unsafeFreezeByteArray brr
#endif
zipBits forall a. Bits a => a -> a -> a
f Vector Bit
xs Vector Bit
ys = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = forall a. Ord a => a -> a -> a
min (forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs) (forall a. Unbox a => Vector a -> Int
U.length Vector Bit
ys)
  MVector s Bit
zs <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
n
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0, Int
wordSize .. Int
n forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i ->
    forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
zs Int
i (forall a. Bits a => a -> a -> a
f (Vector Bit -> Int -> Word
indexWord Vector Bit
xs Int
i) (Vector Bit -> Int -> Word
indexWord Vector Bit
ys Int
i))
  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
zs
{-# INLINABLE zipBits #-}

-- | Map a vectors with the given function.
-- Similar to 'Data.Vector.Unboxed.map',
-- but faster.
--
-- >>> :set -XOverloadedLists
-- >>> import Data.Bits
-- >>> mapBits complement [0,1,1]
-- [1,0,0]
--
-- @since 1.1.0.0
mapBits
  :: (forall a . Bits a => a -> a)
  -> U.Vector Bit
  -> U.Vector Bit
mapBits :: (forall a. Bits a => a -> a) -> Vector Bit -> Vector Bit
mapBits forall a. Bits a => a -> a
f = case (Bit -> Bool
unBit (forall a. Bits a => a -> a
f (Bool -> Bit
Bit Bool
False)), Bit -> Bool
unBit (forall a. Bits a => a -> a
f (Bool -> Bit
Bit Bool
True))) of
  (Bool
False, Bool
False) -> (forall a. Unbox a => Int -> a -> Vector a
`U.replicate` Bool -> Bit
Bit Bool
False) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => Vector a -> Int
U.length
  (Bool
False, Bool
True)  -> forall a. a -> a
id
  (Bool
True, Bool
False)  -> Vector Bit -> Vector Bit
invertBits
  (Bool
True, Bool
True)   -> (forall a. Unbox a => Int -> a -> Vector a
`U.replicate` Bool -> Bit
Bit Bool
True) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Unbox a => Vector a -> Int
U.length
{-# INLINE mapBits #-}

-- | Invert (flip) all bits.
--
-- Users are strongly encouraged to enable the
-- @libgmp@ flag for the ultimate performance of 'invertBits'.
--
-- >>> :set -XOverloadedLists
-- >>> invertBits [0,1,0,1,0]
-- [1,0,1,0,1]
--
-- @since 1.0.1.0
invertBits
  :: U.Vector Bit
  -> U.Vector Bit
invertBits :: Vector Bit -> Vector Bit
invertBits (BitVec Int
_ Int
0 ByteArray
_) = forall a. Unbox a => Vector a
U.empty
#if UseLibGmp
invertBits (BitVec 0 l arg) = runST $ do
  let w = nWords l
  brr <- newByteArray (w `shiftL` gmpLimbShift)
  mpnCom brr arg w
  BitVec 0 l <$> unsafeFreezeByteArray brr
#endif
invertBits Vector Bit
xs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs
  MVector s Bit
ys <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
n
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0, Int
wordSize .. Int
n forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i ->
    forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
ys Int
i (forall a. Bits a => a -> a
complement (Vector Bit -> Int -> Word
indexWord Vector Bit
xs Int
i))
  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
ys

-- | For each set bit of the first argument, deposit
-- the corresponding bit of the second argument
-- to the result. Similar to the
-- [parallel bit deposit instruction (PDEP)](https://en.wikipedia.org/wiki/X86_Bit_manipulation_instruction_set#Parallel_bit_deposit_and_extract).
--
-- >>> :set -XOverloadedLists
-- >>> selectBits [0,1,0,1,1] [1,1,0,0,1]
-- [1,0,1]
--
-- Here is a reference (but slow) implementation:
--
-- > import qualified Data.Vector.Unboxed as U
-- > selectBits mask ws = U.map snd (U.filter (unBit . fst) (U.zip mask ws))
--
-- @since 0.1
selectBits :: U.Vector Bit -> U.Vector Bit -> U.Vector Bit
selectBits :: Vector Bit -> Vector Bit -> Vector Bit
selectBits Vector Bit
is Vector Bit
xs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
xs1 <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
xs
  Int
n   <- forall (m :: * -> *).
PrimMonad m =>
Vector Bit -> MVector (PrimState m) Bit -> m Int
selectBitsInPlace Vector Bit
is MVector s Bit
xs1
  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.take Int
n MVector s Bit
xs1)

-- | For each unset bit of the first argument, deposit
-- the corresponding bit of the second argument
-- to the result.
--
-- >>> :set -XOverloadedLists
-- >>> excludeBits [0,1,0,1,1] [1,1,0,0,1]
-- [1,0]
--
-- Here is a reference (but slow) implementation:
--
-- > import qualified Data.Vector.Unboxed as U
-- > excludeBits mask ws = U.map snd (U.filter (not . unBit . fst) (U.zip mask ws))
--
-- @since 0.1
excludeBits :: U.Vector Bit -> U.Vector Bit -> U.Vector Bit
excludeBits :: Vector Bit -> Vector Bit -> Vector Bit
excludeBits Vector Bit
is Vector Bit
xs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
xs1 <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
xs
  Int
n   <- forall (m :: * -> *).
PrimMonad m =>
Vector Bit -> MVector (PrimState m) Bit -> m Int
excludeBitsInPlace Vector Bit
is MVector s Bit
xs1
  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.take Int
n MVector s Bit
xs1)

-- | Reverse the order of bits.
--
-- >>> :set -XOverloadedLists
-- >>> reverseBits [1,1,0,1,0]
-- [0,1,0,1,1]
--
-- Consider using the [vector-rotcev](https://hackage.haskell.org/package/vector-rotcev) package
-- to reverse vectors in O(1) time.
--
-- @since 1.0.1.0
reverseBits :: U.Vector Bit -> U.Vector Bit
reverseBits :: Vector Bit -> Vector Bit
reverseBits Vector Bit
xs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n    = forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs
  MVector s Bit
ys <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
n

  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0, Int
wordSize .. Int
n forall a. Num a => a -> a -> a
- Int
wordSize] forall a b. (a -> b) -> a -> b
$ \Int
i ->
    forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
ys (Int
n forall a. Num a => a -> a -> a
- Int
i forall a. Num a => a -> a -> a
- Int
wordSize) (Word -> Word
reverseWord (Vector Bit -> Int -> Word
indexWord Vector Bit
xs Int
i))

  let nMod :: Int
nMod = Int -> Int
modWordSize Int
n
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
nMod forall a. Eq a => a -> a -> Bool
/= Int
0) forall a b. (a -> b) -> a -> b
$ do
    let x :: Word
x = Vector Bit -> Int -> Word
indexWord Vector Bit
xs (forall a. Bits a => a -> a
mulWordSize (forall a. Bits a => a -> a
divWordSize Int
n))
    Word
y <- forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector s Bit
ys Int
0
    forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
ys Int
0 (Int -> Word -> Word -> Word
meld Int
nMod (Int -> Word -> Word
reversePartialWord Int
nMod Word
x) Word
y)

  forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
ys

clipLoBits :: Bit -> Int -> Word -> Word
clipLoBits :: Bit -> Int -> Word -> Word
clipLoBits (Bit Bool
True ) Int
k Word
w = Word
w forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
k
clipLoBits (Bit Bool
False) Int
k Word
w = (Word
w forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
k) forall a. Bits a => a -> a -> a
.|. Int -> Word
hiMask (Int
wordSize forall a. Num a => a -> a -> a
- Int
k)

clipHiBits :: Bit -> Int -> Word -> Word
clipHiBits :: Bit -> Int -> Word -> Word
clipHiBits (Bit Bool
True ) Int
k Word
w = Word
w forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
k
clipHiBits (Bit Bool
False) Int
k Word
w = Word
w forall a. Bits a => a -> a -> a
.|. Int -> Word
hiMask Int
k

-- | Return the index of the first bit in the vector
-- with the specified value, if any.
-- Similar to 'Data.Vector.Unboxed.elemIndex', but up to 64x faster.
--
-- >>> :set -XOverloadedLists
-- >>> bitIndex 1 [0,0,1,0,1]
-- Just 2
-- >>> bitIndex 1 [0,0,0,0,0]
-- Nothing
--
-- > bitIndex bit == nthBitIndex bit 1
--
-- One can also use it to reduce a vector with disjunction or conjunction:
--
-- > import Data.Maybe
-- > isAnyBitSet   = isJust    . bitIndex 1
-- > areAllBitsSet = isNothing . bitIndex 0
--
-- @since 1.0.0.0
bitIndex :: Bit -> U.Vector Bit -> Maybe Int
bitIndex :: Bit -> Vector Bit -> Maybe Int
bitIndex Bit
b (BitVec Int
off Int
len ByteArray
arr)
  | Int
len forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Maybe a
Nothing
  | Int
offBits forall a. Eq a => a -> a -> Bool
== Int
0 = case Int -> Int
modWordSize Int
len of
    Int
0    -> Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords Bit
b Int
offWords Int
lWords ByteArray
arr
    Int
nMod -> case Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords Bit
b Int
offWords (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr of
      r :: Maybe Int
r@Just{} -> Maybe Int
r
      Maybe Int
Nothing  -> (forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> a
mulWordSize (Int
lWords forall a. Num a => a -> a -> a
- Int
1)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Word -> Maybe Int
bitIndexInWord
        Bit
b
        (Bit -> Int -> Word -> Word
clipHiBits Bit
b Int
nMod (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1)))
  | Bool
otherwise = case Int -> Int
modWordSize (Int
off forall a. Num a => a -> a -> a
+ Int
len) of
    Int
0 ->
      case
          Bit -> Word -> Maybe Int
bitIndexInWord Bit
b (Bit -> Int -> Word -> Word
clipLoBits Bit
b Int
offBits (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords))
        of
          r :: Maybe Int
r@Just{} -> Maybe Int
r
          Maybe Int
Nothing ->
            (forall a. Num a => a -> a -> a
+ (Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits))
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords Bit
b (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr
    Int
nMod -> case Int
lWords of
      Int
1 -> Bit -> Word -> Maybe Int
bitIndexInWord
        Bit
b
        (Bit -> Int -> Word -> Word
clipHiBits Bit
b Int
len (Bit -> Int -> Word -> Word
clipLoBits Bit
b Int
offBits (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords)))
      Int
_ ->
        case
            Bit -> Word -> Maybe Int
bitIndexInWord
              Bit
b
              (Bit -> Int -> Word -> Word
clipLoBits Bit
b Int
offBits (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords))
          of
            r :: Maybe Int
r@Just{} -> Maybe Int
r
            Maybe Int
Nothing ->
              (forall a. Num a => a -> a -> a
+ (Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits))
                forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords Bit
b (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr of
                      r :: Maybe Int
r@Just{} -> Maybe Int
r
                      Maybe Int
Nothing ->
                        (forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> a
mulWordSize (Int
lWords forall a. Num a => a -> a -> a
- Int
2)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Word -> Maybe Int
bitIndexInWord
                          Bit
b
                          (Bit -> Int -> Word -> Word
clipHiBits
                            Bit
b
                            Int
nMod
                            (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1))
                          )
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits forall a. Num a => a -> a -> a
+ Int
len)

bitIndexInWord :: Bit -> Word -> Maybe Int
bitIndexInWord :: Bit -> Word -> Maybe Int
bitIndexInWord (Bit Bool
True ) = Word -> Maybe Int
ffs
bitIndexInWord (Bit Bool
False) = Word -> Maybe Int
ffs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bits a => a -> a
complement

bitIndexInWords :: Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords :: Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords (Bit Bool
True) !Int
off !Int
len !ByteArray
arr = Int -> Maybe Int
go Int
off
 where
  go :: Int -> Maybe Int
go !Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
off forall a. Num a => a -> a -> a
+ Int
len = forall a. Maybe a
Nothing
    | Bool
otherwise = case Word -> Maybe Int
ffs (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n) of
      Maybe Int
Nothing  -> Int -> Maybe Int
go (Int
n forall a. Num a => a -> a -> a
+ Int
1)
      Just Int
r  -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Bits a => a -> a
mulWordSize (Int
n forall a. Num a => a -> a -> a
- Int
off) forall a. Num a => a -> a -> a
+ Int
r
bitIndexInWords (Bit Bool
False) !Int
off !Int
len !ByteArray
arr = Int -> Maybe Int
go Int
off
 where
  go :: Int -> Maybe Int
go !Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
off forall a. Num a => a -> a -> a
+ Int
len = forall a. Maybe a
Nothing
    | Bool
otherwise = case Word -> Maybe Int
ffs (forall a. Bits a => a -> a
complement (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n)) of
      Maybe Int
Nothing -> Int -> Maybe Int
go (Int
n forall a. Num a => a -> a -> a
+ Int
1)
      Just Int
r  -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Bits a => a -> a
mulWordSize (Int
n forall a. Num a => a -> a -> a
- Int
off) forall a. Num a => a -> a -> a
+ Int
r

-- | Return the index of the @n@-th bit in the vector
-- with the specified value, if any.
-- Here @n@ is 1-based and the index is 0-based.
-- Non-positive @n@ results in an error.
--
-- >>> :set -XOverloadedLists
-- >>> nthBitIndex 1 2 [0,1,0,1,1,1,0] -- 2nd occurence of 1
-- Just 3
-- >>> nthBitIndex 1 5 [0,1,0,1,1,1,0] -- 5th occurence of 1
-- Nothing
--
-- One can use 'nthBitIndex' to implement
-- to implement @select{0,1}@ queries
-- for <https://en.wikipedia.org/wiki/Succinct_data_structure succinct dictionaries>.
--
-- @since 1.0.0.0
nthBitIndex :: Bit -> Int -> U.Vector Bit -> Maybe Int
nthBitIndex :: Bit -> Int -> Vector Bit -> Maybe Int
nthBitIndex Bit
_ Int
k Vector Bit
_ | Int
k forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"nthBitIndex: n must be positive"
nthBitIndex Bit
b Int
k (BitVec Int
off Int
len ByteArray
arr)
  | Int
len forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Maybe a
Nothing
  | Int
offBits forall a. Eq a => a -> a -> Bool
== Int
0 = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ case Int -> Int
modWordSize Int
len of
    Int
0    -> Bit -> Int -> Int -> Int -> ByteArray -> Either Int Int
nthInWords Bit
b Int
k Int
offWords Int
lWords ByteArray
arr
    Int
nMod -> case Bit -> Int -> Int -> Int -> ByteArray -> Either Int Int
nthInWords Bit
b Int
k Int
offWords (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr of
      r :: Either Int Int
r@Right{} -> Either Int Int
r
      Left Int
k'   -> (forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> a
mulWordSize (Int
lWords forall a. Num a => a -> a -> a
- Int
1)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Int -> Word -> Either Int Int
nthInWord
        Bit
b
        Int
k'
        (Bit -> Int -> Word -> Word
clipHiBits Bit
b Int
nMod (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1)))
  | Bool
otherwise = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ case Int -> Int
modWordSize (Int
off forall a. Num a => a -> a -> a
+ Int
len) of
    Int
0 ->
      case Bit -> Int -> Word -> Either Int Int
nthInWord Bit
b Int
k (Bit -> Int -> Word -> Word
clipLoBits Bit
b Int
offBits (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords)) of
        r :: Either Int Int
r@Right{} -> Either Int Int
r
        Left Int
k' ->
          (forall a. Num a => a -> a -> a
+ (Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits))
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Int -> Int -> Int -> ByteArray -> Either Int Int
nthInWords Bit
b Int
k' (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr
    Int
nMod -> case Int
lWords of
      Int
1 -> Bit -> Int -> Word -> Either Int Int
nthInWord
        Bit
b
        Int
k
        (Bit -> Int -> Word -> Word
clipHiBits Bit
b Int
len (Bit -> Int -> Word -> Word
clipLoBits Bit
b Int
offBits (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords)))
      Int
_ ->
        case
            Bit -> Int -> Word -> Either Int Int
nthInWord Bit
b Int
k (Bit -> Int -> Word -> Word
clipLoBits Bit
b Int
offBits (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords))
          of
            r :: Either Int Int
r@Right{} -> Either Int Int
r
            Left Int
k' ->
              (forall a. Num a => a -> a -> a
+ (Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits))
                forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Bit -> Int -> Int -> Int -> ByteArray -> Either Int Int
nthInWords Bit
b Int
k' (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr of
                      r :: Either Int Int
r@Right{} -> Either Int Int
r
                      Left Int
k''  -> (forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> a
mulWordSize (Int
lWords forall a. Num a => a -> a -> a
- Int
2)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Int -> Word -> Either Int Int
nthInWord
                        Bit
b
                        Int
k''
                        (Bit -> Int -> Word -> Word
clipHiBits
                          Bit
b
                          Int
nMod
                          (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1))
                        )
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits forall a. Num a => a -> a -> a
+ Int
len)

nthInWord :: Bit -> Int -> Word -> Either Int Int
nthInWord :: Bit -> Int -> Word -> Either Int Int
nthInWord (Bit Bool
b) Int
k Word
v = if Int
k forall a. Ord a => a -> a -> Bool
> Int
c then forall a b. a -> Either a b
Left (Int
k forall a. Num a => a -> a -> a
- Int
c) else forall a b. b -> Either a b
Right (Int -> Word -> Int
unsafeNthTrueInWord Int
k Word
w)
 where
  w :: Word
w = if Bool
b then Word
v else forall a. Bits a => a -> a
complement Word
v
  c :: Int
c = forall a. Bits a => a -> Int
popCount Word
w

nthInWords :: Bit -> Int -> Int -> Int -> ByteArray -> Either Int Int
nthInWords :: Bit -> Int -> Int -> Int -> ByteArray -> Either Int Int
nthInWords (Bit Bool
True) !Int
k !Int
off !Int
len !ByteArray
arr = Int -> Int -> Either Int Int
go Int
off Int
k
 where
  go :: Int -> Int -> Either Int Int
go !Int
n !Int
l
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
off forall a. Num a => a -> a -> a
+ Int
len = forall a b. a -> Either a b
Left Int
l
    | Bool
otherwise = if Int
l forall a. Ord a => a -> a -> Bool
> Int
c
      then Int -> Int -> Either Int Int
go (Int
n forall a. Num a => a -> a -> a
+ Int
1) (Int
l forall a. Num a => a -> a -> a
- Int
c)
      else forall a b. b -> Either a b
Right (forall a. Bits a => a -> a
mulWordSize (Int
n forall a. Num a => a -> a -> a
- Int
off) forall a. Num a => a -> a -> a
+ Int -> Word -> Int
unsafeNthTrueInWord Int
l Word
w)
   where
    w :: Word
w = forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n
    c :: Int
c = forall a. Bits a => a -> Int
popCount Word
w
nthInWords (Bit Bool
False) !Int
k !Int
off !Int
len !ByteArray
arr = Int -> Int -> Either Int Int
go Int
off Int
k
 where
  go :: Int -> Int -> Either Int Int
go !Int
n !Int
l
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
off forall a. Num a => a -> a -> a
+ Int
len = forall a b. a -> Either a b
Left Int
l
    | Bool
otherwise = if Int
l forall a. Ord a => a -> a -> Bool
> Int
c
      then Int -> Int -> Either Int Int
go (Int
n forall a. Num a => a -> a -> a
+ Int
1) (Int
l forall a. Num a => a -> a -> a
- Int
c)
      else forall a b. b -> Either a b
Right (forall a. Bits a => a -> a
mulWordSize (Int
n forall a. Num a => a -> a -> a
- Int
off) forall a. Num a => a -> a -> a
+ Int -> Word -> Int
unsafeNthTrueInWord Int
l Word
w)
   where
    w :: Word
w = forall a. Bits a => a -> a
complement (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n)
    c :: Int
c = forall a. Bits a => a -> Int
popCount Word
w

unsafeNthTrueInWord :: Int -> Word -> Int
unsafeNthTrueInWord :: Int -> Word -> Int
unsafeNthTrueInWord Int
l Word
w = forall b. FiniteBits b => b -> Int
countTrailingZeros (Word -> Word -> Word
pdep (Word
1 forall a. Bits a => a -> Int -> a
`shiftL` (Int
l forall a. Num a => a -> a -> a
- Int
1)) Word
w)

-- | Return the number of set bits in a vector (population count, popcount).
--
-- Users are strongly encouraged to enable the
-- @libgmp@ flag for the ultimate performance of 'countBits'.
--
-- >>> :set -XOverloadedLists
-- >>> countBits [1,1,0,1,0,1]
-- 4
--
-- One can combine 'countBits' with 'Data.Vector.Unboxed.take'
-- to implement @rank{0,1}@ queries
-- for <https://en.wikipedia.org/wiki/Succinct_data_structure succinct dictionaries>.
--
-- @since 0.1
countBits :: U.Vector Bit -> Int
countBits :: Vector Bit -> Int
countBits (BitVec Int
_ Int
0 ByteArray
_)                      = Int
0
#if UseLibGmp
countBits (BitVec 0 len arr) | modWordSize len == 0 =
  fromIntegral (mpnPopcount arr (divWordSize len))
#endif
countBits (BitVec Int
off Int
len ByteArray
arr) | Int
offBits forall a. Eq a => a -> a -> Bool
== Int
0 = case Int -> Int
modWordSize Int
len of
  Int
0    -> Vector Word -> Int
countBitsInWords (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords Int
lWords ByteArray
arr)
  Int
nMod -> Vector Word -> Int
countBitsInWords (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
    forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1) forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
nMod)
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits forall a. Num a => a -> a -> a
+ Int
len)
countBits (BitVec Int
off Int
len ByteArray
arr) = case Int -> Int
modWordSize (Int
off forall a. Num a => a -> a -> a
+ Int
len) of
  Int
0 -> forall a. Bits a => a -> Int
popCount (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word)
    forall a. Num a => a -> a -> a
+ Vector Word -> Int
countBitsInWords (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
  Int
nMod -> case Int
lWords of
    Int
1 -> forall a. Bits a => a -> Int
popCount
      ((forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits) forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
len)
    Int
_ ->
      forall a. Bits a => a -> Int
popCount (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word)
        forall a. Num a => a -> a -> a
+ Vector Word -> Int
countBitsInWords (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr)
        forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int
popCount (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1) forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
nMod)
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits forall a. Num a => a -> a -> a
+ Int
len)

countBitsInWords :: P.Vector Word -> Int
countBitsInWords :: Vector Word -> Int
countBitsInWords = forall b a. Prim b => (a -> b -> a) -> a -> Vector b -> a
P.foldl' (\Int
acc Word
word -> forall a. Bits a => a -> Int
popCount Word
word forall a. Num a => a -> a -> a
+ Int
acc) Int
0

-- | Return 0-based indices of set bits in a vector.
--
-- >>> :set -XOverloadedLists
-- >>> listBits [1,1,0,1,0,1]
-- [0,1,3,5]
--
-- @since 0.1
listBits :: U.Vector Bit -> [Int]
listBits :: Vector Bit -> [Int]
listBits (BitVec Int
_ Int
0 ByteArray
_)                      = []
listBits (BitVec Int
off Int
len ByteArray
arr) | Int
offBits forall a. Eq a => a -> a -> Bool
== Int
0 = case Int -> Int
modWordSize Int
len of
  Int
0 -> Int -> Vector Word -> [Int] -> [Int]
listBitsInWords Int
0 (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords Int
lWords ByteArray
arr) []
  Int
nMod ->
    Int -> Vector Word -> [Int] -> [Int]
listBitsInWords Int
0 (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
      forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> a
mulWordSize (Int
lWords forall a. Num a => a -> a -> a
- Int
1))
      forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Bits a => a -> Int -> Bool
testBit (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1) :: Word))
               [Int
0 .. Int
nMod forall a. Num a => a -> a -> a
- Int
1]
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits forall a. Num a => a -> a -> a
+ Int
len)
listBits (BitVec Int
off Int
len ByteArray
arr) = case Int -> Int
modWordSize (Int
off forall a. Num a => a -> a -> a
+ Int
len) of
  Int
0 ->
    forall a. (a -> Bool) -> [a] -> [a]
filter
        (forall a. Bits a => a -> Int -> Bool
testBit (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word))
        [Int
0 .. Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits forall a. Num a => a -> a -> a
- Int
1]
      forall a. [a] -> [a] -> [a]
++ Int -> Vector Word -> [Int] -> [Int]
listBitsInWords (Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits)
                         (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
                         []
  Int
nMod -> case Int
lWords of
    Int
1 -> forall a. (a -> Bool) -> [a] -> [a]
filter
      (forall a. Bits a => a -> Int -> Bool
testBit (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word))
      [Int
0 .. Int
len forall a. Num a => a -> a -> a
- Int
1]
    Int
_ ->
      forall a. (a -> Bool) -> [a] -> [a]
filter
          (forall a. Bits a => a -> Int -> Bool
testBit (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word))
          [Int
0 .. Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits forall a. Num a => a -> a -> a
- Int
1]
        forall a. [a] -> [a] -> [a]
++ ( Int -> Vector Word -> [Int] -> [Int]
listBitsInWords (Int
wordSize forall a. Num a => a -> a -> a
- Int
offBits)
                             (forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr)
           forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
+ (forall a. Bits a => a -> a
mulWordSize (Int
lWords forall a. Num a => a -> a -> a
- Int
1) forall a. Num a => a -> a -> a
- Int
offBits))
           forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter
               (forall a. Bits a => a -> Int -> Bool
testBit (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords forall a. Num a => a -> a -> a
+ Int
lWords forall a. Num a => a -> a -> a
- Int
1) :: Word))
               [Int
0 .. Int
nMod forall a. Num a => a -> a -> a
- Int
1]
           )
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits forall a. Num a => a -> a -> a
+ Int
len)

listBitsInWord :: Int -> Word -> [Int]
listBitsInWord :: Int -> Word -> [Int]
listBitsInWord Int
offset Word
word =
  forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
+ Int
offset) forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Bits a => a -> Int -> Bool
testBit Word
word) forall a b. (a -> b) -> a -> b
$ [Int
0 .. Int
wordSize forall a. Num a => a -> a -> a
- Int
1]

listBitsInWords :: Int -> P.Vector Word -> [Int] -> [Int]
listBitsInWords :: Int -> Vector Word -> [Int] -> [Int]
listBitsInWords Int
offset = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall a b. Prim a => (Int -> a -> b -> b) -> b -> Vector a -> b
P.ifoldr
  (\Int
i Word
word [Int]
acc -> Int -> Word -> [Int]
listBitsInWord (Int
offset forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> a
mulWordSize Int
i) Word
word forall a. [a] -> [a] -> [a]
++ [Int]
acc)