{-# LANGUAGE CPP              #-}

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

#ifndef BITVEC_THREADSAFE
module Data.Bit.Mutable
#else
module Data.Bit.MutableTS
#endif
  ( castFromWordsM
  , castToWordsM
  , cloneToWordsM

  , cloneToWords8M

  , zipInPlace
  , mapInPlace

  , invertInPlace
  , selectBitsInPlace
  , excludeBitsInPlace

  , reverseInPlace
  ) where

#include "MachDeps.h"

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
#ifndef BITVEC_THREADSAFE
import Data.Bit.Internal
#else
import Data.Bit.InternalTS
#endif
import Data.Bit.Utils
import Data.Bits
import Data.Primitive.ByteArray
import qualified Data.Vector.Primitive as P
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import Data.Word

#ifdef WORDS_BIGENDIAN
import GHC.Exts
#endif

-- | Cast a vector of words to a vector of bits.
-- Cf. 'Data.Bit.castFromWords'.
castFromWordsM :: MVector s Word -> MVector s Bit
castFromWordsM :: MVector s Word -> MVector s Bit
castFromWordsM (MU.MV_Word (P.MVector off len ws)) =
  Int -> Int -> MutableByteArray s -> MVector s Bit
forall s. Int -> Int -> MutableByteArray s -> MVector s Bit
BitMVec (Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
off) (Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
len) MutableByteArray s
ws

-- | Try to cast a vector of bits to a vector of words.
-- It succeeds if a vector of bits is aligned.
-- Use 'cloneToWordsM' otherwise.
-- Cf. 'Data.Bit.castToWords'.
castToWordsM :: MVector s Bit -> Maybe (MVector s Word)
castToWordsM :: MVector s Bit -> Maybe (MVector s Word)
castToWordsM (BitMVec s n ws)
  | Int -> Bool
aligned Int
s, Int -> Bool
aligned Int
n
  = MVector s Word -> Maybe (MVector s Word)
forall a. a -> Maybe a
Just (MVector s Word -> Maybe (MVector s Word))
-> MVector s Word -> Maybe (MVector s Word)
forall a b. (a -> b) -> a -> b
$ MVector s Word -> MVector s Word
forall s. MVector s Word -> MVector s Word
MU.MV_Word (MVector s Word -> MVector s Word)
-> MVector s Word -> MVector s Word
forall a b. (a -> b) -> a -> b
$ Int -> Int -> MutableByteArray s -> MVector s Word
forall s a. Int -> Int -> MutableByteArray s -> MVector s a
P.MVector (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
s) (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
n) MutableByteArray s
ws
  | Bool
otherwise = Maybe (MVector s Word)
forall a. Maybe a
Nothing

-- | Clone a 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.cloneToWords'.
cloneToWordsM
  :: PrimMonad m
  => MVector (PrimState m) Bit
  -> m (MVector (PrimState m) Word)
cloneToWordsM :: MVector (PrimState m) Bit -> m (MVector (PrimState m) Word)
cloneToWordsM MVector (PrimState m) Bit
v = do
  let lenBits :: Int
lenBits  = MVector (PrimState m) Bit -> Int
forall a s. Unbox a => MVector s a -> Int
MU.length MVector (PrimState m) Bit
v
      lenWords :: Int
lenWords = Int -> Int
nWords Int
lenBits
  w :: MVector (PrimState m) Bit
w@(BitMVec _ _ arr) <- Int -> m (MVector (PrimState m) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.unsafeNew (Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
lenWords)
  MVector (PrimState m) Bit -> MVector (PrimState m) Bit -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
MU.unsafeCopy (Int
-> Int -> MVector (PrimState m) Bit -> MVector (PrimState m) Bit
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
0 Int
lenBits MVector (PrimState m) Bit
w) MVector (PrimState m) Bit
v
  MVector (PrimState m) Bit -> Bit -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MU.set (Int
-> Int -> MVector (PrimState m) Bit -> MVector (PrimState m) Bit
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
lenBits (Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
lenWords Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lenBits) MVector (PrimState m) Bit
w) (Bool -> Bit
Bit Bool
False)
  MVector (PrimState m) Word -> m (MVector (PrimState m) Word)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector (PrimState m) Word -> m (MVector (PrimState m) Word))
-> MVector (PrimState m) Word -> m (MVector (PrimState m) Word)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) Word -> MVector (PrimState m) Word
forall s. MVector s Word -> MVector s Word
MU.MV_Word (MVector (PrimState m) Word -> MVector (PrimState m) Word)
-> MVector (PrimState m) Word -> MVector (PrimState m) Word
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MutableByteArray (PrimState m)
-> MVector (PrimState m) Word
forall s a. Int -> Int -> MutableByteArray s -> MVector s a
P.MVector Int
0 Int
lenWords MutableByteArray (PrimState m)
arr
{-# INLINE cloneToWordsM #-}

-- | Clone a vector of bits to a new unboxed vector of 'Word8'.
-- If the bits don't completely fill the words, the last 'Word8' will be zero-padded.
-- Cf. 'Data.Bit.cloneToWords8'.
cloneToWords8M
  :: PrimMonad m
  => MVector (PrimState m) Bit
  -> m (MVector (PrimState m) Word8)
cloneToWords8M :: MVector (PrimState m) Bit -> m (MVector (PrimState m) Word8)
cloneToWords8M MVector (PrimState m) Bit
v = do
  let lenBits :: Int
lenBits  = MVector (PrimState m) Bit -> Int
forall a s. Unbox a => MVector s a -> Int
MU.length MVector (PrimState m) Bit
v
      -- Take care about big-endian architectures: allocate full words!
      actualLenBytes :: Int
actualLenBytes = (Int
lenBits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
      roundedLenBytes :: Int
roundedLenBytes = Int -> Int
wordsToBytes (Int -> Int
nWords Int
lenBits)
  ws :: MVector (PrimState m) Bit
ws@(BitMVec _ _ arr) <- Int -> m (MVector (PrimState m) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MU.unsafeNew (Int
roundedLenBytes Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3)
  MVector (PrimState m) Bit -> MVector (PrimState m) Bit -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
MU.unsafeCopy (Int
-> Int -> MVector (PrimState m) Bit -> MVector (PrimState m) Bit
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
0 Int
lenBits MVector (PrimState m) Bit
ws) MVector (PrimState m) Bit
v
  MVector (PrimState m) Bit -> Bit -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MU.set (Int
-> Int -> MVector (PrimState m) Bit -> MVector (PrimState m) Bit
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MU.slice Int
lenBits (Int
roundedLenBytes Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
3 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lenBits) MVector (PrimState m) Bit
ws) (Bool -> Bit
Bit Bool
False)

#ifdef WORDS_BIGENDIAN
  forM_ [0..nWords lenBits - 1] $ \i -> do
    W# w <- readByteArray arr i
    writeByteArray arr i (W# (byteSwap# w))
#endif

  MVector (PrimState m) Word8 -> m (MVector (PrimState m) Word8)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MVector (PrimState m) Word8 -> m (MVector (PrimState m) Word8))
-> MVector (PrimState m) Word8 -> m (MVector (PrimState m) Word8)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) Word8 -> MVector (PrimState m) Word8
forall s. MVector s Word8 -> MVector s Word8
MU.MV_Word8 (MVector (PrimState m) Word8 -> MVector (PrimState m) Word8)
-> MVector (PrimState m) Word8 -> MVector (PrimState m) Word8
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MutableByteArray (PrimState m)
-> MVector (PrimState m) Word8
forall s a. Int -> Int -> MutableByteArray s -> MVector s a
P.MVector Int
0 Int
actualLenBytes MutableByteArray (PrimState m)
arr
{-# INLINE cloneToWords8M #-}

-- | Zip two vectors with the given function.
-- rewriting contents of the second argument.
-- Cf. 'Data.Bit.zipBits'.
--
-- >>> :set -XOverloadedLists
-- >>> import Data.Bits
-- >>> Data.Vector.Unboxed.modify (zipInPlace (.&.) [1,1,0]) [0,1,1]
-- [0,1,0]
--
-- __Warning__: if the immutable vector is shorter than the mutable one,
-- it is a caller's responsibility to trim the result:
--
-- >>> :set -XOverloadedLists
-- >>> import Data.Bits
-- >>> Data.Vector.Unboxed.modify (zipInPlace (.&.) [1,1,0]) [0,1,1,1,1,1]
-- [0,1,0,1,1,1] -- note trailing garbage
zipInPlace
  :: forall m.
     PrimMonad m
  => (forall a . Bits a => a -> a -> a)
  -> Vector Bit
  -> MVector (PrimState m) Bit
  -> m ()
zipInPlace :: (forall a. Bits a => a -> a -> a)
-> Vector Bit -> MVector (PrimState m) Bit -> m ()
zipInPlace forall a. Bits a => a -> a -> a
f (BitVec off l xs) (BitMVec off' l' ys) =
  Int -> Int -> Int -> m ()
go (Int
l Int -> Int -> Int
forall a. Ord a => a -> a -> a
`min` Int
l') Int
off Int
off'
  where
    go :: Int -> Int -> Int -> m ()
    go :: Int -> Int -> Int -> m ()
go Int
len Int
offXs Int
offYs
      | Int
shft Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
        Int -> Int -> Int -> m ()
go' Int
len Int
offXs (Int -> Int
forall a. Bits a => a -> a
divWordSize Int
offYs)
      | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
wordSize = do
        Word
y <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
vecYs Int
0
        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
vecYs Int
0 (Word -> Word -> Word
forall a. Bits a => a -> a -> a
f Word
x Word
y)
      | Bool
otherwise = do
        Word
y <- MutableByteArray (PrimState m) -> Int -> m Word
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray (PrimState m)
ys Int
base
        MutableByteArray (PrimState m) -> Int -> Word -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Word -> Word -> m ()
modifyByteArray MutableByteArray (PrimState m)
ys Int
base (Int -> Word
loMask Int
shft) (Word -> Word -> Word
forall a. Bits a => a -> a -> a
f (Word
x Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
shft) Word
y Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Int -> Word
hiMask Int
shft)
        Int -> Int -> Int -> m ()
go' (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
shft) (Int
offXs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
shft) (Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      where
        vecXs :: Vector Bit
vecXs = Int -> Int -> ByteArray -> Vector Bit
BitVec  Int
offXs Int
len ByteArray
xs
        vecYs :: MVector (PrimState m) Bit
vecYs = Int
-> Int
-> MutableByteArray (PrimState m)
-> MVector (PrimState m) Bit
forall s. Int -> Int -> MutableByteArray s -> MVector s Bit
BitMVec Int
offYs Int
len MutableByteArray (PrimState m)
ys
        x :: Word
x     = Vector Bit -> Int -> Word
indexWord Vector Bit
vecXs Int
0
        shft :: Int
shft  = Int -> Int
modWordSize Int
offYs
        base :: Int
base  = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
offYs

    go' :: Int -> Int -> Int -> m ()
    go' :: Int -> Int -> Int -> m ()
go' Int
len Int
offXs Int
offYsW = do
      if Int
shft Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Int -> m ()
loopAligned Int
offYsW
        else Int -> Word -> m ()
loop Int
offYsW (ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
xs Int
base)
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int -> Int
modWordSize Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        let ix :: Int
ix = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
modWordSize Int
len
        let x :: Word
x = Vector Bit -> Int -> Word
indexWord Vector Bit
vecXs Int
ix
        Word
y <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
vecYs Int
ix
        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
vecYs Int
ix (Word -> Word -> Word
forall a. Bits a => a -> a -> a
f Word
x Word
y)

      where

        vecXs :: Vector Bit
vecXs = Int -> Int -> ByteArray -> Vector Bit
BitVec  Int
offXs Int
len ByteArray
xs
        vecYs :: MVector (PrimState m) Bit
vecYs = Int
-> Int
-> MutableByteArray (PrimState m)
-> MVector (PrimState m) Bit
forall s. Int -> Int -> MutableByteArray s -> MVector s Bit
BitMVec (Int -> Int
forall a. Bits a => a -> a
mulWordSize Int
offYsW) Int
len MutableByteArray (PrimState m)
ys
        shft :: Int
shft  = Int -> Int
modWordSize Int
offXs
        shft' :: Int
shft' = Int
wordSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
shft
        base :: Int
base  = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
offXs
        base0 :: Int
base0 = Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offYsW
        base1 :: Int
base1 = Int
base0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        iMax :: Int
iMax  = Int -> Int
forall a. Bits a => a -> a
divWordSize Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offYsW

        loopAligned :: Int -> m ()
        loopAligned :: Int -> m ()
loopAligned !Int
i
          | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
iMax = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Bool
otherwise =  do
            let x :: Word
x = ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
xs (Int
base0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) :: Word
            Word
y <- MutableByteArray (PrimState m) -> Int -> m Word
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray (PrimState m)
ys Int
i
            MutableByteArray (PrimState m) -> Int -> Word -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray (PrimState m)
ys Int
i (Word -> Word -> Word
forall a. Bits a => a -> a -> a
f Word
x Word
y)
            Int -> m ()
loopAligned (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

        loop :: Int -> Word -> m ()
        loop :: Int -> Word -> m ()
loop !Int
i !Word
acc
          | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
iMax = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Bool
otherwise =  do
            let accNew :: Word
accNew = ByteArray -> Int -> Word
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
xs (Int
base1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)
                x :: Word
x = (Word
acc Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
shft) Word -> Word -> Word
forall a. Bits a => a -> a -> a
.|. (Word
accNew Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
shft')
            Word
y <- MutableByteArray (PrimState m) -> Int -> m Word
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray (PrimState m)
ys Int
i
            MutableByteArray (PrimState m) -> Int -> Word -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray (PrimState m)
ys Int
i (Word -> Word -> Word
forall a. Bits a => a -> a -> a
f Word
x Word
y)
            Int -> Word -> m ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Word
accNew

{-# SPECIALIZE zipInPlace :: (forall a. Bits a => a -> a -> a) -> Vector Bit -> MVector s Bit -> ST s () #-}
{-# INLINE zipInPlace #-}

-- | Apply a function to a mutable vector bitwise,
-- rewriting its contents.
-- Cf. 'Data.Bit.mapBits'.
--
-- >>> :set -XOverloadedLists
-- >>> import Data.Bits
-- >>> Data.Vector.Unboxed.modify (mapInPlace complement) [0,1,1]
-- [1,0,0]
mapInPlace
  :: PrimMonad m
  => (forall a . Bits a => a -> a)
  -> U.MVector (PrimState m) Bit
  -> m ()
mapInPlace :: (forall a. Bits a => a -> a) -> MVector (PrimState m) Bit -> m ()
mapInPlace forall a. Bits a => a -> a
f MVector (PrimState m) 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) -> MVector (PrimState m) Bit -> Bit -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MU.set MVector (PrimState m) Bit
xs (Bool -> Bit
Bit Bool
False)
  (Bool
False, Bool
True)  -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  (Bool
True, Bool
False)  -> MVector (PrimState m) Bit -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> m ()
invertInPlace MVector (PrimState m) Bit
xs
  (Bool
True, Bool
True)   -> MVector (PrimState m) Bit -> Bit -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MU.set MVector (PrimState m) Bit
xs (Bool -> Bit
Bit Bool
True)
{-# SPECIALIZE mapInPlace :: (forall a. Bits a => a -> a) -> MVector s Bit -> ST s () #-}
{-# INLINE mapInPlace #-}

-- | Invert (flip) all bits in-place.
--
-- >>> :set -XOverloadedLists
-- >>> Data.Vector.Unboxed.modify invertInPlace [0,1,0,1,0]
-- [1,0,1,0,1]
invertInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> m ()
invertInPlace :: MVector (PrimState m) Bit -> m ()
invertInPlace MVector (PrimState m) Bit
xs = do
  let n :: Int
n = MVector (PrimState m) Bit -> Int
forall a s. Unbox a => MVector s a -> Int
MU.length MVector (PrimState m) Bit
xs
  [Int] -> (Int -> m ()) -> m ()
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 -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    Word
x <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
i
    MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
i (Word -> Word
forall a. Bits a => a -> a
complement Word
x)
{-# SPECIALIZE invertInPlace :: U.MVector s Bit -> ST s () #-}

-- | Same as 'Data.Bit.selectBits', but deposit
-- selected bits in-place. Returns a number of selected bits.
-- It is caller's responsibility to trim the result to this number.
--
-- >>> :set -XOverloadedLists
-- >>> import Control.Monad.ST (runST)
-- >>> import qualified Data.Vector.Unboxed as U
-- >>> runST $ do { vec <- U.unsafeThaw [1,1,0,0,1]; n <- selectBitsInPlace [0,1,0,1,1] vec; U.take n <$> U.unsafeFreeze vec }
-- [1,0,1]
--
selectBitsInPlace
  :: PrimMonad m => U.Vector Bit -> U.MVector (PrimState m) Bit -> m Int
selectBitsInPlace :: Vector Bit -> MVector (PrimState m) Bit -> m Int
selectBitsInPlace Vector Bit
is MVector (PrimState m) Bit
xs = Int -> Int -> m Int
loop Int
0 Int
0
 where
  !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
is) (MVector (PrimState m) Bit -> Int
forall a s. Unbox a => MVector s a -> Int
MU.length MVector (PrimState m) Bit
xs)
  loop :: Int -> Int -> m Int
loop !Int
i !Int
ct
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ct
    | Bool
otherwise = do
      Word
x <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
i
      let !(Int
nSet, Word
x') = Word -> Word -> (Int, Word)
selectWord (Int -> Word -> Word
masked (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) (Vector Bit -> Int -> Word
indexWord Vector Bit
is Int
i)) Word
x
      MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
ct Word
x'
      Int -> Int -> m Int
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wordSize) (Int
ct Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nSet)

-- | Same as 'Data.Bit.excludeBits', but deposit
-- excluded bits in-place. Returns a number of excluded bits.
-- It is caller's responsibility to trim the result to this number.
--
-- >>> :set -XOverloadedLists
-- >>> import Control.Monad.ST (runST)
-- >>> import qualified Data.Vector.Unboxed as U
-- >>> runST $ do { vec <- U.unsafeThaw [1,1,0,0,1]; n <- excludeBitsInPlace [0,1,0,1,1] vec; U.take n <$> U.unsafeFreeze vec }
-- [1,0]
--
excludeBitsInPlace
  :: PrimMonad m => U.Vector Bit -> U.MVector (PrimState m) Bit -> m Int
excludeBitsInPlace :: Vector Bit -> MVector (PrimState m) Bit -> m Int
excludeBitsInPlace Vector Bit
is MVector (PrimState m) Bit
xs = Int -> Int -> m Int
loop Int
0 Int
0
 where
  !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
is) (MVector (PrimState m) Bit -> Int
forall a s. Unbox a => MVector s a -> Int
MU.length MVector (PrimState m) Bit
xs)
  loop :: Int -> Int -> m Int
loop !Int
i !Int
ct
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ct
    | Bool
otherwise = do
      Word
x <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
i
      let !(Int
nSet, Word
x') =
            Word -> Word -> (Int, Word)
selectWord (Int -> Word -> Word
masked (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) (Word -> Word
forall a. Bits a => a -> a
complement (Vector Bit -> Int -> Word
indexWord Vector Bit
is Int
i))) Word
x
      MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
ct Word
x'
      Int -> Int -> m Int
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wordSize) (Int
ct Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nSet)

-- | Reverse the order of bits in-place.
--
-- >>> :set -XOverloadedLists
-- >>> Data.Vector.Unboxed.modify reverseInPlace [1,1,0,1,0]
-- [0,1,0,1,1]
--
-- Consider using @vector-rotcev@ package
-- to reverse vectors in O(1) time.
reverseInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> m ()
reverseInPlace :: MVector (PrimState m) Bit -> m ()
reverseInPlace MVector (PrimState m) Bit
xs
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0  = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise = Int -> m ()
loop Int
0
  where
    len :: Int
len = MVector (PrimState m) Bit -> Int
forall a s. Unbox a => MVector s a -> Int
MU.length MVector (PrimState m) Bit
xs

    loop :: Int -> m ()
loop !Int
i
      | Int
i' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
j' = do
        Word
x <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
i
        Word
y <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
j'

        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
i  (Word -> Word
reverseWord Word
y)
        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
j' (Word -> Word
reverseWord Word
x)

        Int -> m ()
loop Int
i'
      | Int
i' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
j = do
        let w :: Int
w = (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
1
            k :: Int
k = Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w
        Word
x <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
i
        Word
y <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
k

        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
i (Int -> Word -> Word -> Word
meld Int
w (Int -> Word -> Word
reversePartialWord Int
w Word
y) Word
x)
        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
k (Int -> Word -> Word -> Word
meld Int
w (Int -> Word -> Word
reversePartialWord Int
w Word
x) Word
y)

        Int -> m ()
loop Int
i'
      | Bool
otherwise = do
        let w :: Int
w = Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i
        Word
x <- MVector (PrimState m) Bit -> Int -> m Word
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> m Word
readWord MVector (PrimState m) Bit
xs Int
i
        MVector (PrimState m) Bit -> Int -> Word -> m ()
forall (m :: * -> *).
PrimMonad m =>
MVector (PrimState m) Bit -> Int -> Word -> m ()
writeWord MVector (PrimState m) Bit
xs Int
i (Int -> Word -> Word -> Word
meld Int
w (Int -> Word -> Word
reversePartialWord Int
w Word
x) Word
x)
     where
      !j :: Int
j  = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i
      !i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wordSize
      !j' :: Int
j' = Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
wordSize
{-# SPECIALIZE reverseInPlace :: U.MVector s Bit -> ST s () #-}