We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Group by - ProtocolError: cannot decode Object: expected 3 elements, got 2110812590
Hi

```
# this is working
group EntryItem
by .account;

# this is not working
group EntryItem
using accountCode := .account
by accountCode;
-> ProtocolError: cannot decode Object: expected 3 elements, got 2110812590
```

- EdgeDB Version: 2.9
- EdgeDB CLI Version: 2.2.6+7eabbf9
- OS Version: macOS 12.1

Schema:

```
module default {
    type Account {
        required link book -> Book;
        required property code -> str {
            constraint exclusive;
        };

        property displayCode := .code[0:5] ++ ' ' ++ .code[5:];
        required property name -> str;
        required link currency -> Currency;

        constraint exclusive on ((.book, .code));
    }

    type EntryItem {
        required link entry -> Entry;
        required property lineNumber -> int16;

        required link account -> Account;
        ...

        constraint exclusive on ((.entry, .lineNumber))
    }
}
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `edb/edgeql/desugar_group.py`
Content:
```
1 #
2 # This source file is part of the EdgeDB open source project.
3 #
4 # Copyright 2008-present MagicStack Inc. and the EdgeDB authors.
5 #
6 # Licensed under the Apache License, Version 2.0 (the "License");
7 # you may not use this file except in compliance with the License.
8 # You may obtain a copy of the License at
9 #
10 #     http://www.apache.org/licenses/LICENSE-2.0
11 #
12 # Unless required by applicable law or agreed to in writing, software
13 # distributed under the License is distributed on an "AS IS" BASIS,
14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 # See the License for the specific language governing permissions and
16 # limitations under the License.
17 #
18 
19 """Desugar GROUP queries into internal FOR GROUP queries.
20 
21 This code is called by both the model and the real implementation,
22 though if that starts becoming a problem it should just be abandoned.
23 """
24 
25 from __future__ import annotations
26 
27 
28 from typing import Optional, Tuple, AbstractSet, Dict, List
29 
30 from edb.common import ast
31 from edb.common import ordered
32 from edb.common.compiler import AliasGenerator
33 
34 from edb.edgeql import ast as qlast
35 from edb.edgeql.compiler import astutils
36 
37 
38 def key_name(s: str) -> str:
39     return s.split('~')[0]
40 
41 
42 def name_path(name: str) -> qlast.Path:
43     return qlast.Path(steps=[qlast.ObjectRef(name=name)])
44 
45 
46 def make_free_object(els: Dict[str, qlast.Expr]) -> qlast.Shape:
47     return qlast.Shape(
48         expr=None,
49         elements=[
50             qlast.ShapeElement(
51                 expr=qlast.Path(steps=[qlast.Ptr(name=name)]),
52                 compexpr=expr
53             )
54             for name, expr in els.items()
55         ],
56     )
57 
58 
59 def collect_grouping_atoms(
60     els: List[qlast.GroupingElement],
61 ) -> AbstractSet[str]:
62     atoms: ordered.OrderedSet[str] = ordered.OrderedSet()
63 
64     def _collect_atom(el: qlast.GroupingAtom) -> None:
65         if isinstance(el, qlast.GroupingIdentList):
66             for at in el.elements:
67                 _collect_atom(at)
68 
69         else:
70             assert isinstance(el, qlast.ObjectRef)
71             atoms.add(el.name)
72 
73     def _collect_el(el: qlast.GroupingElement) -> None:
74         if isinstance(el, qlast.GroupingSets):
75             for sub in el.sets:
76                 _collect_el(sub)
77         elif isinstance(el, qlast.GroupingOperation):
78             for at in el.elements:
79                 _collect_atom(at)
80         elif isinstance(el, qlast.GroupingSimple):
81             _collect_atom(el.element)
82         else:
83             raise AssertionError('Unknown GroupingElement')
84 
85     for el in els:
86         _collect_el(el)
87 
88     return atoms
89 
90 
91 def desugar_group(
92     node: qlast.GroupQuery,
93     aliases: AliasGenerator,
94 ) -> qlast.InternalGroupQuery:
95     assert not isinstance(node, qlast.InternalGroupQuery)
96     alias_map: Dict[str, Tuple[str, qlast.Expr]] = {}
97 
98     def rewrite_atom(el: qlast.GroupingAtom) -> qlast.GroupingAtom:
99         if isinstance(el, qlast.ObjectRef):
100             return el
101         elif isinstance(el, qlast.Path):
102             assert isinstance(el.steps[0], qlast.Ptr)
103             ptrname = el.steps[0].name
104             if ptrname not in alias_map:
105                 alias = aliases.get(ptrname)
106                 alias_map[ptrname] = (alias, el)
107             alias = alias_map[ptrname][0]
108             return qlast.ObjectRef(name=alias)
109         else:
110             return qlast.GroupingIdentList(
111                 span=el.span,
112                 elements=tuple(rewrite_atom(at) for at in el.elements),
113             )
114 
115     def rewrite(el: qlast.GroupingElement) -> qlast.GroupingElement:
116         if isinstance(el, qlast.GroupingSimple):
117             return qlast.GroupingSimple(
118                 span=el.span, element=rewrite_atom(el.element))
119         elif isinstance(el, qlast.GroupingSets):
120             return qlast.GroupingSets(
121                 span=el.span, sets=[rewrite(s) for s in el.sets])
122         elif isinstance(el, qlast.GroupingOperation):
123             return qlast.GroupingOperation(
124                 span=el.span,
125                 oper=el.oper,
126                 elements=[rewrite_atom(a) for a in el.elements])
127         raise AssertionError
128 
129     for using_clause in (node.using or ()):
130         alias_map[using_clause.alias] = (using_clause.alias, using_clause.expr)
131 
132     using = node.using[:] if node.using else []
133     by = [rewrite(by_el) for by_el in node.by]
134     for alias, path in alias_map.values():
135         using.append(qlast.AliasedExpr(alias=alias, expr=path))
136 
137     actual_keys = collect_grouping_atoms(by)
138 
139     g_alias = aliases.get('g')
140     grouping_alias = aliases.get('grouping')
141     output_dict = {
142         'key': make_free_object({
143             name: name_path(alias)
144             for name, (alias, _) in alias_map.items()
145             if alias in actual_keys
146         }),
147         'grouping': qlast.FunctionCall(
148             func='array_unpack',
149             args=[name_path(grouping_alias)],
150         ),
151         'elements': name_path(g_alias),
152     }
153     output_shape = make_free_object(output_dict)
154 
155     return qlast.InternalGroupQuery(
156         span=node.span,
157         aliases=node.aliases,
158         subject_alias=node.subject_alias,
159         subject=node.subject,
160         # rewritten parts!
161         using=using,
162         by=by,
163         group_alias=g_alias,
164         grouping_alias=grouping_alias,
165         result=output_shape,
166         from_desugaring=True,
167     )
168 
169 
170 def _count_alias_uses(
171     node: qlast.Expr,
172     alias: str,
173 ) -> int:
174     uses = 0
175     for child in ast.find_children(node, qlast.Path):
176         match child:
177             case astutils.alias_view((alias2, _)) if alias == alias2:
178                 uses += 1
179     return uses
180 
181 
182 def try_group_rewrite(
183     node: qlast.Query,
184     aliases: AliasGenerator,
185 ) -> Optional[qlast.Query]:
186     """
187     Try to apply some syntactic rewrites of GROUP expressions so we
188     can generate better code.
189 
190     The two key desugarings are:
191 
192     * Sink a shape into the internal group result
193 
194         SELECT (GROUP ...) <shape>
195         [filter-clause] [order-clause] [other clauses]
196         =>
197         SELECT (
198           FOR GROUP ...
199           UNION <igroup-body> <shape>
200           [filter-clause]
201           [order-clause]
202         ) [other clauses]
203 
204     * Convert a FOR over a group into just an internal group (and
205       a trivial FOR)
206 
207         FOR g in (GROUP ...) UNION <body>
208         =>
209         FOR GROUP ...
210         UNION (
211             FOR g IN (<group-body>)
212             UNION <body>
213         )
214     """
215 
216     # Inline trivial uses of aliases bound to a group and then
217     # immediately used, so that we can apply the other optimizations.
218     match node:
219         case qlast.SelectQuery(
220             aliases=[
221                 *_,
222                 qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)
223             ] as qaliases,
224             result=qlast.Shape(
225                 expr=astutils.alias_view((alias2, [])),
226                 elements=elements,
227             ) as result,
228         ) if alias == alias2 and _count_alias_uses(result, alias) == 1:
229             node = node.replace(
230                 aliases=qaliases[:-1],
231                 result=qlast.Shape(expr=grp, elements=elements),
232             )
233 
234         case qlast.ForQuery(
235             aliases=[
236                 *_,
237                 qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)
238             ] as qaliases,
239             iterator=astutils.alias_view((alias2, [])),
240             result=result,
241         ) if alias == alias2 and _count_alias_uses(result, alias) == 0:
242             node = node.replace(
243                 aliases=qaliases[:-1],
244                 iterator=grp,
245             )
246 
247     # Sink shapes into the GROUP
248     if (
249         isinstance(node, qlast.SelectQuery)
250         and isinstance(node.result, qlast.Shape)
251         and isinstance(node.result.expr, qlast.GroupQuery)
252     ):
253         igroup = desugar_group(node.result.expr, aliases)
254         igroup = igroup.replace(result=qlast.Shape(
255             expr=igroup.result, elements=node.result.elements))
256 
257         # FILTER gets sunk into the body of the FOR GROUP
258         if node.where or node.orderby:
259             igroup = igroup.replace(
260                 # We need to move the result_alias in case
261                 # the FILTER depends on it.
262                 result_alias=node.result_alias,
263                 where=node.where,
264                 orderby=node.orderby,
265             )
266 
267         return node.replace(
268             result=igroup, result_alias=None, where=None, orderby=None)
269 
270     # Eliminate FORs over GROUPs
271     if (
272         isinstance(node, qlast.ForQuery)
273         and isinstance(node.iterator, qlast.GroupQuery)
274     ):
275         igroup = desugar_group(node.iterator, aliases)
276         new_result = qlast.ForQuery(
277             iterator_alias=node.iterator_alias,
278             iterator=igroup.result,
279             result=node.result,
280         )
281         return igroup.replace(result=new_result, aliases=node.aliases)
282 
283     return None
284
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/edb/edgeql/desugar_group.py b/edb/edgeql/desugar_group.py
--- a/edb/edgeql/desugar_group.py
+++ b/edb/edgeql/desugar_group.py
@@ -27,6 +27,8 @@
 
 from typing import Optional, Tuple, AbstractSet, Dict, List
 
+from edb import errors
+
 from edb.common import ast
 from edb.common import ordered
 from edb.common.compiler import AliasGenerator
@@ -126,11 +128,22 @@
                 elements=[rewrite_atom(a) for a in el.elements])
         raise AssertionError
 
+    # The rewrite calls on the grouping elements populate alias_map
+    # with any bindings for pointers the by clause refers to directly.
+    by = [rewrite(by_el) for by_el in node.by]
+
     for using_clause in (node.using or ()):
+        if using_clause.alias in alias_map:
+            # TODO: This would be a great place to allow multiple spans!
+            raise errors.QueryError(
+                f"USING clause binds a variable '{using_clause.alias}' "
+                f"but a property with that name is used directly in the BY "
+                f"clause",
+                span=alias_map[using_clause.alias][1].span,
+            )
         alias_map[using_clause.alias] = (using_clause.alias, using_clause.expr)
 
-    using = node.using[:] if node.using else []
-    by = [rewrite(by_el) for by_el in node.by]
+    using = []
     for alias, path in alias_map.values():
         using.append(qlast.AliasedExpr(alias=alias, expr=path))

{"golden_diff": "diff --git a/edb/edgeql/desugar_group.py b/edb/edgeql/desugar_group.py\n--- a/edb/edgeql/desugar_group.py\n+++ b/edb/edgeql/desugar_group.py\n@@ -27,6 +27,8 @@\n \n from typing import Optional, Tuple, AbstractSet, Dict, List\n \n+from edb import errors\n+\n from edb.common import ast\n from edb.common import ordered\n from edb.common.compiler import AliasGenerator\n@@ -126,11 +128,22 @@\n                 elements=[rewrite_atom(a) for a in el.elements])\n         raise AssertionError\n \n+    # The rewrite calls on the grouping elements populate alias_map\n+    # with any bindings for pointers the by clause refers to directly.\n+    by = [rewrite(by_el) for by_el in node.by]\n+\n     for using_clause in (node.using or ()):\n+        if using_clause.alias in alias_map:\n+            # TODO: This would be a great place to allow multiple spans!\n+            raise errors.QueryError(\n+                f\"USING clause binds a variable '{using_clause.alias}' \"\n+                f\"but a property with that name is used directly in the BY \"\n+                f\"clause\",\n+                span=alias_map[using_clause.alias][1].span,\n+            )\n         alias_map[using_clause.alias] = (using_clause.alias, using_clause.expr)\n \n-    using = node.using[:] if node.using else []\n-    by = [rewrite(by_el) for by_el in node.by]\n+    using = []\n     for alias, path in alias_map.values():\n         using.append(qlast.AliasedExpr(alias=alias, expr=path))\n", "issue": "Group by - ProtocolError: cannot decode Object: expected 3 elements, got 2110812590\nHi\r\n\r\n```\r\n# this is working\r\ngroup EntryItem\r\nby .account;\r\n\r\n# this is not working\r\ngroup EntryItem\r\nusing accountCode := .account\r\nby accountCode;\r\n-> ProtocolError: cannot decode Object: expected 3 elements, got 2110812590\r\n```\r\n\r\n- EdgeDB Version: 2.9\r\n- EdgeDB CLI Version: 2.2.6+7eabbf9\r\n- OS Version: macOS 12.1\r\n\r\nSchema:\r\n\r\n```\r\nmodule default {\r\n    type Account {\r\n        required link book -> Book;\r\n        required property code -> str {\r\n            constraint exclusive;\r\n        };\r\n\r\n        property displayCode := .code[0:5] ++ ' ' ++ .code[5:];\r\n        required property name -> str;\r\n        required link currency -> Currency;\r\n\r\n        constraint exclusive on ((.book, .code));\r\n    }\r\n\r\n    type EntryItem {\r\n        required link entry -> Entry;\r\n        required property lineNumber -> int16;\r\n\r\n        required link account -> Account;\r\n        ...\r\n\r\n        constraint exclusive on ((.entry, .lineNumber))\r\n    }\r\n}\r\n```\r\n\n", "before_files": [{"content": "#\n# This source file is part of the EdgeDB open source project.\n#\n# Copyright 2008-present MagicStack Inc. and the EdgeDB authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n\"\"\"Desugar GROUP queries into internal FOR GROUP queries.\n\nThis code is called by both the model and the real implementation,\nthough if that starts becoming a problem it should just be abandoned.\n\"\"\"\n\nfrom __future__ import annotations\n\n\nfrom typing import Optional, Tuple, AbstractSet, Dict, List\n\nfrom edb.common import ast\nfrom edb.common import ordered\nfrom edb.common.compiler import AliasGenerator\n\nfrom edb.edgeql import ast as qlast\nfrom edb.edgeql.compiler import astutils\n\n\ndef key_name(s: str) -> str:\n    return s.split('~')[0]\n\n\ndef name_path(name: str) -> qlast.Path:\n    return qlast.Path(steps=[qlast.ObjectRef(name=name)])\n\n\ndef make_free_object(els: Dict[str, qlast.Expr]) -> qlast.Shape:\n    return qlast.Shape(\n        expr=None,\n        elements=[\n            qlast.ShapeElement(\n                expr=qlast.Path(steps=[qlast.Ptr(name=name)]),\n                compexpr=expr\n            )\n            for name, expr in els.items()\n        ],\n    )\n\n\ndef collect_grouping_atoms(\n    els: List[qlast.GroupingElement],\n) -> AbstractSet[str]:\n    atoms: ordered.OrderedSet[str] = ordered.OrderedSet()\n\n    def _collect_atom(el: qlast.GroupingAtom) -> None:\n        if isinstance(el, qlast.GroupingIdentList):\n            for at in el.elements:\n                _collect_atom(at)\n\n        else:\n            assert isinstance(el, qlast.ObjectRef)\n            atoms.add(el.name)\n\n    def _collect_el(el: qlast.GroupingElement) -> None:\n        if isinstance(el, qlast.GroupingSets):\n            for sub in el.sets:\n                _collect_el(sub)\n        elif isinstance(el, qlast.GroupingOperation):\n            for at in el.elements:\n                _collect_atom(at)\n        elif isinstance(el, qlast.GroupingSimple):\n            _collect_atom(el.element)\n        else:\n            raise AssertionError('Unknown GroupingElement')\n\n    for el in els:\n        _collect_el(el)\n\n    return atoms\n\n\ndef desugar_group(\n    node: qlast.GroupQuery,\n    aliases: AliasGenerator,\n) -> qlast.InternalGroupQuery:\n    assert not isinstance(node, qlast.InternalGroupQuery)\n    alias_map: Dict[str, Tuple[str, qlast.Expr]] = {}\n\n    def rewrite_atom(el: qlast.GroupingAtom) -> qlast.GroupingAtom:\n        if isinstance(el, qlast.ObjectRef):\n            return el\n        elif isinstance(el, qlast.Path):\n            assert isinstance(el.steps[0], qlast.Ptr)\n            ptrname = el.steps[0].name\n            if ptrname not in alias_map:\n                alias = aliases.get(ptrname)\n                alias_map[ptrname] = (alias, el)\n            alias = alias_map[ptrname][0]\n            return qlast.ObjectRef(name=alias)\n        else:\n            return qlast.GroupingIdentList(\n                span=el.span,\n                elements=tuple(rewrite_atom(at) for at in el.elements),\n            )\n\n    def rewrite(el: qlast.GroupingElement) -> qlast.GroupingElement:\n        if isinstance(el, qlast.GroupingSimple):\n            return qlast.GroupingSimple(\n                span=el.span, element=rewrite_atom(el.element))\n        elif isinstance(el, qlast.GroupingSets):\n            return qlast.GroupingSets(\n                span=el.span, sets=[rewrite(s) for s in el.sets])\n        elif isinstance(el, qlast.GroupingOperation):\n            return qlast.GroupingOperation(\n                span=el.span,\n                oper=el.oper,\n                elements=[rewrite_atom(a) for a in el.elements])\n        raise AssertionError\n\n    for using_clause in (node.using or ()):\n        alias_map[using_clause.alias] = (using_clause.alias, using_clause.expr)\n\n    using = node.using[:] if node.using else []\n    by = [rewrite(by_el) for by_el in node.by]\n    for alias, path in alias_map.values():\n        using.append(qlast.AliasedExpr(alias=alias, expr=path))\n\n    actual_keys = collect_grouping_atoms(by)\n\n    g_alias = aliases.get('g')\n    grouping_alias = aliases.get('grouping')\n    output_dict = {\n        'key': make_free_object({\n            name: name_path(alias)\n            for name, (alias, _) in alias_map.items()\n            if alias in actual_keys\n        }),\n        'grouping': qlast.FunctionCall(\n            func='array_unpack',\n            args=[name_path(grouping_alias)],\n        ),\n        'elements': name_path(g_alias),\n    }\n    output_shape = make_free_object(output_dict)\n\n    return qlast.InternalGroupQuery(\n        span=node.span,\n        aliases=node.aliases,\n        subject_alias=node.subject_alias,\n        subject=node.subject,\n        # rewritten parts!\n        using=using,\n        by=by,\n        group_alias=g_alias,\n        grouping_alias=grouping_alias,\n        result=output_shape,\n        from_desugaring=True,\n    )\n\n\ndef _count_alias_uses(\n    node: qlast.Expr,\n    alias: str,\n) -> int:\n    uses = 0\n    for child in ast.find_children(node, qlast.Path):\n        match child:\n            case astutils.alias_view((alias2, _)) if alias == alias2:\n                uses += 1\n    return uses\n\n\ndef try_group_rewrite(\n    node: qlast.Query,\n    aliases: AliasGenerator,\n) -> Optional[qlast.Query]:\n    \"\"\"\n    Try to apply some syntactic rewrites of GROUP expressions so we\n    can generate better code.\n\n    The two key desugarings are:\n\n    * Sink a shape into the internal group result\n\n        SELECT (GROUP ...) <shape>\n        [filter-clause] [order-clause] [other clauses]\n        =>\n        SELECT (\n          FOR GROUP ...\n          UNION <igroup-body> <shape>\n          [filter-clause]\n          [order-clause]\n        ) [other clauses]\n\n    * Convert a FOR over a group into just an internal group (and\n      a trivial FOR)\n\n        FOR g in (GROUP ...) UNION <body>\n        =>\n        FOR GROUP ...\n        UNION (\n            FOR g IN (<group-body>)\n            UNION <body>\n        )\n    \"\"\"\n\n    # Inline trivial uses of aliases bound to a group and then\n    # immediately used, so that we can apply the other optimizations.\n    match node:\n        case qlast.SelectQuery(\n            aliases=[\n                *_,\n                qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)\n            ] as qaliases,\n            result=qlast.Shape(\n                expr=astutils.alias_view((alias2, [])),\n                elements=elements,\n            ) as result,\n        ) if alias == alias2 and _count_alias_uses(result, alias) == 1:\n            node = node.replace(\n                aliases=qaliases[:-1],\n                result=qlast.Shape(expr=grp, elements=elements),\n            )\n\n        case qlast.ForQuery(\n            aliases=[\n                *_,\n                qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)\n            ] as qaliases,\n            iterator=astutils.alias_view((alias2, [])),\n            result=result,\n        ) if alias == alias2 and _count_alias_uses(result, alias) == 0:\n            node = node.replace(\n                aliases=qaliases[:-1],\n                iterator=grp,\n            )\n\n    # Sink shapes into the GROUP\n    if (\n        isinstance(node, qlast.SelectQuery)\n        and isinstance(node.result, qlast.Shape)\n        and isinstance(node.result.expr, qlast.GroupQuery)\n    ):\n        igroup = desugar_group(node.result.expr, aliases)\n        igroup = igroup.replace(result=qlast.Shape(\n            expr=igroup.result, elements=node.result.elements))\n\n        # FILTER gets sunk into the body of the FOR GROUP\n        if node.where or node.orderby:\n            igroup = igroup.replace(\n                # We need to move the result_alias in case\n                # the FILTER depends on it.\n                result_alias=node.result_alias,\n                where=node.where,\n                orderby=node.orderby,\n            )\n\n        return node.replace(\n            result=igroup, result_alias=None, where=None, orderby=None)\n\n    # Eliminate FORs over GROUPs\n    if (\n        isinstance(node, qlast.ForQuery)\n        and isinstance(node.iterator, qlast.GroupQuery)\n    ):\n        igroup = desugar_group(node.iterator, aliases)\n        new_result = qlast.ForQuery(\n            iterator_alias=node.iterator_alias,\n            iterator=igroup.result,\n            result=node.result,\n        )\n        return igroup.replace(result=new_result, aliases=node.aliases)\n\n    return None\n", "path": "edb/edgeql/desugar_group.py"}], "after_files": [{"content": "#\n# This source file is part of the EdgeDB open source project.\n#\n# Copyright 2008-present MagicStack Inc. and the EdgeDB authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n\"\"\"Desugar GROUP queries into internal FOR GROUP queries.\n\nThis code is called by both the model and the real implementation,\nthough if that starts becoming a problem it should just be abandoned.\n\"\"\"\n\nfrom __future__ import annotations\n\n\nfrom typing import Optional, Tuple, AbstractSet, Dict, List\n\nfrom edb import errors\n\nfrom edb.common import ast\nfrom edb.common import ordered\nfrom edb.common.compiler import AliasGenerator\n\nfrom edb.edgeql import ast as qlast\nfrom edb.edgeql.compiler import astutils\n\n\ndef key_name(s: str) -> str:\n    return s.split('~')[0]\n\n\ndef name_path(name: str) -> qlast.Path:\n    return qlast.Path(steps=[qlast.ObjectRef(name=name)])\n\n\ndef make_free_object(els: Dict[str, qlast.Expr]) -> qlast.Shape:\n    return qlast.Shape(\n        expr=None,\n        elements=[\n            qlast.ShapeElement(\n                expr=qlast.Path(steps=[qlast.Ptr(name=name)]),\n                compexpr=expr\n            )\n            for name, expr in els.items()\n        ],\n    )\n\n\ndef collect_grouping_atoms(\n    els: List[qlast.GroupingElement],\n) -> AbstractSet[str]:\n    atoms: ordered.OrderedSet[str] = ordered.OrderedSet()\n\n    def _collect_atom(el: qlast.GroupingAtom) -> None:\n        if isinstance(el, qlast.GroupingIdentList):\n            for at in el.elements:\n                _collect_atom(at)\n\n        else:\n            assert isinstance(el, qlast.ObjectRef)\n            atoms.add(el.name)\n\n    def _collect_el(el: qlast.GroupingElement) -> None:\n        if isinstance(el, qlast.GroupingSets):\n            for sub in el.sets:\n                _collect_el(sub)\n        elif isinstance(el, qlast.GroupingOperation):\n            for at in el.elements:\n                _collect_atom(at)\n        elif isinstance(el, qlast.GroupingSimple):\n            _collect_atom(el.element)\n        else:\n            raise AssertionError('Unknown GroupingElement')\n\n    for el in els:\n        _collect_el(el)\n\n    return atoms\n\n\ndef desugar_group(\n    node: qlast.GroupQuery,\n    aliases: AliasGenerator,\n) -> qlast.InternalGroupQuery:\n    assert not isinstance(node, qlast.InternalGroupQuery)\n    alias_map: Dict[str, Tuple[str, qlast.Expr]] = {}\n\n    def rewrite_atom(el: qlast.GroupingAtom) -> qlast.GroupingAtom:\n        if isinstance(el, qlast.ObjectRef):\n            return el\n        elif isinstance(el, qlast.Path):\n            assert isinstance(el.steps[0], qlast.Ptr)\n            ptrname = el.steps[0].name\n            if ptrname not in alias_map:\n                alias = aliases.get(ptrname)\n                alias_map[ptrname] = (alias, el)\n            alias = alias_map[ptrname][0]\n            return qlast.ObjectRef(name=alias)\n        else:\n            return qlast.GroupingIdentList(\n                span=el.span,\n                elements=tuple(rewrite_atom(at) for at in el.elements),\n            )\n\n    def rewrite(el: qlast.GroupingElement) -> qlast.GroupingElement:\n        if isinstance(el, qlast.GroupingSimple):\n            return qlast.GroupingSimple(\n                span=el.span, element=rewrite_atom(el.element))\n        elif isinstance(el, qlast.GroupingSets):\n            return qlast.GroupingSets(\n                span=el.span, sets=[rewrite(s) for s in el.sets])\n        elif isinstance(el, qlast.GroupingOperation):\n            return qlast.GroupingOperation(\n                span=el.span,\n                oper=el.oper,\n                elements=[rewrite_atom(a) for a in el.elements])\n        raise AssertionError\n\n    # The rewrite calls on the grouping elements populate alias_map\n    # with any bindings for pointers the by clause refers to directly.\n    by = [rewrite(by_el) for by_el in node.by]\n\n    for using_clause in (node.using or ()):\n        if using_clause.alias in alias_map:\n            # TODO: This would be a great place to allow multiple spans!\n            raise errors.QueryError(\n                f\"USING clause binds a variable '{using_clause.alias}' \"\n                f\"but a property with that name is used directly in the BY \"\n                f\"clause\",\n                span=alias_map[using_clause.alias][1].span,\n            )\n        alias_map[using_clause.alias] = (using_clause.alias, using_clause.expr)\n\n    using = []\n    for alias, path in alias_map.values():\n        using.append(qlast.AliasedExpr(alias=alias, expr=path))\n\n    actual_keys = collect_grouping_atoms(by)\n\n    g_alias = aliases.get('g')\n    grouping_alias = aliases.get('grouping')\n    output_dict = {\n        'key': make_free_object({\n            name: name_path(alias)\n            for name, (alias, _) in alias_map.items()\n            if alias in actual_keys\n        }),\n        'grouping': qlast.FunctionCall(\n            func='array_unpack',\n            args=[name_path(grouping_alias)],\n        ),\n        'elements': name_path(g_alias),\n    }\n    output_shape = make_free_object(output_dict)\n\n    return qlast.InternalGroupQuery(\n        span=node.span,\n        aliases=node.aliases,\n        subject_alias=node.subject_alias,\n        subject=node.subject,\n        # rewritten parts!\n        using=using,\n        by=by,\n        group_alias=g_alias,\n        grouping_alias=grouping_alias,\n        result=output_shape,\n        from_desugaring=True,\n    )\n\n\ndef _count_alias_uses(\n    node: qlast.Expr,\n    alias: str,\n) -> int:\n    uses = 0\n    for child in ast.find_children(node, qlast.Path):\n        match child:\n            case astutils.alias_view((alias2, _)) if alias == alias2:\n                uses += 1\n    return uses\n\n\ndef try_group_rewrite(\n    node: qlast.Query,\n    aliases: AliasGenerator,\n) -> Optional[qlast.Query]:\n    \"\"\"\n    Try to apply some syntactic rewrites of GROUP expressions so we\n    can generate better code.\n\n    The two key desugarings are:\n\n    * Sink a shape into the internal group result\n\n        SELECT (GROUP ...) <shape>\n        [filter-clause] [order-clause] [other clauses]\n        =>\n        SELECT (\n          FOR GROUP ...\n          UNION <igroup-body> <shape>\n          [filter-clause]\n          [order-clause]\n        ) [other clauses]\n\n    * Convert a FOR over a group into just an internal group (and\n      a trivial FOR)\n\n        FOR g in (GROUP ...) UNION <body>\n        =>\n        FOR GROUP ...\n        UNION (\n            FOR g IN (<group-body>)\n            UNION <body>\n        )\n    \"\"\"\n\n    # Inline trivial uses of aliases bound to a group and then\n    # immediately used, so that we can apply the other optimizations.\n    match node:\n        case qlast.SelectQuery(\n            aliases=[\n                *_,\n                qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)\n            ] as qaliases,\n            result=qlast.Shape(\n                expr=astutils.alias_view((alias2, [])),\n                elements=elements,\n            ) as result,\n        ) if alias == alias2 and _count_alias_uses(result, alias) == 1:\n            node = node.replace(\n                aliases=qaliases[:-1],\n                result=qlast.Shape(expr=grp, elements=elements),\n            )\n\n        case qlast.ForQuery(\n            aliases=[\n                *_,\n                qlast.AliasedExpr(alias=alias, expr=qlast.GroupQuery() as grp)\n            ] as qaliases,\n            iterator=astutils.alias_view((alias2, [])),\n            result=result,\n        ) if alias == alias2 and _count_alias_uses(result, alias) == 0:\n            node = node.replace(\n                aliases=qaliases[:-1],\n                iterator=grp,\n            )\n\n    # Sink shapes into the GROUP\n    if (\n        isinstance(node, qlast.SelectQuery)\n        and isinstance(node.result, qlast.Shape)\n        and isinstance(node.result.expr, qlast.GroupQuery)\n    ):\n        igroup = desugar_group(node.result.expr, aliases)\n        igroup = igroup.replace(result=qlast.Shape(\n            expr=igroup.result, elements=node.result.elements))\n\n        # FILTER gets sunk into the body of the FOR GROUP\n        if node.where or node.orderby:\n            igroup = igroup.replace(\n                # We need to move the result_alias in case\n                # the FILTER depends on it.\n                result_alias=node.result_alias,\n                where=node.where,\n                orderby=node.orderby,\n            )\n\n        return node.replace(\n            result=igroup, result_alias=None, where=None, orderby=None)\n\n    # Eliminate FORs over GROUPs\n    if (\n        isinstance(node, qlast.ForQuery)\n        and isinstance(node.iterator, qlast.GroupQuery)\n    ):\n        igroup = desugar_group(node.iterator, aliases)\n        new_result = qlast.ForQuery(\n            iterator_alias=node.iterator_alias,\n            iterator=igroup.result,\n            result=node.result,\n        )\n        return igroup.replace(result=new_result, aliases=node.aliases)\n\n    return None\n", "path": "edb/edgeql/desugar_group.py"}]}