Skip to main content

weir/dispatch_decode/
pack.rs

1//! Shared GPU dispatch-output decoders for Weir via paths.
2//!
3//! Production GPU wrappers must reject malformed backend output instead of
4//! accepting prefixes or silently ignoring trailing bytes. Keep the contract in
5//! one place so every dataflow primitive reports the same failure mode.
6
7use super::{byte_len_for_u32_words, DispatchU32Schema};
8
9/// Encode every input word as little-endian bytes for GPU dispatch.
10#[inline]
11#[must_use]
12pub(crate) fn pack_u32(values: &[u32]) -> Vec<u8> {
13    #[cfg(target_endian = "little")]
14    {
15        bytemuck::cast_slice(values).to_vec()
16    }
17    #[cfg(target_endian = "big")]
18    {
19        let mut bytes = Vec::with_capacity(values.len().saturating_mul(4));
20        for value in values {
21            bytes.extend_from_slice(&value.to_le_bytes());
22        }
23        bytes
24    }
25}
26
27/// Encode every input word as little-endian bytes for GPU dispatch with an
28/// explicit allocator failure contract.
29#[inline]
30pub(crate) fn try_pack_u32(values: &[u32], field: &'static str) -> Result<Vec<u8>, String> {
31    let byte_len = byte_len_for_u32_words(values.len(), field)?;
32    #[cfg(target_endian = "little")]
33    {
34        let bytes = bytemuck::cast_slice(values);
35        if bytes.len() != byte_len {
36            // Cast width mismatch is a cold malformed-input path.
37            return Err(pack_u32_cast_mismatch(
38                field,
39                values.len(),
40                bytes.len(),
41                byte_len,
42            ));
43        }
44        let mut out = crate::staging_reserve::reserved_vec(byte_len, field)?;
45        out.extend_from_slice(bytes);
46        Ok(out)
47    }
48    #[cfg(target_endian = "big")]
49    let mut bytes = crate::staging_reserve::reserved_vec(byte_len, field)?;
50    #[cfg(target_endian = "big")]
51    for value in values {
52        bytes.extend_from_slice(&value.to_le_bytes());
53    }
54    #[cfg(target_endian = "big")]
55    Ok(bytes)
56}
57
58/// Encode one repeated input word as little-endian bytes for GPU dispatch.
59#[inline]
60pub(crate) fn pack_repeated_u32(value: u32, words: usize) -> Result<Vec<u8>, String> {
61    let byte_len = byte_len_for_u32_words(words, "pack_repeated_u32")?;
62    if value == 0 {
63        let mut bytes = Vec::new();
64        try_write_zero_bytes(&mut bytes, byte_len, "repeated zero u32 byte pack")?;
65        return Ok(bytes);
66    }
67    let mut bytes = crate::staging_reserve::reserved_vec(byte_len, "repeated u32 byte pack")?;
68    pack_repeated_u32_into(value, words, &mut bytes)?;
69    Ok(bytes)
70}
71
72/// Encode every input word as little-endian bytes into caller-owned scratch.
73#[inline]
74#[cfg(any(test, feature = "legacy-infallible"))]
75pub fn pack_u32_into(values: &[u32], bytes: &mut Vec<u8>) {
76    try_pack_u32_into(values, bytes, "u32 byte pack")
77        .expect("u32 byte pack allocation failed in legacy infallible caller");
78}
79
80/// Encode every input word as little-endian bytes into caller-owned scratch
81/// with an explicit allocator failure contract.
82#[inline]
83pub(crate) fn try_pack_u32_into(
84    values: &[u32],
85    bytes: &mut Vec<u8>,
86    field: &'static str,
87) -> Result<(), String> {
88    let byte_len = byte_len_for_u32_words(values.len(), field)?;
89    let stable_len = bytes.len() == byte_len;
90    crate::staging_reserve::reserve_vec(bytes, byte_len, field)?;
91    #[cfg(target_endian = "little")]
92    {
93        let packed = bytemuck::cast_slice(values);
94        if packed.len() != byte_len {
95            // Cast width mismatch is a cold malformed-input path.
96            return Err(pack_u32_cast_mismatch(
97                field,
98                values.len(),
99                packed.len(),
100                byte_len,
101            ));
102        }
103        if stable_len {
104            bytes.copy_from_slice(packed);
105        } else {
106            bytes.clear();
107            bytes.extend_from_slice(packed);
108        }
109    }
110    #[cfg(target_endian = "big")]
111    {
112        if stable_len {
113            for (slot, value) in bytes.chunks_exact_mut(4).zip(values) {
114                slot.copy_from_slice(&value.to_le_bytes());
115            }
116        } else {
117            bytes.clear();
118            for value in values {
119                bytes.extend_from_slice(&value.to_le_bytes());
120            }
121        }
122    }
123    Ok(())
124}
125
126/// Encode one repeated input word as little-endian bytes into caller-owned scratch.
127#[inline]
128pub(crate) fn pack_repeated_u32_into(
129    value: u32,
130    words: usize,
131    bytes: &mut Vec<u8>,
132) -> Result<(), String> {
133    let byte_len = byte_len_for_u32_words(words, "pack_repeated_u32_into")?;
134    if value == 0 {
135        return try_write_zero_bytes(bytes, byte_len, "repeated zero u32 byte pack");
136    }
137    let stable_len = bytes.len() == byte_len;
138    crate::staging_reserve::reserve_vec(bytes, byte_len, "repeated u32 byte pack")?;
139    let word = value.to_le_bytes();
140    if stable_len {
141        for slot in bytes.chunks_exact_mut(4) {
142            slot.copy_from_slice(&word);
143        }
144    } else {
145        bytes.clear();
146        for _ in 0..words {
147            bytes.extend_from_slice(&word);
148        }
149    }
150    Ok(())
151}
152
153/// Encode exactly `slots` little-endian words into caller-owned scratch.
154///
155/// # Errors
156///
157/// Returns an actionable ABI error when the caller supplies too few or too
158/// many words.
159#[inline]
160pub(crate) fn pack_exact_u32_slots_into(
161    values: &[u32],
162    slots: usize,
163    name: &str,
164    bytes: &mut Vec<u8>,
165) -> Result<(), String> {
166    DispatchU32Schema::input(name, slots).require_input_words(values.len())?;
167    try_pack_u32_into(values, bytes, "exact u32 slot pack")?;
168    Ok(())
169}
170
171/// Extend a semantically exact byte buffer to a backend dispatch minimum.
172///
173/// This is not semantic padding: callers have already passed
174/// [`pack_exact_u32_slots`]. The extra zero bytes exist only for backends that
175/// require at least one storage word even when the semantic workload is empty.
176#[inline]
177pub(crate) fn pad_dispatch_min_words(
178    bytes: &mut Vec<u8>,
179    dispatch_slots: usize,
180) -> Result<(), String> {
181    try_pad_zero_bytes_to_len(
182        bytes,
183        byte_len_for_u32_words(dispatch_slots, "pad_dispatch_min_words")?,
184        "dispatch minimum word padding",
185    )?;
186    Ok(())
187}
188
189/// Write an exact zero-filled byte buffer into caller-owned scratch with an
190/// explicit allocator failure contract.
191#[inline]
192pub(crate) fn try_write_zero_bytes(
193    bytes: &mut Vec<u8>,
194    byte_len: usize,
195    field: &'static str,
196) -> Result<(), String> {
197    if bytes.len() == byte_len {
198        bytes.fill(0);
199        return Ok(());
200    }
201    crate::staging_reserve::reserve_vec(bytes, byte_len, field)?;
202    bytes.clear();
203    extend_zero_bytes(bytes, byte_len);
204    Ok(())
205}
206
207/// Write an exact zero-filled word buffer into caller-owned scratch.
208#[inline]
209#[cfg(any(test, feature = "legacy-infallible"))]
210pub(crate) fn write_zero_words(words: &mut Vec<u32>, word_len: usize) {
211    try_write_zero_words(words, word_len, "zero word staging")
212        .expect("zero word staging allocation failed in legacy infallible caller");
213}
214
215/// Write an exact zero-filled word buffer into caller-owned scratch with an
216/// explicit allocator failure contract.
217#[inline]
218pub(crate) fn try_write_zero_words(
219    words: &mut Vec<u32>,
220    word_len: usize,
221    field: &'static str,
222) -> Result<(), String> {
223    if words.len() == word_len {
224        words.fill(0);
225        return Ok(());
226    }
227    crate::staging_reserve::reserve_vec(words, word_len, field)?;
228    words.clear();
229    extend_zero_words(words, word_len);
230    Ok(())
231}
232
233/// Append zero bytes up to an exact byte length without disturbing the prefix,
234/// with an explicit allocator failure contract.
235#[inline]
236pub(crate) fn try_pad_zero_bytes_to_len(
237    bytes: &mut Vec<u8>,
238    byte_len: usize,
239    field: &'static str,
240) -> Result<(), String> {
241    crate::staging_reserve::reserve_vec(bytes, byte_len, field)?;
242    if bytes.len() < byte_len {
243        extend_zero_bytes(bytes, byte_len - bytes.len());
244    } else {
245        bytes.truncate(byte_len);
246    }
247    Ok(())
248}
249
250#[inline]
251fn extend_zero_bytes(bytes: &mut Vec<u8>, additional: usize) {
252    bytes.extend(std::iter::repeat_n(0, additional));
253}
254
255#[inline]
256fn extend_zero_words(words: &mut Vec<u32>, additional: usize) {
257    words.extend(std::iter::repeat_n(0, additional));
258}
259
260#[cold]
261fn pack_u32_cast_mismatch(
262    field: &'static str,
263    words: usize,
264    got_bytes: usize,
265    expected_bytes: usize,
266) -> String {
267    format!(
268        "{field} cast {words} u32 words into {got_bytes} bytes, expected {expected_bytes}. Fix: split the Weir GPU dispatch buffer or repair host slice width accounting."
269    )
270}
271
272#[cfg(test)]
273mod packing_tests {
274    use super::{
275        pack_exact_u32_slots_into, pack_repeated_u32_into, pack_u32, pack_u32_into,
276        try_pad_zero_bytes_to_len, try_write_zero_bytes, try_write_zero_words,
277    };
278
279    #[test]
280    fn pack_u32_preserves_little_endian_wire_order() {
281        assert_eq!(
282            pack_u32(&[0x0102_0304, 0xa0b0_c0d0]),
283            vec![0x04, 0x03, 0x02, 0x01, 0xd0, 0xc0, 0xb0, 0xa0]
284        );
285    }
286
287    #[test]
288    fn pack_u32_into_reuses_capacity_and_preserves_wire_order() {
289        let mut bytes = Vec::with_capacity(64);
290        let original_capacity = bytes.capacity();
291        pack_u32_into(&[0x1122_3344], &mut bytes);
292
293        assert_eq!(bytes, vec![0x44, 0x33, 0x22, 0x11]);
294        assert_eq!(
295            bytes.capacity(),
296            original_capacity,
297            "packing into caller-owned scratch must not shrink reusable capacity"
298        );
299    }
300
301    #[test]
302    fn pack_repeated_u32_into_preserves_wire_order_without_word_staging() {
303        let mut bytes = Vec::with_capacity(64);
304        let original_capacity = bytes.capacity();
305        pack_repeated_u32_into(0x1122_3344, 3, &mut bytes)
306            .expect("small repeated u32 pack must fit");
307
308        assert_eq!(
309            bytes,
310            vec![0x44, 0x33, 0x22, 0x11, 0x44, 0x33, 0x22, 0x11, 0x44, 0x33, 0x22, 0x11]
311        );
312        assert_eq!(
313            bytes.capacity(),
314            original_capacity,
315            "repeated packing into caller-owned scratch must not allocate a staging word buffer"
316        );
317    }
318
319    #[test]
320    fn pack_exact_u32_slots_into_reuses_capacity_and_rejects_wrong_width() {
321        let mut bytes = Vec::with_capacity(64);
322        let original_capacity = bytes.capacity();
323
324        pack_exact_u32_slots_into(&[0x1122_3344], 1, "test exact", &mut bytes)
325            .expect("matching semantic width should pack");
326
327        assert_eq!(bytes, vec![0x44, 0x33, 0x22, 0x11]);
328        assert_eq!(bytes.capacity(), original_capacity);
329        let err = pack_exact_u32_slots_into(&[1, 2], 1, "test exact", &mut bytes)
330            .expect_err("mismatched semantic width must be rejected");
331        assert!(
332            err.contains("expected exactly 1"),
333            "unexpected diagnostic: {err}"
334        );
335    }
336
337    #[test]
338    fn zero_staging_extends_after_reserve_without_resize_growth() {
339        let mut bytes = Vec::with_capacity(16);
340        try_write_zero_bytes(&mut bytes, 8, "test zero bytes").expect("zero bytes should stage");
341        assert_eq!(bytes, vec![0; 8]);
342        bytes[0] = 7;
343        try_pad_zero_bytes_to_len(&mut bytes, 12, "test zero pad").expect("zero pad should stage");
344        assert_eq!(bytes[0], 7);
345        assert_eq!(&bytes[8..], &[0, 0, 0, 0]);
346
347        let mut words = Vec::with_capacity(8);
348        try_write_zero_words(&mut words, 3, "test zero words").expect("zero words should stage");
349        assert_eq!(words, vec![0; 3]);
350    }
351
352    #[test]
353    fn dispatch_decode_source_has_no_resize_based_zero_staging() {
354        let source = include_str!("pack.rs");
355        let production = source
356            .split("#[cfg(test)]")
357            .next()
358            .expect("dispatch decode production source must precede tests");
359        assert!(
360            !production.contains(".resize(")
361                && production.contains("fn extend_zero_bytes(")
362                && production.contains("fn extend_zero_words("),
363            "Fix: dispatch decode zero staging must extend after fallible reserve instead of resize-driven growth."
364        );
365        assert!(
366            !production.contains("fn byte_len_for_existing_words")
367                && !production.contains("debug_assert_eq!(bytes.len(), byte_len)")
368                && !production.contains("debug_assert_eq!(packed.len(), byte_len)"),
369            "Fix: dispatch decode byte sizing and cast-width checks must run in release builds."
370        );
371    }
372}
373
374#[cfg(test)]
375mod panic_tests {
376    use super::{pack_u32_into, write_zero_words};
377
378    #[test]
379    #[should_panic(expected = "capacity overflow")]
380    fn pack_u32_into_panics_on_huge_input() {
381        // vec![] itself panics before pack_u32_into is reached; the
382        // important invariant is that huge inputs do not silently succeed.
383        let huge = vec![0u32; usize::MAX / 8 + 1];
384        let mut out = Vec::new();
385        pack_u32_into(&huge, &mut out);
386    }
387
388    #[test]
389    #[should_panic(expected = "zero word staging allocation failed in legacy infallible caller")]
390    fn write_zero_words_panics_on_huge_input() {
391        let mut out = Vec::new();
392        write_zero_words(&mut out, usize::MAX / 8 + 1);
393    }
394}