guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/decorator.py

@@ -0,0 +1,246 @@
+import builtins
+import inspect
+from collections.abc import Callable, Sequence
+from types import FrameType
+from typing import ParamSpec, TypeVar
+
+from hugr import ops
+from hugr import tys as ht
+
+from guppylang.defs import (
+    GuppyDefinition,
+    GuppyFunctionDefinition,
+)
+from guppylang_internals.compiler.core import (
+    CompilerContext,
+    GlobalConstId,
+)
+from guppylang_internals.definition.common import DefId
+from guppylang_internals.definition.custom import (
+    CustomCallChecker,
+    CustomFunctionDef,
+    CustomInoutCallCompiler,
+    DefaultCallChecker,
+    NotImplementedCallCompiler,
+    OpCompiler,
+    RawCustomFunctionDef,
+)
+from guppylang_internals.definition.ty import OpaqueTypeDef, TypeDef
+from guppylang_internals.definition.wasm import RawWasmFunctionDef
+from guppylang_internals.engine import DEF_STORE
+from guppylang_internals.std._internal.checker import WasmCallChecker
+from guppylang_internals.std._internal.compiler.wasm import (
+    WasmModuleCallCompiler,
+    WasmModuleDiscardCompiler,
+    WasmModuleInitCompiler,
+)
+from guppylang_internals.tys.arg import Argument
+from guppylang_internals.tys.builtin import (
+    WasmModuleTypeDef,
+)
+from guppylang_internals.tys.param import Parameter
+from guppylang_internals.tys.subst import Inst
+from guppylang_internals.tys.ty import (
+    FuncInput,
+    FunctionType,
+    InputFlags,
+    NoneType,
+    NumericType,
+)
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def get_calling_frame() -> FrameType:
+    """Finds the first frame that called this function outside the compiler modules."""
+    frame = inspect.currentframe()
+    while frame:
+        module = inspect.getmodule(frame)
+        if module is None:
+            return frame
+        if module.__file__ != __file__:
+            return frame
+        frame = frame.f_back
+    raise RuntimeError("Couldn't obtain stack frame for definition")
+
+
+def custom_function(
+    compiler: CustomInoutCallCompiler | None = None,
+    checker: CustomCallChecker | None = None,
+    higher_order_value: bool = True,
+    name: str = "",
+    signature: FunctionType | None = None,
+) -> Callable[[Callable[P, T]], GuppyFunctionDefinition[P, T]]:
+    """Decorator to add custom typing or compilation behaviour to function decls.
+
+    Optionally, usage of the function as a higher-order value can be disabled. In
+    that case, the function signature can be omitted if a custom call compiler is
+    provided.
+    """
+
+    def dec(f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
+        call_checker = checker or DefaultCallChecker()
+        func = RawCustomFunctionDef(
+            DefId.fresh(),
+            name or f.__name__,
+            None,
+            f,
+            call_checker,
+            compiler or NotImplementedCallCompiler(),
+            higher_order_value,
+            signature,
+        )
+        DEF_STORE.register_def(func, get_calling_frame())
+        return GuppyFunctionDefinition(func)
+
+    return dec
+
+
+def hugr_op(
+    op: Callable[[ht.FunctionType, Inst, CompilerContext], ops.DataflowOp],
+    checker: CustomCallChecker | None = None,
+    higher_order_value: bool = True,
+    name: str = "",
+    signature: FunctionType | None = None,
+) -> Callable[[Callable[P, T]], GuppyFunctionDefinition[P, T]]:
+    """Decorator to annotate function declarations as HUGR ops.
+
+    Args:
+        op: A function that takes an instantiation of the type arguments as well as
+            the inferred input and output types and returns a concrete HUGR op.
+        checker: The custom call checker.
+        higher_order_value: Whether the function may be used as a higher-order
+            value.
+        name: The name of the function.
+    """
+    return custom_function(OpCompiler(op), checker, higher_order_value, name, signature)
+
+
+def extend_type(defn: TypeDef) -> Callable[[type], type]:
+    """Decorator to add new instance functions to a type."""
+
+    def dec(c: type) -> type:
+        for val in c.__dict__.values():
+            if isinstance(val, GuppyDefinition):
+                DEF_STORE.register_impl(defn.id, val.wrapped.name, val.id)
+        return c
+
+    return dec
+
+
+def custom_type(
+    hugr_ty: ht.Type | Callable[[Sequence[Argument], CompilerContext], ht.Type],
+    name: str = "",
+    copyable: bool = True,
+    droppable: bool = True,
+    bound: ht.TypeBound | None = None,
+    params: Sequence[Parameter] | None = None,
+) -> Callable[[type[T]], type[T]]:
+    """Decorator to annotate a class definitions as Guppy types.
+
+    Requires the static Hugr translation of the type. Additionally, the type can be
+    marked as linear. All `@guppy` annotated functions on the class are turned into
+    instance functions.
+
+    For non-generic types, the Hugr representation can be passed as a static value.
+    For generic types, a callable may be passed that takes the type arguments of a
+    concrete instantiation.
+    """
+    mk_hugr_ty = (
+        (lambda args, ctx: hugr_ty) if isinstance(hugr_ty, ht.Type) else hugr_ty
+    )
+
+    def dec(c: type[T]) -> type[T]:
+        defn = OpaqueTypeDef(
+            DefId.fresh(),
+            name or c.__name__,
+            None,
+            params or [],
+            not copyable,
+            not droppable,
+            mk_hugr_ty,
+            bound,
+        )
+        DEF_STORE.register_def(defn, get_calling_frame())
+        for val in c.__dict__.values():
+            if isinstance(val, GuppyDefinition):
+                DEF_STORE.register_impl(defn.id, val.wrapped.name, val.id)
+        # We're pretending to return the class unchanged, but in fact we return
+        # a `GuppyDefinition` that handles the comptime logic
+        return GuppyDefinition(defn)  # type: ignore[return-value]
+
+    return dec
+
+
+def wasm_module(
+    filename: str, filehash: int
+) -> Callable[[builtins.type[T]], GuppyDefinition]:
+    def dec(cls: builtins.type[T]) -> GuppyDefinition:
+        # N.B. Only one module per file and vice-versa
+        wasm_module = WasmModuleTypeDef(
+            DefId.fresh(),
+            cls.__name__,
+            None,
+            filename,
+            filehash,
+        )
+
+        wasm_module_ty = wasm_module.check_instantiate([], None)
+
+        DEF_STORE.register_def(wasm_module, get_calling_frame())
+        for val in cls.__dict__.values():
+            if isinstance(val, GuppyDefinition):
+                DEF_STORE.register_impl(wasm_module.id, val.wrapped.name, val.id)
+        # Add a constructor to the class
+        call_method = CustomFunctionDef(
+            DefId.fresh(),
+            "__new__",
+            None,
+            FunctionType(
+                [FuncInput(NumericType(NumericType.Kind.Nat), flags=InputFlags.Owned)],
+                wasm_module_ty,
+            ),
+            DefaultCallChecker(),
+            WasmModuleInitCompiler(),
+            True,
+            GlobalConstId.fresh(f"{cls.__name__}.__new__"),
+            True,
+        )
+        discard = CustomFunctionDef(
+            DefId.fresh(),
+            "discard",
+            None,
+            FunctionType([FuncInput(wasm_module_ty, InputFlags.Owned)], NoneType()),
+            DefaultCallChecker(),
+            WasmModuleDiscardCompiler(),
+            False,
+            GlobalConstId.fresh(f"{cls.__name__}.__discard__"),
+            True,
+        )
+        DEF_STORE.register_def(call_method, get_calling_frame())
+        DEF_STORE.register_impl(wasm_module.id, "__new__", call_method.id)
+        DEF_STORE.register_def(discard, get_calling_frame())
+        DEF_STORE.register_impl(wasm_module.id, "discard", discard.id)
+
+        return GuppyDefinition(wasm_module)
+
+    return dec
+
+
+def wasm(f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
+    func = RawWasmFunctionDef(
+        DefId.fresh(),
+        f.__name__,
+        None,
+        f,
+        WasmCallChecker(),
+        WasmModuleCallCompiler(f.__name__),
+        True,
+        signature=None,
+    )
+    DEF_STORE.register_def(func, get_calling_frame())
+    return GuppyFunctionDefinition(func)
+
+
+# TODO sphinx hack for custom decorators

guppylang_internals/definition/common.py

@@ -0,0 +1,214 @@
+import ast
+import itertools
+from abc import ABC, abstractmethod
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar, TypeAlias
+
+from hugr.build.dfg import DefinitionBuilder, OpVar
+
+from guppylang_internals.diagnostic import Fatal
+from guppylang_internals.span import SourceMap
+
+if TYPE_CHECKING:
+    from guppylang_internals.checker.core import Globals
+    from guppylang_internals.compiler.core import (
+        CompilerContext,
+        PartiallyMonomorphizedArgs,
+    )
+    from guppylang_internals.tys.param import Parameter
+
+
+RawDef: TypeAlias = "ParsableDef | ParsedDef"
+ParsedDef: TypeAlias = "CheckableDef | CheckedDef"
+CheckedDef: TypeAlias = "CompilableDef | MonomorphizableDef | CompiledDef"
+
+
+@dataclass(frozen=True)
+class DefId:
+    """Unique identifier for definitions across modules.
+
+    This id is persistent across all compilation stages. It can be used to identify a
+    definition at any step in the compilation pipeline.
+
+    Args:
+        id: An integer uniquely identifying the definition.
+        module: The module where the definition was defined.
+    """
+
+    id: int
+
+    _ids: ClassVar[Iterator[int]] = itertools.count()
+
+    @classmethod
+    def fresh(cls) -> "DefId":
+        return DefId(next(cls._ids))
+
+
+@dataclass(frozen=True)
+class Definition(ABC):
+    """Abstract base class for user-defined objects on module-level.
+
+    Each definition is identified by a globally unique id. Furthermore, we store the
+    user-picked name for the defined object and an optional AST node for the definition
+    location.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition.
+        defined_at: The AST node where the definition was defined.
+    """
+
+    id: DefId
+    name: str
+    defined_at: ast.AST | None
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Description of this definition to be used in messages to the user.
+
+        The returned text should fit into messages of the following form: "expected
+        a function, but got {description of this definition} instead".
+        """
+
+
+class ParsableDef(Definition):
+    """Abstract base class for raw definitions that still require parsing.
+
+    For example, raw function definitions first need to parse their signature and check
+    that all types are valid. The result of parsing should be a definition that is ready
+    to be checked.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition.
+        defined_at: The AST node where the definition was defined.
+    """
+
+    @abstractmethod
+    def parse(self, globals: "Globals", sources: SourceMap) -> ParsedDef:
+        """Performs parsing and validation, returning a definition that can be checked.
+
+        The provided globals contain all other raw definitions that have been defined.
+        """
+
+
+class CheckableDef(Definition):
+    """Abstract base class for definitions that still need to be checked.
+
+    The result of checking should be a definition that is ready to be compiled to Hugr.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition.
+        defined_at: The AST node where the definition was defined.
+    """
+
+    @abstractmethod
+    def check(self, globals: "Globals") -> CheckedDef:
+        """Type checks the definition.
+
+        The provided globals contain all other parsed definitions that have been
+        defined.
+
+        Returns a checked version of the definition that can be compiled to Hugr.
+        """
+
+
+class CompilableDef(Definition):
+    """Abstract base class for definitions that still need to be compiled to Hugr.
+
+    The result of compilation should be a `CompiledDef` with a pointer to the Hugr node
+    that was created for this definition.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition.
+        defined_at: The AST node where the definition was defined.
+    """
+
+    @abstractmethod
+    def compile_outer(
+        self, module: DefinitionBuilder[OpVar], ctx: "CompilerContext"
+    ) -> "CompiledDef":
+        """Adds a Hugr node for the definition to the provided Hugr module.
+
+        Note that is not required to fill in the contents of the node. At this point,
+        we don't have access to the globals since they have not all been compiled yet.
+
+        See `CompiledDef.compile_inner()` for the hook to compile the inside of the
+        node. This two-step process enables things like mutual recursion.
+        """
+
+
+class MonomorphizableDef(Definition):
+    """Abstract base class for definitions that require monomorphization when compiling
+    to Hugr.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition.
+        defined_at: The AST node where the definition was defined.
+    """
+
+    @property
+    @abstractmethod
+    def params(self) -> "Sequence[Parameter]":
+        """Generic parameters of this definition."""
+
+    @abstractmethod
+    def monomorphize(
+        self,
+        module: DefinitionBuilder[OpVar],
+        mono_args: "PartiallyMonomorphizedArgs",
+        ctx: "CompilerContext",
+    ) -> "MonomorphizedDef":
+        """Adds a Hugr node for the (partially) monomorphized definition to the provided
+        Hugr module.
+
+        See `MonomorphizedDef.compile_inner()` for the hook to compile the inside of the
+        node. This two-step process enables things like mutual recursion.
+        """
+
+
+class CompiledDef(Definition):
+    """Abstract base class for definitions that have been added to a Hugr.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition.
+        defined_at: The AST node where the definition was defined.
+    """
+
+    def compile_inner(self, ctx: "CompilerContext") -> None:
+        """Optional hook that is called to fill in the content of the Hugr node.
+
+        Opposed to `CompilableDef.compile()`, we have access to all other compiled
+        definitions here, which allows things like mutual recursion.
+        """
+
+
+@dataclass(frozen=True)
+class MonomorphizedDef(CompiledDef):
+    """Abstract base class for definitions that have been added to a Hugr and have been
+    partially monomorphized.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the definition. This will be the same for all monomorphized
+            variants of the definition.
+        defined_at: The AST node where the original polymorphic definition was defined.
+        mono_args: Partial monomorphization of the generic parameters.
+    """
+
+    mono_args: "PartiallyMonomorphizedArgs"
+
+
+@dataclass(frozen=True)
+class UnknownSourceError(Fatal):
+    title: ClassVar[str] = "Cannot find source"
+    message: ClassVar[str] = (
+        "Unable to look up the source code for Python object `{obj}`"
+    )
+    obj: object

guppylang_internals/definition/const.py

@@ -0,0 +1,74 @@
+import ast
+from dataclasses import dataclass, field
+
+from hugr import Node, Wire
+from hugr import val as hv
+from hugr.build.dfg import DefinitionBuilder, OpVar
+
+from guppylang_internals.ast_util import AstNode
+from guppylang_internals.checker.core import Globals
+from guppylang_internals.compiler.core import CompilerContext, DFContainer
+from guppylang_internals.definition.common import CompilableDef, ParsableDef
+from guppylang_internals.definition.value import (
+    CompiledHugrNodeDef,
+    CompiledValueDef,
+    ValueDef,
+)
+from guppylang_internals.span import SourceMap
+from guppylang_internals.tys.parsing import type_from_ast
+
+
+@dataclass(frozen=True)
+class RawConstDef(ParsableDef):
+    """A raw constant definition as provided by the user."""
+
+    type_ast: ast.expr
+    value: hv.Value
+
+    description: str = field(default="constant", init=False)
+
+    def parse(self, globals: Globals, sources: SourceMap) -> "ConstDef":
+        """Parses and checks the user-provided signature of the function."""
+        return ConstDef(
+            self.id,
+            self.name,
+            self.defined_at,
+            type_from_ast(self.type_ast, globals, {}),
+            self.type_ast,
+            self.value,
+        )
+
+
+@dataclass(frozen=True)
+class ConstDef(RawConstDef, ValueDef, CompilableDef):
+    """A constant with a checked type."""
+
+    def compile_outer(
+        self, graph: DefinitionBuilder[OpVar], ctx: CompilerContext
+    ) -> "CompiledConstDef":
+        const_node = graph.add_const(self.value)
+        return CompiledConstDef(
+            self.id,
+            self.name,
+            self.defined_at,
+            self.ty,
+            self.type_ast,
+            self.value,
+            const_node,
+        )
+
+
+@dataclass(frozen=True)
+class CompiledConstDef(ConstDef, CompiledValueDef, CompiledHugrNodeDef):
+    """A constant that has been compiled to a Hugr node."""
+
+    const_node: Node
+
+    @property
+    def hugr_node(self) -> Node:
+        """The Hugr node this definition was compiled into."""
+        return self.const_node
+
+    def load(self, dfg: DFContainer, ctx: CompilerContext, node: AstNode) -> Wire:
+        """Loads the extern value into a local Hugr dataflow graph."""
+        return dfg.builder.load(self.const_node)