XSizeable.hs

{-# LANGUAGE FlexibleInstances, FlexibleContexts, MultiParamTypeClasses, ScopedTypeVariables, OverlappingInstances, DeriveDataTypeable #-}

-- | A resizable array, equivalent to C++'s std::vector. Manipulate them with the Data.Array.MArray class.
--
--   Use with multiple threads has to be synchronized.
module Data.XSizeable (memset, SA, getPtr, Resizable(..)) where

import Foreign.Storable
import Foreign.Ptr
import Foreign.Marshal.Alloc
import Foreign.C.Types
import System.Mem.Weak
import Data.Int
import Data.Word
import Control.Monad
import Data.IORef
import Data.Array.IO
import Data.Array.Base
import Data.Typeable
import System.Win32.Mem

newtype SA i e = SA (IORef (i, i, Int, Ptr e)) deriving (Eq, Typeable)

-- | Returns internal data about the array. But it cannot be counted on to be valid
--   after the next call to 'resize'.
--
--   Also use this to free the array once it is not needed.
getPtr (SA ref) = liftM (\(_, _, sz, p) -> (sz, p)) $ readIORef ref

instance (Storable e) => MArray SA e IO where
	getBounds (SA ref) = liftM (\(l, h, _, _) -> (l, h)) $ readIORef ref
	getNumElements = liftM rangeSize . getBounds
	newArray_ (l, h) = do
		let required = sizeOf (undefined :: e) * rangeSize (l, h)
		p <- mallocBytes required
		memset p 0 (fromIntegral required)
		ref <- newIORef (l, h, required, p)
		return (SA ref)
	unsafeRead (SA ref) i = do
		tup@(_, _, _, p) <- readIORef ref
		peekByteOff p (sizeOf (undefined :: e) * i)
	unsafeWrite (SA ref) i x = do
		tup@(_, _, _, p) <- readIORef ref
		let loc = sizeOf (undefined :: e) * i
		pokeByteOff p loc x

-- A resize that only works on single-dimensional arrays.
resizeBuffer :: forall i e. (Storable e) => (Ix i) => (i, i) -> SA i e -> IO ()
resizeBuffer (l', h') sa@(SA ref) = do
	(l, h, sz, p) <- readIORef ref
	let wasRequired = sizeOf (undefined :: e) * rangeSize (l, h)
	let nowRequired = sizeOf (undefined :: e) * rangeSize (l', h')
	-- I want to maintain about 50% residency in the array.
	(sz, p) <- if 4 * rangeSize (l', h') <= sz then
			-- If the array gets small enough, resize it so it is 50% used.
			let newSz = 2 * rangeSize (l', h') in
			liftM ((,) newSz) $ reallocBytes p newSz
		else if rangeSize (l', h') <= sz then
			return (sz, p)
		else
			-- If the array exceeds its buffer, at least double it in size.
			let newSz = nowRequired `max` (2 * wasRequired) in
			liftM ((,) newSz) $ reallocBytes p newSz
	when (nowRequired > wasRequired) $ memset (plusPtr p wasRequired) 0 (fromIntegral $ nowRequired - wasRequired)
	writeIORef ref (l', h', sz, p)

__resize (l', h') sa = do
	(l, _) <- getBounds sa
	resizeBuffer (l', h') sa
	when (l /= l') $ __resize1 (l', h') sa

__resize1 oldBnd sa = mapM_ (\i -> unsafeRead sa (index oldBnd i) >>= writeArray sa i) (range oldBnd)

-- | Mutable arrays that can in addition be resized.
class (Ix i, MArray a e m) => Resizable a i e m where
	resize :: (i, i) -> a i e -> m ()

instance (Storable e) => Resizable SA Int e IO where
	resize = __resize

instance (Storable e) => Resizable SA Int8 e IO where
	resize = __resize

instance (Storable e) => Resizable SA Int16 e IO where
	resize = __resize

instance (Storable e) => Resizable SA Int32 e IO where
	resize = __resize

instance (Storable e) => Resizable SA Word8 e IO where
	resize = __resize

instance (Storable e) => Resizable SA Word16 e IO where
	resize = __resize

instance (Storable e) => Resizable SA Word32 e IO where
	resize = __resize

instance (Storable e, Ix i) => Resizable SA i e IO where
	resize bnd sa = do
		resizeBuffer bnd sa
		__resize1 bnd sa