diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,120 @@ # Changelog for wireform-core -## 0.1.0.0 -- 2026-05-16 +## 0.2.0.0 -- 2026-07-16 +* `Wireform.Base64` -- RFC 4648 §4 base64 encode / decode. SSSE3 + (via simde) inner loop for the encoder (12 input bytes -> 16 + output chars per iteration) plus an SSE2 pre-scan for the + decoder that rejects any 16-byte window containing a high-bit + byte before scalar sextet extraction. Replaces ad-hoc + `base64-bytestring` usage in downstream packages — the + WebSocket handshake (`wireform-websocket`) is the first + consumer. + +* `Wireform.FFI.fastRandomWord64` -- thread-local xoshiro256++ + PRNG implemented in `cbits/fast_rng.c`. Per-OS-thread 256-bit + state stored in `__thread` storage, seeded on first use from + `getrandom(2)` (`arc4random_buf` on BSDs, `/dev/urandom` + elsewhere). Each call is a single FFI trip + a handful of + register-only XOR / rotate ops — typically ~1 ns including the + FFI boundary, versus ~50 ns for the global `splitmix` `MVar` + generator. Caveat: because Haskell threads are multiplexed + across OS threads, the per-Haskell-thread sequence is not + reproducible; use `System.Random.Stateful` for that. The + intended uses are non-deterministic-randomness needs on hot + paths — WebSocket frame masks, retry jitter, etc. + +* `Wireform.Ring` -- gave `MagicRing` a phantom type parameter `s` + modelled after `Control.Monad.ST.ST`. `withMagicRing` is now + rank-2 (`Int -> (forall s. MagicRing s -> IO a) -> IO a`), which + seals `s` inside the ring's scope. New `RingSlice s` type (a + pointer + length tagged with the ring's `s`) replaces ad-hoc + `ByteString` slicing for callers that want type-system-enforced + safety against dangling references after a refill. `copyRingSlice + :: RingSlice s -> IO ByteString` is the explicit escape hatch. + +* `Wireform.Transport` -- replaced the typed `transportRing :: + MagicRing` field with three raw fields (`transportRingBaseField`, + `transportRingSizeField`, `transportRingMaskField`). This keeps + `Transport` un-parameterised so the new `s` does not have to + cascade through every downstream `Transport`-using package. The + existing `transportRing :: Transport -> MagicRing s` getter + remains for backwards compatibility but is polymorphic in `s`, i.e. + un-scoped from a safety standpoint. Transport constructors must + populate the three raw fields explicitly. + +* `Wireform.Parser` -- `takeBs` and `takeBsCopy` now drain any byte + count, including counts larger than the magic ring's capacity. + Reads that fit in the ring keep their existing fast path + (zero-copy slice or single memcpy). Reads that exceed the ring + allocate a fresh pinned 'ByteString' and chunk-copy through the + ring, calling 'modeCheckpoint' between chunks so the producer has + room to keep refilling. No deadlock, no `ParseRingOverflow` for + these primitives — the caller gets the full byte string they + asked for. Primitives that cannot be drained safely (literal + `byteString` matches, `isolate`, raw `ensureN#`) still surface + `ParseRingOverflow` when their size exceeds the ring. + +* `Wireform.Parser.Driver` / `Wireform.Parser.Error` -- detect the + case where the streaming parser asks (via `ensureN#`, `byteString`, + or `isolate`) for more bytes than the ring can ever hold and + short-circuit with a new `ParseRingOverflow` error variant + (position, requested bytes, ring size). Without this guard the + wait loop spins forever on a no-progress @MoreData@ from the + transport because the producer cannot make room and the consumer + is suspended waiting for it. + +* `Wireform.Parser.Driver` -- fix a latent wrap bug in the + StepSuspend / StepCheckpoint resume paths. The driver used to + compute the resumed eob as @base + (head .&. mask)@, which + collapses onto cur whenever the producer has filled exactly one + ring-worth of bytes since the suspension (because the masked + offsets coincide). The parser would then see zero bytes available + even though `head - cur == ringSize`. Replaced with + @newCur + (head - pos)@, which correctly places eob in the second + mapping when wrap happens. Same fix applied to the initial eob + in `runParserInternal`. + +* `Wireform.Parser.Internal.ParserEnv` -- replaced the immutable + `peStartPos` / `peInitCur` fields with mutable `peAnchorPos` / + `peAnchorCur` cells, plus matching `writeAnchor` / `writeAnchor#` + helpers. Whenever the driver wraps the parser's cur back into the + first mapping (StepCheckpoint or StepSuspend resume), it now also + re-anchors the env so that `curToPos` keeps producing the correct + absolute position even after the wrap. `curToPos` is now + State#-threaded (no behavioural change for ordinary callers — the + existing call sites all already have a State# in scope). Callers + that constructed `ParserEnv` directly (driver, parseByteString, + Stateful) now stack-allocate a 24-byte cells buffer (via + `allocaBytes`) instead of `bracket (mallocBytes 8) free`. Removes + a C `malloc` + `free` round-trip per `runParser` call on the + streaming path. + +## 0.1.0.0 -- 2026 + Initial release. + +* `Wireform.FFI` -- C FFI surface for the per-format encoders / + decoders: + * Packed-varint pre-scan (`countPackedVarints`, + `packedAllSingleByte`). + * SWAR UTF-8 validation (`validateUtf8SWAR`). + * Branchless 8-byte varint decode (`decodeVarintSWAR`). + * Page-boundary relocation helper (`relocatePageBoundary`). + * C-native scalar field encoders (`encodeLengthDelimitedC`, + `encodeVarintFieldC`, `encodeBoolFieldC`). + * SIMD NUL / byte / ASCII / JSON-escape / EDN-whitespace + scanners (`findNul`, `findByte`, `isAscii`, + `findJsonEscape`, `skipWhitespace`). + * SIMD Iceberg partition-bounds comparison (`compareBounds`). + * SIMD Arrow IPC buffer validation + (`validateArrowBuffers`). + * BE / LE endianness helpers (`readBE16H` … `writeLE64H`). + * Text decoding via the `text` library's bundled simdutf + (`decodeTextFast`). +* `Wireform.Encode.Direct` -- shared direct-write encode buffer + primitive (`Buf`) for the per-format encoders. +* `Wireform.Hash` -- SIMD-accelerated hashing helpers. + +C sources live in `cbits/`; the vendored `simde` headers ship in +`include/simde/`. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -41,19 +41,6 @@ duplicating the `__attribute__((target(...)))` / `simde-features.h` plumbing. -## Performance tip - -Compiling with `-fllvm` alongside `-O2` typically yields **20–30% throughput -gains** on the encode/decode hot paths. LLVM produces better instruction -scheduling and loop vectorisation for the unboxed arithmetic in `Wireform.FFI` -and `Wireform.Encode.Direct`. - -``` --- in your package's cabal file, or in cabal.project.local: -package wireform-core - ghc-options: -fllvm -``` - ## License BSD-3-Clause. Vendored `simde` headers carry their own MIT license diff --git a/bench/Common.hs b/bench/Common.hs new file mode 100644 --- /dev/null +++ b/bench/Common.hs @@ -0,0 +1,19 @@ +{-# LANGUAGE StrictData #-} + +module Common where + +import Data.ByteString qualified as B + + +type Name = B.ByteString + + +data Tm + = Var Name + | App Tm Tm + | Lam Name Tm + | Let Name Tm Tm + | Int Int + | Add Tm Tm + | Mul Tm Tm + deriving (Show) diff --git a/bench/FPBasic.hs b/bench/FPBasic.hs new file mode 100644 --- /dev/null +++ b/bench/FPBasic.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +module FPBasic ( + runSexp, + runLongws, + runNumcsv, + runTm, +) where + +import Common +import Data.ByteString qualified as B +import FlatParse.Basic + + +ws, open, close, ident, sexp, src :: Parser () () +ws = skipMany $(switch [|case _ of " " -> pure (); "\n" -> pure ()|]) +open = $(char '(') >> ws +close = $(char ')') >> ws +ident = skipSome (skipSatisfyAscii isLatinLetter) >> ws +sexp = branch open (skipSome sexp >> close) ident +src = sexp >> eof + + +runSexp :: B.ByteString -> Result () () +runSexp = runParser src + + +longw, longws :: Parser () () +longw = $(string "thisisalongkeyword") +longws = skipSome (longw >> ws) >> eof + + +runLongws :: B.ByteString -> Result () () +runLongws = runParser longws + + +numeral, comma, numcsv :: Parser () () +numeral = skipSome (skipSatisfyAscii isDigit) >> ws +comma = $(char ',') >> ws +numcsv = numeral >> skipMany (comma >> numeral) >> eof + + +runNumcsv :: B.ByteString -> Result () () +runNumcsv = runParser numcsv + + +ident' :: Parser () B.ByteString +ident' = byteStringOf (skipSome (skipSatisfyAscii \c -> isLatinLetter c || isDigit c)) <* ws + + +equal, semi, dot, addOp, mulOp, parl, parr :: Parser () () +equal = $(string "=") <* ws +semi = $(string ";") <* ws +dot = $(string ".") <* ws +addOp = $(string "+") <* ws +mulOp = $(string "*") <* ws +parl = $(string "(") <* ws +parr = $(string ")") <* ws + + +add :: Parser () Tm +add = chainl Add mul (addOp *> mul) + + +mul :: Parser () Tm +mul = chainl Mul spine (mulOp *> spine) + + +spine :: Parser () Tm +spine = chainl App atom atom + + +atom :: Parser () Tm +atom = + (Int <$> (anyAsciiDecimalInt <* ws)) + <|> (Var <$> ident') + <|> (parl *> tm <* parr) + + +tm :: Parser () Tm +tm = + $( switch + [| + case _ of + "fun" -> do ws; x <- ident'; dot; t <- tm; pure (Lam x t) + "let" -> do ws; x <- ident'; equal; t <- tm; semi; u <- tm; pure (Let x t u) + _ -> add + |] + ) + + +runTm :: B.ByteString -> Result () Tm +runTm = runParser (ws *> tm <* eof) diff --git a/bench/ParserBench.hs b/bench/ParserBench.hs new file mode 100644 --- /dev/null +++ b/bench/ParserBench.hs @@ -0,0 +1,328 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Main where + +import Criterion.Main +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Builder qualified as BSB +import Data.ByteString.Char8 qualified as BSC +import Data.ByteString.Lazy qualified as LBS +import Data.Int +import Data.Word +import FPBasic qualified +import FlatParse.Basic qualified as FP +import WFBasic qualified +import Wireform.Parser qualified as W +import Wireform.Parser.Driver qualified as W +import Wireform.Parser.Internal (Pure) + + +------------------------------------------------------------------------ +-- Input generation +------------------------------------------------------------------------ + +-- | N copies of a 4-byte big-endian word +mkWord32Input :: Int -> ByteString +mkWord32Input n = + LBS.toStrict . BSB.toLazyByteString $ + mconcat [BSB.word32BE (fromIntegral i) | i <- [0 .. n - 1]] + + +-- | N copies of a single byte +mkByteInput :: Int -> ByteString +mkByteInput n = BS.replicate n 0x42 + + +{- | Length-prefixed messages: 1-byte length + payload +Each message is 8 bytes of payload (length byte = 0x08) +-} +mkLengthPrefixedInput :: Int -> ByteString +mkLengthPrefixedInput n = + LBS.toStrict . BSB.toLazyByteString $ + mconcat + [ BSB.word8 8 <> BSB.byteString (BS.replicate 8 (fromIntegral i)) + | i <- [0 .. n - 1] + ] + + +-- | ASCII decimal numbers separated by newlines +mkAsciiDecimalInput :: Int -> ByteString +mkAsciiDecimalInput n = + LBS.toStrict . BSB.toLazyByteString $ + mconcat [BSB.stringUtf8 (show i) <> BSB.word8 0x0A | i <- [0 .. n - 1]] + + +-- | Alternating tag bytes: 0x01 or 0x02, each followed by a word32be +mkTaggedInput :: Int -> ByteString +mkTaggedInput n = + LBS.toStrict . BSB.toLazyByteString $ + mconcat + [ BSB.word8 (if even i then 0x01 else 0x02) + <> BSB.word32BE (fromIntegral i) + | i <- [0 .. n - 1] + ] + + +-- | UTF-8 text: ASCII (1-byte) characters +mkAsciiTextInput :: Int -> ByteString +mkAsciiTextInput n = BS.replicate n 0x61 -- 'a' + + +-- | UTF-8 text: 2-byte characters (Latin-1 supplement, e.g. é = C3 A9) +mkUtf8_2byteInput :: Int -> ByteString +mkUtf8_2byteInput n = + LBS.toStrict . BSB.toLazyByteString $ + mconcat [BSB.word8 0xC3 <> BSB.word8 0xA9 | _ <- [1 .. n]] + + +------------------------------------------------------------------------ +-- Wireform parsers +------------------------------------------------------------------------ + +type WP = W.Parser Pure () + + +wfWord32s :: Int -> WP () +wfWord32s 0 = pure () +wfWord32s n = do + !_ <- W.anyWord32be + wfWord32s (n - 1) +{-# INLINE wfWord32s #-} + + +wfBytes :: Int -> WP () +wfBytes 0 = pure () +wfBytes n = do + !_ <- W.anyWord8 + wfBytes (n - 1) +{-# INLINE wfBytes #-} + + +wfLengthPrefixed :: Int -> WP () +wfLengthPrefixed 0 = pure () +wfLengthPrefixed n = do + !len <- W.anyWord8 + !_ <- W.takeBs (fromIntegral len) + wfLengthPrefixed (n - 1) +{-# INLINE wfLengthPrefixed #-} + + +wfAsciiDecimals :: Int -> WP () +wfAsciiDecimals 0 = pure () +wfAsciiDecimals n = do + !_ <- W.anyAsciiDecimalWord + W.word8 0x0A + wfAsciiDecimals (n - 1) +{-# INLINE wfAsciiDecimals #-} + + +wfTagged :: Int -> WP () +wfTagged 0 = pure () +wfTagged n = do + (W.word8 0x01 >> W.anyWord32be >> pure ()) + W.<|> (W.word8 0x02 >> W.anyWord32be >> pure ()) + wfTagged (n - 1) +{-# INLINE wfTagged #-} + + +wfAsciiChars :: Int -> WP () +wfAsciiChars 0 = pure () +wfAsciiChars n = do + W.skipSatisfyAscii (\_ -> True) + wfAsciiChars (n - 1) +{-# INLINE wfAsciiChars #-} + + +wfUtf8Chars :: Int -> WP () +wfUtf8Chars 0 = pure () +wfUtf8Chars n = do + !_ <- W.anyChar + wfUtf8Chars (n - 1) +{-# INLINE wfUtf8Chars #-} + + +------------------------------------------------------------------------ +-- FlatParse parsers +------------------------------------------------------------------------ + +type FPP = FP.Parser () + + +fpWord32s :: Int -> FPP () +fpWord32s 0 = pure () +fpWord32s n = do + !_ <- FP.anyWord32be + fpWord32s (n - 1) +{-# INLINE fpWord32s #-} + + +fpBytes :: Int -> FPP () +fpBytes 0 = pure () +fpBytes n = do + !_ <- FP.anyWord8 + fpBytes (n - 1) +{-# INLINE fpBytes #-} + + +fpLengthPrefixed :: Int -> FPP () +fpLengthPrefixed 0 = pure () +fpLengthPrefixed n = do + !len <- FP.anyWord8 + !_ <- FP.take (fromIntegral len) + fpLengthPrefixed (n - 1) +{-# INLINE fpLengthPrefixed #-} + + +fpAsciiDecimals :: Int -> FPP () +fpAsciiDecimals 0 = pure () +fpAsciiDecimals n = do + !_ <- FP.anyAsciiDecimalWord + FP.word8 0x0A + fpAsciiDecimals (n - 1) +{-# INLINE fpAsciiDecimals #-} + + +fpTagged :: Int -> FPP () +fpTagged 0 = pure () +fpTagged n = do + (FP.word8 0x01 >> FP.anyWord32be >> pure ()) + FP.<|> (FP.word8 0x02 >> FP.anyWord32be >> pure ()) + fpTagged (n - 1) +{-# INLINE fpTagged #-} + + +fpAsciiChars :: Int -> FPP () +fpAsciiChars 0 = pure () +fpAsciiChars n = do + FP.skipSatisfyAscii (\_ -> True) + fpAsciiChars (n - 1) +{-# INLINE fpAsciiChars #-} + + +fpUtf8Chars :: Int -> FPP () +fpUtf8Chars 0 = pure () +fpUtf8Chars n = do + !_ <- FP.anyChar + fpUtf8Chars (n - 1) +{-# INLINE fpUtf8Chars #-} + + +------------------------------------------------------------------------ +-- Runners +------------------------------------------------------------------------ + +runWF :: WP a -> ByteString -> a +runWF p bs = case W.parseByteString p bs of + Right a -> a + Left _ -> error "wireform parse failed" +{-# INLINE runWF #-} + + +runFP :: FPP a -> ByteString -> a +runFP p bs = case FP.runParser p bs of + FP.OK a _ -> a + _ -> error "flatparse parse failed" +{-# INLINE runFP #-} + + +------------------------------------------------------------------------ +-- Benchmark harness +------------------------------------------------------------------------ + +main :: IO () +main = do + let !n = 100000 + + -- Pre-generate inputs + let !byteInput = mkByteInput n + !word32Input = mkWord32Input n + !lpInput = mkLengthPrefixedInput n + !decInput = mkAsciiDecimalInput n + !taggedInput = mkTaggedInput n + !asciiInput = mkAsciiTextInput n + !utf8Input = mkUtf8_2byteInput n + + putStrLn $ "Input sizes:" + putStrLn $ " byte: " <> show (BS.length byteInput) <> " bytes" + putStrLn $ " word32: " <> show (BS.length word32Input) <> " bytes" + putStrLn $ " len-pfx: " <> show (BS.length lpInput) <> " bytes" + putStrLn $ " decimal: " <> show (BS.length decInput) <> " bytes" + putStrLn $ " tagged: " <> show (BS.length taggedInput) <> " bytes" + putStrLn $ " ascii: " <> show (BS.length asciiInput) <> " bytes" + putStrLn $ " utf8-2b: " <> show (BS.length utf8Input) <> " bytes" + + defaultMain + [ bgroup + "anyWord8 x100k" + [ bench "wireform" $ nf (runWF (wfBytes n)) byteInput + , bench "flatparse" $ nf (runFP (fpBytes n)) byteInput + ] + , bgroup + "anyWord32be x100k" + [ bench "wireform" $ nf (runWF (wfWord32s n)) word32Input + , bench "flatparse" $ nf (runFP (fpWord32s n)) word32Input + ] + , bgroup + "length-prefixed messages x100k" + [ bench "wireform" $ nf (runWF (wfLengthPrefixed n)) lpInput + , bench "flatparse" $ nf (runFP (fpLengthPrefixed n)) lpInput + ] + , bgroup + "ASCII decimal + newline x100k" + [ bench "wireform" $ nf (runWF (wfAsciiDecimals n)) decInput + , bench "flatparse" $ nf (runFP (fpAsciiDecimals n)) decInput + ] + , bgroup + "tagged alternatives x100k" + [ bench "wireform" $ nf (runWF (wfTagged n)) taggedInput + , bench "flatparse" $ nf (runFP (fpTagged n)) taggedInput + ] + , bgroup + "anyCharASCII x100k" + [ bench "wireform" $ nf (runWF (wfAsciiChars n)) asciiInput + , bench "flatparse" $ nf (runFP (fpAsciiChars n)) asciiInput + ] + , bgroup + "anyChar (2-byte UTF-8) x100k" + [ bench "wireform" $ nf (runWF (wfUtf8Chars n)) utf8Input + , bench "flatparse" $ nf (runFP (fpUtf8Chars n)) utf8Input + ] + , -- Flatparse-equivalent real-world benchmarks + bgroup + "sexp" + [ bench "wireform" $ whnf WFBasic.runSexp sexpInp + , bench "flatparse" $ whnf FPBasic.runSexp sexpInp + ] + , bgroup + "long keyword" + [ bench "wireform" $ whnf WFBasic.runLongws longwsInp + , bench "flatparse" $ whnf FPBasic.runLongws longwsInp + ] + , bgroup + "numeral csv" + [ bench "wireform" $ whnf WFBasic.runNumcsv numcsvInp + , bench "flatparse" $ whnf FPBasic.runNumcsv numcsvInp + ] + , bgroup + "lambda term" + [ bench "wireform" $ whnf WFBasic.runTm tmInp + , bench "flatparse" $ whnf FPBasic.runTm tmInp + ] + ] + where + sexpInp = BS.concat $ "(" : replicate 33333 "(foo (foo (foo ((bar baza)))))" <> [")"] + longwsInp = BS.concat $ replicate 55555 "thisisalongkeyword " + numcsvInp = BSC.pack (concat ("0" : [", " <> show i | i <- [1 .. 100000 :: Int]])) + tmInp = + BSC.pack + ( unlines + ( [ "let x" <> show x <> " = fun f. fun g. fun x. fun y. f (f (f ((g x y + g x y) * g x y * g x y * 13500)));" + | x <- [0 .. 3000 :: Int] + ] + <> [("x1000" :: String)] + ) + ) diff --git a/bench/SendBuilderBench.hs b/bench/SendBuilderBench.hs new file mode 100644 --- /dev/null +++ b/bench/SendBuilderBench.hs @@ -0,0 +1,149 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | Benchmark: sendBuilder (via ByteString) vs sendBuilderDirect (ring sink). + +Uses an in-memory send transport backed by a magic ring so the +benchmark isolates builder-to-ring staging cost without any I/O. +-} +module Main where + +import Criterion.Main +import Data.ByteString qualified as BS +import Data.IORef +import Data.Word +import Foreign.Ptr (Ptr) +import Wireform.Builder +import Wireform.Ring.Internal (ringBase, ringMask, ringSize, withMagicRing) +import Wireform.Transport.Send + + +------------------------------------------------------------------------ +-- In-memory send transport (no I/O, immediate drain) +------------------------------------------------------------------------ + +{- | Build a 'SendTransport' over a magic ring that acts as an +instant drain: every 'sendPublishHead' immediately advances the +tail to match, so the ring never fills. This isolates the +builder staging cost. +-} +mkMemorySendTransport + :: Ptr Word8 + -> Int + -> Int + -> IO SendTransport +mkMemorySendTransport base sz msk = do + headRef <- newIORef (0 :: Word64) + tailRef <- newIORef (0 :: Word64) + let publish h = do + writeIORef headRef h + writeIORef tailRef h + pure + SendTransport + { sendRingBase = base + , sendRingSize = sz + , sendRingMask = msk + , sendLoadTail = readIORef tailRef + , sendLoadHead = readIORef headRef + , sendPublishHead = publish + , sendWaitSpace = \_ -> do + tl <- readIORef tailRef + pure (SendSpaceAvailable tl) + , sendFlush = pure () + , sendShutdownWrite = pure () + , sendClose = pure () + } + + +------------------------------------------------------------------------ +-- Builders of various sizes +------------------------------------------------------------------------ + +-- 64-byte payload: small protocol header / Kafka produce ack +smallBuilder :: Builder +smallBuilder = + mconcat + [ word32BE 0x00000001 + , word32BE 0x00000040 + , byteStringCopy (BS.replicate 56 0xAA) + ] + + +-- 256-byte payload: typical Kafka produce record +mediumBuilder :: Builder +mediumBuilder = + mconcat + [ word32BE 0x00000002 + , word32BE 0x00000100 + , byteStringCopy (BS.replicate 248 0xBB) + ] + + +-- 1 KiB payload: HTTP response header block +kiloBuilder :: Builder +kiloBuilder = + mconcat + [ word32BE 0x00000003 + , word32BE 0x00000400 + , byteStringCopy (BS.replicate 1016 0xCC) + ] + + +-- 4 KiB payload: typical HTTP response body chunk +fourKBuilder :: Builder +fourKBuilder = + mconcat + [ word32BE 0x00000004 + , word32BE 0x00001000 + , byteStringCopy (BS.replicate 4088 0xDD) + ] + + +-- 16 KiB payload: large message, exercises ring overflow +sixteenKBuilder :: Builder +sixteenKBuilder = + mconcat + [ word32BE 0x00000005 + , word32BE 0x00004000 + , byteStringCopy (BS.replicate 16376 0xEE) + ] + + +-- 64 KiB payload: stresses ring chunking +sixtyFourKBuilder :: Builder +sixtyFourKBuilder = + mconcat + [ word32BE 0x00000006 + , word32BE 0x00010000 + , byteStringCopy (BS.replicate 65528 0xFF) + ] + + +------------------------------------------------------------------------ +-- Main +------------------------------------------------------------------------ + +main :: IO () +main = + withMagicRing (256 * 1024) $ \ring -> do + let !base = ringBase ring + !sz = ringSize ring + !msk = ringMask ring + t <- mkMemorySendTransport base sz msk + + let benchPair name builder = + bgroup + name + [ bench "direct (RingSink)" $ nfIO (sendBuilderDirect t builder) + , bench "via ByteString (old)" $ nfIO (sendBuilderViaByteString t builder) + ] + + defaultMain + [ benchPair "64 B" smallBuilder + , benchPair "256 B" mediumBuilder + , benchPair "1 KiB" kiloBuilder + , benchPair "4 KiB" fourKBuilder + , benchPair "16 KiB" sixteenKBuilder + , benchPair "64 KiB" sixtyFourKBuilder + ] diff --git a/bench/WFBasic.hs b/bench/WFBasic.hs new file mode 100644 --- /dev/null +++ b/bench/WFBasic.hs @@ -0,0 +1,155 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | Wireform equivalents of flatparse's benchmark parsers. +These test the same parse patterns (sexp, long keyword, numeric csv, +lambda term) using wireform's combinator API. +-} +module WFBasic ( + runSexp, + runLongws, + runNumcsv, + runTm, +) where + +import Common (Tm (..)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as B +import Wireform.Parser +import Wireform.Parser.Driver (parseByteString) +import Wireform.Parser.Internal (Pure) + + +type P = Parser Pure () + + +------------------------------------------------------------------------ +-- sexp +------------------------------------------------------------------------ + +ws :: P () +ws = skipMany (satisfyAscii (\c -> c == ' ' || c == '\n') *> pure ()) +{-# INLINE ws #-} + + +open :: P () +open = word8 0x28 >> ws -- '(' +{-# INLINE open #-} + + +close :: P () +close = word8 0x29 >> ws -- ')' +{-# INLINE close #-} + + +ident :: P () +ident = skipSome (skipSatisfyAscii isLatinLetter) >> ws +{-# INLINE ident #-} + + +sexp :: P () +sexp = branch open (skipSome sexp >> close) ident + + +src :: P () +src = sexp >> eof + + +runSexp :: ByteString -> Either (ParseError ()) () +runSexp = parseByteString src + + +------------------------------------------------------------------------ +-- long keyword +------------------------------------------------------------------------ + +longw :: P () +longw = byteString "thisisalongkeyword" +{-# INLINE longw #-} + + +longws :: P () +longws = skipSome (longw >> ws) >> eof + + +runLongws :: ByteString -> Either (ParseError ()) () +runLongws = parseByteString longws + + +------------------------------------------------------------------------ +-- numeral csv +------------------------------------------------------------------------ + +numeral :: P () +numeral = skipSome (skipSatisfyAscii isDigit) >> ws +{-# INLINE numeral #-} + + +comma :: P () +comma = word8 0x2C >> ws -- ',' +{-# INLINE comma #-} + + +numcsv :: P () +numcsv = numeral >> skipMany (comma >> numeral) >> eof + + +runNumcsv :: ByteString -> Either (ParseError ()) () +runNumcsv = parseByteString numcsv + + +------------------------------------------------------------------------ +-- lambda term +------------------------------------------------------------------------ + +ident' :: P ByteString +ident' = byteStringOf (skipSome (skipSatisfyAscii \c -> isLatinLetter c || isDigit c)) <* ws +{-# INLINE ident' #-} + + +equal, semi, dot, addOp, mulOp, parl, parr :: P () +equal = byteString "=" >> ws +{-# INLINE equal #-} +semi = byteString ";" >> ws +{-# INLINE semi #-} +dot = byteString "." >> ws +{-# INLINE dot #-} +addOp = byteString "+" >> ws +{-# INLINE addOp #-} +mulOp = byteString "*" >> ws +{-# INLINE mulOp #-} +parl = byteString "(" >> ws +{-# INLINE parl #-} +parr = byteString ")" >> ws +{-# INLINE parr #-} + + +add :: P Tm +add = chainl Add mul (addOp *> mul) + + +mul :: P Tm +mul = chainl Mul spine (mulOp *> spine) + + +spine :: P Tm +spine = chainl App atom atom + + +atom :: P Tm +atom = + (Int <$> (anyAsciiDecimalInt <* ws)) + <|> (Var <$> ident') + <|> (parl *> tm <* parr) + + +tm :: P Tm +tm = + (byteString "fun" *> ws *> do x <- ident'; dot; t <- tm; pure (Lam x t)) + <|> (byteString "let" *> ws *> do x <- ident'; equal; t <- tm; semi; u <- tm; pure (Let x t u)) + <|> add + + +runTm :: ByteString -> Either (ParseError ()) Tm +runTm = parseByteString (ws *> tm <* eof) diff --git a/cbits/base64.c b/cbits/base64.c new file mode 100644 --- /dev/null +++ b/cbits/base64.c @@ -0,0 +1,293 @@ +/* + * RFC 4648 base64 encode / decode. + * + * Two SIMD-accelerated primitives, in line with the rest of + * wireform-core: + * + * * 'hs_base64_encode': SSSE3 (via simde) inner loop, 12 input + * bytes -> 16 output chars per iteration, scalar prologue / + * epilogue for the tail. + * * 'hs_base64_decode': SSE2 (via simde) pre-scan, 16 chars at a + * time, rejects any window containing a high-bit byte (which + * by construction is outside the base64 alphabet); the actual + * sextet extraction then runs through a tight scalar loop on + * a 256-entry decode table. This is the same shape every + * wireform format takes for "validate-then-decode" SIMD: see + * 'hs_proto_validate_utf8_fast' / 'hs_find_byte'. + * + * The encode SIMD path uses the well-known Mula / alfredklomp + * formula (no LUT in memory; ASCII offset is derived from the + * 6-bit value via a single PSHUFB on a 16-byte lookup). + * + * The decoder is kept scalar in the sextet-extraction step + * because the typical wireform call site (SHA-1 = 28 base64 + * chars; proto3 JSON bytes fields = O(100s)) does not benefit + * enough from a Mula-style PSHUFB classifier to be worth the + * table-debugging cost. Encoding -- which is on the hot path + * for big payloads via the proto3 JSON mapping -- stays SSSE3. + */ + +#include +#include +#include +#include + +static const char b64_alphabet[64] = + "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + "abcdefghijklmnopqrstuvwxyz" + "0123456789+/"; + +/* Reverse-lookup table: ASCII -> 6-bit value, or 0xFF for any byte + * outside the standard base64 alphabet. '=' maps to 0xFF too; + * the tail handler checks for it explicitly. */ +static const uint8_t b64_dec_table[256] = { + /* 0x00..0x2A (43 entries) */ + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255, + /* 0x2B '+' = 62, 0x2C ',' invalid, 0x2D '-' invalid, + * 0x2E '.' invalid, 0x2F '/' = 63 */ + 62, 255, 255, 255, 63, + /* 0x30..0x39 '0'..'9' = 52..61 */ + 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, + /* 0x3A..0x40 (7 entries) */ + 255, 255, 255, 255, 255, 255, 255, + /* 0x41..0x5A 'A'..'Z' = 0..25 */ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + 20, 21, 22, 23, 24, 25, + /* 0x5B..0x60 (6 entries) */ + 255, 255, 255, 255, 255, 255, + /* 0x61..0x7A 'a'..'z' = 26..51 */ + 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, + 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, + 46, 47, 48, 49, 50, 51, + /* 0x7B..0xFF (133 entries) */ + 255, 255, 255, 255, 255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255, + 255,255,255,255,255,255,255,255,255,255,255,255,255,255,255,255 +}; + +/* ------------------------------------------------------------------ + * Length helpers (RFC 4648 sec 4). + * + * Encode: 3-byte triplets pad up to the next multiple of 4 chars. + * Decode (upper bound): every 4 chars produce up to 3 bytes; + * trailing '=' chars subtract one byte each. The exact decoded + * length is reported by hs_base64_decode itself; this helper just + * sizes the output buffer. + * ------------------------------------------------------------------ */ + +int hs_base64_encoded_length(int in_len) +{ + if (in_len < 0) return 0; + return ((in_len + 2) / 3) * 4; +} + +int hs_base64_decoded_max_length(int in_len) +{ + if (in_len < 0) return 0; + return (in_len / 4) * 3; +} + +/* ------------------------------------------------------------------ + * Encode + * + * Returns the number of bytes written to @out. Always equals + * hs_base64_encoded_length(in_len) (padded form). + * ------------------------------------------------------------------ */ + +static inline simde__m128i enc_reshuffle(simde__m128i in) +{ + /* Spread each 3-byte triplet into a 4-byte slot so the 6-bit + * fields land in adjacent bytes. */ + in = simde_mm_shuffle_epi8(in, simde_mm_setr_epi8( + 1, 0, 2, 1, + 4, 3, 5, 4, + 7, 6, 8, 7, + 10, 9, 11, 10)); + + simde__m128i t0 = simde_mm_and_si128(in, simde_mm_set1_epi32(0x0fc0fc00)); + simde__m128i t1 = simde_mm_mulhi_epu16(t0, simde_mm_set1_epi32(0x04000040)); + simde__m128i t2 = simde_mm_and_si128(in, simde_mm_set1_epi32(0x003f03f0)); + simde__m128i t3 = simde_mm_mullo_epi16(t2, simde_mm_set1_epi32(0x01000010)); + return simde_mm_or_si128(t1, t3); +} + +static inline simde__m128i enc_translate(simde__m128i in) +{ + /* 6-bit value v -> ASCII char via a 16-entry PSHUFB lookup of + * the offset to add to v. Branchless. */ + simde__m128i lut = simde_mm_setr_epi8( + 65, 71, -4, -4, -4, -4, -4, -4, + -4, -4, -4, -4, -19, -16, 0, 0); + simde__m128i indices = simde_mm_subs_epu8(in, simde_mm_set1_epi8(51)); + simde__m128i mask = simde_mm_cmpgt_epi8(in, simde_mm_set1_epi8(25)); + indices = simde_mm_sub_epi8(indices, mask); + simde__m128i offsets = simde_mm_shuffle_epi8(lut, indices); + return simde_mm_add_epi8(in, offsets); +} + +int hs_base64_encode(const uint8_t *in, int in_len, uint8_t *out) +{ + int i = 0, j = 0; + + /* SIMD: 12 in -> 16 out per iter. Need 16 readable input bytes + * (we load 16 then mask out the upper 4) so guard accordingly. */ + while (in_len - i >= 16) { + simde__m128i chunk = simde_mm_loadu_si128((const simde__m128i *)(in + i)); + simde__m128i v = enc_reshuffle(chunk); + simde__m128i c = enc_translate(v); + simde_mm_storeu_si128((simde__m128i *)(out + j), c); + i += 12; + j += 16; + } + + /* Scalar tail: 3 in -> 4 out. */ + while (in_len - i >= 3) { + uint32_t v = ((uint32_t)in[i] << 16) + | ((uint32_t)in[i+1] << 8) + | (uint32_t)in[i+2]; + out[j+0] = (uint8_t)b64_alphabet[(v >> 18) & 0x3F]; + out[j+1] = (uint8_t)b64_alphabet[(v >> 12) & 0x3F]; + out[j+2] = (uint8_t)b64_alphabet[(v >> 6) & 0x3F]; + out[j+3] = (uint8_t)b64_alphabet[ v & 0x3F]; + i += 3; + j += 4; + } + + /* RFC 4648 sec 4 padding. */ + int rem = in_len - i; + if (rem == 1) { + uint32_t v = (uint32_t)in[i] << 16; + out[j+0] = (uint8_t)b64_alphabet[(v >> 18) & 0x3F]; + out[j+1] = (uint8_t)b64_alphabet[(v >> 12) & 0x3F]; + out[j+2] = (uint8_t)'='; + out[j+3] = (uint8_t)'='; + j += 4; + } else if (rem == 2) { + uint32_t v = ((uint32_t)in[i] << 16) | ((uint32_t)in[i+1] << 8); + out[j+0] = (uint8_t)b64_alphabet[(v >> 18) & 0x3F]; + out[j+1] = (uint8_t)b64_alphabet[(v >> 12) & 0x3F]; + out[j+2] = (uint8_t)b64_alphabet[(v >> 6) & 0x3F]; + out[j+3] = (uint8_t)'='; + j += 4; + } + return j; +} + +/* ------------------------------------------------------------------ + * Decode + * + * Strict RFC 4648 decoder. Input length must be a multiple of 4; + * any out-of-alphabet byte (other than '=' in the trailing + * position) produces -1. Returns the number of output bytes + * written on success. + * ------------------------------------------------------------------ */ + +/* SSE2 fast-path probe: does this 16-byte window contain any byte + * with the high bit set? Any such byte is by construction outside + * the base64 alphabet, so we can short-circuit. */ +static inline int dec_window_high_bits(const uint8_t *in) +{ + simde__m128i chunk = simde_mm_loadu_si128((const simde__m128i *)in); + return simde_mm_movemask_epi8(chunk); +} + +int hs_base64_decode(const uint8_t *in, int in_len, uint8_t *out) +{ + if (in_len < 0 || (in_len & 3) != 0) return -1; + if (in_len == 0) return 0; + + int i = 0, j = 0; + int main_len = in_len - 4; /* tail quartet handled separately */ + + /* SIMD pre-scan + scalar decode of the main body. + * + * The SIMD step rejects any 16-byte window containing a byte + * with the high bit set (cheap PMOVMSKB + branch); the scalar + * step then decodes 4 quartets = 12 bytes from each clean + * window. This gives us the bulk of the SIMD win (early + * rejection of garbage input) without the table-debugging + * overhead of a full PSHUFB classifier. */ + while (main_len - i >= 16) { + if (dec_window_high_bits(in + i) != 0) return -1; + for (int k = 0; k < 16; k += 4) { + uint8_t a = b64_dec_table[in[i+k+0]]; + uint8_t b = b64_dec_table[in[i+k+1]]; + uint8_t c = b64_dec_table[in[i+k+2]]; + uint8_t d = b64_dec_table[in[i+k+3]]; + if ((a | b | c | d) >= 64) return -1; + uint32_t v = ((uint32_t)a << 18) + | ((uint32_t)b << 12) + | ((uint32_t)c << 6) + | (uint32_t)d; + out[j+0] = (uint8_t)(v >> 16); + out[j+1] = (uint8_t)(v >> 8); + out[j+2] = (uint8_t) v; + j += 3; + } + i += 16; + } + + /* Scalar middle: remaining full quartets before the tail. */ + while (i < main_len) { + uint8_t a = b64_dec_table[in[i+0]]; + uint8_t b = b64_dec_table[in[i+1]]; + uint8_t c = b64_dec_table[in[i+2]]; + uint8_t d = b64_dec_table[in[i+3]]; + if ((a | b | c | d) >= 64) return -1; + uint32_t v = ((uint32_t)a << 18) + | ((uint32_t)b << 12) + | ((uint32_t)c << 6) + | (uint32_t)d; + out[j+0] = (uint8_t)(v >> 16); + out[j+1] = (uint8_t)(v >> 8); + out[j+2] = (uint8_t) v; + i += 4; + j += 3; + } + + /* Tail quartet: may contain '=' padding. */ + { + uint8_t a = b64_dec_table[in[i+0]]; + uint8_t b = b64_dec_table[in[i+1]]; + if (a >= 64 || b >= 64) return -1; + if (in[i+2] == (uint8_t)'=') { + if (in[i+3] != (uint8_t)'=') return -1; + uint32_t v = ((uint32_t)a << 18) | ((uint32_t)b << 12); + out[j+0] = (uint8_t)(v >> 16); + j += 1; + } else if (in[i+3] == (uint8_t)'=') { + uint8_t c = b64_dec_table[in[i+2]]; + if (c >= 64) return -1; + uint32_t v = ((uint32_t)a << 18) + | ((uint32_t)b << 12) + | ((uint32_t)c << 6); + out[j+0] = (uint8_t)(v >> 16); + out[j+1] = (uint8_t)(v >> 8); + j += 2; + } else { + uint8_t c = b64_dec_table[in[i+2]]; + uint8_t d = b64_dec_table[in[i+3]]; + if (c >= 64 || d >= 64) return -1; + uint32_t v = ((uint32_t)a << 18) + | ((uint32_t)b << 12) + | ((uint32_t)c << 6) + | (uint32_t)d; + out[j+0] = (uint8_t)(v >> 16); + out[j+1] = (uint8_t)(v >> 8); + out[j+2] = (uint8_t) v; + j += 3; + } + i += 4; + } + (void)i; + return j; +} diff --git a/cbits/fast_rng.c b/cbits/fast_rng.c new file mode 100644 --- /dev/null +++ b/cbits/fast_rng.c @@ -0,0 +1,154 @@ +/* + * xoshiro256++ — a fast, non-cryptographic 64-bit PRNG. + * + * Reference: + * D. Blackman and S. Vigna, "Scrambled Linear Pseudorandom + * Number Generators", ACM Trans. on Math. Software 47:4 + * (2021). https://vigna.di.unimi.it/ftp/papers/ScrambledLinear.pdf + * + * One 256-bit state per pthread, stored in '__thread' storage so + * concurrent generation across many Haskell capabilities never + * touches a shared cache line. Seeded on first use from + * @getrandom(2)@. + * + * Used by 'wireform-websocket' to roll per-frame masking keys + * without going through the global 'splitmix' MVar. Exposed + * generically because the same primitive is what any + * non-cryptographic random need on the hot path should reach + * for. + */ + +#include +#include +#include +#include + +#if defined(__linux__) +#include +#include +static int kernel_random(void *buf, size_t len) +{ + long r = syscall(SYS_getrandom, buf, len, 0); + return r == (long)len ? 0 : -1; +} +#elif defined(__APPLE__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__) +#include +static int kernel_random(void *buf, size_t len) +{ + arc4random_buf(buf, len); + return 0; +} +#else +/* Fall back to /dev/urandom for portability. */ +#include +static int kernel_random(void *buf, size_t len) +{ + int fd = open("/dev/urandom", O_RDONLY); + if (fd < 0) return -1; + size_t off = 0; + while (off < len) { + ssize_t n = read(fd, (char *)buf + off, len - off); + if (n <= 0) { close(fd); return -1; } + off += (size_t)n; + } + close(fd); + return 0; +} +#endif + +/* TLS access model for the per-thread state below. + * + * A file-local 'static __thread' defaults to the local-exec model: a + * single FS/GS-relative MOV, the fastest TLS access (what the comment + * below describes). That is only valid when the object lands in the + * main executable's static TLS block. When the object is built + * position-independent for a shared library — GHC's dynamic + * ('.dyn_o' -> '.so') way, used whenever wireform-core sits in a + * Template-Haskell dependency closure — local-exec / local-dynamic + * emit 'R_X86_64_TPOFF32' / 'R_X86_64_DTPOFF32' relocations that ld + * cannot link into the '.so' on x86_64-linux ("relocation truncated + * to fit"). + * + * '__PIC__' is defined by GCC / Clang exactly when compiling + * position-independent (the '.dyn_o' way) and undefined for the plain + * static '.o'. Keep the fast local-exec model for the static object; + * only step down to initial-exec ('R_X86_64_GOTTPOFF', a single GOT + * indirection that links into a '.so') when actually building PIC. + * initial-exec keeps the symbol file-local (unlike global-dynamic, + * which would require an exported symbol) and its few words of TLS sit + * within glibc's static-TLS surplus when the library is dlopen'd. */ +#if defined(__PIC__) +#define WF_TLS_MODEL __attribute__((tls_model("initial-exec"))) +#else +#define WF_TLS_MODEL +#endif + +/* Per-thread state. 256-bit (4 × uint64_t). '__thread' is the + * GCC / Clang extension; on every modern Unix it lands in the + * TLS block one offset away from FS / GS register, so the + * dispatch is a single MOV (local-exec; see WF_TLS_MODEL above for + * the position-independent / shared-object case). */ +static __thread uint64_t xoshiro_state[4] WF_TLS_MODEL; +static __thread int xoshiro_seeded WF_TLS_MODEL = 0; + +static inline uint64_t rotl64(uint64_t x, int k) +{ + return (x << k) | (x >> (64 - k)); +} + +/* Recover from a degenerate all-zero seed (xoshiro256++ has a + * fixed point at 0). Vigna's reference splitmix64 is the + * recommended way to expand a single 64-bit input into a 256-bit + * xoshiro seed. We only invoke this when 'kernel_random' + * returns all-zero, which on a working kernel never happens. */ +static inline uint64_t splitmix64_step(uint64_t *x) +{ + *x += 0x9E3779B97F4A7C15ULL; + uint64_t z = *x; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); +} + +static void xoshiro_seed(void) +{ + if (kernel_random(xoshiro_state, sizeof(xoshiro_state)) != 0 + || (xoshiro_state[0] | xoshiro_state[1] + | xoshiro_state[2] | xoshiro_state[3]) == 0) + { + /* Last-resort fallback: tid + nanotime through splitmix. */ + uint64_t fallback = (uint64_t)(uintptr_t)&xoshiro_state; + for (int i = 0; i < 4; i++) { + xoshiro_state[i] = splitmix64_step(&fallback); + } + } + xoshiro_seeded = 1; +} + +/* The xoshiro256++ next() generator. Pure register work after + * the (cached) TLS load. */ +uint64_t hs_xoshiro256pp_next(void) +{ + if (!xoshiro_seeded) xoshiro_seed(); + + const uint64_t *s = xoshiro_state; + const uint64_t result = rotl64(s[0] + s[3], 23) + s[0]; + const uint64_t t = xoshiro_state[1] << 17; + + xoshiro_state[2] ^= xoshiro_state[0]; + xoshiro_state[3] ^= xoshiro_state[1]; + xoshiro_state[1] ^= xoshiro_state[2]; + xoshiro_state[0] ^= xoshiro_state[3]; + xoshiro_state[2] ^= t; + xoshiro_state[3] = rotl64(xoshiro_state[3], 45); + + return result; +} + +/* Force reseed. Useful after fork() and for tests that want + * deterministic state — they can call this then immediately + * overwrite the seed via hs_xoshiro256pp_reseed. */ +void hs_xoshiro256pp_reseed(void) +{ + xoshiro_seeded = 0; +} diff --git a/cbits/fast_scan.c b/cbits/fast_scan.c --- a/cbits/fast_scan.c +++ b/cbits/fast_scan.c @@ -49,3 +49,61 @@ return len; } + +/* + * In-place XOR of a 4-byte repeating mask over @buf[0..len). + * Used by wireform-websocket for RFC 6455 sec 5.3 frame masking + * (both directions: the parser un-masks an inbound payload; the + * builder masks an outbound client payload). + * + * The mask repeats in 4-byte network-order chunks; for the SSE2 + * inner loop we broadcast that 32-bit pattern into a 16-byte + * vector once and XOR 16 bytes per iteration. Scalar prologue + + * epilogue handle any sub-16-byte head / tail. + * + * @mask is a uint32_t with the four mask bytes in NETWORK ORDER, + * i.e. byte 0 of the mask in the high byte. The same byte + * ordering the wire format uses. + */ +void hs_ws_mask(uint8_t *buf, int len, uint32_t mask) +{ + /* Build the byte ordering the loop needs: mask[0] at offset 0, + * mask[1] at offset 1, … Easiest done by re-laying the bytes. */ + uint8_t mb[4]; + mb[0] = (uint8_t)(mask >> 24); + mb[1] = (uint8_t)(mask >> 16); + mb[2] = (uint8_t)(mask >> 8); + mb[3] = (uint8_t) mask; + + int i = 0; + + /* Align the scalar prologue so the main loop can use 'movdqa' + * once we hit a 16-byte boundary. The mask byte at offset i + * is mb[i & 3] because the mask repeats every 4 bytes. */ + while (i < len && ((uintptr_t)(buf + i) & 15) != 0) { + buf[i] ^= mb[i & 3]; + i++; + } + + /* SSE2 main loop: 16 bytes per iteration. The 4-byte mask + * tiled into a 16-byte vector once. If i hit alignment at an + * offset where i & 3 != 0, the tile rotates accordingly. */ + if (i + 16 <= len) { + int phase = i & 3; + uint8_t tile[16]; + for (int k = 0; k < 16; k++) { + tile[k] = mb[(phase + k) & 3]; + } + simde__m128i v = simde_mm_loadu_si128((const simde__m128i *)tile); + for (; i + 16 <= len; i += 16) { + simde__m128i chunk = simde_mm_load_si128((const simde__m128i *)(buf + i)); + chunk = simde_mm_xor_si128(chunk, v); + simde_mm_store_si128((simde__m128i *)(buf + i), chunk); + } + } + + /* Scalar epilogue for the trailing <16 bytes. */ + for (; i < len; i++) { + buf[i] ^= mb[i & 3]; + } +} diff --git a/cbits/magic_ring_linux.c b/cbits/magic_ring_linux.c new file mode 100644 --- /dev/null +++ b/cbits/magic_ring_linux.c @@ -0,0 +1,101 @@ +#define _GNU_SOURCE +#include +#include +#include +#include +#include +#include +#include +#include "magic_ring.h" + +static long cached_page_size = 0; + +long hs_page_size(void) { + if (cached_page_size == 0) { + cached_page_size = sysconf(_SC_PAGESIZE); + } + return cached_page_size; +} + +/* Round up to the next power of two >= val, with val >= page_size. */ +static size_t next_pow2(size_t val) { + val--; + val |= val >> 1; + val |= val >> 2; + val |= val >> 4; + val |= val >> 8; + val |= val >> 16; + val |= val >> 32; + val++; + return val; +} + +/* + * Allocate a double-mapped ring of at least `requested` bytes. + * Returns 0 on success, -1 on failure (sets errno). + * On success, out->base and out->size are populated. + */ +int hs_ring_create(size_t requested, struct hs_ring *out) { + long ps = hs_page_size(); + size_t size; + int fd; + void *base, *m1, *m2; + + if (requested < (size_t)ps) + requested = (size_t)ps; + + /* Round up to page-size multiple, then to power of two */ + size = (requested + (size_t)ps - 1) & ~((size_t)ps - 1); + size = next_pow2(size); + + fd = (int)syscall(SYS_memfd_create, "wireform_ring", 1 /* MFD_CLOEXEC */); + if (fd < 0) + return -1; + + if (ftruncate(fd, (off_t)size) != 0) { + close(fd); + return -1; + } + + /* Reserve 2N contiguous virtual addresses with no permissions */ + base = mmap(NULL, size * 2, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (base == MAP_FAILED) { + close(fd); + return -1; + } + + /* Map the first copy */ + m1 = mmap(base, size, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_FIXED, fd, 0); + if (m1 == MAP_FAILED) { + munmap(base, size * 2); + close(fd); + return -1; + } + + /* Map the second copy immediately after */ + m2 = mmap((char *)base + size, size, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_FIXED, fd, 0); + if (m2 == MAP_FAILED) { + munmap(base, size * 2); + close(fd); + return -1; + } + + close(fd); /* mappings keep the memfd alive */ + + /* Pre-fault every page to avoid latency surprises on first touch */ + memset(base, 0, size); + + out->base = base; + out->size = size; + return 0; +} + +void hs_ring_destroy(struct hs_ring *ring) { + if (ring && ring->base) { + munmap(ring->base, ring->size * 2); + ring->base = NULL; + ring->size = 0; + } +} diff --git a/cbits/magic_ring_posix.c b/cbits/magic_ring_posix.c new file mode 100644 --- /dev/null +++ b/cbits/magic_ring_posix.c @@ -0,0 +1,115 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "magic_ring.h" + +static long cached_page_size = 0; + +long hs_page_size(void) { + if (cached_page_size == 0) { + cached_page_size = sysconf(_SC_PAGESIZE); + } + return cached_page_size; +} + +static size_t next_pow2(size_t val) { + val--; + val |= val >> 1; + val |= val >> 2; + val |= val >> 4; + val |= val >> 8; + val |= val >> 16; + val |= val >> 32; + val++; + return val; +} + +#if !defined(__FreeBSD__) || !defined(SHM_ANON) +static _Atomic long shm_counter = 0; +#endif + +int hs_ring_create(size_t requested, struct hs_ring *out) { + long ps = hs_page_size(); + size_t size; + int fd; + void *base, *m1, *m2; + int saved_errno; + + if (requested < (size_t)ps) + requested = (size_t)ps; + + size = (requested + (size_t)ps - 1) & ~((size_t)ps - 1); + size = next_pow2(size); + +#if defined(__FreeBSD__) && defined(SHM_ANON) + fd = shm_open(SHM_ANON, O_RDWR, 0600); +#else + { + char name[64]; + snprintf(name, sizeof(name), "/wfring.%d.%ld", + getpid(), atomic_fetch_add(&shm_counter, 1)); + fd = shm_open(name, O_RDWR | O_CREAT | O_EXCL, 0600); + if (fd >= 0) + shm_unlink(name); + } +#endif + + if (fd < 0) + return -1; + + if (ftruncate(fd, (off_t)size) != 0) { + saved_errno = errno; + close(fd); + errno = saved_errno; + return -1; + } + + base = mmap(NULL, size * 2, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (base == MAP_FAILED) { + saved_errno = errno; + close(fd); + errno = saved_errno; + return -1; + } + + m1 = mmap(base, size, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_FIXED, fd, 0); + if (m1 == MAP_FAILED) { + saved_errno = errno; + munmap(base, size * 2); + close(fd); + errno = saved_errno; + return -1; + } + + m2 = mmap((char *)base + size, size, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_FIXED, fd, 0); + if (m2 == MAP_FAILED) { + saved_errno = errno; + munmap(base, size * 2); + close(fd); + errno = saved_errno; + return -1; + } + + close(fd); + memset(base, 0, size); + + out->base = base; + out->size = size; + return 0; +} + +void hs_ring_destroy(struct hs_ring *ring) { + if (ring && ring->base) { + munmap(ring->base, ring->size * 2); + ring->base = NULL; + ring->size = 0; + } +} diff --git a/cbits/magic_ring_windows.c b/cbits/magic_ring_windows.c new file mode 100644 --- /dev/null +++ b/cbits/magic_ring_windows.c @@ -0,0 +1,103 @@ +#ifdef _WIN32 + +#include +#include +#include +#include "magic_ring.h" + +static long cached_page_size = 0; + +long hs_page_size(void) { + if (cached_page_size == 0) { + SYSTEM_INFO si; + GetSystemInfo(&si); + cached_page_size = (long)si.dwAllocationGranularity; + } + return cached_page_size; +} + +static size_t next_pow2(size_t val) { + val--; + val |= val >> 1; + val |= val >> 2; + val |= val >> 4; + val |= val >> 8; + val |= val >> 16; +#if defined(_WIN64) + val |= val >> 32; +#endif + val++; + return val; +} + +int hs_ring_create(size_t requested, struct hs_ring *out) { + long ps = hs_page_size(); + size_t size; + PVOID placeholder = NULL; + HANDLE section = NULL; + PVOID view1 = NULL, view2 = NULL; + + if (requested < (size_t)ps) + requested = (size_t)ps; + + size = (requested + (size_t)ps - 1) & ~((size_t)ps - 1); + size = next_pow2(size); + + placeholder = VirtualAlloc2(NULL, NULL, size * 2, + MEM_RESERVE | MEM_RESERVE_PLACEHOLDER, PAGE_NOACCESS, NULL, 0); + if (!placeholder) + return -1; + + /* Split the placeholder into two halves */ + if (!VirtualFree(placeholder, size, MEM_RELEASE | MEM_PRESERVE_PLACEHOLDER)) { + VirtualFree(placeholder, 0, MEM_RELEASE); + return -1; + } + + section = CreateFileMappingW(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, + (DWORD)(size >> 32), (DWORD)(size & 0xFFFFFFFF), NULL); + if (!section) { + VirtualFree(placeholder, 0, MEM_RELEASE); + VirtualFree((char *)placeholder + size, 0, MEM_RELEASE); + return -1; + } + + view1 = MapViewOfFile3(section, NULL, placeholder, 0, size, + MEM_REPLACE_PLACEHOLDER, PAGE_READWRITE, NULL, 0); + if (!view1) + goto fail; + + view2 = MapViewOfFile3(section, NULL, (char *)placeholder + size, 0, size, + MEM_REPLACE_PLACEHOLDER, PAGE_READWRITE, NULL, 0); + if (!view2) { + UnmapViewOfFile(view1); + goto fail; + } + + /* Pre-fault */ + memset(view1, 0, size); + + out->base = view1; + out->size = size; + out->section = section; + return 0; + +fail: + CloseHandle(section); + VirtualFree(placeholder, 0, MEM_RELEASE); + VirtualFree((char *)placeholder + size, 0, MEM_RELEASE); + return -1; +} + +void hs_ring_destroy(struct hs_ring *ring) { + if (ring && ring->base) { + UnmapViewOfFile(ring->base); + UnmapViewOfFile((char *)ring->base + ring->size); + CloseHandle(ring->section); + ring->base = NULL; + ring->size = 0; + ring->section = NULL; + } +} + +#endif /* _WIN32 */ diff --git a/include/magic_ring.h b/include/magic_ring.h new file mode 100644 --- /dev/null +++ b/include/magic_ring.h @@ -0,0 +1,18 @@ +#ifndef WIREFORM_MAGIC_RING_H +#define WIREFORM_MAGIC_RING_H + +#include + +struct hs_ring { + void *base; + size_t size; +#ifdef _WIN32 + void *section; /* HANDLE */ +#endif +}; + +long hs_page_size(void); +int hs_ring_create(size_t requested, struct hs_ring *out); +void hs_ring_destroy(struct hs_ring *ring); + +#endif diff --git a/include/simde/simde/arm/neon/st1_lane.h b/include/simde/simde/arm/neon/st1_lane.h --- a/include/simde/simde/arm/neon/st1_lane.h +++ b/include/simde/simde/arm/neon/st1_lane.h @@ -522,3 +522,4 @@ HEDLEY_DIAGNOSTIC_POP #endif /* !defined(SIMDE_ARM_NEON_ST1_LANE_H) */ + diff --git a/src/Wireform/Base64.hs b/src/Wireform/Base64.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Base64.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE BangPatterns #-} + +{- | RFC 4648 \u00a74 base64 encode / decode. + +The implementation lives in @cbits\/base64.c@: an SSSE3 (via +simde) inner loop encoding 12 input bytes per iteration with a +scalar prologue \/ epilogue for the tail. The decoder takes the +same shape and additionally validates the input alphabet +character by character; non-alphabet bytes outside the trailing +@=@ padding produce 'Nothing'. + +These primitives are kept here (rather than in a downstream +package) so every wireform format that needs base64 \u2014 +WebSocket handshakes, BSON binary subtype 0x00, the proto3 JSON +mapping for @bytes@, the Avro JSON encoding, the Iceberg +@avro-data-files@ glue \u2014 shares one SIMD implementation +rather than each pulling its own @base64-bytestring@. + +The encoder produces /padded/ standard base64; there is no +URL-safe variant. Add one in a separate module if needed. +-} +module Wireform.Base64 ( + -- * Encode + encodeBase64, + encodeBase64Length, + + -- * Decode + decodeBase64, + decodeBase64MaxLength, +) where + +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Internal qualified as BSI +import Data.ByteString.Unsafe qualified as BSU +import Data.Word (Word8) +import Foreign.C.Types (CInt (..)) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import System.IO.Unsafe (unsafePerformIO) + + +------------------------------------------------------------------------ +-- FFI +------------------------------------------------------------------------ + +foreign import ccall unsafe "hs_base64_encoded_length" + c_base64_encoded_length :: CInt -> CInt + + +foreign import ccall unsafe "hs_base64_decoded_max_length" + c_base64_decoded_max_length :: CInt -> CInt + + +{- | Encoder \/ decoder run as 'IO' so the buffer-mutating side +effect is sequenced; declaring them as pure @CInt@-returning +functions confuses GHC's strictness analysis and the call can +be silently elided. +-} +foreign import ccall unsafe "hs_base64_encode" + c_base64_encode :: Ptr Word8 -> CInt -> Ptr Word8 -> IO CInt + + +foreign import ccall unsafe "hs_base64_decode" + c_base64_decode :: Ptr Word8 -> CInt -> Ptr Word8 -> IO CInt + + +------------------------------------------------------------------------ +-- Length helpers +------------------------------------------------------------------------ + +{- | Number of base64 characters produced for an input of @n@ bytes +(including @=@ padding). +-} +encodeBase64Length :: Int -> Int +encodeBase64Length = fromIntegral . c_base64_encoded_length . fromIntegral +{-# INLINE encodeBase64Length #-} + + +{- | Upper bound on the number of bytes produced when decoding @n@ +input characters. The actual decoded length is shorter by the +number of trailing @=@ characters. +-} +decodeBase64MaxLength :: Int -> Int +decodeBase64MaxLength = fromIntegral . c_base64_decoded_max_length . fromIntegral +{-# INLINE decodeBase64MaxLength #-} + + +------------------------------------------------------------------------ +-- Encode +------------------------------------------------------------------------ + +{- | Encode a 'ByteString' to RFC 4648 \u00a74 base64 (with @=@ +padding). One allocation of the exact output length, one +SIMD pass. +-} +encodeBase64 :: ByteString -> ByteString +encodeBase64 bs = unsafePerformIO $ do + let !inLen = BS.length bs + !outLen = encodeBase64Length inLen + BSI.create outLen $ \outPtr -> + BSU.unsafeUseAsCStringLen bs $ \(inPtr, _) -> do + _ <- c_base64_encode (castPtr inPtr) (fromIntegral inLen) outPtr + pure () +{-# NOINLINE encodeBase64 #-} + + +------------------------------------------------------------------------ +-- Decode +------------------------------------------------------------------------ + +{- | Decode an RFC 4648 \u00a74 base64 'ByteString'. Returns +'Nothing' if the input is malformed: not a multiple of four +characters, contains an out-of-alphabet byte, or uses padding +in an invalid position. + +Whitespace is /not/ tolerated; strip it before calling if your +input is line-wrapped (PEM, MIME, etc.). +-} +decodeBase64 :: ByteString -> Maybe ByteString +decodeBase64 bs + | BS.null bs = Just BS.empty + | (BS.length bs .&. 3) /= 0 = Nothing + | otherwise = unsafePerformIO $ do + let !inLen = BS.length bs + !maxOut = decodeBase64MaxLength inLen + (inFp, inOff, _) = BSI.toForeignPtr bs + outFp <- BSI.mallocByteString maxOut + actual <- withForeignPtr inFp $ \inBase -> + withForeignPtr outFp $ \outPtr -> do + n <- + c_base64_decode + (inBase `plusPtr` inOff) + (fromIntegral inLen) + outPtr + pure (fromIntegral n :: Int) + if actual < 0 + then pure Nothing + else pure (Just (BSI.fromForeignPtr outFp 0 actual)) +{-# NOINLINE decodeBase64 #-} diff --git a/src/Wireform/Builder.hs b/src/Wireform/Builder.hs --- a/src/Wireform/Builder.hs +++ b/src/Wireform/Builder.hs @@ -55,34 +55,6 @@ @ -} module Wireform.Builder ( - -- * Builder type - -- | A chunk-based byte builder. Combine fragments with @('<>')@. - -- Run with 'toStrictByteString', 'toLazyByteString', or 'hPutBuilder'. - Builder, - - -- * Running builders - -- | Convert a 'Builder' into output bytes or write directly to a handle. - toStrictByteString, - toLazyByteString, - hPutBuilder, - hPutBuilderLen, - hPutBuilderWith, - - -- * Performance tuning - -- | 'rebuild' resets GHC's inlining budget for a builder, which can - -- help when combining many small builders causes excessive code bloat. - rebuild, - - -- * Bounded / fixed primitives - primBounded, - primFixed, - - -- * ByteString to Builder - byteString, - byteStringInsert, - byteStringCopy, - byteStringThreshold, - -- * Single byte word8, int8, @@ -124,10 +96,12 @@ string7, -- * Builder internals (advanced) - -- | 'StreamSink' and 'withStreamTransform' allow interposing a - -- streaming transformation (e.g. compression) between the builder - -- and its output destination. See "Wireform.Builder.FastBuilder" for - -- the full internal API. + + {- | 'StreamSink' and 'withStreamTransform' allow interposing a + streaming transformation (e.g. compression) between the builder + and its output destination. See "Wireform.Builder.FastBuilder" for + the full internal API. + -} module Wireform.Builder.FastBuilder, ) where diff --git a/src/Wireform/Builder/FastBuilder.hs b/src/Wireform/Builder/FastBuilder.hs --- a/src/Wireform/Builder/FastBuilder.hs +++ b/src/Wireform/Builder/FastBuilder.hs @@ -46,6 +46,10 @@ withStreamTransform, runBuilderStreaming, + -- * Ring-direct sink (zero-copy builder → send ring) + RingSinkState (..), + finalizeRingSink, + -- * Basic builders primBounded, primFixed, @@ -69,6 +73,7 @@ import Control.Concurrent.MVar import Control.Exception qualified as E import Control.Monad +import Data.Bits ((.&.)) import Data.ByteString qualified as S import Data.ByteString.Builder.Extra qualified as X import Data.ByteString.Builder.Prim qualified as P @@ -141,15 +146,18 @@ {-# INLINE mempty #-} mappend = (<>) {-# INLINE mappend #-} + + -- Don't be tempted to write `mconcat = fold` here — the @Foldable []@ -- instance defines `fold = mconcat` as an optimisation, so going -- through `fold` would bottom out in a value-level cycle and trip -- GHC's blackhole detection at runtime (<>). mconcat = foldr mappend mempty {-# INLINE mconcat #-} - {-# HLINT ignore "Use fold" #-} +{- HLINT ignore "Use fold" -} + -- | 'fromString' = 'stringUtf8' instance IsString Builder where fromString = builderFromString @@ -162,26 +170,38 @@ DynamicSink !(IORef DynamicSink) | -- | Bytes are accumulated in a contiguous buffer. GrowingBuffer !(IORef (ForeignPtr Word8)) - | -- | Bytes are first accumulated in the 'Queue', then flushed to the - -- 'IO.Handle'. + | {- | Bytes are first accumulated in the 'Queue', then flushed to the + 'IO.Handle'. + -} HandleSink !IO.Handle !Int {-next buffer size-} !(IORef Queue) - | -- | The buffer has a known-exact size. No 'IORef', no growth. - -- Used by 'toStrictByteStringExact' when the output size is - -- known in advance. If the builder overflows, behaviour is - -- undefined (but in practice 'getBytes_' will 'error'). + | {- | The buffer has a known-exact size. No 'IORef', no growth. + Used by 'toStrictByteStringExact' when the output size is + known in advance. If the builder overflows, behaviour is + undefined (but in practice 'getBytes_' will 'error'). + -} FixedBuffer - | -- | Bytes are accumulated in a buffer, then fed to a streaming - -- transform (compression, encryption, etc.) when the buffer fills. + | {- | Bytes are accumulated in a buffer, then fed to a streaming + transform (compression, encryption, etc.) when the buffer fills. + -} StreamingSink !Int {-next buffer size-} !(IORef StreamQueue) + | {- | Bytes are written directly into a double-mapped magic ring + buffer (used by 'Wireform.Transport.Send.sendBuilderDirect'). + When the builder overflows, the already-written bytes are + published to the ring consumer and the builder continues + writing into newly-available ring space — no intermediate + 'S.ByteString' allocation or memcpy. + -} + RingSink !(IORef RingSinkState) -- | Variable-destination cases. data DynamicSink = -- | Bytes are sent to another thread. ThreadedSink !(MVar Request) !(MVar Response) - | -- | Bytes are accumulated in a contiguous buffer until the - -- size limit is reached. After that, the destination switches - -- to a 'ThreadedSink'. + | {- | Bytes are accumulated in a contiguous buffer until the + size limit is reached. After that, the destination switches + to a 'ThreadedSink'. + -} BoundedGrowingBuffer {-# UNPACK #-} !(ForeignPtr Word8) !Int {-bound-} @@ -195,8 +215,9 @@ -} data StreamQueue = StreamQueue { sqBase :: {-# UNPACK #-} !(Ptr Word8) - -- ^ Buffer start, extracted once at init. The hot-path flush - -- reads this instead of dereferencing the 'ForeignPtr'. + {- ^ Buffer start, extracted once at init. The hot-path flush + reads this instead of dereferencing the 'ForeignPtr'. + -} , sqFptr :: !(ForeignPtr Word8) -- ^ Kept alive for GC. Never dereferenced on the hot path. , sqCap :: !Int @@ -208,6 +229,70 @@ } +{- | Mutable state for the ring-direct sink. + +The builder writes directly into the ring's double-mapped memory. +On overflow the handler first checks whether the ring already has +room (deferred publish). Only when the ring is actually full does +it publish ('rsPublishHead') and wait ('rsEnsureRoom'). This +batches writes so that a single 'sendBuilderDirect' call produces +at most one 'sendPublishHead' per ring-size worth of data, matching +the syscall count of the old materialize-then-copy path. + +The callbacks are closed over the 'SendTransport' at construction +time (in "Wireform.Transport.Send"), so 'FastBuilder' does not +depend on the transport types. +-} +data RingSinkState = RingSinkState + { rsHead :: {-# UNPACK #-} !Word64 + {- ^ Logical head: start of the current buffer region. + Tracks the builder's logical write position; may be ahead + of the last published head. + -} + , rsPublished :: {-# UNPACK #-} !Word64 + {- ^ Last position passed to 'rsPublishHead'. Everything + before this has been handed off to the ring consumer. + -} + , rsBase :: {-# UNPACK #-} !(Ptr Word8) + -- ^ Ring base address (first mapping). + , rsMask :: {-# UNPACK #-} !Int + -- ^ @ringSize - 1@. + , rsRingSize :: {-# UNPACK #-} !Int + -- ^ Physical ring size (N, not 2N). + , rsChunkSize :: {-# UNPACK #-} !Int + -- ^ Default chunk offered to the builder on overflow. + , rsPublishHead :: !(Word64 -> IO ()) + -- ^ Publish the head position to the ring consumer. + , rsEnsureRoom :: !(Word64 -> IO ()) + {- ^ Block until the ring has room for head to advance to this + position. Throws on transport failure / close. + -} + , rsLoadTail :: !(IO Word64) + {- ^ Read the consumer's current tail position. Used to check + whether the ring has room without publishing first. + -} + } + + +{- | Publish the final bytes written by a builder into the ring. + +After 'runBuilder' returns a final cursor, call this to publish +the last unpublished bytes. Returns the final head position. +-} +finalizeRingSink :: IORef RingSinkState -> Ptr Word8 -> IO Word64 +finalizeRingSink rsRef finalCur = do + rs <- readIORef rsRef + let !off = fromIntegral (rsHead rs) .&. rsMask rs + !startPtr = rsBase rs `plusPtr` off + !written = finalCur `minusPtr` startPtr + !finalHead = rsHead rs + fromIntegral written + -- Publish everything from the last-published position to + -- the builder's final cursor in a single call. + when (finalHead > rsPublished rs) $ + rsPublishHead rs finalHead + pure finalHead + + {- | A streaming byte transform — middleware that processes chunks as a 'Builder' produces them. Used for compression, encryption, checksumming, or any incremental byte transform. @@ -237,23 +322,25 @@ -} data StreamSink = StreamSink { ssFeedRaw :: !(Ptr Word8 -> Int -> IO ()) - -- ^ Feed a raw pointer region to the transform. The pointer - -- points into the builder's current buffer. The region is - -- valid for the duration of this call only — implementations - -- must consume or copy the data before returning. - -- - -- @ssFeedRaw ptr len@ — @ptr@ is the start of the filled - -- region, @len@ is the number of valid bytes. + {- ^ Feed a raw pointer region to the transform. The pointer + points into the builder's current buffer. The region is + valid for the duration of this call only — implementations + must consume or copy the data before returning. + + @ssFeedRaw ptr len@ — @ptr@ is the start of the filled + region, @len@ is the number of valid bytes. + -} , ssFinish :: !(IO Builder) - -- ^ Flush the transform and return a 'Builder' that emits the - -- transformed output. This builder is run directly inside the - -- outer builder's buffer — no intermediate 'S.ByteString' is - -- materialised for the output either. - -- - -- Typical implementation: collect compressed chunks in an - -- 'IORef', then return @mconcat (map byteStringCopy chunks)@. - -- Or for a single contiguous output buffer, return - -- @unsafeCStringLen (outPtr, outLen)@. + {- ^ Flush the transform and return a 'Builder' that emits the + transformed output. This builder is run directly inside the + outer builder's buffer — no intermediate 'S.ByteString' is + materialised for the output either. + + Typical implementation: collect compressed chunks in an + 'IORef', then return @mconcat (map byteStringCopy chunks)@. + Or for a single contiguous output buffer, return + @unsafeCStringLen (outPtr, outLen)@. + -} } @@ -280,12 +367,14 @@ Error E.SomeException | -- | The builder thread has completed. Done !(Ptr Word8) - | -- | The builder thread has finished generating one chunk, - -- and waits for another request with the specified minimum size. + | {- | The builder thread has finished generating one chunk, + and waits for another request with the specified minimum size. + -} MoreBuffer !(Ptr Word8) !Int - | -- | The builder thread has partially filled the current chunk, - -- and wants to emit the bytestring to be included in the final - -- output. + | {- | The builder thread has partially filled the current chunk, + and wants to emit the bytestring to be included in the final + output. + -} InsertByteString !(Ptr Word8) !S.ByteString deriving (Show) @@ -416,8 +505,8 @@ -- 'BuilderState' is an unlifted tuple, so the @(\\s -> w1 (w0 s))@ / -- @(\\s -> s)@ continuations here can't go through @(.)@ or 'id' — -- those have lifted-kind type parameters and don't apply. -{-# HLINT ignore "Avoid lambda" #-} -{-# HLINT ignore "Use id" #-} +{- HLINT ignore "Avoid lambda" -} +{- HLINT ignore "Use id" -} instance Sem.Semigroup Write where Write s0 w0 <> Write s1 w1 = Write (s0 + s1) (\s -> w1 (w0 s)) @@ -956,6 +1045,50 @@ -- Reset cur to base — buffer is free after flush sq <- io $ readIORef sqRef setCur (sqBase sq) + RingSink rsRef -> do + cur <- getCur + rs <- io $ readIORef rsRef + let !off = fromIntegral (rsHead rs) .&. rsMask rs + !startPtr = rsBase rs `plusPtr` off + !written = cur `minusPtr` startPtr + !newHead = rsHead rs + fromIntegral written + !bsLen = S.length bstr + !ringSz = rsRingSize rs + -- Publish accumulated builder bytes before the BS copy. + when (written > 0) $ io $ do + when (newHead > rsPublished rs) $ + rsPublishHead rs newHead + -- Copy the BS into the ring, chunking at ringSize if + -- the BS is larger than the ring can hold contiguously. + io $ S.unsafeUseAsCString bstr $ \src -> do + let loop !srcOff !hd !remaining + | remaining <= 0 = pure hd + | otherwise = do + let !chunk = min remaining ringSz + !wh = hd + fromIntegral chunk + rsEnsureRoom rs wh + let !dstOff = fromIntegral hd .&. rsMask rs + !dst = rsBase rs `plusPtr` dstOff + copyBytes dst (castPtr src `plusPtr` srcOff) chunk + rsPublishHead rs wh + loop (srcOff + chunk) wh (remaining - chunk) + afterBs <- loop 0 newHead bsLen + let !pub = afterBs + writeIORef + rsRef + rs + { rsHead = afterBs + , rsPublished = pub + } + -- Set up a fresh region for subsequent builder writes. + rs' <- io $ readIORef rsRef + let !contChunk = min ringSz (rsChunkSize rs') + !contWant = rsHead rs' + fromIntegral contChunk + io $ rsEnsureRoom rs' contWant + let !contOff = fromIntegral (rsHead rs') .&. rsMask rs' + !contPtr = rsBase rs' `plusPtr` contOff + setCur contPtr + setEnd (contPtr `plusPtr` contChunk) {-# NOINLINE byteStringInsert_ #-} @@ -1049,6 +1182,55 @@ } setCur newBase setEnd (newBase `plusPtr` I# n) + RingSink rsRef -> do + cur <- getCur + rs <- io $ readIORef rsRef + -- The ring provides at most ringSize contiguous bytes. + -- ensureBytes requests larger than that are a programming + -- error — use byteStringInsert (triggered automatically + -- by byteString / byteStringThreshold) for large payloads. + when (I# n > rsRingSize rs) $ + io $ + error $ + "Wireform.Builder: RingSink getBytes request (" + ++ show (I# n) + ++ ") exceeds ring size (" + ++ show (rsRingSize rs) + ++ "); use byteString/byteStringInsert for large payloads" + let !off = fromIntegral (rsHead rs) .&. rsMask rs + !startPtr = rsBase rs `plusPtr` off + !written = cur `minusPtr` startPtr + !newHead = rsHead rs + fromIntegral written + -- Try to extend without publishing: the ring may already + -- have room beyond what we initially reserved. + let !chunkSz = + min + (rsRingSize rs) + (max (I# n) (rsChunkSize rs)) + !wantHead = newHead + fromIntegral chunkSz + tl <- io $ rsLoadTail rs + let !ringCap = fromIntegral (rsRingSize rs) :: Word64 + published' <- + if wantHead - tl <= ringCap + then + io $ pure (rsPublished rs) + else do + io $ do + when (newHead > rsPublished rs) $ + rsPublishHead rs newHead + rsEnsureRoom rs wantHead + io $ pure newHead + let !newOff = fromIntegral newHead .&. rsMask rs + !newPtr = rsBase rs `plusPtr` newOff + io $ + writeIORef + rsRef + rs + { rsHead = newHead + , rsPublished = published' + } + setCur newPtr + setEnd (newPtr `plusPtr` chunkSz) {-# NOINLINE getBytes_ #-} diff --git a/src/Wireform/Encode/Direct.hs b/src/Wireform/Encode/Direct.hs --- a/src/Wireform/Encode/Direct.hs +++ b/src/Wireform/Encode/Direct.hs @@ -1,69 +1,74 @@ {-# LANGUAGE BangPatterns #-} {-# OPTIONS_GHC -Wno-unused-top-binds #-} --- | Direct-write encoder: bypasses ByteString.Builder entirely. --- --- Uses offset-based writes to a pre-allocated buffer. The offset is --- an unboxed Int threaded through all operations, avoiding any heap --- allocation for the write context. --- --- For field writes (tag + value), uses single C FFI calls that --- handle the entire field in one shot, minimizing Haskell\/C --- call overhead. --- --- @ --- import Proto.Encode.Direct (directEncode, dWord8, dVarint) --- --- let bytes = directEncode 3 (\\p off -> do --- off1 <- dWord8 p off 0x42 --- dVarint p off1 12345) --- @ -module Wireform.Encode.Direct - ( -- * Core - directEncode - -- * Offset-based write primitives (return new offset) - , dWord8 - , dVarint - , dWord32LE - , dWord64LE - , dFloatLE - , dDoubleLE - , dBytes - , dText +{- | Direct-write encoder: bypasses ByteString.Builder entirely. - -- * Field-level writes (tag + value in one call) - , dVarintField - , dBoolField - , dStringField - , dBytesField - , dFixed32Field - , dFixed64Field - , dFloatField - , dDoubleField - ) where +Uses offset-based writes to a pre-allocated buffer. The offset is +an unboxed Int threaded through all operations, avoiding any heap +allocation for the write context. -import Data.Bits ((.&.), (.|.), shiftR) +For field writes (tag + value), uses single C FFI calls that +handle the entire field in one shot, minimizing Haskell\/C +call overhead. + +@ +import Proto.Encode.Direct (directEncode, dWord8, dVarint) + +let bytes = directEncode 3 (\\p off -> do + off1 <- dWord8 p off 0x42 + dVarint p off1 12345) +@ +-} +module Wireform.Encode.Direct ( + -- * Core + directEncode, + + -- * Offset-based write primitives (return new offset) + dWord8, + dVarint, + dWord32LE, + dWord64LE, + dFloatLE, + dDoubleLE, + dBytes, + dText, + + -- * Field-level writes (tag + value in one call) + dVarintField, + dBoolField, + dStringField, + dBytesField, + dFixed32Field, + dFixed64Field, + dFloatField, + dDoubleField, +) where + +import Data.Bits (shiftR, (.&.), (.|.)) import Data.ByteString (ByteString) -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BSI +import Data.ByteString qualified as BS +import Data.ByteString.Internal qualified as BSI import Data.Text (Text) -import qualified Data.Text.Encoding as TE -import Data.Word (Word8, Word32, Word64) +import Data.Text.Encoding qualified as TE +import Data.Word (Word32, Word64, Word8) import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Marshal.Utils (copyBytes) import Foreign.Ptr (Ptr, plusPtr) import Foreign.Storable (pokeByteOff) -import GHC.Float (castFloatToWord32, castDoubleToWord64) +import GHC.Float (castDoubleToWord64, castFloatToWord32) +import Wireform.FFI (encodeBoolFieldC, encodeLengthDelimitedC, encodeVarintFieldC) -import Wireform.FFI (encodeVarintFieldC, encodeBoolFieldC, encodeLengthDelimitedC) --- | Allocate a buffer of exactly @sz@ bytes, run the writer, return ByteString. --- The writer takes (Ptr Word8, offset) and returns the final offset. +{- | Allocate a buffer of exactly @sz@ bytes, run the writer, return ByteString. +The writer takes (Ptr Word8, offset) and returns the final offset. +-} directEncode :: Int -> (Ptr Word8 -> Int -> IO Int) -> ByteString directEncode !sz writer = BSI.unsafeCreate sz $ \ptr -> do _ <- writer ptr 0 pure () {-# INLINE directEncode #-} + -- Offset-based primitives: take (Ptr, offset), return new offset. dWord8 :: Ptr Word8 -> Int -> Word8 -> IO Int @@ -72,26 +77,31 @@ pure $! off + 1 {-# INLINE dWord8 #-} + dWord32LE :: Ptr Word8 -> Int -> Word32 -> IO Int dWord32LE !p !off !v = do pokeByteOff p off v pure $! off + 4 {-# INLINE dWord32LE #-} + dWord64LE :: Ptr Word8 -> Int -> Word64 -> IO Int dWord64LE !p !off !v = do pokeByteOff p off v pure $! off + 8 {-# INLINE dWord64LE #-} + dFloatLE :: Ptr Word8 -> Int -> Float -> IO Int dFloatLE !p !off !v = dWord32LE p off (castFloatToWord32 v) {-# INLINE dFloatLE #-} + dDoubleLE :: Ptr Word8 -> Int -> Double -> IO Int dDoubleLE !p !off !v = dWord64LE p off (castDoubleToWord64 v) {-# INLINE dDoubleLE #-} + -- | Write a varint at offset, return new offset. dVarint :: Ptr Word8 -> Int -> Word64 -> IO Int dVarint !p !off !n @@ -112,14 +122,16 @@ pure $! off + n' {-# INLINE dVarint #-} + -- | Write raw bytes at offset (no length prefix). Returns new offset. dBytes :: Ptr Word8 -> Int -> ByteString -> IO Int dBytes !p !off (BSI.BS fp len) = do withForeignPtr fp $ \src -> - BSI.memcpy (p `plusPtr` off) src len + copyBytes (p `plusPtr` off) src len pure $! off + len {-# INLINE dBytes #-} + -- | Write text (length-prefixed UTF-8 bytes) at offset. dText :: Ptr Word8 -> Int -> Text -> IO Int dText !p !off !t = do @@ -129,6 +141,7 @@ dBytes p off1 bs {-# INLINE dText #-} + -- ============================================================ -- Field-level writes: tag + value in minimal calls -- ============================================================ @@ -140,6 +153,7 @@ pure $! off + n {-# INLINE dVarintField #-} + -- | Bool field: tag_byte + 0/1. Always 2 bytes. dBoolField :: Ptr Word8 -> Int -> Word8 -> Bool -> IO Int dBoolField !p !off !tag !val = do @@ -147,6 +161,7 @@ pure $! off + 2 {-# INLINE dBoolField #-} + -- | String field: tag + varint(len) + UTF-8 bytes. Single C call. dStringField :: Ptr Word8 -> Int -> Word8 -> Text -> IO Int dStringField !p !off !tag !t = do @@ -156,6 +171,7 @@ pure $! off + n {-# INLINE dStringField #-} + -- | Bytes field: tag + varint(len) + raw bytes. Single C call. dBytesField :: Ptr Word8 -> Int -> Word8 -> ByteString -> IO Int dBytesField !p !off !tag (BSI.BS fp len) = do @@ -164,6 +180,7 @@ pure $! off + n {-# INLINE dBytesField #-} + -- | Fixed32 field: tag + 4 LE bytes. Always 5 bytes. dFixed32Field :: Ptr Word8 -> Int -> Word8 -> Word32 -> IO Int dFixed32Field !p !off !tag !val = do @@ -172,6 +189,7 @@ pure $! off + 5 {-# INLINE dFixed32Field #-} + -- | Fixed64 field: tag + 8 LE bytes. Always 9 bytes. dFixed64Field :: Ptr Word8 -> Int -> Word8 -> Word64 -> IO Int dFixed64Field !p !off !tag !val = do @@ -180,10 +198,12 @@ pure $! off + 9 {-# INLINE dFixed64Field #-} + -- | Float field: tag + 4 LE bytes. dFloatField :: Ptr Word8 -> Int -> Word8 -> Float -> IO Int dFloatField !p !off !tag !val = dFixed32Field p off tag (castFloatToWord32 val) {-# INLINE dFloatField #-} + -- | Double field: tag + 8 LE bytes. dDoubleField :: Ptr Word8 -> Int -> Word8 -> Double -> IO Int diff --git a/src/Wireform/FFI.hs b/src/Wireform/FFI.hs --- a/src/Wireform/FFI.hs +++ b/src/Wireform/FFI.hs @@ -50,6 +50,15 @@ findByte, findByteBS, + -- * SIMD 4-byte repeating-key XOR (WebSocket masking, + + -- RFC 6455 sec 5.3) + xorRepeatingKey, + xorRepeatingKeyBS, + + -- * Thread-local xoshiro256++ PRNG + fastRandomWord64, + -- * SIMD ASCII check (general purpose) isAscii, isAsciiBS, @@ -312,6 +321,75 @@ ------------------------------------------------------------------------ +-- 4-byte repeating-key XOR (WebSocket masking) +------------------------------------------------------------------------ + +foreign import ccall unsafe "hs_ws_mask" + c_ws_mask :: Ptr Word8 -> CInt -> Word32 -> IO () + + +{- | In-place XOR of the 4-byte repeating @key@ over @buf[0..len)@. + +The key is interpreted in network byte order: byte 0 of the +key sits in the high byte of @key@. This matches RFC 6455 sec +5.3's masking-key layout. + +Used by 'wireform-websocket' for frame masking; exposed +generically because the same primitive shows up in any protocol +that masks with a periodic 4-byte key. Implemented in +@cbits\/fast_scan.c@ as an SSE2 (via simde) loop processing +16 bytes per iteration with a tiled mask vector. +-} +xorRepeatingKey :: Ptr Word8 -> Int -> Word32 -> IO () +xorRepeatingKey ptr len key = c_ws_mask ptr (fromIntegral len) key +{-# INLINE xorRepeatingKey #-} + + +{- | 'ByteString' variant: XOR the 4-byte repeating key in-place +over the bytes the 'ByteString' references. The 'ByteString' +itself is shared with the caller, so the mutation is visible to +everyone else holding the same 'ByteString'. Callers must own +the backing memory exclusively at this point. +-} +xorRepeatingKeyBS :: ByteString -> Word32 -> IO () +xorRepeatingKeyBS bs key = BSU.unsafeUseAsCStringLen bs $ \(ptr, len) -> + c_ws_mask (castPtr ptr) (fromIntegral len) key +{-# INLINE xorRepeatingKeyBS #-} + + +------------------------------------------------------------------------ +-- Thread-local xoshiro256++ PRNG +------------------------------------------------------------------------ + +foreign import ccall unsafe "hs_xoshiro256pp_next" + c_xoshiro256pp_next :: IO Word64 + + +{- | Pull a non-cryptographically-random 'Word64' from the calling +OS thread's xoshiro256++ generator. + +State is per-OS-thread (@__thread@), seeded from @getrandom(2)@ +(@arc4random_buf@ on BSDs, @\/dev\/urandom@ everywhere else) on +first use. Once seeded, each call is a handful of register-only +arithmetic ops — typically ~1 ns including the FFI boundary, +versus ~50 ns for the global @splitmix@ generator that takes an +@MVar@ on every call. + +Caveat: because Haskell threads are multiplexed across OS +threads, the per-Haskell-thread sequence is not reproducible +(an HS thread may resume on a different capability and draw +from a different RNG stream). This is the right trade-off for +the non-deterministic-randomness needs that motivated the helper +(WebSocket frame masking, retry jitter, …); for reproducible +streams use a 'System.Random.Stateful' generator the caller +owns. +-} +fastRandomWord64 :: IO Word64 +fastRandomWord64 = c_xoshiro256pp_next +{-# INLINE fastRandomWord64 #-} + + +------------------------------------------------------------------------ -- SIMD ASCII check ------------------------------------------------------------------------ @@ -387,14 +465,14 @@ let !escPos = findJsonEscapeBS bs pos !safeLen = escPos - pos in ( if safeLen > 0 - then B.byteString (BSU.unsafeTake safeLen (BSU.unsafeDrop pos bs)) - else mempty + then B.byteString (BSU.unsafeTake safeLen (BSU.unsafeDrop pos bs)) + else mempty ) - <> if escPos >= len - then mempty - else - let !b = BSU.unsafeIndex bs escPos - in escByte b <> go (escPos + 1) + <> if escPos >= len + then mempty + else + let !b = BSU.unsafeIndex bs escPos + in escByte b <> go (escPos + 1) escByte :: Word8 -> B.Builder escByte 0x22 = B.byteString "\\\"" -- " diff --git a/src/Wireform/Hash.hs b/src/Wireform/Hash.hs --- a/src/Wireform/Hash.hs +++ b/src/Wireform/Hash.hs @@ -1,49 +1,54 @@ {-# LANGUAGE BangPatterns #-} --- | Shared C/SIMDe-backed hash and Roaring-bitmap kernels used across the --- wireform format packages. --- --- - 'murmur3_32' / 'bucketLong' — byte-compatible with Apache Iceberg's --- @org.apache.iceberg.util.BucketUtil@ (used both by the Iceberg --- @bucket[N]@ partition transform and by the Parquet bloom filter's --- feeder hash variants in some integrations). --- - 'xxh64' — XXH64 with a caller-supplied seed, byte-exact against the --- upstream xxHash 0.6.x reference. This is the hash function Parquet's --- split-block bloom filter uses (seed @0@) and that Iceberg manifest --- tooling uses for content-addressing. --- - 'roaring*' — portable Roaring 32-bit container kernels used by Iceberg --- V3 deletion vectors and (potentially) Parquet column-index null pages. --- --- The implementations live in @cbits/wireform_hash_simd.c@. All entry --- points are 'unsafePerformIO' / 'unsafeDupablePerformIO' wrapped so --- callers don't think about IO. -module Wireform.Hash - ( -- * Murmur3 32-bit (Iceberg @BucketUtil@) - murmur3_32 - , bucketLong - , bucketBytes - -- * XXH64 - , xxh64 - -- * Roaring 32-bit container - , roaringDecodeArray - , roaringDecodeBitset - , roaringContains - , roaringEncodeArray - , roaringEncodeBitset - , RoaringContainerKind(..) - ) where +{- | Shared C/SIMDe-backed hash and Roaring-bitmap kernels used across the +wireform format packages. + +- 'murmur3_32' / 'bucketLong' — byte-compatible with Apache Iceberg's + @org.apache.iceberg.util.BucketUtil@ (used both by the Iceberg + @bucket[N]@ partition transform and by the Parquet bloom filter's + feeder hash variants in some integrations). +- 'xxh64' — XXH64 with a caller-supplied seed, byte-exact against the + upstream xxHash 0.6.x reference. This is the hash function Parquet's + split-block bloom filter uses (seed @0@) and that Iceberg manifest + tooling uses for content-addressing. +- 'roaring*' — portable Roaring 32-bit container kernels used by Iceberg + V3 deletion vectors and (potentially) Parquet column-index null pages. + +The implementations live in @cbits/wireform_hash_simd.c@. All entry +points are 'unsafePerformIO' / 'unsafeDupablePerformIO' wrapped so +callers don't think about IO. +-} +module Wireform.Hash ( + -- * Murmur3 32-bit (Iceberg @BucketUtil@) + murmur3_32, + bucketLong, + bucketBytes, + + -- * XXH64 + xxh64, + + -- * Roaring 32-bit container + roaringDecodeArray, + roaringDecodeBitset, + roaringContains, + roaringEncodeArray, + roaringEncodeBitset, + RoaringContainerKind (..), +) where + import Data.Bits ((.&.)) -import qualified Data.ByteString as BS -import qualified Data.ByteString.Internal as BSI +import Data.ByteString qualified as BS +import Data.ByteString.Internal qualified as BSI import Data.ByteString.Unsafe (unsafeUseAsCStringLen) import Data.Int (Int32, Int64) -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Storable.Mutable as VSM +import Data.Vector.Storable qualified as VS +import Data.Vector.Storable.Mutable qualified as VSM import Data.Word (Word16, Word32, Word64, Word8) import Foreign.ForeignPtr (withForeignPtr) import Foreign.Ptr (Ptr, castPtr) import System.IO.Unsafe (unsafeDupablePerformIO, unsafePerformIO) + -- ============================================================ -- C entry points -- ============================================================ @@ -51,27 +56,35 @@ foreign import ccall unsafe "hs_wf_murmur3_32" c_murmur3_32 :: Ptr Word8 -> Int32 -> Int32 + foreign import ccall unsafe "hs_wf_bucket_long" c_bucket_long :: Int64 -> Int32 -> Int32 + foreign import ccall unsafe "hs_wf_xxh64" c_xxh64 :: Ptr Word8 -> Int64 -> Word64 -> Word64 + foreign import ccall unsafe "hs_wf_roaring_decode_array" c_roaring_decode_array :: Ptr Word8 -> Int32 -> Word32 -> Ptr Int32 -> Int32 + foreign import ccall unsafe "hs_wf_roaring_decode_bitset" c_roaring_decode_bitset :: Ptr Word8 -> Word32 -> Ptr Int32 -> Int32 + foreign import ccall unsafe "hs_wf_roaring_contains" c_roaring_contains :: Int32 -> Ptr Word8 -> Int32 -> Word16 -> Int32 + foreign import ccall unsafe "hs_wf_roaring_encode_array" c_roaring_encode_array :: Ptr Word16 -> Int32 -> Ptr Word8 -> Int32 + foreign import ccall unsafe "hs_wf_roaring_encode_bitset" c_roaring_encode_bitset :: Ptr Word16 -> Int32 -> Ptr Word8 -> IO () + -- ============================================================ -- Public wrappers -- ============================================================ @@ -83,27 +96,32 @@ pure $! c_murmur3_32 (castPtr p) (fromIntegral len) {-# INLINE murmur3_32 #-} --- | Iceberg @bucket[N]@ partition transform on a 64-bit signed integer. --- The C kernel inlines the 8-byte little-endian serialisation, the --- murmur3 hash, and the @& Integer.MAX_VALUE % N@ reduction. + +{- | Iceberg @bucket[N]@ partition transform on a 64-bit signed integer. +The C kernel inlines the 8-byte little-endian serialisation, the +murmur3 hash, and the @& Integer.MAX_VALUE % N@ reduction. +-} bucketLong :: Int -> Int64 -> Int bucketLong n v - | n <= 0 = 0 + | n <= 0 = 0 | otherwise = fromIntegral (c_bucket_long v (fromIntegral n)) {-# INLINE bucketLong #-} --- | Iceberg @bucket[N]@ on raw bytes (binary / fixed / uuid / decimal / --- string columns once UTF-8-encoded). Hash is murmur3-32 over the bytes --- with seed 0, then @(hash & Integer.MAX_VALUE) % N@. + +{- | Iceberg @bucket[N]@ on raw bytes (binary / fixed / uuid / decimal / +string columns once UTF-8-encoded). Hash is murmur3-32 over the bytes +with seed 0, then @(hash & Integer.MAX_VALUE) % N@. +-} bucketBytes :: Int -> BS.ByteString -> Int bucketBytes n bs - | n <= 0 = 0 + | n <= 0 = 0 | otherwise = - let !h = murmur3_32 bs + let !h = murmur3_32 bs !hu = fromIntegral h .&. (0x7FFFFFFF :: Word32) - in fromIntegral hu `mod` n + in fromIntegral hu `mod` n {-# INLINE bucketBytes #-} + -- | XXH64 with caller-supplied seed. xxh64 :: Word64 -> BS.ByteString -> Word64 xxh64 seed bs = unsafePerformIO $ @@ -111,6 +129,7 @@ pure $! c_xxh64 (castPtr p) (fromIntegral len) seed {-# INLINE xxh64 #-} + -- ============================================================ -- Roaring containers -- ============================================================ @@ -118,12 +137,17 @@ data RoaringContainerKind = ArrayContainer | BitsetContainer deriving (Show, Eq) --- | Decode a Roaring ARRAY container payload into a vector of 32-bit --- positions, OR'd with @hi << 16@. + +{- | Decode a Roaring ARRAY container payload into a vector of 32-bit +positions, OR'd with @hi << 16@. +-} roaringDecodeArray - :: BS.ByteString -- ^ Raw payload (cardinality * 2 bytes). - -> Int -- ^ Cardinality. - -> Word32 -- ^ High 32 bits to OR onto each position. + :: BS.ByteString + -- ^ Raw payload (cardinality * 2 bytes). + -> Int + -- ^ Cardinality. + -> Word32 + -- ^ High 32 bits to OR onto each position. -> VS.Vector Int32 roaringDecodeArray bs card hi = unsafeDupablePerformIO $ do mv <- VSM.unsafeNew card @@ -134,11 +158,15 @@ VS.unsafeFreeze mv {-# INLINE roaringDecodeArray #-} --- | Decode a Roaring BITSET container payload (exactly 8192 bytes) into --- a vector of set positions, OR'd with @hi << 16@. + +{- | Decode a Roaring BITSET container payload (exactly 8192 bytes) into +a vector of set positions, OR'd with @hi << 16@. +-} roaringDecodeBitset - :: BS.ByteString -- ^ 8192 raw bytes. - -> Word32 -- ^ High 32 bits. + :: BS.ByteString + -- ^ 8192 raw bytes. + -> Word32 + -- ^ High 32 bits. -> VS.Vector Int32 roaringDecodeBitset bs hi = unsafeDupablePerformIO $ do mv <- VSM.unsafeNew 65536 @@ -149,11 +177,13 @@ VS.unsafeFreeze (VSM.unsafeSlice 0 n mv) {-# INLINE roaringDecodeBitset #-} + -- | Test whether a 16-bit low value is present in a Roaring container. roaringContains :: RoaringContainerKind -> BS.ByteString - -> Int -- ^ Cardinality (only used for ARRAY). + -> Int + -- ^ Cardinality (only used for ARRAY). -> Word16 -> Bool roaringContains kind bs card v = unsafePerformIO $ @@ -162,22 +192,27 @@ in pure $! c_roaring_contains k (castPtr p) (fromIntegral card) v /= 0 {-# INLINE roaringContains #-} --- | Encode a sorted vector of 16-bit positions as a Roaring ARRAY --- container payload. + +{- | Encode a sorted vector of 16-bit positions as a Roaring ARRAY +container payload. +-} roaringEncodeArray :: VS.Vector Word16 -> BS.ByteString roaringEncodeArray vs = unsafePerformIO $ do let !len = VS.length vs - !sz = len * 2 + !sz = len * 2 fp <- BSI.mallocByteString sz withForeignPtr fp $ \dst -> VS.unsafeWith vs $ \src -> - do _ <- pure $! c_roaring_encode_array src (fromIntegral len) dst - pure () + do + _ <- pure $! c_roaring_encode_array src (fromIntegral len) dst + pure () pure (BSI.BS fp sz) {-# INLINE roaringEncodeArray #-} --- | Encode a sorted vector of 16-bit positions as a Roaring BITSET --- container payload (exactly 8192 bytes). + +{- | Encode a sorted vector of 16-bit positions as a Roaring BITSET +container payload (exactly 8192 bytes). +-} roaringEncodeBitset :: VS.Vector Word16 -> BS.ByteString roaringEncodeBitset vs = unsafePerformIO $ do fp <- BSI.mallocByteString 8192 diff --git a/src/Wireform/Parser.hs b/src/Wireform/Parser.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser.hs @@ -0,0 +1,1221 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +-- | Fast streaming parser with flatparse-grade inner-loop performance. +module Wireform.Parser ( + -- * Parser type + Parser, + + -- * Byte primitives + anyWord8, + anyWord8_, + anyWord16, + anyWord32, + anyWord64, + anyWord16le, + anyWord32le, + anyWord64le, + anyWord16be, + anyWord32be, + anyWord64be, + anyWord16_, + anyWord32_, + anyWord64_, + anyFloatle, + anyFloatbe, + anyDoublele, + anyDoublebe, + + -- * CPS byte primitives (avoid intermediate boxing) + withAnyWord8, + withAnyWord16, + withAnyWord32, + withAnyWord64, + withSatisfyAscii, + + -- * Byte matching + word8, + byteString, + + -- * ByteString operations + takeBs, + takeBsCopy, + skip, + takeRest, + skipBack, + + -- * Signed integers + anyInt8, + anyInt16, + anyInt32, + anyInt64, + anyInt16le, + anyInt32le, + anyInt64le, + anyInt16be, + anyInt32be, + anyInt64be, + + -- * UTF-8 characters + anyChar, + anyChar_, + satisfy, + satisfy_, + satisfyAscii, + satisfyAscii_, + skipSatisfyAscii, + fusedSatisfy, + + -- * UTF-8 characters (additional) + anyAsciiChar, + skipAnyAsciiChar, + + -- * Character classes + isDigit, + isLatinLetter, + isGreekLetter, + + -- * ASCII numeric + anyAsciiDecimalWord, + anyAsciiDecimalInt, + anyAsciiDecimalInteger, + anyAsciiHexWord, + anyAsciiHexInt, + + -- * Control flow + (<|>), + empty, + branch, + lookahead, + fails, + try, + optional, + optional_, + many_, + some_, + many, + some, + skipMany, + skipSome, + withOption, + chainl, + chainr, + notFollowedBy, + isolate, + + -- * Error handling + err, + cut, + cutting, + ensureNOrEof, + withError, + + -- * Position and span + Pos (..), + Span (..), + getPos, + withSpan, + byteStringOf, + spanToByteString, + withByteString, + inSpan, + subPos, + + -- * Marks + Mark, + mark, + restore, + release, + + -- * Low-level + ensureN#, + checkpoint, + eof, + atEnd, + remaining, + + -- * Re-exports + module Wireform.Parser.Error, +) where + +import Control.Monad (when) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Internal qualified as BSI +import Data.Int +import Data.Word +import Foreign.Marshal.Utils (copyBytes) +import GHC.Exts +import GHC.Float (castWord32ToFloat, castWord64ToDouble) +import GHC.ForeignPtr (ForeignPtr (..), mallocPlainForeignPtrBytes, unsafeWithForeignPtr) +import GHC.IO (IO (..)) +import GHC.Word (Word16 (..), Word32 (..), Word64 (..), Word8 (..)) +import System.IO.Unsafe (unsafeDupablePerformIO) +import Wireform.Parser.Error +import Wireform.Parser.Internal +import Wireform.Parser.Mark +import Wireform.Parser.Position + + +------------------------------------------------------------------------ +-- Internal: withEnsure# (bounds check then run continuation) +------------------------------------------------------------------------ + +{- | Ensure @n#@ bytes available, then run the body. + +Fast path (1 register comparison, no memory access). +Semi-fast path: re-read 'peEndPtr' to catch stale @eob@ after a +prior streaming resume; avoids the NOINLINE 'ensureNSlow' call. +Slow path: genuine suspension via 'onEnsureFail'. +-} +withEnsure# :: forall m e a. ParserMode m => Int# -> Parser m e a -> Parser m e a +withEnsure# n# (Parser p) = Parser \env eob s st -> + case n# <=# minusAddr# eob s of + 1# -> p env eob s st + _ -> case readEnd# env st of + (# st', eob' #) -> case n# <=# minusAddr# eob' s of + 1# -> p env eob' s st' + _ -> case onEnsureFail @m env eob' s n# st' of + (# st'', OK# _ s'' #) -> case readEnd# env st'' of + (# st''', eob'' #) -> p env eob'' s'' st''' + (# st'', x #) -> (# st'', unsafeCoerce# x #) +{-# INLINE withEnsure# #-} + + +------------------------------------------------------------------------ +-- CPS byte primitives (match flatparse's withAnyWord8 etc) +------------------------------------------------------------------------ + +{- | Read one byte and pass it to the continuation. Most efficient +form — avoids allocating an intermediate @Word8@ on the heap. +-} +withAnyWord8 :: ParserMode m => (Word8 -> Parser m e r) -> Parser m e r +withAnyWord8 p = withEnsure# 1# $ Parser \env eob s st -> + case indexWord8OffAddr# s 0# of + w# -> runParser# (p (W8# w#)) env eob (plusAddr# s 1#) st +{-# INLINE withAnyWord8 #-} + + +withAnyWord16 :: ParserMode m => (Word16 -> Parser m e r) -> Parser m e r +withAnyWord16 p = withEnsure# 2# $ Parser \env eob s st -> + case indexWord16OffAddr# s 0# of + w# -> runParser# (p (W16# w#)) env eob (plusAddr# s 2#) st +{-# INLINE withAnyWord16 #-} + + +withAnyWord32 :: ParserMode m => (Word32 -> Parser m e r) -> Parser m e r +withAnyWord32 p = withEnsure# 4# $ Parser \env eob s st -> + case indexWord32OffAddr# s 0# of + w# -> runParser# (p (W32# w#)) env eob (plusAddr# s 4#) st +{-# INLINE withAnyWord32 #-} + + +withAnyWord64 :: ParserMode m => (Word64 -> Parser m e r) -> Parser m e r +withAnyWord64 p = withEnsure# 8# $ Parser \env eob s st -> + case indexWord64OffAddr# s 0# of + w# -> runParser# (p (W64# w#)) env eob (plusAddr# s 8#) st +{-# INLINE withAnyWord64 #-} + + +------------------------------------------------------------------------ +-- Non-CPS byte primitives +------------------------------------------------------------------------ + +anyWord8 :: ParserMode m => Parser m e Word8 +anyWord8 = withAnyWord8 pure +{-# INLINE anyWord8 #-} + + +anyWord8_ :: ParserMode m => Parser m e () +anyWord8_ = withEnsure# 1# $ Parser \_ _ s st -> + (# st, OK# () (plusAddr# s 1#) #) +{-# INLINE anyWord8_ #-} + + +anyWord16 :: ParserMode m => Parser m e Word16 +anyWord16 = withAnyWord16 pure +{-# INLINE anyWord16 #-} + + +anyWord32 :: ParserMode m => Parser m e Word32 +anyWord32 = withAnyWord32 pure +{-# INLINE anyWord32 #-} + + +anyWord64 :: ParserMode m => Parser m e Word64 +anyWord64 = withAnyWord64 pure +{-# INLINE anyWord64 #-} + + +-- Skip variants +anyWord16_ :: ParserMode m => Parser m e () +anyWord16_ = withEnsure# 2# $ Parser \_ _ s st -> (# st, OK# () (plusAddr# s 2#) #) +{-# INLINE anyWord16_ #-} + + +anyWord32_ :: ParserMode m => Parser m e () +anyWord32_ = withEnsure# 4# $ Parser \_ _ s st -> (# st, OK# () (plusAddr# s 4#) #) +{-# INLINE anyWord32_ #-} + + +anyWord64_ :: ParserMode m => Parser m e () +anyWord64_ = withEnsure# 8# $ Parser \_ _ s st -> (# st, OK# () (plusAddr# s 8#) #) +{-# INLINE anyWord64_ #-} + +------------------------------------------------------------------------ +-- Endianness variants +------------------------------------------------------------------------ + +#if defined(WORDS_BIGENDIAN) +anyWord16le = withAnyWord16 (pure . byteSwap16) +anyWord32le = withAnyWord32 (pure . byteSwap32) +anyWord64le = withAnyWord64 (pure . byteSwap64) +anyWord16be = anyWord16 +anyWord32be = anyWord32 +anyWord64be = anyWord64 +#else +anyWord16le :: ParserMode m => Parser m e Word16 +anyWord16le = anyWord16 +{-# INLINE anyWord16le #-} +anyWord32le :: ParserMode m => Parser m e Word32 +anyWord32le = anyWord32 +{-# INLINE anyWord32le #-} +anyWord64le :: ParserMode m => Parser m e Word64 +anyWord64le = anyWord64 +{-# INLINE anyWord64le #-} +anyWord16be :: ParserMode m => Parser m e Word16 +anyWord16be = withAnyWord16 (pure . byteSwap16) +{-# INLINE anyWord16be #-} +anyWord32be :: ParserMode m => Parser m e Word32 +anyWord32be = withAnyWord32 (pure . byteSwap32) +{-# INLINE anyWord32be #-} +anyWord64be :: ParserMode m => Parser m e Word64 +anyWord64be = withAnyWord64 (pure . byteSwap64) +{-# INLINE anyWord64be #-} +#endif + + +------------------------------------------------------------------------ +-- Floating-point +------------------------------------------------------------------------ + +anyFloatle :: ParserMode m => Parser m e Float +anyFloatle = castWord32ToFloat <$> anyWord32le +{-# INLINE anyFloatle #-} + + +anyFloatbe :: ParserMode m => Parser m e Float +anyFloatbe = castWord32ToFloat <$> anyWord32be +{-# INLINE anyFloatbe #-} + + +anyDoublele :: ParserMode m => Parser m e Double +anyDoublele = castWord64ToDouble <$> anyWord64le +{-# INLINE anyDoublele #-} + + +anyDoublebe :: ParserMode m => Parser m e Double +anyDoublebe = castWord64ToDouble <$> anyWord64be +{-# INLINE anyDoublebe #-} + + +------------------------------------------------------------------------ +-- Byte matching +------------------------------------------------------------------------ + +word8 :: ParserMode m => Word8 -> Parser m e () +word8 (W8# expected) = withEnsure# 1# $ Parser \_ _ s st -> + case eqWord8# (indexWord8OffAddr# s 0#) expected of + 1# -> (# st, OK# () (plusAddr# s 1#) #) + _ -> (# st, Fail# #) +{-# INLINE word8 #-} + + +-- | Match a literal 'ByteString'. Uses word-at-a-time comparison. +byteString :: ByteString -> Parser m e () +byteString bs = Parser \env eob s st -> + let !(BSI.BS _ (I# len#)) = bs + in case len# <=# minusAddr# eob s of + 1# -> case memcmpAddr# s (bsAddr# bs) len# of + 0# -> (# st, OK# () (plusAddr# s len#) #) + _ -> (# st, Fail# #) + _ -> case ensureNSlow env eob s len# st of + (# st', OK# _ s' #) -> case readEnd# env st' of + (# st'', _ #) -> + case memcmpAddr# s' (bsAddr# bs) len# of + 0# -> (# st'', OK# () (plusAddr# s' len#) #) + _ -> (# st'', Fail# #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE byteString #-} + + +bsAddr# :: ByteString -> Addr# +bsAddr# (BSI.BS (ForeignPtr addr _) _) = addr +{-# INLINE bsAddr# #-} + + +memcmpAddr# :: Addr# -> Addr# -> Int# -> Int# +memcmpAddr# a b n = case unsafeDupablePerformIO (c_memcmp (Ptr a) (Ptr b) (I# n)) of + r -> if r == 0 then 0# else 1# +{-# INLINE memcmpAddr# #-} + + +foreign import ccall unsafe "memcmp" + c_memcmp :: Ptr Word8 -> Ptr Word8 -> Int -> IO Int + + +------------------------------------------------------------------------ +-- ByteString operations +------------------------------------------------------------------------ + +{- | Take @n@ bytes. + +Returns a /zero-copy slice/ into the ring (in streaming mode) or +the input 'ByteString' (in 'parseByteString' mode) whenever @n@ +fits in the current backing buffer in one go. When @n@ exceeds +the magic ring's capacity, falls back to draining the bytes into +a fresh heap allocation: 'takeBs' still hands the caller all @n@ +bytes, the cost is one allocation + memcpy and the result no +longer aliases the ring. + +== Lifetime + +For the zero-copy slice case the returned 'ByteString' holds a +reference to the ring's backing memory. The slice is valid as +long as the ring exists (managed by 'withMagicRing' / 'withRecvTransport'). +If the ring is freed while the 'ByteString' is still alive, the +slice becomes a dangling pointer. In practice this is safe because +'runParser' / 'runParserLoop' run within the transport's scope. + +For the drained-into-a-fresh-allocation case (@n@ > ring size) the +result is an ordinary heap 'ByteString' with no ring dependency +and may outlive 'withMagicRing' freely. + +A type-system-enforced alternative for callers that prefer to +/prove/ slices cannot outlive a refill lives in "Wireform.Ring": +use 'Wireform.Ring.RingSlice' values and 'Wireform.Ring.copyRingSlice' +as the explicit escape hatch. +-} +takeBs :: forall m e. ParserMode m => Int -> Parser m e ByteString +takeBs total + | total <= 0 = pure BS.empty + | otherwise = Parser \env eob s st -> + let !(I# mask#) = peMask env + in case (case total of I# n# -> (n# -# 1#) <=# mask#) of + 1# -> runParser# (takeBsSlice @m total) env eob s st + -- Drain path is Stream-monomorphic so that 'modeCheckpoint' + -- in the loop body resolves to 'checkpoint' at compile time + -- (no per-iteration dictionary indirection). 'coerce' is + -- safe because the 'm' phantom parameter of 'Parser' has + -- role phantom — 'Parser Stream e a' and 'Parser m e a' + -- have identical runtime representation. Pure mode never + -- reaches this branch at runtime (peMask is maxBound for + -- Pure, so the bounds check always picks the fast slice). + _ -> + runParser# + (coerce (takeBsDrainStream @e total) :: Parser m e ByteString) + env + eob + s + st +{-# INLINE takeBs #-} + + +{- | Take @n@ bytes, always copying. Use when you need the result +to outlive the transport's scope. + +Like 'takeBs', accepts any @n@ — for @n@ larger than the ring +size, the bytes are drained chunk-by-chunk through the ring into +the destination, advancing the consumer tail between chunks so +the producer has room to keep filling. +-} +takeBsCopy :: forall m e. ParserMode m => Int -> Parser m e ByteString +takeBsCopy total + | total <= 0 = pure BS.empty + | otherwise = Parser \env eob s st -> + let !(I# mask#) = peMask env + in case (case total of I# n# -> (n# -# 1#) <=# mask#) of + 1# -> runParser# (takeBsCopySingle @m total) env eob s st + -- See 'takeBs' for why the drain path is Stream-monomorphic + -- + 'coerce'd back to 'Parser m'. + _ -> + runParser# + (coerce (takeBsDrainStream @e total) :: Parser m e ByteString) + env + eob + s + st +{-# INLINE takeBsCopy #-} + + +-- | Single-shot zero-copy slice (fast path of 'takeBs'). +takeBsSlice :: forall m e. ParserMode m => Int -> Parser m e ByteString +takeBsSlice (I# n#) = withEnsure# @m n# $ Parser \env _ s st -> + let !bs = BSI.BS (ForeignPtr s (peBackingFp env)) (I# n#) + in (# st, OK# bs (plusAddr# s n#) #) +{-# INLINE takeBsSlice #-} + + +-- | Single-shot memcpy into a fresh allocation (fast path of 'takeBsCopy'). +takeBsCopySingle :: forall m e. ParserMode m => Int -> Parser m e ByteString +takeBsCopySingle (I# n#) = withEnsure# @m n# $ Parser \_ _ s st -> + let !bs = + unsafeDupablePerformIO + (BSI.create (I# n#) \dst -> copyBytes dst (Ptr s) (I# n#)) + in (# st, OK# bs (plusAddr# s n#) #) +{-# INLINE takeBsCopySingle #-} + + +{- | Drain @total@ bytes through the ring into a fresh heap +allocation. Used by 'takeBs' and 'takeBsCopy' when @total@ does +not fit in the ring at once. + +Strategy: + +1. Allocate a pinned 'ForeignPtr' of size @total@. +2. 'checkpoint' — advance tail to the parser's current position + so the producer has the whole ring to refill into. +3. Loop: 'ensureN#' a ring-sized chunk, memcpy from the cursor + into the destination, advance the cursor, 'checkpoint' if + more chunks remain. + +The loop performs roughly @ceil(total / ringSize)@ +producer/consumer round-trips. + +Deliberately monomorphic in @Parser Stream e@ rather than +polymorphic over @ParserMode m@: 'checkpoint' / 'ensureN#' need +to resolve at compile time so the inner loop carries no +dictionary indirection. Pure-mode callers can't reach this +function anyway (their bounds check always picks the fast +slice), so the polymorphic dispatch lives in 'takeBs' / +'takeBsCopy' which 'coerce' the 'Parser Stream' result into +'Parser m' across the (phantom) @m@ parameter. +-} +takeBsDrainStream :: forall e. Int -> Parser Stream e ByteString +takeBsDrainStream total = do + ringSz <- readRingSize + fp <- allocPinnedFp total + checkpoint + drainLoopStream fp 0 total ringSz + pure $! BSI.BS fp total + + +{- | Pull at most 'ringSize' bytes per iteration into @fp@; advance +cur and tail between iterations. + +Bang patterns on every argument so the strictness analyzer +unboxes 'copied' / 'total' / 'ringSz' into 'Int#' across the +recursive call. Without the bang on @ringSz@ in particular the +worker shipped it as a boxed 'Int' and paid one @case I#@ per +iteration to unbox it. +-} +drainLoopStream + :: forall e + . ForeignPtr Word8 + -> Int + -> Int + -> Int + -> Parser Stream e () +drainLoopStream !fp !copied !total !ringSz + | copied >= total = pure () + | otherwise = do + let !(I# c#) = min (total - copied) ringSz + ensureN# @Stream c# + memcpyFromCur fp copied (I# c#) + let !copied' = copied + I# c# + when (copied' < total) checkpoint + drainLoopStream fp copied' total ringSz + + +{- | 'peMask + 1' — the streaming ring's capacity, or 'maxBound' / +arbitrary in whole-input mode (the dispatch in 'takeBs' / 'takeBsCopy' +never enters the drain path in whole-input mode, so the value is +not actually consumed there). +-} +readRingSize :: Parser m e Int +readRingSize = Parser \env _eob s st -> + let !rsz = case peMask env of + m -> if m == maxBound then m else m + 1 + in (# st, OK# rsz s #) +{-# INLINE readRingSize #-} + + +{- | Allocate a pinned 'ForeignPtr' of the given size, threading the +parser's 'State#' so allocation order is well-defined. +-} +allocPinnedFp :: Int -> Parser m e (ForeignPtr Word8) +allocPinnedFp n = Parser \_env _eob s st0 -> + case unIO (mallocPlainForeignPtrBytes n) st0 of + (# st1, fp #) -> (# st1, OK# fp s #) + where + unIO (IO f) = f +{-# INLINE allocPinnedFp #-} + + +{- | Copy @len@ bytes from the parser's cursor into @fp + dstOff@, +advancing the cursor by @len@ bytes. Caller must have called +'ensureN#' (or equivalent) for at least @len@ bytes. + +Uses 'unsafeWithForeignPtr' rather than 'withForeignPtr': the +former skips the per-call 'keepAlive#' barrier and the fresh +closure allocation that goes with it. Safe here because @fp@ +is bound by the outer @do@-block in 'takeBsDrain' and stays live +for the entire drain — the returned @BSI.BS fp _@ keeps it +referenced past the loop's last iteration, and the drain loop +itself only ever reads through the address. +-} +memcpyFromCur :: ForeignPtr Word8 -> Int -> Int -> Parser m e () +memcpyFromCur fp (I# dstOff#) (I# len#) = Parser \_env _eob s st0 -> + case unIO + ( unsafeWithForeignPtr fp \(Ptr dst#) -> + copyBytes + (Ptr (plusAddr# dst# dstOff#)) + (Ptr s) + (I# len#) + ) + st0 of + (# st1, () #) -> (# st1, OK# () (plusAddr# s len#) #) + where + unIO (IO f) = f +{-# INLINE memcpyFromCur #-} + + +skip :: ParserMode m => Int -> Parser m e () +skip (I# n#) = withEnsure# n# $ Parser \_ _ s st -> + (# st, OK# () (plusAddr# s n#) #) +{-# INLINE skip #-} + + +-- | Consume all remaining bytes in the current window. +takeRest :: Parser m e ByteString +takeRest = Parser \env eob s st -> + let !len = I# (minusAddr# eob s) + !bs = + if len <= 0 + then BS.empty + else BSI.BS (ForeignPtr s (peBackingFp env)) len + in (# st, OK# bs eob #) +{-# INLINE takeRest #-} + + +{- | Skip backward @n@ bytes (takes a positive integer). + +In the ring buffer context, you can only skip back to bytes that +haven't been overwritten — i.e., bytes between the consumer tail +and the current position. For 'parseByteString' this is the +entire input. For streaming, it's bounded by whatever the driver +hasn't advanced the tail past (at minimum, the start of the +current 'runParser' invocation). + +Fails if @n@ would move before the most recent anchor (the start +of the current parse run for non-streaming, or the latest +'checkpoint' for streaming). Does NOT check the ring tail +directly — the driver guarantees tail <= anchor. +-} +skipBack :: Int -> Parser m e () +skipBack (I# n#) = Parser \env _ s st -> + let !target = plusAddr# s (negateInt# n#) + !(Ptr anchorPtr#) = peAnchorCur env + in case readAddrOffAddr# anchorPtr# 0# st of + (# st', anchorCur# #) -> + case leAddr# anchorCur# target of + 1# -> (# st', OK# () target #) + _ -> (# st', Fail# #) +{-# INLINE skipBack #-} + + +------------------------------------------------------------------------ +-- Signed integers +------------------------------------------------------------------------ + +anyInt8 :: ParserMode m => Parser m e Int8 +anyInt8 = fromIntegral <$> anyWord8 +{-# INLINE anyInt8 #-} + + +anyInt16 :: ParserMode m => Parser m e Int16 +anyInt16 = fromIntegral <$> anyWord16 +{-# INLINE anyInt16 #-} + + +anyInt32 :: ParserMode m => Parser m e Int32 +anyInt32 = fromIntegral <$> anyWord32 +{-# INLINE anyInt32 #-} + + +anyInt64 :: ParserMode m => Parser m e Int64 +anyInt64 = fromIntegral <$> anyWord64 +{-# INLINE anyInt64 #-} + + +anyInt16le :: ParserMode m => Parser m e Int16 +anyInt16le = fromIntegral <$> anyWord16le +{-# INLINE anyInt16le #-} + + +anyInt32le :: ParserMode m => Parser m e Int32 +anyInt32le = fromIntegral <$> anyWord32le +{-# INLINE anyInt32le #-} + + +anyInt64le :: ParserMode m => Parser m e Int64 +anyInt64le = fromIntegral <$> anyWord64le +{-# INLINE anyInt64le #-} + + +anyInt16be :: ParserMode m => Parser m e Int16 +anyInt16be = fromIntegral <$> anyWord16be +{-# INLINE anyInt16be #-} + + +anyInt32be :: ParserMode m => Parser m e Int32 +anyInt32be = fromIntegral <$> anyWord32be +{-# INLINE anyInt32be #-} + + +anyInt64be :: ParserMode m => Parser m e Int64 +anyInt64be = fromIntegral <$> anyWord64be +{-# INLINE anyInt64be #-} + + +------------------------------------------------------------------------ +-- UTF-8 / ASCII +------------------------------------------------------------------------ + +-- | Parse any UTF-8 character. +anyChar :: ParserMode m => Parser m e Char +anyChar = withEnsure# 1# $ Parser \_ eob s st -> + case indexWord8OffAddr# s 0# of + c1 -> case leWord8# c1 (wordToWord8# 0x7F##) of + 1# -> (# st, OK# (C# (chr# (word2Int# (word8ToWord# c1)))) (plusAddr# s 1#) #) + _ -> case leWord8# c1 (wordToWord8# 0xBF##) of + 1# -> (# st, Fail# #) + _ -> case leWord8# c1 (wordToWord8# 0xDF##) of + 1# -> case 2# <=# minusAddr# eob s of + 1# -> case indexWord8OffAddr# s 1# of + c2 -> + let !cp = + ((word2Int# (word8ToWord# c1) -# 0xC0#) `uncheckedIShiftL#` 6#) + +# (word2Int# (word8ToWord# c2) -# 0x80#) + in (# st, OK# (C# (chr# cp)) (plusAddr# s 2#) #) + _ -> (# st, Fail# #) + _ -> case leWord8# c1 (wordToWord8# 0xEF##) of + 1# -> case 3# <=# minusAddr# eob s of + 1# -> case indexWord8OffAddr# s 1# of + c2 -> + case indexWord8OffAddr# s 2# of + c3 -> + let !cp = + ((word2Int# (word8ToWord# c1) -# 0xE0#) `uncheckedIShiftL#` 12#) + +# ((word2Int# (word8ToWord# c2) -# 0x80#) `uncheckedIShiftL#` 6#) + +# (word2Int# (word8ToWord# c3) -# 0x80#) + in (# st, OK# (C# (chr# cp)) (plusAddr# s 3#) #) + _ -> (# st, Fail# #) + _ -> case 4# <=# minusAddr# eob s of + 1# -> case indexWord8OffAddr# s 1# of + c2 -> + case indexWord8OffAddr# s 2# of + c3 -> + case indexWord8OffAddr# s 3# of + c4 -> + let !cp = + ((word2Int# (word8ToWord# c1) -# 0xF0#) `uncheckedIShiftL#` 18#) + +# ((word2Int# (word8ToWord# c2) -# 0x80#) `uncheckedIShiftL#` 12#) + +# ((word2Int# (word8ToWord# c3) -# 0x80#) `uncheckedIShiftL#` 6#) + +# (word2Int# (word8ToWord# c4) -# 0x80#) + in (# st, OK# (C# (chr# cp)) (plusAddr# s 4#) #) + _ -> (# st, Fail# #) +{-# INLINE anyChar #-} + + +anyChar_ :: ParserMode m => Parser m e () +anyChar_ = () <$ anyChar +{-# INLINE anyChar_ #-} + + +-- | CPS variant: parse an ASCII char and pass to continuation. +withSatisfyAscii :: ParserMode m => (Char -> Bool) -> (Char -> Parser m e r) -> Parser m e r +withSatisfyAscii f p = withEnsure# 1# $ Parser \env eob s st -> + case indexCharOffAddr# s 0# of + c1 -> case leChar# c1 '\x7F'# of + 1# -> + let !ch = C# c1 + in if f ch + then runParser# (p ch) env eob (plusAddr# s 1#) st + else (# st, Fail# #) + _ -> (# st, Fail# #) +{-# INLINE withSatisfyAscii #-} + + +-- | Parse an ASCII character matching a predicate. +satisfyAscii :: ParserMode m => (Char -> Bool) -> Parser m e Char +satisfyAscii f = withSatisfyAscii f pure +{-# INLINE satisfyAscii #-} + + +satisfyAscii_ :: ParserMode m => (Char -> Bool) -> Parser m e () +satisfyAscii_ f = withSatisfyAscii f (\_ -> pure ()) +{-# INLINE satisfyAscii_ #-} + + +skipSatisfyAscii :: ParserMode m => (Char -> Bool) -> Parser m e () +skipSatisfyAscii f = withEnsure# 1# $ Parser \_ _ s st -> + case indexCharOffAddr# s 0# of + c1 -> case leChar# c1 '\x7F'# of + 1# -> + if f (C# c1) + then (# st, OK# () (plusAddr# s 1#) #) + else (# st, Fail# #) + _ -> (# st, Fail# #) +{-# INLINE skipSatisfyAscii #-} + + +-- | Parse any ASCII character (< 0x80). +anyAsciiChar :: ParserMode m => Parser m e Char +anyAsciiChar = satisfyAscii (const True) +{-# INLINE anyAsciiChar #-} + + +-- | Skip one ASCII character. +skipAnyAsciiChar :: ParserMode m => Parser m e () +skipAnyAsciiChar = skipSatisfyAscii (const True) +{-# INLINE skipAnyAsciiChar #-} + + +-- | Parse a UTF-8 character matching a predicate. +satisfy :: ParserMode m => (Char -> Bool) -> Parser m e Char +satisfy f = do + c <- anyChar + if f c then pure c else empty +{-# INLINE satisfy #-} + + +satisfy_ :: ParserMode m => (Char -> Bool) -> Parser m e () +satisfy_ f = () <$ satisfy f +{-# INLINE satisfy_ #-} + + +{- | Fused satisfy: uses the ASCII fast path when the byte is < 0x80, +the full UTF-8 path otherwise. Two separate predicates for each case. +-} +fusedSatisfy :: forall m e. ParserMode m => (Char -> Bool) -> (Word8 -> Bool) -> Parser m e Char +fusedSatisfy charPred _bytePred = withEnsure# 1# $ Parser \env eob s st -> + case indexWord8OffAddr# s 0# of + w -> case leWord8# w (wordToWord8# 0x7F##) of + 1# -> + let !ch = C# (chr# (word2Int# (word8ToWord# w))) + in if charPred ch + then (# st, OK# ch (plusAddr# s 1#) #) + else (# st, Fail# #) + _ -> case runParser# (anyChar @m) env eob s st of + (# st', OK# c s' #) -> if charPred c then (# st', OK# c s' #) else (# st', Fail# #) + x -> x +{-# INLINE fusedSatisfy #-} + + +------------------------------------------------------------------------ +-- Character classes +------------------------------------------------------------------------ + +isDigit :: Char -> Bool +isDigit c = c >= '0' && c <= '9' +{-# INLINE isDigit #-} + + +isLatinLetter :: Char -> Bool +isLatinLetter c = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') +{-# INLINE isLatinLetter #-} + + +------------------------------------------------------------------------ +-- ASCII numeric +------------------------------------------------------------------------ + +anyAsciiDecimalWord :: ParserMode m => Parser m e Word +anyAsciiDecimalWord = withEnsure# 1# $ Parser \_ eob s st -> + case indexWord8OffAddr# s 0# of + c -> case leWord8# (wordToWord8# 0x30##) c of + 0# -> (# st, Fail# #) + _ -> case leWord8# c (wordToWord8# 0x39##) of + 0# -> (# st, Fail# #) + _ -> + let !d0 = W# (word8ToWord# c) - 0x30 + in goDecWord eob (plusAddr# s 1#) st d0 +{-# INLINE anyAsciiDecimalWord #-} + + +goDecWord :: Addr# -> Addr# -> State# RealWorld -> Word -> StRes# e Word +goDecWord eob s st !acc = + case eqAddr# eob s of + 1# -> (# st, OK# acc s #) + _ -> case indexWord8OffAddr# s 0# of + c -> case leWord8# (wordToWord8# 0x30##) c of + 0# -> (# st, OK# acc s #) + _ -> case leWord8# c (wordToWord8# 0x39##) of + 0# -> (# st, OK# acc s #) + _ -> + goDecWord + eob + (plusAddr# s 1#) + st + (acc * 10 + (W# (word8ToWord# c) - 0x30)) +{-# NOINLINE goDecWord #-} + + +anyAsciiDecimalInt :: ParserMode m => Parser m e Int +anyAsciiDecimalInt = fromIntegral <$> anyAsciiDecimalWord +{-# INLINE anyAsciiDecimalInt #-} + + +------------------------------------------------------------------------ +-- Control flow +------------------------------------------------------------------------ + +empty :: Parser m e a +empty = Parser \_ _ _ st -> (# st, Fail# #) +{-# INLINE empty #-} + + +infixr 6 <|> + + +(<|>) :: Parser m e a -> Parser m e a -> Parser m e a +(<|>) (Parser f) (Parser g) = Parser \env eob s st -> + case f env eob s st of + (# st', Fail# #) -> g env eob s st' + x -> x +{-# INLINE [1] (<|>) #-} + + +{-# RULES "wireform/reassoc-alt" forall l m r. (l <|> m) <|> r = l <|> (m <|> r) #-} + + +branch :: Parser m e a -> Parser m e b -> Parser m e b -> Parser m e b +branch (Parser p) (Parser t) (Parser f) = Parser \env eob s st -> + case p env eob s st of + (# st', OK# _ s' #) -> t env eob s' st' + (# st', Fail# #) -> f env eob s st' + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE branch #-} + + +lookahead :: Parser m e a -> Parser m e a +lookahead (Parser f) = Parser \env eob s st -> + case f env eob s st of + (# st', OK# a _ #) -> (# st', OK# a s #) + x -> x +{-# INLINE lookahead #-} + + +fails :: Parser m e a -> Parser m e () +fails (Parser f) = Parser \env eob s st -> + case f env eob s st of + (# st', OK# _ _ #) -> (# st', Fail# #) + (# st', Fail# #) -> (# st', OK# () s #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE fails #-} + + +try :: Parser m e a -> Parser m e a +try (Parser f) = Parser \env eob s st -> + case f env eob s st of + (# st', Err# _ #) -> (# st', Fail# #) + x -> x +{-# INLINE try #-} + + +optional :: Parser m e a -> Parser m e (Maybe a) +optional p = (Just <$> p) <|> pure Nothing +{-# INLINE optional #-} + + +optional_ :: Parser m e a -> Parser m e () +optional_ p = (() <$ p) <|> pure () +{-# INLINE optional_ #-} + + +withOption :: Parser m e a -> (a -> Parser m e b) -> Parser m e b -> Parser m e b +withOption (Parser p) just (Parser nothing) = Parser \env eob s st -> + case p env eob s st of + (# st', OK# a s' #) -> runParser# (just a) env eob s' st' + (# st', Fail# #) -> nothing env eob s st' + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE withOption #-} + + +many_ :: Parser m e a -> Parser m e () +many_ (Parser f) = Parser go + where + go env eob s st = case f env eob s st of + (# st', OK# _ s' #) -> go env eob s' st' + (# st', Fail# #) -> (# st', OK# () s #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE many_ #-} + + +some_ :: Parser m e a -> Parser m e () +some_ p = p *> many_ p +{-# INLINE some_ #-} + + +skipMany :: Parser m e a -> Parser m e () +skipMany = many_ +{-# INLINE skipMany #-} + + +skipSome :: Parser m e a -> Parser m e () +skipSome = some_ +{-# INLINE skipSome #-} + + +-- | Run @p@ zero or more times, collecting results into a list. +many :: Parser m e a -> Parser m e [a] +many (Parser f) = Parser go + where + go env eob s st = case f env eob s st of + (# st', OK# a s' #) -> case go env eob s' st' of + (# st'', OK# as s'' #) -> (# st'', OK# (a : as) s'' #) + x -> x + (# st', Fail# #) -> (# st', OK# [] s #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE many #-} + + +-- | Run @p@ one or more times, collecting results. +some :: Parser m e a -> Parser m e [a] +some p = (:) <$> p <*> many p +{-# INLINE some #-} + + +-- | Succeed iff @p@ fails. Does not consume input. +notFollowedBy :: Parser m e a -> Parser m e () +notFollowedBy = fails +{-# INLINE notFollowedBy #-} + + +{- | Run a parser on exactly @n@ bytes. The inner parser must consume +all @n@ bytes or the overall parse fails. +-} +isolate :: ParserMode m => Int -> Parser m e a -> Parser m e a +isolate (I# n#) (Parser p) = withEnsure# n# $ Parser \env _ s st -> + let !isolEnd = plusAddr# s n# + in case p env isolEnd s st of + (# st', OK# a s' #) -> case eqAddr# s' isolEnd of + 1# -> (# st', OK# a s' #) + _ -> (# st', Fail# #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE isolate #-} + + +------------------------------------------------------------------------ +-- Error handling +------------------------------------------------------------------------ + +-- | Catch an @Err#@ and handle it. +withError :: (e -> Parser m e a) -> Parser m e a -> Parser m e a +withError handler (Parser f) = Parser \env eob s st -> + case f env eob s st of + (# st', Err# e #) -> runParser# (handler e) env eob s st' + x -> x +{-# INLINE withError #-} + + +err :: e -> Parser m e a +err e = Parser \_ _ _ st -> (# st, Err# e #) +{-# INLINE err #-} + + +cut :: Parser m e a -> e -> Parser m e a +cut (Parser f) e = Parser \env eob s st -> + case f env eob s st of + (# st', Fail# #) -> (# st', Err# e #) + x -> x +{-# INLINE cut #-} + + +cutting :: Parser m e a -> e -> (e -> e -> e) -> Parser m e a +cutting (Parser f) e merge = Parser \env eob s st -> + case f env eob s st of + (# st', Fail# #) -> (# st', Err# e #) + (# st', Err# e' #) -> (# st', Err# (merge e' e) #) + x -> x +{-# INLINE cutting #-} + + +------------------------------------------------------------------------ +-- EOF +------------------------------------------------------------------------ + +{- | Succeed iff at end of input. +In streaming mode, this may suspend to check if more data is coming. +Uses ensureN# 1 internally: if ensure fails (no more data), we're at EOF. +-} +eof :: forall m e. ParserMode m => Parser m e () +eof = Parser \env eob s st -> + case eqAddr# eob s of + 1# -> case onEnsureFail @m env eob s 1# st of + (# st', OK# _ _ #) -> (# st', Fail# #) -- data arrived + (# st', Fail# #) -> (# st', OK# () s #) -- genuine EOF + (# st', x #) -> (# st', unsafeCoerce# x #) + _ -> (# st, Fail# #) +{-# INLINE eof #-} + + +{- | Non-consuming check: 'True' if at end of input. +In streaming mode, may suspend to determine. +-} +atEnd :: ParserMode m => Parser m e Bool +atEnd = (eof *> pure True) <|> pure False +{-# INLINE atEnd #-} + + +-- | Number of bytes remaining in the current window (cheap, no suspend). +remaining :: Parser m e Int +remaining = Parser \_ eob s st -> + (# st, OK# (I# (minusAddr# eob s)) s #) +{-# INLINE remaining #-} + + +------------------------------------------------------------------------ +-- Chain combinators +------------------------------------------------------------------------ + +{- | Left-associative chain. +@chainl f p q@ parses @p@ then zero or more @q@, folding left with @f@. +-} +chainl :: (b -> a -> b) -> Parser m e b -> Parser m e a -> Parser m e b +chainl f (Parser p) (Parser q) = Parser \env eob s st -> + case p env eob s st of + (# st', OK# b s' #) -> chainlGo f q env eob s' st' b + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE chainl #-} + + +chainlGo + :: forall e b a + . (b -> a -> b) + -> (ParserEnv -> Addr# -> Addr# -> State# RealWorld -> StRes# e a) + -> ParserEnv + -> Addr# + -> Addr# + -> State# RealWorld + -> b + -> StRes# e b +chainlGo f q env eob s st !acc = + case q env eob s st of + (# st', OK# a s' #) -> chainlGo f q env eob s' st' (f acc a) + (# st', Fail# #) -> (# st', OK# acc s #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# NOINLINE chainlGo #-} + + +{- | Right-associative chain. +@chainr f p q@ parses zero or more @p@, then @q@ at the end, +folding right with @f@. +-} +chainr :: (a -> b -> b) -> Parser m e a -> Parser m e b -> Parser m e b +chainr f (Parser p) (Parser q) = Parser go + where + go env eob s st = + case p env eob s st of + (# st', OK# a s' #) -> case go env eob s' st' of + (# st'', OK# b s'' #) -> (# st'', OK# (f a b) s'' #) + (# st'', x #) -> (# st'', unsafeCoerce# x #) + (# st', Fail# #) -> q env eob s st' + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE chainr #-} + + +------------------------------------------------------------------------ +-- Additional error handling +------------------------------------------------------------------------ + +{- | Like 'ensureN#' but on EOF, errors with a specific error +instead of failing recoverably. +-} +ensureNOrEof :: ParserMode m => Int -> e -> Parser m e () +ensureNOrEof n e = ensureN# (case n of I# n# -> n#) <|> err e +{-# INLINE ensureNOrEof #-} + + +------------------------------------------------------------------------ +-- Additional position combinators +------------------------------------------------------------------------ + +-- | Parse something and return both the result and the bytes consumed. +withByteString :: Parser m e a -> (a -> ByteString -> Parser m e b) -> Parser m e b +withByteString (Parser p) f = Parser \env eob s st -> + case p env eob s st of + (# st', OK# a s' #) -> + let !len = I# (minusAddr# s' s) + !bs = BSI.BS (ForeignPtr s (peBackingFp env)) len + in runParser# (f a bs) env eob s' st' + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE withByteString #-} + + +------------------------------------------------------------------------ +-- Additional character classes +------------------------------------------------------------------------ + +isGreekLetter :: Char -> Bool +isGreekLetter c = (c >= '\x0391' && c <= '\x03A9') || (c >= '\x03B1' && c <= '\x03C9') +{-# INLINE isGreekLetter #-} + + +------------------------------------------------------------------------ +-- Additional ASCII numeric +------------------------------------------------------------------------ + +anyAsciiHexWord :: ParserMode m => Parser m e Word +anyAsciiHexWord = withEnsure# 1# $ Parser \_ eob s st -> + case hexDigit (indexWord8OffAddr# s 0#) of + (# | (# #) #) -> (# st, Fail# #) + (# (# d #) | #) -> goHexWord eob (plusAddr# s 1#) st (W# (word8ToWord# d)) +{-# INLINE anyAsciiHexWord #-} + + +hexDigit :: Word8# -> (# (# Word8# #) | (# #) #) +hexDigit w + | isTrue# (leWord8# (wordToWord8# 0x30##) w) + , isTrue# (leWord8# w (wordToWord8# 0x39##)) = + (# (# wordToWord8# (word8ToWord# w `minusWord#` 0x30##) #) | #) + | isTrue# (leWord8# (wordToWord8# 0x41##) w) + , isTrue# (leWord8# w (wordToWord8# 0x46##)) = + (# (# wordToWord8# (word8ToWord# w `minusWord#` 0x37##) #) | #) + | isTrue# (leWord8# (wordToWord8# 0x61##) w) + , isTrue# (leWord8# w (wordToWord8# 0x66##)) = + (# (# wordToWord8# (word8ToWord# w `minusWord#` 0x57##) #) | #) + | otherwise = (# | (# #) #) +{-# INLINE hexDigit #-} + + +goHexWord :: Addr# -> Addr# -> State# RealWorld -> Word -> StRes# e Word +goHexWord eob s st !acc = + case eqAddr# eob s of + 1# -> (# st, OK# acc s #) + _ -> case hexDigit (indexWord8OffAddr# s 0#) of + (# | (# #) #) -> (# st, OK# acc s #) + (# (# d #) | #) -> + goHexWord + eob + (plusAddr# s 1#) + st + (acc * 16 + W# (word8ToWord# d)) +{-# NOINLINE goHexWord #-} + + +anyAsciiHexInt :: ParserMode m => Parser m e Int +anyAsciiHexInt = fromIntegral <$> anyAsciiHexWord +{-# INLINE anyAsciiHexInt #-} + + +anyAsciiDecimalInteger :: ParserMode m => Parser m e Integer +anyAsciiDecimalInteger = fromIntegral <$> anyAsciiDecimalWord +{-# INLINE anyAsciiDecimalInteger #-} diff --git a/src/Wireform/Parser/Adapter.hs b/src/Wireform/Parser/Adapter.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Adapter.hs @@ -0,0 +1,166 @@ +{-# LANGUAGE BlockArguments #-} + +module Wireform.Parser.Adapter ( + -- * ChunkParser type + ChunkParser (..), + ChunkStep (..), + ChunkFinal (..), + ChunkParseError (..), + + -- * Running + runChunked, + runChunkedLoop, + + -- * Chunk mode + ChunkMode (..), +) where + +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import Data.ByteString.Internal qualified as BSI +import Data.Word (Word64, Word8) +import Foreign.ForeignPtr (newForeignPtr_) +import Foreign.Marshal.Utils (copyBytes) +import Foreign.Ptr (Ptr, plusPtr) +import Wireform.Parser.Driver (LoopControl (..)) +import Wireform.Ring.Internal (ringBase, ringMask, ringSize) +import Wireform.Transport + + +------------------------------------------------------------------------ +-- Types +------------------------------------------------------------------------ + +{- | A chunked-feeding parser. The lingua franca for incremental +parser libraries (attoparsec, binary, cereal, etc.). +-} +data ChunkParser a = ChunkParser + { stepChunk :: !(ByteString -> ChunkStep a) + -- ^ Feed non-empty input. + , stepEof :: !(ChunkFinal a) + -- ^ Signal end of input. + } + + +data ChunkStep a + = -- | Consumed N bytes; needs more. + ChunkConsumed {-# UNPACK #-} !Int !(ChunkParser a) + | -- | Done; consumed N bytes of the fed chunk (rest is leftover). + ChunkDone !a {-# UNPACK #-} !Int + | ChunkFailed !ChunkParseError + + +data ChunkFinal a + = FinalDone !a + | FinalFailed !ChunkParseError + + +data ChunkParseError = ChunkParseError + { chunkErrorMessage :: !String + , chunkErrorContext :: ![String] + , chunkErrorPosition :: {-# UNPACK #-} !Word64 + } + deriving stock (Show) + + +------------------------------------------------------------------------ +-- Chunk mode +------------------------------------------------------------------------ + +data ChunkMode + = -- | Allocate a fresh 'ByteString' per chunk (safe, default). + ChunkCopy + | {- | Reference ring memory directly. Caller must not retain the + 'ByteString' past the @stepChunk@ return. + -} + ChunkZeroCopy + deriving stock (Eq, Show) + + +------------------------------------------------------------------------ +-- Running +------------------------------------------------------------------------ + +-- | Run a 'ChunkParser' against a transport. +runChunked + :: ReceiveTransport + -> ChunkMode + -> ChunkParser a + -> IO (Either ChunkParseError a) +runChunked t mode cp0 = do + startPos <- receiveLoadHead t + loop cp0 startPos startPos startPos + where + ring = receiveRing t + base = ringBase ring + msk = ringMask ring + sz = ringSize ring + + advanceThreshold = sz `div` 4 + + loop cp parserStart parserPos lastTailAdvance = do + h <- receiveLoadHead t + if h > parserPos + then do + let !chunkLen = fromIntegral (min (h - parserPos) (fromIntegral sz)) + !off = fromIntegral parserPos .&. msk + !ptr = base `plusPtr` off + chunk <- makeChunk mode ptr chunkLen + case stepChunk cp chunk of + ChunkDone a consumed -> do + let !newPos = parserPos + fromIntegral consumed + receiveAdvanceTail t newPos + pure (Right a) + ChunkConsumed consumed cp' -> do + let !newPos = parserPos + fromIntegral consumed + if newPos - lastTailAdvance >= fromIntegral advanceThreshold + then do + receiveAdvanceTail t newPos + loop cp' parserStart newPos newPos + else loop cp' parserStart newPos lastTailAdvance + ChunkFailed e -> pure (Left e) + else do + r <- receiveWaitData t parserPos + case r of + ReceiveMoreData _ -> + loop cp parserStart parserPos lastTailAdvance + ReceiveEndOfInput -> + case stepEof cp of + FinalDone a + | parserPos == parserStart -> pure (Right a) + | otherwise -> pure (Right a) + FinalFailed e -> pure (Left e) + ReceiveFailed exc -> + pure (Left (ChunkParseError (show exc) [] parserPos)) + + +-- | Loop variant for repeated parsing. +runChunkedLoop + :: ReceiveTransport + -> ChunkMode + -> ChunkParser a + -> (a -> IO LoopControl) + -> IO (Either ChunkParseError ()) +runChunkedLoop t mode mkParser k = loop + where + loop = do + r <- runChunked t mode mkParser + case r of + Right a -> do + ctl <- k a + case ctl of + Continue -> loop + Stop -> pure (Right ()) + Left e -> pure (Left e) + + +------------------------------------------------------------------------ +-- Helpers +------------------------------------------------------------------------ + +makeChunk :: ChunkMode -> Ptr Word8 -> Int -> IO ByteString +makeChunk ChunkCopy ptr len = + BSI.create len \dst -> copyBytes dst ptr len +makeChunk ChunkZeroCopy ptr len = do + fptr <- newForeignPtr_ ptr + pure (BSI.fromForeignPtr fptr 0 len) diff --git a/src/Wireform/Parser/Driver.hs b/src/Wireform/Parser/Driver.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Driver.hs @@ -0,0 +1,354 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +module Wireform.Parser.Driver ( + runParser, + runParserLoop, + LoopControl (..), + parseByteString, + runParserInternal, + InternalResult (..), +) where + +import Control.Exception (SomeException, mask) +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import Data.ByteString.Internal qualified as BSI +import Data.IORef +import Data.Word (Word64, Word8) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Marshal.Alloc (allocaBytes) +import Foreign.Ptr (Ptr, castPtr, minusPtr, plusPtr) +import Foreign.Storable (poke) +import GHC.Exts +import GHC.ForeignPtr (ForeignPtr (..), ForeignPtrContents (..)) +import GHC.IO (IO (..)) +import System.IO.Unsafe (unsafeDupablePerformIO) +import Wireform.Parser.Error +import Wireform.Parser.Internal +import Wireform.Ring.Internal (ringBase, ringMask, ringSize) +import Wireform.Transport + + +data LoopControl = Continue | Stop + deriving stock (Eq, Show) + + +data InternalResult e a + = IRDone {-# UNPACK #-} !Word64 !a + | IRFail !Word64 + | IRErr !Word64 !e + | IRUnexpectedEof !Word64 !Int + | IRTransportError !SomeException + | IRCleanEof + | -- | position, requested bytes, ring size — see 'ParseRingOverflow'. + IRRingOverflow {-# UNPACK #-} !Word64 {-# UNPACK #-} !Int {-# UNPACK #-} !Int + + +data TransportState + = TSOpen + | TSClosedEof + | TSClosedErr !SomeException + + +------------------------------------------------------------------------ +-- runParser +------------------------------------------------------------------------ + +runParser :: forall e a. ReceiveTransport -> Parser Stream e a -> IO (Either (ParseError e) a) +runParser t p = do + startPos <- receiveLoadHead t + ir <- runParserInternal t p startPos + pure $ case ir of + IRDone _ a -> Right a + IRFail pos -> Left (ParseFail pos) + IRErr pos e -> Left (ParseErr pos e) + IRUnexpectedEof pos n -> Left (ParseUnexpectedEof pos n) + IRTransportError exc -> Left (ParseTransportError exc) + IRCleanEof -> Left (ParseUnexpectedEof 0 0) + IRRingOverflow pos n sz -> Left (ParseRingOverflow pos n sz) + + +runParserInternal :: forall e a. ReceiveTransport -> Parser Stream e a -> Word64 -> IO (InternalResult e a) +runParserInternal t p startPos = mask \restore -> do + let ring = receiveRing t + !base = ringBase ring + !msk = ringMask ring + !sz = ringSize ring + + currentHead <- receiveLoadHead t + let !curOffset = fromIntegral startPos .&. msk + !(Ptr initCur#) = base `plusPtr` curOffset + -- See StepCheckpoint / StepSuspend for the rationale on + -- computing eob from @cur + avail@ rather than + -- @base + (head .&. msk)@: when @head - startPos == ringSize@, + -- the masked offsets coincide and eob collapses onto cur. + -- + -- Clamp to 'sz' (the ring size). A well-behaved transport + -- maintains @head <= tail + ringSize@ so the clamp is a no-op, + -- but a misbehaving transport (e.g. a benchmark fixture that + -- prefills more bytes than the ring holds, or a bug in a + -- recv loop) can claim a much larger head — without the + -- clamp 'eob' would point past the double mapping and the + -- first parser read past 'base + 2*ringSize' would segfault. + !avail = min sz (fromIntegral (currentHead - startPos)) + !(Ptr initEnd#) = (Ptr initCur#) `plusPtr` avail + + -- One stack-allocated buffer holds the three mutable cells the + -- parser shares with the driver: end pointer, anchor pos, anchor + -- cur. The cells live for the duration of the parse run and + -- never escape — neither 'ParserEnv' nor any 'control0#' + -- continuation can outlive the surrounding 'prompt#' frame, which + -- is set up inside this 'allocaBytes' scope. + -- + -- The cells are mutated by the parser-side resume bodies (see + -- 'checkpoint' / 'ensureNSlow' in "Wireform.Parser.Internal") with + -- no lock or fence: those bodies only run while the parser is + -- suspended between 'control0#' and the matching 'prompt#'. See + -- the 'ParserEnv' haddock for the full argument. + allocaBytes 24 \cells -> do + let !endPtr = cells + !anchorPos = cells `plusPtr` 8 + !anchorCur = cells `plusPtr` 16 + poke (castPtr endPtr :: Ptr (Ptr Word8)) (Ptr initEnd#) + poke (castPtr anchorPos :: Ptr Word64) startPos + poke (castPtr anchorCur :: Ptr (Ptr Word8)) (Ptr initCur#) + + highWaterRef <- newIORef startPos + tsRef <- newIORef TSOpen + + -- Allocate tag and construct env inside IO so the tag + -- is available for peTag, then return both. + (env, step0) <- IO \s0 -> case newPromptTag# s0 of + (# s1, (tag :: PromptTag# (Step e a)) #) -> + let !env' = + ParserEnv + { peEndPtr = castPtr endPtr + , peBaseAddr = base + , peMask = msk + , peAnchorPos = castPtr anchorPos + , peAnchorCur = castPtr anchorCur + , peBackingFp = FinalPtr + , peTag = tagToAny tag + } + body :: State# RealWorld -> (# State# RealWorld, Step e a #) + body s = case runParser# p env' initEnd# initCur# s of + (# s', OK# a cur' #) -> + case curToPos env' cur' s' of + (# s'', newPos #) -> (# s'', StepDone newPos a #) + (# s', Fail# #) -> + (# s', StepFail startPos #) + (# s', Err# e #) -> + (# s', StepErr startPos e #) + in case prompt# tag body s1 of + (# s2, step #) -> (# s2, (env', step) #) + + driverLoop restore t env base msk sz startPos highWaterRef tsRef step0 + + +driverLoop + :: forall e a + . (forall x. IO x -> IO x) + -> ReceiveTransport + -> ParserEnv + -> Ptr Word8 + -> Int + -> Int + -> Word64 + -> IORef Word64 + -> IORef TransportState + -> Step e a + -> IO (InternalResult e a) +driverLoop restore t env base msk sz startPos hwRef tsRef = go + where + go :: Step e a -> IO (InternalResult e a) + go step = case step of + StepDone newPos a -> do + receiveAdvanceTail t newPos + pure (IRDone newPos a) + StepErr pos e -> pure (IRErr pos e) + StepFail pos -> do + ts <- readIORef tsRef + hw <- readIORef hwRef + pure $ case ts of + TSOpen -> IRFail pos + TSClosedErr e -> IRTransportError e + TSClosedEof + | hw == startPos -> IRCleanEof + | otherwise -> IRUnexpectedEof pos 0 + StepCheckpoint pos resume -> do + receiveAdvanceTail t pos + h <- receiveLoadHead t + -- Compute eob relative to the new cur so that a wrap (where + -- @pos .&. msk == h .&. msk@ because exactly a ring's worth + -- of bytes are in flight) does not collapse eob onto cur and + -- make the parser see zero bytes when there are actually + -- ringSize bytes available. The double mapping guarantees + -- @newCur + avail@ is addressable for any @avail <= ringSize@. + -- The 'min sz' guards against transports that mis-report + -- head past 'tail + ringSize'; without it the parser could + -- read past the double mapping and segfault. + -- 'resumeContinue' re-anchors the env so 'curToPos' stays + -- correct after the cur wrap. + let !curOff = fromIntegral pos .&. msk + !newCur = base `plusPtr` curOff + !avail = min sz (fromIntegral (h - pos)) + !newEnd = newCur `plusPtr` avail + nextStep <- resumeContinue resume newCur newEnd + go nextStep + StepSuspend pausedAt needed resume + | needed > sz -> + -- The parser asked for more bytes than the entire ring can + -- ever hold. No amount of refilling will satisfy this; + -- waiting would deadlock (producer can't make room because + -- the consumer is suspended waiting for it). Fail loudly + -- instead. + pure (IRRingOverflow pausedAt needed sz) + | otherwise -> do + modifyIORef' hwRef (max pausedAt) + result <- restore (waitUntilAvailable t tsRef pausedAt needed sz) + case result of + WAMoreData newHead -> do + modifyIORef' hwRef (max newHead) + -- See the comment on StepCheckpoint above: compute eob + -- as @newCur + (newHead - pausedAt)@ rather than as + -- @base + (newHead .&. msk)@. The latter collapses to + -- @newCur@ when the producer has filled exactly one + -- ring-worth of bytes since @pausedAt@, hiding all the + -- newly-available bytes from the parser. + -- 'min sz' guards against misbehaving transports — + -- see the matching comment in 'runParserInternal'. + -- 'resumeContinue' re-anchors the env so 'curToPos' + -- stays correct after the cur wrap. + let !newCurOff = fromIntegral pausedAt .&. msk + !newCur = base `plusPtr` newCurOff + !avail = min sz (fromIntegral (newHead - pausedAt)) + !newEnd = newCur `plusPtr` avail + nextStep <- resumeContinue resume newCur newEnd + go nextStep + WAEndOfInput -> do + nextStep <- resumeEof resume + go nextStep + WATransportError exc -> + pure (IRTransportError exc) + + +data WaitAvail + = WAMoreData {-# UNPACK #-} !Word64 + | WAEndOfInput + | WATransportError !SomeException + + +waitUntilAvailable + :: ReceiveTransport + -> IORef TransportState + -> Word64 + -> Int + -> Int + -> IO WaitAvail +waitUntilAvailable t tsRef pos needed _ringSize = do + h0 <- receiveLoadHead t + loop (max pos h0) + where + -- 'waitFrom' is the position we ask the transport to advance past. + -- It advances on each iteration so a transport that delivers data + -- in dribs (a chunked feeder, a TLS layer that yields partial + -- records, an io_uring SQ that hands us one CQE at a time) + -- doesn't spin: each 'receiveWaitData' is asked to wait /past/ + -- the head we already observed, so it must either pull fresh + -- bytes or report ReceiveEndOfInput. + loop !waitFrom = do + h <- receiveLoadHead t + if h - pos >= fromIntegral needed + then pure (WAMoreData h) + else do + r <- receiveWaitData t waitFrom + case r of + ReceiveMoreData h' -> loop (max waitFrom h') + ReceiveEndOfInput -> do writeIORef tsRef TSClosedEof; pure WAEndOfInput + ReceiveFailed exc -> do writeIORef tsRef (TSClosedErr exc); pure (WATransportError exc) + + +------------------------------------------------------------------------ +-- runParserLoop +------------------------------------------------------------------------ + +runParserLoop + :: forall e a + . ReceiveTransport + -> Parser Stream e a + -> (a -> IO LoopControl) + -> IO (Either (ParseError e) ()) +runParserLoop t p k = do + startPos <- receiveLoadHead t + loop startPos + where + loop pos = do + r <- runParserInternal t p pos + case r of + IRDone newPos a -> do + ctl <- k a + case ctl of Continue -> loop newPos; Stop -> pure (Right ()) + IRCleanEof -> pure (Right ()) + IRFail fpos -> pure (Left (ParseFail fpos)) + IRErr fpos e -> pure (Left (ParseErr fpos e)) + IRUnexpectedEof fpos n -> pure (Left (ParseUnexpectedEof fpos n)) + IRTransportError exc -> pure (Left (ParseTransportError exc)) + IRRingOverflow fpos n sz -> pure (Left (ParseRingOverflow fpos n sz)) + + +------------------------------------------------------------------------ +-- parseByteString (non-streaming, flatparse-equivalent) +------------------------------------------------------------------------ + +{- | Run a parser against a whole 'ByteString'. +The hot path is bit-identical to flatparse — no suspension overhead. +-} +parseByteString :: forall e a. Parser Pure e a -> ByteString -> Either (ParseError e) a +parseByteString p b = unsafeDupablePerformIO $ do + -- withForeignPtr keeps the ByteString's backing memory alive + let !(BSI.BS (ForeignPtr buf# fp) (I# len#)) = b + !end# = plusAddr# buf# len# + + withForeignPtr (ForeignPtr buf# fp) \_ -> + allocaBytes 24 \cells -> do + let !endPtr = cells + !anchorPos = cells `plusPtr` 8 + !anchorCur = cells `plusPtr` 16 + poke (castPtr endPtr :: Ptr (Ptr Word8)) (Ptr end#) + poke (castPtr anchorPos :: Ptr Word64) 0 + poke (castPtr anchorCur :: Ptr (Ptr Word8)) (Ptr buf#) + + IO \s0 -> case newPromptTag# s0 of + (# s1, (tag :: PromptTag# (Step e a)) #) -> + let env = + ParserEnv + { peEndPtr = castPtr endPtr + , peBaseAddr = Ptr buf# + , peMask = maxBound + , peAnchorPos = castPtr anchorPos + , peAnchorCur = castPtr anchorCur + , peBackingFp = fp + , peTag = tagToAny tag + } + body :: State# RealWorld -> (# State# RealWorld, Step e a #) + body s = case runParser# p env end# buf# s of + (# s', OK# a cur' #) -> + let !pos = fromIntegral (I# (minusAddr# cur' buf#)) + in (# s', StepDone pos a #) + (# s', Fail# #) -> (# s', StepFail 0 #) + (# s', Err# e #) -> (# s', StepErr 0 e #) + in case prompt# tag body s1 of + (# s2, step #) -> unIO (classifyStep step) s2 + where + classifyStep (StepDone _ a) = pure (Right a) + classifyStep (StepFail pos) = pure (Left (ParseFail pos)) + classifyStep (StepErr pos e) = pure (Left (ParseErr pos e)) + classifyStep (StepSuspend _ _ r) = resumeEof r >>= classifyStep + classifyStep (StepCheckpoint _ r) = resumeEof r >>= classifyStep + unIO (IO f) = f diff --git a/src/Wireform/Parser/Error.hs b/src/Wireform/Parser/Error.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Error.hs @@ -0,0 +1,46 @@ +module Wireform.Parser.Error ( + ParseError (..), + CleanEof (..), +) where + +import Control.Exception (Exception, SomeException) +import Data.Typeable (Typeable) +import Data.Word (Word64) + + +-- | The result of a failed parse. +data ParseError e + = {- | The parser failed (recoverably) at the given absolute position + and the failure was not caught by an enclosing alternative. + -} + ParseFail !Word64 + | -- | The parser failed unrecoverably (via cut\/commit) with user error @e@. + ParseErr !Word64 !e + | {- | Transport reported EOF while the parser was mid-message. + The position is where the parser was; the 'Int' is how many + additional bytes it needed. + -} + ParseUnexpectedEof !Word64 !Int + | -- | Transport-level failure (broken connection, closed socket, etc.). + ParseTransportError !SomeException + | {- | The parser asked for more bytes than the magic ring can ever + hold (or, more precisely, more than the remaining capacity from + the parser's current position). Continuing would deadlock — + the producer cannot make room because the consumer is suspended + waiting for bytes the producer cannot deliver. Fields: parser + position, requested bytes, ring size. + -} + ParseRingOverflow !Word64 !Int !Int + deriving stock (Show, Functor) + + +{- | Internal sentinel: clean EOF before any bytes were consumed for +this parse. Never exposed to users — the driver converts it to +either clean loop termination or 'ParseUnexpectedEof' depending on +context. +-} +data CleanEof = CleanEof + deriving stock (Show, Typeable) + + +instance Exception CleanEof diff --git a/src/Wireform/Parser/Internal.hs b/src/Wireform/Parser/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Internal.hs @@ -0,0 +1,572 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +{- | Core parser types, mirroring flatparse's representation for +identical inner-loop performance. + +The parser threads @Addr#@ pointers and @State# RealWorld@ directly +(no IO wrapper, no boxed Ptr). The result is an unboxed sum +matching flatparse's @Res#@. + +Streaming suspension via @control0#@ is confined to 'ensureNSlow', +which is @NOINLINE@ and never fires on the hot path for whole-input +parsing. +-} +module Wireform.Parser.Internal ( + -- * Parser type + Parser (..), + + -- * Mode types and class + Pure, + Stream, + ParserMode (..), + + -- * Result types + type Res#, + type StRes#, + pattern OK#, + pattern Fail#, + pattern Err#, + + -- * Step / Resume (driver protocol) + Step (..), + Resume (..), + + -- * Parser environment + ParserEnv (..), + curToPos, + curToPosIO, + + -- * Suspension / checkpoint primitives + ensureN#, + ensureNSlow, + checkpoint, + + -- * Tag storage + tagToAny, + pePromptTag, + + -- * End-pointer access + readEnd#, + writeEnd#, + readEnd, + writeEnd, + + -- * Anchor access + writeAnchor#, + writeAnchor, +) where + +import Control.Applicative (Alternative (..)) +import Control.Monad (MonadPlus) +import Data.Kind (Type) +import Data.Word (Word64, Word8) +import GHC.Exts +import GHC.ForeignPtr (ForeignPtrContents (..)) +import GHC.IO (IO (..)) +import GHC.Word (Word64 (..)) + + +------------------------------------------------------------------------ +-- Result types — unboxed, identical to flatparse +------------------------------------------------------------------------ + +{- | Parser result. We drop the State# token from the result type +since it's threaded implicitly by IO. The parser function signature +still takes State# and must thread it, but Res# is just the +unboxed sum of outcomes. +-} +type Res# e a = + (# + (# a, Addr# #) | + (# #) | + (# e #) + #) + + +pattern OK# :: a -> Addr# -> Res# e a +pattern OK# a s = (# (# a, s #) | | #) + + +pattern Fail# :: Res# e a +pattern Fail# = (# | (# #) | #) + + +pattern Err# :: e -> Res# e a +pattern Err# e = (# | | (# e #) #) + + +{-# COMPLETE OK#, Fail#, Err# #-} + + +------------------------------------------------------------------------ +-- Parser mode (compile-time dispatch for ensure slow path) +------------------------------------------------------------------------ + +{- | Whole-input mode. When bounds check fails, just return 'Fail#'. +Identical codegen to flatparse — zero suspension overhead. +-} +data Pure + + +{- | Streaming mode. When bounds check fails, suspend via @control0#@ +to wait for more data from the transport. +-} +data Stream + + +{- | Type class dispatching the ensure slow path. +Resolved at compile time via specialization. +-} +class ParserMode (m :: Type) where + onEnsureFail + :: ParserEnv + -> Addr# + -> Addr# + -> Int# + -> State# RealWorld + -> StRes# e () + + + {- | Mode-polymorphic checkpoint. In streaming mode, advances the + transport's tail to the parser's current position (freeing ring + space behind it). In whole-input mode, a no-op — there is no + ring to refill. + + Used by 'Wireform.Parser.takeBs' / 'takeBsCopy' to drain reads + larger than the ring without deadlocking. + -} + modeCheckpoint :: Parser m e () + + +instance ParserMode Pure where + onEnsureFail _env _eob _s _n st = (# st, Fail# #) + {-# INLINE onEnsureFail #-} + modeCheckpoint = Parser \_env _eob s st -> (# st, OK# () s #) + {-# INLINE modeCheckpoint #-} + + +instance ParserMode Stream where + onEnsureFail = ensureNSlow + {-# INLINE onEnsureFail #-} + modeCheckpoint = checkpoint + {-# INLINE modeCheckpoint #-} + + +------------------------------------------------------------------------ +-- Parser environment (carries mutable end pointer for streaming) +------------------------------------------------------------------------ + +{- | Shared context for a parse run. + +Three mutable cells (pointer-to-Ptr fields) carry state that the +driver updates between parser resumptions: + + * 'peEndPtr' — the producer's head pointer, refreshed after + a 'StepSuspend' / 'StepCheckpoint' round-trip. + * 'peAnchorPos' — the absolute byte offset of the parser's most + recent /anchor/. Initially the parse run's + @startPos@; updated whenever the driver wraps + 'cur' back into the first mapping (currently + on every 'StepCheckpoint' / 'StepSuspend' resume). + * 'peAnchorCur' — the cur pointer that pairs with 'peAnchorPos'. + +Logical position is computed as @anchorPos + (cur - anchorCur)@ +(see 'curToPos'). Re-anchoring on cur-resets keeps that identity +correct even when a single parse run drains more bytes than the +ring buffer can hold (the parser's logical 'cur' advances past +@anchorCur + ringSize@, the driver wraps it back, and the anchor +shifts forward by the wrap amount). + +== Concurrency model + +These cells are mutable but require neither locks nor memory +barriers, because the writer (the driver) and the reader (the +parser) never run at the same time: + + * Parser writes are confined to 'resumeContinue' closures, which + run /while the parser is suspended/ (between 'control0#' + reifying the continuation and 'prompt#' re-entering it). + * The parser cannot observe a cell mid-update because it is not + executing at all during the update. + * The update of all three cells happens-before the @prompt#@ + that resumes the parser, by virtue of being sequenced in the + same @IO@ action. + +So the three writes look atomic to the parser even though they are +three separate stores: suspension is the synchronization +primitive. No GHC fence is needed (single-threaded use; the +transport contract requires that producer and consumer share a +thread for a given parse run). +-} +data ParserEnv = ParserEnv + { peEndPtr :: {-# UNPACK #-} !(Ptr (Ptr Word8)) + -- ^ Mutable end pointer (single deref to read, updated on resume). + , peBaseAddr :: {-# UNPACK #-} !(Ptr Word8) + -- ^ Ring base address. + , peMask :: {-# UNPACK #-} !Int + -- ^ @ringSize - 1@ (or 'maxBound' in whole-input mode). + , peAnchorPos :: {-# UNPACK #-} !(Ptr Word64) + {- ^ Mutable anchor: absolute byte position corresponding to + 'peAnchorCur'. + -} + , peAnchorCur :: {-# UNPACK #-} !(Ptr (Ptr Word8)) + -- ^ Mutable anchor: cur address corresponding to 'peAnchorPos'. + , peBackingFp :: !ForeignPtrContents + -- ^ Keeps the backing memory alive for zero-copy slices. + , peTag :: Any + {- ^ PromptTag# stored as Any via unsafeCoerce#. + Only meaningful for streaming mode. + -} + } + + +-- | Recover the typed PromptTag# from the env. Only call from ensureNSlow. +{-# INLINE pePromptTag #-} +pePromptTag :: forall e r. ParserEnv -> PromptTag# (Step e r) +pePromptTag env = unsafeCoerce# (peTag env) + + +-- | Store a PromptTag# into an Any-typed field. +{-# INLINE tagToAny #-} +tagToAny :: PromptTag# a -> Any +tagToAny t = unsafeCoerce# t + + +{- | Convert a cur address to an absolute logical position. Reads +the mutable anchor cells, so it must be threaded through 'State#'. +-} +{-# INLINE curToPos #-} +curToPos + :: ParserEnv + -> Addr# + -> State# RealWorld + -> (# State# RealWorld, Word64 #) +curToPos env cur s0 = + let !(Ptr posCell) = peAnchorPos env + !(Ptr curCell#) = peAnchorCur env + in case readWord64OffAddr# posCell 0# s0 of + (# s1, ap# #) -> + case readAddrOffAddr# curCell# 0# s1 of + (# s2, ac# #) -> + let !off = I# (minusAddr# cur ac#) + !pos = W64# ap# + fromIntegral off + in (# s2, pos #) + + +-- | Boxed wrapper for IO contexts. +curToPosIO :: ParserEnv -> Ptr Word8 -> IO Word64 +curToPosIO env (Ptr cur#) = IO \s -> curToPos env cur# s +{-# INLINE curToPosIO #-} + + +{- | Write a fresh @(anchorPos, anchorCur)@ pair. Called by the driver +whenever it wraps the parser's cur back into the first mapping +(StepCheckpoint resume, StepSuspend resume). +-} +writeAnchor# + :: ParserEnv + -> Word64 + -> Addr# + -> State# RealWorld + -> State# RealWorld +writeAnchor# env (W64# pos#) cur# s0 = + let !(Ptr posCell) = peAnchorPos env + !(Ptr curCell#) = peAnchorCur env + !s1 = writeWord64OffAddr# posCell 0# pos# s0 + !s2 = writeAddrOffAddr# curCell# 0# cur# s1 + in s2 +{-# INLINE writeAnchor# #-} + + +writeAnchor :: ParserEnv -> Word64 -> Ptr Word8 -> IO () +writeAnchor env pos (Ptr cur#) = IO \s -> (# writeAnchor# env pos cur# s, () #) +{-# INLINE writeAnchor #-} + + +{-# INLINE readEnd# #-} +readEnd# :: ParserEnv -> State# RealWorld -> (# State# RealWorld, Addr# #) +readEnd# env s = + let !(Ptr p) = peEndPtr env + in case readAddrOffAddr# p 0# s of + (# s', a #) -> (# s', a #) + + +{-# INLINE writeEnd# #-} +writeEnd# :: ParserEnv -> Addr# -> State# RealWorld -> State# RealWorld +writeEnd# env val s = + let !(Ptr p) = peEndPtr env + in writeAddrOffAddr# p 0# val s + + +-- Boxed wrappers for driver code +readEnd :: ParserEnv -> IO (Ptr Word8) +readEnd env = IO \s -> case readEnd# env s of + (# s', a #) -> (# s', Ptr a #) +{-# INLINE readEnd #-} + + +writeEnd :: ParserEnv -> Ptr Word8 -> IO () +writeEnd env (Ptr a) = IO \s -> (# writeEnd# env a s, () #) +{-# INLINE writeEnd #-} + + +------------------------------------------------------------------------ +-- Step / Resume (driver protocol) +------------------------------------------------------------------------ + +data Step e r + = StepDone {-# UNPACK #-} !Word64 r + | StepFail {-# UNPACK #-} !Word64 + | StepErr {-# UNPACK #-} !Word64 e + | StepSuspend + {-# UNPACK #-} !Word64 + {-# UNPACK #-} !Int + !(Resume e r) + | StepCheckpoint {-# UNPACK #-} !Word64 !(Resume e r) + + +data Resume e r = Resume + { resumeContinue :: !(Ptr Word8 -> Ptr Word8 -> IO (Step e r)) + , resumeEof :: !(IO (Step e r)) + } + + +------------------------------------------------------------------------ +-- The Parser newtype — flatparse-shaped +------------------------------------------------------------------------ + +{- | A parser consuming bytes from @(cur, eob)@ pointers. + +The representation mirrors flatparse exactly: +@ForeignPtrContents -> Addr# (eob) -> Addr# (cur) -> State# -> Res#@ + +The @ParserEnv@ and @PromptTag#@ are additional parameters for +streaming support. For whole-input parsing ('parseByteString'), +the PromptTag# is allocated but @ensureNSlow@ is never reached. +| Parser result with state token. +-} +type StRes# e a = (# State# RealWorld, Res# e a #) + + +{- | The parser type. Parameterized by mode @m@ ('Pure' or 'Stream'), +error type @e@, and result type @a@. + +Four args on the hot path: env, eob, cur, state. The mode @m@ is +phantom — it controls which 'ParserMode' instance is used for +bounds-check failures, resolved at compile time. +-} +newtype Parser (m :: Type) e a = Parser + { runParser# + :: ParserEnv + -> Addr# -- eob (end of buffer / current end) + -> Addr# -- cur (current position) + -> State# RealWorld + -> StRes# e a + } + + +instance Functor (Parser m e) where + fmap f (Parser g) = Parser \env eob s st -> + case g env eob s st of + (# st', OK# a s' #) -> let !b = f a in (# st', OK# b s' #) + (# st', x #) -> (# st', unsafeCoerce# x #) + {-# INLINE fmap #-} + + + a' <$ Parser g = Parser \env eob s st -> + case g env eob s st of + (# st', OK# _ s' #) -> (# st', OK# a' s' #) + (# st', x #) -> (# st', unsafeCoerce# x #) + {-# INLINE (<$) #-} + + +instance Applicative (Parser m e) where + pure !a = Parser \_ _ s st -> (# st, OK# a s #) + {-# INLINE pure #-} + + + Parser ff <*> Parser fa = Parser \env eob s st -> + case ff env eob s st of + (# st', OK# f s' #) -> case fa env eob s' st' of + (# st'', OK# a s'' #) -> let !b = f a in (# st'', OK# b s'' #) + (# st'', x #) -> (# st'', unsafeCoerce# x #) + (# st', x #) -> (# st', unsafeCoerce# x #) + {-# INLINE (<*>) #-} + + + Parser fa *> Parser fb = Parser \env eob s st -> + case fa env eob s st of + (# st', OK# _ s' #) -> fb env eob s' st' + (# st', x #) -> (# st', unsafeCoerce# x #) + {-# INLINE (*>) #-} + + + Parser fa <* Parser fb = Parser \env eob s st -> + case fa env eob s st of + (# st', OK# a s' #) -> case fb env eob s' st' of + (# st'', OK# _ s'' #) -> (# st'', OK# a s'' #) + (# st'', x #) -> (# st'', unsafeCoerce# x #) + (# st', x #) -> (# st', unsafeCoerce# x #) + {-# INLINE (<*) #-} + + +instance Monad (Parser m e) where + Parser fa >>= f = Parser \env eob s st -> + case fa env eob s st of + (# st', OK# a s' #) -> runParser# (f a) env eob s' st' + (# st', x #) -> (# st', unsafeCoerce# x #) + {-# INLINE (>>=) #-} + + + (>>) = (*>) + {-# INLINE (>>) #-} + + +instance MonadFail (Parser m e) where + fail _ = Parser \_ _ _ st -> (# st, Fail# #) + {-# INLINE fail #-} + + +instance ParserMode m => Alternative (Parser m e) where + empty = Parser \_ _ _ st -> (# st, Fail# #) + {-# INLINE empty #-} + (Parser f) <|> (Parser g) = Parser \env eob s st -> + case f env eob s st of + (# st', Fail# #) -> g env eob s st' + x -> x + {-# INLINE (<|>) #-} + + +instance ParserMode m => MonadPlus (Parser m e) + + +------------------------------------------------------------------------ +-- ensureN#: bounds check + suspension +------------------------------------------------------------------------ + +{- | Require @n#@ bytes available from @cur@. + +Fast path: single register comparison (identical to flatparse). +Semi-fast: re-read 'peEndPtr' to catch stale @eob@ after a +prior streaming resume; no function-call overhead. +Slow path: suspends to the driver via @control0#@. +-} +ensureN# :: forall m e. ParserMode m => Int# -> Parser m e () +ensureN# n# = Parser \env eob s st -> + case n# <=# minusAddr# eob s of + 1# -> (# st, OK# () s #) + _ -> case readEnd# env st of + (# st', eob' #) -> case n# <=# minusAddr# eob' s of + 1# -> (# st', OK# () s #) + _ -> onEnsureFail @m env eob' s n# st' +{-# INLINE ensureN# #-} + + +{- | Streaming slow path. + +Before suspending via @control0#@, re-read the mutable end pointer. +After a prior resume, @>>=@ threads the stale @eob@ that was captured +in its closure, so the fast-path comparison in 'ensureN#' fails even +though new data is already available in the ring. The re-read here +catches that case and avoids a redundant suspension round-trip. +-} +ensureNSlow + :: ParserEnv + -> Addr# + -> Addr# + -> Int# + -> State# RealWorld + -> StRes# e () +ensureNSlow env _eob s n# st = + case readEnd# env st of + (# st0, eob' #) -> + case n# <=# minusAddr# eob' s of + 1# -> (# st0, OK# () s #) + _ -> + let tag :: PromptTag# (Step Any Any) + tag = unsafeCoerce# (peTag env) + in case curToPos env s st0 of + (# st0', pos #) -> + let !needed = I# n# + in control0# + tag + ( \k st1 -> + let resume = + Resume + { resumeContinue = \(Ptr newCur) (Ptr newEnd) -> + IO \st2 -> + -- The driver always wraps newCur back + -- into the first mapping; re-anchor + -- so 'curToPos' stays correct after + -- the wrap. Both stores are safe + -- without a fence: the parser is + -- suspended (between this 'control0#' + -- and the 'prompt#' below) while + -- they execute. + let st3 = writeEnd# env newEnd st2 + st4 = writeAnchor# env pos newCur st3 + in prompt# tag (\st5 -> k (\st6 -> (# st6, OK# () newCur #)) st5) st4 + , resumeEof = + IO \st2 -> prompt# tag (\st3 -> k (\st4 -> (# st4, Fail# #)) st3) st2 + } + in (# st1, StepSuspend pos needed resume #) + ) + st0' +{-# NOINLINE ensureNSlow #-} + + +------------------------------------------------------------------------ +-- checkpoint +------------------------------------------------------------------------ + +-- | Checkpoint is only meaningful in streaming mode. +checkpoint :: Parser Stream e () +checkpoint = Parser \env _ s st0 -> + let tag :: PromptTag# (Step Any Any) + tag = unsafeCoerce# (peTag env) + in case curToPos env s st0 of + (# st0', pos #) -> + control0# + tag + ( \k st1 -> + let resume = + Resume + { resumeContinue = \(Ptr newCur) (Ptr newEnd) -> + IO \st2 -> + -- Re-anchor: the driver has wrapped newCur + -- back to @base + (pos .&. mask)@ which + -- decouples the cur pointer from the original + -- anchor. Without re-anchoring, 'curToPos' + -- would compute the wrong absolute position + -- on the next call. + -- + -- These two stores (end pointer + anchor + -- pair) are not behind any lock or fence: + -- they run inside the resume body, which is + -- only invoked while the parser is suspended + -- between 'control0#' and the matching + -- 'prompt#' below. The parser cannot observe + -- a partial update because it is not running + -- at all during the update — suspension is + -- the synchronization mechanism. + let st3 = writeEnd# env newEnd st2 + st4 = writeAnchor# env pos newCur st3 + in prompt# tag (\st5 -> k (\st6 -> (# st6, OK# () newCur #)) st5) st4 + , resumeEof = + IO \st2 -> prompt# tag (\st3 -> k (\st4 -> (# st4, Fail# #)) st3) st2 + } + in (# st1, StepCheckpoint pos resume #) + ) + st0' diff --git a/src/Wireform/Parser/Mark.hs b/src/Wireform/Parser/Mark.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Mark.hs @@ -0,0 +1,42 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +module Wireform.Parser.Mark ( + Mark (..), + mark, + restore, + release, +) where + +import Data.Bits ((.&.)) +import Data.Word (Word64) +import Foreign.Ptr (plusPtr) +import GHC.Exts +import Wireform.Parser.Internal + + +newtype Mark = Mark {unMark :: Word64} + deriving stock (Eq, Ord, Show) + + +mark :: Parser m e Mark +mark = Parser \env _ s st -> + case curToPos env s st of + (# st', pos #) -> (# st', OK# (Mark pos) s #) +{-# INLINE mark #-} + + +restore :: Mark -> Parser m e () +restore (Mark pos) = Parser \env _ _ st -> + let !offset = fromIntegral pos .&. peMask env + !(Ptr newCur) = peBaseAddr env `plusPtr` offset + in (# st, OK# () newCur #) +{-# INLINE restore #-} + + +release :: Mark -> Parser m e () +release _ = Parser \_ _ s st -> (# st, OK# () s #) +{-# INLINE release #-} diff --git a/src/Wireform/Parser/Position.hs b/src/Wireform/Parser/Position.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Position.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +module Wireform.Parser.Position ( + Pos (..), + Span (..), + subPos, + getPos, + withSpan, + byteStringOf, + spanToByteString, + inSpan, +) where + +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import Data.ByteString.Internal qualified as BSI +import Data.Word (Word64) +import Foreign.Marshal.Utils (copyBytes) +import Foreign.Ptr (plusPtr) +import GHC.Exts +import GHC.ForeignPtr (ForeignPtr (..)) +import System.IO.Unsafe (unsafeDupablePerformIO) +import Wireform.Parser.Internal + + +newtype Pos = Pos {unPos :: Word64} + deriving stock (Eq, Ord, Show) + deriving newtype (Num) + + +data Span = Span !Pos !Pos + deriving stock (Eq, Show) + + +subPos :: Pos -> Pos -> Int +subPos (Pos a) (Pos b) = fromIntegral (a - b) + + +getPos :: Parser m e Pos +getPos = Parser \env _ s st -> + case curToPos env s st of + (# st', pos #) -> (# st', OK# (Pos pos) s #) +{-# INLINE getPos #-} + + +withSpan :: Parser m e a -> (a -> Span -> Parser m e b) -> Parser m e b +withSpan (Parser p) f = Parser \env eob s st -> + case p env eob s st of + (# st', OK# a s' #) -> + case curToPos env s st' of + (# st'', startPos #) -> + case curToPos env s' st'' of + (# st''', endPos #) -> + let !sp = Span (Pos startPos) (Pos endPos) + in runParser# (f a sp) env eob s' st''' + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE withSpan #-} + + +byteStringOf :: Parser m e a -> Parser m e ByteString +byteStringOf (Parser p) = Parser \env eob s st -> + case p env eob s st of + (# st', OK# _ s' #) -> + let !len = I# (minusAddr# s' s) + !bs = BSI.BS (ForeignPtr s (peBackingFp env)) len + in (# st', OK# bs s' #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE byteStringOf #-} + + +spanToByteString :: Span -> Parser m e ByteString +spanToByteString (Span (Pos start) (Pos end)) = Parser \env _ s st -> + let !len = fromIntegral (end - start) + base = peBaseAddr env + mask = peMask env + off = fromIntegral start .&. mask + !(Ptr ptr) = base `plusPtr` off + !bs = unsafeDupablePerformIO (BSI.create len \dst -> copyBytes dst (Ptr ptr) len) + in (# st, OK# bs s #) + + +{- | Run a parser within an explicitly bounded byte window. +Temporarily restricts @eob@ to the span's end position, then +restores it on completion. +-} +inSpan :: Span -> Parser m e a -> Parser m e a +inSpan (Span (Pos start) (Pos end)) (Parser p) = Parser \env _ s st -> + let base = peBaseAddr env + mask = peMask env + !endOff = fromIntegral end .&. mask + !(Ptr spanEnd) = base `plusPtr` endOff + !startOff = fromIntegral start .&. mask + !(Ptr spanStart) = base `plusPtr` startOff + in case p env spanEnd spanStart st of + (# st', OK# a _ #) -> (# st', OK# a s #) + (# st', x #) -> (# st', unsafeCoerce# x #) +{-# INLINE inSpan #-} diff --git a/src/Wireform/Parser/Stateful.hs b/src/Wireform/Parser/Stateful.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Stateful.hs @@ -0,0 +1,252 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +{- | Stateful parser variant with a reader environment @r@ and mutable +state @s@, following flatparse's @FlatParse.Stateful@ design. + +State is stored in a raw @MutVar#@ — no IORef boxing overhead. +Reader is a plain boxed value. +-} +module Wireform.Parser.Stateful ( + ParserS (..), + + -- * Running + runParserS, + parseByteStringS, + + -- * Reader operations + ask, + asks, + local, + + -- * State operations + get, + put, + modify, + modify', + + -- * Lifting basic parsers + liftParser, + + -- * Re-exports + module Wireform.Parser, +) where + +import Data.ByteString.Internal qualified as BSI +import Data.Word (Word64, Word8) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Marshal.Alloc (allocaBytes) +import Foreign.Ptr (castPtr, plusPtr) +import Foreign.Storable (poke) +import GHC.Exts +import GHC.ForeignPtr (ForeignPtr (..)) +import GHC.IO (IO (..)) +import System.IO.Unsafe (unsafeDupablePerformIO) +import Wireform.Parser +import Wireform.Parser.Internal ( + Parser (..), + ParserEnv (..), + Resume (..), + Step (..), + tagToAny, + pattern Err#, + pattern Fail#, + pattern OK#, + type StRes#, + ) + + +------------------------------------------------------------------------ +-- Stateful parser type +------------------------------------------------------------------------ + +{- | Parser with reader @r@ and mutable state @s@ stored in a raw +@MutVar#@. One pointer deref per 'get'/'put', no IORef boxing. +-} +newtype ParserS r s e a = ParserS + { runParserS# + :: forall result + . PromptTag# (Step e result) + -> ParserEnv + -> r + -> MutVar# RealWorld s + -> Addr# + -> Addr# + -> State# RealWorld + -> StRes# e a + } + + +instance Functor (ParserS r s e) where + fmap f (ParserS g) = ParserS \tag env r mv eob s rw -> + case g tag env r mv eob s rw of + (# rw', OK# a s' #) -> let !b = f a in (# rw', OK# b s' #) + (# rw', x #) -> (# rw', unsafeCoerce# x #) + {-# INLINE fmap #-} + + +instance Applicative (ParserS r s e) where + pure !a = ParserS \_tag _env _r _mv _eob s rw -> (# rw, OK# a s #) + {-# INLINE pure #-} + ParserS ff <*> ParserS fa = ParserS \tag env r mv eob s rw -> + case ff tag env r mv eob s rw of + (# rw', OK# f s' #) -> case fa tag env r mv eob s' rw' of + (# rw'', OK# a s'' #) -> let !b = f a in (# rw'', OK# b s'' #) + (# rw'', x #) -> (# rw'', unsafeCoerce# x #) + (# rw', x #) -> (# rw', unsafeCoerce# x #) + {-# INLINE (<*>) #-} + ParserS fa *> ParserS fb = ParserS \tag env r mv eob s rw -> + case fa tag env r mv eob s rw of + (# rw', OK# _ s' #) -> fb tag env r mv eob s' rw' + (# rw', x #) -> (# rw', unsafeCoerce# x #) + {-# INLINE (*>) #-} + + +instance Monad (ParserS r s e) where + ParserS fa >>= f = ParserS \tag env r mv eob s rw -> + case fa tag env r mv eob s rw of + (# rw', OK# a s' #) -> runParserS# (f a) tag env r mv eob s' rw' + (# rw', x #) -> (# rw', unsafeCoerce# x #) + {-# INLINE (>>=) #-} + (>>) = (*>) + {-# INLINE (>>) #-} + + +------------------------------------------------------------------------ +-- Reader operations +------------------------------------------------------------------------ + +ask :: ParserS r s e r +ask = ParserS \_tag _env r _mv _eob s rw -> (# rw, OK# r s #) +{-# INLINE ask #-} + + +asks :: (r -> a) -> ParserS r s e a +asks f = ParserS \_tag _env r _mv _eob s rw -> let !a = f r in (# rw, OK# a s #) +{-# INLINE asks #-} + + +local :: (r -> r) -> ParserS r s e a -> ParserS r s e a +local f (ParserS p) = ParserS \tag env r mv eob s rw -> + p tag env (f r) mv eob s rw +{-# INLINE local #-} + + +------------------------------------------------------------------------ +-- State operations (raw MutVar#, no IORef) +------------------------------------------------------------------------ + +get :: ParserS r s e s +get = ParserS \_tag _env _r mv _eob s rw -> + case readMutVar# mv rw of + (# rw', val #) -> (# rw', OK# val s #) +{-# INLINE get #-} + + +put :: s -> ParserS r s e () +put !v = ParserS \_tag _env _r mv _eob s rw -> + case writeMutVar# mv v rw of + rw' -> (# rw', OK# () s #) +{-# INLINE put #-} + + +modify :: (s -> s) -> ParserS r s e () +modify f = ParserS \_tag _env _r mv _eob s rw -> + case readMutVar# mv rw of + (# rw', val #) -> case writeMutVar# mv (f val) rw' of + rw'' -> (# rw'', OK# () s #) +{-# INLINE modify #-} + + +modify' :: (s -> s) -> ParserS r s e () +modify' f = ParserS \_tag _env _r mv _eob s rw -> + case readMutVar# mv rw of + (# rw', val #) -> + let !val' = f val + in case writeMutVar# mv val' rw' of + rw'' -> (# rw'', OK# () s #) +{-# INLINE modify' #-} + + +------------------------------------------------------------------------ +-- Lifting basic parsers +------------------------------------------------------------------------ + +liftParser :: Parser m e a -> ParserS r s e a +liftParser (Parser p) = ParserS \_tag env _r _mv eob s rw -> + p env eob s rw +{-# INLINE liftParser #-} + + +------------------------------------------------------------------------ +-- Running +------------------------------------------------------------------------ + +runParserS + :: ParserS r s e a + -> r + -> s + -> BSI.ByteString + -> Either (ParseError e) (a, s) +runParserS p r s0 b = unsafeDupablePerformIO $ do + let !(BSI.BS (ForeignPtr buf# fp) (I# len#)) = b + !end# = plusAddr# buf# len# + + withForeignPtr (ForeignPtr buf# fp) \_ -> + allocaBytes 24 \cells -> do + let !endPtr = cells + !anchorPos = cells `plusPtr` 8 + !anchorCur = cells `plusPtr` 16 + poke (castPtr endPtr :: Ptr (Ptr Word8)) (Ptr end#) + poke (castPtr anchorPos :: Ptr Word64) 0 + poke (castPtr anchorCur :: Ptr (Ptr Word8)) (Ptr buf#) + + IO \rw0 -> case newMutVar# s0 rw0 of + (# rw1, mv #) -> case newPromptTag# rw1 of + (# rw2, (tag :: PromptTag# (Step e a)) #) -> + let env = + ParserEnv + { peEndPtr = castPtr endPtr + , peBaseAddr = Ptr buf# + , peMask = maxBound + , peAnchorPos = castPtr anchorPos + , peAnchorCur = castPtr anchorCur + , peBackingFp = fp + , peTag = tagToAny tag + } + body :: State# RealWorld -> (# State# RealWorld, Step e a #) + body rw = case runParserS# p tag env r mv end# buf# rw of + (# rw', OK# a _cur #) -> (# rw', StepDone 0 a #) + (# rw', Fail# #) -> (# rw', StepFail 0 #) + (# rw', Err# e #) -> (# rw', StepErr 0 e #) + in case prompt# tag body rw2 of + (# rw3, step #) -> case readMutVar# mv rw3 of + (# rw4, finalState #) -> + (# rw4, classifyStep step finalState #) + where + classifyStep (StepDone _ a) s = Right (a, s) + classifyStep (StepFail pos) _ = Left (ParseFail pos) + classifyStep (StepErr pos e) _ = Left (ParseErr pos e) + classifyStep (StepSuspend _ _ res) s = + unsafeDupablePerformIO $ + resumeEof res >>= \step -> pure (classifyStep step s) + classifyStep (StepCheckpoint _ res) s = + unsafeDupablePerformIO $ + resumeEof res >>= \step -> pure (classifyStep step s) + + +parseByteStringS + :: ParserS r s e a + -> r + -> s + -> BSI.ByteString + -> Either (ParseError e) (a, s) +parseByteStringS = runParserS diff --git a/src/Wireform/Parser/Switch.hs b/src/Wireform/Parser/Switch.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Parser/Switch.hs @@ -0,0 +1,288 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UnboxedSums #-} +{-# LANGUAGE UnboxedTuples #-} + +-- \$(switch [| case _ of +-- \"if\" -> pure TIf +-- \"else\" -> pure TElse +-- _ -> identifier |]) +-- @ + +{- | Efficient literal branching using Template Haskell. + +Compiles string-literal case branches into a trie of primitive byte +comparisons with grouped bounds checks. + +@ +-} +module Wireform.Parser.Switch ( + switch, + switchWithPost, + + -- * TH literal splices + char, + string, + + -- * Helpers used by generated code (not for direct use) + switchFailed, + switchBranch, + switchAnyWord8Unsafe, + switchPeekWord8Unsafe, + switchSkip1, +) where + +import Control.Monad (forM) +import Data.ByteString.Char8 qualified as BSC +import Data.Char (ord) +import Data.Foldable (foldl') +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as M +import Data.Word (Word8) +import GHC.Exts +import GHC.Word (Word8 (..)) +import Language.Haskell.TH +import Language.Haskell.TH.Syntax (Lift (..)) +import Wireform.Parser (byteString, word8) +import Wireform.Parser.Internal + + +------------------------------------------------------------------------ +-- Helpers called by generated code +------------------------------------------------------------------------ + +-- | Fail without consuming input. +switchFailed :: Parser m e a +switchFailed = Parser \_ _ _ st -> (# st, Fail# #) +{-# INLINE switchFailed #-} + + +-- | Unsafe read — caller must have ensured at least 1 byte. +switchAnyWord8Unsafe :: Parser m e Word8 +switchAnyWord8Unsafe = Parser \_ _ s st -> + case indexWord8OffAddr# s 0# of + w# -> (# st, OK# (W8# w#) (plusAddr# s 1#) #) +{-# INLINE switchAnyWord8Unsafe #-} + + +{- | Peek the next byte without consuming. Used by the trie code +generator so that a wildcard-branch falling back to the +current node's terminal action does not eat a byte that +belongs to whatever runs after the switch. + +(Previously the generator used 'switchAnyWord8Unsafe' for the +byte dispatch, which consumed the byte and then ran the +terminal action /past/ that byte. For input like @\"1,…\"@ on a +switch with literals @\"1\", \"1.0\", \"1.00\"@, that meant the +@\"1\"@ terminal action ran from after the comma — silently +breaking every quality-weighted Accept-* parser whose first +entry used @q=1@.) +-} +switchPeekWord8Unsafe :: Parser m e Word8 +switchPeekWord8Unsafe = Parser \_env _eob s st -> + case indexWord8OffAddr# s 0# of + w# -> (# st, OK# (W8# w#) s #) +{-# INLINE switchPeekWord8Unsafe #-} + + +{- | Skip one byte (used after 'switchPeekWord8Unsafe' has +decided a child branch matches). +-} +switchSkip1 :: Parser m e () +switchSkip1 = Parser \_env _eob s st -> + (# st, OK# () (plusAddr# s 1#) #) +{-# INLINE switchSkip1 #-} + + +-- | Branch on an ensure check: if enough bytes, run @t@; else @f@. +switchBranch :: Int -> Parser m e a -> Parser m e a -> Parser m e a +switchBranch (I# n#) (Parser t) (Parser f) = Parser \env eob s st -> + case n# <=# minusAddr# eob s of + 1# -> t env eob s st + _ -> f env eob s st +{-# INLINE switchBranch #-} + + +------------------------------------------------------------------------ +-- Trie +------------------------------------------------------------------------ + +type Rule = Maybe Int + + +data Trie a = Branch !a !(Map Word8 (Trie a)) + + +nilTrie :: Trie Rule +nilTrie = Branch Nothing mempty + + +insertTrie :: Int -> [Word8] -> Trie Rule -> Trie Rule +insertTrie rule = go + where + go [] (Branch r ts) = Branch (Just $ maybe rule (min rule) r) ts + go (c : cs) (Branch r ts) = + Branch r (M.alter (Just . maybe (go cs nilTrie) (go cs)) c ts) + + +listToTrie :: [(Int, String)] -> Trie Rule +listToTrie = foldl' (\t (!r, !s) -> insertTrie r (charToBytes s) t) nilTrie + + +charToBytes :: String -> [Word8] +charToBytes = concatMap go + where + go c + | n < 0x80 = [fromIntegral n] + | n < 0x800 = + [ fromIntegral (0xC0 + n `div` 64) + , fromIntegral (0x80 + n `mod` 64) + ] + | otherwise = [fromIntegral n] -- simplified; ASCII-only for v1 + where + n = ord c + + +mindepths :: Trie Rule -> Trie (Rule, Int) +mindepths (Branch rule ts) + | M.null ts = Branch (rule, 0) mempty + | otherwise = + let ts' = M.map mindepths ts + d = + minimum $ + M.elems $ + M.map + ( \(Branch (r, depth) _) -> + maybe (depth + 1) (const 1) r + ) + ts' + in Branch (rule, d) ts' + +------------------------------------------------------------------------ +-- TH code generation +------------------------------------------------------------------------ + +#if MIN_VERSION_base(4,15,0) +mkDoE :: [Stmt] -> Exp +mkDoE = DoE Nothing +#else +mkDoE :: [Stmt] -> Exp +mkDoE = DoE +#endif + + +switch :: Q Exp -> Q Exp +switch = switchWithPost Nothing + + +switchWithPost :: Maybe (Q Exp) -> Q Exp -> Q Exp +switchWithPost postAction qexp = do + post <- sequence postAction + (cases, deflt) <- parseSwitchExp qexp + genSwitch post cases deflt + + +parseSwitchExp :: Q Exp -> Q ([(String, Exp)], Maybe Exp) +parseSwitchExp qexp = + qexp >>= \case + CaseE (UnboundVarE _) [] -> fail "switch: empty case list" + CaseE (UnboundVarE _) matches -> do + let (ini, lst) = (init matches, last matches) + cases <- forM ini \case + Match (LitP (StringL s)) (NormalB rhs) [] -> pure (s, rhs) + _ -> fail "switch: expected string literal pattern" + (cases', deflt) <- case lst of + Match (LitP (StringL s)) (NormalB rhs) [] -> pure (cases <> [(s, rhs)], Nothing) + Match WildP (NormalB rhs) [] -> pure (cases, Just rhs) + _ -> fail "switch: expected string literal or wildcard pattern" + pure (cases', deflt) + _ -> fail "switch: expected 'case _ of' expression" + + +genSwitch :: Maybe Exp -> [(String, Exp)] -> Maybe Exp -> Q Exp +genSwitch post cases deflt = do + let indexed = zip [0 :: Int ..] cases + trie = mindepths (listToTrie [(i, s) | (i, (s, _)) <- indexed]) + + let ruleMap = + M.fromList $ + (Nothing, VarE 'switchFailed) + : [(Just i, applyPost post rhs) | (i, (_, rhs)) <- indexed] + defltExp = maybe (VarE 'switchFailed) id deflt + + bindings <- forM (M.toList (M.insert Nothing defltExp ruleMap)) \(k, e) -> do + n <- newName ("r" <> maybe "D" show k) + pure (k, n, e) + + let nameMap = M.fromList [(k, n) | (k, n, _) <- bindings] + + body <- genTrieCode nameMap trie + letE + [valD (varP n) (normalB (pure e)) [] | (_, n, e) <- bindings] + (pure body) + + +genTrieCode :: Map (Maybe Int) Name -> Trie (Rule, Int) -> Q Exp +genTrieCode names = go 0 + where + ruleName r = case M.lookup r names of + Just n -> n + Nothing -> error "switch: missing rule" + + go :: Int -> Trie (Rule, Int) -> Q Exp + go ensured (Branch (rule, depth) ts) + | M.null ts = pure (VarE (ruleName rule)) + | otherwise = do + let need = ensured < 1 + branches <- forM (M.toList ts) \(w, sub) -> do + e <- go (if need then depth - 1 else ensured - 1) sub + pure (w, InfixE (Just (VarE 'switchSkip1)) (VarE '(>>)) (Just e)) + + let defE = VarE (ruleName rule) + body = + mkDoE + [ BindS (VarP (mkName "c")) (VarE 'switchPeekWord8Unsafe) + , NoBindS + ( CaseE + (VarE (mkName "c")) + ( [ Match (LitP (IntegerL (fromIntegral w))) (NormalB e) [] + | (w, e) <- branches + ] + <> [Match WildP (NormalB defE) []] + ) + ) + ] + + if need + then [|switchBranch depth $(pure body) $(pure defE)|] + else pure body + + +applyPost :: Maybe Exp -> Exp -> Exp +applyPost Nothing rhs = rhs +applyPost (Just p) rhs = InfixE (Just p) (VarE '(>>)) (Just rhs) + + +------------------------------------------------------------------------ +-- TH literal splices (flatparse-compatible) +------------------------------------------------------------------------ + +{- | @$(char \'x\')@ compiles to @word8 0xNN@ for ASCII characters, +or the appropriate multi-byte UTF-8 sequence otherwise. +Only ASCII is supported for now. +-} +char :: Char -> Q Exp +char c + | n < 0x80 = [|word8 $(lift (fromIntegral n :: Word8))|] + | otherwise = error ("Wireform.Parser.Switch.char: non-ASCII character: " <> show c) + where + n = ord c + + +-- | @$(string \"foo\")@ compiles to @byteString \"foo\"@. +string :: String -> Q Exp +string s = [|byteString $(lift (BSC.pack s))|] diff --git a/src/Wireform/Ring.hs b/src/Wireform/Ring.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Ring.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE CApiFFI #-} +{-# LANGUAGE RankNTypes #-} + +{- | Double-mapped (\"magic\") ring buffer. + +A buffer of size @N@ (power of two, page-aligned) backed by a single +shared-memory object mapped twice into a contiguous @2N@-byte virtual +region. Any read of up to @N@ bytes starting anywhere in +@[base, base + N)@ is contiguous in virtual memory — the MMU handles +the wrap transparently. + +The parser's pointer-bumping primitives rely on this property: they +never contain wrap logic. + +== Scoping + +'MagicRing' carries a phantom type parameter @s@ that works like the +@s@ on 'Control.Monad.ST.ST'. 'withMagicRing' is rank-2: + +> withMagicRing :: Int -> (forall s. MagicRing s -> IO a) -> IO a + +Slices produced from the ring ('RingSlice') inherit @s@ and therefore +cannot appear in @a@ — they are unable to outlive the buffer that +backs them and would otherwise alias bytes that subsequent refills +overwrite. When a slice has to leave the scope, 'copyRingSlice' +materialises a fresh 'ByteString' by memcpy. +-} +module Wireform.Ring ( + -- * The ring + MagicRing, + withMagicRing, + ringBase, + ringSize, + ringMask, + MagicRingException (..), + + -- * Slices + RingSlice, + ringSlice, + ringSliceAtPos, + ringSliceLength, + withRingSlice, + peekRingSliceByte, + copyRingSlice, +) where + +import Wireform.Ring.Internal + diff --git a/src/Wireform/Ring/Internal.hs b/src/Wireform/Ring/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Ring/Internal.hs @@ -0,0 +1,315 @@ +{-# LANGUAGE CApiFFI #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE UnboxedTuples #-} + +{- | Internal module exposing the unsafe machinery behind 'MagicRing'. + +The phantom @s@ on 'MagicRing' acts like the @s@ on 'Control.Monad.ST.ST': +it is sealed inside 'withMagicRing' with a rank-2 @forall@, so values +that carry the @s@ — most importantly 'RingSlice' — cannot leak out +of the scope in which the ring is alive (and therefore cannot +dangle through a refill that overwrites the bytes they point at). + +Code that genuinely needs the un-scoped variants ('newMagicRing', +'destroyMagicRing', the raw @ringBase@ pointer) lives here. +Day-to-day callers should prefer the safer surface in "Wireform.Ring". +-} +module Wireform.Ring.Internal ( + -- * The ring + MagicRing (..), + newMagicRing, + destroyMagicRing, + withMagicRing, + ringBase, + ringSize, + ringMask, + MagicRingException (..), + + -- * Slices + RingSlice (..), + ringSlice, + ringSliceAtPos, + ringSliceLength, + ringSliceBase, + withRingSlice, + peekRingSliceByte, + copyRingSlice, + + -- * Unsafe internals (no scoping guarantees) + unsafeRingSliceFromPtr, + unsafeRingScope, +) where + +import Control.Exception (Exception, bracket, throwIO) +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import Data.ByteString.Internal qualified as BSI +import Data.Typeable (Typeable) +import Data.Word (Word64, Word8) +import Foreign.C.Types (CInt (..), CLong (..), CSize (..)) +import Foreign.Marshal.Alloc (alloca) +import Foreign.Marshal.Utils (copyBytes) +import Foreign.Ptr (Ptr, nullPtr, plusPtr) +import Foreign.Storable (Storable (..), peek, peekByteOff, poke, pokeByteOff) + + +------------------------------------------------------------------------ +-- MagicRing +------------------------------------------------------------------------ + +{- | Opaque handle to a double-mapped ring buffer. + +The phantom @s@ is /nominal/: it has no run-time content, only the +type-level role of preventing slices from escaping their ring's +scope (see 'withMagicRing' / 'RingSlice'). +-} +data MagicRing s = MagicRing + { mrBase :: {-# UNPACK #-} !(Ptr Word8) + , mrSize :: {-# UNPACK #-} !Int + } + + +type role MagicRing nominal + + +-- C struct layout: { void *base; size_t size; } +-- On Windows there's an extra HANDLE field, but we don't need it from +-- the Haskell side — hs_ring_destroy handles cleanup. +data CRing = CRing + { crBase :: {-# UNPACK #-} !(Ptr Word8) + , crSize :: {-# UNPACK #-} !CSize + } + + +instance Storable CRing where + sizeOf _ = sizeOf (undefined :: Ptr ()) + sizeOf (undefined :: CSize) + + + -- conservative: ignores padding. On all target platforms + -- both fields are pointer-sized and naturally aligned. + alignment _ = alignment (undefined :: Ptr ()) + peek p = do + b <- peekByteOff p 0 + s <- peekByteOff p (sizeOf (undefined :: Ptr ())) + pure (CRing b s) + poke p (CRing b s) = do + pokeByteOff p 0 b + pokeByteOff p (sizeOf (undefined :: Ptr ())) s + + +foreign import capi unsafe "magic_ring.h hs_ring_create" + c_ring_create :: CSize -> Ptr CRing -> IO CInt + + +foreign import capi unsafe "magic_ring.h hs_ring_destroy" + c_ring_destroy :: Ptr CRing -> IO () + + +data MagicRingException + = MagicRingUnavailable !String + deriving stock (Show, Typeable) + + +instance Exception MagicRingException + + +{- | Allocate a magic ring of at least the requested size (in bytes). +The actual size is rounded up to the nearest page-size multiple and +power of two. Throws 'MagicRingUnavailable' on failure. + +The returned ring is polymorphic in @s@: callers either pair it +with 'destroyMagicRing' manually (no scope safety) or, preferably, +use 'withMagicRing' which binds @s@ inside a rank-2 scope. +-} +newMagicRing :: Int -> IO (MagicRing s) +newMagicRing requested = alloca $ \p -> do + poke p (CRing nullPtr 0) + rc <- c_ring_create (fromIntegral requested) p + if rc /= 0 + then + throwIO + ( MagicRingUnavailable $ + "hs_ring_create failed for requested size " <> show requested + ) + else do + cr <- peek p + pure + MagicRing + { mrBase = crBase cr + , mrSize = fromIntegral (crSize cr) + } + + +-- | Release the ring's virtual mappings. Safe to call multiple times. +destroyMagicRing :: MagicRing s -> IO () +destroyMagicRing mr = alloca $ \p -> do + poke p (CRing (mrBase mr) (fromIntegral (mrSize mr))) + c_ring_destroy p + + +{- | Scoped allocation, modelled after 'Control.Monad.ST.runST'. + +The body is rank-2: the @s@ inside the body is a fresh skolem that +cannot unify with anything outside the @forall@. Any value carrying +that @s@ — slices, sub-buffers, anything we add later — therefore +cannot appear in the result type @a@ and cannot escape. + +> withMagicRing 4096 $ \\ring -> do +> let slice = ringSlice ring 0 16 +> copyRingSlice slice -- OK, copy escapes as a fresh ByteString +> pure slice -- type error: RingSlice s escapes +-} +withMagicRing :: Int -> (forall s. MagicRing s -> IO a) -> IO a +withMagicRing n action = bracket (newMagicRing n) destroyMagicRing action + + +-- | Base pointer to the start of the ring (first of the two mappings). +ringBase :: MagicRing s -> Ptr Word8 +ringBase = mrBase +{-# INLINE ringBase #-} + + +-- | Physical size of the ring in bytes (N, not 2N). +ringSize :: MagicRing s -> Int +ringSize = mrSize +{-# INLINE ringSize #-} + + +{- | Bitmask for cheap @pos .&. ringMask@ indexing. +Always @ringSize - 1@ since the size is a power of two. +-} +ringMask :: MagicRing s -> Int +ringMask mr = mrSize mr - 1 +{-# INLINE ringMask #-} + + +------------------------------------------------------------------------ +-- RingSlice +------------------------------------------------------------------------ + +{- | A contiguous slice into a 'MagicRing'. + +'RingSlice' is deliberately /not/ a 'ByteString': the only way to +materialise a 'ByteString' from one is 'copyRingSlice', which performs +a memcpy into freshly-allocated memory. This makes the cost of +escaping the ring's scope visible at the call site and prevents +accidental retention of a pointer that subsequent refills will +overwrite. + +The phantom @s@ ties the slice to its originating ring's scope; +because 'withMagicRing' seals @s@ inside a rank-2 @forall@, a slice +can never outlive the ring it came from when constructed through +the safe API. + +The base pointer lies inside the first mapping (@[base, base + N)@) +but, by virtue of the double mapping, callers may read up to +'ringSliceLength' bytes contiguously from it without wrap logic. +-} +data RingSlice s = RingSlice + { _rsPtr :: {-# UNPACK #-} !(Ptr Word8) + , _rsLen :: {-# UNPACK #-} !Int + } + + +type role RingSlice nominal + + +{- | Build a slice from a byte offset within the ring and a length. + +The offset is taken modulo 'ringSize' so callers can pass an +absolute consumer position directly. The length may exceed +@ringSize - (offset \`mod\` ringSize)@ — the double mapping makes +the read contiguous as long as @len <= ringSize@. +-} +ringSlice :: MagicRing s -> Int -> Int -> RingSlice s +ringSlice mr offset len = + let !msk = mrSize mr - 1 + !off = offset .&. msk + !p = mrBase mr `plusPtr` off + in RingSlice p len +{-# INLINE ringSlice #-} + + +{- | Build a slice from an absolute (monotonic) producer/consumer +position and a length — the form the transport layer naturally +speaks. +-} +ringSliceAtPos :: MagicRing s -> Word64 -> Int -> RingSlice s +ringSliceAtPos mr pos len = + let !msk = mrSize mr - 1 + !off = fromIntegral pos .&. msk + !p = mrBase mr `plusPtr` off + in RingSlice p len +{-# INLINE ringSliceAtPos #-} + + +-- | Slice length in bytes. +ringSliceLength :: RingSlice s -> Int +ringSliceLength (RingSlice _ n) = n +{-# INLINE ringSliceLength #-} + + +{- | Base pointer of the slice. + +Exposed in @.Internal@ only. The pointer is /not/ tagged with @s@, +so once you reach for it you are responsible for not retaining it +past the ring's lifetime. Prefer 'withRingSlice' for scoped access. +-} +ringSliceBase :: RingSlice s -> Ptr Word8 +ringSliceBase (RingSlice p _) = p +{-# INLINE ringSliceBase #-} + + +{- | Borrow the slice's pointer + length for an IO action. The @s@ +on 'RingSlice' guarantees the call site is inside the ring's scope; +the body is conventionally responsible for not stashing the pointer +somewhere it could outlive that scope. +-} +withRingSlice :: RingSlice s -> (Ptr Word8 -> Int -> IO a) -> IO a +withRingSlice (RingSlice p n) k = k p n +{-# INLINE withRingSlice #-} + + +-- | Read a single byte at an offset within the slice. No bounds check. +peekRingSliceByte :: RingSlice s -> Int -> IO Word8 +peekRingSliceByte (RingSlice p _) i = peek (p `plusPtr` i) +{-# INLINE peekRingSliceByte #-} + + +{- | The escape hatch. Allocates a fresh 'ByteString' and memcpys the +slice's bytes into it. The result has no dependency on the ring, +so it can leave the 'withMagicRing' scope freely. +-} +copyRingSlice :: RingSlice s -> IO ByteString +copyRingSlice (RingSlice p n) + | n <= 0 = pure mempty + | otherwise = BSI.create n $ \dst -> copyBytes dst p n +{-# INLINE copyRingSlice #-} + + +------------------------------------------------------------------------ +-- Unsafe internals +------------------------------------------------------------------------ + +{- | Construct a 'RingSlice' directly from a raw pointer + length. + +This is the back door used by the parser internals, which already +track bounds against the live ring. It bypasses the scope check +because the @s@ is freely chosen by the caller; only call this +when you can guarantee, by other means, that the pointer is +inside a live ring. +-} +unsafeRingSliceFromPtr :: Ptr Word8 -> Int -> RingSlice s +unsafeRingSliceFromPtr = RingSlice +{-# INLINE unsafeRingSliceFromPtr #-} + + +{- | Re-tag a ring under a different @s@. Like 'unsafeCoerce' for the +phantom — useful when wrapping a long-lived (unsafely-allocated) +ring in a scoped action. Do not export from "Wireform.Ring". +-} +unsafeRingScope :: MagicRing s -> MagicRing s' +unsafeRingScope (MagicRing b sz) = MagicRing b sz +{-# INLINE unsafeRingScope #-} diff --git a/src/Wireform/Ring/Pool.hs b/src/Wireform/Ring/Pool.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Ring/Pool.hs @@ -0,0 +1,102 @@ +{- | A pool of pre-allocated 'MagicRing' buffers for reuse. + +Allocating a magic ring costs @memfd_create@ + @ftruncate@ + 3× +@mmap@ + @memset@ per ring. This pool keeps a bounded free list of +rings so connection churn doesn't repeat that syscall work. + +Rings are stored as raw @(Ptr Word8, Int)@ pairs — no IORefs, no +per-ring mutable state. The transport-layer cursor state (head, +tail, open\/closed) is always fresh per connection; only the +underlying mmap'd memory is recycled. +-} +module Wireform.Ring.Pool ( + RingPool, + RingPoolConfig (..), + defaultRingPoolConfig, + newRingPool, + closeRingPool, + acquireRing, + releaseRing, +) where + +import Control.Concurrent.MVar +import Control.Exception (SomeException, mask_, try) +import Control.Monad (forM_) +import GHC.Exts (RealWorld) +import Wireform.Ring.Internal (MagicRing (..), destroyMagicRing, newMagicRing) + + +data RingPoolConfig = RingPoolConfig + { ringPoolMaxIdle :: !Int + } + + +defaultRingPoolConfig :: RingPoolConfig +defaultRingPoolConfig = + RingPoolConfig + { ringPoolMaxIdle = 64 + } + + +{- | The pool is just a bounded free list protected by an MVar. +All rings in the list have the same allocated size (the first +allocation's rounded-up size sets the class). +-} +data RingPool = RingPool + { rpMaxIdle :: !Int + , rpFree :: !(MVar [MagicRing RealWorld]) + } + + +newRingPool :: RingPoolConfig -> IO RingPool +newRingPool cfg = do + free <- newMVar [] + pure + RingPool + { rpMaxIdle = ringPoolMaxIdle cfg + , rpFree = free + } + + +{- | Destroy all idle rings. The pool remains usable (acquire +allocates fresh, release destroys immediately). +-} +closeRingPool :: RingPool -> IO () +closeRingPool pool = mask_ $ do + rings <- swapMVar (rpFree pool) [] + forM_ rings $ \r -> + () <$ try @SomeException (destroyMagicRing r) + + +-- | Take a ring from the pool or allocate a fresh one. +acquireRing :: RingPool -> Int -> IO (MagicRing s) +acquireRing pool requested = mask_ $ do + rings <- takeMVar (rpFree pool) + case rings of + (r : rs) | mrSize r >= requested -> do + putMVar (rpFree pool) rs + pure (coerceRing r) + _ -> do + putMVar (rpFree pool) rings + newMagicRing requested + + +{- | Return a ring to the pool. Destroyed immediately if the pool +is full. +-} +releaseRing :: RingPool -> MagicRing s -> IO () +releaseRing pool ring = mask_ $ do + rings <- takeMVar (rpFree pool) + if length rings >= rpMaxIdle pool + then do + putMVar (rpFree pool) rings + destroyMagicRing ring + else + putMVar (rpFree pool) (coerceRing ring : rings) + + +-- MagicRing's phantom s is nominal, but the pool stores them as +-- RealWorld and hands them out polymorphically. The underlying +-- data is just (Ptr Word8, Int) — the coercion is safe. +coerceRing :: MagicRing s -> MagicRing s' +coerceRing (MagicRing p n) = MagicRing p n diff --git a/src/Wireform/Transport.hs b/src/Wireform/Transport.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Transport.hs @@ -0,0 +1,17 @@ +{- | Magic-ring transports — receive and send. + +This module is a thin umbrella over the symmetric pair +'Wireform.Transport.Receive.ReceiveTransport' (consumer of bytes +the producer side wrote into a ring) and +'Wireform.Transport.Send.SendTransport' (producer of bytes the +consumer side will drain from a ring). See the per-direction +modules for the full doc. +-} +module Wireform.Transport ( + module Wireform.Transport.Receive, + module Wireform.Transport.Send, +) where + +import Wireform.Transport.Receive +import Wireform.Transport.Send + diff --git a/src/Wireform/Transport/Capabilities.hs b/src/Wireform/Transport/Capabilities.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Transport/Capabilities.hs @@ -0,0 +1,306 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-local-binds #-} + +module Wireform.Transport.Capabilities + ( SystemCapabilities (..) + , IOUringFeatures (..) + , NumaNodeInfo (..) + , CoreTopology (..) + , Placement (..) + , detectCapabilities + , recommendPlacement + ) where + +import qualified Control.Exception as CE +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IntMap +import Data.IORef +import Foreign.C.Types (CLong (..)) +import System.IO.Unsafe (unsafePerformIO) + +import Wireform.Parser +import Wireform.Parser.Internal (Pure) +import Wireform.Parser.Driver (parseByteString) + +data SystemCapabilities = SystemCapabilities + { capPageSize :: !Int + , capHugePageSizes :: ![Int] + , capHasIOUring :: !Bool + , capIOUringFeatures :: !IOUringFeatures + , capNumaNodes :: ![NumaNodeInfo] + , capIsolatedCores :: ![Int] + , capCoreCount :: !Int + , capCoreTopology :: !CoreTopology + } deriving stock (Show) + +data IOUringFeatures = IOUringFeatures + { ioUringFeatureProvidedBuffers :: !Bool + , ioUringFeatureMultishotRecv :: !Bool + , ioUringFeatureSQPoll :: !Bool + } deriving stock (Show) + +data NumaNodeInfo = NumaNodeInfo + { numaNodeId :: !Int + , numaNodeCores :: ![Int] + , numaNodeMemoryMB :: !Int + } deriving stock (Show) + +data CoreTopology = CoreTopology + { topoCoreToPackage :: !(IntMap Int) + , topoCoreToL3 :: !(IntMap Int) + , topoSiblings :: !(IntMap [Int]) + } deriving stock (Show) + +data Placement = Placement + { placementNumaNode :: !(Maybe Int) + , placementCore :: !(Maybe Int) + } deriving stock (Show) + +{-# NOINLINE cachedCapabilities #-} +cachedCapabilities :: IORef (Maybe SystemCapabilities) +cachedCapabilities = unsafePerformIO (newIORef Nothing) + +detectCapabilities :: IO SystemCapabilities +detectCapabilities = do + cached <- readIORef cachedCapabilities + case cached of + Just c -> pure c + Nothing -> do + c <- detectCapabilitiesUncached + writeIORef cachedCapabilities (Just c) + pure c + +detectCapabilitiesUncached :: IO SystemCapabilities +detectCapabilitiesUncached = do + ps <- getPageSize + hugeSizes <- getHugePageSizes + cores <- getCoreCount + isolated <- getIsolatedCores + hasUring <- getHasIOUring + uringFeats <- getIOUringFeatures + numa <- getNumaNodes + pure SystemCapabilities + { capPageSize = ps + , capHugePageSizes = hugeSizes + , capHasIOUring = hasUring + , capIOUringFeatures = uringFeats + , capNumaNodes = numa + , capIsolatedCores = isolated + , capCoreCount = cores + , capCoreTopology = CoreTopology IntMap.empty IntMap.empty IntMap.empty + } + +recommendPlacement :: Maybe Int -> IO Placement +recommendPlacement _fdNumaNode = pure (Placement Nothing Nothing) + +------------------------------------------------------------------------ +-- Parsers for /proc and /sys files (using our own parser) +------------------------------------------------------------------------ + +type P = Parser Pure () + +-- | Parse a CPU list like "0-3,5,7-9" into [0,1,2,3,5,7,8,9]. +pCpuList :: P [Int] +pCpuList = do + first <- pRange + rest <- many (word8 0x2C *> pRange) -- ',' + pure (concat (first : rest)) + where + pRange :: P [Int] + pRange = do + lo <- anyAsciiDecimalInt + withOption (word8 0x2D *> anyAsciiDecimalInt) -- '-' + (\hi -> pure [lo..hi]) + (pure [lo]) + +-- | Parse a decimal integer, skipping leading whitespace. +pDecimal :: P Int +pDecimal = skipWs *> anyAsciiDecimalInt + +-- | Skip ASCII whitespace. +skipWs :: P () +skipWs = skipMany (satisfyAscii (\c -> c == ' ' || c == '\t' || c == '\n' || c == '\r')) + +-- | Skip to after a keyword, then parse what follows. +pAfterKeyword :: ByteString -> P a -> P a +pAfterKeyword kw p = go + where + go = (byteString kw *> p) <|> (anyWord8 *> go) + +-- | Count lines starting with a given prefix. +pCountPrefix :: ByteString -> P Int +pCountPrefix pfx = go 0 + where + go !n = (byteString pfx *> pSkipLine *> go (n + 1)) + <|> (pSkipLine *> go n) -- pSkipLine fails at EOF, terminating the loop + <|> pure n + +-- | Skip bytes until newline (inclusive). Fails at EOF. +pSkipLine :: P () +pSkipLine = go + where + go = withAnyWord8 \w -> + if w == 0x0A then pure () else go + +-- | Parse "Linux version X.Y.Z ..." from /proc/version. +pKernelVersion :: P (Int, Int) +pKernelVersion = pAfterKeyword "version " $ do + major <- anyAsciiDecimalInt + word8 0x2E -- '.' + minor <- anyAsciiDecimalInt + pure (major, minor) + +-- | Parse "Hugepagesize: 2048 kB" lines from /proc/meminfo. +pHugePageSize :: P [Int] +pHugePageSize = go [] + where + go !acc = + (do byteString "Hugepagesize:" + skipWs + n <- anyAsciiDecimalInt + skipWs + byteString "kB" + pSkipLine + go (n * 1024 : acc)) + <|> (pSkipLine *> go acc) + <|> pure (reverse acc) + +-- | Parse "MemTotal: 12345 kB" from a NUMA node's meminfo. +pNodeMemMB :: P Int +pNodeMemMB = go + where + go = (do byteString "MemTotal:" + skipWs + n <- anyAsciiDecimalInt + pure (n `div` 1024)) + <|> (pSkipLine *> go) + <|> pure 0 + +------------------------------------------------------------------------ +-- File reading + parsing +------------------------------------------------------------------------ + +tryReadFileBS :: FilePath -> IO (Maybe ByteString) +tryReadFileBS path = do + r <- CE.try @CE.SomeException (BS.readFile path) + case r of + Left _ -> pure Nothing + Right s -> pure (Just s) + +runP :: P a -> ByteString -> Maybe a +runP p bs = case parseByteString p bs of + Right a -> Just a + Left _ -> Nothing + +------------------------------------------------------------------------ +-- Platform-specific detection +------------------------------------------------------------------------ + +foreign import ccall unsafe "hs_page_size" + c_page_size :: IO CLong + +getPageSize :: IO Int +getPageSize = fromIntegral <$> c_page_size + +getCoreCount :: IO Int +getCoreCount = do +#if defined(linux_HOST_OS) + r <- tryReadFileBS "/proc/cpuinfo" + case r of + Nothing -> pure 1 + Just bs -> pure $ max 1 (maybe 1 id (runP (pCountPrefix "processor") bs)) +#else + pure 1 +#endif + +getHugePageSizes :: IO [Int] +getHugePageSizes = do +#if defined(linux_HOST_OS) + r <- tryReadFileBS "/proc/meminfo" + case r of + Nothing -> pure [] + Just bs -> pure $ maybe [] id (runP pHugePageSize bs) +#else + pure [] +#endif + +getIsolatedCores :: IO [Int] +getIsolatedCores = do +#if defined(linux_HOST_OS) + r <- tryReadFileBS "/sys/devices/system/cpu/isolated" + case r of + Nothing -> pure [] + Just bs -> pure $ maybe [] id (runP (skipWs *> pCpuList) (BS.filter (/= 0x0A) bs)) +#else + pure [] +#endif + +getHasIOUring :: IO Bool +getHasIOUring = do +#if defined(linux_HOST_OS) + ver <- kernelVersion + pure (ver >= (5, 1)) +#else + pure False +#endif + +getIOUringFeatures :: IO IOUringFeatures +getIOUringFeatures = do +#if defined(linux_HOST_OS) + hasUring <- getHasIOUring + if not hasUring + then pure (IOUringFeatures False False False) + else do + ver <- kernelVersion + pure IOUringFeatures + { ioUringFeatureProvidedBuffers = ver >= (5, 19) + , ioUringFeatureMultishotRecv = ver >= (6, 0) + , ioUringFeatureSQPoll = ver >= (5, 1) + } +#else + pure (IOUringFeatures False False False) +#endif + +getNumaNodes :: IO [NumaNodeInfo] +getNumaNodes = do +#if defined(linux_HOST_OS) + r <- tryReadFileBS "/sys/devices/system/node/online" + case r of + Nothing -> pure [] + Just bs -> do + let nodeIds = maybe [] id (runP (skipWs *> pCpuList) (BS.filter (/= 0x0A) bs)) + mapM getNodeInfo nodeIds +#else + pure [] +#endif + where + getNodeInfo nid = do + cpuList <- tryReadFileBS ("/sys/devices/system/node/node" <> bsShow nid <> "/cpulist") + memInfo <- tryReadFileBS ("/sys/devices/system/node/node" <> bsShow nid <> "/meminfo") + let cores = case cpuList of + Nothing -> [] + Just bs -> maybe [] id (runP pCpuList (BS.filter (/= 0x0A) bs)) + memMB = case memInfo of + Nothing -> 0 + Just bs -> maybe 0 id (runP pNodeMemMB bs) + pure (NumaNodeInfo nid cores memMB) + + bsShow :: Int -> FilePath + bsShow = show + +kernelVersion :: IO (Int, Int) +kernelVersion = do +#if defined(linux_HOST_OS) + r <- tryReadFileBS "/proc/version" + pure $ case r of + Nothing -> (0, 0) + Just bs -> maybe (0, 0) id (runP pKernelVersion bs) +#else + pure (0, 0) +#endif diff --git a/src/Wireform/Transport/Config.hs b/src/Wireform/Transport/Config.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Transport/Config.hs @@ -0,0 +1,221 @@ +module Wireform.Transport.Config ( + -- * Profiles + Profile (..), + profileConfig, + + -- * Full configuration + TransportConfig (..), + defaultTransportConfig, + + -- * Wait policy + WaitPolicy (..), + SpinBudget (..), + + -- * Pinning + PinningPolicy (..), + + -- * Pages + PagePolicy (..), + + -- * NUMA + NumaPolicy (..), + + -- * io_uring + IOUringConfig (..), + SQPollPolicy (..), + CompletionWaitMode (..), + defaultIOUringConfig, + + -- * Capability handling + CapabilityAction (..), +) where + + +------------------------------------------------------------------------ +-- Profiles +------------------------------------------------------------------------ + +{- | High-level performance profile. Most users pick one of these +and never touch 'TransportConfig' directly. +-} +data Profile + = -- | IO-manager-parked. Good citizen on shared systems. + Throughput + | -- | Brief spin before parking, pinned thread, huge pages where available. + LowLatency + | {- | Pure busy-poll, pinned to isolated core, huge pages, mlock. + Burns one CPU core permanently. + -} + UltraLowLatency + deriving stock (Eq, Ord, Show, Read, Bounded, Enum) + + +-- | Expand a profile into a full 'TransportConfig'. +profileConfig :: Profile -> TransportConfig +profileConfig Throughput = defaultTransportConfig +profileConfig LowLatency = + defaultTransportConfig + { waitPolicy = WaitSpinThenPark (SpinNanos 5000) + , pinning = PinNearFd + , pages = PreferHugePages + , memoryLocking = True + , numaPlacement = NumaAutoFromFd + } +profileConfig UltraLowLatency = + defaultTransportConfig + { waitPolicy = WaitBusyPoll + , pinning = PinIsolated + , pages = PreferHugePages + , memoryLocking = True + , numaPlacement = NumaAutoFromFd + , ioUring = + defaultIOUringConfig + { ioUringQueueDepth = 1024 + , ioUringSQPoll = SQPollWithIdle 100 + , ioUringCompletionWait = WaitViaBusyPoll + } + } + + +------------------------------------------------------------------------ +-- Full configuration +------------------------------------------------------------------------ + +data TransportConfig = TransportConfig + { ringSizeHint :: !Int + {- ^ Requested ring size in bytes. Rounded up to page-size multiple + and power of two. Default: 1 MiB. + -} + , waitPolicy :: !WaitPolicy + , pinning :: !PinningPolicy + , pages :: !PagePolicy + , memoryLocking :: !Bool + -- ^ @mlock@ the ring. Requires sufficient RLIMIT_MEMLOCK. + , numaPlacement :: !NumaPolicy + , ioUring :: !IOUringConfig + -- ^ Linux io_uring tuning (silently ignored elsewhere). + , onCapabilityLimit :: !CapabilityAction + -- ^ What to do when a knob cannot be applied on this platform. + } + deriving stock (Show) + + +defaultTransportConfig :: TransportConfig +defaultTransportConfig = + TransportConfig + { ringSizeHint = 1024 * 1024 + , waitPolicy = WaitParkImmediately + , pinning = NoPinning + , pages = StandardPages + , memoryLocking = False + , numaPlacement = NoNumaPreference + , ioUring = defaultIOUringConfig + , onCapabilityLimit = LogAndContinue + } + + +------------------------------------------------------------------------ +-- Wait policy +------------------------------------------------------------------------ + +data WaitPolicy + = -- | Park on the IO manager as soon as caught up. + WaitParkImmediately + | -- | Busy-spin for the budget, then park. + WaitSpinThenPark !SpinBudget + | -- | Never park. Burns a core. + WaitBusyPoll + deriving stock (Show) + + +data SpinBudget + = SpinIterations !Int + | SpinNanos !Int + | SpinUntilPaused + deriving stock (Show) + + +------------------------------------------------------------------------ +-- Pinning +------------------------------------------------------------------------ + +data PinningPolicy + = NoPinning + | PinToCore !Int + | PinToCoreSet ![Int] + | -- | Auto-detect: pin near the NIC / fd's NUMA node. + PinNearFd + | -- | Pick an isolated core (Linux @isolcpus=@); fall back to 'PinNearFd'. + PinIsolated + deriving stock (Show) + + +------------------------------------------------------------------------ +-- Page policy +------------------------------------------------------------------------ + +data PagePolicy + = StandardPages + | -- | Request huge pages; fall back to standard if unavailable. + PreferHugePages + | -- | Fail at ring creation if huge pages are unavailable. + RequireHugePages + deriving stock (Show) + + +------------------------------------------------------------------------ +-- NUMA +------------------------------------------------------------------------ + +data NumaPolicy + = NoNumaPreference + | NumaNode !Int + | NumaAutoFromFd + | NumaAutoFromCurrentCore + deriving stock (Show) + + +------------------------------------------------------------------------ +-- io_uring +------------------------------------------------------------------------ + +data IOUringConfig = IOUringConfig + { ioUringQueueDepth :: !Int + , ioUringSQPoll :: !SQPollPolicy + , ioUringProvidedBuffers :: !Bool + , ioUringCompletionWait :: !CompletionWaitMode + } + deriving stock (Show) + + +defaultIOUringConfig :: IOUringConfig +defaultIOUringConfig = + IOUringConfig + { ioUringQueueDepth = 128 + , ioUringSQPoll = NoSQPoll + , ioUringProvidedBuffers = True + , ioUringCompletionWait = WaitViaEventFd + } + + +data SQPollPolicy + = NoSQPoll + | SQPollWithIdle !Int + deriving stock (Show) + + +data CompletionWaitMode + = WaitViaEventFd + | WaitViaBusyPoll + deriving stock (Show) + + +------------------------------------------------------------------------ +-- Capability action +------------------------------------------------------------------------ + +data CapabilityAction + = SilentlyIgnore + | LogAndContinue + | FailHard + deriving stock (Eq, Ord, Show, Read, Bounded, Enum) diff --git a/src/Wireform/Transport/Receive.hs b/src/Wireform/Transport/Receive.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Transport/Receive.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE RankNTypes #-} + +{- | Receive-side magic-ring transport. + +The parser consumes bytes from a 'ReceiveTransport'. The producer +(a network recv loop, a TLS context, an io_uring backend, an +in-memory test fixture, …) writes into the magic ring at @head@ +and the parser drains @[tail, head)@. + +== Cursor model + +'Word64' positions are monotonic (never wrap during process +lifetime). The byte at logical position @p@ lives at +@ringBase + (p .&. ringMask)@ in the ring. + +== Why the ring is stored as three primitive fields + +'ReceiveTransport' embeds the ring's base pointer, size, and mask +directly rather than carrying a 'Wireform.Ring.Internal.MagicRing'. +That keeps the type free of the ring's phantom @s@ parameter, so +the rank-2 scope safety of 'Wireform.Ring.withMagicRing' is something +callers opt into at the slice layer rather than something that +ripples through every 'ReceiveTransport'-using package. +'receiveRing' reconstructs a polymorphic 'MagicRing' for callers +that still want to call 'ringBase' / 'ringSize' / 'ringMask' on a +record value. +-} +module Wireform.Transport.Receive ( + ReceiveTransport (..), + ReceiveWait (..), + receiveRing, +) where + +import Control.Exception (SomeException) +import Data.Word (Word64, Word8) +import Foreign.Ptr (Ptr) +import Wireform.Ring.Internal (MagicRing (..)) + + +{- | A producer-side cursor + a slot to wait on more data. + +The parser interacts with data exclusively through this record. +Implementations bind one to a socket, TLS context, io_uring +instance, or an in-memory fixture. +-} +data ReceiveTransport = ReceiveTransport + { receiveRingBase :: {-# UNPACK #-} !(Ptr Word8) + -- ^ Base address of the ring's first mapping. + , receiveRingSize :: {-# UNPACK #-} !Int + -- ^ Physical ring size (N, not 2N). Always a power of two. + , receiveRingMask :: {-# UNPACK #-} !Int + -- ^ @ringSize - 1@, cached for cheap @pos .&. mask@. + , receiveLoadHead :: !(IO Word64) + {- ^ Read the producer's current head position + (monotonically increasing). + -} + , receiveAdvanceTail :: !(Word64 -> IO ()) + {- ^ Consumer publishes a new tail position + (monotonically increasing). The transport may use this for + backpressure or buffer reuse. + -} + , receiveWaitData :: !(Word64 -> IO ReceiveWait) + {- ^ Block until head advances past the given position. + Parks on the IO manager when blocking (same thread as the + parser). + -} + , receiveClose :: !(IO ()) + -- ^ Release resources. Idempotent. + } + + +-- | Outcome of waiting for more data. +data ReceiveWait + = -- | New head position. + ReceiveMoreData {-# UNPACK #-} !Word64 + | {- | Clean end (sticky: once returned, all subsequent calls + return this). + -} + ReceiveEndOfInput + | -- | Producer-side failure (sticky). + ReceiveFailed !SomeException + deriving stock (Show) + + +{- | Reconstruct the underlying 'MagicRing'. Polymorphic in the +phantom @s@: the resulting handle does /not/ inherit the scope of +whatever ring originally produced these bytes, so callers should +treat it as un-scoped (i.e. equivalent to the pre-@MagicRing s@ +world). Use 'Wireform.Ring.withMagicRing' directly when you need +the type-system-enforced safety. +-} +receiveRing :: ReceiveTransport -> MagicRing s +receiveRing t = MagicRing (receiveRingBase t) (receiveRingSize t) +{-# INLINE receiveRing #-} diff --git a/src/Wireform/Transport/Send.hs b/src/Wireform/Transport/Send.hs new file mode 100644 --- /dev/null +++ b/src/Wireform/Transport/Send.hs @@ -0,0 +1,451 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} + +{- | Send-side magic-ring transport. + +The dual of 'Wireform.Transport.Receive.ReceiveTransport': + +* Producer = the encoder (a 'Wireform.Builder.Builder', a hand- + written pointer-bumping codec, a 'ByteString' source, …). + Advances @head@ by writing bytes into the ring at + @[head, head + n)@ and publishing. +* Consumer = the wire (a network @sendmsg@ loop, a TLS + @SSL_write@ loop, an io_uring @prep_send@ submission, an in- + memory test sink, …). Advances @tail@ as it drains. + +The encoder never sees a wrap point: a single 'reserveSend' may +span the ring's wrap boundary, but the double mapping makes the +reserved pointer contiguous in virtual memory. +-} +module Wireform.Transport.Send ( + -- * The transport + SendTransport (..), + SendWait (..), + sendRing, + + -- * Encoder-facing API + reserveSend, + withSendReservation, + sendByteString, + sendByteStringMany, + sendBuilder, + sendBuilderDirect, + sendBuilderViaByteString, + + -- * Corking (batch multiple sends into one publish) + withSendCork, + + -- * Exceptions + SendRingFull (..), + SendReservationTooLarge (..), +) where + +import Control.Exception (Exception, SomeException, throwIO) +import Control.Monad (when) +import Data.Bits ((.&.)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Unsafe qualified as BSU +import Data.IORef +import Data.Typeable (Typeable) +import Data.Word (Word64, Word8) +import Foreign.Marshal.Utils (copyBytes) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Wireform.Builder qualified as B +import Wireform.Builder.FastBuilder ( + DataSink (RingSink), + RingSinkState (..), + finalizeRingSink, + runBuilder, + ) +import Wireform.Ring.Internal (MagicRing (..)) + + +------------------------------------------------------------------------ +-- SendTransport +------------------------------------------------------------------------ + +{- | A consumer-side cursor + a slot to wait on more room. + +Encoders interact with the wire exclusively through this record. +Implementations bind one to a socket, a TLS context, an io_uring +instance, an in-memory fixture, etc. +-} +data SendTransport = SendTransport + { sendRingBase :: {-# UNPACK #-} !(Ptr Word8) + -- ^ Base address of the send ring's first mapping. + , sendRingSize :: {-# UNPACK #-} !Int + -- ^ Physical ring size (N, not 2N). Always a power of two. + , sendRingMask :: {-# UNPACK #-} !Int + -- ^ @sendRingSize - 1@, cached for cheap @pos .&. mask@. + , sendLoadTail :: !(IO Word64) + {- ^ Read the consumer's current tail position + (monotonically increasing). + -} + , sendLoadHead :: !(IO Word64) + {- ^ Read the current head position. Encoders maintain this + locally between commits; exposed here so debugging tools + and background workers can observe progress. + -} + , sendPublishHead :: !(Word64 -> IO ()) + {- ^ Producer publishes a new head after writing bytes into the + ring. Implementations may use this as the flush trigger + (e.g. an io_uring SQ submit) or rely on an explicit + 'sendFlush'. + -} + , sendWaitSpace :: !(Word64 -> IO SendWait) + {- ^ Block until 'sendLoadTail >= pos - sendRingSize' — i.e. + until there is room to advance head to @pos@. The dual of + 'Wireform.Transport.Receive.receiveWaitData'. + -} + , sendFlush :: !(IO ()) + {- ^ Request the consumer side drain everything published so + far. For an inline (synchronous @sendmsg@ / @SSL_write@) + consumer this is @pure ()@ because 'sendPublishHead' already + drained. For a background-worker / io_uring consumer this + kicks the SQ / signals the worker. + -} + , sendShutdownWrite :: !(IO ()) + {- ^ Half-close (@shutdown(SHUT_WR)@ on TCP, @close_notify@ + on TLS), keeping the recv side alive so we can still read + the peer's final reply. + -} + , sendClose :: !(IO ()) + {- ^ Full release: tear down the ring + any background worker. + Idempotent. + -} + } + + +-- | Outcome of waiting for more room. +data SendWait + = -- | New tail position. Room is @sendRingSize - (head - tail)@. + SendSpaceAvailable {-# UNPACK #-} !Word64 + | {- | Consumer closed (peer sent RST, FIN-after-shutdown, etc). + Sticky. + -} + SendPeerClosed + | -- | Wire-side failure. Sticky. + SendFailed !SomeException + deriving stock (Show) + + +{- | Reconstruct the underlying send ring as a 'MagicRing'. +Polymorphic in @s@ for the same reason 'receiveRing' is — the +resulting handle does not inherit any scope. +-} +sendRing :: SendTransport -> MagicRing s +sendRing t = MagicRing (sendRingBase t) (sendRingSize t) +{-# INLINE sendRing #-} + + +------------------------------------------------------------------------ +-- Exceptions +------------------------------------------------------------------------ + +{- | The send ring is full and the consumer has not yet drained any +room within the wait policy's budget. Surfaced as a sticky +'SendFailed' (and re-thrown by 'reserveSend' / 'sendByteString') +instead of letting the producer/consumer pair spin forever. +-} +data SendRingFull = SendRingFull + { sendRingFullSize :: !Int + -- ^ Physical ring size (N). + , sendRingFullHead :: !Word64 + , sendRingFullTail :: !Word64 + } + deriving stock (Show, Typeable) + + +instance Exception SendRingFull + + +{- | A 'reserveSend' / 'sendByteString' asked for more bytes in a +single contiguous reservation than the ring physically holds. +The ring is sized at construction time; raise the +'Wireform.Transport.Config.ringSizeHint' for the connection if +you need to stage a larger payload in one shot. +-} +data SendReservationTooLarge = SendReservationTooLarge + { sendReservationRequested :: !Int + , sendReservationRingSize :: !Int + } + deriving stock (Show, Typeable) + + +instance Exception SendReservationTooLarge + + +------------------------------------------------------------------------ +-- Encoder-facing API +------------------------------------------------------------------------ + +{- | Reserve a contiguous span of @n@ bytes at the current head. +Blocks (via 'sendWaitSpace') until there is room. Returns a +pointer into the ring's double mapping (so the encoder can write +@n@ bytes contiguously regardless of wrap) and the new head that +'commitSend' / 'sendPublishHead' should publish once the write +is complete. + +Throws 'SendReservationTooLarge' if @n > sendRingSize@. +Throws the sticky 'SendRingFull' / underlying 'SendFailed' +exception if the transport has been closed (concretely: the +ring's backing mmap may already be gone, so we must not hand +out a pointer into it). +-} +reserveSend + :: SendTransport + -> Int + -> IO (Ptr Word8, Word64) +reserveSend t n + | n < 0 = error "Wireform.Transport.Send.reserveSend: negative length" + | n > sendRingSize t = + throwIO + ( SendReservationTooLarge + { sendReservationRequested = n + , sendReservationRingSize = sendRingSize t + } + ) + | otherwise = do + h <- sendLoadHead t + let !newHead = h + fromIntegral n + ensureRoom t newHead + let !off = fromIntegral h .&. sendRingMask t + !writeP = sendRingBase t `plusPtr` off + pure (writeP, newHead) +{-# INLINE reserveSend #-} + + +{- | Reserve up to @maxLen@ bytes, run the fill callback, publish +head by the number of bytes the callback actually wrote. +Convenient bracket-style wrapper around 'reserveSend' + +'sendPublishHead'. +-} +withSendReservation + :: SendTransport + -> Int + -- ^ max bytes to reserve + -> (Ptr Word8 -> Int -> IO Int) + -- ^ fill callback, returns bytes written + -> IO Int +withSendReservation t maxLen fill + | maxLen <= 0 = pure 0 + | otherwise = do + (p, _) <- reserveSend t maxLen + written <- fill p maxLen + if written < 0 || written > maxLen + then error "Wireform.Transport.Send.withSendReservation: callback returned out-of-range length" + else do + when_ (written > 0) $ do + h <- sendLoadHead t + sendPublishHead t (h + fromIntegral written) + pure written + where + when_ True m = m + when_ False _ = pure () +{-# INLINE withSendReservation #-} + + +{- | Copy a 'ByteString' into the send ring as one contiguous +reservation, then publish head. The bytes are committed by the +time this returns; whether they have hit the wire by then +depends on the transport's consumer (an inline socket consumer +has already drained; an io_uring consumer has at least submitted +the SQE). +-} +sendByteString :: SendTransport -> ByteString -> IO () +sendByteString t bs = do + let !len = BS.length bs + if len == 0 + then pure () + else do + (p, newHead) <- reserveSend t len + BSU.unsafeUseAsCStringLen bs $ \(src, _) -> + copyBytes p (castPtr src) len + sendPublishHead t newHead +{-# INLINE sendByteString #-} + + +{- | Stage many byte strings as one merged reservation, then +publish head once. Lets the consumer coalesce them into a +single @sendmsg@ / io_uring SQE. +-} +sendByteStringMany :: SendTransport -> [ByteString] -> IO () +sendByteStringMany _ [] = pure () +sendByteStringMany t bss = do + let !total = sum (fmap BS.length bss) + if total == 0 + then pure () + else do + (p, newHead) <- reserveSend t total + copyAll p bss + sendPublishHead t newHead + where + copyAll _ [] = pure () + copyAll dst (b : bs) = do + let !n = BS.length b + BSU.unsafeUseAsCStringLen b $ \(src, _) -> + copyBytes dst (castPtr src) n + copyAll (dst `plusPtr` n) bs +{-# INLINEABLE sendByteStringMany #-} + + +{- | Stage a 'B.Builder' into the send ring. Equivalent to +'sendBuilderDirect' -- writes directly into ring memory with +no intermediate 'ByteString'. +-} +sendBuilder :: SendTransport -> B.Builder -> IO () +sendBuilder = sendBuilderDirect +{-# INLINE sendBuilder #-} + + +{- | The pre-optimisation path: materialise to a strict +'ByteString', then memcpy into the ring. Exported so +benchmarks can compare the two strategies side-by-side. +-} +sendBuilderViaByteString :: SendTransport -> B.Builder -> IO () +sendBuilderViaByteString t b = + sendByteString t (B.toStrictByteStringWith 4096 b) +{-# INLINE sendBuilderViaByteString #-} + + +{- | Run a 'B.Builder' directly into the send ring with no +intermediate 'ByteString' allocation or copy. + +The builder writes into the ring's double-mapped memory. +When the current region fills, the already-written bytes are +published and the builder continues into freshly-available ring +space. The final bytes are published before this function +returns. + +For a typical small frame (< 32 KiB) this is one @runBuilder@ +call with zero allocations beyond a single 'IORef'. For larger +payloads the builder publishes in chunks, keeping memory +pressure proportional to the ring size rather than the payload. +-} +sendBuilderDirect :: SendTransport -> B.Builder -> IO () +sendBuilderDirect t b = do + h <- sendLoadHead t + let !chunkSz = min (sendRingSize t) defaultDirectChunk + !wantHead = h + fromIntegral chunkSz + ensureRoom t wantHead + let !off = fromIntegral h .&. sendRingMask t + !ptr = sendRingBase t `plusPtr` off + rsRef <- + newIORef + RingSinkState + { rsHead = h + , rsPublished = h + , rsBase = sendRingBase t + , rsMask = sendRingMask t + , rsRingSize = sendRingSize t + , rsChunkSize = chunkSz + , rsPublishHead = sendPublishHead t + , rsEnsureRoom = ensureRoom t + , rsLoadTail = sendLoadTail t + } + finalCur <- runBuilder b (RingSink rsRef) ptr (ptr `plusPtr` chunkSz) + _ <- finalizeRingSink rsRef finalCur + pure () +{-# INLINE sendBuilderDirect #-} + + +-- 32 KiB: large enough that most protocol frames never overflow, +-- small enough that the initial ensureRoom resolves immediately on +-- a typical 64-256 KiB ring. +defaultDirectChunk :: Int +defaultDirectChunk = 32768 + + +------------------------------------------------------------------------ +-- Corking +------------------------------------------------------------------------ + +{- | Run an action with deferred publishing. Bytes written via +'sendByteString' \/ 'sendBuilderDirect' \/ etc. accumulate in the +ring without triggering the consumer (no @sendmsg@, no io_uring +SQE submission). On scope exit, a single 'sendPublishHead' covers +everything written. + +The cork is demand-driven: 'sendPublishHead' is a no-op inside the +scope, and the cork only breaks when 'sendWaitSpace' detects the +ring is genuinely full and needs a drain to continue. This means +the cork batches as much as the ring can hold — for typical HTTP +responses (headers + body < ringSize) that is one publish total. +For payloads larger than ringSize, publishes happen only when +strictly necessary to free ring space. + +@ +withSendCork transport $ \corked -> do + sendBuilderDirect corked (responseBuilder resp) + sendByteString corked bodyBytes +-- single publish \/ sendmsg here +@ +-} +withSendCork :: SendTransport -> (SendTransport -> IO a) -> IO a +withSendCork t action = do + h0 <- sendLoadHead t + headRef <- newIORef h0 + publishedRef <- newIORef h0 + let corkedPublish h = writeIORef headRef h + corkedWaitSpace pos = do + r <- sendWaitSpace t pos + case r of + SendSpaceAvailable tl + | pos - tl <= fromIntegral (sendRingSize t) -> pure r + | otherwise -> do + -- Ring full — publish accumulated bytes to trigger + -- drain, then retry. + cur <- readIORef headRef + pub <- readIORef publishedRef + when (cur > pub) $ do + sendPublishHead t cur + writeIORef publishedRef cur + sendWaitSpace t pos + _ -> pure r + let corked = + t + { sendPublishHead = corkedPublish + , sendLoadHead = readIORef headRef + , sendWaitSpace = corkedWaitSpace + } + result <- action corked + finalHead <- readIORef headRef + pub <- readIORef publishedRef + when (finalHead > pub) $ sendPublishHead t finalHead + pure result + + +------------------------------------------------------------------------ +-- Internal: wait for room +------------------------------------------------------------------------ + +{- | Block until the ring has room for head to advance to @newHead@. + +Always consults 'sendWaitSpace' once, even when the local cursor +arithmetic says room is already available. This is what surfaces +a sticky closed state (e.g. after 'sendClose') BEFORE the caller +reaches for the ring's base pointer — without this probe, +'reserveSend' on a closed inline transport would happily hand out +a pointer into the ring's already-unmapped backing memory. +-} +ensureRoom :: SendTransport -> Word64 -> IO () +ensureRoom t newHead = loop + where + sz = fromIntegral (sendRingSize t) :: Word64 + loop = do + r <- sendWaitSpace t newHead + case r of + SendSpaceAvailable tl + | newHead - tl <= sz -> pure () + | otherwise -> loop + SendPeerClosed -> + throwIO + ( SendRingFull + { sendRingFullSize = sendRingSize t + , sendRingFullHead = newHead + , sendRingFullTail = newHead -- unknown; report newHead + } + ) + SendFailed e -> throwIO e +{-# INLINE ensureRoom #-} diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,15 @@ +module Main where + +import Test.Syd +import Wireform.Base64.Test qualified as Base64 +import Wireform.Parser.Test qualified as Parser +import Wireform.Ring.Test qualified as Ring +import Wireform.Transport.SendTest qualified as Send + + +main :: IO () +main = sydTest $ do + Ring.spec + Parser.spec + Send.spec + Base64.spec diff --git a/test/Wireform/Base64/Test.hs b/test/Wireform/Base64/Test.hs new file mode 100644 --- /dev/null +++ b/test/Wireform/Base64/Test.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Wireform.Base64.Test (spec) where + +import Data.ByteString qualified as BS +import Data.Word (Word8) +import Test.QuickCheck +import Test.Syd +import Wireform.Base64 + + +spec :: Spec +spec = describe "Wireform.Base64" $ do + describe "RFC 4648 sec 10 test vectors" $ do + it "encodes the canonical vectors" $ do + encodeBase64 "" `shouldBe` "" + encodeBase64 "f" `shouldBe` "Zg==" + encodeBase64 "fo" `shouldBe` "Zm8=" + encodeBase64 "foo" `shouldBe` "Zm9v" + encodeBase64 "foob" `shouldBe` "Zm9vYg==" + encodeBase64 "fooba" `shouldBe` "Zm9vYmE=" + encodeBase64 "foobar" `shouldBe` "Zm9vYmFy" + + it "decodes the canonical vectors" $ do + decodeBase64 "" `shouldBe` Just "" + decodeBase64 "Zg==" `shouldBe` Just "f" + decodeBase64 "Zm8=" `shouldBe` Just "fo" + decodeBase64 "Zm9v" `shouldBe` Just "foo" + decodeBase64 "Zm9vYg==" `shouldBe` Just "foob" + decodeBase64 "Zm9vYmE=" `shouldBe` Just "fooba" + decodeBase64 "Zm9vYmFy" `shouldBe` Just "foobar" + + describe "round-trip" $ do + it "decodes its own encoding for inputs up to 256 bytes" $ + property $ \(bytes :: [Word8Wrap]) -> + let bs = BS.pack (map unwrap (take 256 bytes)) + in decodeBase64 (encodeBase64 bs) === Just bs + + it "drives the SIMD body (>= 16 chars)" $ do + let bs = BS.pack [0 .. 47] + decodeBase64 (encodeBase64 bs) `shouldBe` Just bs + + it "drives the SIMD body and tail (4096 bytes)" $ do + let bs = BS.pack (take 4096 (cycle [0 .. 255])) + decodeBase64 (encodeBase64 bs) `shouldBe` Just bs + + describe "rejection" $ do + it "rejects non-multiple-of-4 inputs" $ + decodeBase64 "Zm9v=" `shouldBe` Nothing + + it "rejects out-of-alphabet bytes" $ + decodeBase64 "Zm9*" `shouldBe` Nothing + + it "rejects bad padding position" $ + decodeBase64 "Z===" `shouldBe` Nothing + + describe "length helpers" $ do + it "encodeBase64Length matches the actual encoded length" $ + property $ \n -> + let n' = abs n `mod` 1024 + bs = BS.replicate n' 0 + in BS.length (encodeBase64 bs) === encodeBase64Length n' + + +-- Small wrapper to bound Arbitrary Word8 in the property test +-- without dragging in QuickCheck-instances. +newtype Word8Wrap = Word8Wrap {unwrap :: Word8} + deriving stock (Show) + + +instance Arbitrary Word8Wrap where + arbitrary = Word8Wrap . fromIntegral <$> chooseInt (0, 255) diff --git a/test/Wireform/Parser/Test.hs b/test/Wireform/Parser/Test.hs new file mode 100644 --- /dev/null +++ b/test/Wireform/Parser/Test.hs @@ -0,0 +1,275 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Wireform.Parser.Test (spec) where + +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.Int +import Data.Word +import Test.QuickCheck +import Test.Syd +import Wireform.Parser +import Wireform.Parser.Driver (parseByteString) +import Wireform.Parser.Error +import Wireform.Parser.Internal (Pure) + + +type P = Parser Pure String + + +ok :: (Show a, Eq a) => Either (ParseError String) a -> a -> Expectation +ok (Right a) expected = a `shouldBe` expected +ok (Left e) _ = expectationFailure ("parse failed: " <> show e) + + +bad :: (Show a) => Either (ParseError String) a -> Expectation +bad (Left _) = pure () +bad (Right a) = expectationFailure ("expected failure, got: " <> show a) + + +spec :: Spec +spec = describe "Parser" $ do + describe "byte primitives" $ do + it "anyWord8" $ + ok (parseByteString (anyWord8 :: P Word8) "\x42") 0x42 + it "anyWord8 on empty" $ + bad (parseByteString (anyWord8 :: P Word8) "") + it "anyWord16 native" $ + ok (parseByteString (anyWord16 :: P Word16) "\x01\x02") (if isLE then 0x0201 else 0x0102) + it "anyWord32be" $ + ok (parseByteString (anyWord32be :: P Word32) "\x00\x00\x01\x00") 256 + it "anyWord32le" $ + ok (parseByteString (anyWord32le :: P Word32) "\x00\x01\x00\x00") 256 + it "anyWord64le" $ + ok (parseByteString (anyWord64le :: P Word64) "\x01\x00\x00\x00\x00\x00\x00\x00") 1 + it "anyWord64be" $ + ok (parseByteString (anyWord64be :: P Word64) "\x00\x00\x00\x00\x00\x00\x00\x01") 1 + + describe "signed integers" $ do + it "anyInt8" $ + ok (parseByteString (anyInt8 :: P Int8) "\xFF") (-1) + it "anyInt16be" $ + ok (parseByteString (anyInt16be :: P Int16) "\xFF\xFE") (-2) + it "anyInt32le" $ + ok (parseByteString (anyInt32le :: P Int32) "\xFE\xFF\xFF\xFF") (-2) + + describe "floating point" $ do + it "anyFloatle parses 1.0" $ + ok (parseByteString (anyFloatle :: P Float) "\x00\x00\x80\x3F") 1.0 + it "anyDoublebe parses 1.0" $ + ok (parseByteString (anyDoublebe :: P Double) "\x3F\xF0\x00\x00\x00\x00\x00\x00") 1.0 + + describe "sequencing" $ do + it "two bytes" $ do + let p = (,) <$> anyWord8 <*> anyWord8 :: P (Word8, Word8) + ok (parseByteString p "\xAA\xBB") (0xAA, 0xBB) + it "insufficient for sequence" $ do + let p = (,) <$> anyWord8 <*> anyWord8 :: P (Word8, Word8) + bad (parseByteString p "\xAA") + it "three-way sequence" $ do + let p = (,,) <$> anyWord8 <*> anyWord16be <*> anyWord32be :: P (Word8, Word16, Word32) + ok (parseByteString p "\x01\x00\x02\x00\x00\x00\x03") (1, 2, 3) + + describe "alternatives" $ do + it "second branch" $ do + let p = (word8 0x01 *> pure "a") <|> (word8 0x02 *> pure "b") :: P String + ok (parseByteString p "\x02") "b" + it "first branch" $ do + let p = (word8 0x01 *> pure "a") <|> (word8 0x02 *> pure "b") :: P String + ok (parseByteString p "\x01") "a" + it "both fail" $ do + let p = (word8 0x01 *> pure "a") <|> (word8 0x02 *> pure "b") :: P String + bad (parseByteString p "\x03") + it "backtracking restores position" $ do + let p = (word8 0x01 *> word8 0x99) <|> (word8 0x01 *> word8 0x02) :: P () + ok (parseByteString p "\x01\x02") () + it "three-way alternative" $ do + let p = word8 0x01 <|> word8 0x02 <|> word8 0x03 :: P () + ok (parseByteString p "\x03") () + + describe "cut/err" $ do + it "cut converts Fail to Err" $ do + let p = cut (word8 0x01 *> anyWord8 *> word8 0xFF) "bad" :: P () + case parseByteString p "\x01\x42\xAA" of + Left (ParseErr _ e) -> e `shouldBe` "bad" + other -> expectationFailure (show other) + it "err produces Err" $ do + let p = err "fatal" :: P () + case parseByteString p "x" of + Left (ParseErr _ e) -> e `shouldBe` "fatal" + other -> expectationFailure (show other) + it "Err bypasses alternatives" $ do + let p = (err "x" :: P ()) <|> pure () + case parseByteString p "x" of + Left (ParseErr _ _) -> pure () + other -> expectationFailure (show other) + it "try converts Err to Fail" $ do + let p = try (err "x" :: P ()) <|> pure () + ok (parseByteString p "x") () + + describe "withError" $ do + it "catches and handles Err" $ do + let p = withError (\e -> pure (e <> "!")) (err "boom") :: P String + ok (parseByteString p "x") "boom!" + + describe "byte matching" $ do + it "byteString match" $ + ok (parseByteString (byteString "hello" :: P ()) "hello") () + it "byteString mismatch" $ + bad (parseByteString (byteString "hello" :: P ()) "hxllo") + it "byteString empty" $ + ok (parseByteString (byteString "" :: P ()) "anything") () + + describe "takeBs" $ do + it "takes n bytes" $ + ok (parseByteString (takeBs 3 :: P ByteString) "abcde") "abc" + it "fails if short" $ + bad (parseByteString (takeBs 10 :: P ByteString) "abc") + it "take 0 bytes" $ + ok (parseByteString (takeBs 0 :: P ByteString) "x") "" + + describe "skip" $ do + it "skip and continue" $ do + let p = skip 2 *> anyWord8 :: P Word8 + ok (parseByteString p "\x00\x00\x42") 0x42 + + describe "takeRest" $ do + it "consumes all" $ + ok (parseByteString (takeRest :: P ByteString) "hello") "hello" + it "empty on empty" $ + ok (parseByteString (takeRest :: P ByteString) "") "" + + describe "eof" $ do + it "succeeds at end" $ + ok (parseByteString (eof :: P ()) "") () + it "fails with remaining" $ + bad (parseByteString (eof :: P ()) "x") + + describe "atEnd / remaining" $ do + it "atEnd true on empty" $ + ok (parseByteString (atEnd :: P Bool) "") True + it "atEnd false with data" $ + ok (parseByteString (atEnd :: P Bool) "x") False + it "remaining counts bytes" $ + ok (parseByteString (remaining :: P Int) "hello") 5 + + describe "UTF-8" $ do + it "ASCII char" $ + ok (parseByteString (anyAsciiChar :: P Char) "A") 'A' + it "rejects non-ASCII" $ + bad (parseByteString (satisfyAscii (const True) :: P Char) "\xC3\xA9") + it "2-byte UTF-8 (é)" $ + ok (parseByteString (anyChar :: P Char) "\xC3\xA9") '\x00E9' + it "3-byte UTF-8 (€)" $ + ok (parseByteString (anyChar :: P Char) "\xE2\x82\xAC") '\x20AC' + it "4-byte UTF-8 (😀)" $ + ok (parseByteString (anyChar :: P Char) "\xF0\x9F\x98\x80") '\x1F600' + + describe "satisfy" $ do + it "satisfy matches" $ + ok (parseByteString (satisfy (== 'A') :: P Char) "A") 'A' + it "satisfy rejects" $ + bad (parseByteString (satisfy (== 'A') :: P Char) "B") + + describe "ASCII decimal" $ do + it "parses number" $ + ok (parseByteString (anyAsciiDecimalWord :: P Word) "12345x") 12345 + it "fails on non-digit" $ + bad (parseByteString (anyAsciiDecimalWord :: P Word) "abc") + it "single digit" $ + ok (parseByteString (anyAsciiDecimalWord :: P Word) "0") 0 + + describe "hex" $ do + it "parses hex" $ + ok (parseByteString (anyAsciiHexWord :: P Word) "FF") 255 + it "mixed case" $ + ok (parseByteString (anyAsciiHexWord :: P Word) "aB") 0xAB + + describe "lookahead and negative lookahead" $ do + it "lookahead does not consume" $ do + let p = lookahead anyWord8 *> anyWord8 :: P Word8 + ok (parseByteString p "\x42") 0x42 + it "fails succeeds" $ do + let p = fails (word8 0x01) :: P () + ok (parseByteString p "\x02") () + it "notFollowedBy" $ do + let p = notFollowedBy (word8 0x01) :: P () + ok (parseByteString p "\x02") () + + describe "many/some/many_/some_" $ do + it "many_ then read" $ do + let p = many_ (word8 0x41) *> anyWord8 :: P Word8 + ok (parseByteString p "\x41\x41\x42") 0x42 + it "many_ on empty" $ + ok (parseByteString (many_ (word8 0x41) *> eof :: P ()) "") () + it "many collects" $ do + let p = many (word8 0x41 *> pure 'A') :: P [Char] + ok (parseByteString p "\x41\x41\x42") "AA" + it "some requires one" $ + bad (parseByteString (some (word8 0x41) :: P [()]) "\x42") + + describe "isolate" $ do + it "isolate consumes exact bytes" $ do + let p = isolate 3 (takeBs 3) :: P ByteString + ok (parseByteString p "abcdef") "abc" + it "isolate fails if inner underconsumed" $ do + let p = isolate 3 (takeBs 2) :: P ByteString + bad (parseByteString p "abcdef") + + describe "chainl" $ do + it "left-associative chain" $ do + let digit = anyAsciiDecimalInt :: P Int + plus = word8 0x2B *> digit + p = chainl (+) digit plus + ok (parseByteString p "1+2+3x") 6 + + describe "position and span" $ do + it "getPos returns 0 at start" $ do + ok (parseByteString (getPos :: P Pos) "hello") (Pos 0) + it "getPos advances" $ do + let p = skip 3 *> getPos :: P Pos + ok (parseByteString p "hello") (Pos 3) + it "byteStringOf captures consumed bytes" $ do + let p = byteStringOf (skip 3) :: P ByteString + ok (parseByteString p "hello") "hel" + it "withSpan captures span" $ do + let p = withSpan (skip 3) (\_ (Span s e) -> pure (subPos e s)) :: P Int + ok (parseByteString p "hello") 3 + + describe "skipBack" $ do + it "skips backward and re-reads" $ do + let p = do + _ <- anyWord8 + _ <- anyWord8 + skipBack 2 + anyWord8 :: P Word8 + ok (parseByteString p "\xAA\xBB") 0xAA + it "fails when skipping past start" $ do + let p = skipBack 1 :: P () + bad (parseByteString p "x") + it "skip forward then back" $ do + let p = do + skip 3 + skipBack 2 + anyWord8 :: P Word8 + -- "hello" -> skip 3 -> at 'l' -> back 2 -> at 'e' + ok (parseByteString p "hello") (fromIntegral (fromEnum 'e')) + + describe "marks" $ do + it "mark and restore" $ do + let p = do + m <- mark + _ <- anyWord8 + _ <- anyWord8 + restore m + anyWord8 :: P Word8 + ok (parseByteString p "\xAA\xBB") 0xAA + where + isLE :: Bool + isLE = + BS.pack [1, 0] + == BS.pack + ( let w = 1 :: Word16 + in [fromIntegral w, fromIntegral (w `div` 256)] + ) diff --git a/test/Wireform/Ring/Test.hs b/test/Wireform/Ring/Test.hs new file mode 100644 --- /dev/null +++ b/test/Wireform/Ring/Test.hs @@ -0,0 +1,214 @@ +module Wireform.Ring.Test (spec) where + +import Control.Exception (SomeException, catch, evaluate) +import Control.Monad (forM_, when) +import Data.Bits qualified +import Data.ByteString qualified as BS +import Data.IORef +import Data.Word (Word8) +import Foreign.Marshal.Array (pokeArray) +import Foreign.Ptr (Ptr, minusPtr, nullPtr, plusPtr) +import Foreign.Storable (peek, poke) +import Test.QuickCheck +import Test.Syd +import Wireform.Ring +import Wireform.Ring.Internal (destroyMagicRing, newMagicRing) +import Wireform.Transport.Capabilities (capCoreCount, capPageSize, detectCapabilities) + + +spec :: Spec +spec = describe "MagicRing" $ do + describe "basic operations" $ do + it "creates and destroys a ring" $ do + withMagicRing 4096 $ \ring -> do + ringSize ring `shouldSatisfy` (>= 4096) + ringMask ring `shouldBe` (ringSize ring - 1) + ringBase ring `shouldSatisfy` (/= nullPtr) + + it "ring size is power of two" $ do + withMagicRing 5000 $ \ring -> + let s = ringSize ring + in (s Data.Bits..&. (s - 1)) `shouldBe` 0 + + it "ring size >= requested" $ do + withMagicRing 12345 $ \ring -> + ringSize ring `shouldSatisfy` (>= 12345) + + describe "double-mapping correctness" $ do + it "write at end, read from second mapping" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + poke (base `plusPtr` (n - 1)) (0xAB :: Word8) + -- Second mapping: base + n points to same physical page + v <- peek (base `plusPtr` (n - 1)) :: IO Word8 + v `shouldBe` 0xAB + + it "contiguous read spanning boundary" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + forM_ [0 .. 19] $ \i -> + poke (base `plusPtr` (n - 10 + i)) (fromIntegral (i + 1) :: Word8) + bytes <- mapM (\i -> peek (base `plusPtr` (n - 10 + i)) :: IO Word8) [0 .. 19] + bytes `shouldBe` [1 .. 20] + + it "write in first mapping, read from second" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + poke (base `plusPtr` 42) (0xCD :: Word8) + v <- peek (base `plusPtr` (n + 42)) :: IO Word8 + v `shouldBe` 0xCD + + it "write in second mapping, read from first" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + poke (base `plusPtr` (n + 100)) (0xEF :: Word8) + v <- peek (base `plusPtr` 100) :: IO Word8 + v `shouldBe` 0xEF + + it "full ring write + wrap-around read" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + forM_ [0 .. n - 1] $ \i -> + poke (base `plusPtr` i) (fromIntegral (i `mod` 256) :: Word8) + forM_ [0 .. n - 1] $ \i -> do + v <- peek (base `plusPtr` (n + i)) :: IO Word8 + v `shouldBe` fromIntegral (i `mod` 256) + + describe "multiple sizes" $ do + it "works with 64KB ring" $ do + withMagicRing (64 * 1024) $ \ring -> do + let n = ringSize ring + base = ringBase ring + poke (base `plusPtr` (n - 1)) (0xFF :: Word8) + v <- peek (base `plusPtr` (n - 1)) :: IO Word8 + v `shouldBe` 0xFF + + it "works with 1MB ring" $ do + withMagicRing (1024 * 1024) $ \ring -> do + let n = ringSize ring + base = ringBase ring + poke (base `plusPtr` (n - 1)) (0x42 :: Word8) + v <- peek (base `plusPtr` (2 * n - 1)) :: IO Word8 + v `shouldBe` 0x42 + + it "minimum size is at least page-sized" $ do + withMagicRing 1 $ \ring -> + ringSize ring `shouldSatisfy` (>= 4096) + + describe "resource management" $ do + it "sequential allocation: no FD/VA leaks (1000 rings)" $ do + forM_ [1 .. 1000 :: Int] $ \_ -> + withMagicRing 4096 $ \ring -> + ringSize ring `shouldSatisfy` (>= 4096) + + it "absurd size throws MagicRingUnavailable" $ do + let absurdSize = 1024 * 1024 * 1024 * 1024 * 1024 + result <- + (newMagicRing absurdSize >> pure False) + `catch` (\(MagicRingUnavailable _) -> pure True) + result `shouldBe` True + + it "destroy is idempotent" $ do + ring <- newMagicRing 4096 + destroyMagicRing ring + destroyMagicRing ring -- should not crash + describe "stress tests" $ do + it "ring used as circular buffer (100k writes)" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + mask = ringMask ring + forM_ [0 .. 99999 :: Int] $ \i -> do + let off = i Data.Bits..&. mask + poke (base `plusPtr` off) (fromIntegral (i `mod` 251) :: Word8) + v <- peek (base `plusPtr` off) :: IO Word8 + when (v /= fromIntegral (i `mod` 251)) $ + error ("mismatch at iteration " <> show i) + + describe "capabilities detection" $ do + it "detects page size > 0" $ do + caps <- detectCapabilities + capPageSize caps `shouldSatisfy` (> 0) + + it "detects at least 1 core" $ do + caps <- detectCapabilities + capCoreCount caps `shouldSatisfy` (>= 1) + + describe "stress tests (continued)" $ do + it "alternating write/read at boundary" $ do + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + forM_ [0 .. 9999 :: Int] $ \i -> do + let off = (n - 4 + (i `mod` 8)) + poke (base `plusPtr` off) (fromIntegral i :: Word8) + v <- peek (base `plusPtr` off) :: IO Word8 + v `shouldBe` fromIntegral i + + describe "RingSlice (runST-like scoping)" $ do + it "ringSliceLength reports the requested length" $ + withMagicRing 4096 $ \ring -> do + let s = ringSlice ring 0 17 + ringSliceLength s `shouldBe` 17 + + it "copyRingSlice produces a ByteString with the slice's bytes" $ + withMagicRing 4096 $ \ring -> do + let base = ringBase ring + forM_ [0 .. 31 :: Int] $ \i -> + poke (base `plusPtr` i) (fromIntegral (i + 1) :: Word8) + copied <- copyRingSlice (ringSlice ring 0 32) + BS.unpack copied `shouldBe` map fromIntegral [1 .. 32 :: Int] + + it "copyRingSlice handles slices that cross the ring boundary" $ + withMagicRing 4096 $ \ring -> do + let n = ringSize ring + base = ringBase ring + -- Lay down a marker pattern that wraps across the boundary by + -- writing through the second mapping (offsets n-8 .. n+7). + forM_ [0 .. 15 :: Int] $ \i -> + poke (base `plusPtr` (n - 8 + i)) (fromIntegral (i + 1) :: Word8) + -- Slice starting at (n-8), length 16 — reads contiguously + -- through the double mapping. + copied <- copyRingSlice (ringSliceAtPos ring (fromIntegral (n - 8)) 16) + BS.unpack copied `shouldBe` map fromIntegral [1 .. 16 :: Int] + + it "copyRingSlice on an empty slice returns mempty" $ + withMagicRing 4096 $ \ring -> do + copied <- copyRingSlice (ringSlice ring 0 0) + copied `shouldBe` BS.empty + + it "copied ByteString outlives the ring's scope" $ do + -- This compiles iff copyRingSlice's result is independent of @s@. + -- The whole point of the scoping mechanism is that a raw + -- RingSlice s cannot escape but a fresh ByteString can. + bs <- withMagicRing 4096 $ \ring -> do + let base = ringBase ring + forM_ [0 .. 7 :: Int] $ \i -> + poke (base `plusPtr` i) (fromIntegral (0xA0 + i) :: Word8) + copyRingSlice (ringSlice ring 0 8) + BS.unpack bs `shouldBe` [0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7] + + it "withRingSlice exposes pointer and length to the body" $ + withMagicRing 4096 $ \ring -> do + let base = ringBase ring + poke (base `plusPtr` 4) (0x42 :: Word8) + let s = ringSlice ring 4 1 + v <- withRingSlice s $ \p n -> do + n `shouldBe` 1 + peek p :: IO Word8 + v `shouldBe` 0x42 + + it "peekRingSliceByte reads at the given offset" $ + withMagicRing 4096 $ \ring -> do + let base = ringBase ring + forM_ [0 .. 15 :: Int] $ \i -> + poke (base `plusPtr` (100 + i)) (fromIntegral (i * 3) :: Word8) + let s = ringSlice ring 100 16 + forM_ [0 .. 15 :: Int] $ \i -> do + v <- peekRingSliceByte s i + v `shouldBe` fromIntegral (i * 3) diff --git a/test/Wireform/Transport/SendTest.hs b/test/Wireform/Transport/SendTest.hs new file mode 100644 --- /dev/null +++ b/test/Wireform/Transport/SendTest.hs @@ -0,0 +1,534 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Wireform.Transport.SendTest (spec) where + +import Control.Exception (catch) +import Data.Bits ((.&.)) +import Data.ByteString qualified as BS +import Data.IORef +import Data.Word (Word64, Word8) +import Foreign.Ptr (castPtr, plusPtr) +import System.IO.Unsafe (unsafePerformIO) +import Test.QuickCheck hiding ((.&.)) +import Test.QuickCheck.Monadic (assert, monadicIO, run) +import Test.Syd +import Wireform.Builder ( + Builder, + byteString, + byteStringCopy, + byteStringInsert, + toStrictByteString, + word32BE, + word64BE, + word8, + ) +import Wireform.Ring.Internal (ringBase, ringMask, ringSize, withMagicRing) +import Wireform.Transport.Send + + +------------------------------------------------------------------------ +-- Helpers: generate non-trivial byte patterns +------------------------------------------------------------------------ + +patternBytes :: Int -> BS.ByteString +patternBytes n = BS.pack (fmap (\i -> fromIntegral (i `mod` 251) :: Word8) [0 .. n - 1]) + + +patternBuilder :: Int -> Builder +patternBuilder n = byteString (patternBytes n) + + +mixedBuilder :: Int -> Builder +mixedBuilder n = + let header = word32BE 0xDEADBEEF <> word64BE (fromIntegral n) + payload = byteString (patternBytes (max 0 (n - 12))) + in header <> payload + + +mixedBuilderBS :: Int -> BS.ByteString +mixedBuilderBS n = toStrictByteString (mixedBuilder n) + + +------------------------------------------------------------------------ +-- Tests +------------------------------------------------------------------------ + +spec :: Spec +spec = describe "SendTransport" $ do + -- ---------------------------------------------------------------- + -- sendBuilderDirect: basic correctness + -- ---------------------------------------------------------------- + describe "sendBuilderDirect" $ do + it "sends empty builder" $ do + (bytes, pubs) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t mempty + readChunks bytes `shouldBe` BS.empty + readIORef pubs `shouldReturn` 0 + + it "sends a single byte" $ do + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (word8 0x42) + readChunks bytes `shouldBe` BS.singleton 0x42 + + it "sends two bytes" $ do + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (word8 0x42 <> word8 0x43) + readChunks bytes `shouldBe` BS.pack [0x42, 0x43] + + it "sends a pattern payload (100 bytes)" $ do + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder 100) + readChunks bytes `shouldBe` patternBytes 100 + + it "sends exactly ringSize bytes" $ do + sz <- ringSizeAt 4096 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder sz) + readChunks bytes `shouldBe` patternBytes sz + + it "sends ringSize - 1 bytes" $ do + sz <- ringSizeAt 4096 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder (sz - 1)) + readChunks bytes `shouldBe` patternBytes (sz - 1) + + it "sends ringSize + 1 bytes (forces overflow)" $ do + sz <- ringSizeAt 16384 + (bytes, _) <- withTestTransport 16384 $ \t -> + sendBuilderDirect t (patternBuilder (sz + 1)) + readChunks bytes `shouldBe` patternBytes (sz + 1) + + it "sends 2x ringSize bytes" $ do + sz <- ringSizeAt 4096 + let n = sz * 2 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder n) + readChunks bytes `shouldBe` patternBytes n + + it "sends 5x ringSize bytes" $ do + sz <- ringSizeAt 4096 + let n = sz * 5 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder n) + readChunks bytes `shouldBe` patternBytes n + + it "sends 10x ringSize bytes" $ do + sz <- ringSizeAt 4096 + let n = sz * 10 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder n) + readChunks bytes `shouldBe` patternBytes n + + it "sends via byteString insert (large BS, > threshold)" $ do + let payload = patternBytes 12000 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (byteString payload) + readChunks bytes `shouldBe` payload + + it "sends via byteString insert larger than ring" $ do + let payload = patternBytes 50000 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (byteString payload) + readChunks bytes `shouldBe` payload + + it "sends mixed builder (header + payload)" $ do + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (mixedBuilder 500) + readChunks bytes `shouldBe` mixedBuilderBS 500 + + it "sends many small builders concatenated" $ do + let builders = fmap (\i -> word8 (fromIntegral (i :: Int))) [0 .. 255] + b = mconcat builders + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t b + readChunks bytes `shouldBe` BS.pack [0 .. 255] + + it "sends builder with explicit byteStringInsert" $ do + let payload = patternBytes 20000 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendBuilderDirect t (byteStringInsert payload) + readChunks bytes `shouldBe` payload + + -- ---------------------------------------------------------------- + -- sendBuilderDirect vs sendBuilderViaByteString equivalence + -- ---------------------------------------------------------------- + describe "sendBuilderDirect vs sendBuilderViaByteString" $ do + it "identical output for small payload" $ do + let b = mixedBuilder 200 + (directBytes, _) <- withTestTransport 65536 $ \t -> + sendBuilderDirect t b + (viaBytes, _) <- withTestTransport 65536 $ \t -> + sendBuilderViaByteString t b + readChunks directBytes `shouldBe` readChunks viaBytes + + it "identical output for medium payload" $ do + let b = mixedBuilder 5000 + (directBytes, _) <- withTestTransport 65536 $ \t -> + sendBuilderDirect t b + (viaBytes, _) <- withTestTransport 65536 $ \t -> + sendBuilderViaByteString t b + readChunks directBytes `shouldBe` readChunks viaBytes + + it "identical output for large payload (multi-overflow)" $ do + let b = mixedBuilder 40000 + (directBytes, _) <- withTestTransport 65536 $ \t -> + sendBuilderDirect t b + (viaBytes, _) <- withTestTransport 65536 $ \t -> + sendBuilderViaByteString t b + readChunks directBytes `shouldBe` readChunks viaBytes + + it "identical output for byteString insert path" $ do + let b = byteString (patternBytes 25000) + (directBytes, _) <- withTestTransport 32768 $ \t -> + sendBuilderDirect t b + (viaBytes, _) <- withTestTransport 32768 $ \t -> + sendBuilderViaByteString t b + readChunks directBytes `shouldBe` readChunks viaBytes + + -- ---------------------------------------------------------------- + -- sendByteString + -- ---------------------------------------------------------------- + describe "sendByteString" $ do + it "sends empty BS (no-op)" $ do + (bytes, pubs) <- withTestTransport 4096 $ \t -> + sendByteString t BS.empty + readChunks bytes `shouldBe` BS.empty + readIORef pubs `shouldReturn` 0 + + it "sends 1 byte" $ do + (bytes, _) <- withTestTransport 4096 $ \t -> + sendByteString t (BS.singleton 0xAB) + readChunks bytes `shouldBe` BS.singleton 0xAB + + it "sends exactly ringSize bytes" $ do + sz <- ringSizeAt 4096 + (bytes, _) <- withTestTransport 4096 $ \t -> + sendByteString t (patternBytes sz) + readChunks bytes `shouldBe` patternBytes sz + + it "sendByteStringMany coalesces" $ do + let bss = [patternBytes 100, patternBytes 200, patternBytes 50] + (bytes, pubs) <- withTestTransport 4096 $ \t -> + sendByteStringMany t bss + readChunks bytes `shouldBe` BS.concat bss + readIORef pubs `shouldReturn` 1 + + -- ---------------------------------------------------------------- + -- sendByteString: reservation too large + -- ---------------------------------------------------------------- + describe "sendByteString limits" $ do + it "throws SendReservationTooLarge for BS > ringSize" $ do + threw <- withMagicRing 4096 $ \ring -> do + let payload = patternBytes (ringSize ring + 1) + headRef' <- newIORef (0 :: Word64) + tailRef' <- newIORef (0 :: Word64) + let t' = + SendTransport + { sendRingBase = ringBase ring + , sendRingSize = ringSize ring + , sendRingMask = ringMask ring + , sendLoadTail = readIORef tailRef' + , sendLoadHead = readIORef headRef' + , sendPublishHead = \h -> writeIORef headRef' h >> writeIORef tailRef' h + , sendWaitSpace = \_ -> SendSpaceAvailable <$> readIORef tailRef' + , sendFlush = pure () + , sendShutdownWrite = pure () + , sendClose = pure () + } + (sendByteString t' payload >> pure False) + `catch` (\(SendReservationTooLarge _ _) -> pure True) + threw `shouldBe` True + + -- ---------------------------------------------------------------- + -- withSendCork: batching + -- ---------------------------------------------------------------- + describe "withSendCork" $ do + it "empty cork: zero publishes, zero bytes" $ do + (bytes, pubs) <- withTestTransport 4096 $ \t -> + withSendCork t $ \_ -> pure () + readChunks bytes `shouldBe` BS.empty + readIORef pubs `shouldReturn` 0 + + it "single send inside cork: one publish" $ do + (bytes, pubs) <- withTestTransport 65536 $ \t -> + withSendCork t $ \corked -> + sendByteString corked (patternBytes 100) + readChunks bytes `shouldBe` patternBytes 100 + readIORef pubs `shouldReturn` 1 + + it "batches two sendByteString into one publish" $ do + let a = patternBytes 100 + b = patternBytes 200 + (bytes, pubs) <- withTestTransport 65536 $ \t -> + withSendCork t $ \corked -> do + sendByteString corked a + sendByteString corked b + readChunks bytes `shouldBe` BS.append a b + readIORef pubs `shouldReturn` 1 + + it "batches three sendByteString into one publish" $ do + let a = patternBytes 100 + b = patternBytes 200 + c = patternBytes 300 + (bytes, pubs) <- withTestTransport 65536 $ \t -> + withSendCork t $ \corked -> do + sendByteString corked a + sendByteString corked b + sendByteString corked c + readChunks bytes `shouldBe` BS.concat [a, b, c] + readIORef pubs `shouldReturn` 1 + + it "batches builder + bytestring into one publish" $ do + let header = word8 0x01 <> word32BE 0x12345678 + body = patternBytes 50 + (bytes, pubs) <- withTestTransport 65536 $ \t -> + withSendCork t $ \corked -> do + sendBuilderDirect corked header + sendByteString corked body + let result = readChunks bytes + BS.take 5 result `shouldBe` BS.pack [0x01, 0x12, 0x34, 0x56, 0x78] + BS.drop 5 result `shouldBe` body + readIORef pubs `shouldReturn` 1 + + it "batches sendBuilderDirect + sendBuilderDirect" $ do + (bytes, pubs) <- withTestTransport 65536 $ \t -> + withSendCork t $ \corked -> do + sendBuilderDirect corked (mixedBuilder 200) + sendBuilderDirect corked (mixedBuilder 300) + readChunks bytes + `shouldBe` BS.append (mixedBuilderBS 200) (mixedBuilderBS 300) + readIORef pubs `shouldReturn` 1 + + it "corked payload exactly ring size" $ do + sz <- ringSizeAt 4096 + (bytes, _) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + sendBuilderDirect corked (patternBuilder sz) + readChunks bytes `shouldBe` patternBytes sz + + -- ---------------------------------------------------------------- + -- withSendCork: backpressure + -- ---------------------------------------------------------------- + describe "withSendCork backpressure" $ do + it "builder payload 2x ring" $ do + sz <- ringSizeAt 4096 + let n = sz * 2 + (bytes, pubs) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + sendBuilderDirect corked (patternBuilder n) + readChunks bytes `shouldBe` patternBytes n + readIORef pubs >>= (`shouldSatisfy` (> 1)) + + it "builder payload 5x ring" $ do + sz <- ringSizeAt 4096 + let n = sz * 5 + (bytes, _) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + sendBuilderDirect corked (patternBuilder n) + readChunks bytes `shouldBe` patternBytes n + + it "builder payload 10x ring" $ do + sz <- ringSizeAt 4096 + let n = sz * 10 + (bytes, _) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + sendBuilderDirect corked (patternBuilder n) + readChunks bytes `shouldBe` patternBytes n + + it "byteString insert > ring inside cork" $ do + sz <- ringSizeAt 4096 + let payload = patternBytes (sz + 1) + (bytes, pubs) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + sendBuilderDirect corked (byteString payload) + readChunks bytes `shouldBe` payload + readIORef pubs >>= (`shouldSatisfy` (> 1)) + + it "many small sends exceeding ring" $ do + sz <- ringSizeAt 4096 + let chunk = patternBytes 500 + numChunks = sz `quot` BS.length chunk + 1 + (bytes, _) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + mapM_ (\_ -> sendByteString corked chunk) [1 .. numChunks] + readChunks bytes `shouldBe` BS.concat (replicate numChunks chunk) + + it "many small builders exceeding ring" $ do + sz <- ringSizeAt 4096 + let numBuilders = sz `quot` 4 + 1 + (bytes, _) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + mapM_ + (sendBuilderDirect corked . word32BE . fromIntegral) + [0 .. numBuilders - 1] + let expected = + toStrictByteString $ + mconcat (fmap (word32BE . fromIntegral) [0 .. numBuilders - 1 :: Int]) + readChunks bytes `shouldBe` expected + + it "interleaved builder + BS exceeding ring" $ do + sz <- ringSizeAt 4096 + let n = sz `quot` 204 + 1 + (bytes, _) <- withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + mapM_ + ( \i -> do + sendBuilderDirect corked (word32BE (fromIntegral i)) + sendByteString corked (patternBytes 200) + ) + [0 .. n - 1] + BS.length (readChunks bytes) `shouldBe` n * (4 + 200) + + -- ---------------------------------------------------------------- + -- Sequential cork + uncork operations + -- ---------------------------------------------------------------- + describe "sequential cork/uncork" $ do + it "cork then uncorked send" $ do + (bytes, _) <- withTestTransport 65536 $ \t -> do + withSendCork t $ \corked -> + sendByteString corked (BS.pack [0x01, 0x02]) + sendByteString t (BS.pack [0x03, 0x04]) + readChunks bytes `shouldBe` BS.pack [0x01, 0x02, 0x03, 0x04] + + it "two sequential corks" $ do + (bytes, pubs) <- withTestTransport 65536 $ \t -> do + withSendCork t $ \corked -> do + sendByteString corked (BS.pack [0x01]) + sendByteString corked (BS.pack [0x02]) + withSendCork t $ \corked -> do + sendByteString corked (BS.pack [0x03]) + sendByteString corked (BS.pack [0x04]) + readChunks bytes `shouldBe` BS.pack [0x01, 0x02, 0x03, 0x04] + readIORef pubs `shouldReturn` 2 + + it "three sequential corks with varying sizes" $ do + (bytes, pubs) <- withTestTransport 65536 $ \t -> do + withSendCork t $ \c -> sendBuilderDirect c (patternBuilder 100) + withSendCork t $ \c -> sendBuilderDirect c (patternBuilder 500) + withSendCork t $ \c -> sendBuilderDirect c (patternBuilder 1000) + readChunks bytes + `shouldBe` BS.concat [patternBytes 100, patternBytes 500, patternBytes 1000] + readIORef pubs `shouldReturn` 3 + + it "cork then many uncorked sends" $ do + (bytes, _) <- withTestTransport 65536 $ \t -> do + withSendCork t $ \c -> + sendByteString c (patternBytes 50) + mapM_ (sendByteString t . BS.singleton . fromIntegral) [0 .. 9 :: Int] + readChunks bytes + `shouldBe` BS.append (patternBytes 50) (BS.pack [0 .. 9]) + + -- ---------------------------------------------------------------- + -- Property-based: arbitrary payload sizes + -- ---------------------------------------------------------------- + describe "property-based" $ do + it "sendBuilderDirect produces correct output for arbitrary sizes" $ + property $ \(Positive (n :: Int)) -> monadicIO $ do + let sz = min n 100000 + (bytes, _) <- run $ withTestTransport 4096 $ \t -> + sendBuilderDirect t (patternBuilder sz) + assert (readChunks bytes == patternBytes sz) + + it "sendBuilderDirect == sendBuilderViaByteString for arbitrary builders" $ + property $ \(Positive (n :: Int)) -> monadicIO $ do + let sz = min n 50000 + b = mixedBuilder sz + (directBytes, _) <- run $ withTestTransport 65536 $ \t -> + sendBuilderDirect t b + (viaBytes, _) <- run $ withTestTransport 65536 $ \t -> + sendBuilderViaByteString t b + assert (readChunks directBytes == readChunks viaBytes) + + it "corked sendBuilderDirect produces correct output for arbitrary sizes" $ + property $ \(Positive (n :: Int)) -> monadicIO $ do + let sz = min n 100000 + (bytes, _) <- run $ withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + sendBuilderDirect corked (patternBuilder sz) + assert (readChunks bytes == patternBytes sz) + + it "corked multi-send produces correct output" $ + property $ \(Positive (n :: Int)) -> monadicIO $ do + let numSends = min n 50 + chunk = patternBytes 100 + (bytes, _) <- run $ withTestTransport 4096 $ \t -> + withSendCork t $ \corked -> + mapM_ (\_ -> sendByteString corked chunk) [1 .. numSends] + assert (readChunks bytes == BS.concat (replicate numSends chunk)) + + +------------------------------------------------------------------------ +-- Test infrastructure +------------------------------------------------------------------------ + +ringSizeAt :: Int -> IO Int +ringSizeAt requested = withMagicRing requested (pure . ringSize) + + +withTestTransport + :: Int + -> (SendTransport -> IO a) + -> IO (IORef [BS.ByteString], IORef Int) +withTestTransport requestedRingSz action = + withMagicRing requestedRingSz $ \ring -> do + let !base = ringBase ring + !sz = ringSize ring + !msk = ringMask ring + chunksRef <- newIORef ([] :: [BS.ByteString]) + publishCount <- newIORef (0 :: Int) + headRef <- newIORef (0 :: Word64) + tailRef <- newIORef (0 :: Word64) + stateRef <- newIORef True + + let drainTo !hd = do + tl <- readIORef tailRef + drainLoop tl + where + drainLoop !cur + | cur >= hd = pure () + | otherwise = do + let !off = fromIntegral cur .&. msk + !want = min (fromIntegral (hd - cur)) (sz - off) + !ptr = base `plusPtr` off + bs <- BS.packCStringLen (castPtr ptr, want) + modifyIORef' chunksRef (bs :) + let !newTail = cur + fromIntegral want + writeIORef tailRef newTail + drainLoop newTail + + let publish h = do + modifyIORef' publishCount (+ 1) + writeIORef headRef h + drainTo h + + let waitSpace _pos = do + isOpen <- readIORef stateRef + if isOpen + then do + tl <- readIORef tailRef + pure (SendSpaceAvailable tl) + else pure SendPeerClosed + + let transport = + SendTransport + { sendRingBase = base + , sendRingSize = sz + , sendRingMask = msk + , sendLoadTail = readIORef tailRef + , sendLoadHead = readIORef headRef + , sendPublishHead = publish + , sendWaitSpace = waitSpace + , sendFlush = readIORef headRef >>= drainTo + , sendShutdownWrite = pure () + , sendClose = writeIORef stateRef False + } + _ <- action transport + pure (chunksRef, publishCount) + + +readChunks :: IORef [BS.ByteString] -> BS.ByteString +readChunks ref = unsafePerformIO $ do + chunks <- readIORef ref + pure (BS.concat (reverse chunks)) +{-# NOINLINE readChunks #-} diff --git a/wireform-core.cabal b/wireform-core.cabal --- a/wireform-core.cabal +++ b/wireform-core.cabal @@ -1,11 +1,24 @@ cabal-version: 3.0 name: wireform-core -version: 0.1.0.0 +version: 0.2.0.0 synopsis: Shared FFI primitives for wireform format packages description: Cross-format SWAR / SIMD-accelerated primitives used by every @wireform-*@ format package. . + Modules: + . + * @Wireform.FFI@ -- C FFI surface: packed-varint pre-scan + decode, + SWAR UTF-8 validation, SIMD NUL / byte / ASCII / whitespace + scanners, SIMD JSON escape scanner, SIMD Iceberg bounds compare, + Arrow IPC buffer validation, endianness helpers. + * @Wireform.Encode.Direct@ -- shared direct-write encode buffer + used by the per-format encoders. + * @Wireform.Hash@ -- SIMD-accelerated hashing helpers. + . + Pulls in C sources (@cbits\/fast_decode.c@, @cbits\/fast_scan.c@, + @cbits\/wireform_hash_simd.c@) and the vendored @simde@ headers. + Contains no format-specific code. homepage: https://github.com/iand675/wireform- bug-reports: https://github.com/iand675/wireform-/issues license: BSD-3-Clause @@ -45,11 +58,13 @@ DeriveAnyClass OverloadedStrings LambdaCase + BlockArguments library import: defaults hs-source-dirs: src exposed-modules: + -- Existing builder modules Wireform.Builder Wireform.Builder.FastBuilder Wireform.Builder.Internal.Prim @@ -57,13 +72,103 @@ Wireform.FFI Wireform.Encode.Direct Wireform.Hash + Wireform.Base64 + + -- Magic ring + Wireform.Ring + Wireform.Ring.Internal + Wireform.Ring.Pool + + -- Transport abstraction (symmetric send / receive) + Wireform.Transport + Wireform.Transport.Receive + Wireform.Transport.Send + Wireform.Transport.Config + Wireform.Transport.Capabilities + + -- Parser + Wireform.Parser + Wireform.Parser.Internal + Wireform.Parser.Stateful + Wireform.Parser.Switch + Wireform.Parser.Position + Wireform.Parser.Mark + Wireform.Parser.Error + + -- Driver + Wireform.Parser.Driver + + -- External-parser adapter + Wireform.Parser.Adapter + c-sources: cbits/fast_decode.c cbits/fast_scan.c cbits/wireform_hash_simd.c + cbits/base64.c + cbits/fast_rng.c include-dirs: include/simde + include + cc-options: -O3 + + if os(linux) + c-sources: cbits/magic_ring_linux.c + if os(darwin) || os(freebsd) || os(netbsd) || os(openbsd) + c-sources: cbits/magic_ring_posix.c + if os(windows) + c-sources: cbits/magic_ring_windows.c + extra-libraries: kernel32 + build-depends: - base >= 4.16 && < 5, + base >= 4.18 && < 5, bytestring >= 0.11 && < 0.13, - ghc-prim >= 0.10 && < 0.12, + containers, + ghc-prim, + template-haskell, text >= 2.0 && < 2.2, vector >= 0.13 && < 0.14 + +test-suite wireform-core-test + import: defaults + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + other-modules: + Wireform.Ring.Test + Wireform.Parser.Test + Wireform.Transport.SendTest + Wireform.Base64.Test + build-depends: + base, + bytestring, + wireform-core, + sydtest, + QuickCheck + +benchmark parser-bench + import: defaults + type: exitcode-stdio-1.0 + hs-source-dirs: bench + main-is: ParserBench.hs + other-modules: + Common + WFBasic + FPBasic + ghc-options: -O2 -rtsopts -threaded + build-depends: + base, + bytestring, + wireform-core, + flatparse >= 0.5 && < 0.6, + criterion >= 1.5 && < 1.7 + +benchmark send-builder-bench + import: defaults + type: exitcode-stdio-1.0 + hs-source-dirs: bench + main-is: SendBuilderBench.hs + ghc-options: -O2 -rtsopts -threaded + build-depends: + base, + bytestring, + wireform-core, + criterion >= 1.5 && < 1.7