Skip to content

Commit 90452a1

Browse files
PoulavBhowmick03PoulavBhowmick03
committed
added helper struct for handling pending updates
1 parent ae3e13a commit 90452a1

File tree

5 files changed

+98
-37
lines changed

5 files changed

+98
-37
lines changed

‎src/iter_arc.rs

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1+
use tree_hash::{TreeHash, TreeHashType};
12
use triomphe::Arc;
23

34
use crate::{
4-
Leaf, Tree, Value,
5+
Error, Leaf, Tree, UpdateMap, Value,
56
utils::{Length, opt_packing_depth},
67
};
78

@@ -20,17 +21,25 @@ pub struct ArcIter<'a, T: Value> {
2021
}
2122

2223
impl<'a, T: Value> ArcIter<'a, T> {
23-
pub fn from_index(index: usize, root: &'a Tree<T>, depth: usize, length: Length) -> Self {
24+
pub fn from_index(
25+
index: usize,
26+
root: &'a Tree<T>,
27+
depth: usize,
28+
length: Length,
29+
) -> Result<Self, Error> {
30+
if <T as TreeHash>::tree_hash_type() == TreeHashType::Basic {
31+
return Err(Error::PackedLeavesNoArc);
32+
}
2433
let mut stack = Vec::with_capacity(depth);
2534
stack.push(root);
2635

27-
ArcIter {
36+
Ok(ArcIter {
2837
stack,
2938
index,
3039
full_depth: depth,
3140
packing_depth: opt_packing_depth::<T>().unwrap_or(0),
3241
length,
33-
}
42+
})
3443
}
3544
}
3645

@@ -99,3 +108,50 @@ impl<'a, T: Value> Iterator for ArcIter<'a, T> {
99108
}
100109

101110
impl<T: Value> ExactSizeIterator for ArcIter<'_, T> {}
111+
#[derive(Debug)]
112+
pub struct ArcInterfaceIter<'a, T: Value, U: UpdateMap<T>> {
113+
tree_iter: ArcIter<'a, T>,
114+
updates: &'a U,
115+
index: usize,
116+
length: usize,
117+
}
118+
119+
impl<'a, T: Value, U: UpdateMap<T>> ArcInterfaceIter<'a, T, U> {
120+
pub fn new(root: &'a Tree<T>, depth: usize, length: Length, updates: &'a U) -> Self {
121+
ArcInterfaceIter {
122+
tree_iter: ArcIter::new(root, depth, length),
123+
updates,
124+
index: 0,
125+
length: length.as_usize(),
126+
}
127+
}
128+
}
129+
130+
impl<'a, T: Value, U: UpdateMap<T>> Iterator for ArcInterfaceIter<'a, T, U> {
131+
type Item = Arc<T>;
132+
133+
fn next(&mut self) -> Option<Self::Item> {
134+
if self.index >= self.length {
135+
return None;
136+
}
137+
let idx = self.index;
138+
self.index += 1;
139+
140+
let backing = self.tree_iter.next();
141+
if let Some(new_val) = self.updates.get(idx) {
142+
Some(
143+
self.updates
144+
.get_arc(idx)
145+
.unwrap_or_else(|| Arc::new(new_val.clone())),
146+
)
147+
} else {
148+
backing.cloned()
149+
}
150+
}
151+
152+
fn size_hint(&self) -> (usize, Option<usize>) {
153+
let rem = self.length.saturating_sub(self.index);
154+
(rem, Some(rem))
155+
}
156+
}
157+
impl<T: Value, U: UpdateMap<T>> ExactSizeIterator for ArcInterfaceIter<'_, T, U> {}

‎src/list.rs

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::builder::Builder;
22
use crate::interface::{ImmList, Interface, MutList};
33
use crate::interface_iter::{InterfaceIter, InterfaceIterCow};
44
use crate::iter::Iter;
5-
use crate::iter_arc::ArcIter;
5+
use crate::iter_arc::{ArcInterfaceIter, ArcIter};
66
use crate::level_iter::{LevelIter, LevelNode};
77
use crate::serde::ListVisitor;
88
use crate::tree::{IntraRebaseAction, RebaseAction};
@@ -137,16 +137,12 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
137137
Ok(self.interface.iter_from(index))
138138
}
139139

140-
pub fn iter_arc(&self) -> Result<impl Iterator<Item = &Arc<T>>, Error> {
141-
if T::tree_hash_type() == TreeHashType::Basic {
142-
// Can't return `Arc`s for packed leaves.
143-
return Err(Error::PackedLeavesNoArc);
144-
}
145-
146-
Ok(ArcIter::new(
140+
pub fn iter_arc(&self) -> Result<impl Iterator<Item = Arc<T>>, Error> {
141+
Ok(ArcInterfaceIter::new(
147142
&self.interface.backing.tree,
148143
self.interface.backing.depth,
149-
self.interface.backing.length,
144+
Length(self.len()),
145+
&self.interface.updates,
150146
))
151147
}
152148

@@ -283,12 +279,7 @@ impl<T: Value, N: Unsigned> ImmList<T> for ListInner<T, N> {
283279
}
284280

285281
fn iter_arc(&self, index: usize) -> Result<ArcIter<'_, T>, Error> {
286-
Ok(ArcIter::from_index(
287-
index,
288-
&self.tree,
289-
self.depth,
290-
self.length,
291-
))
282+
ArcIter::from_index(index, &self.tree, self.depth, self.length)
292283
}
293284
}
294285

‎src/tests/proptest/operations.rs

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -215,12 +215,11 @@ where
215215
}
216216
Op::IterArc => {
217217
if <T as TreeHash>::tree_hash_type() != TreeHashType::Basic {
218-
list.apply_updates().unwrap();
218+
let actual: Vec<T> =
219+
list.iter_arc().unwrap().map(|arc| (*arc).clone()).collect();
220+
let expected: Vec<T> = spec.iter().cloned().collect();
219221

220-
let actual = list.iter_arc().unwrap().map(|arc| arc.as_ref());
221-
let expected = spec.iter();
222-
223-
assert!(actual.eq(expected));
222+
assert_eq!(actual, expected);
224223
}
225224
}
226225

@@ -322,12 +321,11 @@ where
322321
}
323322
Op::IterArc => {
324323
if <T as TreeHash>::tree_hash_type() != TreeHashType::Basic {
325-
vect.apply_updates().unwrap();
326-
327-
let actual = vect.iter_arc().unwrap().map(|arc| arc.as_ref());
328-
let expected = spec.iter();
324+
let actual: Vec<T> =
325+
vect.iter_arc().unwrap().map(|arc| (*arc).clone()).collect();
326+
let expected: Vec<T> = spec.iter().cloned().collect();
329327

330-
assert!(actual.eq(expected));
328+
assert!(actual.eq(&expected));
331329
}
332330
}
333331

‎src/update_map.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use crate::cow::{BTreeCow, Cow, VecCow};
22
use crate::utils::max_btree_index;
33
use std::collections::{BTreeMap, btree_map::Entry};
44
use std::ops::ControlFlow;
5+
use triomphe::Arc;
56
use vec_map::VecMap;
67

78
/// Trait for map types which can be used to store intermediate updates before application
@@ -18,6 +19,8 @@ pub trait UpdateMap<T>: Default + Clone {
1819
F: FnOnce(usize) -> Option<&'a T>,
1920
T: Clone + 'a;
2021

22+
fn get_arc(&self, k: usize) -> Option<Arc<T>>;
23+
2124
fn insert(&mut self, k: usize, value: T) -> Option<T>;
2225

2326
fn for_each_range<F, E>(&self, start: usize, end: usize, f: F) -> Result<(), E>
@@ -72,6 +75,10 @@ impl<T: Clone> UpdateMap<T> for BTreeMap<usize, T> {
7275
Some(Cow::BTree(cow))
7376
}
7477

78+
fn get_arc(&self, k: usize) -> Option<Arc<T>> {
79+
self.get(&k).cloned().map(Arc::new)
80+
}
81+
7582
fn insert(&mut self, idx: usize, value: T) -> Option<T> {
7683
BTreeMap::insert(self, idx, value)
7784
}
@@ -136,6 +143,10 @@ impl<T: Clone> UpdateMap<T> for VecMap<T> {
136143
Some(Cow::Vec(cow))
137144
}
138145

146+
fn get_arc(&self, k: usize) -> Option<Arc<T>> {
147+
self.get(k).cloned().map(Arc::new)
148+
}
149+
139150
fn insert(&mut self, idx: usize, value: T) -> Option<T> {
140151
VecMap::insert(self, idx, value)
141152
}
@@ -202,6 +213,10 @@ where
202213
self.inner.get_cow_with(k, f)
203214
}
204215

216+
fn get_arc(&self, k: usize) -> Option<Arc<T>> {
217+
self.inner.get_arc(k)
218+
}
219+
205220
fn insert(&mut self, k: usize, value: T) -> Option<T> {
206221
if k > self.max_key {
207222
self.max_key = k;

‎src/vector.rs

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crate::interface::{ImmList, Interface, MutList};
22
use crate::interface_iter::InterfaceIter;
33
use crate::iter::Iter;
4+
use crate::iter_arc::ArcInterfaceIter;
45
use crate::level_iter::LevelIter;
56
use crate::tree::{IntraRebaseAction, RebaseAction};
67
use crate::update_map::MaxMap;
@@ -78,8 +79,13 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
7879
self.interface.iter()
7980
}
8081

81-
pub fn iter_arc(&self) -> Result<ArcIter<'_, T>, Error> {
82-
self.interface.iter_arc()
82+
pub fn iter_arc(&self) -> Result<impl Iterator<Item = Arc<T>>, Error> {
83+
Ok(ArcInterfaceIter::new(
84+
&self.interface.backing.tree,
85+
self.interface.backing.depth,
86+
Length(self.len()),
87+
&self.interface.updates,
88+
))
8389
}
8490

8591
pub fn iter_from(&self, index: usize) -> Result<InterfaceIter<'_, T, U>, Error> {
@@ -230,12 +236,7 @@ impl<T: Value, N: Unsigned> ImmList<T> for VectorInner<T, N> {
230236
}
231237

232238
fn iter_arc(&self, index: usize) -> Result<ArcIter<'_, T>, Error> {
233-
Ok(ArcIter::from_index(
234-
index,
235-
&self.tree,
236-
self.depth,
237-
Length(N::to_usize()),
238-
))
239+
ArcIter::from_index(index, &self.tree, self.depth, Length(N::to_usize()))
239240
}
240241
}
241242

0 commit comments

Comments
 (0)