From 41ce62d6ce100d5a66bc1555a3584ba8a48798e4 Mon Sep 17 00:00:00 2001
From: Trevor Manz <trevor.j.manz@gmail.com>
Date: Sun, 4 Jul 2021 11:47:14 -0400
Subject: [PATCH] Add simple cache (#100)

---
 package-lock.json |  5 +++++
 package.json      |  1 +
 src/lru-store.ts  | 36 ++++++++++++++++++++++++++++++++++++
 src/utils.ts      | 23 ++++++++++++++++++-----
 4 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 src/lru-store.ts

diff --git a/package-lock.json b/package-lock.json
index defec8e5..ce1c9eae 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -2741,6 +2741,11 @@
       "integrity": "sha512-N5ZAX4/LxJmF+7wN74pUD6qAh9/wnvdQcjq9TZjevvXzSUo7bfmw91saqMjzGS2xq91/odN2dW/WOl7qQHNDGA==",
       "dev": true
     },
+    "quick-lru": {
+      "version": "6.0.0",
+      "resolved": "https://registry.npmjs.org/quick-lru/-/quick-lru-6.0.0.tgz",
+      "integrity": "sha512-kQ3ACYYIALMi7x8xvMU+qSF6N6baMZms8/q3dn2XFiORhFrLbpR6vTRNGR80HwlakWSoY4RnGoCJWUgnyH6zZQ=="
+    },
     "quickselect": {
       "version": "2.0.0",
       "resolved": "https://registry.npmjs.org/quickselect/-/quickselect-2.0.0.tgz",
diff --git a/package.json b/package.json
index 75f96604..908b564d 100644
--- a/package.json
+++ b/package.json
@@ -9,6 +9,7 @@
     "imjoy-rpc": "^0.2.23",
     "jotai": "^1.0.0",
     "p-map": "^4.0.0",
+    "quick-lru": "^6.0.0",
     "react": "^17.0.1",
     "react-dom": "^17.0.1",
     "reference-spec-reader": "^0.1.1",
diff --git a/src/lru-store.ts b/src/lru-store.ts
new file mode 100644
index 00000000..c183219c
--- /dev/null
+++ b/src/lru-store.ts
@@ -0,0 +1,36 @@
+import type { AsyncStore } from 'zarr/types/storage/types';
+import QuickLRU from 'quick-lru';
+
+export class LRUCacheStore<S extends AsyncStore<ArrayBuffer>> {
+  private cache: QuickLRU<string, ArrayBuffer>;
+
+  constructor(public store: S, public maxSize: number = 100) {
+    this.cache = new QuickLRU({ maxSize });
+  }
+
+  async getItem(...args: Parameters<S['getItem']>) {
+    const [key, opts] = args;
+    if (this.cache.has(key)) {
+      return this.cache.get(key)!;
+    }
+    const value = await this.store.getItem(key, opts);
+    this.cache.set(key, value);
+    return value;
+  }
+
+  async containsItem(key: string) {
+    return this.cache.has(key) || this.store.containsItem(key);
+  }
+
+  async keys() {
+    return [];
+  }
+
+  deleteItem(key: string): never {
+    throw new Error('deleteItem not implemented');
+  }
+
+  setItem(key: string, value: ArrayBuffer): never {
+    throw new Error('setItem not implemented');
+  }
+}
diff --git a/src/utils.ts b/src/utils.ts
index 6279fe60..56d03950 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -1,6 +1,8 @@
 import { ContainsArrayError, HTTPStore, openArray, openGroup, ZarrArray } from 'zarr';
 import type { Group as ZarrGroup } from 'zarr';
+import type { AsyncStore, Store } from 'zarr/types/storage/types';
 import { Matrix4 } from '@math.gl/core/dist/esm';
+import { LRUCacheStore } from './lru-store';
 
 export const MAX_CHANNELS = 6;
 
@@ -17,19 +19,30 @@ 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);
 
-async function normalizeStore(source: string | ZarrArray['store']) {
+async function normalizeStore(source: string | Store) {
   if (typeof source === 'string') {
+    let store: AsyncStore<ArrayBuffer>;
+
     if (source.endsWith('.json')) {
       // import custom store implementation
-      const { ReferenceStore } = await import('reference-spec-reader');
-      return ReferenceStore.fromJSON(await fetch(source).then((res) => res.json()));
+      const [{ ReferenceStore }, json] = await Promise.all([
+        import('reference-spec-reader'),
+        fetch(source).then((res) => res.json()),
+      ]);
+
+      store = new ReferenceStore(json);
+    } else {
+      store = new HTTPStore(source);
     }
-    return new HTTPStore(source);
+
+    // Wrap remote stores in a cache
+    return new LRUCacheStore(store);
   }
+
   return source;
 }
 
-export async function open(source: string | ZarrArray['store']) {
+export async function open(source: string | Store) {
   const store = await normalizeStore(source);
   return openGroup(store).catch((err) => {
     if (err instanceof ContainsArrayError) {