guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/nodes.py

@@ -0,0 +1,427 @@
+"""Custom AST nodes used by Guppy"""
+
+import ast
+from collections.abc import Mapping
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+from guppylang_internals.ast_util import AstNode
+from guppylang_internals.tys.const import Const
+from guppylang_internals.tys.subst import Inst
+from guppylang_internals.tys.ty import FunctionType, StructType, TupleType, Type
+
+if TYPE_CHECKING:
+    from guppylang_internals.cfg.cfg import CFG
+    from guppylang_internals.checker.cfg_checker import CheckedCFG
+    from guppylang_internals.checker.core import Place, Variable
+    from guppylang_internals.definition.common import DefId
+    from guppylang_internals.definition.struct import StructField
+    from guppylang_internals.tys.param import ConstParam
+
+
+class PlaceNode(ast.expr):
+    place: "Place"
+
+    _fields = ("place",)
+
+
+class GlobalName(ast.Name):
+    id: str
+    def_id: "DefId"
+
+    _fields = (
+        "id",
+        "def_id",
+    )
+
+
+class GenericParamValue(ast.Name):
+    id: str
+    param: "ConstParam"
+
+    _fields = (
+        "id",
+        "param",
+    )
+
+
+class LocalCall(ast.expr):
+    func: ast.expr
+    args: list[ast.expr]
+
+    _fields = (
+        "func",
+        "args",
+    )
+
+
+class GlobalCall(ast.expr):
+    def_id: "DefId"
+    args: list[ast.expr]
+    type_args: Inst  # Inferred type arguments
+
+    _fields = (
+        "def_id",
+        "args",
+        "type_args",
+    )
+
+
+class TensorCall(ast.expr):
+    """A call to a tuple of functions. Behaves like a local call, but more
+    unpacking of tuples is required at compilation"""
+
+    func: ast.expr
+    args: list[ast.expr]
+    tensor_ty: FunctionType
+
+    _fields = (
+        "func",
+        "args",
+        "tensor_ty",
+    )
+
+
+class TypeApply(ast.expr):
+    value: ast.expr
+    inst: Inst
+
+    _fields = (
+        "value",
+        "inst",
+    )
+
+
+class PartialApply(ast.expr):
+    """A partial function application.
+
+    This node is emitted when methods are loaded as values, since this requires
+    partially applying the `self` argument.
+    """
+
+    func: ast.expr
+    args: list[ast.expr]
+
+    _fields = (
+        "func",
+        "args",
+    )
+
+
+class FieldAccessAndDrop(ast.expr):
+    """A field access on a struct, dropping all the remaining other fields."""
+
+    value: ast.expr
+    struct_ty: "StructType"
+    field: "StructField"
+
+    _fields = (
+        "value",
+        "struct_ty",
+        "field",
+    )
+
+
+class SubscriptAccessAndDrop(ast.expr):
+    """A subscript element access on an object, dropping all the remaining items."""
+
+    item: "Variable"
+    item_expr: ast.expr
+    getitem_expr: ast.expr
+    original_expr: ast.Subscript
+
+    _fields = ("item", "item_expr", "getitem_expr", "original_expr")
+
+
+class TupleAccessAndDrop(ast.expr):
+    """A subscript element access on a tuple, dropping all the remaining items."""
+
+    value: ast.expr
+    tuple_ty: TupleType
+    index: int
+
+    _fields = ("value", "tuple_ty", "index")
+
+
+class MakeIter(ast.expr):
+    """Creates an iterator using the `__iter__` magic method.
+
+    This node is inserted in `for` loops and list comprehensions.
+    """
+
+    value: ast.expr
+    unwrap_size_hint: bool
+
+    # Node that triggered the creation of this iterator. For example, a for loop stmt.
+    # It is not mentioned in `_fields` so that it is not visible to AST visitors
+    origin_node: ast.AST
+
+    _fields = ("value",)
+
+    def __init__(
+        self, value: ast.expr, origin_node: ast.AST, unwrap_size_hint: bool = True
+    ) -> None:
+        super().__init__(value)
+        self.origin_node = origin_node
+        self.unwrap_size_hint = unwrap_size_hint
+
+
+class IterHasNext(ast.expr):
+    """Checks if an iterator has a next element using the `__hasnext__` magic method.
+
+    This node is inserted in `for` loops and list comprehensions.
+    """
+
+    value: ast.expr
+
+    _fields = ("value",)
+
+
+class IterNext(ast.expr):
+    """Obtains the next element of an iterator using the `__next__` magic method.
+
+    This node is inserted in `for` loops and list comprehensions.
+    """
+
+    value: ast.expr
+
+    _fields = ("value",)
+
+
+class IterEnd(ast.expr):
+    """Finalises an iterator using the `__end__` magic method.
+
+    This node is inserted in `for` loops and list comprehensions. It is needed to
+    consume linear iterators once they are finished.
+    """
+
+    value: ast.expr
+
+    _fields = ("value",)
+
+
+class DesugaredGenerator(ast.expr):
+    """A single desugared generator in a list comprehension.
+
+    Stores assignments of the original generator targets as well as dummy variables for
+    the iterator and hasnext test.
+    """
+
+    iter_assign: ast.Assign
+    next_call: ast.expr
+    iter: ast.expr
+    target: ast.expr
+    ifs: list[ast.expr]
+
+    used_outer_places: "list[Place]"
+
+    _fields = (
+        "iter_assign",
+        "next_call",
+        "iter",
+        "target",
+        "ifs",
+    )
+
+
+class DesugaredGeneratorExpr(ast.expr):
+    """A desugared generator expression."""
+
+    elt: ast.expr
+    generators: list[DesugaredGenerator]
+
+    _fields = (
+        "elt",
+        "generators",
+    )
+
+
+class DesugaredListComp(ast.expr):
+    """A desugared list comprehension."""
+
+    elt: ast.expr
+    generators: list[DesugaredGenerator]
+
+    _fields = (
+        "elt",
+        "generators",
+    )
+
+
+class DesugaredArrayComp(ast.expr):
+    """A desugared array comprehension."""
+
+    elt: ast.expr
+    generator: DesugaredGenerator
+    length: Const
+    elt_ty: Type
+
+    _fields = (
+        "elt",
+        "generator",
+        "length",
+        "elt_ty",
+    )
+
+
+class ComptimeExpr(ast.expr):
+    """A compile-time evaluated `py(...)` expression."""
+
+    value: ast.expr
+
+    _fields = ("value",)
+
+
+class ResultExpr(ast.expr):
+    """A `result(tag, value)` expression."""
+
+    value: ast.expr
+    base_ty: Type
+    #: Array length in case this is an array result, otherwise `None`
+    array_len: Const | None
+    tag: str
+
+    _fields = ("value", "base_ty", "array_len", "tag")
+
+    @property
+    def args(self) -> list[ast.expr]:
+        return [self.value]
+
+
+class ExitKind(Enum):
+    ExitShot = 0  # Exit the current shot
+    Panic = 1  # Panic the program ending all shots
+
+
+class PanicExpr(ast.expr):
+    """A `panic(msg, *args)` or `exit(msg, *args)` expression ."""
+
+    kind: ExitKind
+    signal: int
+    msg: str
+    values: list[ast.expr]
+
+    _fields = ("kind", "signal", "msg", "values")
+
+
+class BarrierExpr(ast.expr):
+    """A `barrier(*args)` expression."""
+
+    args: list[ast.expr]
+    func_ty: FunctionType
+    _fields = ("args", "func_ty")
+
+
+class StateResultExpr(ast.expr):
+    """A `state_result(tag, *args)` expression."""
+
+    tag: str
+    args: list[ast.expr]
+    func_ty: FunctionType
+    #: Array length in case this is an array result, otherwise `None`
+    array_len: Const | None
+    _fields = ("tag", "args", "func_ty", "has_array_input")
+
+
+AnyCall = (
+    LocalCall | GlobalCall | TensorCall | BarrierExpr | ResultExpr | StateResultExpr
+)
+
+
+class InoutReturnSentinel(ast.expr):
+    """An invisible expression corresponding to an implicit use of borrowed vars
+    whenever a function returns."""
+
+    var: "Place | str"
+
+    _fields = ("var",)
+
+
+class UnpackPattern(ast.expr):
+    """The LHS of an unpacking assignment like `a, *bs, c = ...` or
+    `[a, *bs, c] = ...`."""
+
+    #: Patterns occurring on the left of the starred target
+    left: list[ast.expr]
+
+    #: The starred target or `None` if there is none
+    starred: ast.expr | None
+
+    #: Patterns occurring on the right of the starred target. This will be an empty list
+    #: if there is no starred target
+    right: list[ast.expr]
+
+    _fields = ("left", "starred", "right")
+
+
+class TupleUnpack(ast.expr):
+    """The LHS of an unpacking assignment of a tuple."""
+
+    #: The (possibly starred) unpacking pattern
+    pattern: UnpackPattern
+
+    _fields = ("pattern",)
+
+
+class IterableUnpack(ast.expr):
+    """The LHS of an unpacking assignment of an iterable type."""
+
+    #: The (possibly starred) unpacking pattern
+    pattern: UnpackPattern
+
+    #: Comprehension that collects the RHS iterable into an array
+    compr: DesugaredArrayComp
+
+    #: Dummy variable that the RHS should be bound to. This variable is referenced in
+    #: `compr`
+    rhs_var: PlaceNode
+
+    # Don't mention the comprehension in _fields to avoid visitors recursing it
+    _fields = ("pattern",)
+
+    def __init__(
+        self, pattern: UnpackPattern, compr: DesugaredArrayComp, rhs_var: PlaceNode
+    ) -> None:
+        super().__init__(pattern)
+        self.compr = compr
+        self.rhs_var = rhs_var
+
+
+#: Any unpacking operation.
+AnyUnpack = TupleUnpack | IterableUnpack
+
+
+class NestedFunctionDef(ast.FunctionDef):
+    cfg: "CFG"
+    ty: FunctionType
+    docstring: str | None
+
+    def __init__(self, cfg: "CFG", ty: FunctionType, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.cfg = cfg
+        self.ty = ty
+
+
+class CheckedNestedFunctionDef(ast.FunctionDef):
+    def_id: "DefId"
+    cfg: "CheckedCFG[Place]"
+    ty: FunctionType
+
+    #: Mapping from names to variables captured by this function, together with an AST
+    #: node witnessing a use of the captured variable in the function body.
+    captured: Mapping[str, tuple["Variable", AstNode]]
+
+    def __init__(
+        self,
+        def_id: "DefId",
+        cfg: "CheckedCFG[Place]",
+        ty: FunctionType,
+        captured: Mapping[str, tuple["Variable", AstNode]],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.def_id = def_id
+        self.cfg = cfg
+        self.ty = ty
+        self.captured = captured

guppylang_internals/span.py

@@ -0,0 +1,150 @@
+"""Source spans representing locations in the code being compiled."""
+
+import ast
+import linecache
+from dataclasses import dataclass
+from typing import TypeAlias
+
+from guppylang_internals.ast_util import get_file, get_line_offset
+from guppylang_internals.error import InternalGuppyError
+from guppylang_internals.ipython_inspect import normalize_ipython_dummy_files
+
+
+@dataclass(frozen=True, order=True)
+class Loc:
+    """A location in a source file."""
+
+    file: str
+
+    #: Line number starting at 1
+    line: int
+
+    #: Column number starting at 1
+    column: int
+
+    def __str__(self) -> str:
+        """Returns the string representation of this source location."""
+        file = normalize_ipython_dummy_files(self.file)
+        return f"{file}:{self.line}:{self.column}"
+
+    def shift_left(self, cols: int) -> "Loc":
+        """Returns a new location shifted to left by the given number of columns."""
+        assert self.column >= cols
+        return Loc(self.file, self.line, self.column - cols)
+
+    def shift_right(self, cols: int) -> "Loc":
+        """Returns a new location shifted to right by the given number of columns."""
+        return Loc(self.file, self.line, self.column + cols)
+
+
+@dataclass(frozen=True)
+class Span:
+    """A continuous sequence of source code within a file."""
+
+    #: Starting location of the span (inclusive)
+    start: Loc
+
+    # Ending location of the span (exclusive)
+    end: Loc
+
+    def __post_init__(self) -> None:
+        if self.start.file != self.end.file:
+            raise InternalGuppyError("Span: Source spans multiple files")
+        if self.start > self.end:
+            raise InternalGuppyError("Span: Start after end")
+
+    def __contains__(self, x: "Span | Loc") -> bool:
+        """Determines whether another span or location is completely contained in this
+        span."""
+        if self.file != x.file:
+            return False
+        if isinstance(x, Span):
+            return self.start <= x.start <= self.end <= x.end
+        return self.start <= x <= self.end
+
+    def __and__(self, other: "Span") -> "Span | None":
+        """Returns the intersection with the given span or `None` if they don't
+        intersect."""
+        if self.file != other.file:
+            return None
+        if self.start > other.end or other.start > self.end:
+            return None
+        return Span(max(self.start, other.start), min(self.end, other.end))
+
+    def __len__(self) -> int:
+        """Returns the length of a single-line span in columns.
+
+        Querying the length of multiline spans raises an `InternalGuppyError`.
+        """
+        if self.is_multiline:
+            raise InternalGuppyError("Span: Tried to compute length of multi-line span")
+        return self.end.column - self.start.column
+
+    @property
+    def file(self) -> str:
+        """The file containing this span."""
+        return self.start.file
+
+    @property
+    def is_multiline(self) -> bool:
+        """Whether this source sequence spans multiple lines."""
+        return self.start.line != self.end.line
+
+    def shift_left(self, cols: int) -> "Span":
+        """Returns a new span that is shifted to the left by the given number of
+        columns."""
+        return Span(self.start.shift_left(cols), self.end.shift_left(cols))
+
+    def shift_right(self, cols: int) -> "Span":
+        """Returns a new span that is shifted to the right by the given number of
+        columns."""
+        return Span(self.start.shift_right(cols), self.end.shift_right(cols))
+
+
+#: Objects in the compiler that are associated with a source span
+ToSpan: TypeAlias = ast.AST | Span
+
+
+def to_span(x: ToSpan) -> Span:
+    """Extracts a source span from an object."""
+    if isinstance(x, Span):
+        return x
+    file, line_offset = get_file(x), get_line_offset(x)
+    assert file is not None
+    assert line_offset is not None
+    # x.lineno and line_offset both start at 1, so we have to subtract 1
+    start = Loc(file, x.lineno + line_offset - 1, x.col_offset)
+    end = Loc(
+        file,
+        (x.end_lineno or x.lineno) + line_offset - 1,
+        x.end_col_offset or x.col_offset,
+    )
+    return Span(start, end)
+
+
+#: List of source lines in a file
+SourceLines: TypeAlias = list[str]
+
+
+class SourceMap:
+    """Map holding the source code for all files accessed by the compiler.
+
+    Can be used to look up the source code associated with a span.
+    """
+
+    sources: dict[str, SourceLines]
+
+    def __init__(self) -> None:
+        self.sources = {}
+
+    def add_file(self, file: str, content: str | None = None) -> None:
+        """Registers a new source file."""
+        if content is None:
+            self.sources[file] = [line.rstrip() for line in linecache.getlines(file)]
+        else:
+            self.sources[file] = content.splitlines(keepends=False)
+
+    def span_lines(self, span: Span, prefix_lines: int = 0) -> list[str]:
+        return self.sources[span.file][
+            span.start.line - prefix_lines - 1 : span.end.line
+        ]