From a519a6f8ea0c4a4495b71f8ff8aab2d7bab62dcf Mon Sep 17 00:00:00 2001
From: Weixuan Fu <weixuanf@outlook.com>
Date: Thu, 21 Mar 2024 19:27:47 +0800
Subject: [PATCH] better strcuture cod

---
 ui/src/gallery/components/MetaBox/MetaBox.tsx |  1 +
 ui/src/gallery/components/MetaBox/utils.ts    | 37 +++++++++++--------
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/ui/src/gallery/components/MetaBox/MetaBox.tsx b/ui/src/gallery/components/MetaBox/MetaBox.tsx
index 4d0b0c1..c5d412c 100644
--- a/ui/src/gallery/components/MetaBox/MetaBox.tsx
+++ b/ui/src/gallery/components/MetaBox/MetaBox.tsx
@@ -67,6 +67,7 @@ export default function MetaBox({
     if (!imagePrompt) return;
     console.log("imagePrompt", imagePrompt);
     const calcInput = calcInputListRecursive(imagePrompt);
+    console.log("✌️calcInput", calcInput);
     setCalcInputList(calcInput);
   }, [imagePrompt]);
 
diff --git a/ui/src/gallery/components/MetaBox/utils.ts b/ui/src/gallery/components/MetaBox/utils.ts
index 8de596e..c72cb8f 100644
--- a/ui/src/gallery/components/MetaBox/utils.ts
+++ b/ui/src/gallery/components/MetaBox/utils.ts
@@ -91,7 +91,6 @@ export type InputResultItem = {
   formLabel?: string;
 };
 
-let toExecute: ImagePromptNodeItem[] = [];
 export type PromptNodeInputItem = {
   classType: string; // nodeType
   inputName: string;
@@ -99,24 +98,19 @@ export type PromptNodeInputItem = {
   nodeID: string;
   children: string[]; // output links to other node e.g. CLIPTextEncode.text -> KSampler.negative
 };
-const inputList: PromptNodeInputItem[] = [];
+
+let inputList: PromptNodeInputItem[] = [];
+let visitedIDs = new Set<string>();
 
 function dfs(promptNode: ImagePromptNodeItem, prompt: ImagePrompt) {
-  const index = toExecute.findIndex((v) => v.nodeID === promptNode.nodeID);
-  if (index == -1) {
-    return [];
+  if (visitedIDs.has(promptNode.nodeID!)) {
+    return;
   }
-  toExecute = toExecute.filter((v) => v.nodeID !== promptNode.nodeID);
-
+  visitedIDs.add(promptNode.nodeID!);
   Object.entries(promptNode.inputs).forEach(([inputName, value]) => {
     if (Array.isArray(value)) {
       const parentID = value[0];
-      const parentNode = prompt[parentID];
-      if (!parentNode.children) {
-        parentNode.children = [];
-      }
-      parentNode.children.push(inputName);
-      dfs(parentNode, prompt);
+      dfs(prompt[parentID], prompt);
     }
   });
   Object.entries(promptNode.inputs).forEach(([inputName, value]) => {
@@ -136,13 +130,24 @@ function dfs(promptNode: ImagePromptNodeItem, prompt: ImagePrompt) {
 export function calcInputListRecursive(
   prompt: ImagePrompt,
 ): PromptNodeInputItem[] {
-  toExecute = Object.values(prompt);
+  inputList = []; // clear result inputlist
   for (const key of Object.keys(prompt)) {
     prompt[key].nodeID = key;
+    Object.entries(prompt[key].inputs).forEach(([inputName, value]) => {
+      if (Array.isArray(value)) {
+        const parentID = value[0];
+        const parentNode = prompt[parentID];
+        if (!parentNode.children) {
+          parentNode.children = [];
+        }
+        // record children output links
+        parentNode.children.push(inputName);
+      }
+    });
   }
 
-  while (toExecute.length) {
-    dfs(toExecute[0], prompt);
+  for (const key of Object.keys(prompt)) {
+    dfs(prompt[key], prompt);
   }
   console.log("calcInputListRecursive res", inputList);
   return inputList;