Ticket #2589: Vector.hs

{-# LANGUAGE CPP, ForeignFunctionInterface #-}
{-# OPTIONS_GHC -XUnliftedFFITypes -XMagicHash -XUnboxedTuples -XDeriveDataTypeable #-}
-- |
-- Module      : Data.Vector.Shore.Interna
-- License     : BSD-style
-- Maintainer  : Gregory Wright
-- Stability   : experimental
-- Portability : portable
-- 
-- A module containing a low level interface to "Shore Vectors".
-- "Shore vectors" are unboxed, storable vectors whose contents are managed
-- as foreign data.  (Whence the name "shore vector" since they live
-- on the C-side.)
--
-- ATM, this is just a transcription of the code from ByteString, generalized
-- to arbitrary storables. 
--
module Data.Vector.Shore.Internal (

        -- * The @ByteString@ type and representation
        Vector(..),             -- instances: Eq, Ord, Show, Read, Data, Typeable

        -- * Low level introduction and elimination
        create,                 -- :: Int -> (Ptr Word8 -> IO ()) -> IO ByteString
        unsafeCreate,           -- :: Int -> (Ptr Word8 -> IO ()) ->  ByteString
        mallocVectorPayload,    -- :: Int -> IO (ForeignPtr a)

        -- * Conversion to and from ForeignPtrs
        fromForeignPtr,         -- :: ForeignPtr Word8 -> Int -> Int -> ByteString
        toForeignPtr,           -- :: ByteString -> (ForeignPtr Word8, Int, Int)

        -- * Utilities
        inlinePerformIO,        -- :: IO a -> a
        nullForeignPtr,         -- :: ForeignPtr Word8
  ) where

import Foreign.ForeignPtr       (ForeignPtr, withForeignPtr)
import Foreign.Marshal.Array    (advancePtr)
import Foreign.Ptr              (Ptr, FunPtr, plusPtr)
import Foreign.Storable         (Storable(..))
import Foreign.C.Types          (CInt, CSize, CULong)
import Foreign.C.String         (CString)

#ifndef __NHC__
import Control.Exception        (assert)
#endif

import Data.Char                (ord)
import Data.Word                (Word8)

#if defined(__GLASGOW_HASKELL__)
import Data.Generics            (Data(..), Typeable(..))
import GHC.Ptr                  (Ptr(..))
import GHC.Base                 (realWorld#,unsafeChr)
import GHC.IOBase               (IO(IO), RawBuffer)
#if __GLASGOW_HASKELL__ >= 608
import GHC.IOBase               (unsafeDupablePerformIO)
#else
import GHC.IOBase               (unsafePerformIO)
#endif
#else
import Data.Char                (chr)
import System.IO.Unsafe         (unsafePerformIO)
#endif

#if __GLASGOW_HASKELL__ >= 605 && !defined(SLOW_FOREIGN_PTR)
import GHC.ForeignPtr           (mallocPlainForeignPtrBytes)
#else
import Foreign.ForeignPtr       (mallocForeignPtrBytes)
#endif

#if __GLASGOW_HASKELL__>=605
import GHC.ForeignPtr           (ForeignPtr(ForeignPtr))
import GHC.Base                 (nullAddr#)
#else
import Foreign.Ptr              (nullPtr)
#endif

#if __HUGS__
import Hugs.ForeignPtr          (newForeignPtr_)
#elif __GLASGOW_HASKELL__<=604
import Foreign.ForeignPtr       (newForeignPtr_)
#endif

-- CFILES stuff is Hugs only
{-# CFILES cbits/fpstring.c #-}

-- An alternative to Control.Exception (assert) for nhc98
#ifdef __NHC__
#define assert  assertS "__FILE__ : __LINE__"
assertS :: String -> Bool -> a -> a
assertS _ True  = id
assertS s False = error ("assertion failed at "++s)
#endif

-- -----------------------------------------------------------------------------
--
-- Useful macros, until we have bang patterns
--

#define STRICT1(f) f a | a `seq` False = undefined
#define STRICT2(f) f a b | a `seq` b `seq` False = undefined
#define STRICT3(f) f a b c | a `seq` b `seq` c `seq` False = undefined
#define STRICT4(f) f a b c d | a `seq` b `seq` c `seq` d `seq` False = undefined
#define STRICT5(f) f a b c d e | a `seq` b `seq` c `seq` d `seq` e `seq` False = undefined

-- -----------------------------------------------------------------------------

-- | A space-efficient representation of a Storable vector.
--
-- Instances of Eq, Ord, Read, Show, Data, Typeable
--
data Vector a = V {-# UNPACK #-} !(ForeignPtr a) -- payload
                  {-# UNPACK #-} !Int            -- offset
                  {-# UNPACK #-} !Int            -- length

#if defined(__GLASGOW_HASKELL__)
    deriving (Data, Typeable)
#endif

-- | /O(n)/ Converts a 'Vector a' to a '[b]', using a conversion function.
unpackWith :: Storable a => (a -> b) -> Vector a -> [b]
unpackWith _ (V _  _ 0) = []
unpackWith k (V ps s l) = inlinePerformIO $ withForeignPtr ps $ \p ->
        go (p `plusPtr` s) (l - 1) []
    where
        STRICT3(go)
        go p 0 acc = peek p          >>= \e -> return (k e : acc)
        go p n acc = peekElemOff p n >>= \e -> go p (n-1) (k e : acc)
{-# INLINE unpackWith #-}

-- | /O(n)/ Convert a '[a]' into a 'Vector b' using some
-- conversion function
packWith :: Storable b => (a -> b) -> [a] -> Vector b
packWith k str = unsafeCreate (length str) $ \p -> go p str
    where
        STRICT2(go)
        go _ []     = return ()
        go p (x:xs) = poke p (k x) >> go (p `advancePtr` 1) xs -- less space than pokeElemOff
{-# INLINE packWith #-}

------------------------------------------------------------------------

-- | The 0 pointer. Used to indicate the empty Bytestring.
nullForeignPtr :: ForeignPtr a
#if __GLASGOW_HASKELL__>=605
nullForeignPtr = ForeignPtr nullAddr# undefined --TODO: should ForeignPtrContents be strict?
#else
nullForeignPtr = unsafePerformIO $ newForeignPtr_ nullPtr
{-# NOINLINE nullForeignPtr #-}
#endif

-- ---------------------------------------------------------------------
-- Low level constructors

-- | /O(1)/ Build a ByteString from a ForeignPtr
fromForeignPtr :: ForeignPtr a
               -> Int -- ^ Offset
               -> Int -- ^ Length
               -> Vector a
fromForeignPtr fp s l = V fp s l
{-# INLINE fromForeignPtr #-}

-- | /O(1)/ Deconstruct a ForeignPtr from a ByteString
toForeignPtr :: Vector a -> (ForeignPtr a, Int, Int) -- ^ (ptr, offset, length)
toForeignPtr (V ps s l) = (ps, s, l)
{-# INLINE toForeignPtr #-}

-- | A way of creating ByteStrings outside the IO monad. The @Int@
-- argument gives the final size of the ByteString. Unlike
-- 'createAndTrim' the ByteString is not reallocated if the final size
-- is less than the estimated size.
unsafeCreate :: Storable a => Int -> (Ptr a -> IO ()) -> Vector a
unsafeCreate l f = unsafeDupablePerformIO (create l f)
{-# INLINE unsafeCreate #-}

#if !defined(__GLASGOW_HASKELL__) || __GLASGOW_HASKELL__ < 608
-- for Hugs   
unsafeDupablePerformIO :: IO a -> a
unsafeDupablePerformIO = unsafePerformIO
#endif

-- | Create ByteString of size @l@ and use action @f@ to fill it's contents.
create :: Storable a => Int -> (Ptr a -> IO ()) -> IO (Vector a)
create l f = do
    fp <- mallocVectorPayload l
    withForeignPtr fp $ \p -> f p
    return $! V fp 0 l
{-# INLINE create #-}

-- | Wrapper of mallocForeignPtrBytes with faster implementation
-- for GHC 6.5 builds newer than 06/06/06
mallocVectorPayload :: Storable a => Int -> IO (ForeignPtr a)
mallocVectorPayload l = doMalloc undefined l
    where
        doMalloc :: (Storable b) => b -> Int -> IO (ForeignPtr b)
        doMalloc dummy len =
#if __GLASGOW_HASKELL__ >= 605 && !defined(SLOW_FOREIGN_PTR)
            mallocPlainForeignPtrBytes ((sizeOf dummy) * len)
#else
            mallocForeignPtrBytes ((sizeOf dummy) * len)
#endif
{-# INLINE mallocVectorPayload #-}

------------------------------------------------------------------------

-- | Just like unsafePerformIO, but we inline it. Big performance gains as
-- it exposes lots of things to further inlining. /Very unsafe/. In
-- particular, you should do no memory allocation inside an
-- 'inlinePerformIO' block. On Hugs this is just @unsafePerformIO@.
--
{-# INLINE inlinePerformIO #-}
inlinePerformIO :: IO a -> a
#if defined(__GLASGOW_HASKELL__)
inlinePerformIO (IO m) = case m realWorld# of (# _, r #) -> r
#else
inlinePerformIO = unsafePerformIO
#endif