From 0658698a13dd734afc88a16ffd961e253c6f5d9a Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Sun, 16 Feb 2025 21:23:07 -0500
Subject: [PATCH] Selection Overlay (#2592)

---
 package-lock.json                             |  8 +-
 package.json                                  |  2 +-
 src/components/graph/GraphCanvas.vue          |  5 ++
 src/components/graph/SelectionOverlay.vue     | 85 +++++++++++++++++++
 .../functional/useChainCallback.ts            | 16 ++++
 src/scripts/app.ts                            |  8 +-
 6 files changed, 115 insertions(+), 9 deletions(-)
 create mode 100644 src/components/graph/SelectionOverlay.vue
 create mode 100644 src/composables/functional/useChainCallback.ts

diff --git a/package-lock.json b/package-lock.json
index 0b86b26a5f..b2a0255a74 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -11,7 +11,7 @@
       "dependencies": {
         "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
         "@comfyorg/comfyui-electron-types": "^0.4.16",
-        "@comfyorg/litegraph": "^0.8.82",
+        "@comfyorg/litegraph": "^0.8.83",
         "@primevue/forms": "^4.2.5",
         "@primevue/themes": "^4.2.5",
         "@sentry/vue": "^8.48.0",
@@ -1944,9 +1944,9 @@
       "license": "GPL-3.0-only"
     },
     "node_modules/@comfyorg/litegraph": {
-      "version": "0.8.82",
-      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.8.82.tgz",
-      "integrity": "sha512-grwCXpjSKdyH/nVt+PYnv8BLavNNdyIhCAqiAMx9V5cK4bKsFFdHweuyYD8D2ySPR1iPPhOwsvfpi42ud+mbJQ==",
+      "version": "0.8.83",
+      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.8.83.tgz",
+      "integrity": "sha512-4ZfRk0mBcCStY2yRERCrguwFf5v6WajD/6/JEmycD3HnF4OwYgyAspMYrscJcQ/R2MXfnedGe1gi8WXQ955vEQ==",
       "license": "MIT"
     },
     "node_modules/@cspotcode/source-map-support": {
diff --git a/package.json b/package.json
index 66e5670c74..2680547725 100644
--- a/package.json
+++ b/package.json
@@ -84,7 +84,7 @@
   "dependencies": {
     "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
     "@comfyorg/comfyui-electron-types": "^0.4.16",
-    "@comfyorg/litegraph": "^0.8.82",
+    "@comfyorg/litegraph": "^0.8.83",
     "@primevue/forms": "^4.2.5",
     "@primevue/themes": "^4.2.5",
     "@sentry/vue": "^8.48.0",
diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index f5bb67d0a8..e568b0bfb4 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -28,6 +28,10 @@
     class="w-full h-full touch-none"
   />
   <NodeSearchboxPopover />
+  <SelectionOverlay>
+    <!-- Placeholder for selection overlay testing. -->
+    <!-- <div class="w-full h-full bg-red-500"></div> -->
+  </SelectionOverlay>
   <NodeTooltip v-if="tooltipEnabled" />
   <NodeBadge />
 </template>
@@ -40,6 +44,7 @@ import BottomPanel from '@/components/bottomPanel/BottomPanel.vue'
 import GraphCanvasMenu from '@/components/graph/GraphCanvasMenu.vue'
 import NodeBadge from '@/components/graph/NodeBadge.vue'
 import NodeTooltip from '@/components/graph/NodeTooltip.vue'
+import SelectionOverlay from '@/components/graph/SelectionOverlay.vue'
 import TitleEditor from '@/components/graph/TitleEditor.vue'
 import NodeSearchboxPopover from '@/components/searchbox/NodeSearchBoxPopover.vue'
 import SideToolbar from '@/components/sidebar/SideToolbar.vue'
diff --git a/src/components/graph/SelectionOverlay.vue b/src/components/graph/SelectionOverlay.vue
new file mode 100644
index 0000000000..0b9dc1c515
--- /dev/null
+++ b/src/components/graph/SelectionOverlay.vue
@@ -0,0 +1,85 @@
+<!-- This component is used to bound the selected items on the canvas. -->
+<template>
+  <div
+    class="selection-overlay-container pointer-events-none"
+    :style="style"
+    v-show="visible"
+  >
+    <slot></slot>
+  </div>
+</template>
+
+<script setup lang="ts">
+import { createBounds } from '@comfyorg/litegraph'
+import type { LGraphCanvas } from '@comfyorg/litegraph'
+import { ref, watch } from 'vue'
+
+import { useAbsolutePosition } from '@/composables/element/useAbsolutePosition'
+import { useChainCallback } from '@/composables/functional/useChainCallback'
+import { useCanvasStore } from '@/stores/graphStore'
+
+const canvasStore = useCanvasStore()
+const { style, updatePosition } = useAbsolutePosition()
+
+const visible = ref(false)
+
+const positionSelectionOverlay = (canvas: LGraphCanvas) => {
+  const selectedItems = canvas.selectedItems
+  if (!selectedItems.size) {
+    visible.value = false
+    return
+  }
+
+  visible.value = true
+  const bounds = createBounds(selectedItems)
+  updatePosition({
+    pos: [bounds[0], bounds[1]],
+    size: [bounds[2], bounds[3]]
+  })
+}
+
+// Register listener on canvas creation.
+watch(
+  () => canvasStore.canvas,
+  (canvas: LGraphCanvas | null) => {
+    if (!canvas) return
+
+    canvas.onSelectionChange = useChainCallback(canvas.onSelectionChange, () =>
+      positionSelectionOverlay(canvas)
+    )
+  },
+  { immediate: true }
+)
+
+watch(
+  () => {
+    const canvas = canvasStore.canvas
+    if (!canvas) return null
+    return {
+      scale: canvas.ds.state.scale,
+      offset: [canvas.ds.state.offset[0], canvas.ds.state.offset[1]]
+    }
+  },
+  (state) => {
+    if (!state) return
+
+    positionSelectionOverlay(canvasStore.canvas as LGraphCanvas)
+  }
+)
+
+watch(
+  () => canvasStore.canvas?.state?.draggingItems,
+  (draggingItems) => {
+    visible.value = !draggingItems
+    if (draggingItems === false) {
+      positionSelectionOverlay(canvasStore.canvas as LGraphCanvas)
+    }
+  }
+)
+</script>
+
+<style scoped>
+.selection-overlay-container > * {
+  pointer-events: auto;
+}
+</style>
diff --git a/src/composables/functional/useChainCallback.ts b/src/composables/functional/useChainCallback.ts
new file mode 100644
index 0000000000..40ae6a1fd6
--- /dev/null
+++ b/src/composables/functional/useChainCallback.ts
@@ -0,0 +1,16 @@
+/**
+ * Chain multiple callbacks together.
+ *
+ * @param originalCallback - The original callback to chain.
+ * @param callbacks - The callbacks to chain.
+ * @returns A new callback that chains the original callback with the callbacks.
+ */
+export const useChainCallback = <T extends (...args: any[]) => void>(
+  originalCallback: T | undefined,
+  ...callbacks: ((...args: Parameters<T>) => void)[]
+) => {
+  return (...args: Parameters<T>) => {
+    originalCallback?.(...args)
+    callbacks.forEach((callback) => callback(...args))
+  }
+}
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 38a67d6208..dde9275fa6 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -10,7 +10,7 @@ import {
 import type { Rect, Vector2 } from '@comfyorg/litegraph'
 import _ from 'lodash'
 import type { ToastMessageOptions } from 'primevue/toast'
-import { shallowReactive } from 'vue'
+import { reactive } from 'vue'
 
 import { st } from '@/i18n'
 import { useDialogService } from '@/services/dialogService'
@@ -795,10 +795,10 @@ export class ComfyApp {
 
     this.#addAfterConfigureHandler()
 
-    // Make LGraphCanvas.state shallow reactive so that any change on the root
-    // object triggers reactivity.
     this.canvas = new LGraphCanvas(canvasEl, this.graph)
-    this.canvas.state = shallowReactive(this.canvas.state)
+    // Make canvas states reactive so we can observe changes on them.
+    this.canvas.state = reactive(this.canvas.state)
+    this.canvas.ds.state = reactive(this.canvas.ds.state)
 
     this.ctx = canvasEl.getContext('2d')