Rustで単一化を書くスマートな方法を考えた

課題

Rustでシンプルな単一化を書くことを考えます。単一化は主に型推論の実装に用いられます。

ここでは以下の方針で実装します。

• 一階の単一化。
• 変数は非負整数のidで表現し、0から順に付番する。
• 変数以外の項はRustのADTを使った通常の木構造で表現し、ノードの共有は行わない。
• 変数参照の縮約は行わない。

この方針のもと実装すると、およそ以下のようなコードになります。

#[derive(Debug, Clone)]
pub struct UnifyEnv {
    vars: Vec<Option<Type>>,
}

#[derive(Debug)]
pub struct UnificationFailure;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Type {
    Var(usize),
    Constr { constr: u32, args: Vec<Type> },
}

impl Type {
    pub fn unify(&self, other: &Self, u: &mut UnifyEnv) -> Result<(), UnificationFailure> {
        // ここでcloneしないとエラーになる
        // let ty1 = self.resolve(u);
        // let ty2 = other.resolve(u);
        let ty1 = &self.resolve(u).clone();
        let ty2 = &other.resolve(u).clone();
        match (ty1, ty2) {
            (Type::Var(v1), Type::Var(v2)) if v1 == v2 => Ok(()),
            (Type::Var(v), _) => {
                if ty2.has_fv(*v, u) {
                    return Err(UnificationFailure);
                }
                u.vars[*v] = Some(ty2.clone());
                Ok(())
            }
            (_, Type::Var(v)) => {
                if ty1.has_fv(*v, u) {
                    return Err(UnificationFailure);
                }
                u.vars[*v] = Some(ty1.clone());
                Ok(())
            }
            (
                Type::Constr {
                    constr: c1,
                    args: args1,
                },
                Type::Constr {
                    constr: c2,
                    args: args2,
                },
            ) => {
                if c1 != c2 || args1.len() != args2.len() {
                    return Err(UnificationFailure);
                }
                for (arg1, arg2) in args1.iter().zip(args2.iter()) {
                    arg1.unify(arg2, u)?;
                }
                Ok(())
            }
        }
    }

    pub fn resolve<'a>(&'a self, u: &'a UnifyEnv) -> &'a Self {
        let mut ty = self;
        loop {
            match ty {
                Type::Var(v) => {
                    if let Some(resolved) = &u.vars[*v] {
                        ty = resolved;
                        continue;
                    }
                }
                _ => {}
            }
            return ty;
        }
    }

    fn has_fv(&self, v: usize, u: &UnifyEnv) -> bool {
        match self.resolve(u) {
            Type::Var(j) => v == *j,
            Type::Constr { constr: _, args } => args.iter().any(|ty| ty.has_fv(v, u)),
        }
    }
}

このコードでは余計なcloneが発生してしまいます。

なぜcloneが必要になるのか

以下の2種類の借用が競合していると判定されるためです。

• 変数を展開するための読み取り専用借用
• 変数に値を代入するための書き込み借用

実際にはこの2つのアクセスが同じ変数に対して同時に起こることはありません。前者は変数が代入済みのときのみ、そして後者は変数が未代入のときのみ発生するものだからです。

解決方法

この問題を解決する方法は色々あります。すぐ思いつくのは以下のようなパターンです。

• cloneコストを許容する
• ノードをArc等で共有することでdeep cloneのコストを削減する
• Var(_)に内部可変性を導入することで共有書き換えを可能にする

しかしここでは UnifyEnv と Type の定義はそのままで、このアルゴリズムの正当性をコンパイラに納得させることを考えます。こうすることで単一化アルゴリズム以外の場面でユーザーに強いる負担を最小限に抑えることを狙います。

OnceCell

ここで出てくるのがOnceCellです。OnceCellは以下の規約を表現する内部可変性コンテナです:

• 未初期化の場合は、値を1回だけ書き込むことができる。
• 初期化済みの場合は、その値に対する読み取り専用参照を取得できる。

これは単一化アルゴリズムが要求する規約とそっくりです。また、OnceCellは !Sync ですが、これは単一化アルゴリズムが単純並列化できない (並列化すると同じ変数に同時書き込みが発生する可能性がある) ことと整合しています。

ただ、これをそのまま UnifyEnv に導入しても単一化は実装できますが、 UnifyEnv 自体が !Sync になってしまい、ユーザーにとってはやや不便です。

OnceCellとレイアウト

ここでは UnifyEnv をいじらずに単一化アルゴリズム内でのみOnceCellを使うことを考えます。ここでOnceCellの定義を確認すると、以下のようになっていることがわかります。

pub struct OnceCell<T> {
    inner: UnsafeCell<Option<T>>,
}

これに #[repr(transparent)] さえついていれば Option<T> とレイアウト上等価であることがわかります。加えて、 OnceCell<T> のownership invariantは Option<T> のownership invariantと等価です。

※ownership invariant ... Rustのメモリ管理意味論を分離論理で定義する際、型の意味論は2つの異なる述語の組として表される。ひとつは T や &mut T として出現したときの規約をあらわす不変条件 (ownership invariant) で、もうひとつは Arc<T> や &T として出現したときの規約をあらわす不変条件 (shared invariant) である。

このようにレイアウト等価かつownership invariantが等価であれば、 &mut 参照に対して広範なtransmute変換が許容されます。実際、CellにはCell::from_mutとCell::as_slice_of_cellsというメソッドがあり、 &mut [T] から &[Cell<T>] を作り出すような変換が許容されています。

Option<T> を基本の表現として使うことで、単一化以外の場面 (たとえばバックトラック処理の実装など) ではOnceCellの記述の冗長さから解放されます。

OnceCellを使った単一化

以上のことから、OnceCellを改造して Option<T> とレイアウト互換にしたものを用意すれば単一化の実装に使えそうです。

内容自体は簡単なものですが、unsafe実装を繰り返すのはあまりよくないので自作したもの(OptionCell)をpublishしておきました。

https://crates.io/crates/option-cell

https://github.com/qnighy/option-cell

一応軽くテストを書いてMIRIで検証してはいる他、実装は既存のOnceCell/Cellを参考にしているので実装に問題がある可能性は低いと思います。

さて、これを使って先ほどの単一化の実装を書き直すと以下のようになります。ポイントは以下の通りです。

• OnceCell (OptionCell) は公開APIには出現せず、内部的にはOnceCellを持ち回すために専用の再帰メソッド (unify2) が定義されている。
• 変数を展開するresolveメソッドに関しては、公開用のものとは別に単一化アルゴリズム中で利用するためのresolve2が再定義されている。

use option_cell::OptionCell;

#[derive(Debug, Clone)]
pub struct UnifyEnv {
    vars: Vec<Option<Type>>,
}

#[derive(Debug)]
pub struct UnificationFailure;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Type {
    Var(usize),
    Constr { constr: u32, args: Vec<Type> },
}

impl Type {
    pub fn unify(&self, other: &Self, u: &mut UnifyEnv) -> Result<(), UnificationFailure> {
        self.unify2(other, OptionCell::from_mut_slice(&mut u.vars))
    }
    fn unify2(&self, other: &Self, vars: &[OptionCell<Type>]) -> Result<(), UnificationFailure> {
        let ty1 = self.resolve2(vars);
        let ty2 = other.resolve2(vars);
        match (ty1, ty2) {
            (Type::Var(v1), Type::Var(v2)) if v1 == v2 => Ok(()),
            (Type::Var(v), _) => {
                if ty2.has_fv(*v, vars) {
                    return Err(UnificationFailure);
                }
                vars[*v].set(ty2.clone()).unwrap();
                Ok(())
            }
            (_, Type::Var(v)) => {
                if ty1.has_fv(*v, vars) {
                    return Err(UnificationFailure);
                }
                vars[*v].set(ty1.clone()).unwrap();
                Ok(())
            }
            (
                Type::Constr {
                    constr: c1,
                    args: args1,
                },
                Type::Constr {
                    constr: c2,
                    args: args2,
                },
            ) => {
                if c1 != c2 || args1.len() != args2.len() {
                    return Err(UnificationFailure);
                }
                for (arg1, arg2) in args1.iter().zip(args2.iter()) {
                    arg1.unify2(arg2, vars)?;
                }
                Ok(())
            }
        }
    }

    pub fn resolve<'a>(&'a self, u: &'a UnifyEnv) -> &'a Self {
        let mut ty = self;
        loop {
            match ty {
                Type::Var(v) => {
                    if let Some(resolved) = &u.vars[*v] {
                        ty = resolved;
                        continue;
                    }
                }
                _ => {}
            }
            return ty;
        }
    }
    fn resolve2<'a>(&'a self, vars: &'a [OptionCell<Type>]) -> &'a Self {
        let mut ty = self;
        loop {
            match ty {
                Type::Var(v) => {
                    if let Some(resolved) = vars[*v].get() {
                        ty = resolved;
                        continue;
                    }
                }
                _ => {}
            }
            return ty;
        }
    }

    fn has_fv(&self, v: usize, vars: &[OptionCell<Type>]) -> bool {
        match self.resolve2(vars) {
            Type::Var(j) => v == *j,
            Type::Constr { constr: _, args } => args.iter().any(|ty| ty.has_fv(v, vars)),
        }
    }
}

まとめ

単一化アルゴリズムにおけるメモリアクセスの排他性はやや特殊なパターンで、よく知られたRustのプリミティブで表現するのが難しく感じられますが、実はOnceCellを用いることで整理できることを説明しました。また、内部可変性を使っているアルゴリズムであっても、互換なレイアウトと等価なownership invariantを持つ別の型で状態をラップすることで、ユーザー側から内部可変性を隠蔽しつつ、同期のための追加コストも回避することができることを説明しました。