Only reload affected areas when proofreading

CHANGELOG.unreleased.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - In addition to drag and drop, the selected tree(s) in the Skeleton tab can also be moved into another group by right-clicking the target group and selecting "Move selected tree(s) here". [#7005](https://github.com/scalableminds/webknossos/pull/7005)
 
 ### Changed
+- Improved speed of proofreading by only reloading affected areas after a split or merge. [#7050](https://github.com/scalableminds/webknossos/pull/7050)
 
 ### Fixed
 - Fixed that changing a segment color could lead to a crash. [#7000](https://github.com/scalableminds/webknossos/pull/7000)

frontend/javascripts/oxalis/api/api_latest.ts

@@ -2,7 +2,7 @@ import PriorityQueue from "js-priority-queue";
 // @ts-expect-error ts-migrate(7016) FIXME: Could not find a declaration file for module 'twee... Remove this comment to see the full error message
 import TWEEN from "tween.js";
 import _ from "lodash";
-import type { Bucket } from "oxalis/model/bucket_data_handling/bucket";
+import type { Bucket, DataBucket } from "oxalis/model/bucket_data_handling/bucket";
 import { getConstructorForElementClass } from "oxalis/model/bucket_data_handling/bucket";
 import { APICompoundType, APICompoundTypeEnum, ElementClass } from "types/api_flow_types";
 import { InputKeyboardNoLoop } from "libs/input";
@@ -1182,16 +1182,23 @@ class DataApi {
 
   /**
    * Invalidates all downloaded buckets of the given layer so that they are reloaded.
+   * If an additional predicate is passed, each bucket is checked to see whether
+   * it should be reloaded. Note that buckets that are in a REQUESTED state (i.e.,
+   * currently being queued or downloaded) will always be reloaded by cancelling and rescheduling
+   * the request.
    */
-  async reloadBuckets(layerName: string): Promise<void> {
+  async reloadBuckets(
+    layerName: string,
+    predicateFn?: (bucket: DataBucket) => boolean,
+  ): Promise<void> {
     await Promise.all(
       Utils.values(this.model.dataLayers).map(async (dataLayer: DataLayer) => {
         if (dataLayer.name === layerName) {
           if (dataLayer.cube.isSegmentation) {
             await Model.ensureSavedState();
           }
 
-          dataLayer.cube.collectAllBuckets();
+          dataLayer.cube.collectBucketsIf(predicateFn || (() => true));
           dataLayer.layerRenderingManager.refresh();
         }
       }),

frontend/javascripts/oxalis/model/bucket_data_handling/bucket.ts

@@ -121,6 +121,7 @@ export const bucketsAlreadyInUndoState: Set<Bucket> = new Set();
 export function markVolumeTransactionEnd() {
   bucketsAlreadyInUndoState.clear();
 }
+
 export class DataBucket {
   type: "data" = "data";
   elementClass: ElementClass;
@@ -147,6 +148,11 @@ export class DataBucket {
   throttledTriggerLabeled: () => void;
   emitter: Emitter;
   maybeUnmergedBucketLoadedPromise: MaybeUnmergedBucketLoadedPromise;
+  // Especially, for segmentation buckets, it can be interesting to
+  // know whether a certain ID is contained in this bucket. To
+  // speed up such requests a cached set of the contained values
+  // can be stored in cachedValueSet.
+  cachedValueSet: Set<number | BigInt> | null = null;
 
   constructor(
     elementClass: ElementClass,
@@ -232,6 +238,7 @@ export class DataBucket {
     // so that at least the big memory hog is tamed (unfortunately,
     // this doesn't help against references which point directly to this.data)
     this.data = null;
+    this.invalidateValueSet();
     this.trigger("bucketCollected");
     // Remove all event handlers (see https://github.com/ai/nanoevents#remove-all-listeners)
     this.emitter.events = {};
@@ -375,6 +382,7 @@ export class DataBucket {
 
   setData(newData: BucketDataArray, newPendingOperations: Array<(arg0: BucketDataArray) => void>) {
     this.data = newData;
+    this.invalidateValueSet();
     this.pendingOperations = newPendingOperations;
     this.dirty = true;
     this.endDataMutation();
@@ -589,6 +597,7 @@ export class DataBucket {
         } else {
           this.data = data;
         }
+        this.invalidateValueSet();
 
         this.state = BucketStateEnum.LOADED;
         this.trigger("bucketLoaded", data);
@@ -600,6 +609,22 @@ export class DataBucket {
     }
   }
 
+  private invalidateValueSet() {
+    this.cachedValueSet = null;
+  }
+
+  private recomputeValueSet() {
+    // @ts-ignore The Set constructor accepts null and BigUint64Arrays just fine.
+    this.cachedValueSet = new Set(this.data);
+  }
+
+  containsValue(value: number | BigInt): boolean {
+    if (this.cachedValueSet == null) {
+      this.recomputeValueSet();
+    }
+    return this.cachedValueSet!.has(value);
+  }
+
   markAsPushed(): void {
     switch (this.state) {
       case BucketStateEnum.LOADED:

frontend/javascripts/oxalis/model/bucket_data_handling/data_cube.ts

@@ -323,15 +323,27 @@ class DataCube {
   }
 
   collectAllBuckets(): void {
+    this.collectBucketsIf(() => true);
+  }
+
+  collectBucketsIf(predicateFn: (bucket: DataBucket) => boolean): void {
     this.pullQueue.clear();
     this.pullQueue.abortRequests();
 
+    const notCollectedBuckets = [];
     for (const bucket of this.buckets) {
-      this.collectBucket(bucket);
+      // If a bucket is requested, collect it independently of the predicateFn,
+      // because the pullQueue was already cleared (meaning the bucket is in a
+      // requested state, but will never be filled with data).
+      if (bucket.state === "REQUESTED" || predicateFn(bucket)) {
+        this.collectBucket(bucket);
+      } else {
+        notCollectedBuckets.push(bucket);
+      }
     }
 
-    this.buckets = [];
-    this.bucketIterator = 0;
+    this.buckets = notCollectedBuckets;
+    this.bucketIterator = notCollectedBuckets.length;
   }
 
   collectBucket(bucket: DataBucket): void {

frontend/javascripts/oxalis/model/sagas/proofread_saga.ts

@@ -309,8 +309,9 @@ function* splitOrMergeOrMinCutAgglomerate(
   yield* call([Model, Model.ensureSavedState]);
 
   /* Reload the segmentation */
-
-  yield* call([api.data, api.data.reloadBuckets], layerName);
+  yield* call([api.data, api.data.reloadBuckets], layerName, (bucket) =>
+    bucket.containsValue(targetAgglomerateId),
+  );
 
   const [newSourceAgglomerateId, newTargetAgglomerateId] = yield* all([
     call(getDataValue, sourceNodePosition),
@@ -536,8 +537,9 @@ function* handleProofreadMergeOrMinCut(
   yield* call([Model, Model.ensureSavedState]);
 
   /* Reload the segmentation */
-
-  yield* call([api.data, api.data.reloadBuckets], layerName);
+  yield* call([api.data, api.data.reloadBuckets], layerName, (bucket) =>
+    bucket.containsValue(targetAgglomerateId),
+  );
 
   const [newSourceAgglomerateId, newTargetAgglomerateId] = yield* all([
     call(getDataValue, sourcePosition),