Add setByteArray# to Prim
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 26 Jan 2012 21:36:03 +0000 (21:36 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 26 Jan 2012 21:36:03 +0000 (21:36 +0000)
Data/Primitive/Types.hs

index 7ced896..07b267b 100644 (file)
@@ -1,4 +1,5 @@
-{-# LANGUAGE UnboxedTuples, MagicHash, DeriveDataTypeable #-}
+{-# LANGUAGE UnboxedTuples, MagicHash, DeriveDataTypeable,
+             ForeignFunctionInterface, UnliftedFFITypes #-}
 
 -- |
 -- Module      : Data.Primitive.Types
@@ -21,6 +22,7 @@ import Control.Monad.Primitive
 import Data.Primitive.MachDeps
 
 import GHC.Base (
+    unsafeCoerce#,
     Int(..), Char(..),
   )
 import GHC.Float (
@@ -79,6 +81,10 @@ class Prim a where
   -- @a@ rather than in bytes.
   writeByteArray# :: MutableByteArray# s -> Int# -> a -> State# s -> State# s
 
+  -- | Fill a slice of the mutable array with a value. The offset and length
+  -- of the chunk is in elements of type @a@ rather than in bytes.
+  setByteArray# :: MutableByteArray# s -> Int# -> Int# -> a -> State# s -> State# s
+
   -- | Read a value from a memory position given by an address and an offset.
   -- The memory block the address refers to must be immutable. The offset is in
   -- elements of type @a@ rather than in bytes.
@@ -92,7 +98,7 @@ class Prim a where
   -- The offset is in elements of type @a@ rather than in bytes.
   writeOffAddr# :: Addr# -> Int# -> a -> State# s -> State# s
 
-#define derivePrim(ty, ctr, sz, align, idx_arr, rd_arr, wr_arr, idx_addr, rd_addr, wr_addr) \
+#define derivePrim(ty, ctr, sz, align, idx_arr, rd_arr, wr_arr, set_arr, idx_addr, rd_addr, wr_addr) \
 instance Prim ty where {                                        \
   sizeOf# _ = unI# sz                                           \
 ; alignment# _ = unI# align                                     \
@@ -100,6 +106,9 @@ instance Prim ty where {                                        \
 ; readByteArray#  arr# i# s# = case rd_arr arr# i# s# of        \
                         { (# s1#, x# #) -> (# s1#, ctr x# #) }  \
 ; writeByteArray# arr# i# (ctr x#) s# = wr_arr arr# i# x# s#    \
+; setByteArray# arr# i# n# (ctr x#) s#                          \
+    = case internal (set_arr arr# i# n# x#) (unsafeCoerce# s#) of \
+        { (# s1#, _ #) -> unsafeCoerce# s1# }                   \
                                                                 \
 ; indexOffAddr# addr# i# = ctr (idx_addr addr# i#)              \
 ; readOffAddr#  addr# i# s# = case rd_addr addr# i# s# of       \
@@ -110,6 +119,7 @@ instance Prim ty where {                                        \
 ; {-# INLINE indexByteArray# #-}                                \
 ; {-# INLINE readByteArray# #-}                                 \
 ; {-# INLINE writeByteArray# #-}                                \
+; {-# INLINE setByteArray# #-}                                  \
 ; {-# INLINE indexOffAddr# #-}                                  \
 ; {-# INLINE readOffAddr# #-}                                   \
 ; {-# INLINE writeOffAddr# #-}                                  \
@@ -119,45 +129,75 @@ unI# :: Int -> Int#
 unI# (I# n#) = n#
 
 derivePrim(Word, W#, sIZEOF_WORD, aLIGNMENT_WORD,
-           indexWordArray#, readWordArray#, writeWordArray#,
+           indexWordArray#, readWordArray#, writeWordArray#, setWordArray#,
            indexWordOffAddr#, readWordOffAddr#, writeWordOffAddr#)
 derivePrim(Word8, W8#, sIZEOF_WORD8, aLIGNMENT_WORD8,
-           indexWord8Array#, readWord8Array#, writeWord8Array#,
+           indexWord8Array#, readWord8Array#, writeWord8Array#, setWord8Array#,
            indexWord8OffAddr#, readWord8OffAddr#, writeWord8OffAddr#)
 derivePrim(Word16, W16#, sIZEOF_WORD16, aLIGNMENT_WORD16,
-           indexWord16Array#, readWord16Array#, writeWord16Array#,
+           indexWord16Array#, readWord16Array#, writeWord16Array#, setWord16Array#,
            indexWord16OffAddr#, readWord16OffAddr#, writeWord16OffAddr#)
 derivePrim(Word32, W32#, sIZEOF_WORD32, aLIGNMENT_WORD32,
-           indexWord32Array#, readWord32Array#, writeWord32Array#,
+           indexWord32Array#, readWord32Array#, writeWord32Array#, setWord32Array#,
            indexWord32OffAddr#, readWord32OffAddr#, writeWord32OffAddr#)
 derivePrim(Word64, W64#, sIZEOF_WORD64, aLIGNMENT_WORD64,
-           indexWord64Array#, readWord64Array#, writeWord64Array#,
+           indexWord64Array#, readWord64Array#, writeWord64Array#, setWord64Array#,
            indexWord64OffAddr#, readWord64OffAddr#, writeWord64OffAddr#)
 derivePrim(Int, I#, sIZEOF_INT, aLIGNMENT_INT,
-           indexIntArray#, readIntArray#, writeIntArray#,
+           indexIntArray#, readIntArray#, writeIntArray#, setIntArray#,
            indexIntOffAddr#, readIntOffAddr#, writeIntOffAddr#)
 derivePrim(Int8, I8#, sIZEOF_INT8, aLIGNMENT_INT8,
-           indexInt8Array#, readInt8Array#, writeInt8Array#,
+           indexInt8Array#, readInt8Array#, writeInt8Array#, setInt8Array#,
            indexInt8OffAddr#, readInt8OffAddr#, writeInt8OffAddr#)
 derivePrim(Int16, I16#, sIZEOF_INT16, aLIGNMENT_INT16,
-           indexInt16Array#, readInt16Array#, writeInt16Array#,
+           indexInt16Array#, readInt16Array#, writeInt16Array#, setInt16Array#,
            indexInt16OffAddr#, readInt16OffAddr#, writeInt16OffAddr#)
 derivePrim(Int32, I32#, sIZEOF_INT32, aLIGNMENT_INT32,
-           indexInt32Array#, readInt32Array#, writeInt32Array#,
+           indexInt32Array#, readInt32Array#, writeInt32Array#, setInt32Array#,
            indexInt32OffAddr#, readInt32OffAddr#, writeInt32OffAddr#)
 derivePrim(Int64, I64#, sIZEOF_INT64, aLIGNMENT_INT64,
-           indexInt64Array#, readInt64Array#, writeInt64Array#,
+           indexInt64Array#, readInt64Array#, writeInt64Array#, setInt64Array#,
            indexInt64OffAddr#, readInt64OffAddr#, writeInt64OffAddr#)
 derivePrim(Float, F#, sIZEOF_FLOAT, aLIGNMENT_FLOAT,
-           indexFloatArray#, readFloatArray#, writeFloatArray#,
+           indexFloatArray#, readFloatArray#, writeFloatArray#, setFloatArray#,
            indexFloatOffAddr#, readFloatOffAddr#, writeFloatOffAddr#)
 derivePrim(Double, D#, sIZEOF_DOUBLE, aLIGNMENT_DOUBLE,
-           indexDoubleArray#, readDoubleArray#, writeDoubleArray#,
+           indexDoubleArray#, readDoubleArray#, writeDoubleArray#, setDoubleArray#,
            indexDoubleOffAddr#, readDoubleOffAddr#, writeDoubleOffAddr#)
 derivePrim(Char, C#, sIZEOF_CHAR, aLIGNMENT_CHAR,
-           indexWideCharArray#, readWideCharArray#, writeWideCharArray#,
+           indexWideCharArray#, readWideCharArray#, writeWideCharArray#, setWideCharArray#,
            indexWideCharOffAddr#, readWideCharOffAddr#, writeWideCharOffAddr#)
 derivePrim(Addr, Addr, sIZEOF_PTR, aLIGNMENT_PTR,
-           indexAddrArray#, readAddrArray#, writeAddrArray#,
+           indexAddrArray#, readAddrArray#, writeAddrArray#, setAddrArray#,
            indexAddrOffAddr#, readAddrOffAddr#, writeAddrOffAddr#)
 
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word8"
+  setWord8Array# :: MutableByteArray# s -> Int# -> Int# -> Word# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word16"
+  setWord16Array# :: MutableByteArray# s -> Int# -> Int# -> Word# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word32"
+  setWord32Array# :: MutableByteArray# s -> Int# -> Int# -> Word# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word64"
+  setWord64Array# :: MutableByteArray# s -> Int# -> Int# -> Word64# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word"
+  setWordArray# :: MutableByteArray# s -> Int# -> Int# -> Word# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word8"
+  setInt8Array# :: MutableByteArray# s -> Int# -> Int# -> Int# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word16"
+  setInt16Array# :: MutableByteArray# s -> Int# -> Int# -> Int# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word32"
+  setInt32Array# :: MutableByteArray# s -> Int# -> Int# -> Int# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word64"
+  setInt64Array# :: MutableByteArray# s -> Int# -> Int# -> Int64# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Word"
+  setIntArray# :: MutableByteArray# s -> Int# -> Int# -> Int# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Ptr"
+  setAddrArray# :: MutableByteArray# s -> Int# -> Int# -> Addr# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Float"
+  setFloatArray# :: MutableByteArray# s -> Int# -> Int# -> Float# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Double"
+  setDoubleArray# :: MutableByteArray# s -> Int# -> Int# -> Double# -> IO ()
+foreign import ccall unsafe "primitive-memops.h hsprimitive_memset_Char"
+  setWideCharArray# :: MutableByteArray# s -> Int# -> Int# -> Char# -> IO ()
+
+