-- | A parameter's value, which is in memory, lazy, or read from a source as it
-- is sent, and how a value is written.
module Hmssql.Tds.Parameter
  ( ParamValue (..), Source (..), sourceOf, sourceIO, paramValueProblem
  , streamed, writeParamValue, writeParamValueWith
  ) where
import Control.Exception (throwIO)
import Control.Monad (when)
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Lazy as LB
import qualified Data.ByteString.Builder as BB
import Data.Word (Word8)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (plusPtr)
import Foreign.Storable (pokeByteOff)
import Streaming (Stream, Of)
import qualified Streaming.Prelude as S
import Hmssql.Errors (EncodeError (..))
import Hmssql.Tds.TypeInfo
import Hmssql.Tds.PacketWriter

-- | Where a streamed parameter's bytes come from.  It is a function of what
-- consumes the stream, so a resource can be scoped to exactly the send, for
-- example with @withBinaryFile@.  A bare stream couldn't scope one, and a second
-- send would run it from wherever it had got to.
newtype Source a = Source (forall r. (Stream (Of a) IO () -> IO r) -> IO r)
-- | Use an existing stream.  Holding on to the stream yourself keeps whatever
-- it has produced alive, and a stream can be sent only once.
sourceOf :: Stream (Of a) IO () -> Source a
sourceOf :: forall a. Stream (Of a) IO () -> Source a
sourceOf Stream (Of a) IO ()
stream = (forall r. (Stream (Of a) IO () -> IO r) -> IO r) -> Source a
forall a.
(forall r. (Stream (Of a) IO () -> IO r) -> IO r) -> Source a
Source (\Stream (Of a) IO () -> IO r
use -> Stream (Of a) IO () -> IO r
use Stream (Of a) IO ()
stream)
-- | A stream made by this action when sending starts.  For a resource to release
-- after the send, use the 'Source' constructor.
sourceIO :: IO (Stream (Of a) IO ()) -> Source a
sourceIO :: forall a. IO (Stream (Of a) IO ()) -> Source a
sourceIO IO (Stream (Of a) IO ())
open = (forall r. (Stream (Of a) IO () -> IO r) -> IO r) -> Source a
forall a.
(forall r. (Stream (Of a) IO () -> IO r) -> IO r) -> Source a
Source (\Stream (Of a) IO () -> IO r
use -> IO (Stream (Of a) IO ())
open IO (Stream (Of a) IO ()) -> (Stream (Of a) IO () -> IO r) -> IO r
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Stream (Of a) IO () -> IO r
use)

-- | An in-memory value ('Nothing' is NULL), lazy bytes read as they are sent,
-- or a source run as it is sent.  A request holding either of the last two can
-- be sent only once; a query builds a new request each time it is sent.
--
-- The constructors are public so that libraries can inspect and build values.
data ParamValue = ParamBytes !RawBytes | ParamLazy LB.ByteString | ParamStream !(Source B.ByteString)
instance Show ParamValue where
  showsPrec :: Int -> ParamValue -> ShowS
showsPrec Int
n (ParamBytes RawBytes
b) = Bool -> ShowS -> ShowS
showParen (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (String -> ShowS
showString String
"ParamBytes " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RawBytes -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 RawBytes
b)
  showsPrec Int
_ (ParamLazy ByteString
_) = String -> ShowS
showString String
"ParamLazy <bytes>"
  showsPrec Int
_ (ParamStream Source ByteString
_) = String -> ShowS
showString String
"ParamStream <source>"

-- | Whether the value is sent as it is read, so it isn't in memory to check.
streamed :: ParamValue -> Bool
streamed :: ParamValue -> Bool
streamed (ParamBytes RawBytes
_) = Bool
False
streamed ParamValue
_ = Bool
True

-- | Why a value can't be sent in this type, if it can't.  A value sent as it is
-- read has no length ahead of time, so only a PLP type can carry it.
paramValueProblem :: TypeInfo -> ParamValue -> Maybe String
paramValueProblem :: TypeInfo -> ParamValue -> Maybe String
paramValueProblem TypeInfo
ti (ParamBytes RawBytes
value) = TypeInfo -> RawBytes -> Maybe String
valueProblem TypeInfo
ti RawBytes
value
paramValueProblem TypeInfo
ti ParamValue
_ = case TypeInfo -> LengthKind
lengthKind TypeInfo
ti of
  LengthKind
Plp -> TypeInfo -> Maybe String
typeInfoProblem TypeInfo
ti
  LengthKind
_ -> String -> Maybe String
forall a. a -> Maybe a
Just String
"a lazy or streamed parameter requires a PLP (max or xml) type"

-- | Write a parameter's value, which has been checked already
-- ('paramValueProblem').  A lazy or streamed value is written as a PLP value
-- of unknown total length, in chunks of at most 8000 bytes, with the source's
-- empty chunks left out; a UTF-16 source must end on a code unit boundary.
writeParamValue :: PacketWriter -> TypeInfo -> ParamValue -> IO ()
writeParamValue :: PacketWriter -> TypeInfo -> ParamValue -> IO ()
writeParamValue PacketWriter
out = PacketWriter -> Builder -> TypeInfo -> ParamValue -> IO ()
writeParamValueWith PacketWriter
out Builder
forall a. Monoid a => a
mempty

-- | 'writeParamValue', with what comes before the value, such as the parameter's
-- header, in the same Builder.  A BYTE-length scalar is poked straight after the
-- header rather than composed with it; every other value runs the writer once
-- for both.
writeParamValueWith :: PacketWriter -> BB.Builder -> TypeInfo -> ParamValue -> IO ()
writeParamValueWith :: PacketWriter -> Builder -> TypeInfo -> ParamValue -> IO ()
writeParamValueWith PacketWriter
out Builder
prefix TypeInfo
ti ParamValue
value = case ParamValue
value of
    -- Poked, a BYTE-length scalar allocates 147 bytes less per parameter than it
    -- did composed into the header's Builder (a packet of 100 parameters, GHC
    -- 9.12.4, 2026-11-04).  Poking every non-PLP kind the same way was more code
    -- and saved only 98 bytes, so the other kinds keep the Builder.
    ParamBytes RawBytes
bytes
      | LengthKind
ByteLen <- TypeInfo -> LengthKind
lengthKind TypeInfo
ti -> PacketWriter -> Builder -> IO ()
writeBuilder PacketWriter
out Builder
prefix IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> PacketWriter -> RawBytes -> IO ()
writeByteLengthValue PacketWriter
out RawBytes
bytes
      | Bool
otherwise -> PacketWriter -> Builder -> IO ()
writeBuilder PacketWriter
out (Builder
prefix Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> TypeInfo -> RawBytes -> Builder
writeValue TypeInfo
ti RawBytes
bytes)
    ParamLazy ByteString
bytes -> Source ByteString -> IO ()
send (Stream (Of ByteString) IO () -> Source ByteString
forall a. Stream (Of a) IO () -> Source a
sourceOf ([ByteString] -> Stream (Of ByteString) IO ()
forall (m :: * -> *) (f :: * -> *) a.
(Monad m, Foldable f) =>
f a -> Stream (Of a) m ()
S.each (ByteString -> [ByteString]
LB.toChunks ByteString
bytes)))
    ParamStream Source ByteString
source -> Source ByteString -> IO ()
send Source ByteString
source
  where
    -- Checking the odd total before the terminator means a bad value never
    -- completes a message: the failure is handled as a source that threw.
    send :: Source ByteString -> IO ()
send (Source forall r. (Stream (Of ByteString) IO () -> IO r) -> IO r
use) = PacketWriter -> forall a. IO a -> IO a
inSource PacketWriter
out (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      PacketWriter -> Builder -> IO ()
writeBuilder PacketWriter
out (Builder
prefix Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
BB.word64LE Word64
0xFFFFFFFFFFFFFFFE)
      oddBytes <- (Stream (Of ByteString) IO () -> IO Bool) -> IO Bool
forall r. (Stream (Of ByteString) IO () -> IO r) -> IO r
use (Bool -> Stream (Of ByteString) IO () -> IO Bool
go Bool
False)
      when (unicode && oddBytes) (throwIO (EncodeError "UTF-16 stream ended with an odd byte count"))
      writeBuilder out (BB.word32LE 0)
    go :: Bool -> Stream (Of ByteString) IO () -> IO Bool
go !Bool
parity Stream (Of ByteString) IO ()
source = Stream (Of ByteString) IO ()
-> IO (Either () (ByteString, Stream (Of ByteString) IO ()))
forall (m :: * -> *) a r.
Monad m =>
Stream (Of a) m r -> m (Either r (a, Stream (Of a) m r))
S.next Stream (Of ByteString) IO ()
source IO (Either () (ByteString, Stream (Of ByteString) IO ()))
-> (Either () (ByteString, Stream (Of ByteString) IO ())
    -> IO Bool)
-> IO Bool
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Left () -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
parity
      Right (ByteString
bytes, Stream (Of ByteString) IO ()
next) -> do
        ByteString -> IO ()
pieces ByteString
bytes
        Bool -> Stream (Of ByteString) IO () -> IO Bool
go (Bool
parity Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> Bool
forall a. Integral a => a -> Bool
odd (ByteString -> Int
B.length ByteString
bytes)) Stream (Of ByteString) IO ()
next
    -- Input chunks never reach the wire as they are.  An empty one would be the
    -- PLP terminator, and SQL Server's bulk parser refuses chunks over 8000 bytes.
    pieces :: ByteString -> IO ()
pieces ByteString
bytes
      | ByteString -> Bool
B.null ByteString
bytes = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = do
          let (ByteString
chunk, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
8000 ByteString
bytes
          PacketWriter -> Builder -> IO ()
writeBuilder PacketWriter
out (Word32 -> Builder
BB.word32LE (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
chunk)) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
BB.byteString ByteString
chunk)
          ByteString -> IO ()
pieces ByteString
rest
    unicode :: Bool
unicode = case TypeInfo
ti of TINVarChar Word16
_ Collation
_ -> Bool
True; TypeInfo
_ -> Bool
False

-- | Poke a BYTE-length value, checked already, as 'Hmssql.Tds.TypeInfo.writeValue'
-- would write it: its length in one byte, then its bytes.  NULL and an empty
-- value are both a zero length byte.  The payload is an immutable 'B.ByteString',
-- held alive through the copy, and the destination is the writer's private
-- packet storage; neither pointer leaves this function.
writeByteLengthValue :: PacketWriter -> RawBytes -> IO ()
writeByteLengthValue :: PacketWriter -> RawBytes -> IO ()
writeByteLengthValue PacketWriter
writer RawBytes
bytes =
  let payload :: ByteString
payload = ByteString -> (ByteString -> ByteString) -> RawBytes -> ByteString
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ByteString
B.empty ByteString -> ByteString
forall a. a -> a
id RawBytes
bytes
      n :: Int
n = ByteString -> Int
B.length ByteString
payload
   in PacketWriter -> Int -> (Ptr Word8 -> IO ()) -> IO ()
unsafeWriteSized PacketWriter
writer (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> do
        Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
p Int
0 (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Word8)
        case ByteString
payload of
          BI.BS ForeignPtr Word8
src Int
_ -> ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
BI.unsafeWithForeignPtr ForeignPtr Word8
src ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
from -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
1) Ptr Word8
from Int
n