phasm/phasm/type5/typeexpr.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

from .kindexpr import Arrow, KindExpr, Nat, Star


@dataclass
class TypeExpr:
    kind: KindExpr
    name: str

    def __str__(self) -> str:
        return f"{self.name} :: {self.kind}"

@dataclass
class AtomicType(TypeExpr):
    def __init__(self, name: str) -> None:
        super().__init__(Star(), name)

    def __hash__(self) -> int:
        return hash((self.kind, self.name))

@dataclass
class TypeLevelNat(TypeExpr):
    value: int

    def __init__(self, nat: int) -> None:
        assert 0 <= nat

        super().__init__(Nat(), str(nat))

        self.value = nat

@dataclass
class TypeVariable(TypeExpr):
    """
    A placeholder in a type expression
    """
    def __hash__(self) -> int:
        return hash((self.kind, self.name))

@dataclass
class TypeConstructor(TypeExpr):
    def __init__(self, kind: Arrow, name: str) -> None:
        super().__init__(kind, name)

    def __hash__(self) -> int:
        return hash((self.kind, self.name))

@dataclass
class TypeApplication(TypeExpr):
    constructor: TypeConstructor | TypeApplication | TypeVariable
    argument: TypeExpr

    def __init__(
            self,
            constructor: TypeConstructor | TypeApplication | TypeVariable,
            argument: TypeExpr,
        ) -> None:
        if isinstance(constructor.kind, Star):
            raise TypeError("A constructor cannot be a concrete type")

        if isinstance(constructor.kind, Nat):
            raise TypeError("A constructor cannot be a number")

        if constructor.kind.arg_kind != argument.kind:
            raise TypeError("Argument does match construtor's expectations")

        super().__init__(
            constructor.kind.result_kind,
            f"{constructor.name} ({argument.name})",
        )

        self.constructor = constructor
        self.argument = argument

def occurs(lft: TypeVariable, rgt: TypeApplication) -> bool:
    """
    Checks whether the given variable occurs in the given application.
    """
    if lft == rgt.constructor:
        return True

    if lft == rgt.argument:
        return True

    if isinstance(rgt.argument, TypeApplication):
        return occurs(lft, rgt.argument)

    return False

def is_concrete(lft: TypeExpr) -> bool:
    """
    A concrete type has no variables in it.

    This is also known as a monomorphic type or a closed type.
    TODO: I don't know the differen between them yet.
    """
    if isinstance(lft, AtomicType):
        return True

    if isinstance(lft, TypeLevelNat):
        return True

    if isinstance(lft, TypeVariable):
        return False

    if isinstance(lft, TypeConstructor):
        return True

    if isinstance(lft, TypeApplication):
        return is_concrete(lft.constructor) and is_concrete(lft.argument)

    raise NotImplementedError

def is_polymorphic(lft: TypeExpr) -> bool:
    """
    A polymorphic type has one or more variables in it.
    """
    return not is_concrete(lft)

def replace_variable(expr: TypeExpr, var: TypeVariable, rep_expr: TypeExpr) -> TypeExpr:
    assert var.kind == rep_expr.kind, (var, rep_expr, )

    if isinstance(expr, AtomicType):
        # Nothing to replace
        return expr

    if isinstance(expr, TypeLevelNat):
        # Nothing to replace
        return expr

    if isinstance(expr, TypeVariable):
        if expr == var:
            return rep_expr
        return expr

    if isinstance(expr, TypeConstructor):
        # Nothing to replace
        return expr

    if isinstance(expr, TypeApplication):
        new_constructor = replace_variable(expr.constructor, var, rep_expr)

        assert isinstance(new_constructor, TypeConstructor | TypeApplication | TypeVariable)  # type hint

        return TypeApplication(
            constructor=new_constructor,
            argument=replace_variable(expr.argument, var, rep_expr),
        )

    raise NotImplementedError

def instantiate(
        expr: TypeExpr,
        known_map: dict[TypeVariable, TypeVariable],
        make_variable: Callable[[KindExpr, str], TypeVariable],
    ) -> TypeExpr:
    """
    Make a copy of all variables.

    This is need when you use a polymorphic function in two places. In the
    one place, `t a` may refer to `u32[...]` and in another to `f32[4]`.
    These should not be unified since they refer to a different monomorphic
    version of the function.
    """
    if isinstance(expr, AtomicType):
        return expr

    if isinstance(expr, TypeLevelNat):
        return expr

    if isinstance(expr, TypeVariable):
        known_map.setdefault(expr, make_variable(expr.kind, expr.name))
        return known_map[expr]

    if isinstance(expr, TypeConstructor):
        return expr

    if isinstance(expr, TypeApplication):
        new_constructor = instantiate(expr.constructor, known_map, make_variable)

        assert isinstance(new_constructor, TypeConstructor | TypeApplication | TypeVariable)  # type hint

        return TypeApplication(
            constructor=new_constructor,
            argument=instantiate(expr.argument, known_map, make_variable),
        )

    raise NotImplementedError(expr)