Allow reading a store in its onUse callback

AmadeusITGroup · Nov 26, 2024 · 91e4089 · 91e4089
1 parent f4b0efa
commit 91e4089
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 24 deletions.
diff --git a/src/index.spec.ts b/src/index.spec.ts
@@ -1279,6 +1279,14 @@ describe('stores', () => {
       a.set(1);
       expect(scopes).toEqual([undefined]);
     });
+
+    it('should allow reading the store in onUse', () => {
+      const onUseValues: number[] = [];
+      const store = writable(0, () => {
+        onUseValues.push(store());
+      });
+      store();
+    });
   });
 
   describe('asWritable', () => {

diff --git a/src/internal/store.ts b/src/internal/store.ts
@@ -9,11 +9,10 @@ export const enum RawStoreFlags {
   // the following flags are used in RawStoreTrackingUsage and derived classes
   HAS_VISIBLE_ONUSE = 1,
   START_USE_CALLED = 1 << 1,
-  INSIDE_GET = 1 << 2,
-  FLUSH_PLANNED = 1 << 3,
+  FLUSH_PLANNED = 1 << 2,
   // the following flags are used in RawStoreComputedOrDerived and derived classes
-  COMPUTING = 1 << 4,
-  DIRTY = 1 << 5,
+  COMPUTING = 1 << 3,
+  DIRTY = 1 << 4,
 }
 
 export interface BaseLink<T> {

diff --git a/src/internal/storeComputed.ts b/src/internal/storeComputed.ts
@@ -77,7 +77,9 @@ export class RawStoreComputed<T>
         link.producer.registerConsumer(link);
       }
     }
-    this.flags |= RawStoreFlags.DIRTY;
+    if (this.epoch !== epoch) {
+      this.flags |= RawStoreFlags.DIRTY;
+    }
   }
 
   override endUse(): void {

diff --git a/src/internal/storeTrackingUsage.ts b/src/internal/storeTrackingUsage.ts
@@ -29,32 +29,35 @@ export const flushUnused = (): void => {
 };
 
 export abstract class RawStoreTrackingUsage<T> extends RawStoreWritable<T> {
+  private extraUsages = 0;
   abstract startUse(): void;
   abstract endUse(): void;
 
   override updateValue(): void {
+    // Ignoring coverage for the following lines because, unless there is a bug in tansu (which would have to be fixed!)
+    // there should be no way to trigger this error.
+    /* v8 ignore next 3 */
+    if (!(this.flags & RawStoreFlags.START_USE_CALLED)) {
+      throw new Error('assert failed: untracked producer usage');
+    }
+    super.updateValue();
+  }
+
+  override checkUsed(): void {
     const flags = this.flags;
     if (!(flags & RawStoreFlags.START_USE_CALLED)) {
-      // Ignoring coverage for the following lines because, unless there is a bug in tansu (which would have to be fixed!)
-      // there should be no way to trigger this error.
-      /* v8 ignore next 3 */
-      if (!(flags & RawStoreFlags.INSIDE_GET) && !this.consumerLinks?.length) {
-        throw new Error('assert failed: untracked producer usage');
-      }
       this.flags |= RawStoreFlags.START_USE_CALLED;
       untrack(() => this.startUse());
     }
   }
 
   override checkUnused(): void {
     const flags = this.flags;
-    // Ignoring coverage for the following lines because, unless there is a bug in tansu (which would have to be fixed!)
-    // there should be no way to trigger this error.
-    /* v8 ignore next 3 */
-    if (flags & RawStoreFlags.INSIDE_GET) {
-      throw new Error('assert failed: INSIDE_GET flag in checkUnused');
-    }
-    if (flags & RawStoreFlags.START_USE_CALLED && !this.consumerLinks?.length) {
+    if (
+      flags & RawStoreFlags.START_USE_CALLED &&
+      !this.consumerLinks?.length &&
+      !this.extraUsages
+    ) {
       if (inFlushUnused || flags & RawStoreFlags.HAS_VISIBLE_ONUSE) {
         this.flags &= ~RawStoreFlags.START_USE_CALLED;
         untrack(() => this.endUse());
@@ -74,11 +77,9 @@ export abstract class RawStoreTrackingUsage<T> extends RawStoreWritable<T> {
     if (activeConsumer) {
       return activeConsumer.addProducer(this);
     } else {
-      if (this.flags & RawStoreFlags.INSIDE_GET) {
-        throw new Error('recursive computed');
-      }
-      this.flags |= RawStoreFlags.INSIDE_GET;
+      this.extraUsages++;
       try {
+        this.checkUsed();
         this.updateValue();
         // Ignoring coverage for the following lines because, unless there is a bug in tansu (which would have to be fixed!)
         // there should be no way to trigger this error.
@@ -88,8 +89,11 @@ export abstract class RawStoreTrackingUsage<T> extends RawStoreWritable<T> {
         }
         return this.readValue();
       } finally {
-        this.flags &= ~RawStoreFlags.INSIDE_GET;
-        this.checkUnused();
+        const extraUsages = this.extraUsages - 1;
+        this.extraUsages = extraUsages;
+        if (extraUsages === 0) {
+          this.checkUnused();
+        }
       }
     }
   }

diff --git a/src/internal/storeWritable.ts b/src/internal/storeWritable.ts
@@ -79,6 +79,9 @@ export class RawStoreWritable<T> implements RawStore<T, ProducerConsumerLink<T>>
     const indexInProducer = consumerLinks.length;
     link.indexInProducer = indexInProducer;
     consumerLinks[indexInProducer] = link;
+    if (indexInProducer === 0) {
+      this.checkUsed();
+    }
     return link;
   }
 
@@ -102,6 +105,7 @@ export class RawStoreWritable<T> implements RawStore<T, ProducerConsumerLink<T>>
     }
   }
 
+  checkUsed(): void {}
   checkUnused(): void {}
   updateValue(): void {}