Repository: geldata/gel
Pull request: Replace weak namespace stripping with smarter rewriting
Instead of "stripping" weak namespaces from all paths, instead rewrite
weak namespaces to be "absolute" by removing namespaces that don't
appear at the binding site for the path.

Fixes #1381.

Problem statement:
Incorrect result for a value computed in WITH block
Run the following in the existing scope tests dataset (from `test_edgeql_scope.py`):
```
WITH 
    MODULE test,
    avg := math::mean({len(Card.name)})
SELECT Card {name, a := avg};
```
Instead of being the same value across all results `a` is different for each as if the function call was inlined.

Hints from issue discussion:
I can even simplify it further to:
```
WITH 
    MODULE test,
    avg := count({Card.name})
SELECT Card {name, a := avg};
```
Or even
```
WITH 
    MODULE test,
    avg := Card.name
SELECT Card {name, a := avg};
```
@elprans can we fix this for alpha 3?
(if not move it to Alpha 4)
The fix isn't trivial, so Alpha 4.
Also: `WITH X := {1, 2}, Y := X+1 SELECT (X, Y)`.
Should produce a cross-product output with 4 results, actually produces 2.

Commit message: Replace weak namespace stripping with smarter rewriting

Instead of "stripping" weak namespaces from all paths, instead rewrite
weak namespaces to be "absolute" by removing namespaces that don't
appear at the binding site for the path.

Fixes #1381.

Patch to review:
diff --git a/edb/common/ast/visitor.py b/edb/common/ast/visitor.py
index aaa61c6bcd2..830caf758d5 100644
--- a/edb/common/ast/visitor.py
+++ b/edb/common/ast/visitor.py
@@ -170,7 +170,7 @@ def generic_visit(self, node, *, combine_results=None):
         for _field, value in base.iter_fields(node, include_meta=False):
             if typeutils.is_container(value):
                 for item in value:
-                    if base.is_ast_node(item):
+                    if base.is_ast_node(item) or typeutils.is_container(item):
                         res = self.visit(item)
                         if res is not None:
                             field_results.append(res)
diff --git a/edb/edgeql/compiler/func.py b/edb/edgeql/compiler/func.py
index 8249c19d0bb..62a9407c1e7 100644
--- a/edb/edgeql/compiler/func.py
+++ b/edb/edgeql/compiler/func.py
@@ -229,12 +229,12 @@ def compile_FunctionCall(
         nested_path_ids = []
         for n, st in rtype.iter_subtypes(ctx.env.schema):
             elem_path_id = pathctx.get_tuple_indirection_path_id(
-                path_id, n, st, ctx=ctx).strip_weak_namespaces()
+                path_id, n, st, ctx=ctx)
 
             if isinstance(st, s_types.Tuple):
                 nested_path_ids.append([
                     pathctx.get_tuple_indirection_path_id(
-                        elem_path_id, nn, sst, ctx=ctx).strip_weak_namespaces()
+                        elem_path_id, nn, sst, ctx=ctx)
                     for nn, sst in st.iter_subtypes(ctx.env.schema)
                 ])
 
diff --git a/edb/edgeql/compiler/inference/cardinality.py b/edb/edgeql/compiler/inference/cardinality.py
index fb8856b9cc7..b46ffecd97b 100644
--- a/edb/edgeql/compiler/inference/cardinality.py
+++ b/edb/edgeql/compiler/inference/cardinality.py
@@ -206,7 +206,7 @@ def _check_binding_volatility(
         and ir.expr
         and volatility.infer_volatility(ir.expr, env=ctx.env) == VOLATILE
     ):
-        path_id = ir.path_id.strip_weak_namespaces()
+        path_id = ir.path_id
         if path_id not in ctx.bindings:
             return
 
diff --git a/edb/edgeql/compiler/inference/types.py b/edb/edgeql/compiler/inference/types.py
index f04a4dcfa53..abe2acb572b 100644
--- a/edb/edgeql/compiler/inference/types.py
+++ b/edb/edgeql/compiler/inference/types.py
@@ -48,7 +48,8 @@ def amend_empty_set_type(
     alias = es.path_id.target_name_hint.name
     typename = s_name.QualName(module='__derived__', name=alias)
     es.path_id = irast.PathId.from_type(
-        env.schema, t, env=env, typename=typename
+        env.schema, t, env=env, typename=typename,
+        namespace=es.path_id.namespace,
     )
 
 
diff --git a/edb/edgeql/compiler/pathctx.py b/edb/edgeql/compiler/pathctx.py
index 6d14be8241e..6ff109f90d5 100644
--- a/edb/edgeql/compiler/pathctx.py
+++ b/edb/edgeql/compiler/pathctx.py
@@ -168,12 +168,11 @@ def ban_path(
         path_id: irast.PathId, *,
         ctx: context.ContextLevel) -> None:
 
-    ctx.banned_paths.add(path_id.strip_weak_namespaces())
+    ctx.banned_paths.add(path_id)
 
 
 def path_is_banned(
         path_id: irast.PathId, *,
         ctx: context.ContextLevel) -> bool:
 
-    s_path_id = path_id.strip_weak_namespaces()
-    return s_path_id in ctx.banned_paths and ctx.path_scope.is_visible(path_id)
+    return path_id in ctx.banned_paths and ctx.path_scope.is_visible(path_id)
diff --git a/edb/edgeql/compiler/setgen.py b/edb/edgeql/compiler/setgen.py
index 5cfcb5fa616..7580b19b487 100644
--- a/edb/edgeql/compiler/setgen.py
+++ b/edb/edgeql/compiler/setgen.py
@@ -184,7 +184,7 @@ def new_tuple_set(
     for elem in elements:
         elem_path_id = pathctx.get_tuple_indirection_path_id(
             result_path_id, elem.name, get_set_type(elem.val, ctx=ctx),
-            ctx=ctx).strip_weak_namespaces()
+            ctx=ctx)
         final_elems.append(irast.TupleElement(
             name=elem.name,
             val=elem.val,
diff --git a/edb/edgeql/compiler/stmt.py b/edb/edgeql/compiler/stmt.py
index db0c2a5548f..c865c655edc 100644
--- a/edb/edgeql/compiler/stmt.py
+++ b/edb/edgeql/compiler/stmt.py
@@ -337,11 +337,12 @@ def compile_insert_unless_conflict(
         else_ir = dispatch.compile(
             astutils.ensure_qlstmt(else_branch), ctx=ctx)
         assert isinstance(else_ir, irast.Set)
-        else_info = irast.OnConflictElse(select_ir, else_ir)
+        else_info = irast.OnConflictElse(select=select_ir, body=else_ir)
 
     return irast.OnConflictClause(
-        irast.ConstraintRef(id=ex_cnstrs[0].id, module_id=module_id),
-        else_info
+        constraint=irast.ConstraintRef(
+            id=ex_cnstrs[0].id, module_id=module_id),
+        else_ir=else_info
     )
 
 
diff --git a/edb/edgeql/compiler/stmtctx.py b/edb/edgeql/compiler/stmtctx.py
index 1fb41847b84..baab65cda57 100644
--- a/edb/edgeql/compiler/stmtctx.py
+++ b/edb/edgeql/compiler/stmtctx.py
@@ -41,6 +41,8 @@
 from edb.edgeql import qltypes
 from edb.edgeql import parser as qlparser
 
+from edb.common.ast import visitor as ast_visitor
+
 from . import astutils
 from . import context
 from . import dispatch
@@ -133,16 +135,10 @@ def fini_expression(
                 ctx=inf_ctx,
             )
 
-    # Strip weak namespaces
-    for ir_set in ctx.env.set_types:
-        if ir_set.path_id.namespace:
-            ir_set.path_id = ir_set.path_id.strip_weak_namespaces()
+    # Fix up weak namespaces
+    _rewrite_weak_namespaces(ir, ctx)
 
     if ctx.path_scope is not None:
-        for node in ctx.path_scope.path_descendants:
-            if node.path_id.namespace:
-                node.path_id = node.path_id.strip_weak_namespaces()
-
         ctx.path_scope.validate_unique_ids()
 
     if isinstance(ir, irast.Command):
@@ -254,10 +250,106 @@ def fini_expression(
     return result
 
 
+class FindPathScopes(ast_visitor.NodeVisitor):
+    """Visitor to find the enclosing path scope id of sub expressions.
+
+    Sets inherit an effective scope id from enclosing expressions,
+    and this visitor computes those.
+    """
+    def __init__(self, check: bool) -> None:
+        super().__init__()
+        self.path_scope_ids: List[Optional[int]] = [None]
+
+    def visit_Stmt(self, stmt: irast.Stmt) -> Optional[int]:
+        # Sometimes there is sharing, so we want the official scope
+        # for a node to be based on its appearance in the result,
+        # not in a subquery.
+        # I think it might not actually matter, though.
+        self.visit(stmt.bindings)
+        if stmt.iterator_stmt:
+            self.visit(stmt.iterator_stmt)
+        if isinstance(stmt, irast.MutatingStmt):
+            self.visit(stmt.subject)
+        self.visit(stmt.result)
+
+        self.generic_visit(stmt)
+        return None
+
+    def visit_Set(self, node: irast.Set) -> Optional[int]:
+        val = self.path_scope_ids[-1]
+        if node.path_scope_id:
+            self.path_scope_ids.append(node.path_scope_id)
+        if not node.is_binding:
+            val = self.path_scope_ids[-1]
+
+        # Visit sub-trees
+        self.generic_visit(node)
+
+        if node.path_scope_id:
+            self.path_scope_ids.pop()
+
+        return val
+
+
+def _try_namespace_fix(
+    scope: irast.ScopeTreeNode,
+    obj: Union[irast.ScopeTreeNode, irast.Set],
+) -> None:
+    if obj.path_id is None:
+        return
+    for prefix in obj.path_id.iter_prefixes():
+        replacement = scope.find_visible(prefix)
+        if (
+            replacement and replacement.path_id
+            and replacement.path_id != prefix
+        ):
+            new = obj.path_id.replace_prefix(prefix, replacement.path_id)
+            obj.path_id = new
+            break
+
+
+def _rewrite_weak_namespaces(
+    ir: irast.Base, ctx: context.ContextLevel
+) -> None:
+    """Rewrite weak namespaces in path ids to be usable by the backend.
+
+    Weak namespaces in path ids in the frontend are "relative", and
+    their interpretation depends on the current scope tree node and
+    the namespaces on the parent nodes. The IR->pgsql compiler does
+    not do this sort of interpretation, and needs path IDs that are
+    "absolute".
+
+    To accomplish this, we go through all the path ids and rewrite
+    them: using the scope tree, we try to find the binding of the path
+    ID (using a prefix if necessary) and drop all namespace parts that
+    don't appear in the binding.
+    """
+
+    if ctx.path_scope is None:
+        return None
+
+    tree = ctx.path_scope
+
+    for node in tree.strict_descendants:
+        _try_namespace_fix(node, node)
+
+    visitor = FindPathScopes(ctx.env.options.apply_query_rewrites)
+    visitor.visit(ir)
+
+    for ir_set in ctx.env.set_types:
+        path_scope_id: Optional[int] = visitor.memo.get(ir_set)
+        if path_scope_id is not None:
+            # Some entries in set_types are from compiling views
+            # in temporary scopes, so we need to just skip those.
+            if scope := tree.find_by_unique_id(path_scope_id):
+                _try_namespace_fix(scope, ir_set)
+
+
 def _elide_derived_ancestors(
     obj: Union[s_types.InheritingType, s_pointers.Pointer], *,
     ctx: context.ContextLevel
 ) -> None:
+
     """Collapse references to derived objects in bases.
 
     When compiling a schema view expression, make sure we don't
diff --git a/edb/edgeql/compiler/viewgen.py b/edb/edgeql/compiler/viewgen.py
index 91d59bf0480..72e81fdad21 100644
--- a/edb/edgeql/compiler/viewgen.py
+++ b/edb/edgeql/compiler/viewgen.py
@@ -447,6 +447,7 @@ def _normalize_view_ptr_expr(
         )
 
         base_cardinality = _get_base_ptr_cardinality(base_ptrcls, ctx=ctx)
+        base_is_singleton = False
         if base_cardinality is not None and base_cardinality.is_known():
             base_is_singleton = base_cardinality.is_single()
 
diff --git a/edb/ir/ast.py b/edb/ir/ast.py
index dbb3ebaa543..c6822b69ea6 100644
--- a/edb/ir/ast.py
+++ b/edb/ir/ast.py
@@ -745,12 +745,12 @@ class MutatingStmt(Stmt):
     subject: Set
 
 
-class OnConflictElse(typing.NamedTuple):
+class OnConflictElse(Base):
     select: Set
     body: Set
 
 
-class OnConflictClause(typing.NamedTuple):
+class OnConflictClause(Base):
     constraint: typing.Optional[ConstraintRef]
     else_ir: typing.Optional[OnConflictElse]
 
diff --git a/edb/ir/pathid.py b/edb/ir/pathid.py
index b559e0cd5c0..6032ef61eaf 100644
--- a/edb/ir/pathid.py
+++ b/edb/ir/pathid.py
@@ -383,23 +383,6 @@ def merge_namespace(
         else:
             return self
 
-    def strip_weak_namespaces(self) -> PathId:
-        """Return a copy of this ``PathId`` with weak namespace portion
-           removed."""
-        if self._namespace is not None:
-            stripped_ns = {bit for bit in self._namespace
-                           if not isinstance(bit, WeakNamespace)}
-            result = self.replace_namespace(stripped_ns)
-
-            if result._prefix is not None:
-                result._prefix = result._get_minimal_prefix(
-                    result._prefix.strip_weak_namespaces())
-
-        else:
-            result = self
-
-        return result
-
     def strip_namespace(self, namespace: AbstractSet[AnyNamespace]) -> PathId:
         """Return a copy of this ``PathId`` with a given portion of the
            namespace id removed."""
@@ -644,7 +627,7 @@ def replace_prefix(self, prefix: PathId, replacement: PathId) -> PathId:
         if self.startswith(prefix):
             prefix_len = len(prefix)
             if prefix_len < len(self):
-                result = self.__class__()
+                result = self.__class__(self)
                 result._path = replacement._path + self._path[prefix_len:]
                 result._norm_path = \
                     replacement._norm_path + self._norm_path[prefix_len:]
diff --git a/edb/pgsql/compiler/dml.py b/edb/pgsql/compiler/dml.py
index 5b59d126478..c29f93200cd 100644
--- a/edb/pgsql/compiler/dml.py
+++ b/edb/pgsql/compiler/dml.py
@@ -133,7 +133,7 @@ def init_dml_stmt(
     else_cte = None
     if (
         isinstance(ir_stmt, irast.InsertStmt)
-        and ir_stmt.on_conflict and ir_stmt.on_conflict[1] is not None
+        and ir_stmt.on_conflict and ir_stmt.on_conflict.else_ir is not None
     ):
         dml_cte = pgast.CommonTableExpr(
             query=pgast.SelectStmt(),
@@ -660,7 +660,8 @@ def compile_insert_else_body(
     )
 
     if on_conflict.else_ir:
-        else_select, else_branch = on_conflict.else_ir
+        else_select = on_conflict.else_ir.select
+        else_branch = on_conflict.else_ir.body
 
         subject_id = ir_stmt.subject.path_id
 
diff --git a/tests/test_edgeql_scope.py b/tests/test_edgeql_scope.py
index 17e572f9437..da1fa9475f2 100644
--- a/tests/test_edgeql_scope.py
+++ b/tests/test_edgeql_scope.py
@@ -2082,6 +2082,33 @@ async def test_edgeql_scope_with_01(self):
             {'Fire', 'Water', 'Earth', 'Air'},
         )
 
+    async def test_edgeql_scope_with_02(self):
+        # Test a WITH binding that depends on a previous one is still
+        # independent
+        await self.assert_query_result(
+            r"""
+                WITH
+                    MODULE test,
+                    X := {1, 2},
+                    Y := X + 1,
+                SELECT _ := (X, Y) ORDER BY _;
+            """,
+            [[1, 2], [1, 3], [2, 2], [2, 3]]
+        )
+
+    async def test_edgeql_scope_with_03(self):
+        # Test that a WITH binding used in a computable doesn't have its
+        # reference to that type captured
+        await self.assert_query_result(
+            r"""
+                WITH
+                    MODULE test,
+                    a := count({Card.name})
+                SELECT Card {name, a := a} FILTER .name = 'Imp';
+            """,
+            [{"name": "Imp", "a": 9}],
+        )
+
     async def test_edgeql_scope_unused_with_def_01(self):
 
         with self.assertRaisesRegex(
diff --git a/tests/test_edgeql_tree.py b/tests/test_edgeql_tree.py
index a1ef99f7af1..a0bf9c743c6 100644
--- a/tests/test_edgeql_tree.py
+++ b/tests/test_edgeql_tree.py
@@ -784,11 +784,6 @@ async def test_edgeql_tree_update_04(self):
             ]
         )
 
-    @test.xfail('''
-        Fails with a SQL reference to a nonexistant alias.
-        Introduced in #2137, but I think that the issue is #1381
-        (weak namespace stripping being bogus).
-    ''')
     async def test_edgeql_tree_update_05(self):
         # Swap around a tree node and its first child as an atomic operation.
         #


Code review comments:
diff_hunk: @@ -254,10 +250,106 @@ def fini_expression(
     return result
 
 
+class FindPathScopes(ast_visitor.NodeVisitor):
+    """Visitor to find the enclosing path scope id of sub expressions.
+
+    Sets inherit an effective scope id from enclosing expressions,
+    and this visitor computes those.
+    """
+    def __init__(self, check: bool) -> None:
+        super().__init__()
+        self.path_scope_ids: List[Optional[int]] = [None]
+
+    def visit_Stmt(self, stmt: irast.Stmt) -> Optional[int]:
+        # Sometimes there is sharing, so we want the official scope
+        # for a node to be based on its appearance in the result,
+        # not in a subquery.
+        # I think it might not actually matter, though.
+        self.visit(stmt.bindings)
+        if stmt.iterator_stmt:
+            self.visit(stmt.iterator_stmt)
+        if isinstance(stmt, irast.MutatingStmt):
+            self.visit(stmt.subject)
+        self.visit(stmt.result)
+
+        self.generic_visit(stmt)
+        return None
+
+    def visit_Set(self, node: irast.Set) -> Optional[int]:
+        val = self.path_scope_ids[-1]
+        if node.path_scope_id:
+            self.path_scope_ids.append(node.path_scope_id)
+        if not node.is_binding:
+            val = self.path_scope_ids[-1]
+
+        # Visit sub-trees
+        self.generic_visit(node)
+
+        if node.path_scope_id:
+            self.path_scope_ids.pop()
+
+        return val
+
+
+def _try_namespace_fix(
+    scope: irast.ScopeTreeNode,
+    obj: Union[irast.ScopeTreeNode, irast.Set],
+) -> None:
+    if obj.path_id is None:
+        return
+    for prefix in obj.path_id.iter_prefixes():
+        replacement = scope.find_visible(prefix)
+        if (
+            replacement and replacement.path_id
+            and replacement.path_id != prefix
+        ):
+            new = obj.path_id.replace_prefix(prefix, replacement.path_id)
+            obj.path_id = new
+            break
+
+
+def _rewrite_weak_namespaces(
+    ir: irast.Base, ctx: context.ContextLevel
+) -> None:
+    """Rewrite weak namespaces in path ids to be usable by the backend.
+
+    Weak namespaces in path ids in the frontend are "relative", and
+    their interpretation depends on the current scope tree node and
+    the namespaces on the parent nodes. The IR->pgsql compiler does
+    not do this sort of interpretation, and needs path IDs that are
+    "absolute".
+
+    To accomplish this, we go through all the path ids and rewrite
+    them: using the scope tree, we try to find the binding of the path
+    ID (using a prefix if necessary) and drop all namespace parts that
+    don't appear in the binding.
+    """
+
+    if ctx.path_scope is None:
+        return None
+
+    tree = ctx.path_scope
+
+    for node in tree.strict_descendants:
+        _try_namespace_fix(node, node)
+
+    visitor = FindPathScopes(ctx.env.options.apply_query_rewrites)
+    visitor.visit(ir)
+
+    for ir_set in ctx.env.set_types:
+        path_scope_id: Optional[int] = visitor.memo.get(ir_set)
+        if path_scope_id is not None:
+            # Some entries in set_types are from compiling views
+            # in temporary scopes, so we need to just skip those.
+            if scope := tree.find_by_unique_id(path_scope_id):
+                _try_namespace_fix(scope, ir_set)
+
+
 def _elide_derived_ancestors(
     obj: Union[s_types.InheritingType, s_pointers.Pointer], *,
     ctx: context.ContextLevel
 ) -> None:
+
path: edb/edgeql/compiler/stmtctx.py
text: @user1:
Looks like a spurious newline.