Skip to content

Commit 1d7cd42

Browse files
committed
shardtree: make construction of fully-empty Parent nodes an error
Ideally the errors here would all be panics, as construction of such a node represents a programming error; however, it is preferable to return extra contextual information about the circumstance that led to this error where possible, so we model this as an error where possible without altering the public API.
1 parent 561a6dc commit 1d7cd42

File tree

4 files changed

+254
-171
lines changed

4 files changed

+254
-171
lines changed

shardtree/src/batch.rs

Lines changed: 26 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use tracing::trace;
88
use crate::{
99
error::{InsertionError, ShardTreeError},
1010
store::{Checkpoint, ShardStore},
11-
IncompleteAt, LocatedPrunableTree, LocatedTree, RetentionFlags, ShardTree, Tree,
11+
IncompleteAt, LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags, ShardTree, Tree,
1212
};
1313

1414
impl<
@@ -215,7 +215,8 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
215215
// fragments up the stack until we get the largest possible subtree
216216
while let Some((potential_sibling, marked)) = fragments.pop() {
217217
if potential_sibling.root_addr.parent() == subtree.root_addr.parent() {
218-
subtree = unite(potential_sibling, subtree, prune_below);
218+
subtree = unite(potential_sibling, subtree, prune_below)
219+
.expect("subtree is non-Nil");
219220
} else {
220221
// this is not a sibling node, so we push it back on to the stack
221222
// and are done
@@ -238,7 +239,9 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
238239
let minimal_tree_addr =
239240
Address::from(position_range.start).common_ancestor(&last_position.into());
240241
trace!("Building minimal tree at {:?}", minimal_tree_addr);
241-
build_minimal_tree(fragments, minimal_tree_addr, prune_below).map(
242+
build_minimal_tree(fragments, minimal_tree_addr, prune_below)
243+
.expect("construction of minimal tree does not result in creation of invalid parent nodes")
244+
.map(
242245
|(to_insert, contains_marked, incomplete)| BatchInsertionResult {
243246
subtree: to_insert,
244247
contains_marked,
@@ -263,16 +266,18 @@ fn unite<H: Hashable + Clone + PartialEq>(
263266
lroot: LocatedPrunableTree<H>,
264267
rroot: LocatedPrunableTree<H>,
265268
prune_below: Level,
266-
) -> LocatedPrunableTree<H> {
269+
) -> Result<LocatedPrunableTree<H>, Address> {
267270
assert_eq!(lroot.root_addr.parent(), rroot.root_addr.parent());
268-
LocatedTree {
271+
Ok(LocatedTree {
269272
root_addr: lroot.root_addr.parent(),
270273
root: if lroot.root_addr.level() < prune_below {
271-
Tree::unite(lroot.root_addr.level(), None, lroot.root, rroot.root)
274+
PrunableTree::unite(lroot.root_addr.level(), None, lroot.root, rroot.root)
275+
.map_err(|_| lroot.root_addr)?
272276
} else {
273-
Tree::parent(None, lroot.root, rroot.root)
277+
PrunableTree::parent_checked(None, lroot.root, rroot.root)
278+
.map_err(|_| lroot.root_addr)?
274279
},
275-
}
280+
})
276281
}
277282

278283
/// Combines the given subtree with an empty sibling node to obtain the next level
@@ -286,7 +291,7 @@ fn combine_with_empty<H: Hashable + Clone + PartialEq>(
286291
incomplete: &mut Vec<IncompleteAt>,
287292
contains_marked: bool,
288293
prune_below: Level,
289-
) -> LocatedPrunableTree<H> {
294+
) -> Result<LocatedPrunableTree<H>, Address> {
290295
assert_eq!(expect_left_child, root.root_addr.is_left_child());
291296
let sibling_addr = root.root_addr.sibling();
292297
incomplete.push(IncompleteAt {
@@ -313,27 +318,27 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
313318
mut xs: Vec<(LocatedPrunableTree<H>, bool)>,
314319
root_addr: Address,
315320
prune_below: Level,
316-
) -> Option<(LocatedPrunableTree<H>, bool, Vec<IncompleteAt>)> {
321+
) -> Result<Option<(LocatedPrunableTree<H>, bool, Vec<IncompleteAt>)>, Address> {
317322
// First, consume the stack from the right, building up a single tree
318323
// until we can't combine any more.
319324
if let Some((mut cur, mut contains_marked)) = xs.pop() {
320325
let mut incomplete = vec![];
321326
while let Some((top, top_marked)) = xs.pop() {
322327
while cur.root_addr.level() < top.root_addr.level() {
323-
cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below);
328+
cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below)?;
324329
}
325330

326331
if cur.root_addr.level() == top.root_addr.level() {
327332
contains_marked = contains_marked || top_marked;
328333
if cur.root_addr.is_right_child() {
329334
// We have a left child and a right child, so unite them.
330-
cur = unite(top, cur, prune_below);
335+
cur = unite(top, cur, prune_below)?;
331336
} else {
332337
// This is a left child, so we build it up one more level and then
333338
// we've merged as much as we can from the right and need to work from
334339
// the left
335340
xs.push((top, top_marked));
336-
cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below);
341+
cur = combine_with_empty(cur, true, &mut incomplete, top_marked, prune_below)?;
337342
break;
338343
}
339344
} else {
@@ -346,15 +351,15 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
346351

347352
// Ensure we can work from the left in a single pass by making this right-most subtree
348353
while cur.root_addr.level() + 1 < root_addr.level() {
349-
cur = combine_with_empty(cur, true, &mut incomplete, contains_marked, prune_below);
354+
cur = combine_with_empty(cur, true, &mut incomplete, contains_marked, prune_below)?;
350355
}
351356

352357
// push our accumulated max-height right hand node back on to the stack.
353358
xs.push((cur, contains_marked));
354359

355360
// From the stack of subtrees, construct a single sparse tree that can be
356361
// inserted/merged into the existing tree
357-
let res_tree = xs.into_iter().fold(
362+
let res_tree = xs.into_iter().try_fold(
358363
None,
359364
|acc: Option<LocatedPrunableTree<H>>, (next_tree, next_marked)| {
360365
if let Some(mut prev_tree) = acc {
@@ -368,19 +373,19 @@ fn build_minimal_tree<H: Hashable + Clone + PartialEq>(
368373
&mut incomplete,
369374
next_marked,
370375
prune_below,
371-
);
376+
)?;
372377
}
373378

374-
Some(unite(prev_tree, next_tree, prune_below))
379+
Ok::<_, Address>(Some(unite(prev_tree, next_tree, prune_below)?))
375380
} else {
376-
Some(next_tree)
381+
Ok::<_, Address>(Some(next_tree))
377382
}
378383
},
379-
);
384+
)?;
380385

381-
res_tree.map(|t| (t, contains_marked, incomplete))
386+
Ok(res_tree.map(|t| (t, contains_marked, incomplete)))
382387
} else {
383-
None
388+
Ok(None)
384389
}
385390
}
386391

shardtree/src/legacy.rs

Lines changed: 49 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
101101
leaf_addr: Address,
102102
mut filled: impl Iterator<Item = H>,
103103
split_at: Level,
104-
) -> (Self, Option<Self>) {
104+
) -> Result<(Self, Option<Self>), Address> {
105105
// add filled nodes to the subtree; here, we do not need to worry about
106106
// whether or not these nodes can be invalidated by a rewind
107107
let mut addr = leaf_addr;
@@ -113,17 +113,19 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
113113
if let Some(right) = filled.next() {
114114
// once we have a right-hand node, add a parent with the current tree
115115
// as the left-hand sibling
116-
subtree = Tree::parent(
116+
subtree = PrunableTree::parent_checked(
117117
None,
118118
subtree,
119119
Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)),
120-
);
120+
)
121+
.map_err(|_| addr)?;
121122
} else {
122123
break;
123124
}
124125
} else {
125126
// the current address is for a right child, so add an empty left sibling
126-
subtree = Tree::parent(None, Tree::empty(), subtree);
127+
subtree =
128+
PrunableTree::parent_checked(None, Tree::empty(), subtree).map_err(|_| addr)?;
127129
}
128130

129131
addr = addr.parent();
@@ -140,16 +142,22 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
140142
for right in filled {
141143
// build up the right-biased tree until we get a left-hand node
142144
while addr.is_right_child() {
143-
supertree = supertree.map(|t| Tree::parent(None, Tree::empty(), t));
145+
supertree = supertree
146+
.map(|t| PrunableTree::parent_checked(None, Tree::empty(), t))
147+
.transpose()
148+
.map_err(|_| addr)?;
144149
addr = addr.parent();
145150
}
146151

147152
// once we have a left-hand root, add a parent with the current ommer as the right-hand sibling
148-
supertree = Some(Tree::parent(
149-
None,
150-
supertree.unwrap_or_else(PrunableTree::empty),
151-
Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)),
152-
));
153+
supertree = Some(
154+
PrunableTree::parent_checked(
155+
None,
156+
supertree.unwrap_or_else(PrunableTree::empty),
157+
Tree::leaf((right.clone(), RetentionFlags::EPHEMERAL)),
158+
)
159+
.map_err(|_| addr)?,
160+
);
153161
addr = addr.parent();
154162
}
155163

@@ -161,7 +169,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
161169
None
162170
};
163171

164-
(subtree, supertree)
172+
Ok((subtree, supertree))
165173
}
166174

167175
/// Insert the nodes belonging to the given incremental witness to this tree, truncating the
@@ -196,23 +204,30 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
196204
Address::from(witness.witnessed_position()),
197205
witness.filled().iter().cloned(),
198206
self.root_addr.level(),
199-
);
207+
)
208+
.map_err(InsertionError::InputMalformed)?;
200209

201210
// construct subtrees from the `cursor` part of the witness
202-
let cursor_trees = witness.cursor().as_ref().filter(|c| c.size() > 0).map(|c| {
203-
Self::from_frontier_parts(
204-
witness.tip_position(),
205-
c.leaf()
206-
.cloned()
207-
.expect("Cannot have an empty leaf for a non-empty tree"),
208-
c.ommers_iter().cloned(),
209-
&Retention::Checkpoint {
210-
id: checkpoint_id,
211-
marking: Marking::None,
212-
},
213-
self.root_addr.level(),
214-
)
215-
});
211+
let cursor_trees = witness
212+
.cursor()
213+
.as_ref()
214+
.filter(|c| c.size() > 0)
215+
.map(|c| {
216+
Self::from_frontier_parts(
217+
witness.tip_position(),
218+
c.leaf()
219+
.cloned()
220+
.expect("Cannot have an empty leaf for a non-empty tree"),
221+
c.ommers_iter().cloned(),
222+
&Retention::Checkpoint {
223+
id: checkpoint_id,
224+
marking: Marking::None,
225+
},
226+
self.root_addr.level(),
227+
)
228+
})
229+
.transpose()
230+
.map_err(InsertionError::InputMalformed)?;
216231

217232
let (subtree, _) = past_subtree.insert_subtree(future_subtree, true)?;
218233

@@ -253,7 +268,7 @@ mod tests {
253268
frontier::CommitmentTree, witness::IncrementalWitness, Address, Level, Position,
254269
};
255270

256-
use crate::{LocatedPrunableTree, RetentionFlags, Tree};
271+
use crate::{LocatedPrunableTree, PrunableTree, RetentionFlags, Tree};
257272

258273
#[test]
259274
fn insert_witness_nodes() {
@@ -280,19 +295,22 @@ mod tests {
280295

281296
assert_eq!(
282297
c.root,
283-
Tree::parent(
298+
PrunableTree::parent_checked(
284299
None,
285-
Tree::parent(
300+
PrunableTree::parent_checked(
286301
None,
287302
Tree::empty(),
288303
Tree::leaf(("ijklmnop".to_string(), RetentionFlags::EPHEMERAL)),
289-
),
290-
Tree::parent(
304+
)
305+
.unwrap(),
306+
PrunableTree::parent_checked(
291307
None,
292308
Tree::leaf(("qrstuvwx".to_string(), RetentionFlags::EPHEMERAL)),
293309
Tree::empty()
294310
)
311+
.unwrap()
295312
)
313+
.unwrap()
296314
);
297315

298316
assert_eq!(

shardtree/src/lib.rs

Lines changed: 24 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -414,44 +414,44 @@ impl<
414414
fn go<H: Hashable + Clone + PartialEq>(
415415
root_addr: Address,
416416
root: &PrunableTree<H>,
417-
) -> Option<(PrunableTree<H>, Position)> {
417+
) -> Result<Option<(PrunableTree<H>, Position)>, Address> {
418418
match &root.0 {
419419
Node::Parent { ann, left, right } => {
420420
let (l_addr, r_addr) = root_addr
421421
.children()
422422
.expect("has children because we checked `root` is a parent");
423-
go(r_addr, right).map_or_else(
423+
go(r_addr, right)?.map_or_else(
424424
|| {
425-
go(l_addr, left).map(|(new_left, pos)| {
426-
(
427-
Tree::unite(
425+
go(l_addr, left)?
426+
.map(|(new_left, pos)| {
427+
PrunableTree::unite(
428428
l_addr.level(),
429429
ann.clone(),
430430
new_left,
431431
Tree::empty(),
432-
),
433-
pos,
434-
)
435-
})
432+
)
433+
.map_err(|_| l_addr)
434+
.map(|t| (t, pos))
435+
})
436+
.transpose()
436437
},
437438
|(new_right, pos)| {
438-
Some((
439-
Tree::unite(
440-
l_addr.level(),
441-
ann.clone(),
442-
left.as_ref().clone(),
443-
new_right,
444-
),
445-
pos,
446-
))
439+
Tree::unite(
440+
l_addr.level(),
441+
ann.clone(),
442+
left.as_ref().clone(),
443+
new_right,
444+
)
445+
.map_err(|_| l_addr)
446+
.map(|t| Some((t, pos)))
447447
},
448448
)
449449
}
450-
Node::Leaf { value: (h, r) } => Some((
450+
Node::Leaf { value: (h, r) } => Ok(Some((
451451
Tree::leaf((h.clone(), *r | RetentionFlags::CHECKPOINT)),
452452
root_addr.max_position(),
453-
)),
454-
Node::Nil => None,
453+
))),
454+
Node::Nil => Ok(None),
455455
}
456456
}
457457

@@ -469,7 +469,9 @@ impl<
469469
// Update the rightmost subtree to add the `CHECKPOINT` flag to the right-most leaf (which
470470
// need not be a level-0 leaf; it's fine to rewind to a pruned state).
471471
if let Some(subtree) = self.store.last_shard().map_err(ShardTreeError::Storage)? {
472-
if let Some((replacement, pos)) = go(subtree.root_addr, &subtree.root) {
472+
if let Some((replacement, pos)) = go(subtree.root_addr, &subtree.root)
473+
.map_err(|addr| ShardTreeError::Insert(InsertionError::InputMalformed(addr)))?
474+
{
473475
self.store
474476
.put_shard(LocatedTree {
475477
root_addr: subtree.root_addr,

0 commit comments

Comments
 (0)