{-# 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.Mutable as MU
import Data.Word
import Unsafe.Coerce

#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
_ = [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]
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 a vector of bits is aligned.
-- Use 'cloneToWords' otherwise.
-- Cf. 'Data.Bit.castToWordsM'.
--
-- > castToWords (castFromWords v) == Just v
castToWords :: U.Vector Bit -> Maybe (U.Vector Word)
castToWords :: Vector Bit -> Maybe (Vector Word)
castToWords (BitVec s n 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]
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
{-# INLINE cloneToWords #-}

-- | Cast a unboxed vector of 'Word8'
-- to an unboxed vector of bits.
--
-- On big-endian architectures 'castFromWords8'
-- resorts to copying instead of aliasing underlying arrays.
--
-- >>> :set -XOverloadedLists
-- >>> castFromWords8 [123]
-- [1,1,0,1,1,1,1,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
    P.Vector off' len arr' = unsafeCoerce 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
    P.Vector Int
off Int
len ByteArray
arr = Vector Word8 -> Vector Any
forall a b. a -> b
unsafeCoerce Vector Word8
ws
#endif

-- | Try to cast an unboxed vector of bits
-- to an unboxed vector of 'Word8'.
-- It succeeds if a vector of bits is aligned.
-- Use 'Data.Bit.cloneToWords8' otherwise.
--
-- > castToWords8 (castFromWords8 v) == Just v
castToWords8 :: U.Vector Bit -> Maybe (U.Vector Word8)
#ifdef WORDS_BIGENDIAN
castToWords8 = const Nothing
#else
castToWords8 :: Vector Bit -> Maybe (Vector Word8)
castToWords8 (BitVec s n 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 Any -> Vector Word8
forall a b. a -> b
unsafeCoerce (Vector Any -> Vector Word8) -> Vector Any -> Vector Word8
forall a b. (a -> b) -> a -> b
$ Int -> Int -> ByteArray -> Vector Any
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]
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
{-# INLINE 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]
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#"
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. Storable 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 :: (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
-- flag @libgmp@ 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]
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 _ 0 _) Vector Bit
_ = Vector Bit
forall a. Unbox a => Vector a
U.empty
zipBits forall a. Bits a => a -> a -> a
_ Vector Bit
_ (BitVec _ 0 _) = Vector Bit
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 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
forall a. Ord a => a -> a -> a
min (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 -> Word -> Word
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))
  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 #-}

-- | 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]
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 Vector Bit
xs = 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 (Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs) (Bool -> Bit
Bit Bool
False)
  (Bool
False, Bool
True)  -> Vector Bit
xs
  (Bool
True, Bool
False)  -> Vector Bit -> Vector Bit
invertBits Vector Bit
xs
  (Bool
True, Bool
True)   -> Int -> Bit -> Vector Bit
forall a. Unbox a => Int -> a -> Vector a
U.replicate (Vector Bit -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector Bit
xs) (Bool -> Bit
Bit Bool
True)
{-# INLINE mapBits #-}

-- | Invert (flip) all bits.
--
-- Users are strongly encouraged to enable
-- flag @libgmp@ for the ultimate performance of 'invertBits'.
--
-- >>> :set -XOverloadedLists
-- >>> invertBits [0,1,0,1,0]
-- [1,0,1,0,1]
invertBits
  :: U.Vector Bit
  -> U.Vector Bit
invertBits :: Vector Bit -> Vector Bit
invertBits (BitVec _ 0 _) = Vector Bit
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 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, deposit
-- the corresponding bit of the second argument
-- to the result. Similar to the parallel deposit instruction (PDEP).
--
-- >>> :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))
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 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, 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))
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 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 @vector-rotcev@ package
-- to reverse vectors in O(1) time.
reverseBits :: U.Vector Bit -> U.Vector Bit
reverseBits :: Vector Bit -> Vector Bit
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
bitIndex :: Bit -> U.Vector Bit -> Maybe Int
bitIndex :: Bit -> Vector Bit -> Maybe Int
bitIndex Bit
b (BitVec off len 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>.
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"
nthBitIndex Bit
b Int
k (BitVec off len 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).
--
-- Users are strongly encouraged to enable
-- flag @libgmp@ 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>.
countBits :: U.Vector Bit -> Int
countBits :: Vector Bit -> Int
countBits (BitVec _ 0 _)                      = Int
0
#if UseLibGmp
countBits (BitVec 0 len arr) | modWordSize len == 0 =
  fromIntegral (mpnPopcount arr (divWordSize len))
#endif
countBits (BitVec off len 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 off len 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]
listBits :: U.Vector Bit -> [Int]
listBits :: Vector Bit -> [Int]
listBits (BitVec _ 0 _)                      = []
listBits (BitVec off len 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 off len 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)