Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 66 additions & 0 deletions wacore/binary/benches/binary_benchmark.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Node> = (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")
Expand Down Expand Up @@ -243,6 +283,32 @@ fn bench_unmarshal_small(bencher: divan::Bencher) {
});
}

fn setup_ack_marshaled() -> Vec<u8> {
marshal(&create_ack_node()).unwrap()
Comment thread
jlucaso1 marked this conversation as resolved.
}

fn setup_fanout_marshaled() -> Vec<u8> {
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
Expand Down
81 changes: 58 additions & 23 deletions wacore/binary/src/decoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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(())
}
Expand All @@ -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<AttrsRef<'a>> {
if size == 0 {
return Ok(AttrsRef::Empty);
Expand Down
59 changes: 59 additions & 0 deletions wacore/binary/tests/packed_equivalence.rs
Original file line number Diff line number Diff line change
@@ -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}"
);
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.