fix(rln): if only a single insertion, dont use batch_insert (#215)

* fix(rln): if only a single insertion, dont use batch_insert

* fix: more optimizations
This commit is contained in:
Aaryamann Challani 2023-09-14 14:10:53 +05:30 committed by GitHub
parent a86b859b75
commit f73c83b571
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 19 deletions

View File

@ -197,11 +197,13 @@ impl ZerokitMerkleTree for PmTree {
let mut indices = indices.into_iter().collect::<Vec<_>>();
indices.sort();
match (leaves.is_empty(), indices.is_empty()) {
(true, true) => Err(Report::msg("no leaves or indices to be removed")),
(false, true) => self.set_range_with_leaves(start, leaves),
(true, false) => self.remove_indices(indices),
(false, false) => self.remove_indices_and_set_leaves(start, leaves, indices),
match (leaves.len(), indices.len()) {
(0, 0) => Err(Report::msg("no leaves or indices to be removed")),
(1, 0) => self.set(start, leaves[0]),
(0, 1) => self.delete(indices[0]),
(_, 0) => self.set_range_with_leaves(start, leaves),
(0, _) => self.remove_indices(&indices),
(_, _) => self.remove_indices_and_set_leaves(start, leaves, &indices),
}
}
@ -264,18 +266,11 @@ impl PmTree {
.map_err(|e| Report::msg(e.to_string()))
}
fn remove_indices(&mut self, indices: Vec<usize>) -> Result<()> {
fn remove_indices(&mut self, indices: &[usize]) -> Result<()> {
let start = indices[0];
let end = indices.last().unwrap() + 1;
let mut new_leaves: Vec<_> = (start..end)
.map(|i| self.tree.get(i))
.collect::<Result<_, _>>()?;
new_leaves
.iter_mut()
.take(indices.len())
.for_each(|leaf| *leaf = PmTreeHasher::default_leaf());
let new_leaves = (start..end).map(|_| PmTreeHasher::default_leaf());
self.tree
.set_range(start, new_leaves)
@ -286,16 +281,13 @@ impl PmTree {
&mut self,
start: usize,
leaves: Vec<FrOfPmTreeHasher>,
indices: Vec<usize>,
indices: &[usize],
) -> Result<()> {
let min_index = *indices.first().unwrap();
let max_index = start + leaves.len();
// Generated a placeholder with the exact size needed,
// Initiated with default values to be overridden throughout the method
let mut set_values = vec![PmTreeHasher::default_leaf(); max_index - min_index];
// If the index is not in indices list, keep the original value
for i in min_index..start {
if !indices.contains(&i) {
let value = self.tree.get(i)?;
@ -303,7 +295,6 @@ impl PmTree {
}
}
// Insert new leaves after 'start' position
for (i, &leaf) in leaves.iter().enumerate() {
set_values[start - min_index + i] = leaf;
}