【Rust】列挙体ディスパッチのボイラープレート削減

あるトレイトを実装した任意の型を受け取りたい場合はジェネリクスを用いるのが自然ですが，フィールドやVecの要素として複数種類の型をとりたい場合などいくつかの理由で具象型として受け取りたいことがあります．その場合は列挙体やトレイトオブジェクトが利用できます．
例えば以下のスタックとキューを同じ型として使いたい場合を考えてみます．pushとpopができるトレイトPushPopを定義しStackとQueueに実装しています．

use std::collections::VecDeque;

pub trait PushPop<T> {
    fn push(&mut self, value: T);
    fn pop(&mut self) -> Option<T>;
    fn is_empty(&self) -> bool;
}

#[derive(Debug, Default)]
pub struct Stack<T>(Vec<T>);

#[derive(Debug, Default)]
pub struct Queue<T>(VecDeque<T>);

impl<T> Stack<T> {
    pub fn new() -> Self {
        Self(Vec::new())
    }
}

impl<T> Queue<T> {
    pub fn new() -> Self {
        Self(VecDeque::new())
    }
}


impl<T> PushPop<T> for Stack<T> {
    fn push(&mut self, value: T) {
        self.0.push(value);
    }
    fn pop(&mut self) -> Option<T> {
        self.0.pop()
    }
    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

impl<T> PushPop<T> for Queue<T> {
    fn push(&mut self, value: T) {
        self.0.push_back(value);
    }
    fn pop(&mut self) -> Option<T> {
        self.0.pop_front()
    }
    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

ナイーブな実装

StackとQueueを一つの型として扱うために列挙体PushPopCollectionを用意し，それぞれの型をそのバリアントとします．

#[derive(Debug)]
enum PushPopCollection<T> {
    Stack(Stack<T>),
    Queue(Queue<T>),
}

PushPopCollection列挙体は各バリアントに処理を委譲することでPushPopトレイトを実装します．

impl<T> PushPop<T> for PushPopCollection<T> {
    fn push(&mut self, value: T) {
        match self {
            Self::Stack(stack) => stack.push(value),
            Self::Queue(queue) => queue.push(value),
        }
    }
    fn pop(&mut self) -> Option<T> {
        match self {
            Self::Stack(stack) => stack.pop(),
            Self::Queue(queue) => queue.pop(),
        }
    }
    fn is_empty(&self) -> bool {
        match self {
            Self::Stack(stack) => stack.is_empty(),
            Self::Queue(queue) => queue.is_empty(),
        }
    }
}

これによってStackとQueueを同じ型として扱うことができるようになります．

let mut collections: Vec<PushPopCollection<i32>> = vec![
	PushPopCollection::Stack(Stack::new()),
	PushPopCollection::Queue(Queue::new()),
];

collections[0].push(1);
collections[0].push(2);
collections[0].push(3);

collections[0].pop();
println!("stack: {:?}", collections[0]);
// stack: Stack(Stack([1, 2]))

collections[1].push(1);
collections[1].push(2);
collections[1].push(3);

collections[1].pop();
println!("queue: {:?}", collections[1]);
// queue: Queue(Queue([2, 3]))

今回は，この委譲部分のボイラープレートを減らすクレートを紹介します．コード全体はこちら

enum_dispatchを使う

enum_dispatchはこのボイラープレートを減らすためのクレートです．まずPushPopトレイトに以下のようにアトリビュートを付けます．

#[enum_dispatch::enum_dispatch]
pub trait PushPop<T> {
    fn push(&mut self, value: T);
    fn pop(&mut self) -> Option<T>;
    fn is_empty(&self) -> bool;
}

そしてPushPopトレイトを定義したモジュール(ファイル)内で以下のようにPushPopCollection列挙体を定義します．

#[derive(Debug)]
#[enum_dispatch::enum_dispatch(PushPop<X>)]
pub enum PushPopCollection<T> {
    Stack(Stack<T>),
    Queue(Queue<T>),
}

これだけでPushPopCollectionにPushPopが実装されます．

自分の環境ではトレイトと別のモジュールで列挙体を定義するとうまく動きませんでした．別モジュールのトレイトだったり外部クレートのトレイトを使いたい場合は次に説明するambassadorが利用できます．

ambassadorを使う

ambassadorは委譲のボイラープレート削減にフォーカスしたクレートです．構造体のフィールドや列挙体のバリアントに処理を委譲できます．enum_dispatchクレートと同じようにPushPopトレイトにアトリビュートを付けます．

#[ambassador::delegatable_trait]
pub trait PushPop<T> {
    fn push(&mut self, value: T);
    fn pop(&mut self) -> Option<T>;
    fn is_empty(&self) -> bool;
}

トレイトを別モジュールから利用する場合は自動生成されたマクロambassador_impl_PushPopもインポートします．

use ambassador::Delegate;
use my_module::{ambassador_impl_PushPop, PushPop, Queue, Stack};

#[derive(Debug, Delegate)]
#[delegate(PushPop<X>, generics = "X")]
enum PushPopCollection<T> {
    Stack(Stack<T>),
    Queue(Queue<T>),
}

これで列挙体にPushPopが実装されます．

(補足) トレイトオブジェクトを使う

トレイトオブジェクトを利用してもStackとQueueを同一の型として利用できます．

let mut collections: Vec<Box<dyn PushPop<i32>>> =
	vec![Box::new(Stack::new()), Box::new(Queue::new())];

collections[0].push(1);
collections[0].push(2);
collections[0].push(3);

let popped_item = collections[0].pop();
println!("popped item: {popped_item:?}");
// popped item: Some(3)

collections[1].push(1);
collections[1].push(2);
collections[1].push(3);

let popped_item = collections[1].pop();
println!("popped item: {popped_item:?}");
// popped item: Some(1)

トレイトオブジェクトは動的ディスパッチを用いるため列挙体を利用する方法よりパフォーマンスが低下してしまいます．一方で列挙体はバリアントとして型を持つため，新しくPushPopトレイトを実装した型を許容するためにはバリアントを増やさなければなりません．つまり，ユーザーが新しく定義した型を(具象型として)受け取りたい場合はトレイトオブジェクトを使う必要があります．