From b8846228b42397a52d620153c558f6f6121f489a Mon Sep 17 00:00:00 2001
From: huchenlei <chenlei.hu@mail.utoronto.ca>
Date: Sat, 24 Aug 2024 17:21:02 -0400
Subject: [PATCH] Refactor node library drag and drop

---
 src/components/graph/GraphCanvas.vue          |  22 ++--
 src/components/primevueOverride/TreePlus.vue  | 106 ------------------
 .../sidebar/tabs/NodeLibrarySidebarTab.vue    |   6 +-
 .../sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue |  25 ++++-
 src/types/litegraphTypes.ts                   |   5 +
 5 files changed, 45 insertions(+), 119 deletions(-)
 delete mode 100644 src/components/primevueOverride/TreePlus.vue

diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index 92ffae8a..5a2b0d4d 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -20,7 +20,7 @@ import { ref, computed, onUnmounted, watch, onMounted, watchEffect } from 'vue'
 import { app as comfyApp } from '@/scripts/app'
 import { useSettingStore } from '@/stores/settingStore'
 import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
-import { useNodeDefStore } from '@/stores/nodeDefStore'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import {
   LiteGraph,
@@ -113,14 +113,18 @@ onMounted(async () => {
     element: canvasRef.value,
     onDrop: (event) => {
       const loc = event.location.current.input
-      // Add an offset on x to make sure after adding the node, the cursor
-      // is on the node (top left corner)
-      const pos = comfyApp.clientPosToCanvasPos([loc.clientX - 20, loc.clientY])
-      const comfyNodeName = event.source.element.getAttribute(
-        'data-comfy-node-name'
-      )
-      const nodeDef = nodeDefStore.nodeDefsByName[comfyNodeName]
-      comfyApp.addNodeOnGraph(nodeDef, { pos })
+      const dndData = event.source.data
+
+      if (dndData.type === 'add-node') {
+        const nodeDef = dndData.data as ComfyNodeDefImpl
+        // Add an offset on x to make sure after adding the node, the cursor
+        // is on the node (top left corner)
+        const pos = comfyApp.clientPosToCanvasPos([
+          loc.clientX - 20,
+          loc.clientY
+        ])
+        comfyApp.addNodeOnGraph(nodeDef, { pos })
+      }
     }
   })
   emit('ready')
diff --git a/src/components/primevueOverride/TreePlus.vue b/src/components/primevueOverride/TreePlus.vue
deleted file mode 100644
index de163e0c..00000000
--- a/src/components/primevueOverride/TreePlus.vue
+++ /dev/null
@@ -1,106 +0,0 @@
-<!-- Tree with all leaf nodes draggable -->
-<script>
-import Tree from 'primevue/tree'
-import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
-import { h, onMounted, onBeforeUnmount, computed } from 'vue'
-
-export default {
-  name: 'TreePlus',
-  extends: Tree,
-  props: {
-    dragSelector: {
-      type: String,
-      default: '.p-tree-node'
-    },
-    // Explicitly declare all v-model props
-    expandedKeys: {
-      type: Object,
-      default: () => ({})
-    },
-    selectionKeys: {
-      type: Object,
-      default: () => ({})
-    }
-  },
-  emits: ['update:expandedKeys', 'update:selectionKeys'],
-  setup(props, context) {
-    // Create computed properties for each v-model prop
-    const computedExpandedKeys = computed({
-      get: () => props.expandedKeys,
-      set: (value) => context.emit('update:expandedKeys', value)
-    })
-
-    const computedSelectionKeys = computed({
-      get: () => props.selectionKeys,
-      set: (value) => context.emit('update:selectionKeys', value)
-    })
-
-    let observer = null
-
-    const makeDraggable = (element) => {
-      if (!element._draggableCleanup) {
-        element._draggableCleanup = draggable({
-          element
-        })
-      }
-    }
-
-    const observeTreeChanges = (treeElement) => {
-      observer = new MutationObserver((mutations) => {
-        mutations.forEach((mutation) => {
-          if (mutation.type === 'childList') {
-            mutation.addedNodes.forEach((node) => {
-              if (node.nodeType === Node.ELEMENT_NODE) {
-                node.querySelectorAll(props.dragSelector).forEach(makeDraggable)
-              }
-            })
-          }
-        })
-      })
-
-      observer.observe(treeElement, { childList: true, subtree: true })
-
-      // Make existing nodes draggable
-      treeElement.querySelectorAll(props.dragSelector).forEach(makeDraggable)
-    }
-
-    onMounted(() => {
-      const treeElement = document.querySelector('.p-tree')
-      if (treeElement) {
-        observeTreeChanges(treeElement)
-      }
-    })
-
-    onBeforeUnmount(() => {
-      if (observer) {
-        observer.disconnect()
-      }
-      // Clean up draggable instances if necessary
-      const treeElement = document.querySelector('.p-tree')
-      if (treeElement) {
-        treeElement.querySelectorAll(props.dragSelector).forEach((node) => {
-          if (node._draggableCleanup) {
-            node._draggableCleanup()
-          }
-        })
-      }
-    })
-
-    return () =>
-      h(
-        Tree,
-        {
-          ...context.attrs,
-          ...props,
-          expandedKeys: computedExpandedKeys.value,
-          selectionKeys: computedSelectionKeys.value,
-          'onUpdate:expandedKeys': (value) =>
-            (computedExpandedKeys.value = value),
-          'onUpdate:selectionKeys': (value) =>
-            (computedSelectionKeys.value = value)
-        },
-        context.slots
-      )
-  }
-}
-</script>
diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index 4aad2fba..38742f27 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -20,7 +20,7 @@
         @search="handleSearch"
         :placeholder="$t('searchNodes') + '...'"
       />
-      <TreePlus
+      <Tree
         class="node-lib-tree"
         v-model:expandedKeys="expandedKeys"
         selectionMode="single"
@@ -62,7 +62,7 @@
             @toggleBookmark="toggleBookmark(node.data.display_name)"
           />
         </template>
-      </TreePlus>
+      </Tree>
       <div
         v-if="hoveredComfyNode"
         class="node-lib-node-preview"
@@ -89,7 +89,7 @@ import {
 import { computed, ref, nextTick } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 import NodeTreeLeaf from './nodeLibrary/NodeTreeLeaf.vue'
-import TreePlus from '@/components/primevueOverride/TreePlus.vue'
+import Tree from 'primevue/tree'
 import NodePreview from '@/components/node/NodePreview.vue'
 import SearchBox from '@/components/common/SearchBox.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
diff --git a/src/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue b/src/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue
index ec37df17..c33b9ff1 100644
--- a/src/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue
+++ b/src/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue
@@ -1,5 +1,5 @@
 <template>
-  <div class="node-tree-leaf">
+  <div class="node-tree-leaf" ref="container">
     <div class="node-content">
       <Tag
         v-if="node.experimental"
@@ -24,6 +24,9 @@
 import Button from 'primevue/button'
 import Tag from 'primevue/tag'
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+import { onMounted, onUnmounted, ref } from 'vue'
+import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { CanvasDragAndDropData } from '@/types/litegraphTypes'
 
 const props = defineProps<{
   node: ComfyNodeDefImpl
@@ -37,6 +40,26 @@ const emit = defineEmits<{
 const toggleBookmark = () => {
   emit('toggle-bookmark', props.node)
 }
+
+const container = ref<HTMLElement | null>(null)
+let draggableCleanup: () => void
+onMounted(() => {
+  const treeNodeElement = container.value?.closest(
+    '.p-tree-node'
+  ) as HTMLElement
+  draggableCleanup = draggable({
+    element: treeNodeElement,
+    getInitialData() {
+      return {
+        type: 'add-node',
+        data: props.node
+      } as CanvasDragAndDropData<ComfyNodeDefImpl>
+    }
+  })
+})
+onUnmounted(() => {
+  draggableCleanup()
+})
 </script>
 
 <style scoped>
diff --git a/src/types/litegraphTypes.ts b/src/types/litegraphTypes.ts
index 6325badf..ec3d25c0 100644
--- a/src/types/litegraphTypes.ts
+++ b/src/types/litegraphTypes.ts
@@ -73,3 +73,8 @@ export class ConnectingLinkImpl implements ConnectingLink {
     }
   }
 }
+
+export type CanvasDragAndDropData<T = any> = {
+  type: 'add-node'
+  data: T
+}