@@ -327,7 +327,6 @@ pub(super) struct SelectiveDisclosure {
327327struct 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) ]
756746mod tests {
757747 use super :: { TlvStream , SIGNATURE_TYPES } ;
0 commit comments