fix(pyrowave): guard 4:4:4 modes that overflow the rate controller's block index

The vendored rate controller packs its wavelet block index into 16 bits
(RDOperation.block_offset_saving), so a mode whose 32x32-block count exceeds
u16::MAX wraps inside the controller and corrupts the bitstream — ~8K-class
4:4:4 territory. Compute the exact count (`block_count_32x32`, the counting walk
of upstream init_block_meta, pinned against the validated Apple WaveletLayout)
and expose `pyrowave_mode_fits_rdo`; the negotiator downgrades such a session to
4:2:0 before the Welcome (the honest-downgrade channel), and both encoders
refuse outright if one slips through rather than emit a wrapped stream.

Vendor patches: 0002-rdo-saving-clamp (analyze_rate_control.comp clamps the
saving accumulation to the target, same overrun class as 0001; slangmosh.hpp
regenerated), 0003-devel-encode-16bit-read (devel tool y4m 16-bit plane reads;
tool-only, kept so the vendored source stays honest).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

crates/pf-encode/src/enc/linux/pyrowave.rs

@@ -227,6 +227,14 @@ impl PyroWaveEncoder {
         if !chroma.is_444() && (width % 2 != 0 || height % 2 != 0) {
             bail!("pyrowave 4:2:0 needs even dimensions (got {width}x{height})");
         }
+        if chroma.is_444() && !crate::pyrowave_mode_fits_rdo(width, height, true) {
+            // The negotiator downgrades these modes to 4:2:0 pre-Welcome; refuse if one
+            // slips through (e.g. the lab override) rather than wrap the RDO block index.
+            bail!(
+                "pyrowave 4:4:4 at {width}x{height} exceeds the rate controller's 16-bit \
+                 block index (see pyrowave-sys patches/0002 note) — use 4:2:0 at this size"
+            );
+        }
         // SAFETY: `open_inner` only issues Vulkan/pyrowave calls whose preconditions it
         // establishes itself (valid instance/device, correctly-chained create-infos that
         // `DeviceHold` keeps alive); all handles are freshly created and owned by the result.

crates/pf-encode/src/enc/pyrowave_wire.rs

@@ -53,6 +53,34 @@ pub(crate) fn stamp_color_bits(bitstream: &mut [u8], seq_offset: usize, bt2020_p
     }
 }
 
+/// The wavelet block space's total 32x32-block count for a mode — the exact counting walk of
+/// upstream `WaveletBuffers::init_block_meta` (also ported to the Apple `WaveletLayout`, whose
+/// golden tests pin it against real host AUs). Needed because the vendored RDO pass packs the
+/// block index into 16 bits (`RDOperation.block_offset_saving` — see
+/// `patches/0002-rdo-saving-clamp.patch`): a mode whose count exceeds `u16::MAX` would wrap
+/// inside the rate controller, so the host guards such modes out (≈8K 4:4:4 territory).
+pub(crate) fn block_count_32x32(width: u32, height: u32, chroma444: bool) -> u32 {
+    const LEVELS: u32 = 5;
+    let align = |v: u32| ((v + 31) & !31).max(128);
+    let (aw, ah) = (align(width), align(height));
+    let mut count = 0u32;
+    for level in (0..LEVELS).rev() {
+        let lw = (aw / 2) >> level;
+        let lh = (ah / 2) >> level;
+        let blocks_x8 = lw.div_ceil(8);
+        let blocks_y8 = lh.div_ceil(8);
+        let per_band = blocks_x8.div_ceil(4) * blocks_y8.div_ceil(4);
+        let bands = if level == LEVELS - 1 { 4 } else { 3 };
+        for component in 0..3u32 {
+            if level == 0 && component != 0 && !chroma444 {
+                continue;
+            }
+            count += per_band * bands;
+        }
+    }
+    count
+}
+
 /// Frame pyrowave's `packets` (each an `(offset, size)` into `bitstream`) into the wire AU.
 /// `wire_chunk = None` copies the single dense packet; `Some(chunk)` produces the windowed
 /// datagram-aligned AU (a whole number of `chunk`-sized windows).
@@ -189,6 +217,37 @@ mod tests {
         assert_eq!(packet_boundary(None, 777), 777);
     }
 
+    #[test]
+    fn block_count_matches_the_apple_layout_invariant() {
+        // 256x144 (the golden-fixture geometry, aligned 256x160): recompute via the same walk
+        // the validated Apple WaveletLayout uses and pin a few mode-level facts.
+        let manual = |w: u32, h: u32, c444: bool| {
+            let align = |v: u32| ((v + 31) & !31).max(128);
+            let (aw, ah) = (align(w), align(h));
+            let mut n = 0u32;
+            for level in (0..5u32).rev() {
+                let per = ((aw / 2 >> level).div_ceil(8).div_ceil(4))
+                    * ((ah / 2 >> level).div_ceil(8).div_ceil(4));
+                let bands = if level == 4 { 4 } else { 3 };
+                for c in 0..3 {
+                    if level == 0 && c != 0 && !c444 {
+                        continue;
+                    }
+                    n += per * bands;
+                }
+            }
+            n
+        };
+        for (w, h) in [(256, 144), (1920, 1080), (3840, 2160), (7680, 4320)] {
+            assert_eq!(block_count_32x32(w, h, false), manual(w, h, false));
+            assert_eq!(block_count_32x32(w, h, true), manual(w, h, true));
+        }
+        // 4:4:4 fits comfortably at 4K; the 16-bit RDO block index wraps around 8K 4:4:4.
+        assert!(block_count_32x32(3840, 2160, true) <= u16::MAX as u32);
+        assert!(block_count_32x32(7680, 4320, true) > u16::MAX as u32);
+        assert!(block_count_32x32(7680, 4320, false) <= u16::MAX as u32);
+    }
+
     #[test]
     fn stamp_color_bits_sets_range_and_hdr_bits() {
         let mut bs = vec![0u8; 16];

crates/pf-encode/src/enc/windows/pyrowave.rs

@@ -182,6 +182,14 @@ impl PyroWaveEncoder {
         if !chroma444 && (width % 2 != 0 || height % 2 != 0) {
             bail!("pyrowave 4:2:0 needs even dimensions (got {width}x{height})");
         }
+        if chroma444 && !crate::pyrowave_mode_fits_rdo(width, height, true) {
+            // The negotiator downgrades these modes to 4:2:0 pre-Welcome; refuse if one
+            // slips through (e.g. the lab override) rather than wrap the RDO block index.
+            bail!(
+                "pyrowave 4:4:4 at {width}x{height} exceeds the rate controller's 16-bit \
+                 block index (see pyrowave-sys patches/0002 note) — use 4:2:0 at this size"
+            );
+        }
         let fps = fps.max(1);
         // Select pyrowave's device by the SELECTED render adapter's vendor/device-id — NOT by LUID:
         // in Session 0 (the host service context) the Vulkan ICD reports `deviceLUIDValid = false`,

crates/pf-encode/src/lib.rs

@@ -1326,6 +1326,18 @@ mod pyrowave;
 #[path = "enc/pyrowave_wire.rs"]
 mod pyrowave_wire;
 
+/// Whether a PyroWave mode fits the vendored rate controller's packed 16-bit block index
+/// (`patches/0002-rdo-saving-clamp.patch` note): false ≈ 8K-class 4:4:4. The negotiator
+/// downgrades such a session to 4:2:0 before the Welcome; the encoders also refuse outright.
+#[cfg(all(any(target_os = "linux", target_os = "windows"), feature = "pyrowave"))]
+pub fn pyrowave_mode_fits_rdo(width: u32, height: u32, chroma444: bool) -> bool {
+    pyrowave_wire::block_count_32x32(width, height, chroma444) <= u16::MAX as u32
+}
+#[cfg(not(all(any(target_os = "linux", target_os = "windows"), feature = "pyrowave")))]
+pub fn pyrowave_mode_fits_rdo(_width: u32, _height: u32, _chroma444: bool) -> bool {
+    false
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/punktfunk-host/src/native/handshake.rs

@@ -277,6 +277,22 @@ pub(super) async fn negotiate(
     } else {
         crate::encode::ChromaFormat::Yuv420
     };
+    // PyroWave-only mode-size gate: the vendored rate controller packs its block index into
+    // 16 bits (pyrowave-sys patches/0002 note), which ≈8K-class 4:4:4 overflows — downgrade
+    // to 4:2:0 BEFORE the Welcome (the honest-downgrade channel), like every gate above.
+    let chroma = if codec == crate::encode::Codec::PyroWave
+        && chroma.is_444()
+        && !crate::encode::pyrowave_mode_fits_rdo(hello.mode.width, hello.mode.height, true)
+    {
+        tracing::warn!(
+            mode = %format_args!("{}x{}", hello.mode.width, hello.mode.height),
+            "PyroWave 4:4:4 at this mode exceeds the rate controller's block-index range — \
+             negotiating 4:2:0"
+        );
+        crate::encode::ChromaFormat::Yuv420
+    } else {
+        chroma
+    };
     tracing::info!(
         chroma = ?chroma,
         host_wants_444,

crates/pyrowave-sys/patches/0003-devel-encode-16bit-read.patch

@@ -0,0 +1,16 @@
+diff --git a/crates/pyrowave-sys/vendor/pyrowave/encode.cpp b/crates/pyrowave-sys/vendor/pyrowave/encode.cpp
+index 95725551..2e78fada 100644
+--- a/crates/pyrowave-sys/vendor/pyrowave/encode.cpp
++++ b/crates/pyrowave-sys/vendor/pyrowave/encode.cpp
+@@ -201,7 +201,10 @@ static void run_encoder(Device &device, const char *out_path, const char *in_pat
+ 		for (auto &img : inputs.images)
+ 		{
+ 			auto *y = cmd->update_image(*img);
+-			if (!input.read(y, img->get_width() * img->get_height()))
++			// PUNKTFUNK PATCH 0003: size the plane read in BYTES, not texels — 16-bit y4m
++			// inputs otherwise read half a plane and desync the stream after frame 1.
++			if (!input.read(y, size_t(img->get_width()) * img->get_height() *
++			                       YUV4MPEGFile::format_to_bytes_per_component(input.get_format())))
+ 			{
+ 				LOGE("Failed to read plane.\n");
+ 				device.submit_discard(cmd);

crates/pyrowave-sys/vendor/pyrowave/PUNKTFUNK-VENDOR.txt

@@ -9,14 +9,29 @@ Tree is pruned to what the pyrowave-sys standalone build needs
 (see the rm -rf list in the script). All parts are MIT-licensed
 (pyrowave, Granite) or Apache-2.0/MIT (volk, Vulkan-Headers).
 
-Local patches (crates/pyrowave-sys/patches/, re-applied on re-vendor):
+Local patches (crates/pyrowave-sys/patches/, re-applied on re-vendor).
+These are OUR fixes — kept local by decision (we vendor anyway), not filed
+upstream:
   0001-payload-data-444-sizing.patch — encoder payload_data worst-case buffer
     was sized for 4:2:0's 1.5 samples/px; busy 4:4:4 (3 samples/px) overran it
     on the GPU → nondeterministic corrupt bitstreams/crashes at any bitrate.
-    Found + validated 2026-07-18 (RTX 5070 Ti, 1080p/4K, 8/16-bit); to be
-    reported upstream.
-
-  (0002-0003 are reserved by concurrent in-flight work and land separately.)
+    Found + validated 2026-07-18 (RTX 5070 Ti, 1080p/4K, 8/16-bit).
+  0002-rdo-saving-clamp.patch — analyze_rate_control.comp accumulates the full
+    32-bit per-block rate saving into the RDO bucket totals but packs only 16
+    bits into each RDOperation; once a block's saving exceeds 65535 cost units
+    the resolve pass over-credits applied ops and the bitstream can overshoot
+    the hard rate target (the same overrun class as 0001). Clamped to the
+    packed width (conservative direction). Includes the regenerated
+    shaders/slangmosh.hpp (built with the pinned Granite slangmosh, -O --strip
+    per upstream's slangmosh.sh); behavior-neutral at our operating points
+    (byte-identical outputs, validated on the RTX 5070 Ti).
+    NOTE the related UNPATCHED limit: the packed 16-bit block_index wraps when
+    block_count_32x32 > 65535 (~8K 4:4:4) — guarded host-side instead
+    (pf-encode rejects such modes for PyroWave).
+  0003-devel-encode-16bit-read.patch — devel tool encode.cpp read y4m planes
+    with texel-count math (bytes for 8-bit): 16-bit inputs got half-plane
+    reads and a desynced stream after frame 1. Tool-only (our build never
+    compiles the devel tools; kept so the vendored source is honest).
 
   0004-encoder-buffer-pool.patch — pyrowave_c.cpp allocated four Vulkan buffers
     (meta + bitstream, Device + CachedHost) on EVERY encode. At 240 fps with

crates/pyrowave-sys/vendor/pyrowave/encode.cpp

@@ -201,7 +201,10 @@ static void run_encoder(Device &device, const char *out_path, const char *in_pat
 		for (auto &img : inputs.images)
 		{
 			auto *y = cmd->update_image(*img);
-			if (!input.read(y, img->get_width() * img->get_height()))
+			// PUNKTFUNK PATCH 0003: size the plane read in BYTES, not texels — 16-bit y4m
+			// inputs otherwise read half a plane and desync the stream after frame 1.
+			if (!input.read(y, size_t(img->get_width()) * img->get_height() *
+			                       YUV4MPEGFile::format_to_bytes_per_component(input.get_format())))
 			{
 				LOGE("Failed to read plane.\n");
 				device.submit_discard(cmd);

crates/pyrowave-sys/vendor/pyrowave/shaders/analyze_rate_control.comp

@@ -132,6 +132,10 @@ void emit_rdo_operations()
     else if (gl_SubgroupInvocationID < 16)
     {
         uint saving = shared_rate_cost[gl_SubgroupInvocationID - 1] - shared_rate_cost[gl_SubgroupInvocationID];
+        // PUNKTFUNK PATCH 0002: RDOperation packs saving into 16 bits; clamp the bucket
+        // accounting to match, else resolve over-credits applied ops and the produced
+        // bitstream can overshoot the hard rate target (a buffer-overrun class).
+        saving = min(saving, 65535u);
 
         if (saving != 0)
         {