Rust の型でスタック操作の正しさを保証する

要約

Rust の型システムを応用すると、スタックの様な後入れ先出し (LIFO: last in first out) の操作に関して以下の様な点をコンパイル時に検証することができます。

• 空のスタックからは pop できないこと
• 函数内でスタックに要素を push した後、函数から返る前に確実に pop すること

前提課題

Rust の Vec には push と pop メソッドがあるのでそのままスタックとして使用できます。

fn main() {
    let mut stack = Vec::new();
    stack.push(0);
    stack.push(1);
    stack.push(4);
    stack.pop();
    println!("{:?}", stack); // [0, 1]
}

しかし pop を呼ぶタイミングに制約がないので、空のスタックに対して pop を呼ぶことができます。

fn main() {
    let mut stack = Vec::new();
    stack.push(0);
    println!("{}", stack.pop().unwrap()); // 0
    println!("{}", stack.pop().unwrap()); // runtime panic!
}

また、うっかり pop を呼び忘れても気付きにくいです。

fn main() {
    let mut stack = Vec::new();
    {
        stack.push(1);
        println!("{:?}", stack); // [1]
        {
            stack.push(2);
            println!("{:?}", stack); // [1, 2]
            stack.pop();
        }
        println!("{:?}", stack); // [1]
        {
            stack.push(3);
            println!("{:?}", stack); // [1, 3]
            // stack.pop(); // oops!
        }
        stack.pop();
    }
    assert_eq!(stack[..], []); // runtime panic!
}

1 回の push に対して 1 回の pop が確実に対応することをコンパイル時に保証できると良いですね。

ガードオブジェクトを導入する

スタックに要素を push した時にガードオブジェクトが生成され、pop するときにガードオブジェクトが消滅する様にします。すると、要素がスタック内に存在する期間とガードオブジェクトの生存期間が一致する様になるので、push と pop が一対一となることを強制できます。

初めにスタックとガードオブジェクトの構造を定義します。Vec::pop を自由に呼べなくするために Vec を新しい構造体で包みます。ガードオブジェクトはそれへの参照を持ちます。

struct Stack<T>(Vec<T>);
struct Guard<'a, T>(&'a mut Stack<T>);

次に、スタックの生成と push を実装します。Push において新しいガードオブジェクトを生成して返す様にします。

impl<T> Stack<T> {
    fn new() -> Self {
        Stack(Vec::new())
    }
    fn push(&mut self, item: T) -> Guard<'_, T> {
        self.0.push(item);
        Guard(self)
    }
}

さらに、ガードオブジェクトが消滅する時に自動的に pop がされる様に、Drop を実装します。

impl<T> Drop for Guard<'_, T> {
    fn drop(&mut self) {
        self.0.0.pop();
    }
}

最後に、Deref(Mut) でガードオブジェクトからスタックにアクセスできる様にします。

impl<T> Deref for Guard<'_, T> {
    type Target = Stack<T>;
    fn deref(&self) -> &Stack<T> {
        self.0
    }
}
impl<T> DerefMut for Guard<'_, T> {
    fn deref_mut(&mut self) -> &mut Stack<T> {
        self.0
    }
}

使用例

定義したスタックを実際に使ってみます。

fn main() {
    let mut stack = Stack::new();
    {
        let mut guard = stack.push(1);
        println!("{:?}", guard); // [1]
        {
            let guard2 = guard.push(2);
            println!("{:?}", guard2); // [1, 2]
        }
        {
            let guard3 = guard.push(3);
            println!("{:?}", guard3); // [1, 3]
        }
    }
    assert_eq!(stack.0[..], []);
}

最初に push した後にもう一度 push するときはカードオブジェクト経由で push を呼び出すのがミソです。ガードオブジェクトの中に元のスタックへの可変参照があるので、ガードオブジェクトが消滅するまではスタックを直接触ることはできません。

もう間違ったことはできない

Vec::pop を Guard::drop のみから呼ぶ様にしたので、間違ったタイミングで Vec::pop が呼ばれることはもうありません。Guard::drop が呼ばれるには先に Stack::push でガードオブジェクトを生成する必要があります。よって push した後にしか pop は実行されません。

ガードオブジェクトが消滅する時に Vec::pop が呼ばれることは Rust のコンパイラーが保証してくれるので、(ガードオブジェクトをどこか別の場所に保管し続けない限りは) ガードオブジェクトのローカル変数のスコープから処理が抜けるときに確実に pop が実行されます。

明示的な pop

ガードオブジェクトが消滅する時に自動的に pop が呼ばれるので、コード内で明示的に pop を呼ぶ必要はありません。しかし明示的に pop を呼ぶ方が分かりやすく感じられることもあるかもしれません。ガードオブジェクトを受け取って捨てるだけのメソッドを作ると、それらしくなります。(現在 nightly にある std::sync::Mutex::unlock の真似です。)

impl<T> Stack<T> {
    fn pop(guard: Guard<'_, T>) {
        drop(guard)
    }
}

fn main() {
    let mut stack = Stack::new();
    {
        let mut guard = stack.push(1);
        println!("{:?}", guard); // [1]
        Stack::pop(guard);
    }
    assert_eq!(stack.0[..], []);
}