Commit 577e9645 authored by Hero Bird's avatar Hero Bird
Browse files

[pdsl_core] Add Index & IndexMut impls for storage::HashMap

Also add storage::HashMap::get_mut.
parent a791a322
......@@ -360,8 +360,7 @@ where
.map(|slot| slot.index())
}
/// Removes a key from the map,
/// returning the value at the key if the key was previously in the map.
/// Removes a key from the map, returning the value at the key if the key was previously in the map.
///
/// # Note
///
......@@ -387,7 +386,7 @@ where
}
}
/// Returns the value corresponding to the key.
/// Returns an immutable reference to the value corresponding to the key.
///
/// The key may be any borrowed form of the map's key type,
/// but Hash and Eq on the borrowed form must match those for the key type.
......@@ -402,6 +401,21 @@ where
}
}
/// Returns a mutable reference to the value corresponding to the key.
///
/// The key may be any borrowed form of the map's key type,
/// but Hash and Eq on the borrowed form must match those for the key type.
pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
match self.entry_mut(key) {
Some(Entry::Removed) | None => None,
Some(Entry::Occupied(OccupiedEntry{val, ..})) => Some(val),
}
}
/// Returns `true` if there is an entry corresponding to the key in the map.
pub fn contains_key<Q>(&self, key: &Q) -> bool
where
......@@ -414,7 +428,7 @@ where
}
}
/// Returns the entry corresponding to the key.
/// Returns an immutable reference to the entry corresponding to the key.
///
/// The key may be any borrowed form of the map's key type,
/// but Hash and Eq on the borrowed form must match those for the key type.
......@@ -428,4 +442,53 @@ where
}
None
}
/// Returns a mutable reference to the entry corresponding to the key.
///
/// The key may be any borrowed form of the map's key type,
/// but Hash and Eq on the borrowed form must match those for the key type.
fn entry_mut<Q>(&mut self, key: &Q) -> Option<&mut Entry<K, V>>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
if let Some(slot) = self.probe_inspecting(key) {
return self.entries.get_mut(slot)
}
None
}
}
impl<'a, K, Q: ?Sized, V> core::ops::Index<&'a Q> for HashMap<K, V>
where
K: Eq + Hash + Borrow<Q> + parity_codec::Codec,
V: parity_codec::Codec,
Q: Eq + Hash,
{
type Output = V;
fn index(&self, index: &Q) -> &Self::Output {
self
.get(index)
.expect(
"[pdsl_core::HashMap::index] Error: \
expected `index` to be within bounds"
)
}
}
impl<'a, K, Q: ?Sized, V> core::ops::IndexMut<&'a Q> for HashMap<K, V>
where
K: Eq + Hash + Borrow<Q> + parity_codec::Codec,
V: parity_codec::Codec,
Q: Eq + Hash,
{
fn index_mut(&mut self, index: &Q) -> &mut Self::Output {
self
.get_mut(index)
.expect(
"[pdsl_core::HashMap::index_mut] Error: \
expected `index` to be within bounds"
)
}
}
......@@ -37,6 +37,53 @@ fn new_unchecked() {
})
}
#[test]
fn get() {
run_test(|| {
let mut map = new_empty::<String, String>();
// Inserts some elements
assert_eq!(map.insert("Black".into(), "White".into()), None);
assert_eq!(map.insert("Up".into(), "Down".into()), None);
// Check if get returns the right answer
assert_eq!(map.get("Black"), Some(&"White".into()));
assert_eq!(map.get("Up"), Some(&"Down".into()));
assert_eq!(map.get("Forward"), None);
})
}
#[test]
fn index() {
run_test(|| {
let mut map = new_empty::<String, String>();
// Inserts some elements
assert_eq!(map.insert("Black".into(), "White".into()), None);
assert_eq!(map.insert("Up".into(), "Down".into()), None);
// Check if get returns the right answer
assert_eq!(map["Black"], "White");
assert_eq!(map["Up"], "Down");
})
}
#[test]
fn index_repeat() {
run_test(|| {
let mut map = new_empty::<String, String>();
// Inserts some elements
assert_eq!(map.insert("Something".into(), "There it is!".into()), None);
// Check if get returns the right answer repeatedly
assert_eq!(map["Something"], "There it is!");
assert_eq!(map["Something"], "There it is!");
})
}
#[test]
#[should_panic]
fn index_fail0() {
let map = new_empty::<String, String>();
// This will just fail and panic
&map["Won't catch this!"];
}
#[test]
fn insert() {
run_test(|| {
......@@ -51,7 +98,7 @@ fn insert() {
// Should *not* increase len.
assert_eq!(map.len(), 1);
// Should return the new value
assert_eq!(map.get("1"), Some(&", World!".into()));
assert_eq!(map["1"], ", World!");
})
}
......@@ -93,20 +140,6 @@ fn remove() {
})
}
#[test]
fn get() {
run_test(|| {
let mut map = new_empty::<String, String>();
// Inserts some elements
assert_eq!(map.insert("Black".into(), "White".into()), None);
assert_eq!(map.insert("Up".into(), "Down".into()), None);
// Check if get returns the right answer
assert_eq!(map.get("Black"), Some(&"White".into()));
assert_eq!(map.get("Up"), Some(&"Down".into()));
assert_eq!(map.get("Forward"), None);
})
}
#[test]
fn mutate_with() {
run_test(|| {
......@@ -115,8 +148,8 @@ fn mutate_with() {
assert_eq!(map.insert("Dog Breed".into(), "Akita".into()), None); // Shiba Inu
assert_eq!(map.insert("Cat Breed".into(), "Bengal".into()), None); // Burmilla
// Verify the inserted breeds
assert_eq!(map.get("Dog Breed"), Some(&"Akita".into()));
assert_eq!(map.get("Cat Breed"), Some(&"Bengal".into()));
assert_eq!(map["Dog Breed"], "Akita");
assert_eq!(map["Cat Breed"], "Bengal");
// Change the breeds
assert_eq!(
map.mutate_with("Dog Breed", |breed| *breed = "Shiba Inu".into()),
......@@ -127,8 +160,8 @@ fn mutate_with() {
Some(&"Bengal Shorthair".into())
);
// Verify the mutated breeds
assert_eq!(map.get("Dog Breed"), Some(&"Shiba Inu".into()));
assert_eq!(map.get("Cat Breed"), Some(&"Bengal Shorthair".into()));
assert_eq!(map["Dog Breed"], "Shiba Inu");
assert_eq!(map["Cat Breed"], "Bengal Shorthair");
// Mutate for non-existing key
assert_eq!(map.mutate_with("Bird Breed", |breed| *breed = "Parrot".into()), None);
})
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment