Skip to content

Commit 775a061

Browse files
authored
refactor: remove debug assert from proof aggregation (#288)
`compact_attestations` and `extend_proofs_greedily` previously tracked attestations and signature proofs in parallel `Vec`s guarded by a `debug_assert_eq!` on length. Switching to `Vec<(att, proof)>` makes the invariant structural: the two lists can no longer drift out of sync, and the debug assertion is gone. Paired tuples flow end-to-end through `build_block`, unzipping only at the final SSZ boundary. Also eliminates an `AggregatedXMSS::clone()` per child inside `aggregate_mixed` / `aggregate_proofs`: the new `unzip` pattern lets us borrow pubkey slices (required by `xmss_aggregate`'s tuple type) while moving the large aggregate values into the child list. Removes the dead `to_children_refs` helper.
1 parent 2c08f8b commit 775a061

File tree

2 files changed

+159
-95
lines changed

2 files changed

+159
-95
lines changed

crates/blockchain/src/store.rs

Lines changed: 146 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -1170,20 +1170,17 @@ fn union_aggregation_bits(a: &AggregationBits, b: &AggregationBits) -> Aggregati
11701170
/// - Multiple entries: merged into one using recursive proof aggregation
11711171
/// (leanSpec PR #510).
11721172
fn compact_attestations(
1173-
attestations: Vec<AggregatedAttestation>,
1174-
proofs: Vec<AggregatedSignatureProof>,
1173+
entries: Vec<(AggregatedAttestation, AggregatedSignatureProof)>,
11751174
head_state: &State,
1176-
) -> Result<(Vec<AggregatedAttestation>, Vec<AggregatedSignatureProof>), StoreError> {
1177-
debug_assert_eq!(attestations.len(), proofs.len());
1178-
1179-
if attestations.len() <= 1 {
1180-
return Ok((attestations, proofs));
1175+
) -> Result<Vec<(AggregatedAttestation, AggregatedSignatureProof)>, StoreError> {
1176+
if entries.len() <= 1 {
1177+
return Ok(entries);
11811178
}
11821179

11831180
// Group indices by AttestationData, preserving first-occurrence order
11841181
let mut order: Vec<AttestationData> = Vec::new();
11851182
let mut groups: HashMap<AttestationData, Vec<usize>> = HashMap::new();
1186-
for (i, att) in attestations.iter().enumerate() {
1183+
for (i, (att, _)) in entries.iter().enumerate() {
11871184
match groups.entry(att.data.clone()) {
11881185
std::collections::hash_map::Entry::Vacant(e) => {
11891186
order.push(e.key().clone());
@@ -1196,23 +1193,21 @@ fn compact_attestations(
11961193
}
11971194

11981195
// Fast path: no duplicates
1199-
if order.len() == attestations.len() {
1200-
return Ok((attestations, proofs));
1196+
if order.len() == entries.len() {
1197+
return Ok(entries);
12011198
}
12021199

12031200
// Wrap in Option so we can .take() items by index without cloning
12041201
let mut items: Vec<Option<(AggregatedAttestation, AggregatedSignatureProof)>> =
1205-
attestations.into_iter().zip(proofs).map(Some).collect();
1202+
entries.into_iter().map(Some).collect();
12061203

1207-
let mut compacted_atts = Vec::with_capacity(order.len());
1208-
let mut compacted_proofs = Vec::with_capacity(order.len());
1204+
let mut compacted = Vec::with_capacity(order.len());
12091205

12101206
for data in order {
12111207
let indices = &groups[&data];
12121208
if indices.len() == 1 {
1213-
let (att, proof) = items[indices[0]].take().expect("index used once");
1214-
compacted_atts.push(att);
1215-
compacted_proofs.push(proof);
1209+
let item = items[indices[0]].take().expect("index used once");
1210+
compacted.push(item);
12161211
continue;
12171212
}
12181213

@@ -1253,25 +1248,32 @@ fn compact_attestations(
12531248
.map_err(StoreError::SignatureAggregationFailed)?;
12541249

12551250
let merged_proof = AggregatedSignatureProof::new(merged_bits.clone(), merged_proof_data);
1256-
1257-
compacted_proofs.push(merged_proof);
1258-
compacted_atts.push(AggregatedAttestation {
1251+
let merged_att = AggregatedAttestation {
12591252
aggregation_bits: merged_bits,
12601253
data,
1261-
});
1254+
};
1255+
compacted.push((merged_att, merged_proof));
12621256
}
12631257

1264-
Ok((compacted_atts, compacted_proofs))
1258+
Ok(compacted)
12651259
}
12661260

12671261
/// Greedily select proofs maximizing new validator coverage.
12681262
///
12691263
/// For a single attestation data entry, picks proofs that cover the most
1270-
/// uncovered validators. Each selected proof produces one AggregatedAttestation.
1264+
/// uncovered validators. A proof is selected as long as it adds at least
1265+
/// one previously-uncovered validator; partially-overlapping participants
1266+
/// between selected proofs are allowed. `compact_attestations` later feeds
1267+
/// these proofs as children to `aggregate_proofs`, which delegates to
1268+
/// `xmss_aggregate` — that function tracks duplicate pubkeys across
1269+
/// children via its `dup_pub_keys` machinery, so overlap is supported by
1270+
/// the underlying aggregation scheme.
1271+
///
1272+
/// Each selected proof is appended to `selected` paired with its
1273+
/// corresponding AggregatedAttestation.
12711274
fn extend_proofs_greedily(
12721275
proofs: &[AggregatedSignatureProof],
1273-
selected_proofs: &mut Vec<AggregatedSignatureProof>,
1274-
attestations: &mut Vec<AggregatedAttestation>,
1276+
selected: &mut Vec<(AggregatedAttestation, AggregatedSignatureProof)>,
12751277
att_data: &AttestationData,
12761278
) {
12771279
if proofs.is_empty() {
@@ -1309,16 +1311,16 @@ fn extend_proofs_greedily(
13091311
.filter(|vid| !covered.contains(vid))
13101312
.collect();
13111313

1312-
attestations.push(AggregatedAttestation {
1314+
let att = AggregatedAttestation {
13131315
aggregation_bits: proof.participants.clone(),
13141316
data: att_data.clone(),
1315-
});
1316-
selected_proofs.push(proof.clone());
1317+
};
13171318

13181319
metrics::inc_pq_sig_aggregated_signatures();
13191320
metrics::inc_pq_sig_attestations_in_aggregated_signatures(new_covered.len() as u64);
13201321

13211322
covered.extend(new_covered);
1323+
selected.push((att, proof.clone()));
13221324
remaining_indices.remove(&best_idx);
13231325
}
13241326
}
@@ -1339,8 +1341,7 @@ fn build_block(
13391341
known_block_roots: &HashSet<H256>,
13401342
aggregated_payloads: &HashMap<H256, (AttestationData, Vec<AggregatedSignatureProof>)>,
13411343
) -> Result<(Block, Vec<AggregatedSignatureProof>, PostBlockCheckpoints), StoreError> {
1342-
let mut aggregated_attestations: Vec<AggregatedAttestation> = Vec::new();
1343-
let mut aggregated_signatures: Vec<AggregatedSignatureProof> = Vec::new();
1344+
let mut selected: Vec<(AggregatedAttestation, AggregatedSignatureProof)> = Vec::new();
13441345

13451346
if !aggregated_payloads.is_empty() {
13461347
// Genesis edge case: when building on genesis (slot 0),
@@ -1382,21 +1383,18 @@ fn build_block(
13821383
processed_data_roots.insert(*data_root);
13831384
found_new = true;
13841385

1385-
extend_proofs_greedily(
1386-
proofs,
1387-
&mut aggregated_signatures,
1388-
&mut aggregated_attestations,
1389-
att_data,
1390-
);
1386+
extend_proofs_greedily(proofs, &mut selected, att_data);
13911387
}
13921388

13931389
if !found_new {
13941390
break;
13951391
}
13961392

13971393
// Check if justification advanced
1398-
let attestations: AggregatedAttestations = aggregated_attestations
1399-
.clone()
1394+
let attestations: AggregatedAttestations = selected
1395+
.iter()
1396+
.map(|(att, _)| att.clone())
1397+
.collect::<Vec<_>>()
14001398
.try_into()
14011399
.expect("attestation count exceeds limit");
14021400
let candidate = Block {
@@ -1421,8 +1419,10 @@ fn build_block(
14211419

14221420
// Compact: merge proofs sharing the same AttestationData via recursive
14231421
// aggregation so each AttestationData appears at most once (leanSpec #510).
1424-
let (aggregated_attestations, aggregated_signatures) =
1425-
compact_attestations(aggregated_attestations, aggregated_signatures, head_state)?;
1422+
let compacted = compact_attestations(selected, head_state)?;
1423+
1424+
let (aggregated_attestations, aggregated_signatures): (Vec<_>, Vec<_>) =
1425+
compacted.into_iter().unzip();
14261426

14271427
// Build final block
14281428
let attestations: AggregatedAttestations = aggregated_attestations
@@ -1891,28 +1891,28 @@ mod tests {
18911891
let bits_a = make_bits(&[0]);
18921892
let bits_b = make_bits(&[1]);
18931893

1894-
let atts = vec![
1895-
AggregatedAttestation {
1896-
aggregation_bits: bits_a.clone(),
1897-
data: data_a.clone(),
1898-
},
1899-
AggregatedAttestation {
1900-
aggregation_bits: bits_b.clone(),
1901-
data: data_b.clone(),
1902-
},
1903-
];
1904-
let proofs = vec![
1905-
AggregatedSignatureProof::empty(bits_a),
1906-
AggregatedSignatureProof::empty(bits_b),
1894+
let entries = vec![
1895+
(
1896+
AggregatedAttestation {
1897+
aggregation_bits: bits_a.clone(),
1898+
data: data_a.clone(),
1899+
},
1900+
AggregatedSignatureProof::empty(bits_a),
1901+
),
1902+
(
1903+
AggregatedAttestation {
1904+
aggregation_bits: bits_b.clone(),
1905+
data: data_b.clone(),
1906+
},
1907+
AggregatedSignatureProof::empty(bits_b),
1908+
),
19071909
];
19081910

19091911
let state = State::from_genesis(1000, vec![]);
1910-
let (out_atts, out_proofs) =
1911-
compact_attestations(atts.clone(), proofs.clone(), &state).unwrap();
1912-
assert_eq!(out_atts.len(), 2);
1913-
assert_eq!(out_proofs.len(), 2);
1914-
assert_eq!(out_atts[0].data, data_a);
1915-
assert_eq!(out_atts[1].data, data_b);
1912+
let out = compact_attestations(entries, &state).unwrap();
1913+
assert_eq!(out.len(), 2);
1914+
assert_eq!(out[0].0.data, data_a);
1915+
assert_eq!(out[1].0.data, data_b);
19161916
}
19171917

19181918
#[test]
@@ -1925,32 +1925,36 @@ mod tests {
19251925
let bits_1 = make_bits(&[1]);
19261926
let bits_2 = make_bits(&[2]);
19271927

1928-
let atts = vec![
1929-
AggregatedAttestation {
1930-
aggregation_bits: bits_0.clone(),
1931-
data: data_a.clone(),
1932-
},
1933-
AggregatedAttestation {
1934-
aggregation_bits: bits_1.clone(),
1935-
data: data_b.clone(),
1936-
},
1937-
AggregatedAttestation {
1938-
aggregation_bits: bits_2.clone(),
1939-
data: data_c.clone(),
1940-
},
1941-
];
1942-
let proofs = vec![
1943-
AggregatedSignatureProof::empty(bits_0),
1944-
AggregatedSignatureProof::empty(bits_1),
1945-
AggregatedSignatureProof::empty(bits_2),
1928+
let entries = vec![
1929+
(
1930+
AggregatedAttestation {
1931+
aggregation_bits: bits_0.clone(),
1932+
data: data_a.clone(),
1933+
},
1934+
AggregatedSignatureProof::empty(bits_0),
1935+
),
1936+
(
1937+
AggregatedAttestation {
1938+
aggregation_bits: bits_1.clone(),
1939+
data: data_b.clone(),
1940+
},
1941+
AggregatedSignatureProof::empty(bits_1),
1942+
),
1943+
(
1944+
AggregatedAttestation {
1945+
aggregation_bits: bits_2.clone(),
1946+
data: data_c.clone(),
1947+
},
1948+
AggregatedSignatureProof::empty(bits_2),
1949+
),
19461950
];
19471951

19481952
let state = State::from_genesis(1000, vec![]);
1949-
let (out_atts, _) = compact_attestations(atts, proofs, &state).unwrap();
1950-
assert_eq!(out_atts.len(), 3);
1951-
assert_eq!(out_atts[0].data, data_a);
1952-
assert_eq!(out_atts[1].data, data_b);
1953-
assert_eq!(out_atts[2].data, data_c);
1953+
let out = compact_attestations(entries, &state).unwrap();
1954+
assert_eq!(out.len(), 3);
1955+
assert_eq!(out[0].0.data, data_a);
1956+
assert_eq!(out[1].0.data, data_b);
1957+
assert_eq!(out[2].0.data, data_c);
19541958
}
19551959

19561960
#[test]
@@ -2036,4 +2040,65 @@ mod tests {
20362040
"Expected DuplicateAttestationData, got: {result:?}"
20372041
);
20382042
}
2043+
2044+
/// A partially-overlapping proof is still selected as long as it adds at
2045+
/// least one previously-uncovered validator. The greedy prefers the
2046+
/// largest proof first, then picks additional proofs whose coverage
2047+
/// extends `covered`. The resulting overlap is handled downstream by
2048+
/// `aggregate_proofs` → `xmss_aggregate` (which tracks duplicate pubkeys
2049+
/// across children via its `dup_pub_keys` machinery).
2050+
#[test]
2051+
fn extend_proofs_greedily_allows_overlap_when_it_adds_coverage() {
2052+
let data = make_att_data(1);
2053+
2054+
// Distinct sizes to avoid tie-breaking ambiguity (HashSet iteration
2055+
// order differs between debug/release):
2056+
// A = {0, 1, 2, 3} (4 validators — largest, picked first)
2057+
// B = {2, 3, 4} (overlaps A on {2,3} but adds validator 4)
2058+
// C = {1, 2} (subset of A — adds nothing, must be skipped)
2059+
let proof_a = AggregatedSignatureProof::empty(make_bits(&[0, 1, 2, 3]));
2060+
let proof_b = AggregatedSignatureProof::empty(make_bits(&[2, 3, 4]));
2061+
let proof_c = AggregatedSignatureProof::empty(make_bits(&[1, 2]));
2062+
2063+
let mut selected = Vec::new();
2064+
extend_proofs_greedily(&[proof_a, proof_b, proof_c], &mut selected, &data);
2065+
2066+
assert_eq!(
2067+
selected.len(),
2068+
2,
2069+
"A and B selected (B adds validator 4); C adds nothing and is skipped"
2070+
);
2071+
2072+
let covered: HashSet<u64> = selected
2073+
.iter()
2074+
.flat_map(|(_, p)| p.participant_indices())
2075+
.collect();
2076+
assert_eq!(covered, HashSet::from([0, 1, 2, 3, 4]));
2077+
2078+
// Attestation bits mirror the proof's participants for each entry.
2079+
for (att, proof) in &selected {
2080+
assert_eq!(att.aggregation_bits, proof.participants);
2081+
assert_eq!(att.data, data);
2082+
}
2083+
}
2084+
2085+
/// When no proof contributes new coverage (subset of a previously selected
2086+
/// proof), greedy terminates without selecting it.
2087+
#[test]
2088+
fn extend_proofs_greedily_stops_when_no_new_coverage() {
2089+
let data = make_att_data(1);
2090+
2091+
// B's participants are a subset of A's. After picking A, B offers zero
2092+
// new coverage and must not be selected (its inclusion would also
2093+
// violate the disjoint invariant).
2094+
let proof_a = AggregatedSignatureProof::empty(make_bits(&[0, 1, 2, 3]));
2095+
let proof_b = AggregatedSignatureProof::empty(make_bits(&[1, 2]));
2096+
2097+
let mut selected = Vec::new();
2098+
extend_proofs_greedily(&[proof_a, proof_b], &mut selected, &data);
2099+
2100+
assert_eq!(selected.len(), 1);
2101+
let covered: HashSet<u64> = selected[0].1.participant_indices().collect();
2102+
assert_eq!(covered, HashSet::from([0, 1, 2, 3]));
2103+
}
20392104
}

crates/common/crypto/src/lib.rs

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -144,8 +144,14 @@ pub fn aggregate_mixed(
144144

145145
ensure_prover_ready();
146146

147-
let deserialized = deserialize_children(children)?;
148-
let children_refs = to_children_refs(&deserialized);
147+
// Split deserialized children into parallel Vecs so we can borrow pubkey
148+
// slices (required by xmss_aggregate's tuple type) while moving the large
149+
// AggregatedXMSS values into the children list without cloning. `pks_list`
150+
// must outlive `children_refs`.
151+
let (pks_list, aggs): (Vec<Vec<LeanSigPubKey>>, Vec<AggregatedXMSS>) =
152+
deserialize_children(children)?.into_iter().unzip();
153+
let children_refs: Vec<(&[LeanSigPubKey], AggregatedXMSS)> =
154+
pks_list.iter().map(Vec::as_slice).zip(aggs).collect();
149155

150156
let raw_xmss: Vec<(LeanSigPubKey, LeanSigSignature)> = raw_public_keys
151157
.into_iter()
@@ -178,8 +184,11 @@ pub fn aggregate_proofs(
178184

179185
ensure_prover_ready();
180186

181-
let deserialized = deserialize_children(children)?;
182-
let children_refs = to_children_refs(&deserialized);
187+
// See `aggregate_mixed` for why this unzip-and-rezip dance is needed.
188+
let (pks_list, aggs): (Vec<Vec<LeanSigPubKey>>, Vec<AggregatedXMSS>) =
189+
deserialize_children(children)?.into_iter().unzip();
190+
let children_refs: Vec<(&[LeanSigPubKey], AggregatedXMSS)> =
191+
pks_list.iter().map(Vec::as_slice).zip(aggs).collect();
183192

184193
let (_sorted_pubkeys, aggregate) = xmss_aggregate(&children_refs, vec![], &message.0, slot, 2);
185194

@@ -204,16 +213,6 @@ fn deserialize_children(
204213
.collect()
205214
}
206215

207-
/// Build the reference slice that `xmss_aggregate` expects.
208-
fn to_children_refs(
209-
deserialized: &[(Vec<LeanSigPubKey>, AggregatedXMSS)],
210-
) -> Vec<(&[LeanSigPubKey], AggregatedXMSS)> {
211-
deserialized
212-
.iter()
213-
.map(|(pks, agg)| (pks.as_slice(), agg.clone()))
214-
.collect()
215-
}
216-
217216
/// Serialize an `AggregatedXMSS` into the `ByteListMiB` wire format.
218217
fn serialize_aggregate(aggregate: AggregatedXMSS) -> Result<ByteListMiB, AggregationError> {
219218
let serialized = aggregate.serialize();

0 commit comments

Comments
 (0)