Zenn
🗓

[Rust] 型推論実装してみた

2025/02/16に公開
Rust
型推論
自作言語
typeinference
tech

しばらく間が空きましたが、 Rust で作るプログラミング言語シリーズです。

https://www.amazon.co.jp/dp/4297141922

モダンな言語の多くは型推論の機能を備えています。型推論はコード量を抑えるだけではなく、リファクタリング時に必要な変更を減らしてくれる有用な機能です。静的型付け言語では特に効果が高いといえます。

実装は Mascal 言語に対して行っています。

https://github.com/msakuta/mascal?tab=readme-ov-file

型推論とは

型推論とは省略した型宣言を自動的に補完する機能です。 Rust プログラマなら日常的に使っているものです。最も簡単な例は次のようなものでしょう。

var i: i32 = 0;
var j = i;

ここで i の型は i32 と宣言されていますが、 j は省略されています。ところが、 j は i で初期化されているため、型は同じであると推論できます。また、 i を初期化している数値リテラル 0 ですが、この型も i32 と i64 のどちらでもあり得ます。こちらも i から i32 と推論できます。

var i: i32 = 0i32;
var j: i32 = i;

さらに、型宣言の順番が前後しても、一意に決定できれば推論できます。

var i = 0;
var j: i32 = i;

基本的には、型の制約を満たすように型の選択肢を絞り込むのが型推論だと言えます。数独パズルを解くときに考えることに似ています。各マスに入れても整合性が保てる選択肢を絞り込んでいくのです。

関数型言語では有名な Hindley-Milner 型推論というものがありますが、ここでは扱いません[1]。

また、関数を跨いだ推論は行わず、関数宣言の型は完全であることを必要とします。これも Rust と同じです[2]。

動的型付けや Duck typing との違い

型推論は動的型付けや Duck typing とは似て非なるものです。似てもいないかもしれませんが、「型を書かなくてもよい」という意味では初心者は混同してしまいそうです。

動的型付けは変数がどんな型でも取れるという意味で型を書かなくても良いのですが、実行時に予期しない型によるエラーが生じえます。大規模な開発では特にエラーの原因がわかりにくくなります。型推論はあくまでも静的型付けの補佐であり、型は一意に定まります。

Duck typing は静的型付けでも使われることがありますが、複数の型を許すという意味で型推論とは異なります。典型例が C++ のテンプレートです。動的型付けと違い、実行時エラーにはなりませんが、具体的な型を代入した時(C++ テンプレートではインスタンス化と呼ばれます)、その型によってコンパイルエラーになることがあります。例えば下記の C++ テンプレートでは、 T という型の値 val に 2 を足して返しています。数値型であればどんな T でも動作しますが、数値リテラルとの operator+() がオーバーロードされていない型である場合はコンパイルエラーになります。

template<typename T>
T add_2(T val) {
    return val + 2;
}

これは T の型に依存し、このテンプレート自体ではコンパイル可能かどうかは判断できません。これは多数の型パラメータを持った複雑なテンプレートでは問題になりやすく、エラーメッセージが非常に読みにくくなる原因にもなります。

ポリモーフィックとモノモーフィック

ポリモーフィック(多相)な型推論は、型の選択肢を絞り込んだ時に2つ以上の選択肢があっても許すものです。関数型言語では一般的です。複数の選択肢が残った場合は、 Rust でいうジェネリックスや C++ でいうテンプレートに相当します。しかし、 Rust も C++ もモノモーフィックな型推論しかしていません[3]。実装は難しい傾向にあります。気取った言い方をするとパラメトリック多相となります。

モノモーフィック(単相)な型推論は、型を絞り込んだ結果がただ一つの可能性である場合のみを許すものです。ネイティブコンパイル言語では厳密な型を制御するためにあえてモノモーフィックな型推論をするものが多い印象です。ただし、 Rust や C++ ではジェネリックスやテンプレートで明示的にポリモーフィックにすることは可能であり、型推論の対象ではないというだけです。本稿ではモノモーフィック型推論を対象にします。

実装方針

型推論は型チェッカーの拡張と言えます。型チェッカーは AST を変更はせず、宣言された型に整合性があることを確認するだけですが、型推論は整合性が保てるように型の穴を埋めるアルゴリズムと言えます。埋める型が一つに定まらないときは型チェックエラーとなります。このため、前から実装していた型チェッカーに手を加えて型推論器にすることにします。

このため、型推論に慣れていない方は、まずは型チェッカーの実装をお勧めします。型推論のほうが難易度は数段高いです。型チェッカーの実装方法については、拙著を参考にしてください(宣伝)。

型集合

型推論を実装する前に、型が持ちうる可能性を表現する必要があります。単純に考えると、次のようにそれぞれの型のフラグを集めた構造体で表現できそうです。

pub struct TypeSet {
    pub i32: bool,
    pub i64: bool,
    pub f32: bool,
    pub f64: bool,
    pub void: bool,
    pub string: bool,
}

void が有効な型の一つとして扱われていることに注目してください。これは値を返さない関数 fn () -> void の返り値を表現するために必要です。

型を省略した場合は TypeSet::any() として全てのフラグをオンにした型集合で表現できそうですが、実はそう簡単ではありません。

型集合の問題点

上記のデータ構造は単純なプリミティブ型には十分ですが、入れ子に構造化された型が出現すると崩壊します。例えば、 TypSet に次のような配列型を追加したと考えてみてください。

pub struct TypeSet {
    // ...
    pub array: Option<Box<TypeSet>>,
}

これだと配列の要素の型が一つしか持てません。任意の型 TypeSet::any() を表現するには、配列の要素についても複数の可能性を表現できるようにする必要があります。

では、次のように複数の要素型を持てるようにすればよいでしょうか。

pub struct TypeSet {
    // ...
    pub arrays: Vec<TypeSet>,
}

しかしこれでもダメです。配列の要素もまた TypeSet であるため、その any() を再帰的に呼び出すと、無限再帰を起こします。

「任意の型」を表現するには、美しくはないですが次のように Any 型を特別に扱う必要があります。

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub enum TypeSet {
    /// "Any" type set is used to avoid infinite recursion of nested array types.
    #[default]
    Any,
    Set(TypeSetFlags),
}

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct TypeSetFlags {
    pub i32: bool,
    pub i64: bool,
    pub f32: bool,
    pub f64: bool,
    pub void: bool,
    pub string: bool,
    pub array: Option<Box<TypeSet>>,
}

ここで、 array の要素型は一種類に限定しています。これは、型推論器の内部で複数の配列要素型を考慮する部分がないからです(TypeScript のように Array<a | b> のような型宣言はできません)。もし複数の競合する要素型が型推論に現れたら、その時点でエラーを返すことにします。もちろん、これは後ほど型システムを改善すれば変わる可能性はあります。

省力された型宣言の型付け

さて、ここまででデータ構造の準備ができたので、パーサをそれに対応させます。 TypeDecl は既存の型ですが、省略されたときは TypeDecl::Any を返すことにします。

pub(crate) fn type_spec(input: Span) -> IResult<Span, TypeDecl> {
    let (r, type_) = opt(delimited(ws(char(':')), type_decl, multispace0))(input)?;
    Ok((r, type_.unwrap_or(TypeDecl::Any)))
}

TypeDecl は TypeSet との棲み分けが分かりにくいですが、「宣言された型」を表します。例えば「i32 と i64 のどちらか」のような型宣言はできないので、それは表現できないようになっています。宣言が省略された場合は Any となります。

#[derive(Debug, PartialEq, Clone)]
#[repr(u8)]
pub enum TypeDecl {
    Any,
    F64,
    F32,
    I64,
    I32,
    Str,
    Array(Box<TypeDecl>, ArraySize),
    Tuple(Vec<TypeDecl>),
}

また、構文木においては、数値リテラルは型集合を持てるようにします。

#[derive(Debug, PartialEq, Clone)]
pub(crate) enum ExprEnum<'a> {
    NumLiteral(Value, TypeSet),
    // ...
}

数値リテラル向けに int() と float() という TypeSet のコンストラクタを用意しておきます。これは2つの型のフラグをセットした型集合を構築します。

impl TypeSet {
    pub fn int() -> Self {
        Self::Set(TypeSetFlags {
            i32: true,
            i64: true,
            ..TypeSetFlags::default()
        })
    }

    pub fn float() -> Self {
        Self::Set(TypeSetFlags {
            f32: true,
            f64: true,
            ..TypeSetFlags::default()
        })
    }
}

数値リテラルの値は I64 または F64 で保持しますが、型宣言は int() と float() を使います。

fn decimal_value(i: Span) -> IResult<Span, (Value, Span)> {
    let (r, v) = recognize(pair(opt(one_of("+-")), decimal))(i)?;
    let parsed = v.parse().map_err(/* ... */)?;
    Ok((r, (Value::I64(parsed), v)))
}

fn float_value(i: Span) -> IResult<Span, (Value, Span)> {
    let (r, v) = float(i)?;
    let parsed = v.parse().map_err(/* ... */)?;
    Ok((r, (Value::F64(parsed), v)))
}

fn double_expr(input: Span) -> IResult<Span, Expression> {
    let (r, (value, value_span)) = alt((float_value, decimal_value))(input)?;
    let ts = match value {
        Value::I64(_) => TypeSet::int(),
        _ => TypeSet::float(),
    };
    Ok((
        r,
        Expression::new(ExprEnum::NumLiteral(value, ts), value_span),
    ))
}

変数の型集合の記憶

型は変数に結び付けられますが、型推論によって宣言時に明示しなくても良いことになっているので、型集合として扱う必要があります。
例えば、下記の s と i の型は宣言時には決定せず (TypeSet::Any)、後で決定します。

var s;
var i;
s = "str";
i = 0i32;

型と変数の関連は TypeCheckContext に HashMap として覚えておきます。これも型チェッカーと同じです。

pub struct TypeCheckContext<'src, 'native, 'ctx> {
    /// Variables table for type checking.
    variables: HashMap<&'src str, TypeSet>,
    // ...
}

型の伝搬

次に、 AST を辿って型を伝搬させるロジックを実装します。型チェッカーでも同じように実装しましたが、型推論では型集合を伝搬させていくところが異なります。

ここで、2つの AST ツリーの型集合の積集合を求める演算が頻繁に登場するので、 TypeSet に次のようなメソッドを用意しておきます。実装はかなり複雑なので省略します。詳細はコードをご覧ください。

impl TypeSet {
    pub fn try_intersect(&self, rhs: &Self) -> Result<Self, String>;
}

論理的には BitAnd トレイトを実装することで typeset_a & typeset_b のように簡潔に書けるのですが、型の積集合が空集合だったり、許されていない組み合わせだった場合は、エラーとして早期リターンするため、メソッドとしています。使い方は typeset_a.try_intersect(&typeset_b)? のようになります。

型の伝搬は次のような関数で再帰的に行います。 tc_expr_forward は式の型伝搬、 tc_stmt_forward は文の式伝搬を行い、お互いに再帰呼び出しをしながら AST を辿りますが、詳細は長いわりに本質的でないので省略します。

fn tc_expr_forward<'src, 'b, 'native>(
    e: &'b Expression<'src>,
    ctx: &mut TypeCheckContext<'src, 'native, '_>,
) -> Result<TypeSet, TypeCheckError<'src>>
where
    'native: 'src,
{
    Ok(match &e.expr {
        // ...
    })
}

fn tc_stmt_forward<'src, 'ast, 'native>(
    stmt: &'ast Statement<'src>,
    ctx: &mut TypeCheckContext<'src, 'native, '_>,
) -> Result<TypeSet, TypeCheckError<'src>>
where
    'native: 'src,
{
    let mut res = TypeSet::all();
    match stmt {
        // ...
    }
}

型の逆伝搬

型チェッカーとは異なるのが、「逆伝搬」のフェーズです。冒頭に挙げた通り、型は後での「使用」に基づいて絞り込める場合がありますので、 AST を逆方向に辿る必要があります。この過程で、型が絞り込めたら結果を AST に書き戻すため、可変参照を引数に取ります。もちろん AST のコピーを構築して返すというインターフェースにすることもできますが、 AST のほとんどはそのままコピーで、型集合をちょっといじるだけなので、かなり無駄なコピーになります。

pub fn tc_stmt_reverse<'src, 'ast, 'native>(
    stmt: &'ast mut Statement<'src>,
    ts: &TypeSet,
    ctx: &mut TypeCheckContext<'src, 'native, '_>,
) -> Result<(), TypeCheckError<'src>>
where
    'native: 'src,
{
    match stmt {
        // ...
    }
}

fn tc_expr_reverse<'src, 'b, 'native>(
    e: &'b mut Expression<'src>,
    ts: &TypeSet,
    ctx: &mut TypeCheckContext<'src, 'native, '_>,
) -> Result<(), TypeCheckError<'src>>
where
    'native: 'src,
{
    let span = e.span;
    match &mut e.expr {
        // ...
    }
}

エントリポイントの定義

最後に、型推論のエントリポイントとして type_check 関数を定義します。これは元々型チェッカーを改造しているので名前はそのままになっていますが、引数に AST への可変参照を取ります。

この関数では何度か AST をスキャンします。まずは宣言された関数の型を拾うために一度スキャンします。次に、関数の内部のコードに型推論を適用します。最後に、トップレベルのコードに型推論を適用します。

pub fn type_check<'src, 'ast, 'native>(
    stmts: &'ast mut Vec<Statement<'src>>,
    ctx: &mut TypeCheckContext<'src, 'native, '_>,
) -> Result<(), TypeCheckError<'src>>
where
    'native: 'src,
{
    for stmt in stmts.iter_mut() {
        match stmt {
            Statement::FnDecl {
                // ...
            } => {
                ctx.functions.insert(
                    // ...
                );
            }
        }
    }

    for stmt in stmts.iter_mut() {
        match stmt {
            Statement::FnDecl {
                // ...
            } => {
                let mut inferer = TypeCheckContext::push_stack(ctx);
            }
        }
    }

    tc_stmts_forward(stmts, ctx)?;
    tc_stmts_reverse(stmts, &TypeSet::all(), ctx)?;
}

代入演算子と左辺値

上記の前進と後退の2パススキャンはとてもよく機能しますが、問題が一つあります。それは集成体型への代入です。(集成体型については前に書いた記事を参照してください)

たとえば、次のコードでは a を配列とし、その要素を整数として型制約を掛けたいです。

a[0] = 1;

ところが、前進のスキャンだけではその型制約を a に届かせることができません。これは代入の左辺が次のように構文解析され、変数名の識別子が構文木のサブツリーに押し込められてしまうためです。

((a)[0]) = 1

左辺はネストした式になりうるため、任意の複雑度を持ちます。

a[0][1] = 1;

さらに左辺値を取る式は次のようにタプルや関数呼び出しがチェインする可能性もあります。

f().a.0[1] = 1;

後退のパスだけでもうまくいきません。右辺も任意の式になりうるからです。

a[0] = 1 + 2;

解決策は、前述の記事のように左辺値を辿る関数を別に定義して代入演算子の左辺に適用することです。

enum VarRef<'src> {
    Variable(&'src str),
    Array(Box<VarRef<'src>>),
    Tuple(Box<VarRef<'src>>, usize),
}

fn forward_lvalue<'src, 'b, 'native>(
    ex: &'b Expression<'src>,
    ctx: &mut TypeCheckContext<'src, 'native, '_>,
) -> Result<Option<(VarRef<'src>, TypeSet)>, String>;

VarRef 構造体は集成体型の要素へのパスとして機能します。これを右辺を前進のパス (tc_expr_forward) で評価した結果と try_intersect し、結果を変数テーブルに反映します。

fn tc_expr_forward<'src, 'b, 'native>( /* ... */ ) {
    Ok(match &e.expr {
        // ...
        ExprEnum::VarAssign(lhs, rhs) => {
            let span = e.span;
            if let Some((var_ref, decl_ty)) = forward_lvalue(lhs, ctx)
                .map_err(|err| TypeCheckError::new(err, span, ctx.source_file))?
            {
                let rhs_ty = tc_expr_forward(rhs, ctx)?;
                let ty = decl_ty
                    .clone()
                    .try_intersect(&rhs_ty)
                    .map_err(|err| TypeCheckError::new(err, span, ctx.source_file))?;
                ctx.set_var_type(&var_ref, &ty);
            }

            binary_op(&lhs, &rhs, e.span, ctx, "Assignment")?
        }
        // ...
    })
}

関数境界を跨いだ型推論

本稿では実装しませんが、関数境界を跨いだ型推論をしようと思ったらさらに複雑になります。前進と後退の2パスでは不十分で、関数間の依存関係を元に依存先を最初に推論するように順番を調整する必要があります。並び変えはトポロジカルソートでできますが、この依存関係グラフは再帰呼び出しを含むときは循環となり、依存関係順に並び変えられない可能性もあります。

ここに至ると Hindley-Milner 型推論が必要になると思います。しかし、 Rust でこれを実装するのはまた骨が折れます。 Hindley-Milner では型制約を任意の構文木のノード間の関係として抽出し(型環境と呼ばれます)、解決するので、構文木の任意のノードへの可変参照を持てるようにする必要があります。これは借用ができないので Rc<RefCell<_>> などで包む必要があります。しかし Rc<RefCell<_>> はパフォーマンスの面ではあまり魅力的ではありません。

このため、ノードを配列に確保しインデックスで参照する相対アドレッシングにデータ構造を変更したほうがよいかもしれません。このアイデアについては こちらのリポジトリで単純な構文について実験していますが、これは Mascal のようなスケールの言語に適用するのはなかなか大きな改造になりそうです。

型推論結果のフィードバック

型推論器はチェックするだけではなく、推論結果として AST へ変更を加えます。これを可視化するため、コマンドラインツールに -a を指定した時は次のように前後の推論結果を表示するようにしました。

fn fact(n: i64) -> i64 {
  if (n < 1(i32|i64)) {
    1(i32|i64)
  } else {
    (n * fact((n - 1(i32|i64))))
  }
}
print(fact(5(i32|i64)));
AST after type inference:
fn fact(n: i64) -> i64 {
  if (n < 1i64) {
    1i64
  } else {
    (n * fact((n - 1i64)))
  }
}
print(fact(5(i32|i64)));

また、 Wasmバージョン では「Type Infer/Check」ボタンで結果を表示します。

img

最終的には LSP を実装してエディタで inlay hint として表示したいところです。

まとめと感想

  • 基本的な型推論を実装することができました。
  • 型チェッカーの自然な拡張として実装することができます。
  • 難易度は型チェッカーよりも数段高いです。
  • 前進と後退の2パスで実装しましたが、部分的に前後を切り替えないと行けない部分があり、 AST を無駄に辿ることがあります。これは今後の最適化を考慮する価値があります。
  • 関数境界を跨いだ型推論をしようと思ったら、さらに困難になります。
  • これ以上は Hindley-Milner 型推論を考慮すべきと思います。
脚注

  1. Hindley-Milner はポリモーフィック(多相)であることと、関数をまたいだ型推論が可能であることが特徴です。関数型言語では特に人気のある型推論手法です。しかしながら、ここでは Rust レベルの単純な型推論を実装します。とはいえ、考え方はここで示した「型制約を満たす型を絞り込む」というもので変わりありません。 ↩︎

  2. 厳密に言うと、 Rust でもラムダや return impl 構文で型の名前を明示しない関数宣言が可能な場所はありますが、基本的な思想としては関数の境界のようなインターフェースでは型を明示するという言語設計になっています。 ↩︎

  3. C++ はかなりベーシックな型推論しかできませんが、時折 Rust よりも簡潔に記述できるときがあります。式のコンテキストが構造体だった場合、構造体名を省略した aggregate initializer を記述することができたりします。 (e.g. vetor.push_back({x, y}))) また、 auto 宣言することで大抵の型を省力できます。これは複雑な型のイテレータを宣言するときなどに重宝します。 ↩︎

GitHubで編集を提案

Discussion