#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
#else
#define safe
#endif
module Data.Vector.Unboxed.Mutable.Bit
( module Data.Bit
, module U
, wordSize
, wordLength
, cloneFromWords
, cloneToWords
, readWord
, writeWord
, mapMInPlaceWithIndex
, mapInPlaceWithIndex
, mapMInPlace
, mapInPlace
, zipInPlace
, unionInPlace
, intersectionInPlace
, differenceInPlace
, symDiffInPlace
, invertInPlace
, selectBitsInPlace
, excludeBitsInPlace
, countBits
, listBits
, and
, or
, any
, anyBits
, all
, allBits
, reverseInPlace
) where
import safe Control.Monad
import Control.Monad.Primitive
import safe Data.Bit
import safe Data.Bit.Internal
import safe Data.Bits
import qualified Data.Vector.Generic.Mutable as MV
import safe qualified Data.Vector.Generic.Safe as V
import safe qualified Data.Vector.Unboxed.Safe as U (Vector)
import safe Data.Vector.Unboxed.Mutable.Safe as U
import Data.Vector.Unboxed.Bit.Internal
import safe Data.Word
import safe Prelude as P
hiding (and, or, any, all, reverse)
wordLength :: U.MVector s Bit -> Int
wordLength = nWords . MV.length
cloneFromWords :: PrimMonad m => Int -> U.MVector (PrimState m) Word -> m (U.MVector (PrimState m) Bit)
cloneFromWords n ws = do
let wordsNeeded = nWords n
wordsGiven = MV.length ws
fillNeeded = wordsNeeded wordsGiven
v <- MV.new wordsNeeded
if fillNeeded > 0
then do
MV.copy (MV.slice 0 wordsGiven v) ws
MV.set (MV.slice wordsGiven fillNeeded v) 0
else do
MV.copy v (MV.slice 0 wordsNeeded ws)
return (BitMVec 0 n v)
cloneToWords :: PrimMonad m => U.MVector (PrimState m) Bit -> m (U.MVector (PrimState m) Word)
cloneToWords v@(BitMVec s n ws)
| aligned s = do
ws <- MV.clone (MV.slice (divWordSize s) (nWords n) ws)
when (not (aligned n)) $ do
readWord v (alignDown n) >>= MV.write ws (divWordSize n)
return ws
| otherwise = cloneWords v
mapMInPlaceWithIndex ::
PrimMonad m =>
(Int -> Word -> m Word)
-> U.MVector (PrimState m) Bit -> m ()
mapMInPlaceWithIndex f xs@(BitMVec 0 n v) = loop 0 0
where
!n_ = alignDown (MV.length xs)
loop !i !j
| i >= n_ = when (n_ /= MV.length xs) $ do
readWord xs i >>= f i >>= writeWord xs i
| otherwise = do
MV.read v j >>= f i >>= MV.write v j
loop (i + wordSize) (j + 1)
mapMInPlaceWithIndex f xs = loop 0
where
!n = MV.length xs
loop !i
| i >= n = return ()
| otherwise = do
readWord xs i >>= f i >>= writeWord xs i
loop (i + wordSize)
mapInPlaceWithIndex ::
PrimMonad m =>
(Int -> Word -> Word)
-> U.MVector (PrimState m) Bit -> m ()
mapInPlaceWithIndex f = mapMInPlaceWithIndex g
where
g i x = return $! f i x
mapMInPlace :: PrimMonad m => (Word -> m Word) -> U.MVector (PrimState m) Bit -> m ()
mapMInPlace f = mapMInPlaceWithIndex (const f)
mapInPlace :: PrimMonad m => (Word -> Word) -> U.MVector (PrimState m) Bit -> m ()
mapInPlace f = mapMInPlaceWithIndex (\_ x -> return (f x))
zipInPlace :: PrimMonad m => (Word -> Word -> Word) -> U.MVector (PrimState m) Bit -> U.Vector Bit -> m ()
zipInPlace f xs ys@(BitVec 0 n2 v) =
mapInPlaceWithIndex g (MV.basicUnsafeSlice 0 n xs)
where
!n = min (MV.length xs) (V.length ys)
g !i !x =
let !w = masked (n2 i) (v V.! divWordSize i)
in f x w
zipInPlace f xs ys =
mapInPlaceWithIndex g (MV.basicUnsafeSlice 0 n xs)
where
!n = min (MV.length xs) (V.length ys)
g !i !x =
let !w = indexWord ys i
in f x w
unionInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> U.Vector Bit -> m ()
unionInPlace = zipInPlace (.|.)
intersectionInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> U.Vector Bit -> m ()
intersectionInPlace = zipInPlace (.&.)
differenceInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> U.Vector Bit -> m ()
differenceInPlace = zipInPlace diff
symDiffInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> U.Vector Bit -> m ()
symDiffInPlace = zipInPlace xor
invertInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> m ()
invertInPlace = mapInPlace complement
selectBitsInPlace :: PrimMonad m => U.Vector Bit -> U.MVector (PrimState m) Bit -> m Int
selectBitsInPlace is xs = loop 0 0
where
!n = min (V.length is) (MV.length xs)
loop !i !ct
| i >= n = return ct
| otherwise = do
x <- readWord xs i
let !(nSet, x') = selectWord (masked (n i) (indexWord is i)) x
writeWord xs ct x'
loop (i + wordSize) (ct + nSet)
excludeBitsInPlace :: PrimMonad m => U.Vector Bit -> U.MVector (PrimState m) Bit -> m Int
excludeBitsInPlace is xs = loop 0 0
where
!n = min (V.length is) (MV.length xs)
loop !i !ct
| i >= n = return ct
| otherwise = do
x <- readWord xs i
let !(nSet, x') = selectWord (masked (n i) (complement (indexWord is i))) x
writeWord xs ct x'
loop (i + wordSize) (ct + nSet)
countBits :: PrimMonad m => U.MVector (PrimState m) Bit -> m Int
countBits v = loop 0 0
where
!n = alignUp (MV.length v)
loop !s !i
| i >= n = return s
| otherwise = do
x <- readWord v i
loop (s + popCount x) (i + wordSize)
listBits :: PrimMonad m => U.MVector (PrimState m) Bit -> m [Int]
listBits v = loop id 0
where
!n = MV.length v
loop bs !i
| i >= n = return $! bs []
| otherwise = do
w <- readWord v i
loop (bs . bitsInWord i w) (i + wordSize)
and :: PrimMonad m => U.MVector (PrimState m) Bit -> m Bool
and v = loop 0
where
!n = MV.length v
loop !i
| i >= n = return True
| otherwise = do
y <- readWord v i
if y == mask (n i)
then loop (i + wordSize)
else return False
or :: PrimMonad m => U.MVector (PrimState m) Bit -> m Bool
or v = loop 0
where
!n = MV.length v
loop !i
| i >= n = return False
| otherwise = do
y <- readWord v i
if y /= 0
then return True
else loop (i + wordSize)
all :: PrimMonad m => (Bit -> Bool) -> U.MVector (PrimState m) Bit -> m Bool
all p = case (p 0, p 1) of
(False, False) -> return . MV.null
(False, True) -> allBits 1
(True, False) -> allBits 0
(True, True) -> flip seq (return True)
any :: PrimMonad m => (Bit -> Bool) -> U.MVector (PrimState m) Bit -> m Bool
any p = case (p 0, p 1) of
(False, False) -> flip seq (return False)
(False, True) -> anyBits 1
(True, False) -> anyBits 0
(True, True) -> return . not . MV.null
allBits, anyBits :: PrimMonad m => Bit -> U.MVector (PrimState m) Bit -> m Bool
allBits 0 = liftM not . or
allBits 1 = and
anyBits 0 = liftM not . and
anyBits 1 = or
reverseInPlace :: PrimMonad m => U.MVector (PrimState m) Bit -> m ()
reverseInPlace xs = loop 0 (MV.length xs)
where
loop !i !j
| i' <= j' = do
x <- readWord xs i
y <- readWord xs j'
writeWord xs i (reverseWord y)
writeWord xs j' (reverseWord x)
loop i' j'
| i' < j = do
let w = (j i) `shiftR` 1
k = j w
x <- readWord xs i
y <- readWord xs k
writeWord xs i (meld w (reversePartialWord w y) x)
writeWord xs k (meld w (reversePartialWord w x) y)
loop i' j'
| i < j = do
let w = j i
x <- readWord xs i
writeWord xs i (meld w (reversePartialWord w x) x)
| otherwise = return ()
where
!i' = i + wordSize
!j' = j wordSize