diff --git a/wacore/binary/benches/binary_benchmark.rs b/wacore/binary/benches/binary_benchmark.rs index 3367b91aa..1903a8d0d 100644 --- a/wacore/binary/benches/binary_benchmark.rs +++ b/wacore/binary/benches/binary_benchmark.rs @@ -22,6 +22,46 @@ fn create_small_node() -> Node { .build() } +/// The shape the wire is mostly made of: a short ack carrying a JID, a +/// nibble-packed number and a hex-packed id at once. `create_small_node` +/// already reaches `read_jid_pair` and the nibble half of `read_packed`, but +/// nothing reached the hex half, since the large fixture's lowercase `abcdef` +/// encodes as raw bytes rather than `HEX_8`. +fn create_ack_node() -> Node { + NodeBuilder::new("ack") + .attr("to", "5511999990000@s.whatsapp.net") + .attr("id", "3EB0A1B2C3D4E5F60718") + .attr("class", "message") + .build() +} + +/// A device fanout. Device-qualified JIDs encode as `AD_JID`, so what repeats +/// per child is the packed decode and the AD path, not `read_jid_pair`. +fn create_fanout_node() -> Node { + let devices: Vec = (0..8) + .map(|i| { + NodeBuilder::new("to") + .attr("jid", format!("5511999990000:{i}@s.whatsapp.net")) + .children(vec![ + NodeBuilder::new("enc") + .attr("v", "2") + .attr("type", "msg") + .bytes(vec![0xAB; 128]) + .build(), + ]) + .build() + }) + .collect(); + NodeBuilder::new("message") + .attr("to", "5511999990000@g.us") + .attr("id", "3EB0A1B2C3D4E5F60718") + .attr("type", "text") + .children(vec![ + NodeBuilder::new("participants").children(devices).build(), + ]) + .build() +} + fn create_large_node() -> Node { NodeBuilder::new("iq") .attr("to", "server@s.whatsapp.net") @@ -243,6 +283,32 @@ fn bench_unmarshal_small(bencher: divan::Bencher) { }); } +fn setup_ack_marshaled() -> Vec { + marshal(&create_ack_node()).unwrap() +} + +fn setup_fanout_marshaled() -> Vec { + marshal(&create_fanout_node()).unwrap() +} + +#[divan::bench] +fn bench_unmarshal_ack(bencher: divan::Bencher) { + bencher + .with_inputs(setup_ack_marshaled) + .bench_refs(|marshaled| { + black_box(unmarshal_ref(black_box(&marshaled[1..])).unwrap()); + }); +} + +#[divan::bench] +fn bench_unmarshal_fanout(bencher: divan::Bencher) { + bencher + .with_inputs(setup_fanout_marshaled) + .bench_refs(|marshaled| { + black_box(unmarshal_ref(black_box(&marshaled[1..])).unwrap()); + }); +} + #[divan::bench] fn bench_unmarshal_large(bencher: divan::Bencher) { bencher diff --git a/wacore/binary/src/decoder.rs b/wacore/binary/src/decoder.rs index 5cd447ff1..7b32189b7 100644 --- a/wacore/binary/src/decoder.rs +++ b/wacore/binary/src/decoder.rs @@ -15,6 +15,44 @@ fn jid_ref_to_compact(j: &JidRef<'_>) -> CompactString { s } +/// Each byte's two output characters, so unpacking is one load and one 2-byte +/// store per input byte instead of two shifts, two lookups and two bounds +/// checks. Packed values on the wire (a 13-digit phone number, a 20-character +/// id) are shorter than the SIMD chunk above, so this is the path that runs. +static HEX_PAIRS: [[u8; 2]; 256] = { + const HEX: [u8; 16] = *b"0123456789ABCDEF"; + let mut table = [[0u8; 2]; 256]; + let mut i = 0; + while i < 256 { + table[i] = [HEX[i >> 4], HEX[i & 0x0F]]; + i += 1; + } + table +}; + +/// Nibble values 12, 13 and 14 encode nothing, so they are marked and the byte +/// carrying one falls back to the scalar path, which reports which half was +/// bad. `NIBBLE_INVALID` cannot collide with an output character. +const NIBBLE_INVALID: u8 = 0xFF; +static NIBBLE_PAIRS: [[u8; 2]; 256] = { + const fn glyph(nibble: usize) -> u8 { + match nibble { + 0..=9 => b'0' + nibble as u8, + 10 => b'-', + 11 => b'.', + 15 => 0, + _ => NIBBLE_INVALID, + } + } + let mut table = [[0u8; 2]; 256]; + let mut i = 0; + while i < 256 { + table[i] = [glyph(i >> 4), glyph(i & 0x0F)]; + i += 1; + } + table +}; + /// Node-nesting cap rejecting deep-`LIST` frames that would overflow the stack via /// unbounded `read_node_ref` recursion (real WA trees are well under 20 levels). const MAX_NODE_DEPTH: usize = 128; @@ -357,14 +395,14 @@ impl<'a> Decoder<'a> { remainder }; - for &byte in packed_data { - let high = (byte & 0xF0) >> 4; - let low = byte & 0x0F; - out[*pos] = Self::unpack_hex(high); - *pos += 1; - out[*pos] = Self::unpack_hex(low); - *pos += 1; + let written = packed_data.len() * 2; + for (slot, &byte) in out[*pos..*pos + written] + .chunks_exact_mut(2) + .zip(packed_data) + { + slot.copy_from_slice(&HEX_PAIRS[byte as usize]); } + *pos += written; } #[inline] @@ -415,14 +453,20 @@ impl<'a> Decoder<'a> { remainder }; - for &byte in packed_data { - let high = (byte & 0xF0) >> 4; - let low = byte & 0x0F; - out[*pos] = Self::unpack_nibble(high)?; - *pos += 1; - out[*pos] = Self::unpack_nibble(low)?; - *pos += 1; + let written = packed_data.len() * 2; + for (slot, &byte) in out[*pos..*pos + written] + .chunks_exact_mut(2) + .zip(packed_data) + { + let pair = NIBBLE_PAIRS[byte as usize]; + if pair[0] == NIBBLE_INVALID || pair[1] == NIBBLE_INVALID { + // Report the bad half in the order the byte carries it. + Self::unpack_nibble((byte & 0xF0) >> 4)?; + Self::unpack_nibble(byte & 0x0F)?; + } + slot.copy_from_slice(&pair); } + *pos += written; Ok(()) } @@ -438,15 +482,6 @@ impl<'a> Decoder<'a> { } } - #[inline(always)] - fn unpack_hex(value: u8) -> u8 { - match value { - 0..=9 => b'0' + value, - 10..=15 => b'A' + value - 10, - _ => unreachable!("hex nibble validated by 4-bit mask"), - } - } - fn read_attributes(&mut self, size: usize) -> Result> { if size == 0 { return Ok(AttrsRef::Empty); diff --git a/wacore/binary/tests/packed_equivalence.rs b/wacore/binary/tests/packed_equivalence.rs new file mode 100644 index 000000000..69fb0767f --- /dev/null +++ b/wacore/binary/tests/packed_equivalence.rs @@ -0,0 +1,59 @@ +//! The packed-value tables must decode exactly what the scalar nibble/hex +//! walk did, including where an invalid nibble reports the failure. +use wacore_binary::builder::NodeBuilder; +use wacore_binary::marshal::{marshal, unmarshal_ref}; + +fn roundtrip_attr(value: &str) -> String { + let node = NodeBuilder::new("t").attr("v", value).build(); + let bytes = marshal(&node).unwrap(); + let decoded = unmarshal_ref(&bytes[1..]).unwrap(); + match decoded.get_attr("v").unwrap() { + wacore_binary::node::ValueRef::String(s) => s.to_string(), + other => panic!("unexpected value {other:?}"), + } +} + +#[test] +fn packed_values_round_trip_across_the_simd_boundary() { + // Hex and nibble, at lengths below, at, and above the 16-byte chunk the + // SIMD path handles, plus odd lengths that exercise the half-byte trim. + for len in [1, 2, 3, 15, 16, 17, 31, 32, 33, 63, 100, 127] { + let hex: String = (0..len) + .map(|i| b"0123456789ABCDEF"[i % 16] as char) + .collect(); + assert_eq!(roundtrip_attr(&hex), hex, "hex len {len}"); + + let nibble: String = (0..len).map(|i| b"0123456789-."[i % 12] as char).collect(); + assert_eq!(roundtrip_attr(&nibble), nibble, "nibble len {len}"); + } +} + +/// Nibble 12, 13 and 14 encode nothing. Frames are built by hand, since the +/// encoder would never emit one. Both lengths matter: one byte stays in the +/// scalar remainder, while 20 bytes puts the bad nibble inside the first +/// 16-byte chunk, which is where the SIMD build takes its validation +/// fallback. +#[test] +fn an_invalid_nibble_is_still_rejected() { + for bad in [0x0Cu8, 0x0D, 0x0E] { + // NIBBLE_8, length 1, one byte whose low half is invalid. + let short = [248u8, 2, 1, 255, 1, 0xF0 | bad]; + let err = unmarshal_ref(&short).unwrap_err(); + assert!( + format!("{err}").contains(&bad.to_string()), + "short frame must name the bad nibble {bad}: {err}" + ); + + // NIBBLE_8, length 20, with the bad nibble at byte 5 so it lands in + // the chunk rather than the remainder. + let mut long = vec![248u8, 2, 1, 255, 20]; + for i in 0..20u8 { + long.push(if i == 5 { 0xF0 | bad } else { 0x12 }); + } + let err = unmarshal_ref(&long).unwrap_err(); + assert!( + format!("{err}").contains(&bad.to_string()), + "chunked frame must name the bad nibble {bad}: {err}" + ); + } +}