Source code
Revision control
Copy as Markdown
Other Tools
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// Copyright by contributors to this project.
// SPDX-License-Identifier: (Apache-2.0 OR MIT)
use alloc::vec::Vec;
use core::{
hash::Hash,
ops::{Deref, DerefMut},
};
use map_impl::SmallMapInner;
pub use map_impl::{LargeMap, LargeMapEntry, SmallMap};
use mls_rs_codec::{MlsDecode, MlsEncode, MlsSize};
#[cfg(feature = "std")]
mod map_impl {
use core::hash::Hash;
use std::collections::{hash_map::Entry, HashMap};
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct SmallMap<K: Hash + Eq, V>(pub(super) HashMap<K, V>);
pub type LargeMap<K, V> = SmallMap<K, V>;
pub(super) type SmallMapInner<K, V> = HashMap<K, V>;
pub type LargeMapEntry<'a, K, V> = Entry<'a, K, V>;
}
#[cfg(not(feature = "std"))]
mod map_impl {
use core::hash::Hash;
use alloc::{
collections::{btree_map::Entry, BTreeMap},
vec::Vec,
};
#[cfg(feature = "by_ref_proposal")]
use itertools::Itertools;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SmallMap<K: Hash + Eq, V>(pub(super) Vec<(K, V)>);
pub type LargeMap<K, V> = BTreeMap<K, V>;
pub(super) type SmallMapInner<K, V> = Vec<(K, V)>;
pub type LargeMapEntry<'a, K, V> = Entry<'a, K, V>;
#[cfg(feature = "by_ref_proposal")]
impl<K: Hash + Eq, V> SmallMap<K, V> {
pub fn get(&self, key: &K) -> Option<&V> {
self.find(key).map(|i| &self.0[i].1)
}
pub fn insert(&mut self, key: K, value: V) {
match self.0.iter_mut().find(|(k, _)| (k == &key)) {
Some((_, v)) => *v = value,
None => self.0.push((key, value)),
}
}
pub fn remove(&mut self, key: &K) -> Option<V> {
self.find(key).map(|i| self.0.remove(i).1)
}
fn find(&self, key: &K) -> Option<usize> {
self.0.iter().position(|(k, _)| k == key)
}
}
}
impl<K: Hash + Eq, V> Default for SmallMap<K, V> {
fn default() -> Self {
Self(SmallMapInner::new())
}
}
impl<K: Hash + Eq, V> Deref for SmallMap<K, V> {
type Target = SmallMapInner<K, V>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<K: Hash + Eq, V> DerefMut for SmallMap<K, V> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl<K, V> MlsDecode for SmallMap<K, V>
where
K: Hash + Eq + MlsEncode + MlsDecode + MlsSize,
V: MlsEncode + MlsDecode + MlsSize,
{
fn mls_decode(reader: &mut &[u8]) -> Result<Self, mls_rs_codec::Error> {
SmallMapInner::mls_decode(reader).map(Self)
}
}
impl<K, V> MlsSize for SmallMap<K, V>
where
K: Hash + Eq + MlsEncode + MlsDecode + MlsSize,
V: MlsEncode + MlsDecode + MlsSize,
{
fn mls_encoded_len(&self) -> usize {
self.0.mls_encoded_len()
}
}
impl<K, V> MlsEncode for SmallMap<K, V>
where
K: Hash + Eq + MlsEncode + MlsDecode + MlsSize,
V: MlsEncode + MlsDecode + MlsSize,
{
fn mls_encode(&self, writer: &mut Vec<u8>) -> Result<(), mls_rs_codec::Error> {
self.0.mls_encode(writer)
}
}