diff --git a/.gitignore b/.gitignore index 07243e00f7..d83cb0a983 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,6 @@ tsconfig.tsbuildinfo -/docs/python/api/index.rst /docs/python/api/*.rst /.vite +/.idea +/.local +/.env diff --git a/docs/index.rst b/docs/index.rst index aba9088945..cf3b08d20f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -7,6 +7,7 @@ Neuroglancer :caption: User Guide user-guide/navigation + user-guide/voxel_annotation .. toctree:: :hidden: diff --git a/docs/user-guide/voxel_annotation.rst b/docs/user-guide/voxel_annotation.rst new file mode 100644 index 0000000000..bc654a64c1 --- /dev/null +++ b/docs/user-guide/voxel_annotation.rst @@ -0,0 +1,130 @@ +.. _voxel-annotation: + +Voxel Annotation +================ + +Voxel annotation allows for direct painting and editing of volumetric data +within Neuroglancer. This feature is available for both :ref:`image-layer` and +:ref:`segmentation-layer`. + +Enabling Voxel Editing +---------------------- + +To enable voxel editing, you must first have a writable volume source. + +1. Open the **Source** tab of an Image or Segmentation layer. +2. Locate a volume source and click the **Write** checkbox next to it. + +.. note:: + Only one source can be writable at a time within a layer. + +**Known limitations**: + +- Only 3D volumes are supported (2D volumes are not). +- Float32 data type is not supported. +- Multi-resolution datasets must have a strict many-to-one hierarchy. See `About multi-resolution datasets`_ for more details. +- On image layers, the value ``0`` (``VOXEL_EMPTY_VALUE``) cannot be used as a paint value, as it is reserved to represent empty (unedited) voxels in the overlay and is rendered as transparent. + +The first time you attempt a drawing operation (like a brush stroke) after enabling writing, a confirmation dialog will appear. Note that this initial operation will be canceled; you can resume drawing once you have confirmed. + +.. note:: + For the segmentation layer, it is recommended to deactivate the **Highlight on hover** option under the **Render** tab. When enabled, painted voxels become highlighted as the mouse moves over them, which can be visually distracting during annotation. + +Supported Storage and Formats +----------------------------- + +Voxel editing is currently supported for the following configurations: + +**Storage**: + +- Amazon S3 or any S3 compatible storage. + +.. note:: + Write operations require that the S3 bucket's CORS policy allows ``PUT`` and ``DELETE`` methods. See :ref:`s3-kvstore` for a reference CORS policy. + +**Data Format**: + +- Zarr v2 and Zarr v3 (including OME-Zarr), with the following Compression/Encoding: + - None (Raw) + - Blosc + - Gzip + +.. note:: + Writing is not supported for arrays whose codec chain includes + ``sharding_indexed`` (the Zarr v3 sharded format) or an array-to-array codec + such as ``transpose``. + +Tools +----- + +Voxel editing provides several tools for different annotation tasks. You can +bind these tools from the **Draw** tab to use/activate them. + +.. _voxel-brush-tool: + +Brush Tool +~~~~~~~~~~ + +The Brush tool allows you to paint voxels by clicking and dragging. + +- **Paint**: Hold :kbd:`Control` + :kbd:`Left Click` and drag. +- **Erase**: Hold :kbd:`Control` + :kbd:`Shift` + :kbd:`Left Click` and drag. +- **Quick Flood Fill**: Hold :kbd:`Control` + :kbd:`Right Click` to trigger a + flood fill at the current position. Hold :kbd:`Shift` as well to erase. + +Settings: + - **Brush size**: Adjust the radius of the brush. + - **Brush shape**: Choose between **Disk** and **Sphere** shapes. + +.. _voxel-flood-fill-tool: + +Flood Fill Tool +~~~~~~~~~~~~~~~ + +The Flood Fill tool fills a connected region of voxels on the current 2D plane. + +- **Fill**: Hold :kbd:`Control` + :kbd:`Left Click`. +- **Clear**: Hold :kbd:`Control` + :kbd:`Shift` + :kbd:`Left Click`. + +Settings: + - **Max fill voxels**: Limits the maximum number of voxels to fill to + prevent accidental large-scale changes. If the limit is exceeded, the + operation will be canceled. + +.. note:: + The flood fill will automatically fill small gaps in the connected region, proportionally to the number of voxels in the region. This feature may sometimes leave unpainted voxels in tight corners of the region. + +.. _voxel-value-picker-tool: + +Value Picker +~~~~~~~~~~~~ + +The Value Picker tool allows you to adopt the voxel value at the current mouse +position as your active Paint Value. + +Common Controls +--------------- + +The **Draw** tab provides several common controls: + +- **Erase only selected value**: When enabled, the erase action only affects + voxels that match the current **Paint Value**. This feature will slow down + painting performance when erasing. +- **Undo / Redo**: Revert or re-apply recent changes. +- **Paint Value**: Manually specify the segment ID or intensity value to paint. +- **New Random Value**: Generates a new random segment ID or intensity value. + +Stamina System +-------------- + +When you perform many edits quickly, a stamina bar will appear below your cursor. This bar represents the amount of remaining work before all of your edits are processed and saved. **If you reload the page while the stamina bar is visible, you will lose some edits**. If the bar gets emptied painting will be halted until the system is able to catch up, this prevents neuroglancer from crashing due to too many edits in a short period of time. + +About multi-resolution datasets +------------------------------- + +Any multi-resolution dataset that has many-to-1 chunk mapping (i.e. one child chunk cannot have multiple parents) can be used for voxel annotation. + +Although voxel annotation supports multi-resolution, any drawing operation will be performed on the highest resolution level, no matter what the current view is. Once an operation is completed, a downsampling pipeline will be triggered to update the lower resolution levels. + +.. note:: + Because of the 3D nature of the datasets, the downsampling may cause visual artifacts: when zoomed out you may see annotations that then disappear when zoomed in, those "invisible" annotations will be found on nearby slices. diff --git a/src/async_computation/encode_blosc.ts b/src/async_computation/encode_blosc.ts new file mode 100644 index 0000000000..443b166cbf --- /dev/null +++ b/src/async_computation/encode_blosc.ts @@ -0,0 +1,25 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { encodeBlosc } from "#src/async_computation/encode_blosc_request.js"; +import { registerAsyncComputation } from "#src/async_computation/handler.js"; + +registerAsyncComputation(encodeBlosc, async (data, config) => { + const { default: Blosc } = await import("numcodecs/blosc"); + const codec = Blosc.fromConfig({ id: "blosc", ...config }); + const result = await codec.encode(data); + return { value: result, transfer: [result.buffer] }; +}); diff --git a/src/async_computation/encode_blosc_request.ts b/src/async_computation/encode_blosc_request.ts new file mode 100644 index 0000000000..c1151f0819 --- /dev/null +++ b/src/async_computation/encode_blosc_request.ts @@ -0,0 +1,22 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { asyncComputation } from "#src/async_computation/index.js"; + +export const encodeBlosc = + asyncComputation<(data: Uint8Array, config: any) => Uint8Array>( + "encodeBlosc", + ); diff --git a/src/chunk_manager/backend.ts b/src/chunk_manager/backend.ts index df56bb02bb..e5c4c5d68c 100644 --- a/src/chunk_manager/backend.ts +++ b/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 = 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 }) { diff --git a/src/chunk_manager/base.ts b/src/chunk_manager/base.ts index d727085373..ec7f096910 100644 --- a/src/chunk_manager/base.ts +++ b/src/chunk_manager/base.ts @@ -100,6 +100,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"; diff --git a/src/chunk_manager/frontend.ts b/src/chunk_manager/frontend.ts index 3df76ce075..258e18eaf0 100644 --- a/src/chunk_manager/frontend.ts +++ b/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,30 @@ export class ChunkSource extends SharedObject { this.chunks.delete(key); } + invalidateChunks(keys: string[], options?: { lazy?: boolean }): void { + // When `lazy` is set, the frontend chunk is kept on display while the + // backend cache is invalidated and refetched; the fresh data replaces it + // in place on arrival, avoiding the lower-resolution fallback flicker. + const lazy = options?.lazy ?? false; + const validKeys: string[] = []; + for (const key of keys) { + const chunk = this.chunks.get(key); + if (chunk) { + validKeys.push(key); + if (!lazy) 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); } diff --git a/src/chunk_worker.bundle.js b/src/chunk_worker.bundle.js index a463171535..1ce9c822d1 100644 --- a/src/chunk_worker.bundle.js +++ b/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/backend.js"; diff --git a/src/datasource/index.ts b/src/datasource/index.ts index 38d2332bc5..4213b19abf 100644 --- a/src/datasource/index.ts +++ b/src/datasource/index.ts @@ -132,6 +132,8 @@ export interface DataSubsource { singleMesh?: SingleMeshSource; segmentPropertyMap?: SegmentPropertyMap; segmentationGraph?: SegmentationGraphSource; + // Specifies whether the datasource & kvstore implementations support writing. + supportsWriting?: boolean; } export interface CompleteUrlOptionsBase extends Partial { @@ -216,6 +218,7 @@ export interface DataSourceWithRedirectInfo extends DataSource { export interface DataSubsourceSpecification { enabled?: boolean; + writingEnabled?: boolean; } export interface DataSourceSpecification { diff --git a/src/datasource/zarr/async_computation.ts b/src/datasource/zarr/async_computation.ts index db56b860ff..678e1d6bfc 100644 --- a/src/datasource/zarr/async_computation.ts +++ b/src/datasource/zarr/async_computation.ts @@ -1,2 +1,3 @@ import "#src/async_computation/decode_blosc.js"; import "#src/async_computation/decode_zstd.js"; +import "#src/async_computation/encode_blosc.js"; diff --git a/src/datasource/zarr/backend.ts b/src/datasource/zarr/backend.ts index 7370f7af0e..389bac8e1a 100644 --- a/src/datasource/zarr/backend.ts +++ b/src/datasource/zarr/backend.ts @@ -19,6 +19,10 @@ import "#src/datasource/zarr/codec/zstd/decode.js"; import "#src/datasource/zarr/codec/bytes/decode.js"; import "#src/datasource/zarr/codec/crc32c/decode.js"; +import "#src/datasource/zarr/codec/bytes/encode.js"; +import "#src/datasource/zarr/codec/gzip/encode.js"; +import "#src/datasource/zarr/codec/blosc/encode.js"; + import { WithParameters } from "#src/chunk_manager/backend.js"; import { VolumeChunkSourceParameters } from "#src/datasource/zarr/base.js"; import { @@ -28,11 +32,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() @@ -46,11 +54,9 @@ export class ZarrVolumeChunkSource extends WithParameters( this.sharedKvStoreContext.kvStoreContext.getKvStore(this.parameters.url), ); - async download(chunk: VolumeChunk, signal: AbortSignal) { - chunk.chunkDataSize = this.spec.chunkDataSize; - const { parameters } = this; - const { chunkGridPosition } = chunk; - const { metadata } = parameters; + // Sharded stores use structured (non-string) keys, hence the unknown type. + private getChunkStoreKey(chunkGridPosition: Float32Array): unknown { + const { metadata } = this.parameters; let baseKey = ""; const rank = this.spec.rank; const { physicalToLogicalDimension } = metadata.codecs.layoutInfo[0]; @@ -83,9 +89,15 @@ export class ZarrVolumeChunkSource extends WithParameters( baseKey += `${sep}${keyCoords[i]}`; sep = metadata.dimensionSeparator; } + return this.chunkKvStore.getChunkKey(chunkGridPosition, baseKey); + } + + async download(chunk: VolumeChunk, signal: AbortSignal) { + chunk.chunkDataSize = this.spec.chunkDataSize; + const { chunkGridPosition } = chunk; const { chunkKvStore } = this; const response = await chunkKvStore.kvStore.read( - chunkKvStore.getChunkKey(chunkGridPosition, baseKey), + this.getChunkStoreKey(chunkGridPosition), { signal }, ); if (response !== undefined) { @@ -97,4 +109,101 @@ export class ZarrVolumeChunkSource extends WithParameters( await postProcessRawData(chunk, signal, decoded); } } + + async writeChunk(chunk: VolumeChunk): Promise { + const { kvStore, decodeCodecs } = this.chunkKvStore; + 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; + + // For segmentation layers (uint32/uint64), neuroglancer automatically + // transcodes chunks into its own compressed segmentation format for GPU + // efficiency (see makeDefaultVolumeChunkSpecifications in + // sliceview/volume/base.ts). chunk.data is therefore stored in that + // compressed format in memory. Before writing back to the zarr store we + // must decompress it so the on-disk data stays as raw integers. + 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; + // A zero-length compressed buffer for a non-empty chunk means the + // underlying ArrayBuffer was detached (transferred to the frontend). + // Proceeding would silently decode to a full-size zero-filled array and + // durably overwrite the stored chunk with zeros. + if (compressedData.length === 0 && numElements > 0) { + throw new Error( + "ZarrVolumeChunkSource.writeChunk: refusing to write chunk from a " + + "zero-length (detached?) compressed buffer.", + ); + } + 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; + } + } + + // Never write a body whose element count does not match the chunk. A + // zero-length view here means the underlying buffer was detached (its + // ArrayBuffer transferred to the frontend); writing it would durably + // corrupt the stored object (e.g. a 20-byte gzip of an empty stream). + const writeChunkDataSize = chunk.chunkDataSize; + if (!writeChunkDataSize) { + throw new Error("ZarrVolumeChunkSource.writeChunk: unknown chunk size"); + } + const expectedElements = writeChunkDataSize.reduce((a, b) => a * b, 1); + const actualElements = (dataToWrite as unknown as { length: number }) + .length; + if (expectedElements === 0 || actualElements !== expectedElements) { + throw new Error( + `ZarrVolumeChunkSource.writeChunk: refusing to write chunk with ` + + `${actualElements} elements (expected ${expectedElements}); ` + + `buffer detached or chunk size invalid.`, + ); + } + + const encoded = await encodeArray( + decodeCodecs, + dataToWrite as ArrayBufferView, + new AbortController().signal, + ); + + const key = this.getChunkStoreKey(chunk.chunkGridPosition); + const arrayBuffer = new Uint8Array(encoded).buffer; + await kvStore.write!(key, arrayBuffer); + } } diff --git a/src/datasource/zarr/codec/blosc/encode.ts b/src/datasource/zarr/codec/blosc/encode.ts new file mode 100644 index 0000000000..58871abd7c --- /dev/null +++ b/src/datasource/zarr/codec/blosc/encode.ts @@ -0,0 +1,40 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { encodeBlosc } from "#src/async_computation/encode_blosc_request.js"; +import { requestAsyncComputation } from "#src/async_computation/request.js"; +import type { Configuration } from "#src/datasource/zarr/codec/blosc/resolve.js"; +import { registerCodec } from "#src/datasource/zarr/codec/encode.js"; +import { CodecKind } from "#src/datasource/zarr/codec/index.js"; + +registerCodec({ + name: "blosc", + kind: CodecKind.bytesToBytes, + encode( + configuration: Configuration, + decoded: Uint8Array, + signal: AbortSignal, + ): Promise { + configuration; + return requestAsyncComputation( + encodeBlosc, + signal, + [decoded.buffer], + decoded, + {}, + ); + }, +}); diff --git a/src/datasource/zarr/codec/bytes/encode.ts b/src/datasource/zarr/codec/bytes/encode.ts new file mode 100644 index 0000000000..b1fa5d127d --- /dev/null +++ b/src/datasource/zarr/codec/bytes/encode.ts @@ -0,0 +1,42 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { Configuration } from "#src/datasource/zarr/codec/bytes/resolve.js"; +import { registerCodec } from "#src/datasource/zarr/codec/encode.js"; +import { + type CodecArrayInfo, + CodecKind, +} from "#src/datasource/zarr/codec/index.js"; +import { DATA_TYPE_BYTES } from "#src/sliceview/base.js"; +import { convertEndian } from "#src/util/endian.js"; + +registerCodec({ + name: "bytes", + kind: CodecKind.arrayToBytes, + async encode( + configuration: Configuration, + encodedArrayInfo: CodecArrayInfo, + decoded: ArrayBufferView, + ): Promise { + const bytesPerElement = DATA_TYPE_BYTES[encodedArrayInfo.dataType]; + convertEndian(decoded, configuration.endian, bytesPerElement); + return new Uint8Array( + decoded.buffer, + decoded.byteOffset, + decoded.byteLength, + ); + }, +}); diff --git a/src/datasource/zarr/codec/decode.ts b/src/datasource/zarr/codec/decode.ts index 2cc3ea9184..519be93a15 100644 --- a/src/datasource/zarr/codec/decode.ts +++ b/src/datasource/zarr/codec/decode.ts @@ -18,16 +18,16 @@ import type { ChunkManager } from "#src/chunk_manager/backend.js"; import type { CodecArrayInfo, CodecChainSpec, + Codec, } from "#src/datasource/zarr/codec/index.js"; import { CodecKind } from "#src/datasource/zarr/codec/index.js"; -import type { KvStoreWithPath, ReadableKvStore } from "#src/kvstore/index.js"; +import type { + KvStore, + KvStoreWithPath, + ReadableKvStore, +} from "#src/kvstore/index.js"; import type { RefCounted } from "#src/util/disposable.js"; -export interface Codec { - name: string; - kind: CodecKind; -} - export interface ArrayToArrayCodec extends Codec { kind: CodecKind.arrayToArray; decode( @@ -145,14 +145,14 @@ export function applySharding( codecs: CodecChainSpec, baseKvStore: KvStoreWithPath, ): { - kvStore: ReadableKvStore; + kvStore: KvStore; getChunkKey: ( chunkGridPosition: ArrayLike, baseKey: string, ) => unknown; decodeCodecs: CodecChainSpec; } { - let kvStore: ReadableKvStore = baseKvStore.store; + let kvStore: KvStore = baseKvStore.store; let curCodecs = codecs; while (true) { const { shardingInfo } = curCodecs; diff --git a/src/datasource/zarr/codec/encode.ts b/src/datasource/zarr/codec/encode.ts new file mode 100644 index 0000000000..38bc755672 --- /dev/null +++ b/src/datasource/zarr/codec/encode.ts @@ -0,0 +1,93 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { + CodecChainSpec, + Codec, + CodecArrayInfo, +} from "#src/datasource/zarr/codec/index.js"; +import { CodecKind } from "#src/datasource/zarr/codec/index.js"; + +interface ArrayToBytesCodec extends Codec { + kind: CodecKind.arrayToBytes; + encode( + configuration: Configuration, + encodedArrayInfo: CodecArrayInfo, + decoded: ArrayBufferView, + signal: AbortSignal, + ): Promise; +} + +interface BytesToBytesCodec extends Codec { + kind: CodecKind.bytesToBytes; + encode( + configuration: Configuration, + decoded: Uint8Array, + signal: AbortSignal, + ): Promise; +} + +const codecRegistry = { + [CodecKind.arrayToBytes]: new Map(), + [CodecKind.bytesToBytes]: new Map(), +}; + +export function registerCodec( + codec: ArrayToBytesCodec | BytesToBytesCodec, +) { + codecRegistry[codec.kind].set(codec.name, codec as any); +} + +export async function encodeArray( + codecs: CodecChainSpec, + decoded: ArrayBufferView, + signal: AbortSignal, +): Promise { + if (codecs[CodecKind.arrayToArray].length > 0) { + throw new Error("array -> array codecs are not supported for writing."); + } + + const arrayToBytesCodecSpec = codecs[CodecKind.arrayToBytes]; + const arrayToBytesImpl = codecRegistry[CodecKind.arrayToBytes].get( + arrayToBytesCodecSpec.name, + ); + if (!arrayToBytesImpl) { + throw new Error( + `Unsupported array -> bytes codec for writing: ${arrayToBytesCodecSpec.name}`, + ); + } + const arrayInfo = codecs.arrayInfo[codecs.arrayInfo.length - 1]; + let data = await arrayToBytesImpl.encode( + arrayToBytesCodecSpec.configuration, + arrayInfo, + decoded, + signal, + ); + + for (const codecSpec of codecs[CodecKind.bytesToBytes]) { + const bytesToBytesImpl = codecRegistry[CodecKind.bytesToBytes].get( + codecSpec.name, + ); + if (!bytesToBytesImpl) { + throw new Error( + `Unsupported bytes -> bytes codec for writing: ${codecSpec.name}`, + ); + } + data = await bytesToBytesImpl.encode(codecSpec.configuration, data, signal); + } + + return data; +} diff --git a/src/datasource/zarr/codec/gzip/encode.ts b/src/datasource/zarr/codec/gzip/encode.ts new file mode 100644 index 0000000000..613d25bf4c --- /dev/null +++ b/src/datasource/zarr/codec/gzip/encode.ts @@ -0,0 +1,38 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { registerCodec } from "#src/datasource/zarr/codec/encode.js"; +import type { Configuration } from "#src/datasource/zarr/codec/gzip/resolve.js"; +import { CodecKind } from "#src/datasource/zarr/codec/index.js"; +import { encodeGzip } from "#src/util/gzip.js"; + +for (const [name, compressionFormat] of [ + ["gzip", "gzip"], + ["zlib", "deflate"], +] as const) { + registerCodec({ + name, + kind: CodecKind.bytesToBytes, + async encode( + configuration: Configuration, + decoded: Uint8Array, + ): Promise { + configuration; + const result = await encodeGzip(decoded, compressionFormat); + return new Uint8Array(result); + }, + }); +} diff --git a/src/datasource/zarr/codec/index.ts b/src/datasource/zarr/codec/index.ts index bb99f48ebb..7bafc22849 100644 --- a/src/datasource/zarr/codec/index.ts +++ b/src/datasource/zarr/codec/index.ts @@ -16,6 +16,11 @@ import type { DataType } from "#src/util/data_type.js"; +export interface Codec { + name: string; + kind: CodecKind; +} + export enum CodecKind { arrayToArray = 0, arrayToBytes = 1, @@ -44,6 +49,13 @@ export interface ShardingInfo { subChunkCodecs: CodecChainSpec; } +export function codecChainSupportsWriting(codecs: CodecChainSpec): boolean { + return ( + codecs.shardingInfo === undefined && + codecs[CodecKind.arrayToArray].length === 0 + ); +} + export interface CodecArrayInfo { dataType: DataType; // Specifies the chunk shape, indexed by logical dimension. diff --git a/src/datasource/zarr/frontend.ts b/src/datasource/zarr/frontend.ts index c501eebec5..bc9d637a9d 100644 --- a/src/datasource/zarr/frontend.ts +++ b/src/datasource/zarr/frontend.ts @@ -33,6 +33,7 @@ import type { } from "#src/datasource/index.js"; import { getKvStorePathCompletions } from "#src/datasource/kvstore_completions.js"; import { VolumeChunkSourceParameters } from "#src/datasource/zarr/base.js"; +import { codecChainSupportsWriting } from "#src/datasource/zarr/codec/index.js"; import "#src/datasource/zarr/codec/bytes/resolve.js"; import "#src/datasource/zarr/codec/crc32c/resolve.js"; import "#src/datasource/zarr/codec/gzip/resolve.js"; @@ -485,6 +486,7 @@ export class ZarrDataSource implements KvStoreBasedDataSourceProvider { this.zarrVersion === undefined ? "" : ` v${this.zarrVersion}`; return `Zarr${versionStr} data source`; } + get(options: GetKvStoreBasedDataSourceOptions): Promise { let { kvStoreUrl, additionalPath, fragment } = resolveUrl(options); kvStoreUrl = kvstoreEnsureDirectoryPipelineUrl( @@ -547,6 +549,12 @@ export class ZarrDataSource implements KvStoreBasedDataSourceProvider { sharedKvStoreContext, multiscaleInfo, ); + const supportsWriting = + sharedKvStoreContext.kvStoreContext.getKvStore(kvStoreUrl).store + .write !== undefined && + multiscaleInfo.scales.every((scale) => + codecChainSupportsWriting(scale.metadata.codecs), + ); return { canonicalUrl: `${kvStoreUrl}|zarr${metadata.zarrVersion}:`, modelTransform: makeIdentityTransform(volume.modelSpace), @@ -556,7 +564,7 @@ export class ZarrDataSource implements KvStoreBasedDataSourceProvider { id: "default", default: true, url: undefined, - subsource: { volume }, + subsource: { volume, supportsWriting }, }, { id: "bounds", diff --git a/src/kvstore/index.ts b/src/kvstore/index.ts index e94f870efb..159e6caa7c 100644 --- a/src/kvstore/index.ts +++ b/src/kvstore/index.ts @@ -91,7 +91,15 @@ export interface ListableKvStore { list?: (prefix: string, options: DriverListOptions) => Promise; } -export interface KvStore extends ReadableKvStore, ListableKvStore { +export interface WritableKvStore { + write?: (key: Key, value: ArrayBuffer) => Promise; + delete?: (key: Key) => Promise; +} + +export interface KvStore + extends ReadableKvStore, + ListableKvStore, + WritableKvStore { // Indicates that the only valid key is the empty string. singleKey?: boolean; } diff --git a/src/kvstore/s3/backend.ts b/src/kvstore/s3/backend.ts index 7577a34618..971c3c9dc7 100644 --- a/src/kvstore/s3/backend.ts +++ b/src/kvstore/s3/backend.ts @@ -17,9 +17,9 @@ import type { SharedKvStoreContextCounterpart } from "#src/kvstore/backend.js"; import type { DriverListOptions, ListResponse } from "#src/kvstore/index.js"; import { proxyList } from "#src/kvstore/proxy.js"; -import { ReadableS3KvStore } from "#src/kvstore/s3/common.js"; +import { S3KvStoreBase } from "#src/kvstore/s3/common.js"; -export class S3KvStore extends ReadableS3KvStore { +export class S3KvStore extends S3KvStoreBase { list(prefix: string, options: DriverListOptions): Promise { return proxyList(this.sharedKvStoreContext, this.getUrl(prefix), options); } diff --git a/src/kvstore/s3/common.ts b/src/kvstore/s3/common.ts index e1e19ca06f..bd3fc40866 100644 --- a/src/kvstore/s3/common.ts +++ b/src/kvstore/s3/common.ts @@ -35,10 +35,21 @@ import { } from "#src/kvstore/s3/list.js"; import { joinBaseUrlAndPath } from "#src/kvstore/url.js"; import type { FetchOk } from "#src/util/http_request.js"; -import { fetchOk } from "#src/util/http_request.js"; +import { HttpError, fetchOk } from "#src/util/http_request.js"; import { ProgressSpan } from "#src/util/progress_listener.js"; +import { getRandomHexString } from "#src/util/random.js"; -export class ReadableS3KvStore< +function validateObjectKey(key: string) { + for (const component of key.split("/")) { + if (component === "." || component === "..") { + throw new Error( + `Invalid S3 object key ${JSON.stringify(key)}: "." and ".." path components are not supported`, + ); + } + } +} + +export class S3KvStoreBase< SharedKvStoreContext extends SharedKvStoreContextBase, > implements KvStore { @@ -50,17 +61,28 @@ export class ReadableS3KvStore< protected fetchOkImpl: FetchOk = fetchOk, ) {} + private getBaseObjectUrl(key: string): string { + validateObjectKey(key); + return joinBaseUrlAndPath(this.baseUrl, key); + } + + // Random query parameter (ignored by S3) so cached responses are never + // used — same rationale as GcsKvStore.getObjectUrl: stale ACAO headers on + // 304s (S3's CORS headers also vary with the Origin) and staleness after + // this or another session writes to the bucket. + private getObjectUrl(key: string): string { + return this.getBaseObjectUrl(key) + `?neuroglancer=${getRandomHexString()}`; + } + stat(key: string, options: StatOptions): Promise { - const url = joinBaseUrlAndPath(this.baseUrl, key); - return stat(this, key, url, options, this.fetchOkImpl); + return stat(this, key, this.getObjectUrl(key), options, this.fetchOkImpl); } read( key: string, options: DriverReadOptions, ): Promise { - const url = joinBaseUrlAndPath(this.baseUrl, key); - return read(this, key, url, options, this.fetchOkImpl); + return read(this, key, this.getObjectUrl(key), options, this.fetchOkImpl); } list(prefix: string, options: DriverListOptions): Promise { @@ -88,6 +110,32 @@ export class ReadableS3KvStore< ); } + async write(key: string, value: ArrayBuffer): Promise { + const url = this.getBaseObjectUrl(key); + try { + await this.fetchOkImpl(url, { + method: "PUT", + body: value, + }); + } catch (e) { + throw new Error(`Failed to write to ${url}.`, { cause: e }); + } + } + + async delete(key: string): Promise { + const url = this.getBaseObjectUrl(key); + try { + await this.fetchOkImpl(url, { + method: "DELETE", + }); + } catch (e) { + if (e instanceof HttpError && e.status === 404) { + return; + } + throw new Error(`Failed to delete ${url}.`, { cause: e }); + } + } + getUrl(path: string) { return joinBaseUrlAndPath(this.baseUrlForDisplay, path); } @@ -104,7 +152,7 @@ function amazonS3Provider< SharedKvStoreContext extends SharedKvStoreContextBase, >( sharedKvStoreContext: SharedKvStoreContext, - s3KvStoreClass: typeof ReadableS3KvStore, + s3KvStoreClass: typeof S3KvStoreBase, ): BaseKvStoreProvider { return { scheme: "s3", @@ -131,7 +179,7 @@ function amazonS3Provider< function s3Provider( sharedKvStoreContext: SharedKvStoreContext, httpScheme: "http" | "https", - s3KvStoreClass: typeof ReadableS3KvStore, + s3KvStoreClass: typeof S3KvStoreBase, ): BaseKvStoreProvider { return { scheme: `s3+${httpScheme}`, @@ -161,7 +209,7 @@ export function registerProviders< SharedKvStoreContext extends SharedKvStoreContextBase, >( registry: KvStoreProviderRegistry, - s3KvStoreClass: typeof ReadableS3KvStore, + s3KvStoreClass: typeof S3KvStoreBase, ) { registry.registerBaseKvStoreProvider((context) => amazonS3Provider(context, s3KvStoreClass), diff --git a/src/kvstore/s3/frontend.ts b/src/kvstore/s3/frontend.ts index b9c4fee91d..ef0ed81111 100644 --- a/src/kvstore/s3/frontend.ts +++ b/src/kvstore/s3/frontend.ts @@ -16,7 +16,7 @@ import type { SharedKvStoreContext } from "#src/kvstore/frontend.js"; import type { DriverListOptions, ListResponse } from "#src/kvstore/index.js"; -import { ReadableS3KvStore } from "#src/kvstore/s3/common.js"; +import { S3KvStoreBase } from "#src/kvstore/s3/common.js"; import { getS3BucketListing, listS3CompatibleUrl, @@ -24,7 +24,7 @@ import { import { joinBaseUrlAndPath } from "#src/kvstore/url.js"; import { ProgressSpan } from "#src/util/progress_listener.js"; -export class S3KvStore extends ReadableS3KvStore { +export class S3KvStore extends S3KvStoreBase { list(prefix: string, options: DriverListOptions): Promise { const { progressListener } = options; using _span = diff --git a/src/kvstore/s3/index.rst b/src/kvstore/s3/index.rst index 01cbfc4a42..7483326b7d 100644 --- a/src/kvstore/s3/index.rst +++ b/src/kvstore/s3/index.rst @@ -72,3 +72,31 @@ such as the following: "MaxAgeSeconds": 3000 } ] + +If the bucket also needs to support write operations (e.g. for :ref:`voxel-annotation`), ``PUT`` and ``DELETE`` must be added to ``AllowedMethods``. When allowing write methods, do not use a wildcard ``AllowedOrigins``; explicitly list the origins that should be permitted to write: + +.. code-block:: json + + [ + { + "AllowedHeaders": [ + "*" + ], + "AllowedMethods": [ + "GET", + "HEAD", + "PUT", + "DELETE" + ], + "AllowedOrigins": [ + "https://example.com" + ], + "ExposeHeaders": [ + "ETag", + "Content-Range", + "Content-Encoding", + "Content-Length" + ], + "MaxAgeSeconds": 3000 + } + ] diff --git a/src/layer/image/index.ts b/src/layer/image/index.ts index 241b69853e..b70ebdcf65 100644 --- a/src/layer/image/index.ts +++ b/src/layer/image/index.ts @@ -34,7 +34,10 @@ import { UserLayer, } from "#src/layer/index.js"; import type { LoadedDataSubsource } from "#src/layer/layer_data_source.js"; +import { registerVoxelLayerControls } from "#src/layer/voxel_annotation/controls.js"; +import { UserLayerWithVoxelEditingMixin } from "#src/layer/voxel_annotation/index.js"; import { Overlay } from "#src/overlay.js"; +import type { RenderLayerTransformOrError } from "#src/render_coordinate_transform.js"; import { getChannelSpace } from "#src/render_coordinate_transform.js"; import { RenderScaleHistogram, @@ -48,21 +51,25 @@ import { ImageRenderLayer, } from "#src/sliceview/volume/image_renderlayer.js"; import { trackableAlphaValue } from "#src/trackable_alpha.js"; -import { trackableBlendModeValue } from "#src/trackable_blend.js"; +import { BLEND_MODES, trackableBlendModeValue } from "#src/trackable_blend.js"; import { TrackableBoolean } from "#src/trackable_boolean.js"; import { trackableFiniteFloat } from "#src/trackable_finite_float.js"; import type { WatchableValueInterface } from "#src/trackable_value.js"; import { makeCachedDerivedWatchableValue, makeCachedLazyDerivedWatchableValue, + makeDerivedWatchableValue, registerNested, + TrackableValue, WatchableValue, } from "#src/trackable_value.js"; import { UserLayerWithAnnotationsMixin } from "#src/ui/annotations.js"; +import { registerVoxelTools } from "#src/ui/voxel_annotations.js"; import { setClipboard } from "#src/util/clipboard.js"; import type { Borrowed } from "#src/util/disposable.js"; import { makeValueOrError } from "#src/util/error.js"; -import { verifyOptionalObjectProperty } from "#src/util/json.js"; +import { verifyFloat01, verifyOptionalObjectProperty } from "#src/util/json.js"; +import { TrackableEnum } from "#src/util/trackable_enum.js"; import { trackableShaderModeValue, VolumeRenderingModes, @@ -72,6 +79,7 @@ import { VOLUME_RENDERING_DEPTH_SAMPLES_DEFAULT_VALUE, VolumeRenderingRenderLayer, } from "#src/volume_rendering/volume_render_layer.js"; +import { VOXEL_EMPTY_VALUE } from "#src/voxel_annotation/base.js"; import type { ParameterizedShaderGetterResult } from "#src/webgl/dynamic_shader.js"; import { makeWatchableShaderError } from "#src/webgl/dynamic_shader.js"; import type { ShaderControlsBuilderState } from "#src/webgl/shader_ui_controls.js"; @@ -119,7 +127,9 @@ export interface ImageLayerSelectionState extends UserLayerSelectionState { value: any; } -const Base = UserLayerWithAnnotationsMixin(UserLayer); +const Base = UserLayerWithVoxelEditingMixin( + UserLayerWithAnnotationsMixin(UserLayer), +); const [ volumeRenderingDepthSamplesOriginLogScale, volumeRenderingDepthSamplesMaxLogScale, @@ -132,6 +142,7 @@ export class ImageUserLayer extends Base { shaderError = makeWatchableShaderError(); dataType = new WatchableValue(undefined); sliceViewRenderScaleHistogram = new RenderScaleHistogram(); + voxelOverlayRenderScaleHistogram = new RenderScaleHistogram(); sliceViewRenderScaleTarget = trackableRenderScaleTarget(1); volumeRenderingGain = trackableFiniteFloat(0); volumeRenderingChunkResolutionHistogram = new RenderScaleHistogram( @@ -188,6 +199,50 @@ export class ImageUserLayer extends Base { }; } + _createVoxelOverlayRenderLayer( + source: MultiscaleVolumeChunkSource, + transform: WatchableValueInterface, + ): ImageRenderLayer { + const wrappedFragmentMain = makeDerivedWatchableValue( + (originalShader: string) => ` +#define main userMain +${originalShader} +#undef main + +void main() { + // VOXEL_EMPTY_VALUE is transparent in the overlay so the underlying data + // shows through. This means it cannot be used as a paint value on image layers. + if (toRaw(getDataValue()) == ${VOXEL_EMPTY_VALUE}n) { + emitTransparent(); + return; + } + userMain(); +} +`, + this.fragmentMain, + ); + this.registerDisposer(wrappedFragmentMain); + + const shaderControlState = new ShaderControlState( + wrappedFragmentMain, + this.shaderControlState.dataContext, + this.channelCoordinateSpaceCombiner, + ); + this.registerDisposer(shaderControlState); + + return new ImageRenderLayer(source, { + opacity: new TrackableValue(1.0, verifyFloat01), + blendMode: new TrackableEnum(BLEND_MODES, BLEND_MODES.ADDITIVE), + shaderControlState: shaderControlState, + shaderError: this.shaderError, + transform: transform, + renderScaleTarget: this.sliceViewRenderScaleTarget, + renderScaleHistogram: this.voxelOverlayRenderScaleHistogram, + localPosition: this.localPosition, + channelCoordinateSpace: this.channelCoordinateSpace, + }); + } + addCoordinateSpace( coordinateSpace: WatchableValueInterface, ) { @@ -246,21 +301,20 @@ export class ImageUserLayer extends Base { } dataType = volume.dataType; loadedSubsource.activate((context) => { - loadedSubsource.addRenderLayer( - new ImageRenderLayer(volume, { - opacity: this.opacity, - blendMode: this.blendMode, - shaderControlState: this.shaderControlState, - shaderError: this.shaderError, - transform: loadedSubsource.getRenderLayerTransform( - this.channelCoordinateSpace, - ), - renderScaleTarget: this.sliceViewRenderScaleTarget, - renderScaleHistogram: this.sliceViewRenderScaleHistogram, - localPosition: this.localPosition, - channelCoordinateSpace: this.channelCoordinateSpace, - }), - ); + const imageRenderLayer = new ImageRenderLayer(volume, { + opacity: this.opacity, + blendMode: this.blendMode, + shaderControlState: this.shaderControlState, + shaderError: this.shaderError, + transform: loadedSubsource.getRenderLayerTransform( + this.channelCoordinateSpace, + ), + renderScaleTarget: this.sliceViewRenderScaleTarget, + renderScaleHistogram: this.sliceViewRenderScaleHistogram, + localPosition: this.localPosition, + channelCoordinateSpace: this.channelCoordinateSpace, + }); + loadedSubsource.addRenderLayer(imageRenderLayer); const volumeRenderLayer = context.registerDisposer( new VolumeRenderingRenderLayer({ gain: this.volumeRenderingGain, @@ -290,6 +344,18 @@ export class ImageUserLayer extends Base { }, this.volumeRenderingMode), ); this.shaderError.changed.dispatch(); + context.registerDisposer( + registerNested((context, writingEnabled) => { + this.initializeVoxelEditingForSubsource( + loadedSubsource, + imageRenderLayer, + writingEnabled, + ); + context.registerDisposer(() => { + this.deinitializeVoxelEditingForSubsource(loadedSubsource); + }); + }, loadedSubsource.writingEnabled), + ); }); } this.dataType.value = dataType; @@ -588,6 +654,8 @@ class ShaderCodeOverlay extends Overlay { } registerLayerType(ImageUserLayer); +registerVoxelTools(ImageUserLayer); +registerVoxelLayerControls(ImageUserLayer); registerVolumeLayerType(VolumeType.IMAGE, ImageUserLayer); // Use ImageUserLayer as a fallback layer type if there is a `volume` subsource. registerLayerTypeDetector((subsource) => { diff --git a/src/layer/index.ts b/src/layer/index.ts index f4ea769103..5d51f23296 100644 --- a/src/layer/index.ts +++ b/src/layer/index.ts @@ -85,6 +85,7 @@ import { LayerToolBinder, SelectedLegacyTool } from "#src/ui/tool.js"; import { gatherUpdate } from "#src/util/array.js"; import type { Borrowed, Owned } from "#src/util/disposable.js"; import { invokeDisposers, RefCounted } from "#src/util/disposable.js"; +import type { vec3 } from "#src/util/geom.js"; import { emptyToUndefined, parseArray, @@ -1137,6 +1138,7 @@ export class MouseSelectionState implements PickState { unsnappedPosition: Float32Array = kEmptyFloat32Vec; active = false; displayDimensions: DisplayDimensions | undefined = undefined; + planeNormal: vec3 | undefined = undefined; pickedRenderLayer: RenderLayer | null = null; pickedValue = 0n; pickedOffset = 0; diff --git a/src/layer/layer_data_source.ts b/src/layer/layer_data_source.ts index 63d07a7746..fc1a08679a 100644 --- a/src/layer/layer_data_source.ts +++ b/src/layer/layer_data_source.ts @@ -36,6 +36,7 @@ import { makeEmptyDataSourceSpecification } from "#src/datasource/index.js"; import type { UserLayer } from "#src/layer/index.js"; import { getWatchableRenderLayerTransform } from "#src/render_coordinate_transform.js"; import type { RenderLayer } from "#src/renderlayer.js"; +import { TrackableBoolean } from "#src/trackable_boolean.js"; import type { WatchableValueInterface } from "#src/trackable_value.js"; import { arraysEqual } from "#src/util/array.js"; import type { Borrowed, Owned } from "#src/util/disposable.js"; @@ -63,6 +64,11 @@ export function parseDataSubsourceSpecificationFromJson( verifyObject(json); return { enabled: verifyOptionalObjectProperty(json, "enabled", verifyBoolean), + writingEnabled: verifyOptionalObjectProperty( + json, + "writingEnabled", + verifyBoolean, + ), }; } @@ -108,7 +114,11 @@ export function layerDataSourceSpecificationFromJson( } function dataSubsourceSpecificationToJson(spec: DataSubsourceSpecification) { - return spec.enabled; + const { enabled, writingEnabled } = spec; + if (writingEnabled === undefined) { + return enabled; + } + return { enabled, writingEnabled }; } export function layerDataSourceSpecificationToJson( @@ -146,6 +156,7 @@ export class LoadedDataSubsource { subsourceToModelSubspaceTransform: Float32Array; modelSubspaceDimensionIndices: number[]; enabled: boolean; + writingEnabled: TrackableBoolean; activated: RefCounted | undefined = undefined; guardValues: any[] = []; messages = new MessageList(); @@ -178,6 +189,13 @@ export class LoadedDataSubsource { ), } = subsourceEntry; this.enabled = enabled; + this.writingEnabled = new TrackableBoolean( + subsourceSpec?.writingEnabled ?? false, + false, + ); + this.writingEnabled.changed.add( + loadedDataSource.layer.dataSourcesChanged.dispatch, + ); this.subsourceToModelSubspaceTransform = subsourceToModelSubspaceTransform; this.modelSubspaceDimensionIndices = modelSubspaceDimensionIndices; this.isActiveChanged.add( @@ -451,6 +469,7 @@ export class LayerDataSource extends RefCounted { if (refCounted.wasDisposed) return; this.loadState_ = { error }; this.messages.clearMessages(); + this.messages.addMessage({ severity: MessageSeverity.error, message: formatErrorMessage(error), @@ -491,6 +510,9 @@ export class LayerDataSource extends RefCounted { loadedSubsource.enabled !== defaultEnabledValue ? loadedSubsource.enabled : undefined, + writingEnabled: loadedSubsource.writingEnabled.value + ? true + : undefined, }, ]; }), diff --git a/src/layer/segmentation/index.ts b/src/layer/segmentation/index.ts index 43a21b05cd..70aa9b755c 100644 --- a/src/layer/segmentation/index.ts +++ b/src/layer/segmentation/index.ts @@ -35,12 +35,15 @@ import type { LoadedDataSubsource } from "#src/layer/layer_data_source.js"; import { layerDataSourceSpecificationFromJson } from "#src/layer/layer_data_source.js"; import * as json_keys from "#src/layer/segmentation/json_keys.js"; import { registerLayerControls } from "#src/layer/segmentation/layer_controls.js"; +import { registerVoxelLayerControls } from "#src/layer/voxel_annotation/controls.js"; +import { UserLayerWithVoxelEditingMixin } from "#src/layer/voxel_annotation/index.js"; import { MeshLayer, MeshSource, MultiscaleMeshLayer, MultiscaleMeshSource, } from "#src/mesh/frontend.js"; +import type { RenderLayerTransformOrError } from "#src/render_coordinate_transform.js"; import { RenderScaleHistogram, trackableRenderScaleTarget, @@ -92,6 +95,7 @@ import type { WatchableValueInterface, } from "#src/trackable_value.js"; import { + registerNested, IndirectTrackableValue, IndirectWatchableValue, makeCachedDerivedWatchableValue, @@ -106,6 +110,7 @@ import { SegmentDisplayTab } from "#src/ui/segment_list.js"; import { registerSegmentSelectTools } from "#src/ui/segment_select_tools.js"; import { registerSegmentSplitMergeTools } from "#src/ui/segment_split_merge_tools.js"; import { DisplayOptionsTab } from "#src/ui/segmentation_display_options_tab.js"; +import { registerVoxelTools } from "#src/ui/voxel_annotations.js"; import { Uint64Map } from "#src/uint64_map.js"; import { Uint64OrderedSet } from "#src/uint64_ordered_set.js"; import { Uint64Set } from "#src/uint64_set.js"; @@ -129,6 +134,11 @@ import { verifyString, } from "#src/util/json.js"; import { Signal } from "#src/util/signal.js"; +import { + SEG_ERASE_SENTINEL, + VOXEL_EMPTY_VALUE, + type VoxelValueGetter, +} from "#src/voxel_annotation/base.js"; import { makeWatchableShaderError } from "#src/webgl/dynamic_shader.js"; import type { DependentViewContext } from "#src/widget/dependent_view_widget.js"; import { registerLayerShaderControlsTool } from "#src/widget/shader_controls.js"; @@ -579,9 +589,12 @@ interface SegmentationActionContext extends LayerActionContext { segmentationToggleSegmentState?: boolean | undefined; } -const Base = UserLayerWithAnnotationsMixin(UserLayer); +const Base = UserLayerWithVoxelEditingMixin( + UserLayerWithAnnotationsMixin(UserLayer), +); export class SegmentationUserLayer extends Base { sliceViewRenderScaleHistogram = new RenderScaleHistogram(); + voxelOverlayRenderScaleHistogram = new RenderScaleHistogram(); sliceViewRenderScaleTarget = trackableRenderScaleTarget(1); codeVisible = new TrackableBoolean(true); @@ -606,6 +619,26 @@ export class SegmentationUserLayer extends Base { ); }; + _createVoxelOverlayRenderLayer( + source: MultiscaleVolumeChunkSource, + transform: WatchableValueInterface, + ): SegmentationRenderLayer { + return new SegmentationRenderLayer(source, { + ...this.displayState, + transform: transform, + renderScaleTarget: this.sliceViewRenderScaleTarget, + renderScaleHistogram: this.voxelOverlayRenderScaleHistogram, + localPosition: this.localPosition, + }); + } + + getVoxelPaintValue(erase: boolean): VoxelValueGetter { + return (isPreview) => { + if (erase) return isPreview ? SEG_ERASE_SENTINEL : VOXEL_EMPTY_VALUE; + return this.paintValue.value; + }; + } + filterBySegmentLabel = (id: bigint) => { const augmented = augmentSegmentId(this.displayState, id); const { label } = augmented; @@ -774,19 +807,28 @@ export class SegmentationUserLayer extends Base { continue; } hasVolume = true; - loadedSubsource.activate( - () => - loadedSubsource.addRenderLayer( - new SegmentationRenderLayer(volume, { - ...this.displayState, - transform: loadedSubsource.getRenderLayerTransform(), - renderScaleTarget: this.sliceViewRenderScaleTarget, - renderScaleHistogram: this.sliceViewRenderScaleHistogram, - localPosition: this.localPosition, - }), - ), - this.displayState.segmentationGroupState.value, - ); + loadedSubsource.activate((context) => { + const segmentationRenderLayer = new SegmentationRenderLayer(volume, { + ...this.displayState, + transform: loadedSubsource.getRenderLayerTransform(), + renderScaleTarget: this.sliceViewRenderScaleTarget, + renderScaleHistogram: this.sliceViewRenderScaleHistogram, + localPosition: this.localPosition, + }); + loadedSubsource.addRenderLayer(segmentationRenderLayer); + context.registerDisposer( + registerNested((context, writingEnabled) => { + this.initializeVoxelEditingForSubsource( + loadedSubsource, + segmentationRenderLayer, + writingEnabled, + ); + context.registerDisposer(() => { + this.deinitializeVoxelEditingForSubsource(loadedSubsource); + }); + }, loadedSubsource.writingEnabled), + ); + }, this.displayState.segmentationGroupState.value); } else if (mesh !== undefined) { loadedSubsource.activate(() => { const displayState = { @@ -1391,7 +1433,8 @@ export class SegmentationUserLayer extends Base { } registerLayerControls(SegmentationUserLayer); - +registerVoxelTools(SegmentationUserLayer); +registerVoxelLayerControls(SegmentationUserLayer); registerLayerType(SegmentationUserLayer); registerVolumeLayerType(VolumeType.SEGMENTATION, SegmentationUserLayer); registerLayerTypeDetector((subsource) => { diff --git a/src/layer/voxel_annotation/controls.ts b/src/layer/voxel_annotation/controls.ts new file mode 100644 index 0000000000..c14422f0db --- /dev/null +++ b/src/layer/voxel_annotation/controls.ts @@ -0,0 +1,326 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { UserLayerConstructor } from "#src/layer/index.js"; +import { LayerActionContext } from "#src/layer/index.js"; +import type { + UserLayerWithVoxelEditing, + VoxelEditingContext, +} from "#src/layer/voxel_annotation/index.js"; +import type { RenderedDataPanel } from "#src/rendered_data_panel.js"; +import { SliceViewPanel } from "#src/sliceview/panel.js"; +import { observeWatchable } from "#src/trackable_value.js"; +import { mat3, vec3 } from "#src/util/geom.js"; +import { + BRUSH_TOOL_ID, + FLOODFILL_MAX_POSSIBLE_VOXELS, + FLOODFILL_MIN_POSSIBLE_VOXELS, + FLOODFILL_TOOL_ID, + getBasisFromNormal, + VALUE_PICKER_TOOL_ID, +} from "#src/voxel_annotation/base.js"; +import type { LayerControlDefinition } from "#src/widget/layer_control.js"; +import { registerLayerControl } from "#src/widget/layer_control.js"; +import { buttonLayerControl } from "#src/widget/layer_control_button.js"; +import { checkboxLayerControl } from "#src/widget/layer_control_checkbox.js"; +import { enumLayerControl } from "#src/widget/layer_control_enum.js"; +import { rangeLayerControl } from "#src/widget/layer_control_range.js"; +import type { RangeWidget } from "#src/widget/range.js"; + +export function getEditingContext( + layer: UserLayerWithVoxelEditing, +): VoxelEditingContext | undefined { + const it = layer.editingContexts.values(); + let ctx: VoxelEditingContext; + while ((ctx = it.next().value) !== undefined) { + if (ctx.writingEnabled) return ctx; + } + return undefined; +} + +export function drawBrushCursor( + layer: UserLayerWithVoxelEditing, + panel: RenderedDataPanel, + ctx: CanvasRenderingContext2D, +): { radiusX: number; radiusY: number } { + const context = getEditingContext(layer); + const radiusXY = { radiusX: -1, radiusY: -1 }; + if (context === undefined || !(panel instanceof SliceViewPanel)) { + return radiusXY; + } + + const { projectionParameters } = panel.sliceView; + const { displayDimensionRenderInfo, viewMatrix } = projectionParameters.value; + const { displayRank } = displayDimensionRenderInfo; + + if (displayRank < 2) { + return radiusXY; + } + + const chunkTransform = context.getChunkTransform(); + if (!chunkTransform) { + return radiusXY; + } + const { chunkToLayerTransform, layerRank } = chunkTransform; + const { globalToRenderLayerDimensions } = chunkTransform.modelTransform; + const stride = layerRank + 1; + + const nWorld = + projectionParameters.value.viewportNormalInCanonicalCoordinates; + const nChunk = context.transformGlobalToVoxelNormal(nWorld); + + const { u: uChunk, v: vChunk } = getBasisFromNormal(nChunk); + + const radius = layer.brushRadius.value - 0.5; + vec3.scale(uChunk, uChunk, radius); + vec3.scale(vChunk, vChunk, radius); + + const chunkToCam3 = mat3.create(); + + // manually creating chunkToCam3 matrix to avoid any unwanted scaling + for (let row = 0; row < 3; ++row) { + for (let col = 0; col < 3; ++col) { + let sum = 0; + for (let globalDim = 0; globalDim < 3; ++globalDim) { + const layerDim = globalToRenderLayerDimensions[globalDim]; + if (layerDim !== -1) { + const viewVal = viewMatrix[globalDim * 4 + row]; + const layerVal = chunkToLayerTransform[col * stride + layerDim]; + sum += viewVal * layerVal; + } + } + chunkToCam3[col * 3 + row] = sum; + } + } + + const uCam = vec3.create(); + const vCam = vec3.create(); + vec3.transformMat3(uCam, uChunk, chunkToCam3); + vec3.transformMat3(vCam, vChunk, chunkToCam3); + + const uScrX = uCam[0]; + const uScrY = uCam[1]; + const vScrX = vCam[0]; + const vScrY = vCam[1]; + + const Q11 = uScrX * uScrX + vScrX * vScrX; + const Q12 = uScrX * uScrY + vScrX * vScrY; + const Q22 = uScrY * uScrY + vScrY * vScrY; + + const trace = Q11 + Q22; + const det = Q11 * Q22 - Q12 * Q12; + + const D_sq = trace * trace - 4 * det; + const D = D_sq < 0 ? 0 : Math.sqrt(D_sq); + + const lambda1 = (trace + D) / 2; + const lambda2 = (trace - D) / 2; + + const radiusX = Math.sqrt(lambda1); + const radiusY = Math.sqrt(lambda2); + + const rotation = Math.atan2(lambda1 - Q11, Q12); + + if (radiusX > 0 && radiusY > 0) { + ctx.save(); + ctx.beginPath(); + ctx.ellipse( + panel.mouseX, + panel.mouseY, + radiusX, + radiusY, + rotation, + 0, + 2 * Math.PI, + ); + ctx.restore(); + + const isEraser = layer.shouldErase() || layer.cursorInEraseMode.value; + const color = "white"; + ctx.fillStyle = isEraser ? "red" : color; + ctx.globalAlpha = 0.2; + ctx.fill(); + ctx.globalAlpha = 1; + ctx.strokeStyle = isEraser ? "rgb(255,136,136)" : "rgba(255, 255, 255, 1)"; + ctx.lineWidth = 4; + ctx.stroke(); + ctx.strokeStyle = isEraser ? "rgb(97,0,0)" : "rgba(0, 0, 0, 1)"; + ctx.lineWidth = 1.5; + ctx.stroke(); + + return { radiusX, radiusY }; + } + + return radiusXY; +} + +export type VoxelTabElement = + | { type: "header"; label: string } + | { type: "tool-row"; tools: { toolId: string; label: string }[] } + | LayerControlDefinition; + +const TOOL_SPECIFIC_CONTROLS: LayerControlDefinition[] = + [ + { + label: "Brush size", + toolJson: { type: "vox-brush-size" }, + ...(() => { + const control = rangeLayerControl( + (layer: UserLayerWithVoxelEditing) => ({ + value: layer.brushRadius, + options: { min: 1, max: 64, step: 1 }, + }), + ); + const originalActivateTool = control.activateTool; + return { + ...control, + activateTool: (activation, controlContext) => { + originalActivateTool(activation, controlContext as RangeWidget); + + const layer = activation.tool.layer as UserLayerWithVoxelEditing; + const trigger = () => { + for (const panel of layer.manager.root.display.panels) { + if (panel instanceof SliceViewPanel) { + panel.scheduleOverlayRedraw(); + } + } + }; + + trigger(); + activation.registerDisposer( + layer.manager.root.layerSelectedValues.mouseState.changed.add( + trigger, + ), + ); + activation.registerDisposer(layer.brushRadius.changed.add(trigger)); + activation.registerDisposer(() => { + trigger(); + }); + }, + }; + })(), + }, + { + label: "Brush shape", + toolJson: { type: "vox-brush-shape" }, + ...enumLayerControl( + (layer: UserLayerWithVoxelEditing) => layer.brushShape, + ), + }, + { + label: "Max fill voxels", + toolJson: { type: "vox-flood-max-voxels" }, + ...rangeLayerControl((layer) => ({ + value: layer.floodMaxVoxels, + options: { + min: FLOODFILL_MIN_POSSIBLE_VOXELS, + max: FLOODFILL_MAX_POSSIBLE_VOXELS, + step: FLOODFILL_MIN_POSSIBLE_VOXELS, + }, + })), + }, + ]; + +const COMMON_CONTROLS: VoxelTabElement[] = [ + { type: "header", label: "Settings" }, + { + label: "Erase only selected value", + toolJson: { type: "vox-erase-mode" }, + ...checkboxLayerControl((layer) => layer.lockToSelectedValue), + }, + { type: "header", label: "Actions" }, + { + label: "Undo", + toolJson: { type: "vox-undo" }, + ...buttonLayerControl({ + text: "Undo", + onClick: (layer) => + layer.handleVoxAction("undo", new LayerActionContext()), + }), + }, + { + label: "Redo", + toolJson: { type: "vox-redo" }, + ...buttonLayerControl({ + text: "Redo", + onClick: (layer) => + layer.handleVoxAction("redo", new LayerActionContext()), + }), + }, + { + label: "Paint Value", + toolJson: { type: "vox-paint-value" }, + makeControl: (layer, context) => { + const control = document.createElement("input"); + control.type = "text"; + control.title = "Specify segment ID or intensity value to paint"; + control.addEventListener("change", () => { + try { + layer.setVoxelPaintValue(control.value); + } catch { + control.value = layer.paintValue.value.toString(); + } + }); + context.registerDisposer( + observeWatchable((value) => { + control.value = value.toString(); + }, layer.paintValue), + ); + control.value = layer.paintValue.value.toString(); + return { control, controlElement: control, parent: context }; + }, + activateTool: () => {}, + }, + { + label: "New Random Value", + toolJson: { type: "vox-random-value" }, + ...buttonLayerControl({ + text: "Random", + onClick: (layer) => + layer.handleVoxAction( + "randomize-paint-value", + new LayerActionContext(), + ), + }), + }, +]; + +export const VOXEL_LAYER_CONTROLS: LayerControlDefinition[] = + [...TOOL_SPECIFIC_CONTROLS, ...COMMON_CONTROLS].filter( + (x): x is LayerControlDefinition => + !("type" in x) || (x.type !== "header" && x.type !== "tool-row"), + ); + +export const VOXEL_TAB_LAYOUT: VoxelTabElement[] = [ + { type: "header", label: "Tools" }, + { + type: "tool-row", + tools: [ + { toolId: BRUSH_TOOL_ID, label: "Brush" }, + { toolId: FLOODFILL_TOOL_ID, label: "Flood Fill" }, + { toolId: VALUE_PICKER_TOOL_ID, label: "Value Picker" }, + ], + }, + ...COMMON_CONTROLS, +]; + +export function registerVoxelLayerControls( + layerType: UserLayerConstructor, +) { + for (const control of VOXEL_LAYER_CONTROLS) { + registerLayerControl(layerType, control); + } +} diff --git a/src/layer/voxel_annotation/draw_tab.ts b/src/layer/voxel_annotation/draw_tab.ts new file mode 100644 index 0000000000..e3a5894e9e --- /dev/null +++ b/src/layer/voxel_annotation/draw_tab.ts @@ -0,0 +1,109 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { VOXEL_TAB_LAYOUT } from "#src/layer/voxel_annotation/controls.js"; +import type { UserLayerWithVoxelEditing } from "#src/layer/voxel_annotation/index.js"; +import { observeWatchable } from "#src/trackable_value.js"; +import { makeToolButton } from "#src/ui/tool.js"; +import type { VoxelEditController } from "#src/voxel_annotation/frontend.js"; +import { DependentViewWidget } from "#src/widget/dependent_view_widget.js"; +import { addLayerControlToOptionsTab } from "#src/widget/layer_control.js"; +import { Tab } from "#src/widget/tab_view.js"; + +export class VoxToolTab extends Tab { + constructor(public layer: UserLayerWithVoxelEditing) { + super(); + const { element } = this; + + const toolbox = document.createElement("div"); + + for (const elementDef of VOXEL_TAB_LAYOUT) { + if ("type" in elementDef && elementDef.type === "header") { + const title = document.createElement("div"); + title.textContent = elementDef.label; + title.style.fontWeight = "600"; + toolbox.appendChild(title); + } else if ("type" in elementDef && elementDef.type === "tool-row") { + const toolButtonsContainer = document.createElement("div"); + toolButtonsContainer.style.display = "flex"; + toolButtonsContainer.style.gap = "8px"; + + for (const tool of elementDef.tools) { + const button = makeToolButton(this, layer.toolBinder, { + toolJson: tool.toolId, + label: tool.label, + }); + toolButtonsContainer.appendChild(button); + } + toolbox.appendChild(toolButtonsContainer); + } else { + const controlDef = elementDef as any; + const controlElement = addLayerControlToOptionsTab( + this, + this.layer, + this.visibility, + controlDef, + ); + + if ( + controlDef.toolJson.type === "vox-undo" || + controlDef.toolJson.type === "vox-redo" + ) { + const button = controlElement.querySelector("button"); + if (button) { + this.registerDisposer( + new DependentViewWidget( + { + changed: this.layer.layersChanged, + get value() { + return ( + layer.editingContexts.values().next().value + ?._controller ?? undefined + ); + }, + }, + ( + controller: VoxelEditController | undefined, + _parent, + context, + ) => { + if (!controller) { + button.disabled = true; + return; + } + const watchable = + controlDef.toolJson.type === "vox-undo" + ? controller.undoCount + : controller.redoCount; + context.registerDisposer( + observeWatchable((count) => { + button.disabled = count === 0; + }, watchable), + ); + }, + this.visibility, + ), + ); + } + } + + toolbox.appendChild(controlElement); + } + } + + element.appendChild(toolbox); + } +} diff --git a/src/layer/voxel_annotation/index.browser_test.ts b/src/layer/voxel_annotation/index.browser_test.ts new file mode 100644 index 0000000000..edda068d8c --- /dev/null +++ b/src/layer/voxel_annotation/index.browser_test.ts @@ -0,0 +1,432 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import type { ChunkManager } from "#src/chunk_manager/frontend.js"; +import type { CoordinateSpaceTransform } from "#src/coordinate_transform.js"; +import { + makeCoordinateSpace, + makeIdentityTransform, +} from "#src/coordinate_transform.js"; +import { getDefaultCredentialsManager } from "#src/credentials_provider/default_manager.js"; +import { SharedCredentialsManager } from "#src/credentials_provider/shared.js"; +import { DataManagementContext } from "#src/data_management_context.js"; +import { + DataSourceRegistry, + makeEmptyDataSourceSpecification, +} from "#src/datasource/index.js"; +import { DisplayContext } from "#src/display_context.js"; +import { SharedKvStoreContext } from "#src/kvstore/frontend.js"; +import { + LayerManager, + LayerSelectedValues, + ManagedUserLayer, + MouseSelectionState, + SelectedLayerState, + TopLevelLayerListSpecification, + TrackableDataSelectionState, +} from "#src/layer/index.js"; +import { + LayerDataSource, + LoadedDataSubsource, + LoadedLayerDataSource, +} from "#src/layer/layer_data_source.js"; +import { SegmentationUserLayer } from "#src/layer/segmentation/index.js"; +import { Position } from "#src/navigation_state.js"; +import { + DataType, + VolumeType, + makeVolumeChunkSpecification, +} from "#src/sliceview/volume/base.js"; +import { + InMemoryVolumeChunkSource, + MultiscaleVolumeChunkSource, +} from "#src/sliceview/volume/frontend.js"; +import { SegmentationRenderLayer } from "#src/sliceview/volume/segmentation_renderlayer.js"; +import { WatchableValue } from "#src/trackable_value.js"; +import { GlobalToolBinder } from "#src/ui/tool.js"; +import { mat4, vec3 } from "#src/util/geom.js"; +import "#src/sliceview/uncompressed_chunk_format.js"; + +class TestMultiscaleSource extends MultiscaleVolumeChunkSource { + constructor( + chunkManager: ChunkManager, + public dataType: DataType, + ) { + super(chunkManager); + } + get volumeType() { + return VolumeType.SEGMENTATION; + } + get rank() { + return 3; + } + getSources(options: any) { + void options; + const spec = makeVolumeChunkSpecification({ + dataType: this.dataType, + chunkDataSize: Uint32Array.from([32, 32, 32]), + lowerVoxelBound: Float32Array.from([0, 0, 0]), + upperVoxelBound: Float32Array.from([100, 100, 100]), + rank: 3, + }); + const chunkSource = this.chunkManager.getChunkSource( + InMemoryVolumeChunkSource, + { spec }, + ); + return [ + [ + { + chunkSource, + chunkToMultiscaleTransform: new Float32Array([ + 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, + ]), + }, + ], + ]; + } +} + +describe("Voxel Editing Utilities", () => { + let display: DisplayContext; + let dataContext: DataManagementContext; + + const createLayer = ( + dataType: DataType = DataType.UINT64, + modelTransform?: CoordinateSpaceTransform, + ) => { + const credentialsManager = new SharedCredentialsManager( + getDefaultCredentialsManager(), + dataContext.rpc, + ); + const kvStoreContext = new SharedKvStoreContext( + dataContext.chunkManager, + credentialsManager, + ); + const dataSourceProvider = new DataSourceRegistry(kvStoreContext); + + const layerManager = new LayerManager(); + const layerSelectedValues = new LayerSelectedValues( + layerManager, + new MouseSelectionState(), + ); + const coordinateSpace = new WatchableValue( + makeCoordinateSpace({ + names: ["x", "y", "z"], + units: ["m", "m", "m"], + scales: Float64Array.of(1, 1, 1), + }), + ); + const selectionState = new TrackableDataSelectionState( + coordinateSpace, + layerSelectedValues, + ); + const selectedLayer = new SelectedLayerState(layerManager); + const globalPosition = new Position(coordinateSpace); + const toolBinder = new GlobalToolBinder(() => {}, {} as any); + + const layerSpecification = new TopLevelLayerListSpecification( + display, + dataSourceProvider, + layerManager, + dataContext.chunkManager, + selectionState, + selectedLayer, + coordinateSpace, + globalPosition, + toolBinder, + ); + + const managedLayer = new ManagedUserLayer("test-layer", layerSpecification); + const userLayer = new SegmentationUserLayer(managedLayer); + managedLayer.layer = userLayer; + + const multiscaleSource = new TestMultiscaleSource( + dataContext.chunkManager, + dataType, + ); + + const dataSource = new LayerDataSource(userLayer); + const loadedLayerDataSource = new LoadedLayerDataSource( + dataSource, + { + canonicalUrl: "test", + modelTransform: + modelTransform ?? + makeIdentityTransform( + makeCoordinateSpace({ + names: ["x", "y", "z"], + scales: Float64Array.of(1, 1, 1), + units: ["m", "m", "m"], + }), + ), + subsources: [], + } as any, + makeEmptyDataSourceSpecification(), + ); + + userLayer.addCoordinateSpace(loadedLayerDataSource.transform.outputSpace); + + const subsourceEntry = { + id: "default", + default: true, + subsource: { volume: multiscaleSource }, + }; + + const loadedSubsource = new LoadedDataSubsource( + loadedLayerDataSource, + subsourceEntry, + undefined, + 0, + true, + ); + + let renderLayer: SegmentationRenderLayer | undefined; + + loadedSubsource.activate(() => { + const transform = loadedSubsource.getRenderLayerTransform(); + renderLayer = new SegmentationRenderLayer(multiscaleSource, { + ...userLayer.displayState, + transform, + renderScaleTarget: userLayer.sliceViewRenderScaleTarget, + renderScaleHistogram: userLayer.sliceViewRenderScaleHistogram, + localPosition: userLayer.localPosition, + }); + loadedSubsource.addRenderLayer(renderLayer); + + loadedSubsource.writingEnabled.value = true; + userLayer.initializeVoxelEditingForSubsource( + loadedSubsource, + renderLayer, + ); + }); + + if (!renderLayer) throw new Error("Failed to create renderLayer"); + + return { userLayer, loadedSubsource, renderLayer }; + }; + + beforeEach(() => { + display = new DisplayContext(document.createElement("div")); + dataContext = new DataManagementContext(display.gl, display); + }); + + afterEach(() => { + display.dispose(); + dataContext.dispose(); + }); + + describe("getVoxelPositionFromMouse", () => { + it("Success: returns mapped voxel position", () => { + const space = makeCoordinateSpace({ + names: ["x", "y", "z"], + scales: Float64Array.of(1, 1, 1), + units: ["m", "m", "m"], + }); + const transform = new Float32Array(16); + mat4.identity(transform as unknown as mat4); + mat4.translate( + transform as unknown as mat4, + transform as unknown as mat4, + [10, 5, 0], + ); + mat4.scale( + transform as unknown as mat4, + transform as unknown as mat4, + [2, 0.5, 1], + ); + + const modelTransform = { + inputSpace: space, + outputSpace: space, + transform: transform as unknown as Float64Array, + rank: 3, + sourceRank: 3, + }; + + const { userLayer, loadedSubsource } = createLayer( + DataType.UINT64, + modelTransform, + ); + const context = userLayer.editingContexts.get(loadedSubsource)!; + + const mouseState = new MouseSelectionState(); + mouseState.unsnappedPosition = Float32Array.of(20, 10, 5); + + const result = context.getVoxelPositionFromMouse(mouseState); + + expect(result).toBeDefined(); + expect(result![0]).toBeCloseTo(5); + expect(result![1]).toBeCloseTo(10); + expect(result![2]).toBeCloseTo(5); + }); + + it("Transform Error: returns undefined", () => { + const { userLayer, loadedSubsource, renderLayer } = createLayer( + DataType.UINT64, + ); + const context = userLayer.editingContexts.get(loadedSubsource)!; + + renderLayer.transform.value = { error: "Transform error" } as any; + + const mouseState = new MouseSelectionState(); + mouseState.unsnappedPosition = Float32Array.of(10, 10, 10); + + const result = context.getVoxelPositionFromMouse(mouseState); + expect(result).toBeUndefined(); + }); + + it("Out of Bounds: returns coordinate", () => { + const { userLayer, loadedSubsource } = createLayer(DataType.UINT64); + const context = userLayer.editingContexts.get(loadedSubsource)!; + + const mouseState = new MouseSelectionState(); + mouseState.unsnappedPosition = Float32Array.of(1000, 2000, 3000); + + const result = context.getVoxelPositionFromMouse(mouseState); + expect(result).toBeDefined(); + expect(result![0]).toBeCloseTo(1000); + expect(result![1]).toBeCloseTo(2000); + expect(result![2]).toBeCloseTo(3000); + }); + }); + + describe("transformGlobalToVoxelNormal", () => { + it("Identity Transform: returns same vector", () => { + const { userLayer, loadedSubsource } = createLayer(DataType.UINT64); + const context = userLayer.editingContexts.get(loadedSubsource)!; + + const mouseState = new MouseSelectionState(); + mouseState.unsnappedPosition = Float32Array.of(10, 10, 10); + context.getVoxelPositionFromMouse(mouseState); + + const globalNormal = vec3.fromValues(1, 0, 0); + const result = context.transformGlobalToVoxelNormal(globalNormal); + + expect(result).toEqual(globalNormal); + }); + + it("Rotation/Permutation: transforms vector", () => { + const space = makeCoordinateSpace({ + names: ["x", "y", "z"], + scales: Float64Array.of(1, 1, 1), + units: ["m", "m", "m"], + }); + + const transform = new Float64Array([ + 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, + ]); + + const modelTransform = { + inputSpace: space, + outputSpace: space, + transform: transform, + rank: 3, + sourceRank: 3, + }; + + const { userLayer, loadedSubsource } = createLayer( + DataType.UINT64, + modelTransform, + ); + const context = userLayer.editingContexts.get(loadedSubsource)!; + + const mouseState = new MouseSelectionState(); + mouseState.unsnappedPosition = Float32Array.of(10, 10, 10); + context.getVoxelPositionFromMouse(mouseState); + + const globalNormal = vec3.fromValues(1, 0, 0); + const result = context.transformGlobalToVoxelNormal(globalNormal); + + expect(result[0]).toBeCloseTo(0); + expect(result[1]).toBeCloseTo(1); + expect(result[2]).toBeCloseTo(0); + }); + + it("Non-aligned Normal with Scaling: correctly transforms and normalizes", () => { + const space = makeCoordinateSpace({ + names: ["x", "y", "z"], + scales: Float64Array.of(1, 1, 1), + units: ["m", "m", "m"], + }); + + const transform = new Float64Array([ + 2, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, + ]); + + const modelTransform = { + inputSpace: space, + outputSpace: space, + transform: transform, + rank: 3, + sourceRank: 3, + }; + + const { userLayer, loadedSubsource } = createLayer( + DataType.UINT64, + modelTransform, + ); + const context = userLayer.editingContexts.get(loadedSubsource)!; + + const mouseState = new MouseSelectionState(); + mouseState.unsnappedPosition = Float32Array.of(10, 10, 10); + context.getVoxelPositionFromMouse(mouseState); + + const inputLen = Math.sqrt(2); + const globalNormal = vec3.fromValues(1 / inputLen, 1 / inputLen, 0); + + const result = context.transformGlobalToVoxelNormal(globalNormal); + + // transform(2,1,1) * (1,1,0) -> (2,1,0). Normalized -> (2,1,0)/sqrt(5). + const expectedX = 1 / Math.sqrt(5); + const expectedY = 2 / Math.sqrt(5); + + expect(result[0]).toBeCloseTo(expectedX); + expect(result[1]).toBeCloseTo(expectedY); + expect(result[2]).toBeCloseTo(0); + }); + }); + + describe("setVoxelPaintValue", () => { + it("UINT8: Clamps and wraps", () => { + const { userLayer } = createLayer(DataType.UINT8); + expect(userLayer.setVoxelPaintValue(255)).toBe(255n); + expect(userLayer.setVoxelPaintValue(256)).toBe(0n); + expect(userLayer.setVoxelPaintValue(-1)).toBe(255n); + }); + + it("INT8: Signed wrapping", () => { + const { userLayer } = createLayer(DataType.INT8); + expect(userLayer.setVoxelPaintValue(127)).toBe(127n); + expect(userLayer.setVoxelPaintValue(128)).toBe(-128n); + expect(userLayer.setVoxelPaintValue(-129)).toBe(127n); + }); + + it("UINT64: Handles BigInts", () => { + const { userLayer } = createLayer(DataType.UINT64); + const bigVal = BigInt(Number.MAX_SAFE_INTEGER) + 10n; + expect(userLayer.setVoxelPaintValue(bigVal)).toBe(bigVal); + }); + + it("No Context: Fails", () => { + const { userLayer } = createLayer(DataType.UINT64); + userLayer.editingContexts.clear(); + expect(() => userLayer.setVoxelPaintValue(1)).toThrow( + "No voxel editing context available", + ); + }); + }); +}); diff --git a/src/layer/voxel_annotation/index.ts b/src/layer/voxel_annotation/index.ts new file mode 100644 index 0000000000..d46a2bd5cf --- /dev/null +++ b/src/layer/voxel_annotation/index.ts @@ -0,0 +1,887 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { + LayerActionContext, + MouseSelectionState, +} from "#src/layer/index.js"; +import { UserLayer } from "#src/layer/index.js"; +import type { LoadedDataSubsource } from "#src/layer/layer_data_source.js"; +import { + drawBrushCursor, + getEditingContext, +} from "#src/layer/voxel_annotation/controls.js"; +import { VoxToolTab } from "#src/layer/voxel_annotation/draw_tab.js"; +import type { + ChunkTransformParameters, + RenderLayerTransformOrError, +} from "#src/render_coordinate_transform.js"; +import { + getChunkPositionFromCombinedGlobalLocalPositions, + getChunkTransformParameters, +} from "#src/render_coordinate_transform.js"; +import type { RenderedDataPanel } from "#src/rendered_data_panel.js"; +import type { + SliceViewSourceOptions, + SliceViewRenderLayer, +} from "#src/sliceview/base.js"; +import { DataType } from "#src/sliceview/base.js"; +import { SliceViewPanel } from "#src/sliceview/panel.js"; +import type { MultiscaleVolumeChunkSource } from "#src/sliceview/volume/frontend.js"; +import type { ImageRenderLayer } from "#src/sliceview/volume/image_renderlayer.js"; +import { SegmentationRenderLayer } from "#src/sliceview/volume/segmentation_renderlayer.js"; +import { StatusMessage } from "#src/status.js"; +import { TrackableBoolean } from "#src/trackable_boolean.js"; +import type { WatchableValueInterface } from "#src/trackable_value.js"; +import { + makeDerivedWatchableValue, + TrackableValue, + WatchableValue, +} from "#src/trackable_value.js"; +import type { UserLayerWithAnnotations } from "#src/ui/annotations.js"; +import { randomUint64 } from "#src/util/bigint.js"; +import { RefCounted } from "#src/util/disposable.js"; +import { vec3 } from "#src/util/geom.js"; +import { + parseUint64, + verifyFiniteFloat, + verifyInt, + verifyOptionalObjectProperty, +} from "#src/util/json.js"; +import { TrackableEnum } from "#src/util/trackable_enum.js"; +import { VoxelPreviewMultiscaleSource } from "#src/voxel_annotation/preview_multiscale_chunk_source.js"; +import type { + VoxelEditControllerHost, + VoxelValueGetter, +} from "#src/voxel_annotation/base.js"; +import { + VOXEL_EDIT_STAMINA, + VOXEL_EMPTY_VALUE, + BRUSH_TOOL_ID, + BrushShape, + MAX_VOXEL_EDIT_STAMINA, +} from "#src/voxel_annotation/base.js"; +import { VoxelEditController } from "#src/voxel_annotation/frontend.js"; + +const BRUSH_SIZE_JSON_KEY = "brushSize"; +const ERASE_SELECTED_MODE_JSON_KEY = "eraseSelectedMode"; +const BRUSH_SHAPE_JSON_KEY = "brushShape"; +const FLOOD_FILL_MAX_VOXELS_JSON_KEY = "floodFillMaxVoxels"; +const FLOOD_FILL_MORPHOLOGICAL_JSON_KEY = "floodFillMorphological"; +const PAINT_VALUE_JSON_KEY = "paintValue"; + +const DATA_TYPE_BIT_INFO = { + [DataType.UINT8]: { bits: 8, signed: false }, + [DataType.INT8]: { bits: 8, signed: true }, + [DataType.UINT16]: { bits: 16, signed: false }, + [DataType.INT16]: { bits: 16, signed: true }, + [DataType.UINT32]: { bits: 32, signed: false }, + [DataType.INT32]: { bits: 32, signed: true }, + [DataType.UINT64]: { bits: 64, signed: false }, +}; + +export class VoxelEditingContext + extends RefCounted + implements VoxelEditControllerHost +{ + private readonly _controller: VoxelEditController | undefined = undefined; + private _pendingPermissionPromise: Promise | undefined; + private hasUserConfirmedWriting = false; + + private cachedChunkTransform: ChunkTransformParameters | undefined; + private cachedTransformGeneration: number = -1; + private cachedVoxelPosition: Float32Array = new Float32Array(3); + optimisticRenderLayer: + | ImageRenderLayer + | SegmentationRenderLayer + | undefined = undefined; + previewSource: VoxelPreviewMultiscaleSource | undefined = undefined; + + private localLoadEstimate = new WatchableValue(0); + public totalPending: WatchableValueInterface; + + constructor( + public hostLayer: UserLayerWithVoxelEditing, + public primarySource: MultiscaleVolumeChunkSource, + public primaryRenderLayer: ImageRenderLayer | SegmentationRenderLayer, + public writingEnabled: boolean, + public dataSourceUrl: string | undefined, + ) { + super(); + + if (!writingEnabled) return; + + // The following checks are in place due to limitations in the implementation, and could be removed if support for the checked constraint is added. + if (primarySource.rank !== 3) { + throw new Error( + `Voxel annotation only supports rank 3 volumes (got ${primarySource.rank}).`, + ); + } + if (primarySource.dataType === DataType.FLOAT32) { + throw new Error(`Voxel annotation does not support Float32 datasets.`); + } + this.validateHierarchy(primarySource); + + this.previewSource = new VoxelPreviewMultiscaleSource( + this.hostLayer.manager.chunkManager, + primarySource, + ); + + this.optimisticRenderLayer = this.hostLayer._createVoxelOverlayRenderLayer( + this.previewSource, + primaryRenderLayer.transform, + ); + + if ( + this.primaryRenderLayer instanceof SegmentationRenderLayer && + this.optimisticRenderLayer instanceof SegmentationRenderLayer + ) { + this.optimisticRenderLayer.forceHiddenFromMainRenderLoop = true; + this.primaryRenderLayer.setVoxelPreviewLayer(this.optimisticRenderLayer); + } + + // since we only allow drawing at max res, we can lock the optimistic render layer to it + ( + this.optimisticRenderLayer as SliceViewRenderLayer + ).getForcedSourceIndexOverride = () => 0; + + this.hostLayer.addRenderLayer(this.optimisticRenderLayer); + + this._controller = new VoxelEditController(this); + + this.totalPending = this.registerDisposer( + makeDerivedWatchableValue( + (local, backend) => local + backend, + this.localLoadEstimate, + this._controller.pendingOpCount, + ), + ); + } + + private async checkPermission(): Promise { + if (this.hasUserConfirmedWriting) { + return true; + } + if (this._pendingPermissionPromise) { + return false; // this._pendingPermissionPromise; + } + + this._pendingPermissionPromise = new Promise((resolve) => { + const msg = new StatusMessage(/*delay=*/ false, /*modal=*/ true); + msg.element.textContent = `Are you sure you want to write to ${this.dataSourceUrl} `; + + const yes = document.createElement("button"); + yes.textContent = "Yes"; + yes.onclick = () => { + this.hasUserConfirmedWriting = true; + msg.dispose(); + resolve(true); + }; + const no = document.createElement("button"); + no.textContent = "No"; + no.onclick = () => { + msg.dispose(); + resolve(false); + }; + msg.element.appendChild(yes); + msg.element.appendChild(no); + msg.setVisible(true); + }).then((result) => { + this._pendingPermissionPromise = undefined; + return result; + }); + + return false; // this._pendingPermissionPromise; + } + + private async withCost( + cost: number, + op: () => Promise, + ): Promise { + if (this.localLoadEstimate.value >= MAX_VOXEL_EDIT_STAMINA) return; + this.localLoadEstimate.value += cost; + try { + if (await this.checkPermission()) { + return await op(); + } + } finally { + this.localLoadEstimate.value -= cost; + } + } + + beginStroke(): number { + if (!this._controller) + throw new Error("Cannot use beginStroke without a controller"); + return this._controller.beginStroke(); + } + + rollbackStroke(seq: number): void { + this._controller?.rollbackStroke(seq); + } + + async applyBrushPreview( + points: Float32Array[], + radiusCanonical: number, + value: VoxelValueGetter, + shape: BrushShape, + basis: { u: Float32Array; v: Float32Array }, + seq: number, + filterValue?: bigint, + ) { + if (!this._controller) + throw new Error("Cannot use applyBrushPreview without a controller"); + if (!(await this.checkPermission())) return; + await this._controller.applyBrushPreview( + points, + radiusCanonical, + value, + shape, + basis, + seq, + filterValue, + ); + } + + async dispatchBrushStroke( + centers: Float32Array[], + radiusCanonical: number, + value: VoxelValueGetter, + shape: BrushShape, + basis: { u: Float32Array; v: Float32Array }, + seq: number, + filterValue?: bigint, + ) { + if (!this._controller) + throw new Error("Cannot use dispatchBrushStroke without a controller"); + const cost = + VOXEL_EDIT_STAMINA.brush( + shape, + radiusCanonical, + filterValue !== undefined, + ) * centers.length; + // The stroke's previews already tagged overlay chunks with `seq`. If the + // dispatch does not reach the backend — stamina or permission refusal in + // withCost, or an RPC failure — those edits will never be written and no + // reload would ever clear them: roll the stroke's overlay chunks back. + let dispatched = false; + try { + await this.withCost(cost, async () => { + await this._controller!.dispatchBrushStroke( + centers, + radiusCanonical, + value, + shape, + basis, + seq, + filterValue, + ); + dispatched = true; + }); + } finally { + if (!dispatched) this._controller.rollbackStroke(seq); + } + } + + async floodFillPlane2D( + startPositionCanonical: Float32Array, + fillValue: VoxelValueGetter, + maxVoxels: number, + basis: { u: Float32Array; v: Float32Array }, + filterValue?: bigint, + morphological = true, + ) { + if (!this._controller) + throw new Error("Cannot use floodFillPlane2D without a controller"); + const cost = VOXEL_EDIT_STAMINA.floodFill(maxVoxels); + await this.withCost(cost, () => + this._controller!.floodFillPlane2D( + startPositionCanonical, + fillValue, + maxVoxels, + basis, + filterValue, + morphological, + ), + ); + } + + async undo() { + if (!this._controller) + throw new Error("Cannot use undo without a controller"); + await this.withCost(VOXEL_EDIT_STAMINA.undoRedo(), () => + this._controller!.undo(), + ); + } + + async redo() { + if (!this._controller) + throw new Error("Cannot use redo without a controller"); + await this.withCost(VOXEL_EDIT_STAMINA.undoRedo(), () => + this._controller!.redo(), + ); + } + + get rpc() { + return this.hostLayer.manager.chunkManager.rpc!; + } + + disposed() { + if (this._controller) this._controller.dispose(); + if (this.optimisticRenderLayer) { + if ( + this.primaryRenderLayer instanceof SegmentationRenderLayer && + this.optimisticRenderLayer instanceof SegmentationRenderLayer + ) { + this.primaryRenderLayer.setVoxelPreviewLayer(undefined); + } + this.hostLayer.removeRenderLayer(this.optimisticRenderLayer); + } + super.disposed(); + } + + /** + * Verifies that the size of a parent chunk is an integer multiple + * of the size of a child chunk. + */ + private validateHierarchy(primarySource: MultiscaleVolumeChunkSource) { + const rank = primarySource.rank; + + const identityOptions = this.hostLayer.getIdentitySliceViewSourceOptions(); + const scales = primarySource.getSources(identityOptions)[0]; + + if (!scales || scales.length < 2) return; + + const getPhysicalChunkExtent = (lodIndex: number) => { + const source = scales[lodIndex]; + const transform = source.chunkToMultiscaleTransform; + const chunkVoxels = source.chunkSource.spec.chunkDataSize; + + const extent = new Float32Array(rank); + + for (let i = 0; i < rank; i++) { + let sumSq = 0; + for (let row = 0; row < rank; row++) { + const val = transform[i * (rank + 1) + row]; + sumSq += val * val; + } + const scaleFactor = Math.sqrt(sumSq); + extent[i] = chunkVoxels[i] * scaleFactor; + } + return extent; + }; + + for (let i = 0; i < scales.length - 1; i++) { + const childExtents = getPhysicalChunkExtent(i); + const parentExtents = getPhysicalChunkExtent(i + 1); + + for (let d = 0; d < rank; d++) { + const ratio = parentExtents[d] / childExtents[d]; + const isInteger = Math.abs(ratio - Math.round(ratio)) < 0.001; + + if (!isInteger) { + throw new Error( + `Hierarchy mismatch between LOD ${i} and ${i + 1}. ` + + `Parent chunk must contain a whole number of child chunks. ` + + `Ratio dim ${d}: ${ratio.toFixed(3)}`, + ); + } + } + } + } + + getChunkTransform(): ChunkTransformParameters | undefined { + const renderLayer = this.primaryRenderLayer; + const renderLayerTransform = renderLayer.transform.value; + if (renderLayerTransform.error !== undefined) { + return undefined; + } + + const transformGeneration = renderLayer.transform.changed.count; + if (this.cachedTransformGeneration !== transformGeneration) { + this.cachedChunkTransform = undefined; + try { + this.cachedChunkTransform = getChunkTransformParameters( + renderLayerTransform, + this.primarySource.getSources( + this.hostLayer.getIdentitySliceViewSourceOptions(), + )[0][0]!.chunkToMultiscaleTransform, + ); + this.cachedTransformGeneration = transformGeneration; + } catch (e) { + this.cachedTransformGeneration = -1; + console.error("Error computing chunk transform parameters:", e); + return undefined; + } + } + return this.cachedChunkTransform; + } + + getVoxelPositionFromMouse( + mouseState: MouseSelectionState, + ): Float32Array | undefined { + const chunkTransform = this.getChunkTransform(); + if (chunkTransform === undefined) return undefined; + + if ( + this.cachedVoxelPosition.length !== + chunkTransform.modelTransform.unpaddedRank + ) { + this.cachedVoxelPosition = new Float32Array( + chunkTransform.modelTransform.unpaddedRank, + ); + } + + const ok = getChunkPositionFromCombinedGlobalLocalPositions( + this.cachedVoxelPosition, + mouseState.unsnappedPosition, + this.hostLayer.localPosition.value, + chunkTransform.layerRank, + chunkTransform.combinedGlobalLocalToChunkTransform, + ); + if (!ok) return undefined; + return this.cachedVoxelPosition; + } + + transformGlobalToVoxelNormal(globalNormal: vec3): vec3 { + const chunkTransform = this.getChunkTransform(); + if (chunkTransform === undefined) + throw new Error("Chunk transform not computed"); + const { modelTransform, layerToChunkTransform, layerRank } = chunkTransform; + const { globalToRenderLayerDimensions } = modelTransform; + const globalRank = globalToRenderLayerDimensions.length; + const voxelNormal = vec3.create(); + + for (let chunkDim = 0; chunkDim < 3; ++chunkDim) { + let sum = 0; + for ( + let globalDim = 0; + globalDim < Math.min(globalRank, 3); + ++globalDim + ) { + const layerDim = globalToRenderLayerDimensions[globalDim]; + if (layerDim !== -1) { + sum += + layerToChunkTransform[chunkDim + layerDim * (layerRank + 1)] * + globalNormal[globalDim]; + } + } + voxelNormal[chunkDim] = sum; + } + vec3.normalize(voxelNormal, voxelNormal); + return voxelNormal; + } +} + +export declare abstract class UserLayerWithVoxelEditing extends UserLayer { + hasSubsourcesWithWritingEnabled: WatchableValue; + + brushRadius: TrackableValue; + lockToSelectedValue: TrackableBoolean; + brushShape: TrackableEnum; + floodMaxVoxels: TrackableValue; + floodMorphological: TrackableBoolean; + paintValue: TrackableValue; + cursorInEraseMode: TrackableBoolean; + + editingContexts: Map; + + abstract _createVoxelOverlayRenderLayer( + source: MultiscaleVolumeChunkSource, + transform: WatchableValueInterface, + ): ImageRenderLayer | SegmentationRenderLayer; + abstract getVoxelPaintValue(erase: boolean): VoxelValueGetter; + abstract setVoxelPaintValue(value: any): bigint; + setEraseState(erase: boolean): void; + shouldErase(): boolean; + scheduleOverlayRedraw(): void; + + initializeVoxelEditingForSubsource( + loadedSubsource: LoadedDataSubsource, + renderlayer: SegmentationRenderLayer | ImageRenderLayer, + ): void; + deinitializeVoxelEditingForSubsource( + loadedSubsource: LoadedDataSubsource, + ): void; + updateHasSubsourcesWithWritingEnabled(): void; + getIdentitySliceViewSourceOptions(): SliceViewSourceOptions; + handleVoxAction(action: string, context: LayerActionContext): void; +} + +export function UserLayerWithVoxelEditingMixin< + TBase extends { new (...args: any[]): UserLayerWithAnnotations }, +>(Base: TBase) { + abstract class C extends Base implements UserLayerWithVoxelEditing { + editingContexts = new Map(); + hasSubsourcesWithWritingEnabled = new WatchableValue(false); + paintValue = new TrackableValue(1n, (x) => parseUint64(x)); + + // Brush properties + brushRadius = new TrackableValue(3, verifyInt); + lockToSelectedValue = new TrackableBoolean(false); + brushShape = new TrackableEnum(BrushShape, BrushShape.DISK); + floodMaxVoxels = new TrackableValue(10000, verifyFiniteFloat); + floodMorphological = new TrackableBoolean(true); + cursorInEraseMode = new TrackableBoolean(false, false); + + private _isInEraseState = false; + + constructor(...args: any[]) { + super(...args); + this.registerDisposer(() => { + for (const context of this.editingContexts.values()) { + context.dispose(); + } + this.editingContexts.clear(); + }); + this.brushRadius.changed.add(this.specificationChanged.dispatch); + this.lockToSelectedValue.changed.add(this.specificationChanged.dispatch); + this.brushShape.changed.add(this.specificationChanged.dispatch); + this.floodMaxVoxels.changed.add(this.specificationChanged.dispatch); + this.floodMorphological.changed.add(this.specificationChanged.dispatch); + this.paintValue.changed.add(this.specificationChanged.dispatch); + + this.bindOverlayToPanels(); + this.registerDisposer( + this.manager.root.display.updateStarted.add(() => + this.bindOverlayToPanels(), + ), + ); + + this.brushRadius.changed.add(this.scheduleOverlayRedraw); + this.manager.root.layerSelectedValues.mouseState.changed.add( + this.scheduleOverlayRedraw, + ); + this.cursorInEraseMode.changed.add(this.scheduleOverlayRedraw); + + this.layersChanged.add(() => { + const ctx = getEditingContext(this); + if (ctx) { + this.registerDisposer( + ctx.totalPending.changed.add(this.scheduleOverlayRedraw), + ); + } + }); + this.scheduleOverlayRedraw(); + + this.tabs.add("Draw", { + label: "Draw", + order: 5, + hidden: makeDerivedWatchableValue( + (editable) => !editable, + this.hasSubsourcesWithWritingEnabled, + ), + getter: () => new VoxToolTab(this), + }); + } + + scheduleOverlayRedraw = () => { + for (const panel of this.manager.root.display.panels) { + if (panel instanceof SliceViewPanel) { + panel.scheduleOverlayRedraw(); + } + } + }; + + private boundPanelCleanups = new Map void>(); + private bindOverlayToPanels() { + for (const panel of this.manager.root.display.panels) { + if ( + panel instanceof SliceViewPanel && + !this.boundPanelCleanups.has(panel) + ) { + const rm = panel.overlayDraw.add((ctx, _w, _h, p) => + this.handleOverlayDraw(ctx, p), + ); + this.boundPanelCleanups.set(panel, rm); + } + } + for (const [panel, cleanup] of this.boundPanelCleanups) { + if (!this.manager.root.display.panels.has(panel)) { + cleanup(); + this.boundPanelCleanups.delete(panel); + } + } + } + + private handleOverlayDraw( + ctx: CanvasRenderingContext2D, + panel: RenderedDataPanel, + ) { + if ( + panel.mouseX < 0 || + panel.mouseY < 0 || + !this.hasSubsourcesWithWritingEnabled.value + ) { + return; + } + const globalToolBinder = this.manager.root.toolBinder; + const activation = globalToolBinder.activeTool_; + let radiusY = 0; + + if ( + activation && + activation.tool.localBinder === this.toolBinder && + activation.tool.toJSON() === BRUSH_TOOL_ID + ) { + const radiusXY = drawBrushCursor(this, panel, ctx); + if (radiusXY.radiusY > 0) { + radiusY = radiusXY.radiusY; + } + } + + const editContext = getEditingContext(this); + const pending = editContext?.totalPending.value || 0; + this.drawStaminaBar(ctx, panel.mouseX, panel.mouseY, pending, radiusY); + } + + private drawStaminaBar( + ctx: CanvasRenderingContext2D, + x: number, + y: number, + pendingCount: number, + brushOffset: number, + ) { + const ratio = Math.max(0, 1.0 - pendingCount / MAX_VOXEL_EDIT_STAMINA); + if (ratio > 0.99) { + return; + } + + const w = 42; + const h = 5; + const radius = h / 2; + + const yOffset = 16 + brushOffset; + const xOffset = 1; + const bx = x - w / 2 + xOffset; + const by = y + yOffset; + + ctx.save(); + + if (ratio === 0) { + const [text, textX, textY] = ["WAIT!", x + xOffset, by + 8]; + ctx.fillStyle = "#ff0000"; + ctx.textAlign = "center"; + ctx.font = "bold 16px monospace"; + ctx.lineWidth = 3; + ctx.strokeStyle = "rgba(0, 0, 0, 0.8)"; + ctx.lineJoin = "round"; + ctx.strokeText(text, textX, textY); + ctx.fillText(text, textX, textY); + } else { + ctx.beginPath(); + ctx.moveTo(bx + radius, by); + ctx.lineTo(bx + w - radius, by); + ctx.arcTo(bx + w, by, bx + w, by + h, radius); + ctx.arcTo(bx + w, by + h, bx, by + h, radius); + ctx.arcTo(bx, by + h, bx, by, radius); + ctx.arcTo(bx, by, bx + w, by, radius); + ctx.closePath(); + + ctx.fillStyle = "rgba(159,159,159,0.5)"; + ctx.fill(); + + ctx.clip(); + + ctx.fillStyle = ratio > 0.3 ? "#5e5e5e" : "#FF0000"; + ctx.fillRect(bx, by, w * ratio, h); + } + ctx.restore(); + } + + setEraseState(erase: boolean): void { + this._isInEraseState = erase; + } + + shouldErase(): boolean { + return this._isInEraseState; + } + + toJSON() { + const json = super.toJSON(); + json[BRUSH_SIZE_JSON_KEY] = this.brushRadius.toJSON(); + json[ERASE_SELECTED_MODE_JSON_KEY] = this.lockToSelectedValue.toJSON(); + json[BRUSH_SHAPE_JSON_KEY] = this.brushShape.toJSON(); + json[FLOOD_FILL_MAX_VOXELS_JSON_KEY] = this.floodMaxVoxels.toJSON(); + json[FLOOD_FILL_MORPHOLOGICAL_JSON_KEY] = + this.floodMorphological.toJSON(); + const pv = this.paintValue.toJSON(); + json[PAINT_VALUE_JSON_KEY] = pv === undefined ? undefined : pv.toString(); + return json; + } + + restoreState(specification: any) { + super.restoreState(specification); + verifyOptionalObjectProperty(specification, BRUSH_SIZE_JSON_KEY, (v) => + this.brushRadius.restoreState(v), + ); + verifyOptionalObjectProperty( + specification, + ERASE_SELECTED_MODE_JSON_KEY, + (v) => this.lockToSelectedValue.restoreState(v), + ); + verifyOptionalObjectProperty(specification, BRUSH_SHAPE_JSON_KEY, (v) => + this.brushShape.restoreState(v), + ); + verifyOptionalObjectProperty( + specification, + FLOOD_FILL_MAX_VOXELS_JSON_KEY, + (v) => this.floodMaxVoxels.restoreState(v), + ); + verifyOptionalObjectProperty( + specification, + FLOOD_FILL_MORPHOLOGICAL_JSON_KEY, + (v) => this.floodMorphological.restoreState(v), + ); + verifyOptionalObjectProperty(specification, PAINT_VALUE_JSON_KEY, (v) => + this.paintValue.restoreState(v), + ); + } + + getVoxelPaintValue(erase: boolean): VoxelValueGetter { + return (_isPreview: boolean) => + erase ? VOXEL_EMPTY_VALUE : this.paintValue.value; + } + + setVoxelPaintValue(x: any) { + const editContext = this.editingContexts.values().next().value; + if (!editContext) throw new Error("No voxel editing context available"); + const dataType = editContext.primarySource.dataType; + if (dataType === DataType.FLOAT32) { + throw new Error("Voxel annotation does not support Float32 datasets."); + } + + const value = BigInt(x); + const info = + DATA_TYPE_BIT_INFO[dataType as keyof typeof DATA_TYPE_BIT_INFO]; + if (!info) { + this.paintValue.value = value; + return value; + } + + const { bits, signed } = info; + const mask = (1n << BigInt(bits)) - 1n; + let truncated = value & mask; + + if (signed) { + const signBit = 1n << BigInt(bits - 1); + if ((truncated & signBit) !== 0n) { + truncated -= 1n << BigInt(bits); + } + } + + this.paintValue.value = truncated; + return truncated; + } + + abstract _createVoxelOverlayRenderLayer( + source: MultiscaleVolumeChunkSource, + transform: WatchableValueInterface, + ): ImageRenderLayer | SegmentationRenderLayer; + + updateHasSubsourcesWithWritingEnabled(): void { + this.hasSubsourcesWithWritingEnabled.value = this.editingContexts + .entries() + .some((value) => value[1].writingEnabled); + } + + initializeVoxelEditingForSubsource( + loadedSubsource: LoadedDataSubsource, + renderlayer: SegmentationRenderLayer | ImageRenderLayer, + writingEnabled: boolean = true, + ): void { + if (writingEnabled) { + for (const [otherSubsource, _] of this.editingContexts) { + if ( + otherSubsource !== loadedSubsource && + otherSubsource.writingEnabled.value + ) { + otherSubsource.writingEnabled.value = false; + } + } + } + + if (this.editingContexts.has(loadedSubsource)) return; + + const primarySource = loadedSubsource.subsourceEntry.subsource + .volume as MultiscaleVolumeChunkSource; + + try { + const context = new VoxelEditingContext( + this, + primarySource, + renderlayer, + writingEnabled, + loadedSubsource.loadedDataSource.dataSource.canonicalUrl, + ); + this.editingContexts.set(loadedSubsource, context); + this.updateHasSubsourcesWithWritingEnabled(); + this.setVoxelPaintValue(this.paintValue.value); + } catch (e) { + if (writingEnabled) { + loadedSubsource.writingEnabled.value = false; + this.updateHasSubsourcesWithWritingEnabled(); + const msg = e instanceof Error ? e.message : String(e); + console.warn("Failed to initialize voxel editing:", msg); + StatusMessage.showTemporaryMessage(msg, 5000); + } + } + } + + deinitializeVoxelEditingForSubsource(loadedSubsource: LoadedDataSubsource) { + const context = this.editingContexts.get(loadedSubsource); + if (context) { + context.dispose(); + this.editingContexts.delete(loadedSubsource); + } + this.updateHasSubsourcesWithWritingEnabled(); + } + + getIdentitySliceViewSourceOptions(): SliceViewSourceOptions { + const rank = this.localCoordinateSpace.value.rank; + const displayRank = rank; + const multiscaleToViewTransform = new Float32Array(displayRank * rank); + for (let chunkDim = 0; chunkDim < rank; ++chunkDim) { + for (let displayDim = 0; displayDim < displayRank; ++displayDim) { + multiscaleToViewTransform[displayRank * chunkDim + displayDim] = + chunkDim === displayDim ? 1 : 0; + } + } + return { + displayRank, + multiscaleToViewTransform, + modelChannelDimensionIndices: [], + }; + } + + handleVoxAction(action: string, _context: LayerActionContext): void { + const firstContext = this.editingContexts.values().next() + .value as VoxelEditingContext; + if (!firstContext) return; + switch (action) { + case "undo": + void firstContext.undo(); + break; + case "redo": + void firstContext.redo(); + break; + case "randomize-paint-value": + this.setVoxelPaintValue(randomUint64()); + break; + } + } + } + return C; +} diff --git a/src/rendered_data_panel.ts b/src/rendered_data_panel.ts index fd8a82ff01..6cdc1d394d 100644 --- a/src/rendered_data_panel.ts +++ b/src/rendered_data_panel.ts @@ -40,6 +40,7 @@ import { KeyboardEventBinder } from "#src/util/keyboard_bindings.js"; import * as matrix from "#src/util/matrix.js"; import { MouseEventBinder } from "#src/util/mouse_bindings.js"; import { startRelativeMouseDrag } from "#src/util/mouse_drag.js"; +import { Signal } from "#src/util/signal.js"; import type { TouchPinchInfo, TouchTranslateInfo, @@ -165,6 +166,9 @@ export abstract class RenderedDataPanel extends RenderedPanel { */ pickRequestPending = false; + private overlay_canvas: HTMLCanvasElement; + private overlay_context: CanvasRenderingContext2D; + private mouseStateForcer = () => this.blockOnPickRequest(); protected isMovingToMousePosition: boolean = false; @@ -812,6 +816,52 @@ export abstract class RenderedDataPanel extends RenderedPanel { } }, ); + + this.overlay_canvas = document.createElement("canvas"); + this.overlay_canvas.style.position = "absolute"; + this.overlay_canvas.style.top = "0"; + this.overlay_canvas.style.left = "0"; + this.overlay_canvas.style.width = "100%"; + this.overlay_canvas.style.height = "100%"; + this.overlay_canvas.style.pointerEvents = "none"; + this.overlay_canvas.style.zIndex = "10"; + this.element.appendChild(this.overlay_canvas); + this.overlay_context = this.overlay_canvas.getContext("2d")!; + + this.boundsUpdated.add(() => { + this.overlay_canvas.width = this.renderViewport.logicalWidth; + this.overlay_canvas.height = this.renderViewport.logicalHeight; + }); + } + + overlayDraw = new Signal< + ( + ctx: CanvasRenderingContext2D, + width: number, + height: number, + panel: RenderedDataPanel, + ) => void + >(); + + scheduleOverlayRedraw() { + if (this.visible) { + requestAnimationFrame(this.drawOverlayInternal.bind(this)); + } + } + + private drawOverlayInternal() { + this.overlay_context.clearRect( + 0, + 0, + this.renderViewport.logicalWidth, + this.renderViewport.logicalHeight, + ); + this.overlayDraw.dispatch( + this.overlay_context, + this.renderViewport.logicalWidth, + this.renderViewport.logicalHeight, + this, + ); } abstract translateDataPointByViewportPixels( diff --git a/src/renderlayer.ts b/src/renderlayer.ts index 583e2026ae..c880c7dc43 100644 --- a/src/renderlayer.ts +++ b/src/renderlayer.ts @@ -71,6 +71,7 @@ export class RenderLayer extends RefCounted { layerChanged = new NullarySignal(); redrawNeeded = new NullarySignal(); layerChunkProgressInfo = new LayerChunkProgressInfo(); + forceHiddenFromMainRenderLoop = false; handleAction(_action: string) { // Do nothing by default. diff --git a/src/sliceview/README.md b/src/sliceview/README.md index 8851e64279..8939962e17 100644 --- a/src/sliceview/README.md +++ b/src/sliceview/README.md @@ -2,7 +2,7 @@ This directory contains the code for `SliceView`, which provides the cross-secti # Architecture -A volume is divided into a regular grid of 3-d chunks. Each chunk has voxel dimensions `chunkDataSize` (a 3-d vector of positive integers). All chunks have the same dimensions, except at the the upper bound of the volume in each dimension, where the chunks are allowed to be truncated to fit within the volume dimensions. +A volume is divided into a regular grid of 3-d chunks. Each chunk has voxel dimensions `chunkDataSize` (a 3-d vector of positive integers). All chunks have the same dimensions, except at the upper bound of the volume in each dimension, where the chunks are allowed to be truncated to fit within the volume dimensions. Chunks are the unit at which portions of the volume are queued, retrieved, transcoded (if necessary), copied to the GPU, and rendered: diff --git a/src/sliceview/base.ts b/src/sliceview/base.ts index b6d061cc07..855ae6d7da 100644 --- a/src/sliceview/base.ts +++ b/src/sliceview/base.ts @@ -169,6 +169,14 @@ export interface SliceViewRenderLayer { localPosition: WatchableValueInterface; renderScaleTarget: WatchableValueInterface; + /** + * If implemented by a render layer, return a non-negative integer scale index to override + * automatic multiscale selection. When defined, the sliceview must use only the specified + * scale from the current orientation. Implementations must ensure the index is valid for + * their multiscale source; this function should return undefined when no override is desired. + */ + getForcedSourceIndexOverride?(): number | undefined; + filterVisibleSources( sliceView: SliceViewBase, sources: readonly TransformedSource[], @@ -686,6 +694,22 @@ export function* filterVisibleSources( renderLayer: SliceViewRenderLayer, sources: readonly TransformedSource[], ): Iterable { + // allows a render layer to force a specific multiscale + const forcedIndex = renderLayer.getForcedSourceIndexOverride?.(); + if (forcedIndex !== undefined) { + if ( + !Number.isInteger(forcedIndex) || + forcedIndex < 0 || + forcedIndex >= sources.length + ) { + throw new Error( + `filterVisibleSources: forced source index ${forcedIndex} is out of range [0, ${sources.length - 1}]`, + ); + } + yield sources[forcedIndex]; + return; + } + // Increase pixel size by a small margin. const pixelSize = sliceView.projectionParameters.value.pixelSize * 1.1; // At the smallest scale, all alternative sources must have the same voxel size, which is diff --git a/src/sliceview/chunk_base.ts b/src/sliceview/chunk_base.ts new file mode 100644 index 0000000000..27146c992a --- /dev/null +++ b/src/sliceview/chunk_base.ts @@ -0,0 +1,31 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { ChunkState } from "#src/chunk_manager/base.js"; +import { Chunk } from "#src/chunk_manager/frontend.js"; +import type { SliceViewChunkSource } from "#src/sliceview/frontend.js"; +import type { vec3 } from "#src/util/geom.js"; + +export class SliceViewChunk extends Chunk { + chunkGridPosition: vec3; + declare source: SliceViewChunkSource; + + constructor(source: SliceViewChunkSource, x: any) { + super(source); + this.chunkGridPosition = x.chunkGridPosition; + this.state = ChunkState.SYSTEM_MEMORY; + } +} diff --git a/src/sliceview/compressed_segmentation/chunk_format.ts b/src/sliceview/compressed_segmentation/chunk_format.ts index b025c49720..db2fc6ac05 100644 --- a/src/sliceview/compressed_segmentation/chunk_format.ts +++ b/src/sliceview/compressed_segmentation/chunk_format.ts @@ -26,7 +26,7 @@ import type { ChunkFormatHandler, VolumeChunkSource, } from "#src/sliceview/volume/frontend.js"; -import { registerChunkFormatHandler } from "#src/sliceview/volume/frontend.js"; +import { registerChunkFormatHandler } from "#src/sliceview/volume/registry.js"; import { RefCounted } from "#src/util/disposable.js"; import { vec3, vec3Key } from "#src/util/geom.js"; import type { GL } from "#src/webgl/context.js"; diff --git a/src/sliceview/frontend.ts b/src/sliceview/frontend.ts index b87475efb6..ea6ff2a111 100644 --- a/src/sliceview/frontend.ts +++ b/src/sliceview/frontend.ts @@ -19,8 +19,9 @@ import { ChunkState } from "#src/chunk_manager/base.js"; import type { ChunkManager, ChunkRequesterState, + Chunk, } from "#src/chunk_manager/frontend.js"; -import { Chunk, ChunkSource } from "#src/chunk_manager/frontend.js"; +import { ChunkSource } from "#src/chunk_manager/frontend.js"; import { applyRenderViewportToProjectionMatrix } from "#src/display_context.js"; import type { LayerManager } from "#src/layer/index.js"; import type { @@ -60,6 +61,7 @@ import { SliceViewBase, SliceViewProjectionParameters, } from "#src/sliceview/base.js"; +import type { SliceViewChunk } from "#src/sliceview/chunk_base.js"; import { ChunkLayout } from "#src/sliceview/chunk_layout.js"; import type { SliceViewerState } from "#src/sliceview/panel.js"; import { SliceViewRenderLayer } from "#src/sliceview/renderlayer.js"; @@ -89,6 +91,8 @@ import { getSquareCornersBuffer } from "#src/webgl/square_corners_buffer.js"; import type { RPC } from "#src/worker_rpc.js"; import { registerSharedObjectOwner } from "#src/worker_rpc.js"; +export { SliceViewChunk } from "#src/sliceview/chunk_base.js"; + export type GenericChunkKey = string; class FrontendSliceViewBase extends SliceViewBase< @@ -420,7 +424,9 @@ export class SliceView extends Base { visibleLayerList.length = 0; for (const renderLayer of this.layerManager.readyRenderLayers()) { if (renderLayer instanceof SliceViewRenderLayer) { - visibleLayerList.push(renderLayer); + if (!renderLayer.forceHiddenFromMainRenderLoop) { + visibleLayerList.push(renderLayer); + } let layerInfo = visibleLayers.get(renderLayer); if (layerInfo === undefined) { const disposers: Disposer[] = []; @@ -708,16 +714,11 @@ export interface SliceViewChunkSource { getChunk(x: any): any; } +/* export class SliceViewChunk extends Chunk { - chunkGridPosition: vec3; - declare source: SliceViewChunkSource; - - constructor(source: SliceViewChunkSource, x: any) { - super(source); - this.chunkGridPosition = x.chunkGridPosition; - this.state = ChunkState.SYSTEM_MEMORY; - } + // MOVED to chunk_base.ts to avoid import loop } +*/ /** * Helper for rendering a SliceView that has been pre-rendered to a texture. diff --git a/src/sliceview/panel.ts b/src/sliceview/panel.ts index adb9c8cc41..751e83b7d5 100644 --- a/src/sliceview/panel.ts +++ b/src/sliceview/panel.ts @@ -279,6 +279,12 @@ export class SliceViewPanel extends RenderedDataPanel { ); } + handleMouseMove(clientX: number, clientY: number) { + super.handleMouseMove(clientX, clientY); + this.viewer.mouseState.planeNormal = + this.sliceView.projectionParameters.value.viewportNormalInCanonicalCoordinates; + } + translateByViewportPixels(deltaX: number, deltaY: number): void { const { pose } = this.viewer.navigationState; pose.updateDisplayPosition((pos) => { diff --git a/src/sliceview/single_texture_chunk_format.ts b/src/sliceview/single_texture_chunk_format.ts index 53468aedfd..81958fa953 100644 --- a/src/sliceview/single_texture_chunk_format.ts +++ b/src/sliceview/single_texture_chunk_format.ts @@ -14,11 +14,11 @@ * limitations under the License. */ +import { VolumeChunk } from "#src/sliceview/volume/chunk.js"; import type { VolumeChunkSource, ChunkFormat, } from "#src/sliceview/volume/frontend.js"; -import { VolumeChunk } from "#src/sliceview/volume/frontend.js"; import type { TypedArray } from "#src/util/array.js"; import type { DataType } from "#src/util/data_type.js"; import type { Disposable } from "#src/util/disposable.js"; @@ -146,6 +146,28 @@ export abstract class SingleTextureVolumeChunk< gl.bindTexture(textureTarget, null); } + updateFromCpuData(gl: GL) { + if (this.data == null) return; + + if (this.texture == null) { + this.copyToGPU(gl); + return; + } + + const textureTarget = + textureTargetForSamplerType[this.chunkFormat.shaderSamplerType]; + gl.bindTexture(textureTarget, this.texture); + try { + this.chunkFormat.setTextureData( + gl, + this.textureLayout!, + this.data as unknown as TypedArray, + ); + } finally { + gl.bindTexture(textureTarget, null); + } + } + freeGPUMemory(gl: GL) { super.freeGPUMemory(gl); if (this.data === null) return; diff --git a/src/sliceview/uncompressed_chunk_format.ts b/src/sliceview/uncompressed_chunk_format.ts index 2f1eadbd71..09195806e1 100644 --- a/src/sliceview/uncompressed_chunk_format.ts +++ b/src/sliceview/uncompressed_chunk_format.ts @@ -26,7 +26,7 @@ import type { ChunkFormatHandler, VolumeChunkSource, } from "#src/sliceview/volume/frontend.js"; -import { registerChunkFormatHandler } from "#src/sliceview/volume/frontend.js"; +import { registerChunkFormatHandler } from "#src/sliceview/volume/registry.js"; import type { TypedArray, TypedNumberArrayConstructor, diff --git a/src/sliceview/volume/backend.spec.ts b/src/sliceview/volume/backend.spec.ts new file mode 100644 index 0000000000..ba5a35c0f9 --- /dev/null +++ b/src/sliceview/volume/backend.spec.ts @@ -0,0 +1,349 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { ChunkState } from "#src/chunk_manager/base.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 { HttpError } from "#src/util/http_request.js"; +import type { RPC } from "#src/worker_rpc.js"; + +vi.mock("#src/sliceview/compressed_segmentation/decode_uint64.js", () => ({ + decodeChannel: vi.fn((out) => out.fill(5n)), +})); +vi.mock("#src/sliceview/compressed_segmentation/encode_uint64.js", () => ({ + encodeChannel: vi.fn((builder) => { + builder.data = new Uint32Array([888]); + }), +})); +vi.mock("#src/sliceview/compressed_segmentation/decode_uint32.js", () => ({ + decodeChannel: vi.fn((out) => out.fill(5)), +})); +vi.mock("#src/sliceview/compressed_segmentation/encode_uint32.js", () => ({ + encodeChannel: vi.fn((builder) => { + builder.data = new Uint32Array([444]); + }), +})); + +vi.mock("#src/sliceview/volume/registry.js", () => ({ + getChunkFormatHandler: vi.fn().mockReturnValue({ + chunkFormat: { dataType: 0 }, + dispose: vi.fn(), + getChunk: (source: any, x: any) => new source.chunkConstructor(source, x), + }), +})); + +class MockBackendSource extends VolumeChunkSource { + public serverStorage = new Map(); + + async download(chunk: VolumeChunk) { + const key = chunk.chunkGridPosition.join(","); + if (this.serverStorage.has(key)) { + chunk.data = new Uint8Array(this.serverStorage.get(key)!.slice(0)); + } + } + + async writeChunk(chunk: VolumeChunk) { + const key = chunk.chunkGridPosition.join(","); + this.serverStorage.set(key, chunk.data!.buffer.slice(0) as ArrayBuffer); + } +} + +describe("VolumeChunkSource: applyEdits", () => { + let mockRpc: RPC; + let source: MockBackendSource; + + const BASE_SPEC = { + rank: 3, + dataType: DataType.UINT64, + chunkDataSize: Uint32Array.from([2, 2, 2]), + upperVoxelBound: Float32Array.from([10, 10, 10]), + baseVoxelOffset: Float32Array.from([0, 0, 0]), + compressedSegmentationBlockSize: undefined, + }; + + beforeEach(() => { + const mockQueueManager = { + sources: new Set(), + adjustCapacitiesForChunk: vi.fn(), + updateChunkState: vi.fn(), + scheduleUpdate: vi.fn(), + invalidateCachedChunks: vi.fn(), + moveChunkToFrontend: vi.fn(), + markRecentlyUsed: vi.fn(), + gl: {}, + }; + + const mockChunkManager = { + queueManager: mockQueueManager, + chunkQueueManager: mockQueueManager, + rpc: null, + memoize: { get: (_k: string, fn: Function) => fn() }, + }; + + mockRpc = { + newId: () => 0, + set: vi.fn(), + get: vi.fn().mockReturnValue(mockChunkManager), + invoke: vi.fn(), + promiseInvoke: vi.fn(), + } as unknown as RPC; + + source = new MockBackendSource(mockRpc, { + spec: { ...BASE_SPEC }, + chunkManager: 0, + }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe("Input Validation", () => { + it("should throw if indices and values lengths mismatch", async () => { + await expect(source.applyEdits("0,0,0", [1], [1n, 2n])).rejects.toThrow( + /length mismatch/, + ); + }); + + it("should throw for invalid chunk keys (wrong rank)", async () => { + await expect(source.applyEdits("0,0", [0], [1n])).rejects.toThrow( + /invalid chunk key/, + ); + }); + + it("should throw for invalid chunk keys (NaN)", async () => { + await expect(source.applyEdits("0,NaN,0", [0], [1n])).rejects.toThrow( + /invalid chunk key/, + ); + }); + }); + + describe("Uncompressed Data (UINT64)", () => { + it("should initialize data if missing", async () => { + const writeSpy = vi.spyOn(source, "writeChunk"); + const result = await source.applyEdits("0,0,0", [0], [100n]); + + expect(writeSpy).toHaveBeenCalled(); + const written = writeSpy.mock.calls[0][0] as VolumeChunk; + expect(written.data).toBeInstanceOf(BigUint64Array); + expect((written.data as BigUint64Array)[0]).toBe(100n); + expect(result.newValues[0]).toBe(100n); + }); + + it("should update existing data", async () => { + const chunk = source.getChunk(new Float32Array([0, 0, 0])) as VolumeChunk; + chunk.data = new BigUint64Array(8); + chunk.state = ChunkState.SYSTEM_MEMORY; + (chunk.data as BigUint64Array)[0] = 50n; + + const result = await source.applyEdits("0,0,0", [0], [100n]); + + expect(result.oldValues[0]).toBe(50n); + expect(result.newValues[0]).toBe(100n); + }); + + it("should throw on out-of-bounds index", async () => { + await expect(source.applyEdits("0,0,0", [9], [1n])).rejects.toThrow( + /index 9 out of bounds/, + ); + }); + }); + + describe("Uncompressed Data (UINT32)", () => { + it("should handle edits correctly", async () => { + const uint32Spec = { ...BASE_SPEC, dataType: DataType.UINT32 }; + const uint32Source = new MockBackendSource(mockRpc, { + spec: uint32Spec, + chunkManager: 0, + }); + + const writeSpy = vi.spyOn(uint32Source, "writeChunk"); + const result = await uint32Source.applyEdits("0,0,0", [0], [123]); + + const written = writeSpy.mock.calls[0][0] as VolumeChunk; + expect(written.data).toBeInstanceOf(Uint32Array); + expect((written.data as Uint32Array)[0]).toBe(123); + expect(result.newValues[0]).toBe(123); + }); + }); + + describe("Compressed Segmentation", () => { + it("should handle UINT64 compressed segmentation", async () => { + const compressedSpec = { + ...BASE_SPEC, + compressedSegmentationBlockSize: Uint32Array.from([2, 2, 1]), + }; + const compressedSource = new MockBackendSource(mockRpc, { + spec: compressedSpec, + chunkManager: 0, + }); + + const chunk = compressedSource.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + chunk.data = new Uint32Array([123]); + chunk.state = ChunkState.SYSTEM_MEMORY_WORKER; + + const writeSpy = vi.spyOn(compressedSource, "writeChunk"); + const result = await compressedSource.applyEdits("0,0,0", [0], [99n]); + + expect(result.oldValues[0]).toBe(5n); + const written = writeSpy.mock.calls[0][0] as VolumeChunk; + expect((written.data as Uint32Array)[0]).toBe(888); + // The shared cache entry is invalidated, never mutated in place. + expect((chunk.data as Uint32Array)[0]).toBe(123); + }); + + it("should handle UINT32 compressed segmentation", async () => { + const compressedSpec = { + ...BASE_SPEC, + dataType: DataType.UINT32, + compressedSegmentationBlockSize: Uint32Array.from([2, 2, 1]), + }; + const compressedSource = new MockBackendSource(mockRpc, { + spec: compressedSpec, + chunkManager: 0, + }); + + const chunk = compressedSource.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + chunk.data = new Uint32Array([123]); + chunk.state = ChunkState.SYSTEM_MEMORY_WORKER; + + const writeSpy = vi.spyOn(compressedSource, "writeChunk"); + const result = await compressedSource.applyEdits("0,0,0", [0], [77]); + + expect(result.oldValues[0]).toBe(5); + expect(result.newValues[0]).toBe(77); + const written = writeSpy.mock.calls[0][0] as VolumeChunk; + expect((written.data as Uint32Array)[0]).toBe(444); + }); + + it("should handle zero-offset compressed data (empty/new)", async () => { + const compressedSpec = { + ...BASE_SPEC, + compressedSegmentationBlockSize: Uint32Array.from([2, 2, 1]), + }; + const compressedSource = new MockBackendSource(mockRpc, { + spec: compressedSpec, + chunkManager: 0, + }); + + const chunk = compressedSource.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + chunk.data = new Uint32Array([]); + chunk.state = ChunkState.SYSTEM_MEMORY_WORKER; + + const writeSpy = vi.spyOn(compressedSource, "writeChunk"); + await compressedSource.applyEdits("0,0,0", [0], [50n]); + const written = writeSpy.mock.calls[0][0] as VolumeChunk; + expect((written.data as Uint32Array)[0]).toBe(888); + }); + + it("should handle zero-offset compressed data for UINT32", async () => { + const compressedSpec = { + ...BASE_SPEC, + dataType: DataType.UINT32, + compressedSegmentationBlockSize: Uint32Array.from([2, 2, 1]), + }; + const compressedSource = new MockBackendSource(mockRpc, { + spec: compressedSpec, + chunkManager: 0, + }); + + const chunk = compressedSource.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + chunk.data = new Uint32Array([]); + chunk.state = ChunkState.SYSTEM_MEMORY_WORKER; + + const writeSpy = vi.spyOn(compressedSource, "writeChunk"); + await compressedSource.applyEdits("0,0,0", [0], [50]); + const written = writeSpy.mock.calls[0][0] as VolumeChunk; + expect((written.data as Uint32Array)[0]).toBe(444); + }); + }); + + describe("Error Handling & Bounds", () => { + it("should throw if chunk size cannot be determined", async () => { + const computeSpy = vi + .spyOn(source, "computeChunkBounds") + .mockImplementation(() => new Float32Array()); + + const chunk = source.getChunk(new Float32Array([0, 0, 0])) as VolumeChunk; + chunk.chunkDataSize = null; + + await expect(source.applyEdits("0,0,0", [0], [1n])).rejects.toThrow( + /size is unknown/, + ); + + computeSpy.mockRestore(); + }); + + it("should retry on 500 errors and eventually succeed", async () => { + vi.useFakeTimers(); + const writeSpy = vi + .spyOn(source, "writeChunk") + .mockRejectedValueOnce(new HttpError("", 500, "")) + .mockRejectedValueOnce(new HttpError("", 503, "")) + .mockResolvedValue(undefined); + + const promise = source.applyEdits("0,0,0", [0], [1n]); + await vi.runAllTimersAsync(); + await promise; + + expect(writeSpy).toHaveBeenCalledTimes(3); + vi.useRealTimers(); + }); + + it("should stop retrying at one point", async () => { + vi.useFakeTimers(); + + vi.spyOn(source, "writeChunk").mockRejectedValue( + new Error("Fatal DB Error"), + ); + + const promise = source.applyEdits("0,0,0", [0], [1n]); + + const assertRejection = expect(promise).rejects.toThrow( + /Failed to write chunk/, + ); + + await vi.runAllTimersAsync(); + await assertRejection; + + vi.useRealTimers(); + }); + + it("should NOT retry on 400 errors", async () => { + vi.useFakeTimers(); + const writeSpy = vi + .spyOn(source, "writeChunk") + .mockRejectedValue(new HttpError("Bad Request", 400, "")); + + const promise = source.applyEdits("0,0,0", [0], [1n]); + + await expect(promise).rejects.toThrow(/Failed to write chunk/); + + expect(writeSpy).toHaveBeenCalledTimes(1); + vi.useRealTimers(); + }); + }); +}); diff --git a/src/sliceview/volume/backend.ts b/src/sliceview/volume/backend.ts index 4c7c39b8de..0e0b9b776b 100644 --- a/src/sliceview/volume/backend.ts +++ b/src/sliceview/volume/backend.ts @@ -15,21 +15,31 @@ */ import type { Chunk } from "#src/chunk_manager/backend.js"; +import { ChunkState } from "#src/chunk_manager/base.js"; import { SliceViewChunk, SliceViewChunkSourceBackend, } from "#src/sliceview/backend.js"; -import type { - DataType, - SliceViewChunkSpecification, -} from "#src/sliceview/base.js"; +import type { SliceViewChunkSpecification } from "#src/sliceview/base.js"; +import { DataType } from "#src/sliceview/base.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 { encodeChannel as encodeChannelUint32 } from "#src/sliceview/compressed_segmentation/encode_uint32.js"; +import { encodeChannel as encodeChannelUint64 } from "#src/sliceview/compressed_segmentation/encode_uint64.js"; import type { VolumeChunkSource as VolumeChunkSourceInterface, VolumeChunkSpecification, } from "#src/sliceview/volume/base.js"; +import { IN_MEMORY_VOLUME_CHUNK_SOURCE_RPC_ID } from "#src/sliceview/volume/base.js"; +import type { TypedArray } from "#src/util/array.js"; +import { TypedArrayBuilder } from "#src/util/array.js"; +import { DATA_TYPE_ARRAY_CONSTRUCTOR } from "#src/util/data_type.js"; import type { vec3 } from "#src/util/geom.js"; +import { HttpError } from "#src/util/http_request.js"; import * as vector from "#src/util/vector.js"; +import type { VoxelChange } from "#src/voxel_annotation/base.js"; import type { RPC } from "#src/worker_rpc.js"; +import { registerSharedObject } from "#src/worker_rpc.js"; export class VolumeChunk extends SliceViewChunk { source: VolumeChunkSource | null = null; @@ -138,6 +148,16 @@ export function computeChunkBounds( return chunkPosition; } +// Recognizes deterministic decode failures proving the stored bytes themselves +// are corrupt (e.g. the empty-body objects written by an earlier bug). Only +// these may be repaired by overwriting the stored chunk. Any UNRECOGNIZED +// error is treated as potentially transient (network, server, cancellation, +// codec OOM, wrapped kvstore errors) and must fail the edit instead: repairing +// on a false positive would durably overwrite valid data (fail-closed). +function isCorruptStoredChunkError(e: unknown): boolean { + return e instanceof Error && /Raw-format chunk is \d+ bytes/.test(e.message); +} + export class VolumeChunkSource extends SliceViewChunkSourceBackend implements VolumeChunkSourceInterface @@ -145,6 +165,7 @@ export class VolumeChunkSource declare spec: VolumeChunkSpecification; tempChunkDataSize: Uint32Array; tempChunkPosition: Float32Array; + constructor(rpc: RPC, options: any) { super(rpc, options); const rank = this.spec.rank; @@ -155,5 +176,240 @@ export class VolumeChunkSource computeChunkBounds(chunk: VolumeChunk) { return computeChunkBounds(this, chunk); } + + /** + * Returns a chunk for `chunkGridPosition` that is NOT registered in the shared + * chunk cache. The queue manager never sees it, so concurrent invalidation or + * eviction cannot dispose it (`chunk.source = null`) while an edit-path + * download or write is in flight on it. Discard after use; it must never be + * added to `this.chunks`. + */ + getIsolatedChunk(chunkGridPosition: Float32Array): VolumeChunk { + const chunk = new (this.chunkConstructor as new () => VolumeChunk)(); + chunk.source = this; + chunk.initializeVolumeChunk( + chunkGridPosition.join(), + chunkGridPosition as vec3, + ); + return chunk; + } + + // Override in data source backends to actually persist the chunk. + // Default throws to ensure write capability is explicitly implemented. + async writeChunk(_chunk: VolumeChunk): Promise { + throw new Error( + "VolumeChunkSource.writeChunk not implemented for this datasource", + ); + } + + async applyEdits( + chunkKey: string, + indices: ArrayLike, + values: ArrayLike, + ): Promise { + if (indices.length !== values.length) { + throw new Error("applyEdits: indices and values length mismatch"); + } + const chunkGridPosition = new Float32Array(chunkKey.split(",").map(Number)); + if ( + chunkGridPosition.length !== this.spec.rank || + chunkGridPosition.some((v) => !Number.isFinite(v)) + ) { + throw new Error(`applyEdits: invalid chunk key ${chunkKey}`); + } + + // The whole read-modify-write runs on an isolated chunk the queue manager + // cannot see: mutating or encoding a shared cache entry races with + // concurrent invalidation (dispose nulls `chunk.source` mid-download) and + // with promotion/serialization (the buffer transfer to the frontend + // detaches it mid-write). If the shared entry has resident data, copy it + // synchronously (a detached source buffer makes `slice()` throw, failing + // the edit loudly instead of writing garbage); otherwise download. The + // shared entry itself is invalidated after the successful write below. + const resident = this.chunks.get(chunkKey) as VolumeChunk | undefined; + const chunk = this.getIsolatedChunk(chunkGridPosition); + if ( + resident !== undefined && + resident.state <= ChunkState.SYSTEM_MEMORY_WORKER && + resident.data + ) { + chunk.chunkDataSize = resident.chunkDataSize; + chunk.data = (resident.data as TypedArray).slice(); + } else { + try { + await this.download(chunk, new AbortController().signal); + } catch (e) { + if (!isCorruptStoredChunkError(e)) throw e; + // The stored object itself is corrupt: treat the chunk as absent so + // the edit proceeds on fill data and the write below repairs it. + console.warn( + `applyEdits: stored chunk ${chunkKey} is unreadable; ` + + `treating it as empty and repairing it on write.`, + e, + ); + } + } + + if (!chunk.chunkDataSize) { + this.computeChunkBounds(chunk); + } + if (!chunk.chunkDataSize) { + throw new Error( + `applyEdits: Cannot create new chunk ${chunkKey} because its size is unknown.`, + ); + } + + if (!chunk.data) { + const numElements = chunk.chunkDataSize.reduce((a, b) => a * b, 1); + let Ctor; + if (this.spec.compressedSegmentationBlockSize !== undefined) { + Ctor = Uint32Array; + } else { + Ctor = DATA_TYPE_ARRAY_CONSTRUCTOR[this.spec.dataType]; + } + chunk.data = new (Ctor as any)(numElements) as TypedArray; + } + + const ArrayCtor = DATA_TYPE_ARRAY_CONSTRUCTOR[this.spec.dataType] as any; + const indicesCopy = new Uint32Array(indices); + const newValuesArray = new ArrayCtor(values.length); + for (let i = 0; i < values.length; ++i) { + newValuesArray[i] = + this.spec.dataType === DataType.UINT64 + ? values[i]! + : Number(values[i]!); + } + const oldValuesArray = new ArrayCtor(indices.length); + + if (this.spec.compressedSegmentationBlockSize !== undefined) { + const compressedData = chunk.data as Uint32Array; + const { chunkDataSize } = chunk; + const numElements = + chunkDataSize[0] * chunkDataSize[1] * chunkDataSize[2]; + const { dataType, compressedSegmentationBlockSize: subchunkSize } = + this.spec; + const baseOffset = compressedData.length > 0 ? compressedData[0] : 0; + + let uncompressedData: Uint32Array | BigUint64Array; + if (dataType === DataType.UINT32) { + uncompressedData = new Uint32Array(numElements); + if (baseOffset !== 0) { + decodeChannelUint32( + uncompressedData, + compressedData, + baseOffset, + chunkDataSize, + subchunkSize!, + ); + } + } else { + uncompressedData = new BigUint64Array(numElements); + if (baseOffset !== 0) { + decodeChannelUint64( + uncompressedData, + compressedData, + baseOffset, + chunkDataSize, + subchunkSize!, + ); + } + } + + for (let i = 0; i < indices.length; ++i) { + const idx = indices[i]!; + if (idx < 0 || idx >= uncompressedData.length) { + throw new Error( + `applyEdits: index ${idx} out of bounds for chunk ${chunkKey}`, + ); + } + oldValuesArray[i] = uncompressedData[idx]; + if (dataType === DataType.UINT32) { + (uncompressedData as Uint32Array)[idx] = Number(values[i]!); + newValuesArray[i] = Number(values[i]!); + } else { + (uncompressedData as BigUint64Array)[idx] = values[i]! as bigint; + newValuesArray[i] = values[i]! as bigint; + } + } + + const outputBuilder = new TypedArrayBuilder(Uint32Array); + outputBuilder.resize(1); + outputBuilder.data[0] = 1; + + if (dataType === DataType.UINT32) { + encodeChannelUint32( + outputBuilder, + subchunkSize!, + uncompressedData as Uint32Array, + chunkDataSize, + ); + } else { + encodeChannelUint64( + outputBuilder, + subchunkSize!, + uncompressedData as BigUint64Array, + chunkDataSize, + ); + } + + chunk.data = outputBuilder.view; + } else { + const data = chunk.data as TypedArray; + for (let i = 0; i < indices.length; ++i) { + const idx = indices[i]!; + if (idx < 0 || idx >= data.length) { + throw new Error( + `applyEdits: index ${idx} out of bounds for chunk ${chunkKey}`, + ); + } + oldValuesArray[i] = data[idx]; + data[idx] = newValuesArray[i]; + } + } + + const maxRetries = 3; + let lastError: Error | undefined; + + for (let i = 0; i < maxRetries; i++) { + try { + await this.writeChunk(chunk); + // The edit was written from an isolated chunk, so the shared cache + // entry (if any) still holds pre-edit data — or has a pre-edit + // download in flight that would otherwise later land as "fresh" and + // even be reused by the next edit, durably erasing this write. + // Invalidate it backend-side: routing through the frontend RPC would + // miss chunks the frontend does not hold (e.g. still DOWNLOADING). + const { queueManager } = this.chunkManager; + queueManager.invalidateCachedChunks(this, [chunkKey]); + queueManager.scheduleUpdate(); + return { + indices: indicesCopy, + oldValues: oldValuesArray, + newValues: newValuesArray, + }; + } catch (e) { + lastError = e as Error; + if (e instanceof HttpError && e.status < 500 && e.status !== 429) { + break; + } + await new Promise((resolve) => + setTimeout(resolve, 250 * Math.pow(2, i)), + ); + } + } + throw new Error( + `Failed to write chunk ${chunkKey} after ${maxRetries} attempts.`, + { cause: lastError }, + ); + } } + +@registerSharedObject(IN_MEMORY_VOLUME_CHUNK_SOURCE_RPC_ID) +export class InMemoryVolumeChunkSourceBackend extends VolumeChunkSource { + async download(chunk: VolumeChunk, _signal: AbortSignal): Promise { + chunk.data = null; + return new Promise((_resolve) => {}); + } +} + VolumeChunkSource.prototype.chunkConstructor = VolumeChunk; diff --git a/src/sliceview/volume/base.ts b/src/sliceview/volume/base.ts index 583d8653e7..a272452d51 100644 --- a/src/sliceview/volume/base.ts +++ b/src/sliceview/volume/base.ts @@ -153,6 +153,22 @@ export function makeVolumeChunkSpecification( }; } +export function computeChunkGridPosition( + chunkGridPosition: Float32Array, + positionWithinChunk: Uint32Array, + voxelCoord: Float32Array, + chunkDataSize: Uint32Array | Float32Array, +) { + const rank = chunkGridPosition.length; + for (let i = 0; i < rank; ++i) { + const voxel = voxelCoord[i]; + const size = chunkDataSize[i]; + const chunkIndex = Math.floor(voxel / size); + chunkGridPosition[i] = chunkIndex; + positionWithinChunk[i] = Math.floor(voxel - size * chunkIndex); + } +} + function shouldTranscodeToCompressedSegmentation( options: VolumeChunkSpecificationDefaultCompressionOptions & VolumeChunkSpecificationOptions & @@ -310,3 +326,5 @@ export interface VolumeChunkSource extends SliceViewChunkSource { } export const VOLUME_RPC_ID = "volume"; +export const IN_MEMORY_VOLUME_CHUNK_SOURCE_RPC_ID = + "sliceview/volume/InMemoryChunkSource"; diff --git a/src/sliceview/volume/chunk.ts b/src/sliceview/volume/chunk.ts new file mode 100644 index 0000000000..b764cdfabd --- /dev/null +++ b/src/sliceview/volume/chunk.ts @@ -0,0 +1,39 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { SliceViewChunk } from "#src/sliceview/chunk_base.js"; +import type { + ChunkFormat, + VolumeChunkSource, +} from "#src/sliceview/volume/frontend.js"; +import type { GL } from "#src/webgl/context.js"; + +export abstract class VolumeChunk extends SliceViewChunk { + declare source: VolumeChunkSource; + chunkDataSize: Uint32Array; + declare CHUNK_FORMAT_TYPE: ChunkFormat; + + get chunkFormat(): this["CHUNK_FORMAT_TYPE"] { + return this.source.chunkFormat; + } + + constructor(source: VolumeChunkSource, x: any) { + super(source, x); + this.chunkDataSize = x.chunkDataSize || source.spec.chunkDataSize; + } + abstract getValueAt(dataPosition: Uint32Array): any; + abstract updateFromCpuData(gl: GL): void; +} diff --git a/src/sliceview/volume/frontend.spec.ts b/src/sliceview/volume/frontend.spec.ts new file mode 100644 index 0000000000..892c8c8eb0 --- /dev/null +++ b/src/sliceview/volume/frontend.spec.ts @@ -0,0 +1,230 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { describe, it, expect, vi, beforeEach, beforeAll } from "vitest"; +import { DataType } from "#src/util/data_type.js"; + +class MockChunk { + data: any = null; + chunkDataSize: Uint32Array; + chunkGridPosition: Float32Array; + state = 0; + + constructor( + public source: any, + x: any, + ) { + this.chunkGridPosition = x.chunkGridPosition; + this.chunkDataSize = source.spec.chunkDataSize; + } + + initializeVolumeChunk() {} + dispose() {} + + updateFromCpuData = vi.fn(); + freeGPUMemory = vi.fn(); +} + +vi.mock("#src/sliceview/volume/registry.js", () => ({ + getChunkFormatHandler: () => ({ + chunkFormat: { + dataType: 0, + }, + getChunk: (source: any, x: any) => { + return new MockChunk(source, x); + }, + dispose: () => {}, + }), +})); + +describe("InMemoryVolumeChunkSource", () => { + let InMemoryVolumeChunkSource: any; + let chunkManagerMock: any; + let glMock: any; + let visibleChunksChangedMock: any; + + beforeAll(async () => { + if (typeof WebGL2RenderingContext === "undefined") { + global.WebGL2RenderingContext = class { + static VERTEX_SHADER = 35633; + static FRAGMENT_SHADER = 35632; + static ARRAY_BUFFER = 34962; + static STATIC_DRAW = 35044; + static TEXTURE_2D = 3553; + static TEXTURE_3D = 32879; + } as any; + } + if (typeof WebGLTexture === "undefined") { + global.WebGLTexture = class {} as any; + } + + const mod = await import("#src/sliceview/volume/frontend.js"); + InMemoryVolumeChunkSource = mod.InMemoryVolumeChunkSource; + }); + + beforeEach(() => { + glMock = { name: "mockGL" }; + visibleChunksChangedMock = { dispatch: vi.fn() }; + chunkManagerMock = { + chunkQueueManager: { + gl: glMock, + visibleChunksChanged: visibleChunksChangedMock, + sources: { add: () => {}, delete: () => {} }, + }, + rpc: { + newId: () => 0, + invoke: () => {}, + register: () => {}, + delete: () => {}, + get: () => {}, + set: () => {}, + }, + }; + }); + + const createSource = (dataType: DataType) => { + const spec: any = { + rank: 3, + chunkDataSize: Uint32Array.from([2, 2, 2]), + dataType, + upperVoxelBound: Float32Array.from([10, 10, 10]), + lowerVoxelBound: Float32Array.from([0, 0, 0]), + baseVoxelOffset: Float32Array.from([0, 0, 0]), + }; + return new InMemoryVolumeChunkSource(chunkManagerMock, { spec }); + }; + + it("Lazy Creation: creates a chunk if it does not exist", () => { + const source = createSource(DataType.UINT64); + const edits = new Map(); + edits.set("0,0,0", { indices: [0], value: 123n }); + + expect(source.chunks.size).toBe(0); + source.applyLocalEdits(edits); + expect(source.chunks.size).toBe(1); + expect(source.chunks.has("0,0,0")).toBe(true); + }); + + it("Lazy Allocation: allocates data buffer if null", () => { + const source = createSource(DataType.UINT64); + const edits = new Map(); + edits.set("0,0,0", { indices: [0], value: 123n }); + + source.applyLocalEdits(edits); + const chunk = source.chunks.get("0,0,0") as unknown as MockChunk; + expect(chunk.data).toBeInstanceOf(BigUint64Array); + expect(chunk.data).toHaveLength(8); // 2*2*2 + expect(chunk.data[0]).toBe(123n); + }); + + it("Data Type Handling: UINT32", () => { + const source = createSource(DataType.UINT32); + const edits = new Map(); + edits.set("0,0,0", { indices: [1], value: 456n }); + + source.applyLocalEdits(edits); + const chunk = source.chunks.get("0,0,0") as unknown as MockChunk; + expect(chunk.data).toBeInstanceOf(Uint32Array); + expect(chunk.data[1]).toBe(456); + }); + + it("Data Type Handling: UINT8", () => { + const source = createSource(DataType.UINT8); + const edits = new Map(); + edits.set("0,0,0", { indices: [2], value: 255n }); + + source.applyLocalEdits(edits); + const chunk = source.chunks.get("0,0,0") as unknown as MockChunk; + expect(chunk.data).toBeInstanceOf(Uint8Array); + expect(chunk.data[2]).toBe(255); + }); + + it("GPU Trigger: calls updateFromCpuData and dispatches change", () => { + const source = createSource(DataType.UINT64); + const edits = new Map(); + edits.set("0,0,0", { indices: [0], value: 123n }); + + source.applyLocalEdits(edits); + const chunk = source.chunks.get("0,0,0") as unknown as MockChunk; + + expect(chunk.updateFromCpuData).toHaveBeenCalledWith(glMock); + expect(visibleChunksChangedMock.dispatch).toHaveBeenCalled(); + }); + + it("Updates existing chunk data", () => { + const source = createSource(DataType.UINT64); + source.applyLocalEdits(new Map([["0,0,0", { indices: [0], value: 123n }]])); + + const chunk = source.chunks.get("0,0,0") as unknown as MockChunk; + chunk.updateFromCpuData.mockClear(); + visibleChunksChangedMock.dispatch.mockClear(); + + source.applyLocalEdits(new Map([["0,0,0", { indices: [0], value: 456n }]])); + + expect(chunk.data[0]).toBe(456n); + expect(chunk.updateFromCpuData).toHaveBeenCalledWith(glMock); + expect(visibleChunksChangedMock.dispatch).toHaveBeenCalled(); + }); + + it("Uses provided chunkGridPosition instead of parsing the key", () => { + const source = createSource(DataType.UINT64); + const chunkGridPosition = Float32Array.of(4, 5, 6); + + source.applyLocalEdits( + new Map([["4,5,6", { indices: [0], value: 123n, chunkGridPosition }]]), + ); + + const chunk = source.chunks.get("4,5,6") as unknown as MockChunk; + expect(chunk.chunkGridPosition).toBe(chunkGridPosition); + expect(chunk.data[0]).toBe(123n); + }); + + it("Overlay seqs: default to 0, are set per key and filtered by seq", () => { + const source = createSource(DataType.UINT64); + expect(source.getOverlaySeq("0,0,0")).toBe(0); + + source.setOverlaySeq("0,0,0", 3); + source.setOverlaySeq("1,0,0", 3); + source.setOverlaySeq("2,0,0", 4); + + expect(source.getOverlaySeq("0,0,0")).toBe(3); + expect(source.keysWithOverlaySeq(3).sort()).toEqual(["0,0,0", "1,0,0"]); + expect(source.keysWithOverlaySeq(4)).toEqual(["2,0,0"]); + expect(source.keysWithOverlaySeq(5)).toEqual([]); + }); + + it("Overlay seqs: purged when the chunk is deleted", () => { + const source = createSource(DataType.UINT64); + source.applyLocalEdits(new Map([["0,0,0", { indices: [0], value: 1n }]])); + source.setOverlaySeq("0,0,0", 7); + + source.invalidateChunks(["0,0,0"]); + + expect(source.getOverlaySeq("0,0,0")).toBe(0); + expect(source.keysWithOverlaySeq(7)).toEqual([]); + }); + + it("Applies contiguous index ranges with typed-array fills", () => { + const source = createSource(DataType.UINT64); + + source.applyLocalEdits( + new Map([["0,0,0", { indices: [], indexRanges: [1, 3], value: 99n }]]), + ); + + const chunk = source.chunks.get("0,0,0") as unknown as MockChunk; + expect(Array.from(chunk.data)).toEqual([0n, 99n, 99n, 99n, 0n, 0n, 0n, 0n]); + }); +}); diff --git a/src/sliceview/volume/frontend.ts b/src/sliceview/volume/frontend.ts index a373b5644a..d4c96411a2 100644 --- a/src/sliceview/volume/frontend.ts +++ b/src/sliceview/volume/frontend.ts @@ -14,29 +14,55 @@ * limitations under the License. */ -import type { ChunkManager } from "#src/chunk_manager/frontend.js"; +import { ChunkState } from "#src/chunk_manager/base.js"; +import type { Chunk, ChunkManager } from "#src/chunk_manager/frontend.js"; import type { ChunkChannelAccessParameters } from "#src/render_coordinate_transform.js"; -import type { - DataType, - SliceViewChunkSpecification, -} from "#src/sliceview/base.js"; +import type { SliceViewChunkSpecification } from "#src/sliceview/base.js"; +import { DataType } from "#src/sliceview/base.js"; +import type { SliceViewChunk } from "#src/sliceview/frontend.js"; import { MultiscaleSliceViewChunkSource, - SliceViewChunk, SliceViewChunkSource, } from "#src/sliceview/frontend.js"; +import type { UncompressedVolumeChunk } from "#src/sliceview/uncompressed_chunk_format.js"; import type { VolumeChunkSource as VolumeChunkSourceInterface, VolumeChunkSpecification, VolumeSourceOptions, VolumeType, } from "#src/sliceview/volume/base.js"; +import { + computeChunkGridPosition, + IN_MEMORY_VOLUME_CHUNK_SOURCE_RPC_ID, +} from "#src/sliceview/volume/base.js"; +import { VolumeChunk } from "#src/sliceview/volume/chunk.js"; +import { getChunkFormatHandler } from "#src/sliceview/volume/registry.js"; +import type { TypedArray } from "#src/util/array.js"; +import { DATA_TYPE_ARRAY_CONSTRUCTOR } from "#src/util/data_type.js"; import type { Disposable } from "#src/util/disposable.js"; import type { GL } from "#src/webgl/context.js"; import type { ShaderBuilder, ShaderProgram } from "#src/webgl/shader.js"; import { getShaderType, glsl_mixLinear } from "#src/webgl/shader_lib.js"; +import { registerSharedObjectOwner } from "#src/worker_rpc.js"; + +export interface LocalVolumeEdit { + indices: number[]; + indexRanges?: number[]; + value: bigint; + chunkGridPosition?: Float32Array; +} -export type VolumeChunkKey = string; +function parseChunkGridPositionKey(key: string): Float32Array { + const pos = new Float32Array(3); + let component = 0; + let start = 0; + for (let i = 0; i <= key.length && component < 3; ++i) { + if (i !== key.length && key.charCodeAt(i) !== 44) continue; + pos[component++] = Number(key.slice(start, i)); + start = i + 1; + } + return pos; +} export interface ChunkFormat { shaderKey: string; @@ -154,27 +180,6 @@ export interface ChunkFormatHandler extends Disposable { getChunk(source: SliceViewChunkSource, x: any): SliceViewChunk; } -export type ChunkFormatHandlerFactory = ( - gl: GL, - spec: VolumeChunkSpecification, -) => ChunkFormatHandler | null; - -const chunkFormatHandlers = new Array(); - -export function registerChunkFormatHandler(factory: ChunkFormatHandlerFactory) { - chunkFormatHandlers.push(factory); -} - -export function getChunkFormatHandler(gl: GL, spec: VolumeChunkSpecification) { - for (const handler of chunkFormatHandlers) { - const result = handler(gl, spec); - if (result != null) { - return result; - } - } - throw new Error("No chunk format handler found."); -} - export class VolumeChunkSource extends SliceViewChunkSource implements VolumeChunkSourceInterface @@ -196,6 +201,18 @@ export class VolumeChunkSource this.tempPositionWithinChunk = new Uint32Array(rank); } + addChunk(key: string, chunk: Chunk) { + // A `new` chunk update can replace a chunk kept on display by a lazy + // invalidation. Free the replaced chunk's texture right before the swap: + // it happens in a single tick with no intermediate frame, and the plain + // `chunks.set` would otherwise orphan the texture. + const existing = this.chunks.get(key); + if (existing !== undefined && existing.state === ChunkState.GPU_MEMORY) { + existing.freeGPUMemory(this.gl); + } + super.addChunk(key, chunk); + } + static encodeSpec(spec: SliceViewChunkSpecification) { const s = spec as VolumeChunkSpecification; return { @@ -212,6 +229,22 @@ export class VolumeChunkSource return this.chunkFormatHandler.chunkFormat; } + computeChunkIndices(voxelCoord: Float32Array): { + chunkGridPosition: Float32Array; + positionWithinChunk: Uint32Array; + } { + computeChunkGridPosition( + this.tempChunkGridPosition, + this.tempPositionWithinChunk, + voxelCoord, + this.spec.chunkDataSize, + ); + return { + chunkGridPosition: this.tempChunkGridPosition, + positionWithinChunk: this.tempPositionWithinChunk, + }; + } + getValueAt( chunkPosition: Float32Array, channelAccess: ChunkChannelAccessParameters, @@ -267,20 +300,123 @@ export class VolumeChunkSource } } -export abstract class VolumeChunk extends SliceViewChunk { - declare source: VolumeChunkSource; - chunkDataSize: Uint32Array; - declare CHUNK_FORMAT_TYPE: ChunkFormat; +@registerSharedObjectOwner(IN_MEMORY_VOLUME_CHUNK_SOURCE_RPC_ID) +export class InMemoryVolumeChunkSource extends VolumeChunkSource { + // Stroke seq of the last local edit that touched each chunk. A reload may + // clear a chunk only once its write coverage reaches this seq. Entries live + // and die with the chunk (purged in deleteChunk), so the map never outgrows + // the set of live overlay chunks. + private overlaySeqs = new Map(); - get chunkFormat(): this["CHUNK_FORMAT_TYPE"] { - return this.source.chunkFormat; + setOverlaySeq(key: string, seq: number): void { + this.overlaySeqs.set(key, seq); } - constructor(source: VolumeChunkSource, x: any) { - super(source, x); - this.chunkDataSize = x.chunkDataSize || source.spec.chunkDataSize; + getOverlaySeq(key: string): number { + return this.overlaySeqs.get(key) ?? 0; + } + + clearOverlaySeq(key: string): void { + this.overlaySeqs.delete(key); + } + + keysWithOverlaySeq(seq: number): string[] { + const keys: string[] = []; + for (const [key, s] of this.overlaySeqs) { + if (s === seq) keys.push(key); + } + return keys; + } + + deleteChunk(key: string) { + this.overlaySeqs.delete(key); + super.deleteChunk(key); + } + + constructor( + chunkManager: ChunkManager, + options: { spec: VolumeChunkSpecification }, + ) { + super(chunkManager, options); + this.initializeCounterpart(this.chunkManager.rpc!, {}); + } + + private invalidateGpuData(chunks: Set): void { + if (chunks.size === 0) return; + for (const chunk of chunks) { + chunk.updateFromCpuData(this.chunkManager.chunkQueueManager.gl); + } + this.chunkManager.chunkQueueManager.visibleChunksChanged.dispatch(); + } + + invalidateChunks(keys: string[], _options?: { lazy?: boolean }): void { + // Signature matches the base `ChunkSource.invalidateChunks`, but `lazy` does + // not apply here: an in-memory source has no backend refetch to swap in, so + // there is nothing to keep the stale chunk on screen for. Deletion is always + // immediate; the crossfade with the real data is timed by the caller (the + // overlay is dropped when its pending swap resolves, or as a rollback on + // write failure or an undispatched stroke), not by a blind delay here. + 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.chunkManager.chunkQueueManager.visibleChunksChanged.dispatch(); + } + } + + applyLocalEdits(edits: Map): void { + const chunksToUpdate = new Set(); + const { dataType } = this.spec; + const isUint64 = dataType === DataType.UINT64; + + for (const [key, edit] of edits.entries()) { + const chunkGridPosition = + edit.chunkGridPosition ?? parseChunkGridPositionKey(key); + + let chunk = this.chunks.get(key) as UncompressedVolumeChunk | undefined; + if (chunk === undefined) { + chunk = this.getChunk({ + chunkGridPosition, + }) as UncompressedVolumeChunk; + this.addChunk(key, chunk); + } + + if (chunk.data == undefined) { + const numElements = chunk.chunkDataSize.reduce((a, b) => a * b, 1); + const Ctor = DATA_TYPE_ARRAY_CONSTRUCTOR[dataType]; + chunk.data = new (Ctor as any)(numElements) as TypedArray; + } + chunksToUpdate.add(chunk); + + const cpuArray = chunk.data!; + const fillValue = isUint64 ? edit.value : Number(edit.value); + + const { indexRanges } = edit; + if (indexRanges !== undefined) { + for (let i = 0; i < indexRanges.length; i += 2) { + const start = indexRanges[i]!; + const length = indexRanges[i + 1]!; + (cpuArray as any).fill(fillValue, start, start + length); + } + } + + for (const index of edit.indices) { + if (isUint64) { + (cpuArray as BigUint64Array)[index] = edit.value; + } else { + cpuArray[index] = fillValue as number; + } + } + } + + this.invalidateGpuData(chunksToUpdate); } - abstract getValueAt(dataPosition: Uint32Array): any; } export abstract class MultiscaleVolumeChunkSource extends MultiscaleSliceViewChunkSource< @@ -290,3 +426,5 @@ export abstract class MultiscaleVolumeChunkSource extends MultiscaleSliceViewChu abstract dataType: DataType; abstract volumeType: VolumeType; } + +export { VolumeChunk }; diff --git a/src/sliceview/volume/registry.ts b/src/sliceview/volume/registry.ts new file mode 100644 index 0000000000..59a8ae6e64 --- /dev/null +++ b/src/sliceview/volume/registry.ts @@ -0,0 +1,40 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { VolumeChunkSpecification } from "#src/sliceview/volume/base.js"; +import type { ChunkFormatHandler } from "#src/sliceview/volume/frontend.js"; +import type { GL } from "#src/webgl/context.js"; + +export type ChunkFormatHandlerFactory = ( + gl: GL, + spec: VolumeChunkSpecification, +) => ChunkFormatHandler | null; + +const chunkFormatHandlers = new Array(); + +export function registerChunkFormatHandler(factory: ChunkFormatHandlerFactory) { + chunkFormatHandlers.push(factory); +} + +export function getChunkFormatHandler(gl: GL, spec: VolumeChunkSpecification) { + for (const handler of chunkFormatHandlers) { + const result = handler(gl, spec); + if (result != null) { + return result; + } + } + throw new Error("No chunk format handler found."); +} diff --git a/src/sliceview/volume/segmentation_renderlayer.ts b/src/sliceview/volume/segmentation_renderlayer.ts index c4a6c44b16..d2dd6b5760 100644 --- a/src/sliceview/volume/segmentation_renderlayer.ts +++ b/src/sliceview/volume/segmentation_renderlayer.ts @@ -35,6 +35,7 @@ import type { SliceView, SliceViewSingleResolutionSource, } from "#src/sliceview/frontend.js"; +import type { SliceViewRenderContext } from "#src/sliceview/renderlayer.js"; import type { MultiscaleVolumeChunkSource, VolumeChunkSource, @@ -43,6 +44,7 @@ import type { RenderLayerBaseOptions } from "#src/sliceview/volume/renderlayer.j import { SliceViewVolumeRenderLayer } from "#src/sliceview/volume/renderlayer.js"; import type { WatchableValueInterface } from "#src/trackable_value.js"; import { + WatchableValue, AggregateWatchableValue, makeCachedDerivedWatchableValue, } from "#src/trackable_value.js"; @@ -85,6 +87,7 @@ interface ShaderParameters { hideSegmentZero: boolean; hasSegmentDefaultColor: boolean; hasHighlightColor: boolean; + isForOptimisticPreview: boolean; } const HAS_SELECTED_SEGMENT_FLAG = 1; @@ -108,11 +111,14 @@ export class SegmentationRenderLayer extends SliceViewVolumeRenderLayer; constructor( multiscaleSource: MultiscaleVolumeChunkSource, public displayState: SliceViewSegmentationDisplayState, ) { + const isForOptimisticPreview = new WatchableValue(false); super(multiscaleSource, { shaderParameters: new AggregateWatchableValue((refCounted) => ({ hasEquivalences: refCounted.registerDisposer( @@ -163,12 +169,14 @@ export class SegmentationRenderLayer extends SliceViewVolumeRenderLayer + + + + + + + + + + + + + + + diff --git a/src/ui/layer_data_sources_tab.css b/src/ui/layer_data_sources_tab.css index a18860239f..e4df9dcaa3 100644 --- a/src/ui/layer_data_sources_tab.css +++ b/src/ui/layer_data_sources_tab.css @@ -19,6 +19,10 @@ flex-direction: column; } +.neuroglancer-layer-data-source-writable-label { + margin-right: 5px; +} + .neuroglancer-layer-data-sources-container { overflow-y: auto; display: flex; diff --git a/src/ui/layer_data_sources_tab.ts b/src/ui/layer_data_sources_tab.ts index 7b53214290..0cbf3e280c 100644 --- a/src/ui/layer_data_sources_tab.ts +++ b/src/ui/layer_data_sources_tab.ts @@ -37,9 +37,15 @@ import { createImageLayerAsMultiChannel } from "#src/layer/multi_channel_setup.j import { MeshSource, MultiscaleMeshSource } from "#src/mesh/frontend.js"; import { SkeletonSource } from "#src/skeleton/frontend.js"; import { MultiscaleVolumeChunkSource } from "#src/sliceview/volume/frontend.js"; -import { TrackableBooleanCheckbox } from "#src/trackable_boolean.js"; +import { + ElementVisibilityFromTrackableBoolean, + TrackableBooleanCheckbox, +} from "#src/trackable_boolean.js"; import type { WatchableValueInterface } from "#src/trackable_value.js"; -import { WatchableValue } from "#src/trackable_value.js"; +import { + makeCachedDerivedWatchableValue, + WatchableValue, +} from "#src/trackable_value.js"; import type { DebouncedFunction } from "#src/util/animation_frame_debounce.js"; import { animationFrameDebounce } from "#src/util/animation_frame_debounce.js"; import { DataType } from "#src/util/data_type.js"; @@ -193,22 +199,58 @@ export class DataSourceSubsourceView extends RefCounted { this.registerDisposer( loadedSource.enabledSubsourcesChanged.add(updateActiveAttribute), ); + const enabledState: WatchableValueInterface & { + set value(v: boolean); + } = { + get value() { + return loadedSubsource.enabled; + }, + set value(value: boolean) { + if (loadedSubsource.enabled === value) return; + loadedSubsource.enabled = value; + loadedSource.enableDefaultSubsources = false; + loadedSource.enabledSubsourcesChanged.dispatch(); + }, + changed: loadedSource.enabledSubsourcesChanged, + }; const enabledCheckbox = this.registerDisposer( - new TrackableBooleanCheckbox({ - get value() { - return loadedSubsource.enabled; - }, - set value(value: boolean) { - loadedSubsource.enabled = value; - loadedSource.enableDefaultSubsources = false; - loadedSource.enabledSubsourcesChanged.dispatch(); - }, - changed: loadedSource.enabledSubsourcesChanged, - }), + new TrackableBooleanCheckbox(enabledState), ); sourceInfoLine.classList.add("neuroglancer-layer-data-sources-info-line"); sourceInfoLine.appendChild(enabledCheckbox.element); + if ( + loadedSubsource.subsourceEntry.subsource.volume instanceof + MultiscaleVolumeChunkSource + ) { + const writableCheckbox = this.registerDisposer( + new TrackableBooleanCheckbox(loadedSubsource.writingEnabled), + ); + writableCheckbox.element.title = "Enable voxel editing for this source"; + const writableLabel = document.createElement("label"); + writableLabel.className = "neuroglancer-layer-data-source-writable-label"; + writableLabel.appendChild(writableCheckbox.element); + writableLabel.appendChild(document.createTextNode("[Enable writing?]")); + + this.registerDisposer( + new ElementVisibilityFromTrackableBoolean( + makeCachedDerivedWatchableValue( + (enabled, supportsWriting) => enabled && supportsWriting, + [ + enabledState, + new WatchableValue( + loadedSubsource.subsourceEntry.subsource.supportsWriting ?? + false, + ), + ], + ), + writableLabel, + ), + ); + + sourceInfoLine.appendChild(writableLabel); + } + const sourceId = document.createElement("span"); sourceId.classList.add("neuroglancer-layer-data-sources-source-id"); const { id } = loadedSubsource.subsourceEntry; diff --git a/src/ui/value_picker_cursor.svg b/src/ui/value_picker_cursor.svg new file mode 100644 index 0000000000..5eb7cf0f04 --- /dev/null +++ b/src/ui/value_picker_cursor.svg @@ -0,0 +1,14 @@ + + + + + + + + + + + + + + diff --git a/src/ui/voxel_annotations.css b/src/ui/voxel_annotations.css new file mode 100644 index 0000000000..f7ce9a849d --- /dev/null +++ b/src/ui/voxel_annotations.css @@ -0,0 +1,29 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +.neuroglancer-voxel-tool-options-body { + display: flex; + flex-direction: row; + flex-wrap: wrap; + align-items: center; + gap: 15px; +} + +.neuroglancer-tool-activation-status-header { + white-space: nowrap; + font-weight: bold; + margin: 5px; +} diff --git a/src/ui/voxel_annotations.ts b/src/ui/voxel_annotations.ts new file mode 100644 index 0000000000..071a5b6e16 --- /dev/null +++ b/src/ui/voxel_annotations.ts @@ -0,0 +1,735 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import "#src/ui/voxel_annotations.css"; + +import type { MouseSelectionState } from "#src/layer/index.js"; +import { + getEditingContext, + VOXEL_LAYER_CONTROLS, +} from "#src/layer/voxel_annotation/controls.js"; +import type { + UserLayerWithVoxelEditing, + VoxelEditingContext, +} from "#src/layer/voxel_annotation/index.js"; +import type { ChunkChannelAccessParameters } from "#src/render_coordinate_transform.js"; +import { StatusMessage } from "#src/status.js"; +import { TrackableBoolean } from "#src/trackable_boolean.js"; +import { linkWatchableValue } from "#src/trackable_value.js"; +import svg_floodFillCursor from "#src/ui/flood_fill_cursor.svg?raw"; +import { + LayerTool, + makeToolActivationStatusMessageWithHeader, + registerTool, + ToolBindingWidget, + type ToolActivation, +} from "#src/ui/tool.js"; +import svg_valuePicker from "#src/ui/value_picker_cursor.svg?raw"; +import { vec3 } from "#src/util/geom.js"; +import type { ActionEvent } from "#src/util/mouse_bindings.js"; +import { EventActionMap } from "#src/util/mouse_bindings.js"; +import { startRelativeMouseDrag } from "#src/util/mouse_drag.js"; +import { WatchableVisibilityPriority } from "#src/visibility_priority/frontend.js"; +import { + BRUSH_TOOL_ID, + type BrushShape, + FLOODFILL_TOOL_ID, + getBasisFromNormal, + VALUE_PICKER_TOOL_ID, + type VoxelValueGetter, +} from "#src/voxel_annotation/base.js"; + +const BRUSH_INPUT_MAP = EventActionMap.fromObject({ + ["at:control+mousedown0"]: "paint-voxels", + ["at:control+shift+mousedown0"]: "erase-voxels", + ["at:control+shift?+mousedown1"]: "flood-fill-shortcut", +}); + +const FLOOD_INPUT_MAP = EventActionMap.fromObject({ + ["at:control+mousedown0"]: "paint-voxels", + ["at:control+shift+mousedown0"]: "erase-voxels", +}); + +const CONTROLS_FOR_TOOL = new Map([ + [BRUSH_TOOL_ID, ["vox-brush-size", "vox-brush-shape"]], + [FLOODFILL_TOOL_ID, ["vox-flood-max-voxels"]], +]); + +function getFloodFillCursor(erase: boolean) { + const lightColor = erase ? "#FF8888" : "#FFFFFF"; + const darkColor = erase ? "#610000" : "#000000"; + + const floodFillSVG = svg_floodFillCursor + .replaceAll("LIGHT_COLOR", lightColor) + .replaceAll("DARK_COLOR", darkColor) + .replace(/\s\s+/g, " "); + + return `url('data:image/svg+xml;utf8,${encodeURIComponent(floodFillSVG)}') 24 24, crosshair`; +} + +abstract class BaseVoxelTool extends LayerTool { + protected latestMouseState: MouseSelectionState | null = null; + private lastNormal: vec3 | undefined = undefined; + protected cursorEraseMode = new TrackableBoolean(false); + + protected getPoint(mouseState: MouseSelectionState): Int32Array | undefined { + const editContext = getEditingContext(this.layer); + if (editContext === undefined) return undefined; + const vox = editContext.getVoxelPositionFromMouse(mouseState) as + | Float32Array + | undefined; + if (!mouseState?.active || !vox) return undefined; + if (!mouseState.planeNormal) return; + this.lastNormal = editContext.transformGlobalToVoxelNormal( + mouseState.planeNormal, + ); + const CHUNK_POSITION_EPSILON = 1e-3; + const shiftedVox = new Float32Array(3); + for (let i = 0; i < 3; ++i) { + shiftedVox[i] = + vox[i] + CHUNK_POSITION_EPSILON * Math.abs(this.lastNormal[i]); + } + return new Int32Array([ + Math.floor(shiftedVox[0]), + Math.floor(shiftedVox[1]), + Math.floor(shiftedVox[2]), + ]); + } + + protected linePoints(a: Int32Array, b: Int32Array): Float32Array[] { + const dx = b[0] - a[0]; + const dy = b[1] - a[1]; + const dz = b[2] - a[2]; + const steps = Math.max(Math.abs(dx), Math.abs(dy), Math.abs(dz)); + const out: Float32Array[] = []; + if (steps <= 0) return out; + let lastX = a[0], + lastY = a[1], + lastZ = a[2]; + for (let s = 1; s <= steps; ++s) { + const x = Math.round(a[0] + (dx * s) / steps); + const y = Math.round(a[1] + (dy * s) / steps); + const z = Math.round(a[2] + (dz * s) / steps); + if (x !== lastX || y !== lastY || z !== lastZ) { + out.push(new Float32Array([x, y, z])); + lastX = x; + lastY = y; + lastZ = z; + } + } + return out; + } + + abstract bindToolInput(activation: ToolActivation): void; + + activate(activation: ToolActivation): boolean { + if (!this.layer.hasSubsourcesWithWritingEnabled.value) { + StatusMessage.showTemporaryMessage( + 'Voxel editing is not available. Please select a writable volume source in the "Source" tab.', + 5000, + ); + activation.cancel(); + return false; + } + this.showToolOptionsBar(activation); + this.bindToolInput(activation); + + const updateCursorState = (e: KeyboardEvent | MouseEvent) => { + this.cursorEraseMode.value = e.ctrlKey && e.shiftKey; + }; + activation.registerEventListener(window, "keydown", updateCursorState); + activation.registerEventListener(window, "keyup", updateCursorState); + activation.registerEventListener(window, "mousemove", updateCursorState); + + const paintCallback = + (erasing: boolean) => (event: ActionEvent) => { + event.stopPropagation(); + this.layer.setEraseState(erasing); + this.activationCallback(activation); + startRelativeMouseDrag( + event.detail as MouseEvent, + () => { + this.latestMouseState = this.mouseState; + }, + () => { + this.deactivationCallback(activation); + this.layer.setEraseState(false); + }, + ); + return true; + }; + + activation.bindAction("paint-voxels", paintCallback(false)); + activation.bindAction("erase-voxels", paintCallback(true)); + return true; + } + + private showToolOptionsBar(activation: ToolActivation) { + const toolId = this.toJSON(); + const controlTypes = CONTROLS_FOR_TOOL.get(toolId); + + const { header, body } = + makeToolActivationStatusMessageWithHeader(activation); + header.textContent = `${this.layer.managedLayer.name} - ${this.description}`; + header.classList.add("neuroglancer-tool-activation-status-header"); + body.classList.add("neuroglancer-voxel-tool-options-body"); + + if (!controlTypes) return; + + const visibility = new WatchableVisibilityPriority( + WatchableVisibilityPriority.VISIBLE, + ); + + for (const type of controlTypes) { + const def = VOXEL_LAYER_CONTROLS.find( + (c) => c.toolJson && c.toolJson.type === type, + ); + if (!def) continue; + + const controlContainer = document.createElement("label"); + controlContainer.classList.add("neuroglancer-layer-control-container"); + controlContainer.addEventListener("mousedown", (event) => { + event.stopPropagation(); + }); + + const labelContainer = document.createElement("div"); + labelContainer.classList.add( + "neuroglancer-layer-control-label-container", + ); + controlContainer.appendChild(labelContainer); + + const label = document.createElement("div"); + label.classList.add("neuroglancer-layer-control-label"); + if (def.title) { + label.title = def.title; + } + labelContainer.appendChild(label); + + const labelTextContainer = document.createElement("div"); + labelTextContainer.classList.add( + "neuroglancer-layer-control-label-text-container", + ); + labelTextContainer.textContent = def.label; + label.appendChild(labelTextContainer); + + const { controlElement } = def.makeControl(this.layer, activation, { + labelContainer, + labelTextContainer, + display: this.layer.manager.root.display, + visibility, + }); + controlElement.classList.add("neuroglancer-layer-control-control"); + controlContainer.appendChild(controlElement); + + if (def.toolJson) { + const widget = new ToolBindingWidget( + this.layer.toolBinder, + def.toolJson, + undefined, + ); + activation.registerDisposer(widget); + label.prepend(widget.element); + } + + body.appendChild(controlContainer); + } + } + + protected async performFloodFill(erasing: boolean): Promise { + const editContext = getEditingContext(this.layer); + if (editContext === undefined) { + StatusMessage.showTemporaryMessage( + 'Voxel editing is not available. Please select a writable volume source in the "Source" tab.', + 5000, + ); + return; + } + const seed = this.getPoint(this.mouseState); + // getBasis is only meaningful after a successful getPoint (which records + // the slice normal it reads). + const basis = seed ? this.getBasis() : undefined; + if (!seed || !basis) { + StatusMessage.showTemporaryMessage( + "Unable to retrieve mouse position. Please try again.", + 5000, + ); + return; + } + try { + const value = this.layer.getVoxelPaintValue(erasing); + const max = Number(this.layer.floodMaxVoxels.value); + if (!Number.isFinite(max) || max <= 0) { + throw new Error("Invalid max fill voxels setting"); + } + + const filterValue = + this.layer.lockToSelectedValue.value && erasing + ? this.layer.getVoxelPaintValue(false)(false) + : undefined; + + void editContext + .floodFillPlane2D( + new Float32Array(seed), + value, + Math.floor(max), + basis, + filterValue, + this.layer.floodMorphological.value, + ) + .catch((e: any) => + StatusMessage.showTemporaryMessage(String(e?.message ?? e)), + ); + } catch (e: any) { + StatusMessage.showTemporaryMessage(String(e?.message ?? e)); + } + } + + abstract activationCallback(activation: ToolActivation): void; + abstract deactivationCallback(activation: ToolActivation): void; + + protected setCursor(cursor: string) { + for (const panel of this.layer.manager.root.display.panels) { + panel.element.style.setProperty("cursor", cursor, "important"); + } + } + + protected resetCursor() { + for (const panel of this.layer.manager.root.display.panels) { + panel.element.style.removeProperty("cursor"); + } + } + + protected getBasis() { + const n = this.lastNormal; + if (!n) { + console.error("getBasis: Unexpected behavior: lastNormal is undefined"); + return undefined; + } + return getBasisFromNormal(n); + } +} + +export class VoxelBrushTool extends BaseVoxelTool { + private isDrawing = false; + private lastPoint: Int32Array | undefined; + private mouseDisposer: (() => void) | undefined; + private animationFrameHandle: number | null = null; + private cursorResetTimer: number | null = null; + private accumulatedCenters: Float32Array[] = []; + private activeStroke: + | { + radius: number; + shape: BrushShape; + basis: { u: Float32Array; v: Float32Array }; + value: VoxelValueGetter; + filterValue: bigint | undefined; + seq: number; + // Snapshotted so a stroke abandoned without dispatch can always be + // rolled back, even if the layer's editing context is gone by then. + context: VoxelEditingContext; + } + | undefined = undefined; + + activate(activation: ToolActivation): boolean { + if (!super.activate(activation)) return false; + + activation.registerDisposer( + linkWatchableValue(this.cursorEraseMode, this.layer.cursorInEraseMode), + ); + + activation.registerDisposer(() => { + if (this.cursorResetTimer !== null) { + clearTimeout(this.cursorResetTimer); + this.cursorResetTimer = null; + } + this.layer.cursorInEraseMode.value = false; + this.resetCursor(); + this.layer.scheduleOverlayRedraw(); + }); + + activation.registerDisposer( + this.mouseState.changed.add(this.layer.scheduleOverlayRedraw), + ); + this.layer.scheduleOverlayRedraw(); + + activation.bindAction( + "flood-fill-shortcut", + (event: ActionEvent) => { + event.stopPropagation(); + this.triggerFloodFill(event.detail.shiftKey); + }, + ); + return true; + } + + private triggerFloodFill(erasing: boolean) { + const wasErasing = this.layer.shouldErase(); + this.layer.setEraseState(erasing); + + if (this.cursorResetTimer !== null) clearTimeout(this.cursorResetTimer); + this.setCursor(getFloodFillCursor(erasing)); + + this.performFloodFill(erasing).finally(() => { + if (this.cursorResetTimer !== null) clearTimeout(this.cursorResetTimer); + this.cursorResetTimer = window.setTimeout(() => { + this.layer.setEraseState(wasErasing); + this.resetCursor(); + this.cursorResetTimer = null; + }, 1000); + }); + } + + activationCallback(_activation: ToolActivation): void { + if (getEditingContext(this.layer) === undefined) { + StatusMessage.showTemporaryMessage( + 'Voxel editing is not available. Please select a writable volume source in the "Source" tab.', + 5000, + ); + this.stopDrawing(); + return; + } + this.startDrawing(this.mouseState); + } + + deactivationCallback(_activation: ToolActivation): void { + this.stopDrawing(); + } + + constructor(layer: UserLayerWithVoxelEditing) { + super(layer, /*toggle=*/ true); + } + + toJSON() { + return BRUSH_TOOL_ID; + } + + get description() { + return "Brush tool"; + } + + bindToolInput(activation: ToolActivation) { + activation.bindInputEventMap(BRUSH_INPUT_MAP); + } + + private drawLoop = (): void => { + if (!this.isDrawing) { + this.animationFrameHandle = null; + return; + } + if (this.latestMouseState === null) { + this.animationFrameHandle = requestAnimationFrame(this.drawLoop); + return; + } + const cur = this.getPoint(this.latestMouseState); + this.latestMouseState = null; + if (cur) { + const last = this.lastPoint; + if ( + last && + (cur[0] !== last[0] || cur[1] !== last[1] || cur[2] !== last[2]) + ) { + const points = this.linePoints(last, cur); + if (points.length > 0) { + this.paintPoints(points); + } + } + this.lastPoint = cur; + } + this.animationFrameHandle = requestAnimationFrame(this.drawLoop); + }; + + private startDrawing(mouseState: MouseSelectionState) { + if (this.isDrawing) return; + const editContext = getEditingContext(this.layer); + if (editContext === undefined) { + throw new Error("editContext is undefined"); + } + // getPoint must run before the stroke snapshot below: it also records the + // slice normal that getBasis() reads. + const start = this.getPoint(mouseState); + if (!start) { + throw new Error( + "startDrawing: could not compute a starting voxel position from mouse", + ); + } + this.isDrawing = true; + this.accumulatedCenters = []; + this.activeStroke = { + radius: this.layer.brushRadius.value, + shape: this.layer.brushShape.value, + basis: this.getBasis()!, + value: this.layer.getVoxelPaintValue(this.layer.shouldErase()), + filterValue: + this.layer.lockToSelectedValue.value && this.layer.shouldErase() + ? this.layer.getVoxelPaintValue(false)(false) + : undefined, + seq: editContext.beginStroke(), + context: editContext, + }; + + this.paintPoints([new Float32Array([start[0], start[1], start[2]])]); + this.lastPoint = start; + this.latestMouseState = mouseState; + + this.mouseDisposer = mouseState.changed.add(() => { + this.latestMouseState = mouseState; + }); + + if (this.animationFrameHandle === null) { + this.animationFrameHandle = requestAnimationFrame(this.drawLoop); + } + } + + private stopDrawing() { + if (!this.isDrawing) return; + this.isDrawing = false; + this.lastPoint = undefined; + if (this.animationFrameHandle !== null) { + cancelAnimationFrame(this.animationFrameHandle); + this.animationFrameHandle = null; + } + if (this.mouseDisposer) { + this.mouseDisposer(); + this.mouseDisposer = undefined; + } + + const stroke = this.activeStroke!; + const centers = this.accumulatedCenters; + this.accumulatedCenters = []; + if (centers.length === 0) { + // Nothing to dispatch: whatever the previews tagged must be rolled + // back, or it would wait forever for a write that never comes. + stroke.context.rollbackStroke(stroke.seq); + return; + } + + const editContext = getEditingContext(this.layer); + if (editContext === undefined) { + stroke.context.rollbackStroke(stroke.seq); + return; + } + + void editContext.dispatchBrushStroke( + centers, + stroke.radius, + stroke.value, + stroke.shape, + stroke.basis, + stroke.seq, + stroke.filterValue, + ); + } + + private paintPoints(points: Float32Array[]) { + const editContext = getEditingContext(this.layer); + if (editContext === undefined) { + throw new Error("editContext is undefined"); + } + const stroke = this.activeStroke!; + + for (const p of points) { + this.accumulatedCenters.push(p); + } + + void editContext.applyBrushPreview( + points, + stroke.radius, + stroke.value, + stroke.shape, + stroke.basis, + stroke.seq, + stroke.filterValue, + ); + } +} + +export class VoxelFloodFillTool extends BaseVoxelTool { + activate(activation: ToolActivation) { + if (!super.activate(activation)) return false; + this.setCursor(getFloodFillCursor(this.cursorEraseMode.value)); + activation.registerDisposer( + this.cursorEraseMode.changed.add(() => { + this.setCursor(getFloodFillCursor(this.cursorEraseMode.value)); + }), + ); + activation.registerDisposer(() => { + this.resetCursor(); + }); + return true; + } + + activationCallback(_activation: ToolActivation): void { + this.performFloodFill(this.layer.shouldErase()); + } + + bindToolInput(activation: ToolActivation) { + activation.bindInputEventMap(FLOOD_INPUT_MAP); + } + + deactivationCallback(_activation: ToolActivation): void { + return; + } + + constructor(layer: UserLayerWithVoxelEditing) { + super(layer, /*toggle=*/ true); + } + + toJSON() { + return FLOODFILL_TOOL_ID; + } + + get description() { + return "Flood fill tool"; + } +} + +const pickerCursor = `url('data:image/svg+xml;utf8,${encodeURIComponent(svg_valuePicker)}') 24 24, crosshair`; + +export class AdoptVoxelValueTool extends LayerTool { + private lastPickPosition: Float32Array | undefined; + private lastCheckedSourceIndex = -1; + + readonly singleChannelAccess: ChunkChannelAccessParameters = { + numChannels: 1, + channelSpaceShape: new Uint32Array([]), + chunkChannelDimensionIndices: [], + chunkChannelCoordinates: new Uint32Array([0]), + }; + + constructor(layer: UserLayerWithVoxelEditing) { + super(layer, /*toggle=*/ false); + } + + protected setCursor(cursor: string) { + for (const panel of this.layer.manager.root.display.panels) { + panel.element.style.setProperty("cursor", cursor, "important"); + } + } + + protected resetCursor() { + for (const panel of this.layer.manager.root.display.panels) { + panel.element.style.removeProperty("cursor"); + } + } + + toJSON() { + return VALUE_PICKER_TOOL_ID; + } + + get description() { + return "Picking tool"; + } + + activate(activation: ToolActivation): void { + if (!this.layer.hasSubsourcesWithWritingEnabled.value) { + StatusMessage.showTemporaryMessage( + 'Voxel editing is not available. Please select a writable volume source in the "Source" tab.', + 5000, + ); + activation.cancel(); + return; + } + if (!this.mouseState?.active) return; + this.setCursor(pickerCursor); + activation.registerDisposer(() => { + this.resetCursor(); + }); + + const currentPosition = this.mouseState.position.slice() as vec3; + + if ( + this.lastPickPosition === undefined || + !vec3.equals(this.lastPickPosition as vec3, currentPosition) + ) { + this.lastPickPosition = currentPosition; + this.lastCheckedSourceIndex = -1; + } + + const allContexts = Array.from(this.layer.editingContexts.values()); + + if (allContexts.length === 0) { + StatusMessage.showTemporaryMessage( + "No volume sources found in this layer.", + 3000, + ); + return; + } + + const numSources = allContexts.length; + const startIndex = this.lastCheckedSourceIndex + 1; + + const checkNextSource = async () => { + for (let i = 0; i < numSources; ++i) { + const sourceIndex = (startIndex + i) % numSources; + const context = allContexts[sourceIndex]!; + + const voxelCoord = context.getVoxelPositionFromMouse(this.mouseState); + if (voxelCoord === undefined) continue; + + const source = context.primarySource.getSources( + this.layer.getIdentitySliceViewSourceOptions(), + )[0][0]!.chunkSource; + + const valueResult = source.getValueAt( + voxelCoord, + this.singleChannelAccess, + ); + const value = Array.isArray(valueResult) ? valueResult[0] : valueResult; + const bigValue = BigInt(value || 0); + + if (bigValue !== 0n) { + this.layer.setVoxelPaintValue(bigValue); + this.lastCheckedSourceIndex = sourceIndex; + StatusMessage.showTemporaryMessage( + `Adopted value: ${bigValue} (from source ${sourceIndex + 1}/${numSources})`, + 3000, + ); + return; + } + } + + this.lastCheckedSourceIndex = -1; + StatusMessage.showTemporaryMessage( + "No further segments found at this position.", + 3000, + ); + }; + + StatusMessage.forPromise(checkNextSource(), { + initialMessage: "Picking voxel value...", + delay: true, + errorPrefix: "Error picking value: ", + }); + } +} + +export function registerVoxelTools(LayerCtor: any) { + registerTool( + LayerCtor, + BRUSH_TOOL_ID, + (layer: UserLayerWithVoxelEditing) => new VoxelBrushTool(layer), + ); + registerTool( + LayerCtor, + FLOODFILL_TOOL_ID, + (layer: UserLayerWithVoxelEditing) => new VoxelFloodFillTool(layer), + ); + registerTool( + LayerCtor, + VALUE_PICKER_TOOL_ID, + (layer: UserLayerWithVoxelEditing) => new AdoptVoxelValueTool(layer), + ); +} diff --git a/src/util/gzip.ts b/src/util/gzip.ts index 1e74c856a9..2af3043897 100644 --- a/src/util/gzip.ts +++ b/src/util/gzip.ts @@ -69,3 +69,18 @@ export async function maybeDecompressGzip( } return byteView; } + +export async function encodeGzip( + data: Uint8Array | ArrayBuffer, + format: CompressionFormat, +): Promise { + const readableStream = new ReadableStream({ + start(controller) { + controller.enqueue(data); + controller.close(); + }, + }); + const compressionStream = new CompressionStream(format); + const compressedStream = readableStream.pipeThrough(compressionStream); + return await new Response(compressedStream).arrayBuffer(); +} diff --git a/src/voxel_annotation/backend.spec.ts b/src/voxel_annotation/backend.spec.ts new file mode 100644 index 0000000000..6813340ec8 --- /dev/null +++ b/src/voxel_annotation/backend.spec.ts @@ -0,0 +1,1904 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import type { VolumeChunk } from "#src/sliceview/volume/backend.js"; +import { VolumeChunkSource } from "#src/sliceview/volume/backend.js"; +import { DATA_TYPE_ARRAY_CONSTRUCTOR, DataType } from "#src/util/data_type.js"; +import { mat4 } from "#src/util/geom.js"; +import { HttpError } from "#src/util/http_request.js"; +import { VoxelEditController } from "#src/voxel_annotation/backend.js"; +import { + makeVoxChunkKey, + VOXEL_EMPTY_VALUE, + VOX_EDIT_FAILURE_RPC_ID, + VOX_EDIT_HISTORY_UPDATE_RPC_ID, + VoxelOperationType, + BrushShape, +} from "#src/voxel_annotation/base.js"; +import type { RPC } from "#src/worker_rpc.js"; + +const mockQueueManager = { + sources: new Set(), + adjustCapacitiesForChunk: vi.fn(), + updateChunkState: vi.fn(), + scheduleUpdate: vi.fn(), + moveChunkToFrontend: vi.fn(), + markRecentlyUsed: vi.fn(), + invalidateCachedChunks: vi.fn(), + gl: {}, +}; + +const mockChunkManager = { + queueManager: mockQueueManager, + chunkQueueManager: mockQueueManager, + rpc: null, + memoize: { get: (_k: string, fn: Function) => fn() }, +}; + +const mockRpc = { + get: vi.fn(), + invoke: vi.fn(), + newId: () => 0, + register: vi.fn(), + set: vi.fn(), + delete: vi.fn(), +} as unknown as RPC; + +const MOCK_SPEC = { + rank: 3, + chunkDataSize: new Uint32Array([2, 2, 2]), + dataType: 0, + lowerVoxelBound: new Float32Array([0, 0, 0]), + upperVoxelBound: new Float32Array([100, 100, 100]), + baseVoxelOffset: new Float32Array([0, 0, 0]), + fillValue: 0, +}; + +class MockBackendSource extends VolumeChunkSource { + public serverStorage = new Map(); + + async download(chunk: VolumeChunk) { + const key = chunk.chunkGridPosition.join(","); + if (this.serverStorage.has(key)) { + const buffer = this.serverStorage.get(key)!; + const Ctor = DATA_TYPE_ARRAY_CONSTRUCTOR[this.spec.dataType]; + chunk.data = new Ctor(buffer.slice(0)); + } + } + + async writeChunk(chunk: VolumeChunk) { + const key = chunk.chunkGridPosition.join(","); + if (chunk.data) { + this.serverStorage.set(key, chunk.data.buffer.slice(0) as ArrayBuffer); + } + } +} + +const createMockSource = (specOverride: any = {}) => { + const source = new MockBackendSource(mockRpc, { + spec: { ...MOCK_SPEC, ...specOverride }, + chunkManager: 0, + }); + vi.spyOn(source, "getChunk"); + vi.spyOn(source, "applyEdits"); + vi.spyOn(source, "download"); + vi.spyOn(source, "writeChunk"); + return source; +}; + +const resConfig = ( + lod: number, + scale: [number, number, number], + chunkSize: [number, number, number], + translation: [number, number, number] = [0, 0, 0], +) => { + const transform = new Float32Array(16); + mat4.identity(transform as unknown as mat4); + mat4.translate( + transform as unknown as mat4, + transform as unknown as mat4, + translation as any, + ); + mat4.scale( + transform as unknown as mat4, + transform as unknown as mat4, + scale as any, + ); + return { + lodIndex: lod, + transform: Array.from(transform), + chunkSize, + sourceRpc: 100 + lod, + }; +}; + +type Grid3D = (number | bigint)[][][]; // Z -> Y -> X + +function flattenGrid(grid: Grid3D, Ctor: any = Uint32Array) { + const d = grid.length; + const h = grid[0].length; + const w = grid[0][0].length; + const data = new Ctor(w * h * d); + + for (let z = 0; z < d; z++) { + for (let y = 0; y < h; y++) { + for (let x = 0; x < w; x++) { + data[x + y * w + z * w * h] = grid[z][y][x]; + } + } + } + return { data, size: [w, h, d] as [number, number, number] }; +} + +describe("VoxelEditController: _calculateParentUpdate", () => { + let controller: VoxelEditController; + let runDownsample: Function; + + beforeEach(() => { + vi.resetAllMocks(); + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + const source = createMockSource(); + source.rpcId = id; + return source; + }); + }); + + const runScenario = ( + scale: [number, number, number], + parentChunkSize: [number, number, number], + inputGrid: Grid3D, + expectedUpdates: { + x: number; + y: number; + z: number; + val: number | bigint; + }[], + childChunkOffset: { x: number; y: number; z: number } = { + x: 0, + y: 0, + z: 0, + }, + dataCtor: any = Uint32Array, + translation: [number, number, number] = [0, 0, 0], + ) => { + const { data: childData, size: childSize } = flattenGrid( + inputGrid, + dataCtor, + ); + + const childRes = resConfig(0, [1, 1, 1], childSize); + const parentRes = resConfig(1, scale, parentChunkSize, translation); + controller = new VoxelEditController(mockRpc, { + resolutions: [childRes, parentRes], + }); + runDownsample = (controller as any)._calculateParentUpdate.bind(controller); + + const result = runDownsample( + childData, + (controller as any).resolutions.get(0), + (controller as any).resolutions.get(1), + childChunkOffset, + childSize, + ); + + const [pw, ph, pd] = parentChunkSize; + const actualUpdatesMap = new Map(); + + for (let i = 0; i < result.indices.length; i++) { + const idx = result.indices[i]; + const val = result.values[i]; + const maxIdx = pw * ph * pd; + expect(idx).toBeLessThan(maxIdx); + expect(idx).toBeGreaterThanOrEqual(0); + + const pz = Math.floor(idx / (pw * ph)); + const rem = idx % (pw * ph); + const py = Math.floor(rem / pw); + const px = rem % pw; + actualUpdatesMap.set(`${px},${py},${pz}`, BigInt(val)); + } + + for (const { x, y, z, val } of expectedUpdates) { + const key = `${x},${y},${z}`; + const actual = actualUpdatesMap.get(key); + expect(actual, `Missing update at Parent(${x},${y},${z})`).toBeDefined(); + expect(actual, `Incorrect value at Parent(${x},${y},${z})`).toBe( + BigInt(val), + ); + actualUpdatesMap.delete(key); + } + + if (actualUpdatesMap.size > 0) { + const extras = Array.from(actualUpdatesMap.entries()) + .map(([k, v]) => `(${k}): ${v}`) + .join(", "); + throw new Error(`Unexpected updates at: ${extras}`); + } + }; + + it("Standard 2x2x2 Downsampling", () => { + runScenario( + [2, 2, 2], + [2, 2, 2], + [ + [ + [1, 1, 2, 2], + [1, 1, 2, 3], + [4, 4, 0, 0], + [4, 4, 0, 0], + ], + [ + [1, 1, 2, 2], + [1, 1, 3, 3], + [4, 4, 0, 0], + [4, 4, 0, 0], + ], + [ + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + ], + [ + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + ], + ], + [ + { x: 0, y: 0, z: 0, val: 1 }, + { x: 1, y: 0, z: 0, val: 2 }, + { x: 0, y: 1, z: 0, val: 4 }, + { x: 1, y: 1, z: 0, val: 0 }, + { x: 0, y: 0, z: 1, val: 0 }, + { x: 1, y: 0, z: 1, val: 0 }, + { x: 0, y: 1, z: 1, val: 0 }, + { x: 1, y: 1, z: 1, val: 0 }, + ], + ); + }); + + it("Anisotropic 1x2x1", () => { + runScenario( + [1, 2, 1], + [2, 2, 2], + [ + [ + [5, 5], + [5, 5], + [6, 6], + [7, 7], + ], + [ + [0, 0], + [0, 0], + [0, 0], + [0, 0], + ], + ], + [ + { x: 0, y: 0, z: 0, val: 5 }, + { x: 1, y: 0, z: 0, val: 5 }, + { x: 0, y: 1, z: 0, val: 6 }, + { x: 1, y: 1, z: 0, val: 6 }, + { x: 0, y: 0, z: 1, val: 0 }, + { x: 1, y: 0, z: 1, val: 0 }, + { x: 0, y: 1, z: 1, val: 0 }, + { x: 1, y: 1, z: 1, val: 0 }, + ], + ); + }); + + it("Odd Factors: 3x2x5", () => { + runScenario( + [3, 2, 5], + [3, 3, 1], + [ + [ + [1, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + ], + [ + [0, 0, 0, 2], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + ], + [ + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 3, 0, 0], + [0, 0, 0, 0], + ], + [ + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0], + [0, 0, 4, 0], + ], + ], + [ + { x: 0, y: 0, z: 0, val: 1 }, + { x: 1, y: 0, z: 0, val: 2 }, + { x: 0, y: 1, z: 0, val: 3 }, + { x: 1, y: 1, z: 0, val: 0 }, + ], + ); + }); + + it("Erasure: Single non-zero pixel cleared", () => { + runScenario( + [2, 2, 1], + [1, 1, 2], + [ + [ + [0, 0], + [0, 0], + ], + [ + [0, 0], + [0, 0], + ], + ], + [ + { x: 0, y: 0, z: 0, val: 0 }, + { x: 0, y: 0, z: 1, val: 0 }, + ], + ); + }); + + it("Offset Parent: Update at z=6", () => { + runScenario( + [2, 2, 2], + [4, 4, 8], + [ + [ + [99, 99], + [99, 99], + ], + [ + [99, 99], + [99, 99], + ], + ], + [{ x: 2, y: 2, z: 6, val: 99 }], + { x: 2, y: 2, z: 6 }, + ); + }); + + it("BigUint64Array: Supports large integers > 2^53", () => { + const bigVal = BigInt(Number.MAX_SAFE_INTEGER) + 50n; + runScenario( + [2, 2, 2], + [2, 2, 2], + [ + [ + [bigVal, bigVal], + [bigVal, bigVal], + ], + ], + [{ x: 0, y: 0, z: 0, val: bigVal }], + { x: 0, y: 0, z: 0 }, + BigUint64Array, + ); + }); + + it("Uint8Array: Supports lower precision types", () => { + runScenario( + [2, 2, 2], + [2, 2, 2], + [ + [ + [255, 255], + [255, 255], + ], + ], + [{ x: 0, y: 0, z: 0, val: 255 }], + { x: 0, y: 0, z: 0 }, + Uint8Array, + ); + }); + + it("Tie Breaking: Lowest value wins when counts are equal", () => { + runScenario( + [2, 2, 1], + [2, 2, 1], + [ + [ + [10, 10], + [5, 5], + ], + ], + [{ x: 0, y: 0, z: 0, val: 5 }], + ); + }); + + it("Non-Zero Dominance: 0 only wins if all values are 0", () => { + runScenario( + [2, 2, 1], + [2, 2, 1], + [ + [ + [0, 0], + [0, 9], + ], + ], + [{ x: 0, y: 0, z: 0, val: 9 }], + ); + }); + + it("Matrix Translation: Handles misaligned grids", () => { + runScenario( + [1, 1, 1], + [4, 1, 1], + [[[42]]], + [{ x: 2, y: 0, z: 0, val: 42 }], + { x: 0, y: 0, z: 0 }, + Uint32Array, + [-2, 0, 0], + ); + }); + + it("Fractional Scale: Aggregates across fractional boundaries (2.5x)", () => { + runScenario( + [2.5, 1, 1], + [2, 1, 1], + [[[1, 1, 2, 2, 2]]], + [ + { x: 0, y: 0, z: 0, val: 1 }, + { x: 1, y: 0, z: 0, val: 2 }, + ], + ); + }); +}); + +describe("VoxelEditController: _getParentChunkInfo", () => { + let controller: VoxelEditController; + + const setupController = (resConfigs: any[]) => { + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + const source = createMockSource(); + source.rpcId = id; + return source; + }); + controller = new VoxelEditController(mockRpc, { resolutions: resConfigs }); + return controller; + }; + + it("Standard Alignment: 2x scaling", () => { + const childRes = resConfig(0, [1, 1, 1], [4, 4, 4]); + const parentRes = resConfig(1, [2, 2, 2], [4, 4, 4]); + + setupController([childRes, parentRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + let res = getInfo(makeVoxChunkKey("0,0,0", 0), childRes); + expect(res.chunkKey).toBe("0,0,0"); + expect(res.parentKey).toBe(makeVoxChunkKey("0,0,0", 1)); + + res = getInfo(makeVoxChunkKey("1,0,0", 0), childRes); + expect(res.chunkKey).toBe("0,0,0"); + + res = getInfo(makeVoxChunkKey("2,0,0", 0), childRes); + expect(res.chunkKey).toBe("1,0,0"); + }); + + it("Matrix Translation: Parent Origin Shift", () => { + const childRes = resConfig(0, [1, 1, 1], [4, 4, 4]); + const parentRes = resConfig(1, [1, 1, 1], [4, 4, 4], [-4, 0, 0]); + + setupController([childRes, parentRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + const res = getInfo(makeVoxChunkKey("0,0,0", 0), childRes); + expect(res.chunkKey).toBe("1,0,0"); + }); + + it("Negative Coordinates", () => { + const childRes = resConfig(0, [1, 1, 1], [4, 4, 4]); + const parentRes = resConfig(1, [1, 1, 1], [4, 4, 4]); + + setupController([childRes, parentRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + const res = getInfo(makeVoxChunkKey("-1,-1,-1", 0), childRes); + expect(res.chunkKey).toBe("-1,-1,-1"); + }); + + it("Max LOD Boundary", () => { + const childRes = resConfig(0, [1, 1, 1], [4, 4, 4]); + setupController([childRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + const res = getInfo(makeVoxChunkKey("0,0,0", 0), childRes); + expect(res).toBeNull(); + }); + + it("Odd Integer Scale (3x)", () => { + const childRes = resConfig(0, [1, 1, 1], [2, 2, 2]); + const parentRes = resConfig(1, [3, 3, 3], [2, 2, 2]); + + setupController([childRes, parentRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + const res = getInfo(makeVoxChunkKey("3,0,0", 0), childRes); + expect(res.chunkKey).toBe("1,0,0"); + }); + + it("Fractional Scale (2.5x)", () => { + const childRes = resConfig(0, [1, 1, 1], [10, 10, 10]); + const parentRes = resConfig(1, [2.5, 2.5, 2.5], [10, 10, 10]); + + setupController([childRes, parentRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + let res = getInfo(makeVoxChunkKey("2,0,0", 0), childRes); + expect(res.chunkKey).toBe("0,0,0"); + + res = getInfo(makeVoxChunkKey("3,0,0", 0), childRes); + expect(res.chunkKey).toBe("1,0,0"); + }); + + it("Anisotropic Scale (1x, 2x, 5x)", () => { + const childRes = resConfig(0, [1, 1, 1], [10, 10, 10]); + const parentRes = resConfig(1, [1, 2, 5], [10, 10, 10]); + + setupController([childRes, parentRes]); + const getInfo = (controller as any)._getParentChunkInfo.bind(controller); + + const res = getInfo(makeVoxChunkKey("1,1,1", 0), childRes); + expect(res.chunkKey).toBe("1,0,0"); + }); +}); + +describe("VoxelEditController: Downsampling Integration", () => { + let controller: VoxelEditController; + let childSource: any; + let parentSource: any; + let grandParentSource: any; + + const setupIntegration = (numLevels: number = 2) => { + // Chunk data is injected through the mock server storage (the download + // path) rather than by stubbing `getChunk`: edit-path reads now download + // into isolated chunks and never consult `getChunk`. + childSource = createMockSource(); + childSource.serverStorage.set( + "0,0,0", + new Uint8Array(8).fill(1).buffer, // dataType 0 = UINT8 + ); + + parentSource = createMockSource(); + // Parent/grandparent are absent from server storage: applyEdits creates + // them as empty (zero-filled) chunks, matching the old zero-filled stubs. + + grandParentSource = createMockSource(); + + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + if (id === 100) return childSource; + if (id === 101) return parentSource; + if (id === 102) return grandParentSource; + if (id === 999) return { value: 0 }; + return null; + }); + + const resolutions = [ + resConfig(0, [1, 1, 1], [2, 2, 2]), // Child + resConfig(1, [2, 2, 2], [2, 2, 2]), // Parent (2x scale) + ]; + + if (numLevels > 2) { + resolutions.push(resConfig(2, [4, 4, 4], [2, 2, 2])); // Grandparent (4x scale) + } + + controller = new VoxelEditController(mockRpc, { + resolutions, + pendingOpCount: 999, + }); + + vi.spyOn(controller as any, "callChunkReload"); + }; + + it("Single Step Flow: Writes to parent and notifies frontend", async () => { + setupIntegration(2); + const key = makeVoxChunkKey("0,0,0", 0); + + (controller as any).enqueueDownsample(key); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(childSource.download).toHaveBeenCalled(); + + expect(parentSource.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.any(Array), + expect.arrayContaining([1n]), + ); + + // The real parent is reloaded lazily; when it reaches the GPU it clears the + // originating LOD-0 overlay (swap-on-arrival), passed as a { real -> overlay } + // map along with the origin's covered dispatch seq (0 here: no seq-tagged + // edit was flushed for the origin). The old eager preview-clear call no + // longer exists. + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [makeVoxChunkKey("0,0,0", 1)], + false, // isForPreviewChunks + { [makeVoxChunkKey("0,0,0", 1)]: makeVoxChunkKey("0,0,0", 0) }, + { [makeVoxChunkKey("0,0,0", 1)]: 0 }, + ); + }); + + it("Coverage echo: flush and cascade reloads carry the max flushed dispatch seq", async () => { + setupIntegration(2); + const key = makeVoxChunkKey("0,0,0", 0); + + // Two dispatched strokes touch the chunk before the flush runs: the write + // covers both, so the echoed coverage must be the max seq (7). + controller.commitVoxels([ + { key, indices: [0], value: 1n, seq: 3 }, + { key, indices: [1], value: 1n, seq: 7 }, + ]); + await (controller as any).flushPending(); + + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [key], + false, + undefined, + { [key]: 7 }, + ); + + // The cascade reload claims the origin's flushed coverage for the parent. + await new Promise((resolve) => setTimeout(resolve, 0)); + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [makeVoxChunkKey("0,0,0", 1)], + false, + { [makeVoxChunkKey("0,0,0", 1)]: key }, + { [makeVoxChunkKey("0,0,0", 1)]: 7 }, + ); + }); + + it("Coverage pruning: lastFlushedSeq entry dropped once the cascade completes", async () => { + setupIntegration(2); + const key = makeVoxChunkKey("0,0,0", 0); + + controller.commitVoxels([{ key, indices: [0], value: 1n, seq: 3 }]); + await (controller as any).flushPending(); + // Queued for cascade: still readable for the chain-start capture. + expect((controller as any).lastFlushedSeq.has(key)).toBe(true); + + await new Promise((resolve) => setTimeout(resolve, 0)); + expect((controller as any).lastFlushedSeq.has(key)).toBe(false); + }); + + it("Coverage pruning: a later flush still covers older overlay tags", async () => { + setupIntegration(2); + const key = makeVoxChunkKey("0,0,0", 0); + + controller.commitVoxels([{ key, indices: [0], value: 1n, seq: 3 }]); + await (controller as any).flushPending(); + await new Promise((resolve) => setTimeout(resolve, 0)); // cascade done, pruned + + // Dispatch seqs are globally monotonic: the recomputed coverage (7) + // exceeds anything an old tag could require. + controller.commitVoxels([{ key, indices: [1], value: 1n, seq: 7 }]); + await (controller as any).flushPending(); + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [key], + false, + undefined, + { [key]: 7 }, + ); + }); + + it("Coverage echo: a chain claims its start-of-chain coverage even if a flush lands mid-chain", async () => { + setupIntegration(3); + const originKey = makeVoxChunkKey("0,0,0", 0); + + (controller as any).lastFlushedSeq.set(originKey, 1); + + // Mid-chain (during the L0->L1 write), a newer flush of the same origin + // completes (seq 5). The data propagated by the running chain derives from + // a read made before that write, so the L1->L2 reload must keep claiming + // seq 1 — claiming 5 would clear the overlay over a parent lacking those + // edits. + parentSource.applyEdits.mockImplementation(async () => { + parentSource.serverStorage.set("0,0,0", new Uint8Array(8).fill(1).buffer); + (controller as any).lastFlushedSeq.set(originKey, 5); + }); + + (controller as any).enqueueDownsample(originKey); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(grandParentSource.applyEdits).toHaveBeenCalled(); + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [makeVoxChunkKey("0,0,0", 2)], + false, + { [makeVoxChunkKey("0,0,0", 2)]: originKey }, + { [makeVoxChunkKey("0,0,0", 2)]: 1 }, + ); + }); + + it("Recursive Propagation: L0 -> L1 -> L2", async () => { + setupIntegration(3); + + // After the L0->L1 write, the L1 chunk must read back as 1s so the L1->L2 + // step propagates. Simulate by publishing it to the mock server storage. + parentSource.applyEdits.mockImplementation(async () => { + parentSource.serverStorage.set("0,0,0", new Uint8Array(8).fill(1).buffer); + }); + + const key = makeVoxChunkKey("0,0,0", 0); + (controller as any).enqueueDownsample(key); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(parentSource.applyEdits).toHaveBeenCalled(); + + expect(grandParentSource.applyEdits).toHaveBeenCalled(); + + const reloadCalls = (controller as any).callChunkReload.mock.calls; + const keysReloaded = reloadCalls.flatMap((c: any) => c[0]); + + expect(keysReloaded).toContain(makeVoxChunkKey("0,0,0", 1)); + expect(keysReloaded).toContain(makeVoxChunkKey("0,0,0", 2)); + }); + + it("Queue Deduplication: Processes same key once per batch", async () => { + setupIntegration(2); + const key = makeVoxChunkKey("0,0,0", 0); + + (controller as any).enqueueDownsample(key); + (controller as any).enqueueDownsample(key); + (controller as any).enqueueDownsample(key); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(parentSource.applyEdits).toHaveBeenCalledTimes(1); + }); + + it("Lazy Loading: Downloads child chunk if missing", async () => { + setupIntegration(2); + + // Chunk absent from server storage; the isolated-chunk download fills it. + childSource.serverStorage.delete("0,0,0"); + childSource.download.mockImplementation(async (chunk: any) => { + chunk.data = new Uint32Array(8).fill(1); + }); + + const key = makeVoxChunkKey("0,0,0", 0); + (controller as any).enqueueDownsample(key); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(childSource.download).toHaveBeenCalled(); + expect(parentSource.applyEdits).toHaveBeenCalled(); + }); + + it("Error Handling: Child download failure aborts chain gracefully", async () => { + setupIntegration(2); + const consoleSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + + childSource.download.mockRejectedValue(new Error("Network Error")); + + const key = makeVoxChunkKey("0,0,0", 0); + (controller as any).enqueueDownsample(key); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(parentSource.applyEdits).not.toHaveBeenCalled(); + + consoleSpy.mockRestore(); + }); + + it("Error Handling: Parent write failure reports error and stops recursion", async () => { + setupIntegration(3); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + parentSource.applyEdits.mockRejectedValue(new Error("Write Failed")); + + const key = makeVoxChunkKey("0,0,0", 0); + (controller as any).enqueueDownsample(key); + + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(parentSource.applyEdits).toHaveBeenCalled(); + + expect(grandParentSource.applyEdits).not.toHaveBeenCalled(); + + expect(mockRpc.invoke).toHaveBeenCalledWith( + "vox.edit.failure", + expect.objectContaining({ + voxChunkKeys: [makeVoxChunkKey("0,0,0", 1)], + }), + ); + + consoleSpy.mockRestore(); + }); +}); + +describe("VoxelEditController: flushPending", () => { + let controller: VoxelEditController; + let mockSource0: any; + let mockSource1: any; + + beforeEach(() => { + vi.useFakeTimers(); + vi.clearAllMocks(); + + mockSource0 = createMockSource(); + vi.spyOn(mockSource0, "applyEdits").mockResolvedValue({ + indices: new Uint32Array([]), + oldValues: new BigUint64Array([]), + newValues: new BigUint64Array([]), + }); + + mockSource1 = createMockSource(); + vi.spyOn(mockSource1, "applyEdits").mockResolvedValue({}); + + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + if (id === 100) return mockSource0; + if (id === 101) return mockSource1; + if (id === 999) return { value: 0 }; + return null; + }); + + controller = new VoxelEditController(mockRpc, { + resolutions: [ + resConfig(0, [1, 1, 1], [2, 2, 2]), + resConfig(1, [2, 2, 2], [2, 2, 2]), + ], + pendingOpCount: 999, + }); + + vi.spyOn(controller as any, "enqueueDownsample").mockImplementation( + () => {}, + ); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("Batching: Aggregates multiple edits to the same chunk into one write", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + + controller.commitVoxels([ + { key, indices: [1], value: 50n }, + { key, indices: [2], value: 60n }, + ]); + + const otherKey = makeVoxChunkKey("1,0,0", 0); + controller.commitVoxels([{ key: otherKey, indices: [5], value: 99n }]); + + controller.commitVoxels([ + { key, indices: [1], value: 42n }, + { key, indices: [3], value: 70n }, + ]); + + await vi.runAllTimersAsync(); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "0,0,0", + [1, 2, 3], + [42n, 60n, 70n], + ); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith("1,0,0", [5], [99n]); + }); + + it("History: Updates stacks and notifies frontend correctly", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + + (controller as any).redoStack.push({ + changes: new Map(), + timestamp: 0, + description: "dummy", + }); + expect((controller as any).redoStack.length).toBe(1); + + controller.commitVoxels([{ key, indices: [1], value: 50n }]); + await vi.runAllTimersAsync(); + + expect((controller as any).undoStack.length).toBe(1); + + expect((controller as any).redoStack.length).toBe(0); + + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_HISTORY_UPDATE_RPC_ID, + expect.objectContaining({ + undoCount: 1, + redoCount: 0, + }), + ); + }); + + it("Partial Failure: Succeeds for valid chunks even if one chunk fails", async () => { + const validKey = makeVoxChunkKey("0,0,0", 0); + const failKey = makeVoxChunkKey("1,0,0", 0); + + mockSource0.applyEdits.mockImplementation((chunkKey: string) => { + if (chunkKey === "1,0,0") { + return Promise.reject(new Error("Network Error")); + } + return Promise.resolve({ + indices: new Uint32Array([1]), + oldValues: new BigUint64Array([VOXEL_EMPTY_VALUE]), + newValues: new BigUint64Array([50n]), + }); + }); + + controller.commitVoxels([ + { key: validKey, indices: [1], value: 50n }, + { key: failKey, indices: [1], value: 50n }, + ]); + + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + await vi.runAllTimersAsync(); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.anything(), + expect.anything(), + ); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "1,0,0", + expect.anything(), + expect.anything(), + ); + + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_FAILURE_RPC_ID, + expect.objectContaining({ + voxChunkKeys: [failKey], + }), + ); + + const undoStack = (controller as any).undoStack; + expect(undoStack.length).toBe(1); + expect(undoStack[0].changes.has(validKey)).toBe(true); + expect(undoStack[0].changes.has(failKey)).toBe(false); + + errorSpy.mockRestore(); + }); + + it("Invalid Data: Handles malformed keys gracefully without crashing", async () => { + const validKey = makeVoxChunkKey("0,0,0", 0); + const badKey = "invalid_format_key"; + + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + controller.commitVoxels([ + { key: validKey, indices: [1], value: 50n }, + { key: badKey, indices: [1], value: 50n }, + ]); + + await vi.runAllTimersAsync(); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.anything(), + expect.anything(), + ); + + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_FAILURE_RPC_ID, + expect.objectContaining({ + voxChunkKeys: [badKey], + }), + ); + + errorSpy.mockRestore(); + }); + + it("Downsample Trigger: Enqueues modified keys for processing", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const enqueueSpy = vi.spyOn(controller as any, "enqueueDownsample"); + + controller.commitVoxels([{ key, indices: [1], value: 50n }]); + await vi.runAllTimersAsync(); + + expect(enqueueSpy).toHaveBeenCalledWith(key); + }); +}); + +describe("VoxelEditController: Undo/Redo", () => { + let controller: VoxelEditController; + let mockSource0: any; + + beforeEach(() => { + vi.clearAllMocks(); + + mockSource0 = createMockSource(); + vi.spyOn(mockSource0, "applyEdits").mockResolvedValue({ + indices: new Uint32Array([]), + oldValues: new BigUint64Array([]), + newValues: new BigUint64Array([]), + }); + + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + if (id === 100) return mockSource0; + if (id === 999) return { value: 0 }; + return null; + }); + + controller = new VoxelEditController(mockRpc, { + resolutions: [resConfig(0, [1, 1, 1], [2, 2, 2])], + pendingOpCount: 999, + }); + + vi.spyOn(controller as any, "callChunkReload"); + vi.spyOn(controller as any, "enqueueDownsample").mockImplementation( + () => {}, + ); + (mockRpc.invoke as any).mockClear(); + }); + + it("A flush triggered during an undo waits for the undo to finish", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const order: string[] = []; + let releaseUndo!: () => void; + const undoGate = new Promise((resolve) => (releaseUndo = resolve)); + mockSource0.applyEdits.mockImplementation(async () => { + if (order.length === 0) { + order.push("undo-start"); + await undoGate; + order.push("undo-end"); + } else { + order.push("flush"); + } + return { + indices: new Uint32Array([]), + oldValues: new BigUint64Array([]), + newValues: new BigUint64Array([]), + }; + }); + (controller as any).undoStack.push({ + changes: new Map([ + [ + key, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([0n]), + newValues: new BigUint64Array([10n]), + }, + ], + ]), + timestamp: Date.now(), + description: "stroke", + }); + + const undoDone = controller.undo(); + await new Promise((resolve) => setTimeout(resolve, 0)); + // A stroke committed mid-undo: its flush must queue behind the undo. + controller.commitVoxels([{ key, indices: [1], value: 42n, seq: 2 }]); + const flushDone = (controller as any).runExclusive(() => + (controller as any).flushPendingLocked(), + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(order).toEqual(["undo-start"]); + + releaseUndo(); + await undoDone; + await flushDone; + expect(order).toEqual(["undo-start", "undo-end", "flush"]); + }); + + it("Concurrent undos are serialized and pop in order", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const makeAction = (label: string) => ({ + changes: new Map([ + [ + key, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([0n]), + newValues: new BigUint64Array([10n]), + }, + ], + ]), + timestamp: Date.now(), + description: label, + }); + const action1 = makeAction("stroke 1"); + const action2 = makeAction("stroke 2"); + (controller as any).undoStack.push(action1, action2); + + let releaseFirst!: () => void; + const firstGate = new Promise((resolve) => (releaseFirst = resolve)); + let calls = 0; + mockSource0.applyEdits.mockImplementation(async () => { + if (++calls === 1) await firstGate; + return { + indices: new Uint32Array([]), + oldValues: new BigUint64Array([]), + newValues: new BigUint64Array([]), + }; + }); + + const undo1 = controller.undo(); + const undo2 = controller.undo(); + await new Promise((resolve) => setTimeout(resolve, 0)); + // The second undo must not have started while the first is writing. + expect(calls).toBe(1); + + releaseFirst(); + await undo1; + await undo2; + expect((controller as any).undoStack.length).toBe(0); + expect((controller as any).redoStack).toEqual([action2, action1]); + }); + + it("Undo flushes pending edits first, so it targets the latest stroke", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const action1 = { + changes: new Map([ + [ + key, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([0n]), + newValues: new BigUint64Array([10n]), + }, + ], + ]), + timestamp: Date.now(), + description: "stroke 1", + }; + (controller as any).undoStack.push(action1); + + // A second stroke is still pending (debounce not fired) when undo runs: + // it must be flushed and become the undo target, not stroke 1. + controller.commitVoxels([{ key, indices: [1], value: 42n, seq: 2 }]); + await controller.undo(); + + expect((controller as any).undoStack).toEqual([action1]); + expect((controller as any).redoStack.length).toBe(1); + expect((controller as any).redoStack[0]).not.toBe(action1); + }); + + it("Successful Undo and Redo Lifecycle", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const editAction = { + changes: new Map([ + [ + key, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([10n]), + newValues: new BigUint64Array([20n]), + }, + ], + ]), + timestamp: Date.now(), + description: "Test Action", + }; + + (controller as any).undoStack.push(editAction); + + await controller.undo(); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.any(Uint32Array), + expect.any(BigUint64Array), + ); + const undoCallArgs = mockSource0.applyEdits.mock.calls[0]; + expect(undoCallArgs[2][0]).toBe(10n); + + // A single rollback reload: the frontend purges the overlay tags and the + // swap clears on first arrival. + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [key], + false, + undefined, + undefined, + true, + ); + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_HISTORY_UPDATE_RPC_ID, + expect.objectContaining({ undoCount: 0, redoCount: 1 }), + ); + + expect((controller as any).undoStack.length).toBe(0); + expect((controller as any).redoStack.length).toBe(1); + + mockSource0.applyEdits.mockClear(); + (mockRpc.invoke as any).mockClear(); + + await controller.redo(); + + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.any(Uint32Array), + expect.any(BigUint64Array), + ); + const redoCallArgs = mockSource0.applyEdits.mock.calls[0]; + expect(redoCallArgs[2][0]).toBe(20n); + + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + [key], + false, + undefined, + undefined, + true, + ); + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_HISTORY_UPDATE_RPC_ID, + expect.objectContaining({ undoCount: 1, redoCount: 0 }), + ); + + expect((controller as any).redoStack.length).toBe(0); + expect((controller as any).undoStack.length).toBe(1); + }); + + it("Empty Stack Behavior", async () => { + await expect(controller.undo()).rejects.toThrow(/Nothing to undo/); + await expect(controller.redo()).rejects.toThrow(/Nothing to redo/); + }); + + it("Undo Failure Handling", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const editAction = { + changes: new Map([ + [ + key, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([10n]), + newValues: new BigUint64Array([20n]), + }, + ], + ]), + timestamp: Date.now(), + description: "Test Action", + }; + (controller as any).undoStack.push(editAction); + + mockSource0.applyEdits.mockRejectedValue(new Error("Backend Write Failed")); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + await controller.undo(); + + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_FAILURE_RPC_ID, + expect.objectContaining({ + voxChunkKeys: [key], + message: "Undo failed.", + }), + ); + + expect((controller as any).undoStack.length).toBe(1); + expect((controller as any).redoStack.length).toBe(0); + + expect((controller as any).callChunkReload).not.toHaveBeenCalled(); + + consoleSpy.mockRestore(); + }); + + it("Redo Failure Handling", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + const editAction = { + changes: new Map([ + [ + key, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([10n]), + newValues: new BigUint64Array([20n]), + }, + ], + ]), + timestamp: Date.now(), + description: "Test Action", + }; + (controller as any).redoStack.push(editAction); + + mockSource0.applyEdits.mockRejectedValue(new Error("Backend Write Failed")); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + await controller.redo(); + + expect(mockRpc.invoke).toHaveBeenCalledWith( + VOX_EDIT_FAILURE_RPC_ID, + expect.objectContaining({ + voxChunkKeys: [key], + message: "Redo failed.", + }), + ); + + expect((controller as any).redoStack.length).toBe(1); + expect((controller as any).undoStack.length).toBe(0); + + consoleSpy.mockRestore(); + }); + + it("Multi-Chunk Action Consistency", async () => { + const key1 = makeVoxChunkKey("0,0,0", 0); + const key2 = makeVoxChunkKey("1,0,0", 0); + + const editAction = { + changes: new Map([ + [ + key1, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([1n]), + newValues: new BigUint64Array([2n]), + }, + ], + [ + key2, + { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([3n]), + newValues: new BigUint64Array([4n]), + }, + ], + ]), + timestamp: Date.now(), + description: "Multi Chunk Action", + }; + + (controller as any).undoStack.push(editAction); + + await controller.undo(); + + expect(mockSource0.applyEdits).toHaveBeenCalledTimes(2); + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.anything(), + expect.anything(), + ); + expect(mockSource0.applyEdits).toHaveBeenCalledWith( + "1,0,0", + expect.anything(), + expect.anything(), + ); + + expect((controller as any).callChunkReload).toHaveBeenCalledWith( + expect.arrayContaining([key1, key2]), + false, + undefined, + undefined, + true, + ); + + expect((controller as any).undoStack.length).toBe(0); + expect((controller as any).redoStack.length).toBe(1); + }); +}); + +describe("VoxelEditController: Tool Operations", () => { + let controller: VoxelEditController; + let mockSource: MockBackendSource; + + beforeEach(() => { + vi.useFakeTimers(); + vi.clearAllMocks(); + + const spec = { + rank: 3, + chunkDataSize: new Uint32Array([10, 10, 10]), + dataType: DataType.UINT64, + lowerVoxelBound: new Float32Array([0, 0, 0]), + upperVoxelBound: new Float32Array([100, 100, 100]), + baseVoxelOffset: new Float32Array([0, 0, 0]), + fillValue: VOXEL_EMPTY_VALUE, + }; + mockSource = createMockSource({ ...spec }); + vi.spyOn(mockSource, "applyEdits").mockResolvedValue({ + indices: new Uint32Array([]), + oldValues: new BigUint64Array([]), + newValues: new BigUint64Array([]), + }); + + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + if (id === 100) return mockSource; + if (id === 999) return { value: 0 }; + return null; + }); + + controller = new VoxelEditController(mockRpc, { + resolutions: [resConfig(0, [1, 1, 1], [10, 10, 10])], + pendingOpCount: 999, + }); + + vi.spyOn(controller as any, "enqueueDownsample").mockImplementation( + () => {}, + ); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("paintBrushWithShape: 3D Sphere", async () => { + const center = new Float32Array([5, 5, 5]); + const radius = 3; + const value = 5n; + + await controller.performOperation({ + type: VoxelOperationType.BRUSH, + centers: [center], + radius, + value, + shape: BrushShape.SPHERE, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + }); + + await vi.runAllTimersAsync(); + + expect(mockSource.applyEdits).toHaveBeenCalledWith( + "0,0,0", + expect.any(Array), + expect.any(Array), + ); + + const call = (mockSource.applyEdits as any).mock.calls[0]; + const indices = call[1]; + const values = call[2]; + const indexSet = new Set(indices); + + const getIdx = (x: number, y: number, z: number) => z * 100 + y * 10 + x; + expect(indexSet.has(getIdx(5, 5, 5))).toBe(true); + expect(indexSet.has(getIdx(7, 5, 5))).toBe(true); + expect(indexSet.has(getIdx(8, 5, 5))).toBe(false); + expect(values[0]).toBe(5n); + }); + + it("paintBrushWithShape: 2D Disk", async () => { + const center = new Float32Array([5, 5, 5]); + const radius = 3; + const value = 3n; + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await controller.performOperation({ + type: VoxelOperationType.BRUSH, + centers: [center], + radius, + value, + shape: BrushShape.DISK, + basis, + }); + + await vi.runAllTimersAsync(); + + const call = (mockSource.applyEdits as any).mock.calls[0]; + const indices = call[1]; + const values = call[2]; + const indexSet = new Set(indices); + const getIdx = (x: number, y: number, z: number) => z * 100 + y * 10 + x; + + for (const idx of indices) { + const z = Math.floor(idx / 100); + expect(z).toBe(5); + } + expect(indexSet.has(getIdx(5, 5, 5))).toBe(true); + expect(indexSet.has(getIdx(7, 5, 5))).toBe(true); + expect(values[0]).toBe(3n); + }); + + it("floodFillPlane2D: Bounded region (Bucket)", async () => { + const data = new BigUint64Array(1000); + for (let x = 3; x <= 7; x++) { + data[0 * 100 + 3 * 10 + x] = 1n; // y=3 + data[0 * 100 + 7 * 10 + x] = 1n; // y=7 + } + for (let y = 3; y <= 7; y++) { + data[0 * 100 + y * 10 + 3] = 1n; // x=3 + data[0 * 100 + y * 10 + 7] = 1n; // x=7 + } + mockSource.serverStorage.set("0,0,0", data.buffer); + + const seed = new Float32Array([5, 5, 0]); + const fillValue = 5n; + const maxVoxels = 100; + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed, + value: fillValue, + maxVoxels, + basis, + }); + + await vi.runAllTimersAsync(); + + const call = (mockSource.applyEdits as any).mock.calls[0]; + const indices = call[1]; + expect(indices.length).toBe(9); + }); + + it("floodFillPlane2D: Plane constraint", async () => { + const data = new BigUint64Array(1000); + const z = 5; + for (let x = 3; x <= 7; x++) { + data[z * 100 + 3 * 10 + x] = 1n; + data[z * 100 + 7 * 10 + x] = 1n; + } + for (let y = 3; y <= 7; y++) { + data[z * 100 + y * 10 + 3] = 1n; + data[z * 100 + y * 10 + 7] = 1n; + } + mockSource.serverStorage.set("0,0,0", data.buffer); + + const seed = new Float32Array([5, 5, 5]); + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed, + value: 2n, + maxVoxels: 100, + basis, + }); + + await vi.runAllTimersAsync(); + + const call = (mockSource.applyEdits as any).mock.calls[0]; + const indices = call[1]; + expect(indices.length).toBe(9); + for (const idx of indices) { + const cz = Math.floor(idx / 100); + expect(cz).toBe(5); + } + }); + + it("floodFillPlane2D: Max voxels exceeded", async () => { + const seed = new Float32Array([5, 5, 0]); + const maxVoxels = 5; + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await expect( + controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed, + value: 9n, + maxVoxels, + basis, + }), + ).rejects.toThrow("Flood fill failed: too many voxels filled."); + }); + + it("floodFillPlane2D: Seed value equals fill value", async () => { + const data = new BigUint64Array(1000); + const seedIdx = 0 * 100 + 5 * 10 + 5; + data[seedIdx] = 5n; + mockSource.serverStorage.set("0,0,0", data.buffer); + + const seed = new Float32Array([5, 5, 0]); + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed, + value: 5n, + maxVoxels: 100, + basis, + }); + + await vi.runAllTimersAsync(); + + expect(mockSource.applyEdits).not.toHaveBeenCalled(); + }); + + it("floodFillPlane2D: Leak prevention (morphological)", async () => { + (controller as any).morphologicalConfig = { + growthThresholds: [{ count: 5, size: 3 }], + maxSize: 9, + }; + + const data = new BigUint64Array(1000); + // Box 0..9 in X, Y, at Z=0. + for (let x = 0; x <= 9; x++) { + if (x !== 0) data[0 * 100 + 0 * 10 + x] = 1n; // y=0 + data[0 * 100 + 9 * 10 + x] = 1n; // y=9 + } + for (let y = 0; y <= 9; y++) { + if (y !== 5) data[0 * 100 + y * 10 + 0] = 1n; // x=0 (hole at y=5) + data[0 * 100 + y * 10 + 9] = 1n; // x=9 + } + mockSource.serverStorage.set("0,0,0", data.buffer); + + const seed = new Float32Array([5, 5, 0]); + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed, + value: 2n, + maxVoxels: 1000, + basis, + }); + + await vi.runAllTimersAsync(); + + expect(mockSource.applyEdits).toHaveBeenCalled(); + const indices = (mockSource.applyEdits as any).mock.calls[0][1]; + expect(indices.length).toBeLessThan(100); + expect(indices.length).toBeGreaterThan(50); + }); + + it("floodFillPlane2D: morphological=false fills through thin channels", async () => { + (controller as any).morphologicalConfig = { + growthThresholds: [{ count: 5, size: 3 }], + maxSize: 9, + }; + + const data = new BigUint64Array(1000); + // Closed box 0..9 at z=0, split by a wall column at x=5 with a 1-voxel + // hole at y=5. + for (let i = 0; i <= 9; i++) { + data[0 * 100 + 0 * 10 + i] = 1n; + data[0 * 100 + 9 * 10 + i] = 1n; + data[0 * 100 + i * 10 + 0] = 1n; + data[0 * 100 + i * 10 + 9] = 1n; + if (i !== 5) data[0 * 100 + i * 10 + 5] = 1n; + } + mockSource.serverStorage.set("0,0,0", data.buffer); + + await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed: new Float32Array([2, 5, 0]), + value: 2n, + maxVoxels: 1000, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + morphological: false, + }); + + await vi.runAllTimersAsync(); + + // Left chamber (4×8) + hole + right chamber (3×8). + const indices = (mockSource.applyEdits as any).mock.calls[0][1]; + expect(indices.length).toBe(57); + }); + + it("performOperation: flood fill returns the covered vox chunk keys", async () => { + const data = new BigUint64Array(1000); + for (let x = 3; x <= 7; x++) { + data[0 * 100 + 3 * 10 + x] = 1n; + data[0 * 100 + 7 * 10 + x] = 1n; + } + for (let y = 3; y <= 7; y++) { + data[0 * 100 + y * 10 + 3] = 1n; + data[0 * 100 + y * 10 + 7] = 1n; + } + mockSource.serverStorage.set("0,0,0", data.buffer); + + const covered = await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed: new Float32Array([5, 5, 0]), + value: 5n, + maxVoxels: 100, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + }); + + expect(covered).toEqual(["lod0#0,0,0"]); + await vi.runAllTimersAsync(); + }); + + it("performOperation: refused flood fill returns empty coverage", async () => { + const data = new BigUint64Array(1000); + data[0 * 100 + 5 * 10 + 5] = 5n; + mockSource.serverStorage.set("0,0,0", data.buffer); + + const covered = await controller.performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seed: new Float32Array([5, 5, 0]), + value: 5n, + maxVoxels: 100, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + }); + + expect(covered).toEqual([]); + expect(mockSource.applyEdits).not.toHaveBeenCalled(); + }); + + it("performOperation: brush returns the covered vox chunk keys", async () => { + const covered = await controller.performOperation({ + type: VoxelOperationType.BRUSH, + centers: [new Float32Array([5, 5, 5])], + radius: 3, + value: 5n, + shape: BrushShape.SPHERE, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + }); + + expect(covered).toEqual(["lod0#0,0,0"]); + await vi.runAllTimersAsync(); + }); + + it("performOperation: filtered brush counts already-written voxels as covered", async () => { + const data = new BigUint64Array(1000).fill(5n); + mockSource.serverStorage.set("0,0,0", data.buffer); + + const covered = await controller.performOperation({ + type: VoxelOperationType.BRUSH, + centers: [new Float32Array([5, 5, 5])], + radius: 3, + value: 5n, + shape: BrushShape.SPHERE, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + filterValue: 1n, + }); + + expect(covered).toEqual(["lod0#0,0,0"]); + expect(mockSource.applyEdits).not.toHaveBeenCalled(); + }); + + it("performOperation: filtered brush leaves filtered-out chunks uncovered", async () => { + const data = new BigUint64Array(1000).fill(7n); + mockSource.serverStorage.set("0,0,0", data.buffer); + + const covered = await controller.performOperation({ + type: VoxelOperationType.BRUSH, + centers: [new Float32Array([5, 5, 5])], + radius: 3, + value: 5n, + shape: BrushShape.SPHERE, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + filterValue: 1n, + }); + + expect(covered).toEqual([]); + expect(mockSource.applyEdits).not.toHaveBeenCalled(); + }); +}); + +describe("VolumeChunkSource.applyEdits: unreadable stored chunk", () => { + it("treats a corrupt (undecodable) chunk as empty and repairs it on write", async () => { + const source = createMockSource(); + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + (source.download as any).mockRejectedValue( + new Error("Raw-format chunk is 0 bytes, but 8 * 1 = 8 bytes expected."), + ); + + const change = await source.applyEdits("0,0,0", [3], [7n]); + + expect(source.writeChunk).toHaveBeenCalled(); + const written = new Uint8Array(source.serverStorage.get("0,0,0")!); + expect(Array.from(written)).toEqual([0, 0, 0, 7, 0, 0, 0, 0]); + expect(Number((change.oldValues as any)[0])).toBe(0); + expect(warnSpy).toHaveBeenCalled(); + warnSpy.mockRestore(); + }); + + it("propagates transient (network) failures without writing", async () => { + const source = createMockSource(); + (source.download as any).mockRejectedValue( + new HttpError("http://store/chunk", 503, "Service Unavailable"), + ); + + await expect(source.applyEdits("0,0,0", [3], [7n])).rejects.toThrow( + "HTTP error 503", + ); + expect(source.writeChunk).not.toHaveBeenCalled(); + }); + + it("fails closed on unrecognized read errors instead of repairing", async () => { + const source = createMockSource(); + (source.download as any).mockRejectedValue( + new Error("some wrapped kvstore failure"), + ); + + await expect(source.applyEdits("0,0,0", [3], [7n])).rejects.toThrow( + "some wrapped kvstore failure", + ); + expect(source.writeChunk).not.toHaveBeenCalled(); + }); + + it("invalidates the shared backend cache entry after a successful write", async () => { + const source = createMockSource(); + mockQueueManager.invalidateCachedChunks.mockClear(); + source.serverStorage.set("0,0,0", new Uint8Array(8).fill(2).buffer); + + await source.applyEdits("0,0,0", [3], [7n]); + + expect(mockQueueManager.invalidateCachedChunks).toHaveBeenCalledWith( + source, + ["0,0,0"], + ); + }); +}); + +describe("BackendVoxelAccessor: mid-flight invalidation", () => { + let controller: VoxelEditController; + let source: any; + + beforeEach(() => { + vi.clearAllMocks(); + source = createMockSource(); + (mockRpc.get as any).mockImplementation((id: number) => { + if (id === 0) return mockChunkManager; + if (id === 100) return source; + if (id === 999) return { value: 0 }; + return null; + }); + controller = new VoxelEditController(mockRpc, { + resolutions: [resConfig(0, [1, 1, 1], [2, 2, 2])], + pendingOpCount: 999, + }); + }); + + it("a load in flight across a write does not re-cache pre-write data", async () => { + source.serverStorage.set("0,0,0", new Uint8Array(8).fill(1).buffer); + let release!: () => void; + const gate = new Promise((resolve) => (release = resolve)); + source.download.mockImplementation(async (chunk: any) => { + // Snapshot at request time, deliver after the gate: simulates a read + // whose response predates a write that lands while it is in flight. + const buffer = source.serverStorage.get("0,0,0")!.slice(0); + await gate; + chunk.data = new Uint8Array(buffer); + }); + + const accessor = (controller as any).getAccessor(0); + const read = accessor.getValue(0, 0, 0); + + source.serverStorage.set("0,0,0", new Uint8Array(8).fill(9).buffer); + accessor.invalidate("0,0,0"); + release(); + + expect(await read).toBe(9n); + expect(await accessor.getValue(0, 0, 0)).toBe(9n); + }); + + it("callers keep sharing the in-flight load across an invalidation", async () => { + source.serverStorage.set("0,0,0", new Uint8Array(8).fill(1).buffer); + const accessor = (controller as any).getAccessor(0); + const p1 = accessor.getOrLoadChunkContext("0,0,0", 0, 0, 0); + accessor.invalidate("0,0,0"); + const p2 = accessor.getOrLoadChunkContext("0,0,0", 0, 0, 0); + expect(p2).toBe(p1); + await p1; + }); + + it("flood fill flushes pending edits before reading", async () => { + const key = makeVoxChunkKey("0,0,0", 0); + controller.commitVoxels([{ key, indices: [0], value: 5n, seq: 1 }]); + + // Fill value equals the seed's current value: early return right after + // the flush, keeping the test cheap. + await (controller as any).performOperation({ + type: VoxelOperationType.FLOOD_FILL, + seq: 2, + seed: new Float32Array([0, 0, 0]), + value: 5n, + maxVoxels: 8, + basis: { u: new Float32Array([1, 0, 0]), v: new Float32Array([0, 1, 0]) }, + }); + + expect(source.applyEdits).toHaveBeenCalledWith( + "0,0,0", + [0], + expect.arrayContaining([5n]), + ); + }); +}); diff --git a/src/voxel_annotation/backend.ts b/src/voxel_annotation/backend.ts new file mode 100644 index 0000000000..22d3ee77f3 --- /dev/null +++ b/src/voxel_annotation/backend.ts @@ -0,0 +1,1816 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { ChunkState } from "#src/chunk_manager/base.js"; +import type { SharedWatchableValue } from "#src/shared_watchable_value.js"; +import { DataType } from "#src/sliceview/base.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, + VolumeChunkSource, +} from "#src/sliceview/volume/backend.js"; +import type { TypedArray } from "#src/util/array.js"; +import { mat4, vec3 } from "#src/util/geom.js"; +import * as matrix from "#src/util/matrix.js"; +import type { + VoxelLayerResolution, + EditAction, + VoxelChange, + VoxelOperation, + BrushOperation, + FloodFillOperation, +} from "#src/voxel_annotation/base.js"; +import { + VOXEL_EMPTY_VALUE, + VOXEL_EDIT_STAMINA, + VOX_EDIT_BACKEND_RPC_ID, + VOX_EDIT_COMMIT_VOXELS_RPC_ID, + VOX_RELOAD_CHUNKS_RPC_ID, + VOX_EDIT_FAILURE_RPC_ID, + VOX_EDIT_UNDO_RPC_ID, + VOX_EDIT_REDO_RPC_ID, + VOX_EDIT_HISTORY_UPDATE_RPC_ID, + VOX_EDIT_OPERATION_RPC_ID, + VoxelOperationType, + BrushShape, + makeVoxChunkKey, + parseVoxChunkKey, + makeChunkKey, +} from "#src/voxel_annotation/base.js"; +import type { RPC } from "#src/worker_rpc.js"; +import { + registerPromiseRPC, + SharedObject, + registerRPC, + registerSharedObject, + initializeSharedObjectCounterpart, +} from "#src/worker_rpc.js"; + +const OFFSETS_26_CONNECTED_BACKEND: number[][] = []; +for (let z = -1; z <= 1; z++) { + for (let y = -1; y <= 1; y++) { + for (let x = -1; x <= 1; x++) { + if (x === 0 && y === 0 && z === 0) continue; + OFFSETS_26_CONNECTED_BACKEND.push([x, y, z]); + } + } +} + +function getFlatChunkData( + chunk: VolumeChunk, + spec: any, +): Uint32Array | BigUint64Array | TypedArray | null { + if (!chunk.data) return null; + + if (!spec.compressedSegmentationBlockSize) { + return chunk.data as TypedArray; + } + + const size = chunk.chunkDataSize!; + const numElements = size[0] * size[1] * size[2]; + const compressedData = chunk.data as Uint32Array; + const baseOffset = compressedData.length > 0 ? compressedData[0] : 0; + const subchunkSize = spec.compressedSegmentationBlockSize; + + if (spec.dataType === DataType.UINT32) { + const out = new Uint32Array(numElements); + if (baseOffset !== 0) { + decodeChannelUint32(out, compressedData, baseOffset, size, subchunkSize); + } + return out; + } else { + const out = new BigUint64Array(numElements); + if (baseOffset !== 0) { + decodeChannelUint64(out, compressedData, baseOffset, size, subchunkSize); + } + return out; + } +} + +interface ChunkContext { + key: string; + data: TypedArray | null; + minX: number; + maxX: number; + minY: number; + maxY: number; + minZ: number; + maxZ: number; + strideY: number; + strideZ: number; +} + +class BackendVoxelAccessor { + private chunkContexts = new Map(); + private pendingLoads = new Map>(); + + // 64 chunks * (64^3 voxels * 8 bytes) ~= 134 MB max per LOD level if chunks are uncompressed 64^3 UINT64. + private readonly MAX_CACHE_SIZE = 64; + + private readonly volMin: Float32Array; + private readonly volMax: Float32Array; + private readonly chunkDimension: Uint32Array; + private readonly fillValue: bigint; + + constructor(private source: VolumeChunkSource) { + const spec = source.spec; + this.volMin = spec.lowerVoxelBound; + this.volMax = spec.upperVoxelBound; + this.chunkDimension = spec.chunkDataSize; + const fv = spec.fillValue; + this.fillValue = typeof fv === "bigint" ? fv : BigInt(fv); + } + + // Bumped when a key is invalidated while its load is in flight, so the + // load reloads instead of re-caching pre-write data. Pruned on settle. + private loadGenerations = new Map(); + + public invalidate(key: string) { + if (this.pendingLoads.has(key)) { + this.loadGenerations.set(key, (this.loadGenerations.get(key) ?? 0) + 1); + } + this.chunkContexts.delete(key); + } + + async getValue(x: number, y: number, z: number): Promise { + if ( + x < this.volMin[0] || + x >= this.volMax[0] || + y < this.volMin[1] || + y >= this.volMax[1] || + z < this.volMin[2] || + z >= this.volMax[2] + ) { + return null; + } + + const cx = Math.floor(x / this.chunkDimension[0]); + const cy = Math.floor(y / this.chunkDimension[1]); + const cz = Math.floor(z / this.chunkDimension[2]); + const key = `${cx},${cy},${cz}`; + + let ctx = this.chunkContexts.get(key); + if (ctx) { + this.chunkContexts.delete(key); + this.chunkContexts.set(key, ctx); + } else { + ctx = await this.getOrLoadChunkContext(key, cx, cy, cz); + } + + return this.readLocal(ctx, x, y, z); + } + + public getOrLoadChunkContext( + key: string, + cx: number, + cy: number, + cz: number, + ): Promise { + let promise = this.pendingLoads.get(key); + if (promise) return promise; + + promise = (async () => { + try { + while (true) { + const generation = this.loadGenerations.get(key); + const ctx = await this.loadChunkContext(key, cx, cy, cz); + // Invalidated mid-load: the data may predate the write — reload. + if (this.loadGenerations.get(key) !== generation) continue; + if (this.chunkContexts.size >= this.MAX_CACHE_SIZE) { + const oldestKey = this.chunkContexts.keys().next().value; + if (oldestKey) { + this.chunkContexts.delete(oldestKey); + } + } + this.chunkContexts.set(key, ctx); + return ctx; + } + } finally { + this.pendingLoads.delete(key); + this.loadGenerations.delete(key); + } + })(); + this.pendingLoads.set(key, promise); + return promise; + } + + private readLocal( + ctx: ChunkContext, + x: number, + y: number, + z: number, + ): bigint { + if (!ctx.data) return this.fillValue; + + const lx = x - ctx.minX; + const ly = y - ctx.minY; + const lz = z - ctx.minZ; + const index = lz * ctx.strideZ + ly * ctx.strideY + lx; + + const val = ctx.data[index]; + return typeof val === "bigint" ? val : BigInt(val); + } + + private async loadChunkContext( + key: string, + cx: number, + cy: number, + cz: number, + ): Promise { + // Same pattern as `applyEdits`: only use the shared cache entry when its + // data is already resident (synchronous read, safe). Otherwise download + // into an isolated chunk the queue manager cannot dispose mid-download. + let chunk = this.source.chunks.get(key) as VolumeChunk | undefined; + if ( + chunk === undefined || + chunk.state > ChunkState.SYSTEM_MEMORY_WORKER || + !chunk.data + ) { + chunk = this.source.getIsolatedChunk(new Float32Array([cx, cy, cz])); + try { + await this.source.download(chunk, new AbortController().signal); + } catch { + return this.createContext(key, null, cx, cy, cz, chunk); + } + } + + if (!chunk.chunkDataSize) { + this.source.computeChunkBounds(chunk); + } + + const flatData = getFlatChunkData(chunk, this.source.spec); + return this.createContext(key, flatData, cx, cy, cz, chunk); + } + + private createContext( + key: string, + data: TypedArray | null, + cx: number, + cy: number, + cz: number, + chunk: VolumeChunk, + ): ChunkContext { + const size = chunk.chunkDataSize || this.chunkDimension; + const minX = cx * this.chunkDimension[0]; + const minY = cy * this.chunkDimension[1]; + const minZ = cz * this.chunkDimension[2]; + + return { + key, + data, + minX, + maxX: minX + size[0], + minY, + maxY: minY + size[1], + minZ, + maxZ: minZ + size[2], + strideY: size[0], + strideZ: size[0] * size[1], + }; + } +} + +type Skipper = (x: number, y: number, z: number) => boolean; + +class BrushOptimizationCache { + private active = false; + + private val: bigint = 0n; + private shape: BrushShape = BrushShape.SPHERE; + private cx = 0; + private cy = 0; + private cz = 0; + private r2 = 0; + + private ux = 0; + private uy = 0; + private uz = 0; + private vx = 0; + private vy = 0; + private vz = 0; + + reset() { + this.active = false; + } + + buildSkipper( + newValue: bigint, + newShape: BrushShape, + newBasis?: { u: Float32Array; v: Float32Array }, + ): Skipper { + if (!this.active || this.val !== newValue || this.shape !== newShape) { + return () => false; + } + + const { cx, cy, cz, r2 } = this; + + if (this.shape === BrushShape.SPHERE) { + return (x: number, y: number, z: number) => { + const dx = x - cx; + const dy = y - cy; + const dz = z - cz; + return dx * dx + dy * dy + dz * dz <= r2; + }; + } + + if (this.shape === BrushShape.DISK && newBasis) { + const dotU = + this.ux * newBasis.u[0] + + this.uy * newBasis.u[1] + + this.uz * newBasis.u[2]; + const dotV = + this.vx * newBasis.v[0] + + this.vy * newBasis.v[1] + + this.vz * newBasis.v[2]; + + if (dotU < 0.9999 || dotV < 0.9999) { + return () => false; + } + + const { ux, uy, uz, vx, vy, vz } = this; + + return (x: number, y: number, z: number) => { + const dx = x - cx; + const dy = y - cy; + const dz = z - cz; + const distU = dx * ux + dy * uy + dz * uz; + const distV = dx * vx + dy * vy + dz * vz; + return distU * distU + distV * distV <= r2; + }; + } + + return () => false; + } + + update( + center: { x: number; y: number; z: number }, + radius: number, + value: bigint, + shape: BrushShape, + basis?: { u: Float32Array; v: Float32Array }, + ) { + this.active = true; + this.cx = center.x; + this.cy = center.y; + this.cz = center.z; + this.r2 = radius * radius; + this.val = value; + this.shape = shape; + + if (basis) { + this.ux = basis.u[0]; + this.uy = basis.u[1]; + this.uz = basis.u[2]; + this.vx = basis.v[0]; + this.vy = basis.v[1]; + this.vz = basis.v[2]; + } + } +} + +@registerSharedObject(VOX_EDIT_BACKEND_RPC_ID) +export class VoxelEditController extends SharedObject { + private sources = new Map(); + private resolutions = new Map< + number, + VoxelLayerResolution & { invTransform: mat4 } + >(); + + private pendingEdits: { + key: string; + indices: number[] | Uint32Array; + value?: bigint; + values?: ArrayLike; + size?: number[]; + seq?: number; + }[] = []; + + // Per LOD-0 vox key: the highest frontend dispatch seq whose edits have been + // durably written to that chunk. Echoed in reload messages (including + // downsample cascade reloads, keyed by origin) so the frontend can tell + // whether the refetched data covers everything its overlay represents. + // Pruned via maybePruneFlushedSeq: bounded by the cascades in flight. + private lastFlushedSeq = new Map(); + + // Running cascade chains per origin (a counter: re-enqueueing can put two + // chains in flight). + private activeChainCounts = new Map(); + + // A future flush recomputes a higher max (dispatch seqs are globally + // monotonic); the only reader early pruning breaks is a cascade capturing + // at chain start — hence the two guards. + private maybePruneFlushedSeq(key: string): void { + if (this.activeChainCounts.has(key)) return; + if (this.downsampleQueueSet.has(key)) return; + this.lastFlushedSeq.delete(key); + } + public pendingOpCount: SharedWatchableValue; + + private updatePendingCount() { + const editedVoxel = this.pendingEdits.reduce( + (a, { indices }) => indices.length + a, + 0, + ); + const pendingEdits = VOXEL_EDIT_STAMINA.pendingEdits(editedVoxel); + const downsampling = VOXEL_EDIT_STAMINA.downsamplingJobs( + this.downsampleQueue.length, + this.resolutions.size, + ); + this.pendingOpCount.value = pendingEdits + downsampling; + } + + private commitDebounceTimer: number | undefined; + private readonly commitDebounceDelayMs: number = 300; + // Serializes the LOD-0 writers (debounced flush, undo/redo): two + // concurrent applyEdits on the same chunk are isolated read-modify-writes, + // so the last one would silently drop the other's voxels. Brush commits + // stay a synchronous append; only their flush enters the chain. + private opChain: Promise = Promise.resolve(); + + private runExclusive(op: () => Promise): Promise { + const run = this.opChain.then(op, op); + this.opChain = run.catch(() => {}); + return run; + } + + // Flushes edits committed before a discrete operation. Caller must hold + // the exclusion chain. + private async flushPendingLocked(): Promise { + if (this.commitDebounceTimer !== undefined) { + clearTimeout(this.commitDebounceTimer); + this.commitDebounceTimer = undefined; + } + if (this.pendingEdits.length > 0) await this.flushPending(); + } + + // Undo/redo history + private undoStack: EditAction[] = []; + private redoStack: EditAction[] = []; + private readonly MAX_HISTORY_SIZE: number = 100; + + private downsampleQueue: string[] = []; + private downsampleQueueSet: Set = new Set(); + private isProcessingDownsampleQueue: boolean = false; + private activeDownsamples = 0; + private readonly MAX_CONCURRENT_DOWNSAMPLES = 16; + private downsampleChunkLocks = new Map>(); + + private brushCache = new BrushOptimizationCache(); + private accessors = new Map(); + + private morphologicalConfig = { + growthThresholds: [ + { count: 100, size: 1 }, + { count: 1000, size: 3 }, + { count: 10000, size: 5 }, + { count: 100000, size: 7 }, + ], + maxSize: 9, + }; + + constructor(rpc: RPC, options: any) { + super(); + this.pendingOpCount = rpc.get(options.pendingOpCount); + initializeSharedObjectCounterpart(this, rpc, options); + + const passedResolutions = options?.resolutions as + | VoxelLayerResolution[] + | undefined; + if (passedResolutions === undefined || !Array.isArray(passedResolutions)) { + throw new Error( + "VoxelEditBackend: missing required 'resolutions' array during initialization", + ); + } + + for (const res of passedResolutions) { + const rank = res.chunkSize.length; + const invTransform = new Float32Array((rank + 1) ** 2); + matrix.inverse( + invTransform, + rank + 1, + new Float32Array(res.transform), + rank + 1, + rank + 1, + ); + this.resolutions.set(res.lodIndex, { + ...res, + invTransform: invTransform as mat4, + }); + const resolved = rpc.get(res.sourceRpc) as VolumeChunkSource | undefined; + if (!resolved) { + throw new Error( + `VoxelEditBackend: failed to resolve VolumeChunkSource for LOD ${res.lodIndex}`, + ); + } + this.sources.set(res.lodIndex, resolved); + } + + this.notifyHistoryChanged(); + } + + private async flushPending(): Promise { + this.brushCache.reset(); + const edits = this.pendingEdits; + this.pendingEdits = []; + this.commitDebounceTimer = undefined; + if (edits.length === 0) { + // Even if nothing to flush, history sizes may not have changed. + this.notifyHistoryChanged(); + return; + } + + const editsByVoxKey = new Map>(); + const maxSeqByVoxKey = new Map(); + + for (const edit of edits) { + let chunkMap = editsByVoxKey.get(edit.key); + if (!chunkMap) { + chunkMap = new Map(); + editsByVoxKey.set(edit.key, chunkMap); + } + if (edit.seq !== undefined) { + maxSeqByVoxKey.set( + edit.key, + Math.max(maxSeqByVoxKey.get(edit.key) ?? 0, edit.seq), + ); + } + + const inds = edit.indices as ArrayLike; + if (edit.values) { + // Handle array of values + const vals = Array.from(edit.values); + if (vals.length !== inds.length) { + throw new Error("flushPending: values length mismatch with indices"); + } + for (let i = 0; i < inds.length; ++i) { + chunkMap.set(inds[i]!, vals[i]!); + } + } else if (edit.value !== undefined) { + // Handle single value for all indices + for (let i = 0; i < inds.length; ++i) { + chunkMap.set(inds[i]!, edit.value); + } + } else { + throw new Error("flushPending: edit missing value(s)"); + } + } + + const failedVoxChunkKeys: string[] = []; + let firstErrorMessage: string | undefined = undefined; + + const newAction: EditAction = { + changes: new Map(), + timestamp: Date.now(), + description: "Voxel Edit", + }; + + for (const [voxKey, chunkEdits] of editsByVoxKey.entries()) { + try { + const parsedKey = parseVoxChunkKey(voxKey); + if (!parsedKey) { + const msg = `flushPending: Failed to parse vox chunk key: ${voxKey}`; + console.error(msg); + failedVoxChunkKeys.push(voxKey); + if (firstErrorMessage === undefined) firstErrorMessage = msg; + continue; + } + const source = this.sources.get(parsedKey.lodIndex); + if (!source) { + const msg = `flushPending: No source found for LOD index ${parsedKey.lodIndex}`; + console.error(msg); + failedVoxChunkKeys.push(voxKey); + if (firstErrorMessage === undefined) firstErrorMessage = msg; + continue; + } + + const indices = Array.from(chunkEdits.keys()); + const values = Array.from(chunkEdits.values()); + + const change = await source.applyEdits( + parsedKey.chunkKey, + indices, + values, + ); + const accessor = this.getAccessor(parsedKey.lodIndex); + accessor.invalidate(parsedKey.chunkKey); + + const flushedSeq = maxSeqByVoxKey.get(voxKey); + if (flushedSeq !== undefined) { + this.lastFlushedSeq.set( + voxKey, + Math.max(this.lastFlushedSeq.get(voxKey) ?? 0, flushedSeq), + ); + } + + newAction.changes.set(voxKey, change); + } catch (e) { + const msg = e instanceof Error ? e.message : String(e); + console.error(`Failed to write chunk ${voxKey}:`, e); + failedVoxChunkKeys.push(voxKey); + if (firstErrorMessage === undefined) firstErrorMessage = msg; + } + } + + // Failed keys are excluded: their store was not modified, so a real + // reload would refetch unchanged data and race with the failure + // rollback below. + const flushedKeys = editsByVoxKey + .keys() + .toArray() + .filter((voxKey) => !failedVoxChunkKeys.includes(voxKey)); + const coveredSeqs: Record = {}; + for (const voxKey of flushedKeys) { + coveredSeqs[voxKey] = this.lastFlushedSeq.get(voxKey) ?? 0; + } + if (flushedKeys.length > 0) { + this.callChunkReload(flushedKeys, false, undefined, coveredSeqs); + } + + if (newAction.changes.size > 0) { + this.undoStack.push(newAction); + if (this.undoStack.length > this.MAX_HISTORY_SIZE) { + this.undoStack.shift(); + } + this.redoStack.length = 0; + } + + // Notify frontend of history changes after any commit attempt + this.notifyHistoryChanged(); + + if (failedVoxChunkKeys.length > 0) { + this.rpc?.invoke(VOX_EDIT_FAILURE_RPC_ID, { + rpcId: this.rpcId, + voxChunkKeys: failedVoxChunkKeys, + message: firstErrorMessage ?? "Voxel edit commit failed.", + }); + } + + const hasDownsampling = this.resolutions.size > 1; + if (hasDownsampling) { + for (const [voxKey, _] of editsByVoxKey.entries()) { + if (failedVoxChunkKeys.includes(voxKey)) continue; + this.enqueueDownsample(voxKey); + } + } + // The overlay is cleared by the swap-on-arrival registered in the real + // callChunkReload above (not here): on the no-downsampling path the max-res + // overlay is dropped when the refetched real chunk arrives. + + // With downsampling, the queue guard defers pruning to the chain-end hook. + for (const voxKey of editsByVoxKey.keys()) { + this.maybePruneFlushedSeq(voxKey); + } + + this.updatePendingCount(); + } + + private getAccessor(lodIndex: number): BackendVoxelAccessor { + let accessor = this.accessors.get(lodIndex); + if (!accessor) { + const source = this.sources.get(lodIndex); + if (!source) throw new Error(`No source for LOD ${lodIndex}`); + accessor = new BackendVoxelAccessor(source); + this.accessors.set(lodIndex, accessor); + } + return accessor; + } + + commitVoxels( + edits: { + key: string; + indices: number[] | Uint32Array; + value?: bigint; + values?: ArrayLike; + size?: number[]; + seq?: number; + }[], + ) { + for (const e of edits) { + if (!e || !e.key || !e.indices) { + throw new Error( + "VoxelEditController.commitVoxels: invalid edit payload", + ); + } + this.pendingEdits.push(e); + } + this.updatePendingCount(); + if (this.commitDebounceTimer !== undefined) + clearTimeout(this.commitDebounceTimer); + this.commitDebounceTimer = setTimeout(() => { + void this.runExclusive(() => this.flushPending()); + }, this.commitDebounceDelayMs) as unknown as number; + } + + // `overlayKeysToClear[realKey]` is the overlay vox key (LOD 0) to drop once the + // real chunk `realKey` reaches the GPU. A real key absent from the map defaults + // to itself on the frontend; downsampled parents map to the originating LOD-0 + // key so the visible (forced LOD-0) overlay is cleared as soon as any covering + // real LOD arrives. Using a keyed map (not parallel arrays) keeps real and + // overlay keys aligned regardless of ordering or partial population. + // `coveredSeqs[realKey]` is the highest frontend dispatch seq whose edits + // are guaranteed present in the stored data this reload refetches (for + // downsampled parents: the origin chunk's flushed seq at the time the child + // was read). The frontend clears the matching overlay only if this covers + // the last dispatched stroke that touched it. + // `isRollback` marks a state rollback (undo/redo): the frontend purges the + // overlay tags so the swap clears on the first arrival, whatever it covers. + callChunkReload( + voxChunkKeys: string[], + isForPreviewChunks = false, + overlayKeysToClear?: Record, + coveredSeqs?: Record, + isRollback = false, + ) { + this.rpc?.invoke(VOX_RELOAD_CHUNKS_RPC_ID, { + rpcId: this.rpcId, + voxChunkKeys: voxChunkKeys, + isForPreviewChunks, + overlayKeysToClear, + coveredSeqs, + isRollback, + }); + } + + // --- Start of Downsampling Logic --- + + /** + * NOTE: Architecture Limitation + * The current downsampling architecture assumes a Many-to-1 (or 1-to-1) mapping between + * child chunks and parent chunks. It calculates a single parent chunk key for a given + * child chunk. + */ + + private enqueueDownsample(key: string): void { + if (key.length === 0) return; + if (!this.downsampleQueueSet.has(key)) { + this.downsampleQueueSet.add(key); + this.downsampleQueue.push(key); + } + if (!this.isProcessingDownsampleQueue) { + this.isProcessingDownsampleQueue = true; + Promise.resolve().then(() => this.processDownsampleQueue()); + } + } + + private async processDownsampleQueue(): Promise { + this.isProcessingDownsampleQueue = true; + + const scheduleNext = () => { + while ( + this.downsampleQueue.length > 0 && + this.activeDownsamples < this.MAX_CONCURRENT_DOWNSAMPLES + ) { + const key = this.downsampleQueue.shift() as string; + this.downsampleQueueSet.delete(key); + this.activeDownsamples++; + this.activeChainCounts.set( + key, + (this.activeChainCounts.get(key) ?? 0) + 1, + ); + + this.processDownsampleChain(key).finally(() => { + this.activeDownsamples--; + const remaining = (this.activeChainCounts.get(key) ?? 1) - 1; + if (remaining <= 0) this.activeChainCounts.delete(key); + else this.activeChainCounts.set(key, remaining); + this.maybePruneFlushedSeq(key); + scheduleNext(); + }); + } + + if (this.activeDownsamples === 0 && this.downsampleQueue.length === 0) { + this.isProcessingDownsampleQueue = false; + } + }; + + scheduleNext(); + } + + private async processDownsampleChain(key: string): Promise { + // Coverage is captured once for the whole chain, before the first child + // read: every step propagates data derived from the LOD-0 content read at + // chain start, so a flush completing mid-chain (higher seq) is NOT + // included in what the later steps write — claiming it would clear the + // overlay over a parent that lacks those edits. Under-claiming is safe: + // that flush enqueues its own chain, which re-claims with its seq. + const chainCoveredSeq = this.lastFlushedSeq.get(key) ?? 0; + let currentKey: string | null = key; + + while (currentKey !== null) { + currentKey = await this.downsampleStep(currentKey, key, chainCoveredSeq); + } + + // Note: the overlay is no longer cleared eagerly here. Each parent's real + // reload (above) drops the originating LOD-0 overlay once it reaches the + // GPU, so the visible overlay is never removed before its replacement is + // rendered — even when zoomed out. + this.updatePendingCount(); + } + + private async withChunkLock( + key: string, + op: () => Promise, + ): Promise { + const prev = this.downsampleChunkLocks.get(key) || Promise.resolve(); + + const current = (async () => { + try { + await prev; + } catch { + // + } + return op(); + })(); + + const nextPromise = current.then( + () => {}, + () => {}, + ); + this.downsampleChunkLocks.set(key, nextPromise); + + nextPromise.then(() => { + if (this.downsampleChunkLocks.get(key) === nextPromise) { + this.downsampleChunkLocks.delete(key); + } + }); + + return current; + } + + /** + * Performs a single downsampling step from a child chunk to its parent. + * @returns The key of the parent chunk that was updated, or null if the cascade should stop. + */ + private async downsampleStep( + childKey: string, + originKey: string, + originCoveredSeq: number, + ): Promise { + const childInfo = parseVoxChunkKey(childKey); + if (childInfo === null) { + console.error(`[Downsample] Invalid child key format: ${childKey}`); + return null; + } + const childRes = this.resolutions.get(childInfo.lodIndex)!; + + const parentInfo = this._getParentChunkInfo(childKey, childRes); + if (parentInfo === null) { + return null; + } + const { parentKey, parentSource, parentRes } = parentInfo; + + // The child read and the parent-update computation must run under the + // parent lock along with the write: with only the write serialized, two + // concurrent chains can compute from reads taken at different times and + // land in the wrong order, durably overwriting the fresher parent with a + // stale result. Serializing the compute alongside is accepted: + // applyEdits' network I/O dominates this lock anyway, and parallelism + // across distinct parents is unaffected. + return this.withChunkLock(parentKey, async () => { + const childAccessor = this.getAccessor(childInfo.lodIndex); + const childCtx = await childAccessor.getOrLoadChunkContext( + childInfo.chunkKey, + childInfo.x, + childInfo.y, + childInfo.z, + ); + + if (!childCtx.data) { + return null; + } + + const childChunkData = childCtx.data as Uint32Array | BigUint64Array; + + const childActualSize = [ + childCtx.maxX - childCtx.minX, + childCtx.maxY - childCtx.minY, + childCtx.maxZ - childCtx.minZ, + ]; + + const update = this._calculateParentUpdate( + childChunkData, + childRes, + parentRes, + childInfo, + childActualSize, + ); + if (update.indices.length === 0) { + return parentKey; + } + + try { + await parentSource.applyEdits( + parentInfo.chunkKey, + update.indices, + update.values, + ); + // Reload the real parent lazily; when it reaches the GPU, clear the + // originating LOD-0 overlay (the visible one when zoomed out) — if + // the propagated data covers every stroke dispatched to the origin. + this.callChunkReload( + [parentKey], + false, + { [parentKey]: originKey }, + { [parentKey]: originCoveredSeq }, + ); + const parentAccessor = this.getAccessor(parentRes.lodIndex); + parentAccessor.invalidate(parentInfo.chunkKey); + } catch (e) { + console.error( + `[Downsample] Failed to apply edits to parent chunk ${parentKey}:`, + e, + ); + this.rpc?.invoke(VOX_EDIT_FAILURE_RPC_ID, { + rpcId: this.rpcId, + voxChunkKeys: [parentKey], + message: `Downsampling to ${parentKey} failed.`, + }); + return null; + } + + return parentKey; + }); + } + + /** + * Helper to find and describe the parent chunk. + */ + private _getParentChunkInfo( + childKey: string, + childRes: VoxelLayerResolution, + ) { + const childInfo = parseVoxChunkKey(childKey)!; + const parentLodIndex = childInfo.lodIndex + 1; + const parentRes = this.resolutions.get(parentLodIndex); + if (parentRes === undefined) return null; // No parent LOD exists + + const parentSource = this.sources.get(parentLodIndex)!; + const rank = childRes.chunkSize.length; + + // Find the world coordinate of the child chunk's origin + const childVoxelOrigin = new Float32Array(rank); + childVoxelOrigin.set([ + childInfo.x * childRes.chunkSize[0], + childInfo.y * childRes.chunkSize[1], + childInfo.z * childRes.chunkSize[2], + ]); + const childPhysOrigin = new Float32Array(rank); + matrix.transformPoint( + childPhysOrigin, + new Float32Array(childRes.transform), + rank + 1, + childVoxelOrigin, + rank, + ); + + // Transform that world coordinate into the parent's voxel space + const parentVoxelCoordOfChildOrigin = new Float32Array(rank); + matrix.transformPoint( + parentVoxelCoordOfChildOrigin, + parentRes.invTransform, + rank + 1, + childPhysOrigin, + rank, + ); + + // Determine the parent chunk's grid position + const parentX = Math.floor( + parentVoxelCoordOfChildOrigin[0] / parentRes.chunkSize[0], + ); + const parentY = Math.floor( + parentVoxelCoordOfChildOrigin[1] / parentRes.chunkSize[1], + ); + const parentZ = Math.floor( + parentVoxelCoordOfChildOrigin[2] / parentRes.chunkSize[2], + ); + + const parentChunkKey = makeChunkKey(parentX, parentY, parentZ); + const parentKey = makeVoxChunkKey(parentChunkKey, parentLodIndex); + return { parentKey, chunkKey: parentChunkKey, parentRes, parentSource }; + } + + /** + * Calculates the downsampled voxel values for a region of a parent chunk. + * This is the core aggregation logic. + */ + private _calculateParentUpdate( + childChunkData: Uint32Array | BigUint64Array, + childRes: VoxelLayerResolution & { invTransform: mat4 }, + parentRes: VoxelLayerResolution & { invTransform: mat4 }, + childInfo: { x: number; y: number; z: number }, + childActualSize: number[], + ) { + const indices: number[] = []; + const values: bigint[] = []; + const rank = childRes.chunkSize.length; + const childChunkSize = childRes.chunkSize; + const [childDataW, childDataH, childDataD] = childActualSize; + const parentChunkSize = parentRes.chunkSize; + + const childDataLength = childChunkData.length; + + // Transform to map a point in parent-voxel-space to a point in child-voxel-space. + const parentVoxelToChildVoxelTransform = mat4.multiply( + mat4.create(), + childRes.invTransform, + new Float32Array(parentRes.transform) as mat4, + ); + + // Calculate the child chunk's origin and extent in absolute child-voxel-space + const childChunkOrigin = new Float32Array([ + childInfo.x * childChunkSize[0], + childInfo.y * childChunkSize[1], + childInfo.z * childChunkSize[2], + ]); + const childChunkMax = new Float32Array([ + (childInfo.x + 1) * childChunkSize[0], + (childInfo.y + 1) * childChunkSize[1], + (childInfo.z + 1) * childChunkSize[2], + ]); + + // Transform child chunk bounds to physical space + const childPhysOrigin = new Float32Array(rank); + matrix.transformPoint( + childPhysOrigin, + new Float32Array(childRes.transform), + rank + 1, + childChunkOrigin, + rank, + ); + const childPhysMax = new Float32Array(rank); + matrix.transformPoint( + childPhysMax, + new Float32Array(childRes.transform), + rank + 1, + childChunkMax, + rank, + ); + + // Transform to parent-voxel-space to find the affected region + const parentVoxelMin = new Float32Array(rank); + matrix.transformPoint( + parentVoxelMin, + parentRes.invTransform, + rank + 1, + childPhysOrigin, + rank, + ); + const parentVoxelMax = new Float32Array(rank); + matrix.transformPoint( + parentVoxelMax, + parentRes.invTransform, + rank + 1, + childPhysMax, + rank, + ); + + // Determine which parent chunk this corresponds to (should match _getParentChunkInfo) + const parentChunkGridX = Math.floor(parentVoxelMin[0] / parentChunkSize[0]); + const parentChunkGridY = Math.floor(parentVoxelMin[1] / parentChunkSize[1]); + const parentChunkGridZ = Math.floor(parentVoxelMin[2] / parentChunkSize[2]); + + // Calculate the parent chunk's origin in absolute parent-voxel-space + const parentChunkOriginInParentVoxels = new Float32Array([ + parentChunkGridX * parentChunkSize[0], + parentChunkGridY * parentChunkSize[1], + parentChunkGridZ * parentChunkSize[2], + ]); + + // Calculate the region to iterate over in the parent chunk's LOCAL coordinate space (0 to chunkSize) + const parentLocalMin = new Float32Array(rank); + const parentLocalMax = new Float32Array(rank); + for (let i = 0; i < rank; ++i) { + parentLocalMin[i] = Math.max( + 0, + Math.floor(parentVoxelMin[i] - parentChunkOriginInParentVoxels[i]), + ); + parentLocalMax[i] = Math.min( + parentChunkSize[i], + Math.ceil(parentVoxelMax[i] - parentChunkOriginInParentVoxels[i]), + ); + } + + const [startX, startY, startZ] = parentLocalMin; + const [endX, endY, endZ] = parentLocalMax; + + const corners = new Array(8).fill(0).map(() => vec3.create()); + const transformedCorners = new Array(8).fill(0).map(() => vec3.create()); + const sourceVoxels: bigint[] = []; + const [parentW, parentH] = parentChunkSize; + + // Iterate over each voxel in the affected region of the parent chunk (in local coordinates) + for (let pz = startZ; pz < endZ; ++pz) { + for (let py = startY; py < endY; ++py) { + for (let px = startX; px < endX; ++px) { + // Convert from parent-chunk-local to absolute parent-voxel-space + const absParentX = parentChunkOriginInParentVoxels[0] + px; + const absParentY = parentChunkOriginInParentVoxels[1] + py; + const absParentZ = parentChunkOriginInParentVoxels[2] + pz; + + // Define the 8 corners of the current parent voxel in absolute parent-voxel-space + vec3.set(corners[0], absParentX, absParentY, absParentZ); + vec3.set(corners[1], absParentX + 1, absParentY, absParentZ); + vec3.set(corners[2], absParentX, absParentY + 1, absParentZ); + vec3.set(corners[3], absParentX + 1, absParentY + 1, absParentZ); + vec3.set(corners[4], absParentX, absParentY, absParentZ + 1); + vec3.set(corners[5], absParentX + 1, absParentY, absParentZ + 1); + vec3.set(corners[6], absParentX, absParentY + 1, absParentZ + 1); + vec3.set(corners[7], absParentX + 1, absParentY + 1, absParentZ + 1); + + // Transform corners to absolute child-voxel-space + for (let i = 0; i < 8; ++i) { + vec3.transformMat4( + transformedCorners[i], + corners[i], + parentVoxelToChildVoxelTransform, + ); + } + + // Find bounding box in absolute child-voxel-space + const childMin = vec3.clone(transformedCorners[0]); + const childMax = vec3.clone(transformedCorners[0]); + for (let i = 1; i < 8; ++i) { + vec3.min(childMin, childMin, transformedCorners[i]); + vec3.max(childMax, childMax, transformedCorners[i]); + } + + // Convert to child-chunk-local coordinates for array indexing + const localChildMin = vec3.create(); + const localChildMax = vec3.create(); + vec3.subtract(localChildMin, childMin, childChunkOrigin as any); + vec3.subtract(localChildMax, childMax, childChunkOrigin as any); + + // Collect all child voxels within this bounding box (in local coordinates) + sourceVoxels.length = 0; + const cStartX = Math.max(0, Math.floor(localChildMin[0])); + const cEndX = Math.min(childDataW, Math.ceil(localChildMax[0])); + const cStartY = Math.max(0, Math.floor(localChildMin[1])); + const cEndY = Math.min(childDataH, Math.ceil(localChildMax[1])); + const cStartZ = Math.max(0, Math.floor(localChildMin[2])); + const cEndZ = Math.min(childDataD, Math.ceil(localChildMax[2])); + + for (let cz = cStartZ; cz < cEndZ; ++cz) { + for (let cy = cStartY; cy < cEndY; ++cy) { + for (let cx = cStartX; cx < cEndX; ++cx) { + const srcIndex = + cz * (childDataW * childDataH) + cy * childDataW + cx; + + if (srcIndex >= 0 && srcIndex < childDataLength) { + const val = childChunkData[srcIndex]; + if (val !== undefined) { + sourceVoxels.push(BigInt(val)); + } + } + } + } + } + + if (sourceVoxels.length > 0) { + const mode = this._calculateMode(sourceVoxels); + // Use local coordinates for the parent chunk index + const parentIndex = pz * (parentW * parentH) + py * parentW + px; + indices.push(parentIndex); + values.push(mode); + } + } + } + } + + return { indices, values }; + } + + private _calculateMode(values: (bigint | number)[]): bigint { + if (values.length === 0) return VOXEL_EMPTY_VALUE; + const counts = new Map(); + let maxCount = 0; + let mode = 0n; + for (const v of values) { + const bigV = BigInt(v); + if (bigV === VOXEL_EMPTY_VALUE) continue; + const c = (counts.get(bigV) ?? 0) + 1; + counts.set(bigV, c); + if (c > maxCount) { + maxCount = c; + mode = bigV; + } else if (c === maxCount && bigV < mode) { + mode = bigV; + } + } + return mode; + } + private notifyHistoryChanged(): void { + this.rpc?.invoke(VOX_EDIT_HISTORY_UPDATE_RPC_ID, { + rpcId: this.rpcId, + undoCount: this.undoStack.length, + redoCount: this.redoStack.length, + }); + } + + private performUndoRedo( + sourceStack: EditAction[], + targetStack: EditAction[], + useOldValues: boolean, + actionDescription: "undo" | "redo", + ): Promise { + return this.runExclusive(() => + this.performUndoRedoLocked( + sourceStack, + targetStack, + useOldValues, + actionDescription, + ), + ); + } + + private async performUndoRedoLocked( + sourceStack: EditAction[], + targetStack: EditAction[], + useOldValues: boolean, + actionDescription: "undo" | "redo", + ): Promise { + await this.flushPendingLocked(); + + if (sourceStack.length === 0) { + throw new Error(`Nothing to ${actionDescription}.`); + } + + const action = sourceStack.pop()!; + + const chunksToReload = new Set(); + let success = true; + + for (const [voxKey, change] of action.changes.entries()) { + const parsedKey = parseVoxChunkKey(voxKey); + if (!parsedKey) continue; + const source = this.sources.get(parsedKey.lodIndex); + if (!source) continue; + + const valuesToApply = useOldValues ? change.oldValues : change.newValues; + try { + await source.applyEdits( + parsedKey.chunkKey, + change.indices, + valuesToApply, + ); + chunksToReload.add(voxKey); + } catch (e) { + success = false; + console.error( + `performUndoRedo: failed to apply edits for ${voxKey}`, + e, + ); + this.rpc?.invoke(VOX_EDIT_FAILURE_RPC_ID, { + rpcId: this.rpcId, + voxChunkKeys: [voxKey], + message: useOldValues ? "Undo failed." : "Redo failed.", + }); + break; + } + } + + if (success) { + targetStack.push(action); + } else { + // Known limitation: chunks reverted before the failure stay reverted + // while the action returns to the stack (no transactional rollback); + // the failure was already surfaced via VOX_EDIT_FAILURE. + sourceStack.push(action); + } + + if (chunksToReload.size > 0 && success) { + const keys = Array.from(chunksToReload); + const hasDownsampling = this.resolutions.size > 1; + if (hasDownsampling) { + for (const key of chunksToReload) { + this.enqueueDownsample(key); + } + } + // Rollback reload: the overlay keeps showing the undone strokes until + // real data arrives. Clearing immediately would reveal older data and + // make the stroke blink back when a pre-undo refetch lands. + this.callChunkReload(keys, false, undefined, undefined, true); + } + + this.notifyHistoryChanged(); + } + + public async undo(): Promise { + await this.performUndoRedo(this.undoStack, this.redoStack, true, "undo"); + } + + public async redo(): Promise { + await this.performUndoRedo(this.redoStack, this.undoStack, false, "redo"); + } + + // Resolves with the vox chunk keys whose stored data will contain the + // operation's overlay content once its edits flush ("covered" chunks). + async performOperation(operation: VoxelOperation): Promise { + switch (operation.type) { + case VoxelOperationType.BRUSH: + return this.performBrush(operation); + case VoxelOperationType.FLOOD_FILL: + return this.performFloodFill(operation); + default: + throw new Error( + `Unknown voxel operation type: ${(operation as any).type}`, + ); + } + } + + private async performBrush(op: BrushOperation): Promise { + const { centers, radius, value, shape, basis, filterValue, seq } = op; + const voxelSize = 1; // Hardcoded LOD 0 + const sourceIndex = 0; + const accessor = this.getAccessor(sourceIndex); + + const covered = new Set(); + // Voxels skipped because the store already holds the brush value are + // covered without being rewritten: clearing their overlay would flash + // pre-write data until the prior write's own reload lands. + const coveredSpec = this.sources.get(sourceIndex)?.spec; + const csizeX = coveredSpec?.chunkDataSize[0] ?? 1; + const csizeY = coveredSpec?.chunkDataSize[1] ?? 1; + const csizeZ = coveredSpec?.chunkDataSize[2] ?? 1; + let lastCX = -Infinity; + let lastCY = -Infinity; + let lastCZ = -Infinity; + const markCovered = (x: number, y: number, z: number) => { + const cx = Math.floor(x / csizeX); + const cy = Math.floor(y / csizeY); + const cz = Math.floor(z / csizeZ); + if (cx === lastCX && cy === lastCY && cz === lastCZ) return; + lastCX = cx; + lastCY = cy; + lastCZ = cz; + covered.add(`lod${sourceIndex}#${cx},${cy},${cz}`); + }; + + let r = Math.round(radius / voxelSize); + if (r <= 0) throw new Error(`Brush radius must be positive.`); + r -= 1; + const rr = r * r; + + // This capacity should ensure we never get out of bounds + const maxCapacity = Math.ceil((2 * r + 1) ** 3); + const voxelBuffer = new Int32Array(maxCapacity * 3); + + for (const center of centers) { + const cx = Math.round((center[0] ?? 0) / voxelSize); + const cy = Math.round((center[1] ?? 0) / voxelSize); + const cz = Math.round((center[2] ?? 0) / voxelSize); + + let voxelCount = 0; + const bufferEnqueue = (x: number, y: number, z: number) => { + // don;t need to check bounds as the capacity is assumed to be large enough + const base = voxelCount * 3; + voxelBuffer[base] = x; + voxelBuffer[base + 1] = y; + voxelBuffer[base + 2] = z; + voxelCount++; + }; + + const shouldSkip = this.brushCache.buildSkipper(value, shape, basis); + + const toAwait = new Set>(); + const pushIf = (x: number, y: number, z: number) => { + if (shouldSkip(x, y, z)) { + markCovered(x, y, z); + return; + } + if (filterValue == undefined) { + bufferEnqueue(x, y, z); + return; + } + toAwait.add( + accessor.getValue(x, y, z).then((v) => { + if (v == null) return; + if (v === value) { + markCovered(x, y, z); + return; + } + if (filterValue !== undefined && v !== filterValue) return; + bufferEnqueue(x, y, z); + }), + ); + }; + + if (shape !== BrushShape.DISK) { + for (let dz = -r; dz <= r; ++dz) { + for (let dy = -r; dy <= r; ++dy) { + for (let dx = -r; dx <= r; ++dx) { + if (dx * dx + dy * dy + dz * dz <= rr) + pushIf(cx + dx, cy + dy, cz + dz); + } + } + } + } else { + if (basis === undefined) + throw new Error("Brush shape requires a basis."); + const { u, v } = basis as { u: vec3; v: vec3 }; + const ux = u[0], + uy = u[1], + uz = u[2]; + const vx = v[0], + vy = v[1], + vz = v[2]; + + for (let j = -r; j <= r; ++j) { + const j2 = j * j; + const vPartX = vx * j; + const vPartY = vy * j; + const vPartZ = vz * j; + + for (let i = -r; i <= r; ++i) { + if (i * i + j2 <= rr) { + const px = Math.round(cx + ux * i + vPartX); + const py = Math.round(cy + uy * i + vPartY); + const pz = Math.round(cz + uz * i + vPartZ); + pushIf(px, py, pz); + } + } + } + } + + await Promise.all(toAwait); + this.brushCache.update({ x: cx, y: cy, z: cz }, r, value, shape, basis); + + if (voxelCount === 0) continue; + + let buffer: Int32Array = voxelBuffer; + let count = voxelCount; + if (basis && shape === BrushShape.DISK) { + const result = this.fillPlaneAliasingGaps( + voxelBuffer, + voxelCount, + basis, + center, + ); + buffer = result.buffer; + count = result.count; + } + for (const key of this.processBackendEdits( + buffer, + count, + value, + sourceIndex, + seq, + )) { + covered.add(key); + } + } + return Array.from(covered); + } + + private async performFloodFill(op: FloodFillOperation): Promise { + // The fill walk reads the store, which does not see pending edits: flush + // them first so a fill right after a stroke sees that stroke. The walk + // itself stays outside the chain (reads only; its edits flush later). + await this.runExclusive(() => this.flushPendingLocked()); + const { seed, value: fillValue, maxVoxels, basis, filterValue, seq } = op; + const morphological = op.morphological !== false; + const sourceIndex = 0; + const accessor = this.getAccessor(sourceIndex); + + const startVoxelLod = vec3.round(vec3.create(), seed as vec3); + const originalValue = await accessor.getValue( + startVoxelLod[0], + startVoxelLod[1], + startVoxelLod[2], + ); + + if (originalValue === null) return []; + if (filterValue !== undefined && originalValue !== filterValue) return []; + if (originalValue === fillValue) return []; + + const visited = new Set(); + const queue: [number, number][] = []; + let filledCount = 0; + const voxelBuffer = new Int32Array(maxVoxels * 3 + 200); // +200 since fillBorderRegion may exceed the maxVoxelCount without failure + + const map2dTo3d = (u: number, v: number): vec3 => { + const point = vec3.clone(startVoxelLod); + vec3.scaleAndAdd(point, point, basis.u as vec3, u); + vec3.scaleAndAdd(point, point, basis.v as vec3, v); + return vec3.round(vec3.create(), point); + }; + + const isFillable = async (p: vec3): Promise => { + const val = await accessor.getValue(p[0], p[1], p[2]); + if (val === null) return false; + if (originalValue === VOXEL_EMPTY_VALUE) return val === VOXEL_EMPTY_VALUE; + return val === originalValue; + }; + + const getCurrentThickness = (): number => { + if (!morphological) return 1; + let thickness = 1; + for (const threshold of this.morphologicalConfig.growthThresholds) { + if (filledCount >= threshold.count) { + thickness = Math.max(thickness, threshold.size); + } + } + return Math.min(thickness, this.morphologicalConfig.maxSize); + }; + + const hasThickEnoughChannel = async ( + u: number, + v: number, + nu: number, + nv: number, + requiredThickness: number, + ): Promise => { + if (requiredThickness <= 1) return true; + const halfThickness = Math.floor(requiredThickness / 2); + const du = nu - u; + const dv = nv - v; + const perpU = -dv; + const perpV = du; + + for (let offset = -halfThickness; offset <= halfThickness; ++offset) { + const testU = nu + perpU * offset; + const testV = nv + perpV * offset; + const pointToTest = map2dTo3d(testU, testV); + if (!(await isFillable(pointToTest))) return false; + } + return true; + }; + + const fillBorderRegion = async ( + startU: number, + startV: number, + requiredThickness: number, + ) => { + const subQueue: [number, number][] = []; + const halfSize = requiredThickness * 2; + const startKey = `${startU},${startV}`; + if (visited.has(startKey)) return; + + subQueue.push([startU, startV]); + visited.add(startKey); + + while (subQueue.length > 0) { + if (filledCount >= maxVoxels) return; + const [u, v] = subQueue.shift()!; + const currentPoint = map2dTo3d(u, v); + const base = filledCount * 3; + voxelBuffer[base] = currentPoint[0]; + voxelBuffer[base + 1] = currentPoint[1]; + voxelBuffer[base + 2] = currentPoint[2]; + filledCount++; + + const neighbors2d: [number, number][] = [ + [u + 1, v], + [u - 1, v], + [u, v + 1], + [u, v - 1], + ]; + for (const [nu, nv] of neighbors2d) { + const du = nu - startU; + const dv = nv - startV; + if (du * du + dv * dv > halfSize * halfSize) continue; + const neighborKey = `${nu},${nv}`; + if (visited.has(neighborKey)) continue; + if (await isFillable(map2dTo3d(nu, nv))) { + visited.add(neighborKey); + subQueue.push([nu, nv]); + } + } + } + }; + + queue.push([0, 0]); + visited.add("0,0"); + + while (queue.length > 0) { + if (filledCount >= maxVoxels) + throw new Error(`Flood fill failed: too many voxels filled.`); + const [u, v] = queue.shift()!; + const currentPoint = map2dTo3d(u, v); + const base = filledCount * 3; + voxelBuffer[base] = currentPoint[0]; + voxelBuffer[base + 1] = currentPoint[1]; + voxelBuffer[base + 2] = currentPoint[2]; + filledCount++; + + const requiredThickness = getCurrentThickness(); + const neighbors2d: [number, number][] = [ + [u + 1, v], + [u - 1, v], + [u, v + 1], + [u, v - 1], + ]; + + for (const [nu, nv] of neighbors2d) { + const k = `${nu},${nv}`; + if (visited.has(k)) continue; + const neighborPoint = map2dTo3d(nu, nv); + if (await isFillable(neighborPoint)) { + if (await hasThickEnoughChannel(u, v, nu, nv, requiredThickness)) { + visited.add(k); + queue.push([nu, nv]); + } else { + await fillBorderRegion(nu, nv, requiredThickness); + } + } + } + } + + const result = this.fillPlaneAliasingGaps( + voxelBuffer, + filledCount, + basis, + seed, + ); + return this.processBackendEdits( + result.buffer, + result.count, + fillValue, + sourceIndex, + seq, + ); + } + + private fillPlaneAliasingGaps( + inputBuffer: Int32Array, + inputCount: number, + basis: { u: Float32Array; v: Float32Array }, + center: Float32Array, + ): { buffer: Int32Array; count: number } { + const u = basis.u as vec3; + const v = basis.v as vec3; + const normal = vec3.create(); + vec3.cross(normal, u, v); + vec3.normalize(normal, normal); + + const SKIP_THRESHOLD = 0.99; + if ( + Math.abs(normal[0]) > SKIP_THRESHOLD || + Math.abs(normal[1]) > SKIP_THRESHOLD || + Math.abs(normal[2]) > SKIP_THRESHOLD + ) { + return { buffer: inputBuffer, count: inputCount }; + } + + const d = -vec3.dot(normal, center as vec3); + const DISTANCE_THRESHOLD = + Math.abs(normal[0]) + Math.abs(normal[1]) + Math.abs(normal[2]) + 1e-5; + + const voxelSet = new Set(); + + let outputBuffer = inputBuffer; + let outputCount = inputCount; + if (inputCount * 6 > inputBuffer.length) { + const newBuf = new Int32Array(inputCount * 6); + newBuf.set(inputBuffer.subarray(0, inputCount * 3)); + outputBuffer = newBuf; + } + + for (let i = 0; i < inputCount; i++) { + const base = i * 3; + const px = outputBuffer[base]; + const py = outputBuffer[base + 1]; + const pz = outputBuffer[base + 2]; + voxelSet.add(`${px},${py},${pz}`); + } + + for (let i = 0; i < inputCount; i++) { + const base = i * 3; + const px = inputBuffer[base]; + const py = inputBuffer[base + 1]; + const pz = inputBuffer[base + 2]; + + for (const [ox, oy, oz] of OFFSETS_26_CONNECTED_BACKEND) { + const nx = px + ox; + const ny = py + oy; + const nz = pz + oz; + + const dist = Math.abs( + normal[0] * nx + normal[1] * ny + normal[2] * nz + d, + ); + if (dist <= DISTANCE_THRESHOLD) { + const key = `${nx},${ny},${nz}`; + if (!voxelSet.has(key)) { + voxelSet.add(key); + + if (outputCount * 3 + 3 > outputBuffer.length) { + const newBuf = new Int32Array(outputBuffer.length * 2); + newBuf.set(outputBuffer); + outputBuffer = newBuf; + } + const outBase = outputCount * 3; + outputBuffer[outBase] = nx; + outputBuffer[outBase + 1] = ny; + outputBuffer[outBase + 2] = nz; + outputCount++; + } + } + } + } + return { buffer: outputBuffer, count: outputCount }; + } + + private processBackendEdits( + voxelBuffer: Int32Array, + voxelCount: number, + value: bigint, + lodIndex: number, + seq?: number, + ): string[] { + const source = this.sources.get(lodIndex); + if (!source) return []; + + const { chunkDataSize } = source.spec; + const indicesByVoxKey = new Map(); + + let lastGridX = -Infinity; + let lastGridY = -Infinity; + let lastGridZ = -Infinity; + let currentIndicesList: number[] | undefined; + + const sizeX = chunkDataSize[0]; + const sizeY = chunkDataSize[1]; + const sizeZ = chunkDataSize[2]; + const strideY = sizeX; + const strideZ = sizeX * sizeY; + + for (let i = 0; i < voxelCount; i++) { + const base = i * 3; + const vx = voxelBuffer[base]; + const vy = voxelBuffer[base + 1]; + const vz = voxelBuffer[base + 2]; + + const cx = Math.floor(vx / sizeX); + const cy = Math.floor(vy / sizeY); + const cz = Math.floor(vz / sizeZ); + + if (cx !== lastGridX || cy !== lastGridY || cz !== lastGridZ) { + lastGridX = cx; + lastGridY = cy; + lastGridZ = cz; + + const voxKey = `lod${lodIndex}#${cx},${cy},${cz}`; + + currentIndicesList = indicesByVoxKey.get(voxKey); + if (!currentIndicesList) { + currentIndicesList = []; + indicesByVoxKey.set(voxKey, currentIndicesList); + } + } + + const lx = Math.floor(vx - sizeX * cx); + const ly = Math.floor(vy - sizeY * cy); + const lz = Math.floor(vz - sizeZ * cz); + + const index = lz * strideZ + ly * strideY + lx; + + currentIndicesList!.push(index); + } + + const backendEdits = []; + for (const [voxKey, indices] of indicesByVoxKey.entries()) { + backendEdits.push({ key: voxKey, indices, value, seq }); + } + this.commitVoxels(backendEdits); + return Array.from(indicesByVoxKey.keys()); + } +} + +registerRPC(VOX_EDIT_COMMIT_VOXELS_RPC_ID, function (x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + obj.commitVoxels(Array.isArray(x.edits) ? x.edits : []); +}); + +registerPromiseRPC(VOX_EDIT_UNDO_RPC_ID, async function (this: RPC, x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + await obj.undo(); + return { value: undefined }; +}); + +registerPromiseRPC(VOX_EDIT_REDO_RPC_ID, async function (this: RPC, x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + await obj.redo(); + return { value: undefined }; +}); + +registerPromiseRPC( + VOX_EDIT_OPERATION_RPC_ID, + async function (this: RPC, x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + const coveredVoxKeys = await obj.performOperation(x.operation); + return { value: coveredVoxKeys }; + }, +); diff --git a/src/voxel_annotation/base.ts b/src/voxel_annotation/base.ts new file mode 100644 index 0000000000..7d0093bfe9 --- /dev/null +++ b/src/voxel_annotation/base.ts @@ -0,0 +1,216 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { MultiscaleVolumeChunkSource } from "#src/sliceview/volume/frontend.js"; +import { vec3 } from "#src/util/geom.js"; +import type { VoxelPreviewMultiscaleSource } from "#src/voxel_annotation/preview_multiscale_chunk_source.js"; +import type { RPC } from "#src/worker_rpc.js"; + +export const VOX_RELOAD_CHUNKS_RPC_ID = "vox.chunk.reload"; +export const VOX_EDIT_BACKEND_RPC_ID = "vox.EditBackend"; +export const VOX_EDIT_COMMIT_VOXELS_RPC_ID = "vox.edit.commitVoxels"; +export const VOX_EDIT_FAILURE_RPC_ID = "vox.edit.failure"; +export const VOX_EDIT_UNDO_RPC_ID = "vox.edit.undo"; +export const VOX_EDIT_REDO_RPC_ID = "vox.edit.redo"; +export const VOX_EDIT_HISTORY_UPDATE_RPC_ID = "vox.edit.historyUpdate"; + +export const VOX_EDIT_OPERATION_RPC_ID = "vox.edit.operation"; + +export const FLOODFILL_MAX_POSSIBLE_VOXELS = 1000000; +export const FLOODFILL_MIN_POSSIBLE_VOXELS = 1000; + +export enum VoxelOperationType { + BRUSH = 0, + FLOOD_FILL = 1, +} + +export interface VoxelOperationBase { + type: VoxelOperationType; + // Frontend dispatch sequence number. Echoed back per written chunk (as the + // max seq the write covers) in reload messages, so the frontend clears an + // overlay chunk only once the stored data covers every dispatched stroke + // that touched it. + seq?: number; +} + +export interface BrushOperation extends VoxelOperationBase { + type: VoxelOperationType.BRUSH; + centers: Float32Array[]; + radius: number; + value: bigint; + shape: BrushShape; + basis: { u: Float32Array; v: Float32Array }; + filterValue?: bigint; +} + +export interface FloodFillOperation extends VoxelOperationBase { + type: VoxelOperationType.FLOOD_FILL; + seed: Float32Array; + value: bigint; + maxVoxels: number; + basis: { u: Float32Array; v: Float32Array }; + filterValue?: bigint; + // Defaults to true; false disables the channel-thickness gating so the fill + // is a plain 4-connected walk. + morphological?: boolean; +} + +export type VoxelOperation = BrushOperation | FloodFillOperation; + +export const BRUSH_TOOL_ID = "vox-brush"; +export const FLOODFILL_TOOL_ID = "vox-flood-fill"; +export const VALUE_PICKER_TOOL_ID = "vox-value-picker"; + +export const VOXEL_EMPTY_VALUE = 0n; + +// Special value used to indicate to the optimistic renderer that a voxel has been erased +export const SEG_ERASE_SENTINEL = ~1n; + +export const VOXEL_EDIT_STAMINA = { + pendingEdits: (voxelCount: number) => Math.round(voxelCount * 0.0004), + downsamplingJobs: (count: number, downsamplingSteps: number) => + Math.round(count * 10 * downsamplingSteps), + brush: (shape: BrushShape, radius: number, hasFiltering: boolean) => { + const FILTERING = hasFiltering ? 7 : 1; + if (shape === BrushShape.DISK) { + return Math.round(0.00035 * Math.pow(radius, 2) * FILTERING); + } else { + return Math.round(0.00012 * Math.pow(radius, 3) * FILTERING); + } + }, + floodFill: (maxVoxels: number) => Math.round(maxVoxels * 0.005), + undoRedo: () => 20, +}; +export const MAX_VOXEL_EDIT_STAMINA = 10000; + +export type VoxelValueGetter = (isPreview: boolean) => bigint; + +export interface VoxelLayerResolution { + lodIndex: number; + transform: number[]; + chunkSize: number[]; + sourceRpc: number; +} + +export type VoxelChangeValues = Uint32Array | BigUint64Array; + +export interface VoxelChange { + indices: Uint32Array; + oldValues: VoxelChangeValues; + newValues: VoxelChangeValues; +} + +export interface EditAction { + changes: Map; + timestamp: number; + description: string; +} + +export function makeVoxChunkKey(chunkKey: string, lodIndex: number) { + return `lod${lodIndex}#${chunkKey}`; +} + +export function makeChunkKey(x: number, y: number, z: number) { + return `${x},${y},${z}`; +} + +export function parseVoxChunkKey(key: string) { + const parts = [ + Number(key.split("#")[0].substring(3)), + ...key.split("#")[1].split(",").map(Number), + ]; + if (parts.length !== 4 || parts.some(isNaN)) { + console.warn(`Invalid chunk key format: ${key}`); + return null; + } + return { + lodIndex: parts[0], + x: parts[1], + y: parts[2], + z: parts[3], + chunkKey: key.split("#")[1], + }; +} + +export function getBasisFromNormal(n: vec3) { + const u = vec3.create(); + const tempVec = + Math.abs(vec3.dot(n, vec3.fromValues(1, 0, 0))) < 0.9 + ? vec3.fromValues(1, 0, 0) + : vec3.fromValues(0, 1, 0); + vec3.cross(u, tempVec, n); + vec3.normalize(u, u); + const v = vec3.cross(vec3.create(), n, u); + vec3.normalize(v, v); + return { u, v }; +} + +export enum BrushShape { + DISK = 0, + SPHERE = 1, +} + +const sphereRowRangesKernelCache = new Map(); + +export function getSphereRowRangesKernel(radius: number): Int16Array { + let kernel = sphereRowRangesKernelCache.get(radius); + if (kernel !== undefined) return kernel; + if (!Number.isInteger(radius) || radius < 0) { + throw new Error(`Invalid sphere radius: ${radius}`); + } + const rr = radius * radius; + const ranges: number[] = []; + for (let dz = -radius; dz <= radius; ++dz) { + const dz2 = dz * dz; + for (let dy = -radius; dy <= radius; ++dy) { + const remaining = rr - dz2 - dy * dy; + if (remaining < 0) continue; + const maxDx = Math.floor(Math.sqrt(remaining)); + ranges.push(dy, dz, -maxDx, maxDx + 1); + } + } + kernel = Int16Array.from(ranges); + sphereRowRangesKernelCache.set(radius, kernel); + return kernel; +} + +const diskStencilKernelCache = new Map(); + +export function getDiskStencilKernel(radius: number): Int16Array { + let kernel = diskStencilKernelCache.get(radius); + if (kernel !== undefined) return kernel; + if (!Number.isInteger(radius) || radius < 0) { + throw new Error(`Invalid disk radius: ${radius}`); + } + const rr = radius * radius; + const pairs: number[] = []; + for (let j = -radius; j <= radius; ++j) { + for (let i = -radius; i <= radius; ++i) { + if (i * i + j * j <= rr) { + pairs.push(i, j); + } + } + } + kernel = Int16Array.from(pairs); + diskStencilKernelCache.set(radius, kernel); + return kernel; +} + +export interface VoxelEditControllerHost { + primarySource: MultiscaleVolumeChunkSource; + previewSource?: VoxelPreviewMultiscaleSource; + rpc: RPC; +} diff --git a/src/voxel_annotation/frontend.spec.ts b/src/voxel_annotation/frontend.spec.ts new file mode 100644 index 0000000000..a0649dc738 --- /dev/null +++ b/src/voxel_annotation/frontend.spec.ts @@ -0,0 +1,288 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { ChunkState } from "#src/chunk_manager/base.js"; +import { NullarySignal } from "#src/util/signal.js"; +import { makeVoxChunkKey } from "#src/voxel_annotation/base.js"; +import { VoxelEditController } from "#src/voxel_annotation/frontend.js"; +import type { RPC } from "#src/worker_rpc.js"; + +const mockRpc = { + get: vi.fn(), + invoke: vi.fn(), + newId: () => 0, + register: vi.fn(), + set: vi.fn(), + delete: vi.fn(), +} as unknown as RPC; + +// Real chunk source mock: `fireFreshChunk` replaces the chunk with a new +// object in GPU_MEMORY state, like applyChunkUpdate does for an `update.new` +// followed by its GPU promotion. +function createRealSourceMock() { + return { + rpcId: 1, + spec: { chunkDataSize: new Uint32Array([2, 2, 2]) }, + chunks: new Map(), + invalidateChunks: vi.fn(), + fireFreshChunk(key: string, state = ChunkState.GPU_MEMORY) { + this.chunks.set(key, { state }); + }, + }; +} + +// Overlay source mock mirroring InMemoryVolumeChunkSource's stroke-seq tags: +// invalidateChunks purges the tag, like the real deleteChunk does. +function createOverlaySourceMock() { + const overlaySeqs = new Map(); + return { + overlaySeqs, + setOverlaySeq: (key: string, seq: number) => overlaySeqs.set(key, seq), + getOverlaySeq: (key: string) => overlaySeqs.get(key) ?? 0, + keysWithOverlaySeq: (seq: number) => + [...overlaySeqs.entries()].filter(([, s]) => s === seq).map(([k]) => k), + clearOverlaySeq: (key: string) => overlaySeqs.delete(key), + invalidateChunks: vi.fn((keys: string[]) => { + for (const key of keys) overlaySeqs.delete(key); + }), + }; +} + +describe("VoxelEditController.callChunkReload: overlay swap observation", () => { + let realSources: ReturnType[]; + let overlaySources: ReturnType[]; + let visibleChunksChanged: NullarySignal; + let controller: VoxelEditController; + + beforeEach(() => { + vi.clearAllMocks(); + realSources = [createRealSourceMock(), createRealSourceMock()]; + overlaySources = [createOverlaySourceMock(), createOverlaySourceMock()]; + visibleChunksChanged = new NullarySignal(); + const makeMultiscale = (sources: unknown[]) => ({ + rank: 3, + chunkManager: { chunkQueueManager: { visibleChunksChanged } }, + getSources: () => [ + sources.map((chunkSource) => ({ + chunkSource, + chunkToMultiscaleTransform: Float32Array.of( + 1, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 1, + ), + })), + ], + }); + const host = { + rpc: mockRpc, + primarySource: makeMultiscale(realSources), + previewSource: makeMultiscale(overlaySources), + }; + controller = new VoxelEditController(host as any); + }); + + it("allocates monotonically increasing stroke seqs", () => { + expect(controller.beginStroke()).toBe(1); + expect(controller.beginStroke()).toBe(2); + expect(controller.beginStroke()).toBe(3); + }); + + it("clears the overlay once the refetched chunk replaces the stale one on the GPU", () => { + const seq = controller.beginStroke(); + overlaySources[0].setOverlaySeq("0,0,0", seq); + // The stale chunk is on display when the reload arrives. + realSources[0].fireFreshChunk("0,0,0"); + + const voxKey = makeVoxChunkKey("0,0,0", 0); + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: seq }); + + expect(realSources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"], { + lazy: true, + }); + + // Signal fires while the stale chunk is still displayed: no clear. + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + + // The refetched chunk (a new object) reaches the GPU. + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"]); + }); + + it("waits until the refetched chunk actually reaches the GPU", () => { + const seq = controller.beginStroke(); + overlaySources[0].setOverlaySeq("0,0,0", seq); + + const voxKey = makeVoxChunkKey("0,0,0", 0); + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: seq }); + + // Refetched data arrived in system memory only: keep the overlay. + realSources[0].fireFreshChunk("0,0,0", ChunkState.SYSTEM_MEMORY); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + + // Promotion to the GPU resolves the swap. + realSources[0].chunks.get("0,0,0")!.state = ChunkState.GPU_MEMORY; + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"]); + }); + + it("skips the clear when the write does not cover the last stroke", () => { + controller.beginStroke(); // seq 1, written + const seq2 = controller.beginStroke(); // seq 2, dispatched but unwritten + overlaySources[0].setOverlaySeq("0,0,0", seq2); + + // The reload for stroke 1's flush only covers seq 1 < 2. + const voxKey = makeVoxChunkKey("0,0,0", 0); + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: 1 }); + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + + // Stroke 2's own flush covers seq 2: its reload performs the clear. + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: 2 }); + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"]); + }); + + it("reads the tag at swap time: a stroke touching the chunk after arming blocks the clear", () => { + const seq1 = controller.beginStroke(); + overlaySources[0].setOverlaySeq("0,0,0", seq1); + + const voxKey = makeVoxChunkKey("0,0,0", 0); + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: seq1 }); + + // A new stroke's preview touches the chunk before the refetch lands: + // the arriving data cannot contain it. + const seq2 = controller.beginStroke(); + overlaySources[0].setOverlaySeq("0,0,0", seq2); + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + }); + + it("skips the clear when a reload carries no coverage but a stroke tagged the chunk", () => { + // e.g. an undo/redo reload: without echoed coverage it must not clear an + // overlay that a dispatched-but-unwritten stroke still owns. + overlaySources[0].setOverlaySeq("0,0,0", controller.beginStroke()); + + controller.callChunkReload([makeVoxChunkKey("0,0,0", 0)], false); + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + }); + + it("clears without coverage info when no stroke ever tagged the chunk", () => { + controller.callChunkReload([makeVoxChunkKey("0,0,0", 0)], false); + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"]); + }); + + it("guards downsampled-parent reloads with the origin chunk's coverage", () => { + controller.beginStroke(); // seq 1, written + const seq2 = controller.beginStroke(); // seq 2, unwritten + overlaySources[0].setOverlaySeq("1,2,3", seq2); + + const parentKey = makeVoxChunkKey("0,0,0", 1); + const originKey = makeVoxChunkKey("1,2,3", 0); + + // Cascade reload from stroke 1's flush: parent data only covers seq 1. + controller.callChunkReload( + [parentKey], + false, + { [parentKey]: originKey }, + { [parentKey]: 1 }, + ); + expect(realSources[1].invalidateChunks).toHaveBeenCalledWith(["0,0,0"], { + lazy: true, + }); + realSources[1].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + + // Cascade re-run after stroke 2's flush covers seq 2. + controller.callChunkReload( + [parentKey], + false, + { [parentKey]: originKey }, + { [parentKey]: 2 }, + ); + realSources[1].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["1,2,3"]); + }); + + it("a rollback reload clears the overlay on first arrival regardless of tags", () => { + // An undone stroke's tag can never be covered by a future write; the + // rollback purges it so the swap resolves unconditionally. + overlaySources[0].setOverlaySeq("0,0,0", controller.beginStroke()); + + const voxKey = makeVoxChunkKey("0,0,0", 0); + controller.callChunkReload([voxKey], false, undefined, undefined, true); + + // Not cleared before data arrives: the overlay keeps showing the stroke. + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"]); + }); + + it("a newer reload overwrites the pending swap for the same chunk", () => { + // Two reloads arm before any refetch arrives (covered 1 then 2, tag at + // 2): only the newest entry remains, so the single arrival clears once, + // with the newest coverage. + controller.beginStroke(); + const seq2 = controller.beginStroke(); + overlaySources[0].setOverlaySeq("0,0,0", seq2); + + const voxKey = makeVoxChunkKey("0,0,0", 0); + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: 1 }); + controller.callChunkReload([voxKey], false, undefined, { [voxKey]: 2 }); + + realSources[0].fireFreshChunk("0,0,0"); + visibleChunksChanged.dispatch(); + visibleChunksChanged.dispatch(); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledTimes(1); + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledWith(["0,0,0"]); + }); + + it("rollbackStroke drops exactly the overlay chunks tagged by that stroke", () => { + const seq1 = controller.beginStroke(); + overlaySources[0].setOverlaySeq("0,0,0", seq1); + const seq2 = controller.beginStroke(); + overlaySources[0].setOverlaySeq("1,0,0", seq2); + overlaySources[0].setOverlaySeq("2,0,0", seq2); + + controller.rollbackStroke(seq2); + + expect(overlaySources[0].invalidateChunks).toHaveBeenCalledTimes(1); + const [rolledBack] = overlaySources[0].invalidateChunks.mock.calls[0]; + expect([...rolledBack].sort()).toEqual(["1,0,0", "2,0,0"]); + // The other stroke's chunk is untouched and still tagged. + expect(overlaySources[0].getOverlaySeq("0,0,0")).toBe(seq1); + }); + + it("rollbackStroke with no tagged chunks is a no-op", () => { + const seq = controller.beginStroke(); + controller.rollbackStroke(seq); + expect(overlaySources[0].invalidateChunks).not.toHaveBeenCalled(); + }); +}); diff --git a/src/voxel_annotation/frontend.ts b/src/voxel_annotation/frontend.ts new file mode 100644 index 0000000000..04d3a82c3e --- /dev/null +++ b/src/voxel_annotation/frontend.ts @@ -0,0 +1,845 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { ChunkState } from "#src/chunk_manager/base.js"; +import type { ChunkChannelAccessParameters } from "#src/render_coordinate_transform.js"; +import { SharedWatchableValue } from "#src/shared_watchable_value.js"; +import type { + InMemoryVolumeChunkSource, + LocalVolumeEdit, + VolumeChunkSource, +} from "#src/sliceview/volume/frontend.js"; +import { StatusMessage } from "#src/status.js"; +import { WatchableValue } from "#src/trackable_value.js"; +import type { vec3 } from "#src/util/geom.js"; +import type { + VoxelEditControllerHost, + VoxelLayerResolution, + VoxelOperation, + VoxelValueGetter, +} from "#src/voxel_annotation/base.js"; +import { + BrushShape, + getDiskStencilKernel, + getSphereRowRangesKernel, + parseVoxChunkKey, + VOX_EDIT_BACKEND_RPC_ID, + VOX_EDIT_FAILURE_RPC_ID, + VOX_EDIT_HISTORY_UPDATE_RPC_ID, + VOX_EDIT_OPERATION_RPC_ID, + VOX_EDIT_REDO_RPC_ID, + VOX_EDIT_UNDO_RPC_ID, + VOX_RELOAD_CHUNKS_RPC_ID, + VoxelOperationType, +} from "#src/voxel_annotation/base.js"; +import { + registerRPC, + registerSharedObjectOwner, + SharedObject, +} from "#src/worker_rpc.js"; + +@registerSharedObjectOwner(VOX_EDIT_BACKEND_RPC_ID) +export class VoxelEditController extends SharedObject { + public undoCount = new WatchableValue(0); + public redoCount = new WatchableValue(0); + public pendingOpCount: SharedWatchableValue; + + // Monotonic counter identifying each stroke/operation dispatched to the + // backend. + private dispatchSeq = 0; + + // Overlay swaps awaiting their refetched real chunk, keyed by real vox + // chunk key and resolved by observing visibleChunksChanged: the real + // chunk's object identity distinguishes the lazily kept stale chunk + // (recorded at arming) from the refetched one, since a `new` chunk update + // always builds a fresh Chunk object. A newer reload for the same chunk + // overwrites its entry, so the map never grows beyond the set of chunks + // awaiting a swap. + private pendingOverlaySwaps = new Map< + string, + { + source: VolumeChunkSource; + chunkKey: string; + staleChunk: unknown; + overlaySource: InMemoryVolumeChunkSource; + overlayChunkKey: string; + coveredSeq: number; + } + >(); + + private processPendingOverlaySwaps(): void { + if (this.pendingOverlaySwaps.size === 0) return; + for (const [voxKey, swap] of this.pendingOverlaySwaps) { + const chunk = swap.source.chunks.get(swap.chunkKey); + // Still the stale chunk (or gone): the refetch has not landed yet. + if (chunk === undefined || chunk === swap.staleChunk) continue; + // Refetched but not yet displayed: keep waiting. + if (chunk.state !== ChunkState.GPU_MEMORY) continue; + this.pendingOverlaySwaps.delete(voxKey); + // The overlay tag is read now, at swap time: a stroke that touched the + // chunk since arming raised it above coveredSeq, keeping the overlay on + // screen; the covering write's own reload re-arms the swap. + if ( + swap.coveredSeq < swap.overlaySource.getOverlaySeq(swap.overlayChunkKey) + ) { + continue; + } + swap.overlaySource.invalidateChunks([swap.overlayChunkKey]); + } + } + + // Allocates a stroke's seq before its first preview. Previews tag the + // overlay chunks they touch with it and the dispatch carries the same + // value, so what the overlay shows and what the write covers cannot + // diverge. + beginStroke(): number { + return ++this.dispatchSeq; + } + + private getOverlaySource(): InMemoryVolumeChunkSource | undefined { + return this.host.previewSource?.getSources( + this.getIdentitySliceViewSourceOptions(), + )[0]?.[0]?.chunkSource as InMemoryVolumeChunkSource | undefined; + } + + // For a stroke whose edits will never be written: stamina/permission + // refusal, empty stroke, dispatch failure. + rollbackStroke(seq: number): void { + this.reconcileStroke(seq, []); + } + + // Drops the overlay chunks a stroke tagged that its backend write does not + // cover: nothing will be written there, so no reload would ever clear them + // and the real data beneath is already correct. Chunks re-tagged by a newer + // stroke no longer match `seq` and are left untouched. + reconcileStroke(seq: number, coveredVoxKeys: string[]): void { + const overlaySource = this.getOverlaySource(); + if (overlaySource === undefined) return; + const tagged = overlaySource.keysWithOverlaySeq(seq); + if (tagged.length === 0) return; + let stale = tagged; + if (coveredVoxKeys.length > 0) { + const covered = new Set(); + for (const voxKey of coveredVoxKeys) { + const parsed = parseVoxChunkKey(voxKey); + if (parsed !== null && parsed.lodIndex === 0) { + covered.add(parsed.chunkKey); + } + } + stale = tagged.filter((key) => !covered.has(key)); + } + if (stale.length > 0) overlaySource.invalidateChunks(stale); + } + + constructor(private host: VoxelEditControllerHost) { + super(); + const rpc = this.host.rpc; + if (!rpc) { + throw new Error( + "VoxelEditController: Missing RPC from multiscale chunk manager.", + ); + } + + const sourcesByScale = this.host.primarySource.getSources( + this.getIdentitySliceViewSourceOptions(), + ); + const sources = sourcesByScale[0]; + if (!sources) { + throw new Error( + "VoxelEditController: Could not retrieve sources from multiscale object.", + ); + } + + const resolutions: VoxelLayerResolution[] = []; + + for (let i = 0; i < sources.length; ++i) { + const source = sources[i]!.chunkSource; + const rpcId = source.rpcId; + if (rpcId == null) { + throw new Error( + `VoxelEditController: Source at LOD index ${i} has null rpcId during initialization.`, + ); + } + resolutions.push({ + lodIndex: i, + transform: Array.from(sources[i]!.chunkToMultiscaleTransform), + chunkSize: Array.from(source.spec.chunkDataSize), + sourceRpc: rpcId, + }); + } + + this.pendingOpCount = this.registerDisposer( + SharedWatchableValue.make(this.host.rpc, 0), + ); + + this.initializeCounterpart(rpc, { + resolutions, + pendingOpCount: this.pendingOpCount.rpcId, + }); + + // Pending overlay swaps are resolved by observing chunk changes rather + // than by hooks inside the chunk manager: the signal fires after every + // applied batch of chunk updates, and the check below is a cheap scan of + // the (small) pending map. + this.registerDisposer( + this.host.primarySource.chunkManager.chunkQueueManager.visibleChunksChanged.add( + () => this.processPendingOverlaySwaps(), + ), + ); + } + + private async dispatchOperation( + operation: VoxelOperation, + ): Promise { + if (!this.rpc) throw new Error("RPC unavailable"); + const coveredVoxKeys = await this.rpc.promiseInvoke( + VOX_EDIT_OPERATION_RPC_ID, + { + rpcId: this.rpcId, + operation, + }, + ); + return Array.isArray(coveredVoxKeys) ? coveredVoxKeys : []; + } + + readonly singleChannelAccess: ChunkChannelAccessParameters = { + numChannels: 1, + channelSpaceShape: new Uint32Array([]), + chunkChannelDimensionIndices: [], + chunkChannelCoordinates: new Uint32Array([0]), + }; + + private getIdentitySliceViewSourceOptions() { + const rank = this.host.primarySource.rank as number | undefined; + if (!Number.isInteger(rank) || (rank as number) <= 0) { + throw new Error("VoxelEditController: Invalid multiscale rank."); + } + const r = rank as number; + const displayRank = r; + const multiscaleToViewTransform = new Float32Array(displayRank * r); + for (let chunkDim = 0; chunkDim < r; ++chunkDim) { + for (let displayDim = 0; displayDim < displayRank; ++displayDim) { + multiscaleToViewTransform[displayRank * chunkDim + displayDim] = + chunkDim === displayDim ? 1 : 0; + } + } + return { + displayRank, + multiscaleToViewTransform, + modelChannelDimensionIndices: [], + } as const; + } + + async applyBrushPreview( + points: Float32Array[], + radiusCanonical: number, + valueGetter: VoxelValueGetter, + shape: BrushShape, + basis: { u: Float32Array; v: Float32Array }, + seq: number, + filterValue?: bigint, + ) { + if (!this.host.previewSource) return; + + const voxelSize = 1; // Assuming LOD 0 + let r = Math.round(radiusCanonical / voxelSize); + if (r <= 0) { + throw new Error("Brush radius must be positive."); + } + r -= 1; + + const previewSource = this.host.previewSource.getSources( + this.getIdentitySliceViewSourceOptions(), + )[0][0].chunkSource as InMemoryVolumeChunkSource; + + const { chunkDataSize } = previewSource.spec; + const sizeX = chunkDataSize[0]; + const sizeY = chunkDataSize[1]; + const sizeZ = chunkDataSize[2]; + const strideY = sizeX; + const strideZ = sizeX * sizeY; + + // Deduplicate centers that round to the same voxel position + const centers: [number, number, number][] = []; + if (points.length <= 1) { + const c = points[0]!; + centers.push([ + Math.round((c[0] ?? 0) / voxelSize), + Math.round((c[1] ?? 0) / voxelSize), + Math.round((c[2] ?? 0) / voxelSize), + ]); + } else { + const seen = new Set(); + for (const c of points) { + const cx = Math.round((c[0] ?? 0) / voxelSize); + const cy = Math.round((c[1] ?? 0) / voxelSize); + const cz = Math.round((c[2] ?? 0) / voxelSize); + const key = `${cx},${cy},${cz}`; + if (seen.has(key)) continue; + seen.add(key); + centers.push([cx, cy, cz]); + } + } + + // filterValue setup (slow path only) + let baseSource: VolumeChunkSource | undefined; + const tempPos = new Float32Array(3); + if (filterValue !== undefined) { + const sourcesByScale = this.host.primarySource.getSources( + this.getIdentitySliceViewSourceOptions(), + ); + baseSource = sourcesByScale[0][0].chunkSource as VolumeChunkSource; + } + + const passesFilter = (x: number, y: number, z: number): boolean => { + tempPos[0] = x; + tempPos[1] = y; + tempPos[2] = z; + const val = baseSource!.getValueAt(tempPos, this.singleChannelAccess); + if (val == null) return true; + const bigVal = typeof val === "bigint" ? val : BigInt(val); + return bigVal === filterValue; + }; + + type EditEntry = LocalVolumeEdit & { + indices: number[]; + indexRanges: number[]; + }; + const edits = new Map(); + const previewValue = valueGetter(true); + + const getOrCreateEdit = ( + chunkX: number, + chunkY: number, + chunkZ: number, + ): EditEntry => { + const key = `${chunkX},${chunkY},${chunkZ}`; + let entry = edits.get(key); + if (entry !== undefined) return entry; + entry = { + indices: [], + indexRanges: [], + value: previewValue, + chunkGridPosition: Float32Array.of(chunkX, chunkY, chunkZ), + }; + edits.set(key, entry); + return entry; + }; + + // Append a single voxel to the correct chunk edit (used by slow paths) + const appendVoxel = (x: number, y: number, z: number) => { + const chunkX = Math.floor(x / sizeX); + const chunkY = Math.floor(y / sizeY); + const chunkZ = Math.floor(z / sizeZ); + const lx = x - chunkX * sizeX; + const ly = y - chunkY * sizeY; + const lz = z - chunkZ * sizeZ; + getOrCreateEdit(chunkX, chunkY, chunkZ).indices.push( + lz * strideZ + ly * strideY + lx, + ); + }; + + // Append a contiguous x-run [xStart, xEndExcl) at (y, z), splitting at chunk boundaries. + // Merges adjacent ranges within the same chunk entry. + const appendRange = ( + xStart: number, + xEndExcl: number, + y: number, + z: number, + ) => { + const chunkY = Math.floor(y / sizeY); + const chunkZ = Math.floor(z / sizeZ); + const ly = y - chunkY * sizeY; + const lz = z - chunkZ * sizeZ; + const baseIndex = lz * strideZ + ly * strideY; + + let currentX = xStart; + while (currentX < xEndExcl) { + const chunkX = Math.floor(currentX / sizeX); + const segEnd = Math.min(xEndExcl, (chunkX + 1) * sizeX); + const startIndex = baseIndex + (currentX - chunkX * sizeX); + const length = segEnd - currentX; + + const entry = getOrCreateEdit(chunkX, chunkY, chunkZ); + if (length === 1) { + entry.indices.push(startIndex); + } else { + const ranges = entry.indexRanges; + const last = ranges.length - 1; + if (last >= 1 && ranges[last - 1]! + ranges[last]! === startIndex) { + ranges[last] = ranges[last]! + length; + } else { + ranges.push(startIndex, length); + } + } + currentX = segEnd; + } + }; + + if (shape === BrushShape.SPHERE) { + const kernel = getSphereRowRangesKernel(r); + if (filterValue === undefined) { + // Fast path: contiguous row ranges → TypedArray.fill() + for (const [cx, cy, cz] of centers) { + for (let i = 0; i < kernel.length; i += 4) { + appendRange( + cx + kernel[i + 2]!, + cx + kernel[i + 3]!, + cy + kernel[i]!, + cz + kernel[i + 1]!, + ); + } + } + } else { + // Slow path: per-voxel filter check + for (const [cx, cy, cz] of centers) { + for (let i = 0; i < kernel.length; i += 4) { + const dy = kernel[i]!; + const dz = kernel[i + 1]!; + const xStart = kernel[i + 2]!; + const xEndExcl = kernel[i + 3]!; + for (let dx = xStart; dx < xEndExcl; ++dx) { + const x = cx + dx; + const y = cy + dy; + const z = cz + dz; + if (passesFilter(x, y, z)) appendVoxel(x, y, z); + } + } + } + } + } else { + // DISK: project stencil through basis vectors, no axis-aligned runs possible + const { u: uVec, v: vVec } = basis as { u: vec3; v: vec3 }; + const ux = uVec[0], + uy = uVec[1], + uz = uVec[2]; + const vx = vVec[0], + vy = vVec[1], + vz = vVec[2]; + const stencil = getDiskStencilKernel(r); + + for (const [cx, cy, cz] of centers) { + for (let i = 0; i < stencil.length; i += 2) { + const si = stencil[i]!; + const sj = stencil[i + 1]!; + const x = Math.round(cx + ux * si + vx * sj); + const y = Math.round(cy + uy * si + vy * sj); + const z = Math.round(cz + uz * si + vz * sj); + if (filterValue !== undefined && !passesFilter(x, y, z)) continue; + appendVoxel(x, y, z); + } + } + } + + if (edits.size > 0) { + previewSource.applyLocalEdits(edits); + for (const key of edits.keys()) previewSource.setOverlaySeq(key, seq); + } + } + + async dispatchBrushStroke( + centers: Float32Array[], + radiusCanonical: number, + valueGetter: VoxelValueGetter, + shape: BrushShape, + basis: { u: Float32Array; v: Float32Array }, + seq: number, + filterValue?: bigint, + ) { + if (centers.length === 0) { + // Nothing will be written for this stroke; drop whatever its previews + // tagged so the overlay does not wait for a write that never comes. + this.rollbackStroke(seq); + return; + } + const storageValue = valueGetter(false); + const coveredVoxKeys = await this.dispatchOperation({ + type: VoxelOperationType.BRUSH, + seq, + centers, + radius: radiusCanonical, + value: storageValue, + shape, + basis, + filterValue, + }); + this.reconcileStroke(seq, coveredVoxKeys); + } + + async floodFillPlane2D( + startPositionCanonical: Float32Array, + fillValueGetter: VoxelValueGetter, + maxVoxels: number, + basis: { u: Float32Array; v: Float32Array }, + filterValue?: bigint, + morphological = true, + ) { + const seq = this.beginStroke(); + const previewValue = fillValueGetter(true); + const sourcesByScale = this.host.primarySource.getSources( + this.getIdentitySliceViewSourceOptions(), + ); + const primaryChunkSource = sourcesByScale[0][0] + .chunkSource as VolumeChunkSource; + + const previewMultiscale = this.host.previewSource; + if (!previewMultiscale) return; + const previewChunkSource = previewMultiscale.getSources( + this.getIdentitySliceViewSourceOptions(), + )[0][0].chunkSource as InMemoryVolumeChunkSource; + + const startX = Math.round(startPositionCanonical[0]); + const startY = Math.round(startPositionCanonical[1]); + const startZ = Math.round(startPositionCanonical[2]); + + const tempPos = new Float32Array(3); + + const getValue = (x: number, y: number, z: number): bigint | null => { + tempPos[0] = x; + tempPos[1] = y; + tempPos[2] = z; + + const previewVal = previewChunkSource.getValueAt( + tempPos, + this.singleChannelAccess, + ); + if (previewVal != null) { + return typeof previewVal === "bigint" ? previewVal : BigInt(previewVal); + } + + const primaryVal = primaryChunkSource.getValueAt( + tempPos, + this.singleChannelAccess, + ); + if (primaryVal != null) { + return typeof primaryVal === "bigint" ? primaryVal : BigInt(primaryVal); + } + return null; + }; + + const originalValue = getValue(startX, startY, startZ); + if (originalValue === null) { + // Chunk not yet loaded on the frontend — skip preview and dispatch directly + // to the backend, which will load the chunk itself. + await this.dispatchOperation({ + type: VoxelOperationType.FLOOD_FILL, + seq, + seed: startPositionCanonical, + value: fillValueGetter(false), + maxVoxels, + basis, + filterValue, + morphological, + }); + return; + } + if (filterValue !== undefined && originalValue !== filterValue) return; + if (originalValue === previewValue) return; + + const visited = new Set(); + const queue: [number, number][] = [[0, 0]]; + visited.add("0,0"); + + const edits = new Map(); + const { chunkDataSize } = previewChunkSource.spec; + const sizeX = chunkDataSize[0]; + const sizeY = chunkDataSize[1]; + const sizeZ = chunkDataSize[2]; + const strideY = sizeX; + const strideZ = sizeX * sizeY; + + let filledCount = 0; + const start = Date.now(); + const MAX_TIME_MS = 50; + const ux = basis.u[0], + uy = basis.u[1], + uz = basis.u[2]; + const vx = basis.v[0], + vy = basis.v[1], + vz = basis.v[2]; + + while (queue.length > 0 && filledCount < maxVoxels) { + if ((filledCount & 63) === 0 && Date.now() - start > MAX_TIME_MS) break; + + const [u, v] = queue.shift()!; + const x = Math.round(startX + ux * u + vx * v); + const y = Math.round(startY + uy * u + vy * v); + const z = Math.round(startZ + uz * u + vz * v); + + const cx = Math.floor(x / sizeX); + const cy = Math.floor(y / sizeY); + const cz = Math.floor(z / sizeZ); + const lx = x - cx * sizeX; + const ly = y - cy * sizeY; + const lz = z - cz * sizeZ; + + const key = `${cx},${cy},${cz}`; + let entry = edits.get(key); + if (!entry) { + entry = { indices: [], value: previewValue }; + edits.set(key, entry); + } + const index = lz * strideZ + ly * strideY + lx; + entry.indices.push(index); + filledCount++; + + const neighbors: [number, number][] = [ + [u + 1, v], + [u - 1, v], + [u, v + 1], + [u, v - 1], + ]; + + for (const [nu, nv] of neighbors) { + const k = `${nu},${nv}`; + if (visited.has(k)) continue; + visited.add(k); + + const nx = Math.round(startX + ux * nu + vx * nv); + const ny = Math.round(startY + uy * nu + vy * nv); + const nz = Math.round(startZ + uz * nu + vz * nv); + + const val = getValue(nx, ny, nz); + if (val !== null && val === originalValue) { + queue.push([nu, nv]); + } + } + } + + if (filledCount > 0) { + previewChunkSource.applyLocalEdits(edits); + for (const key of edits.keys()) { + previewChunkSource.setOverlaySeq(key, seq); + } + } + + const storageValue = fillValueGetter(false); + try { + const coveredVoxKeys = await this.dispatchOperation({ + type: VoxelOperationType.FLOOD_FILL, + seq, + seed: startPositionCanonical, + value: storageValue, + maxVoxels, + basis, + filterValue, + morphological, + }); + this.reconcileStroke(seq, coveredVoxKeys); + } catch (e) { + this.rollbackStroke(seq); + throw e; + } + } + + callChunkReload( + voxChunkKeys: string[], + isForPreviewChunks: boolean, + overlayKeysToClear?: Record, + coveredSeqs?: Record, + isRollback = false, + ) { + if (!Array.isArray(voxChunkKeys) || voxChunkKeys.length === 0) return; + const multiscaleSource = isForPreviewChunks + ? this.host.previewSource + : this.host.primarySource; + if (!multiscaleSource) { + throw new Error( + "VoxelEditController.callChunkReload: ERROR Missing source", + ); + } + const sources = multiscaleSource.getSources( + this.getIdentitySliceViewSourceOptions(), + )[0]; + if (!sources) { + throw new Error( + "VoxelEditController.callChunkReload: Missing base source", + ); + } + + const chunksToInvalidateBySource = new Map(); + + if (!isForPreviewChunks) { + // Real chunks: invalidate lazily so the current GPU chunk stays on screen, + // and clear the matching overlay chunk only once the refetched real data + // has actually arrived on the GPU (swap-on-arrival), never on a timer. + // + // The overlay to clear defaults to the same key/LOD (max-res edits). For + // downsampled parents the backend passes the originating LOD-0 key, so the + // visible (forced LOD-0) overlay is cleared as soon as the real chunk of + // whatever LOD is on screen arrives. A never-swapped overlay simply stays + // displayed (correct) at worst leaking a little memory — far better than a + // timer that would clear it with nothing to show. + const previewSources = this.host.previewSource?.getSources( + this.getIdentitySliceViewSourceOptions(), + )[0]; + // Known, accepted race: a refetch predating the write we are reloading + // for may still sit in the frontend's pending-update queue at arming + // time; its arrival is indistinguishable from the fresh one and can + // clear the overlay over pre-write data. The window requires the queue + // to lag behind RPC processing (heavy load only), and the backend's + // in-flight download cancellation on write guarantees a correct + // refetch follows within one round trip, so the effect is a rare, + // self-healing flicker — the trade-off for keeping the chunk manager + // free of voxel-specific hooks. + for (const voxKey of voxChunkKeys) { + const parsed = parseVoxChunkKey(voxKey); + if (!parsed) continue; + const source = sources[parsed.lodIndex]?.chunkSource as + | VolumeChunkSource + | undefined; + if (!source) continue; + const { chunkKey } = parsed; + + const overlayParsed = parseVoxChunkKey( + overlayKeysToClear?.[voxKey] ?? voxKey, + ); + const overlaySource = overlayParsed + ? (previewSources?.[overlayParsed.lodIndex]?.chunkSource as + | InMemoryVolumeChunkSource + | undefined) + : undefined; + if (overlaySource) { + // The backend echoes, per reloaded chunk, the highest stroke seq + // its write covers. The overlay chunk carries the seq of the last + // stroke whose preview touched it (including a stroke still under + // the mouse — its seq is allocated before its first preview), read + // when the swap resolves. Clearing only when coverage reaches that + // tag guarantees the arriving data contains everything the overlay + // shows. A skipped clear is re-armed by the covering write's own + // reload; a stroke that never gets written is rolled back + // explicitly (rollbackStroke) instead of waited for. + if (isRollback) { + // Undo/redo: whatever arrives next is the truth — purge the tag + // so the swap clears on first arrival. An in-progress stroke's + // chunk would lose its preview until its dispatch rewrites it + // (Ctrl+Z mid-drag, accepted). + overlaySource.clearOverlaySeq(overlayParsed!.chunkKey); + } + this.pendingOverlaySwaps.set(voxKey, { + source, + chunkKey, + staleChunk: source.chunks.get(chunkKey), + overlaySource, + overlayChunkKey: overlayParsed!.chunkKey, + coveredSeq: coveredSeqs?.[voxKey] ?? 0, + }); + } + let arr = chunksToInvalidateBySource.get(source); + if (!arr) { + arr = []; + chunksToInvalidateBySource.set(source, arr); + } + arr.push(chunkKey); + } + + for (const [source, keys] of chunksToInvalidateBySource.entries()) { + if (keys.length > 0) { + source.invalidateChunks(keys, { lazy: true }); + } + } + return; + } + + // Preview chunks: clear the optimistic overlay immediately (write-failure + // rollback, and downsampled-overlay cleanup). + for (const voxKey of voxChunkKeys) { + const parsed = parseVoxChunkKey(voxKey); + if (!parsed) continue; + const source = sources[parsed.lodIndex]?.chunkSource as + | VolumeChunkSource + | undefined; + if (source) { + let arr = chunksToInvalidateBySource.get(source); + if (!arr) { + arr = []; + chunksToInvalidateBySource.set(source, arr); + } + arr.push(parsed.chunkKey); + } + } + + for (const [source, keys] of chunksToInvalidateBySource.entries()) { + if (keys.length > 0) { + source.invalidateChunks(keys); + } + } + } + + handleCommitFailure(voxChunkKeys: string[], message: string): void { + try { + this.callChunkReload(voxChunkKeys, true); + } finally { + StatusMessage.showTemporaryMessage(message); + } + } + + public async undo() { + if (!this.rpc) + throw new Error("VoxelEditController.undo: RPC not initialized."); + await this.rpc + .promiseInvoke(VOX_EDIT_UNDO_RPC_ID, { rpcId: this.rpcId }) + .catch((error: unknown) => { + const message = error instanceof Error ? error.message : String(error); + StatusMessage.showTemporaryMessage(`Undo failed: ${message}`, 3000); + }); + } + + public async redo() { + if (!this.rpc) + throw new Error("VoxelEditController.redo: RPC not initialized."); + await this.rpc + .promiseInvoke(VOX_EDIT_REDO_RPC_ID, { rpcId: this.rpcId }) + .catch((error: unknown) => { + const message = error instanceof Error ? error.message : String(error); + StatusMessage.showTemporaryMessage(`Redo failed: ${message}`, 3000); + }); + } +} + +function asRecordOrUndefined(x: unknown): Record | undefined { + return x !== null && typeof x === "object" + ? (x as Record) + : undefined; +} + +registerRPC(VOX_RELOAD_CHUNKS_RPC_ID, function (x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + const keys: string[] = Array.isArray(x.voxChunkKeys) ? x.voxChunkKeys : []; + obj.callChunkReload( + keys, + x.isForPreviewChunks, + asRecordOrUndefined(x.overlayKeysToClear), + asRecordOrUndefined(x.coveredSeqs), + x.isRollback === true, + ); +}); + +registerRPC(VOX_EDIT_FAILURE_RPC_ID, function (x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + const keys: string[] = Array.isArray(x.voxChunkKeys) ? x.voxChunkKeys : []; + const message: string = + typeof x.message === "string" ? x.message : "Voxel edit failed."; + obj.handleCommitFailure(keys, message); +}); + +registerRPC(VOX_EDIT_HISTORY_UPDATE_RPC_ID, function (x: any) { + const obj = this.get(x.rpcId) as VoxelEditController; + const undoCount = typeof x.undoCount === "number" ? x.undoCount : 0; + const redoCount = typeof x.redoCount === "number" ? x.redoCount : 0; + obj.undoCount.value = undoCount; + obj.redoCount.value = redoCount; +}); diff --git a/src/voxel_annotation/preview_multiscale_chunk_source.ts b/src/voxel_annotation/preview_multiscale_chunk_source.ts new file mode 100644 index 0000000000..43f3b63153 --- /dev/null +++ b/src/voxel_annotation/preview_multiscale_chunk_source.ts @@ -0,0 +1,73 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { ChunkManager } from "#src/chunk_manager/frontend.js"; +import type { SliceViewSingleResolutionSource } from "#src/sliceview/frontend.js"; +import type { + VolumeChunkSpecification, + VolumeSourceOptions, + DataType, + VolumeType, +} from "#src/sliceview/volume/base.js"; +import { + InMemoryVolumeChunkSource, + MultiscaleVolumeChunkSource, + type VolumeChunkSource, +} from "#src/sliceview/volume/frontend.js"; + +export class VoxelPreviewMultiscaleSource extends MultiscaleVolumeChunkSource { + dataType: DataType; + volumeType: VolumeType; + rank: number; + + constructor( + chunkManager: ChunkManager, + public primarySource: MultiscaleVolumeChunkSource, + ) { + super(chunkManager); + this.dataType = primarySource.dataType; + this.volumeType = primarySource.volumeType; + this.rank = primarySource.rank; + } + + getSources( + options: VolumeSourceOptions, + ): SliceViewSingleResolutionSource[][] { + const sourcesByScale = this.primarySource.getSources(options); + + return sourcesByScale.map((orientation) => { + return orientation.map((primaryResSource) => { + const spec = primaryResSource.chunkSource.spec; + + const previewSpec: VolumeChunkSpecification = { + ...spec, + compressedSegmentationBlockSize: undefined, + }; + + const previewSource = this.chunkManager.getChunkSource( + InMemoryVolumeChunkSource, + { spec: previewSpec }, + ); + + return { + chunkSource: previewSource, + chunkToMultiscaleTransform: + primaryResSource.chunkToMultiscaleTransform, + }; + }); + }); + } +} diff --git a/src/voxel_annotation/stamina_calibration.benchmark.ts b/src/voxel_annotation/stamina_calibration.benchmark.ts new file mode 100644 index 0000000000..db18c2e8f0 --- /dev/null +++ b/src/voxel_annotation/stamina_calibration.benchmark.ts @@ -0,0 +1,311 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { describe, bench, vi } from "vitest"; +import type { VolumeChunk } from "#src/sliceview/volume/backend.js"; +import { VolumeChunkSource } from "#src/sliceview/volume/backend.js"; +import { DataType, DATA_TYPE_ARRAY_CONSTRUCTOR } from "#src/util/data_type.js"; +import { vec3 } from "#src/util/geom.js"; +import { VoxelEditController } from "#src/voxel_annotation/backend.js"; +import { + BrushShape, + VOXEL_EMPTY_VALUE, + VoxelOperationType, + makeVoxChunkKey, +} from "#src/voxel_annotation/base.js"; + +const NETWORK_LATENCY = 0; +class RealisticInMemorySource extends VolumeChunkSource { + public storage = new Map(); + + async download(chunk: VolumeChunk) { + if (NETWORK_LATENCY) + await new Promise((resolve) => setTimeout(resolve, NETWORK_LATENCY)); + + if (!chunk.chunkDataSize) { + this.computeChunkBounds(chunk); + } + + const numElements = chunk.chunkDataSize!.reduce((a, b) => a * b, 1); + const Ctor = DATA_TYPE_ARRAY_CONSTRUCTOR[this.spec.dataType]; + const key = chunk.chunkGridPosition.join(","); + + if (this.storage.has(key)) { + chunk.data = new (Ctor as any)(this.storage.get(key)!); + } else { + chunk.data = new (Ctor as any)(numElements); + } + } + + async writeChunk(chunk: VolumeChunk) { + if (NETWORK_LATENCY) + await new Promise((resolve) => setTimeout(resolve, NETWORK_LATENCY)); + + const key = chunk.chunkGridPosition.join(","); + const Ctor = DATA_TYPE_ARRAY_CONSTRUCTOR[this.spec.dataType]; + if (chunk.data) { + this.storage.set(key, new (Ctor as any)(chunk.data)); + } + } +} + +const createResConfig = (lod: number, chunkSize: number) => ({ + lodIndex: lod, + transform: [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1], + chunkSize: [chunkSize, chunkSize, chunkSize], + sourceRpc: 100 + lod, +}); + +const CHUNK_SIZE = 64; + +const spec = { + rank: 3, + chunkDataSize: new Uint32Array([CHUNK_SIZE, CHUNK_SIZE, CHUNK_SIZE]), + dataType: DataType.UINT64, + lowerVoxelBound: new Float32Array([0, 0, 0]), + upperVoxelBound: new Float32Array([10000, 10000, 10000]), + baseVoxelOffset: new Float32Array([0, 0, 0]), + fillValue: VOXEL_EMPTY_VALUE, +}; + +const mockChunkQueueManager = { sources: new Set() }; +const mockChunkManager = { + queueManager: mockChunkQueueManager, + memoize: { get: (_k: string, f: Function) => f() }, +}; + +const rpcObjects = new Map([ + [0, mockChunkManager], + [999, { value: 0 }], +]); + +const rpcHandler = { + get: (id: number) => rpcObjects.get(id) ?? null, + invoke: () => {}, + newId: () => 0, + register: () => {}, + set: () => {}, + promiseInvoke: async () => {}, +} as any; + +const mockSource0 = new RealisticInMemorySource(rpcHandler, { + spec, + chunkManager: 0, +}); +rpcObjects.set(100, mockSource0); + +const controller = new VoxelEditController(rpcHandler, { + resolutions: [createResConfig(0, CHUNK_SIZE)], + pendingOpCount: 999, +}); + +vi.spyOn(controller as any, "enqueueDownsample").mockImplementation(() => {}); + +// Pre-populate the downsample chunk with worst-case data once. +const _initChunk = mockSource0.getChunk( + new Float32Array([0, 0, 0]), +) as VolumeChunk; +await mockSource0.download(_initChunk); +const _initData = _initChunk.data as BigUint64Array; +for (let i = 0; i < _initData.length; i++) { + _initData[i] = BigInt(i % 5); +} + +// -------------------------------------------------------------------------- +// 1. SYSTEM OVERHEAD (Commit) +// -------------------------------------------------------------------------- + +describe("Commit", () => { + const voxelsCounts = [1000, 10000, 100000]; + const numOfEdits = [1, 10, 50]; + + for (const numOfEdit of numOfEdits) { + for (const count of voxelsCounts) { + // Pre-compute data at collection time, not during the bench run. + const indices = new Uint32Array(count); + const values = new BigUint64Array(count); + for (let i = 0; i < count; i++) { + indices[i] = i; + values[i] = BigInt(i); + } + const edits = [{ key: "lod0#0,0,0", indices, values }]; + for (let i = 1; i < numOfEdit; i++) + edits.push({ key: `lod0#0,0,${i}`, indices, values }); + + bench(`voxels=${count} edits=${numOfEdit}`, async () => { + (controller as any).pendingEdits.push(...edits); + await (controller as any).flushPending(); + }); + } + } +}); + +// -------------------------------------------------------------------------- +// 2. DOWNSAMPLING +// -------------------------------------------------------------------------- + +describe("Downsample", () => { + bench(`inputVoxels=${CHUNK_SIZE ** 3}`, async () => { + const originKey = makeVoxChunkKey("0,0,0", 0); + await (controller as any).downsampleStep(originKey, originKey, 0); + }); +}); + +// -------------------------------------------------------------------------- +// 3. BRUSH STROKES +// -------------------------------------------------------------------------- + +const STROKE_LEN = 20; +const runStroke = async ( + shape: BrushShape, + radius: number, + useFilter: boolean, +) => { + (controller as any).brushCache.reset(); + const center = new Float32Array(3); + const basis = { u: vec3.fromValues(1, 0, 0), v: vec3.fromValues(0, 1, 0) }; + + if (useFilter) { + const chunk = mockSource0.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + if (!chunk.data) await mockSource0.download(chunk); + } + const filterVal = useFilter ? 999n : undefined; + + for (let i = 0; i < STROKE_LEN; i++) { + center[0] = 30 + i; + center[1] = 30; + center[2] = 30; + + await (controller as any).performBrush({ + type: VoxelOperationType.BRUSH, + centers: [center], + radius, + value: 1n, + shape, + basis, + filterValue: filterVal, + }); + } +}; + +describe("Brush SPHERE", () => { + for (const r of [4, 8, 12, 16, 20, 24, 28, 32]) { + bench(`r=${r} no filter`, async () => { + await runStroke(BrushShape.SPHERE, r, false); + }); + + bench(`r=${r} with filter`, async () => { + await runStroke(BrushShape.SPHERE, r, true); + }); + } +}); + +describe("Brush DISK", () => { + for (const r of [16, 24, 32, 40, 48, 56, 64]) { + bench(`r=${r} no filter`, async () => { + await runStroke(BrushShape.DISK, r, false); + }); + + bench(`r=${r} with filter`, async () => { + await runStroke(BrushShape.DISK, r, true); + }); + } +}); + +// -------------------------------------------------------------------------- +// 4. FLOOD FILL +// -------------------------------------------------------------------------- + +describe("FloodFill", () => { + for (const size of [1000, 5000, 10000, 25000, 50000]) { + bench(`maxVoxels=${size}`, async () => { + // Reset chunk data each iteration so the flood fill always starts clean. + const chunk = mockSource0.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + (chunk.data as BigUint64Array).fill(VOXEL_EMPTY_VALUE); + + try { + await (controller as any).performFloodFill({ + type: VoxelOperationType.FLOOD_FILL, + seed: new Float32Array([32, 32, 32]), + value: 1n, + maxVoxels: size, + basis: { u: vec3.fromValues(1, 0, 0), v: vec3.fromValues(0, 1, 0) }, + }); + } catch (e: any) { + if (!e.message.includes("too many voxels")) throw e; + } + }); + } +}); + +// -------------------------------------------------------------------------- +// 5. UNDO +// Each iteration performs the setup (pushing to the undo stack) and the undo +// itself, so reported timings include both. The setup cost is minimal +// compared to the undo write-back. +// -------------------------------------------------------------------------- + +describe("Undo (chunks)", () => { + for (const count of [1, 10, 50, 100]) { + bench(`chunks=${count}`, async () => { + const changes = new Map(); + for (let i = 0; i < count; i++) { + const chunk = mockSource0.getChunk( + new Float32Array([i, 0, 0]), + ) as VolumeChunk; + if (!chunk.data) await mockSource0.download(chunk); + changes.set(`lod0#${i},0,0`, { + indices: new Uint32Array([0]), + oldValues: new BigUint64Array([1n]), + newValues: new BigUint64Array([2n]), + }); + } + (controller as any).undoStack.push({ + changes, + timestamp: 0, + description: "bench", + }); + await controller.undo(); + }); + } +}); + +describe("Undo (voxels, 1 chunk)", () => { + for (const count of [1000, 10000, 100000]) { + const indices = new Uint32Array(count); + const vals = new BigUint64Array(count).fill(1n); + + bench(`voxels=${count}`, async () => { + const chunk = mockSource0.getChunk( + new Float32Array([0, 0, 0]), + ) as VolumeChunk; + if (!chunk.data) await mockSource0.download(chunk); + + (controller as any).undoStack.push({ + changes: new Map([ + [`lod0#0,0,0`, { indices, oldValues: vals, newValues: vals }], + ]), + timestamp: 0, + description: "bench", + }); + await controller.undo(); + }); + } +}); diff --git a/src/widget/layer_control_button.ts b/src/widget/layer_control_button.ts new file mode 100644 index 0000000000..273162bd62 --- /dev/null +++ b/src/widget/layer_control_button.ts @@ -0,0 +1,37 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { UserLayer } from "#src/layer/index.js"; +import type { LayerControlFactory } from "#src/widget/layer_control.js"; + +export function buttonLayerControl(options: { + text: string; + onClick: (layer: LayerType) => void; +}): LayerControlFactory { + return { + makeControl: (layer, context) => { + const control = document.createElement("button"); + control.textContent = options.text; + context.registerEventListener(control, "click", () => + options.onClick(layer), + ); + return { control, controlElement: control }; + }, + activateTool: (activation) => { + options.onClick(activation.tool.layer as LayerType); + }, + }; +} diff --git a/tests/kvstore/s3.spec.ts b/tests/kvstore/s3.spec.ts index eeb1fd3690..75e2eb9c54 100644 --- a/tests/kvstore/s3.spec.ts +++ b/tests/kvstore/s3.spec.ts @@ -83,6 +83,34 @@ describe("http:// path-style URL", () => { testKvStore(constantFixture(`https://s3.amazonaws.com/${BUCKET}/`)); }); +describe("dot segment key components", () => { + test.for(["..", "../b", "a/../b", "a/./b"])( + "read rejects %s", + async (key) => { + const context = await sharedKvStoreContext(); + await expect( + context.kvStoreContext.read(`s3://${BUCKET}/${key}`, { + throwIfMissing: true, + }), + ).rejects.toThrow(/path components are not supported/); + }, + ); + test("write rejects a/../b", async () => { + const context = await sharedKvStoreContext(); + const { store } = context.kvStoreContext.getKvStore(`s3://${BUCKET}/`); + await expect(store.write!("a/../b", new ArrayBuffer(0))).rejects.toThrow( + /path components are not supported/, + ); + }); + test("delete rejects a/../b", async () => { + const context = await sharedKvStoreContext(); + const { store } = context.kvStoreContext.getKvStore(`s3://${BUCKET}/`); + await expect(store.delete!("a/../b")).rejects.toThrow( + /path components are not supported/, + ); + }); +}); + describe("special characters", () => { test.for(SPECIAL_CHAR_CODES)("charCode=%s", async (charCode) => { const context = await sharedKvStoreContext(); diff --git a/tests/voxel_annotation/pipeline_zarr_s3.browser_test.ts b/tests/voxel_annotation/pipeline_zarr_s3.browser_test.ts new file mode 100644 index 0000000000..dbcabe4cdc --- /dev/null +++ b/tests/voxel_annotation/pipeline_zarr_s3.browser_test.ts @@ -0,0 +1,769 @@ +/** + * @license + * Copyright 2025 Google Inc. + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import "#src/datasource/zarr/register_default.js"; +import "#src/kvstore/s3/register_frontend.js"; +import "#src/sliceview/uncompressed_chunk_format.js"; +import "#src/layer/segmentation/index.js"; +import "#src/layer/image/index.js"; + +import { http, HttpResponse } from "msw"; +import { test, beforeEach, afterEach } from "vitest"; +import { DisplayContext } from "#src/display_context.js"; +import { makeLayer } from "#src/layer/index.js"; +import type { + VoxelEditingContext, + UserLayerWithVoxelEditing, +} from "#src/layer/voxel_annotation/index.js"; +import { Viewer } from "#src/viewer.js"; +import { mswFixture } from "#tests/fixtures/msw"; + +const msw = mswFixture(); +let viewer: Viewer | undefined; +const storage = new Map(); +const baseUrl = "http://localhost:9000"; + +beforeEach(async () => { + storage.clear(); + const display = new DisplayContext(document.createElement("div")); + viewer = new Viewer(display, { + showLayerDialog: false, + resetStateWhenEmpty: false, + }); + + (await msw()).use( + http.put(`${baseUrl}/*`, async ({ request }) => { + const parsed = parseBucketKey(request.url); + if (!parsed) return new HttpResponse(null, { status: 400 }); + const storageKey = `${parsed.bucket}/${parsed.key}`; + const buffer = await request.arrayBuffer(); + storage.set(storageKey, buffer); + return new HttpResponse(null, { status: 200 }); + }), + http.get(`${baseUrl}/*`, ({ request }) => { + const parsed = parseBucketKey(request.url); + const storageKey = parsed ? `${parsed.bucket}/${parsed.key}` : ""; + const data = storage.get(storageKey); + if (!data) return new HttpResponse(null, { status: 404 }); + return new HttpResponse(data); + }), + http.head(`${baseUrl}/*`, ({ request }) => { + const parsed = parseBucketKey(request.url); + const storageKey = parsed ? `${parsed.bucket}/${parsed.key}` : ""; + const data = storage.get(storageKey); + if (!data) return new HttpResponse(null, { status: 404 }); + return new HttpResponse(null, { + status: 200, + headers: { + "Content-Length": data.byteLength.toString(), + }, + }); + }), + ); +}); + +afterEach(() => { + if (viewer) { + viewer.dispose(); + viewer = undefined; + } +}); + +async function poll( + condition: () => boolean | Promise, + what: string, + timeout = 5000, +) { + const start = Date.now(); + while (Date.now() - start < timeout) { + if (await condition()) return; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + throw new Error("Timeout polling condition: " + what); +} + +function parseBucketKey( + urlStr: string, +): { bucket: string; key: string } | null { + const url = new URL(urlStr); + const path = url.pathname.slice(1); + const parts = path.split("/"); + if (parts.length < 2) return null; + return { bucket: parts[0], key: parts.slice(1).join("/") }; +} + +async function waitForEditingContext() { + if (!viewer) throw new Error("Viewer not initialized"); + await poll(() => { + const userLayer = viewer!.layerManager.managedLayers[0] + ?.layer as UserLayerWithVoxelEditing; + return userLayer?.editingContexts?.size > 0; + }, "Wait for Editing Context"); + const userLayer = viewer.layerManager.managedLayers[0] + .layer as UserLayerWithVoxelEditing; + const context = userLayer.editingContexts.values().next() + .value as VoxelEditingContext; + (context as any).hasUserConfirmedWriting = true; + return { userLayer, context }; +} + +test("Pipeline: Zarr V2 (UINT8) Undo/Redo with Brush", async () => { + const BUCKET = "test-v2-uint8"; + const zarray = JSON.stringify({ + zarr_format: 2, + shape: [64, 64, 64], + chunks: [32, 32, 32], + dtype: "|u1", + fill_value: 0, + order: "C", + dimension_separator: ".", + compressor: null, + }); + storage.set( + `${BUCKET}/data.zarr/.zarray`, + new TextEncoder().encode(zarray).buffer, + ); + storage.set( + `${BUCKET}/data.zarr/.zgroup`, + new TextEncoder().encode("{}").buffer, + ); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "image", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const center = new Float32Array([16, 16, 16]); + await context.dispatchBrushStroke( + [center], + 5, + (_) => 100n, + 0 /* DISK */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + const chunkKey = `${BUCKET}/data.zarr/0.0.0`; + + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new Uint8Array(data); + return arr.some((v) => v === 100); + }, "Verify painted chunk"); + + await context.undo(); + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return true; + return new Uint8Array(data).every((v) => v === 0); + }, "Verify undo"); + + await context.redo(); + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return false; + return new Uint8Array(data).some((v) => v === 100); + }, "Verify redo"); +}); + +test("Pipeline: Zarr V3 (UINT64) Brush", async () => { + const BUCKET = "test-v3-uint64"; + const zarrJson = JSON.stringify({ + zarr_format: 3, + node_type: "array", + shape: [64, 64, 64], + data_type: "uint64", + chunk_grid: { + name: "regular", + configuration: { chunk_shape: [32, 32, 32] }, + }, + chunk_key_encoding: { + name: "default", + configuration: { separator: "/" }, + }, + codecs: [{ name: "bytes", configuration: { endian: "little" } }], + fill_value: 0, + attributes: {}, + }); + + storage.set( + `${BUCKET}/data.zarr/zarr.json`, + new TextEncoder().encode(zarrJson).buffer, + ); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "segmentation", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr|zarr3:`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const center = new Float32Array([16, 16, 16]); + const paintVal = 123456789n; + await context.dispatchBrushStroke( + [center], + 2, + (_) => paintVal, + 0 /* DISK */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + const chunkKey = `${BUCKET}/data.zarr/c/0/0/0`; + + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new BigUint64Array(data); + return arr.some((v) => v === paintVal); + }, "Verify painted chunk (UINT64)"); +}); + +test("Pipeline: Repaint over existing chunk data (Zarr V3 UINT64)", async () => { + const BUCKET = "test-v3-repaint"; + const zarrJson = JSON.stringify({ + zarr_format: 3, + node_type: "array", + shape: [64, 64, 64], + data_type: "uint64", + chunk_grid: { + name: "regular", + configuration: { chunk_shape: [32, 32, 32] }, + }, + chunk_key_encoding: { + name: "default", + configuration: { separator: "/" }, + }, + codecs: [{ name: "bytes", configuration: { endian: "little" } }], + fill_value: 0, + attributes: {}, + }); + + storage.set( + `${BUCKET}/data.zarr/zarr.json`, + new TextEncoder().encode(zarrJson).buffer, + ); + + // Pre-seed the chunk with existing data: the brush must merge into it, not + // replace it (regression test for edits only working on empty chunks). + const existingVal = 7n; + const existing = new BigUint64Array(32 * 32 * 32).fill(existingVal); + const chunkKey = `${BUCKET}/data.zarr/c/0/0/0`; + storage.set(chunkKey, existing.buffer.slice(0) as ArrayBuffer); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "segmentation", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr|zarr3:`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const paintVal = 123n; + await context.dispatchBrushStroke( + [new Float32Array([16, 16, 16])], + 2, + (_) => paintVal, + 0 /* DISK */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new BigUint64Array(data); + return ( + arr.some((v) => v === paintVal) && arr.some((v) => v === existingVal) + ); + }, "Verify stroke applied AND pre-existing data preserved"); +}); + +test("Pipeline: Repaint over existing gzip-compressed chunk (Zarr V3 UINT64)", async () => { + const BUCKET = "test-v3-repaint-gzip"; + const zarrJson = JSON.stringify({ + zarr_format: 3, + node_type: "array", + shape: [64, 64, 64], + data_type: "uint64", + chunk_grid: { + name: "regular", + configuration: { chunk_shape: [32, 32, 32] }, + }, + chunk_key_encoding: { + name: "default", + configuration: { separator: "/" }, + }, + codecs: [ + { name: "bytes", configuration: { endian: "little" } }, + { name: "gzip", configuration: { level: 1 } }, + ], + fill_value: 0, + attributes: {}, + }); + + storage.set( + `${BUCKET}/data.zarr/zarr.json`, + new TextEncoder().encode(zarrJson).buffer, + ); + + const gzip = async (data: ArrayBuffer): Promise => { + const cs = new CompressionStream("gzip"); + const writer = cs.writable.getWriter(); + void writer.write(new Uint8Array(data)); + void writer.close(); + return await new Response(cs.readable).arrayBuffer(); + }; + const gunzip = async (data: ArrayBuffer): Promise => { + const ds = new DecompressionStream("gzip"); + const writer = ds.writable.getWriter(); + void writer.write(new Uint8Array(data)); + void writer.close(); + return await new Response(ds.readable).arrayBuffer(); + }; + + const existingVal = 7n; + const existing = new BigUint64Array(32 * 32 * 32).fill(existingVal); + const chunkKey = `${BUCKET}/data.zarr/c/0/0/0`; + storage.set(chunkKey, await gzip(existing.buffer as ArrayBuffer)); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "segmentation", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr|zarr3:`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const paintVal = 123n; + await context.dispatchBrushStroke( + [new Float32Array([16, 16, 16])], + 2, + (_) => paintVal, + 0 /* DISK */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + await poll(async () => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new BigUint64Array(await gunzip(data)); + return ( + arr.some((v) => v === paintVal) && arr.some((v) => v === existingVal) + ); + }, "Verify gzip repaint applied AND pre-existing data preserved"); +}); + +test("Pipeline: Repaint over dense large-label segmentation data (Zarr V3 UINT64)", async () => { + const BUCKET = "test-v3-repaint-dense"; + const zarrJson = JSON.stringify({ + zarr_format: 3, + node_type: "array", + shape: [64, 64, 64], + data_type: "uint64", + chunk_grid: { + name: "regular", + configuration: { chunk_shape: [32, 32, 32] }, + }, + chunk_key_encoding: { + name: "default", + configuration: { separator: "/" }, + }, + codecs: [ + { name: "bytes", configuration: { endian: "little" } }, + { name: "gzip", configuration: { level: 1 } }, + ], + fill_value: 0, + attributes: {}, + }); + + storage.set( + `${BUCKET}/data.zarr/zarr.json`, + new TextEncoder().encode(zarrJson).buffer, + ); + + const gzip = async (data: ArrayBuffer): Promise => { + const cs = new CompressionStream("gzip"); + const writer = cs.writable.getWriter(); + void writer.write(new Uint8Array(data)); + void writer.close(); + return await new Response(cs.readable).arrayBuffer(); + }; + const gunzip = async (data: ArrayBuffer): Promise => { + const ds = new DecompressionStream("gzip"); + const writer = ds.writable.getWriter(); + void writer.write(new Uint8Array(data)); + void writer.close(); + return await new Response(ds.readable).arrayBuffer(); + }; + + // Worst case for the compressed-segmentation in-memory transcode: every + // voxel a distinct large (>2^63) label, as produced by real segmentation + // pipelines — unlike the small uniform values of the other tests. + const existing = new BigUint64Array(32 * 32 * 32); + for (let i = 0; i < existing.length; ++i) { + existing[i] = 0x8000000000000000n + BigInt(i); + } + const sentinel = existing[0]; + const chunkKey = `${BUCKET}/data.zarr/c/0/0/0`; + storage.set(chunkKey, await gzip(existing.buffer as ArrayBuffer)); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "segmentation", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr|zarr3:`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const paintVal = 0x9999999999999999n; + await context.dispatchBrushStroke( + [new Float32Array([16, 16, 16])], + 2, + (_) => paintVal, + 0 /* DISK */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + await poll(async () => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new BigUint64Array(await gunzip(data)); + return arr.some((v) => v === paintVal) && arr[0] === sentinel; + }, "Verify dense-label repaint applied AND existing labels preserved"); +}); + +test("Pipeline: Multiscale repaint with NON-EMPTY downsample parent (OME zarr3)", async () => { + const BUCKET = "test-v3-multiscale"; + const groupJson = JSON.stringify({ + zarr_format: 3, + node_type: "group", + attributes: { + multiscales: [ + { + version: "0.5", + axes: [ + { name: "x", type: "space", unit: "nanometer" }, + { name: "y", type: "space", unit: "nanometer" }, + { name: "z", type: "space", unit: "nanometer" }, + ], + datasets: [ + { + path: "s0", + coordinateTransformations: [ + { type: "scale", scale: [32, 32, 32] }, + ], + }, + { + path: "s1", + coordinateTransformations: [ + { type: "scale", scale: [64, 64, 64] }, + ], + }, + ], + name: "test-multiscale", + }, + ], + }, + }); + const arrayJson = (shape: number) => + JSON.stringify({ + zarr_format: 3, + node_type: "array", + shape: [shape, shape, shape], + data_type: "uint64", + chunk_grid: { + name: "regular", + configuration: { chunk_shape: [32, 32, 32] }, + }, + chunk_key_encoding: { + name: "default", + configuration: { separator: "/" }, + }, + codecs: [ + { name: "bytes", configuration: { endian: "little" } }, + { name: "gzip", configuration: { level: 1 } }, + ], + fill_value: 0, + attributes: {}, + }); + + storage.set( + `${BUCKET}/data.zarr/zarr.json`, + new TextEncoder().encode(groupJson).buffer, + ); + storage.set( + `${BUCKET}/data.zarr/s0/zarr.json`, + new TextEncoder().encode(arrayJson(64)).buffer, + ); + storage.set( + `${BUCKET}/data.zarr/s1/zarr.json`, + new TextEncoder().encode(arrayJson(32)).buffer, + ); + + const gzip = async (data: ArrayBuffer): Promise => { + const cs = new CompressionStream("gzip"); + const writer = cs.writable.getWriter(); + void writer.write(new Uint8Array(data)); + void writer.close(); + return await new Response(cs.readable).arrayBuffer(); + }; + const gunzip = async (data: ArrayBuffer): Promise => { + const ds = new DecompressionStream("gzip"); + const writer = ds.writable.getWriter(); + void writer.write(new Uint8Array(data)); + void writer.close(); + return await new Response(ds.readable).arrayBuffer(); + }; + + // Both the edited chunk AND its downsample parent pre-exist with data — + // replicating painting over pipeline-produced regions (the untested case; + // blank regions have empty parents). + const existingVal = 7n; + const parentVal = 9n; + const child = new BigUint64Array(32 * 32 * 32).fill(existingVal); + const parent = new BigUint64Array(32 * 32 * 32).fill(parentVal); + const childKey = `${BUCKET}/data.zarr/s0/c/0/0/0`; + const parentKey = `${BUCKET}/data.zarr/s1/c/0/0/0`; + storage.set(childKey, await gzip(child.buffer as ArrayBuffer)); + storage.set(parentKey, await gzip(parent.buffer as ArrayBuffer)); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "segmentation", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr|zarr3:`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const paintVal = 123n; + await context.dispatchBrushStroke( + [new Float32Array([16, 16, 16])], + 4, + (_) => paintVal, + 1 /* SPHERE */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + await poll(async () => { + const data = storage.get(childKey); + if (!data) return false; + const arr = new BigUint64Array(await gunzip(data)); + return ( + arr.some((v) => v === paintVal) && arr.some((v) => v === existingVal) + ); + }, "Verify s0 stroke applied AND existing s0 data preserved"); + + await poll(async () => { + const data = storage.get(parentKey); + if (!data) return false; + const arr = new BigUint64Array(await gunzip(data)); + return arr.some((v) => v === paintVal) && arr.some((v) => v === parentVal); + }, "Verify s1 downsample applied AND existing s1 data preserved"); +}); + +test("Pipeline: Zarr V2 (UINT32) with Slash Separator", async () => { + const BUCKET = "test-v2-sep"; + const zarray = JSON.stringify({ + zarr_format: 2, + shape: [64, 64, 64], + chunks: [32, 32, 32], + dtype: "new TextEncoder().encode(zarray).buffer, + ); + storage.set( + `${BUCKET}/data.zarr/.zgroup`, + new TextEncoder().encode("{}").buffer, + ); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "segmentation", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const center = new Float32Array([10, 10, 10]); + const paintVal = 42n; + await context.dispatchBrushStroke( + [center], + 2, + (_) => paintVal, + 0 /* DISK */, + { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }, + context.beginStroke(), + ); + + const chunkKey = `${BUCKET}/data.zarr/0/0/0`; + + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new Uint32Array(data); + return arr.some((v) => v === 42); + }, "Verify painted chunk with slash separator"); +}); + +test("Pipeline: Flood Fill (Zarr V2 UINT8 on img layer)", async () => { + const BUCKET = "test-flood-fill"; + const CHUNK_SIZE = 32; + const zarray = JSON.stringify({ + zarr_format: 2, + shape: [64, 64, 64], + chunks: [CHUNK_SIZE, CHUNK_SIZE, CHUNK_SIZE], + dtype: "|u1", + fill_value: 0, + order: "C", + dimension_separator: ".", + compressor: null, + }); + storage.set( + `${BUCKET}/data.zarr/.zarray`, + new TextEncoder().encode(zarray).buffer, + ); + storage.set( + `${BUCKET}/data.zarr/.zgroup`, + new TextEncoder().encode("{}").buffer, + ); + + const chunkData = new Uint8Array(CHUNK_SIZE * CHUNK_SIZE * CHUNK_SIZE); + // Create a hollow box from 5,5,0 to 25,25,0 in z=0 slice + for (let y = 5; y <= 25; y++) { + for (let x = 5; x <= 25; x++) { + if (x === 5 || x === 25 || y === 5 || y === 25) { + const index = 0 * CHUNK_SIZE * CHUNK_SIZE + y * CHUNK_SIZE + x; + chunkData[index] = 255; + } + } + } + const chunkKey = `${BUCKET}/data.zarr/0.0.0`; + storage.set(chunkKey, chunkData.buffer); + + const layer = makeLayer(viewer!.layerSpecification, "volume", { + type: "image", + source: { + url: `s3+http://localhost:9000/${BUCKET}/data.zarr`, + subsources: { default: { enabled: true, writingEnabled: true } }, + enableDefaultSubsources: false, + }, + }); + viewer!.layerSpecification.add(layer); + + const { context } = await waitForEditingContext(); + + const seed = new Float32Array([15, 15, 0]); + const fillValue = 128n; + const maxVoxels = 1000; + const basis = { + u: new Float32Array([1, 0, 0]), + v: new Float32Array([0, 1, 0]), + }; + + await poll( + async () => { + try { + await context.floodFillPlane2D( + seed, + (_) => fillValue, + maxVoxels, + basis, + ); + return true; + } catch (e: any) { + if (e.message.includes("unloaded")) { + return false; + } + throw e; + } + }, + "Execute flood fill", + 5000, + ); + + await poll(() => { + const data = storage.get(chunkKey); + if (!data) return false; + const arr = new Uint8Array(data); + const insideIndex = 0 * CHUNK_SIZE * CHUNK_SIZE + 15 * CHUNK_SIZE + 15; + const outsideIndex = 0 * CHUNK_SIZE * CHUNK_SIZE + 2 * CHUNK_SIZE + 2; + return arr[insideIndex] === 128 && arr[outsideIndex] === 0; + }, "Verify flood fill result"); +});