From 772e1208542ba046bb7b80ac6a0cae3c3a373927 Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Mon, 17 Jun 2024 11:28:26 -0700
Subject: [PATCH] Fix ai::to_context duplication (#7464)

The problem is that when we use the new inlined `body` field of
`FunctionCall`, the backend compiler still always compiles the
arguments, and since the arguments weren't registered in the scope
tree, the compiled arguments might not get picked up.

Fix this by always registering arguments when we produce an inlined body.
(We could *not* compile the arguments unconditionally for inlined
functions like this, but we'd still need to register the sets in case
the argument is used twice).

Fixes #7463.
---
 edb/edgeql/compiler/func.py                   | 12 ++++++++++
 edb/edgeql/compiler/inference/multiplicity.py |  3 +++
 tests/schemas/ext_ai.esdl                     |  2 +-
 tests/test_ext_ai.py                          | 22 +++++++++++++++++++
 4 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/edb/edgeql/compiler/func.py b/edb/edgeql/compiler/func.py
index 452c28c71cf..879ec815870 100644
--- a/edb/edgeql/compiler/func.py
+++ b/edb/edgeql/compiler/func.py
@@ -348,6 +348,18 @@ def compile_FunctionCall(
     else:
         res = fcall
 
+    if isinstance(res, irast.FunctionCall) and res.body:
+        # If we are generating a special-cased inlined function call,
+        # make sure to register all the arguments in the scope tree
+        # to ensure that the compiled arguments get picked up when
+        # compiling the body.
+        for arg in res.args.values():
+            pathctx.register_set_in_scope(
+                arg.expr,
+                optional=arg.param_typemod == ft.TypeModifier.OptionalType,
+                ctx=ctx,
+            )
+
     ir_set = setgen.ensure_set(res, typehint=rtype, path_id=path_id, ctx=ctx)
     return stmt.maybe_add_view(ir_set, ctx=ctx)
 
diff --git a/edb/edgeql/compiler/inference/multiplicity.py b/edb/edgeql/compiler/inference/multiplicity.py
index 6c0f9854efb..935d0191875 100644
--- a/edb/edgeql/compiler/inference/multiplicity.py
+++ b/edb/edgeql/compiler/inference/multiplicity.py
@@ -357,6 +357,9 @@ def __infer_func_call(
         for g_arg in ir.global_args:
             _infer_set(g_arg, scope_tree=scope_tree, ctx=ctx)
 
+    if ir.body:
+        infer_multiplicity(ir.body, scope_tree=scope_tree, ctx=ctx)
+
     if card.is_single():
         return UNIQUE
     elif str(ir.func_shortname) == 'std::assert_distinct':
diff --git a/tests/schemas/ext_ai.esdl b/tests/schemas/ext_ai.esdl
index f596e7544ed..9c8320cda8c 100644
--- a/tests/schemas/ext_ai.esdl
+++ b/tests/schemas/ext_ai.esdl
@@ -36,7 +36,7 @@ type Astronomy {
 type Stuff extending Astronomy {
     content2: str;
     deferred index ext::ai::index(embedding_model := 'text-embedding-test')
-        on (.content ++ .content2);
+        on ({.content} ++ {.content2});
 };
 
 type Star extending Astronomy;
diff --git a/tests/test_ext_ai.py b/tests/test_ext_ai.py
index e56a31d1b97..835d1762b92 100644
--- a/tests/test_ext_ai.py
+++ b/tests/test_ext_ai.py
@@ -209,6 +209,17 @@ async def test_ext_ai_indexing_02(self):
             variables=dict(qv=qv),
         )
 
+        await self.assert_query_result(
+            '''
+            select _ := ext::ai::to_context((select Stuff))
+            order by _
+            ''',
+            [
+                'Skies on Earth are blue',
+                'Skies on Mars are red',
+            ],
+        )
+
         async for tr in self.try_until_succeeds(
             ignore=(AssertionError,),
             timeout=10.0,
@@ -261,6 +272,17 @@ async def test_ext_ai_indexing_03(self):
             """,
         )
 
+        await self.assert_query_result(
+            '''
+            select _ := ext::ai::to_context((select Star))
+            order by _
+            ''',
+            [
+                'Skies on Earth are blue',
+                'Skies on Mars are red',
+            ],
+        )
+
         async for tr in self.try_until_succeeds(
             ignore=(AssertionError,),
             timeout=10.0,