From ec04cd05de1bee6dc70416d4843c3e56487ea878 Mon Sep 17 00:00:00 2001
From: Trevor Manz <trevor.j.manz@gmail.com>
Date: Fri, 19 Mar 2021 10:26:32 -0400
Subject: [PATCH] Add experimental FileReferenceStore (#82)

* Try out FileReferenceStore

* Pass abort signal from Viv
---
 src/io.ts      |  3 +--
 src/storage.ts | 73 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/utils.ts   | 12 ++++++---
 3 files changed, 83 insertions(+), 5 deletions(-)
 create mode 100644 src/storage.ts

diff --git a/src/io.ts b/src/io.ts
index d57ceaf5..6bb9f212 100644
--- a/src/io.ts
+++ b/src/io.ts
@@ -59,7 +59,6 @@ function loadSingleChannel(config: SingleChannelConfig, data: ZarrPixelSource<st
 function loadMultiChannel(config: MultichannelConfig, data: ZarrPixelSource<string[]>[], max: number): SourceData {
   const { names, channel_axis, name, model_matrix, opacity = 1, colormap = '' } = config;
   let { contrast_limits, visibilities, colors } = config;
-
   const n = data[0].shape[channel_axis as number];
   for (const channelProp of [contrast_limits, visibilities, names, colors]) {
     if (channelProp && channelProp.length !== n) {
@@ -101,7 +100,7 @@ function loadMultiChannel(config: MultichannelConfig, data: ZarrPixelSource<stri
   return {
     loader: data,
     name,
-    channel_axis: channel_axis as number,
+    channel_axis: Number(channel_axis as number),
     colors,
     names: names ?? range(n).map((i) => `channel_${i}`),
     contrast_limits: contrast_limits ?? Array(n).fill([0, max]),
diff --git a/src/storage.ts b/src/storage.ts
new file mode 100644
index 00000000..dbdf38b8
--- /dev/null
+++ b/src/storage.ts
@@ -0,0 +1,73 @@
+import type { AsyncStore } from 'zarr/types/storage/types';
+import { KeyError, HTTPError } from 'zarr';
+
+type Ref = string | [url: string] | [url: string, offset: number, length: number];
+
+const encoder = new TextEncoder();
+
+export class FileReferenceStore implements AsyncStore<ArrayBuffer> {
+  constructor(public ref: Map<string, Ref>) {}
+
+  static async fromUrl(url: string) {
+    const json: Record<string, Ref> = await fetch(url).then((res) => res.json());
+    if ('version' in json) {
+      throw Error('Only v0 ReferenceFileSystem description is currently supported!');
+    }
+    const ref = new Map(Object.entries(json));
+    return new FileReferenceStore(ref);
+  }
+
+  _url(url: string) {
+    const [protocol, _path] = url.split('://');
+    if (protocol === 'https' || protocol === 'http') {
+      return url;
+    }
+    throw Error('Protocol not supported, got: ' + JSON.stringify(protocol));
+  }
+
+  _fetch({ url, offset, size }: { url: string; offset?: number; size?: number }, opts: RequestInit) {
+    if (offset !== undefined && size !== undefined) {
+      // add range headers to request options
+      opts = { ...opts, headers: { ...opts.headers, Range: `bytes=${offset}-${offset + size - 1}` } };
+    }
+    return fetch(this._url(url), opts);
+  }
+
+  async getItem(key: string, opts: RequestInit = {}) {
+    const entry = this.ref.get(key);
+
+    if (!entry) {
+      throw new KeyError(key);
+    }
+
+    if (typeof entry === 'string') {
+      // JSON data entry in reference
+      return encoder.encode(entry).buffer;
+    }
+
+    const [url, offset, size] = entry;
+    const res = await this._fetch({ url, offset, size }, opts);
+
+    if (res.status === 200 || res.status === 206) {
+      return res.arrayBuffer();
+    }
+
+    throw new HTTPError(`Request unsuccessful for key ${key}. Response status: ${res.status}.`);
+  }
+
+  async containsItem(key: string) {
+    return this.ref.has(key);
+  }
+
+  keys() {
+    return Promise.resolve([...this.ref.keys()]);
+  }
+
+  setItem(key: string, value: ArrayBuffer): never {
+    throw Error('FileReferenceStore.setItem is not implemented.');
+  }
+
+  deleteItem(key: string): never {
+    throw Error('FileReferenceStore.deleteItem is not implemented.');
+  }
+}
diff --git a/src/utils.ts b/src/utils.ts
index fe99f220..ecd902ce 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -2,6 +2,8 @@ import { ContainsArrayError, HTTPStore, openArray, openGroup, ZarrArray } from '
 import type { Group as ZarrGroup } from 'zarr';
 import { Matrix4 } from '@math.gl/core/dist/esm';
 
+import { FileReferenceStore } from './storage';
+
 export const MAX_CHANNELS = 6;
 
 export const COLORS = {
@@ -17,16 +19,20 @@ export const MAGENTA_GREEN = [COLORS.magenta, COLORS.green];
 export const RGB = [COLORS.red, COLORS.green, COLORS.blue];
 export const CYMRGB = Object.values(COLORS).slice(0, -2);
 
-function normalizeStore(source: string | ZarrArray['store']) {
+async function normalizeStore(source: string | ZarrArray['store']) {
   if (typeof source === 'string') {
+    if (source.endsWith('.json')) {
+      const store = await FileReferenceStore.fromUrl(source);
+      return { store };
+    }
     const [root, path] = source.split('.zarr');
     return { store: new HTTPStore(root + '.zarr'), path };
   }
-  return { store: source, path: '' };
+  return { store: source };
 }
 
 export async function open(source: string | ZarrArray['store']) {
-  const { store, path } = normalizeStore(source);
+  const { store, path } = await normalizeStore(source);
   return openGroup(store, path).catch((err) => {
     if (err instanceof ContainsArrayError) {
       return openArray({ store, path });