From 319a2cff046446eb1e8568c3ce197a447cee05dc Mon Sep 17 00:00:00 2001
From: Carson Full <carsonfull@gmail.com>
Date: Tue, 3 Sep 2024 18:34:38 -0500
Subject: [PATCH] Fix ancestor types missing their concrete descendant type
 names

---
 integration-tests/lts/dbschema/default.esdl   |  6 +++-
 .../dbschema/migrations/00027-m173dds.edgeql  |  8 +++++
 integration-tests/lts/select.test.ts          | 22 +++++++++++-
 .../src/edgeql-js/generateObjectTypes.ts      | 35 +++++++++++++------
 4 files changed, 58 insertions(+), 13 deletions(-)
 create mode 100644 integration-tests/lts/dbschema/migrations/00027-m173dds.edgeql

diff --git a/integration-tests/lts/dbschema/default.esdl b/integration-tests/lts/dbschema/default.esdl
index ae8c30b0f..83e7da7d0 100644
--- a/integration-tests/lts/dbschema/default.esdl
+++ b/integration-tests/lts/dbschema/default.esdl
@@ -26,7 +26,11 @@ module default {
     property character_name -> str;
   }
 
-  abstract type Person {
+  abstract type LivingThing {
+    age: int32;
+  }
+
+  abstract type Person extending LivingThing {
     required property name -> str {
       constraint exclusive;
     };
diff --git a/integration-tests/lts/dbschema/migrations/00027-m173dds.edgeql b/integration-tests/lts/dbschema/migrations/00027-m173dds.edgeql
new file mode 100644
index 000000000..1b2dad03c
--- /dev/null
+++ b/integration-tests/lts/dbschema/migrations/00027-m173dds.edgeql
@@ -0,0 +1,8 @@
+CREATE MIGRATION m173ddshjvgy5ampp7rzi2g7cwwdkaljvlifemr5vugnkimtbpp6ca
+    ONTO m1wb2dgjeppqex272zwvqnsdfzdvvppub4iwa5vaxu3xxigyjlruka
+{
+  CREATE ABSTRACT TYPE default::LivingThing {
+      CREATE PROPERTY age: std::int32;
+  };
+  ALTER TYPE default::Person EXTENDING default::LivingThing LAST;
+};
diff --git a/integration-tests/lts/select.test.ts b/integration-tests/lts/select.test.ts
index 3f184bcfd..dda4f4163 100644
--- a/integration-tests/lts/select.test.ts
+++ b/integration-tests/lts/select.test.ts
@@ -348,16 +348,18 @@ describe("select", () => {
       ...e.is(e.Hero, e.Hero["*"]),
       name: true,
     }));
+    type result = $infer<typeof q>;
 
     // 'id' is filtered out since it is not valid in a polymorphic expr
     tc.assert<
       tc.IsExact<
-        $infer<typeof q>,
+        result,
         ({ name: string } & (
           | { __typename: "default::Villain" }
           | {
               __typename: "default::Hero";
               height: string | null;
+              age: number | null;
               isAdult: boolean | null;
               number_of_movies: number | null;
               secret_identity: string | null;
@@ -374,6 +376,24 @@ describe("select", () => {
     tc.assert<tc.IsExact<typeof name, "default::Hero">>(true);
   });
 
+  test("polymorphic type names", () => {
+    tc.assert<
+      tc.IsExact<
+        typeof e.LivingThing.__element__.__polyTypenames__,
+        "default::Hero" | "default::Villain"
+      >
+    >(true);
+    tc.assert<
+      tc.IsExact<
+        typeof e.Person.__element__.__polyTypenames__,
+        "default::Hero" | "default::Villain"
+      >
+    >(true);
+    tc.assert<
+      tc.IsExact<typeof e.Hero.__element__.__polyTypenames__, "default::Hero">
+    >(true);
+  });
+
   test("limit/offset inference", () => {
     const testSet = e.set(1, 2, 3);
 
diff --git a/packages/generate/src/edgeql-js/generateObjectTypes.ts b/packages/generate/src/edgeql-js/generateObjectTypes.ts
index 815554ca3..957bbeeb9 100644
--- a/packages/generate/src/edgeql-js/generateObjectTypes.ts
+++ b/packages/generate/src/edgeql-js/generateObjectTypes.ts
@@ -175,17 +175,7 @@ export const getStringRepresentation: (
 export const generateObjectTypes = (params: GeneratorParams) => {
   const { dir, types } = params;
 
-  const descendents = new Map<string, Set<string>>();
-  for (const type of types.values()) {
-    if (type.kind === "object" && !type.is_abstract) {
-      for (const base of type.bases) {
-        if (!descendents.has(base.id)) {
-          descendents.set(base.id, new Set());
-        }
-        descendents.get(base.id)!.add(type.name);
-      }
-    }
-  }
+  const descendents = generatePolyTypenames(types);
 
   for (const type of types.values()) {
     if (type.kind !== "object") {
@@ -393,3 +383,26 @@ export const generateObjectTypes = (params: GeneratorParams) => {
     body.addToDefaultExport(literal, name);
   }
 };
+
+function generatePolyTypenames(types: $.introspect.Types) {
+  const descendents = new Map<string, Set<string>>();
+
+  const visit = (current: $.introspect.Type, descendent: $.introspect.Type) => {
+    if (current.kind !== "object") {
+      return;
+    }
+    for (const base of current.bases) {
+      if (!descendents.has(base.id)) {
+        descendents.set(base.id, new Set());
+      }
+      descendents.get(base.id)!.add(descendent.name);
+      visit(types.get(base.id), descendent);
+    }
+  };
+  for (const type of types.values()) {
+    if (type.kind === "object" && !type.is_abstract) {
+      visit(type, type);
+    }
+  }
+  return descendents;
+}