{-# LANGUAGE CPP                  #-}

{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE BinaryLiterals       #-}
{-# 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 UseSIMD
import Data.Bit.SIMD
#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

-- | Note: For '(.&.)', '(.|.)' and 'xor',
-- if one input is larger than the other, the remaining bits will be ignored.
-- 'bitSize' is undefined (throws an exception).
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 a -> a -> a
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 a -> a -> a
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 a -> a -> a
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
_ = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"bitSize is undefined"
  bitSizeMaybe :: Vector Bit -> Maybe Int
bitSizeMaybe Vector Bit
_ = Maybe Int
forall a. Maybe a
Nothing
  isSigned :: Vector Bit -> Bool
isSigned Vector Bit
_ = Bool
False
  zeroBits :: Vector Bit
zeroBits = Vector Bit
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v = Bool
False
    | Bool
otherwise = Bit -> Bool
unBit (Vector Bit -> Int -> Bit
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v = Vector Bit
v
    | Bool
otherwise = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
      MVector s Bit
u <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
v
      MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MU.unsafeWrite MVector s Bit
MVector (PrimState (ST s)) Bit
u Int
n (Bool -> Bit
Bit Bool
True)
      MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
MVector (PrimState (ST s)) Bit
u

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

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

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

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

  rotate :: Vector Bit -> Int -> Vector Bit
rotate Vector Bit
v Int
n'
    | Vector Bit -> Bool
forall a. Unbox a => Vector a -> Bool
U.null Vector Bit
v = Vector Bit
v
    | Bool
otherwise = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
      let l :: Int
l = Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
v
          n :: Int
n = Int
n' Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
l
      MVector s Bit
u <- Int -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
l
      MVector (PrimState (ST s)) Bit -> Vector Bit -> ST s ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
U.copy (Int -> MVector s Bit -> MVector s Bit
forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.drop Int
n MVector s Bit
u) (Int -> Vector Bit -> Vector Bit
forall a. Unbox a => Int -> Vector a -> Vector a
U.take (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Vector Bit
v)
      MVector (PrimState (ST s)) Bit -> Vector Bit -> ST s ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
U.copy (Int -> MVector s Bit -> MVector s Bit
forall a s. Unbox a => Int -> MVector s a -> MVector s a
MU.take Int
n MVector s Bit
u) (Int -> Vector Bit -> Vector Bit
forall a. Unbox a => Int -> Vector a -> Vector a
U.drop (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Vector Bit
v)
      MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
MVector (PrimState (ST 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 (Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
off) (Int -> Int
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 =
    Vector Word -> Maybe (Vector Word)
forall a. a -> Maybe a
Just (Vector Word -> Maybe (Vector Word))
-> Vector Word -> Maybe (Vector Word)
forall a b. (a -> b) -> a -> b
$ Vector Word -> Vector Word
fromPrimVector (Vector Word -> Vector Word) -> Vector Word -> Vector Word
forall a b. (a -> b) -> a -> b
$ Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
s) (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
n) ByteArray
ws
  | Bool
otherwise = Maybe (Vector Word)
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 s. ST s (Vector Word)) -> Vector Word
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Word)) -> Vector Word)
-> (forall s. ST s (Vector Word)) -> Vector Word
forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
v' <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
v
  MVector s Word
w  <- MVector (PrimState (ST s)) Bit
-> ST s (MVector (PrimState (ST s)) Word)
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> m (MVector (PrimState m) Word)
cloneToWordsM MVector s Bit
MVector (PrimState (ST s)) Bit
v'
  MVector (PrimState (ST s)) Word -> ST s (Vector Word)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Word
MVector (PrimState (ST 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 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3) (Int
len Int -> Int -> Int
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 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0, Int
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
  = Vector Word8 -> Maybe (Vector Word8)
forall a. a -> Maybe a
Just (Vector Word8 -> Maybe (Vector Word8))
-> Vector Word8 -> Maybe (Vector Word8)
forall a b. (a -> b) -> a -> b
$ Vector Word8 -> Vector Word8
UB.V_Word8 (Vector Word8 -> Vector Word8) -> Vector Word8 -> Vector Word8
forall a b. (a -> b) -> a -> b
$ Int -> Int -> ByteArray -> Vector Word8
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
s Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) (Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3) ByteArray
ws
  | Bool
otherwise = Maybe (Vector Word8)
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 s. ST s (Vector Word8)) -> Vector Word8
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Word8)) -> Vector Word8)
-> (forall s. ST s (Vector Word8)) -> Vector Word8
forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
v' <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw Vector Bit
v
  MVector s Word8
w  <- MVector (PrimState (ST s)) Bit
-> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> m (MVector (PrimState m) Word8)
cloneToWords8M MVector s Bit
MVector (PrimState (ST s)) Bit
v'
  MVector (PrimState (ST s)) Word8 -> ST s (Vector Word8)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Word8
MVector (PrimState (ST 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
  (Vector Word8 -> Vector Bit)
-> (ByteString -> Vector Word8) -> ByteString -> Vector Bit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word8 -> Vector Word8
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
U.convert
  (Vector Word8 -> Vector Word8)
-> (ByteString -> Vector Word8) -> ByteString -> Vector Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ForeignPtr Word8 -> Int -> Int -> Vector Word8)
-> (ForeignPtr Word8, Int, Int) -> Vector Word8
forall a b c d. (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 ForeignPtr Word8 -> Int -> Int -> Vector Word8
forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
S.unsafeFromForeignPtr
  ((ForeignPtr Word8, Int, Int) -> Vector Word8)
-> (ByteString -> (ForeignPtr Word8, Int, Int))
-> ByteString
-> Vector Word8
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
  = (ForeignPtr Word8 -> Int -> Int -> ByteString)
-> (ForeignPtr Word8, Int, Int) -> ByteString
forall a b c d. (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 ForeignPtr Word8 -> Int -> Int -> ByteString
BS.fromForeignPtr
  ((ForeignPtr Word8, Int, Int) -> ByteString)
-> (Vector Bit -> (ForeignPtr Word8, Int, Int))
-> Vector Bit
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word8 -> (ForeignPtr Word8, Int, Int)
forall a. Vector a -> (ForeignPtr a, Int, Int)
S.unsafeToForeignPtr
  (Vector Word8 -> (ForeignPtr Word8, Int, Int))
-> (Vector Bit -> Vector Word8)
-> Vector Bit
-> (ForeignPtr Word8, Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word8 -> Vector Word8
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
U.convert
  (Vector Word8 -> Vector Word8)
-> (Vector Bit -> Vector Word8) -> Vector Bit -> Vector Word8
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 3500x (!) faster.
--
-- Note: If one input is larger than the other, the remaining bits will be ignored.
--
-- For sufficiently dense sets, represented as bitmaps,
-- 'zipBits' is up to 64x faster than
-- 'Data.IntSet.union', 'Data.IntSet.intersection', etc.
--
-- The function passed to zipBits may only use the following
-- 'Bits' methods:
--
-- '.&.', '.|.', 'xor', 'complement', 'zeroBits', and (likely uselessly)
-- 'bitSizeMaybe' and 'isSigned'.
--
-- >>> :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
f = \Vector Bit
xs Vector Bit
ys -> case (Vector Bit
xs, Vector Bit
ys) of
  (BitVec Int
_ Int
0 ByteArray
_, !Vector Bit
_) -> Vector Bit
forall a. Unbox a => Vector a
U.empty
  (Vector Bit
_, BitVec Int
_ Int
0 ByteArray
_) -> Vector Bit
forall a. Unbox a => Vector a
U.empty
#if UseSIMD
  (BitVec Int
0 Int
l1 ByteArray
arg1, BitVec Int
0 Int
l2 ByteArray
arg2) -> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
    let
        l :: Int
l = Int -> Int -> Int
noinlineMin Int
l1 Int
l2
        w :: Int
w = Int -> Int
nWords Int
l
        b :: Int
b = Int -> Int
wordsToBytes Int
w
    MutableByteArray s
brr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
b
    -- We used to calculate (f False False, f False True, f True False, f True True).
    -- Now we calculate all those in one go by passing all four possibilities within
    -- a word.
    case Word
0b1111 Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. (Bitsy -> Word
unBitsy (Bitsy -> Word) -> Bitsy -> Word
forall a b. (a -> b) -> a -> b
$ Bitsy -> Bitsy -> Bitsy
forall a. Bits a => a -> a -> a
f (Word -> Bitsy
Bitsy Word
0b0011) (Word -> Bitsy
Bitsy Word
0b0101)) of
     Word
0b0000 -> MutableByteArray (PrimState (ST s))
-> Int -> Int -> Word -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> Int -> a -> m ()
setByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
brr Int
0 Int
w (Word
forall a. Bits a => a
zeroBits :: Word)
     Word
0b0001 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompAnd  MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b0010 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompAndn MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b0011 -> MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
brr Int
0 ByteArray
arg1 Int
0 Int
b
     Word
0b0100 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompAndn MutableByteArray s
brr ByteArray
arg2 ByteArray
arg1 Int
b
     Word
0b0101 -> MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
brr Int
0 ByteArray
arg2 Int
0 Int
b
     Word
0b0110 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompXor  MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b0111 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompIor  MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b1000 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompNior MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b1001 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompXnor MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b1010 -> MutableByteArray s -> ByteArray -> Int -> ST s ()
forall s. MutableByteArray s -> ByteArray -> Int -> ST s ()
ompCom  MutableByteArray s
brr ByteArray
arg2      Int
b
     Word
0b1011 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompIorn MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
0b1100 -> MutableByteArray s -> ByteArray -> Int -> ST s ()
forall s. MutableByteArray s -> ByteArray -> Int -> ST s ()
ompCom  MutableByteArray s
brr ByteArray
arg1      Int
b
     Word
0b1101 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompIorn MutableByteArray s
brr ByteArray
arg2 ByteArray
arg1 Int
b
     Word
0b1110 -> MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
forall s.
MutableByteArray s -> ByteArray -> ByteArray -> Int -> ST s ()
ompNand MutableByteArray s
brr ByteArray
arg1 ByteArray
arg2 Int
b
     Word
_0b1111 -> MutableByteArray (PrimState (ST s))
-> Int -> Int -> Word -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> Int -> a -> m ()
setByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
brr Int
0 Int
w (Word -> Word
forall a. Bits a => a -> a
complement Word
forall a. Bits a => a
zeroBits :: Word)
    Int -> Int -> ByteArray -> Vector Bit
BitVec Int
0 Int
l (ByteArray -> Vector Bit) -> ST s ByteArray -> ST s (Vector Bit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
brr
#endif
  (Vector Bit, Vector Bit)
_ -> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
    let n :: Int
n = Int -> Int -> Int
noinlineMin (Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs) (Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
ys)
    MVector s Bit
zs <- Int -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
n
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0, Int
wordSize .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      MVector (PrimState (ST s)) Bit -> Int -> Word -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
MVector (PrimState (ST s)) Bit
zs Int
i (Word -> ST s ()) -> (Bitsy -> Word) -> Bitsy -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bitsy -> Word
unBitsy (Bitsy -> ST s ()) -> Bitsy -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bitsy -> Bitsy -> Bitsy
forall a. Bits a => a -> a -> a
f (Word -> Bitsy
Bitsy (Word -> Bitsy) -> Word -> Bitsy
forall a b. (a -> b) -> a -> b
$ Vector Bit -> Int -> Word
indexWord Vector Bit
xs Int
i) (Word -> Bitsy
Bitsy (Word -> Bitsy) -> Word -> Bitsy
forall a b. (a -> b) -> a -> b
$ Vector Bit -> Int -> Word
indexWord Vector Bit
ys Int
i)
    MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
MVector (PrimState (ST s)) Bit
zs
{-# INLINE zipBits #-}

-- | This is hideous, but it keeps the code size down in applications of
-- 'zipBits'. Otherwise we end up taking different code paths depending
-- on how the comparison goes in the min calculation, and the Core gets
-- seriously ugly. Ugh!
noinlineMin :: Int -> Int -> Int
noinlineMin :: Int -> Int -> Int
noinlineMin = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
{-# NOINLINE noinlineMin #-}

-- | A version of 'Word' that only supports operations that make sense in
-- zipBits. This ensures that if someone does something overly silly in the function
-- they pass to zipBits, then they'll get a helpful (albeit run-time) error rather than just
-- weird garbage results.
newtype Bitsy = Bitsy {Bitsy -> Word
unBitsy :: Word}
instance Eq Bitsy where
  Bitsy
_ == :: Bitsy -> Bitsy -> Bool
== Bitsy
_ = [Char] -> Bool
forall a. [Char] -> a
notBitsy [Char]
"=="
instance Bits Bitsy where
  Bitsy Word
x .&. :: Bitsy -> Bitsy -> Bitsy
.&. Bitsy Word
y = Word -> Bitsy
Bitsy (Word
x Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
y)
  Bitsy Word
x .|. :: Bitsy -> Bitsy -> Bitsy
.|. Bitsy Word
y = Word -> Bitsy
Bitsy (Word
x Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. Word
y)
  Bitsy Word
x xor :: Bitsy -> Bitsy -> Bitsy
`xor` Bitsy Word
y = Word -> Bitsy
Bitsy (Word
x Word -> Word -> Word
forall a. Bits a => a -> a -> a
`xor` Word
y)
  complement :: Bitsy -> Bitsy
complement (Bitsy Word
x) = Word -> Bitsy
Bitsy (Word -> Word
forall a. Bits a => a -> a
complement Word
x)
  zeroBits :: Bitsy
zeroBits = Word -> Bitsy
Bitsy Word
forall a. Bits a => a
zeroBits
  bitSizeMaybe :: Bitsy -> Maybe Int
bitSizeMaybe Bitsy
_ = Maybe Int
forall a. Maybe a
Nothing
  isSigned :: Bitsy -> Bool
isSigned Bitsy
_ = Bool
False  -- Not useful, but not harmful
  {-# INLINE (.&.) #-}
  {-# INLINE (.|.) #-}
  {-# INLINE xor #-}
  {-# INLINE complement #-}
  {-# INLINE zeroBits #-}

  shiftL :: Bitsy -> Int -> Bitsy
shiftL Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"shiftL"
  shiftR :: Bitsy -> Int -> Bitsy
shiftR Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"shiftR"
  shift :: Bitsy -> Int -> Bitsy
shift Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"shift"
  unsafeShiftL :: Bitsy -> Int -> Bitsy
unsafeShiftL Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"unsafeShiftL"
  unsafeShiftR :: Bitsy -> Int -> Bitsy
unsafeShiftR Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"unsafeShiftR"
  rotateL :: Bitsy -> Int -> Bitsy
rotateL Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"rotateL"
  rotateR :: Bitsy -> Int -> Bitsy
rotateR Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"rotateR"
  rotate :: Bitsy -> Int -> Bitsy
rotate Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"rotate"
  bitSize :: Bitsy -> Int
bitSize Bitsy
_ = [Char] -> Int
forall a. [Char] -> a
notBitsy [Char]
"bitSize"
  testBit :: Bitsy -> Int -> Bool
testBit Bitsy
_ Int
_ = [Char] -> Bool
forall a. [Char] -> a
notBitsy [Char]
"testBit"
  bit :: Int -> Bitsy
bit Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"bit"
  setBit :: Bitsy -> Int -> Bitsy
setBit Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"setBit"
  clearBit :: Bitsy -> Int -> Bitsy
clearBit Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"clearBit"
  complementBit :: Bitsy -> Int -> Bitsy
complementBit Bitsy
_ Int
_ = [Char] -> Bitsy
forall a. [Char] -> a
notBitsy [Char]
"complementBit"
  popCount :: Bitsy -> Int
popCount Bitsy
_ = [Char] -> Int
forall a. [Char] -> a
notBitsy [Char]
"popCount"

{-# NOINLINE notBitsy #-}
notBitsy :: String -> a
notBitsy :: forall a. [Char] -> a
notBitsy [Char]
fun = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$
  [Char]
"The function passed to zipBits may only use\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++
  [Char]
".&., .|., xor, complement, zeroBits, bitSizeMaybe, and isSigned.\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++
  [Char]
"You used " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
fun

-- | 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 (Bit -> Bit
forall a. Bits a => a -> a
f (Bool -> Bit
Bit Bool
False)), Bit -> Bool
unBit (Bit -> Bit
forall a. Bits a => a -> a
f (Bool -> Bit
Bit Bool
True))) of
  (Bool
False, Bool
False) -> (Int -> Bit -> Vector Bit
forall a. Unbox a => Int -> a -> Vector a
`U.replicate` Bool -> Bit
Bit Bool
False) (Int -> Vector Bit)
-> (Vector Bit -> Int) -> Vector Bit -> Vector Bit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length
  (Bool
False, Bool
True)  -> Vector Bit -> Vector Bit
forall a. a -> a
id
  (Bool
True, Bool
False)  -> Vector Bit -> Vector Bit
invertBits
  (Bool
True, Bool
True)   -> (Int -> Bit -> Vector Bit
forall a. Unbox a => Int -> a -> Vector a
`U.replicate` Bool -> Bit
Bit Bool
True) (Int -> Vector Bit)
-> (Vector Bit -> Int) -> Vector Bit -> Vector Bit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length
{-# INLINE mapBits #-}

-- | Invert (flip) all bits.
--
-- >>> :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
_) = Vector Bit
forall a. Unbox a => Vector a
U.empty
#if UseSIMD
invertBits (BitVec Int
0 Int
l ByteArray
arg) = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  let w :: Int
w = Int -> Int
nWords Int
l
      b :: Int
b = Int -> Int
wordsToBytes Int
w
  MutableByteArray s
brr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
b
  MutableByteArray s -> ByteArray -> Int -> ST s ()
forall s. MutableByteArray s -> ByteArray -> Int -> ST s ()
ompCom MutableByteArray s
brr ByteArray
arg Int
b
  Int -> Int -> ByteArray -> Vector Bit
BitVec Int
0 Int
l (ByteArray -> Vector Bit) -> ST s ByteArray -> ST s (Vector Bit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
brr
#endif
invertBits Vector Bit
xs = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs
  MVector s Bit
ys <- Int -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
n
  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0, Int
wordSize .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    MVector (PrimState (ST s)) Bit -> Int -> Word -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
MVector (PrimState (ST s)) Bit
ys Int
i (Word -> Word
forall a. Bits a => a -> a
complement (Vector Bit -> Int -> Word
indexWord Vector Bit
xs Int
i))
  MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze MVector s Bit
MVector (PrimState (ST s)) Bit
ys

-- | For each set bit of the first argument, extract
-- the corresponding bit of the second argument
-- to the result. Similar to the
-- [parallel bit extract instruction (PEXT)](https://en.wikipedia.org/wiki/X86_Bit_manipulation_instruction_set#Parallel_bit_deposit_and_extract).
--
-- Note: If one input is larger than the other, the remaining bits will be ignored.
--
-- >>> :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
#ifdef UseSIMD
selectBits :: Vector Bit -> Vector Bit -> Vector Bit
selectBits (BitVec Int
0 Int
iLen ByteArray
iArr) (BitVec Int
0 Int
xLen ByteArray
xArr) | Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
marr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
len Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3)
  Int
n <- MutableByteArray s
-> ByteArray -> ByteArray -> Int -> Bool -> ST s Int
forall s.
MutableByteArray s
-> ByteArray -> ByteArray -> Int -> Bool -> ST s Int
selectBitsC MutableByteArray s
marr ByteArray
xArr ByteArray
iArr (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
len) Bool
False
  Int -> Int -> ByteArray -> Vector Bit
BitVec Int
0 Int
n (ByteArray -> Vector Bit) -> ST s ByteArray -> ST s (Vector Bit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
marr
 where
  len :: Int
len = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
iLen Int
xLen
#endif
selectBits Vector Bit
is Vector Bit
xs = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
xs1 <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
xs
  Int
n   <- Vector Bit -> MVector (PrimState (ST s)) Bit -> ST s Int
forall (m :: * -> *).
PrimMonad m =>
Vector Bit -> MVector (PrimState m) Bit -> m Int
selectBitsInPlace Vector Bit
is MVector s Bit
MVector (PrimState (ST s)) Bit
xs1
  MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (Int -> MVector s Bit -> MVector s Bit
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, extract
-- the corresponding bit of the second argument
-- to the result.
--
-- Note: If one input is larger than the other, the remaining bits will be ignored.
--
-- >>> :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
#ifdef UseSIMD
excludeBits :: Vector Bit -> Vector Bit -> Vector Bit
excludeBits (BitVec Int
0 Int
iLen ByteArray
iArr) (BitVec Int
0 Int
xLen ByteArray
xArr) | Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
marr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
len Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3)
  Int
n <- MutableByteArray s
-> ByteArray -> ByteArray -> Int -> Bool -> ST s Int
forall s.
MutableByteArray s
-> ByteArray -> ByteArray -> Int -> Bool -> ST s Int
selectBitsC MutableByteArray s
marr ByteArray
xArr ByteArray
iArr (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
len) Bool
True
  Int -> Int -> ByteArray -> Vector Bit
BitVec Int
0 Int
n (ByteArray -> Vector Bit) -> ST s ByteArray -> ST s (Vector Bit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
marr
 where
  len :: Int
len = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
iLen Int
xLen
#endif
excludeBits Vector Bit
is Vector Bit
xs = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  MVector s Bit
xs1 <- Vector Bit -> ST s (MVector (PrimState (ST s)) Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector Bit
xs
  Int
n   <- Vector Bit -> MVector (PrimState (ST s)) Bit -> ST s Int
forall (m :: * -> *).
PrimMonad m =>
Vector Bit -> MVector (PrimState m) Bit -> m Int
excludeBitsInPlace Vector Bit
is MVector s Bit
MVector (PrimState (ST s)) Bit
xs1
  MVector (PrimState (ST s)) Bit -> ST s (Vector Bit)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (Int -> MVector s Bit -> MVector s Bit
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
#ifdef UseSIMD
reverseBits :: Vector Bit -> Vector Bit
reverseBits (BitVec Int
0 Int
len ByteArray
arr) | Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
marr <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
len Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3)
  MutableByteArray s -> ByteArray -> Int -> ST s ()
forall s. MutableByteArray s -> ByteArray -> Int -> ST s ()
reverseBitsC MutableByteArray s
marr ByteArray
arr (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
len)
  Int -> Int -> ByteArray -> Vector Bit
BitVec Int
0 Int
len (ByteArray -> Vector Bit) -> ST s ByteArray -> ST s (Vector Bit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
marr
#endif
reverseBits Vector Bit
xs = (forall s. ST s (Vector Bit)) -> Vector Bit
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Bit)) -> Vector Bit)
-> (forall s. ST s (Vector Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n    = Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs
  MVector s Bit
ys <- Int -> ST s (MVector (PrimState (ST s)) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.new Int
n

  [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0, Int
wordSize .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
wordSize] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    MVector (PrimState (ST s)) Bit -> Int -> Word -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
MVector (PrimState (ST s)) Bit
ys (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
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
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
nMod Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    let x :: Word
x = Vector Bit -> Int -> Word
indexWord Vector Bit
xs (Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
n))
    Word
y <- MVector (PrimState (ST s)) Bit -> Int -> ST s Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector s Bit
MVector (PrimState (ST s)) Bit
ys Int
0
    MVector (PrimState (ST s)) Bit -> Int -> Word -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector s Bit
MVector (PrimState (ST s)) Bit
ys Int
0 (Int -> Word -> Word -> Word
meld Int
nMod (Int -> Word -> Word
reversePartialWord Int
nMod Word
x) Word
y)

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

clipLoBits :: Bit -> Int -> Word -> Word
clipLoBits :: Bit -> Int -> Word -> Word
clipLoBits (Bit Bool
True ) Int
k Word
w = Word
w Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
k
clipLoBits (Bit Bool
False) Int
k Word
w = (Word
w Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
k) Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. Int -> Word
hiMask (Int
wordSize Int -> Int -> Int
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 Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
k
clipHiBits (Bit Bool
False) Int
k Word
w = Word
w Word -> Word -> Word
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
#if UseSIMD
bitIndex :: Bit -> Vector Bit -> Maybe Int
bitIndex (Bit Bool
b) (BitVec Int
0 Int
len ByteArray
arr) | Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
  let res :: Int
res = ByteArray -> Int -> Bool -> Int
bitIndexC ByteArray
arr (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
len) Bool
b
  in if Int
res Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
res
#endif
bitIndex Bit
b (BitVec Int
off Int
len ByteArray
arr)
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Maybe Int
forall a. Maybe a
Nothing
  | Int
offBits Int -> Int -> Bool
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr of
      r :: Maybe Int
r@Just{} -> Maybe Int
r
      Maybe Int
Nothing  -> (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (Int -> Int) -> Maybe Int -> Maybe Int
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 (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))
  | Bool
otherwise = case Int -> Int
modWordSize (Int
off Int -> Int -> Int
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 (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords))
        of
          r :: Maybe Int
r@Just{} -> Maybe Int
r
          Maybe Int
Nothing ->
            (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits))
              (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords Bit
b (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
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 (ByteArray -> Int -> Word
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 (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords))
          of
            r :: Maybe Int
r@Just{} -> Maybe Int
r
            Maybe Int
Nothing ->
              (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits))
                (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Bit -> Int -> Int -> ByteArray -> Maybe Int
bitIndexInWords Bit
b (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr of
                      r :: Maybe Int
r@Just{} -> Maybe Int
r
                      Maybe Int
Nothing ->
                        (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)) (Int -> Int) -> Maybe Int -> Maybe Int
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
                            (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
                          )
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits Int -> Int -> Int
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 (Word -> Maybe Int) -> (Word -> Word) -> Word -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Word
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len = Maybe Int
forall a. Maybe a
Nothing
    | Bool
otherwise = case Word -> Maybe Int
ffs (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n) of
      Maybe Int
Nothing  -> Int -> Maybe Int
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Just Int
r  -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off) Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len = Maybe Int
forall a. Maybe a
Nothing
    | Bool
otherwise = case Word -> Maybe Int
ffs (Word -> Word
forall a. Bits a => a -> a
complement (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n)) of
      Maybe Int
Nothing -> Int -> Maybe Int
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Just Int
r  -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off) Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [Char] -> Maybe Int
forall a. HasCallStack => [Char] -> a
error [Char]
"nthBitIndex: n must be positive"
#if UseSIMD
nthBitIndex (Bit Bool
b) Int
n (BitVec Int
0 Int
len ByteArray
arr) | Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
  let res :: Int
res = ByteArray -> Int -> Bool -> Int -> Int
nthBitIndexC ByteArray
arr (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
len) Bool
b Int
n
  in if Int
res Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
res
#endif
nthBitIndex Bit
b Int
k (BitVec Int
off Int
len ByteArray
arr)
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Maybe Int
forall a. Maybe a
Nothing
  | Int
offBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (Int -> Maybe Int)
-> (Int -> Maybe Int) -> Either Int Int -> Maybe Int
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe Int -> Int -> Maybe Int
forall a b. a -> b -> a
const Maybe Int
forall a. Maybe a
Nothing) Int -> Maybe Int
forall a. a -> Maybe a
Just (Either Int Int -> Maybe Int) -> Either Int Int -> Maybe Int
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 Int -> Int -> Int
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'   -> (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (Int -> Int) -> Either Int Int -> Either Int Int
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 (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))
  | Bool
otherwise = (Int -> Maybe Int)
-> (Int -> Maybe Int) -> Either Int Int -> Maybe Int
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe Int -> Int -> Maybe Int
forall a b. a -> b -> a
const Maybe Int
forall a. Maybe a
Nothing) Int -> Maybe Int
forall a. a -> Maybe a
Just (Either Int Int -> Maybe Int) -> Either Int Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ case Int -> Int
modWordSize (Int
off Int -> Int -> Int
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 (ByteArray -> Int -> Word
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' ->
          (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits))
            (Int -> Int) -> Either Int Int -> Either Int Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
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 (ByteArray -> Int -> Word
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 (ByteArray -> Int -> Word
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' ->
              (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits))
                (Int -> Int) -> Either Int Int -> Either Int Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
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''  -> (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)) (Int -> Int) -> Either Int Int -> Either Int Int
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
                          (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
                        )
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
c then Int -> Either Int Int
forall a b. a -> Either a b
Left (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c) else Int -> Either Int Int
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 Word -> Word
forall a. Bits a => a -> a
complement Word
v
  c :: Int
c = Word -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len = Int -> Either Int Int
forall a b. a -> Either a b
Left Int
l
    | Bool
otherwise = if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
c
      then Int -> Int -> Either Int Int
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c)
      else Int -> Either Int Int
forall a b. b -> Either a b
Right (Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Word -> Int
unsafeNthTrueInWord Int
l Word
w)
   where
    w :: Word
w = ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n
    c :: Int
c = Word -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len = Int -> Either Int Int
forall a b. a -> Either a b
Left Int
l
    | Bool
otherwise = if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
c
      then Int -> Int -> Either Int Int
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c)
      else Int -> Either Int Int
forall a b. b -> Either a b
Right (Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Word -> Int
unsafeNthTrueInWord Int
l Word
w)
   where
    w :: Word
w = Word -> Word
forall a. Bits a => a -> a
complement (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n)
    c :: Int
c = Word -> Int
forall a. Bits a => a -> Int
popCount Word
w

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

-- | Return the number of set bits in a vector (population count, popcount).
--
-- >>> :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 UseSIMD
countBits (BitVec Int
0 Int
len ByteArray
arr) | Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
  ByteArray -> Int -> Int
ompPopcount ByteArray
arr (Int
len Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
5)
#endif
countBits (BitVec Int
off Int
len ByteArray
arr) | Int
offBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = case Int -> Int
modWordSize Int
len of
  Int
0    -> Vector Word -> Int
countBitsInWords (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords Int
lWords ByteArray
arr)
  Int
nMod -> Vector Word -> Int
countBitsInWords (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
    Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word -> Int
forall a. Bits a => a -> Int
popCount (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
nMod)
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
countBits (BitVec Int
off Int
len ByteArray
arr) = case Int -> Int
modWordSize (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) of
  Int
0 -> Word -> Int
forall a. Bits a => a -> Int
popCount (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word)
    Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Word -> Int
countBitsInWords (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
  Int
nMod -> case Int
lWords of
    Int
1 -> Word -> Int
forall a. Bits a => a -> Int
popCount
      ((ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits) Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
len)
    Int
_ ->
      Word -> Int
forall a. Bits a => a -> Int
popCount (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word)
        Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Word -> Int
countBitsInWords (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr)
        Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word -> Int
forall a. Bits a => a -> Int
popCount (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Int -> Word
loMask Int
nMod)
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)

countBitsInWords :: P.Vector Word -> Int
countBitsInWords :: Vector Word -> Int
countBitsInWords = (Int -> Word -> Int) -> Int -> Vector Word -> Int
forall b a. Prim b => (a -> b -> a) -> a -> Vector b -> a
P.foldl' (\Int
acc Word
word -> Word -> Int
forall a. Bits a => a -> Int
popCount Word
word Int -> Int -> Int
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 Int -> Int -> Bool
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 (Int -> Int -> ByteArray -> Vector Word
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 (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector Int
offWords (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
      ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
      ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) :: Word))
               [Int
0 .. Int
nMod Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
listBits (BitVec Int
off Int
len ByteArray
arr) = case Int -> Int
modWordSize (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) of
  Int
0 ->
    (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter
        (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word))
        [Int
0 .. Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> Vector Word -> [Int] -> [Int]
listBitsInWords (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits)
                         (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ByteArray
arr)
                         []
  Int
nMod -> case Int
lWords of
    Int
1 -> (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter
      (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word))
      [Int
0 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    Int
_ ->
      (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter
          (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
offWords Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
offBits :: Word))
          [Int
0 .. Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ ( Int -> Vector Word -> [Int] -> [Int]
listBitsInWords (Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits)
                             (Int -> Int -> ByteArray -> Vector Word
forall a. Int -> Int -> ByteArray -> Vector a
P.Vector (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) ByteArray
arr)
           ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int -> Int
forall a. Bits a => a -> a
mulWordSize (Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offBits))
           ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter
               (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr (Int
offWords Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) :: Word))
               [Int
0 .. Int
nMod Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
           )
 where
  offBits :: Int
offBits  = Int -> Int
modWordSize Int
off
  offWords :: Int
offWords = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
off
  lWords :: Int
lWords   = Int -> Int
nWords (Int
offBits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)

listBitsInWord :: Int -> Word -> [Int]
listBitsInWord :: Int -> Word -> [Int]
listBitsInWord Int
offset Word
word =
  (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word
word) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Int
0 .. Int
wordSize Int -> Int -> Int
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 = ([Int] -> Vector Word -> [Int]) -> Vector Word -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (([Int] -> Vector Word -> [Int]) -> Vector Word -> [Int] -> [Int])
-> ([Int] -> Vector Word -> [Int]) -> Vector Word -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Word -> [Int] -> [Int]) -> [Int] -> Vector Word -> [Int]
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
i) Word
word [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
acc)