From 95dc6ff5dee23c7212dc1ea50c5102484251ab66 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Fri, 9 Aug 2024 17:53:42 -0400
Subject: [PATCH] Refactor nodeDefStore.nodeTree (#351)

---
 src/stores/nodeDefStore.ts | 34 ++++++----------------------------
 1 file changed, 6 insertions(+), 28 deletions(-)

diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index dda58ca58..c754ba021 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -4,6 +4,7 @@ import { defineStore } from 'pinia'
 import { Type, Transform, plainToClass } from 'class-transformer'
 import { ComfyWidgetConstructor } from '@/scripts/widgets'
 import { TreeNode } from 'primevue/treenode'
+import { buildTree } from '@/utils/treeUtil'
 
 export class BaseInputSpec<T = any> {
   name: string
@@ -262,34 +263,11 @@ export const useNodeDefStore = defineStore('nodeDef', {
     nodeSearchService(state) {
       return new NodeSearchService(Object.values(state.nodeDefsByName))
     },
-    nodeTree(state): TreeNode {
-      const root: TreeNode = {
-        key: 'root',
-        label: 'Nodes',
-        leaf: false,
-        children: []
-      }
-      for (const nodeDef of Object.values(state.nodeDefsByName)) {
-        const path = nodeDef.category.split('/')
-        let current = root
-        let key = 'root'
-        for (const part of path) {
-          key += `/${part}`
-          let next = current.children.find((child) => child.label === part)
-          if (!next) {
-            next = { key, label: part, children: [], leaf: false }
-            current.children.push(next)
-          }
-          current = next
-        }
-        current.children.push({
-          label: nodeDef.display_name,
-          data: nodeDef,
-          key: `${key}/${nodeDef.name}`,
-          leaf: true
-        })
-      }
-      return root
+    nodeTree(): TreeNode {
+      return buildTree(this.nodeDefs, (nodeDef: ComfyNodeDefImpl) => [
+        ...nodeDef.category.split('/'),
+        nodeDef.display_name
+      ])
     },
     sortedNodeTree(): TreeNode {
       return sortedTree(this.nodeTree)