[WIP][POC] Add useOnyxState multi-key derived state hook

lib/OnyxSubscriptionManager.ts

@@ -16,6 +16,14 @@ type Listener<TKey extends OnyxKey = OnyxKey> = (value: OnyxValue<TKey>, key: TK
  */
 type GenericListener = (value: unknown, key: OnyxKey) => void;
 
+/** Listener fired when any of a state listener's declared dependency keys changes. Used by `useOnyxState`. */
+type StateListenerCallback = () => void;
+
+type StateListenerEntry = {
+    listener: StateListenerCallback;
+    deps: Set<OnyxKey>;
+};
+
 const reportedCacheMisses = new Set<OnyxKey>();
 
 function reportCacheMissIfStorageHasValue(key: OnyxKey): void {
@@ -43,17 +51,20 @@ function reportCacheMissIfStorageHasValue(key: OnyxKey): void {
 }
 
 /**
- * OnyxSubscriptionManager is a registry for Onyx subscriptions.
- * Subscriptions are stored in `keyListeners`, a flat map keyed by OnyxKey.
- * Subscribers are notified on a per-key basis:
- *    - `notifyKey` for individual keys or collection members
- *    - `notifyCollection` for batch updates to collections
+ * OnyxSubscriptionManager is a registry for Onyx subscriptions, backed by two indexes:
+ *    - keyListeners:        exact-key listeners, keyed by OnyxKey (a single key, a collection object,
+ *                           or a collection member). Notified via `notifyKey`/`notifyCollection`.
+ *    - stateListenersByDep: listeners that re-evaluate when any of their declared dependency keys
+ *                           change, indexed by dep key. Used by `useOnyxState`.
  */
 class OnyxSubscriptionManager {
     private keyListeners: Map<OnyxKey, Set<GenericListener>>;
 
+    private stateListenersByDep: Map<OnyxKey, Set<StateListenerEntry>>;
+
     constructor() {
         this.keyListeners = new Map();
+        this.stateListenersByDep = new Map();
     }
 
     /**
@@ -103,13 +114,45 @@ class OnyxSubscriptionManager {
         };
     }
 
+    /**
+     * Subscribe to state-tree changes. The listener fires when any of the declared dependency keys
+     * changes. A collection key counts as a single dep: any member change fires it. Used by `useOnyxState`.
+     *
+     * Returns an unsubscribe function.
+     */
+    subscribeState(listener: StateListenerCallback, deps: readonly OnyxKey[]): () => void {
+        const entry: StateListenerEntry = {listener, deps: new Set(deps)};
+        for (const dep of entry.deps) {
+            let set = this.stateListenersByDep.get(dep);
+            if (!set) {
+                set = new Set();
+                this.stateListenersByDep.set(dep, set);
+            }
+            set.add(entry);
+        }
+
+        return () => {
+            for (const dep of entry.deps) {
+                const set = this.stateListenersByDep.get(dep);
+                if (!set) {
+                    continue;
+                }
+                set.delete(entry);
+                if (set.size === 0) {
+                    this.stateListenersByDep.delete(dep);
+                }
+            }
+        };
+    }
+
     /**
      * Notify of a single-key write.
      *
      * Dispatch:
      *   1. keyListeners.get(key): exact-key subscribers (always fires).
      *   2. If key is a collection member, keyListeners.get(collectionKey): collection
      *      listeners for the parent collection.
+     *   3. State listeners whose deps include `key` or its collection key.
      */
     notifyKey<TKey extends OnyxKey>(key: TKey, value: OnyxValue<TKey>): void {
         // 1. Exact-key listeners
@@ -134,6 +177,13 @@ class OnyxSubscriptionManager {
                 }
             }
         }
+
+        // 3. State listeners, fired at most once even when both the member and its collection key match.
+        const fired = new Set<StateListenerEntry>();
+        this.fireStateListenersForDep(key, fired);
+        if (isCollectionMemberWrite) {
+            this.fireStateListenersForDep(collectionKey, fired);
+        }
     }
 
     /**
@@ -143,6 +193,7 @@ class OnyxSubscriptionManager {
      *   1. keyListeners.get(collectionKey): fires once with the new collection object.
      *   2. keyListeners.get(memberKey): fires per changed member whose value differs from
      *      the previous, preserving ref-equality on unchanged members.
+     *   3. State listeners affected by the collection key or any changed member key, each fired at most once.
      */
     notifyCollection<TKey extends CollectionKeyBase>(
         collectionKey: TKey,
@@ -186,13 +237,21 @@ class OnyxSubscriptionManager {
                 this.safeInvoke(() => listener(value, memberKey), memberKey);
             }
         }
+
+        // 3. State listeners, each affected entry fired at most once across the collection and its members.
+        const fired = new Set<StateListenerEntry>();
+        this.fireStateListenersForDep(collectionKey, fired);
+        for (const memberKey of changedKeys) {
+            this.fireStateListenersForDep(memberKey, fired);
+        }
     }
 
     /**
      * Wipe all subscriptions. Used by tests and `Onyx.clear()` follow-on.
      */
     clearAll(): void {
         this.keyListeners.clear();
+        this.stateListenersByDep.clear();
         this.resetCacheMissReports();
     }
 
@@ -216,6 +275,22 @@ class OnyxSubscriptionManager {
         return false;
     }
 
+    /** Fires each state listener registered on `depKey`, skipping any already fired this dispatch. */
+    private fireStateListenersForDep(depKey: OnyxKey, alreadyFired: Set<StateListenerEntry>): void {
+        const set = this.stateListenersByDep.get(depKey);
+        if (!set || set.size === 0) {
+            return;
+        }
+
+        for (const entry of [...set]) {
+            if (alreadyFired.has(entry)) {
+                continue;
+            }
+            alreadyFired.add(entry);
+            this.safeInvoke(entry.listener, depKey);
+        }
+    }
+
     /**
      * Runs a listener, catching and logging any throw so one failing listener can't stop the rest.
      */
@@ -231,4 +306,4 @@ class OnyxSubscriptionManager {
 const onyxSubscriptionManager = new OnyxSubscriptionManager();
 
 export default onyxSubscriptionManager;
-export type {Listener};
+export type {Listener, StateListenerCallback};

lib/index.ts

@@ -20,10 +20,12 @@ import type {
 } from './types';
 import type {FetchStatus, ResultMetadata, UseOnyxResult, UseOnyxOptions} from './useOnyx';
 import useOnyx from './useOnyx';
+import type {OnyxStateView, UseOnyxStateOptions, UseOnyxStateSelector} from './useOnyxState';
+import useOnyxState from './useOnyxState';
 import type {OnyxSQLiteKeyValuePair} from './storage/providers/SQLiteProvider';
 
 export default Onyx;
-export {useOnyx};
+export {useOnyx, useOnyxState};
 export type {
     ConnectOptions,
     CustomTypeOptions,
@@ -48,5 +50,8 @@ export type {
     UseOnyxResult,
     Connection,
     UseOnyxOptions,
+    OnyxStateView,
+    UseOnyxStateOptions,
+    UseOnyxStateSelector,
     OnyxSQLiteKeyValuePair,
 };