module Hasql.PreparedStatementRegistry
  ( PreparedStatementRegistry,
    new,
    update,
    LocalKey (..),
  )
where

import ByteString.StrictBuilder qualified as B
import Data.HashTable.IO qualified as A
import Hasql.LibPq14 qualified as Pq
import Hasql.Prelude hiding (lookup)

data PreparedStatementRegistry
  = PreparedStatementRegistry !(A.BasicHashTable LocalKey ByteString) !(IORef Word)

{-# INLINEABLE new #-}
new :: IO PreparedStatementRegistry
new :: IO PreparedStatementRegistry
new =
  HashTable RealWorld LocalKey ByteString
-> IORef Word -> PreparedStatementRegistry
BasicHashTable LocalKey ByteString
-> IORef Word -> PreparedStatementRegistry
PreparedStatementRegistry (HashTable RealWorld LocalKey ByteString
 -> IORef Word -> PreparedStatementRegistry)
-> IO (HashTable RealWorld LocalKey ByteString)
-> IO (IORef Word -> PreparedStatementRegistry)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (HashTable RealWorld LocalKey ByteString)
IO (BasicHashTable LocalKey ByteString)
forall (h :: * -> * -> * -> *) k v.
HashTable h =>
IO (IOHashTable h k v)
A.new IO (IORef Word -> PreparedStatementRegistry)
-> IO (IORef Word) -> IO PreparedStatementRegistry
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Word -> IO (IORef Word)
forall a. a -> IO (IORef a)
newIORef Word
0

{-# INLINEABLE update #-}
update :: LocalKey -> (ByteString -> IO (Bool, a)) -> (ByteString -> IO a) -> PreparedStatementRegistry -> IO a
update :: forall a.
LocalKey
-> (ByteString -> IO (Bool, a))
-> (ByteString -> IO a)
-> PreparedStatementRegistry
-> IO a
update LocalKey
localKey ByteString -> IO (Bool, a)
onNewRemoteKey ByteString -> IO a
onOldRemoteKey (PreparedStatementRegistry BasicHashTable LocalKey ByteString
table IORef Word
counter) =
  IO (Maybe ByteString)
lookup IO (Maybe ByteString) -> (Maybe ByteString -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO a -> (ByteString -> IO a) -> Maybe ByteString -> IO a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO a
new ByteString -> IO a
old
  where
    lookup :: IO (Maybe ByteString)
lookup =
      BasicHashTable LocalKey ByteString
-> LocalKey -> IO (Maybe ByteString)
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> IO (Maybe v)
A.lookup BasicHashTable LocalKey ByteString
table LocalKey
localKey
    new :: IO a
new =
      IORef Word -> IO Word
forall a. IORef a -> IO a
readIORef IORef Word
counter IO Word -> (Word -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word -> IO a
onN
      where
        onN :: Word -> IO a
onN Word
n =
          do
            (Bool
save, a
result) <- ByteString -> IO (Bool, a)
onNewRemoteKey ByteString
remoteKey
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
save (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
              BasicHashTable LocalKey ByteString
-> LocalKey -> ByteString -> IO ()
forall (h :: * -> * -> * -> *) k v.
(HashTable h, Eq k, Hashable k) =>
IOHashTable h k v -> k -> v -> IO ()
A.insert BasicHashTable LocalKey ByteString
table LocalKey
localKey ByteString
remoteKey
              IORef Word -> Word -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Word
counter (Word -> Word
forall a. Enum a => a -> a
succ Word
n)
            return a
result
          where
            remoteKey :: ByteString
remoteKey =
              Builder -> ByteString
B.builderBytes (Builder -> ByteString) -> (Word -> Builder) -> Word -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word -> Builder
forall a. Integral a => a -> Builder
B.asciiIntegral (Word -> ByteString) -> Word -> ByteString
forall a b. (a -> b) -> a -> b
$ Word
n
    old :: ByteString -> IO a
old =
      ByteString -> IO a
onOldRemoteKey

-- |
-- Local statement key.
data LocalKey
  = LocalKey !ByteString ![Pq.Oid]
  deriving (Int -> LocalKey -> ShowS
[LocalKey] -> ShowS
LocalKey -> String
(Int -> LocalKey -> ShowS)
-> (LocalKey -> String) -> ([LocalKey] -> ShowS) -> Show LocalKey
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LocalKey -> ShowS
showsPrec :: Int -> LocalKey -> ShowS
$cshow :: LocalKey -> String
show :: LocalKey -> String
$cshowList :: [LocalKey] -> ShowS
showList :: [LocalKey] -> ShowS
Show, LocalKey -> LocalKey -> Bool
(LocalKey -> LocalKey -> Bool)
-> (LocalKey -> LocalKey -> Bool) -> Eq LocalKey
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LocalKey -> LocalKey -> Bool
== :: LocalKey -> LocalKey -> Bool
$c/= :: LocalKey -> LocalKey -> Bool
/= :: LocalKey -> LocalKey -> Bool
Eq)

instance Hashable LocalKey where
  {-# INLINE hashWithSalt #-}
  hashWithSalt :: Int -> LocalKey -> Int
hashWithSalt Int
salt (LocalKey ByteString
template [Oid]
types) =
    Int -> ByteString -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt ByteString
template