Stop fully inlining of packed field encoders. (#101)

We want to inline the field number because it is often a constant. But by forcing the entire fold to be INLINE, not just the field number, we force an impossible choice on proto3-suite: inline nothing or everything. Therefore we leave the actual folds merely INLINABLE. (If that is insufficient control then we may have to expose a more complex API.) Also give the compiler more discretion about whether to inline varints.
awakesecurity · Oct 11, 2023 · 9385232 · 9385232
1 parent 6dcc557
commit 9385232
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 28 deletions.
diff --git a/src/Proto3/Wire/Encode.hs b/src/Proto3/Wire/Encode.hs
@@ -535,9 +535,9 @@ shortByteString num = embedded num . MessageBuilder . RB.shortByteString
 -- >>> packedVarints 1 [1, 2, 3]
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\ETX\SOH\STX\ETX"
 packedVarints :: Foldable f => FieldNumber -> f Word64 -> MessageBuilder
-packedVarints num =
-    etaMessageBuilder
-      (embedded num . foldMap (liftBoundedPrim . base128Varint64))
+packedVarints num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = foldMap (liftBoundedPrim . base128Varint64)
 {-# INLINE packedVarints #-}
 
 -- | A faster but more specialized variant of:
@@ -548,8 +548,9 @@ packedVarints num =
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\ETX\SOH\STX\ETX"
 packedVarintsV ::
   Vector v a => (a -> Word64) -> FieldNumber -> v a -> MessageBuilder
-packedVarintsV f num =
-    embedded num . vectorMessageBuilder (liftBoundedPrim . base128Varint64 . f)
+packedVarintsV f num = embedded num . payload
+  where
+    payload = vectorMessageBuilder (liftBoundedPrim . base128Varint64 . f)
 {-# INLINE packedVarintsV #-}
 
 -- | A faster but more specialized variant of:
@@ -560,10 +561,9 @@ packedVarintsV f num =
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\STX\SOH\NUL"
 packedBoolsV ::
   Vector v a => (a -> Bool) -> FieldNumber -> v a -> MessageBuilder
-packedBoolsV f num =
-    embedded num . MessageBuilder . Prim.vectorFixedPrim op
+packedBoolsV f num = embedded num . MessageBuilder . payload
   where
-    op = Prim.word8 . fromIntegral . fromEnum . f
+    payload = Prim.vectorFixedPrim (Prim.word8 . fromIntegral . fromEnum . f)
 {-# INLINE packedBoolsV #-}
 
 -- | Encode fixed-width Word32s in the space-efficient packed format.
@@ -574,8 +574,9 @@ packedBoolsV f num =
 -- >>> packedFixed32 1 [1, 2, 3]
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\f\SOH\NUL\NUL\NUL\STX\NUL\NUL\NUL\ETX\NUL\NUL\NUL"
 packedFixed32 :: Foldable f => FieldNumber -> f Word32 -> MessageBuilder
-packedFixed32 num =
-    etaMessageBuilder (embedded num . foldMap (MessageBuilder . RB.word32LE))
+packedFixed32 num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = foldMap (MessageBuilder . RB.word32LE)
 {-# INLINE packedFixed32 #-}
 
 -- | A faster but more specialized variant of:
@@ -586,8 +587,9 @@ packedFixed32 num =
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\f\SOH\NUL\NUL\NUL\STX\NUL\NUL\NUL\ETX\NUL\NUL\NUL"
 packedFixed32V ::
   Vector v a => (a -> Word32) -> FieldNumber -> v a -> MessageBuilder
-packedFixed32V f num =
-    embedded num . MessageBuilder . Prim.vectorFixedPrim (Prim.word32LE . f)
+packedFixed32V f num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = MessageBuilder . Prim.vectorFixedPrim (Prim.word32LE . f)
 {-# INLINE packedFixed32V #-}
 
 -- | Encode fixed-width Word64s in the space-efficient packed format.
@@ -598,8 +600,9 @@ packedFixed32V f num =
 -- >>> packedFixed64 1 [1, 2, 3]
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\CAN\SOH\NUL\NUL\NUL\NUL\NUL\NUL\NUL\STX\NUL\NUL\NUL\NUL\NUL\NUL\NUL\ETX\NUL\NUL\NUL\NUL\NUL\NUL\NUL"
 packedFixed64 :: Foldable f => FieldNumber -> f Word64 -> MessageBuilder
-packedFixed64 num =
-    etaMessageBuilder (embedded num . foldMap (MessageBuilder . RB.word64LE))
+packedFixed64 num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = foldMap (MessageBuilder . RB.word64LE)
 {-# INLINE packedFixed64 #-}
 
 -- | A faster but more specialized variant of:
@@ -610,8 +613,9 @@ packedFixed64 num =
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\CAN\SOH\NUL\NUL\NUL\NUL\NUL\NUL\NUL\STX\NUL\NUL\NUL\NUL\NUL\NUL\NUL\ETX\NUL\NUL\NUL\NUL\NUL\NUL\NUL"
 packedFixed64V ::
   Vector v a => (a -> Word64) -> FieldNumber -> v a -> MessageBuilder
-packedFixed64V f num =
-    embedded num . MessageBuilder . Prim.vectorFixedPrim (Prim.word64LE . f)
+packedFixed64V f num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = MessageBuilder . Prim.vectorFixedPrim (Prim.word64LE . f)
 {-# INLINE packedFixed64V #-}
 
 -- | Encode floats in the space-efficient packed format.
@@ -620,8 +624,9 @@ packedFixed64V f num =
 -- >>> 1 `packedFloats` [1, 2, 3]
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\f\NUL\NUL\128?\NUL\NUL\NUL@\NUL\NUL@@"
 packedFloats :: Foldable f => FieldNumber -> f Float -> MessageBuilder
-packedFloats num =
-    etaMessageBuilder (embedded num . foldMap (MessageBuilder . RB.floatLE))
+packedFloats num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = foldMap (MessageBuilder . RB.floatLE)
 {-# INLINE packedFloats #-}
 
 -- | A faster but more specialized variant of:
@@ -632,8 +637,9 @@ packedFloats num =
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\f\NUL\NUL\128?\NUL\NUL\NUL@\NUL\NUL@@"
 packedFloatsV ::
   Vector v a => (a -> Float) -> FieldNumber -> v a -> MessageBuilder
-packedFloatsV f num =
-    embedded num . MessageBuilder . Prim.vectorFixedPrim (Prim.floatLE . f)
+packedFloatsV f num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = MessageBuilder . Prim.vectorFixedPrim (Prim.floatLE . f)
 {-# INLINE packedFloatsV #-}
 
 -- | Encode doubles in the space-efficient packed format.
@@ -642,8 +648,9 @@ packedFloatsV f num =
 -- >>> 1 `packedDoubles` [1, 2, 3]
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\CAN\NUL\NUL\NUL\NUL\NUL\NUL\240?\NUL\NUL\NUL\NUL\NUL\NUL\NUL@\NUL\NUL\NUL\NUL\NUL\NUL\b@"
 packedDoubles :: Foldable f => FieldNumber -> f Double -> MessageBuilder
-packedDoubles num =
-    etaMessageBuilder (embedded num . foldMap (MessageBuilder . RB.doubleLE))
+packedDoubles num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = foldMap (MessageBuilder . RB.doubleLE)
 {-# INLINE packedDoubles #-}
 
 -- | A faster but more specialized variant of:
@@ -654,8 +661,9 @@ packedDoubles num =
 -- Proto3.Wire.Encode.unsafeFromLazyByteString "\n\CAN\NUL\NUL\NUL\NUL\NUL\NUL\240?\NUL\NUL\NUL\NUL\NUL\NUL\NUL@\NUL\NUL\NUL\NUL\NUL\NUL\b@"
 packedDoublesV ::
   Vector v a => (a -> Double) -> FieldNumber -> v a -> MessageBuilder
-packedDoublesV f num =
-    embedded num . MessageBuilder . Prim.vectorFixedPrim (Prim.doubleLE . f)
+packedDoublesV f num = etaMessageBuilder (embedded num . payload)
+  where
+    payload = MessageBuilder . Prim.vectorFixedPrim (Prim.doubleLE . f)
 {-# INLINE packedDoublesV #-}
 
 -- | Encode an embedded message.

diff --git a/src/Proto3/Wire/Reverse/Prim.hs b/src/Proto3/Wire/Reverse/Prim.hs
@@ -711,8 +711,13 @@ wordBase128LEVar_inline (W# w) = word64Base128LEVar_inline (W64# w)
 -- | The bounded primitive implementing
 -- `Proto3.Wire.Reverse.word32Base128LEVar`.
 word32Base128LEVar :: Word32 -> BoundedPrim 5
-word32Base128LEVar = word32Base128LEVar_inline
-{-# INLINE word32Base128LEVar #-}
+word32Base128LEVar (W32# x0) =
+  ( wordBase128LEVar_choose 1 wordBase128LE_p1 $
+    wordBase128LEVar_choose 2 wordBase128LE_p2 $
+    wordBase128LEVar_choose 3 wordBase128LE_p3 $
+    wordBase128LEVar_choose 4 wordBase128LE_p4 $
+    (\x -> liftFixedPrim (wordBase128LE_p5 0## x))
+  ) x0
 
 -- | Like 'word32Base128LEVar' but inlined, which currently means
 -- that it is just the same as 'word32Base128LEVar', which we inline.
@@ -798,14 +803,13 @@ word64Base128LEVar = \(W64# x) ->
     pif (W64# x <= fromIntegral (maxBound :: Word32))
           (word32Base128LEVar (fromIntegral (W64# x)))
           (word64Base128LEVar_big x)
-{-# INLINE word64Base128LEVar #-}
 
 -- | Like 'word64Base128LEVar' but inlined, possibly bloating your code.  On
 -- the other hand, inlining an application to a constant may shrink your code.
 word64Base128LEVar_inline :: Word64 -> BoundedPrim 10
 word64Base128LEVar_inline = \(W64# x) ->
     pif (W64# x <= fromIntegral (maxBound :: Word32))
-          (word32Base128LEVar (fromIntegral (W64# x)))
+          (word32Base128LEVar_inline (fromIntegral (W64# x)))
           (inline (word64Base128LEVar_big x))
 {-# INLINE word64Base128LEVar_inline #-}