From 5b5d19c5807ba900b36ce6e581211c3ba74b0bc4 Mon Sep 17 00:00:00 2001
From: Carson Full <carsonfull@gmail.com>
Date: Sun, 1 Sep 2024 15:38:01 -0500
Subject: [PATCH] Require opt-in to prevent it being a breaking change

---
 integration-tests/lts/select.test.ts       | 34 ++++++++-------
 packages/generate/src/syntax/typesystem.ts | 50 ++++++++++++++++++++++
 2 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/integration-tests/lts/select.test.ts b/integration-tests/lts/select.test.ts
index 375c0db6c..0610455d1 100644
--- a/integration-tests/lts/select.test.ts
+++ b/integration-tests/lts/select.test.ts
@@ -5,6 +5,12 @@ import * as $ from "../../packages/generate/src/syntax/reflection";
 import e, { type $infer } from "./dbschema/edgeql-js";
 import { setupTests, teardownTests, tc, type TestData } from "./setupTeardown";
 
+declare module "./dbschema/edgeql-js/typesystem" {
+  export interface TypesystemOptions {
+    polymorphismAsDistributedUnions: true;
+  }
+}
+
 let client: edgedb.Client;
 let data: TestData;
 
@@ -239,7 +245,7 @@ describe("select", () => {
     assert.deepEqual(query.__element__.__kind__, $.TypeKind.object);
     assert.equal(query.__element__.__name__, "default::Person");
 
-    type result = $.BaseTypeToTsType<(typeof query)["__element__"]>;
+    type result = $infer<typeof query>[number];
     tc.assert<
       tc.IsExact<
         result,
@@ -294,21 +300,17 @@ describe("select", () => {
       }),
     }));
 
-    type q = $.setToTsType<typeof q>;
-    tc.assert<
-      tc.IsExact<
-        q,
-        ({
-          id: string;
-        } & (
-          | { __typename: "default::Hero"; secret_identity: string | null }
-          | {
-              __typename: "default::Villain";
-              nemesis: { id: string; computable: 1234 } | null;
-            }
-        ))[]
-      >
-    >(true);
+    type actual = $infer<typeof q>[number];
+    type expected = {
+      id: string;
+    } & (
+      | { __typename: "default::Hero"; secret_identity: string | null }
+      | {
+          __typename: "default::Villain";
+          nemesis: { id: string; computable: 1234 } | null;
+        }
+    );
+    tc.assert<tc.IsExact<actual, expected>>(true);
   });
 
   test("parent type props in polymorphic", () => {
diff --git a/packages/generate/src/syntax/typesystem.ts b/packages/generate/src/syntax/typesystem.ts
index e02b34be7..791d0574d 100644
--- a/packages/generate/src/syntax/typesystem.ts
+++ b/packages/generate/src/syntax/typesystem.ts
@@ -12,6 +12,11 @@ import type { cardutil } from "./cardinality";
 import type { Range, MultiRange } from "edgedb";
 import type { $Shape, normaliseShape } from "./select";
 
+// eslint-disable-next-line @typescript-eslint/no-empty-object-type
+export interface TypesystemOptions {
+  // polymorphismAsDistributedUnions: boolean;
+}
+
 //////////////////
 // BASETYPE
 //////////////////
@@ -372,6 +377,51 @@ export type $expr_PolyShapeElement<
 };
 
 export type computeObjectShape<
+  Pointers extends ObjectTypePointers,
+  Shape,
+  TypeName extends string = string,
+> = TypesystemOptions extends { polymorphismAsDistributedUnions: true }
+  ? computeObjectShapeNew<Pointers, Shape, TypeName>
+  : computeObjectShapeLegacy<Pointers, Shape, TypeName>;
+
+type computeObjectShapeLegacy<
+  Pointers extends ObjectTypePointers,
+  Shape,
+  TypeName extends string,
+> = typeutil.flatten<
+  keyof Shape extends never
+    ? { id: string }
+    : typeutil.stripNever<{
+        [k in keyof Shape]: Shape[k] extends $expr_PolyShapeElement<
+          infer PolyType,
+          infer ShapeEl
+        >
+          ? [k] extends [keyof PolyType["__element__"]["__pointers__"]]
+            ? shapeElementToTs<
+                PolyType["__element__"]["__pointers__"][k],
+                ShapeEl,
+                k extends "__type__" ? TypeName : null
+              > | null
+            : never
+          : 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],
+                  k extends "__type__" ? TypeName : null
+                >
+              : never;
+      }>
+>;
+
+type computeObjectShapeNew<
   Pointers extends ObjectTypePointers,
   Shape,
   TypeName extends string,