From 695e464255348e87289824d730baa152eef86f98 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 17 Jun 2024 17:53:30 +1000
Subject: [PATCH] feat(ui): document bounds overlay

---
 .../components/StageComponent.tsx             |  7 +++
 .../features/controlLayers/konva/events.ts    |  5 +-
 .../konva/renderers/previewLayer.ts           | 54 +++++++++++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx b/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
index c634a355c1..8b56b9598f 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/StageComponent.tsx
@@ -11,6 +11,7 @@ import {
   debouncedRenderers,
   renderers as normalRenderers,
 } from 'features/controlLayers/konva/renderers/layers';
+import { renderDocumentBoundsOverlay } from 'features/controlLayers/konva/renderers/previewLayer';
 import { renderLayers } from 'features/controlLayers/konva/renderers/rasterLayer';
 import { renderRegions } from 'features/controlLayers/konva/renderers/rgLayer';
 import {
@@ -298,6 +299,7 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
         scale: stage.scaleX(),
       });
       renderBackgroundLayer(stage);
+      renderDocumentBoundsOverlay(stage, $document.get);
     };
 
     const resizeObserver = new ResizeObserver(fitStageToContainer);
@@ -328,6 +330,7 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
   }, [
     asPreview,
     currentFill,
+    document,
     isDrawing,
     isMouseDown,
     lastCursorPos,
@@ -372,6 +375,10 @@ const useStageRenderer = (stage: Konva.Stage, container: HTMLDivElement | null,
     renderControlAdapters(stage, controlAdapters, getImageDTO);
   }, [controlAdapters, stage]);
 
+  useLayoutEffect(() => {
+    renderDocumentBoundsOverlay(stage, $document.get);
+  }, [stage, document]);
+
   useLayoutEffect(() => {
     arrangeEntities(stage, layers, controlAdapters, regions);
   }, [layers, controlAdapters, regions, stage]);
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/events.ts b/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
index 825a2dcb0c..80edb1fcbd 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/events.ts
@@ -1,5 +1,5 @@
 import { renderBackgroundLayer } from 'features/controlLayers/konva/renderers/background';
-import { scaleToolPreview } from 'features/controlLayers/konva/renderers/previewLayer';
+import { renderDocumentBoundsOverlay, scaleToolPreview } from 'features/controlLayers/konva/renderers/previewLayer';
 import { getScaledFlooredCursorPosition } from 'features/controlLayers/konva/util';
 import type {
   BrushLineAddedArg,
@@ -469,6 +469,7 @@ export const setStageEventHandlers = ({
       setStageAttrs({ ...newPos, width: stage.width(), height: stage.height(), scale: newScale });
       renderBackgroundLayer(stage);
       scaleToolPreview(stage, getToolState());
+      renderDocumentBoundsOverlay(stage, getDocument);
     }
   });
 
@@ -482,6 +483,7 @@ export const setStageEventHandlers = ({
       scale: stage.scaleX(),
     });
     renderBackgroundLayer(stage);
+    renderDocumentBoundsOverlay(stage, getDocument);
   });
 
   //#region dragend
@@ -526,6 +528,7 @@ export const setStageEventHandlers = ({
       setStageAttrs({ x, y, width, height, scale });
       scaleToolPreview(stage, getToolState());
       renderBackgroundLayer(stage);
+      renderDocumentBoundsOverlay(stage, getDocument);
     }
   };
   window.addEventListener('keydown', onKeyDown);
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts
index eff1e48bc3..a3b8f7a893 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts
@@ -1,3 +1,4 @@
+import { getArbitraryBaseColor } from '@invoke-ai/ui-library';
 import { rgbaColorToString } from 'common/util/colorCodeTransformers';
 import { roundToMultiple, roundToMultipleMin } from 'common/util/roundDownToMultiple';
 import {
@@ -449,3 +450,56 @@ export const scaleToolPreview = (stage: Konva.Stage, toolState: CanvasV2State['t
     ?.findOne<Konva.Circle>(`#${PREVIEW_BRUSH_BORDER_OUTER_ID}`)
     ?.setAttrs({ strokeWidth: BRUSH_ERASER_BORDER_WIDTH / scale, radius: radius + BRUSH_ERASER_BORDER_WIDTH / scale });
 };
+
+const getDocumentOverlayGroup = (stage: Konva.Stage): Konva.Group => {
+  const previewLayer = getPreviewLayer(stage);
+  let documentOverlayGroup = previewLayer.findOne<Konva.Group>('#document_overlay_group');
+  if (documentOverlayGroup) {
+    return documentOverlayGroup;
+  }
+
+  documentOverlayGroup = new Konva.Group({ id: 'document_overlay_group', listening: false });
+  const documentOverlayOuterRect = new Konva.Rect({
+    id: 'document_overlay_outer_rect',
+    listening: false,
+    fill: getArbitraryBaseColor(10),
+    opacity: 0.7,
+  });
+  const documentOverlayInnerRect = new Konva.Rect({
+    id: 'document_overlay_inner_rect',
+    listening: false,
+    fill: 'white',
+    globalCompositeOperation: 'destination-out',
+  });
+  documentOverlayGroup.add(documentOverlayOuterRect);
+  documentOverlayGroup.add(documentOverlayInnerRect);
+  previewLayer.add(documentOverlayGroup);
+  return documentOverlayGroup;
+};
+
+export const renderDocumentBoundsOverlay = (stage: Konva.Stage, getDocument: () => CanvasV2State['document']): void => {
+  const document = getDocument();
+  const documentOverlayGroup = getDocumentOverlayGroup(stage);
+
+  documentOverlayGroup.zIndex(0);
+
+  const x = stage.x();
+  const y = stage.y();
+  const width = stage.width();
+  const height = stage.height();
+  const scale = stage.scaleX();
+
+  documentOverlayGroup.findOne<Konva.Rect>('#document_overlay_outer_rect')?.setAttrs({
+    offsetX: x / scale,
+    offsetY: y / scale,
+    width: width / scale,
+    height: height / scale,
+  });
+
+  documentOverlayGroup.findOne<Konva.Rect>('#document_overlay_inner_rect')?.setAttrs({
+    x: 0,
+    y: 0,
+    width: document.width,
+    height: document.height,
+  });
+};