Line data    Source code

       1              : import * as React from 'react';
       2              : 
       3              : /// A state container outside of React.
       4              : ///
       5              : /// React may invoke the updater of a state hook more than once, e.g. when it rebases updates or
       6              : /// double-invokes them in strict mode. The registries free Wasm objects and schedule storage writes
       7              : /// while they reduce, so their updates must run exactly once. A store applies every update right
       8              : /// away, synchronously and in call order, and notifies its subscribers afterwards.
       9            1 : export class Store<T> {
      10              :     /// The current state
      11            1 :     current: T;
      12              :     /// The subscribers
      13            1 :     private readonly listeners: Set<() => void>;
      14              : 
      15            1 :     constructor(initial: T) {
      16            1 :         this.current = initial;
      17            1 :         this.listeners = new Set();
      18            1 :     }
      19              : 
      20              :     /// Subscribe to state changes
      21            1 :     subscribe = (listener: () => void): (() => void) => {
      22            1 :         this.listeners.add(listener);
      23            1 :         return () => { this.listeners.delete(listener); };
      24            1 :     };
      25              : 
      26              :     /// Get the current state
      27            1 :     get = (): T => this.current;
      28              : 
      29              :     /// Replace the state or derive it from the current one. The updater runs exactly once.
      30            1 :     update = (action: React.SetStateAction<T>): void => {
      31            1 :         const next = typeof action === 'function'
      32            1 :             ? (action as (prev: T) => T)(this.current)
      33            1 :             : action;
      34            1 :         if (Object.is(next, this.current)) return;
      35            1 :         this.current = next;
      36            1 :         for (const listener of this.listeners) {
      37            1 :             listener();
      38              :         }
      39            1 :     };
      40            1 : }
      41              : 
      42              : /// Batches the updates of a store and applies them in a single store update, in the order they were pushed.
      43              : ///
      44              : /// Updates are applied at the end of the current task. Updates that may wait for the next frame, e.g.
      45              : /// the progress of a result stream, are applied once per frame. An update that must not wait applies
      46              : /// the queued updates right away to preserve their order.
      47            1 : export class UpdateQueue<T> {
      48              :     /// The store
      49            1 :     private readonly store: Store<T>;
      50              :     /// The pending updates
      51            1 :     private pending: Array<(state: T) => T>;
      52              :     /// The scheduled flush
      53            1 :     private scheduled: { type: 'microtask' } | { type: 'frame'; handle: number } | null;
      54              : 
      55            1 :     constructor(store: Store<T>) {
      56            1 :         this.store = store;
      57            1 :         this.pending = [];
      58            1 :         this.scheduled = null;
      59            1 :     }
      60              : 
      61              :     /// Queue an update
      62            1 :     push(update: (state: T) => T, waitForFrame: boolean = false): void {
      63            1 :         this.pending.push(update);
      64            1 :         if (waitForFrame) {
      65            1 :             if (this.scheduled == null) {
      66            1 :                 this.scheduled = { type: 'frame', handle: requestAnimationFrame(this.flush) };
      67              :             }
      68            1 :         } else if (this.scheduled?.type === 'frame') {
      69            1 :             cancelAnimationFrame(this.scheduled.handle);
      70            1 :             this.flush();
      71            1 :         } else if (this.scheduled == null) {
      72            1 :             this.scheduled = { type: 'microtask' };
      73            1 :             queueMicrotask(this.flush);
      74            1 :         }
      75            1 :     }
      76              : 
      77              :     /// Apply the pending updates. An update that throws is skipped, the others still apply.
      78            1 :     flush = (): void => {
      79            1 :         this.scheduled = null;
      80            1 :         const pending = this.pending;
      81            1 :         if (pending.length === 0) return;
      82            1 :         this.pending = [];
      83            1 :         this.store.update(state => pending.reduce((s, update) => {
      84            1 :             try {
      85            1 :                 return update(s);
      86            1 :             } catch (e: unknown) {
      87            1 :                 console.error('Store update failed', e);
      88            1 :                 return s;
      89            1 :             }
      90            1 :         }, state));
      91            1 :     };
      92              : 
      93              :     /// Drop the pending updates
      94            1 :     dispose(): void {
      95            1 :         if (this.scheduled?.type === 'frame') {
      96            1 :             cancelAnimationFrame(this.scheduled.handle);
      97              :         }
      98            1 :         this.scheduled = null;
      99            1 :         this.pending = [];
     100            1 :     }
     101            1 : }
     102              : 
     103              : /// Create the update queue of a store that lives as long as the component
     104            1 : export function useUpdateQueue<T>(store: Store<T>): UpdateQueue<T> {
     105            1 :     const [queue] = React.useState(() => new UpdateQueue(store));
     106            1 :     React.useEffect(() => () => queue.dispose(), [queue]);
     107            1 :     return queue;
     108            1 : }
     109              : 
     110              : /// Create a store that lives as long as the component
     111            1 : export function useStore<T>(initial: () => T): Store<T> {
     112            1 :     const [store] = React.useState(() => new Store(initial()));
     113            1 :     return store;
     114            1 : }
     115              : 
     116              : /// Subscribe to the state of a store
     117            1 : export function useStoreState<T>(store: Store<T>): T {
     118            1 :     return React.useSyncExternalStore(store.subscribe, store.get);
     119            1 : }