feat: refactor to share chunk data across multiple layers/viewports (#531)

### Summary
This is picking up where #514 left off, basically just that PR rebased
on `main`.

This PR introduces a `ChunkStore`/`ChunkStoreView` architecture to
enable efficient chunk sharing across multiple viewports. Currently,
`ChunkManagerSource` conflates two responsibilities: managing a
canonical chunk data store (for a given source) and determining
viewport-specific chunk visibility. This design prevents resource
sharing when multiple viewports render overlapping portions of the same
data source.

The implementation splits `ChunkManagerSource` into two complementary
classes:
- `ChunkStore`: Manages the canonical data store and global state for
all chunks for a single source by aggregating view state across
associated `ChunkStoreView` objects
- `ChunkStoreView`: Represents a layer/viewport-specific "view" into
that store, calculating LOD, visibility, and prefetch requirements based
on its camera and slice position

Multiple layers (including in separate viewports) can now share the same
`ChunkStore` while each maintains its own `ChunkStoreView`. The
`ChunkStore` aggregates chunk requirements across all views: chunks are
kept loaded if any view needs them, and disposed when no views reference
them. Priority for loading is determined by taking the minimum priority
across all views.

`ChunkStoreView` maintains an efficient sparse `Map<Chunk,
ChunkViewState>` rather than tracking all chunks. View state aggregation
happens synchronously during the update cycle, with `maybeForgetChunk()`
called inline to clean up no-longer-visible chunks within a view.

`chunkManager.update()` is now called *after* rendering, as chunk view
states are updated as the runtime iterates over viewports to render.
This performs the view state aggregation and ultimately chunk queueing
and discarding.

### Related Issue
For architecture diagrams and class relationships, see the [multiple
viewports tech
spec](https://docs.google.com/document/d/1OwGL7KW34hxFAENyOm2WiQ71E58XGu9Y7jKkJlyRx84/edit?tab=t.0).

This is the final PR of the multi-viewport implementation plan (see tech
spec):
  - #407: WebGL infrastructure
  - #419: Viewport class
  - #451: Viewport events
  - #506: Idetik API integration
  - #531 (this PR): `ChunkManagerSource` refactor

### Tests & Checks
- Existing chunk management tests updated to use new
`ChunkStore`/`ChunkStoreView` API
  - All tests pass
  - Existing examples work without modification
- Added a second `ChunkedImageLayer` to the multi-viewport example and
confirmed it behaves as expected, with chunk memory sharing confirmed
(according to stats.js)
- Validated that chunks are properly disposed when no views need them
(memory goes back down)

Author: aganders3

packages/core/examples/chunk_streaming/chunk_info_overlay.ts

@@ -18,21 +18,26 @@ export class ChunkInfoOverlay {
 
   public update(idetik: Idetik, _timestamp?: DOMHighResTimeStamp): void {
     if (this.textDiv_.style.display === "none") return;
-    const chunkManagerSource = this.imageLayer_.chunkManagerSource;
-    if (!chunkManagerSource) {
-      this.textDiv_.textContent = "No chunk manager source";
+    const chunkStore = this.imageLayer_.chunkStoreView?.store;
+    const chunkStoreView = this.imageLayer_.chunkStoreView;
+    if (!chunkStore || !chunkStoreView) {
+      this.textDiv_.textContent = "No chunk store";
       return;
     }
 
-    const chunksAtCurrentTime = chunkManagerSource.getChunksAtCurrentTime();
+    const chunksAtCurrentTime = chunkStore.getChunksAtTime(
+      chunkStore.getTimeIndex(this.imageLayer_.sliceCoords)
+    );
     if (!chunksAtCurrentTime) {
       this.textDiv_.textContent = "No chunks available";
       return;
     }
 
     const chunkDetails: string[] = [];
-    const currentLOD = chunkManagerSource.currentLOD;
-    const renderedChunks = chunkManagerSource.getChunks();
+    const currentLOD = chunkStoreView.currentLOD;
+    const renderedChunks = chunkStoreView.getChunksToRender(
+      this.imageLayer_.sliceCoords
+    );
     const totalChunks = chunksAtCurrentTime.length;
 
     let loadedChunks = 0;
@@ -41,7 +46,7 @@ export class ChunkInfoOverlay {
       visible: number;
       rendered: number;
       prefetched: number;
-    }[] = Array.from({ length: chunkManagerSource.lodCount }, () => ({
+    }[] = Array.from({ length: chunkStore.lodCount }, () => ({
       visible: 0,
       rendered: 0,
       prefetched: 0,

packages/core/examples/multi_viewport/index.html

@@ -18,6 +18,17 @@
         height: 100vh;
       }
 
+      #viewport-left {
+        flex: 1;
+        border: 3px solid black;
+      }
+
+      #right-panel {
+        flex: 1;
+        display: flex;
+        flex-direction: column;
+      }
+
       .viewport {
         flex: 1;
         border: 3px solid black;
@@ -39,7 +50,10 @@
   <body>
     <div id="container">
       <div class="viewport" id="viewport-left"></div>
-      <div class="viewport" id="viewport-right"></div>
+      <div id="right-panel">
+        <div class="viewport" id="viewport-top-right"></div>
+        <div class="viewport" id="viewport-bottom-right"></div>
+      </div>
     </div>
     <canvas id="canvas"></canvas>
     <script type="module" src="./main.ts"></script>

packages/core/examples/multi_viewport/main.ts

@@ -26,21 +26,29 @@ const zMin = z.translate;
 const zMax = z.translate + z.scale * z.shape - z.scale;
 const zRange = { min: zMin, max: zMax };
 
+// shared source between viewports
 const source = new OmeZarrImageSource(url);
-const sliceCoords = { t: 400, z: 200, c: 0 };
-const camera2D = new OrthographicCamera(left, right, top, bottom);
-const imageLayer = new ChunkedImageLayer({
+
+const camera3D = new PerspectiveCamera();
+const volumeLayer = new VolumeLayer();
+
+const sliceCoords1 = { t: 400, z: 200, c: 0 };
+const camera2D1 = new OrthographicCamera(left, right, top, bottom);
+const imageLayer1 = new ChunkedImageLayer({
   source,
-  sliceCoords,
+  sliceCoords: sliceCoords1,
   policy: createExplorationPolicy(),
   channelProps: [{ contrastLimits: [0, 200] }],
 });
 
-// TODO: the reason this example works is that the volume viewport uses a perspective camera
-// otherwise the ChunkManager update causes interference between the two viewports
-// this example can be updated to include multiple views of the same volume once we have
-// better handling of multiple viewports using the same source
-const camera3D = new PerspectiveCamera();
+const sliceCoords2 = { t: 400, z: 300, c: 0 };
+const camera2D2 = new OrthographicCamera(left, right, top, bottom);
+const imageLayer2 = new ChunkedImageLayer({
+  source,
+  sliceCoords: sliceCoords2,
+  policy: createExplorationPolicy(),
+  channelProps: [{ contrastLimits: [0, 200] }],
+});
 
 new Idetik({
   canvas: document.querySelector<HTMLCanvasElement>("#canvas")!,
@@ -50,25 +58,50 @@ new Idetik({
       element: document.querySelector<HTMLDivElement>("#viewport-left")!,
       camera: camera3D,
       cameraControls: new OrbitControls(camera3D, { radius: 3 }),
-      layers: [new VolumeLayer()],
+      layers: [volumeLayer],
+    },
+    {
+      id: "slice1",
+      element: document.querySelector<HTMLDivElement>("#viewport-top-right")!,
+      camera: camera2D1,
+      cameraControls: new PanZoomControls(camera2D1),
+      layers: [imageLayer1],
     },
     {
-      id: "slice",
-      element: document.querySelector<HTMLDivElement>("#viewport-right")!,
-      camera: camera2D,
-      cameraControls: new PanZoomControls(camera2D),
-      layers: [imageLayer],
+      id: "slice2",
+      element: document.querySelector<HTMLDivElement>(
+        "#viewport-bottom-right"
+      )!,
+      camera: camera2D2,
+      cameraControls: new PanZoomControls(camera2D2),
+      layers: [imageLayer2],
     },
   ],
   showStats: true,
 }).start();
 
 const gui = new GUI({ width: 300 });
+
+const topRightViewportFolder = gui.addFolder("Top Right Viewport (Slice 1)");
+addDimensionSlider({
+  gui: topRightViewportFolder,
+  sliceCoords: sliceCoords1,
+  dimensionName: "z",
+  minValue: zRange.min,
+  maxValue: zRange.max,
+  stepValue: z.scale,
+});
+topRightViewportFolder.open();
+
+const bottomRightViewportFolder = gui.addFolder(
+  "Bottom Right Viewport (Slice 2)"
+);
 addDimensionSlider({
-  gui,
-  sliceCoords,
+  gui: bottomRightViewportFolder,
+  sliceCoords: sliceCoords2,
   dimensionName: "z",
   minValue: zRange.min,
   maxValue: zRange.max,
   stepValue: z.scale,
 });
+bottomRightViewportFolder.open();

packages/core/src/core/chunk_manager.ts

@@ -1,76 +1,181 @@
-import { ChunkSource, SliceCoordinates } from "../data/chunk";
-import { OrthographicCamera } from "../objects/cameras/orthographic_camera";
+import { Logger } from "../utilities/logger";
+import { Chunk, ChunkSource } from "../data/chunk";
 import { ChunkQueue } from "../data/chunk_queue";
-import { ChunkManagerSource } from "./chunk_manager_source";
+import { ChunkStore } from "./chunk_store";
+import { ChunkStoreView } from "./chunk_store_view";
 import { ImageSourcePolicy } from "./image_source_policy";
 
 export class ChunkManager {
-  private readonly sources_ = new Map<ChunkSource, ChunkManagerSource>();
-  private readonly pendingSources_ = new Map<
-    ChunkSource,
-    Promise<ChunkManagerSource>
-  >();
+  private readonly stores_ = new Map<ChunkSource, ChunkStore>();
+  private readonly pendingStores_ = new Map<ChunkSource, Promise<ChunkStore>>();
+  private readonly views_ = new Map<ChunkStore, ChunkStoreView[]>();
   private readonly queue_ = new ChunkQueue();
 
-  public async addSource(
+  public async addView(
     source: ChunkSource,
-    sliceCoords: SliceCoordinates,
     policy: ImageSourcePolicy
-  ) {
+  ): Promise<ChunkStoreView> {
+    const store = await this.addSource(source);
+    const view = new ChunkStoreView(store, policy);
+    this.views_.set(store, (this.views_.get(store) ?? []).concat(view));
+    return view;
+  }
+
+  public removeView(view: ChunkStoreView): void {
+    const store = view.store;
+    const source = this.getSourceForStore(store);
+
+    const views = this.views_.get(store);
+    if (!views) {
+      throw new Error("Cannot remove view: store not managed by ChunkManager");
+    }
+
+    const index = views.indexOf(view);
+    if (index === -1) {
+      throw new Error(
+        `Cannot remove view: view not found in store's view list (source: ${source})`
+      );
+    }
+
+    const affectedChunks = Array.from(view.chunkViewStates.keys());
+    views.splice(index, 1);
+
+    for (const chunk of affectedChunks) {
+      this.aggregateChunkViewStates(chunk, store);
+    }
+
+    if (views.length === 0) {
+      this.stores_.delete(source);
+      this.views_.delete(store);
+    }
+  }
+
+  private async addSource(source: ChunkSource): Promise<ChunkStore> {
     const existingOrPending =
-      this.sources_.get(source) ?? this.pendingSources_.get(source);
+      this.stores_.get(source) ?? this.pendingStores_.get(source);
     if (existingOrPending) {
       return existingOrPending;
     }
 
-    const initializeSource = async () => {
+    const initializeStore = async () => {
       const loader = await source.open();
-      const chunkManagerSource = new ChunkManagerSource(
-        loader,
-        sliceCoords,
-        policy
-      );
-      this.sources_.set(source, chunkManagerSource);
-      this.pendingSources_.delete(source);
-      return chunkManagerSource;
+      const store = new ChunkStore(loader);
+      this.stores_.set(source, store);
+      this.pendingStores_.delete(source);
+      return store;
     };
 
-    const pending = initializeSource();
-    this.pendingSources_.set(source, pending);
+    const pending = initializeStore();
+    this.pendingStores_.set(source, pending);
     return pending;
   }
 
-  public update(camera: OrthographicCamera, bufferWidth: number) {
-    if (this.sources_.size === 0) return;
-
-    if (camera.type !== "OrthographicCamera") {
-      throw new Error(
-        "ChunkManager currently supports only orthographic cameras. " +
-          "Update the implementation before using a perspective camera."
-      );
+  private getSourceForStore(store: ChunkStore): ChunkSource {
+    for (const [source, s] of this.stores_) {
+      if (s === store) {
+        return source;
+      }
     }
+    throw new Error("Source not found for the given store.");
+  }
 
-    const viewBounds2D = camera.getWorldViewRect();
-    const virtualWidth = Math.abs(viewBounds2D.max[0] - viewBounds2D.min[0]);
-    const virtualUnitsPerScreenPixel = virtualWidth / bufferWidth;
-    const lodFactor = Math.log2(1 / virtualUnitsPerScreenPixel);
+  public update() {
+    for (const [_, store] of this.stores_) {
+      const updatedChunks = this.updateAndCollectChunkChanges(store);
 
-    for (const [_, source] of this.sources_) {
-      const updatedChunks = source.updateAndCollectChunkChanges(
-        lodFactor,
-        viewBounds2D
-      );
       for (const chunk of updatedChunks) {
         if (chunk.priority === null) {
           this.queue_.cancel(chunk);
         } else if (chunk.state === "queued") {
           this.queue_.enqueue(chunk, (signal) =>
-            source.loadChunkData(chunk, signal)
+            store.loadChunkData(chunk, signal)
           );
         }
       }
     }
 
     this.queue_.flush();
   }
+
+  private updateAndCollectChunkChanges(store: ChunkStore): Set<Chunk> {
+    const views = this.views_.get(store);
+    if (!views) return new Set<Chunk>();
+    const affectedChunks = new Set<Chunk>();
+    for (const view of views) {
+      for (const [chunk, _viewState] of view.chunkViewStates) {
+        affectedChunks.add(chunk);
+      }
+    }
+
+    for (const chunk of affectedChunks) {
+      this.aggregateChunkViewStates(chunk, store);
+    }
+
+    return affectedChunks;
+  }
+
+  private aggregateChunkViewStates(chunk: Chunk, store: ChunkStore): void {
+    const views = this.views_.get(store);
+    if (!views) return;
+    let anyVisible = false;
+    let anyPrefetch = false;
+    let minPriority: number | null = null;
+    let orderKeyForMinPriority: number | null = null;
+
+    for (const view of views) {
+      const viewState = view.chunkViewStates.get(chunk);
+      if (!viewState) continue;
+
+      if (viewState.visible) anyVisible = true;
+      if (viewState.prefetch) anyPrefetch = true;
+
+      if (viewState.priority !== null) {
+        if (minPriority === null || viewState.priority < minPriority) {
+          minPriority = viewState.priority;
+          orderKeyForMinPriority = viewState.orderKey;
+        }
+      }
+
+      if (
+        !viewState.visible &&
+        !viewState.prefetch &&
+        viewState.priority === null
+      ) {
+        view.maybeForgetChunk(chunk);
+      }
+    }
+
+    chunk.visible = anyVisible;
+    chunk.prefetch = anyPrefetch;
+    chunk.priority = minPriority;
+    chunk.orderKey = orderKeyForMinPriority;
+
+    const shouldEnqueueChunk =
+      chunk.priority !== null && chunk.state === "unloaded";
+    if (shouldEnqueueChunk) {
+      chunk.state = "queued";
+      return;
+    }
+
+    const shouldCancelQueuedChunk =
+      chunk.priority === null && chunk.state === "queued";
+    if (shouldCancelQueuedChunk) {
+      chunk.state = "unloaded";
+      return;
+    }
+
+    const shouldDisposeChunk =
+      chunk.state === "loaded" && !chunk.visible && !chunk.prefetch;
+    if (shouldDisposeChunk) {
+      chunk.data = undefined;
+      chunk.state = "unloaded";
+      chunk.priority = null;
+      chunk.orderKey = null;
+      chunk.prefetch = false;
+      Logger.debug(
+        "ChunkManager",
+        `Disposing chunk ${JSON.stringify(chunk.chunkIndex)} in LOD ${chunk.lod}`
+      );
+    }
+  }
 }