RustでUnion Typesを実現する型レベルプログラミングの技術

はじめに

この記事では、RustでUnion Typesのように型の集合を表現する方法について解説します。型レベルリストを用いることで型の集合を表現し、Trait Resolutionを使って特定の型が集合に含まれるかどうか、またある集合が別の集合を含むかどうかを判定することができます。

実装の核となるのは、再帰の深さを表す型パラメータを使った型レベルリストの再帰的処理です。このトリックを使うことで、型レベルリストに対する様々な処理が実現できます。

筆者はRustの初心者です。もし間違いや用語の誤用などがあれば、お気軽にご指摘ください。

モチベーション

Union TypesとEnum

TypeScript、Scalaなどの言語のUnion Typesは、複数の型を1つの型として扱うことができる機能です。この記事では、例えば、TypeScriptにおいて string | number は、文字列または数値のどちらかを示します。

function printValue(value: string | number) {
    console.log(value); // valueは string または number
}

printValue("hello"); // OK
printValue(42);      // OK

RustにはUnion Typesが存在しません。多くの場合、Enumを使うことで似たようなことを実現できます。文字列または数値のどちらかを表すUnion TypesをEnumで表現すると以下のようになります。

enum Value {
    String(String),
    Number(i32),
}

fn print_value(value: Value) {
    match value {
        Value::String(s) => println!("{}", s),
        Value::Number(n) => println!("{}", n),
    }
}

fn main() {
    print_value(Value::String("hello".to_string())); // OK
    print_value(Value::Number(42));                  // OK
}

この例では、Value というEnumを定義し、String と Number の2つのケースを持つようにしています。print_value 関数は Value に対してパターンマッチを行い、それぞれのケースに応じて処理を行うことができます。

Enumの問題点

この2つの機能は似ていますが、Union Typesでは表現できるもののEnumでは表現できないケースがあります。それが、特定のケースだけを受け付ける関数を書くことです。例として、クロスプラットフォームなグラフィックエンジンを考えます。

このレンダリングエンジンは、丸(Circle)・四角(Square)・テキスト(Text)の3つの要素(Shape)をレンダリングすることができます。

abstract class Shape {}

class Circle extends Shape {
    constructor(public radius: number) {
        super();
    }
}

class Square extends Shape {
    constructor(public size: number) {
        super();
    }
}

class Text extends Shape {
    constructor(public text: string) {
        super();
    }
}

このレンダラは複数のシステムで利用されます。

• GeometryOSというシステムは、幾何学模様の描画に特化しており、CircleとSquareのみをサポートしています。
• TextOSというシステムは、テキストの描画に特化しており、Textのみをサポートしています。
• MightyOSというシステムは、Circle, Square, Textの全てをサポートしています。

TypeScriptでは、Union Typesを使って、特定の形状のみを受け付ける関数を書くことができます。

function renderOnGeometryOS(shape: Circle | Square) {
    // CircleとSquareのみを受け付ける
}

function renderOnTextOS(shape: Text) {
    // Textのみを受け付ける
}

function renderOnMightyOS(shape: Circle | Square | Text) {
    // Circle, Square, Textの全てを受け付ける
}

サポートされていない形状を渡すとタイプエラーとなります。

renderOnGeometryOS(new Circle(10));    // OK
renderOnGeometryOS(new Text("hello")); // Error

Enumではどうでしょうか。ShapeというEnumを定義し、Circle, Square, Textの3つのケースを持つようにします。

enum Shape {
    Circle(radius: f64),
    Square(size: f64),
    Text(text: String),
}

しかし、Enumでは指定したケースのみを受け付ける関数を書くことができません。

fn render_on_geometry_os(shape: ???) {
    // CircleとSquareのみを受け付けたいが、型レベルで制限することができない
}

パターンマッチを使って、サポートされていないケースでプログラムをクラッシュさせることはできますが、実行時にエラーが発生する可能性があります。このようなケースで、コンパイル時に型エラーを検出することが、この記事のモチベーションです。

型レベルリストで型の集合を表現する

Union Typesを実現するためには、(A) 型の集合を表現する型を定義する (B) 特定の型が集合に含まれるかどうかを判定する という2つのステップが必要です。

型の集合を表現する型は、型レベルリストを使って実現することができます。具体的には、以下のようにHListというトレイトを定義し、HNilとHConsという2つの型を実装します。Cons Listが要素とそのあとに続くリストを持つように、型レベルリストも要素(Head)とそのあとに続く型レベルリスト(Tail)を持つようにします。HNil は空のリストを表します。

struct HNil;
struct HCons<Head, Tail>(Head, Tail);

trait HList {}

impl HList for HNil {}
impl<Head, Tail> HList for HCons<Head, Tail> {}

このように定義することで、以下のように型レベルリストを表現することができます。

// 形状ごとに型 (struct) を定義する
struct Circle {
    radius: f64,
}
struct Square {
    size: f64,
}
struct Text {
    text: String,
}


// すべての形状
type AllShapes = HCons<Circle, HCons<Square, HCons<Text, HNil>>>;
// 幾何学模様の形状
type GeometryShapes = HCons<Circle, HCons<Square, HNil>>;

HCons / HNil を繰り返す表現は冗長であるため、これを簡潔にするためのマクロを定義します。

macro_rules! HList {
    () => { HNil };
    ($head:ty $(,)*) => { HCons<$head, HNil> };
    ($head:ty, $($tail:tt)*) => { HCons<$head, HCons!($($tail)*)> };
}

このマクロを使うことで、以下のように型レベルリストを定義することができます。

type AllShapes = HList!(Circle, Square, Text);
type GeometryShapes = HList!(Circle, Square);

以下、このマクロを使って型の集合を定義します。このマクロは表現を簡潔にするためのものであり、Union Typesを実現するために必須ではありません。

型が集合に含まれるかどうかを判定する

Union Typesを実現するためには、特定の型が集合に含まれるかを判定し、含まれない場合には型エラーとなるようにする必要があります。

「特定の集合に属す型である」ことを示すTraitを定義できれば、Genericsと組み合わせて以下のように「集合のうち1つの型を取る関数」を定義することができます。

type GeometryShapes = HList!(Circle, Square);

fn render_on_geometry_os<S>(shape: S) where S: Member<GeometryShapes> {
    // shape は Circle もしくは Square
}

そこで、この Member トレイトを定義することを考えます。 L を任意の HList とするとき、その全ての要素に Member<L> を impl することがゴールとなります。

これを実現するには、再帰的に Member トレイトを実装する必要があります。具体的には、

• ベースケース: Head は Member<HCons<Head, Tail>>
• 再帰ケース: T が Member<Tail> のとき、T は Member<HCons<Head, Tail>>

の2ケースを実装する必要があります。素直に上記の2つのルールをトレイト実装として書くと以下のようになります。

trait Member<H> {}
impl<Head, Tail>    Member<HCons<Head, Tail>> for Head {}
impl<T, Head, Tail> Member<Tail>              for T
    where T: Member<HCons<Head, Tail>> {}

残念ながら、このコードをコンパイルすると以下のエラーとなります。

error[E0119]: conflicting implementations of trait `Member<HCons<_, _>>`
  --> src/main.rs:44:1
   |
43 | impl<Head, Tail> Member<HCons<Head, Tail>> for Head {}
   | --------------------------------------------------- first implementation here
44 | impl<T, Head, Tail> Member<Tail> for T where T: Member<HCons<Head, Tail>> {}
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ conflicting implementation

このエラーは、同じ型に対して複数のトレイト実装が存在するために発生します。探している型が Head であるとき、ベースケースは Head に Member<HCons<Head, Tail>> を実装します。再帰ケースにおける Member<Tail> がこれと重なることがあるため、コンパイラが2つのトレイト実装を区別できません。

Specializationを使うとこの問題を解決できるかもしれませんが、この記事ではあくまでStable Rustを使うことを前提としています。

しかし、この問題を解決する方法があります。それが、Gentle Intro to Type-level Recursion in Rust: From Zero to HList Sculpting で紹介されているトリックで、ここでは Index トリックと呼びます。トリックの肝は、ベースケースと再帰ケースを区別するために、Member トレイトに新たな型パラメータを追加し、常に異なるトレイトを実装するようにすることです。

まず、Member トレイトに新たな型パラメータ Index を追加します。

trait Member<H, Index> {}

次に、ベースケースと再帰ケースを区別するための Index として使う2つの型を定義します。

struct Here;
struct There<Index>(Index);

Here はベースケース、There は再帰ケースを表します。ベースケースは Member<H, Here> を実装し、再帰ケースは Member<T, There<Index>> を実装します。Member<H, Here> と Member<T, There<Index>> は異なるトレイトとして扱われるため、コンパイラがエラーを出さなくなります。

impl<Head, Tail: HList> Member<HCons<Head, Tail>, Here> for Head {}
impl<T, Head, Tail, Index> Member<HCons<Head, Tail>, There<Index>> for T where T: Member<Tail, Index>
{}

Index とは、再帰の深さを表す型パラメータです。Here は再帰の深さが0であることを示し、これ以上の再帰を行いません。There<Index> は、さらに Index 回 (合計Index + 1回)再帰を行うことを示します。このようにすることで、コンパイラが異なるトレイト実装として扱うことができます。

例えば、

type AllShapes = HList!(Circle, Square, Text);

のとき、

• Circle は Member<AllShapes, Here> を
• Square は Member<AllShapes, There<Here>> を
• Text は Member<AllShapes, There<There<Here>>> を

実装します。

さて、render_on_geometry_os 関数はどのように定義できるでしょうか。Member トレイトには Index が必要です。幸い、Rustの型推論を用いて、Index をタイプチェッカーに推論させることができます。以下のように render_on_geometry_os 関数を定義します。

type GeometryShapes = HList!(Circle, Square);

fn render_on_geometry_os<S, Index>(shape: S)
where
    S: Member<GeometryShapes, Index>,
{
    // ...
}

このとき、Circle または Square を引数に渡すと、コンパイラは Index を推論し、Member<GeometryShapes, Here> または Member<GeometryShapes, There<Here>> として解釈します。一方、Text を引数に渡すと、S: Member<GeometryShapes, Index> の条件を満たす Index を推論することができないため、タイプエラーとなります。

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=45b93e8f3c7a0349cb3eb3cdee5c2439

部分集合の判定

Here と There を使うことで、型レベルリストから再帰的にMember トレイトの実装を行うことができることを示しました。この Index トリックは型レベルリストに対して再帰処理を行う汎用的な方法です。例えば、Member トレイトの応用として、ある集合が別の集合の部分集合であるかどうかを判定する Subset トレイトを定義することができます。

具体的なモチベーションとして、以下のようなケースを考えます。render_multiple_on_geometry_os 関数は、複数の幾何学模様の形状のみを受け付けます。

fn render_multiple_on_geometry_os<L>(shapes: L)
{
    // ...
}

shapes は HList の値であり、その全ての要素の型が GeometryShapes に含まれる場合にのみ受け付けることができます。

// OK
render_multiple_on_geometry_os(HCons(
    Circle::new(),
    HCons(Square::new(), HNil),
));
// NG
render_multiple_on_geometry_os(HCons(Text::new(), HNil));

まず、Member と同じように、どのように Subset トレイトを定義するかを考えます。

fn render_multiple_on_geometry_os<L>(shapes: L)
where
    L: Subset<GeometryShapes> 
{
    // ...
}

このように書くことができれば、L が GeometryShapes の部分集合であるかどうかチェックしたい意図が明確になります。

では、Subset トレイトはどのように定義できるでしょうか。S1 が S2 の部分集合であるとは、S1 の全ての要素が S2 に含まれることを意味します。そのため、Subset<S2> は、リストの全ての要素が Member<S2> である場合に impl される必要があります。Member と同様に、ベースケースと再帰ケースを考えます。

• ベースケース: HNil はどの集合にも部分集合であるため、任意の S について Subset<S> を impl する。
• 再帰ケース: HCons<Head, Tail> について、Head が Member<S2> であり、Tail が Subset<S2> のとき、HCons<Head, Tail> は Subset<S2> である。

これを素直に実装すると Member と同様に実装の競合が発生するため、Index トリックを使います。Member はその他の再帰を行わないため、一つの Index (自然数) で再帰の深さを表すことができました。一方、Subset は Head に対して Member の再帰を行いつつ、Tail に対して Subset の再帰を行うため、2つの Index が必要です。これらは Subset の再帰が深くなるごとにツリー構造を形成するため、Subset の型パラメータを追加することでは対応できません。そこで、2つの Index を合体させるための Pair という型を定義します。

struct Pair<Index1, Index2>(Index1, Index2);

Pair は2つの Index を組み合わせたものです。Pair を使って、Subset トレイトを定義します。

trait Subset<S, Index> {}

// ベースケース: HNil はどの集合にも部分集合である
impl<S> Subset<S, Here> for HNil {}

// 再帰ケース: Head が Member<S, Index1> であり、Tail が Subset<S, Index2> のとき、HCons<Head, Tail> は Subset<S, Pair<Index1, Index2>> である
impl<Head, Tail, S, Index1, Index2> Subset<S, Pair<Index1, Index2>> for HCons<Head, Tail>
where
    Head: Member<S, Index1>,
    Tail: Subset<S, Index2>,
{}

最後に、Member と同様に Subset の Index を推論させるようにrender_multiple_on_geometry_os 関数を定義します。

fn render_multiple_on_geometry_os<L>(shapes: L)
where
    L: Subset<GeometryShapes> 
{
    // ...
}

こうすることで、render_multiple_on_geometry_os 関数は、幾何学模様の形状のみで構成されるリストを受け付けることができます。

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=d9ab5ef42557f4ce0f1aeedf04d9aff0

where でトレイト境界を指定する場合、複数の境界を指定することができ、また型パラメータを右辺に持つことができます。この特性を利用すると、両方向の部分集合の判定を行うことで、2つの集合が等しいかどうかを判定することもできます。

まとめ

この記事では、RustのTrait解決を使ってUnion Typesのように型の集合を表現する方法について解説しました。型レベルリストを用いることで型の集合を表現し、Trait Resolutionを使って特定の型が集合に含まれるかどうか、またある集合が別の集合を含むかどうかを判定することができます。Index トリックを使うことで、型レベルリストに対する再帰処理を行うことができ、元の判定から、部分集合の判定や等価性の判定など、様々な処理を実現することができます。

この記事では取り上げませんでしたが、型レベルリストに対するhigher-order functionを実装することで、型リストに対するfoldを処理を実装することも可能です。こうしたトリックを用いることで、Rustの型システムを活用し、DSLや型安全なAPIを実現することができます。興味がある方は、お気軽にご連絡ください。

最後に、このテクニックを用いてChoRusというRustで振り付けプログラミングをするライブラリを開発しています。振り付けプログラミングでは、分散システムにおける複数のノードの振る舞いを Choreography という1つのプログラムで記述し、Endpoint Projection という手法で各ノードのプログラムを生成します。こうして記述されたプログラムは、通信のミスマッチやデッドロックが発生しないことが知られており、安全な分散システムの構築に役立ちます。ChoRusは、このような振り付けプログラミングをRustで実現するためのライブラリです。興味がある方は、以下のリンクから詳細をご覧ください。

https://lsd-ucsc.github.io/ChoRus/

https://github.com/lsd-ucsc/ChoRus