Tweak how we set IR source spans (#8531)

Author: aljazerzen

edb/edgeql/compiler/stmt.py

@@ -30,9 +30,11 @@
 
 from collections import defaultdict
 import textwrap
+import itertools
 
 from edb import errors
 from edb.common import ast
+from edb.common import span as edb_span
 from edb.common.typeutils import not_none
 
 from edb.ir import ast as irast
@@ -1585,6 +1587,15 @@ def compile_query_subject(
     if (shape is not None or view_scls is not None) and len(set.path_id) == 1:
         ctx.class_view_overrides[set.path_id.target.id] = set_stype
 
+    if shape:
+        # make sure that an applied shape expands the span of the set
+        set.span = edb_span.merge_spans(
+            itertools.chain(
+                (s.span for s in [set] if s.span),
+                (el.span for el in shape if el.span)
+            )
+        )
+
     return set
 
 

edb/edgeql/compiler/viewgen.py

@@ -98,6 +98,7 @@ class EarlyShapePtr(NamedTuple):
     ptrcls: s_pointers.Pointer
     target_set: Optional[irast.Set]
     shape_origin: qlast.ShapeOrigin
+    span: Optional[parsing.Span]
 
 
 class ShapePtr(NamedTuple):
@@ -106,6 +107,7 @@ class ShapePtr(NamedTuple):
     ptrcls: s_pointers.Pointer
     shape_op: qlast.ShapeOp
     target_set: Optional[irast.Set]
+    span: Optional[parsing.Span]
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True)
@@ -497,7 +499,8 @@ def make_error(
             )
 
             pointers[pointer] = EarlyShapePtr(
-                pointer, ptr_set, shape_el_desc.ql.origin)
+                pointer, ptr_set, shape_el_desc.ql.origin, shape_el_desc.ql.span
+            )
 
     # If we are not defining a shape (so we might care about
     # materialization), look through our parent view (if one exists)
@@ -534,7 +537,8 @@ def make_error(
             )
 
         pointers[pointer] = EarlyShapePtr(
-            pointer, ptr_set, qlast.ShapeOrigin.MATERIALIZATION)
+            pointer, ptr_set, qlast.ShapeOrigin.MATERIALIZATION, None
+        )
 
     specified_ptrs = {
         ptrcls.get_local_name(ctx.env.schema) for ptrcls in pointers
@@ -571,7 +575,7 @@ def make_error(
     set_shape = []
     shape_ptrs: list[ShapePtr] = []
 
-    for ptrcls, ptr_set, _ in pointers.values():
+    for ptrcls, ptr_set, _, span in pointers.values():
         source: s_types.Type | s_pointers.PointerLike
 
         if ptrcls.is_link_property(ctx.env.schema):
@@ -592,16 +596,15 @@ def make_error(
             continue
 
         ctx.env.view_shapes[source].append((ptrcls, shape_op))
-        shape_ptrs.append(ShapePtr(ir_set, ptrcls, shape_op, ptr_set))
+        shape_ptrs.append(ShapePtr(ir_set, ptrcls, shape_op, ptr_set, span))
 
     rptrcls = view_rptr.ptrcls if view_rptr else None
     shape_ptrs = _get_early_shape_configuration(
         ir_set, shape_ptrs, rptrcls=rptrcls, ctx=ctx)
 
     # Produce the shape. The main thing here is that we need to fixup
     # all of the rptrs to properly point back at ir_set.
-    for _, ptrcls, shape_op, ptr_set in shape_ptrs:
-        ptr_span = None
+    for _, ptrcls, shape_op, ptr_set, ptr_span in shape_ptrs:
         if ptrcls in ctx.env.pointer_specified_info:
             _, _, ptr_span = ctx.env.pointer_specified_info[ptrcls]
 
@@ -636,7 +639,7 @@ def make_error(
                 ir_set,
                 ptrcls,
                 same_computable_scope=True,
-                span=ptr_span or span,
+                span=ptr_span,
                 ctx=ctx,
             )
 
@@ -716,6 +719,7 @@ def _shape_el_ql_to_shape_el_desc(
                 lexpr,
             ],
             partial=True,
+            span=shape_el.span,
         )
     else:
         path_ql = qlast.Path(
@@ -725,6 +729,7 @@ def _shape_el_ql_to_shape_el_desc(
                 qlast.TypeIntersection(type=target_typexpr),
             ],
             partial=True,
+            span=shape_el.span,
         )
 
     return ShapeElementDesc(
@@ -890,7 +895,8 @@ def _gen_pointers_from_defaults(
             )
 
             result.append(EarlyShapePtr(
-                pointer, ptr_set, qlast.ShapeOrigin.DEFAULT))
+                pointer, ptr_set, qlast.ShapeOrigin.DEFAULT, None
+            ))
 
     schema = ctx.env.schema
 
@@ -900,11 +906,11 @@ def _gen_pointers_from_defaults(
     # We cannot check or preprocess this at migration time, because some
     # defaults may not be used for some inserts.
     pointer_indexes = {}
-    for (index, (pointer, _, _)) in enumerate(result):
+    for (index, (pointer, _, _, _)) in enumerate(result):
         p = pointer.get_nearest_non_derived_parent(schema)
         pointer_indexes[p.get_name(schema).name] = index
     graph = {}
-    for (index, (_, irset, _)) in enumerate(result):
+    for (index, (_, irset, _, _)) in enumerate(result):
         assert irset
         dep_pointers = ast.find_children(irset, irast.Pointer)
         dep_rptrs = (
@@ -1249,7 +1255,7 @@ def _compile_rewrites_for_stype(
                 ctx=scopectx,
             )
             res[pn] = EarlyShapePtr(
-                pointer, ptr_set, qlast.ShapeOrigin.DEFAULT
+                pointer, ptr_set, qlast.ShapeOrigin.DEFAULT, None
             )
     return res
 
@@ -2254,7 +2260,8 @@ def _inline_type_computable(
             ctx.env.pointer_specified_info[ptr] = (None, None, None)
         view_shape.insert(0, (ptr, qlast.ShapeOp.ASSIGN))
         shape_ptrs.insert(
-            0, ShapePtr(ir_set, ptr, qlast.ShapeOp.ASSIGN, ptr_set))
+            0, ShapePtr(ir_set, ptr, qlast.ShapeOp.ASSIGN, ptr_set, None)
+        )
 
 
 def _get_shape_configuration_inner(
@@ -2267,7 +2274,8 @@ def _get_shape_configuration_inner(
 ) -> None:
     is_objtype = ir_set.path_id.is_objtype_path()
     all_materialize = all(
-        op == qlast.ShapeOp.MATERIALIZE for _, _, op, _ in shape_ptrs)
+        op == qlast.ShapeOp.MATERIALIZE for _, _, op, _, _ in shape_ptrs
+    )
 
     if is_objtype:
         assert isinstance(stype, s_objtypes.ObjectType)
@@ -2306,7 +2314,8 @@ def _get_shape_configuration_inner(
                     view_shape.insert(0, (ptr, implicit_op))
                     shape_metadata.has_implicit_id = True
                     shape_ptrs.insert(
-                        0, ShapePtr(ir_set, ptr, implicit_op, None))
+                        0, ShapePtr(ir_set, ptr, implicit_op, None, None)
+                    )
                 break
 
     is_mutation = parent_view_type in {
@@ -2409,7 +2418,8 @@ def _get_late_shape_configuration(
 
     for source in sources:
         for ptr, shape_op in ctx.env.view_shapes[source]:
-            shape_ptrs.append(ShapePtr(ir_set, ptr, shape_op, None))
+            shape_ptrs.append(ShapePtr(ir_set, ptr, shape_op, None, None)
+        )
 
     _get_shape_configuration_inner(
         ir_set, shape_ptrs, stype, parent_view_type=parent_view_type, ctx=ctx)
@@ -2506,16 +2516,16 @@ def _late_compile_view_shapes_in_set(
             return
 
         shape = []
-        for path_tip, ptr, shape_op, _ in shape_ptrs:
-            span = None
+        for path_tip, ptr, shape_op, _, ptr_span in shape_ptrs:
+            ptr_span = None
             if ptr in ctx.env.pointer_specified_info:
-                _, _, span = ctx.env.pointer_specified_info[ptr]
+                _, _, ptr_span = ctx.env.pointer_specified_info[ptr]
 
             element = setgen.extend_path(
                 path_tip,
                 ptr,
                 same_computable_scope=True,
-                span=span,
+                span=ptr_span,
                 ctx=ctx,
             )
 

tests/test_edgeql_explain.py

@@ -72,7 +72,12 @@ async def test_edgeql_explain_simple_01(self):
         ''')
         self.assert_plan(res['fine_grained'], {
             "contexts": [
-                {"buffer_idx": 0, "end": 32, "start": 28, "text": "User"}
+                {
+                    "buffer_idx": 0,
+                    "start": 28,
+                    "end": 43,
+                    "text": "User { id, name"
+                }
             ],
             "pipeline": [
                 {
@@ -356,9 +361,9 @@ async def test_edgeql_explain_multi_link_01(self):
             "contexts": [
                 {
                     "start": 28,
-                    "end": 32,
+                    "end": 60,
                     "buffer_idx": 0,
-                    "text": "User",
+                    "text": "User { name, todo: {name, number",
                 },
             ],
             "pipeline": [
@@ -413,9 +418,9 @@ async def test_edgeql_explain_multi_link_01(self):
                     "contexts": [
                         {
                             "start": 41,
-                            "end": 45,
+                            "end": 60,
                             "buffer_idx": 0,
-                            "text": "todo",
+                            "text": "todo: {name, number",
                         },
                     ],
                     "pipeline": tb.bag([
@@ -606,9 +611,9 @@ async def test_edgeql_explain_computed_backlink_01(self):
         shape = {
             "contexts": [{
                 "start": 28,
-                "end": 32,
+                "end": 68,
                 "buffer_idx": 0,
-                "text": "User"
+                "text": "User { name, owned_issues: {name, number"
             }],
             "pipeline": [{
                 "plan_rows": 1,
@@ -663,9 +668,9 @@ async def test_edgeql_explain_computed_backlink_01(self):
                     "text": '.<owner[is default::Issue]'
                 }, {
                     "start": 41,
-                    "end": 53,
+                    "end": 68,
                     "buffer_idx": 0,
-                    "text": "owned_issues"
+                    "text": "owned_issues: {name, number"
                 }]),
                 "pipeline": tb.bag([
                     {
@@ -1020,9 +1025,13 @@ async def test_edgeql_explain_type_intersect_01(self):
                     "contexts": [
                         {
                             "start": 28,
-                            "end": 32,
+                            "end": 93,
                             "buffer_idx": 0,
-                            "text": 'Text',
+                            'text': (
+                                'Text {\n'
+                                '                body,\n'
+                                '                z := [is Issue].name'
+                            )
                         },
                     ],
                     "pipeline": [