guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/__init__.py

@@ -0,0 +1,3 @@
+# This is updated by our release-please workflow, triggered by this
+# annotation: x-release-please-version
+__version__ = "0.21.0"

guppylang_internals/ast_util.py

@@ -0,0 +1,350 @@
+import ast
+import textwrap
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, cast
+
+if TYPE_CHECKING:
+    from guppylang_internals.tys.ty import Type
+
+AstNode = (
+    ast.AST
+    | ast.operator
+    | ast.expr
+    | ast.arg
+    | ast.stmt
+    | ast.Name
+    | ast.keyword
+    | ast.FunctionDef
+)
+
+T = TypeVar("T", covariant=True)
+
+
+class AstVisitor(Generic[T]):
+    """
+    Note: This class is based on the implementation of `ast.NodeVisitor` but
+    allows extra arguments to be passed to the `visit` functions.
+
+    Original documentation:
+
+    A node visitor base class that walks the abstract syntax tree and calls a
+    visitor function for every node found.  This function may return a value
+    which is forwarded by the `visit` method.
+
+    This class is meant to be subclassed, with the subclass adding visitor
+    methods.
+
+    Per default the visitor functions for the nodes are ``'visit_'`` +
+    class name of the node.  So a `TryFinally` node visit function would
+    be `visit_TryFinally`.  This behavior can be changed by overriding
+    the `visit` method.  If no visitor function exists for a node
+    (return value `None`) the `generic_visit` visitor is used instead.
+
+    Don't use the `NodeVisitor` if you want to apply changes to nodes during
+    traversing.  For this a special visitor exists (`NodeTransformer`) that
+    allows modifications.
+    """
+
+    def visit(self, node: Any, *args: Any, **kwargs: Any) -> T:
+        """Visit a node."""
+        method = "visit_" + node.__class__.__name__
+        visitor = getattr(self, method, self.generic_visit)
+        return visitor(node, *args, **kwargs)
+
+    def generic_visit(self, node: Any, *args: Any, **kwargs: Any) -> T:
+        """Called if no explicit visitor function exists for a node."""
+        raise NotImplementedError(f"visit_{node.__class__.__name__} is not implemented")
+
+
+class AstSearcher(ast.NodeVisitor):
+    """Visitor that searches for occurrences of specific nodes in an AST."""
+
+    matcher: Callable[[ast.AST], bool]
+    dont_recurse_into: set[type[ast.AST]]
+    found: list[ast.AST]
+    is_first_node: bool
+
+    def __init__(
+        self,
+        matcher: Callable[[ast.AST], bool],
+        dont_recurse_into: set[type[ast.AST]] | None = None,
+    ) -> None:
+        self.matcher = matcher
+        self.dont_recurse_into = dont_recurse_into or set()
+        self.found = []
+        self.is_first_node = True
+
+    def generic_visit(self, node: ast.AST) -> None:
+        if self.matcher(node):
+            self.found.append(node)
+        if self.is_first_node or type(node) not in self.dont_recurse_into:
+            self.is_first_node = False
+            super().generic_visit(node)
+
+
+def find_nodes(
+    matcher: Callable[[ast.AST], bool],
+    node: ast.AST,
+    dont_recurse_into: set[type[ast.AST]] | None = None,
+) -> list[ast.AST]:
+    """Returns all nodes in the AST that satisfy the matcher."""
+    v = AstSearcher(matcher, dont_recurse_into)
+    v.visit(node)
+    return v.found
+
+
+def name_nodes_in_ast(node: Any) -> list[ast.Name]:
+    """Returns all `Name` nodes occurring in an AST."""
+    found = find_nodes(lambda n: isinstance(n, ast.Name), node)
+    return cast(list[ast.Name], found)
+
+
+def return_nodes_in_ast(node: Any) -> list[ast.Return]:
+    """Returns all `Return` nodes occurring in an AST."""
+    found = find_nodes(lambda n: isinstance(n, ast.Return), node, {ast.FunctionDef})
+    return cast(list[ast.Return], found)
+
+
+def breaks_in_loop(node: Any) -> list[ast.Break]:
+    """Returns all `Break` nodes occurring in a loop.
+
+    Note that breaks in nested loops are excluded.
+    """
+    found = find_nodes(
+        lambda n: isinstance(n, ast.Break), node, {ast.For, ast.While, ast.FunctionDef}
+    )
+    return cast(list[ast.Break], found)
+
+
+class ContextAdjuster(ast.NodeTransformer):
+    """Updates the `ast.Context` indicating if expressions occur on the LHS or RHS."""
+
+    ctx: ast.expr_context
+
+    def __init__(self, ctx: ast.expr_context) -> None:
+        self.ctx = ctx
+
+    def visit(self, node: ast.AST) -> ast.AST:
+        return cast(ast.AST, super().visit(node))
+
+    def visit_Name(self, node: ast.Name) -> ast.Name:
+        return with_loc(node, ast.Name(id=node.id, ctx=self.ctx))
+
+    def visit_Starred(self, node: ast.Starred) -> ast.Starred:
+        return with_loc(node, ast.Starred(value=self.visit(node.value), ctx=self.ctx))
+
+    def visit_Tuple(self, node: ast.Tuple) -> ast.Tuple:
+        return with_loc(
+            node, ast.Tuple(elts=[self.visit(elt) for elt in node.elts], ctx=self.ctx)
+        )
+
+    def visit_List(self, node: ast.List) -> ast.List:
+        return with_loc(
+            node, ast.List(elts=[self.visit(elt) for elt in node.elts], ctx=self.ctx)
+        )
+
+    def visit_Subscript(self, node: ast.Subscript) -> ast.Subscript:
+        # Don't adjust the slice!
+        return with_loc(
+            node,
+            ast.Subscript(value=self.visit(node.value), slice=node.slice, ctx=self.ctx),
+        )
+
+    def visit_Attribute(self, node: ast.Attribute) -> ast.Attribute:
+        return with_loc(
+            node,
+            ast.Attribute(value=self.visit(node.value), attr=node.attr, ctx=self.ctx),
+        )
+
+
+@dataclass(frozen=True, eq=False)
+class TemplateReplacer(ast.NodeTransformer):
+    """Replaces nodes in a template."""
+
+    replacements: Mapping[str, ast.AST | Sequence[ast.AST]]
+    default_loc: ast.AST
+
+    def _get_replacement(self, x: str) -> ast.AST | Sequence[ast.AST]:
+        if x not in self.replacements:
+            msg = f"No replacement for `{x}` is given"
+            raise ValueError(msg)
+        return self.replacements[x]
+
+    def visit_Name(self, node: ast.Name) -> ast.AST:
+        repl = self._get_replacement(node.id)
+        if not isinstance(repl, ast.expr):
+            msg = f"Replacement for `{node.id}` must be an expression"
+            raise TypeError(msg)
+
+        # Update the context
+        adjuster = ContextAdjuster(node.ctx)
+        return with_loc(repl, adjuster.visit(repl))
+
+    def visit_Expr(self, node: ast.Expr) -> ast.AST | Sequence[ast.AST]:
+        if isinstance(node.value, ast.Name):
+            repl = self._get_replacement(node.value.id)
+            repls = [repl] if not isinstance(repl, Sequence) else repl
+            # Wrap expressions to turn them into statements
+            return [
+                with_loc(r, ast.Expr(value=r)) if isinstance(r, ast.expr) else r
+                for r in repls
+            ]
+        return self.generic_visit(node)
+
+    def generic_visit(self, node: ast.AST) -> ast.AST:
+        # Insert the default location
+        node = super().generic_visit(node)
+        return with_loc(self.default_loc, node)
+
+
+def template_replace(
+    template: str, default_loc: ast.AST, **kwargs: ast.AST | Sequence[ast.AST]
+) -> list[ast.stmt]:
+    """Turns a template into a proper AST by substituting all placeholders."""
+    nodes = ast.parse(textwrap.dedent(template)).body
+    replacer = TemplateReplacer(kwargs, default_loc)
+    new_nodes = []
+    for n in nodes:
+        new = replacer.visit(n)
+        if isinstance(new, list):
+            new_nodes.extend(new)
+        else:
+            new_nodes.append(new)
+    return new_nodes
+
+
+def line_col(node: ast.AST) -> tuple[int, int]:
+    """Returns the line and column of an ast node."""
+    return node.lineno, node.col_offset
+
+
+def set_location_from(node: ast.AST, loc: ast.AST) -> None:
+    """Copy source location from one AST node to the other."""
+    node.lineno = loc.lineno
+    node.col_offset = loc.col_offset
+    node.end_lineno = loc.end_lineno
+    node.end_col_offset = loc.end_col_offset
+
+    source, file, line_offset = get_source(loc), get_file(loc), get_line_offset(loc)
+    assert source is not None
+    assert file is not None
+    assert line_offset is not None
+    annotate_location(node, source, file, line_offset)
+
+
+def annotate_location(
+    node: ast.AST, source: str, file: str, line_offset: int, recurse: bool = True
+) -> None:
+    node.line_offset = line_offset  # type: ignore[attr-defined]
+    node.file = file  # type: ignore[attr-defined]
+    node.source = source  # type: ignore[attr-defined]
+
+    if recurse:
+        for _field, value in ast.iter_fields(node):
+            if isinstance(value, list):
+                for item in value:
+                    if isinstance(item, ast.AST):
+                        annotate_location(item, source, file, line_offset, recurse)
+            elif isinstance(value, ast.AST):
+                annotate_location(value, source, file, line_offset, recurse)
+
+
+def shift_loc(node: ast.AST, delta_lineno: int, delta_col_offset: int) -> None:
+    """Shifts all line and column number in the AST node by the given amount."""
+    if hasattr(node, "lineno"):
+        node.lineno += delta_lineno
+    if hasattr(node, "end_lineno") and node.end_lineno is not None:
+        node.end_lineno += delta_lineno
+    if hasattr(node, "col_offset"):
+        node.col_offset += delta_col_offset
+    if hasattr(node, "end_col_offset") and node.end_col_offset is not None:
+        node.end_col_offset += delta_col_offset
+    for _, value in ast.iter_fields(node):
+        if isinstance(value, list):
+            for item in value:
+                if isinstance(item, ast.AST):
+                    shift_loc(item, delta_lineno, delta_col_offset)
+        elif isinstance(value, ast.AST):
+            shift_loc(value, delta_lineno, delta_col_offset)
+
+
+def get_file(node: AstNode) -> str | None:
+    """Tries to retrieve a file annotation from an AST node."""
+    try:
+        file = node.file  # type: ignore[union-attr]
+        return file if isinstance(file, str) else None
+    except AttributeError:
+        return None
+
+
+def get_source(node: AstNode) -> str | None:
+    """Tries to retrieve a source annotation from an AST node."""
+    try:
+        source = node.source  # type: ignore[union-attr]
+        return source if isinstance(source, str) else None
+    except AttributeError:
+        return None
+
+
+def get_line_offset(node: AstNode) -> int | None:
+    """Tries to retrieve a line offset annotation from an AST node."""
+    try:
+        line_offset = node.line_offset  # type: ignore[union-attr]
+        return line_offset if isinstance(line_offset, int) else None
+    except AttributeError:
+        return None
+
+
+A = TypeVar("A", bound=ast.AST)
+
+
+def with_loc(loc: ast.AST, node: A) -> A:
+    """Copy source location from one AST node to the other."""
+    set_location_from(node, loc)
+    return node
+
+
+def with_type(ty: "Type", node: A) -> A:
+    """Annotates an AST node with a type."""
+    node.type = ty  # type: ignore[attr-defined]
+    return node
+
+
+def get_type_opt(node: AstNode) -> Optional["Type"]:
+    """Tries to retrieve a type annotation from an AST node."""
+    from guppylang_internals.tys.ty import Type, TypeBase
+
+    try:
+        ty = node.type  # type: ignore[union-attr]
+        return cast(Type, ty) if isinstance(ty, TypeBase) else None
+    except AttributeError:
+        return None
+
+
+def get_type(node: AstNode) -> "Type":
+    """Retrieve a type annotation from an AST node.
+
+    Fails if the node is not annotated.
+    """
+    ty = get_type_opt(node)
+    assert ty is not None
+    return ty
+
+
+def has_empty_body(func_ast: ast.FunctionDef) -> bool:
+    """Returns `True` if the body of a function definition is empty.
+
+    This is the case if the body only contains a single `pass` statement or an ellipsis
+    `...` expression.
+    """
+    if len(func_ast.body) == 0:
+        return True
+    if len(func_ast.body) > 1:
+        return False
+    [n] = func_ast.body
+    return (
+        isinstance(n, ast.Expr)
+        and isinstance(n.value, ast.Constant)
+        and n.value.value == Ellipsis
+    )

guppylang_internals/cfg/analysis.py

@@ -0,0 +1,230 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from typing import Generic, TypeVar
+
+from guppylang_internals.cfg.bb import BB, VariableStats, VId
+
+# Type variable for the lattice domain
+T = TypeVar("T")
+
+# Analysis result is a mapping from basic blocks to lattice values
+Result = dict[BB, T]
+
+
+class Analysis(Generic[T], ABC):
+    """Abstract base class for a program analysis pass over the lattice `T`"""
+
+    def eq(self, t1: T, t2: T, /) -> bool:
+        """Equality on lattice values"""
+        return t1 == t2
+
+    @abstractmethod
+    def include_unreachable(self) -> bool:
+        """Whether unreachable BBs and jumps should be taken into account for the
+        analysis."""
+
+    @abstractmethod
+    def initial(self) -> T:
+        """Initial lattice value"""
+
+    @abstractmethod
+    def join(self, *ts: T) -> T:
+        """Lattice join operation"""
+
+    @abstractmethod
+    def run(self, bbs: Iterable[BB]) -> Result[T]:
+        """Runs the analysis pass.
+
+        Returns a mapping from basic blocks to lattice values at the start of each BB.
+        """
+
+
+class ForwardAnalysis(Generic[T], Analysis[T], ABC):
+    """Abstract base class for a program analysis pass running in forward direction."""
+
+    @abstractmethod
+    def apply_bb(self, val_before: T, bb: BB, /) -> T:
+        """Transformation a basic block applies to a lattice value"""
+
+    def run(self, bbs: Iterable[BB]) -> Result[T]:
+        """Runs the analysis pass.
+
+        Returns a mapping from basic blocks to lattice values at the start of each BB.
+        """
+        if not self.include_unreachable():
+            bbs = [bb for bb in bbs if bb.reachable]
+        vals_before = {bb: self.initial() for bb in bbs}  # return value
+        vals_after = {bb: self.apply_bb(vals_before[bb], bb) for bb in bbs}  # cache
+        queue = set(bbs)
+        while len(queue) > 0:
+            bb = queue.pop()
+            preds = (
+                bb.predecessors + bb.dummy_predecessors
+                if self.include_unreachable()
+                else bb.predecessors
+            )
+            vals_before[bb] = self.join(*(vals_after[pred] for pred in preds))
+            val_after = self.apply_bb(vals_before[bb], bb)
+            if not self.eq(val_after, vals_after[bb]):
+                vals_after[bb] = val_after
+                queue.update(bb.successors)
+        return vals_before
+
+
+class BackwardAnalysis(Generic[T], Analysis[T], ABC):
+    """Abstract base class for a program analysis pass running in backward direction."""
+
+    @abstractmethod
+    def apply_bb(self, val_after: T, bb: BB, /) -> T:
+        """Transformation a basic block applies to a lattice value"""
+
+    def run(self, bbs: Iterable[BB]) -> Result[T]:
+        """Runs the analysis pass.
+
+        Returns a mapping from basic blocks to lattice values at the start of each BB.
+        """
+        vals_before = {bb: self.initial() for bb in bbs}
+        queue = set(bbs)
+        while len(queue) > 0:
+            bb = queue.pop()
+            succs = (
+                bb.successors + bb.dummy_successors
+                if self.include_unreachable()
+                else bb.successors
+            )
+            val_after = self.join(*(vals_before[succ] for succ in succs))
+            val_before = self.apply_bb(val_after, bb)
+            if not self.eq(vals_before[bb], val_before):
+                vals_before[bb] = val_before
+                queue.update(bb.predecessors)
+        return vals_before
+
+
+# For live variable analysis, we also store a BB in which a use occurs as evidence of
+# liveness.
+LivenessDomain = dict[VId, BB]
+
+
+class LivenessAnalysis(Generic[VId], BackwardAnalysis[LivenessDomain[VId]]):
+    """Live variable analysis pass.
+
+    Computes the variables that are live before the execution of each BB. The analysis
+    runs over the lattice of mappings from variable names to BBs containing a use.
+    """
+
+    stats: dict[BB, VariableStats[VId]]
+
+    def __init__(
+        self,
+        stats: dict[BB, VariableStats[VId]],
+        initial: LivenessDomain[VId] | None = None,
+        include_unreachable: bool = False,
+    ) -> None:
+        self.stats = stats
+        self._initial = initial or {}
+        self._include_unreachable = include_unreachable
+
+    def eq(self, live1: LivenessDomain[VId], live2: LivenessDomain[VId]) -> bool:
+        # Only check that both contain the same variables. We don't care about the BB
+        # in which the use occurs, we just need any one, to report to the user.
+        return live1.keys() == live2.keys()
+
+    def initial(self) -> LivenessDomain[VId]:
+        return self._initial
+
+    def include_unreachable(self) -> bool:
+        return self._include_unreachable
+
+    def join(self, *ts: LivenessDomain[VId]) -> LivenessDomain[VId]:
+        res: LivenessDomain[VId] = {}
+        for t in ts:
+            res |= t
+        return res
+
+    def apply_bb(self, live_after: LivenessDomain[VId], bb: BB) -> LivenessDomain[VId]:
+        stats = self.stats[bb]
+        return {x: bb for x in stats.used} | {
+            x: b for x, b in live_after.items() if x not in stats.assigned
+        }
+
+
+# Set of variables that are definitely assigned at the start of a BB
+DefAssignmentDomain = set[VId]
+
+# Set of variables that are assigned on (at least) some paths to a BB. Definitely
+# assigned variables are a subset of this
+MaybeAssignmentDomain = set[VId]
+
+# For assignment analysis, we do definite- and maybe-assignment in one pass
+AssignmentDomain = tuple[DefAssignmentDomain[VId], MaybeAssignmentDomain[VId]]
+
+
+class AssignmentAnalysis(Generic[VId], ForwardAnalysis[AssignmentDomain[VId]]):
+    """Assigned variable analysis pass.
+
+    Computes the set of variables (i.e. `V`s) that are definitely assigned at the start
+    of a BB. Additionally, we compute the set of variables that are assigned on (at
+    least) some paths to a BB (the definitely assigned variables are a subset of this).
+    """
+
+    stats: dict[BB, VariableStats[VId]]
+    all_vars: set[VId]
+    ass_before_entry: set[VId]
+    maybe_ass_before_entry: set[VId]
+
+    def __init__(
+        self,
+        stats: dict[BB, VariableStats[VId]],
+        ass_before_entry: set[VId],
+        maybe_ass_before_entry: set[VId],
+        include_unreachable: bool = False,
+    ) -> None:
+        """Constructs an `AssignmentAnalysis` pass for a CFG.
+
+        Also takes a set variables that are definitely assigned before the entry of the
+        CFG (for example function arguments).
+        """
+        assert ass_before_entry.issubset(maybe_ass_before_entry)
+        self.stats = stats
+        self.ass_before_entry = ass_before_entry
+        self.maybe_ass_before_entry = maybe_ass_before_entry
+        self.all_vars = (
+            set.union(*(set(stat.assigned.keys()) for stat in stats.values()))
+            | ass_before_entry
+        )
+        self._include_unreachable = include_unreachable
+
+    def initial(self) -> AssignmentDomain[VId]:
+        # Note that definite assignment must start with `all_vars` instead of only
+        # `ass_before_entry` since we want to compute the *greatest* fixpoint.
+        return self.all_vars, self.maybe_ass_before_entry
+
+    def include_unreachable(self) -> bool:
+        return self._include_unreachable
+
+    def join(self, *ts: AssignmentDomain[VId]) -> AssignmentDomain[VId]:
+        # We always include the variables that are definitely assigned before the entry,
+        # even if the join is empty
+        if len(ts) == 0:
+            return self.ass_before_entry, self.ass_before_entry
+
+        def_ass = set.intersection(*(def_ass for def_ass, _ in ts))
+        maybe_ass = set.union(*(maybe_ass for _, maybe_ass in ts))
+        return def_ass, maybe_ass
+
+    def apply_bb(
+        self, val_before: AssignmentDomain[VId], bb: BB
+    ) -> AssignmentDomain[VId]:
+        stats = self.stats[bb]
+        def_ass_before, maybe_ass_before = val_before
+        return (
+            def_ass_before | stats.assigned.keys(),
+            maybe_ass_before | stats.assigned.keys(),
+        )
+
+    def run_unpacked(
+        self, bbs: Iterable[BB]
+    ) -> tuple[Result[DefAssignmentDomain[VId]], Result[MaybeAssignmentDomain[VId]]]:
+        """Runs the analysis and unpacks the definite- and maybe-assignment results."""
+        res = self.run(bbs)
+        return {bb: res[bb][0] for bb in res}, {bb: res[bb][1] for bb in res}