DataRows decoder

src/Database/PostgreSQL/Protocol/DataRows.hs

@@ -2,6 +2,8 @@ module Database.PostgreSQL.Protocol.DataRows
     ( loopExtractDataRows
     , countDataRows
     , flattenDataRows
+    , decodeManyRows
+    , decodeOneRow
     ) where
 
 import Data.Monoid ((<>))
@@ -9,10 +11,15 @@ import Data.Word    (Word8, byteSwap32)
 import Foreign      (peek, peekByteOff, castPtr)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Unsafe as B
+import qualified Data.Vector as V
+import qualified Data.Vector.Mutable as MV
+import Data.Foldable
+import System.IO.Unsafe
 
 import Database.PostgreSQL.Driver.Error
 import Database.PostgreSQL.Protocol.Types
 import Database.PostgreSQL.Protocol.Parsers
+import Database.PostgreSQL.Protocol.Store.Decode
 import Database.PostgreSQL.Protocol.Utils
 
 -- Optimized loop for extracting chunks of DataRows.
@@ -123,6 +130,32 @@ loopExtractDataRows readMoreAction callback = go "" Empty
             w <- byteSwap32 <$> peekByteOff (castPtr ptr) 1
             pure $ Header b $ fromIntegral (w - 4)
 
+----
+-- Decoding 
+-----
+
+--  It is better that Decode throws exception on invalid input
+decodeOneRow :: Decode a -> DataRows -> a
+decodeOneRow dec Empty                         = snd $ runDecode dec ""
+decodeOneRow dec (DataRows (DataChunk _ bs) _) = snd $ runDecode dec bs
+
+decodeManyRows :: Decode a -> DataRows -> V.Vector a
+decodeManyRows dec dr = unsafePerformIO $ do
+    vec <- MV.unsafeNew count
+    go vec 0 dr
+    V.unsafeFreeze vec
+  where
+    go vec startInd Empty = pure ()
+    go vec startInd (DataRows (DataChunk len bs) nextDr) = do
+        let endInd = startInd + fromIntegral len
+        runDecodeIO 
+            (traverse_ (writeDec vec) [startInd .. (endInd  -1)]) 
+            bs
+        go vec endInd nextDr
+
+    count = fromIntegral $ countDataRows dr
+    writeDec vec pos = dec >>= embedIO . MV.unsafeWrite vec pos
+
 ---
 -- Utils
 -- 

src/Database/PostgreSQL/Protocol/Store/Decode.hs

@@ -23,6 +23,18 @@ runDecode (Decode dec) bs =
     let (offset,v ) = decodeExPortionWith dec bs
     in (B.drop offset bs, v)
 
+{-# INLINE runDecodeIO #-}
+runDecodeIO :: Decode a -> B.ByteString -> IO (B.ByteString, a)
+runDecodeIO (Decode dec) bs = do
+    (offset, v) <- decodeIOPortionWith dec bs
+    pure (B.drop offset bs, v)
+
+{-# INLINE embedIO #-}
+embedIO :: IO a -> Decode a
+embedIO action = Decode $ Peek $ \_ ptr -> do
+    v <- action
+    return (ptr, v)
+
 {-# INLINE prim #-}
 prim :: Int -> (Ptr Word8 -> IO a) -> Decode a
 prim len f = Decode $ Peek $ \ps ptr -> do

tests/Driver.hs

@@ -228,13 +228,9 @@ testCorrectDatarows = withConnection $ \c -> do
     case r of
         Left e -> error $ show e
         Right rows -> do
-            let bs = flattenDataRows rows
-            map (BS.pack . show ) [1 .. 1000] @=? go bs
+            map (BS.pack . show ) [1 .. 1000] @=? V.toList (decodeManyRows decodeDataRow rows)
             countDataRows rows @=? 1000
   where
-    go bs | B.null bs = []
-          | otherwise = let (rest, v) = runDecode decodeDataRow bs
-                        in v : go rest
     -- TODO Right parser later
     decodeDataRow :: Decode B.ByteString
     decodeDataRow = do