るいもの戯れ言

Rustで、Multimap(1つのキーに複数の値を登録できるMap)を実装してみる。後でキーをソート順の取り出したいので、BtreeMapを使って、値にVecを入れる。


use std::collections::BTreeMap;

struct Bag {
    map: BTreeMap<i32, Vec<i32>>
}

impl Bag {
    fn add(&mut self, key: i32, value: i32) {
        match self.map.get_mut(&key) {
            None => {
                self.map.insert(key, vec!(value));
            },
            Some(vec) => vec.push(value)
        }
    }
}

でも、これはコンパイルが通らない。


error[E0499]: cannot borrow `self.map` as mutable more than once at a time
  --> /Users/shanai/.emacs.d/rust-playground/at-2017-01-22-230525/snippet.rs:32:17
   |
30 |         match self.map.get_mut(&key) {
   |               -------- first mutable borrow occurs here
31 |             None => {
32 |                 self.map.insert(key, vec!(value));
   |                 ^^^^^^^^ second mutable borrow occurs here
...
35 |         }
   |         - first borrow ends here

mutable borrowが2回起きているという警告。まずここでself.mapに対するmutable参照を取得している。


        match self.map.get_mut(&key) {

で、このmatchが終了する前に、ここでもう一度self.mapに対するmutable参照を取得している(insertは、mutableなselfが必要)。


            None => {
                self.map.insert(key, vec!(value));
            },

つまりmutableな参照が、2つ存在してしまう瞬間がある。これはRustでは禁止されている。以下がエラーになるのと同じ。


    let mut i = 0;
    let pi = &mut i;
    let pi2 = &mut i;

スコープを分けて、mutable参照が2つ同時に存在する瞬間を無くしてしまえば良い。


impl Bag {
    fn add(&mut self, key: i32, value: i32) {
        {
            match self.map.get_mut(&key) {
                None => {},
                Some(vec) => {
                    vec.push(value);
                    return
                }
            }
        }
        
        self.map.insert(key, vec!(value));
    }
}

でも、BtreeMapにはentry()という素敵なメソッドがあるので、これを使うのが簡単。


impl Bag {
    fn add(&mut self, key: i32, value: i32) {
        self.map.entry(key).or_insert(vec!()).push(value)
    }
}

entryはキーを引数に取り、Entryというenumを返す。指定されたキーが無ければVacant、あればOccupiedが返る。or_insert()を呼び出すと、Vacantの場合は、引数に指定された値をmapに格納した上で、その値へのmutable参照を返す。Ocupiedの場合は、値へのmutable参照を返す。なので上のように書けば目的が達成できる。


fn main() {
    let mut bag = Bag { map: BTreeMap::new() };
    bag.add(1, 10);
    bag.add(1, 20);
    bag.add(2, 100);
    println!("map = {:?}", bag.map);
}


map = {1: [10, 20], 2: [100]}