{-|

Module      :  Data.Array.BitArray.IO
Copyright   :  (c) Claude Heiland-Allen 2012,2018
License     :  BSD3

Maintainer  :  claude@mathr.co.uk
Stability   :  unstable
Portability :  portable

Unboxed mutable bit arrays in the 'IO' monad.

-}
module Data.Array.BitArray.IO
  ( IOBitArray()
  -- * MArray-like interface.
  , getBounds
  , newArray
  , newArray_
  , newListArray
  , readArray
  , writeArray
  , mapArray
  , mapIndices
  , getElems
  , getAssocs
  -- * Conversion to/from immutable bit arrays.
  , freeze
  , thaw
  -- * Construction
  , copy
  , fill
  -- * Short-circuiting reductions.
  , or
  , and
  , isUniform
  , elemIndex
  -- * Aggregate operations.
  , fold
  , map
  , zipWith
  , popCount
  -- * Unsafe.
  , unsafeReadArray
  , unsafeGetElems
  , unsafeFreeze
  , unsafeThaw
  ) where

import Prelude hiding (and, or, map, zipWith)

import Control.Monad (forM_, when)
import Data.Bits (shiftL, shiftR, testBit, setBit, clearBit, (.&.), complement)
import qualified Data.Bits
import Data.Ix (Ix, index, inRange, range, rangeSize)
import Data.List (foldl1')
import Data.Word (Word8, Word64)
import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr)
import Foreign.Ptr (Ptr, plusPtr, castPtr)
import Foreign.Storable (poke, pokeByteOff, pokeElemOff, peekByteOff, peekElemOff)
import System.IO.Unsafe (unsafeInterleaveIO)

import Data.Bits.Bitwise (packWord8LE, mask)
import qualified Data.Bits.Bitwise as Bitwise

import Data.Array.BitArray.Internal
  ( IOBitArray(..)
  , getBounds
  , newArray_
  , freeze
  , unsafeFreeze
  , thaw
  , unsafeThaw
  , copy
  )

-- | Create a new array filled with an initial value.
{-# INLINE newArray #-}
newArray :: Ix i => (i, i) {- ^ bounds -} -> Bool {- ^ initial value -} -> IO (IOBitArray i)
newArray :: forall i. Ix i => (i, i) -> Bool -> IO (IOBitArray i)
newArray (i, i)
bs Bool
b = do
  IOBitArray i
a <- (i, i) -> IO (IOBitArray i)
forall i. Ix i => (i, i) -> IO (IOBitArray i)
newArray_ (i, i)
bs
  IOBitArray i -> Bool -> IO ()
forall i. Ix i => IOBitArray i -> Bool -> IO ()
fill IOBitArray i
a Bool
b
  IOBitArray i -> IO (IOBitArray i)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IOBitArray i
a

-- | Create a new array filled with values from a list.
{-# INLINE newListArray #-}
newListArray :: Ix i => (i, i) {- ^ bounds -} -> [Bool] {- ^ elems -} -> IO (IOBitArray i)
newListArray :: forall i. Ix i => (i, i) -> [Bool] -> IO (IOBitArray i)
newListArray (i, i)
bs [Bool]
es = do
  IOBitArray i
a <- (i, i) -> IO (IOBitArray i)
forall i. Ix i => (i, i) -> IO (IOBitArray i)
newArray_ (i, i)
bs
  let byteBits :: Int
byteBits = Int
8
      writeBytes :: Ptr Word8 -> [Bool] -> IO ()
      writeBytes :: Ptr Word8 -> [Bool] -> IO ()
writeBytes Ptr Word8
p (Bool
b0:Bool
b1:Bool
b2:Bool
b3:Bool
b4:Bool
b5:Bool
b6:Bool
b7:[Bool]
rest) = do
        Ptr Word8 -> Word8 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Word8
p (Bool
-> Bool -> Bool -> Bool -> Bool -> Bool -> Bool -> Bool -> Word8
packWord8LE Bool
b0 Bool
b1 Bool
b2 Bool
b3 Bool
b4 Bool
b5 Bool
b6 Bool
b7)
        Ptr Word8 -> [Bool] -> IO ()
writeBytes (Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
p Int
1) [Bool]
rest
      writeBytes Ptr Word8
_ [] = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      writeBytes Ptr Word8
p [Bool]
rest = Ptr Word8 -> [Bool] -> IO ()
writeBytes Ptr Word8
p (Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take Int
byteBits ([Bool]
rest [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ Bool -> [Bool]
forall a. a -> [a]
repeat Bool
False))
  ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> do
    Ptr Word8 -> [Bool] -> IO ()
writeBytes (Ptr Word64 -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr Word64
p) (Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take (Int
byteBits Int -> Int -> Int
forall a. Num a => a -> a -> a
* IOBitArray i -> Int
forall i. IOBitArray i -> Int
iobBytes IOBitArray i
a) [Bool]
es)
  IOBitArray i -> IO (IOBitArray i)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IOBitArray i
a

-- | Read from an array at an index.
{-# INLINE readArray #-}
readArray :: Ix i => IOBitArray i -> i -> IO Bool
readArray :: forall i. Ix i => IOBitArray i -> i -> IO Bool
readArray IOBitArray i
a i
i = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not ((i, i) -> i -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (i, i)
bs i
i)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"array index out of bounds"
  IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a ((i, i) -> i -> Int
forall a. Ix a => (a, a) -> a -> Int
index (i, i)
bs i
i)

-- | Read from an array at an index without bounds checking.  Unsafe.
{-# INLINE unsafeReadArray #-}
unsafeReadArray :: Ix i => IOBitArray i -> i -> IO Bool
unsafeReadArray :: forall i. Ix i => IOBitArray i -> i -> IO Bool
unsafeReadArray IOBitArray i
a i
i = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a ((i, i) -> i -> Int
forall a. Ix a => (a, a) -> a -> Int
index (i, i)
bs i
i)

{-# INLINE readArrayRaw #-}
readArrayRaw :: Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw :: forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a Int
n = do
  let byte :: Int
byte = Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
      bit :: Int
bit = Int
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7
  ForeignPtr Word64 -> (Ptr Word64 -> IO Bool) -> IO Bool
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO Bool) -> IO Bool)
-> (Ptr Word64 -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> do
    Word8
b0 <- Ptr Word64 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word64
p Int
byte
    Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Word8 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (Word8
b0 :: Word8) Int
bit)

-- | Write to an array at an index.
{-# INLINE writeArray #-}
writeArray :: Ix i => IOBitArray i -> i -> Bool -> IO ()
writeArray :: forall i. Ix i => IOBitArray i -> i -> Bool -> IO ()
writeArray IOBitArray i
a i
i Bool
b = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not ((i, i) -> i -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (i, i)
bs i
i)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"array index out of bounds"
  let n :: Int
n = (i, i) -> i -> Int
forall a. Ix a => (a, a) -> a -> Int
index (i, i)
bs i
i
      byte :: Int
byte = Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
      bit :: Int
bit = Int
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
7
  ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> do
    Word8
b0 <- Ptr Word64 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word64
p Int
byte
    let b1 :: Word8
b1 = (if Bool
b then Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
setBit else Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
clearBit) (Word8
b0 :: Word8) Int
bit
    Ptr Word64 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word64
p Int
byte Word8
b1

-- | Alias for 'map'.
{-# INLINE mapArray #-}
mapArray :: Ix i => (Bool -> Bool) -> IOBitArray i -> IO (IOBitArray i)
mapArray :: forall i.
Ix i =>
(Bool -> Bool) -> IOBitArray i -> IO (IOBitArray i)
mapArray = (Bool -> Bool) -> IOBitArray i -> IO (IOBitArray i)
forall i.
Ix i =>
(Bool -> Bool) -> IOBitArray i -> IO (IOBitArray i)
map

-- unsafeInterleaveIO is used to avoid having to create the whole list in
-- memory before the function can return, but need to keep the ForeignPtr
-- alive to avoid GC stealing our data.
interleavedMapMThenTouch :: Ix i => IOBitArray i -> (a -> IO b) -> [a] -> IO [b]
interleavedMapMThenTouch :: forall i a b. Ix i => IOBitArray i -> (a -> IO b) -> [a] -> IO [b]
interleavedMapMThenTouch IOBitArray i
a a -> IO b
_ [] = ForeignPtr Word64 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) IO () -> IO [b] -> IO [b]
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [b] -> IO [b]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return []
interleavedMapMThenTouch IOBitArray i
a a -> IO b
f (a
x:[a]
xs) = IO [b] -> IO [b]
forall a. IO a -> IO a
unsafeInterleaveIO (IO [b] -> IO [b]) -> IO [b] -> IO [b]
forall a b. (a -> b) -> a -> b
$ do
  b
y <- a -> IO b
f a
x
  [b]
ys <- IOBitArray i -> (a -> IO b) -> [a] -> IO [b]
forall i a b. Ix i => IOBitArray i -> (a -> IO b) -> [a] -> IO [b]
interleavedMapMThenTouch IOBitArray i
a a -> IO b
f [a]
xs
  [b] -> IO [b]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (b
yb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
ys)

-- | Create a new array by reading from another.
{-# INLINE mapIndices #-}
mapIndices :: (Ix i, Ix j) => (i, i) {- ^ new bounds -} -> (i -> j) {- ^ index transformation -} -> IOBitArray j {- ^ source array -} -> IO (IOBitArray i)
mapIndices :: forall i j.
(Ix i, Ix j) =>
(i, i) -> (i -> j) -> IOBitArray j -> IO (IOBitArray i)
mapIndices (i, i)
bs i -> j
h IOBitArray j
a = (i, i) -> [Bool] -> IO (IOBitArray i)
forall i. Ix i => (i, i) -> [Bool] -> IO (IOBitArray i)
newListArray (i, i)
bs ([Bool] -> IO (IOBitArray i)) -> IO [Bool] -> IO (IOBitArray i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IOBitArray j -> (i -> IO Bool) -> [i] -> IO [Bool]
forall i a b. Ix i => IOBitArray i -> (a -> IO b) -> [a] -> IO [b]
interleavedMapMThenTouch IOBitArray j
a (IOBitArray j -> j -> IO Bool
forall i. Ix i => IOBitArray i -> i -> IO Bool
readArray IOBitArray j
a (j -> IO Bool) -> (i -> j) -> i -> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> j
h) ((i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range (i, i)
bs)

-- | Get a list of all elements of an array.
{-# INLINE getElems #-}
getElems :: Ix i => IOBitArray i -> IO [Bool]
getElems :: forall i. Ix i => IOBitArray i -> IO [Bool]
getElems IOBitArray i
a = IOBitArray i -> IO [Bool]
forall i. Ix i => IOBitArray i -> IO [Bool]
unsafeGetElems (IOBitArray i -> IO [Bool]) -> IO (IOBitArray i) -> IO [Bool]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IOBitArray i -> IO (IOBitArray i)
forall i. Ix i => IOBitArray i -> IO (IOBitArray i)
copy IOBitArray i
a

-- | Get a list of all elements of an array.  Unsafe when the source
--   array can be modified later.
{-# INLINE unsafeGetElems #-}
unsafeGetElems :: Ix i => IOBitArray i -> IO [Bool]
unsafeGetElems :: forall i. Ix i => IOBitArray i -> IO [Bool]
unsafeGetElems IOBitArray i
a' = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a'
  let r :: Int
r = (i, i) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (i, i)
bs
      count :: Int
count = (Int
r 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
  Ptr Word64
p <- ForeignPtr Word64
-> (Ptr Word64 -> IO (Ptr Word64)) -> IO (Ptr Word64)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a') ((Ptr Word64 -> IO (Ptr Word64)) -> IO (Ptr Word64))
-> (Ptr Word64 -> IO (Ptr Word64)) -> IO (Ptr Word64)
forall a b. (a -> b) -> a -> b
$ Ptr Word64 -> IO (Ptr Word64)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
  [Word8]
bytes <- IOBitArray i -> (Int -> IO Word8) -> [Int] -> IO [Word8]
forall i a b. Ix i => IOBitArray i -> (a -> IO b) -> [a] -> IO [b]
interleavedMapMThenTouch IOBitArray i
a' (Ptr Word64 -> Int -> IO Word8
forall b. Ptr b -> Int -> IO Word8
forall a b. Storable a => Ptr b -> Int -> IO a
peekByteOff Ptr Word64
p) [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  [Bool] -> IO [Bool]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Bool] -> IO [Bool])
-> ([Word8] -> [Bool]) -> [Word8] -> IO [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take Int
r ([Bool] -> [Bool]) -> ([Word8] -> [Bool]) -> [Word8] -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word8 -> [Bool]) -> [Word8] -> [Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Word8 -> [Bool]
forall b. Bits b => b -> [Bool]
Bitwise.toListLE ([Word8] -> IO [Bool]) -> [Word8] -> IO [Bool]
forall a b. (a -> b) -> a -> b
$ ([Word8]
bytes :: [Word8])

-- | Get a list of all (index, element) pairs.
{-# INLINE getAssocs #-}
getAssocs :: Ix i => IOBitArray i -> IO [(i, Bool)]
getAssocs :: forall i. Ix i => IOBitArray i -> IO [(i, Bool)]
getAssocs IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  [i] -> [Bool] -> [(i, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range (i, i)
bs) ([Bool] -> [(i, Bool)]) -> IO [Bool] -> IO [(i, Bool)]
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` IOBitArray i -> IO [Bool]
forall i. Ix i => IOBitArray i -> IO [Bool]
getElems IOBitArray i
a
  
-- | Fill an array with a uniform value.
{-# INLINE fill #-}
fill :: Ix i => IOBitArray i -> Bool -> IO ()
fill :: forall i. Ix i => IOBitArray i -> Bool -> IO ()
fill IOBitArray i
a Bool
b = do
  let count :: Int
count = IOBitArray i -> Int
forall i. IOBitArray i -> Int
iobBytes IOBitArray i
a Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
      word :: Word64
      word :: Word64
word = if Bool
b then Word64 -> Word64
forall a. Bits a => a -> a
complement Word64
0 else Word64
0
  ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p ->
    [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      Ptr Word64 -> Int -> Word64 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word64
p Int
i Word64
word

-- | Short-circuit bitwise reduction: True when any bit is True.
{-# INLINE or #-}
or :: Ix i => IOBitArray i -> IO Bool
or :: forall i. Ix i => IOBitArray i -> IO Bool
or IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  let total :: Int
total = (i, i) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (i, i)
bs
      full :: Int
full = Int
total Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement (Int -> Int
forall b. (Num b, Bits b) => Int -> b
mask Int
6)
      count :: Int
count = Int
full Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
      loop :: Ptr Word64 -> Int -> IO Bool
      loop :: Ptr Word64 -> Int -> IO Bool
loop Ptr Word64
p Int
n
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
count = do
            Word64
w <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
p Int
n
            if Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= (Word64
0 :: Word64) then Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True else Ptr Word64 -> Int -> IO Bool
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int -> IO Bool
rest Int
full
      rest :: Int -> IO Bool
rest Int
m
        | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
total = do
            Bool
b <- IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a Int
m
            if Bool
b then Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True else Int -> IO Bool
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  ForeignPtr Word64 -> (Ptr Word64 -> IO Bool) -> IO Bool
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO Bool) -> IO Bool)
-> (Ptr Word64 -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> Ptr Word64 -> Int -> IO Bool
loop Ptr Word64
p Int
0

-- | Short-circuit bitwise reduction: False when any bit is False.
{-# INLINE and #-}
and :: Ix i => IOBitArray i -> IO Bool
and :: forall i. Ix i => IOBitArray i -> IO Bool
and IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  let total :: Int
total = (i, i) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (i, i)
bs
      full :: Int
full = Int
total Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement (Int -> Int
forall b. (Num b, Bits b) => Int -> b
mask Int
6)
      count :: Int
count = Int
full Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
      loop :: Ptr Word64 -> Int -> IO Bool
      loop :: Ptr Word64 -> Int -> IO Bool
loop Ptr Word64
p Int
n
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
count = do
            Word64
w <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
p Int
n
            if Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= (Word64 -> Word64
forall a. Bits a => a -> a
complement Word64
0 :: Word64) then Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False else Ptr Word64 -> Int -> IO Bool
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int -> IO Bool
rest Int
full
      rest :: Int -> IO Bool
rest Int
m
        | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
total = do
            Bool
b <- IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a Int
m
            if Bool -> Bool
not Bool
b then Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False else Int -> IO Bool
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
  ForeignPtr Word64 -> (Ptr Word64 -> IO Bool) -> IO Bool
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO Bool) -> IO Bool)
-> (Ptr Word64 -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> Ptr Word64 -> Int -> IO Bool
loop Ptr Word64
p Int
0

-- | Short-circuit bitwise reduction: 'Nothing' when any bits differ,
--   'Just' when all bits are the same.
{-# INLINE isUniform #-}
isUniform :: Ix i => IOBitArray i -> IO (Maybe Bool)
isUniform :: forall i. Ix i => IOBitArray i -> IO (Maybe Bool)
isUniform IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  let total :: Int
total = (i, i) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (i, i)
bs
      full :: Int
full = Int
total Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement (Int -> Int
forall b. (Num b, Bits b) => Int -> b
mask Int
6)
      count :: Int
count = Int
full Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
      loop :: Ptr Word64 -> Int -> Bool -> Bool -> IO (Maybe Bool)
      loop :: Ptr Word64 -> Int -> Bool -> Bool -> IO (Maybe Bool)
loop Ptr Word64
p Int
n Bool
st Bool
sf
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
count = do
            Word64
w <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
p Int
n
            let t :: Bool
t = Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= (Word64
0 :: Word64)  Bool -> Bool -> Bool
|| Bool
st
                f :: Bool
f = Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= (Word64 -> Word64
forall a. Bits a => a -> a
complement Word64
0) Bool -> Bool -> Bool
|| Bool
sf
            if Bool
t Bool -> Bool -> Bool
&& Bool
f then Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Bool
forall a. Maybe a
Nothing else Ptr Word64 -> Int -> Bool -> Bool -> IO (Maybe Bool)
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
t Bool
f
        | Bool
otherwise = Int -> Bool -> Bool -> IO (Maybe Bool)
rest Int
full Bool
st Bool
sf
      rest :: Int -> Bool -> Bool -> IO (Maybe Bool)
rest Int
m Bool
st Bool
sf
        | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
total = do
            Bool
b <- IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a Int
m
            let t :: Bool
t =     Bool
b Bool -> Bool -> Bool
|| Bool
st
                f :: Bool
f = Bool -> Bool
not Bool
b Bool -> Bool -> Bool
|| Bool
sf
            if Bool
t Bool -> Bool -> Bool
&& Bool
f then Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Bool
forall a. Maybe a
Nothing else Int -> Bool -> Bool -> IO (Maybe Bool)
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool
t Bool
f
        | Bool
st Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
sf = Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True)
        | Bool -> Bool
not Bool
st Bool -> Bool -> Bool
&& Bool
sf = Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False)
        | Bool
otherwise = Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Bool
forall a. Maybe a
Nothing
  ForeignPtr Word64
-> (Ptr Word64 -> IO (Maybe Bool)) -> IO (Maybe Bool)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO (Maybe Bool)) -> IO (Maybe Bool))
-> (Ptr Word64 -> IO (Maybe Bool)) -> IO (Maybe Bool)
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> Ptr Word64 -> Int -> Bool -> Bool -> IO (Maybe Bool)
loop Ptr Word64
p Int
0 Bool
False Bool
False

-- | Look up index of first matching bit.
--
--   Note that the index type is limited to Int because there
--   is no 'unindex' method in the 'Ix' class.
{-# INLINE elemIndex #-}
elemIndex :: Bool -> IOBitArray Int -> IO (Maybe Int)
elemIndex :: Bool -> IOBitArray Int -> IO (Maybe Int)
elemIndex Bool
which IOBitArray Int
a = do
  (Int, Int)
bs <- IOBitArray Int -> IO (Int, Int)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray Int
a
  let skip :: Word64
      skip :: Word64
skip | Bool
which = Word64
0
           | Bool
otherwise = Word64 -> Word64
forall a. Bits a => a -> a
complement Word64
0
      total :: Int
total = (Int, Int) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (Int, Int)
bs
      full :: Int
full = Int
total Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement (Int -> Int
forall b. (Num b, Bits b) => Int -> b
mask Int
6)
      count :: Int
count = Int
full Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
      loop :: Ptr Word64 -> Int -> IO (Maybe Int)
      loop :: Ptr Word64 -> Int -> IO (Maybe Int)
loop Ptr Word64
p Int
n
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
count = do
            Word64
w <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
p Int
n
            if Word64
w Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
skip then Int -> IO (Maybe Int)
rest (Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
6) else Ptr Word64 -> Int -> IO (Maybe Int)
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Int -> IO (Maybe Int)
rest Int
full
      rest :: Int -> IO (Maybe Int)
rest Int
m
        | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
total = do
            Bool
b <- IOBitArray Int -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray Int
a Int
m
            if Bool
b Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
which then Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Maybe Int
forall a. a -> Maybe a
Just ((Int, Int) -> Int
forall a b. (a, b) -> a
fst (Int, Int)
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m)) else Int -> IO (Maybe Int)
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing
  ForeignPtr Word64
-> (Ptr Word64 -> IO (Maybe Int)) -> IO (Maybe Int)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray Int -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray Int
a) ((Ptr Word64 -> IO (Maybe Int)) -> IO (Maybe Int))
-> (Ptr Word64 -> IO (Maybe Int)) -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> Ptr Word64 -> Int -> IO (Maybe Int)
loop Ptr Word64
p Int
0

-- | Bitwise reduction with an associative commutative boolean operator.
--   Implementation lifts from 'Bool' to 'Bits' and folds large chunks
--   at a time.  Each bit is used as a source exactly once.
{-# INLINE fold #-}
fold :: Ix i => (Bool -> Bool -> Bool) {- ^ operator -} -> IOBitArray i -> IO (Maybe Bool)
fold :: forall i.
Ix i =>
(Bool -> Bool -> Bool) -> IOBitArray i -> IO (Maybe Bool)
fold Bool -> Bool -> Bool
f IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  let g :: Word64 -> Word64 -> Word64
g = (Bool -> Bool -> Bool) -> Word64 -> Word64 -> Word64
forall b. Bits b => (Bool -> Bool -> Bool) -> b -> b -> b
Bitwise.zipWith Bool -> Bool -> Bool
f
      total :: Int
total = (i, i) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (i, i)
bs
      full :: Int
full = Int
total Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement (Int -> Int
forall b. (Num b, Bits b) => Int -> b
mask Int
6)
      count :: Int
count = Int
full Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
      loop :: Ptr Word64 -> Int -> Maybe Word64 -> IO (Maybe Bool)
      loop :: Ptr Word64 -> Int -> Maybe Word64 -> IO (Maybe Bool)
loop Ptr Word64
p Int
n Maybe Word64
mw
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
count = do
            Word64
w <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
p Int
n
            case Maybe Word64
mw of
              Maybe Word64
Nothing -> Ptr Word64 -> Int -> Maybe Word64 -> IO (Maybe Bool)
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (Word64 -> Maybe Word64) -> Word64 -> Maybe Word64
forall a b. (a -> b) -> a -> b
$!      Word64
w)
              Just Word64
w0 -> Ptr Word64 -> Int -> Maybe Word64 -> IO (Maybe Bool)
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (Word64 -> Maybe Word64) -> Word64 -> Maybe Word64
forall a b. (a -> b) -> a -> b
$! Word64 -> Word64 -> Word64
g Word64
w0 Word64
w)
        | Bool
otherwise =
            case Maybe Word64
mw of
              Maybe Word64
Nothing -> Int -> Maybe Bool -> IO (Maybe Bool)
rest Int
full Maybe Bool
forall a. Maybe a
Nothing
              Just Word64
w0 -> Int -> Maybe Bool -> IO (Maybe Bool)
rest Int
full (Bool -> Maybe Bool
forall a. a -> Maybe a
Just (Bool -> Maybe Bool) -> Bool -> Maybe Bool
forall a b. (a -> b) -> a -> b
$! (Bool -> Bool -> Bool) -> [Bool] -> Bool
forall a. HasCallStack => (a -> a -> a) -> [a] -> a
foldl1' Bool -> Bool -> Bool
f (Word64 -> [Bool]
forall b. Bits b => b -> [Bool]
Bitwise.toListLE Word64
w0))
      rest :: Int -> Maybe Bool -> IO (Maybe Bool)
rest Int
m Maybe Bool
mb
        | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
total = do
            Bool
b <- IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a Int
m
            case Maybe Bool
mb of
              Maybe Bool
Nothing -> Int -> Maybe Bool -> IO (Maybe Bool)
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Bool -> Maybe Bool
forall a. a -> Maybe a
Just (Bool -> Maybe Bool) -> Bool -> Maybe Bool
forall a b. (a -> b) -> a -> b
$!      Bool
b)
              Just Bool
b0 -> Int -> Maybe Bool -> IO (Maybe Bool)
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Bool -> Maybe Bool
forall a. a -> Maybe a
Just (Bool -> Maybe Bool) -> Bool -> Maybe Bool
forall a b. (a -> b) -> a -> b
$! Bool -> Bool -> Bool
f Bool
b0 Bool
b)
        | Bool
otherwise = Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Bool
mb
  ForeignPtr Word64
-> (Ptr Word64 -> IO (Maybe Bool)) -> IO (Maybe Bool)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO (Maybe Bool)) -> IO (Maybe Bool))
-> (Ptr Word64 -> IO (Maybe Bool)) -> IO (Maybe Bool)
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> Ptr Word64 -> Int -> Maybe Word64 -> IO (Maybe Bool)
loop Ptr Word64
p Int
0 Maybe Word64
forall a. Maybe a
Nothing

-- | Bitwise map.  Implementation lifts from 'Bool' to 'Bits' and maps
--   large chunks at a time.
{-# INLINE map #-}
map :: Ix i => (Bool -> Bool) -> IOBitArray i -> IO (IOBitArray i)
map :: forall i.
Ix i =>
(Bool -> Bool) -> IOBitArray i -> IO (IOBitArray i)
map Bool -> Bool
f IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  IOBitArray i
b <- (i, i) -> IO (IOBitArray i)
forall i. Ix i => (i, i) -> IO (IOBitArray i)
newArray_ (i, i)
bs
  IOBitArray i -> (Bool -> Bool) -> IOBitArray i -> IO ()
forall i.
Ix i =>
IOBitArray i -> (Bool -> Bool) -> IOBitArray i -> IO ()
mapTo IOBitArray i
b Bool -> Bool
f IOBitArray i
a
  IOBitArray i -> IO (IOBitArray i)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IOBitArray i
b

{-# INLINE mapTo #-}
mapTo :: Ix i => IOBitArray i -> (Bool -> Bool) -> IOBitArray i -> IO ()
mapTo :: forall i.
Ix i =>
IOBitArray i -> (Bool -> Bool) -> IOBitArray i -> IO ()
mapTo IOBitArray i
dst Bool -> Bool
f IOBitArray i
src = do
  -- {
  (i, i)
sbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
src
  (i, i)
dbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
dst
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((i, i)
sbs (i, i) -> (i, i) -> Bool
forall a. Eq a => a -> a -> Bool
/= (i, i)
dbs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"mapTo mismatched bounds"
  -- }
  let count :: Int
count = IOBitArray i -> Int
forall i. IOBitArray i -> Int
iobBytes IOBitArray i
dst Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
      g :: Word64 -> Word64
      g :: Word64 -> Word64
g = (Bool -> Bool) -> Word64 -> Word64
forall b. Bits b => (Bool -> Bool) -> b -> b
Bitwise.map Bool -> Bool
f
  ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
src) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
sp ->
    ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
dst) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
dp ->
      [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
n -> do
        Ptr Word64 -> Int -> Word64 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word64
dp Int
n (Word64 -> IO ()) -> (Word64 -> Word64) -> Word64 -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Word64
g (Word64 -> IO ()) -> IO Word64 -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
sp Int
n

-- | Bitwise zipWith.  Implementation lifts from 'Bool' to 'Bits' and
--   combines large chunks at a time.
--
--   The bounds of the source arrays must be identical.
{-# INLINE zipWith #-}
zipWith :: Ix i => (Bool -> Bool -> Bool) -> IOBitArray i -> IOBitArray i -> IO (IOBitArray i)
zipWith :: forall i.
Ix i =>
(Bool -> Bool -> Bool)
-> IOBitArray i -> IOBitArray i -> IO (IOBitArray i)
zipWith Bool -> Bool -> Bool
f IOBitArray i
l IOBitArray i
r = do
  (i, i)
lbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
l
  (i, i)
rbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
r
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((i, i)
lbs (i, i) -> (i, i) -> Bool
forall a. Eq a => a -> a -> Bool
/= (i, i)
rbs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"zipWith mismatched bounds"
  IOBitArray i
c <- (i, i) -> IO (IOBitArray i)
forall i. Ix i => (i, i) -> IO (IOBitArray i)
newArray_ (i, i)
lbs
  IOBitArray i
-> (Bool -> Bool -> Bool) -> IOBitArray i -> IOBitArray i -> IO ()
forall i.
Ix i =>
IOBitArray i
-> (Bool -> Bool -> Bool) -> IOBitArray i -> IOBitArray i -> IO ()
zipWithTo IOBitArray i
c Bool -> Bool -> Bool
f IOBitArray i
l IOBitArray i
r
  IOBitArray i -> IO (IOBitArray i)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IOBitArray i
c

{-# INLINE zipWithTo #-}
zipWithTo :: Ix i => IOBitArray i -> (Bool -> Bool -> Bool) -> IOBitArray i -> IOBitArray i -> IO ()
zipWithTo :: forall i.
Ix i =>
IOBitArray i
-> (Bool -> Bool -> Bool) -> IOBitArray i -> IOBitArray i -> IO ()
zipWithTo IOBitArray i
dst Bool -> Bool -> Bool
f IOBitArray i
l IOBitArray i
r = do
  (i, i)
lbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
l
  (i, i)
rbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
r
  (i, i)
dbs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
dst
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((i, i)
lbs (i, i) -> (i, i) -> Bool
forall a. Eq a => a -> a -> Bool
/= (i, i)
rbs Bool -> Bool -> Bool
|| (i, i)
dbs (i, i) -> (i, i) -> Bool
forall a. Eq a => a -> a -> Bool
/= (i, i)
lbs Bool -> Bool -> Bool
|| (i, i)
dbs (i, i) -> (i, i) -> Bool
forall a. Eq a => a -> a -> Bool
/= (i, i)
rbs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"zipWithTo mismatched bounds"
  let count :: Int
count = IOBitArray i -> Int
forall i. IOBitArray i -> Int
iobBytes IOBitArray i
dst Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
3
      g :: Word64 -> Word64 -> Word64
      g :: Word64 -> Word64 -> Word64
g = (Bool -> Bool -> Bool) -> Word64 -> Word64 -> Word64
forall b. Bits b => (Bool -> Bool -> Bool) -> b -> b -> b
Bitwise.zipWith Bool -> Bool -> Bool
f
  ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
l) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
lp ->
    ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
r) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
rp ->
      ForeignPtr Word64 -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
dst) ((Ptr Word64 -> IO ()) -> IO ()) -> (Ptr Word64 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
dp ->
          [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
n -> do
            Word64
p <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
lp Int
n
            Word64
q <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
rp Int
n
            Ptr Word64 -> Int -> Word64 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word64
dp Int
n (Word64 -> Word64 -> Word64
g Word64
p Word64
q)

-- | Count set bits.
{-# INLINE popCount #-}
popCount :: Ix i => IOBitArray i -> IO Int
popCount :: forall i. Ix i => IOBitArray i -> IO Int
popCount IOBitArray i
a = do
  (i, i)
bs <- IOBitArray i -> IO (i, i)
forall i. Ix i => IOBitArray i -> IO (i, i)
getBounds IOBitArray i
a
  let total :: Int
total = (i, i) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (i, i)
bs
      full :: Int
full = Int
total Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement (Int -> Int
forall b. (Num b, Bits b) => Int -> b
mask Int
6)
      count :: Int
count = Int
full Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
      loop :: Ptr Word64 -> Int -> Int -> IO Int
      loop :: Ptr Word64 -> Int -> Int -> IO Int
loop Ptr Word64
p Int
n Int
acc
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
count = Int
acc Int -> IO Int -> IO Int
forall a b. a -> b -> b
`seq` do
            Word64
w <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word64
p Int
n
            Ptr Word64 -> Int -> Int -> IO Int
loop Ptr Word64
p (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word64 -> Int
forall a. Bits a => a -> Int
Data.Bits.popCount Word64
w)
        | Bool
otherwise = Int -> Int -> IO Int
rest Int
full Int
acc
      rest :: Int -> Int -> IO Int
rest Int
m Int
acc
        | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
total = Int
acc Int -> IO Int -> IO Int
forall a b. a -> b -> b
`seq` do
            Bool
b <- IOBitArray i -> Int -> IO Bool
forall i. Ix i => IOBitArray i -> Int -> IO Bool
readArrayRaw IOBitArray i
a Int
m
            Int -> Int -> IO Int
rest (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Bool -> Int
forall a. Enum a => a -> Int
fromEnum Bool
b)
        | Bool
otherwise = Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
acc
  ForeignPtr Word64 -> (Ptr Word64 -> IO Int) -> IO Int
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IOBitArray i -> ForeignPtr Word64
forall i. IOBitArray i -> ForeignPtr Word64
iobData IOBitArray i
a) ((Ptr Word64 -> IO Int) -> IO Int)
-> (Ptr Word64 -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Ptr Word64
p -> Ptr Word64 -> Int -> Int -> IO Int
loop Ptr Word64
p Int
0 Int
0