Rustで作る自作言語(8) Let多相と向き合う

型推論のバグ

このようなバグがあることに気付きました。

>>let f(x) = {let g(y) = [x, y]; g;};
>>:t f
(a6) -> (a7) -> Vector[a7]

本来、fの型はa7 -> a7 -> Vector[a7]のようになるはずですが、１つ目の引数と２つ目の引数の型が異なってしまっています。そのため、型安全性が破壊されてしまっています。

>>f(1, 1==1)
Ok(VVector([VInt(1), VBool(true)]))

これは、現状だとlet束縛する時にすべての型変数を多相化してしまっているためです。具体的に動作を追ってみると、

let f(x) = ...

まず、fに型変数_a1(数字に意味はない)を割り当てます。_から始まる型変数は全称量化されてない型変数を意味します。そして、xに型変数_a2を割り当てます。

let f(x) = {
    let g(y) = ...

gに型変数_a3を割り当てます。yに型変数_a4を割り当てます。

    let g(y) = [x, y];

[x, y]はVectorなので、まず最初に型変数_a5を生成し、Vectorの中身のx, yの型に対して単一化します。すると、_a5=_a2=_a4となります。
gがletで束縛された時点で、型変数はgeneralizeされるので、gの型はa1->Vector[a1]となります。

let f(x) = {
    ...
    g;
};

gが使われたので、gの型a->Vector[a]がinstantiateされ、_a6->Vector[_a6]となります。よってfの型は、_a2->_a6->Vector[_a6]がgeneralizeされた、a2->a3->Vector[a3]となります。

これを解消するために、levelという概念を導入します。

level base多相型推論

ここから先は次の記事の焼き直しですので、詳しくは以下の記事を読むことをおすすめします。
https://rhysd.hatenablog.com/entry/2017/12/16/002048
levelとは、どれだけletの内側にいるかという値です。

//level 0
let f(x) =
    //level 1
    {
    let g(y) =
        //level 2
        [x, y];
    //level 1
    g;
    };

ある型変数が作られた時、その型変数が作られた時のlevelが型変数のパラメータとして付与されます。例えば次のようなシンプルな多相関数について考えてみます。

let id(x) = x;

ここで、まずidに_a1(0) (括弧の内側はlevelを表すとする)を割り当てます。xには_a2(1)を割り当てます。なぜxはlevel1になるかというと、

let id = fn x -> x;

let id(x) = x;は上の構文糖衣だからです。
そして、letでidに束縛する時に型がgeneralizeされるのですが、この時全称量化されるのは、letのある場所のlevelよりも高いlevelの型変数のみです。
fn x -> xの型は_a2(1) -> _a2(1)で、letがあるのはlevel 0なので、generalizeされ、idの型はa->aとなります。

次は、元のプログラムでバグが発生していた例で、どのように改善されるのかを見ていきましょう。

let f(x) = {
    let g(y) = ...

fには_a1(0)が、xには_a2(1)が、gには_a3(1)が、yには_a4(2)が割り当てられます。

    let g(y) = [x, y];

[x, y]では、Vectorなので新しい型変数_a5(2)が作られ、x, yの型と単一化されます。
つまり、_a5(2)=_a2(1)=_a4(2)となります。ここで、level base多相型推論での単一化のルールに、型変数同士で単一化する時は型変数のlevelを低い方に合わせるというものがあるので、_a2(1)のレベルに合わせて、_a5(2)は_a5(1)に、_a4(2)は_a4(1)になります。これによって、letでgに束縛する時、letはlevel 1の場所にあるので、gの型_a4(1) -> Vector[_a4(1)]はgeneralizeされません。

let f(x) = {
    ...
    g;
};

最後に、fの型_a2(1)->_a2(1)->Vector[_a2(1)] (_a5(1)=_a2(1)=_a4(1)であるため)は、letのある場所よりlevelが高いのでgeneralizeされ、a -> a -> Vector[a]となり、正しい型となります。

実装

    TVar(u64, Rc<RefCell<u64>>, Rc<RefCell<Option<Type>>>),

TVarにlevelを持たせます。Rc<RefCell<u64>>となっているのは、後にunifyでlevelを変更することがあるからです。

pub struct TypeInfer {
    env: Rc<RefCell<TypeEnv>>,
    unassigned_num: u64,
    level: u64,
}

同様に、TypeInferの方にもlevelを追加しておきます。

impl TypeInfer {
    pub fn newTVar(&mut self) -> Type {
        let ty = Type::TVar(
            self.unassigned_num,
            Rc::new(RefCell::new(self.level)),
            Rc::new(RefCell::new(None)),
        );
        self.unassigned_num += 1;
        ty
    }
}

newTvarで新しい型変数を作る時に、TypeInferのlevelをTVarに持たせます。

impl TypeInfer {
    fn typeinfer_expr_levelup(&mut self, ast: &Expr) -> Result<Type, TypeInferError> {
        self.level += 1;
        let ty = self.typeinfer_expr(ast)?;
        self.level -= 1;
        Ok(ty)
    }
}

let束縛の右辺の式の型推論に用いられます。

fn unify(t1: &Type, t2: &Type) -> Result<(), TypeInferError> {
    match (t1.simplify(), t2.simplify()) {
        ...
	(Type::TVar(n1, level1, t1), t2) => {
            if occur(n1, &t2) {
                Err(TypeInferError::OccurError(n1, t2))
            } else {
                level_balance(&Type::TVar(n1, level1, Rc::clone(&t1)), &t2);
                *(*t1).borrow_mut() = Some(t2);
                Ok(())
            }
        }
        ...
    }
}

level_balance関数で、左右の型が両方とも型変数である時に、levelを低い方に合わせています。
level_balance関数の定義はこんな感じ。

fn level_balance(t1: &Type, t2: &Type) {
    match (t1.simplify(), t2.simplify()) {
        (Type::TVar(_, level1, _), Type::TVar(_, level2, _)) => {
            let level1num: u64 = *Rc::clone(&level1).borrow();
            let level2num: u64 = *Rc::clone(&level2).borrow();
            if level1num > level2num {
                *level1.borrow_mut() = level2num;
            } else {
                *level2.borrow_mut() = level1num;
            }
        }
        (_, _) => (),
    }
}