Ownership is theft? TockのTakeCellについて

tockというrust製の組み込み向けOSがあります． このtockの作成者らが2015年にOwnership is Theft: Experiences Building an Embedded OS in Rust (PLOS 2015)という論文を発表しました．そこでは著者がtockを開発する上で嵌ったownershipに関する問題と，それを解決するための言語の修正アプローチが述べられています．ただその後rustの開発者とのやりとりもあり，言語に修正を加えることなく当初実装したかったことが実装できたようです．The Case for Writing a Kernel in Rust (APSys'17)にそのことが簡単に書かれています．

Ownershipと問題

ご存知の通りrustにはメモリ安全を保証するためにownershipという仕組みがあり，たとえシングルスレッドでも同一データに対するmutableなreferenceを複数持つことができません．といってもmutableなreferenceが複数欲し場合は往々にしてあります．論文では以下のような構成の乱数生成器呼び出しを例に挙げています．

SysCallDispatcher <--> SimpleRNG <---> RNG

pub struct SimpleRNG {
    pub busy: bool,
}

impl SimpleRNG {
    pub fn command(&mut self) {
        self.busy = true;
        //...
    }

    pub fn deliver(&mut self, rand: u32) {
        self.busy = false;
        //...
    }
}

pub struct SysCallDispatcher<'a> {
    pub srng: &'a mut SimpleRNG,
}

pub struct RNG<'a> {
    pub srng: &'a mut SimpleRNG,
}

impl<'a> SysCallDispatcher<'a> {
    pub fn dispatch(&mut self) {
        self.srng.command();
    }
}

impl<'a> RNG<'a> {
    pub fn done(&mut self, num: u32) {
        self.srng.deliver(num);
    }
}

ここではSysCallDispatcherとRNGがともにSimpleRNGのreferenceを所有する構成を考えていますが，このコードを実際に使用する場合はborrow checkに引っかかりコンパイルできません．

let mut srng = SimpleRNG { busy: false };
let mut dispathcer = SysCallDispatcher { srng: &mut srng };
// エラー: cannot borrow `srng` as mutable more than once at a time
let mut rng = RNG { srng: &mut srng };

解決策1. Cell

このような場合の解決策の一つはCellを使うことです．

use std::cell::Cell;

pub struct SimpleRNG {
    pub busy: Cell<bool>,
}

impl SimpleRNG {
    pub fn command(&self) {
        self.busy.set(true);
    }

    pub fn deliver(&self, rand: u32) {
        self.busy.set(false);
    }
}

pub struct SysCallDispatcher<'a> {
    pub srng: &'a SimpleRNG,
}

pub struct RNG<'a> {
    pub srng: &'a SimpleRNG,
}

impl<'a> SysCallDispatcher<'a> {
    pub fn dispatch(&self) {
        self.srng.command();
    }
}

impl<'a> RNG<'a> {
    pub fn done(&self, num: u32) {
        self.srng.deliver(num);
    }
}

こうするとコンパイルが通るようになります．

let mut srng = SimpleRNG { busy: false };
let mut dispathcer = SysCallDispatcher { srng: &srng };
let mut rng = RNG { srng: &srng };

ポイントとしては今回SysCallDispatcherやRNGは&mut Tではなく&Tを保有している点です．imutableなreferenceなのでborrow checkerにひかかりません．Cellは内部的にunsafeなコードを利用することで値を変更します．get()は値のコピーを返し，set()ではstd::mem::replace()を利用して値を書き換えます．

CellはCopyを実装している型しか使えません．またCellはSyncを実装していないため，Cellをスレッド間で共有するようなコードはコンパイルエラーになります．そのためデータ競合が生じることはありません．

Copyを実装している型しか使えないというのは大きな制約です．プリミティブ型はCellでいいですが，バッファ領域などを共有することができません(&mut TはCopyを実装していません)．

解決策2. TakeCell

バッファ領域などを共有するために，tockではTakeCellというものを利用しています．TakeCellはCellと似ていますが，以下のデータ構造を利用してreferenceを保持します．

pub struct TakeCell<'a, T: 'a + ?Sized> {
    val: UnsafeCell<Option<&'a mut T>>,
}

take()を利用してTakeCellの値を取得することができます．もし仮にTakeCellの中身が別から取得されている場合はNoneが返ります．

    pub fn take(&self) -> Option<&'a mut T> {
        unsafe {
            let inner = &mut *self.val.get();
            inner.take()
        }
    }

また，クロージャからTakeCellのデータを簡単に利用できるようにmap()というメソッドを提供しています．

    pub fn map<F, R>(&self, closure: F) -> Option<R>
        where F: FnOnce(&mut T) -> R
    {
        let maybe_val = self.take();
        maybe_val.map(|mut val| {
            let res = closure(&mut val);
            self.replace(val);
            res
        })
    }

take()するとTakeCellのデータはNoneになります．そこでデータをTakeCellに戻したい場合はmap()のコードにあるように明示的に戻す必要があります．

TakeCellを使うと例えば以下のようなコードが書けます．

pub struct DMAChannel {
    pub buffer: TakeCell<'static, [u8]>,
}

impl DMAChannel {
    pub fn foo(&mut self) {
        self.buffer.map(|b| {
            //...
        });
    }
}

static mut buffer: [u8; 128] = [0; 128];

fn foo(){
    let mut chan = unsafe {
        DMAChannel {
            buffer: TakeCell::new(&mut buffer),
        }
    };
    chan.foo();
}

staticなborrowを作るためにはunsafeなコードが必要です．

RefCellとtry_borrow_mut

rustを書いたことがある人ならTakeCellなんか利用しなくてもRefCellでいいのでは?と思うと思います． RefCellはget(), set()の代わりにborrow()及びborrow_mut()を提供し，値のreferenceを操作することができます．RefCellはCopyを実装していなくても使用することができます．

RefCell最大の特徴はコンパイル時ではなく実行時にborrow checkをするということです．もし仮に複数のmutableなreferenceを使用とした場合，実行時にpanicします．RefCellを使って安全なコードを書くのはプログラマの責任です．カーネル内でpanicしてしまうのは大変望ましくありません．TakeCellの場合はもしtake()されていたらNoneが返ります (map()の場合は何もしないで終わります). ただし，RefCellにはtry_borrow_mut()というメソッドがあり，もし仮にすでにmutableなrreferenceが取得されていた場合Errが返ります．

TakeCellを利用しなくても，try_borrow_mut()でそれと同等な機能ができるような気がします． それではなぜtockがTakeCellを使っているのかというと... このredditのコメントによると TakeCellができたのはtry_borrow_mut()が導入されるよりも前 *1 というのが最大の理由のようです．

ちなみに，tockにはTakeCellとすごく似ているMapCellというデータ構造もあります．MapCellはreferenceではなく実際の値を保持するようになっており，実際に値を保持しているかはoccupiedと呼ばれるフィールドで保持しています．一方TakeCellは値をreferenceに限定することでOptionのnon-null optimizationを利用してoccupied分のデータ構造を節約しています*2．

また，普通だとCellやRefCellはRcと組み合わせて利用することが多いと思いますが今回カーネル内で動的にメモリを確保することは考えていないのでRcの利用はオプション外です．tockだと基本的にTakeCellでは'staticなバッファ領域を共有するようです．