From f6a203951f061759a3d6a9847b865dd9c7d77a38 Mon Sep 17 00:00:00 2001
From: "Regis Gaughan, III" <regis.gaughan@gmail.com>
Date: Mon, 27 May 2024 14:05:51 -0400
Subject: [PATCH] Extend core snapToGrid to LiteGraph Groups. (#3393)

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 dc534d6e..aac01774 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;
+		};
 	},
 });