🐍

Pythonのコンパイラを作りたい #4 - 関数定義とスコープ

に公開
Python
LLVM
compiler
tech

こんにちは。前回（「#3 - 数値演算と型システム」）では、Python の int をネイティブの i32 として処理し、x + y などの数値演算を高速に実行する方法についてお話ししました。
このアプローチによって、Python的には動的型付けの要素があるのに、あえて静的に型 (i32) を決め打ちすることでコンパイラ構造を大幅に単純化し、かつ LLVM の最適化恩恵を得られるようにしていました。

今回のテーマは、関数定義とスコープの扱いについてです。
Python の関数といえば、動的言語ならではの柔軟さが大きな特徴ですが、このプロジェクトでは「静的な型注釈に基づいて LLVM IR の関数定義を作る」というアプローチを取ることで、シンプルな実装を目指しています。「この関数は i32 を受け取って i32 を返す」とわかっていれば、ほぼ C 言語の int foo(int x) のように扱えるわけです。

1. 関数定義の AST ノード

1-1. FunctionDef ノードの構造

Python で以下のような関数を書くとしましょう。

def foo(x: int, y: int) -> int:
    return x + y

Python の標準ライブラリ ast を使って ast.parse() すると、FunctionDef ノードとして解析され、典型的には下記のような構造になっています。

FunctionDef(
  name='foo',
  args=Arguments(
    args=[
      arg(arg='x', annotation=Name(id='int')),
      arg(arg='y', annotation=Name(id='int'))
    ],
    vararg=None,
    kwarg=None,
    ...
  ),
  returns=Name(id='int'),
  body=[
    Return(value=BinOp(left=Name(id='x'), op=Add(), right=Name(id='y')))
  ],
  ...
)

ここで

  • name='foo' は関数名
  • args.args に引数 x と y が並んでおり、それぞれ annotation=Name(id='int')
  • returns=Name(id='int') で戻り値が int であることを示している
  • body に Return(...) などのステートメントが入っている

という構造になっています。

1-2. 型アノテーションの重要性

前回お話ししたとおり、このコンパイラ pyc では「引数・戻り値が int である」とアノテートされている場合、それを「ネイティブ i32」としてマッピングします。
そのおかげで、foo(x, y) は「i32 を受け取り、i32 を返す関数」と見なせるため、静的型付き言語に近い形で LLVM IR の関数を組み立てられるわけです。
逆にアノテーションが書かれていなかったり、未知の型が指定されている場合は、現状では「動的型 (ptr) に落とす」「ビルドエラー扱いにする」などの制限が発生します。今後の拡張によっては、より柔軟な対応が可能になるかもしれませんが、まずは単純化を優先しています。

2. シンボルテーブルとスコープ管理

2-1. シンボルテーブルの概要

Python に限らず、多くのコンパイラでは「変数名や型情報などをどこに記録しておくか」が大きな問題になります。本プロジェクト pyc では、BaseVisitor や StmtVisitor と呼ばれるクラス内に、以下のような簡易シンボルテーブルを用意しています。

  • self.symbol_table = {}
  • 辞書のキーが変数名、値が型文字列 ("i32", "ptr", など)
  • 関数定義やブロックを訪問するときに、必要に応じてシンボルテーブルを更新する

こうした仕組みを利用すると、たとえば「引数 x は i32、y は i32、関数内で使うローカル変数 tmp も i32」といった情報を簡単に管理できます。

2-2. 関数単位でのスコープ管理

Python は本来、グローバルスコープ・ローカルスコープ・関数のネスト（クロージャ）など、より複雑なスコープがあります。しかし、pyc では初期段階でそこまで全部再現するのは大変なので、「関数定義ごとにローカルスコープを1セット作る」程度の簡易的対応にとどまっています。

  1. visit_FunctionDef(node) を呼び出したら、新しいスコープを用意する (あるいはシンボルテーブルをシャドウする)
  2. 引数やローカル変数があれば、symbol_table[var_name] = var_type として登録する
  3. 関数の body を再帰的に訪問する
  4. 関数定義が終わったら、上位スコープに戻る

この基本だけで、簡単なプログラム (グローバルに関数をいくつか定義し、それぞれの中でローカル変数を持つレベル) なら十分にコンパイル可能になります。
もちろん、グローバル変数やクラス変数などは未対応の場合が多いので今後実装を進めていきます。

3. LLVM IR の関数定義を組み立てる

3-1. 引数と戻り値のマッピング

FunctionDef ノードから取得できる name (関数名) 、args (引数のリスト) 、returns (戻り値型の注釈) をもとに、LLVM IR の関数シグネチャを組み立てます。

例えば、

def foo(x: int, y: int) -> int:
    return x + y

なら、以下のような IR にしたいわけです。

define i32 @foo(i32 %x, i32 %y) {
entry:
  ...
  ret i32 ...
}

ここでポイントになるのは、

  • define i32 @foo: 戻り値が i32、関数名が foo
  • (i32 %x, i32 %y): 引数もそれぞれ i32
  • { ... }: 関数本体の命令列
  • entry:: LLVM IR でしばしば使われるエントリブロック用ラベル

この宣言部分を文字列として生成するために、Visitor は arg_str = ", ".join(...) のように引数を連結して組み立て、IRBuilder に emit(f"define i32 @{func_name}({arg_str}) {{") と書き込みます。

3-2. 関数本体の処理

次に、node.body に含まれるステートメントを順に Visitor で処理し、命令列を出力します。具体的には

  1. self.builder.emit("entry:") でエントリブロックのラベルを宣言
  2. for stmt in node.body: self.visit(stmt) で文ごとに再帰的に処理
    • Assign, Return, If, などに応じて異なる IR が生成される
  3. visit_Return(...) が呼ばれれば、ret i32 や ret ptr といった命令を出す

もし return が一度も出現しなかった場合、Python 的には暗黙の None を返しますが、このコンパイラでは「戻り値型が i32 なら ret i32 0」などのデフォルト動作にしていることが多いです。これはかなり単純化した仕様ですが、小規模な実装では十分使い物になります。

3-3. シンプルな例：加算関数

先ほどの加算関数 add(x: int, y: int) -> int なら、以下のような IR が生成されるイメージです。

define i32 @add(i32 %x, i32 %y) {
entry:
  %t0 = add i32 %x, %y
  ret i32 %t0
}
  • %t0 = add i32 %x, %y は前回記事で解説した BinOp 処理 (x + y) が visit_BinOp 経由で生成する部分
  • 最終行 ret i32 %t0 は return x + y の結果を返す行

Python らしさが薄いほどにシンプルですが、こうすることで「C 言語とほぼ同等のオーバーヘッド」で数値演算を扱えるわけです。

4. Return 文の扱い

4-1. ast.Return ノードの処理

Python AST における return は、Return(value=...) というノードで表されます。
実際に visit_Return(node) を書くなら、以下のような流れになることが多いです。

  1. ret_val = self.visit(node.value) で戻り値となる式を再帰的に評価し、その型とレジスタ名を取得する (例: TypedValue("%t0", "i32"))
  2. 関数の期待される戻り値型と ret_val.type_ を照合する
    • もし合わなければコンパイラエラーにするなど
  3. IRBuilder に emit(f"ret i32 {ret_val.llvm_value}") のように書き込み、関数を終了させる

4-2. 型不一致時の対応

たとえば関数が「-> int」と宣言されているのに return "Hello" してしまったら矛盾が生じます。Python の動的型風に考えればそれ自体はエラーではありませんが、本コンパイラはあくまで「静的型っぽくコンパイルする」方針なので、その場でコンパイルエラーにしています。
このように、「Python の柔軟性をどこまで許容するか？」と「静的型コンパイラとしての整合性」は常にトレードオフを伴います。将来的には、型注釈が無い場合は遅い動的ディスパッチを使うなどの高度な仕組みを検討するかもしれませんが、現状では単純化を優先している状態です。

4-3. デフォルトの戻り値

Python では、return 文が存在しない関数は暗黙的に None を返します。しかし pyc が目指す静的型モデルでは「戻り値が i32 か ptr か」を明示しないと扱いづらいため、とりあえず

  • 戻り値型が i32 の関数で return が無ければ ret i32 0
  • 戻り値型が ptr (例えば -> None とアノテートしているケース) の関数で return が無ければ ret ptr null

というような簡易実装で落とし込んでいます。実行時の挙動は Python とは少々異なりますが、最低限「クラッシュしないで動く」レベルにはなります。

5. 簡易的な静的型付き関数モデルの今後

5-1. Python の高度な機能とのギャップ

これまでの説明のとおり、pyc が採用している「静的型付き関数モデル」は、Python 本来の柔軟性をかなり削っています。
例えば以下のような機能は未サポートまたは制限付きです:

  • デフォルト引数: def foo(x=10) のようなデフォルト値
  • 可変長引数: *args や **kwargs
  • クロージャ: 関数内関数が外側の変数をキャプチャする
  • デコレータ: 関数定義に対して動的に処理を付加する
  • クラスメソッド・静的メソッド: @classmethod, @staticmethod など

これらを対応しようとすると、スコープ管理や呼び出し規約が一気に複雑化していきます。今後の課題として、部分的にサポートしたり、ある程度書き換えを促したりする可能性があります。

5-2. メリット: シンプルで高速

一方で、このシンプルさのおかげで、ある程度定型的な関数 (int の引数をとり、int を返す) を書く分にはコンパイル生成物が軽量かつ速く動きます。C 言語のコンパイラに近い生成物が期待できるため、ループ演算や数値計算タスクを静的型で書きたいケースには有効です。
「Python としての柔軟な機能は捨てるが、低レイヤーの高速処理部分だけ独自に書きたい」というニーズに応えられるかもしれません。

5-3. 今後の拡張

このプロジェクト自体は実験的なもので、今後「型アノテーションなしの関数をどう扱うか」「部分的に動的型を導入するか」「クラスや継承をどう再現するか」など、多数の発展課題があります。とはいえ、関数定義をこのようにシンプルにマップできるだけでも、ある程度面白いプログラムは書ける状態になっています。

6. まとめ & 次回の話

今回は、関数定義 (FunctionDef) や return 文、スコープ管理をどのように LLVM IR に落とし込んでいるかを解説しました。
主なポイントとしては、

  • AST 上の FunctionDef から名前・引数型・戻り値型を読み取り、C 言語に似た define i32 @funcName(i32 %x, ...) { ... } を生成する
  • シンボルテーブルを使ってスコープ内の変数の型を管理し、Visitor で各ステートメントを巡回して命令列を出力
  • return x + y のような単純な処理は、add i32 / ret i32 というシンプルな IR に落とせる
  • 今のところ、Python の高度な動的機能は大幅に制限されているが、その分高速かつコンパイラとしては実装しやすい

という流れです。

次回は、「リストと辞書 (PyList, PyDict) の実装」にフォーカスして、Python 的に異なる型の要素が混在できるコレクションを C 言語でどう再現しているかをご紹介します。

  • リスト: 可変長配列を GC_malloc で動的拡張するしくみ
  • 辞書: 連想配列をハッシュテーブルで管理するしくみ
    など、やや低レイヤー寄りの話題が増えますが、Python らしいデータ構造を支えるための工夫をお見せできればと思います。

次回:
https://zenn.dev/t3tra/articles/10c356a2560e5a

Discussion