{-# 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

{- HLINT ignore "Reduce duplication"  -}

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             -- vector
  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 -- position
  let i :: Int
i = Int
p Int -> Word64 -> Int
forall a. Shift a => a -> Word64 -> a
.>. Word64
6                     -- index into vector
  let o :: Int
o = Int
p Int -> Int -> Int
forall a. BitWise a => a -> a -> a
.&. Int
0x3f                  -- offset within a word
  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             -- vector
  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 -- position
  let i :: Int
i = Int
p Int -> Word64 -> Int
forall a. Shift a => a -> Word64 -> a
.>. Word64
6                     -- index into vector
  let o :: Int
o = Int
p Int -> Int -> Int
forall a. BitWise a => a -> a -> a
.&. Int
0x3f                  -- offset within a word
  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)  -- masked word
  let v :: MVector s Word64
v = Writer s -> MVector s Word64
forall s. Writer s -> MVector s Word64
vector Writer s
writer                       -- vector
  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         -- position
  let i :: Int
i = Int
p Int -> Word64 -> Int
forall a. Shift a => a -> Word64 -> a
.>. Word64
6                             -- index into vector for lo part
  let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1                               -- index into vector for hi part
  let o :: Int
o = Int
p Int -> Int -> Int
forall a. BitWise a => a -> a -> a
.&. Int
0x3f                          -- offset within a word
  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         -- position
  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 #-}