Use Vector for column table for constant element access

README.md

@@ -12,4 +12,4 @@
 - [ ] Implement JSON decoding
 - [ ] Implement `ByteString` decoding (`bytea`)
   - Can we make the fromField instance choose whether it wants binary or text?
-- [ ] Clean up and document column table stuff
+- [x] Clean up and document column table stuff

flake.nix

@@ -21,6 +21,7 @@
             postgresql-libpq
             text
             transformers
+            vector
           ]))
 
           pkgs.haskell-language-server

lib/Database/PostgreSQL/Opium.hs

@@ -10,11 +10,13 @@
 {-# LANGUAGE TypeOperators #-}
 
 module Database.PostgreSQL.Opium
-  ( Error (..)
+  ( ColumnTable
+  , Error (..)
   , ErrorPosition (..)
   , FromField (..)
   , FromRow (..)
   , fetch_
+  , toListColumnTable
   )
   where
 
@@ -24,6 +26,7 @@ import Data.ByteString (ByteString)
 import Data.IORef (IORef, modifyIORef', newIORef, readIORef)
 import Data.Proxy (Proxy (..))
 import Data.Text (Text)
+import Data.Vector (Vector)
 import Database.PostgreSQL.LibPQ
   ( Column
   , Connection
@@ -36,6 +39,7 @@ import GHC.TypeLits (KnownSymbol, symbolVal)
 
 import qualified Data.Text as Text
 import qualified Data.Text.Encoding as Encoding
+import qualified Data.Vector as Vector
 import qualified Database.PostgreSQL.LibPQ as LibPQ
 
 import Database.PostgreSQL.Opium.Error (Error (..), ErrorPosition (..))
@@ -57,20 +61,26 @@ execParams conn query = do
 fetch_ :: forall a. FromRow a => Connection -> ByteString -> IO (Either Error [a])
 fetch_ conn query = runExceptT $ do
   result <- execParams conn query
-  -- TODO: Use unboxed array for columnTable
   columnTable <- ExceptT $ getColumnTable @a Proxy result
   nRows <- liftIO $ LibPQ.ntuples result
   mapM (ExceptT . fromRow result columnTable) [0..nRows - 1]
 
-type ColumnTable = [(Column, Oid)]
+newtype ColumnTable = ColumnTable (Vector (Column, Oid))
+  deriving (Eq, Show)
+
+newColumnTable :: [(Column, Oid)] -> ColumnTable
+newColumnTable = ColumnTable . Vector.fromList
 
 indexColumnTable :: ColumnTable -> Int -> (Column, Oid)
-indexColumnTable = (!!)
+indexColumnTable (ColumnTable v) i = v `Vector.unsafeIndex` i
+
+toListColumnTable :: ColumnTable -> [(Column, Oid)]
+toListColumnTable (ColumnTable v) = Vector.toList v
 
 class FromRow a where
   getColumnTable :: Proxy a -> Result -> IO (Either Error ColumnTable)
   default getColumnTable :: (Generic a, GetColumnTable' (Rep a)) => Proxy a -> Result -> IO (Either Error ColumnTable)
-  getColumnTable Proxy = runExceptT . getColumnTable' @(Rep a) Proxy
+  getColumnTable Proxy = runExceptT . fmap newColumnTable . getColumnTable' @(Rep a) Proxy
 
   fromRow :: Result -> ColumnTable -> Row -> IO (Either Error a)
   default fromRow :: (Generic a, FromRow' (Rep a)) => Result -> ColumnTable -> Row -> IO (Either Error a)
@@ -79,7 +89,7 @@ class FromRow a where
     runExceptT $ to <$> fromRow' (FromRowCtx result columnTable iRef) row
 
 class GetColumnTable' f where
-  getColumnTable' :: Proxy (f p) -> Result -> ExceptT Error IO ColumnTable
+  getColumnTable' :: Proxy (f p) -> Result -> ExceptT Error IO [(Column, Oid)]
 
 instance GetColumnTable' f => GetColumnTable' (M1 D c f) where
   getColumnTable' Proxy = getColumnTable' @f Proxy
@@ -91,7 +101,7 @@ instance (GetColumnTable' f, GetColumnTable' g) => GetColumnTable' (f :*: g) whe
   getColumnTable' Proxy result =
     (++) <$> getColumnTable' @f Proxy result <*> getColumnTable' @g Proxy result
 
-checkColumn :: forall a. FromField a => Proxy a -> String -> Result -> ExceptT Error IO ColumnTable
+checkColumn :: forall a. FromField a => Proxy a -> String -> Result -> ExceptT Error IO [(Column, Oid)]
 checkColumn Proxy nameStr result = do
   column <- ExceptT $ maybe (Left $ ErrorMissingColumn nameText) Right <$> LibPQ.fnumber result name
   oid <- liftIO $ LibPQ.ftype result column

opium.cabal

@@ -79,7 +79,8 @@ library
         containers,
         postgresql-libpq,
         text,
-        transformers
+        transformers,
+        vector
 
     -- Directories containing source files.
     hs-source-dirs:   lib

test/Database/PostgreSQL/OpiumSpec.hs

@@ -5,6 +5,7 @@
 
 module Database.PostgreSQL.OpiumSpec (spec) where
 
+import Data.ByteString (ByteString)
 import Data.Proxy (Proxy (..))
 import Data.Text (Text)
 import Database.PostgreSQL.LibPQ (Connection)
@@ -41,22 +42,35 @@ isLeft :: Either a b -> Bool
 isLeft (Left _) = True
 isLeft _ = False
 
+shouldHaveColumns
+  :: Opium.FromRow a
+  => Proxy a
+  -> Connection
+  -> ByteString
+  -> [LibPQ.Column]
+  -> IO ()
+shouldHaveColumns proxy conn query expectedColumns = do
+  Just result <- LibPQ.execParams conn query [] LibPQ.Text
+  columnTable <- Opium.getColumnTable proxy result
+  let actualColumns = fmap (map fst . Opium.toListColumnTable) columnTable
+  actualColumns `shouldBe` Right expectedColumns
+
 spec :: SpecWith Connection
 spec = do
   describe "getColumnTable" $ do
     it "Gets the column table for a result" $ \conn -> do
-      Just result <- LibPQ.execParams conn "SELECT name, age FROM person" [] LibPQ.Text
-      columnTable <- Opium.getColumnTable @Person Proxy result
-      fmap (map fst) columnTable `shouldBe` Right [0, 1]
+      shouldHaveColumns @Person Proxy conn
+        "SELECT name, age FROM person"
+        [0, 1]
 
     it "Gets the numbers right for funky configurations" $ \conn -> do
-      Just result0 <- LibPQ.execParams conn "SELECT age, name FROM person" [] LibPQ.Text
-      columnTable0 <- Opium.getColumnTable @Person Proxy result0
-      fmap (map fst) columnTable0 `shouldBe` Right [1, 0]
+      shouldHaveColumns @Person Proxy conn
+        "SELECT age, name FROM person"
+        [1, 0]
 
-      Just result1 <- LibPQ.execParams conn "SELECT 0 AS a, 1 AS b, 2 AS c, age, 4 AS d, name FROM person" [] LibPQ.Text
-      columnTable1 <- Opium.getColumnTable @Person Proxy result1
-      fmap (map fst) columnTable1 `shouldBe` Right [5, 3]
+      shouldHaveColumns @Person Proxy conn
+        "SELECT 0 AS a, 1 AS b, 2 AS c, age, 4 AS d, name FROM person"
+        [5, 3]
 
     it "Fails for missing columns" $ \conn -> do
       Just result <- LibPQ.execParams conn "SELECT 0 AS a FROM person" [] LibPQ.Text