Building a Python Compiler #4: Function Definitions and Scope

Hello. Last time (“#3 - Numerical Operations and Type System”), we discussed how to treat Python int as native i32 to execute numerical operations like x + y at high speed.
By intentionally deciding types (i32) statically despite Python's dynamic nature, we significantly simplified the compiler structure and gained the benefits of LLVM optimization.

The theme for this installment is handling function definitions and scopes.
While Python functions are characterized by their flexibility as a dynamic language, this project aims for a simple implementation by taking the approach of "creating LLVM IR function definitions based on static type annotations." If we know that "this function takes an i32 and returns an i32," we can treat it almost exactly like int foo(int x) in C.

1. AST Nodes for Function Definitions

1-1. Structure of the FunctionDef Node

Let's say we write a function like the following in Python:

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

When parsed with ast.parse() using Python's standard ast library, it is analyzed as a FunctionDef node, which typically has a structure like this:

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')))
  ],
  ...
)

In this structure:

• name='foo' is the function name.
• The arguments x and y are listed in args.args, each with annotation=Name(id='int').
• returns=Name(id='int') indicates that the return value is an int.
• The body contains statements such as Return(...).

1-2. Importance of Type Annotations

As mentioned last time, this compiler pyc maps cases where "arguments and return values are annotated as int" to "native i32."
Thanks to this, foo(x, y) can be regarded as a "function that receives i32 and returns i32," allowing the LLVM IR function to be assembled in a manner close to statically typed languages.
Conversely, if annotations are missing or an unknown type is specified, limitations currently occur, such as "falling back to a dynamic type (ptr)" or "treating it as a build error." While more flexible handling may be possible in future extensions, we are prioritizing simplification for now.

2. Symbol Table and Scope Management

2-1. Overview of the Symbol Table

In many compilers, not just Python, the question of "where to record information such as variable names and types" is a major concern. In this project pyc, a simple symbol table like the one below is maintained within classes called BaseVisitor and StmtVisitor.

• self.symbol_table = {}
• The dictionary keys are variable names, and the values are type strings ("i32", "ptr", etc.).
• When visiting function definitions or blocks, the symbol table is updated as necessary.

By using such a mechanism, it becomes easy to manage information such as "argument x is i32, y is i32, and the local variable tmp used inside the function is also i32."

2-2. Scope Management on a Function Basis

Python inherently has more complex scopes, such as global scope, local scope, and nested functions (closures). However, since reproducing all of them in the early stages of pyc would be difficult, we have stuck to a simple approach of "creating one set of local scope for each function definition."

1. When visit_FunctionDef(node) is called, prepare a new scope (or shadow the symbol table).
2. If there are arguments or local variables, register them as symbol_table[var_name] = var_type.
3. Recursively visit the body of the function.
4. Once the function definition is finished, return to the upper scope.

With this basic approach alone, simple programs (at a level where several functions are defined globally, each having its own local variables) become sufficiently compilable.
Of course, since global variables and class variables are often unsupported, we will proceed with their implementation in the future.

3. Assembling LLVM IR Function Definitions

3-1. Mapping Arguments and Return Values

We assemble the LLVM IR function signature based on the name (function name), args (list of arguments), and returns (return type annotation) obtained from the FunctionDef node.

For example:

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

We want to turn this into IR like the following:

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

The key points here are:

• define i32 @foo: The return value is i32, and the function name is foo.
• (i32 %x, i32 %y): Each argument is also i32.
• { ... }: The sequence of instructions for the function body.
• entry:: A label for the entry block often used in LLVM IR.

To generate this declaration part as a string, the Visitor assembles the arguments using something like arg_str = ", ".join(...) and writes it to the IRBuilder with emit(f"define i32 @{func_name}({arg_str}) {{").

3-2. Processing the Function Body

Next, the statements contained in node.body are processed in order by the Visitor to output the instruction sequence. Specifically:

1. Declare the entry block label with self.builder.emit("entry:").
2. Recursively process each statement using for stmt in node.body: self.visit(stmt).
   • Different IR is generated depending on Assign, Return, If, etc.
3. If visit_Return(...) is called, output instructions such as ret i32 or ret ptr.

If a return never appears, Python implicitly returns None. However, in this compiler, we often set the default behavior to something like "if the return type is i32, ret i32 0." While this is a heavily simplified specification, it is quite usable for small-scale implementations.

3-3. Simple Example: Addition Function

For the addition function add(x: int, y: int) -> int mentioned earlier, the generated IR would look something like this:

define i32 @add(i32 %x, i32 %y) {
entry:
  %t0 = add i32 %x, %y
  ret i32 %t0
}

• %t0 = add i32 %x, %y is the part generated via visit_BinOp for the BinOp processing (x + y) explained in the previous article.
• The final line ret i32 %t0 is the line that returns the result of return x + y.

It is so simple that the "Python-ness" is faint, but by doing this, numerical operations can be handled with "overhead nearly equivalent to C."

4. Handling Return Statements

4-1. Processing ast.Return Nodes

In the Python AST, return is represented by a Return(value=...) node.
If we were to write visit_Return(node), the flow would typically look like this:

1. Recursively evaluate the expression to be returned with ret_val = self.visit(node.value) and obtain its type and register name (e.g., TypedValue("%t0", "i32")).
2. Compare the expected return type of the function with ret_val.type_.
   • If they don't match, raise a compiler error, etc.
3. Write to the IRBuilder with something like emit(f"ret i32 {ret_val.llvm_value}") to terminate the function.

4-2. Handling Type Mismatches

For example, if a function is declared with -> int but executes return "Hello", a contradiction occurs. In terms of Python's dynamic typing, this isn't an error in itself, but since this compiler follows the policy of "compiling as if it were statically typed," it results in a compile error on the spot.
In this way, there is always a trade-off between "how much of Python's flexibility to allow" and "consistency as a statically typed compiler." In the future, we may consider advanced mechanisms such as using slow dynamic dispatch when type annotations are missing, but for now, we are prioritizing simplification.

4-3. Default Return Values

In Python, a function without a return statement implicitly returns None. However, in the static type model pyc aims for, it is difficult to handle without specifying whether the return value is i32 or ptr. Therefore, we have implemented a simple fallback:

• If a function with a return type of i32 has no return, it performs ret i32 0.
• If a function with a return type of ptr (e.g., cases annotated with -> None) has no return, it performs ret ptr null.

Although the runtime behavior differs slightly from Python, it reaches a level where it "works without crashing" at a minimum.

5. Future of the Simple Statically Typed Function Model

5-1. Gap with Advanced Python Features

As explained so far, the "statically typed function model" adopted by pyc significantly trims down Python's original flexibility. For instance, the following features are unsupported or limited:

• Default arguments: Default values like def foo(x=10)
• Variable-length arguments: *args and **kwargs
• Closures: Inner functions capturing variables from the outer scope
• Decorators: Adding processing dynamically to function definitions
• Class and static methods: @classmethod, @staticmethod, etc.

Attempting to support these would drastically complicate scope management and calling conventions. As a future challenge, we might support them partially or encourage a certain amount of code rewriting.

5-2. Advantage: Simple and Fast

On the other hand, thanks to this simplicity, when writing standard functions (taking int arguments and returning an int), the compiled artifact is lightweight and runs fast. Since we can expect output close to that of a C compiler, it is effective for cases where you want to write loop operations or numerical calculation tasks with static types.

It might meet the need of those who say, "I'll discard Python's flexible features, but I want to write high-speed low-layer processing parts myself."

5-3. Future Extensions

This project itself is experimental, and there are many future challenges, such as "how to handle functions without type annotations," "whether to partially introduce dynamic types," and "how to reproduce classes and inheritance." Nevertheless, even just being able to map function definitions simply like this allows us to write reasonably interesting programs.

6. Summary & Next Topic

In this installment, we explained how function definitions (FunctionDef), return statements, and scope management are translated into LLVM IR.
The main points are:

• Reading the name, argument types, and return type from FunctionDef on the AST and generating define i32 @funcName(i32 %x, ...) { ... } similar to C.
• Managing the types of variables within the scope using a symbol table, and traversing each statement with a Visitor to output the instruction sequence.
• Simple processes like return x + y can be boiled down to simple IR like add i32 / ret i32.
• For now, Python's advanced dynamic features are heavily restricted, but in exchange, it's fast and easier to implement as a compiler.

Next time, we will focus on the "Implementation of Lists and Dictionaries (PyList, PyDict)" and introduce how to reproduce Python-style collections, where elements of different types can coexist, in C.

• Lists: A mechanism for dynamically expanding variable-length arrays using GC_malloc.
• Dictionaries: A mechanism for managing associative arrays with hash tables.
  The discussion will lean more toward the lower layers, but I hope to show you the ingenuity behind supporting Pythonic data structures.

Next time:
https://zenn.dev/t3tra/articles/10c356a2560e5a