perf: binary - #279
Conversation
📝 WalkthroughWalkthroughAdds a ByteWriter abstraction and JID/string-hint-aware encoder with precomputed marshaling plans and new marshal APIs (auto/exact/vec/ref variants); expands benchmarks; and optimizes decoder packed-string decoding with unchecked/SIMD paths and tag-dispatch entry points. Public API gains new marshal re-exports. Changes
Sequence DiagramsequenceDiagram
participant Caller as Caller
participant Marshal as marshal_auto / marshal_exact
participant Planner as build_marshaled_node_plan
participant Cache as StringHintCache
participant Encoder as Encoder::new_slice / new_vec
participant Writer as SliceByteWriter / VecByteWriter / IoByteWriter
Caller->>Marshal: request marshal(node)
alt exact-size path
Marshal->>Planner: build_marshaled_node_plan(node)
Planner->>Cache: classify_string_hint / collect hints
Cache-->>Planner: hints
Planner-->>Marshal: MarshaledSizePlan {size, hints}
Marshal->>Encoder: Encoder::new_slice(buf, hints)
else auto-reserve path
Marshal->>Marshal: estimate_capacity_node / reserve Vec
Marshal->>Encoder: Encoder::new_vec(buf)
end
Encoder->>Writer: write_u8 / write_bytes (use hints / write_jid_*)
Writer-->>Encoder: bytes_written
Encoder-->>Marshal: encoded bytes
Marshal-->>Caller: payload
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
No actionable comments were generated in the recent review. 🎉 🧹 Recent nitpick comments
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Bencher Report![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
wacore/binary/src/encoder.rs (1)
264-314:⚠️ Potential issue | 🔴 Critical
parse_jid_metasilently drops text after the first:or_even when the suffix is not a valid number.When
find(':')matches butdevice_part.parse::<u16>()fails,ua(the part before:) is still used for subsequent underscore scanning and determinesuser_end. Likewise, whenfind('_')matches but the agent suffix is not a validu16,user_endis set tounderscore_idx, truncating the user.Concrete example: the string
"foo:bar@baz.com"(≤ 48 chars, contains@) will be classified asStringHint::Jidwithuser_end = 3.write_jid_from_metathen encodes user"foo"+ server"baz.com"viaJID_PAIR, and the decoder reconstructs"foo@baz.com"— silently dropping":bar".The same data-loss applies to underscores:
"hello_world@s.whatsapp.net"→ encoded as user"hello", decoded as"hello@s.whatsapp.net".When the suffix doesn't parse as a number, fall back to the full
user_combinedso no data is lost:Proposed fix
fn parse_jid_meta(input: &str) -> Option<ParsedJidMeta> { let sep_idx = input.find('@')?; let server_start = sep_idx + 1; let server = &input[server_start..]; let user_combined = &input[..sep_idx]; - let (user_agent, device) = if let Some(colon_idx) = user_combined.find(':') { - let ua = &user_combined[..colon_idx]; - let device_part = &user_combined[colon_idx + 1..]; - let parsed_device = if device_part.is_empty() { - None - } else { - device_part.parse::<u16>().ok() - }; - (ua, parsed_device) + let (user_agent, device) = if let Some(colon_idx) = user_combined.find(':') + && let Ok(d) = user_combined[colon_idx + 1..].parse::<u16>() + { + (&user_combined[..colon_idx], Some(d)) } else { (user_combined, None) }; - let (user_end, agent_override) = if let Some(underscore_idx) = user_agent.find('_') { - let agent_part = &user_agent[underscore_idx + 1..]; - ( - underscore_idx, - if agent_part.is_empty() { - None - } else { - agent_part.parse::<u16>().ok() - }, - ) + let (user_end, agent_override) = if let Some(underscore_idx) = user_agent.find('_') + && let Ok(a) = user_agent[underscore_idx + 1..].parse::<u16>() + { + (underscore_idx, Some(a)) } else { (user_agent.len(), None) };
🤖 Fix all issues with AI agents
In `@wacore/binary/src/marshal.rs`:
- Around line 45-51: marshal_ref currently allocates payload using plan.size and
returns it without truncating trailing zeros, mirroring the same trailing-zeros
risk as marshal; update marshal_ref so after writing with Encoder::new_slice and
encoder.write_node it slices/truncates payload to encoder.bytes_written() before
returning (use build_marshaled_node_ref_plan, plan.size,
encoder.bytes_written(), Encoder::new_slice and encoder.write_node to locate the
logic) to ensure the returned Vec<u8> contains only the actual encoded bytes.
- Around line 27-32: The vector payload is allocated with plan.size which may
overestimate the actual encoded length and leave trailing zeros; after calling
Encoder::write_node (using Encoder::new_slice with plan.hints) retrieve the
actual number of bytes via encoder.bytes_written() and truncate the payload to
that length (e.g., payload.truncate(bytes_written)) before returning so the
returned Vec<u8> contains only the real encoded bytes; update the return path in
the same function that calls build_marshaled_node_plan and uses
Encoder::new_slice/encoder.write_node to perform this truncation.
🧹 Nitpick comments (4)
wacore/binary/src/decoder.rs (1)
306-310: Unsafefrom_utf8_uncheckedis sound but removes a safety net.The SAFETY argument is correct: both
HEX_LOOKUPandNIBBLE_LOOKUPproduce only ASCII bytes (including\x00padding), so the output is always valid UTF-8. The SIMD validation for NIBBLE_8 (lines 271-284) and the scalar fallback viaunpack_byte(lines 295-300) further ensure only expected nibble values reach the lookup.That said,
from_utf8_uncheckedturns future lookup-table bugs into UB rather than a cleanInvalidUtf8error. Given the bytes are ASCII,from_utf8should already be near-zero cost (fast-path short-circuit). Consider whether the marginal speedup justifies losing that defense layer — if it does, this is fine as-is.wacore/binary/src/encoder.rs (3)
633-669:write_string_uncachedduplicateswrite_string_with_hint— delegate instead.The match arms are identical.
write_string_uncachedcan just classify and delegate:Proposed simplification
fn write_string_uncached(&mut self, s: &str) -> Result<()> { - match classify_string_hint(s) { - 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(()) + self.write_string_with_hint(s, classify_string_hint(s)) }
188-202:StrKeyrelies on pointer identity — sound here but fragile across allocations.
StrKeycaches based on(ptr, len), so it only matches the exact same string allocation. This is fine when the cache is built and consumed during a single immutable borrow of the same node, but a staleStringHintCachereused with a different node could silently miss every entry (benign) or, after allocator reuse, hit the wrong hint (UB-adjacent correctness bug). Consider adding a doc comment toStringHintCachenoting that it must not outlive the strings it was built from.
692-732:device as u8truncates silently for devices > 255.Lines 676, 697, and 719 all cast
device(au16) tou8. If a device value exceeds 255, the high byte is silently dropped. The protocol'sAD_JIDformat encodes only one byte for device, so this is likely a known protocol constraint — but a debug assertion or comment documenting this invariant would prevent confusion.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@wacore/binary/src/encoder.rs`:
- Around line 346-366: The comment points out a mismatch between the JID length
heuristic and the test comment: in classify_string_hint the JID gate uses
is_likely_jid = s.len() <= 48, not 256; update the test comment to reflect the
actual 48-byte threshold (or change the heuristic if intended) so they match.
Locate classify_string_hint and the test around the comment (test line
referencing "256 bytes or less") and either adjust the test comment to say "48
bytes or less" or change the heuristic to use 256 if that was the original
intent; ensure any related tests still pass after the change.
🧹 Nitpick comments (6)
wacore/binary/src/decoder.rs (2)
113-121: Prefer token constants over magic numbers for consistency.Lines 117–118 use raw
248/249while the rest of the file (and the encoder) consistently referencestoken::LIST_8/token::LIST_16. Using the named constants improves readability and reduces the risk of silent divergence if those values ever change.Suggested fix
fn read_list_size(&mut self, tag: u8) -> Result<usize> { 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)), } }
331-373: SIMD nibble validation fallback usesunreachable!— will panic on hypothetical SIMD mismatch instead of returning an error.If the SIMD validation (
simd_le/simd_eq) ever disagrees with the scalarunpack_nibblelogic (e.g., due to a portable_simd bug on an unusual target), line 355 panics in release builds. The current validation logic is correct, so this is unlikely, but replacingunreachable!withunsafe { core::hint::unreachable_unchecked() }(to match the perf-orientedunsafestyle elsewhere) or simply removing it and falling through to the scalar output would be more resilient.That said, the SIMD mask (
<= 11 || == 15) provably matches the scalar match arms, so this is a very low-risk concern.wacore/binary/src/marshal.rs (1)
162-232: Capacity estimation functions are near-identical forNodeandNodeRef— consider a shared helper.
estimate_capacity_nodeandestimate_capacity_node_refhave the same structure and constants, differing only in how they access content (&node.contentvsnode.content.as_deref()). A trait or closure-based approach could eliminate the duplication, but given the structural match differences, this is acceptable as-is.wacore/binary/src/encoder.rs (3)
212-226:StrKeyuses pointer identity — relies on node immutability between planning and encoding passes.This is safe in the current
marshal_exactflow wherebuild_marshaled_node_planandEncoder::new_sliceoperate on the same node in the same function scope. If this cache were ever reused across different node instances or after mutation/reallocation, lookups would silently miss.A brief
// SAFETY:comment documenting this invariant (same node, no mutation between passes) would help future maintainers.
676-714:write_string_uncachedandwrite_string_with_hintshare identical match arms.Both functions dispatch on the same
StringHintvariants with the same writing logic. The only difference is thatwrite_string_with_hintreceives a pre-computed hint whilewrite_string_uncachedcomputes it viaclassify_string_hint. Consider extracting the shared match body into a singlewrite_string_from_hintand having both callers resolve the hint first.Suggested consolidation
- fn write_string_uncached(&mut self, s: &str) -> Result<()> { - match classify_string_hint(s) { - StringHint::Empty => { ... } - // ... identical arms ... - } - Ok(()) - } - - fn write_string_with_hint(&mut self, s: &str, hint: StringHint) -> Result<()> { - match hint { - StringHint::Empty => { ... } - // ... identical arms ... - } - Ok(()) - } + fn write_string_uncached(&mut self, s: &str) -> Result<()> { + let hint = classify_string_hint(s); + self.write_string_with_hint(s, hint) + }
252-280: Linear scan inStringHintCacheis appropriate for bounded size but could become a bottleneck for deeply nested nodes.
hint_forandhint_or_insertdo O(n) scans on every string lookup, capped atMAX_HINT_ENTRIES = 96. For the current workloads this is fine and likely faster than aHashMap. Just noting that ifMAX_HINT_ENTRIESis ever raised significantly, switching to aHashMap(or a fixed-size hash table) would be warranted.
1 participant
Summary by CodeRabbit
New Features
Refactor
Performance
Tests / Benchmarks