allow output lengths to be set

chaiNNer-org · Jul 1, 2024 · 8f45730 · 8f45730
1 parent 7c20a13
commit 8f45730
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 8 deletions.
diff --git a/src/common/nodes/TypeState.ts b/src/common/nodes/TypeState.ts
@@ -65,7 +65,7 @@ export class TypeState {
     static create(
         nodesMap: ReadonlyMap<string, Node<NodeData>>,
         rawEdges: readonly Edge<EdgeData>[],
-        outputNarrowing: ReadonlyMap<string, ReadonlyMap<OutputId, Type>>,
+        outputNarrowing: ReadonlyMap<string, ReadonlyMap<OutputId | 'length', Type>>,
         functionDefinitions: ReadonlyMap<SchemaId, FunctionDefinition>,
         passthrough?: PassthroughMap,
         previousTypeState?: TypeState

diff --git a/src/common/types/function.ts b/src/common/types/function.ts
@@ -591,7 +591,7 @@ export class FunctionInstance {
     static fromPartialInputs(
         definition: FunctionDefinition,
         partialInputs: (inputId: InputId) => NonNeverType | undefined,
-        outputNarrowing: ReadonlyMap<OutputId, Type> = EMPTY_MAP,
+        outputNarrowing: ReadonlyMap<OutputId | 'length', Type> = EMPTY_MAP,
         passthrough?: PassthroughInfo
     ): FunctionInstance {
         const inputErrors: FunctionInputAssignmentError[] = [];
@@ -705,9 +705,31 @@ export class FunctionInstance {
 
                 if (item.type === 'Output') {
                     outputs.set(item.output.id, type);
+                    scope.assignParameter(getOutputParamName(item.output.id), type);
                 } else {
                     for (const id of item.iterOutput.outputs) {
-                        outputLengths.set(id, type);
+                        if (item.iterOutput.outputs.includes(id)) {
+                            const predeterminedLength = outputNarrowing.get('length');
+                            if (predeterminedLength && predeterminedLength.type !== 'never') {
+                                const sequenceType = evaluate(
+                                    fromJson(
+                                        `Sequence { length: ${predeterminedLength.toString()} }`
+                                    ),
+                                    scope
+                                );
+                                if (sequenceType.type !== 'never') {
+                                    outputLengths.set(id, sequenceType);
+                                }
+                            } else {
+                                const lengthType = evaluate(
+                                    fromJson(item.iterOutput.sequenceType),
+                                    scope
+                                );
+                                if (lengthType.type !== 'never') {
+                                    outputLengths.set(id, lengthType);
+                                }
+                            }
+                        }
                     }
                 }
                 scope.assignParameter(item.param, type);

diff --git a/src/renderer/contexts/ExecutionContext.tsx b/src/renderer/contexts/ExecutionContext.tsx
@@ -262,6 +262,7 @@ export const ExecutionProvider = memo(({ children }: React.PropsWithChildren<{}>
         let broadcastData;
         let types;
         let progress;
+        let expectedLength;
 
         for (const { type, data } of events) {
             if (type === 'node-start') {
@@ -272,6 +273,7 @@ export const ExecutionProvider = memo(({ children }: React.PropsWithChildren<{}>
             } else if (type === 'node-broadcast') {
                 broadcastData = data.data;
                 types = data.types;
+                expectedLength = data.expectedLength;
             } else {
                 progress = data;
             }
@@ -313,6 +315,11 @@ export const ExecutionProvider = memo(({ children }: React.PropsWithChildren<{}>
                 }
             }
         }
+
+        if (expectedLength) {
+            const type = evaluate(fromJson(expectedLength), getChainnerScope());
+            setManualOutputType(nodeId, 'length', type);
+        }
     };
     const nodeEventBacklog = useEventBacklog({
         process: (events: NodeEvents[]) => {

diff --git a/src/renderer/contexts/GlobalNodeState.tsx b/src/renderer/contexts/GlobalNodeState.tsx
@@ -137,7 +137,11 @@ interface Global {
     setZoom: SetState<number>;
     exportViewportScreenshot: () => void;
     exportViewportScreenshotToClipboard: () => void;
-    setManualOutputType: (nodeId: string, outputId: OutputId, type: Expression | undefined) => void;
+    setManualOutputType: (
+        nodeId: string,
+        outputId: OutputId | 'length',
+        type: Expression | undefined
+    ) => void;
     clearManualOutputTypes: (nodes: Iterable<string>) => void;
     typeStateRef: Readonly<React.MutableRefObject<TypeState>>;
     chainLineageRef: Readonly<React.MutableRefObject<ChainLineage>>;
@@ -204,10 +208,14 @@ export const GlobalProvider = memo(
         );
 
         const [manualOutputTypes, setManualOutputTypes] = useState(() => ({
-            map: new Map<string, Map<OutputId, Type>>(),
+            map: new Map<string, Map<OutputId | 'length', Type>>(),
         }));
         const setManualOutputType = useCallback(
-            (nodeId: string, outputId: OutputId, expression: Expression | undefined): void => {
+            (
+                nodeId: string,
+                outputId: OutputId | 'length',
+                expression: Expression | undefined
+            ): void => {
                 const getType = () => {
                     if (expression === undefined) {
                         return undefined;
@@ -262,7 +270,7 @@ export const GlobalProvider = memo(
                 // remove manual overrides of nodes that no longer exist
                 if (manualOutputTypes.map.size > 0) {
                     const ids = [...manualOutputTypes.map.keys()];
-                    for (const id of ids.filter((key) => !nodeMap.has(key))) {
+                    for (const id of ids.filter((key) => !nodeMap.has(key) && key !== 'length')) {
                         // use interior mutability to not cause updates
                         manualOutputTypes.map.delete(id);
                     }

diff --git a/src/renderer/helpers/nodeState.ts b/src/renderer/helpers/nodeState.ts
@@ -1,3 +1,4 @@
+import { NonNeverType } from '@chainner/navi';
 import { useMemo } from 'react';
 import { useContext, useContextSelector } from 'use-context-selector';
 import {
@@ -24,13 +25,14 @@ import { useMemoObject } from '../hooks/useMemo';
 export interface TypeInfo {
     readonly instance: FunctionInstance | undefined;
     readonly connectedInputs: ReadonlySet<InputId>;
+    readonly iteratedInputLengths?: ReadonlyMap<InputId, NonNeverType>;
+    readonly iteratedOutputLengths?: ReadonlyMap<OutputId, NonNeverType>;
 }
 
 const useTypeInfo = (id: string): TypeInfo => {
     const instance = useContextSelector(GlobalVolatileContext, (c) =>
         c.typeState.functions.get(id)
     );
-
     const connectedInputsString = useContextSelector(GlobalVolatileContext, (c) => {
         const connected = c.typeState.edges.byTarget.get(id);
         return IdSet.from(connected?.map((connection) => connection.inputId) ?? EMPTY_ARRAY);