Is there an `Entry` mechanism for `BTreeMap` that allows returning an immutable reference?

user_1932452 0
Viewed 56

Is there an Entry mechanism for BTreeMap that allows returning an immutable reference?

Take the following code, which implements a kind of sparse vector. Essentially, it's a vector that assumes the elements are zero unless specifically set and manipulated:

// External libraries
use std::cell::{Ref, RefCell, RefMut};
use std::collections::BTreeMap;

// Stores a sparse vector
#[derive(Debug)]
struct MyVec {
    data: RefCell<BTreeMap<usize, f64>>,
}
impl MyVec {
    // Returns a reference to an element
    pub fn get<'a>(&'a self, i: usize) -> Ref <'a,f64> {
        // Insert a missing element
        let mut data = self.data.replace(BTreeMap::new());
        match data.get(&i) {
            Some(_) => (),
            None => {data.insert(i,0.0); ()}
        }
        self.data.replace(data);

        // Return the element
        Ref::map(self.data.borrow(), |data| data.get(&i).unwrap())
    }

    // Returns a mutable reference to an element
    pub fn get_mut<'a>(&'a self, i: usize) -> RefMut<'a, f64> {
        RefMut::map(self.data.borrow_mut(), |data| data.entry(i).or_insert(0.0))
    }

    // Create a new sparse vector
    pub fn new() -> MyVec {
        MyVec {
            data: RefCell::new(BTreeMap::new()),
        }
    }
}

fn main() {
    // Create a vector
    let x = MyVec::new();
    *x.get_mut(3) = 4.0;
    *x.get_mut(5) += 0.5;
    let x2 = *x.get(2);
    let x3 = *x.get(3);
    println!("x: {:?}, x[2]: {}, x[3]: {}", x, x2, x3);
}

This returns the result:

x: MyVec { data: RefCell { value: {2: 0.0, 3: 4.0, 5: 0.5} } }, x[2]: 0, x[3]: 4

My issue here is that the function get requires three look-ups into the map. It needs one look-up to determine if the data exists, one lookup to insert the entry if it does not, and a third look-up to return the data. This contrasts from the function get_mut which can do a single look-up for the element, insert a missing zero if necessary, and then return a mutable reference.

Is there a better way to insert a missing element if required and then return an immutable reference to the result?

rust data-structures
1 Answers

You can use the same entry + or_insert chain you have in get_mut in get to combine the first lookup and insertion.

Let's get rid of all the 'a lifetimes, too. Rust's inferred lifetimes are sufficient.

pub fn get(&self, i: usize) -> Ref<'_, f64> {
    self.data.borrow_mut().entry(i).or_insert(0.0);
    Ref::map(self.data.borrow(), |data| data.get(&i).unwrap())
}

(I'm not sure how to get it down to a single lookup. The issue is that we need to borrow the cell mutably to insert the default value, yet we need to return the mutable borrow before the function ends. If there's a way to do it all with a single lookup I can't come up with it.)

Now let's talk about the API. It's not good design to return Refs and RefMuts: it exposes that your struct uses RefCell, which ought to be an implementation detail. It's better to return f64 from get and &mut f64 from get_mut.

Also, get_mut allows the data structure to be modified so it ought to take &mut self. You shouldn't be allowing mutations unless the struct is mutable.

impl MyVec {
    pub fn get(&self, i: usize) -> f64 {
        *self.data.borrow_mut().entry(i).or_insert(0.0)
    }

    pub fn get_mut(&mut self, i: usize) -> &mut f64 {
        self.data.get_mut().entry(i).or_insert(0.0)
    }
}

let mut x = MyVec::new();
*x.get_mut(3) = 4.0;
*x.get_mut(5) += 0.5;
let x2 = x.get(2);
let x3 = x.get(3);

assert_eq!(x2, 0.0);
assert_eq!(x3, 4.0);
assert_eq!(x.get(5), 0.5);

Playground

Notice that x is now declared let mut x to allow the get_mut calls, and that the get calls no longer need to be dereferenced. Also we can condense get to a single lookup since we're not returning a reference (thanks @loganfsmyth).

You may wish that get returned references instead of copies of the values. Unfortunately, it is difficult to do so safely. Those references can be invalidated any time the tree is modified. get_mut can safely return references because there can only be one mutable borrow at a time. get can't because regular borrows have no such restriction.

user_685870
Related