Skip to content

Commit fb8c68c

Browse files
Reduce allocations in payer proof merkle operations
Remove nonce_hash from TlvMerkleData, saving 32 bytes per TLV record. The nonce hash was stored for every TLV but only consumed by included TLVs when building leaf_hashes, so collect those directly during the TlvStream iteration instead. Inline reconstruct_positions_from_records into reconstruct_merkle_root so that TreeNodes are built directly during position reconstruction, eliminating the intermediate Vec<bool>. Pre-allocate missing_with_types and needs_hash with with_capacity(num_omitted) to avoid repeated reallocations during tree traversal. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 6ca5d54 commit fb8c68c

2 files changed

Lines changed: 63 additions & 66 deletions

File tree

lightning/src/offers/merkle.rs

Lines changed: 56 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,6 @@ pub(super) struct SelectiveDisclosure {
327327
struct TlvMerkleData {
328328
tlv_type: u64,
329329
per_tlv_hash: sha256::Hash,
330-
nonce_hash: sha256::Hash,
331330
is_included: bool,
332331
}
333332

@@ -358,27 +357,23 @@ pub(super) fn compute_selective_disclosure(
358357
let branch_tag = tagged_hash_engine(sha256::Hash::hash("LnBranch".as_bytes()));
359358

360359
let mut tlv_data: Vec<TlvMerkleData> = Vec::new();
360+
let mut leaf_hashes: Vec<sha256::Hash> = Vec::new();
361361
for record in tlv_stream.filter(|r| !SIGNATURE_TYPES.contains(&r.r#type)) {
362362
let leaf_hash = tagged_hash_from_engine(leaf_tag.clone(), record.record_bytes);
363363
let nonce_hash = tagged_hash_from_engine(nonce_tag.clone(), record.type_bytes);
364364
let per_tlv_hash =
365365
tagged_branch_hash_from_engine(branch_tag.clone(), leaf_hash, nonce_hash);
366366

367367
let is_included = included_types.contains(&record.r#type);
368-
tlv_data.push(TlvMerkleData {
369-
tlv_type: record.r#type,
370-
per_tlv_hash,
371-
nonce_hash,
372-
is_included,
373-
});
368+
if is_included {
369+
leaf_hashes.push(nonce_hash);
370+
}
371+
tlv_data.push(TlvMerkleData { tlv_type: record.r#type, per_tlv_hash, is_included });
374372
}
375373

376374
if tlv_data.is_empty() {
377375
return Err(SelectiveDisclosureError::EmptyTlvStream);
378376
}
379-
380-
let leaf_hashes: Vec<_> =
381-
tlv_data.iter().filter(|d| d.is_included).map(|d| d.nonce_hash).collect();
382377
let omitted_markers = compute_omitted_markers(&tlv_data);
383378
let (merkle_root, missing_hashes) = build_tree_with_disclosure(&tlv_data, &branch_tag);
384379

@@ -438,6 +433,8 @@ fn build_tree_with_disclosure(
438433
let num_nodes = tlv_data.len();
439434
debug_assert!(num_nodes > 0, "TLV stream must contain at least one record");
440435

436+
let num_omitted = tlv_data.iter().filter(|d| !d.is_included).count();
437+
441438
let mut nodes: Vec<TreeNode> = tlv_data
442439
.iter()
443440
.map(|data| TreeNode {
@@ -447,7 +444,7 @@ fn build_tree_with_disclosure(
447444
})
448445
.collect();
449446

450-
let mut missing_with_types: Vec<(u64, sha256::Hash)> = Vec::new();
447+
let mut missing_with_types: Vec<(u64, sha256::Hash)> = Vec::with_capacity(num_omitted);
451448

452449
for level in 0.. {
453450
let step = 2 << level;
@@ -527,32 +524,65 @@ pub(super) fn reconstruct_merkle_root<'a>(
527524
return Err(SelectiveDisclosureError::LeafHashCountMismatch);
528525
}
529526

530-
let positions = reconstruct_positions_from_records(included_records, omitted_markers);
531-
532-
let num_nodes = positions.len();
533-
534527
let leaf_tag = tagged_hash_engine(sha256::Hash::hash("LnLeaf".as_bytes()));
535528
let branch_tag = tagged_hash_engine(sha256::Hash::hash("LnBranch".as_bytes()));
536529

530+
// Build TreeNode vec directly by interleaving included/omitted positions,
531+
// eliminating the intermediate Vec<bool> from reconstruct_positions_from_records.
532+
let num_nodes = 1 + included_records.len() + omitted_markers.len();
537533
let mut nodes: Vec<TreeNode> = Vec::with_capacity(num_nodes);
538-
let mut leaf_hash_idx = 0;
539-
for (i, &incl) in positions.iter().enumerate() {
540-
let hash = if incl {
541-
let (_, record_bytes) = included_records[leaf_hash_idx];
534+
535+
// TLV0 is always omitted
536+
nodes.push(TreeNode { hash: None, included: false, min_type: 0 });
537+
538+
let mut inc_idx = 0;
539+
let mut mrk_idx = 0;
540+
let mut prev_marker: u64 = 0;
541+
let mut node_idx: u64 = 1;
542+
543+
while inc_idx < included_records.len() || mrk_idx < omitted_markers.len() {
544+
if mrk_idx >= omitted_markers.len() {
545+
// No more markers, remaining positions are included
546+
let (_, record_bytes) = included_records[inc_idx];
542547
let leaf_hash = tagged_hash_from_engine(leaf_tag.clone(), record_bytes);
543-
let nonce_hash = leaf_hashes[leaf_hash_idx];
544-
leaf_hash_idx += 1;
545-
Some(tagged_branch_hash_from_engine(branch_tag.clone(), leaf_hash, nonce_hash))
548+
let nonce_hash = leaf_hashes[inc_idx];
549+
let hash = tagged_branch_hash_from_engine(branch_tag.clone(), leaf_hash, nonce_hash);
550+
nodes.push(TreeNode { hash: Some(hash), included: true, min_type: node_idx });
551+
inc_idx += 1;
552+
} else if inc_idx >= included_records.len() {
553+
// No more included types, remaining positions are omitted
554+
nodes.push(TreeNode { hash: None, included: false, min_type: node_idx });
555+
prev_marker = omitted_markers[mrk_idx];
556+
mrk_idx += 1;
546557
} else {
547-
None
548-
};
549-
nodes.push(TreeNode { hash, included: incl, min_type: i as u64 });
558+
let marker = omitted_markers[mrk_idx];
559+
let (inc_type, _) = included_records[inc_idx];
560+
561+
if marker == prev_marker + 1 {
562+
// Continuation of current run -> omitted position
563+
nodes.push(TreeNode { hash: None, included: false, min_type: node_idx });
564+
prev_marker = marker;
565+
mrk_idx += 1;
566+
} else {
567+
// Jump detected -> included position comes first
568+
let (_, record_bytes) = included_records[inc_idx];
569+
let leaf_hash = tagged_hash_from_engine(leaf_tag.clone(), record_bytes);
570+
let nonce_hash = leaf_hashes[inc_idx];
571+
let hash =
572+
tagged_branch_hash_from_engine(branch_tag.clone(), leaf_hash, nonce_hash);
573+
nodes.push(TreeNode { hash: Some(hash), included: true, min_type: node_idx });
574+
prev_marker = inc_type;
575+
inc_idx += 1;
576+
}
577+
}
578+
node_idx += 1;
550579
}
551580

552581
// First pass: walk the tree to discover which positions need missing hashes.
553582
// We mutate nodes[].included and nodes[].min_type directly since the second
554583
// pass only reads nodes[].hash, making this safe without a separate allocation.
555-
let mut needs_hash: Vec<(u64, usize)> = Vec::new();
584+
let num_omitted = omitted_markers.len() + 1; // +1 for implicit TLV0
585+
let mut needs_hash: Vec<(u64, usize)> = Vec::with_capacity(num_omitted);
556586

557587
for level in 0.. {
558588
let step = 2 << level;
@@ -712,46 +742,6 @@ fn reconstruct_positions(included_types: &[u64], omitted_markers: &[u64]) -> Vec
712742
positions
713743
}
714744

715-
/// Like `reconstruct_positions`, but extracts types directly from included records,
716-
/// avoiding a separate Vec allocation for the types.
717-
fn reconstruct_positions_from_records(
718-
included_records: &[(u64, &[u8])], omitted_markers: &[u64],
719-
) -> Vec<bool> {
720-
let total = 1 + included_records.len() + omitted_markers.len();
721-
let mut positions = Vec::with_capacity(total);
722-
positions.push(false); // TLV0 is always omitted
723-
724-
let mut inc_idx = 0;
725-
let mut mrk_idx = 0;
726-
let mut prev_marker: u64 = 0;
727-
728-
while inc_idx < included_records.len() || mrk_idx < omitted_markers.len() {
729-
if mrk_idx >= omitted_markers.len() {
730-
positions.push(true);
731-
inc_idx += 1;
732-
} else if inc_idx >= included_records.len() {
733-
positions.push(false);
734-
prev_marker = omitted_markers[mrk_idx];
735-
mrk_idx += 1;
736-
} else {
737-
let marker = omitted_markers[mrk_idx];
738-
let (inc_type, _) = included_records[inc_idx];
739-
740-
if marker == prev_marker + 1 {
741-
positions.push(false);
742-
prev_marker = marker;
743-
mrk_idx += 1;
744-
} else {
745-
positions.push(true);
746-
prev_marker = inc_type;
747-
inc_idx += 1;
748-
}
749-
}
750-
}
751-
752-
positions
753-
}
754-
755745
#[cfg(test)]
756746
mod tests {
757747
use super::{TlvStream, SIGNATURE_TYPES};

lightning/src/offers/payer_proof.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -466,6 +466,13 @@ impl AsRef<[u8]> for PayerProof {
466466
}
467467
}
468468

469+
// TODO: This uses manual TLV parsing rather than the standard `ParsedMessage` +
470+
// `tlv_stream!` pattern (used by Offer, InvoiceRequest, Bolt12Invoice) because payer
471+
// proofs have a hybrid structure: a dynamic set of included invoice TLV records
472+
// (preserved as raw bytes for merkle reconstruction) plus payer-proof-specific TLVs
473+
// (240-250) with non-standard encodings (BigSize lists, concatenated hashes).
474+
// Possible improvements: separate parsing from semantic validation into two layers,
475+
// extract helpers for repeated cursor+read patterns (read_tlv_value, read_hash_list).
469476
impl TryFrom<Vec<u8>> for PayerProof {
470477
type Error = crate::offers::parse::Bolt12ParseError;
471478

0 commit comments

Comments
 (0)