forbidden-fruit-0.1.0: src/Data/Vector/Dynamic.hs

{-# OPTIONS_GHC -w #-}
{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE TypeFamilies          #-}
module Data.Vector.Dynamic
( DynamicVector
, newDyn
, newDyn'
, readDyn
, writeDyn
, pushDyn
, popDyn
, shiftDyn
, unshiftDyn
, sizeDyn
, toListDyn
) where
import           Control.Imperative.Vector.Base (VectorElem, VectorEntity)
import           Control.Monad
import           Control.Monad.Primitive        (PrimMonad, PrimState)
import           Data.Bits
import           Data.Int
import           Data.Primitive.MutVar
import qualified Data.Vector.Generic.Mutable    as GMV
import qualified Data.Vector.Mutable            as MV
import qualified Data.Vector.Unboxed            as UV
import           Data.Word

data DynamicVector m a = Dynamic
  { logSizeVar :: MutVar (PrimState m) Int
  , frontVar   :: MutVar (PrimState m) Int
  , rearVar    :: MutVar (PrimState m) Int
  , vecVar     :: MutVar (PrimState m) (VectorEntity a (PrimState m) a)
  }

toPow2 :: Int -> Int
toPow2 n
  | n <= 0 = 1
  | otherwise = go (n-1) 1
    where
      go 0 !c = c
      go n !c = go (shiftR n 1) (shiftL c 1)
{-# INLINE toPow2 #-}

newDyn :: (PrimMonad m, VectorElem a) => Int -> m (DynamicVector m a)
newDyn logSize = liftM4 Dynamic
  (newMutVar logSize)
  (newMutVar 0)
  (newMutVar (max 0 (logSize-1)))
  (GMV.unsafeNew (toPow2 logSize) >>= newMutVar)
{-# INLINE newDyn #-}

newDyn' :: (PrimMonad m, VectorElem a) => Int -> a -> m (DynamicVector m a)
newDyn' logSize x = do
  d@(Dynamic {vecVar}) <- newDyn logSize
  v <- readMutVar vecVar
  w <- GMV.replicate logSize x
  unsafeMemCopy 0 v 0 (max 0 logSize) w
  return d
{-# INLINE newDyn' #-}

unsafeMemCopy
  :: (PrimMonad m, GMV.MVector v a)
  => Int -- dst offset
  -> v (PrimState m) a -- dst
  -> Int -- src offset
  -> Int -- length
  -> v (PrimState m) a -- src
  -> m ()
unsafeMemCopy i dst j n src = do
  let sdst = GMV.unsafeSlice i n dst
      ssrc = GMV.unsafeSlice j n src
  GMV.unsafeCopy sdst ssrc
{-# INLINE unsafeMemCopy #-}

readDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> Int -> m a
readDyn d@(Dynamic {logSizeVar, frontVar, rearVar, vecVar}) logIndex = do
  logSize  <- readMutVar logSizeVar
  when (logSize <= logIndex || logIndex < 0) $
    error $ "index out of bounds " ++ show (logIndex, logSize)
  phyIndex <- log2phy d logIndex
  vec <- readMutVar vecVar
  GMV.unsafeRead vec phyIndex
{-# INLINE readDyn #-}

log2phy :: (PrimMonad m, VectorElem a) => DynamicVector m a -> Int -> m Int
log2phy (Dynamic {frontVar, vecVar}) i = do
  front <- readMutVar frontVar
  vec   <- readMutVar vecVar
  let realCap = GMV.length vec
  return $ (front + i) .&. (realCap - 1)
{-# INLINE log2phy #-}

writeDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> Int -> a -> m ()
writeDyn d@(Dynamic {logSizeVar, rearVar, vecVar}) logIndex x = do
  logSize <- readMutVar logSizeVar
  when (logIndex < 0) $
    error $ "index out of bounds " ++ show (logIndex, logSize)
  when (logSize <= logIndex) $ do
    resizeDyn d (logIndex + 1)
    log2phy d logIndex >>= writeMutVar rearVar
  phyIndex <- log2phy d logIndex
  vec <- readMutVar vecVar
  GMV.unsafeWrite vec phyIndex x
{-# INLINE writeDyn #-}

resizeDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> Int -> m ()
resizeDyn (Dynamic {logSizeVar, frontVar, rearVar, vecVar}) newLogSize = do
  vec <- readMutVar vecVar
  let realCap = GMV.length vec
  when (realCap < newLogSize) $ do
    let newRealCap = until (newLogSize <=) (*2) realCap
        diff = newRealCap - realCap
    front <- readMutVar frontVar
    rear <- readMutVar rearVar
    newVec <- if front > rear
      then do
        let n = realCap - front
        nv <- GMV.unsafeNew newRealCap
        unsafeMemCopy 0 nv 0 (rear+1) vec
        unsafeMemCopy (front + diff) nv front n vec
        writeMutVar frontVar (front + diff)
        return nv
      else GMV.unsafeGrow vec diff
    writeMutVar vecVar newVec
  writeMutVar logSizeVar newLogSize
{-# INLINE resizeDyn #-}

pushDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> a -> m ()
pushDyn d@(Dynamic {logSizeVar, rearVar}) x = do
  logSize <- readMutVar logSizeVar
  writeDyn d logSize x
{-# INLINE pushDyn #-}

popDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> m a
popDyn d@(Dynamic {logSizeVar, frontVar, rearVar}) = do
  logSize <- readMutVar logSizeVar
  when (logSize < 1) $ error "Couldn't pop on an empty vector"
  x <- readDyn d (logSize-1)
  when (logSize > 1) $
    log2phy d (logSize-2) >>= writeMutVar rearVar
  modifyMutVar' logSizeVar pred
  return x
{-# INLINE popDyn #-}

shiftDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> m a
shiftDyn d@(Dynamic {logSizeVar, frontVar, rearVar, vecVar}) = do
  logSize <- readMutVar logSizeVar
  when (logSize < 1) $ error "Couldn't shift on an empty vector"
  x <- readDyn d 0
  when (logSize > 1) $
    log2phy d 1 >>= writeMutVar frontVar
  modifyMutVar' logSizeVar pred
  return x
{-# INLINE shiftDyn #-}

unshiftDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> a -> m ()
unshiftDyn d@(Dynamic {logSizeVar, frontVar, rearVar, vecVar}) x = do
  logSize <- readMutVar logSizeVar
  realCap <- liftM GMV.length $ readMutVar vecVar
  if logSize == realCap
    then resizeDyn d (logSize+1)
    else modifyMutVar' logSizeVar succ
  phyNewHead <- log2phy d (-1)
  vec <- readMutVar vecVar
  GMV.unsafeWrite vec phyNewHead x
  writeMutVar frontVar phyNewHead
  when (logSize==0) $ writeMutVar rearVar phyNewHead
{-# INLINE unshiftDyn #-}

toListDyn :: (PrimMonad m, VectorElem a) => DynamicVector m a -> m [a]
toListDyn (Dynamic {logSizeVar, frontVar, rearVar, vecVar}) = do
  logSize <- readMutVar logSizeVar
  if logSize == 0
    then return []
    else do
      front <- readMutVar frontVar
      rear <- readMutVar rearVar
      vec  <- readMutVar vecVar
      let realCap = GMV.length vec
          indicies = if front <= rear then [front..rear] else [front..realCap-1] ++ [0..rear]
      forM indicies $ GMV.unsafeRead vec
{-# INLINE toListDyn #-}

sizeDyn :: PrimMonad m => DynamicVector m a -> m Int
sizeDyn (Dynamic {logSizeVar}) = readMutVar logSizeVar
{-# INLINE sizeDyn #-}