Skip to content

Commit

Permalink
Add mc-util-vec-map which exposes a heapless map type (mobilecoinfoun…
Browse files Browse the repository at this point in the history
…dation#2810)

* Add mc-util-vec-map which exposes a heapless map type

A heapless map type is needed for things like hardware wallets.

This was pulled out of mobilecoinfoundation#2725 to ease review.

* Update util/vec-map/src/lib.rs

Co-authored-by: Nick Santana <[email protected]>

* Update util/vec-map/README.md

Co-authored-by: Nick Santana <[email protected]>

* fixup trait bounds per suggestions

* address more review and test comments

* Update util/vec-map/src/lib.rs

Co-authored-by: Remoun Metyas <[email protected]>

* more review comments

* fix test build

Co-authored-by: Nick Santana <[email protected]>
Co-authored-by: Remoun Metyas <[email protected]>
3 people authored Nov 3, 2022

Partially verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
We cannot verify signatures from co-authors, and some of the co-authors attributed to this commit require their commits to be signed.
1 parent 3b25b55 commit 6b17c72
Showing 5 changed files with 503 additions and 3 deletions.
175 changes: 172 additions & 3 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
@@ -160,6 +160,8 @@ members = [
"util/test-helper",
"util/test-vector",
"util/uri",
"util/vec-map",
"util/zip-exact",
"wasm-test",
"watcher",
"watcher/api",
11 changes: 11 additions & 0 deletions util/vec-map/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
[package]
name = "mc-util-vec-map"
version = "2.1.0-pre0"
authors = ["MobileCoin"]
edition = "2021"
description = "A map object based on heapless Vec"
readme = "README.md"

[dependencies]
displaydoc = "0.2"
heapless = { version = "0.7", default-features = false }
16 changes: 16 additions & 0 deletions util/vec-map/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
mc-util-vec-map
===============

This is a map (container of key-value pairs) whose storage is arranged as two
`Heapless::Vec` objects, one for keys and one for values.

The motivation is somewhat connected to crates.io crates:
* https://crates.io/crates/vector-map
* https://docs.rs/vec_map/latest/vec_map/struct.VecMap.html
* https://docs.rs/vec-collections/latest/vec_collections/

However, these crates use `std::vec::Vec` rather than a `no_std` friendly object.

We are using `heapless` because we want this to be friendly for hardware wallets.

This will be much smaller on the stack than some kind of hash table.
302 changes: 302 additions & 0 deletions util/vec-map/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,302 @@
// Copyright (c) 2018-2022 The MobileCoin Foundation

//! Provides a map (key-value store) interface backed by `heapless::Vec`
#![no_std]
#![deny(missing_docs)]

use displaydoc::Display;
use heapless::Vec;

/// An error which can occur when using VecMap
#[derive(Clone, Debug, Display, Eq, PartialEq)]
pub enum Error {
/// VecMap capacity exceeded
CapacityExceeded,
}

/// This is a mini version of VecMap that is no_std compatible and uses
/// heapless::Vec instead alloc::Vec.
///
/// It may be better to patch upstream for no_std compatibility and use that,
/// but that crate has other issues -- it relies on an
/// experimental "contracts" crate that causes a dependency on rand crate.
/// Porting to Heapless would be a breaking change to the API.
///
/// TBD what the best path is: https://github.com/p-avital/vec-map-rs/blob/master/src/lib.rs
#[derive(Clone, Debug)]
pub struct VecMap<K, V, const N: usize> {
keys: Vec<K, N>,
values: Vec<V, N>,
}

impl<K, V, const N: usize> Default for VecMap<K, V, N> {
fn default() -> Self {
Self {
keys: Default::default(),
values: Default::default(),
}
}
}

impl<K, V, const N: usize> VecMap<K, V, N> {
/// Check if the map is empty
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}

/// Get the number of items in the map
#[inline]
pub fn len(&self) -> usize {
debug_assert!(self.keys.len() == self.values.len());
self.keys.len()
}
}

impl<K: Eq, V, const N: usize> VecMap<K, V, N> {
/// Get the value associated to a key, if present
pub fn get(&self, key: &K) -> Option<&V> {
self.keys
.iter()
.position(|k| k == key)
.map(|idx| &self.values[idx])
}

/// Get a mutable reference to the value associated to a key, if present
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
self.keys
.iter()
.position(|k| k == key)
.map(|idx| &mut self.values[idx])
}
}

impl<K: Clone + Eq, V, const N: usize> VecMap<K, V, N> {
/// Get a mutable reference to the value associated to a key, if present,
/// or else insert such a value produced by given callback,
/// and then return a mutable reference
///
/// Returns an error if the heapless::Vec capacity was exceeded
pub fn get_mut_or_insert_with(
&mut self,
key: &K,
val_fn: impl FnOnce() -> V,
) -> Result<&mut V, Error> {
if let Some(idx) = self.keys.iter().position(|k| k == key) {
Ok(&mut self.values[idx])
} else {
let idx = self.keys.len();
debug_assert_eq!(idx, self.values.len());
self.keys
.push(key.clone())
.map_err(|_| Error::CapacityExceeded)?;
self.values
.push(val_fn())
.map_err(|_| panic!("Length invariant violated"))?;
Ok(&mut self.values[idx])
}
}

/// Get an iterator over the pairs in the VecMap
pub fn iter(&self) -> IterVecMap<K, V, N> {
IterVecMap::new(self)
}
}

// Sorting is possible when keys are ordered, and keys and values are cloneable
impl<K: Clone + Ord, V: Clone, const N: usize> VecMap<K, V, N> {
/// Sort the key-value pairs of the VecMap
pub fn sort(&mut self) {
// First compute the order that would sort the set of keys
let mut indices: Vec<usize, N> = (0..self.keys.len()).collect();
indices.sort_by_key(|&i| &self.keys[i]);
// Make new key and val sets
let mut new_keys = Vec::<K, N>::default();
let mut new_vals = Vec::<V, N>::default();
// Push items into the new sets in appropriate order
for idx in indices {
// Safety: This is okay because indices
// has length at most n, so we are pushing at most n
// things into new_keys and new_vals.
unsafe { new_keys.push_unchecked(self.keys[idx].clone()) };
unsafe { new_vals.push_unchecked(self.values[idx].clone()) };
}
// Overwrite old sets
self.keys = new_keys;
self.values = new_vals;
}
}

/// An iterator over a VecMap
pub struct IterVecMap<'a, K: Clone + Eq, V, const N: usize> {
src: &'a VecMap<K, V, N>,
idx: usize,
}

impl<'a, K: Clone + Eq, V, const N: usize> IterVecMap<'a, K, V, N> {
fn new(src: &'a VecMap<K, V, N>) -> Self {
Self { src, idx: 0 }
}
}

impl<'a, K: Clone + Eq, V, const N: usize> Iterator for IterVecMap<'a, K, V, N> {
type Item = (&'a K, &'a V);

fn next(&mut self) -> Option<Self::Item> {
let result = self.src.keys.get(self.idx).and_then(|key_ref| {
self.src
.values
.get(self.idx)
.map(|value_ref| (key_ref, value_ref))
});
self.idx += 1;
result
}
}

#[cfg(test)]
mod tests {
use super::*;
extern crate alloc;
use alloc::vec;

#[test]
fn test_get_mut_or_insert_with() {
let mut vec_map = VecMap::<u32, u64, 4>::default();
assert_eq!(*vec_map.get_mut_or_insert_with(&3, || 5).unwrap(), 5);
assert_eq!(*vec_map.get_mut_or_insert_with(&4, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&7, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&8, || 10).unwrap(), 10);
assert_eq!(
vec_map.get_mut_or_insert_with(&9, || 11),
Err(Error::CapacityExceeded)
);
assert_eq!(vec_map.len(), 4);
assert_eq!(*vec_map.get_mut_or_insert_with(&7, || 10).unwrap(), 6);
assert_eq!(vec_map.len(), 4);
}

#[test]
fn test_get() {
let mut vec_map = VecMap::<u32, u64, 4>::default();
assert_eq!(*vec_map.get_mut_or_insert_with(&3, || 5).unwrap(), 5);
assert_eq!(*vec_map.get_mut_or_insert_with(&4, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&7, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&8, || 10).unwrap(), 10);

assert_eq!(vec_map.get(&3), Some(&5));
assert_eq!(vec_map.get(&4), Some(&6));
assert_eq!(vec_map.get(&5), None);
assert_eq!(vec_map.get(&6), None);
assert_eq!(vec_map.get(&7), Some(&6));
assert_eq!(vec_map.get(&8), Some(&10));
}

#[test]
fn test_get_mut() {
let mut vec_map = VecMap::<u32, u64, 4>::default();
assert_eq!(*vec_map.get_mut_or_insert_with(&3, || 5).unwrap(), 5);
assert_eq!(*vec_map.get_mut_or_insert_with(&4, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&7, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&8, || 10).unwrap(), 10);

assert_eq!(vec_map.get_mut(&3), Some(5).as_mut());
assert_eq!(vec_map.get_mut(&4), Some(6).as_mut());
assert_eq!(vec_map.get_mut(&5), None);
assert_eq!(vec_map.get_mut(&6), None);
assert_eq!(vec_map.get_mut(&7), Some(6).as_mut());
assert_eq!(vec_map.get_mut(&8), Some(10).as_mut());
}

#[test]
fn test_iter() {
let mut vec_map = VecMap::<u32, u64, 4>::default();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![]);

assert_eq!(*vec_map.get_mut_or_insert_with(&3, || 5).unwrap(), 5);

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&3, &5)]);

assert_eq!(*vec_map.get_mut_or_insert_with(&4, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&7, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&8, || 10).unwrap(), 10);

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&3, &5), (&4, &6), (&7, &6), (&8, &10)]);
}

#[test]
fn test_sort_empty() {
let mut vec_map = VecMap::<u32, u64, 4>::default();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![]);

vec_map.sort();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![]);
}

#[test]
fn test_sort_one() {
let mut vec_map = VecMap::<u32, u64, 4>::default();

assert_eq!(*vec_map.get_mut_or_insert_with(&9, || 5).unwrap(), 5);

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&9, &5)]);

vec_map.sort();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&9, &5)]);
}

#[test]
fn test_sort_two() {
let mut vec_map = VecMap::<u32, u64, 4>::default();

assert_eq!(*vec_map.get_mut_or_insert_with(&9, || 5).unwrap(), 5);
assert_eq!(*vec_map.get_mut_or_insert_with(&3, || 6).unwrap(), 6);

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&9, &5), (&3, &6)]);

vec_map.sort();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&3, &6), (&9, &5)]);

vec_map.sort();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&3, &6), (&9, &5)]);
}

#[test]
fn test_sort_at_capacity() {
let mut vec_map = VecMap::<u32, u64, 4>::default();
assert_eq!(*vec_map.get_mut_or_insert_with(&9, || 5).unwrap(), 5);
assert_eq!(*vec_map.get_mut_or_insert_with(&3, || 6).unwrap(), 6);
assert_eq!(*vec_map.get_mut_or_insert_with(&7, || 7).unwrap(), 7);
assert_eq!(*vec_map.get_mut_or_insert_with(&1, || 10).unwrap(), 10);

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&9, &5), (&3, &6), (&7, &7), (&1, &10)]);

vec_map.sort();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&1, &10), (&3, &6), (&7, &7), (&9, &5)]);

vec_map.sort();

let seq: alloc::vec::Vec<_> = vec_map.iter().collect();
assert_eq!(seq, vec![(&1, &10), (&3, &6), (&7, &7), (&9, &5)]);
}
}

0 comments on commit 6b17c72

Please sign in to comment.