Type Checking — Static, Dynamic and Hindley-Milner

class TypeEnvironment:
    def __init__(self, parent=None):
        self.bindings = {}
        self.parent = parent

    def lookup(self, name):
        if name in self.bindings:
            return self.bindings[name]
        if self.parent:
            return self.parent.lookup(name)
        return None

    def bind(self, name, type_name):
        self.bindings[name] = type_name

class StaticTypeChecker:
    def __init__(self):
        self.env = TypeEnvironment()
        self.errors = []

    def check(self, ast):
        method = f"check_{type(ast).__name__}"
        checker = getattr(self, method, None)
        if checker:
            return checker(ast)
        raise TypeError(f"No checker for {type(ast).__name__}")

    def check_NumNode(self, node):
        return "int"

    def check_FloatNode(self, node):
        return "float"

    def check_StringNode(self, node):
        return "string"

    def check_BoolNode(self, node):
        return "bool"

    def check_VarNode(self, node):
        type_name = self.env.lookup(node.name)
        if not type_name:
            self.errors.append(f"Undeclared variable: {node.name}")
            return "error"
        return type_name

    def check_BinOpNode(self, node):
        left_type = self.check(node.left)
        right_type = self.check(node.right)
        if left_type == "error" or right_type == "error":
            return "error"
        if node.op in ("+", "-", "*", "/"):
            if left_type == "int" and right_type == "int":
                return "int"
            if left_type == "float" or right_type == "float":
                if left_type in ("int", "float") and right_type in ("int", "float"):
                    return "float"
            self.errors.append(f"Type mismatch: {left_type} {node.op} {right_type}")
            return "error"
        return left_type

    def check_AssignNode(self, node):
        value_type = self.check(node.value)
        existing = self.env.lookup(node.name)
        if existing and existing != value_type:
            self.errors.append(
                f"Type mismatch: cannot assign {value_type} to {node.name}:{existing}"
            )
        self.env.bind(node.name, value_type)
        return value_type

    def check_IfNode(self, node):
        cond_type = self.check(node.condition)
        if cond_type != "bool":
            self.errors.append("If condition must be boolean")
        self.check(node.then_branch)
        if node.else_branch:
            self.check(node.else_branch)

checker = StaticTypeChecker()
checker.check(AssignNode("x", NumNode(42)))
checker.check(AssignNode("y", StringNode("hello")))
checker.check(BinOpNode(VarNode("x"), "+", VarNode("y")))
for err in checker.errors:
    print(err)

class HMTypeInferrer:
    def __init__(self):
        self.type_vars = 0
        self.constraints = []
        self.substitutions = {}

    def fresh_type_var(self):
        self.type_vars += 1
        return f"T{self.type_vars}"

    def infer(self, expr, env):
        if isinstance(expr, NumNode):
            return "int"
        elif isinstance(expr, BoolNode):
            return "bool"
        elif isinstance(expr, VarNode):
            if expr.name in env:
                return env[expr.name]
            raise TypeError(f"Unbound variable: {expr.name}")
        elif hasattr(expr, 'param') and hasattr(expr, 'body'):
            param_type = self.fresh_type_var()
            new_env = {**env, expr.param: param_type}
            body_type = self.infer(expr.body, new_env)
            return ("->", param_type, body_type)
        elif hasattr(expr, 'fun') and hasattr(expr, 'arg'):
            fun_type = self.infer(expr.fun, env)
            arg_type = self.infer(expr.arg, env)
            result_type = self.fresh_type_var()
            self.constraints.append((fun_type, ("->", arg_type, result_type)))
            return result_type
        raise TypeError(f"Cannot infer type for {type(expr).__name__}")

    def unify(self, t1, t2):
        t1 = self.apply(t1)
        t2 = self.apply(t2)
        if t1 == t2:
            return
        if isinstance(t1, str) and t1.startswith("T"):
            self.substitutions[t1] = t2
        elif isinstance(t2, str) and t2.startswith("T"):
            self.substitutions[t2] = t1
        elif isinstance(t1, tuple) and isinstance(t2, tuple) and t1[0] == "->" and t2[0] == "->":
            self.unify(t1[1], t2[1])
            self.unify(t1[2], t2[2])
        else:
            raise TypeError(f"Cannot unify {t1} with {t2}")

    def apply(self, t):
        if isinstance(t, str):
            while t in self.substitutions:
                t = self.substitutions[t]
            return t
        if isinstance(t, tuple):
            return tuple(self.apply(x) if isinstance(x, (str, tuple)) else x for x in t)
        return t

class LambdaNode:
    def __init__(self, param, body):
        self.param = param
        self.body = body

inferrer = HMTypeInferrer()
expr = LambdaNode("x", BinOpNode(VarNode("x"), "+", NumNode(1)))
result = inferrer.infer(expr, {})
print(f"Inferred type: {inferrer.apply(result)}")