From f4958c9a721221a42de5c1c73793f05c05d52674 Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Fri, 26 Jul 2024 16:49:54 +0200
Subject: [PATCH] frontend-plugin-api: update createExtension to use
 union-based verification of factory outputs

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>
---
 .../src/wiring/createExtension.test.ts        | 112 ++++++++++--------
 .../src/wiring/createExtension.ts             |  49 ++++----
 2 files changed, 86 insertions(+), 75 deletions(-)

diff --git a/packages/frontend-plugin-api/src/wiring/createExtension.test.ts b/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
index febc7ec133..bbf2e539d1 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtension.test.ts
@@ -357,7 +357,68 @@ describe('createExtension', () => {
     }).toThrow("Missing required value at 'foo'");
   });
 
-  it('should new form of inputs and outputs', () => {
+  it('should support new form of outputs', () => {
+    // @ts-expect-error
+    createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      output: [stringDataRef, numberDataRef],
+      factory() {
+        return []; // Missing all outputs
+      },
+    });
+
+    // @ts-expect-error
+    createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      output: [stringDataRef, numberDataRef],
+      factory() {
+        return [stringDataRef('hello')]; // Missing number output
+      },
+    });
+
+    // Duplicate output, we won't attempt to handle this a compile time and instead error out at runtime
+    createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      output: [stringDataRef],
+      factory() {
+        return [stringDataRef('hello'), stringDataRef('hello')];
+      },
+    });
+
+    createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      output: [stringDataRef, numberDataRef],
+      factory() {
+        return [stringDataRef('hello'), numberDataRef(4)];
+      },
+    });
+
+    createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      output: [stringDataRef, numberDataRef.optional()],
+      factory() {
+        return [stringDataRef('hello'), numberDataRef(4)];
+      },
+    });
+
+    createExtension({
+      namespace: 'test',
+      attachTo: { id: 'root', input: 'default' },
+      output: [stringDataRef, numberDataRef.optional()],
+      factory() {
+        return [stringDataRef('hello')]; // Missing number output, but it's optional so that's allowed
+      },
+    });
+
+    expect(true).toBe(true);
+  });
+
+  it('should support new form of inputs', () => {
     createExtension({
       namespace: 'test',
       attachTo: { id: 'root', input: 'default' },
@@ -389,55 +450,6 @@ describe('createExtension', () => {
       },
     });
 
-    // Double output
-    createExtension({
-      namespace: 'test',
-      attachTo: { id: 'root', input: 'default' },
-      output: [stringDataRef],
-      // @ts-expect-error
-      factory() {
-        return [stringDataRef('hello'), stringDataRef('hello')];
-      },
-    });
-
-    // Missing output
-    createExtension({
-      namespace: 'test',
-      attachTo: { id: 'root', input: 'default' },
-      output: [stringDataRef, numberDataRef],
-      // @ts-expect-error
-      factory() {
-        return [stringDataRef('hello')];
-      },
-    });
-
-    createExtension({
-      namespace: 'test',
-      attachTo: { id: 'root', input: 'default' },
-      output: [stringDataRef, numberDataRef],
-      factory() {
-        return [stringDataRef('hello'), numberDataRef(4)];
-      },
-    });
-
-    createExtension({
-      namespace: 'test',
-      attachTo: { id: 'root', input: 'default' },
-      output: [stringDataRef, numberDataRef.optional()],
-      factory() {
-        return [stringDataRef('hello'), numberDataRef(4)];
-      },
-    });
-
-    createExtension({
-      namespace: 'test',
-      attachTo: { id: 'root', input: 'default' },
-      output: [stringDataRef, numberDataRef.optional()],
-      factory() {
-        return [stringDataRef('hello')];
-      },
-    });
-
     expect(true).toBe(true);
   });
 });
diff --git a/packages/frontend-plugin-api/src/wiring/createExtension.ts b/packages/frontend-plugin-api/src/wiring/createExtension.ts
index c514750a84..cfce13d494 100644
--- a/packages/frontend-plugin-api/src/wiring/createExtension.ts
+++ b/packages/frontend-plugin-api/src/wiring/createExtension.ts
@@ -143,27 +143,6 @@ export interface LegacyCreateExtensionOptions<
   }): Expand<ExtensionDataValues<TOutput>>;
 }
 
-// TODO(Rugvip): This creates a permutation of all possible output tuples,
-// taking the optionality into account. It's not optimal, since it might hurt
-// performance for large outputs, doesn't provide a very clear error message,
-// and doesn't allow us to refactor to allow for a generator.
-/** @public */
-type ExtensionFactoryOutput<
-  URef extends AnyExtensionDataRef,
-  TPickRef extends AnyExtensionDataRef = URef,
-> = [URef] extends [never]
-  ? []
-  : TPickRef extends ExtensionDataRef<infer IData, infer IId, infer IConfig>
-  ?
-      | (IConfig['optional'] extends true
-          ? ExtensionFactoryOutput<Exclude<URef, TPickRef>>
-          : never)
-      | [
-          ExtensionDataValue<IData, IId>,
-          ...ExtensionFactoryOutput<Exclude<URef, TPickRef>>,
-        ]
-  : never;
-
 /** @public */
 export type CreateExtensionOptions<
   UOutput extends AnyExtensionDataRef,
@@ -171,6 +150,7 @@ export type CreateExtensionOptions<
   TConfig,
   TConfigInput,
   TConfigSchema extends { [key: string]: (zImpl: typeof z) => z.ZodType },
+  UFactoryOutput extends ExtensionDataValue<any, any>,
 > = {
   kind?: string;
   namespace?: string;
@@ -195,8 +175,23 @@ export type CreateExtensionOptions<
             >;
           });
     inputs: Expand<ResolvedExtensionInputs<TInputs>>;
-  }): ExtensionFactoryOutput<UOutput>;
-};
+  }): Iterable<UFactoryOutput>;
+} & ((
+  UOutput extends any
+    ? UOutput['config']['optional'] extends true
+      ? never
+      : UOutput['id']
+    : never
+) extends infer IRequiredOutputIds
+  ? [IRequiredOutputIds] extends [UFactoryOutput['id']]
+    ? {}
+    : {
+        'Error: The extension factory is missing the following outputs': Exclude<
+          IRequiredOutputIds,
+          UFactoryOutput['id']
+        >;
+      }
+  : never);
 
 /** @public */
 export interface ExtensionDefinition<TConfig, TConfigInput = TConfig> {
@@ -256,13 +251,15 @@ export function createExtension<
   TConfig,
   TConfigInput,
   TConfigSchema extends { [key: string]: (zImpl: typeof z) => z.ZodType },
+  UFactoryOutput extends ExtensionDataValue<any, any>,
 >(
   options: CreateExtensionOptions<
     UOutput,
     TInputs,
     TConfig,
     TConfigInput,
-    TConfigSchema
+    TConfigSchema,
+    UFactoryOutput
   >,
 ): ExtensionDefinition<
   TConfig &
@@ -320,6 +317,7 @@ export function createExtension<
   TConfig,
   TConfigInput,
   TConfigSchema extends { [key: string]: (zImpl: typeof z) => z.ZodType },
+  UFactoryOutput extends ExtensionDataValue<any, any>,
 >(
   options:
     | CreateExtensionOptions<
@@ -327,7 +325,8 @@ export function createExtension<
         TInputs,
         TConfig,
         TConfigInput,
-        TConfigSchema
+        TConfigSchema,
+        UFactoryOutput
       >
     | LegacyCreateExtensionOptions<
         AnyExtensionDataMap,