Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: abstract common functions in Keys and KeysRange #989

Merged
merged 3 commits into from
Dec 6, 2022
Merged

refactor: abstract common functions in Keys and KeysRange #989

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1a6d19c
Replace `TreeMap->Keys` with `TreeMap->KeysRange`
Dec 3, 2022
d2e61eb
merge master
Dec 3, 2022
1351cad
Add `Keys` back and abstract common functions
Dec 3, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
212 changes: 15 additions & 197 deletions near-sdk/src/store/tree_map/iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ where
V: BorshSerialize,
H: ToKey,
{
keys: Keys<'a, K>,
keys: KeysRange<'a, K>,
values: &'a LookupMap<K, V, H>,
}

Expand All @@ -56,7 +56,10 @@ where
H: ToKey,
{
pub(super) fn new(map: &'a TreeMap<K, V, H>) -> Self {
Self { keys: Keys::new(&map.tree), values: &map.values }
Self {
keys: KeysRange::new(&map.tree, (Bound::Unbounded, Bound::Unbounded)),
values: &map.values,
}
}
}

Expand Down Expand Up @@ -148,7 +151,7 @@ where
H: ToKey,
{
/// Values iterator which contains empty and filled cells.
keys: Keys<'a, K>,
keys: KeysRange<'a, K>,
/// Exclusive reference to underlying map to lookup values with `keys`.
values: &'a mut LookupMap<K, V, H>,
}
Expand All @@ -160,7 +163,10 @@ where
H: ToKey,
{
pub(super) fn new(map: &'a mut TreeMap<K, V, H>) -> Self {
Self { keys: Keys::new(&map.tree), values: &mut map.values }
Self {
keys: KeysRange::new(&map.tree, (Bound::Unbounded, Bound::Unbounded)),
values: &mut map.values,
}
}
}

Expand Down Expand Up @@ -294,198 +300,6 @@ where
}
}

/// An iterator over the keys of a [`TreeMap`], in sorted order.
///
/// This `struct` is created by the `keys` method on [`TreeMap`].
pub struct Keys<'a, K: 'a>
where
K: BorshSerialize + BorshDeserialize + Ord,
{
tree: &'a Tree<K>,
length: u32,
min: FindUnbounded,
max: FindUnbounded,
//The last element in the stack is the latest value returned by the iterator
stack_asc: Vec<FreeListIndex>,
stack_desc: Vec<FreeListIndex>,
}

impl<'a, K> Keys<'a, K>
where
K: BorshSerialize + BorshDeserialize + Ord,
{
pub(super) fn new(tree: &'a Tree<K>) -> Self {
Self {
tree,
length: tree.nodes.len(),
min: FindUnbounded::First,
max: FindUnbounded::First,
stack_asc: Vec::new(),
stack_desc: Vec::new(),
}
}
}

impl<'a, K> Iterator for Keys<'a, K>
where
K: BorshSerialize + BorshDeserialize + Ord,
{
type Item = &'a K;

fn next(&mut self) -> Option<&'a K> {
if self.length == 0 {
// Short circuit if all elements have been iterated.
return None;
}

let next = match self.min {
FindUnbounded::First => self.find_min(self.tree.root.as_ref()),
FindUnbounded::Next => self.find_next_asc(),
};

if next.is_some() {
// Update minimum bound.
self.min = FindUnbounded::Next;

// Decrease count of potential elements
self.length -= 1;
} else {
// No more elements to iterate, set length to 0 to avoid duplicate lookups.
// Bounds can never be updated manually once initialized, so this can be done.
self.length = 0;
}

next
}

fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.length as usize;
(len, Some(len))
}

fn count(self) -> usize {
self.length as usize
}
}

impl<'a, K> ExactSizeIterator for Keys<'a, K> where K: BorshSerialize + BorshDeserialize + Ord {}
impl<'a, K> FusedIterator for Keys<'a, K> where K: BorshSerialize + BorshDeserialize + Ord {}

impl<'a, K> DoubleEndedIterator for Keys<'a, K>
where
K: BorshSerialize + Ord + BorshDeserialize,
{
fn next_back(&mut self) -> Option<&'a K> {
if self.length == 0 {
// Short circuit if all elements have been iterated.
return None;
}

let next = match self.max {
FindUnbounded::First => self.find_max(self.tree.root.as_ref()),
FindUnbounded::Next => self.find_next_desc(),
};

if next.is_some() {
// Update maximum bound.
self.max = FindUnbounded::Next;

// Decrease count of potential elements
self.length -= 1;
} else {
// No more elements to iterate, set length to 0 to avoid duplicate lookups.
// Bounds can never be updated manually once initialized, so this can be done.
self.length = 0;
}

next
}
}

impl<'a, K> Keys<'a, K>
where
K: BorshSerialize + BorshDeserialize + Ord,
{
fn find_min(&mut self, root: Option<&FreeListIndex>) -> Option<&'a K> {
let mut curr = root;
let mut seen: Option<&K> = None;

while let Some(curr_idx) = curr {
if let Some(node) = self.tree.node(*curr_idx) {
seen = Some(&node.key);
self.stack_asc.push(*curr_idx);
curr = node.lft.as_ref();
} else {
curr = None
}
}
seen
}

//The last element in the stack is the last returned key.
//Find the next key to the last item in the stack
fn find_next_asc(&mut self) -> Option<&'a K> {
let last_key_idx = self.stack_asc.pop();
let mut seen: Option<&K> = None;
if let Some(last_idx) = last_key_idx {
if let Some(node) = self.tree.node(last_idx) {
//If the last returned key has right node then return minimum key from the
//tree where the right node is the root.
seen = match node.rgt {
Some(rgt) => self.find_min(Some(&rgt)),
None => None,
}
}
}
//If the last returned key does not have right node then return the
//last value in the stack.
if seen.is_none() && !self.stack_asc.is_empty() {
if let Some(result_idx) = self.stack_asc.last() {
seen = self.tree.node(*result_idx).map(|f| &f.key);
}
}
seen
}
fn find_max(&mut self, root: Option<&FreeListIndex>) -> Option<&'a K> {
let mut curr = root;
let mut seen: Option<&K> = None;

while let Some(curr_idx) = curr {
if let Some(node) = self.tree.node(*curr_idx) {
seen = Some(&node.key);
self.stack_desc.push(*curr_idx);
curr = node.rgt.as_ref();
} else {
curr = None
}
}
seen
}

fn find_next_desc(&mut self) -> Option<&'a K> {
let last_key_idx = self.stack_desc.pop();
let mut seen: Option<&K> = None;
if let Some(last_idx) = last_key_idx {
if let Some(node) = self.tree.node(last_idx) {
//If the last returned key has left node then return maximum key from the
//tree where the left node is the root.
seen = match node.lft {
Some(lft) => self.find_max(Some(&lft)),
None => None,
}
}
}
//If the last returned key does not have left node then return the
//last value in the stack.
if seen.is_none() && !self.stack_desc.is_empty() {
if let Some(result_idx) = self.stack_desc.last() {
seen = self.tree.node(*result_idx).map(|f| &f.key);
}
}
seen
}
}

/// An iterator over the keys of a [`TreeMap`], in sorted order.
///
/// This `struct` is created by the `keys` method on [`TreeMap`].
Expand Down Expand Up @@ -584,7 +398,11 @@ where

fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.length as usize;
(0, Some(len))
(len, Some(len))
askoa marked this conversation as resolved.
Show resolved Hide resolved
}

fn count(self) -> usize {
self.length as usize
}
}

Expand Down
8 changes: 4 additions & 4 deletions near-sdk/src/store/tree_map/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,10 @@ use crate::store::LookupMap;
use crate::{env, IntoStorageKey};
use borsh::{BorshDeserialize, BorshSerialize};
pub use entry::Entry;
pub use iter::{Iter, IterMut, Keys, Range, RangeMut, Values, ValuesMut};
pub use iter::{Iter, IterMut, KeysRange, Range, RangeMut, Values, ValuesMut};
askoa marked this conversation as resolved.
Show resolved Hide resolved
use std::borrow::Borrow;
use std::fmt;
use std::ops::RangeBounds;
use std::ops::{Bound, RangeBounds};

type NodeAndIndex<'a, K> = (FreeListIndex, &'a Node<K>);

Expand Down Expand Up @@ -825,11 +825,11 @@ where

/// An iterator visiting all keys in arbitrary order.
/// The iterator element type is `&'a K`.
pub fn keys(&self) -> Keys<K>
pub fn keys(&self) -> KeysRange<K>
where
K: BorshDeserialize,
{
Keys::new(&self.tree)
KeysRange::new(&self.tree, (Bound::Unbounded, Bound::Unbounded))
}

/// An iterator visiting all values in arbitrary order.
Expand Down