vector-functorlazy-0.0.1: src/Data/Vector/FunctorLazy/Mutable.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables,TypeFamilies,FlexibleInstances,MultiParamTypeClasses #-}

-- | Mutable functor-lazy vectors are like mutable boxed vectors, but support mapping a function onto all elements in constant time.  All vector operations (except slicing) are fully supported.  See <http://github.com/mikeizbicki/functor-lazy> for more details.
module Data.Vector.FunctorLazy.Mutable
    (
    -- * Mutable functorlazy vectors
    MVector(..), IOVector, STVector,

    forceElement, Data.Vector.FunctorLazy.Mutable.mapM,

    -- * Accessors

    -- ** Length information
    VGM.length, VGM.null,

--     -- ** Extracting subvectors
--     slice, init, tail, take, drop, splitAt,
--     unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,

--     -- ** Overlapping
--     overlaps,

    -- * Construction

    -- ** Initialisation
    VGM.new, VGM.unsafeNew, VGM.replicate, VGM.replicateM, VGM.clone,

    -- ** Growing
    VGM.grow, VGM.unsafeGrow,

    -- ** Restricting memory usage
    VGM.clear,

    -- * Accessing individual elements
    VGM.read, VGM.write, VGM.swap,
    VGM.unsafeRead, VGM.unsafeWrite, VGM.unsafeSwap,

    -- * Modifying vectors

    -- ** Filling and copying
    VGM.set, VGM.copy, VGM.move, VGM.unsafeCopy, VGM.unsafeMove
    )
    where

import Data.Monoid hiding (Any)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Unboxed.Mutable as VUM
import Data.Vector.Unboxed.Deriving
import Data.Primitive.Array
import Data.Primitive.ByteArray

import Control.Monad.ST
import Control.Monad.Primitive
import Unsafe.Coerce
import System.IO.Unsafe
import GHC.Prim

import Data.Vector.FunctorLazy.Common

-------------------------------------------------------------------------------
-- data types 

data MVector s a = MVector 
    { mvecAny :: !(MutableArray s Any)
    , mvecInt :: !(MutableByteArray s)
    , mlen :: !Int
    , mcontrol :: !LazyController
    }

type IOVector = MVector RealWorld
type STVector s = MVector s

uninitialized :: a
uninitialized = error "Data.Vector.FunctorLazy: uninitialized element"

-------------------------------------------------------------------------------
-- vector instances 

instance VGM.MVector MVector a where
    {-# INLINE basicLength #-}
    basicLength (MVector va bi l c) = l

    {-# INLINE basicUnsafeNew #-}
    basicUnsafeNew len = do
        mvecAny <- newArray len uninitialized
        mvecInt <- newByteArray (len*8)
        setByteArray mvecInt 0 len (0::Int)
        return $ MVector
            { mvecAny = mvecAny
            , mvecInt = mvecInt
            , mlen = len
            , mcontrol = mempty
            }

--     {-# INLINE basicUnsafeRead #-}
    basicUnsafeRead (MVector va vi len (LazyController fl fc)) i = do 
        any <- readArray va i
        count :: Int <- readByteArray vi i
        let val = unsafeCoerce any
        if fc == count
            then return val
            else return $ forceElement (MVector va vi len (LazyController fl fc)) i 
                
--     {-# INLINE basicUnsafeWrite #-}
    basicUnsafeWrite (MVector va vi len (LazyController fl fc)) i a = do 
        writeArray va i (unsafeCoerce a)
        writeByteArray vi i fc

    basicOverlaps = error "Data.Vector.FunctorLazy.MVector: basicOverlaps not supported"

--     basicUnsafeSlice s t v = error "Data.Vector.FunctorLazy.Mvector: basicUnsafeSlice" 
    basicUnsafeSlice s len v = unsafePerformIO $ do
        v' :: MVector RealWorld a <- VGM.basicUnsafeNew len
        do_copy s v'
        return $ unsafeCoerce v' 
        where
            do_copy i dst 
                | i < s+len = do
                    x <- VGM.basicUnsafeRead (unsafeCoerce v :: MVector RealWorld a) (s+i)
                    VGM.basicUnsafeWrite dst i x
                    do_copy (i+1) dst
                | otherwise = return ()

    basicUnsafeGrow v by = do
        v' <- VGM.basicUnsafeNew (n+by)
        VGM.basicUnsafeCopy v' v
        return v'
        where
            n = VGM.basicLength v 

-------------------------------------------------------------------------------
-- functions unique to functor lazy vectors

-- | forces all queued functions to be applied at a given index; this does not actually evaluate the functions, however, only stores the appropriate thunk in the index
{-# NOINLINE forceElement #-}
forceElement :: MVector s a -> Int -> a
forceElement (MVector va vi len (LazyController fl fc)) i = unsafePerformIO $ do
    any <- readArray (unsafeCoerce va) i
    count :: Int <- readByteArray (unsafeCoerce vi) i
    let count' = fc 
    let any' = appList any (take (fc - count) fl) :: a
    writeArray (unsafeCoerce va) i (unsafeCoerce any')
    writeByteArray (unsafeCoerce vi) i (count')
    return any'

-------------------------------------------------------------------------------
-- more efficient functions 

-- | map a function onto all elements in the vector; uses time O(1)
{-# INLINE mapM #-}
mapM :: (Monad m) => (a -> b) -> MVector s a -> m (MVector s b)
mapM f v = return $ v { mcontrol = LazyController
    { funcL = (unsafeCoerce f):(funcL $ mcontrol v)
    , funcC = 1+(funcC $ mcontrol v)
    }}