From 7067c650e33025fe9b9dc77259fde0c2814c8630 Mon Sep 17 00:00:00 2001
From: Carlos Precioso <github@precioso.design>
Date: Fri, 14 Apr 2023 15:15:50 +0200
Subject: [PATCH] feat: add caches and keyed version

---
 src/cache/keyed.ts           | 16 ++++++++++++++++
 src/cache/single-value.ts    | 15 +++++++++++++++
 src/index.ts                 |  5 ++++-
 src/{ => lib}/cache-ref.ts   |  0
 src/{ => lib}/suspend.ts     |  8 ++++++--
 src/single-value.ts          | 32 --------------------------------
 src/suspense/keyed.ts        | 14 ++++++++++++++
 src/suspense/single-value.ts | 12 ++++++++++++
 8 files changed, 67 insertions(+), 35 deletions(-)
 create mode 100644 src/cache/keyed.ts
 create mode 100644 src/cache/single-value.ts
 rename src/{ => lib}/cache-ref.ts (100%)
 rename src/{ => lib}/suspend.ts (71%)
 delete mode 100644 src/single-value.ts
 create mode 100644 src/suspense/keyed.ts
 create mode 100644 src/suspense/single-value.ts

diff --git a/src/cache/keyed.ts b/src/cache/keyed.ts
new file mode 100644
index 0000000..092665a
--- /dev/null
+++ b/src/cache/keyed.ts
@@ -0,0 +1,16 @@
+import { CacheValue } from "../lib/cache-ref";
+import { suspendOnPromise } from "../lib/suspend";
+
+export interface KeyedCacheStorage<K, T> {
+  get(key: K): CacheValue<T> | undefined | null;
+  set(key: K, value: CacheValue<T>): void;
+}
+
+export const createKeyedCache = <K, T>(
+  storage: KeyedCacheStorage<K, T> = new Map()
+) => {
+  const use = (key: K, fn: () => Promise<T>) =>
+    suspendOnPromise(fn, storage.get(key), (value) => storage.set(key, value));
+
+  return { use, storage };
+};
diff --git a/src/cache/single-value.ts b/src/cache/single-value.ts
new file mode 100644
index 0000000..1951b18
--- /dev/null
+++ b/src/cache/single-value.ts
@@ -0,0 +1,15 @@
+import { CacheValue } from "../lib/cache-ref";
+import { suspendOnPromise } from "../lib/suspend";
+
+export const createSingleValueCache = <T>() => {
+  const storage = {
+    current: null as CacheValue<T> | null,
+  };
+
+  const use = (fn: () => Promise<T>) =>
+    suspendOnPromise(fn, storage.current, (value) => {
+      storage.current = value;
+    });
+
+  return { use, storage };
+};
diff --git a/src/index.ts b/src/index.ts
index 2564bc1..6f0fbe0 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1 +1,4 @@
-export { createSingleValueSuspense } from "./single-value";
+export { createKeyedCache } from "./cache/keyed";
+export { createSingleValueCache } from "./cache/single-value";
+export { createKeyedSuspense } from "./suspense/keyed";
+export { createSingleValueSuspense } from "./suspense/single-value";
diff --git a/src/cache-ref.ts b/src/lib/cache-ref.ts
similarity index 100%
rename from src/cache-ref.ts
rename to src/lib/cache-ref.ts
diff --git a/src/suspend.ts b/src/lib/suspend.ts
similarity index 71%
rename from src/suspend.ts
rename to src/lib/suspend.ts
index 45bb0e0..6bf4e10 100644
--- a/src/suspend.ts
+++ b/src/lib/suspend.ts
@@ -1,12 +1,16 @@
 import { CacheValue } from "./cache-ref";
 
+// We wrap the `fn` in an async function to avoid
+// throwing sync errors.
+const callPromiseFn = async <T>(fn: () => Promise<T>) => await fn();
+
 export const suspendOnPromise = <T>(
   fn: () => Promise<T>,
-  cachedValue: CacheValue<T> | null,
+  cachedValue: CacheValue<T> | undefined | null,
   storeInCache: (value: CacheValue<T>) => void
 ) => {
   if (!cachedValue) {
-    const promise = fn().then(
+    const promise = callPromiseFn(fn).then(
       (value) => storeInCache({ status: "fulfilled", value }),
       (error) => storeInCache({ status: "rejected", error })
     );
diff --git a/src/single-value.ts b/src/single-value.ts
deleted file mode 100644
index 80567b6..0000000
--- a/src/single-value.ts
+++ /dev/null
@@ -1,32 +0,0 @@
-import { useSyncExternalStore } from "react";
-import { CacheValue } from "./cache-ref";
-import { suspendOnPromise } from "./suspend";
-
-const subscribe = (cb: () => void) => {
-  cb();
-  return () => {};
-};
-
-const makeGetSnapshot = <T>(getValue: () => Promise<T>) => {
-  let value: CacheValue<T> | null;
-
-  const getSnapshot = () =>
-    suspendOnPromise(getValue, value, (newValue) => {
-      value = newValue;
-    });
-
-  return getSnapshot;
-};
-
-export const createSingleValueSuspense = <T>(
-  getClientValue: () => Promise<T>,
-  getServerValue?: () => Promise<T>
-) => {
-  const getClientSnapshot = makeGetSnapshot(getClientValue);
-  const getServerSnapshot = getServerValue && makeGetSnapshot(getServerValue);
-
-  const useSingleValueSuspense = () =>
-    useSyncExternalStore(subscribe, getClientSnapshot, getServerSnapshot);
-
-  return useSingleValueSuspense;
-};
diff --git a/src/suspense/keyed.ts b/src/suspense/keyed.ts
new file mode 100644
index 0000000..466e683
--- /dev/null
+++ b/src/suspense/keyed.ts
@@ -0,0 +1,14 @@
+import { createKeyedCache } from "../cache/keyed";
+
+export const createKeyedSuspense = <K extends string, T>(
+  fn: (key: K) => Promise<T>
+) => {
+  const cache = createKeyedCache<K, T>();
+
+  const useKeyedSuspense = (key: K) => {
+    const value = cache.use(key, () => fn(key));
+    return { value, cache };
+  };
+
+  return useKeyedSuspense;
+};
diff --git a/src/suspense/single-value.ts b/src/suspense/single-value.ts
new file mode 100644
index 0000000..e05c204
--- /dev/null
+++ b/src/suspense/single-value.ts
@@ -0,0 +1,12 @@
+import { createSingleValueCache } from "../cache/single-value";
+
+export const createSingleValueSuspense = <T>(fn: () => Promise<T>) => {
+  const cache = createSingleValueCache<T>();
+
+  const useSingleValueSuspense = () => {
+    const value = cache.use(fn);
+    return { value, cache };
+  };
+
+  return useSingleValueSuspense;
+};