fix(segmentation): undoable overlap labelmap strokes, faster eager stroke resolution, brush hover guard

packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapEditTransaction.ts

@@ -13,10 +13,20 @@ import {
   getSegmentsOnLabelmap,
 } from './labelmapSegmentBindings';
 import { moveSegmentToPrivateLabelmap as defaultMoveSegmentToPrivateLabelmap } from './privateLabelmap';
+import type { LabelmapRestoreStep } from '../../../utilities/segmentation/createLabelmapMemo';
+
+/** A batch of cross-layer erases performed by one strategy fill call, with
+ *  everything needed to reverse/replay it on undo/redo. */
+type CrossLayerEraseRecord = {
+  voxelManager: Types.IVoxelManager<number>;
+  labelValue: number;
+  indices: number[];
+};
 
 type MoveSegmentToPrivateLabelmap = (
   segmentation: Segmentation,
-  segmentIndex: number
+  segmentIndex: number,
+  options?: { moveStepCallback?: (step: LabelmapRestoreStep) => void }
 ) => LabelmapLayer | undefined;
 
 type BeginLabelmapEditTransactionOptions = {
@@ -39,6 +49,10 @@ type LabelmapEditTransaction = {
   protectedSegmentIndices: number[];
   crossLayerEraseBindings: SegmentLabelmapBindingState[];
   movedSegment: boolean;
+  /** Undo/redo step for a segment move performed by this transaction
+   *  (layer registration + bulk voxel move + binding change), for recording
+   *  on the stroke's memo. */
+  moveStep?: LabelmapRestoreStep;
 };
 
 type ResolveLabelmapLayerEditTargetOptions = {
@@ -62,6 +76,9 @@ type EraseLabelmapEditTransactionOptions = {
   isInObject: (point: Types.Point3) => boolean;
   isInObjectBoundsIJK?: Types.BoundsIJK;
   imageId?: string;
+  /** When provided, receives the erased voxels of each touched layer so the
+   *  caller can record them for undo/redo. */
+  crossLayerEraseCallback?: (record: CrossLayerEraseRecord) => void;
 };
 
 function getProtectedSegmentIndicesForLayer(
@@ -197,13 +214,16 @@ function beginLabelmapEditTransaction(
     }
   );
 
+  let moveStep: LabelmapRestoreStep | undefined;
+
   if (shouldMoveSegment) {
     const moveSegmentToPrivateLabelmap =
       options.moveSegmentToPrivateLabelmap ??
       defaultMoveSegmentToPrivateLabelmap;
     const privateLayer = moveSegmentToPrivateLabelmap(
       segmentation,
-      segmentIndex
+      segmentIndex,
+      { moveStepCallback: (step) => (moveStep = step) }
     );
     const privateBinding = getSegmentBinding(segmentation, segmentIndex);
 
@@ -234,6 +254,7 @@ function beginLabelmapEditTransaction(
       overwriteSegmentIndices
     ),
     movedSegment,
+    moveStep,
   };
 }
 
@@ -329,6 +350,7 @@ function eraseVolumeLayer(
     return;
   }
 
+  const erasedIndices: number[] = [];
   volume.voxelManager.forEach(
     ({ value, index, pointIJK }) => {
       if (value !== binding.labelValue) {
@@ -343,13 +365,22 @@ function eraseVolumeLayer(
       }
 
       volume.voxelManager.setAtIndex(index, 0);
+      erasedIndices.push(index);
     },
     {
       imageData: volume.imageData,
       boundsIJK: options.isInObjectBoundsIJK,
     }
   );
 
+  if (erasedIndices.length) {
+    options.crossLayerEraseCallback?.({
+      voxelManager: volume.voxelManager as Types.IVoxelManager<number>,
+      labelValue: binding.labelValue,
+      indices: erasedIndices,
+    });
+  }
+
   volume.voxelManager
     ?.getArrayOfModifiedSlices?.()
     ?.forEach((sliceIndex) => modifiedSlices.add(sliceIndex));
@@ -373,6 +404,7 @@ function eraseStackLayer(
     return;
   }
 
+  const erasedIndices: number[] = [];
   voxelManager.forEach(
     ({ value, index, pointIJK }) => {
       if (value !== binding.labelValue) {
@@ -387,13 +419,22 @@ function eraseStackLayer(
       }
 
       voxelManager.setAtIndex(index, 0);
+      erasedIndices.push(index);
     },
     {
       imageData: options.referenceImageData,
       boundsIJK: options.isInObjectBoundsIJK,
     }
   );
 
+  if (erasedIndices.length) {
+    options.crossLayerEraseCallback?.({
+      voxelManager,
+      labelValue: binding.labelValue,
+      indices: erasedIndices,
+    });
+  }
+
   const currentSlice = stackViewport.getCurrentImageIdIndex?.();
   if (typeof currentSlice === 'number') {
     modifiedSlices.add(currentSlice);
@@ -440,6 +481,7 @@ export {
 };
 export type {
   BeginLabelmapEditTransactionOptions,
+  CrossLayerEraseRecord,
   EraseLabelmapEditTransactionOptions,
   LabelmapEditTransaction,
   LabelmapLayerEditTarget,

packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts

@@ -215,6 +215,41 @@ describe('LabelmapImageReferenceResolver', () => {
       );
       expect(result).toBe('lm-a-0');
     });
+
+    it('replays the cached result for the matching viewport scan key', () => {
+      const segmentation = makeSegmentation();
+      getSegmentation.mockReturnValue(segmentation);
+
+      const viewport1 = makeStackViewport({
+        id: 'vp-1',
+        getCurrentImageId: jest.fn(() => 'ref-0'),
+        isReferenceViewable: jest.fn(
+          (ref) => ref.referencedImageId === 'lm-a-0'
+        ),
+      });
+      const viewport2 = makeStackViewport({
+        id: 'vp-2',
+        getCurrentImageId: jest.fn(() => 'ref-0'),
+        isReferenceViewable: jest.fn(
+          (ref) => ref.referencedImageId === 'lm-a-1'
+        ),
+      });
+      getEnabledElementByViewportId.mockImplementation((viewportId) => ({
+        viewport: viewportId === 'vp-1' ? viewport1 : viewport2,
+      }));
+
+      expect(
+        resolver.updateLabelmapSegmentationImageReferences('vp-1', 'seg-1')
+      ).toBe('lm-a-0');
+      expect(
+        resolver.updateLabelmapSegmentationImageReferences('vp-2', 'seg-1')
+      ).toBe('lm-a-1');
+      expect(
+        resolver.updateLabelmapSegmentationImageReferences('vp-1', 'seg-1')
+      ).toBe('lm-a-0');
+      expect(viewport1.isReferenceViewable).toHaveBeenCalledTimes(2);
+      expect(viewport2.isReferenceViewable).toHaveBeenCalledTimes(2);
+    });
   });
 
   describe('getCurrentLabelmapImageIdsForViewport', () => {

packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts

@@ -31,6 +31,17 @@ class LabelmapImageReferenceResolver {
   // labelmapImageIdReferenceMap, so removeSegmentation can purge entries in
   // O(k) without scanning the whole map.
   private readonly keysBySegmentationId = new Map<string, Set<string>>();
+  // Memo of already-performed isReferenceViewable scans. Each scan resolves
+  // every labelmap imageId against the viewport's current image via full
+  // metadata resolution (image plane, slice basis), which costs tens of ms
+  // for multi-layer labelmaps - and the segmentation data-modified pipeline
+  // re-enters this on EVERY brush/erase drag event. Scan keys include the
+  // viewport id and a signature of the labelmap imageId set, so adding/removing
+  // a layer (which changes the set) naturally forces a fresh scan.
+  private readonly completedScanResultsBySegmentationId = new Map<
+    string,
+    Map<string, string | undefined>
+  >();
 
   constructor(getSegmentation: GetSegmentation) {
     this.getSegmentation = getSegmentation;
@@ -40,10 +51,12 @@ class LabelmapImageReferenceResolver {
     this.stackLabelmapImageIdReferenceMap.clear();
     this.labelmapImageIdReferenceMap.clear();
     this.keysBySegmentationId.clear();
+    this.completedScanResultsBySegmentationId.clear();
   }
 
   removeSegmentation(segmentationId: string): void {
     this.stackLabelmapImageIdReferenceMap.delete(segmentationId);
+    this.completedScanResultsBySegmentationId.delete(segmentationId);
     const keys = this.keysBySegmentationId.get(segmentationId);
     if (keys) {
       for (const key of keys) {
@@ -142,11 +155,71 @@ class LabelmapImageReferenceResolver {
       return;
     }
 
-    return this.updateLabelmapSegmentationReferences(
+    const viewport = enabledElement.viewport as Types.IStackViewport;
+    const scanKey = this.generateScanKey('current', viewport, labelmapImageIds);
+    if (scanKey && this.hasCompletedScan(segmentationId, scanKey)) {
+      return this.getCompletedScanResult(segmentationId, scanKey);
+    }
+
+    const result = this.updateLabelmapSegmentationReferences(
       segmentationId,
-      enabledElement.viewport as Types.IStackViewport,
+      viewport,
       labelmapImageIds
     );
+
+    if (scanKey) {
+      this.markCompletedScan(segmentationId, scanKey, result);
+    }
+
+    return result;
+  }
+
+  /** Signature of one isReferenceViewable scan: the viewport's current image
+   *  plus a cheap fingerprint of the labelmap imageId set (count + endpoints),
+   *  so layer additions/removals invalidate naturally. */
+  private generateScanKey(
+    kind: 'current' | 'all',
+    viewport: Types.IStackViewport,
+    labelmapImageIds: string[]
+  ): string | undefined {
+    const referenceImageId =
+      typeof viewport.getCurrentImageId === 'function'
+        ? viewport.getCurrentImageId()
+        : undefined;
+    if (!referenceImageId) {
+      return;
+    }
+    return `${kind}|${viewport.id}|${referenceImageId}|${
+      labelmapImageIds.length
+    }|${labelmapImageIds[0]}|${labelmapImageIds[labelmapImageIds.length - 1]}`;
+  }
+
+  private hasCompletedScan(segmentationId: string, scanKey: string): boolean {
+    return !!this.completedScanResultsBySegmentationId
+      .get(segmentationId)
+      ?.has(scanKey);
+  }
+
+  private getCompletedScanResult(
+    segmentationId: string,
+    scanKey: string
+  ): string | undefined {
+    return this.completedScanResultsBySegmentationId
+      .get(segmentationId)
+      ?.get(scanKey);
+  }
+
+  private markCompletedScan(
+    segmentationId: string,
+    scanKey: string,
+    result?: string
+  ): void {
+    let results = this.completedScanResultsBySegmentationId.get(segmentationId);
+    if (!results) {
+      results = new Map();
+      this.completedScanResultsBySegmentationId.set(segmentationId, results);
+    }
+    results.set(scanKey, result);
   }
 
   getCurrentLabelmapImageIdsForViewport(
@@ -382,6 +455,18 @@ class LabelmapImageReferenceResolver {
 
     const stackViewport = enabledElement.viewport as Types.IStackViewport;
 
+    const scanKey = this.generateScanKey(
+      'all',
+      stackViewport,
+      labelmapImageIds
+    );
+    if (scanKey && this.hasCompletedScan(segmentationId, scanKey)) {
+      return;
+    }
+    if (scanKey) {
+      this.markCompletedScan(segmentationId, scanKey);
+    }
+
     this.updateLabelmapSegmentationReferences(
       segmentationId,
       stackViewport,