Decode.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE UnboxedTuples #-}

{- | Low-level, high-performance wire format decoding primitives.

The decoder uses unboxed sums for the result type, avoiding heap
allocation for intermediate decode results. Each decode operation
returns @(# (# a, Int# #) | DecodeError #)@ — either a value with
a new offset (on the stack, not heap) or an error.

This approach is more robust than CPS: it doesn't depend on GHC
successfully inlining all continuations, and the result type can
be unboxed by GHC. On modern CPUs the branch prediction for the
success/failure case split is near-perfect since decoding almost
always succeeds.
-}
module Proto.Internal.Wire.Decode (
  -- * Decode result
  DecodeResult (..),
  DecodeError (..),

  -- * Varint decoding
  getVarint,
  getVarintSigned,
  getSVarint32,
  getSVarint64,

  -- * Fixed-width decoding
  getFixed32,
  getFixed64,
  getFloat,
  getDouble,

  -- * Length-delimited
  getLengthDelimited,
  getByteString,
  getText,

  -- * Tags
  getTag,
  getTagOr,

  -- * Skipping unknown fields
  skipField,

  -- * Running a decoder
  runDecoder,
  Decoder (..),

  -- * ZigZag
  unZigZag32,
  unZigZag64,

  -- * Low-level access (for generated code)
  runDecoder',

  -- * Unboxed internal variants (zero-allocation hot path)
  UMaybe (UJust, UNothing),
  umaybe,
  getTagOrU,

  -- * Three-way tag result (flattened unboxed sum for the decode loop)
  TagResult#,
  withTag,

  -- * Monadic CPS tag dispatch (zero Tag allocation, for generated code)
  withTagM,
  skipWireType,

  -- * Non-throwing UTF-8 validation
  validateUtf8,
) where

import Control.DeepSeq (NFData (..))
import Data.Bits (shiftL, shiftR, xor, (.&.), (.|.))
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.Int (Int32, Int64)
import Data.Text (Text)
import Data.Text.Encoding qualified as TE
import Data.Word (Word32, Word64)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import Foreign.Storable (peek)
import GHC.Exts (Int (I#), Int#, isTrue#, (+#), (>=#))
import GHC.Float (castWord32ToFloat, castWord64ToDouble)
import Proto.Internal.Wire (Tag (..), WireType (..), decodeTag)
import System.IO.Unsafe (unsafeDupablePerformIO)
import Wireform.FFI (validateUtf8SWAR)


-- | Errors that can occur during protobuf wire-format decoding.
data DecodeError
  = UnexpectedEnd
  | InvalidVarint
  | InvalidTag !Word64
  | InvalidWireType !Int
  | InvalidUtf8
  | NegativeLength
  | ExtraBytes
  | SubMessageError !DecodeError
  | CustomError !String
  deriving stock (Show, Eq)


instance NFData DecodeError where
  rnf UnexpectedEnd = ()
  rnf InvalidVarint = ()
  rnf (InvalidTag w) = rnf w
  rnf (InvalidWireType i) = rnf i
  rnf InvalidUtf8 = ()
  rnf NegativeLength = ()
  rnf ExtraBytes = ()
  rnf (SubMessageError e) = rnf e
  rnf (CustomError s) = rnf s


-- | Legacy result type, kept for compatibility with runDecoder'.
data DecodeResult a
  = DecodeOK !a {-# UNPACK #-} !Int
  | DecodeFail !DecodeError
  deriving stock (Show)


{- | Decoder monad using unboxed sums for the result.
Returns either (value, new_offset) or an error, with no heap
allocation for the result envelope.
-}
newtype Decoder a = Decoder
  { runDecoder# :: ByteString -> Int# -> (# (# a, Int# #) | DecodeError #)
  }


instance Functor Decoder where
  fmap f (Decoder g) = Decoder $ \bs off -> case g bs off of
    (# (# a, off' #) | #) -> (# (# f a, off' #) | #)
    (# | e #) -> (# | e #)
  {-# INLINE fmap #-}


instance Applicative Decoder where
  pure a = Decoder $ \_ off -> (# (# a, off #) | #)
  {-# INLINE pure #-}
  Decoder f <*> Decoder g = Decoder $ \bs off -> case f bs off of
    (# (# fab, off' #) | #) -> case g bs off' of
      (# (# a, off'' #) | #) -> (# (# fab a, off'' #) | #)
      (# | e #) -> (# | e #)
    (# | e #) -> (# | e #)
  {-# INLINE (<*>) #-}


instance Monad Decoder where
  Decoder g >>= f = Decoder $ \bs off -> case g bs off of
    (# (# a, off' #) | #) -> runDecoder# (f a) bs off'
    (# | e #) -> (# | e #)
  {-# INLINE (>>=) #-}


-- | Run a decoder on a ByteString, producing Either.
runDecoder :: Decoder a -> ByteString -> Either DecodeError a
runDecoder (Decoder f) bs =
  case f bs 0# of
    (# (# a, off' #) | #)
      | I# off' == BS.length bs -> Right a
      | otherwise -> Left ExtraBytes
    (# | e #) -> Left e
{-# INLINE runDecoder #-}


-- | Run a decoder returning the legacy DecodeResult (for internal use).
runDecoder' :: Decoder a -> ByteString -> Int -> DecodeResult a
runDecoder' (Decoder f) bs (I# off) =
  case f bs off of
    (# (# a, off' #) | #) -> DecodeOK a (I# off')
    (# | e #) -> DecodeFail e
{-# INLINE runDecoder' #-}


-- Helpers for offset arithmetic
bsLen :: ByteString -> Int#
bsLen bs = case BS.length bs of I# n -> n
{-# INLINE bsLen #-}


{- | Decode a varint. Inline fast path for 1-4 byte varints.

Hyperpb insight: most varints are 1-3 bytes (tags, small integers,
enum values). Inlining up to 4 bytes covers field tags up to
field number ~500k and values up to 2^28, hitting the slow path
only for genuinely large values.
-}
getVarint :: Decoder Word64
getVarint = Decoder $ \bs off ->
  let len = bsLen bs
  in if isTrue# (off >=# len)
       then (# | UnexpectedEnd #)
       else
         let !b0 = fromIntegral (BSU.unsafeIndex bs (I# off)) :: Word64
         in if b0 < 0x80
              then (# (# b0, off +# 1# #) | #)
              else
                let off1 = off +# 1#
                in if isTrue# (off1 >=# len)
                     then (# | UnexpectedEnd #)
                     else
                       let !b1 = fromIntegral (BSU.unsafeIndex bs (I# off1)) :: Word64
                       in if b1 < 0x80
                            then (# (# (b0 .&. 0x7F) .|. (b1 `shiftL` 7), off +# 2# #) | #)
                            else
                              let off2 = off +# 2#
                              in if isTrue# (off2 >=# len)
                                   then (# | UnexpectedEnd #)
                                   else
                                     let !b2 = fromIntegral (BSU.unsafeIndex bs (I# off2)) :: Word64
                                     in if b2 < 0x80
                                          then
                                            (#
                                              (#
                                                (b0 .&. 0x7F) .|. ((b1 .&. 0x7F) `shiftL` 7) .|. (b2 `shiftL` 14)
                                                , off +# 3#
                                              #) |
                                            #)
                                          else
                                            let off3 = off +# 3#
                                            in if isTrue# (off3 >=# len)
                                                 then (# | UnexpectedEnd #)
                                                 else
                                                   let !b3 = fromIntegral (BSU.unsafeIndex bs (I# off3)) :: Word64
                                                   in if b3 < 0x80
                                                        then
                                                          (#
                                                            (#
                                                              (b0 .&. 0x7F)
                                                                .|. ((b1 .&. 0x7F) `shiftL` 7)
                                                                .|. ((b2 .&. 0x7F) `shiftL` 14)
                                                                .|. (b3 `shiftL` 21)
                                                              , off +# 4#
                                                            #) |
                                                          #)
                                                        else getVarintSlow bs off
{-# INLINE getVarint #-}


getVarintSlow :: ByteString -> Int# -> (# (# Word64, Int# #) | DecodeError #)
getVarintSlow bs = go 0 0
  where
    len = bsLen bs
    go :: Word64 -> Int -> Int# -> (# (# Word64, Int# #) | DecodeError #)
    go !acc !shift !pos
      | shift > 63 = (# | InvalidVarint #)
      | isTrue# (pos >=# len) = (# | UnexpectedEnd #)
      | otherwise =
          let !b = BSU.unsafeIndex bs (I# pos)
              !val = acc .|. ((fromIntegral b .&. 0x7F) `shiftL` shift)
          in if b < 0x80
               then (# (# val, pos +# 1# #) | #)
               else go val (shift + 7) (pos +# 1#)
{-# INLINE getVarintSlow #-}


-- | Decode a varint as a signed 'Int64'.
getVarintSigned :: Decoder Int64
getVarintSigned = fromIntegral <$> getVarint
{-# INLINE getVarintSigned #-}


-- | Decode a ZigZag-encoded sint32 value.
getSVarint32 :: Decoder Int32
getSVarint32 = unZigZag32 . fromIntegral <$> getVarint
{-# INLINE getSVarint32 #-}


-- | Decode a ZigZag-encoded sint64 value.
getSVarint64 :: Decoder Int64
getSVarint64 = unZigZag64 <$> getVarint
{-# INLINE getSVarint64 #-}


-- | Decode a ZigZag-encoded 32-bit value back to a signed 'Int32'.
unZigZag32 :: Word32 -> Int32
unZigZag32 n = fromIntegral ((n `shiftR` 1) `xor` negate (n .&. 1))
{-# INLINE unZigZag32 #-}


-- | Decode a ZigZag-encoded 64-bit value back to a signed 'Int64'.
unZigZag64 :: Word64 -> Int64
unZigZag64 n = fromIntegral ((n `shiftR` 1) `xor` negate (n .&. 1))
{-# INLINE unZigZag64 #-}


{- | Decode a fixed32 (little-endian) via a single aligned/unaligned
word load.  On x86_64 and aarch64-LE this compiles to one MOV.
-}
getFixed32 :: Decoder Word32
getFixed32 = Decoder $ \bs off ->
  if I# (off +# 4#) > BS.length bs
    then (# | UnexpectedEnd #)
    else
      let !val = readWord32LE bs (I# off)
      in (# (# val, off +# 4# #) | #)
{-# INLINE getFixed32 #-}


-- | Decode a fixed64 (little-endian) via a single word load.
getFixed64 :: Decoder Word64
getFixed64 = Decoder $ \bs off ->
  if I# (off +# 8#) > BS.length bs
    then (# | UnexpectedEnd #)
    else
      let !val = readWord64LE bs (I# off)
      in (# (# val, off +# 8# #) | #)
{-# INLINE getFixed64 #-}


-- Direct word-sized reads from a ByteString.
-- On little-endian platforms (all targets we care about: x86_64, aarch64-LE)
-- this is a single unaligned load instruction, replacing 4 or 8 separate
-- byte reads + shifts + ORs.

readWord32LE :: ByteString -> Int -> Word32
readWord32LE (BSI.BS fp _) off = unsafeDupablePerformIO $
  withForeignPtr fp $ \ptr ->
    peek (castPtr (ptr `plusPtr` off) :: Ptr Word32)
{-# INLINE readWord32LE #-}


readWord64LE :: ByteString -> Int -> Word64
readWord64LE (BSI.BS fp _) off = unsafeDupablePerformIO $
  withForeignPtr fp $ \ptr ->
    peek (castPtr (ptr `plusPtr` off) :: Ptr Word64)
{-# INLINE readWord64LE #-}


-- | Decode a 32-bit IEEE 754 float from a fixed32 wire value.
getFloat :: Decoder Float
getFloat = castWord32ToFloat <$> getFixed32
{-# INLINE getFloat #-}


-- | Decode a 64-bit IEEE 754 double from a fixed64 wire value.
getDouble :: Decoder Double
getDouble = castWord64ToDouble <$> getFixed64
{-# INLINE getDouble #-}


-- | Decode a length-delimited field. Zero-copy ByteString slice.
getLengthDelimited :: Decoder ByteString
getLengthDelimited = Decoder $ \bs off ->
  case runDecoder# getVarint bs off of
    (# (# lenW, off' #) | #) ->
      let !len = fromIntegral lenW :: Int
      in if len < 0
           then (# | NegativeLength #)
           else
             if I# off' + len > BS.length bs
               then (# | UnexpectedEnd #)
               else
                 let !i = I# off'
                 in (# (# BSU.unsafeTake len (BSU.unsafeDrop i bs), case i + len of I# r -> r #) | #)
    (# | e #) -> (# | e #)
{-# INLINE getLengthDelimited #-}


-- | Decode a length-delimited bytes field (alias for 'getLengthDelimited').
getByteString :: Decoder ByteString
getByteString = getLengthDelimited
{-# INLINE getByteString #-}


{- | Decode a text field.

We rely on text >= 2.0's simdutf-powered 'decodeUtf8'' for UTF-8
validation and decoding in a single pass. The text library uses
AVX2/NEON internally, which is faster than a separate pre-check + decode.
-}
getText :: Decoder Text
getText = Decoder $ \bs off ->
  case runDecoder# getLengthDelimited bs off of
    (# (# bytes, off' #) | #) ->
      case TE.decodeUtf8' bytes of
        Right t -> (# (# t, off' #) | #)
        Left _ -> (# | InvalidUtf8 #)
    (# | e #) -> (# | e #)
{-# INLINE getText #-}


-- | Decode a field tag (field number + wire type). Fails on invalid tags.
getTag :: Decoder Tag
getTag = Decoder $ \bs off ->
  case runDecoder# getVarint bs off of
    (# (# w, off' #) | #) ->
      case decodeTag w of
        Just tag -> (# (# tag, off' #) | #)
        Nothing -> (# | InvalidTag w #)
    (# | e #) -> (# | e #)
{-# INLINE getTag #-}


-- | Try to decode a tag, returning Nothing at end-of-input.
getTagOr :: Decoder (Maybe Tag)
getTagOr = Decoder $ \bs off ->
  if isTrue# (off >=# bsLen bs)
    then (# (# Nothing, off #) | #)
    else case runDecoder# getTag bs off of
      (# (# tag, off' #) | #) -> (# (# Just tag, off' #) | #)
      (# | e #) -> (# | e #)
{-# INLINE getTagOr #-}


-- | Skip over a field value based on its wire type.
skipField :: WireType -> Decoder ()
skipField = \case
  WireVarint -> skipVarint
  Wire64Bit -> skip 8
  WireLengthDelimited -> Decoder $ \bs off ->
    case runDecoder# getVarint bs off of
      (# (# lenW, off' #) | #) ->
        let !len = fromIntegral lenW :: Int
        in if I# off' + len > BS.length bs
             then (# | UnexpectedEnd #)
             else (# (# (), case I# off' + len of I# r -> r #) | #)
      (# | e #) -> (# | e #)
  WireStartGroup -> skipGroup
  WireEndGroup -> pure ()
  Wire32Bit -> skip 4


skip :: Int -> Decoder ()
skip (I# n) = Decoder $ \bs off ->
  if I# (off +# n) > BS.length bs
    then (# | UnexpectedEnd #)
    else (# (# (), off +# n #) | #)
{-# INLINE skip #-}


skipVarint :: Decoder ()
skipVarint = Decoder $ \bs off0 ->
  let len = bsLen bs
      go !pos
        | isTrue# (pos >=# len) = (# | UnexpectedEnd #)
        | BSU.unsafeIndex bs (I# pos) < 0x80 = (# (# (), pos +# 1# #) | #)
        | otherwise = go (pos +# 1#)
  in go off0
{-# INLINE skipVarint #-}


skipGroup :: Decoder ()
skipGroup = Decoder $ \bs off ->
  case runDecoder# getTagOrU bs off of
    (# (# mt, off' #) | #) -> case mt of
      UNothing -> (# | UnexpectedEnd #)
      UJust (Tag _ WireEndGroup) -> (# (# (), off' #) | #)
      UJust (Tag _ wt) -> runDecoder# (skipField wt >> skipGroup) bs off'
    (# | e #) -> (# | e #)


-- | Unboxed optional for zero-allocation tag-or-EOF in the decode loop.
data UMaybe a = UMaybe (# (# #) | a #)


-- | A 'UMaybe' containing a value.
pattern UJust :: a -> UMaybe a
pattern UJust a = UMaybe (# | a #)


-- | An empty 'UMaybe'.
pattern UNothing :: UMaybe a
pattern UNothing = UMaybe (# (# #) | #)


{-# COMPLETE UJust, UNothing #-}


-- | Eliminate a 'UMaybe': supply a default for 'UNothing' and a function for 'UJust'.
umaybe :: b -> (a -> b) -> UMaybe a -> b
umaybe def f (UMaybe x) = case x of
  (# (# #) | #) -> def
  (# | a #) -> f a
{-# INLINE umaybe #-}


-- | Like 'getTagOr' but returns 'UMaybe' to avoid allocating a boxed Maybe.
getTagOrU :: Decoder (UMaybe Tag)
getTagOrU = Decoder $ \bs off ->
  if isTrue# (off >=# bsLen bs)
    then (# (# UNothing, off #) | #)
    else case runDecoder# getTag bs off of
      (# (# tag, off' #) | #) -> (# (# UJust tag, off' #) | #)
      (# | e #) -> (# | e #)
{-# INLINE getTagOrU #-}


{- | Three-way unboxed result for the tag-or-EOF operation.
Flattens what would otherwise be two nested unboxed sums
(Decoder result × UMaybe) into a single three-way split.

* @(# (# #) | _ | _ #)@ — end of input (offset unchanged)
* @(# _ | (# Int#, Int#, Int# #) | _ #)@ — got a tag: field number, wire type, new offset
* @(# _ | _ | DecodeError #)@ — decode error
-}
type TagResult# = (# (# #) | (# Int#, Int#, Int# #) | DecodeError #)


{- | CPS interface to the three-way tag result, specialized for decoder results.
Avoids constructing any intermediate value — the continuation is applied
directly to the unboxed field number and wire type.
-}
withTag
  :: ByteString
  -> Int#
  -> (Int# -> (# (# a, Int# #) | DecodeError #))
  -> (Int# -> Int# -> Int# -> (# (# a, Int# #) | DecodeError #))
  -> (DecodeError -> (# (# a, Int# #) | DecodeError #))
  -> (# (# a, Int# #) | DecodeError #)
withTag bs off kEOF kTag kErr =
  if isTrue# (off >=# bsLen bs)
    then kEOF off
    else case runDecoder# getVarint bs off of
      (# (# w, off' #) | #) ->
        case decodeTagParts w of
          (# fn, wt #) -> kTag fn wt off'
      (# | e #) -> kErr e
{-# INLINE withTag #-}


-- | Decode tag into unboxed field number and wire type.
decodeTagParts :: Word64 -> (# Int#, Int# #)
decodeTagParts w =
  let fn = fromIntegral (w `shiftR` 3) :: Int
      wt = fromIntegral (w .&. 0x07) :: Int
  in case fn of
       I# fn# -> case wt of
         I# wt# -> (# fn#, wt# #)
{-# INLINE decodeTagParts #-}


{- | Monadic CPS tag dispatch for generated decoders.

At end-of-input: calls @kEOF@.
On a valid tag: calls @kTag fieldNumber wireType@, where both are
unboxed 'Int' values (no Tag constructor allocated).
On error: propagates the decode error.

This is the monadic counterpart to 'withTag', intended for generated
code. Avoids allocating the 'Tag' record that 'getTagOrU' produces.
-}
withTagM
  :: Decoder a
  -- ^ kEOF: continuation at end-of-input
  -> (Int -> Int -> Decoder a)
  -- ^ kTag: continuation with (fieldNumber, wireType)
  -> Decoder a
withTagM (Decoder kEOF) kTag = Decoder $ \bs off ->
  if isTrue# (off >=# bsLen bs)
    then kEOF bs off
    else case runDecoder# getVarint bs off of
      (# (# w, off' #) | #) ->
        case decodeTagParts w of
          (# fn#, wt# #) -> runDecoder# (kTag (I# fn#) (I# wt#)) bs off'
      (# | e #) -> (# | e #)
{-# INLINE withTagM #-}


-- | Skip a field given its wire type as an 'Int' (for use with 'withTagM').
skipWireType :: Int -> Decoder ()
skipWireType wt = case wt of
  0 -> skipVarint
  1 -> skip 8
  2 -> Decoder $ \bs off ->
    case runDecoder# getVarint bs off of
      (# (# lenW, off' #) | #) ->
        let !len = fromIntegral lenW :: Int
        in if I# off' + len > BS.length bs
             then (# | UnexpectedEnd #)
             else (# (# (), case I# off' + len of I# r -> r #) | #)
      (# | e #) -> (# | e #)
  5 -> skip 4
  _ -> Decoder $ \_ _ -> (# | InvalidWireType wt #)
{-# INLINE skipWireType #-}


{- | Validate UTF-8 without exceptions.

Uses the SWAR-accelerated C validator. Useful for paths that need
validation without decoding (e.g. conformance checks).

Note: 'getText' and 'decodeTextFast' do /not/ use this — they rely on
text >= 2.0's simdutf-powered 'TE.decodeUtf8'' which validates and
decodes in a single pass.
-}
validateUtf8 :: ByteString -> Bool
validateUtf8 = validateUtf8SWAR
{-# INLINE validateUtf8 #-}