bug fix

Author: marcello3d

Diff for: src/cache.test.ts

@@ -15,7 +15,13 @@ describe('getOrAdd', () => {
     const fn = jest.fn(() => Promise.resolve(5));
     await expect(getOrAdd(cache, 'foo', fn)).resolves.toBe(5);
     expect(getOrAdd(cache, 'foo', fn)).toBe(5);
+    expect(fn).toHaveBeenCalledTimes(1);
+
+    const bar = getOrAdd(cache, 'bar', fn);
+    expect(bar).toBeInstanceOf(Promise);
+    expect(getOrAdd(cache, 'bar', fn)).toBe(bar);
     await expect(getOrAdd(cache, 'bar', fn)).resolves.toBe(5);
+    expect(getOrAdd(cache, 'bar', fn)).toBe(5);
     expect(fn).toHaveBeenCalledTimes(2);
   });
 });

Diff for: src/cache.ts

@@ -8,29 +8,28 @@
 
 import { PromiseOrValue } from './type';
 
-export type PromiseOrValueMap<K, V> = CacheLike<K, PromiseOrValue<V>>;
+export type PromiseOrValueMapLike<K, V> = {
+  get(key: K): PromiseOrValue<V> | undefined;
+  set(key: K, value: PromiseOrValue<V>): void;
+};
 
 export function getOrAdd<K, V>(
-  cache: PromiseOrValueMap<K, V>,
+  cache: PromiseOrValueMapLike<K, V>,
   key: K,
   compute: (key: K) => PromiseOrValue<V>,
 ): PromiseOrValue<V> {
   const existing = cache.get(key);
   if (existing) {
     return existing;
   }
-  const resultOrPromise = compute(key);
-  cache.set(key, resultOrPromise);
+  let resultOrPromise = compute(key);
   if (resultOrPromise instanceof Promise) {
-    return resultOrPromise.then((result) => {
+    // Return a promise that resolves after the cache has been updated
+    resultOrPromise = resultOrPromise.then((result) => {
       cache.set(key, result);
       return result;
     });
   }
+  cache.set(key, resultOrPromise);
   return resultOrPromise;
 }
-
-export type CacheLike<K, V> = {
-  get(key: K): V | undefined;
-  set(key: K, value: V): void;
-};