Typed Lua (Old)

Some random thoughts and ideas about types

There's a lot of stuff that should be formalized and notation that is abused. Mainly to get the point across. This is not an academic paper, we're trying to be readable here. 😉

We distinguish between regular Lua functions $f:\prod_i(x_i:\tau_i)\mapsto (y:\tau)$ which take in some parameters of unknown type $\tau_i$ and outputs a value of type $\tau$ versus a function which takes in types as parameters and outputs a type $F:\prod_i\tau_i\mapsto\tau$ (breaking from the traditional lambda form to represent functions).

SSA Type Annotations

For this, we assume the influence of external function calls. We also ignore metatables. Type inference/deduction usually is as follows:

1. Type constructors prove variable types

2. Types are then propagated through operations

Both of which are complicated by function calls (unintended side effects) and metamethods (unspecified and unknown operation type deductions). So we ignore those for now.

Consider this Lua code:

t = {}
t.x = 10
t.y = t.x + 10
t.x = "hello"

We can construct an SSA representation like this:

%1 = NEWTABLE()
%2 = NEWINDEX(%1, x, 10)
%3 = INDEX(%2, x)
%4 = ADD(%3, 10)
%5 = NEWINDEX(%2, y, %4)
%6 = NEWSTRING("hello")
%7 = NEWINDEX(%5, x, %6)

Then the type deductions are as follows:

%1: {}
%2: typeof(%1) & { x: number }
%3: typeof(%2.x)
%4: number
%5: typeof(%2) & { y: typeof(%4) }
%6: string
%7: typeof(%5) & { x: typeof(%6) }

Notice a few key concepts:

• NEWINDEX is called even when replacing the value at a table index. This is due to the semantics of Lua's __newindex metamethod.

• Adding a new entry to a table is considered a new static assignment

• This new static assignment gets a new type

• Other type deductions are done normally

  • The deduction for %4 arises because typeof(%3) is a number so number + number gives a number. However this is not trivial if metatables are involved.

• We use type intersection to refine the type of t as the program executes. In fact, %1, %2 and %5 are aliased to the same table reference. In Lua, tables are always pass-by-reference-value.

• Type intersection will replace table entry types if they exist

We extend this type algebra to consider PHI nodes also. Consider:

a = 10
t = {}
if b then
    a = a + 1
    t.x = a
else
    a = "string"
    t.y = 10
end

Where in the SSA representation, we expect some PHI nodes to be generated:

.lbb0:
    %1 = 10
    %2 = NEWTABLE()
    %3 = GETGLOBAL(b)
    JUMP_IF_FALSE(%3, .lbb2)
.lbb1:
    %4 = ADD(%1, 1)
    %5 = NEWINDEX(%2, x, %4)
.lbb2:
    %6 = NEWSTRING("string")
    %7 = NEWINDEX(%2, y, 10)
.lbb3:
    -- PHI cond true_val false_val
    %8 = PHI(%3, %4, %6)
    %9 = PHI(%3, %5, %7)

When we type annotate it (only the interesting annotations are shown):

%8: typeof(%4) | typeof(%6)
%9: typeof(%5) | typeof(%7)

Where the other types are derived as usual. Just like how NEWINDEX forces a type intersection, the PHI node forces a type union.

This naturally gives rise to an algorithm to process the code in SSA form:

1. Transform into SSA as usual, using virtual register assigns (such as NEWINDEX)

2. Treat each SSA instruction as a function $F:\tau_1\times\cdots\times\tau_n\to\tau_{n+1}$ and if no such function legally exists, throw a type error.
3. Perform DCE on the no-return instructions by removing the virtual assignment (treat them as simply side-effect instructions).

4. Perform the usual SSA optimizations.

Type Inference with Unknown Parameters

A critical step to realizing type inferencing with function calls is to be able to predict types (as much as possible) from unknown parameters.

Idea: Consider function $f:\prod_i(x_i:\tau_i)\mapsto (y:\tau_i)$​ then there must exist a function $F:\prod_i\tau_i\mapsto\tau$ which takes in the parameter types (i.e., typeof(xi)) and outputs the deduced type. Moreover, this function is trivial to compute as the typical type deduction techniques (in previous section) applies.

Problem: Evaluating the output of the generic function F is as complex as parsing the function f itself. Most of this complexity comes from metamethods as operators may have undefined type signatures! Completely unknown at compile-time! Moreover, this is way too slow. For instance,

function f(x, y)
    return (x+y)*x/y
end

Becomes the generic version

F(t1, t2) => opdiv(opmul(opadd(t1, t2), t1), t2)

Where we try to match the type signatures for the op functions (if we even can) for each unique invocation of f. Now, given an average of k calls per function of depth d we have a runtime of $O(k^d)$​ which is highly undesirable. Basically it's (way) too slow.

Idea: Only infer function signatures when we have to. This gives rise to the following algorithm:

1. For each function, infer the most general type F of f. Note we should get the most general return type of f from this.

2. Start at the top-level and assume we don't have any untyped parameters/locals/globals. Then for each function invocation, use the most-generic return type for continued inference. Add the function invocation signature.

3. Then during codegen, we can take the most frequently called function signatures and specialize it into typed IR. Otherwise, we leave the function as generic bytecode for the standard Lua interpreter to use.

This leads to another problem: The more polymorphic the written code is, the less optimal the AOT generated code becomes. This is an acceptable trade-off.