Is there something similar to try_retain()? - rust

Question

Is it possible to apply FnMut(&mut V) -> Result<bool> for collection elements and if result is:
Ok(false) -- remove element and continue
Ok(true) -- leave element in collection and continue
Err(_) -- stop and return the error
Basically, how to code a try_retain(), an equivalent of this C++ code:
struct S {
// ...
};
// Updates `s` (and possibly some external state), returns true if `s` is no longer needed
// throws on error
bool bar(S& s);
void foo(std::map<int, S>& m) {
try
{
for(auto it = m.begin(), it_end = m.end(); it != it_end; ) {
if (bar(it->second))
m.erase(it++);
else
++it;
}
}
catch(...)
{
printf("bailed early\n");
throw;
}
}
For some reason I keep running into this pattern and I can't figure out how to do it in Rust without traversing collection twice or using additional memory...

Answer 1

This won't short-circuit, but you can do it by storing the global result in a mutable variable that's updated by the closure you pass to retain. Something like this:
use std::collections::HashMap;
fn try_retain<K, V, E, F>(m: &mut HashMap<K, V>, mut f: F) -> Result<(), E>
where
F: FnMut(&K, &mut V) -> Result<bool, E>,
{
let mut result = Ok(());
m.retain(|k, v| match f(k, v) {
Ok(b) => b,
Err(e) => {
result = Err(e);
true
}
});
return result;
}
Playground

Answer 2

If you're fine with switching to hashbrown::HashMap (which is what std::collections::HashMap uses as it's inner implementation), you could copy the implementation of fn retain and slightly change it to use the ? operator to return the first error:
fn try_retain<K, V, E, F>(m: &mut hashbrown::HashMap<K, V>, mut f: F) -> Result<(), E>
where
F: FnMut(&K, &mut V) -> Result<bool, E>,
{
let mut raw_table = m.raw_table();
// Here we only use `iter` as a temporary, preventing use-after-free
unsafe {
for item in raw_table.iter() {
let &mut (ref key, ref mut value) = item.as_mut();
if !f(key, value)? {
raw_table.erase(item);
}
}
}
Ok(())
}
This requires hashbrown's raw feature to be enabled, so you can access the underlying raw table.
This stops at the first error, and the elements are visited in a random order, so it's non-deterministic which elements are removed if an error occurs.

Answer 3

We could use try_for_each and remember the keys to remove (not to retain) till the end of the iteration because of the borrow checker:
fn try_retain(
map: &mut HashMap<i32, i32>,
to_retain: fn(k: i32, v: i32) -> bool,
) -> Result<bool, String> {
let mut to_remove: Vec<i32> = vec![];
map.iter().try_for_each(|(k, v)| {
if has_error() {
return Err("bailed early".to_string());
}
if !to_retain(*k, *v) {
to_remove.push(*k);
}
Ok::<_, String>(())
})?; // short-circuiting
if to_remove.is_empty() {
Ok(true)
} else {
(*map).retain(|&k, _| !to_remove.contains(&k));
Ok(false)
}
}
The return value results from whether there are keys to remove.
Can short-circuiting.
Playground
If the extra space for the to-remove-items is still too much, here is an unstable solution with drain_filter:
fn try_retain(
map: &mut HashMap<i32, i32>,
to_retain: fn(k: i32, v: i32) -> bool,
) -> bool {
let mut rslt = true;
map.drain_filter(|k, v| {
if !to_retain(*k, *v) {
rslt = false;
true
} else {
false
}
});
rslt
}
Short-circuiting is only possible through panic.
Playground