guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/checker/func_checker.py

@@ -0,0 +1,269 @@
+"""Type checking code for top-level and nested function definitions.
+
+For top-level functions, we take the `DefinedFunction` containing the `ast.FunctionDef`
+node straight from the Python AST. We build a CFG, check it, and return a
+`CheckedFunction` containing a `CheckedCFG` with type annotations.
+"""
+
+import ast
+import sys
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+from guppylang_internals.ast_util import return_nodes_in_ast, with_loc
+from guppylang_internals.cfg.bb import BB
+from guppylang_internals.cfg.builder import CFGBuilder
+from guppylang_internals.checker.cfg_checker import CheckedCFG, check_cfg
+from guppylang_internals.checker.core import Context, Globals, Place, Variable
+from guppylang_internals.checker.errors.generic import UnsupportedError
+from guppylang_internals.definition.common import DefId
+from guppylang_internals.diagnostic import Error, Help, Note
+from guppylang_internals.engine import DEF_STORE, ENGINE
+from guppylang_internals.error import GuppyError
+from guppylang_internals.experimental import check_capturing_closures_enabled
+from guppylang_internals.nodes import CheckedNestedFunctionDef, NestedFunctionDef
+from guppylang_internals.tys.parsing import parse_function_io_types
+from guppylang_internals.tys.ty import FunctionType, InputFlags, NoneType
+
+if sys.version_info >= (3, 12):
+    from guppylang_internals.tys.parsing import parse_parameter
+
+if TYPE_CHECKING:
+    from guppylang_internals.tys.param import Parameter
+
+
+@dataclass(frozen=True)
+class IllegalAssignError(Error):
+    title: ClassVar[str] = "Illegal assignment"
+    span_label: ClassVar[str] = (
+        "Variable `{var}` may not be assigned to since `{var}` is captured from an "
+        "outer scope"
+    )
+    var: str
+
+    @dataclass(frozen=True)
+    class DefHint(Note):
+        span_label: ClassVar[str] = "`{var}` defined here"
+        var: str
+
+
+@dataclass(frozen=True)
+class MissingArgAnnotationError(Error):
+    title: ClassVar[str] = "Missing type annotation"
+    span_label: ClassVar[str] = "Argument requires a type annotation"
+
+
+@dataclass(frozen=True)
+class MissingReturnAnnotationError(Error):
+    title: ClassVar[str] = "Missing type annotation"
+    span_label: ClassVar[str] = "Return type must be annotated"
+
+    @dataclass(frozen=True)
+    class ReturnNone(Help):
+        message: ClassVar[str] = (
+            "Looks like `{func}` doesn't return anything. Consider annotating it with "
+            "`-> None`."
+        )
+        func: str
+
+
+def check_global_func_def(
+    func_def: ast.FunctionDef, ty: FunctionType, globals: Globals
+) -> CheckedCFG[Place]:
+    """Type checks a top-level function definition."""
+    args = func_def.args.args
+    returns_none = isinstance(ty.output, NoneType)
+    assert ty.input_names is not None
+
+    cfg = CFGBuilder().build(func_def.body, returns_none, globals)
+    inputs = [
+        Variable(x, inp.ty, loc, inp.flags, is_func_input=True)
+        for x, inp, loc in zip(ty.input_names, ty.inputs, args, strict=True)
+        # Comptime inputs are turned into generic args, so are not included here
+        if InputFlags.Comptime not in inp.flags
+    ]
+    generic_params = {
+        param.name: param.with_idx(i) for i, param in enumerate(ty.params)
+    }
+    return check_cfg(cfg, inputs, ty.output, generic_params, func_def.name, globals)
+
+
+def check_nested_func_def(
+    func_def: NestedFunctionDef, bb: BB, ctx: Context
+) -> CheckedNestedFunctionDef:
+    """Type checks a local (nested) function definition."""
+    func_ty = check_signature(func_def, ctx.globals)
+    assert func_ty.input_names is not None
+
+    if func_ty.parametrized:
+        raise GuppyError(
+            UnsupportedError(func_def, "Nested generic function definitions")
+        )
+
+    # We've already built the CFG for this function while building the CFG of the
+    # enclosing function
+    cfg = func_def.cfg
+
+    # Find captured variables
+    parent_cfg = bb.containing_cfg
+    def_ass_before = set(func_ty.input_names) | ctx.locals.keys()
+    maybe_ass_before = def_ass_before | parent_cfg.maybe_ass_before[bb]
+    inout_vars = inout_var_names(func_ty)
+    cfg.analyze(def_ass_before, maybe_ass_before, inout_vars)
+    captured = {
+        x: (ctx.locals[x], using_bb.vars.used[x])
+        for x, using_bb in cfg.live_before[cfg.entry_bb].items()
+        if x not in func_ty.input_names and x in ctx.locals
+    }
+
+    # Capturing closures are an experimental features since they aren't supported
+    # further down the stack yet
+    if captured:
+        _, loc = captured[next(iter(captured.keys()))]
+        check_capturing_closures_enabled(loc)
+
+    # Captured variables may never be assigned to
+    for bb in cfg.bbs:
+        for v, _ in captured.values():
+            x = v.name
+            if x in bb.vars.assigned:
+                err = IllegalAssignError(bb.vars.assigned[x], x)
+                err.add_sub_diagnostic(IllegalAssignError.DefHint(v.defined_at, x))
+                raise GuppyError(err)
+
+    # Construct inputs for checking the body CFG
+    inputs = [v for v, _ in captured.values()] + [
+        Variable(x, inp.ty, func_def.args.args[i], inp.flags, is_func_input=True)
+        for i, (x, inp) in enumerate(
+            zip(func_ty.input_names, func_ty.inputs, strict=True)
+        )
+        # Comptime inputs are turned into generic args, so are not included here
+        if InputFlags.Comptime not in inp.flags
+    ]
+    def_id = DefId.fresh()
+    globals = ctx.globals
+
+    # Check if the body contains a free (recursive) occurrence of the function name.
+    # By checking if the name is free at the entry BB, we avoid false positives when
+    # a user shadows the name with a local variable
+    if func_def.name in cfg.live_before[cfg.entry_bb]:
+        if not captured:
+            # If there are no captured vars, we treat the function like a global name
+            from guppylang.defs import GuppyDefinition
+            from guppylang_internals.definition.function import ParsedFunctionDef
+
+            func = ParsedFunctionDef(def_id, func_def.name, func_def, func_ty, None)
+            DEF_STORE.register_def(func, None)
+            ENGINE.parsed[def_id] = func
+            globals.f_locals[func_def.name] = GuppyDefinition(func)
+        else:
+            # Otherwise, we treat it like a local name
+            inputs.append(Variable(func_def.name, func_def.ty, func_def))
+
+    checked_cfg = check_cfg(cfg, inputs, func_ty.output, {}, func_def.name, globals)
+    checked_def = CheckedNestedFunctionDef(
+        def_id,
+        checked_cfg,
+        func_ty,
+        captured,
+        name=func_def.name,
+        args=func_def.args,
+        body=func_def.body,
+        decorator_list=func_def.decorator_list,
+        returns=func_def.returns,
+        type_comment=func_def.type_comment,
+    )
+    return with_loc(func_def, checked_def)
+
+
+def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType:
+    """Checks the signature of a function definition and returns the corresponding
+    Guppy type."""
+    if len(func_def.args.posonlyargs) != 0:
+        raise GuppyError(
+            UnsupportedError(func_def.args.posonlyargs[0], "Positional-only parameters")
+        )
+    if len(func_def.args.kwonlyargs) != 0:
+        raise GuppyError(
+            UnsupportedError(func_def.args.kwonlyargs[0], "Keyword-only parameters")
+        )
+    if func_def.args.vararg is not None:
+        raise GuppyError(UnsupportedError(func_def.args.vararg, "Variadic args"))
+    if func_def.args.kwarg is not None:
+        raise GuppyError(UnsupportedError(func_def.args.kwarg, "Keyword args"))
+    if func_def.args.defaults:
+        raise GuppyError(
+            UnsupportedError(func_def.args.defaults[0], "Default arguments")
+        )
+    if func_def.returns is None:
+        err = MissingReturnAnnotationError(func_def)
+        # TODO: Error location is incorrect
+        if all(r.value is None for r in return_nodes_in_ast(func_def)):
+            err.add_sub_diagnostic(
+                MissingReturnAnnotationError.ReturnNone(None, func_def.name)
+            )
+        raise GuppyError(err)
+
+    # Prepopulate parameter mapping when using Python 3.12 style generic syntax
+    param_var_mapping: dict[str, Parameter] = {}
+    if sys.version_info >= (3, 12):
+        for i, param_node in enumerate(func_def.type_params):
+            param = parse_parameter(param_node, i, globals)
+            param_var_mapping[param.name] = param
+
+    input_nodes = []
+    input_names = []
+    for inp in func_def.args.args:
+        ty_ast = inp.annotation
+        if ty_ast is None:
+            raise GuppyError(MissingArgAnnotationError(inp))
+        input_nodes.append(ty_ast)
+        input_names.append(inp.arg)
+    inputs, output = parse_function_io_types(
+        input_nodes,
+        func_def.returns,
+        input_names,
+        func_def,
+        globals,
+        param_var_mapping,
+        True,
+    )
+    return FunctionType(
+        inputs,
+        output,
+        input_names,
+        sorted(param_var_mapping.values(), key=lambda v: v.idx),
+    )
+
+
+def parse_function_with_docstring(
+    func_ast: ast.FunctionDef,
+) -> tuple[ast.FunctionDef, str | None]:
+    """Check if the first line of a function is a docstring.
+
+    If it is, return the function with the docstring removed, plus the docstring.
+    Else, return the original function and `None`
+    """
+    docstring = None
+    match func_ast.body:
+        case [doc, *xs]:
+            if (
+                isinstance(doc, ast.Expr)
+                and isinstance(doc.value, ast.Constant)
+                and isinstance(doc.value.value, str)
+            ):
+                docstring = doc.value.value
+                func_ast.body = xs
+        case _:
+            pass
+    return func_ast, docstring
+
+
+def inout_var_names(func_ty: FunctionType) -> list[str]:
+    """Returns the names of all borrowed arguments in a function type."""
+    assert func_ty.input_names is not None
+    return [
+        x
+        for inp, x in zip(func_ty.inputs, func_ty.input_names, strict=True)
+        if InputFlags.Inout in inp.flags
+    ]