Implement copy from/to pointer via primops

Author: Bodigrim

src/Data/Text/Array.hs

@@ -31,6 +31,8 @@ module Data.Text.Array
     , shrinkM
     , copyM
     , copyI
+    , copyFromPointer
+    , copyToPointer
     , empty
     , equal
     , compare
@@ -250,6 +252,40 @@ copyI count@(I# count#) (MutableByteArray dst#) dstOff@(I# dstOff#) (ByteArray s
       s2# -> (# s2#, () #)
 {-# INLINE copyI #-}
 
+-- | Copy from pointer.
+copyFromPointer
+  :: MArray s               -- ^ Destination
+  -> Int                    -- ^ Destination offset
+  -> Ptr Word8              -- ^ Source
+  -> Int                    -- ^ Count
+  -> ST s ()
+copyFromPointer (MutableByteArray dst#) dstOff@(I# dstOff#) (Ptr src#) count@(I# count#)
+#if defined(ASSERTS)
+  | count < 0 = error $
+    "copyFromPointer: count must be >= 0, but got " ++ show count
+#endif
+  | otherwise = ST $ \s1# ->
+    case copyAddrToByteArray# src# dst# dstOff# count# s1# of
+      s2# -> (# s2#, () #)
+{-# INLINE copyFromPointer #-}
+
+-- | Copy to pointer.
+copyToPointer
+  :: Array                  -- ^ Source
+  -> Int                    -- ^ Source offset
+  -> Ptr Word8              -- ^ Destination
+  -> Int                    -- ^ Count
+  -> ST s ()
+copyToPointer (ByteArray src#) srcOff@(I# srcOff#) (Ptr dst#) count@(I# count#)
+#if defined(ASSERTS)
+  | count < 0 = error $
+    "copyToPointer: count must be >= 0, but got " ++ show count
+#endif
+  | otherwise = ST $ \s1# ->
+    case copyByteArrayToAddr# src# srcOff# dst# count# s1# of
+      s2# -> (# s2#, () #)
+{-# INLINE copyToPointer #-}
+
 -- | Compare portions of two arrays for equality.  No bounds checking
 -- is performed.
 equal :: Array -> Int -> Array -> Int -> Int -> Bool

src/Data/Text/Foreign.hs

@@ -34,7 +34,7 @@ module Data.Text.Foreign
 #if defined(ASSERTS)
 import Control.Exception (assert)
 #endif
-import Control.Monad.ST.Unsafe (unsafeIOToST)
+import Control.Monad.ST.Unsafe (unsafeSTToIO)
 import Data.ByteString.Unsafe (unsafePackCStringLen, unsafeUseAsCStringLen)
 import Data.Text.Encoding (decodeUtf8, encodeUtf8)
 import Data.Text.Internal (Text(..), empty)
@@ -44,8 +44,7 @@ import Data.Word (Word8)
 import Foreign.C.String (CStringLen)
 import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtrArray)
 import Foreign.Marshal.Alloc (allocaBytes)
-import Foreign.Ptr (Ptr, castPtr, plusPtr)
-import Foreign.Storable (peek, poke)
+import Foreign.Ptr (Ptr, castPtr)
 import qualified Data.Text.Array as A
 
 -- $interop
@@ -68,20 +67,11 @@ newtype I8 = I8 Int
 fromPtr :: Ptr Word8           -- ^ source array
         -> I8                  -- ^ length of source array (in 'Word8' units)
         -> IO Text
-fromPtr _   (I8 0)   = return empty
-fromPtr ptr (I8 len) =
-#if defined(ASSERTS)
-    assert (len > 0) $
-#endif
-    return $! Text arr 0 len
-  where
-    arr = A.run (A.new len >>= copy)
-    copy marr = loop ptr 0
-      where
-        loop !p !i | i == len = return marr
-                   | otherwise = do
-          A.unsafeWrite marr i =<< unsafeIOToST (peek p)
-          loop (p `plusPtr` 1) (i + 1)
+fromPtr ptr (I8 len) = unsafeSTToIO $ do
+  dst <- A.new len
+  A.copyFromPointer dst 0 ptr len
+  arr <- A.unsafeFreeze dst
+  return $! Text arr 0 len
 
 -- $lowlevel
 --
@@ -130,13 +120,7 @@ splitAtWord8 (I8 n) t@(Text arr off len)
 -- | /O(n)/ Copy a 'Text' to an array.  The array is assumed to be big
 -- enough to hold the contents of the entire 'Text'.
 unsafeCopyToPtr :: Text -> Ptr Word8 -> IO ()
-unsafeCopyToPtr (Text arr off len) ptr = loop ptr off
-  where
-    end = off + len
-    loop !p !i | i == end  = return ()
-               | otherwise = do
-      poke p (A.unsafeIndex arr i)
-      loop (p `plusPtr` 1) (i + 1)
+unsafeCopyToPtr (Text arr off len) ptr = unsafeSTToIO $ A.copyToPointer arr off ptr len
 
 -- | /O(n)/ Perform an action on a temporary, mutable copy of a
 -- 'Text'.  The copy is freed as soon as the action returns.