refactor(ui): create entity to konva node map abstraction (wip)

Instead of chaining konva `find` and `findOne` methods, all konva nodes are added to a mapping object. Finding and manipulating them is much simpler.

Done for regions and layers, wip for control adapters.

invokeai/frontend/web/src/features/controlLayers/konva/konvaMap.ts

@@ -0,0 +1,104 @@
+import type { BrushLine, EraserLine, ImageObject, RectShape } from 'features/controlLayers/store/types';
+import type Konva from 'konva';
+
+export type BrushLineEntry = {
+  id: string;
+  type: BrushLine['type'];
+  konvaLine: Konva.Line;
+  konvaLineGroup: Konva.Group;
+};
+
+export type EraserLineEntry = {
+  id: string;
+  type: EraserLine['type'];
+  konvaLine: Konva.Line;
+  konvaLineGroup: Konva.Group;
+};
+
+export type RectShapeEntry = {
+  id: string;
+  type: RectShape['type'];
+  konvaRect: Konva.Rect;
+};
+
+export type ImageEntry = {
+  id: string;
+  type: ImageObject['type'];
+  konvaImage: Konva.Image | null; // The image is loaded asynchronously, so it may not be available immediately
+  konvaGroup: Konva.Group;
+};
+
+type Entry = BrushLineEntry | EraserLineEntry | RectShapeEntry | ImageEntry;
+
+export class EntityToKonvaMap {
+  mappings: Record<string, EntityToKonvaMapping>;
+
+  constructor() {
+    this.mappings = {};
+  }
+
+  addMapping(id: string, konvaLayer: Konva.Layer, konvaObjectGroup: Konva.Group): EntityToKonvaMapping {
+    const mapping = new EntityToKonvaMapping(id, konvaLayer, konvaObjectGroup);
+    this.mappings[id] = mapping;
+    return mapping;
+  }
+
+  getMapping(id: string): EntityToKonvaMapping | undefined {
+    return this.mappings[id];
+  }
+
+  getMappings(): EntityToKonvaMapping[] {
+    return Object.values(this.mappings);
+  }
+
+  destroyMapping(id: string): void {
+    const mapping = this.getMapping(id);
+    if (!mapping) {
+      return;
+    }
+    mapping.konvaObjectGroup.destroy();
+    delete this.mappings[id];
+  }
+}
+
+export class EntityToKonvaMapping {
+  id: string;
+  konvaLayer: Konva.Layer;
+  konvaObjectGroup: Konva.Group;
+  konvaNodeEntries: Record<string, Entry>;
+
+  constructor(id: string, konvaLayer: Konva.Layer, konvaObjectGroup: Konva.Group) {
+    this.id = id;
+    this.konvaLayer = konvaLayer;
+    this.konvaObjectGroup = konvaObjectGroup;
+    this.konvaNodeEntries = {};
+  }
+
+  addEntry<T extends Entry>(entry: T): T {
+    this.konvaNodeEntries[entry.id] = entry;
+    return entry;
+  }
+
+  getEntry<T extends Entry>(id: string): T | undefined {
+    return this.konvaNodeEntries[id] as T | undefined;
+  }
+
+  getEntries(): Entry[] {
+    return Object.values(this.konvaNodeEntries);
+  }
+
+  destroyEntry(id: string): void {
+    const entry = this.getEntry(id);
+    if (!entry) {
+      return;
+    }
+    if (entry.type === 'brush_line' || entry.type === 'eraser_line') {
+      entry.konvaLineGroup.destroy();
+    } else if (entry.type === 'rect_shape') {
+      entry.konvaRect.destroy();
+    } else if (entry.type === 'image') {
+      entry.konvaGroup.destroy();
+    }
+    delete this.konvaNodeEntries[id];
+  }
+}

invokeai/frontend/web/src/features/controlLayers/konva/renderers/caLayer.ts

@@ -1,4 +1,5 @@
 import { LightnessToAlphaFilter } from 'features/controlLayers/konva/filters';
+import type { EntityToKonvaMap } from 'features/controlLayers/konva/konvaMap';
 import { CA_LAYER_IMAGE_NAME, CA_LAYER_NAME, getCAImageId } from 'features/controlLayers/konva/naming';
 import type { ControlAdapterEntity } from 'features/controlLayers/store/types';
 import Konva from 'konva';
@@ -34,6 +35,7 @@ const createCALayerImage = (konvaLayer: Konva.Layer, imageEl: HTMLImageElement):
   const konvaImage = new Konva.Image({
     name: CA_LAYER_IMAGE_NAME,
     image: imageEl,
+    listening: false,
   });
   konvaLayer.add(konvaImage);
   return konvaImage;
@@ -128,6 +130,7 @@ const updateCALayerImageAttrs = (stage: Konva.Stage, konvaImage: Konva.Image, ca
  */
 export const renderCALayer = (
   stage: Konva.Stage,
+  controlAdapterMap: EntityToKonvaMap,
   ca: ControlAdapterEntity,
   getImageDTO: (imageName: string) => Promise<ImageDTO | null>
 ): void => {
@@ -157,16 +160,17 @@ export const renderCALayer = (
 
 export const renderControlAdapters = (
   stage: Konva.Stage,
+  controlAdapterMap: EntityToKonvaMap,
   controlAdapters: ControlAdapterEntity[],
   getImageDTO: (imageName: string) => Promise<ImageDTO | null>
 ): void => {
   // Destroy nonexistent layers
-  for (const konvaLayer of stage.find<Konva.Layer>(`.${CA_LAYER_NAME}`)) {
+  for (const mapping of controlAdapterMap.getMappings()) {
-    if (!controlAdapters.find((ca) => ca.id === konvaLayer.id())) {
+    if (!controlAdapters.find((ca) => ca.id === mapping.id)) {
-      konvaLayer.destroy();
+      controlAdapterMap.destroyMapping(mapping.id);
     }
   }
   for (const ca of controlAdapters) {
-    renderCALayer(stage, ca, getImageDTO);
+    renderCALayer(stage, controlAdapterMap, ca, getImageDTO);
   }
 };

invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts

@@ -1,4 +1,11 @@
 import { rgbaColorToString } from 'common/util/colorCodeTransformers';
+import type {
+  BrushLineEntry,
+  EntityToKonvaMapping,
+  EraserLineEntry,
+  ImageEntry,
+  RectShapeEntry,
+} from 'features/controlLayers/konva/konvaMap';
 import {
   getLayerBboxId,
   getObjectGroupId,
@@ -23,18 +30,17 @@ import { v4 as uuidv4 } from 'uuid';
  * @param layerObjectGroup The konva layer's object group to add the line to
  * @param name The konva name for the line
  */
-export const getBrushLine = (brushLine: BrushLine, layerObjectGroup: Konva.Group, name: string): Konva.Line => {
+export const getBrushLine = (mapping: EntityToKonvaMapping, brushLine: BrushLine, name: string): BrushLineEntry => {
-  let konvaLineGroup = layerObjectGroup.findOne<Konva.Group>(`#${brushLine.id}_group`);
+  let entry = mapping.getEntry<BrushLineEntry>(brushLine.id);
-  let konvaLine = konvaLineGroup?.findOne<Konva.Line>(`#${brushLine.id}`);
+  if (entry) {
-  if (konvaLine) {
+    return entry;
-    return konvaLine;
   }
 
-  konvaLineGroup = new Konva.Group({
+  const konvaLineGroup = new Konva.Group({
-    id: `${brushLine.id}_group`,
+    clip: brushLine.clip,
-    // clip: brushLine.clip,
+    listening: false,
   });
-  konvaLine = new Konva.Line({
+  const konvaLine = new Konva.Line({
     id: brushLine.id,
     name,
     strokeWidth: brushLine.strokeWidth,
@@ -47,8 +53,9 @@ export const getBrushLine = (brushLine: BrushLine, layerObjectGroup: Konva.Group
     stroke: rgbaColorToString(brushLine.color),
   });
   konvaLineGroup.add(konvaLine);
-  layerObjectGroup.add(konvaLineGroup);
+  mapping.konvaObjectGroup.add(konvaLineGroup);
-  return konvaLine;
+  entry = mapping.addEntry({ id: brushLine.id, type: 'brush_line', konvaLine, konvaLineGroup });
+  return entry;
 };
 
 /**
@@ -57,10 +64,18 @@ export const getBrushLine = (brushLine: BrushLine, layerObjectGroup: Konva.Group
  * @param layerObjectGroup The konva layer's object group to add the line to
  * @param name The konva name for the line
  */
-export const getEraserLine = (eraserLine: EraserLine, layerObjectGroup: Konva.Group, name: string): Konva.Line => {
+export const getEraserLine = (mapping: EntityToKonvaMapping, eraserLine: EraserLine, name: string): EraserLineEntry => {
+  let entry = mapping.getEntry<EraserLineEntry>(eraserLine.id);
+  if (entry) {
+    return entry;
+  }
+
+  const konvaLineGroup = new Konva.Group({
+    clip: eraserLine.clip,
+    listening: false,
+  });
   const konvaLine = new Konva.Line({
     id: eraserLine.id,
-    key: eraserLine.id,
     name,
     strokeWidth: eraserLine.strokeWidth,
     tension: 0,
@@ -72,8 +87,10 @@ export const getEraserLine = (eraserLine: EraserLine, layerObjectGroup: Konva.Gr
     stroke: rgbaColorToString(DEFAULT_RGBA_COLOR),
     clip: eraserLine.clip,
   });
-  layerObjectGroup.add(konvaLine);
+  konvaLineGroup.add(konvaLine);
-  return konvaLine;
+  mapping.konvaObjectGroup.add(konvaLineGroup);
+  entry = mapping.addEntry({ id: eraserLine.id, type: 'eraser_line', konvaLine, konvaLineGroup });
+  return entry;
 };
 
 /**
@@ -82,7 +99,11 @@ export const getEraserLine = (eraserLine: EraserLine, layerObjectGroup: Konva.Gr
  * @param layerObjectGroup The konva layer's object group to add the rect to
  * @param name The konva name for the rect
  */
-export const createRectShape = (rectShape: RectShape, layerObjectGroup: Konva.Group, name: string): Konva.Rect => {
+export const getRectShape = (mapping: EntityToKonvaMapping, rectShape: RectShape, name: string): RectShapeEntry => {
+  let entry = mapping.getEntry<RectShapeEntry>(rectShape.id);
+  if (entry) {
+    return entry;
+  }
   const konvaRect = new Konva.Rect({
     id: rectShape.id,
     key: rectShape.id,
@@ -94,8 +115,9 @@ export const createRectShape = (rectShape: RectShape, layerObjectGroup: Konva.Gr
     listening: false,
     fill: rgbaColorToString(rectShape.color),
   });
-  layerObjectGroup.add(konvaRect);
+  mapping.konvaObjectGroup.add(konvaRect);
-  return konvaRect;
+  entry = mapping.addEntry({ id: rectShape.id, type: 'rect_shape', konvaRect });
+  return entry;
 };
 
 /**
@@ -112,6 +134,7 @@ const createImagePlaceholderGroup = (
     fill: 'hsl(220 12% 45% / 1)', // 'base.500'
     width,
     height,
+    listening: false,
   });
   const konvaPlaceholderText = new Konva.Text({
     name: 'image-placeholder-text',
@@ -150,14 +173,20 @@ const createImagePlaceholderGroup = (
  * @returns A promise that resolves to the konva group for the image object
  */
 export const createImageObjectGroup = async (
+  mapping: EntityToKonvaMapping,
   imageObject: ImageObject,
-  layerObjectGroup: Konva.Group,
   name: string
-): Promise<Konva.Group> => {
+): Promise<ImageEntry> => {
+  let entry = mapping.getEntry<ImageEntry>(imageObject.id);
+  if (entry) {
+    return entry;
+  }
   const konvaImageGroup = new Konva.Group({ id: imageObject.id, name, listening: false });
   const placeholder = createImagePlaceholderGroup(imageObject);
   konvaImageGroup.add(placeholder.konvaPlaceholderGroup);
-  layerObjectGroup.add(konvaImageGroup);
+  mapping.konvaObjectGroup.add(konvaImageGroup);
+
+  entry = mapping.addEntry({ id: imageObject.id, type: 'image', konvaGroup: konvaImageGroup, konvaImage: null });
   getImageDTO(imageObject.image.name).then((imageDTO) => {
     if (!imageDTO) {
       placeholder.onError();
@@ -173,6 +202,7 @@ export const createImageObjectGroup = async (
       });
       placeholder.onLoaded();
       konvaImageGroup.add(konvaImage);
+      entry.konvaImage = konvaImage;
     };
     imageEl.onerror = () => {
       placeholder.onError();
@@ -180,7 +210,7 @@ export const createImageObjectGroup = async (
     imageEl.id = imageObject.id;
     imageEl.src = imageDTO.image_url;
   });
-  return konvaImageGroup;
+  return entry;
 };
 
 /**

invokeai/frontend/web/src/features/controlLayers/konva/renderers/rasterLayer.ts

@@ -1,3 +1,4 @@
+import type { EntityToKonvaMap, EntityToKonvaMapping } from 'features/controlLayers/konva/konvaMap';
 import {
   RASTER_LAYER_BRUSH_LINE_NAME,
   RASTER_LAYER_ERASER_LINE_NAME,
@@ -9,11 +10,11 @@ import {
 import {
   createImageObjectGroup,
   createObjectGroup,
-  createRectShape,
   getBrushLine,
   getEraserLine,
+  getRectShape,
 } from 'features/controlLayers/konva/renderers/objects';
-import { mapId, selectRasterObjects } from 'features/controlLayers/konva/util';
+import { mapId } from 'features/controlLayers/konva/util';
 import type { CanvasEntity, LayerEntity, PosChangedArg, Tool } from 'features/controlLayers/store/types';
 import Konva from 'konva';
 
@@ -27,11 +28,16 @@ import Konva from 'konva';
  * @param layerState The raster layer state
  * @param onPosChanged Callback for when the layer's position changes
  */
-const createRasterLayer = (
+const getLayer = (
   stage: Konva.Stage,
+  layerMap: EntityToKonvaMap,
   layerState: LayerEntity,
   onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
-): Konva.Layer => {
+): EntityToKonvaMapping => {
+  let mapping = layerMap.getMapping(layerState.id);
+  if (mapping) {
+    return mapping;
+  }
   // This layer hasn't been added to the konva state yet
   const konvaLayer = new Konva.Layer({
     id: layerState.id,
@@ -48,9 +54,11 @@ const createRasterLayer = (
     });
   }
 
+  const konvaObjectGroup = createObjectGroup(konvaLayer, RASTER_LAYER_OBJECT_GROUP_NAME);
+  konvaLayer.add(konvaObjectGroup);
   stage.add(konvaLayer);
-
+  mapping = layerMap.addMapping(layerState.id, konvaLayer, konvaObjectGroup);
-  return konvaLayer;
+  return mapping;
 };
 
 /**
@@ -62,63 +70,51 @@ const createRasterLayer = (
  */
 export const renderRasterLayer = async (
   stage: Konva.Stage,
+  layerMap: EntityToKonvaMap,
   layerState: LayerEntity,
   tool: Tool,
   onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
 ) => {
-  const konvaLayer =
+  const mapping = getLayer(stage, layerMap, layerState, onPosChanged);
-    stage.findOne<Konva.Layer>(`#${layerState.id}`) ?? createRasterLayer(stage, layerState, onPosChanged);
 
   // Update the layer's position and listening state
-  konvaLayer.setAttrs({
+  mapping.konvaLayer.setAttrs({
     listening: tool === 'move', // The layer only listens when using the move tool - otherwise the stage is handling mouse events
     x: Math.floor(layerState.x),
     y: Math.floor(layerState.y),
   });
 
-  const konvaObjectGroup =
-    konvaLayer.findOne<Konva.Group>(`.${RASTER_LAYER_OBJECT_GROUP_NAME}`) ??
-    createObjectGroup(konvaLayer, RASTER_LAYER_OBJECT_GROUP_NAME);
-
   const objectIds = layerState.objects.map(mapId);
-  // Destroy any objects that are no longer in the redux state
+  // Destroy any objects that are no longer in state
-  // TODO(psyche): `konvaObjectGroup.getChildren()` seems to return a stale array of children, but find is never stale.
+  for (const entry of mapping.getEntries()) {
-  // Should report upstream
+    if (!objectIds.includes(entry.id)) {
-  for (const objectNode of konvaObjectGroup.find(selectRasterObjects)) {
+      mapping.destroyEntry(entry.id);
-    if (!objectIds.includes(objectNode.id())) {
-      objectNode.destroy();
     }
   }
 
   for (const obj of layerState.objects) {
     if (obj.type === 'brush_line') {
-      const konvaBrushLine = getBrushLine(obj, konvaObjectGroup, RASTER_LAYER_BRUSH_LINE_NAME);
+      const entry = getBrushLine(mapping, obj, RASTER_LAYER_BRUSH_LINE_NAME);
       // Only update the points if they have changed.
-      if (konvaBrushLine.points().length !== obj.points.length) {
+      if (entry.konvaLine.points().length !== obj.points.length) {
-        konvaBrushLine.points(obj.points);
+        entry.konvaLine.points(obj.points);
       }
     } else if (obj.type === 'eraser_line') {
-      const konvaEraserLine =
+      const entry = getEraserLine(mapping, obj, RASTER_LAYER_ERASER_LINE_NAME);
-        konvaObjectGroup.findOne<Konva.Line>(`#${obj.id}`) ??
-        getEraserLine(obj, konvaObjectGroup, RASTER_LAYER_ERASER_LINE_NAME);
       // Only update the points if they have changed.
-      if (konvaEraserLine.points().length !== obj.points.length) {
+      if (entry.konvaLine.points().length !== obj.points.length) {
-        konvaEraserLine.points(obj.points);
+        entry.konvaLine.points(obj.points);
       }
     } else if (obj.type === 'rect_shape') {
-      if (!konvaObjectGroup.findOne<Konva.Rect>(`#${obj.id}`)) {
+      getRectShape(mapping, obj, RASTER_LAYER_RECT_SHAPE_NAME);
-        createRectShape(obj, konvaObjectGroup, RASTER_LAYER_RECT_SHAPE_NAME);
-      }
     } else if (obj.type === 'image') {
-      if (!konvaObjectGroup.findOne<Konva.Group>(`#${obj.id}`)) {
+      createImageObjectGroup(mapping, obj, RASTER_LAYER_IMAGE_NAME);
-        createImageObjectGroup(obj, konvaObjectGroup, RASTER_LAYER_IMAGE_NAME);
-      }
     }
   }
 
   // Only update layer visibility if it has changed.
-  if (konvaLayer.visible() !== layerState.isEnabled) {
+  if (mapping.konvaLayer.visible() !== layerState.isEnabled) {
-    konvaLayer.visible(layerState.isEnabled);
+    mapping.konvaLayer.visible(layerState.isEnabled);
   }
 
   // const bboxRect = konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(layerState, konvaLayer);
@@ -139,22 +135,23 @@ export const renderRasterLayer = async (
   //   bboxRect.visible(false);
   // }
 
-  konvaObjectGroup.opacity(layerState.opacity);
+  mapping.konvaObjectGroup.opacity(layerState.opacity);
 };
 
 export const renderLayers = (
   stage: Konva.Stage,
+  layerMap: EntityToKonvaMap,
   layers: LayerEntity[],
   tool: Tool,
   onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
 ): void => {
   // Destroy nonexistent layers
-  for (const konvaLayer of stage.find<Konva.Layer>(`.${RASTER_LAYER_NAME}`)) {
+  for (const mapping of layerMap.getMappings()) {
-    if (!layers.find((l) => l.id === konvaLayer.id())) {
+    if (!layers.find((l) => l.id === mapping.id)) {
-      konvaLayer.destroy();
+      layerMap.destroyMapping(mapping.id);
     }
   }
   for (const layer of layers) {
-    renderRasterLayer(stage, layer, tool, onPosChanged);
+    renderRasterLayer(stage, layerMap, layer, tool, onPosChanged);
   }
 };

invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts

@@ -4,6 +4,7 @@ import { logger } from 'app/logging/logger';
 import { $isDebugging } from 'app/store/nanostores/isDebugging';
 import type { RootState } from 'app/store/store';
 import { setStageEventHandlers } from 'features/controlLayers/konva/events';
+import { EntityToKonvaMap } from 'features/controlLayers/konva/konvaMap';
 import { renderBackgroundLayer } from 'features/controlLayers/konva/renderers/background';
 import { updateBboxes } from 'features/controlLayers/konva/renderers/bbox';
 import { renderControlAdapters } from 'features/controlLayers/konva/renderers/caLayer';
@@ -281,6 +282,10 @@ export const initializeRenderer = (
   // the entire state over when needed.
   const debouncedUpdateBboxes = debounce(updateBboxes, 300);
 
+  const regionMap = new EntityToKonvaMap();
+  const layerMap = new EntityToKonvaMap();
+  const controlAdapterMap = new EntityToKonvaMap();
+
   const renderCanvas = () => {
     const { canvasV2 } = store.getState();
 
@@ -298,7 +303,7 @@ export const initializeRenderer = (
       canvasV2.tool.selected !== prevCanvasV2.tool.selected
     ) {
       logIfDebugging('Rendering layers');
-      renderLayers(stage, canvasV2.layers, canvasV2.tool.selected, onPosChanged);
+      renderLayers(stage, layerMap, canvasV2.layers, canvasV2.tool.selected, onPosChanged);
     }
 
     if (
@@ -310,6 +315,7 @@ export const initializeRenderer = (
       logIfDebugging('Rendering regions');
       renderRegions(
         stage,
+        regionMap,
         canvasV2.regions,
         canvasV2.settings.maskOpacity,
         canvasV2.tool.selected,
@@ -320,7 +326,7 @@ export const initializeRenderer = (
 
     if (isFirstRender || canvasV2.controlAdapters !== prevCanvasV2.controlAdapters) {
       logIfDebugging('Rendering control adapters');
-      renderControlAdapters(stage, canvasV2.controlAdapters, getImageDTO);
+      renderControlAdapters(stage, controlAdapterMap, canvasV2.controlAdapters, getImageDTO);
     }
 
     if (isFirstRender || canvasV2.document !== prevCanvasV2.document) {

invokeai/frontend/web/src/features/controlLayers/konva/renderers/rgLayer.ts

@@ -1,8 +1,7 @@
 import { rgbColorToString } from 'common/util/colorCodeTransformers';
-import { BBOX_SELECTED_STROKE } from 'features/controlLayers/konva/constants';
+import type { EntityToKonvaMap, EntityToKonvaMapping } from 'features/controlLayers/konva/konvaMap';
 import {
   COMPOSITING_RECT_NAME,
-  LAYER_BBOX_NAME,
   RG_LAYER_BRUSH_LINE_NAME,
   RG_LAYER_ERASER_LINE_NAME,
   RG_LAYER_NAME,
@@ -11,13 +10,12 @@ import {
 } from 'features/controlLayers/konva/naming';
 import { getLayerBboxFast } from 'features/controlLayers/konva/renderers/bbox';
 import {
-  createBboxRect,
   createObjectGroup,
-  createRectShape,
   getBrushLine,
   getEraserLine,
+  getRectShape,
 } from 'features/controlLayers/konva/renderers/objects';
-import { mapId, selectVectorMaskObjects } from 'features/controlLayers/konva/util';
+import { mapId } from 'features/controlLayers/konva/util';
 import type {
   CanvasEntity,
   CanvasEntityIdentifier,
@@ -47,17 +45,22 @@ const createCompositingRect = (konvaLayer: Konva.Layer): Konva.Rect => {
 /**
  * Creates a regional guidance layer.
  * @param stage The konva stage
- * @param rg The regional guidance layer state
+ * @param region The regional guidance layer state
  * @param onLayerPosChanged Callback for when the layer's position changes
  */
-const createRGLayer = (
+const getRegion = (
   stage: Konva.Stage,
-  rg: RegionEntity,
+  regionMap: EntityToKonvaMap,
+  region: RegionEntity,
   onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
-): Konva.Layer => {
+): EntityToKonvaMapping => {
+  let mapping = regionMap.getMapping(region.id);
+  if (mapping) {
+    return mapping;
+  }
   // This layer hasn't been added to the konva state yet
   const konvaLayer = new Konva.Layer({
-    id: rg.id,
+    id: region.id,
     name: RG_LAYER_NAME,
     draggable: true,
     dragDistance: 0,
@@ -67,117 +70,114 @@ const createRGLayer = (
   // the position - we do not need to call this on the `dragmove` event.
   if (onPosChanged) {
     konvaLayer.on('dragend', function (e) {
-      onPosChanged({ id: rg.id, x: Math.floor(e.target.x()), y: Math.floor(e.target.y()) }, 'regional_guidance');
+      onPosChanged({ id: region.id, x: Math.floor(e.target.x()), y: Math.floor(e.target.y()) }, 'regional_guidance');
     });
   }
 
-  stage.add(konvaLayer);
+  const konvaObjectGroup = createObjectGroup(konvaLayer, RG_LAYER_OBJECT_GROUP_NAME);
 
-  return konvaLayer;
+  konvaLayer.add(konvaObjectGroup);
+  stage.add(konvaLayer);
+  mapping = regionMap.addMapping(region.id, konvaLayer, konvaObjectGroup);
+  return mapping;
 };
 
 /**
  * Renders a raster layer.
  * @param stage The konva stage
- * @param rg The regional guidance layer state
+ * @param region The regional guidance layer state
  * @param globalMaskLayerOpacity The global mask layer opacity
  * @param tool The current tool
  * @param onPosChanged Callback for when the layer's position changes
  */
 export const renderRGLayer = (
   stage: Konva.Stage,
-  rg: RegionEntity,
+  regionMap: EntityToKonvaMap,
+  region: RegionEntity,
   globalMaskLayerOpacity: number,
   tool: Tool,
   selectedEntityIdentifier: CanvasEntityIdentifier | null,
   onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
 ): void => {
-  const konvaLayer = stage.findOne<Konva.Layer>(`#${rg.id}`) ?? createRGLayer(stage, rg, onPosChanged);
+  const mapping = getRegion(stage, regionMap, region, onPosChanged);
 
   // Update the layer's position and listening state
-  konvaLayer.setAttrs({
+  mapping.konvaLayer.setAttrs({
     listening: tool === 'move', // The layer only listens when using the move tool - otherwise the stage is handling mouse events
-    x: Math.floor(rg.x),
+    x: Math.floor(region.x),
-    y: Math.floor(rg.y),
+    y: Math.floor(region.y),
   });
 
   // Convert the color to a string, stripping the alpha - the object group will handle opacity.
-  const rgbColor = rgbColorToString(rg.fill);
+  const rgbColor = rgbColorToString(region.fill);
-
-  const konvaObjectGroup =
-    konvaLayer.findOne<Konva.Group>(`.${RG_LAYER_OBJECT_GROUP_NAME}`) ??
-    createObjectGroup(konvaLayer, RG_LAYER_OBJECT_GROUP_NAME);
 
   // We use caching to handle "global" layer opacity, but caching is expensive and we should only do it when required.
   let groupNeedsCache = false;
 
-  const objectIds = rg.objects.map(mapId);
+  const objectIds = region.objects.map(mapId);
-  // Destroy any objects that are no longer in the redux state
+  // Destroy any objects that are no longer in state
-  for (const objectNode of konvaObjectGroup.find(selectVectorMaskObjects)) {
+  for (const entry of mapping.getEntries()) {
-    if (!objectIds.includes(objectNode.id())) {
+    if (!objectIds.includes(entry.id)) {
-      objectNode.destroy();
+      mapping.destroyEntry(entry.id);
       groupNeedsCache = true;
     }
   }
 
-  for (const obj of rg.objects) {
+  for (const obj of region.objects) {
     if (obj.type === 'brush_line') {
-      const konvaBrushLine =
+      const entry = getBrushLine(mapping, obj, RG_LAYER_BRUSH_LINE_NAME);
-        stage.findOne<Konva.Line>(`#${obj.id}`) ?? getBrushLine(obj, konvaObjectGroup, RG_LAYER_BRUSH_LINE_NAME);
 
       // Only update the points if they have changed. The point values are never mutated, they are only added to the
       // array, so checking the length is sufficient to determine if we need to re-cache.
-      if (konvaBrushLine.points().length !== obj.points.length) {
+      if (entry.konvaLine.points().length !== obj.points.length) {
-        konvaBrushLine.points(obj.points);
+        entry.konvaLine.points(obj.points);
         groupNeedsCache = true;
       }
       // Only update the color if it has changed.
-      if (konvaBrushLine.stroke() !== rgbColor) {
+      if (entry.konvaLine.stroke() !== rgbColor) {
-        konvaBrushLine.stroke(rgbColor);
+        entry.konvaLine.stroke(rgbColor);
         groupNeedsCache = true;
       }
     } else if (obj.type === 'eraser_line') {
-      const konvaEraserLine =
+      const entry = getEraserLine(mapping, obj, RG_LAYER_ERASER_LINE_NAME);
-        stage.findOne<Konva.Line>(`#${obj.id}`) ?? getEraserLine(obj, konvaObjectGroup, RG_LAYER_ERASER_LINE_NAME);
 
       // Only update the points if they have changed. The point values are never mutated, they are only added to the
       // array, so checking the length is sufficient to determine if we need to re-cache.
-      if (konvaEraserLine.points().length !== obj.points.length) {
+      if (entry.konvaLine.points().length !== obj.points.length) {
-        konvaEraserLine.points(obj.points);
+        entry.konvaLine.points(obj.points);
         groupNeedsCache = true;
       }
       // Only update the color if it has changed.
-      if (konvaEraserLine.stroke() !== rgbColor) {
+      if (entry.konvaLine.stroke() !== rgbColor) {
-        konvaEraserLine.stroke(rgbColor);
+        entry.konvaLine.stroke(rgbColor);
         groupNeedsCache = true;
       }
     } else if (obj.type === 'rect_shape') {
-      const konvaRectShape =
+      const entry = getRectShape(mapping, obj, RG_LAYER_RECT_SHAPE_NAME);
-        stage.findOne<Konva.Rect>(`#${obj.id}`) ?? createRectShape(obj, konvaObjectGroup, RG_LAYER_RECT_SHAPE_NAME);
 
       // Only update the color if it has changed.
-      if (konvaRectShape.fill() !== rgbColor) {
+      if (entry.konvaRect.fill() !== rgbColor) {
-        konvaRectShape.fill(rgbColor);
+        entry.konvaRect.fill(rgbColor);
         groupNeedsCache = true;
       }
     }
   }
 
   // Only update layer visibility if it has changed.
-  if (konvaLayer.visible() !== rg.isEnabled) {
+  if (mapping.konvaLayer.visible() !== region.isEnabled) {
-    konvaLayer.visible(rg.isEnabled);
+    mapping.konvaLayer.visible(region.isEnabled);
     groupNeedsCache = true;
   }
 
-  if (konvaObjectGroup.getChildren().length === 0) {
+  if (mapping.konvaObjectGroup.getChildren().length === 0) {
     // No objects - clear the cache to reset the previous pixel data
-    konvaObjectGroup.clearCache();
+    mapping.konvaObjectGroup.clearCache();
     return;
   }
 
   const compositingRect =
-    konvaLayer.findOne<Konva.Rect>(`.${COMPOSITING_RECT_NAME}`) ?? createCompositingRect(konvaLayer);
+    mapping.konvaLayer.findOne<Konva.Rect>(`.${COMPOSITING_RECT_NAME}`) ?? createCompositingRect(mapping.konvaLayer);
-  const isSelected = selectedEntityIdentifier?.id === rg.id;
+  const isSelected = selectedEntityIdentifier?.id === region.id;
 
   /**
    * When the group is selected, we use a rect of the selected preview color, composited over the shapes. This allows
@@ -192,54 +192,56 @@ export const renderRGLayer = (
    */
   if (isSelected && tool !== 'move') {
     // We must clear the cache first so Konva will re-draw the group with the new compositing rect
-    if (konvaObjectGroup.isCached()) {
+    if (mapping.konvaObjectGroup.isCached()) {
-      konvaObjectGroup.clearCache();
+      mapping.konvaObjectGroup.clearCache();
     }
     // The user is allowed to reduce mask opacity to 0, but we need the opacity for the compositing rect to work
-    konvaObjectGroup.opacity(1);
+    mapping.konvaObjectGroup.opacity(1);
 
     compositingRect.setAttrs({
       // The rect should be the size of the layer - use the fast method if we don't have a pixel-perfect bbox already
-      ...(!rg.bboxNeedsUpdate && rg.bbox ? rg.bbox : getLayerBboxFast(konvaLayer)),
+      ...(!region.bboxNeedsUpdate && region.bbox ? region.bbox : getLayerBboxFast(mapping.konvaLayer)),
       fill: rgbColor,
       opacity: globalMaskLayerOpacity,
       // Draw this rect only where there are non-transparent pixels under it (e.g. the mask shapes)
       globalCompositeOperation: 'source-in',
       visible: true,
       // This rect must always be on top of all other shapes
-      zIndex: konvaObjectGroup.getChildren().length,
+      zIndex: mapping.konvaObjectGroup.getChildren().length,
     });
   } else {
     // The compositing rect should only be shown when the layer is selected.
     compositingRect.visible(false);
     // Cache only if needed - or if we are on this code path and _don't_ have a cache
-    if (groupNeedsCache || !konvaObjectGroup.isCached()) {
+    if (groupNeedsCache || !mapping.konvaObjectGroup.isCached()) {
-      konvaObjectGroup.cache();
+      mapping.konvaObjectGroup.cache();
     }
     // Updating group opacity does not require re-caching
-    konvaObjectGroup.opacity(globalMaskLayerOpacity);
+    mapping.konvaObjectGroup.opacity(globalMaskLayerOpacity);
   }
 
-  const bboxRect = konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(rg, konvaLayer);
+  // const bboxRect =
+  //   regionMap.konvaLayer.findOne<Konva.Rect>(`.${LAYER_BBOX_NAME}`) ?? createBboxRect(rg, regionMap.konvaLayer);
 
-  if (rg.bbox) {
+  // if (rg.bbox) {
-    const active = !rg.bboxNeedsUpdate && isSelected && tool === 'move';
+  //   const active = !rg.bboxNeedsUpdate && isSelected && tool === 'move';
-    bboxRect.setAttrs({
+  //   bboxRect.setAttrs({
-      visible: active,
+  //     visible: active,
-      listening: active,
+  //     listening: active,
-      x: rg.bbox.x,
+  //     x: rg.bbox.x,
-      y: rg.bbox.y,
+  //     y: rg.bbox.y,
-      width: rg.bbox.width,
+  //     width: rg.bbox.width,
-      height: rg.bbox.height,
+  //     height: rg.bbox.height,
-      stroke: isSelected ? BBOX_SELECTED_STROKE : '',
+  //     stroke: isSelected ? BBOX_SELECTED_STROKE : '',
-    });
+  //   });
-  } else {
+  // } else {
-    bboxRect.visible(false);
+  //   bboxRect.visible(false);
-  }
+  // }
 };
 
 export const renderRegions = (
   stage: Konva.Stage,
+  regionMap: EntityToKonvaMap,
   regions: RegionEntity[],
   maskOpacity: number,
   tool: Tool,
@@ -247,12 +249,12 @@ export const renderRegions = (
   onPosChanged?: (arg: PosChangedArg, entityType: CanvasEntity['type']) => void
 ): void => {
   // Destroy nonexistent layers
-  for (const konvaLayer of stage.find<Konva.Layer>(`.${RG_LAYER_NAME}`)) {
+  for (const mapping of regionMap.getMappings()) {
-    if (!regions.find((rg) => rg.id === konvaLayer.id())) {
+    if (!regions.find((rg) => rg.id === mapping.id)) {
-      konvaLayer.destroy();
+      regionMap.destroyMapping(mapping.id);
     }
   }
   for (const rg of regions) {
-    renderRGLayer(stage, rg, maskOpacity, tool, selectedEntityIdentifier, onPosChanged);
+    renderRGLayer(stage, regionMap, rg, maskOpacity, tool, selectedEntityIdentifier, onPosChanged);
   }
 };