From c4bc0e84305de6bb0e5f00bc4ad9b47eb708d723 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Tue, 20 Aug 2024 11:01:05 -0400
Subject: [PATCH] Auto expand tree on search in node library tab (#558)

* Add custom nodelib searchbox

* Auto expand on search

* Support alphabetical sort in filtered tree
---
 .../SearchBox.vue}                            | 34 +++++++-----
 .../dialog/content/SettingDialogContent.vue   |  5 +-
 .../sidebar/tabs/NodeLibrarySidebarTab.vue    | 54 ++++++++++++++++---
 src/i18n.ts                                   |  2 +
 src/stores/nodeDefStore.ts                    | 24 ---------
 src/utils/treeUtil.ts                         | 21 ++++++++
 6 files changed, 96 insertions(+), 44 deletions(-)
 rename src/components/{dialog/content/setting/SettingSearchBox.vue => common/SearchBox.vue} (58%)

diff --git a/src/components/dialog/content/setting/SettingSearchBox.vue b/src/components/common/SearchBox.vue
similarity index 58%
rename from src/components/dialog/content/setting/SettingSearchBox.vue
rename to src/components/common/SearchBox.vue
index e5420db36..bc3bb93cf 100644
--- a/src/components/dialog/content/setting/SettingSearchBox.vue
+++ b/src/components/common/SearchBox.vue
@@ -1,35 +1,45 @@
 <template>
   <IconField :class="props.class">
-    <InputIcon class="pi pi-search" />
+    <InputIcon :class="props.icon" />
     <InputText
       class="search-box-input"
       @input="handleInput"
       :modelValue="props.modelValue"
-      :placeholder="$t('searchSettings') + '...'"
+      :placeholder="props.placeholder"
     />
   </IconField>
 </template>
 
 <script setup lang="ts">
+import { debounce } from 'lodash'
 import IconField from 'primevue/iconfield'
 import InputIcon from 'primevue/inputicon'
 import InputText from 'primevue/inputtext'
-import { debounce } from 'lodash'
 
-const props = defineProps<{
+interface Props {
   class?: string
   modelValue: string
-}>()
-const emit = defineEmits(['update:modelValue', 'search'])
-const emitSearch = debounce((event: KeyboardEvent) => {
-  const target = event.target as HTMLInputElement
-  emit('search', target.value)
-}, 300)
+  placeholder?: string
+  icon?: string
+  debounceTime?: number
+}
 
-const handleInput = (event: KeyboardEvent) => {
+const props = withDefaults(defineProps<Props>(), {
+  placeholder: 'Search...',
+  icon: 'pi pi-search',
+  debounceTime: 300
+})
+
+const emit = defineEmits(['update:modelValue', 'search'])
+
+const emitSearch = debounce((value: string) => {
+  emit('search', value)
+}, props.debounceTime)
+
+const handleInput = (event: Event) => {
   const target = event.target as HTMLInputElement
   emit('update:modelValue', target.value)
-  emitSearch(event)
+  emitSearch(target.value)
 }
 </script>
 
diff --git a/src/components/dialog/content/SettingDialogContent.vue b/src/components/dialog/content/SettingDialogContent.vue
index a1441401c..90ab80009 100644
--- a/src/components/dialog/content/SettingDialogContent.vue
+++ b/src/components/dialog/content/SettingDialogContent.vue
@@ -1,10 +1,11 @@
 <template>
   <div class="settings-container">
     <div class="settings-sidebar">
-      <SettingSearchBox
+      <SearchBox
         class="settings-search-box"
         v-model:modelValue="searchQuery"
         @search="handleSearch"
+        :placeholder="$t('searchSettings') + '...'"
       />
       <Listbox
         v-model="activeCategory"
@@ -66,7 +67,7 @@ import Divider from 'primevue/divider'
 import { SettingTreeNode, useSettingStore } from '@/stores/settingStore'
 import { SettingParams } from '@/types/settingTypes'
 import SettingGroup from './setting/SettingGroup.vue'
-import SettingSearchBox from './setting/SettingSearchBox.vue'
+import SearchBox from '@/components/common/SearchBox.vue'
 import NoResultsPlaceholder from '@/components/common/NoResultsPlaceholder.vue'
 import { flattenTree } from '@/utils/treeUtil'
 
diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index 9be23a6c6..598ba85da 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -14,13 +14,17 @@
       </ToggleButton>
     </template>
     <template #body>
+      <SearchBox
+        class="node-lib-search-box"
+        v-model:modelValue="searchQuery"
+        @search="handleSearch"
+        :placeholder="$t('searchNodes') + '...'"
+      />
       <TreePlus
         class="node-lib-tree"
         v-model:expandedKeys="expandedKeys"
         selectionMode="single"
         :value="renderedRoot.children"
-        :filter="true"
-        filterMode="lenient"
         dragSelector=".p-tree-node-leaf"
         :pt="{
           nodeLabel: 'node-lib-tree-node-label',
@@ -85,10 +89,11 @@ import { computed, ref, nextTick } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 import TreePlus from '@/components/primevueOverride/TreePlus.vue'
 import NodePreview from '@/components/node/NodePreview.vue'
+import SearchBox from '@/components/common/SearchBox.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import { useSettingStore } from '@/stores/settingStore'
 import { app } from '@/scripts/app'
-import { ComfyNodeDef } from '@/types/apiTypes'
+import { buildTree, sortedTree } from '@/utils/treeUtil'
 
 const nodeDefStore = useNodeDefStore()
 const alphabeticalSort = ref(false)
@@ -101,6 +106,7 @@ const hoveredComfyNode = computed<ComfyNodeDefImpl | null>(() => {
   return nodeDefStore.nodeDefsByName[hoveredComfyNodeName.value] || null
 })
 const previewRef = ref<InstanceType<typeof NodePreview> | null>(null)
+const searchQuery = ref<string>('')
 
 const settingStore = useSettingStore()
 const sidebarLocation = computed<'left' | 'right'>(() =>
@@ -112,9 +118,10 @@ const nodePreviewStyle = ref<Record<string, string>>({
   left: '0px'
 })
 
-const root = computed(() =>
-  alphabeticalSort.value ? nodeDefStore.sortedNodeTree : nodeDefStore.nodeTree
-)
+const root = computed(() => {
+  const root = filteredRoot.value || nodeDefStore.nodeTree
+  return alphabeticalSort.value ? sortedTree(root) : root
+})
 const renderedRoot = computed(() => {
   return fillNodeInfo(root.value)
 })
@@ -176,6 +183,35 @@ const toggleNode = (id: string) => {
 const insertNode = (nodeDef: ComfyNodeDefImpl) => {
   app.addNodeOnGraph(nodeDef, { pos: app.getCanvasCenter() })
 }
+
+const filteredRoot = ref<TreeNode | null>(null)
+const handleSearch = (query: string) => {
+  if (query.length < 3) {
+    filteredRoot.value = null
+    expandedKeys.value = {}
+    return
+  }
+
+  const matchedNodes = nodeDefStore.nodeSearchService.searchNode(query, [], {
+    limit: 64
+  })
+
+  filteredRoot.value = buildTree(matchedNodes, (nodeDef: ComfyNodeDefImpl) => [
+    ...nodeDef.category.split('/'),
+    nodeDef.display_name
+  ])
+  expandNode(filteredRoot.value)
+}
+
+const expandNode = (node: TreeNode) => {
+  if (node.children && node.children.length) {
+    expandedKeys.value[node.key] = true
+
+    for (let child of node.children) {
+      expandNode(child)
+    }
+  }
+}
 </script>
 
 <style>
@@ -185,3 +221,9 @@ const insertNode = (nodeDef: ComfyNodeDefImpl) => {
   margin-left: var(--p-tree-node-gap);
 }
 </style>
+
+<style scoped>
+:deep(.node-lib-search-box) {
+  @apply mx-4 mt-4;
+}
+</style>
diff --git a/src/i18n.ts b/src/i18n.ts
index 5c7638a16..c3e4cb364 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -10,6 +10,7 @@ const messages = {
     loadWorkflow: 'Load Workflow',
     settings: 'Settings',
     searchSettings: 'Search Settings',
+    searchNodes: 'Search Nodes',
     noResultsFound: 'No Results Found',
     searchFailedMessage:
       "We couldn't find any settings matching your search. Try adjusting your search terms.",
@@ -35,6 +36,7 @@ const messages = {
     loadWorkflow: '加载工作流',
     settings: '设置',
     searchSettings: '搜索设置',
+    searchNodes: '搜索节点',
     noResultsFound: '未找到结果',
     noTasksFound: '未找到任务',
     noTasksFoundMessage: '队列中没有任务。',
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 6ca8de839..bbac6e043 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -226,27 +226,6 @@ export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
   }
 }
 
-function sortedTree(node: TreeNode): TreeNode {
-  // Create a new node with the same label and data
-  const newNode: TreeNode = {
-    ...node
-  }
-
-  if (node.children) {
-    // Sort the children of the current node
-    const sortedChildren = [...node.children].sort((a, b) =>
-      a.label.localeCompare(b.label)
-    )
-    // Recursively sort the children and add them to the new node
-    newNode.children = []
-    for (const child of sortedChildren) {
-      newNode.children.push(sortedTree(child))
-    }
-  }
-
-  return newNode
-}
-
 interface State {
   nodeDefsByName: Record<string, ComfyNodeDefImpl>
   widgets: Record<string, ComfyWidgetConstructor>
@@ -269,9 +248,6 @@ export const useNodeDefStore = defineStore('nodeDef', {
         ...nodeDef.category.split('/'),
         nodeDef.display_name
       ])
-    },
-    sortedNodeTree(): TreeNode {
-      return sortedTree(this.nodeTree)
     }
   },
   actions: {
diff --git a/src/utils/treeUtil.ts b/src/utils/treeUtil.ts
index 0bf54b7d3..3269a1ff7 100644
--- a/src/utils/treeUtil.ts
+++ b/src/utils/treeUtil.ts
@@ -47,3 +47,24 @@ export function flattenTree<T>(tree: TreeNode): T[] {
   }
   return result
 }
+
+export function sortedTree(node: TreeNode): TreeNode {
+  // Create a new node with the same label and data
+  const newNode: TreeNode = {
+    ...node
+  }
+
+  if (node.children) {
+    // Sort the children of the current node
+    const sortedChildren = [...node.children].sort((a, b) =>
+      a.label.localeCompare(b.label)
+    )
+    // Recursively sort the children and add them to the new node
+    newNode.children = []
+    for (const child of sortedChildren) {
+      newNode.children.push(sortedTree(child))
+    }
+  }
+
+  return newNode
+}