Merge pull request #1 from robx/binary-cstring

Avoid copies when passing binary parameters

Author: robx

CHANGELOG.md

@@ -1,3 +1,5 @@
+- Avoid copies when passing binary parameters
+
 0.9.4.3
 -------
 

src/Database/PostgreSQL/LibPQ.hsc

@@ -746,6 +746,70 @@ 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 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 unsafeUseParamAsCString) values $ \c_values ->
+            withArray c_values $ \vs ->
+                withArray c_lengths $ \ls ->
+                    withArrayLen formats $ \n fs ->
+                        action (toEnum n) ts vs ls fs
+  where
+    (oids, values, c_lengths, formats) =
+        foldl' accum ([],[],[],[]) $ reverse params
+
+    accum (!a,!b,!c,!d) Nothing = ( invalidOid:a
+                                  , Nothing:b
+                                  , 0:c
+                                  , 0:d
+                                  )
+    accum (!a,!b,!c,!d) (Just (t,v,f)) = ( t:a
+                                         , (Just (v,f)):b
+                                         , (toEnum $ B.length v):c
+                                         , (toEnum $ fromEnum f):d
+                                         )
+
+-- | 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 unsafeUseParamAsCString) values $ \c_values ->
+        withArray c_values $ \vs ->
+            withArray c_lengths $ \ls ->
+                withArrayLen formats $ \n fs ->
+                    action (toEnum n) vs ls fs
+  where
+    (values, c_lengths, formats) = foldl' accum ([],[],[]) $ reverse params
+
+    accum (!a,!b,!c) Nothing       = ( Nothing:a
+                                     , 0:b
+                                     , 0:c
+                                     )
+    accum (!a,!b,!c) (Just (v, f)) = ( (Just (v,f)):a
+                                     , (toEnum $ B.length v):b
+                                     , (toEnum $ fromEnum f):c
+                                     )
+
+
 -- | Submits a command to the server and waits for the result.
 --
 -- Returns a 'Result' or possibly 'Nothing'. A 'Result' will generally
@@ -812,31 +876,12 @@ execParams :: Connection                          -- ^ connection
            -> Format                              -- ^ result format
            -> IO (Maybe Result)                   -- ^ result
 execParams connection statement params rFmt =
-    do let (oids, values, lengths, formats) =
-               foldl' accum ([],[],[],[]) $ reverse params
-           !c_lengths = map toEnum lengths :: [CInt]
-           !n = toEnum $ length params
-           !f = toEnum $ fromEnum rFmt
-       resultFromConn connection $ \c ->
-           B.useAsCString statement $ \s ->
-               withArray oids $ \ts ->
-                   withMany (maybeWith B.useAsCString) values $ \c_values ->
-                       withArray c_values $ \vs ->
-                           withArray c_lengths $ \ls ->
-                               withArray formats $ \fs ->
-                                   c_PQexecParams c s n ts vs ls fs f
-
-    where
-      accum (!a,!b,!c,!d) Nothing = ( invalidOid:a
-                                    , Nothing:b
-                                    , 0:c
-                                    , 0:d
-                                    )
-      accum (!a,!b,!c,!d) (Just (t,v,f)) = ( t:a
-                                           , (Just v):b
-                                           , (B.length v):c
-                                           , (toEnum $ fromEnum f):d
-                                           )
+    resultFromConn connection $ \c ->
+        B.useAsCString statement $ \s ->
+            withParams params $ \n ts vs ls fs ->
+                c_PQexecParams c s n ts vs ls fs f
+  where
+    !f = toEnum $ fromEnum rFmt
 
 
 -- | Submits a request to create a prepared statement with the given
@@ -911,28 +956,13 @@ execPrepared :: Connection                     -- ^ connection
              -> [Maybe (B.ByteString, Format)] -- ^ parameters
              -> Format                         -- ^ result format
              -> IO (Maybe Result)              -- ^ result
-execPrepared connection stmtName mPairs rFmt =
-    do let (values, lengths, formats) = foldl' accum ([],[],[]) $ reverse mPairs
-           !c_lengths = map toEnum lengths :: [CInt]
-           !n = toEnum $ length mPairs
-           !f = toEnum $ fromEnum rFmt
-       resultFromConn connection $ \c ->
-           B.useAsCString stmtName $ \s ->
-               withMany (maybeWith B.useAsCString) values $ \c_values ->
-                   withArray c_values $ \vs ->
-                       withArray c_lengths $ \ls ->
-                           withArray formats $ \fs ->
-                               c_PQexecPrepared c s n vs ls fs f
-
+execPrepared connection stmtName params rFmt =
+    resultFromConn connection $ \c ->
+        B.useAsCString stmtName $ \s ->
+            withParamsPrepared params $ \n vs ls fs ->
+                c_PQexecPrepared c s n vs ls fs f
     where
-      accum (!a,!b,!c) Nothing       = ( Nothing:a
-                                       , 0:b
-                                       , 0:c
-                                       )
-      accum (!a,!b,!c) (Just (v, f)) = ( (Just v):a
-                                       , (B.length v):b
-                                       , (toEnum $ fromEnum f):c
-                                       )
+      !f = toEnum $ fromEnum rFmt
 
 
 -- | Submits a request to obtain information about the specified
@@ -1676,31 +1706,13 @@ sendQueryParams :: Connection
                 -> Format
                 -> IO Bool
 sendQueryParams connection statement params rFmt =
-    do let (oids, values, lengths, formats) =
-               foldl' accum ([],[],[],[]) $ reverse params
-           !c_lengths = map toEnum lengths :: [CInt]
-           !n = toEnum $ length params
-           !f = toEnum $ fromEnum rFmt
-       enumFromConn connection $ \c ->
-           B.useAsCString statement $ \s ->
-               withArray oids $ \ts ->
-                   withMany (maybeWith B.useAsCString) values $ \c_values ->
-                       withArray c_values $ \vs ->
-                           withArray c_lengths $ \ls ->
-                               withArray formats $ \fs ->
-                                   c_PQsendQueryParams c s n ts vs ls fs f
+    enumFromConn connection $ \c ->
+        B.useAsCString statement $ \s ->
+            withParams params $ \n ts vs ls fs ->
+                c_PQsendQueryParams c s n ts vs ls fs f
 
     where
-      accum (!a,!b,!c,!d) Nothing = ( invalidOid:a
-                                    , Nothing:b
-                                    , 0:c
-                                    , 0:d
-                                    )
-      accum (!a,!b,!c,!d) (Just (t,v,f)) = ( t:a
-                                           , (Just v):b
-                                           , (B.length v):c
-                                           , (toEnum $ fromEnum f):d
-                                           )
+      !f = toEnum $ fromEnum rFmt
 
 
 -- | Sends a request to create a prepared statement with the given
@@ -1726,28 +1738,14 @@ sendQueryPrepared :: Connection
                   -> [Maybe (B.ByteString, Format)]
                   -> Format
                   -> IO Bool
-sendQueryPrepared connection stmtName mPairs rFmt =
-    do let (values, lengths, formats) = foldl' accum ([],[],[]) $ reverse mPairs
-           !c_lengths = map toEnum lengths :: [CInt]
-           !n = toEnum $ length mPairs
-           !f = toEnum $ fromEnum rFmt
-       enumFromConn connection $ \c ->
-           B.useAsCString stmtName $ \s ->
-               withMany (maybeWith B.useAsCString) values $ \c_values ->
-                   withArray c_values $ \vs ->
-                       withArray c_lengths $ \ls ->
-                           withArray formats $ \fs ->
-                               c_PQsendQueryPrepared c s n vs ls fs f
+sendQueryPrepared connection stmtName params rFmt =
+    enumFromConn connection $ \c ->
+        B.useAsCString stmtName $ \s ->
+            withParamsPrepared params $ \n vs ls fs ->
+                c_PQsendQueryPrepared c s n vs ls fs f
 
     where
-      accum (!a,!b,!c) Nothing       = ( Nothing:a
-                                       , 0:b
-                                       , 0:c
-                                       )
-      accum (!a,!b,!c) (Just (v, f)) = ( (Just v):a
-                                       , (B.length v):b
-                                       , (toEnum $ fromEnum f):c
-                                       )
+      !f = toEnum $ fromEnum rFmt
 
 
 -- | Submits a request to obtain information about the specified