diff --git a/wacore/binary/benches/binary_benchmark.rs b/wacore/binary/benches/binary_benchmark.rs index f26ac5c40..d07d26ca5 100644 --- a/wacore/binary/benches/binary_benchmark.rs +++ b/wacore/binary/benches/binary_benchmark.rs @@ -7,7 +7,10 @@ use iai_callgrind::{ use std::hint::black_box; use std::io::Write; use wacore_binary::builder::NodeBuilder; -use wacore_binary::marshal::{marshal, marshal_ref, marshal_to, unmarshal_ref}; +use wacore_binary::marshal::{ + marshal, marshal_auto, marshal_exact, marshal_ref, marshal_ref_auto, marshal_ref_exact, + marshal_to, marshal_to_vec, unmarshal_ref, +}; use wacore_binary::node::Node; use wacore_binary::util::unpack; @@ -60,7 +63,7 @@ fn create_attr_node() -> Node { } // Creates a node with long string content to test the JID parsing optimization. -// Long strings (> 256 chars) should skip JID parsing for better performance. +// Long strings (> 48 chars) should skip JID parsing for better performance. fn create_long_string_node() -> Node { // Generate a 500+ character string that contains '@' but is NOT a valid JID. // Without the optimization, parse_jid would scan the entire string. @@ -124,6 +127,28 @@ fn create_jid_heavy_node() -> Node { .build() } +fn create_huge_bytes_node() -> Node { + NodeBuilder::new("message") + .attr("to", "server@s.whatsapp.net") + .attr("id", "huge-binary") + .bytes(vec![0x5A; 256 * 1024]) + .build() +} + +fn create_many_children_node() -> Node { + NodeBuilder::new("iq") + .attr("to", "server@s.whatsapp.net") + .attr("id", "many-children") + .children((0..2048).map(|i| { + NodeBuilder::new("item") + .attr("index", i.to_string()) + .attr("type", "entry") + .string_content("ok") + .build() + })) + .build() +} + // Marshal benchmarks - self-contained, no setup needed #[library_benchmark] fn bench_marshal_allocating() -> Vec { @@ -131,6 +156,18 @@ fn bench_marshal_allocating() -> Vec { black_box(marshal(black_box(&node)).unwrap()) } +#[library_benchmark] +fn bench_marshal_auto_allocating() -> Vec { + let node = create_large_node(); + black_box(marshal_auto(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_exact_allocating() -> Vec { + let node = create_large_node(); + black_box(marshal_exact(black_box(&node)).unwrap()) +} + #[library_benchmark] fn bench_marshal_reusing_buffer() -> Vec { let node = create_large_node(); @@ -139,6 +176,14 @@ fn bench_marshal_reusing_buffer() -> Vec { black_box(buffer) } +#[library_benchmark] +fn bench_marshal_reusing_buffer_vec_writer() -> Vec { + let node = create_large_node(); + let mut buffer = Vec::with_capacity(4096); + marshal_to_vec(black_box(&node), &mut buffer).unwrap(); + black_box(buffer) +} + // Benchmark for marshaling nodes with long string content. // This demonstrates the JID parsing optimization: long strings skip parse_jid. #[library_benchmark] @@ -147,6 +192,54 @@ fn bench_marshal_long_string() -> Vec { black_box(marshal(black_box(&node)).unwrap()) } +#[library_benchmark] +fn bench_marshal_auto_long_string() -> Vec { + let node = create_long_string_node(); + black_box(marshal_auto(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_exact_long_string() -> Vec { + let node = create_long_string_node(); + black_box(marshal_exact(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_huge_bytes_allocating() -> Vec { + let node = create_huge_bytes_node(); + black_box(marshal(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_auto_huge_bytes_allocating() -> Vec { + let node = create_huge_bytes_node(); + black_box(marshal_auto(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_exact_huge_bytes_allocating() -> Vec { + let node = create_huge_bytes_node(); + black_box(marshal_exact(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_many_children_allocating() -> Vec { + let node = create_many_children_node(); + black_box(marshal(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_auto_many_children_allocating() -> Vec { + let node = create_many_children_node(); + black_box(marshal_auto(black_box(&node)).unwrap()) +} + +#[library_benchmark] +fn bench_marshal_exact_many_children_allocating() -> Vec { + let node = create_many_children_node(); + black_box(marshal_exact(black_box(&node)).unwrap()) +} + // Setup functions for unmarshal benchmarks - pre-compute marshaled data // Note: marshal() adds a flag byte at position 0, unmarshal_ref expects data without it fn setup_small_marshaled() -> Vec { @@ -215,6 +308,24 @@ fn bench_roundtrip(marshaled: Vec) -> Vec { black_box(marshal_ref(&node_ref).unwrap()) } +#[library_benchmark] +#[bench::small(setup = setup_small_marshaled)] +#[bench::large(setup = setup_large_marshaled)] +fn bench_roundtrip_auto(marshaled: Vec) -> Vec { + // Skip the flag byte at position 0 + let node_ref = unmarshal_ref(black_box(&marshaled[1..])).unwrap(); + black_box(marshal_ref_auto(&node_ref).unwrap()) +} + +#[library_benchmark] +#[bench::small(setup = setup_small_marshaled)] +#[bench::large(setup = setup_large_marshaled)] +fn bench_roundtrip_exact(marshaled: Vec) -> Vec { + // Skip the flag byte at position 0 + let node_ref = unmarshal_ref(black_box(&marshaled[1..])).unwrap(); + black_box(marshal_ref_exact(&node_ref).unwrap()) +} + // Child iteration benchmark: tests get_children_by_tag performance // Simulates the recursive traversal pattern used in usync parsing #[library_benchmark] @@ -242,7 +353,21 @@ fn bench_get_children_by_tag() { library_benchmark_group!( name = marshal_group; - benchmarks = bench_marshal_allocating, bench_marshal_reusing_buffer, bench_marshal_long_string + benchmarks = + bench_marshal_allocating, + bench_marshal_auto_allocating, + bench_marshal_exact_allocating, + bench_marshal_reusing_buffer, + bench_marshal_reusing_buffer_vec_writer, + bench_marshal_long_string, + bench_marshal_auto_long_string, + bench_marshal_exact_long_string, + bench_marshal_huge_bytes_allocating, + bench_marshal_auto_huge_bytes_allocating, + bench_marshal_exact_huge_bytes_allocating, + bench_marshal_many_children_allocating, + bench_marshal_auto_many_children_allocating, + bench_marshal_exact_many_children_allocating ); library_benchmark_group!( @@ -262,7 +387,7 @@ library_benchmark_group!( library_benchmark_group!( name = roundtrip_group; - benchmarks = bench_roundtrip + benchmarks = bench_roundtrip, bench_roundtrip_auto, bench_roundtrip_exact ); library_benchmark_group!( diff --git a/wacore/binary/src/decoder.rs b/wacore/binary/src/decoder.rs index 9908c5f81..e80eeae47 100644 --- a/wacore/binary/src/decoder.rs +++ b/wacore/binary/src/decoder.rs @@ -23,6 +23,7 @@ impl<'a> Decoder<'a> { self.data.len() - self.position } + #[inline(always)] fn check_eos(&self, len: usize) -> Result<()> { if self.bytes_left() >= len { Ok(()) @@ -31,50 +32,76 @@ impl<'a> Decoder<'a> { } } + #[inline(always)] fn read_u8(&mut self) -> Result { self.check_eos(1)?; - let value = self.data[self.position]; + let position = self.position; self.position += 1; + // SAFETY: `check_eos(1)` guarantees that `position` is a valid index. + let value = unsafe { *self.data.get_unchecked(position) }; Ok(value) } + #[inline(always)] fn read_u16_be(&mut self) -> Result { self.check_eos(2)?; - let value = u16::from_be_bytes([self.data[self.position], self.data[self.position + 1]]); + let position = self.position; self.position += 2; + // SAFETY: `check_eos(2)` guarantees both indexes are in bounds. + let value = unsafe { + u16::from_be_bytes([ + *self.data.get_unchecked(position), + *self.data.get_unchecked(position + 1), + ]) + }; Ok(value) } + #[inline(always)] fn read_u20_be(&mut self) -> Result { self.check_eos(3)?; - let bytes = [ - self.data[self.position], - self.data[self.position + 1], - self.data[self.position + 2], - ]; + let position = self.position; self.position += 3; + // SAFETY: `check_eos(3)` guarantees all indexes are in bounds. + let bytes = unsafe { + [ + *self.data.get_unchecked(position), + *self.data.get_unchecked(position + 1), + *self.data.get_unchecked(position + 2), + ] + }; Ok(((bytes[0] as u32 & 0x0F) << 16) | ((bytes[1] as u32) << 8) | (bytes[2] as u32)) } + #[inline(always)] fn read_u32_be(&mut self) -> Result { self.check_eos(4)?; - let value = u32::from_be_bytes([ - self.data[self.position], - self.data[self.position + 1], - self.data[self.position + 2], - self.data[self.position + 3], - ]); + let position = self.position; self.position += 4; + // SAFETY: `check_eos(4)` guarantees all indexes are in bounds. + let value = unsafe { + u32::from_be_bytes([ + *self.data.get_unchecked(position), + *self.data.get_unchecked(position + 1), + *self.data.get_unchecked(position + 2), + *self.data.get_unchecked(position + 3), + ]) + }; Ok(value) } + #[inline(always)] fn read_bytes(&mut self, len: usize) -> Result<&'a [u8]> { self.check_eos(len)?; - let slice = &self.data[self.position..self.position + len]; - self.position += len; + let start = self.position; + let end = start + len; + self.position = end; + // SAFETY: `check_eos(len)` guarantees `[start..end]` is in bounds. + let slice = unsafe { self.data.get_unchecked(start..end) }; Ok(slice) } + #[inline(always)] fn read_string(&mut self, len: usize) -> Result> { let bytes = self.read_bytes(len)?; match std::str::from_utf8(bytes) { @@ -83,11 +110,12 @@ impl<'a> Decoder<'a> { } } + #[inline(always)] fn read_list_size(&mut self, tag: u8) -> Result { match tag { token::LIST_EMPTY => Ok(0), - 248 => self.read_u8().map(|v| v as usize), - 249 => self.read_u16_be().map(|v| v as usize), + token::LIST_8 => self.read_u8().map(|v| v as usize), + token::LIST_16 => self.read_u16_be().map(|v| v as usize), _ => Err(BinaryError::InvalidToken(tag)), } } @@ -166,6 +194,11 @@ impl<'a> Decoder<'a> { fn read_value_as_string(&mut self) -> Result>> { let tag = self.read_u8()?; + self.read_value_as_string_from_tag(tag) + } + + #[inline(always)] + fn read_value_as_string_from_tag(&mut self, tag: u8) -> Result>> { match tag { token::LIST_EMPTY => Ok(None), token::BINARY_8 => { @@ -252,14 +285,56 @@ impl<'a> Decoder<'a> { let packed_data = self.read_bytes(len)?; let mut unpacked_bytes = Vec::with_capacity(raw_len); - const NIBBLE_LOOKUP: [u8; 16] = *b"0123456789-.\x00\x00\x00\x00"; + match tag { + token::HEX_8 => Self::decode_packed_hex(packed_data, &mut unpacked_bytes), + token::NIBBLE_8 => Self::decode_packed_nibble(packed_data, &mut unpacked_bytes)?, + _ => return Err(BinaryError::InvalidToken(tag)), + } + + if is_half_byte { + unpacked_bytes.pop(); + } + + // SAFETY: unpacked bytes are built exclusively from protocol lookup tables: + // - HEX_8 produces ASCII '0'..'9' and 'A'..'F' + // - NIBBLE_8 produces ASCII '0'..'9', '-', '.', or '\0' padding + // All generated bytes are valid UTF-8 scalar values. + Ok(unsafe { String::from_utf8_unchecked(unpacked_bytes) }) + } + + #[inline] + fn decode_packed_hex(packed_data: &[u8], unpacked_bytes: &mut Vec) { const HEX_LOOKUP: [u8; 16] = *b"0123456789ABCDEF"; - let lookup_table = Simd::from_array(if tag == token::NIBBLE_8 { - NIBBLE_LOOKUP - } else { - HEX_LOOKUP - }); + let lookup_table = Simd::from_array(HEX_LOOKUP); + let low_mask = Simd::splat(0x0F); + + let (chunks, remainder) = packed_data.as_chunks::<16>(); + for chunk in chunks { + let data = u8x16::from_array(*chunk); + let high_nibbles = (data >> 4) & low_mask; + let low_nibbles = data & low_mask; + let high_chars = lookup_table.swizzle_dyn(high_nibbles); + let low_chars = lookup_table.swizzle_dyn(low_nibbles); + let (lo, hi) = Simd::interleave(high_chars, low_chars); + unpacked_bytes.extend_from_slice(lo.as_array()); + unpacked_bytes.extend_from_slice(hi.as_array()); + } + + for &byte in remainder { + let high = (byte & 0xF0) >> 4; + let low = byte & 0x0F; + unpacked_bytes.push(Self::unpack_hex(high)); + unpacked_bytes.push(Self::unpack_hex(low)); + } + } + + #[inline] + fn decode_packed_nibble(packed_data: &[u8], unpacked_bytes: &mut Vec) -> Result<()> { + const NIBBLE_LOOKUP: [u8; 16] = *b"0123456789-.\x00\x00\x00\x00"; + let lookup_table = Simd::from_array(NIBBLE_LOOKUP); let low_mask = Simd::splat(0x0F); + let le11 = Simd::splat(11); + let f15 = Simd::splat(15); let (chunks, remainder) = packed_data.as_chunks::<16>(); for chunk in chunks { @@ -268,25 +343,27 @@ impl<'a> Decoder<'a> { let high_nibbles = (data >> 4) & low_mask; let low_nibbles = data & low_mask; - if tag == token::NIBBLE_8 { - let le11 = Simd::splat(11); - let f15 = Simd::splat(15); - let hi_valid = high_nibbles.simd_le(le11) | high_nibbles.simd_eq(f15); - let lo_valid = low_nibbles.simd_le(le11) | low_nibbles.simd_eq(f15); - if !(hi_valid & lo_valid).all() { - for byte in *chunk { - let high = (byte & 0xF0) >> 4; - let low = byte & 0x0F; - Self::unpack_byte(tag, high)?; - Self::unpack_byte(tag, low)?; - } - unreachable!("SIMD validation should match scalar validation"); + let hi_valid = high_nibbles.simd_le(le11) | high_nibbles.simd_eq(f15); + let lo_valid = low_nibbles.simd_le(le11) | low_nibbles.simd_eq(f15); + if !(hi_valid & lo_valid).all() { + // Validate first, then decode scalar as a conservative fallback. + for byte in *chunk { + let high = (byte & 0xF0) >> 4; + let low = byte & 0x0F; + Self::unpack_nibble(high)?; + Self::unpack_nibble(low)?; + } + for byte in *chunk { + let high = (byte & 0xF0) >> 4; + let low = byte & 0x0F; + unpacked_bytes.push(Self::unpack_nibble(high)?); + unpacked_bytes.push(Self::unpack_nibble(low)?); } + continue; } let high_chars = lookup_table.swizzle_dyn(high_nibbles); let low_chars = lookup_table.swizzle_dyn(low_nibbles); - let (lo, hi) = Simd::interleave(high_chars, low_chars); unpacked_bytes.extend_from_slice(lo.as_array()); unpacked_bytes.extend_from_slice(hi.as_array()); @@ -295,32 +372,30 @@ impl<'a> Decoder<'a> { for &byte in remainder { let high = (byte & 0xF0) >> 4; let low = byte & 0x0F; - unpacked_bytes.push(Self::unpack_byte(tag, high)? as u8); - unpacked_bytes.push(Self::unpack_byte(tag, low)? as u8); + unpacked_bytes.push(Self::unpack_nibble(high)?); + unpacked_bytes.push(Self::unpack_nibble(low)?); } - if is_half_byte { - unpacked_bytes.pop(); - } + Ok(()) + } - String::from_utf8(unpacked_bytes).map_err(|e| BinaryError::InvalidUtf8(e.utf8_error())) + #[inline(always)] + fn unpack_nibble(value: u8) -> Result { + match value { + 0..=9 => Ok(b'0' + value), + 10 => Ok(b'-'), + 11 => Ok(b'.'), + 15 => Ok(0), + _ => Err(BinaryError::InvalidToken(value)), + } } - fn unpack_byte(tag: u8, value: u8) -> Result { - match tag { - token::NIBBLE_8 => match value { - 0..=9 => Ok((b'0' + value) as char), - 10 => Ok('-'), - 11 => Ok('.'), - 15 => Ok('\x00'), - _ => Err(BinaryError::InvalidToken(value)), - }, - token::HEX_8 => match value { - 0..=9 => Ok((b'0' + value) as char), - 10..=15 => Ok((b'A' + value - 10) as char), - _ => Err(BinaryError::InvalidToken(value)), - }, - _ => Err(BinaryError::InvalidToken(tag)), + #[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"), } } @@ -341,6 +416,11 @@ impl<'a> Decoder<'a> { fn read_content(&mut self) -> Result>> { let tag = self.read_u8()?; + self.read_content_from_tag(tag) + } + + #[inline(always)] + fn read_content_from_tag(&mut self, tag: u8) -> Result>> { match tag { token::LIST_EMPTY => Ok(None), @@ -370,8 +450,7 @@ impl<'a> Decoder<'a> { } _ => { - self.position -= 1; - let string_content = self.read_value_as_string()?; + let string_content = self.read_value_as_string_from_tag(tag)?; match string_content { Some(s) => Ok(Some(NodeContentRef::String(s))), diff --git a/wacore/binary/src/encoder.rs b/wacore/binary/src/encoder.rs index 2a5d7ea59..b1a7ccf12 100644 --- a/wacore/binary/src/encoder.rs +++ b/wacore/binary/src/encoder.rs @@ -4,11 +4,106 @@ use core::simd::Select; use core::simd::prelude::*; use core::simd::{Simd, u8x16}; -use crate::error::Result; +use crate::error::{BinaryError, Result}; use crate::jid::{self, Jid, JidRef}; use crate::node::{Node, NodeContent, NodeContentRef, NodeRef, NodeValue, ValueRef}; use crate::token; +pub(crate) trait ByteWriter { + fn write_u8(&mut self, value: u8) -> Result<()>; + fn write_bytes(&mut self, bytes: &[u8]) -> Result<()>; +} + +pub(crate) struct IoByteWriter { + writer: W, +} + +impl IoByteWriter { + fn new(writer: W) -> Self { + Self { writer } + } +} + +impl ByteWriter for IoByteWriter { + #[inline] + fn write_u8(&mut self, value: u8) -> Result<()> { + self.writer.write_all(&[value])?; + Ok(()) + } + + #[inline] + fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> { + self.writer.write_all(bytes)?; + Ok(()) + } +} + +pub(crate) struct VecByteWriter<'a> { + buffer: &'a mut Vec, +} + +impl<'a> VecByteWriter<'a> { + fn new(buffer: &'a mut Vec) -> Self { + Self { buffer } + } +} + +impl ByteWriter for VecByteWriter<'_> { + #[inline] + fn write_u8(&mut self, value: u8) -> Result<()> { + self.buffer.push(value); + Ok(()) + } + + #[inline] + fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> { + self.buffer.extend_from_slice(bytes); + Ok(()) + } +} + +pub(crate) struct SliceByteWriter<'a> { + buffer: &'a mut [u8], + position: usize, +} + +impl<'a> SliceByteWriter<'a> { + fn new(buffer: &'a mut [u8]) -> Self { + Self { + buffer, + position: 0, + } + } + + #[inline] + fn bytes_written(&self) -> usize { + self.position + } +} + +impl ByteWriter for SliceByteWriter<'_> { + #[inline] + fn write_u8(&mut self, value: u8) -> Result<()> { + if self.position >= self.buffer.len() { + return Err(BinaryError::UnexpectedEof); + } + self.buffer[self.position] = value; + self.position += 1; + Ok(()) + } + + #[inline] + fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> { + let end = self.position + bytes.len(); + if end > self.buffer.len() { + return Err(BinaryError::UnexpectedEof); + } + self.buffer[self.position..end].copy_from_slice(bytes); + self.position = end; + Ok(()) + } +} + /// Trait for encoding node structures (both owned Node and borrowed NodeRef). /// All encoding logic lives in the trait implementation, keeping /// the Encoder simple and focused on low-level byte writing. @@ -18,10 +113,10 @@ pub(crate) trait EncodeNode { fn has_content(&self) -> bool; /// Encode all attributes to the encoder - fn encode_attrs(&self, encoder: &mut Encoder) -> Result<()>; + fn encode_attrs<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()>; /// Encode content (string, bytes, or child nodes) to the encoder - fn encode_content(&self, encoder: &mut Encoder) -> Result<()>; + fn encode_content<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()>; } impl EncodeNode for Node { @@ -37,7 +132,7 @@ impl EncodeNode for Node { self.content.is_some() } - fn encode_attrs(&self, encoder: &mut Encoder) -> Result<()> { + fn encode_attrs<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()> { for (k, v) in &self.attrs { encoder.write_string(k)?; match v { @@ -48,7 +143,7 @@ impl EncodeNode for Node { Ok(()) } - fn encode_content(&self, encoder: &mut Encoder) -> Result<()> { + fn encode_content<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()> { if let Some(content) = &self.content { match content { NodeContent::String(s) => encoder.write_string(s)?, @@ -78,7 +173,7 @@ impl EncodeNode for NodeRef<'_> { self.content.is_some() } - fn encode_attrs(&self, encoder: &mut Encoder) -> Result<()> { + fn encode_attrs<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()> { for (k, v) in &self.attrs { encoder.write_string(k)?; match v { @@ -89,7 +184,7 @@ impl EncodeNode for NodeRef<'_> { Ok(()) } - fn encode_content(&self, encoder: &mut Encoder) -> Result<()> { + fn encode_content<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()> { if let Some(content) = self.content.as_deref() { match content { NodeContentRef::String(s) => encoder.write_string(s)?, @@ -106,36 +201,121 @@ impl EncodeNode for NodeRef<'_> { } } -struct ParsedJid<'a> { - user: &'a str, - server: &'a str, +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ParsedJidMeta { + user_end: usize, + server_start: usize, domain_type: u8, - device: Option, + device: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct StrKey { + ptr: usize, + len: usize, +} + +impl StrKey { + #[inline] + fn from_str(s: &str) -> Self { + Self { + ptr: s.as_ptr() as usize, + len: s.len(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StringHint { + Empty, + SingleToken(u8), + DoubleToken { dict: u8, token: u8 }, + PackedNibble, + PackedHex, + Jid(ParsedJidMeta), + RawBytes, +} + +#[derive(Debug)] +pub(crate) struct StringHintCache { + // Keys use (ptr, len) identity, so this cache is only valid while encoding + // the same immutable node/strings it was built from. + hints: Vec<(StrKey, StringHint)>, +} + +impl Default for StringHintCache { + fn default() -> Self { + Self { + hints: Vec::with_capacity(32), + } + } } -fn parse_jid(input: &str) -> Option> { +impl StringHintCache { + const MAX_HINT_ENTRIES: usize = 96; + + #[inline] + fn hint_for(&self, s: &str) -> Option { + let key = StrKey::from_str(s); + self.hints + .iter() + .find_map(|(cached_key, hint)| (*cached_key == key).then_some(*hint)) + } + + #[inline] + fn hint_or_insert(&mut self, s: &str) -> StringHint { + let key = StrKey::from_str(s); + if let Some(existing) = self + .hints + .iter() + .find_map(|(cached_key, hint)| (*cached_key == key).then_some(*hint)) + { + existing + } else { + let hint = classify_string_hint(s); + if self.hints.len() < Self::MAX_HINT_ENTRIES { + self.hints.push((key, hint)); + } + hint + } + } +} + +#[derive(Debug)] +pub(crate) struct MarshaledSizePlan { + pub(crate) size: usize, + pub(crate) hints: StringHintCache, +} + +fn parse_jid_meta(input: &str) -> Option { let sep_idx = input.find('@')?; - let server = &input[sep_idx + 1..]; + let server_start = sep_idx + 1; + let server = &input[server_start..]; let user_combined = &input[..sep_idx]; - let (user_agent, device) = match user_combined.split_once(':') { - Some((ua, device_part)) => { - let parsed_device = if device_part.is_empty() { - None - } else { - device_part.parse::().ok() - }; - (ua, parsed_device) + let (user_agent, device) = if let Some(colon_idx) = user_combined.find(':') { + let device_part = &user_combined[colon_idx + 1..]; + if let Ok(parsed_device) = device_part.parse::() { + (&user_combined[..colon_idx], Some(parsed_device)) + } else { + (user_combined, None) } - None => (user_combined, None), + } else { + (user_combined, None) }; - let (user, agent_override) = match user_agent.split_once('_') { - Some((u, agent_part)) => (u, agent_part.parse::().ok()), - None => (user_agent, None), + let (user_end, agent_override) = if let Some(underscore_idx) = user_agent.find('_') { + let agent_part = &user_agent[underscore_idx + 1..]; + if let Ok(parsed_agent) = agent_part.parse::() { + (underscore_idx, Some(parsed_agent)) + } else { + (user_agent.len(), None) + } + } else { + (user_agent.len(), None) }; - let agent_byte = agent_override.unwrap_or(0) as u8; + let agent_byte = agent_override.unwrap_or(0); let domain_type = if server == jid::HIDDEN_USER_SERVER { 1 } else if server == jid::HOSTED_SERVER { @@ -146,55 +326,325 @@ fn parse_jid(input: &str) -> Option> { agent_byte }; - Some(ParsedJid { - user, - server, + Some(ParsedJidMeta { + user_end, + server_start, domain_type, device, }) } -pub(crate) struct Encoder { +#[inline] +fn split_jid_from_meta(input: &str, meta: ParsedJidMeta) -> (&str, &str) { + (&input[..meta.user_end], &input[meta.server_start..]) +} + +#[inline] +fn classify_string_hint(s: &str) -> StringHint { + if s.is_empty() { + return StringHint::Empty; + } + + let is_likely_jid = s.len() <= 48; + + if let Some(token) = token::index_of_single_token(s) { + StringHint::SingleToken(token) + } else if let Some((dict, token)) = token::index_of_double_byte_token(s) { + StringHint::DoubleToken { dict, token } + } else if validate_nibble(s) { + StringHint::PackedNibble + } else if validate_hex(s) { + StringHint::PackedHex + } else if is_likely_jid { + parse_jid_meta(s).map_or(StringHint::RawBytes, StringHint::Jid) + } else { + StringHint::RawBytes + } +} + +pub(crate) fn build_marshaled_node_plan(node: &Node) -> MarshaledSizePlan { + let mut hints = StringHintCache::default(); + let size = 1 + node_encoded_size_with_cache(node, &mut hints); + MarshaledSizePlan { size, hints } +} + +pub(crate) fn build_marshaled_node_ref_plan(node: &NodeRef<'_>) -> MarshaledSizePlan { + let mut hints = StringHintCache::default(); + let size = 1 + node_ref_encoded_size_with_cache(node, &mut hints); + MarshaledSizePlan { size, hints } +} + +#[inline] +fn list_start_encoded_size(len: usize) -> usize { + if len == 0 { + 1 + } else if len < 256 { + 2 + } else { + 3 + } +} + +#[inline] +fn binary_len_prefix_size(len: usize) -> usize { + if len < 256 { + 2 + } else if len < (1 << 20) { + 4 + } else { + 5 + } +} + +#[inline] +fn bytes_with_len_encoded_size(len: usize) -> usize { + binary_len_prefix_size(len) + len +} + +#[inline] +fn packed_encoded_size(value_len: usize) -> usize { + 2 + value_len.div_ceil(2) +} + +fn node_encoded_size_with_cache(node: &Node, hints: &mut StringHintCache) -> usize { + let content_len = usize::from(node.content.is_some()); + let list_len = 1 + (node.attrs.len() * 2) + content_len; + + let attrs_size: usize = node + .attrs + .iter() + .map(|(k, v)| { + let value_size = match v { + NodeValue::String(s) => string_encoded_size_with_cache(s, hints), + NodeValue::Jid(jid) => owned_jid_encoded_size_with_cache(jid, hints), + }; + string_encoded_size_with_cache(k, hints) + value_size + }) + .sum(); + + let content_size = match &node.content { + Some(NodeContent::String(s)) => string_encoded_size_with_cache(s, hints), + Some(NodeContent::Bytes(b)) => bytes_with_len_encoded_size(b.len()), + Some(NodeContent::Nodes(nodes)) => { + list_start_encoded_size(nodes.len()) + + nodes + .iter() + .map(|child| node_encoded_size_with_cache(child, hints)) + .sum::() + } + None => 0, + }; + + list_start_encoded_size(list_len) + + string_encoded_size_with_cache(node.tag.as_str(), hints) + + attrs_size + + content_size +} + +fn node_ref_encoded_size_with_cache(node: &NodeRef<'_>, hints: &mut StringHintCache) -> usize { + let content_len = usize::from(node.content.is_some()); + let list_len = 1 + (node.attrs.len() * 2) + content_len; + + let attrs_size: usize = node + .attrs + .iter() + .map(|(k, v)| { + let value_size = match v { + ValueRef::String(s) => string_encoded_size_with_cache(s, hints), + ValueRef::Jid(jid) => jid_ref_encoded_size_with_cache(jid, hints), + }; + string_encoded_size_with_cache(k, hints) + value_size + }) + .sum(); + + let content_size = match node.content.as_deref() { + Some(NodeContentRef::String(s)) => string_encoded_size_with_cache(s, hints), + Some(NodeContentRef::Bytes(b)) => bytes_with_len_encoded_size(b.len()), + Some(NodeContentRef::Nodes(nodes)) => { + list_start_encoded_size(nodes.len()) + + nodes + .iter() + .map(|child| node_ref_encoded_size_with_cache(child, hints)) + .sum::() + } + None => 0, + }; + + list_start_encoded_size(list_len) + + string_encoded_size_with_cache(node.tag.as_ref(), hints) + + attrs_size + + content_size +} + +#[inline] +fn string_encoded_size_with_cache(s: &str, hints: &mut StringHintCache) -> usize { + let hint = hints.hint_or_insert(s); + string_encoded_size_from_hint_with_cache(s, hint, hints) +} + +#[inline] +fn string_encoded_size_from_hint_with_cache( + s: &str, + hint: StringHint, + hints: &mut StringHintCache, +) -> usize { + match hint { + StringHint::Empty => 2, + StringHint::SingleToken(_) => 1, + StringHint::DoubleToken { .. } => 2, + StringHint::PackedNibble | StringHint::PackedHex => packed_encoded_size(s.len()), + StringHint::RawBytes => bytes_with_len_encoded_size(s.len()), + StringHint::Jid(meta) => parsed_jid_encoded_size_with_cache(s, meta, hints), + } +} + +#[inline] +fn parsed_jid_encoded_size_with_cache( + jid: &str, + meta: ParsedJidMeta, + hints: &mut StringHintCache, +) -> usize { + let (user, server) = split_jid_from_meta(jid, meta); + if meta.device.is_some() { + 3 + string_encoded_size_with_cache(user, hints) + } else { + let user_size = if user.is_empty() { + 1 + } else { + string_encoded_size_with_cache(user, hints) + }; + 1 + user_size + string_encoded_size_with_cache(server, hints) + } +} + +#[inline] +fn owned_jid_encoded_size_with_cache(jid: &Jid, hints: &mut StringHintCache) -> usize { + if jid.device > 0 { + 3 + string_encoded_size_with_cache(&jid.user, hints) + } else { + let user_size = if jid.user.is_empty() { + 1 + } else { + string_encoded_size_with_cache(&jid.user, hints) + }; + 1 + user_size + string_encoded_size_with_cache(&jid.server, hints) + } +} + +#[inline] +fn jid_ref_encoded_size_with_cache(jid: &JidRef<'_>, hints: &mut StringHintCache) -> usize { + if jid.device > 0 { + 3 + string_encoded_size_with_cache(&jid.user, hints) + } else { + let user_size = if jid.user.is_empty() { + 1 + } else { + string_encoded_size_with_cache(&jid.user, hints) + }; + 1 + user_size + string_encoded_size_with_cache(&jid.server, hints) + } +} + +#[inline] +fn validate_nibble(value: &str) -> bool { + if value.len() > token::PACKED_MAX as usize { + return false; + } + value + .as_bytes() + .iter() + .all(|&b| b.is_ascii_digit() || b == b'-' || b == b'.') +} + +#[inline] +fn validate_hex(value: &str) -> bool { + if value.len() > token::PACKED_MAX as usize { + return false; + } + value + .as_bytes() + .iter() + .all(|&b| b.is_ascii_digit() || (b'A'..=b'F').contains(&b)) +} + +pub(crate) struct Encoder<'a, W: ByteWriter> { writer: W, + string_hints: Option<&'a StringHintCache>, } -impl Encoder { +impl Encoder<'static, IoByteWriter> { pub(crate) fn new(writer: W) -> Result { - let mut enc = Self { writer }; + let mut enc = Self { + writer: IoByteWriter::new(writer), + string_hints: None, + }; + enc.write_u8(0)?; + Ok(enc) + } +} + +impl<'v> Encoder<'static, VecByteWriter<'v>> { + pub(crate) fn new_vec(buffer: &'v mut Vec) -> Result { + let mut enc = Self { + writer: VecByteWriter::new(buffer), + string_hints: None, + }; enc.write_u8(0)?; Ok(enc) } +} +impl<'a> Encoder<'a, SliceByteWriter<'a>> { + pub(crate) fn new_slice( + buffer: &'a mut [u8], + string_hints: Option<&'a StringHintCache>, + ) -> Result { + let mut enc = Self { + writer: SliceByteWriter::new(buffer), + string_hints, + }; + enc.write_u8(0)?; + Ok(enc) + } + + #[inline] + pub(crate) fn bytes_written(&self) -> usize { + self.writer.bytes_written() + } +} + +impl<'a, W: ByteWriter> Encoder<'a, W> { + #[inline(always)] fn write_u8(&mut self, val: u8) -> Result<()> { - self.writer.write_all(&[val])?; - Ok(()) + self.writer.write_u8(val) } + #[inline(always)] fn write_u16_be(&mut self, val: u16) -> Result<()> { - self.writer.write_all(&val.to_be_bytes())?; - Ok(()) + self.writer.write_bytes(&val.to_be_bytes()) } + #[inline(always)] fn write_u32_be(&mut self, val: u32) -> Result<()> { - self.writer.write_all(&val.to_be_bytes())?; - Ok(()) + self.writer.write_bytes(&val.to_be_bytes()) } + #[inline(always)] fn write_u20_be(&mut self, value: u32) -> Result<()> { let bytes = [ ((value >> 16) & 0x0F) as u8, ((value >> 8) & 0xFF) as u8, (value & 0xFF) as u8, ]; - self.writer.write_all(&bytes)?; - Ok(()) + self.writer.write_bytes(&bytes) } + #[inline(always)] fn write_raw_bytes(&mut self, bytes: &[u8]) -> Result<()> { - self.writer.write_all(bytes)?; - Ok(()) + self.writer.write_bytes(bytes) } + #[inline(always)] fn write_bytes_with_len(&mut self, bytes: &[u8]) -> Result<()> { let len = bytes.len(); if len < 256 { @@ -210,49 +660,57 @@ impl Encoder { self.write_raw_bytes(bytes) } + #[inline(always)] fn write_string(&mut self, s: &str) -> Result<()> { - // Empty strings must be encoded as BINARY_8 + 0 - if s.is_empty() { - self.write_u8(token::BINARY_8)?; - self.write_u8(0)?; - return Ok(()); + if let Some(string_hints) = self.string_hints + && let Some(hint) = string_hints.hint_for(s) + { + return self.write_string_with_hint(s, hint); } + self.write_string_uncached(s) + } - // Optimization: JID formats are tightly bounded (max ~41 chars for user+agent+device - // with domain). Use a small headroom threshold to avoid scanning long text payloads. - let is_likely_jid = s.len() <= 48; - - if let Some(token) = token::index_of_single_token(s) { - self.write_u8(token)?; - } else if let Some((dict, token)) = token::index_of_double_byte_token(s) { - self.write_u8(token::DICTIONARY_0 + dict)?; - self.write_u8(token)?; - } else if Self::validate_nibble(s) { - self.write_packed_bytes(s, token::NIBBLE_8)?; - } else if Self::validate_hex(s) { - self.write_packed_bytes(s, token::HEX_8)?; - } else if is_likely_jid && let Some(jid) = parse_jid(s) { - self.write_jid(&jid)?; - } else { - self.write_bytes_with_len(s.as_bytes())?; + #[inline(always)] + fn write_string_uncached(&mut self, s: &str) -> Result<()> { + self.write_string_with_hint(s, classify_string_hint(s)) + } + + #[inline(always)] + fn write_string_with_hint(&mut self, s: &str, hint: StringHint) -> Result<()> { + match hint { + StringHint::Empty => { + self.write_u8(token::BINARY_8)?; + self.write_u8(0)?; + } + StringHint::SingleToken(token) => self.write_u8(token)?, + StringHint::DoubleToken { dict, token } => { + self.write_u8(token::DICTIONARY_0 + dict)?; + self.write_u8(token)?; + } + StringHint::PackedNibble => self.write_packed_bytes(s, token::NIBBLE_8)?, + StringHint::PackedHex => self.write_packed_bytes(s, token::HEX_8)?, + StringHint::Jid(meta) => self.write_jid_from_meta(s, meta)?, + StringHint::RawBytes => self.write_bytes_with_len(s.as_bytes())?, } Ok(()) } - fn write_jid(&mut self, jid: &ParsedJid<'_>) -> Result<()> { - if let Some(device) = jid.device { + #[inline(always)] + fn write_jid_from_meta(&mut self, jid: &str, meta: ParsedJidMeta) -> Result<()> { + let (user, server) = split_jid_from_meta(jid, meta); + if let Some(device) = meta.device { self.write_u8(token::AD_JID)?; - self.write_u8(jid.domain_type)?; - self.write_u8(device as u8)?; - self.write_string(jid.user)?; + self.write_u8(meta.domain_type)?; + self.write_u8(device)?; + self.write_string(user)?; } else { self.write_u8(token::JID_PAIR)?; - if jid.user.is_empty() { + if user.is_empty() { self.write_u8(token::LIST_EMPTY)?; } else { - self.write_string(jid.user)?; + self.write_string(user)?; } - self.write_string(jid.server)?; + self.write_string(server)?; } Ok(()) } @@ -262,9 +720,12 @@ impl Encoder { fn write_jid_ref(&mut self, jid: &JidRef<'_>) -> Result<()> { if jid.device > 0 { // AD_JID format: agent/domain_type, device, user + let device = u8::try_from(jid.device).map_err(|_| { + BinaryError::AttrParse(format!("AD_JID device id out of range: {}", jid.device)) + })?; self.write_u8(token::AD_JID)?; self.write_u8(jid.agent)?; - self.write_u8(jid.device as u8)?; + self.write_u8(device)?; self.write_string(&jid.user)?; } else { // JID_PAIR format: user, server @@ -284,9 +745,12 @@ impl Encoder { fn write_jid_owned(&mut self, jid: &Jid) -> Result<()> { if jid.device > 0 { // AD_JID format: agent/domain_type, device, user + let device = u8::try_from(jid.device).map_err(|_| { + BinaryError::AttrParse(format!("AD_JID device id out of range: {}", jid.device)) + })?; self.write_u8(token::AD_JID)?; self.write_u8(jid.agent)?; - self.write_u8(jid.device as u8)?; + self.write_u8(device)?; self.write_string(&jid.user)?; } else { // JID_PAIR format: user, server @@ -301,44 +765,29 @@ impl Encoder { Ok(()) } - fn validate_nibble(value: &str) -> bool { - if value.len() > token::PACKED_MAX as usize { - return false; - } - value - .chars() - .all(|c| c.is_ascii_digit() || c == '-' || c == '.') - } - - fn pack_nibble(value: char) -> u8 { + #[inline(always)] + fn pack_nibble(value: u8) -> u8 { match value { - '-' => 10, - '.' => 11, - '\x00' => 15, - c if c.is_ascii_digit() => c as u8 - b'0', + b'-' => 10, + b'.' => 11, + 0 => 15, + c if c.is_ascii_digit() => c - b'0', _ => panic!("Invalid char for nibble packing: {value}"), } } - fn validate_hex(value: &str) -> bool { - if value.len() > token::PACKED_MAX as usize { - return false; - } - value - .chars() - .all(|c| c.is_ascii_hexdigit() && (c.is_ascii_uppercase() || c.is_ascii_digit())) - } - - fn pack_hex(value: char) -> u8 { + #[inline(always)] + fn pack_hex(value: u8) -> u8 { match value { - c if c.is_ascii_digit() => c as u8 - b'0', - c if ('A'..='F').contains(&c) => 10 + (c as u8 - b'A'), - '\x00' => 15, + c if c.is_ascii_digit() => c - b'0', + c if (b'A'..=b'F').contains(&c) => 10 + (c - b'A'), + 0 => 15, _ => panic!("Invalid char for hex packing: {value}"), } } - fn pack_byte_pair(&self, packer: fn(char) -> u8, part1: char, part2: char) -> u8 { + #[inline(always)] + fn pack_byte_pair(packer: fn(u8) -> u8, part1: u8, part2: u8) -> u8 { (packer(part1) << 4) | packer(part2) } @@ -357,43 +806,58 @@ impl Encoder { let mut input_bytes = value.as_bytes(); - while input_bytes.len() >= 16 { - let (chunk, rest) = input_bytes.split_at(16); - let input = u8x16::from_slice(chunk); + if data_type == token::NIBBLE_8 { + const NIBBLE_LOOKUP: [u8; 16] = + [10, 11, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 255, 255, 255]; + let lookup = Simd::from_array(NIBBLE_LOOKUP); + let nibble_base = Simd::splat(b'-'); - let nibbles = if data_type == token::NIBBLE_8 { - let indices = input.saturating_sub(Simd::splat(b'-')); - const LOOKUP: [u8; 16] = [10, 11, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 255, 255, 255]; - Simd::from_array(LOOKUP).swizzle_dyn(indices) - } else { - let ascii_0 = Simd::splat(b'0'); - let ascii_a = Simd::splat(b'A'); - let ten = Simd::splat(10); + while input_bytes.len() >= 16 { + let (chunk, rest) = input_bytes.split_at(16); + let input = u8x16::from_slice(chunk); + let indices = input.saturating_sub(nibble_base); + let nibbles = lookup.swizzle_dyn(indices); + + let (evens, odds) = nibbles.deinterleave(nibbles.rotate_elements_left::<1>()); + let packed: Simd = (evens << Simd::splat(4)) | odds; + let packed_bytes = packed.to_array(); + self.write_raw_bytes(&packed_bytes[..8])?; + + input_bytes = rest; + } + + let mut bytes_iter = input_bytes.iter().copied(); + while let Some(part1) = bytes_iter.next() { + let part2 = bytes_iter.next().unwrap_or(0); + self.write_u8(Self::pack_byte_pair(Self::pack_nibble, part1, part2))?; + } + } else { + let ascii_0 = Simd::splat(b'0'); + let ascii_a = Simd::splat(b'A'); + let ten = Simd::splat(10); + + while input_bytes.len() >= 16 { + let (chunk, rest) = input_bytes.split_at(16); + let input = u8x16::from_slice(chunk); let digit_vals = input - ascii_0; let letter_vals = input - ascii_a + ten; let is_letter = input.simd_ge(ascii_a); - is_letter.select(letter_vals, digit_vals) - }; - - let (evens, odds) = nibbles.deinterleave(nibbles.rotate_elements_left::<1>()); - let packed: Simd = (evens << Simd::splat(4)) | odds; - let packed_bytes = packed.to_array(); - self.write_raw_bytes(&packed_bytes[..8])?; + let nibbles = is_letter.select(letter_vals, digit_vals); - input_bytes = rest; - } + let (evens, odds) = nibbles.deinterleave(nibbles.rotate_elements_left::<1>()); + let packed: Simd = (evens << Simd::splat(4)) | odds; + let packed_bytes = packed.to_array(); + self.write_raw_bytes(&packed_bytes[..8])?; - let packer: fn(char) -> u8 = if data_type == token::NIBBLE_8 { - Self::pack_nibble - } else { - Self::pack_hex - }; + input_bytes = rest; + } - let mut chars = core::str::from_utf8(input_bytes)?.chars(); - while let Some(part1) = chars.next() { - let part2 = chars.next().unwrap_or('\x00'); - self.write_u8(self.pack_byte_pair(packer, part1, part2))?; + let mut bytes_iter = input_bytes.iter().copied(); + while let Some(part1) = bytes_iter.next() { + let part2 = bytes_iter.next().unwrap_or(0); + self.write_u8(Self::pack_byte_pair(Self::pack_hex, part1, part2))?; + } } Ok(()) } @@ -520,42 +984,42 @@ mod tests { #[test] fn test_hex_validation() { // Valid hex strings (uppercase A-F, digits 0-9) - assert!(Encoder::>::validate_hex("0123456789ABCDEF")); - assert!(Encoder::>::validate_hex("DEADBEEF")); - assert!(Encoder::>::validate_hex("1234")); + assert!(validate_hex("0123456789ABCDEF")); + assert!(validate_hex("DEADBEEF")); + assert!(validate_hex("1234")); // Invalid: lowercase letters - assert!(!Encoder::>::validate_hex("abcdef")); - assert!(!Encoder::>::validate_hex("DeadBeef")); + assert!(!validate_hex("abcdef")); + assert!(!validate_hex("DeadBeef")); // Invalid: special characters - assert!(!Encoder::>::validate_hex("-")); - assert!(!Encoder::>::validate_hex(".")); - assert!(!Encoder::>::validate_hex(" ")); + assert!(!validate_hex("-")); + assert!(!validate_hex(".")); + assert!(!validate_hex(" ")); // Empty string is valid (but will be encoded as regular string) - assert!(Encoder::>::validate_hex("")); + assert!(validate_hex("")); } /// Test nibble packing validation #[test] fn test_nibble_validation() { // Valid nibble strings: digits, dash, dot - assert!(Encoder::>::validate_nibble("0123456789")); - assert!(Encoder::>::validate_nibble("-")); - assert!(Encoder::>::validate_nibble(".")); - assert!(Encoder::>::validate_nibble("123-456.789")); + assert!(validate_nibble("0123456789")); + assert!(validate_nibble("-")); + assert!(validate_nibble(".")); + assert!(validate_nibble("123-456.789")); // Invalid: letters - assert!(!Encoder::>::validate_nibble("abc")); - assert!(!Encoder::>::validate_nibble("123abc")); + assert!(!validate_nibble("abc")); + assert!(!validate_nibble("123abc")); // Invalid: uppercase letters - assert!(!Encoder::>::validate_nibble("ABC")); + assert!(!validate_nibble("ABC")); // Invalid: special characters other than - and . - assert!(!Encoder::>::validate_nibble("123!456")); - assert!(!Encoder::>::validate_nibble("@")); + assert!(!validate_nibble("123!456")); + assert!(!validate_nibble("@")); } /// Test BINARY_8, BINARY_20, BINARY_32 boundary transitions @@ -624,11 +1088,11 @@ mod tests { fn test_packed_max_boundary() { // Exactly PACKED_MAX characters should be valid for packing let max_nibble = "0".repeat(token::PACKED_MAX as usize); - assert!(Encoder::>::validate_nibble(&max_nibble)); + assert!(validate_nibble(&max_nibble)); // One more than PACKED_MAX should NOT be packed let over_max = "0".repeat(token::PACKED_MAX as usize + 1); - assert!(!Encoder::>::validate_nibble(&over_max)); + assert!(!validate_nibble(&over_max)); } /// Test empty string encoding - should be BINARY_8 + 0, not just 0 @@ -699,7 +1163,7 @@ mod tests { use crate::decoder::Decoder; use crate::token; - // Short JID: should be encoded as JID token (256 bytes or less) + // Short JID: should be encoded as a JID token (48 bytes or less) let short_jid = "user@s.whatsapp.net"; let mut buffer = Vec::new(); let mut encoder = Encoder::new(Cursor::new(&mut buffer))?; @@ -712,7 +1176,7 @@ mod tests { "Short JID should be encoded as JID_PAIR token" ); - // Long string (> 256 chars): should be encoded as raw bytes, not as JID + // Long string (> 48 chars): should be encoded as raw bytes, not as JID let long_text = "x".repeat(300) + "@s.whatsapp.net"; let mut buffer = Vec::new(); let mut encoder = Encoder::new(Cursor::new(&mut buffer))?; @@ -750,4 +1214,52 @@ mod tests { Ok(()) } + + #[test] + fn test_jid_parser_preserves_non_numeric_device_suffix() -> TestResult { + use crate::decoder::Decoder; + + let value = "foo:bar@s.whatsapp.net"; + let node = Node::new( + "msg", + Attrs::new(), + Some(NodeContent::String(value.to_string())), + ); + + let mut buffer = Vec::new(); + let mut encoder = Encoder::new(Cursor::new(&mut buffer))?; + encoder.write_node(&node)?; + + let mut decoder = Decoder::new(&buffer[1..]); + let decoded = decoder.read_node_ref()?.to_owned(); + match decoded.content { + Some(NodeContent::String(s)) => assert_eq!(s, value), + other => panic!("Expected string content, got {:?}", other), + } + Ok(()) + } + + #[test] + fn test_jid_parser_preserves_non_numeric_agent_suffix() -> TestResult { + use crate::decoder::Decoder; + + let value = "hello_world@s.whatsapp.net"; + let node = Node::new( + "msg", + Attrs::new(), + Some(NodeContent::String(value.to_string())), + ); + + let mut buffer = Vec::new(); + let mut encoder = Encoder::new(Cursor::new(&mut buffer))?; + encoder.write_node(&node)?; + + let mut decoder = Decoder::new(&buffer[1..]); + let decoded = decoder.read_node_ref()?.to_owned(); + match decoded.content { + Some(NodeContent::String(s)) => assert_eq!(s, value), + other => panic!("Expected string content, got {:?}", other), + } + Ok(()) + } } diff --git a/wacore/binary/src/lib.rs b/wacore/binary/src/lib.rs index d3c4697fe..6b4e2066d 100644 --- a/wacore/binary/src/lib.rs +++ b/wacore/binary/src/lib.rs @@ -14,5 +14,8 @@ pub mod util; pub use attrs::{AttrParser, AttrParserRef}; pub use error::{BinaryError, Result}; -pub use marshal::{marshal, marshal_ref, marshal_ref_to, marshal_to}; +pub use marshal::{ + marshal, marshal_auto, marshal_exact, marshal_ref, marshal_ref_auto, marshal_ref_exact, + marshal_ref_to, marshal_ref_to_vec, marshal_to, marshal_to_vec, +}; pub use node::{Node, NodeRef, NodeValue}; diff --git a/wacore/binary/src/marshal.rs b/wacore/binary/src/marshal.rs index d44eba6e6..e30a7ae2e 100644 --- a/wacore/binary/src/marshal.rs +++ b/wacore/binary/src/marshal.rs @@ -1,6 +1,21 @@ use std::io::Write; -use crate::{BinaryError, Node, NodeRef, Result, decoder::Decoder, encoder::Encoder}; +use crate::{ + BinaryError, Node, NodeRef, Result, + decoder::Decoder, + encoder::{Encoder, build_marshaled_node_plan, build_marshaled_node_ref_plan}, + node::{NodeContent, NodeContentRef}, +}; + +const DEFAULT_MARSHAL_CAPACITY: usize = 1024; +const AUTO_RESERVE_ATTRS_THRESHOLD: usize = 24; +const AUTO_RESERVE_CHILDREN_THRESHOLD: usize = 64; +const AUTO_RESERVE_SCALAR_THRESHOLD: usize = 8 * 1024; +const AUTO_CHILD_SAMPLE_LIMIT: usize = 32; +const AUTO_MAX_HINT_CAPACITY: usize = 512 * 1024; +const AUTO_ATTR_ESTIMATE: usize = 24; +const AUTO_CHILD_ESTIMATE: usize = 96; +const AUTO_GRANDCHILD_ESTIMATE: usize = 40; pub fn unmarshal_ref(data: &[u8]) -> Result> { let mut decoder = Decoder::new(data); @@ -19,9 +34,44 @@ pub fn marshal_to(node: &Node, writer: &mut impl Write) -> Result<()> { Ok(()) } +/// Serialize an owned node directly into a `Vec` using the fast vec writer path. +pub fn marshal_to_vec(node: &Node, output: &mut Vec) -> Result<()> { + let mut encoder = Encoder::new_vec(output)?; + encoder.write_node(node)?; + Ok(()) +} + pub fn marshal(node: &Node) -> Result> { - let mut payload = Vec::with_capacity(1024); - marshal_to(node, &mut payload)?; + let mut payload = Vec::with_capacity(DEFAULT_MARSHAL_CAPACITY); + marshal_to_vec(node, &mut payload)?; + Ok(payload) +} + +/// Serialize a `Node` using a conservative auto strategy. +/// +/// This keeps the fast one-pass path for typical payloads and only uses +/// a lightweight preallocation hint for obviously larger payload shapes. +pub fn marshal_auto(node: &Node) -> Result> { + if should_auto_reserve_node(node) { + marshal_with_capacity(node, estimate_capacity_node(node)) + } else { + marshal(node) + } +} + +/// Serialize a `Node` using a two-pass strategy: +/// 1) compute exact encoded size +/// 2) write directly into a fixed-size output buffer +/// +/// This avoids output buffer growth/copies and can be beneficial for large/variable payloads. +pub fn marshal_exact(node: &Node) -> Result> { + let plan = build_marshaled_node_plan(node); + let mut payload = vec![0; plan.size]; + let mut encoder = Encoder::new_slice(payload.as_mut_slice(), Some(&plan.hints))?; + encoder.write_node(node)?; + let written = encoder.bytes_written(); + debug_assert_eq!(written, payload.len(), "plan size mismatch for Node"); + payload.truncate(written); Ok(payload) } @@ -33,10 +83,356 @@ pub fn marshal_ref_to(node: &NodeRef<'_>, writer: &mut impl Write) -> Result<()> Ok(()) } +/// Serialize a borrowed node directly into a `Vec` using the fast vec writer path. +pub fn marshal_ref_to_vec(node: &NodeRef<'_>, output: &mut Vec) -> Result<()> { + let mut encoder = Encoder::new_vec(output)?; + encoder.write_node(node)?; + Ok(()) +} + /// Zero-copy serialization of a `NodeRef` to a new `Vec`. /// Prefer `marshal_ref_to` with a reusable buffer for best performance. pub fn marshal_ref(node: &NodeRef<'_>) -> Result> { - let mut payload = Vec::with_capacity(1024); - marshal_ref_to(node, &mut payload)?; + let mut payload = Vec::with_capacity(DEFAULT_MARSHAL_CAPACITY); + marshal_ref_to_vec(node, &mut payload)?; + Ok(payload) +} + +/// Serialize a `NodeRef` using the same conservative auto strategy as `marshal_auto`. +pub fn marshal_ref_auto(node: &NodeRef<'_>) -> Result> { + if should_auto_reserve_node_ref(node) { + marshal_ref_with_capacity(node, estimate_capacity_node_ref(node)) + } else { + marshal_ref(node) + } +} + +/// Serialize a `NodeRef` using a two-pass exact-size strategy. +/// +/// This avoids output buffer growth/copies and preserves zero-copy input semantics. +pub fn marshal_ref_exact(node: &NodeRef<'_>) -> Result> { + let plan = build_marshaled_node_ref_plan(node); + let mut payload = vec![0; plan.size]; + let mut encoder = Encoder::new_slice(payload.as_mut_slice(), Some(&plan.hints))?; + encoder.write_node(node)?; + let written = encoder.bytes_written(); + debug_assert_eq!(written, payload.len(), "plan size mismatch for NodeRef"); + payload.truncate(written); + Ok(payload) +} + +#[inline] +fn marshal_with_capacity(node: &Node, capacity: usize) -> Result> { + let mut payload = Vec::with_capacity(capacity); + marshal_to_vec(node, &mut payload)?; Ok(payload) } + +#[inline] +fn marshal_ref_with_capacity(node: &NodeRef<'_>, capacity: usize) -> Result> { + let mut payload = Vec::with_capacity(capacity); + marshal_ref_to_vec(node, &mut payload)?; + Ok(payload) +} + +#[inline] +fn should_auto_reserve_node(node: &Node) -> bool { + if node.attrs.len() >= AUTO_RESERVE_ATTRS_THRESHOLD { + return true; + } + + match &node.content { + Some(NodeContent::Bytes(bytes)) => bytes.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, + Some(NodeContent::String(text)) => text.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, + Some(NodeContent::Nodes(children)) => children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD, + None => false, + } +} + +#[inline] +fn should_auto_reserve_node_ref(node: &NodeRef<'_>) -> bool { + if node.attrs.len() >= AUTO_RESERVE_ATTRS_THRESHOLD { + return true; + } + + match node.content.as_deref() { + Some(NodeContentRef::Bytes(bytes)) => bytes.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, + Some(NodeContentRef::String(text)) => text.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, + Some(NodeContentRef::Nodes(children)) => children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD, + None => false, + } +} + +#[inline] +fn estimate_capacity_node(node: &Node) -> usize { + let mut estimate = DEFAULT_MARSHAL_CAPACITY + 16; + estimate += node.tag.len(); + estimate += node.attrs.len() * AUTO_ATTR_ESTIMATE; + + match &node.content { + Some(NodeContent::Bytes(bytes)) => { + estimate += bytes.len() + 8; + } + Some(NodeContent::String(text)) => { + estimate += text.len() + 8; + } + Some(NodeContent::Nodes(children)) => { + estimate += children.len() * AUTO_CHILD_ESTIMATE; + for child in children.iter().take(AUTO_CHILD_SAMPLE_LIMIT) { + estimate += child.tag.len() + child.attrs.len() * AUTO_ATTR_ESTIMATE; + match &child.content { + Some(NodeContent::Bytes(bytes)) => estimate += bytes.len() + 8, + Some(NodeContent::String(text)) => estimate += text.len() + 8, + Some(NodeContent::Nodes(grand_children)) => { + estimate += grand_children.len() * AUTO_GRANDCHILD_ESTIMATE; + } + None => {} + } + if estimate >= AUTO_MAX_HINT_CAPACITY { + return AUTO_MAX_HINT_CAPACITY; + } + } + } + None => {} + } + + estimate.clamp(DEFAULT_MARSHAL_CAPACITY, AUTO_MAX_HINT_CAPACITY) +} + +#[inline] +fn estimate_capacity_node_ref(node: &NodeRef<'_>) -> usize { + let mut estimate = DEFAULT_MARSHAL_CAPACITY + 16; + estimate += node.tag.len(); + estimate += node.attrs.len() * AUTO_ATTR_ESTIMATE; + + match node.content.as_deref() { + Some(NodeContentRef::Bytes(bytes)) => { + estimate += bytes.len() + 8; + } + Some(NodeContentRef::String(text)) => { + estimate += text.len() + 8; + } + Some(NodeContentRef::Nodes(children)) => { + estimate += children.len() * AUTO_CHILD_ESTIMATE; + for child in children.iter().take(AUTO_CHILD_SAMPLE_LIMIT) { + estimate += child.tag.len() + child.attrs.len() * AUTO_ATTR_ESTIMATE; + match child.content.as_deref() { + Some(NodeContentRef::Bytes(bytes)) => estimate += bytes.len() + 8, + Some(NodeContentRef::String(text)) => estimate += text.len() + 8, + Some(NodeContentRef::Nodes(grand_children)) => { + estimate += grand_children.len() * AUTO_GRANDCHILD_ESTIMATE; + } + None => {} + } + if estimate >= AUTO_MAX_HINT_CAPACITY { + return AUTO_MAX_HINT_CAPACITY; + } + } + } + None => {} + } + + estimate.clamp(DEFAULT_MARSHAL_CAPACITY, AUTO_MAX_HINT_CAPACITY) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::jid::Jid; + use crate::node::{Attrs, NodeContent, NodeValue}; + + type TestResult = crate::error::Result<()>; + + fn fixture_node() -> Node { + let mut attrs = Attrs::with_capacity(4); + attrs.push("id".to_string(), "ABC123"); + attrs.push("to".to_string(), "123456789@s.whatsapp.net"); + attrs.push( + "participant".to_string(), + NodeValue::Jid("15551234567@s.whatsapp.net".parse::().unwrap()), + ); + attrs.push("hex".to_string(), "DEADBEEF"); + + let child = Node::new( + "item", + Attrs::new(), + Some(NodeContent::Bytes(vec![1, 2, 3, 4, 5, 6, 7, 8])), + ); + + Node::new( + "message", + attrs, + Some(NodeContent::Nodes(vec![ + child, + Node::new( + "text", + Attrs::new(), + Some(NodeContent::String("hello".repeat(40))), + ), + ])), + ) + } + + fn large_binary_fixture() -> Node { + Node::new( + "message", + Attrs::new(), + Some(NodeContent::Bytes(vec![ + 0xAB; + AUTO_RESERVE_SCALAR_THRESHOLD + 2048 + ])), + ) + } + + #[test] + fn test_marshaled_node_size_matches_output() -> TestResult { + let node = fixture_node(); + let plan = build_marshaled_node_plan(&node); + let payload = marshal(&node)?; + assert_eq!(payload.len(), plan.size); + Ok(()) + } + + #[test] + fn test_marshaled_node_ref_size_matches_output() -> TestResult { + let node = fixture_node(); + let node_ref = node.as_node_ref(); + let plan = build_marshaled_node_ref_plan(&node_ref); + let payload = marshal_ref(&node_ref)?; + assert_eq!(payload.len(), plan.size); + Ok(()) + } + + #[test] + fn test_marshal_matches_marshal_to_bytes() -> TestResult { + let node = fixture_node(); + + let payload_alloc = marshal(&node)?; + + let mut payload_writer = Vec::new(); + marshal_to(&node, &mut payload_writer)?; + + assert_eq!(payload_alloc, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_ref_matches_marshal_ref_to_bytes() -> TestResult { + let node = fixture_node(); + let node_ref = node.as_node_ref(); + + let payload_alloc = marshal_ref(&node_ref)?; + + let mut payload_writer = Vec::new(); + marshal_ref_to(&node_ref, &mut payload_writer)?; + + assert_eq!(payload_alloc, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_to_vec_matches_marshal_to() -> TestResult { + let node = fixture_node(); + + let mut payload_vec_writer = Vec::new(); + marshal_to_vec(&node, &mut payload_vec_writer)?; + + let mut payload_writer = Vec::new(); + marshal_to(&node, &mut payload_writer)?; + + assert_eq!(payload_vec_writer, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_ref_to_vec_matches_marshal_ref_to() -> TestResult { + let node = fixture_node(); + let node_ref = node.as_node_ref(); + + let mut payload_vec_writer = Vec::new(); + marshal_ref_to_vec(&node_ref, &mut payload_vec_writer)?; + + let mut payload_writer = Vec::new(); + marshal_ref_to(&node_ref, &mut payload_writer)?; + + assert_eq!(payload_vec_writer, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_exact_matches_marshal_to_bytes() -> TestResult { + let node = fixture_node(); + + let payload_exact = marshal_exact(&node)?; + + let mut payload_writer = Vec::new(); + marshal_to(&node, &mut payload_writer)?; + + assert_eq!(payload_exact, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_ref_exact_matches_marshal_ref_to_bytes() -> TestResult { + let node = fixture_node(); + let node_ref = node.as_node_ref(); + + let payload_exact = marshal_ref_exact(&node_ref)?; + + let mut payload_writer = Vec::new(); + marshal_ref_to(&node_ref, &mut payload_writer)?; + + assert_eq!(payload_exact, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_auto_matches_marshal_to_bytes() -> TestResult { + let node = fixture_node(); + let payload_auto = marshal_auto(&node)?; + + let mut payload_writer = Vec::new(); + marshal_to(&node, &mut payload_writer)?; + + assert_eq!(payload_auto, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_ref_auto_matches_marshal_ref_to_bytes() -> TestResult { + let node = fixture_node(); + let node_ref = node.as_node_ref(); + let payload_auto = marshal_ref_auto(&node_ref)?; + + let mut payload_writer = Vec::new(); + marshal_ref_to(&node_ref, &mut payload_writer)?; + + assert_eq!(payload_auto, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_auto_large_binary_matches_marshal_to_bytes() -> TestResult { + let node = large_binary_fixture(); + let payload_auto = marshal_auto(&node)?; + + let mut payload_writer = Vec::new(); + marshal_to(&node, &mut payload_writer)?; + + assert_eq!(payload_auto, payload_writer); + Ok(()) + } + + #[test] + fn test_marshal_ref_auto_large_binary_matches_marshal_ref_to_bytes() -> TestResult { + let node = large_binary_fixture(); + let node_ref = node.as_node_ref(); + let payload_auto = marshal_ref_auto(&node_ref)?; + + let mut payload_writer = Vec::new(); + marshal_ref_to(&node_ref, &mut payload_writer)?; + + assert_eq!(payload_auto, payload_writer); + Ok(()) + } +}