Feature/voxel annotation

.gitignore

@@ -26,3 +26,6 @@ tsconfig.tsbuildinfo
 -/docs/python/api/index.rst
 /docs/python/api/*.rst
 /.vite
+/.idea
+/.local
+/.env

package.json

@@ -485,6 +485,37 @@
       "neuroglancer/kvstore/s3:disabled": "./src/util/false.ts",
       "default": "./src/kvstore/s3/register_backend.ts"
     },
+    "#kvstore/opfs/register_frontend": {
+      "neuroglancer/kvstore/opfs:enabled": "./src/kvstore/opfs/register_frontend.ts",
+      "neuroglancer/kvstore:none_by_default": "./src/util/false.ts",
+      "neuroglancer/kvstore/opfs:disabled": "./src/util/false.ts",
+      "default": "./src/kvstore/opfs/register_frontend.ts"
+    },
+    "#kvstore/opfs/register_backend": {
+      "neuroglancer/kvstore/opfs:enabled": "./src/kvstore/opfs/register_backend.ts",
+      "neuroglancer/kvstore:none_by_default": "./src/util/false.ts",
+      "neuroglancer/kvstore/opfs:disabled": "./src/util/false.ts",
+      "default": "./src/kvstore/opfs/register_backend.ts"
+    },
+    "#kvstore/ssa_s3/register_credentials_provider": {
+      "neuroglancer/python": "./src/util/false.ts",
+      "neuroglancer/kvstore/ssa_s3:enabled": "./src/kvstore/ssa_s3/register_credentials_provider.ts",
+      "neuroglancer/kvstore:none_by_default": "./src/util/false.ts",
+      "neuroglancer/kvstore/ssa_s3:disabled": "./src/util/false.ts",
+      "default": "./src/kvstore/ssa_s3/register_credentials_provider.ts"
+    },
+    "#kvstore/ssa_s3/register_frontend": {
+      "neuroglancer/kvstore/ssa_s3:enabled": "./src/kvstore/ssa_s3/register_frontend.ts",
+      "neuroglancer/kvstore:none_by_default": "./src/util/false.ts",
+      "neuroglancer/kvstore/ssa_s3:disabled": "./src/util/false.ts",
+      "default": "./src/kvstore/ssa_s3/register_frontend.ts"
+    },
+    "#kvstore/ssa_s3/register_backend": {
+      "neuroglancer/kvstore/ssa_s3:enabled": "./src/kvstore/ssa_s3/register_backend.ts",
+      "neuroglancer/kvstore:none_by_default": "./src/util/false.ts",
+      "neuroglancer/kvstore/ssa_s3:disabled": "./src/util/false.ts",
+      "default": "./src/kvstore/ssa_s3/register_backend.ts"
+    },
     "#kvstore/zip/register_frontend": {
       "neuroglancer/kvstore/zip:enabled": "./src/kvstore/zip/register_frontend.ts",
       "neuroglancer/kvstore:none_by_default": "./src/util/false.ts",

src/chunk_manager/backend.ts

@@ -20,6 +20,7 @@ import type {
   LayerChunkProgressInfo,
 } from "#src/chunk_manager/base.js";
 import {
+  CHUNK_SOURCE_INVALIDATE_CHUNKS_RPC_ID,
   CHUNK_LAYER_STATISTICS_RPC_ID,
   CHUNK_MANAGER_RPC_ID,
   CHUNK_QUEUE_MANAGER_RPC_ID,
@@ -1110,8 +1111,10 @@ export class ChunkQueueManager extends SharedObjectCounterpart {
     }
   }
 
-  invalidateSourceCache(source: ChunkSource) {
-    for (const chunk of source.chunks.values()) {
+  invalidateCachedChunks(source: ChunkSource, keys: string[]) {
+    for (const key of keys) {
+      const chunk = source.chunks.get(key);
+      if (!chunk) continue;
       switch (chunk.state) {
         case ChunkState.DOWNLOADING:
           cancelChunkDownload(chunk);
@@ -1123,6 +1126,10 @@ export class ChunkQueueManager extends SharedObjectCounterpart {
       // Note: After calling this, chunk may no longer be valid.
       this.updateChunkState(chunk, ChunkState.QUEUED);
     }
+  }
+
+  invalidateSourceCache(source: ChunkSource) {
+    this.invalidateCachedChunks(source, [...source.chunks.keys()]);
     this.rpc!.invoke("Chunk.update", { source: source.rpcId });
     this.scheduleUpdate();
   }
@@ -1378,6 +1385,12 @@ registerRPC(CHUNK_SOURCE_INVALIDATE_RPC_ID, function (x) {
   source.chunkManager.queueManager.invalidateSourceCache(source);
 });
 
+registerRPC(CHUNK_SOURCE_INVALIDATE_CHUNKS_RPC_ID, function (x) {
+  const source = <ChunkSource>this.get(x.id);
+  source.chunkManager.queueManager.invalidateCachedChunks(source, x.keys);
+  source.chunkManager.queueManager.scheduleUpdate();
+});
+
 registerPromiseRPC(
   REQUEST_CHUNK_STATISTICS_RPC_ID,
   function (x: { queue: number }) {

src/chunk_manager/base.ts

@@ -97,6 +97,8 @@ export const PREFETCH_PRIORITY_MULTIPLIER = 1e13;
 export const CHUNK_QUEUE_MANAGER_RPC_ID = "ChunkQueueManager";
 export const CHUNK_MANAGER_RPC_ID = "ChunkManager";
 export const CHUNK_SOURCE_INVALIDATE_RPC_ID = "ChunkSource.invalidate";
+export const CHUNK_SOURCE_INVALIDATE_CHUNKS_RPC_ID =
+  "ChunkSource.invalidateChunks";
 
 export const REQUEST_CHUNK_STATISTICS_RPC_ID =
   "ChunkQueueManager.requestChunkStatistics";

src/chunk_manager/frontend.ts

@@ -19,6 +19,7 @@ import type {
   LayerChunkProgressInfo,
 } from "#src/chunk_manager/base.js";
 import {
+  CHUNK_SOURCE_INVALIDATE_CHUNKS_RPC_ID,
   CHUNK_LAYER_STATISTICS_RPC_ID,
   CHUNK_MANAGER_RPC_ID,
   CHUNK_QUEUE_MANAGER_RPC_ID,
@@ -463,6 +464,26 @@ export class ChunkSource extends SharedObject {
     this.chunks.delete(key);
   }
 
+  invalidateChunks(keys: string[]): void {
+    const validKeys: string[] = [];
+    for (const key of keys) {
+      const chunk = this.chunks.get(key);
+      if (chunk) {
+        validKeys.push(key);
+        this.deleteChunk(key);
+      }
+    }
+
+    if (validKeys.length > 0) {
+      this.rpc!.invoke(CHUNK_SOURCE_INVALIDATE_CHUNKS_RPC_ID, {
+        id: this.rpcId,
+        keys: validKeys,
+      });
+
+      this.chunkManager.chunkQueueManager.visibleChunksChanged.dispatch();
+    }
+  }
+
   addChunk(key: string, chunk: Chunk) {
     this.chunks.set(key, chunk);
   }

src/chunk_worker.bundle.js

@@ -12,3 +12,4 @@ import "#src/annotation/backend.js";
 import "#src/datasource/enabled_backend_modules.js";
 import "#src/kvstore/enabled_backend_modules.js";
 import "#src/worker_rpc_context.js";
+import "#src/voxel_annotation/edit_backend.js";

src/datasource/index.ts

@@ -132,6 +132,7 @@ export interface DataSubsource {
   singleMesh?: SingleMeshSource;
   segmentPropertyMap?: SegmentPropertyMap;
   segmentationGraph?: SegmentationGraphSource;
+  isPotentiallyWritable?: boolean;
 }
 
 export interface CompleteUrlOptionsBase extends Partial<ProgressOptions> {
@@ -216,6 +217,7 @@ export interface DataSourceWithRedirectInfo extends DataSource {
 
 export interface DataSubsourceSpecification {
   enabled?: boolean;
+  writable?: boolean;
 }
 
 export interface DataSourceSpecification {

src/datasource/zarr/backend.ts

@@ -28,11 +28,15 @@ import {
 import "#src/datasource/zarr/codec/gzip/decode.js";
 import "#src/datasource/zarr/codec/sharding_indexed/decode.js";
 import "#src/datasource/zarr/codec/transpose/decode.js";
+import { encodeArray } from "#src/datasource/zarr/codec/encode.js";
 import { ChunkKeyEncoding } from "#src/datasource/zarr/metadata/index.js";
 import { WithSharedKvStoreContextCounterpart } from "#src/kvstore/backend.js";
 import { postProcessRawData } from "#src/sliceview/backend_chunk_decoders/postprocess.js";
+import { decodeChannel as decodeChannelUint32 } from "#src/sliceview/compressed_segmentation/decode_uint32.js";
+import { decodeChannel as decodeChannelUint64 } from "#src/sliceview/compressed_segmentation/decode_uint64.js";
 import type { VolumeChunk } from "#src/sliceview/volume/backend.js";
 import { VolumeChunkSource } from "#src/sliceview/volume/backend.js";
+import { DataType } from "#src/util/data_type.js";
 import { registerSharedObject } from "#src/worker_rpc.js";
 
 @registerSharedObject()
@@ -97,4 +101,102 @@ export class ZarrVolumeChunkSource extends WithParameters(
       await postProcessRawData(chunk, signal, decoded);
     }
   }
+
+  async writeChunk(chunk: VolumeChunk): Promise<void> {
+    const { kvStore, getChunkKey, decodeCodecs } = this.chunkKvStore as any;
+    if (!kvStore.write) {
+      throw new Error(
+        "ZarrVolumeChunkSource.writeChunk: underlying kvStore is not writable",
+      );
+    }
+    if (!chunk.data) {
+      throw new Error("ZarrVolumeChunkSource.writeChunk: missing chunk.data");
+    }
+    let dataToWrite = chunk.data;
+
+    const { compressedSegmentationBlockSize } = this.spec;
+    if (compressedSegmentationBlockSize !== undefined) {
+      const compressedData = chunk.data as Uint32Array;
+      const { chunkDataSize } = chunk;
+      if (!chunkDataSize) {
+        throw new Error("Cannot write chunk with unknown size.");
+      }
+      const numElements =
+        chunkDataSize[0] * chunkDataSize[1] * chunkDataSize[2];
+      const { dataType } = this.spec;
+      const baseOffset = compressedData.length > 0 ? compressedData[0] : 0;
+
+      if (dataType === DataType.UINT32) {
+        const uncompressedData = new Uint32Array(numElements);
+        if (baseOffset !== 0) {
+          decodeChannelUint32(
+            uncompressedData,
+            compressedData,
+            baseOffset,
+            chunkDataSize,
+            compressedSegmentationBlockSize,
+          );
+        }
+        dataToWrite = uncompressedData;
+      } else {
+        const uncompressedData = new BigUint64Array(numElements);
+        if (baseOffset !== 0) {
+          decodeChannelUint64(
+            uncompressedData,
+            compressedData,
+            baseOffset,
+            chunkDataSize,
+            compressedSegmentationBlockSize,
+          );
+        }
+        dataToWrite = uncompressedData;
+      }
+    }
+
+    const encoded = await encodeArray(
+      decodeCodecs,
+      dataToWrite as ArrayBufferView<ArrayBufferLike>,
+      new AbortController().signal,
+    );
+
+    const { parameters } = this;
+    const { chunkGridPosition } = chunk;
+    const { metadata } = parameters;
+    let baseKey = "";
+    const rank = this.spec.rank;
+    const { physicalToLogicalDimension } = metadata.codecs.layoutInfo[0];
+    let sep: string;
+    if (metadata.chunkKeyEncoding === ChunkKeyEncoding.DEFAULT) {
+      baseKey += "c";
+      sep = metadata.dimensionSeparator;
+    } else {
+      sep = "";
+      if (rank === 0) {
+        baseKey += "0";
+      }
+    }
+    const keyCoords = new Array<number>(rank);
+    const { readChunkShape } = metadata.codecs.layoutInfo[0];
+    const { chunkShape } = metadata;
+    for (
+      let fOrderPhysicalDim = 0;
+      fOrderPhysicalDim < rank;
+      ++fOrderPhysicalDim
+    ) {
+      const decodedDim =
+        physicalToLogicalDimension[rank - 1 - fOrderPhysicalDim];
+      keyCoords[decodedDim] = Math.floor(
+        (chunkGridPosition[fOrderPhysicalDim] * readChunkShape[decodedDim]) /
+          chunkShape[decodedDim],
+      );
+    }
+    for (let i = 0; i < rank; ++i) {
+      baseKey += `${sep}${keyCoords[i]}`;
+      sep = metadata.dimensionSeparator;
+    }
+
+    const key = getChunkKey(chunkGridPosition, baseKey) as string | unknown;
+    const arrayBuffer = new Uint8Array(encoded).buffer;
+    await kvStore.write!(key as any, arrayBuffer);
+  }
 }

src/datasource/zarr/codec/encode.ts

@@ -0,0 +1,25 @@
+/**
+ * Minimal Zarr encode pipeline to persist chunks.
+ * Supports only the common case of raw bytes (no transpose/compression/sharding).
+ */
+import type { CodecChainSpec } from "#src/datasource/zarr/codec/index.js";
+import { CodecKind } from "#src/datasource/zarr/codec/index.js";
+
+export async function encodeArray(
+  codecs: CodecChainSpec,
+  typed: ArrayBufferView<ArrayBufferLike>,
+  _signal: AbortSignal,
+): Promise<Uint8Array<ArrayBufferLike>> {
+  // Only support simple "bytes" encoding with no array-to-array and no bytes-to-bytes codecs.
+  const hasArrayToArray = codecs[CodecKind.arrayToArray].length > 0;
+  const hasBytesToBytes = codecs[CodecKind.bytesToBytes].length > 0;
+  const arrayToBytes = codecs[CodecKind.arrayToBytes];
+  if (hasArrayToArray || hasBytesToBytes || arrayToBytes.name !== "bytes") {
+    throw new Error(
+      `encodeArray: Unsupported codec chain; only raw 'bytes' without additional codecs is supported. Got arrayToArray=${hasArrayToArray}, bytesToBytes=${hasBytesToBytes}, arrayToBytes=${arrayToBytes.name}`,
+    );
+  }
+  // For raw bytes, we can write the underlying buffer.
+  const { buffer, byteOffset, byteLength } = typed;
+  return new Uint8Array(buffer, byteOffset, byteLength);
+}

src/datasource/zarr/frontend.ts

@@ -556,7 +556,7 @@ export class ZarrDataSource implements KvStoreBasedDataSourceProvider {
               id: "default",
               default: true,
               url: undefined,
-              subsource: { volume },
+              subsource: { volume, isPotentiallyWritable: true },
             },
             {
               id: "bounds",

src/kvstore/enabled_backend_modules.ts

@@ -9,3 +9,5 @@ import "#kvstore/ngauth/register";
 import "#kvstore/ocdbt/register_backend";
 import "#kvstore/s3/register_backend";
 import "#kvstore/zip/register_backend";
+import "#kvstore/ssa_s3/register_backend";
+import "#kvstore/opfs/register_backend";

src/kvstore/enabled_frontend_modules.ts

@@ -10,4 +10,7 @@ import "#kvstore/ngauth/register";
 import "#kvstore/ngauth/register_credentials_provider";
 import "#kvstore/ocdbt/register_frontend";
 import "#kvstore/s3/register_frontend";
+import "#kvstore/ssa_s3/register_credentials_provider";
+import "#kvstore/ssa_s3/register_frontend";
 import "#kvstore/zip/register_frontend";
+import "#kvstore/opfs/register_frontend";

src/kvstore/index.ts

@@ -91,7 +91,15 @@ export interface ListableKvStore {
   list?: (prefix: string, options: DriverListOptions) => Promise<ListResponse>;
 }
 
-export interface KvStore extends ReadableKvStore, ListableKvStore {
+export interface WritableKvStore {
+  write?: (key: string, value: ArrayBuffer) => Promise<void>;
+  delete?: (key: string) => Promise<void>;
+}
+
+export interface KvStore
+  extends ReadableKvStore,
+    ListableKvStore,
+    WritableKvStore {
   // Indicates that the only valid key is the empty string.
   singleKey?: boolean;
 }