-- |
-- Module:      Data.BitStream
-- Copyright:   (c) 2017 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Lazy, infinite, compact stream of 'Bool' with O(1) indexing.
-- Most useful for memoization of predicates.
--
-- __Example 1__
--
-- Consider following predicate:
--
-- > isOdd :: Word -> Bool
-- > isOdd 0 = False
-- > isOdd n = not (isOdd (n - 1))
--
-- Its computation is expensive, so we'd like to memoize its values into
-- 'BitStream' using 'tabulate' and access this stream via 'index'
-- instead of recalculation of @isOdd@:
--
-- > isOddBS :: BitStream
-- > isOddBS = tabulate isOdd
-- >
-- > isOdd' :: Word -> Bool
-- > isOdd' = index isOddBS
--
-- We can do even better by replacing part of recursive calls to @isOdd@
-- by indexing memoized values. Write @isOddF@
-- such that @isOdd = 'fix' isOddF@:
--
-- > isOddF :: (Word -> Bool) -> Word -> Bool
-- > isOddF _ 0 = False
-- > isOddF f n = not (f (n - 1))
--
-- and use 'tabulateFix':
--
-- > isOddBS :: BitStream
-- > isOddBS = tabulateFix isOddF
-- >
-- > isOdd' :: Word -> Bool
-- > isOdd' = index isOddBS
--
-- __Example 2__
--
-- Define a predicate, which checks whether its argument is
-- a prime number by trial division.
--
-- > isPrime :: Word -> Bool
-- > isPrime n
-- >   | n < 2     = False
-- >   | n < 4     = True
-- >   | even n    = False
-- >   | otherwise = and [ n `rem` d /= 0 | d <- [3, 5 .. ceiling (sqrt (fromIntegral n))], isPrime d]
--
-- Convert it to unfixed form:
--
-- > isPrimeF :: (Word -> Bool) -> Word -> Bool
-- > isPrimeF f n
-- >   | n < 2     = False
-- >   | n < 4     = True
-- >   | even n    = False
-- >   | otherwise = and [ n `rem` d /= 0 | d <- [3, 5 .. ceiling (sqrt (fromIntegral n))], f d]
--
-- Create its memoized version for faster evaluation:
--
-- > isPrimeBS :: BitStream
-- > isPrimeBS = tabulateFix isPrimeF
-- >
-- > isPrime' :: Word -> Bool
-- > isPrime' = index isPrimeBS

{-# LANGUAGE ScopedTypeVariables #-}

module Data.BitStream
  ( BitStream
  , tabulate
  , tabulateFix
  , tabulateM
  , tabulateFixM
  , index

  , mapWithKey
  , traverseWithKey
  , not

  , zipWithKey
  , zipWithKeyM
  , and
  , or
  ) where

import Prelude hiding ((^), (*), div, mod, fromIntegral, not, and, or)
import Data.Bits
import Data.Foldable hiding (and, or)
import Data.Function (fix)
import Data.Functor.Identity
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
import Unsafe.Coerce

-- | Compact representation of infinite stream of 'Bool'.
--
-- It spends one bit (1/8 byte) for one 'Bool' in store.
-- Compare it to at least 24 bytes per element in @[Bool]@,
-- approximately 2 bytes per element in 'IntSet'
-- and 1 byte per element in unboxed @Vector Bool@.
--
-- It also offers indexing in constant time.
-- Compare it to linear time for lists and logarithmic time for sets.
--
-- Moreover, it is lazy: querying n-th element triggers computation
-- of first @max(64, 2 ^ ceiling (logBase 2 n))@ elements only. On contrary,
-- sets and unboxed vectors are completely strict.
newtype BitStream = BitStream { BitStream -> Vector (Vector Word)
_unBitStream :: V.Vector (U.Vector Word) }

word2int :: Word -> Int
word2int :: Word -> Int
word2int = Word -> Int
forall a b. a -> b
unsafeCoerce

int2word :: Int -> Word
int2word :: Int -> Word
int2word = Int -> Word
forall a b. a -> b
unsafeCoerce

bits :: Int
bits :: Int
bits = Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word)

bitsLog :: Int
bitsLog :: Int
bitsLog = Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int -> Word
int2word Int
bits)

-- | Create a bit stream from the predicate.
-- The predicate must be well-defined for any value of argument
-- and should not return 'error' / 'undefined'.
tabulate :: (Word -> Bool) -> BitStream
tabulate :: (Word -> Bool) -> BitStream
tabulate Word -> Bool
f = Identity BitStream -> BitStream
forall a. Identity a -> a
runIdentity (Identity BitStream -> BitStream)
-> Identity BitStream -> BitStream
forall a b. (a -> b) -> a -> b
$ (Word -> Identity Bool) -> Identity BitStream
forall (m :: * -> *). Monad m => (Word -> m Bool) -> m BitStream
tabulateM (Bool -> Identity Bool
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Identity Bool) -> (Word -> Bool) -> Word -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Bool
f)

-- | Create a bit stream from the monadic predicate.
-- The predicate must be well-defined for any value of argument
-- and should not return 'error' / 'undefined'.
tabulateM :: forall m. Monad m => (Word -> m Bool) -> m BitStream
tabulateM :: forall (m :: * -> *). Monad m => (Word -> m Bool) -> m BitStream
tabulateM Word -> m Bool
f = do
  Word
z  <- Int -> m Word
tabulateW Int
0
  Vector (Vector Word)
zs <- Int -> (Int -> m (Vector Word)) -> m (Vector (Vector Word))
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Vector a)
V.generateM (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bitsLog) Int -> m (Vector Word)
tabulateU
  BitStream -> m BitStream
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (BitStream -> m BitStream) -> BitStream -> m BitStream
forall a b. (a -> b) -> a -> b
$ Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> Vector (Vector Word) -> BitStream
forall a b. (a -> b) -> a -> b
$ Word -> Vector Word
forall a. Unbox a => a -> Vector a
U.singleton Word
z Vector Word -> Vector (Vector Word) -> Vector (Vector Word)
forall a. a -> Vector a -> Vector a
`V.cons` Vector (Vector Word)
zs
  where
    tabulateU :: Int -> m (U.Vector Word)
    tabulateU :: Int -> m (Vector Word)
tabulateU Int
i = Int -> (Int -> m Word) -> m (Vector Word)
forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> (Int -> m a) -> m (Vector a)
U.generateM Int
ii (\Int
j -> Int -> m Word
tabulateW (Int
ii Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j))
      where
        ii :: Int
ii = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
i

    tabulateW :: Int -> m Word
    tabulateW :: Int -> m Word
tabulateW Int
j = (Word -> Int -> m Word) -> Word -> [Int] -> m Word
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Word -> Int -> m Word
forall {b}. Bits b => b -> Int -> m b
go Word
forall a. Bits a => a
zeroBits [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      where
        jj :: Int
jj = Int
j Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
bitsLog
        go :: b -> Int -> m b
go b
acc Int
k = do
          Bool
b <- Word -> m Bool
f (Int -> Word
int2word (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ Int
jj Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)
          b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ if Bool
b then b
acc b -> Int -> b
forall a. Bits a => a -> Int -> a
`setBit` Int
k else b
acc
{-# SPECIALIZE tabulateM :: (Word -> Identity Bool) -> Identity BitStream #-}

-- | Create a bit stream from the unfixed predicate.
-- The predicate must be well-defined for any value of argument
-- and should not return 'error' / 'undefined'.
tabulateFix :: ((Word -> Bool) -> Word -> Bool) -> BitStream
tabulateFix :: ((Word -> Bool) -> Word -> Bool) -> BitStream
tabulateFix (Word -> Bool) -> Word -> Bool
uf = Identity BitStream -> BitStream
forall a. Identity a -> a
runIdentity (Identity BitStream -> BitStream)
-> Identity BitStream -> BitStream
forall a b. (a -> b) -> a -> b
$ ((Word -> Identity Bool) -> Word -> Identity Bool)
-> Identity BitStream
forall (m :: * -> *).
Monad m =>
((Word -> m Bool) -> Word -> m Bool) -> m BitStream
tabulateFixM ((Bool -> Identity Bool
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Identity Bool) -> (Word -> Bool) -> Word -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Word -> Bool) -> Word -> Identity Bool)
-> ((Word -> Identity Bool) -> Word -> Bool)
-> (Word -> Identity Bool)
-> Word
-> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> Bool) -> Word -> Bool
uf ((Word -> Bool) -> Word -> Bool)
-> ((Word -> Identity Bool) -> Word -> Bool)
-> (Word -> Identity Bool)
-> Word
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Identity Bool -> Bool
forall a. Identity a -> a
runIdentity (Identity Bool -> Bool) -> (Word -> Identity Bool) -> Word -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.))

-- | Create a bit stream from the unfixed monadic predicate.
-- The predicate must be well-defined for any value of argument
-- and should not return 'error' / 'undefined'.
tabulateFixM :: forall m. Monad m => ((Word -> m Bool) -> Word -> m Bool) -> m BitStream
tabulateFixM :: forall (m :: * -> *).
Monad m =>
((Word -> m Bool) -> Word -> m Bool) -> m BitStream
tabulateFixM (Word -> m Bool) -> Word -> m Bool
uf = m BitStream
bs
  where
    bs :: m BitStream
    bs :: m BitStream
bs = do
      Word
z  <- (Word -> m Bool) -> Int -> m Word
tabulateW (((Word -> m Bool) -> Word -> m Bool) -> Word -> m Bool
forall a. (a -> a) -> a
fix (Word -> m Bool) -> Word -> m Bool
uf) Int
0
      Vector (Vector Word)
zs <- Int -> (Int -> m (Vector Word)) -> m (Vector (Vector Word))
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Vector a)
V.generateM (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bitsLog) Int -> m (Vector Word)
tabulateU
      BitStream -> m BitStream
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (BitStream -> m BitStream) -> BitStream -> m BitStream
forall a b. (a -> b) -> a -> b
$ Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> Vector (Vector Word) -> BitStream
forall a b. (a -> b) -> a -> b
$ Word -> Vector Word
forall a. Unbox a => a -> Vector a
U.singleton Word
z Vector Word -> Vector (Vector Word) -> Vector (Vector Word)
forall a. a -> Vector a -> Vector a
`V.cons` Vector (Vector Word)
zs

    tabulateU :: Int -> m (U.Vector Word)
    tabulateU :: Int -> m (Vector Word)
tabulateU Int
i = Int -> (Int -> m Word) -> m (Vector Word)
forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> (Int -> m a) -> m (Vector a)
U.generateM Int
ii (\Int
j -> (Word -> m Bool) -> Int -> m Word
tabulateW ((Word -> m Bool) -> Word -> m Bool
uf Word -> m Bool
f) (Int
ii Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j))
      where
        ii :: Int
ii = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
i
        iii :: Int
iii = Int
ii Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
bitsLog
        f :: Word -> m Bool
f Word
k = do
          BitStream
bs' <- m BitStream
bs
          if Word
k Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Word
int2word Int
iii then Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (BitStream -> Word -> Bool
index BitStream
bs' Word
k) else (Word -> m Bool) -> Word -> m Bool
uf Word -> m Bool
f Word
k

    tabulateW :: (Word -> m Bool) -> Int -> m Word
    tabulateW :: (Word -> m Bool) -> Int -> m Word
tabulateW Word -> m Bool
f Int
j = (Word -> Int -> m Word) -> Word -> [Int] -> m Word
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Word -> Int -> m Word
forall {b}. Bits b => b -> Int -> m b
go Word
forall a. Bits a => a
zeroBits [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      where
        jj :: Int
jj = Int
j Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
bitsLog
        go :: b -> Int -> m b
go b
acc Int
k = do
          Bool
b <- Word -> m Bool
f (Int -> Word
int2word (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ Int
jj Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)
          b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ if Bool
b then b
acc b -> Int -> b
forall a. Bits a => a -> Int -> a
`setBit` Int
k else b
acc
{-# SPECIALIZE tabulateFixM :: ((Word -> Identity Bool) -> Word -> Identity Bool) -> Identity BitStream #-}

-- | Convert a bit stream back to predicate.
-- Indexing itself works in O(1) time, but triggers evaluation and allocation
-- of surrounding elements of the stream, if they were not computed before.
index :: BitStream -> Word -> Bool
index :: BitStream -> Word -> Bool
index (BitStream Vector (Vector Word)
vus) Word
i =
  if Int
sgm Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then Vector Word -> Int -> Bool
indexU (Vector (Vector Word) -> Vector Word
forall a. Vector a -> a
V.unsafeHead Vector (Vector Word)
vus) (Word -> Int
word2int Word
i)
  else Vector Word -> Int -> Bool
indexU (Vector (Vector Word)
vus Vector (Vector Word) -> Int -> Vector Word
forall a. Vector a -> Int -> a
`V.unsafeIndex` (Int
sgm Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Word -> Int
word2int (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$ Word
i Word -> Word -> Word
forall a. Num a => a -> a -> a
- Int -> Word
int2word Int
bits Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` Int
sgm)
  where
    sgm :: Int
    sgm :: Int
sgm = Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize Word
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bitsLog Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Word
i

    indexU :: U.Vector Word -> Int -> Bool
    indexU :: Vector Word -> Int -> Bool
indexU Vector Word
vec Int
j = Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit (Vector Word
vec Vector Word -> Int -> Word
forall a. Unbox a => Vector a -> Int -> a
`U.unsafeIndex` Int
jHi) Int
jLo
      where
        jHi :: Int
jHi = Int
j Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
bitsLog
        jLo :: Int
jLo = Int
j Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | Element-wise 'not'.
not :: BitStream -> BitStream
not :: BitStream -> BitStream
not (BitStream Vector (Vector Word)
vus) = Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> Vector (Vector Word) -> BitStream
forall a b. (a -> b) -> a -> b
$ (Vector Word -> Vector Word)
-> Vector (Vector Word) -> Vector (Vector Word)
forall a b. (a -> b) -> Vector a -> Vector b
V.map ((Word -> Word) -> Vector Word -> Vector Word
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
U.map (Word
forall a. Bounded a => a
maxBound Word -> Word -> Word
forall a. Num a => a -> a -> a
-)) Vector (Vector Word)
vus

-- | Map over all indices and respective elements in the stream.
mapWithKey :: (Word -> Bool -> Bool) -> BitStream -> BitStream
mapWithKey :: (Word -> Bool -> Bool) -> BitStream -> BitStream
mapWithKey Word -> Bool -> Bool
f = Identity BitStream -> BitStream
forall a. Identity a -> a
runIdentity (Identity BitStream -> BitStream)
-> (BitStream -> Identity BitStream) -> BitStream -> BitStream
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> Bool -> Identity Bool) -> BitStream -> Identity BitStream
forall (m :: * -> *).
Monad m =>
(Word -> Bool -> m Bool) -> BitStream -> m BitStream
traverseWithKey ((Bool -> Identity Bool
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Identity Bool) -> (Bool -> Bool) -> Bool -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Bool -> Bool) -> Bool -> Identity Bool)
-> (Word -> Bool -> Bool) -> Word -> Bool -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Bool -> Bool
f)

-- | Traverse over all indices and respective elements in the stream.
traverseWithKey :: forall m. Monad m => (Word -> Bool -> m Bool) -> BitStream -> m BitStream
traverseWithKey :: forall (m :: * -> *).
Monad m =>
(Word -> Bool -> m Bool) -> BitStream -> m BitStream
traverseWithKey Word -> Bool -> m Bool
f (BitStream Vector (Vector Word)
bs) = Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> m (Vector (Vector Word)) -> m BitStream
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Vector Word -> m (Vector Word))
-> Vector (Vector Word) -> m (Vector (Vector Word))
forall (m :: * -> *) a b.
Monad m =>
(Int -> a -> m b) -> Vector a -> m (Vector b)
V.imapM Int -> Vector Word -> m (Vector Word)
g Vector (Vector Word)
bs
  where
    g :: Int -> U.Vector Word -> m (U.Vector Word)
    g :: Int -> Vector Word -> m (Vector Word)
g Int
0         = (Int -> Word -> m Word) -> Vector Word -> m (Vector Word)
forall (m :: * -> *) a b.
(Monad m, Unbox a, Unbox b) =>
(Int -> a -> m b) -> Vector a -> m (Vector b)
U.imapM Int -> Word -> m Word
h
    g Int
logOffset = (Int -> Word -> m Word) -> Vector Word -> m (Vector Word)
forall (m :: * -> *) a b.
(Monad m, Unbox a, Unbox b) =>
(Int -> a -> m b) -> Vector a -> m (Vector b)
U.imapM (Int -> Word -> m Word
h (Int -> Word -> m Word) -> (Int -> Int) -> Int -> Word -> m Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
bitsLog) (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset))
      where
        offset :: Int
offset = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` (Int
logOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

    h :: Int -> Word -> m Word
    h :: Int -> Word -> m Word
h Int
offset Word
w = (Word -> Int -> m Word) -> Word -> [Int] -> m Word
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Word -> Int -> m Word
forall {b}. Bits b => b -> Int -> m b
go Word
forall a. Bits a => a
zeroBits [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      where
        go :: b -> Int -> m b
go b
acc Int
k = do
          Bool
b <- Word -> Bool -> m Bool
f (Int -> Word
int2word (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word
w Int
k)
          b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ if Bool
b then b
acc b -> Int -> b
forall a. Bits a => a -> Int -> a
`setBit` Int
k else b
acc
{-# SPECIALIZE traverseWithKey :: (Word -> Bool -> Identity Bool) -> BitStream -> Identity BitStream #-}

-- | Element-wise 'and'.
and :: BitStream -> BitStream -> BitStream
and :: BitStream -> BitStream -> BitStream
and (BitStream Vector (Vector Word)
vus) (BitStream Vector (Vector Word)
wus) = Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> Vector (Vector Word) -> BitStream
forall a b. (a -> b) -> a -> b
$ (Vector Word -> Vector Word -> Vector Word)
-> Vector (Vector Word)
-> Vector (Vector Word)
-> Vector (Vector Word)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith ((Word -> Word -> Word) -> Vector Word -> Vector Word -> Vector Word
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
U.zipWith Word -> Word -> Word
forall a. Bits a => a -> a -> a
(.&.)) Vector (Vector Word)
vus Vector (Vector Word)
wus

-- | Element-wise 'or'.
or  :: BitStream -> BitStream -> BitStream
or :: BitStream -> BitStream -> BitStream
or (BitStream Vector (Vector Word)
vus) (BitStream Vector (Vector Word)
wus) = Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> Vector (Vector Word) -> BitStream
forall a b. (a -> b) -> a -> b
$ (Vector Word -> Vector Word -> Vector Word)
-> Vector (Vector Word)
-> Vector (Vector Word)
-> Vector (Vector Word)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith ((Word -> Word -> Word) -> Vector Word -> Vector Word -> Vector Word
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
U.zipWith Word -> Word -> Word
forall a. Bits a => a -> a -> a
(.|.)) Vector (Vector Word)
vus Vector (Vector Word)
wus

-- | Zip two streams with the function, which is provided with an index and respective elements of both streams.
zipWithKey :: (Word -> Bool -> Bool -> Bool) -> BitStream -> BitStream -> BitStream
zipWithKey :: (Word -> Bool -> Bool -> Bool)
-> BitStream -> BitStream -> BitStream
zipWithKey Word -> Bool -> Bool -> Bool
f = (Identity BitStream -> BitStream
forall a. Identity a -> a
runIdentity (Identity BitStream -> BitStream)
-> (BitStream -> Identity BitStream) -> BitStream -> BitStream
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((BitStream -> Identity BitStream) -> BitStream -> BitStream)
-> (BitStream -> BitStream -> Identity BitStream)
-> BitStream
-> BitStream
-> BitStream
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> Bool -> Bool -> Identity Bool)
-> BitStream -> BitStream -> Identity BitStream
forall (m :: * -> *).
Monad m =>
(Word -> Bool -> Bool -> m Bool)
-> BitStream -> BitStream -> m BitStream
zipWithKeyM (((Bool -> Identity Bool
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Identity Bool) -> (Bool -> Bool) -> Bool -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Bool -> Bool) -> Bool -> Identity Bool)
-> (Bool -> Bool -> Bool) -> Bool -> Bool -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Bool -> Bool -> Bool) -> Bool -> Bool -> Identity Bool)
-> (Word -> Bool -> Bool -> Bool)
-> Word
-> Bool
-> Bool
-> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Bool -> Bool -> Bool
f)

-- | Zip two streams with the monadic function, which is provided with an index and respective elements of both streams.
zipWithKeyM :: forall m. Monad m => (Word -> Bool -> Bool -> m Bool) -> BitStream -> BitStream -> m BitStream
zipWithKeyM :: forall (m :: * -> *).
Monad m =>
(Word -> Bool -> Bool -> m Bool)
-> BitStream -> BitStream -> m BitStream
zipWithKeyM Word -> Bool -> Bool -> m Bool
f (BitStream Vector (Vector Word)
bs1) (BitStream Vector (Vector Word)
bs2) = Vector (Vector Word) -> BitStream
BitStream (Vector (Vector Word) -> BitStream)
-> m (Vector (Vector Word)) -> m BitStream
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Vector Word -> Vector Word -> m (Vector Word))
-> Vector (Vector Word)
-> Vector (Vector Word)
-> m (Vector (Vector Word))
forall (m :: * -> *) a b c.
Monad m =>
(Int -> a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
V.izipWithM Int -> Vector Word -> Vector Word -> m (Vector Word)
g Vector (Vector Word)
bs1 Vector (Vector Word)
bs2
  where
    g :: Int -> U.Vector Word -> U.Vector Word -> m (U.Vector Word)
    g :: Int -> Vector Word -> Vector Word -> m (Vector Word)
g Int
0         = (Int -> Word -> Word -> m Word)
-> Vector Word -> Vector Word -> m (Vector Word)
forall (m :: * -> *) a b c.
(Monad m, Unbox a, Unbox b, Unbox c) =>
(Int -> a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
U.izipWithM Int -> Word -> Word -> m Word
h
    g Int
logOffset = (Int -> Word -> Word -> m Word)
-> Vector Word -> Vector Word -> m (Vector Word)
forall (m :: * -> *) a b c.
(Monad m, Unbox a, Unbox b, Unbox c) =>
(Int -> a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
U.izipWithM (Int -> Word -> Word -> m Word
h (Int -> Word -> Word -> m Word)
-> (Int -> Int) -> Int -> Word -> Word -> m Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
bitsLog) (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset))
      where
        offset :: Int
offset = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` (Int
logOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

    h :: Int -> Word -> Word -> m Word
    h :: Int -> Word -> Word -> m Word
h Int
offset Word
w1 Word
w2 = (Word -> Int -> m Word) -> Word -> [Int] -> m Word
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Word -> Int -> m Word
forall {b}. Bits b => b -> Int -> m b
go Word
forall a. Bits a => a
zeroBits [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      where
        go :: b -> Int -> m b
go b
acc Int
k = do
          Bool
b <- Word -> Bool -> Bool -> m Bool
f (Int -> Word
int2word (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word
w1 Int
k) (Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word
w2 Int
k)
          b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ if Bool
b then b
acc b -> Int -> b
forall a. Bits a => a -> Int -> a
`setBit` Int
k else b
acc
{-# SPECIALIZE zipWithKeyM :: (Word -> Bool -> Bool -> Identity Bool) -> BitStream -> BitStream -> Identity BitStream #-}