Toy language type inference with generics and tuples: Frontier Interview's solution

Question: Toy language type inference with generics and tuples

Companies: OpenAI

Difficulty: Medium

Problem statement

A small typed language has three kinds of type:

Primitives: int, float, char.
Generics: placeholders named T followed by a number (T1, T2, ...). A generic can stand for any type, including a tuple.
Tuples: an ordered list of types, which may be nested, e.g. [int,[T1,char]].

A function has a list of parameter types and a single return type.

Types are given to you already built as objects. You do not need to parse strings. Model a type as a Node (a named leaf, or a tuple with child nodes) and a function as Function(params, ret).

Part 1: printing

Implement to_str() for both classes:

A primitive or generic prints its name.
A tuple prints its elements comma-separated inside square brackets, with no spaces: [int,[T1,char]].
A function prints its parameters comma-separated inside parentheses, then ->, then the return type: (int,[T1,float]) -> T1.

Part 2: return type inference

Implement infer_return(fn, args). args is a list of concrete argument types, one for each parameter. Return the function's return type with every generic replaced by the type it must stand for. Raise an error when:

the number of arguments differs from the number of parameters;
a primitive parameter receives a different type;
a tuple parameter receives a non-tuple, or a tuple of a different length;
one generic would have to stand for two different types;
the return type uses a generic that no parameter binds.

Examples

Function (T1, [T1,T2]) -> [T2,T1] called with (float, [float,[int,int]]): T1 binds to float and T2 binds to [int,int], so the result is [[int,int],float].

Function (T1, T1) -> T1 called with (int, char) raises an error, because T1 cannot be both int and char.

Approach

The language has three kinds of type, and they nest:

Primitive: int, float, char.
Generic: a placeholder such as T1 or T2. It can stand for any type, including a tuple.
Tuple: an ordered list of types, which may contain other tuples.

A function has a list of parameter types and one return type. Two tasks are usually asked, in this order:

to_str on a type and on a function, which shows you can walk the tree.
infer_return(fn, args), which works out the concrete return type when the function is called with arguments of known types.

Representation. Use one Node class. A leaf holds a name (int, T1), and a tuple holds a list of children. A name is generic exactly when it is not a primitive. This rule (or a regex like T\d+) is the one place where a leaf's kind gets decided. Give Node structural equality so that two separately built [int,char] trees compare equal.

Printing. This is a plain recursive walk. A leaf prints its name. A tuple prints [ + its children joined by , + ]. A function prints ( + its parameters joined by , + ) -> + its return type. Settle the exact format with the interviewer before you write code. Separators and spacing are the usual source of failed asserts.

Brute force for inference. One idea is to substitute candidate types for each generic and check which choice matches. That search has no bound, because a generic can be any tuple. Don't do it. The arguments already tell you what each generic must be.

Optimal: one-sided unification. Walk each parameter type alongside its argument type, and fill a bindings map from generic name to Node:

Parameter is a generic T: if T has no binding yet, bind it to the whole argument subtree. Otherwise the argument must equal the existing binding, or you raise a conflict.
Parameter is a primitive: the argument must be the same primitive.
Parameter is a tuple: the argument must be a tuple of the same length. Recurse pairwise.

Next, check that the number of arguments matches the number of parameters. Then rebuild the return type, replacing each bound generic with a copy of its binding. A generic that appears only in the return type has nothing to bind it, so raise an error. (Some interviewers want it left as is. Ask.)

Why it works. Matching is structural and deterministic. Each position in a parameter tree has exactly one argument subtree under it, so there is never a choice to make. The first binding of a generic is forced, and every later occurrence only checks against it. This is unification where only the pattern side has variables. The full two-sided version is a follow-up.

Invariants.

After a parameter is processed, every generic in it is bound, and every occurrence of that generic matched the same subtree.
bindings only grows. A binding never changes once it is set.
The output contains no generic that came from the function's own signature.

Code

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Optional

PRIMITIVES = {"int", "float", "char"}


class TypeInferenceError(Exception):
    pass


@dataclass(frozen=True)
class Node:
    """A leaf (base is a primitive or generic name) or a tuple (children)."""
    base: Optional[str] = None
    children: tuple = field(default_factory=tuple)

    @staticmethod
    def leaf(name: str) -> "Node":
        return Node(base=name)

    @staticmethod
    def tup(*items: "Node") -> "Node":
        return Node(children=tuple(items))

    def is_tuple(self) -> bool:
        return self.base is None

    def is_generic(self) -> bool:
        return self.base is not None and self.base not in PRIMITIVES

    def to_str(self) -> str:
        if self.is_tuple():
            return "[" + ",".join(c.to_str() for c in self.children) + "]"
        return self.base


@dataclass(frozen=True)
class Function:
    params: List[Node]
    ret: Node

    def to_str(self) -> str:
        return "(" + ",".join(p.to_str() for p in self.params) + ") -> " + self.ret.to_str()


def _bind(param: Node, arg: Node, bindings: Dict[str, Node]) -> None:
    if param.is_generic():
        bound = bindings.get(param.base)
        if bound is None:
            bindings[param.base] = arg
        elif bound != arg:
            raise TypeInferenceError(
                f"conflict for {param.base}: {bound.to_str()} vs {arg.to_str()}")
        return
    if not param.is_tuple():  # primitive
        if arg != param:
            raise TypeInferenceError(f"expected {param.to_str()}, got {arg.to_str()}")
        return
    if not arg.is_tuple() or len(arg.children) != len(param.children):
        raise TypeInferenceError(f"expected {param.to_str()}, got {arg.to_str()}")
    for p, a in zip(param.children, arg.children):
        _bind(p, a, bindings)


def _substitute(node: Node, bindings: Dict[str, Node]) -> Node:
    if node.is_tuple():
        return Node.tup(*(_substitute(c, bindings) for c in node.children))
    if node.is_generic():
        if node.base not in bindings:
            raise TypeInferenceError(f"unbound generic {node.base} in return type")
        return bindings[node.base]
    return node


def infer_return(fn: Function, args: List[Node]) -> Node:
    if len(args) != len(fn.params):
        raise TypeInferenceError(f"expected {len(fn.params)} args, got {len(args)}")
    bindings: Dict[str, Node] = {}
    for p, a in zip(fn.params, args):
        _bind(p, a, bindings)
    return _substitute(fn.ret, bindings)

Nodes are immutable, so sharing a bound subtree in the result is safe and no deep copy is needed.

Complexity

Let P be the total node count of the parameter types, A that of the arguments, and R that of the return type.

to_str: O(n) time to visit the nodes, plus the output length. Joining strings at each level copies text once per level, so very deep trees cost O(n * depth). O(depth) recursion stack.
infer_return: binding visits each parameter node once. Each check of a repeated generic costs O(size of the bound subtree) for the equality test. The worst case is O(P * A). With hash-consing or cached hashes it becomes O(P + A) in practice. Substitution is O(R) new nodes, because bound subtrees are shared, not copied. Space is O(number of generics + R + depth).

Test cases

Notation: int, T1 are leaves; [a,b] is a tuple.

to_str of [int,[T1,char]] → "[int,[T1,char]]".
to_str of a function with params (int, [T1,float]) and return T1 → "(int,[T1,float]) -> T1".
Function (T1, T2) -> [T2,T1], args (int, char) → [char,int].
Function (T1, [T1,T2]) -> T2, args (float, [float,[int,int]]) → [int,int] (a generic bound to a tuple).
Function (T1, T1) -> T1, args (int, char) → error: conflict for T1.
Function (T1, T1) -> T1, args ([int,char], [int,char]) → [int,char] (tuples compare structurally).
Function ([int,T1]) -> T1, args ([int,char,float]) → error: tuple length mismatch.
Function (int) -> int, args (float) → error: primitive mismatch.
Function ([T1]) -> T1, args (int) → error: expected a tuple.
Function (T1) -> T2, args (int) → error: unbound generic T2.
Function (int, char) -> float, args (int,) → error: arity mismatch.
Function () -> int, args () → int.

I, F, C = Node.leaf("int"), Node.leaf("float"), Node.leaf("char")
T1, T2 = Node.leaf("T1"), Node.leaf("T2")
assert infer_return(Function([T1, T2], Node.tup(T2, T1)), [I, C]).to_str() == "[char,int]"
assert infer_return(Function([T1, Node.tup(T1, T2)], T2),
                    [F, Node.tup(F, Node.tup(I, I))]).to_str() == "[int,int]"

Follow-ups

Parse the string form. Write parse("(int,[T1,char]) -> T1") that returns a Function. Use a recursive-descent parser over tokens ([, ], ,, (, ), ->, names), with one function for types and one for the signature. Check that to_str(parse(s)) == s as a round trip.
Arguments that contain generics (two-sided unification). Callers may pass types like [T3,int]. Rename the function's generics apart so they don't collide with the caller's, then run real unification with a union-find or a substitution map that both sides can bind. Add an occurs check so that T1 = [T1,int] is rejected.
Function types as values (higher-order). Add a FuncType node, so a parameter can be (T1) -> T2. Matching it against an argument function type unifies the parameters pairwise and the return types. The binding code gains one more case, and the rest stays the same.
Better errors. Report the path to the mismatch, for example arg 2 → element 1: expected int, got char. Pass a path list through _bind and add it to the message.
Function composition or a call chain. Given f and g, infer the type of g(f(x)), or check that a pipeline is well typed. This reuses infer_return with the output of one step as the input of the next. If generics are left free, it needs the two-sided version.

Interviewer expectations

Mid-level (L4)

Models types as a small recursive tree. Gets to_str right, including nested tuples and the agreed separators.
Implements binding with a map, and catches the conflict and primitive-mismatch cases.
Tests by hand with at least one tuple example and one error example.

Senior (L5)

Handles every failure mode without being prompted: arity, tuple length, tuple vs leaf, unbound return generic. Uses a dedicated exception type.
Uses structural equality, so tuple bindings compare correctly. Explains why there is no search: each parameter position decides its binding.
States the complexity, including the cost of repeated subtree comparisons, and keeps the code split into clear bind and substitute steps.

Staff (L6+)

Names the technique as one-sided unification. Explains what changes for two-sided unification: renaming apart, union-find, the occurs check.
Extends the design cleanly to function-typed parameters and to parsing, without rewriting the core.
Discusses API choices: immutable nodes vs copying, error messages with paths, and whether unbound return generics should be an error or stay polymorphic.