From 7d75cc99baaf2ffdb88850fe8c5c275456713979 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Mon, 29 Jul 2024 12:39:54 -0400
Subject: [PATCH] Add sort button in node library sidebar tab (#259)

* Add sort button on node library

* tab template
---
 .../sidebar/tabs/NodeLibrarySideBarTab.vue    | 116 +++++++++++-------
 .../sidebar/tabs/SideBarTabTemplate.vue       |  65 ++++++++++
 src/i18n.ts                                   |  10 +-
 src/stores/nodeDefStore.ts                    |  24 ++++
 4 files changed, 166 insertions(+), 49 deletions(-)
 create mode 100644 src/components/sidebar/tabs/SideBarTabTemplate.vue

diff --git a/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue b/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue
index 0697cf67b..4925e4945 100644
--- a/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySideBarTab.vue
@@ -1,63 +1,83 @@
 <template>
-  <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',
-      nodeChildren: ({ props }) => ({
-        'data-comfy-node-name': props.node?.data?.name,
-        onMouseenter: (event: MouseEvent) => {
-          hoveredComfyNodeName = props.node?.data?.name
+  <SideBarTabTemplate :title="$t('sideToolBar.nodeLibrary')">
+    <template #tool-buttons>
+      <ToggleButton
+        v-model:model-value="alphabeticalSort"
+        on-icon="pi pi-sort-alpha-down"
+        off-icon="pi pi-sort-alt"
+        aria-label="Sort"
+        :pt="{
+          label: { style: { display: 'none' } }
+        }"
+        v-tooltip="$t('sideToolBar.nodeLibraryTab.sortOrder')"
+      >
+      </ToggleButton>
+    </template>
+    <template #body>
+      <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',
+          nodeChildren: ({ props }) => ({
+            'data-comfy-node-name': props.node?.data?.name,
+            onMouseenter: (event: MouseEvent) => {
+              hoveredComfyNodeName = props.node?.data?.name
 
-          const hoverTarget = event.target as HTMLElement
-          const targetRect = hoverTarget.getBoundingClientRect()
-          nodePreviewStyle.top = `${targetRect.top - 40}px`
-          nodePreviewStyle.left = `${targetRect.right}px`
-        },
-        onMouseleave: () => {
-          hoveredComfyNodeName = null
-        }
-      })
-    }"
-  >
-    <template #folder="{ node }">
-      <span class="folder-label">{{ node.label }}</span>
-      <Badge
-        :value="node.totalNodes"
-        severity="secondary"
-        :style="{ marginLeft: '0.5rem' }"
-      ></Badge>
+              const hoverTarget = event.target as HTMLElement
+              const targetRect = hoverTarget.getBoundingClientRect()
+              nodePreviewStyle.top = `${targetRect.top - 40}px`
+              nodePreviewStyle.left = `${targetRect.right}px`
+            },
+            onMouseleave: () => {
+              hoveredComfyNodeName = null
+            }
+          })
+        }"
+      >
+        <template #folder="{ node }">
+          <span class="folder-label">{{ node.label }}</span>
+          <Badge
+            :value="node.totalNodes"
+            severity="secondary"
+            :style="{ marginLeft: '0.5rem' }"
+          ></Badge>
+        </template>
+        <template #node="{ node }">
+          <span class="node-label">{{ node.label }}</span>
+        </template>
+      </TreePlus>
+      <div
+        v-if="hoveredComfyNode"
+        class="node-lib-node-preview"
+        :style="nodePreviewStyle"
+      >
+        <NodePreview
+          :key="hoveredComfyNode.name"
+          :nodeDef="hoveredComfyNode"
+        ></NodePreview>
+      </div>
     </template>
-    <template #node="{ node }">
-      <span class="node-label">{{ node.label }}</span>
-    </template>
-  </TreePlus>
-  <div
-    v-if="hoveredComfyNode"
-    class="node-lib-node-preview"
-    :style="nodePreviewStyle"
-  >
-    <NodePreview
-      :key="hoveredComfyNode.name"
-      :nodeDef="hoveredComfyNode"
-    ></NodePreview>
-  </div>
+  </SideBarTabTemplate>
 </template>
 
 <script setup lang="ts">
 import Badge from 'primevue/badge'
+import ToggleButton from 'primevue/togglebutton'
 import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
 import { computed, ref } from 'vue'
 import { TreeNode } from 'primevue/treenode'
 import TreePlus from '@/components/primevueOverride/TreePlus.vue'
 import NodePreview from '@/components/NodePreview.vue'
+import SideBarTabTemplate from '@/components/sidebar/tabs/SideBarTabTemplate.vue'
 
 const nodeDefStore = useNodeDefStore()
+const alphabeticalSort = ref(false)
 const expandedKeys = ref({})
 const hoveredComfyNodeName = ref<string | null>(null)
 const hoveredComfyNode = computed<ComfyNodeDefImpl | null>(() => {
@@ -72,7 +92,9 @@ const nodePreviewStyle = ref<Record<string, string>>({
   left: '0px'
 })
 
-const root = computed(() => nodeDefStore.nodeTree)
+const root = computed(() =>
+  alphabeticalSort.value ? nodeDefStore.sortedNodeTree : nodeDefStore.nodeTree
+)
 const renderedRoot = computed(() => {
   return fillNodeInfo(root.value)
 })
diff --git a/src/components/sidebar/tabs/SideBarTabTemplate.vue b/src/components/sidebar/tabs/SideBarTabTemplate.vue
new file mode 100644
index 000000000..ff0450515
--- /dev/null
+++ b/src/components/sidebar/tabs/SideBarTabTemplate.vue
@@ -0,0 +1,65 @@
+<template>
+  <div class="comfy-vue-side-bar-container">
+    <Toolbar class="comfy-vue-side-bar-header">
+      <template #start>
+        <span class="comfy-vue-side-bar-header-span">{{
+          props.title.toUpperCase()
+        }}</span>
+      </template>
+      <template #end>
+        <slot name="tool-buttons"></slot>
+      </template>
+    </Toolbar>
+    <div class="comfy-vue-side-bar-body">
+      <slot name="body"></slot>
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Toolbar from 'primevue/toolbar'
+
+const props = defineProps({
+  title: {
+    type: String,
+    required: true
+  }
+})
+</script>
+
+<style scoped>
+.comfy-vue-side-bar-container {
+  display: flex;
+  flex-direction: column;
+  height: 100%;
+  overflow: hidden;
+}
+
+.comfy-vue-side-bar-header {
+  flex-shrink: 0;
+  border-left: none;
+  border-right: none;
+  border-top: none;
+  border-radius: 0;
+  padding: 0.25rem 1rem;
+}
+
+.comfy-vue-side-bar-header-span {
+  font-size: small;
+}
+
+.comfy-vue-side-bar-body {
+  flex-grow: 1;
+  overflow: auto;
+  scrollbar-width: thin;
+  scrollbar-color: transparent transparent;
+}
+
+.comfy-vue-side-bar-body::-webkit-scrollbar {
+  width: 1px;
+}
+
+.comfy-vue-side-bar-body::-webkit-scrollbar-thumb {
+  background-color: transparent;
+}
+</style>
diff --git a/src/i18n.ts b/src/i18n.ts
index d51f4996e..cadcc2dbf 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -6,7 +6,10 @@ const messages = {
       settings: 'Settings',
       themeToggle: 'Toggle Theme',
       queue: 'Queue',
-      nodeLibrary: 'Node Library'
+      nodeLibrary: 'Node Library',
+      nodeLibraryTab: {
+        sortOrder: 'Sort Order'
+      }
     }
   },
   zh: {
@@ -14,7 +17,10 @@ const messages = {
       settings: '设置',
       themeToggle: '主题切换',
       queue: '队列',
-      nodeLibrary: '节点库'
+      nodeLibrary: '节点库',
+      nodeLibraryTab: {
+        sortOrder: '排序顺序'
+      }
     }
   }
   // TODO: Add more languages
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 2a50c448c..499fb218a 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -218,6 +218,27 @@ export const SYSTEM_NODE_DEFS: 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>
@@ -261,6 +282,9 @@ export const useNodeDefStore = defineStore('nodeDef', {
         })
       }
       return root
+    },
+    sortedNodeTree(): TreeNode {
+      return sortedTree(this.nodeTree)
     }
   },
   actions: {