From cadf3a2fecb5637d3818bb348db3fc91db10994d Mon Sep 17 00:00:00 2001
From: "Regis Gaughan, III" <regis.gaughan@gmail.com>
Date: Thu, 2 May 2024 21:15:24 -0400
Subject: [PATCH] Extend core snapToGrid to LiteGraph Groups.

Extends the core Comfy.SnapToGrid behavior for nodes to apply to LiteGraph's LGraphGroup with the same behavior. Also, pulls out redundant rounding code into util function.
---
 web/extensions/core/snapToGrid.js | 96 ++++++++++++++++++++++++++++---
 1 file changed, 89 insertions(+), 7 deletions(-)

diff --git a/web/extensions/core/snapToGrid.js b/web/extensions/core/snapToGrid.js
index dc534d6edf..aac0177484 100644
--- a/web/extensions/core/snapToGrid.js
+++ b/web/extensions/core/snapToGrid.js
@@ -2,6 +2,13 @@ import { app } from "../../scripts/app.js";
 
 // Shift + drag/resize to snap to grid
 
+/** Rounds a Vector2 in-place to the current CANVAS_GRID_SIZE. */
+function roundVectorToGrid(vec) {
+	vec[0] = LiteGraph.CANVAS_GRID_SIZE * Math.round(vec[0] / LiteGraph.CANVAS_GRID_SIZE);
+	vec[1] = LiteGraph.CANVAS_GRID_SIZE * Math.round(vec[1] / LiteGraph.CANVAS_GRID_SIZE);
+	return vec;
+}
+
 app.registerExtension({
 	name: "Comfy.SnapToGrid",
 	init() {
@@ -43,10 +50,7 @@ app.registerExtension({
 			const onResize = node.onResize;
 			node.onResize = function () {
 				if (app.shiftDown) {
-					const w = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.size[0] / LiteGraph.CANVAS_GRID_SIZE);
-					const h = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.size[1] / LiteGraph.CANVAS_GRID_SIZE);
-					node.size[0] = w;
-					node.size[1] = h;
+					roundVectorToGrid(node.size);
 				}
 				return onResize?.apply(this, arguments);
 			};
@@ -57,9 +61,7 @@ app.registerExtension({
 		const origDrawNode = LGraphCanvas.prototype.drawNode;
 		LGraphCanvas.prototype.drawNode = function (node, ctx) {
 			if (app.shiftDown && this.node_dragged && node.id in this.selected_nodes) {
-				const x = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.pos[0] / LiteGraph.CANVAS_GRID_SIZE);
-				const y = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.pos[1] / LiteGraph.CANVAS_GRID_SIZE);
-
+				const [x, y] = roundVectorToGrid([...node.pos]);
 				const shiftX = x - node.pos[0];
 				let shiftY = y - node.pos[1];
 
@@ -85,5 +87,85 @@ app.registerExtension({
 
 			return origDrawNode.apply(this, arguments);
 		};
+
+		
+		
+		/**
+  		 * The currently moving, selected group only. Set after the `selected_group` has actually started
+		 * moving.
+   		 */
+		let selectedAndMovingGroup = null;
+
+		/**
+  		 * Handles moving a group; tracking when a group has been moved (to show the ghost in `drawGroups`
+     		 * below) as well as handle the last move call from LiteGraph's `processMouseUp`.
+		 */
+		const groupMove = LGraphGroup.prototype.move;
+		LGraphGroup.prototype.move = function(deltax, deltay, ignore_nodes) {
+			const v = groupMove.apply(this, arguments);
+			// When we've started moving, set `selectedAndMovingGroup` as LiteGraph sets `selected_group`
+			// too eagerly and we don't want to behave like we're moving until we get a delta.
+			if (!selectedAndMovingGroup && app.canvas.selected_group === this && (deltax || deltay)) {
+				selectedAndMovingGroup = this;
+			}
+
+			// LiteGraph will call group.move both on mouse-move as well as mouse-up though we only want
+			// to snap on a mouse-up which we can determine by checking if `app.canvas.last_mouse_dragging`
+			// has been set to `false`. Essentially, this check here is the equivilant to calling an
+			// `LGraphGroup.prototype.onNodeMoved` if it had existed.
+			if (app.canvas.last_mouse_dragging === false && app.shiftDown) {
+				// After moving a group (while app.shiftDown), snap all the child nodes and, finally,
+				//  align the group itself.
+				this.recomputeInsideNodes();
+				for (const node of this._nodes) {
+					node.alignToGrid();
+				}
+				LGraphNode.prototype.alignToGrid.apply(this);
+			}
+			return v;
+		};
+
+		/**
+  		 * Handles drawing a group when, snapping the size when one is actively being resized tracking and/or
+     		 * drawing a ghost box when one is actively being moved. This mimics the node snapping behavior for
+		 * both.
+		 */
+		const drawGroups = LGraphCanvas.prototype.drawGroups;
+		LGraphCanvas.prototype.drawGroups = function (canvas, ctx) {
+			if (this.selected_group && app.shiftDown) {
+				if (this.selected_group_resizing) {
+					roundVectorToGrid(this.selected_group.size);
+				} else if (selectedAndMovingGroup) {
+					const [x, y] = roundVectorToGrid([...selectedAndMovingGroup.pos]);
+					const f = ctx.fillStyle;
+					const s = ctx.strokeStyle;
+					ctx.fillStyle = "rgba(100, 100, 100, 0.33)";
+					ctx.strokeStyle = "rgba(100, 100, 100, 0.66)";
+					ctx.rect(x, y, ...selectedAndMovingGroup.size);
+					ctx.fill();
+					ctx.stroke();
+					ctx.fillStyle = f;
+					ctx.strokeStyle = s;
+				}
+			} else if (!this.selected_group) {
+				selectedAndMovingGroup = null;
+			}
+			return drawGroups.apply(this, arguments);
+		};
+
+
+		/** Handles adding a group in a snapping-enabled state. */
+		const onGroupAdd = LGraphCanvas.onGroupAdd;
+		LGraphCanvas.onGroupAdd = function() {
+			const v = onGroupAdd.apply(app.canvas, arguments);
+			if (app.shiftDown) {
+				const lastGroup = app.graph._groups[app.graph._groups.length - 1];
+				if (lastGroup) {
+					roundVectorToGrid(lastGroup.pos);
+					roundVectorToGrid(lastGroup.size);
+				}
+			}
+			return v;
+		};
 	},
 });