Infer Generic Return Types

Infer Generic Return Types

In statically typed and generic programming systems, compilers determine return types at call sites using unification and substitution. You are asked to build a core type-inference function for a simplified generic type system.

A type is represented by a TypeNode object:

class TypeNode:
    def __init__(self, name: str, args: list['TypeNode'] = None, is_var: bool = False):
        self.name = name          # e.g., "T", "int", "List", "Dict"
        self.args = args or []    # Nested types, e.g. for List[int], args=[TypeNode("int")]
        self.is_var = is_var      # True if this node represents a generic type variable (e.g., T, U)

Write a function infer_return_type that accepts:

• param_types: A list of TypeNode representing the formal parameter types of a function, which may contain type variables (is_var=True).
• arg_types: A list of TypeNode representing the actual argument types passed at the call site (all is_var=False).
• return_type: A TypeNode representing the declared return type of the function, which may contain type variables.

Your function should unify param_types with arg_types to create a substitution mapping from each generic variable name to a concrete TypeNode. Then, apply this substitution to return_type and return the instantiated TypeNode.

If unification fails (e.g., mismatched argument counts, structural mismatches, conflicting variable assignments, or if any type variable in return_type remains unmapped), return None.

Function Signature

from typing import Optional

def infer_return_type(
    param_types: list[TypeNode],
    arg_types: list[TypeNode],
    return_type: TypeNode
) -> Optional[TypeNode]:
    pass

Examples

Example 1:

# Function signature: foo(x: List[T], y: Dict[T, U]) -> Pair[U, T]
# Arguments: x: List[int], y: Dict[int, str]

T = TypeNode("T", is_var=True)
U = TypeNode("U", is_var=True)

param_types = [
    TypeNode("List", [T]),
    TypeNode("Dict", [T, U])
]
return_type = TypeNode("Pair", [U, T])

arg_types = [
    TypeNode("List", [TypeNode("int")]),
    TypeNode("Dict", [TypeNode("int"), TypeNode("str")])
]

# T unifies to int, U unifies to str.
# Result: Pair[str, int]
result = infer_return_type(param_types, arg_types, return_type)
# result.name == "Pair", result.args[0].name == "str", result.args[1].name == "int"

Example 2:

# Signature: choose(a: T, b: T) -> T
# Arguments: a: int, b: str
# T cannot be both int and str; unification fails.
result = infer_return_type([T, T], [TypeNode("int"), TypeNode("str")], T)
# result == None

Constraints

• 0 <= len(param_types), len(arg_types) <= 50
• Maximum depth of nested TypeNode structures is 20.
• All concrete types in arg_types have is_var = False.
• Node names are non-empty alphanumeric strings.