From c63404d52bc888bbc74808686d46184bf29ca0f3 Mon Sep 17 00:00:00 2001
From: Scott Trinh <scott@edgedb.com>
Date: Mon, 17 Jun 2024 15:47:40 -0400
Subject: [PATCH] Use given cardinality in select expression (#1043)

Previously, if you made a select expression with a constant or coalescing
expression, we always calculated the cardinality of a object key based on the
pointer's cardinality, rather than allowing you to return something assignable
to that cardinality. This change will check if the assignment is possible, and
override the cardinality when converting the query builder expression to a TS
type.
---
 integration-tests/lts/bench.ts             |  7 +++
 integration-tests/lts/select.test.ts       | 55 ++++++++++++++++++++++
 packages/generate/src/syntax/typesystem.ts | 14 ++++--
 3 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/integration-tests/lts/bench.ts b/integration-tests/lts/bench.ts
index 7eac3e769..10b3b9a77 100644
--- a/integration-tests/lts/bench.ts
+++ b/integration-tests/lts/bench.ts
@@ -225,6 +225,13 @@ bench("select: with offset", () => {
   return {} as typeof query;
 }).types([8891, "instantiations"]);
 
+bench("select: with pointer override", () => {
+  const query = e.select(e.Hero, (h) => ({
+    height: e.decimal("10.0"),
+  }));
+  return {} as typeof query;
+}).types([2160, "instantiations"]);
+
 bench("params select", () => {
   const query = e.params({ name: e.str }, (params) =>
     e.select(e.Hero, (hero) => ({
diff --git a/integration-tests/lts/select.test.ts b/integration-tests/lts/select.test.ts
index 0533626bf..b575b1159 100644
--- a/integration-tests/lts/select.test.ts
+++ b/integration-tests/lts/select.test.ts
@@ -1170,6 +1170,61 @@ SELECT __scope_0_defaultPerson {
     assert.equal(res4, 1);
   });
 
+  test("overriding pointers", async () => {
+    const q = e.select(e.Hero, () => ({
+      height: e.decimal("10.0"),
+    }));
+    type Height = $infer<typeof q.__element__.__shape__.height>;
+    type Q = $infer<typeof q>;
+    tc.assert<
+      tc.IsExact<
+        Q,
+        {
+          height: Height;
+        }[]
+      >
+    >(true);
+
+    await assert.rejects(
+      async () =>
+        e
+          // @ts-expect-error cannot assign multi cardinality if point is single
+          .select(e.Hero, () => ({
+            height: e.set(e.decimal("10.0"), e.decimal("11.0")),
+          }))
+          .run(client),
+      (err) => {
+        if (
+          err &&
+          typeof err === "object" &&
+          "message" in err &&
+          typeof err.message === "string"
+        ) {
+          assert.match(err.message, /possibly more than one element returned/);
+          return true;
+        } else {
+          assert.fail("Expected error to be an object with a message");
+        }
+      },
+    );
+  });
+
+  test("coalescing pointers", () => {
+    const q = e.select(e.Hero, (h) => ({
+      height: e.op(h.height, "??", e.decimal("10.0")),
+    }));
+    type Height = $infer<typeof q.__element__.__shape__.height>;
+    type Q = $infer<typeof q>;
+    tc.assert<
+      tc.IsExact<
+        Q,
+        {
+          height: Height;
+        }[]
+      >
+    >(true);
+  });
+
   test("nested matching scopes", async () => {
     const q = e.select(e.Hero, (h) => ({
       name: h.name,
diff --git a/packages/generate/src/syntax/typesystem.ts b/packages/generate/src/syntax/typesystem.ts
index 3f7f57c0b..21ac4347a 100644
--- a/packages/generate/src/syntax/typesystem.ts
+++ b/packages/generate/src/syntax/typesystem.ts
@@ -372,10 +372,16 @@ export type computeObjectShape<
                 ShapeEl
               > | null
             : never
-          : [k] extends [keyof Pointers]
-            ? shapeElementToTs<Pointers[k], Shape[k]>
-            : Shape[k] extends TypeSet
-              ? setToTsType<Shape[k]>
+          : Shape[k] extends TypeSet
+            ? [k] extends [keyof Pointers]
+              ? Shape[k]["__cardinality__"] extends cardutil.assignable<
+                  Pointers[k]["cardinality"]
+                >
+                ? setToTsType<Shape[k]>
+                : never
+              : setToTsType<Shape[k]>
+            : [k] extends [keyof Pointers]
+              ? shapeElementToTs<Pointers[k], Shape[k]>
               : never;
       }
 >;