{-# LANGUAGE ScopedTypeVariables #-}
module HaskellWorks.Data.Bits.Writer.Storable where
import Control.Monad.ST
import Data.Word
import HaskellWorks.Data.Bits.BitWise
import qualified Data.STRef as ST
import qualified Data.Vector.Storable.Mutable as DVSM
data Writer s = Writer
{ Writer s -> MVector s Word64
vector :: DVSM.MVector s Word64
, Writer s -> STRef s Int
position :: ST.STRef s Int
}
full :: Writer s -> ST s Bool
full :: Writer s -> ST s Bool
full Writer s
writer = do
Int
p <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
ST.readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer
Bool -> ST s Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> ST s Bool) -> Bool -> ST s Bool
forall a b. (a -> b) -> a -> b
$ Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
64 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= MVector s Word64 -> Int
forall a s. Storable a => MVector s a -> Int
DVSM.length (Writer s -> MVector s Word64
forall s. Writer s -> MVector s Word64
vector Writer s
writer)
{-# INLINE full #-}
newWriter :: Int -> ST s (Writer s)
newWriter :: Int -> ST s (Writer s)
newWriter Int
size = do
MVector s Word64
v <- Int -> ST s (MVector (PrimState (ST s)) Word64)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.new Int
size
STRef s Int
p <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
ST.newSTRef Int
0
Writer s -> ST s (Writer s)
forall (m :: * -> *) a. Monad m => a -> m a
return (Writer s -> ST s (Writer s)) -> Writer s -> ST s (Writer s)
forall a b. (a -> b) -> a -> b
$ MVector s Word64 -> STRef s Int -> Writer s
forall s. MVector s Word64 -> STRef s Int -> Writer s
Writer MVector s Word64
v STRef s Int
p
{-# INLINE newWriter #-}
unsafeWriteBit :: Writer s -> Word64 -> ST s ()
unsafeWriteBit :: Writer s -> Word64 -> ST s ()
unsafeWriteBit Writer s
writer Word64
w = do
let v :: MVector s Word64
v = Writer s -> MVector s Word64
forall s. Writer s -> MVector s Word64
vector Writer s
writer
Int
p <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
ST.readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer
let i :: Int
i = Int
p Int -> Word64 -> Int
forall a. Shift a => a -> Word64 -> a
.>. Word64
6
let o :: Int
o = Int
p Int -> Int -> Int
forall a. BitWise a => a -> a -> a
.&. Int
0x3f
Word64
e <- MVector (PrimState (ST s)) Word64 -> Int -> ST s Word64
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
DVSM.unsafeRead MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
i
MVector (PrimState (ST s)) Word64 -> Int -> Word64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
i (((Word64
w Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64
1) Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
o) Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. Word64
e)
STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
ST.writeSTRef (Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer) (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE unsafeWriteBit #-}
unsafeWriteLoBits :: Writer s -> Int -> Word64 -> ST s ()
unsafeWriteLoBits :: Writer s -> Int -> Word64 -> ST s ()
unsafeWriteLoBits Writer s
writer Int
c Word64
w = do
let u :: Word64
u = Word64
w Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. ((Word64
1 Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
let v :: MVector s Word64
v = Writer s -> MVector s Word64
forall s. Writer s -> MVector s Word64
vector Writer s
writer
Int
p <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
ST.readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer
let i :: Int
i = Int
p Int -> Word64 -> Int
forall a. Shift a => a -> Word64 -> a
.>. Word64
6
let o :: Int
o = Int
p Int -> Int -> Int
forall a. BitWise a => a -> a -> a
.&. Int
0x3f
Word64
lo <- MVector (PrimState (ST s)) Word64 -> Int -> ST s Word64
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
DVSM.unsafeRead MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
i
MVector (PrimState (ST s)) Word64 -> Int -> Word64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
i (Word64 -> ST s ()) -> Word64 -> ST s ()
forall a b. (a -> b) -> a -> b
$ Word64
lo Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. (Word64
u Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
o)
STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
ST.writeSTRef (Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer) (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c)
{-# INLINE unsafeWriteLoBits #-}
unsafeWriteBits :: Writer s -> Int -> Word64 -> ST s ()
unsafeWriteBits :: Writer s -> Int -> Word64 -> ST s ()
unsafeWriteBits Writer s
writer Int
c Word64
w = do
let u :: Word64
u = Word64
w Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. ((Word64
1 Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1)
let v :: MVector s Word64
v = Writer s -> MVector s Word64
forall s. Writer s -> MVector s Word64
vector Writer s
writer
Int
p <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
ST.readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer
let i :: Int
i = Int
p Int -> Word64 -> Int
forall a. Shift a => a -> Word64 -> a
.>. Word64
6
let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
let o :: Int
o = Int
p Int -> Int -> Int
forall a. BitWise a => a -> a -> a
.&. Int
0x3f
Word64
lo <- MVector (PrimState (ST s)) Word64 -> Int -> ST s Word64
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
DVSM.unsafeRead MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
i
MVector (PrimState (ST s)) Word64 -> Int -> Word64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
i (Word64 -> ST s ()) -> Word64 -> ST s ()
forall a b. (a -> b) -> a -> b
$ Word64
lo Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. (Word64
u Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
o)
Word64
hi <- MVector (PrimState (ST s)) Word64 -> Int -> ST s Word64
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
DVSM.unsafeRead MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
j
MVector (PrimState (ST s)) Word64 -> Int -> Word64 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector s Word64
MVector (PrimState (ST s)) Word64
v Int
j (Word64 -> ST s ()) -> Word64 -> ST s ()
forall a b. (a -> b) -> a -> b
$ Word64
hi Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. (Word64
u Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.>. Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
64 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o))
STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
ST.writeSTRef (Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer) (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
c)
{-# INLINE unsafeWriteBits #-}
written :: Writer s -> ST s (DVSM.MVector s Word64)
written :: Writer s -> ST s (MVector s Word64)
written Writer s
writer = do
Int
p <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
ST.readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Writer s -> STRef s Int
forall s. Writer s -> STRef s Int
position Writer s
writer
MVector s Word64 -> ST s (MVector s Word64)
forall (m :: * -> *) a. Monad m => a -> m a
return (MVector s Word64 -> ST s (MVector s Word64))
-> MVector s Word64 -> ST s (MVector s Word64)
forall a b. (a -> b) -> a -> b
$ Int -> MVector s Word64 -> MVector s Word64
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take ((Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
63) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
64) (Writer s -> MVector s Word64
forall s. Writer s -> MVector s Word64
vector Writer s
writer)
{-# INLINE written #-}