Use datastore segment index when loading ad-hoc meshes

frontend/javascripts/admin/rest_api.ts

@@ -80,6 +80,7 @@ import type { DatasourceConfiguration } from "types/schemas/datasource.types";
 import type { AnnotationTypeFilterEnum, LOG_LEVELS, Vector3 } from "viewer/constants";
 import Constants, { ControlModeEnum, AnnotationStateFilterEnum } from "viewer/constants";
 import type BoundingBox from "viewer/model/bucket_data_handling/bounding_box";
+import type { LayerSourceInfo } from "viewer/model/bucket_data_handling/wkstore_adapter";
 import {
   parseProtoAnnotation,
   parseProtoListOfLong,
@@ -855,13 +856,28 @@ export function hasSegmentIndexInDataStore(
   );
 }
 
+export const hasSegmentIndexInDataStoreCached = _.memoize(hasSegmentIndexInDataStore, (...args) =>
+  args.join("::"),
+);
+
+export function getVolumeRequestUrl(layerSourceInfo: LayerSourceInfo) {
+  const { dataset, annotation, visibleSegmentationLayer, tracingId } = layerSourceInfo;
+  if (annotation == null || tracingId == null) {
+    return `${dataset.dataStore.url}/data/datasets/${dataset.id}/layers/${visibleSegmentationLayer?.name}`;
+  } else {
+    const tracingStoreHost = annotation?.tracingStore.url;
+    return `${tracingStoreHost}/tracings/volume/${tracingId}`;
+  }
+}
+
 export function getSegmentVolumes(
-  requestUrl: string,
+  layerSourceInfo: LayerSourceInfo,
   mag: Vector3,
   segmentIds: Array<number>,
   additionalCoordinates: AdditionalCoordinate[] | undefined | null,
   mappingName: string | null | undefined,
 ): Promise<number[]> {
+  const requestUrl = getVolumeRequestUrl(layerSourceInfo);
   return doWithToken((token) =>
     Request.sendJSONReceiveJSON(`${requestUrl}/segmentStatistics/volume?token=${token}`, {
       data: { additionalCoordinates, mag, segmentIds, mappingName },
@@ -871,12 +887,13 @@ export function getSegmentVolumes(
 }
 
 export function getSegmentBoundingBoxes(
-  requestUrl: string,
+  layerSourceInfo: LayerSourceInfo,
   mag: Vector3,
   segmentIds: Array<number>,
   additionalCoordinates: AdditionalCoordinate[] | undefined | null,
   mappingName: string | null | undefined,
 ): Promise<Array<{ topLeft: Vector3; width: number; height: number; depth: number }>> {
+  const requestUrl = getVolumeRequestUrl(layerSourceInfo);
   return doWithToken((token) =>
     Request.sendJSONReceiveJSON(`${requestUrl}/segmentStatistics/boundingBox?token=${token}`, {
       data: { additionalCoordinates, mag, segmentIds, mappingName },
@@ -1880,12 +1897,13 @@ type MeshRequest = {
 };
 
 export function computeAdHocMesh(
-  requestUrl: string,
+  dataSourceInfo: LayerSourceInfo,
   meshRequest: MeshRequest,
 ): Promise<{
   buffer: ArrayBuffer;
   neighbors: Array<number>;
 }> {
+  const requestUrl = getVolumeRequestUrl(dataSourceInfo);
   const {
     positionWithPadding,
     additionalCoordinates,
@@ -1927,23 +1945,25 @@ export function computeAdHocMesh(
 }
 
 export function getBucketPositionsForAdHocMesh(
-  tracingStoreUrl: string,
-  tracingId: string,
+  dataSourceInfo: LayerSourceInfo,
   segmentId: number,
   cubeSize: Vector3,
   mag: Vector3,
   additionalCoordinates: AdditionalCoordinate[] | null | undefined,
+  mappingName: string | null | undefined,
 ): Promise<Vector3[]> {
+  const requestUrl = getVolumeRequestUrl(dataSourceInfo);
   return doWithToken(async (token) => {
     const params = new URLSearchParams();
     params.set("token", token);
     const positions = await Request.sendJSONReceiveJSON(
-      `${tracingStoreUrl}/tracings/volume/${tracingId}/segmentIndex/${segmentId}?${params}`,
+      `${requestUrl}/segmentIndex/${segmentId}?${params}`,
       {
         data: {
           cubeSize,
           mag,
           additionalCoordinates,
+          mappingName,
         },
         method: "POST",
       },

frontend/javascripts/viewer/model/bucket_data_handling/wkstore_adapter.ts

@@ -5,7 +5,12 @@ import { parseMaybe } from "libs/utils";
 import WebworkerPool from "libs/webworker_pool";
 import window from "libs/window";
 import _ from "lodash";
-import type { AdditionalCoordinate } from "types/api_types";
+import type {
+  APIDataLayer,
+  APIDataset,
+  APISegmentationLayer,
+  AdditionalCoordinate,
+} from "types/api_types";
 import type { BucketAddress, Vector3 } from "viewer/constants";
 import constants, { MappingStatusEnum } from "viewer/constants";
 import {
@@ -22,14 +27,22 @@ import type { DataBucket } from "viewer/model/bucket_data_handling/bucket";
 import { bucketPositionToGlobalAddress } from "viewer/model/helpers/position_converter";
 import type { UpdateActionWithoutIsolationRequirement } from "viewer/model/sagas/volume/update_actions";
 import { updateBucket } from "viewer/model/sagas/volume/update_actions";
-import type { DataLayerType, VolumeTracing } from "viewer/store";
+import type { DataLayerType, StoreAnnotation, VolumeTracing } from "viewer/store";
 import Store from "viewer/store";
 import ByteArraysToLz4Base64Worker from "viewer/workers/byte_arrays_to_lz4_base64.worker";
 import { createWorker } from "viewer/workers/comlink_wrapper";
 import DecodeFourBitWorker from "viewer/workers/decode_four_bit.worker";
 import { getGlobalDataConnectionInfo } from "../data_connection_info";
 import type { MagInfo } from "../helpers/mag_info";
 
+export type LayerSourceInfo = {
+  dataset: APIDataset;
+  annotation: StoreAnnotation | null;
+  tracingId: string | undefined;
+  visibleSegmentationLayer: APISegmentationLayer | APIDataLayer | undefined | null;
+  forceUsingDataStore?: boolean | undefined | null;
+};
+
 const decodeFourBit = createWorker(DecodeFourBitWorker);
 
 // For 32-bit buckets with 32^3 voxels, a COMPRESSION_BATCH_SIZE of

frontend/javascripts/viewer/model/sagas/meshes/ad_hoc_mesh_saga.ts

@@ -1,6 +1,7 @@
 import {
   computeAdHocMesh,
   getBucketPositionsForAdHocMesh,
+  hasSegmentIndexInDataStoreCached,
   sendAnalyticsEvent,
 } from "admin/rest_api";
 import ThreeDMap from "libs/ThreeDMap";
@@ -41,13 +42,14 @@ import {
   type LoadAdHocMeshAction,
   loadPrecomputedMeshAction,
 } from "viewer/model/actions/segmentation_actions";
+import type { LayerSourceInfo } from "viewer/model/bucket_data_handling/wkstore_adapter";
 import type DataLayer from "viewer/model/data_layer";
 import type { MagInfo } from "viewer/model/helpers/mag_info";
 import { zoomedAddressToAnotherZoomStepWithInfo } from "viewer/model/helpers/position_converter";
 import type { Saga } from "viewer/model/sagas/effect-generators";
 import { select } from "viewer/model/sagas/effect-generators";
 import { Model } from "viewer/singletons";
-import Store from "viewer/store";
+import Store, { type StoreDataset, type VolumeTracing } from "viewer/store";
 import { getAdditionalCoordinatesAsString } from "../../accessors/flycam_accessor";
 import { ensureSceneControllerReady, ensureWkReady } from "../ready_sagas";
 
@@ -293,6 +295,25 @@ function removeMeshWithoutVoxels(
   }
 }
 
+function* getUsePositionsFromSegmentIndex(
+  volumeTracing: VolumeTracing | null | undefined,
+  dataset: StoreDataset,
+  layerName: string,
+  maybeTracingId?: string | null,
+): Saga<boolean> {
+  if (volumeTracing == null) {
+    return yield* call(
+      hasSegmentIndexInDataStoreCached,
+      dataset.dataStore.url,
+      dataset.id,
+      layerName,
+    );
+  }
+  return (
+    volumeTracing?.hasSegmentIndex && !volumeTracing.hasEditableMapping && maybeTracingId != null
+  );
+}
+
 function* loadFullAdHocMesh(
   layer: DataLayer,
   segmentId: number,
@@ -320,36 +341,48 @@ function* loadFullAdHocMesh(
   yield* put(startedLoadingMeshAction(layer.name, segmentId));
 
   const cubeSize = marchingCubeSizeInTargetMag();
-  const tracingStoreHost = yield* select((state) => state.annotation.tracingStore.url);
+  const dataset = yield* select((state) => state.dataset);
   const mag = magInfo.getMagByIndexOrThrow(zoomStep);
 
   const volumeTracing = yield* select((state) => getActiveSegmentationTracing(state));
+  const annotation = yield* select((state) => state.annotation);
   const visibleSegmentationLayer = yield* select((state) => getVisibleSegmentationLayer(state));
   // Fetch from datastore if no volumetracing ...
-  let useDataStore = volumeTracing == null || visibleSegmentationLayer?.tracingId == null;
+  let forceUsingDataStore = volumeTracing == null || visibleSegmentationLayer?.tracingId == null;
   if (meshExtraInfo.useDataStore != null) {
     // ... except if the caller specified whether to use the data store ...
-    useDataStore = meshExtraInfo.useDataStore;
+    forceUsingDataStore = meshExtraInfo.useDataStore;
   } else if (volumeTracing?.hasEditableMapping) {
     // ... or if an editable mapping is active.
-    useDataStore = false;
+    forceUsingDataStore = false;
   }
 
-  // Segment stats can only be used for volume tracings that have a segment index
+  // Segment stats can only be used for segmentation layers that have a segment index
   // and that don't have editable mappings.
-  const usePositionsFromSegmentIndex =
-    volumeTracing?.hasSegmentIndex &&
-    !volumeTracing.hasEditableMapping &&
-    visibleSegmentationLayer?.tracingId != null;
+  const usePositionsFromSegmentIndex = getUsePositionsFromSegmentIndex(
+    volumeTracing,
+    dataset,
+    layer.name,
+    visibleSegmentationLayer?.tracingId,
+  );
+
+  const dataSourceInfo: LayerSourceInfo = {
+    dataset,
+    annotation,
+    tracingId: visibleSegmentationLayer?.tracingId,
+    visibleSegmentationLayer,
+    forceUsingDataStore,
+  };
+
   let positionsToRequest = usePositionsFromSegmentIndex
     ? yield* getChunkPositionsFromSegmentIndex(
-        tracingStoreHost,
-        layer,
+        dataSourceInfo,
         segmentId,
         cubeSize,
         mag,
         clippedPosition,
         additionalCoordinates,
+        mappingName,
       )
     : [clippedPosition];
 
@@ -373,7 +406,7 @@ function* loadFullAdHocMesh(
       magInfo,
       isInitialRequest,
       removeExistingMesh && isInitialRequest,
-      useDataStore,
+      dataSourceInfo,
       !usePositionsFromSegmentIndex,
     );
     isInitialRequest = false;
@@ -390,22 +423,22 @@ function* loadFullAdHocMesh(
 }
 
 function* getChunkPositionsFromSegmentIndex(
-  tracingStoreHost: string,
-  layer: DataLayer,
+  dataSourceInfo: LayerSourceInfo,
   segmentId: number,
   cubeSize: Vector3,
   mag: Vector3,
   clippedPosition: Vector3,
   additionalCoordinates: AdditionalCoordinate[] | null | undefined,
+  mappingName: string | null | undefined,
 ) {
   const targetMagPositions = yield* call(
     getBucketPositionsForAdHocMesh,
-    tracingStoreHost,
-    layer.name,
+    dataSourceInfo,
     segmentId,
     cubeSize,
     mag,
     additionalCoordinates,
+    mappingName,
   );
   const mag1Positions = targetMagPositions.map((pos) => V3.scale3(pos, mag));
   return sortByDistanceTo(mag1Positions, clippedPosition) as Vector3[];
@@ -424,7 +457,7 @@ function* maybeLoadMeshChunk(
   magInfo: MagInfo,
   isInitialRequest: boolean,
   removeExistingMesh: boolean,
-  useDataStore: boolean,
+  dataSourceInfo: LayerSourceInfo,
   findNeighbors: boolean,
 ): Saga<Vector3[]> {
   const additionalCoordinates = yield* select((state) => state.flycam.additionalCoordinates);
@@ -445,17 +478,10 @@ function* maybeLoadMeshChunk(
   batchCounterPerSegment[segmentId]++;
   threeDMap.set(paddedPositionWithinLayer, true);
   const scaleFactor = yield* select((state) => state.dataset.dataSource.scale.factor);
-  const dataStoreHost = yield* select((state) => state.dataset.dataStore.url);
-  const datasetId = yield* select((state) => state.dataset.id);
-  const tracingStoreHost = yield* select((state) => state.annotation.tracingStore.url);
-  const dataStoreUrl = `${dataStoreHost}/data/datasets/${datasetId}/layers/${
-    layer.fallbackLayer != null ? layer.fallbackLayer : layer.name
-  }`;
-  const tracingStoreUrl = `${tracingStoreHost}/tracings/volume/${layer.name}`;
 
   if (isInitialRequest) {
     sendAnalyticsEvent("request_isosurface", {
-      mode: useDataStore ? "view" : "annotation",
+      mode: dataSourceInfo.forceUsingDataStore ? "view" : "annotation",
     });
   }
 
@@ -472,7 +498,7 @@ function* maybeLoadMeshChunk(
           context: null,
           fn: computeAdHocMesh,
         },
-        useDataStore ? dataStoreUrl : tracingStoreUrl,
+        dataSourceInfo,
         {
           positionWithPadding: paddedPositionWithinLayer,
           additionalCoordinates: additionalCoordinates || undefined,

frontend/javascripts/viewer/view/context_menu.tsx

@@ -146,10 +146,7 @@ import type {
 import { deleteNodeAsUserAction } from "viewer/model/actions/skeletontracing_actions_with_effects";
 import { type MutableNode, type Tree, TreeMap } from "viewer/model/types/tree_types";
 import Store from "viewer/store";
-import {
-  getVolumeRequestUrl,
-  withMappingActivationConfirmation,
-} from "viewer/view/right-border-tabs/segments_tab/segments_view_helper";
+import { withMappingActivationConfirmation } from "viewer/view/right-border-tabs/segments_tab/segments_view_helper";
 import { LayoutEvents, layoutEmitter } from "./layouting/layout_persistence";
 import { LoadMeshMenuItemLabel } from "./right-border-tabs/segments_tab/load_mesh_menu_item_label";
 
@@ -1690,26 +1687,26 @@ function ContextMenuInner() {
       if (visibleSegmentationLayer == null || !isSegmentIndexAvailable) return [];
       const tracingId = volumeTracing?.tracingId;
       const additionalCoordinates = flycam.additionalCoordinates;
-      const requestUrl = getVolumeRequestUrl(
+      const infoForRequestUrl = {
         dataset,
         annotation,
         tracingId,
         visibleSegmentationLayer,
-      );
+      };
       const magInfo = getMagInfo(visibleSegmentationLayer.resolutions);
       const layersFinestMag = magInfo.getFinestMag();
       const voxelSize = dataset.dataSource.scale;
 
       try {
         const [segmentSize] = await getSegmentVolumes(
-          requestUrl,
+          infoForRequestUrl,
           layersFinestMag,
           [clickedSegmentOrMeshId],
           additionalCoordinates,
           mappingName,
         );
         const [boundingBoxInRequestedMag] = await getSegmentBoundingBoxes(
-          requestUrl,
+          infoForRequestUrl,
           layersFinestMag,
           [clickedSegmentOrMeshId],
           additionalCoordinates,