phasm/phasm/type3/functions.py

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Hashable, Iterable, List

if TYPE_CHECKING:
    from .typeclasses import Type3Class
    from .types import Type3


class TypeVariableApplication_Base[T: Hashable, S: Hashable]:
    """
    Records the constructor and arguments used to create this type.

    Nullary types, or types of kind *, have both arguments set to None.
    """
    constructor: T
    arguments: S

    def __init__(self, constructor: T, arguments: S) -> None:
        self.constructor = constructor
        self.arguments = arguments

    def __hash__(self) -> int:
        return hash((self.constructor, self.arguments, ))

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, TypeVariableApplication_Base):
            raise NotImplementedError

        return (self.constructor == other.constructor  # type: ignore[no-any-return]
            and self.arguments == other.arguments)

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}({self.constructor!r}, {self.arguments!r})'

class TypeVariable:
    """
    Types variable are used in function definition.

    They are used in places where you don't know the exact type.
    They are different from PlaceholderForType, as those are instanced
    during type checking. These type variables are used solely in the
    function's definition
    """
    __slots__ = ('name', 'application', )

    name: str
    application: TypeVariableApplication_Base[Any, Any]

    def __init__(self, name: str, application: TypeVariableApplication_Base[Any, Any]) -> None:
        self.name = name
        self.application = application

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

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, TypeVariable):
            raise NotImplementedError

        return (self.name == other.name
            and self.application == other.application)

    def __repr__(self) -> str:
        return f'TypeVariable({repr(self.name)})'

class TypeVariableApplication_Nullary(TypeVariableApplication_Base[None, None]):
    """
    For the type for this function argument it's not relevant if it was constructed.
    """

def make_typevar(name: str) -> TypeVariable:
    """
    Helper function to make a type variable for a non-constructed type.
    """
    return TypeVariable(name, TypeVariableApplication_Nullary(None, None))

class TypeConstructorVariable:
    """
    Types constructor variable are used in function definition.

    They are a lot like TypeVariable, except that they represent a
    type constructor rather than a type directly.

    For now, we only have type constructor variables for kind
    * -> *.
    """
    __slots__ = ('name', )

    name: str

    def __init__(self, name: str) -> None:
        self.name = name

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

    def __eq__(self, other: Any) -> bool:
        if other is None:
            return False

        if not isinstance(other, TypeConstructorVariable):
            raise NotImplementedError(other)

        return (self.name == other.name)

    def __call__(self, tvar: TypeVariable) -> 'TypeVariable':
        return TypeVariable(
            self.name + ' ' + tvar.name,
            TypeVariableApplication_Unary(self, tvar)
        )

    def __repr__(self) -> str:
        return f'TypeConstructorVariable({self.name!r})'

class TypeVariableApplication_Unary(TypeVariableApplication_Base[TypeConstructorVariable, TypeVariable]):
    """
    The type for this function argument should be constructed from a type constructor.

    And we need to know what construtor that was, since that's the one we support.
    """

class ConstraintBase:
    __slots__ = ()

class Constraint_TypeClassInstanceExists(ConstraintBase):
    __slots__ = ('type_class3', 'types', )

    type_class3: 'Type3Class'
    types: list[TypeVariable]

    def __init__(self, type_class3: 'Type3Class', types: Iterable[TypeVariable]) -> None:
        self.type_class3 = type_class3
        self.types = list(types)

        # Sanity check. AFAIK, if you have a multi-parameter type class,
        # you can only add a constraint by supplying types for all variables
        assert len(self.type_class3.args) == len(self.types)

    def __str__(self) -> str:
        return self.type_class3.name + ' ' + ' '.join(x.name for x in self.types)

    def __repr__(self) -> str:
        return f'Constraint_TypeClassInstanceExists({self.type_class3.name}, {self.types!r})'

class TypeVariableContext:
    __slots__ = ('constraints', )

    constraints: list[ConstraintBase]

    def __init__(self, constraints: Iterable[ConstraintBase] = ()) -> None:
        self.constraints = list(constraints)

    def __copy__(self) -> 'TypeVariableContext':
        return TypeVariableContext(self.constraints)

    def __str__(self) -> str:
        if not self.constraints:
            return ''

        return '(' + ', '.join(str(x) for x in self.constraints) + ') => '

    def __repr__(self) -> str:
        return f'TypeVariableContext({self.constraints!r})'

class FunctionArgument:
    __slots__ = ('args', 'name', )

    args: list[Type3 | TypeVariable]
    name: str

    def __init__(self, args: list[Type3 | TypeVariable]) -> None:
        self.args = args

        self.name = '(' + ' -> '.join(x.name for x in args) + ')'

class FunctionSignature:
    __slots__ = ('context', 'args', )

    context: TypeVariableContext
    args: List[Type3 | TypeVariable | FunctionArgument]

    def __init__(self, context: TypeVariableContext, args: Iterable[Type3 | TypeVariable | list[Type3 | TypeVariable]]) -> None:
        self.context = context.__copy__()
        self.args = list(
            FunctionArgument(x) if isinstance(x, list) else x
            for x in args
        )

    def __str__(self) -> str:
        return str(self.context) + ' -> '.join(x.name for x in self.args)

    def __repr__(self) -> str:
        return f'FunctionSignature({self.context!r}, {self.args!r})'