feat(iterators): allow key-only iteration

[rymnc]

Linked Issues/PRs

• Optimize iteration over the keys in the database #1997
• chore(iterators): replace iter_all with iter_all_keys where required #2076

Description

Implements iter_store_keys and iter_keys appropriately for in memory stores and rocksdb. The previous PR made them call it, and this one changes the underlying implementation to avoid value allocations.

Checklist

• Breaking changes are clearly marked as such in the PR description and changelog (no breaking changes)
• New behavior is reflected in tests (tests which test this behaviour under the hood still pass)
• The specification matches the implemented behavior (no changes to spec)

Before requesting review

• I have reviewed the code myself
• I have created follow-up issues caused by this PR and linked them here

After merging, notify other teams

[Add or remove entries as needed]

• Rust SDK
• Sway compiler
• Platform documentation (for out-of-organization contributors, the person merging the PR will do this)
• Someone else?

[rymnc]

taken mostly as-is from the iter_store fn defn, but only uses keys and uses _iter_all_keys under the hood.

I have an implementation which abstracts the functionality of using key prefixes and start index, but it seems verbose and might not provide much value.

[rymnc]

Favouring "owning" the IteratorMode because our keys are relatively small, and to avoid all the lifetime annotations in places which use this to cut cognitive load.

[xgreenx]

You could copy/paste everything from the function above and just remove value.clone()=)

[rymnc]

sure, addressed in 06bd0e9
thanks

[xgreenx]

[nit] We could avoid value.clone() if we implemented that in some another way=) Maybe we could use generics to minimize the code duplciation

[rymnc]

yeah, maybe we should change the fn signature of iter_all_keys in IteratorOverTable to include the column, prefix, start key. it is currently -

    fn iter_all_keys<M>(
        &self,
        direction: Option<IterDirection>,
    ) -> BoxedIter<super::Result<M::OwnedKey>>
    where
        M: Mappable,
        Self: IterableTable<M>,
    {
        self.iter_all_filtered_keys::<M, [u8; 0]>(None, None, direction)
    }

[rymnc]

nevermind, fixed in 1665b86

thanks!

[xgreenx]

The reverse_prefix_iter function is pretty complex on its own, and having two complex functions will be hard to support=D Maybe we could use the generic, which will define what type to return and defines the closure to get that type?

[rymnc]

totally agree, however, since the types we want to implement for, namely KVItem and KeyItem are a tuple and a distinct type, writing a generic for this becomes more involved. We will have to define a trait for both KVItem and KeyItem which provide the following helper fn -

trait StorageItem {
  type Key = Vec<u8>;

  fn key() -> Self::Key;
  fn key_starts_with(prefix: Self::Key) -> bool; 
}

which makes it easier to write the function, in the following way -

fn reverse_prefix_iter_generic<T, F>(
        &self,
        prefix: &[u8],
        column: Description::Column,
        iter_store_fn: F,
        iter_all_fn: impl Fn(
            &Self,
            Description::Column,
            ReadOptions,
            OwnedIteratorMode,
        ) -> BoxedIter<StorageResult<T>>,
    ) -> BoxedIter<StorageResult<T>>
    where
        T: 'static + StorageItem,
        F: Fn(
            &Self,
            Description::Column,
            Option<&[u8]>,
            Option<&[u8]>,
            IterDirection,
        ) -> BoxedIter<'static, StorageResult<T>>,
    {
        let maybe_next_item = next_prefix(prefix.to_vec())
            .and_then(|next_prefix| {
                iter_store_fn(
                    self,
                    column,
                    Some(next_prefix.as_slice()),
                    None,
                    IterDirection::Forward,
                )
                .next()
            })
            .and_then(|res| res.ok());

        let prefix_vec = prefix.to_vec(); 
        
        match maybe_next_item {
            Some(key_or_item) => {   
                let next_start_key_vec = key_or_item.key().to_vec();

                let iter_mode = OwnedIteratorMode::From(
                    next_start_key_vec.to_vec(),
                    rocksdb::Direction::Reverse,
                );

                iter_all_fn(self, column, self.read_options(), iter_mode)
                    .skip(1)
                    .take_while(move |item| match item {
                        Ok(key_or_item) => key_or_item.key_starts_with(&prefix_vec), 
                        Err(_) => true,
                    })
                    .into_boxed()
            }
            None => {
                iter_all_fn(self, column, self.read_options(), OwnedIteratorMode::End)
                    .take_while(move |item| match item {
                        Ok(key_or_item) => key_or_item.key_starts_with(&prefix_vec), 
                        Err(_) => true,
                    })
                    .into_boxed()
            }
        }
    }

the important line of code in the above fn being let next_start_key_vec = key_or_item.key().to_vec();

to avoid destructuring the tuple since the generic T can either be a KVItem (a tuple) or a KeyItem (Vec)

perhaps this can be addressed in a follow up PR to avoid a large diff :)

what are your thoughts?

[rymnc]

@xgreenx check out the diff in e449766, I added it there and if we want to include as a follow up pr, i can create it

[xgreenx]

I'm curious. Is it possible to return a slice instead of the vector? In most cases we decode key and don't use the vector itself

[rymnc]

not without adding additional lifetime constraints :/

[xgreenx]

Okay, we can try to play with it in follow up PR=)

[rymnc]

#2097 (comment)

[Dentosal]

After merge with master, this is in the wrong section

[rymnc]

addressed in ad29d89

[MitchTurner]

This probably doesn't need to be to_vec you can just use the slice below.

[MitchTurner]

It looks like it will be painful trying to get the lifetimes to align here. It's probable that we can just use a constant sized array in the future, rather than a vec, in which case this shouldn't be a problem.

Not really a problem now either, but worth optimizing for sure.