feat: Update CFG checker to use new diagnostics

examples/demo.ipynb

@@ -203,13 +203,17 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Guppy compilation failed. Error in file <In [7]>:7\n",
+      "Error: Variable not defined (at <In [7]>:7:11)\n",
+      "  | \n",
+      "5 |     if b:\n",
+      "6 |         x = 4\n",
+      "7 |     return x  # x not defined if b is False\n",
+      "  |            ^ `x` might be undefined ...\n",
+      "  | \n",
+      "5 |     if b:\n",
+      "  |        - ... if this expression is `False`\n",
       "\n",
-      "5:    if b:\n",
-      "6:        x = 4\n",
-      "7:    return x  # x not defined if b is False\n",
-      "             ^\n",
-      "GuppyError: Variable `x` is not defined on all control-flow paths.\n"
+      "Guppy compilation failed due to 1 previous error\n"
      ]
     }
    ],
@@ -247,13 +251,20 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Guppy compilation failed. Error in file <In [8]>:9\n",
+      "Error: Different types (at <In [8]>:9:15)\n",
+      "  | \n",
+      "7 |     else:\n",
+      "8 |         x = True\n",
+      "9 |     return int(x)  # x has different types depending on b\n",
+      "  |                ^ Variable `x` may refer to different types\n",
+      "  | \n",
+      "6 |         x = 4\n",
+      "  |         - This is of type `int`\n",
+      "  | \n",
+      "8 |         x = True\n",
+      "  |         - This is of type `bool`\n",
       "\n",
-      "7:    else:\n",
-      "8:        x = True\n",
-      "9:    return int(x)  # x has different types depending on b\n",
-      "                 ^\n",
-      "GuppyError: Variable `x` can refer to different types: `int` (at 6:8) vs `bool` (at 8:8)\n"
+      "Guppy compilation failed due to 1 previous error\n"
      ]
     }
    ],

guppylang/cfg/builder.py

@@ -177,7 +177,7 @@ def visit_For(self, node: ast.For, bb: BB, jumps: Jumps) -> BB | None:
         b = make_var(next(tmp_vars), node.iter)
         new_nodes = template_replace(
             template,
-            node,
+            node.iter,
             it=it,
             b=b,
             x=node.target,

guppylang/cfg/cfg.py

@@ -1,3 +1,5 @@
+from collections import deque
+from collections.abc import Iterator
 from typing import Generic, TypeVar
 
 from guppylang.cfg.analysis import (
@@ -37,6 +39,18 @@ def __init__(
         self.ass_before = {}
         self.maybe_ass_before = {}
 
+    def ancestors(self, *bbs: T) -> Iterator[T]:
+        """Returns an iterator over all ancestors of the given BBs in BFS order."""
+        queue = deque(bbs)
+        visited = set()
+        while queue:
+            bb = queue.popleft()
+            if bb in visited:
+                continue
+            visited.add(bb)
+            yield bb
+            queue += bb.predecessors
+
 
 class CFG(BaseCFG[BB]):
     """A control-flow graph of unchecked basic blocks."""

guppylang/checker/cfg_checker.py

@@ -4,17 +4,19 @@
 `CheckedBB`s with inferred type signatures.
 """
 
+import ast
 import collections
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass, field
-from typing import Generic, TypeVar
+from typing import ClassVar, Generic, TypeVar
 
 from guppylang.ast_util import line_col
 from guppylang.cfg.bb import BB
 from guppylang.cfg.cfg import CFG, BaseCFG
 from guppylang.checker.core import Context, Globals, Locals, Place, V, Variable
 from guppylang.checker.expr_checker import ExprSynthesizer, to_bool
 from guppylang.checker.stmt_checker import StmtChecker
+from guppylang.diagnostic import Error, Note
 from guppylang.error import GuppyError
 from guppylang.tys.ty import InputFlags, Type
 
@@ -129,6 +131,44 @@ def check_cfg(
     return linearity_checked_cfg
 
 
+@dataclass(frozen=True)
+class VarNotDefinedError(Error):
+    title: ClassVar[str] = "Variable not defined"
+    span_label: ClassVar[str] = "`{var}` is not defined"
+    var: str
+
+
+@dataclass(frozen=True)
+class VarMaybeNotDefinedError(Error):
+    title: ClassVar[str] = "Variable not defined"
+    var: str
+
+    @dataclass(frozen=True)
+    class BadBranch(Note):
+        span_label: ClassVar[str] = "... if this expression is `{truth_value}`"
+        var: str
+        truth_value: bool
+
+    @property
+    def rendered_span_label(self) -> str:
+        s = f"`{self.var}` might be undefined"
+        if self.children:
+            s += " ..."
+        return s
+
+
+@dataclass(frozen=True)
+class BranchTypeError(Error):
+    title: ClassVar[str] = "Different types"
+    span_label: ClassVar[str] = "{ident} may refer to different types"
+    ident: str
+
+    @dataclass(frozen=True)
+    class TypeHint(Note):
+        span_label: ClassVar[str] = "This is of type `{ty}`"
+        ty: Type
+
+
 def check_bb(
     bb: BB,
     checked_cfg: CheckedCFG[Variable],
@@ -144,7 +184,7 @@ def check_bb(
         assert len(bb.predecessors) == 0
         for x, use in bb.vars.used.items():
             if x not in cfg.ass_before[bb] and x not in globals:
-                raise GuppyError(f"Variable `{x}` is not defined", use)
+                raise GuppyError(VarNotDefinedError(use, x))
 
     # Check the basic block
     ctx = Context(globals, Locals({v.name: v for v in inputs}))
@@ -163,14 +203,15 @@ def check_bb(
                 # If the variable is defined on *some* paths, we can give a more
                 # informative error message
                 if x in cfg.maybe_ass_before[use_bb]:
-                    # TODO: This should be "Variable x is not defined when coming
-                    #  from {bb}". But for this we need a way to associate BBs with
-                    #  source locations.
-                    raise GuppyError(
-                        f"Variable `{x}` is not defined on all control-flow paths.",
-                        use_bb.vars.used[x],
-                    )
-                raise GuppyError(f"Variable `{x}` is not defined", use_bb.vars.used[x])
+                    err = VarMaybeNotDefinedError(use_bb.vars.used[x], x)
+                    if bad_branch := diagnose_maybe_undefined(use_bb, x, cfg):
+                        branch_expr, truth_value = bad_branch
+                        note = VarMaybeNotDefinedError.BadBranch(
+                            branch_expr, x, truth_value
+                        )
+                        err.add_sub_diagnostic(note)
+                    raise GuppyError(err)
+                raise GuppyError(VarNotDefinedError(use_bb.vars.used[x], x))
 
     # Finally, we need to compute the signature of the basic block
     outputs = [
@@ -209,12 +250,39 @@ def check_rows_match(row1: Row[Variable], row2: Row[Variable], bb: BB) -> None:
             # We shouldn't mention temporary variables (starting with `%`)
             # in error messages:
             ident = "Expression" if v1.name.startswith("%") else f"Variable `{v1.name}`"
-            raise GuppyError(
-                f"{ident} can refer to different types: "
-                f"`{v1.ty}` (at {{}}) vs `{v2.ty}` (at {{}})",
-                bb.containing_cfg.live_before[bb][v1.name].vars.used[v1.name],
-                [v1.defined_at, v2.defined_at],
-            )
+            use = bb.containing_cfg.live_before[bb][v1.name].vars.used[v1.name]
+            err = BranchTypeError(use, ident)
+            err.add_sub_diagnostic(BranchTypeError.TypeHint(v1.defined_at, v1.ty))
+            err.add_sub_diagnostic(BranchTypeError.TypeHint(v2.defined_at, v2.ty))
+            raise GuppyError(err)
+
+
+def diagnose_maybe_undefined(
+    bb: BB, x: str, cfg: BaseCFG[BB]
+) -> tuple[ast.expr, bool] | None:
+    """Given a BB and a variable `x`, tries to find a branch where one of the successors
+    leads to an assignment of `x` while the other one does not.
+
+    Returns the branch condition and a flag whether the value being `True` leads to the
+    undefined path. Returns `None` if no such branch can be found.
+    """
+    assert x in cfg.maybe_ass_before[bb]
+    # Find all BBs that can reach this BB and which ones of those assign `x`
+    ancestors = list(cfg.ancestors(bb))
+    assigns = [anc for anc in ancestors if x in anc.vars.assigned]
+    # Compute which ancestors can possibly reach an assignment
+    reaches_assignment = set(cfg.ancestors(*assigns))
+    # Try to find a branching BB where one of paths can reach an assignment, while the
+    # other one cannot
+    for anc in ancestors:
+        match anc.successors:
+            case [true_succ, false_succ]:
+                assert anc.branch_pred is not None
+                true_reaches_assignment = true_succ in reaches_assignment
+                false_reaches_assignment = false_succ in reaches_assignment
+                if true_reaches_assignment != false_reaches_assignment:
+                    return anc.branch_pred, true_reaches_assignment
+    return None
 
 
 T = TypeVar("T")

guppylang/definition/function.py

@@ -241,7 +241,13 @@ def parse_py_func(f: PyFunc, sources: SourceMap) -> tuple[ast.FunctionDef, str |
             defn = find_ipython_def(func_ast.name)
             if defn is not None:
                 file = f"<{defn.cell_name}>"
-        sources.add_file(file, source)
+                sources.add_file(file, defn.cell_source)
+            else:
+                # If we couldn't find the defining cell, just use the source code we
+                # got from inspect. Line numbers will be wrong, but that's the best we
+                # can do.
+                sources.add_file(file, source)
+                line_offset = 1
     else:
         file = inspect.getsourcefile(f)
         if file is None:

guppylang/diagnostic.py

@@ -208,10 +208,16 @@ def render_diagnostic(self, diag: Diagnostic) -> None:
         else:
             span = to_span(diag.span)
             level = self.level_str(diag.level)
+            all_spans = [
+                span,
+                *(to_span(child.span) for child in diag.children if child.span),
+            ]
+            max_lineno = max(s.end.line for s in all_spans)
             self.buffer.append(f"{level}: {diag.rendered_title} (at {span.start})")
             self.render_snippet(
                 span,
                 diag.rendered_span_label,
+                max_lineno,
                 is_primary=True,
                 prefix_lines=self.PREFIX_CONTEXT_LINES,
             )
@@ -221,6 +227,7 @@ def render_diagnostic(self, diag: Diagnostic) -> None:
                     self.render_snippet(
                         to_span(sub_diag.span),
                         sub_diag.rendered_span_label,
+                        max_lineno,
                         is_primary=False,
                     )
             if diag.message:
@@ -238,7 +245,12 @@ def render_diagnostic(self, diag: Diagnostic) -> None:
                 )
 
     def render_snippet(
-        self, span: Span, label: str | None, is_primary: bool, prefix_lines: int = 0
+        self,
+        span: Span,
+        label: str | None,
+        max_lineno: int,
+        is_primary: bool,
+        prefix_lines: int = 0,
     ) -> None:
         """Renders the source associated with a span together with an optional label.
 
@@ -272,7 +284,7 @@ def render_snippet(
         additional context.
         """
         # Check how much space we need to reserve for the leading line numbers
-        ll_length = len(str(span.end.line))
+        ll_length = len(str(max_lineno))
         highlight_char = "^" if is_primary else "-"
 
         def render_line(line: str, line_number: int | None = None) -> None:

tests/diagnostics/snapshots/test_two_spans_different_lineno_lens.txt

@@ -0,0 +1,9 @@
+Error: Can't compare apples with oranges (at <unknown>:99:6)
+    | 
+ 97 | apple == orange
+ 98 | apple == orange
+ 99 | apple == orange
+    |       ^^ Comparison attempted here
+    | 
+100 | apple == orange
+    | ----- This is an apple

tests/diagnostics/test_diagnostics_rendering.py

@@ -146,6 +146,19 @@ def test_justify_line_number(snapshot, request):
     run_test(source, diagnostic, snapshot, request)
 
 
+def test_two_spans_different_lineno_lens(snapshot, request):
+    @dataclass(frozen=True)
+    class MySubDiagnostic(Note):
+        span_label: ClassVar[str] = "This is an apple"
+
+    source = 100 * "apple == orange\n"
+    span1 = Span(Loc(file, 99, 6), Loc(file, 99, 8))
+    span2 = Span(Loc(file, 100, 0), Loc(file, 100, 5))
+    diagnostic = MyError(span1)
+    diagnostic.add_sub_diagnostic(MySubDiagnostic(span2))
+    run_test(source, diagnostic, snapshot, request)
+
+
 def test_indented(snapshot, request):
     source = " " * 50 + "super_apple := apple ** 2\n"
     source += " " * 50 + "    lemon := orange - apple\n"

tests/error/errors_on_usage/and_not_defined.err

@@ -1,7 +1,11 @@
-Guppy compilation failed. Error in file $FILE:9
+Error: Variable not defined (at $FILE:9:15)
+  | 
+7 |         return z
+8 |     else:
+9 |         return z
+  |                ^ `z` might be undefined ...
+  | 
+6 |     if x and (z := y + 1):
+  |        - ... if this expression is `False`
 
-7:        return z
-8:    else:
-9:        return z
-                 ^
-GuppyError: Variable `z` is not defined on all control-flow paths.
+Guppy compilation failed due to 1 previous error

tests/error/errors_on_usage/else_expr_not_defined.err

@@ -1,7 +1,11 @@
-Guppy compilation failed. Error in file $FILE:7
+Error: Variable not defined (at $FILE:7:11)
+  | 
+5 | def foo(x: bool) -> int:
+6 |     (y := 1) if x else (z := 2)
+7 |     return z
+  |            ^ `z` might be undefined ...
+  | 
+6 |     (y := 1) if x else (z := 2)
+  |                 - ... if this expression is `True`
 
-5:    def foo(x: bool) -> int:
-6:        (y := 1) if x else (z := 2)
-7:        return z
-                 ^
-GuppyError: Variable `z` is not defined on all control-flow paths.
+Guppy compilation failed due to 1 previous error

tests/error/errors_on_usage/else_expr_type_change.err

@@ -1,7 +1,14 @@
-Guppy compilation failed. Error in file $FILE:8
+Error: Different types (at $FILE:8:11)
+  | 
+6 |     y = 3
+7 |     (y := y + 1) if x else (y := True)
+8 |     return y
+  |            ^ Variable `y` may refer to different types
+  | 
+7 |     (y := y + 1) if x else (y := True)
+  |      - This is of type `int`
+  | 
+7 |     (y := y + 1) if x else (y := True)
+  |                             - This is of type `bool`
 
-6:    y = 3
-7:    (y := y + 1) if x else (y := True)
-8:    return y
-             ^
-GuppyError: Variable `y` can refer to different types: `int` (at 7:5) vs `bool` (at 7:28)
+Guppy compilation failed due to 1 previous error

tests/error/errors_on_usage/else_not_defined.err

@@ -1,7 +1,11 @@
-Guppy compilation failed. Error in file $FILE:10
+Error: Variable not defined (at $FILE:10:11)
+   | 
+ 8 |     else:
+ 9 |         z = 2
+10 |     return z
+   |            ^ `z` might be undefined ...
+   | 
+ 6 |     if x:
+   |        - ... if this expression is `True`
 
-8:     else:
-9:         z = 2
-10:    return z
-              ^
-GuppyError: Variable `z` is not defined on all control-flow paths.
+Guppy compilation failed due to 1 previous error