{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}
module Data.Chimera.Bool
( Chimera
, index
, trueIndices
, falseIndices
, tabulate
, tabulateFix
, tabulateM
, tabulateFixM
, 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 Data.Word
import Data.Chimera.Compat
import Data.Chimera.FromIntegral
newtype Chimera = Chimera { _unChimera :: V.Vector (U.Vector Word) }
bits :: Int
bits = fbs (0 :: Word)
bitsLog :: Int
bitsLog = bits - 1 - word2int (clz (int2word bits))
tabulate :: (Word -> Bool) -> Chimera
tabulate f = runIdentity $ tabulateM (return . f)
tabulateM :: forall m. Monad m => (Word -> m Bool) -> m Chimera
tabulateM f = do
z <- tabulateW 0
zs <- V.generateM (bits - bitsLog) tabulateU
return $ Chimera $ U.singleton z `V.cons` zs
where
tabulateU :: Int -> m (U.Vector Word)
tabulateU i = U.generateM ii (\j -> tabulateW (ii + j))
where
ii = 1 `shiftL` i
tabulateW :: Int -> m Word
tabulateW j = foldlM go 0 [0 .. bits - 1]
where
jj = j `shiftL` bitsLog
go acc k = do
b <- f (int2word $ jj + k)
return $ if b then acc `setBit` k else acc
{-# SPECIALIZE tabulateM :: (Word -> Identity Bool) -> Identity Chimera #-}
tabulateFix :: ((Word -> Bool) -> Word -> Bool) -> Chimera
tabulateFix uf = runIdentity $ tabulateFixM ((return .) . uf . (runIdentity .))
tabulateFixM :: forall m. Monad m => ((Word -> m Bool) -> Word -> m Bool) -> m Chimera
tabulateFixM uf = bs
where
bs :: m Chimera
bs = do
z <- tabulateW (fix uf) 0
zs <- V.generateM (bits - bitsLog) tabulateU
return $ Chimera $ U.singleton z `V.cons` zs
tabulateU :: Int -> m (U.Vector Word)
tabulateU i = U.generateM ii (\j -> tabulateW (uf f) (ii + j))
where
ii = 1 `shiftL` i
iii = ii `shiftL` bitsLog
f k = do
bs' <- bs
if k < int2word iii then return (index bs' k) else uf f k
tabulateW :: (Word -> m Bool) -> Int -> m Word
tabulateW f j = foldlM go 0 [0 .. bits - 1]
where
jj = j `shiftL` bitsLog
go acc k = do
b <- f (int2word $ jj + k)
return $ if b then acc `setBit` k else acc
{-# SPECIALIZE tabulateFixM :: ((Word -> Identity Bool) -> Word -> Identity Bool) -> Identity Chimera #-}
index :: Chimera -> Word -> Bool
index (Chimera vus) i =
if sgm < 0 then indexU (V.unsafeHead vus) (word2int i)
else indexU (vus `V.unsafeIndex` (sgm + 1)) (word2int $ i - int2word bits `shiftL` sgm)
where
sgm :: Int
sgm = fbs i - 1 - bitsLog - word2int (clz i)
indexU :: U.Vector Word -> Int -> Bool
indexU vec j = testBit (vec `U.unsafeIndex` jHi) jLo
where
jHi = j `shiftR` bitsLog
jLo = j .&. (bits - 1)
trueIndices :: Chimera -> [Word]
trueIndices bs = someIndices True bs
falseIndices :: Chimera -> [Word]
falseIndices bs = someIndices False bs
someIndices :: Bool -> Chimera -> [Word]
someIndices bool (Chimera b) = V.ifoldr goU [] b
where
goU :: Int -> U.Vector Word -> [Word] -> [Word]
goU i vec rest = U.ifoldr (\j -> goW (ii + j)) rest vec
where
ii = case i of
0 -> 0
_ -> 1 `shiftL` (i - 1)
goW :: Int -> Word -> [Word] -> [Word]
goW j w rest
= map (\k -> int2word $ jj + k)
(filter (\bt -> testBit w bt == bool) [0 .. bits - 1])
++ rest
where
jj = j `shiftL` bitsLog
{-# INLINE someIndices #-}
not :: Chimera -> Chimera
not (Chimera vus) = Chimera $ V.map (U.map (maxBound -)) vus
mapWithKey :: (Word -> Bool -> Bool) -> Chimera -> Chimera
mapWithKey f = runIdentity . traverseWithKey ((return .) . f)
traverseWithKey :: forall m. Monad m => (Word -> Bool -> m Bool) -> Chimera -> m Chimera
traverseWithKey f (Chimera bs) = do
bs' <- V.imapM g bs
return $ Chimera bs'
where
g :: Int -> U.Vector Word -> m (U.Vector Word)
g 0 = U.imapM h
g logOffset = U.imapM (h . (`shiftL` bitsLog) . (+ offset))
where
offset = 1 `shiftL` (logOffset - 1)
h :: Int -> Word -> m Word
h offset w = foldlM go 0 [0 .. bits - 1]
where
go acc k = do
b <- f (int2word $ offset + k) (testBit w k)
return $ if b then acc `setBit` k else acc
{-# SPECIALIZE traverseWithKey :: (Word -> Bool -> Identity Bool) -> Chimera -> Identity Chimera #-}
and :: Chimera -> Chimera -> Chimera
and (Chimera vus) (Chimera wus) = Chimera $ V.zipWith (U.zipWith (.&.)) vus wus
or :: Chimera -> Chimera -> Chimera
or (Chimera vus) (Chimera wus) = Chimera $ V.zipWith (U.zipWith (.|.)) vus wus
zipWithKey :: (Word -> Bool -> Bool -> Bool) -> Chimera -> Chimera -> Chimera
zipWithKey f = (runIdentity .) . zipWithKeyM (((return .) .) . f)
zipWithKeyM :: forall m. Monad m => (Word -> Bool -> Bool -> m Bool) -> Chimera -> Chimera -> m Chimera
zipWithKeyM f (Chimera bs1) (Chimera bs2) = do
bs' <- V.izipWithM g bs1 bs2
return $ Chimera bs'
where
g :: Int -> U.Vector Word -> U.Vector Word -> m (U.Vector Word)
g 0 = U.izipWithM h
g logOffset = U.izipWithM (h . (`shiftL` bitsLog) . (+ offset))
where
offset = 1 `shiftL` (logOffset - 1)
h :: Int -> Word -> Word -> m Word
h offset w1 w2 = foldlM go 0 [0 .. bits - 1]
where
go acc k = do
b <- f (int2word $ offset + k) (testBit w1 k) (testBit w2 k)
return $ if b then acc `setBit` k else acc
{-# SPECIALIZE zipWithKeyM :: (Word -> Bool -> Bool -> Identity Bool) -> Chimera -> Chimera -> Identity Chimera #-}