{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}
module Data.Chimera.Unboxed
( Chimera
, index
, toList
, tabulate
, tabulateFix
, tabulateM
, tabulateFixM
, mapWithKey
, traverseWithKey
, zipWithKey
, zipWithKeyM
) where
import Prelude hiding ((^), (*), div, mod, fromIntegral, not, and, or, iterate)
import Data.Bits
import Data.Foldable hiding (and, or, toList)
import Data.Function (fix)
import Data.Functor.Identity
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Data.Word
import Data.Chimera.Compat
import Data.Chimera.FromIntegral
newtype Chimera a = Chimera { _unChimera :: V.Vector (U.Vector a) }
bits :: Int
bits = fbs (0 :: Word)
tabulate :: U.Unbox a => (Word -> a) -> Chimera a
tabulate f = runIdentity $ tabulateM (return . f)
tabulateM :: forall m a. (Monad m, U.Unbox a) => (Word -> m a) -> m (Chimera a)
tabulateM f = do
z <- f 0
zs <- V.generateM bits tabulateU
return $ Chimera $ U.singleton z `V.cons` zs
where
tabulateU :: Int -> m (U.Vector a)
tabulateU i = U.generateM ii (\j -> f (int2word (ii + j)))
where
ii = 1 `shiftL` i
{-# SPECIALIZE tabulateM :: U.Unbox a => (Word -> Identity a) -> Identity (Chimera a) #-}
tabulateFix :: U.Unbox a => ((Word -> a) -> Word -> a) -> Chimera a
tabulateFix uf = runIdentity $ tabulateFixM ((return .) . uf . (runIdentity .))
tabulateFixM :: forall m a. (Monad m, U.Unbox a) => ((Word -> m a) -> Word -> m a) -> m (Chimera a)
tabulateFixM uf = bs
where
bs :: m (Chimera a)
bs = do
z <- fix uf 0
zs <- V.generateM bits tabulateU
return $ Chimera $ U.singleton z `V.cons` zs
tabulateU :: Int -> m (U.Vector a)
tabulateU i = U.generateM ii (\j -> uf f (int2word (ii + j)))
where
ii = 1 `shiftL` i
f k = do
bs' <- bs
if k < int2word ii then return (index bs' k) else uf f k
{-# SPECIALIZE tabulateFixM :: U.Unbox a => ((Word -> Identity a) -> Word -> Identity a) -> Identity (Chimera a) #-}
index :: U.Unbox a => Chimera a -> Word -> a
index (Chimera vus) 0 = U.unsafeHead (V.unsafeHead vus)
index (Chimera vus) i = U.unsafeIndex (vus `V.unsafeIndex` (sgm + 1)) (word2int $ i - 1 `shiftL` sgm)
where
sgm :: Int
sgm = fbs i - 1 - word2int (clz i)
toList :: U.Unbox a => Chimera a -> [a]
toList (Chimera vus) = foldMap U.toList vus
mapWithKey :: (U.Unbox a, U.Unbox b) => (Word -> a -> b) -> Chimera a -> Chimera b
mapWithKey f = runIdentity . traverseWithKey ((return .) . f)
traverseWithKey :: forall m a b. (Monad m, U.Unbox a, U.Unbox b) => (Word -> a -> m b) -> Chimera a -> m (Chimera b)
traverseWithKey f (Chimera bs) = do
bs' <- V.imapM g bs
return $ Chimera bs'
where
g :: Int -> U.Vector a -> m (U.Vector b)
g 0 = U.imapM (f . int2word)
g logOffset = U.imapM (f . int2word . (+ offset))
where
offset = 1 `shiftL` (logOffset - 1)
{-# SPECIALIZE traverseWithKey :: U.Unbox a => (Word -> a -> Identity a) -> Chimera a -> Identity (Chimera a) #-}
zipWithKey :: (U.Unbox a, U.Unbox b, U.Unbox c) => (Word -> a -> b -> c) -> Chimera a -> Chimera b -> Chimera c
zipWithKey f = (runIdentity .) . zipWithKeyM (((return .) .) . f)
zipWithKeyM :: forall m a b c. (Monad m, U.Unbox a, U.Unbox b, U.Unbox c) => (Word -> a -> b -> m c) -> Chimera a -> Chimera b -> m (Chimera c)
zipWithKeyM f (Chimera bs1) (Chimera bs2) = do
bs' <- V.izipWithM g bs1 bs2
return $ Chimera bs'
where
g :: Int -> U.Vector a -> U.Vector b -> m (U.Vector c)
g 0 = U.izipWithM (f . int2word)
g logOffset = U.izipWithM (f . int2word . (+ offset))
where
offset = 1 `shiftL` (logOffset - 1)
{-# SPECIALIZE zipWithKeyM :: U.Unbox a => (Word -> a -> a -> Identity a) -> Chimera a -> Chimera a -> Identity (Chimera a) #-}