Pass binary parameter data without copying

Author: robx

src/Database/PostgreSQL/LibPQ.hsc

@@ -745,13 +745,26 @@ newtype Oid = Oid CUInt deriving (Eq, Ord, Read, Show, Storable, Typeable)
 invalidOid :: Oid
 invalidOid = Oid (#const InvalidOid)
 
+-- | Prepare the given parameter bytestring for passing on to libpq,
+-- without copying for binary parameters.
+--
+-- This is safe to use to pass parameters to libpq considering:
+-- * libpq treats the parameter data as read-only
+-- * 'ByteString' uses pinned memory
+-- * the reference to the 'CString' doesn't escape
+unsafeUseParamAsCString :: (B.ByteString, Format) -> (CString -> IO a) -> IO a
+unsafeUseParamAsCString (bs, format) =
+    case format of
+        Binary -> B.unsafeUseAsCString bs
+        Text   -> B.useAsCString bs
+
 -- | Convert a list of parameters to the format expected by libpq FFI calls.
 withParams :: [Maybe (Oid, B.ByteString, Format)]
            -> (CInt -> Ptr Oid -> Ptr CString -> Ptr CInt -> Ptr CInt -> IO a)
            -> IO a
 withParams params action =
     withArray oids $ \ts ->
-        withMany (maybeWith B.useAsCString) values $ \c_values ->
+        withMany (maybeWith unsafeUseParamAsCString) values $ \c_values ->
             withArray c_values $ \vs ->
                 withArray c_lengths $ \ls ->
                     withArray formats $ \fs ->
@@ -770,7 +783,7 @@ withParams params action =
     accum (Just (t,v,f)) ~(AccumParams xs ys zs ws)  =
         let z = intToCInt (B.length v)
             w = formatToCInt f
-        in z `seq` w `seq` AccumParams (t : xs) (Just v : ys) (z : zs) (w : ws)
+        in z `seq` w `seq` AccumParams (t : xs) (Just (v, f) : ys) (z : zs) (w : ws)
 
 formatToCInt :: Format -> CInt
 formatToCInt Text   = 0
@@ -779,15 +792,15 @@ formatToCInt Binary = 1
 intToCInt :: Int -> CInt
 intToCInt = toEnum
 
-data AccumParams = AccumParams ![Oid] ![Maybe B.ByteString] ![CInt] ![CInt]
+data AccumParams = AccumParams ![Oid] ![Maybe (B.ByteString, Format)] ![CInt] ![CInt]
 
 -- | Convert a list of parameters to the format expected by libpq FFI calls,
 -- prepared statement variant.
 withParamsPrepared :: [Maybe (B.ByteString, Format)]
                    -> (CInt -> Ptr CString -> Ptr CInt -> Ptr CInt -> IO a)
                    -> IO a
 withParamsPrepared params action =
-    withMany (maybeWith B.useAsCString) values $ \c_values ->
+    withMany (maybeWith unsafeUseParamAsCString) values $ \c_values ->
         withArray c_values $ \vs ->
             withArray c_lengths $ \ls ->
                 withArray formats $ \fs ->
@@ -806,9 +819,9 @@ withParamsPrepared params action =
     accum (Just (v, f)) ~(AccumPrepParams xs ys zs) =
         let y = intToCInt (B.length v)
             z = formatToCInt f
-        in y `seq` z `seq` AccumPrepParams (Just v : xs) (y : ys) (z : zs)
+        in y `seq` z `seq` AccumPrepParams (Just (v, f) : xs) (y : ys) (z : zs)
 
-data AccumPrepParams = AccumPrepParams ![Maybe B.ByteString] ![CInt] ![CInt]
+data AccumPrepParams = AccumPrepParams ![Maybe (B.ByteString, Format)] ![CInt] ![CInt]
 
 -- | Submits a command to the server and waits for the result.
 --