Computed values

Author: janicduplessis

lib/Onyx.ts

@@ -8,8 +8,10 @@ import Storage from './storage';
 import utils from './utils';
 import DevTools from './DevTools';
 import type {
+    AnyComputedKey,
     Collection,
     CollectionKeyBase,
+    ComputedKey,
     ConnectOptions,
     InitOptions,
     KeyValueMapping,
@@ -63,6 +65,15 @@ function init({
     Promise.all([OnyxUtils.addAllSafeEvictionKeysToRecentlyAccessedList(), OnyxUtils.initializeWithDefaultKeyStates()]).then(deferredInitTask.resolve);
 }
 
+function computeAndSendData(mapping: Mapping<AnyComputedKey>, dependencies: Record<string, unknown>) {
+    let val = cache.getValue(mapping.key.cacheKey);
+    if (val === undefined) {
+        val = mapping.key.compute(dependencies);
+        cache.set(mapping.key.cacheKey, val);
+    }
+    OnyxUtils.sendDataToConnection(mapping, val, mapping.key.cacheKey, true);
+}
+
 /**
  * Subscribes a react component's state directly to a store key
  *
@@ -91,12 +102,54 @@ function init({
  * Note that it will not cause the component to have the loading prop set to true.
  * @returns an ID to use when calling disconnect
  */
-function connect<TKey extends OnyxKey>(mapping: ConnectOptions<TKey>): number {
+function connect<TKey extends OnyxKey | AnyComputedKey>(mapping: ConnectOptions<TKey>): number {
     const connectionID = lastConnectionID++;
     const callbackToStateMapping = OnyxUtils.getCallbackToStateMapping();
     callbackToStateMapping[connectionID] = mapping;
     callbackToStateMapping[connectionID].connectionID = connectionID;
 
+    const mappingKey = mapping.key;
+    if (OnyxUtils.isComputedKey(mappingKey)) {
+        deferredInitTask.promise
+            .then(() => OnyxUtils.addKeyToRecentlyAccessedIfNeeded(mapping))
+            .then(() => {
+                const mappingDependencies = mappingKey.dependencies || {};
+                const dependenciesCount = _.size(mappingDependencies);
+                if (dependenciesCount === 0) {
+                    // If we have no dependencies we can send the computed value immediately.
+                    computeAndSendData(mapping as Mapping<AnyComputedKey>, {});
+                } else {
+                    callbackToStateMapping[connectionID].dependencyConnections = [];
+
+                    const dependencyValues: Record<string, unknown> = {};
+                    _.each(mappingDependencies, (dependency, dependencyKey) => {
+                        // Create a mapping of dependent cache keys so when a key changes, all dependent keys
+                        // can also be cleared from the cache.
+                        const cacheKey = OnyxUtils.getCacheKey(dependency);
+                        OnyxUtils.addDependentCacheKey(cacheKey, mappingKey.cacheKey);
+
+                        // Connect to dependencies.
+                        const dependencyConnection = connect({
+                            key: dependency,
+                            waitForCollectionCallback: true,
+                            callback: (value) => {
+                                dependencyValues[dependencyKey] = value;
+
+                                // Once all dependencies are ready, compute the value and send it to the connection.
+                                if (_.size(dependencyValues) === dependenciesCount) {
+                                    computeAndSendData(mapping as Mapping<AnyComputedKey>, dependencyValues);
+                                }
+                            },
+                        });
+
+                        // Store dependency connections so we can disconnect them later.
+                        callbackToStateMapping[connectionID].dependencyConnections.push(dependencyConnection);
+                    });
+                }
+            });
+        return connectionID;
+    }
+
     if (mapping.initWithStoredValues === false) {
         return connectionID;
     }
@@ -108,24 +161,24 @@ function connect<TKey extends OnyxKey>(mapping: ConnectOptions<TKey>): number {
             // Performance improvement
             // If the mapping is connected to an onyx key that is not a collection
             // we can skip the call to getAllKeys() and return an array with a single item
-            if (Boolean(mapping.key) && typeof mapping.key === 'string' && !mapping.key.endsWith('_') && cache.storageKeys.has(mapping.key)) {
-                return new Set([mapping.key]);
+            if (Boolean(mappingKey) && typeof mappingKey === 'string' && !mappingKey.endsWith('_') && cache.storageKeys.has(mappingKey)) {
+                return new Set([mappingKey]);
             }
             return OnyxUtils.getAllKeys();
         })
         .then((keys) => {
             // We search all the keys in storage to see if any are a "match" for the subscriber we are connecting so that we
             // can send data back to the subscriber. Note that multiple keys can match as a subscriber could either be
             // subscribed to a "collection key" or a single key.
-            const matchingKeys = Array.from(keys).filter((key) => OnyxUtils.isKeyMatch(mapping.key, key));
+            const matchingKeys = Array.from(keys).filter((key) => OnyxUtils.isKeyMatch(mappingKey, key));
 
             // If the key being connected to does not exist we initialize the value with null. For subscribers that connected
             // directly via connect() they will simply get a null value sent to them without any information about which key matched
             // since there are none matched. In withOnyx() we wait for all connected keys to return a value before rendering the child
             // component. This null value will be filtered out so that the connected component can utilize defaultProps.
             if (matchingKeys.length === 0) {
-                if (mapping.key && !OnyxUtils.isCollectionKey(mapping.key)) {
-                    cache.set(mapping.key, null);
+                if (mappingKey && !OnyxUtils.isCollectionKey(mappingKey)) {
+                    cache.set(mappingKey, null);
                 }
 
                 // Here we cannot use batching because the null value is expected to be set immediately for default props
@@ -138,7 +191,7 @@ function connect<TKey extends OnyxKey>(mapping: ConnectOptions<TKey>): number {
             // into an object and just make a single call. The latter behavior is enabled by providing a waitForCollectionCallback key
             // combined with a subscription to a collection key.
             if (typeof mapping.callback === 'function') {
-                if (OnyxUtils.isCollectionKey(mapping.key)) {
+                if (OnyxUtils.isCollectionKey(mappingKey)) {
                     if (mapping.waitForCollectionCallback) {
                         OnyxUtils.getCollectionDataAndSendAsObject(matchingKeys, mapping);
                         return;
@@ -147,26 +200,26 @@ function connect<TKey extends OnyxKey>(mapping: ConnectOptions<TKey>): number {
                     // We did not opt into using waitForCollectionCallback mode so the callback is called for every matching key.
                     // eslint-disable-next-line @typescript-eslint/prefer-for-of
                     for (let i = 0; i < matchingKeys.length; i++) {
-                        OnyxUtils.get(matchingKeys[i]).then((val) => OnyxUtils.sendDataToConnection(mapping, val, matchingKeys[i] as TKey, true));
+                        OnyxUtils.get(matchingKeys[i]).then((val) => OnyxUtils.sendDataToConnection(mapping, val, matchingKeys[i], true));
                     }
                     return;
                 }
 
                 // If we are not subscribed to a collection key then there's only a single key to send an update for.
-                OnyxUtils.get(mapping.key).then((val) => OnyxUtils.sendDataToConnection(mapping, val, mapping.key, true));
+                OnyxUtils.get(mappingKey).then((val) => OnyxUtils.sendDataToConnection(mapping, val, mappingKey, true));
                 return;
             }
 
             // If we have a withOnyxInstance that means a React component has subscribed via the withOnyx() HOC and we need to
             // group collection key member data into an object.
             if (mapping.withOnyxInstance) {
-                if (OnyxUtils.isCollectionKey(mapping.key)) {
+                if (OnyxUtils.isCollectionKey(mappingKey)) {
                     OnyxUtils.getCollectionDataAndSendAsObject(matchingKeys, mapping);
                     return;
                 }
 
                 // If the subscriber is not using a collection key then we just send a single value back to the subscriber
-                OnyxUtils.get(mapping.key).then((val) => OnyxUtils.sendDataToConnection(mapping, val, mapping.key, true));
+                OnyxUtils.get(mappingKey).then((val) => OnyxUtils.sendDataToConnection(mapping, val, mappingKey, true));
                 return;
             }
 
@@ -197,6 +250,10 @@ function disconnect(connectionID: number, keyToRemoveFromEvictionBlocklist?: Ony
         OnyxUtils.removeFromEvictionBlockList(keyToRemoveFromEvictionBlocklist, connectionID);
     }
 
+    if (callbackToStateMapping[connectionID].dependencyConnections) {
+        callbackToStateMapping[connectionID].dependencyConnections.forEach((id: number) => disconnect(id));
+    }
+
     delete callbackToStateMapping[connectionID];
 }
 

lib/OnyxUtils.d.ts

@@ -1,6 +1,6 @@
 import {Component} from 'react';
 import * as Logger from './Logger';
-import {CollectionKey, CollectionKeyBase, DeepRecord, KeyValueMapping, NullishDeep, OnyxCollection, OnyxEntry, OnyxKey, Selector} from './types';
+import {AnyComputedKey, ComputedKey, CollectionKey, CollectionKeyBase, DeepRecord, KeyValueMapping, NullishDeep, OnyxCollection, OnyxEntry, OnyxKey, Selector} from './types';
 
 declare const METHOD: {
     readonly SET: 'set';
@@ -275,6 +275,21 @@ declare function applyMerge(existingValue: OnyxValue<OnyxKey>, changes: Array<On
  */
 declare function initializeWithDefaultKeyStates(): Promise<void>;
 
+/**
+ * Returns a string cache key for a possible computed key.
+ */
+declare function getCacheKey(key: OnyxKey | AnyComputedKey): string;
+
+/**
+ * Returns if a key is a computed key.
+ */
+declare function isComputedKey(key: OnyxKey | AnyComputedKey): key is AnyComputedKey;
+
+/**
+ * Adds an entry in the dependent cache key map.
+ */
+declare function addDependentCacheKey(key: OnyxKey, dependentKey: OnyxKey): void;
+
 const OnyxUtils = {
     METHOD,
     getMergeQueue,
@@ -315,6 +330,9 @@ const OnyxUtils = {
     prepareKeyValuePairsForStorage,
     applyMerge,
     initializeWithDefaultKeyStates,
+    getCacheKey,
+    isComputedKey,
+    addDependentCacheKey,
 } as const;
 
 export default OnyxUtils;

lib/OnyxUtils.js

@@ -26,6 +26,9 @@ const mergeQueuePromise = {};
 // Holds a mapping of all the react components that want their state subscribed to a store key
 const callbackToStateMapping = {};
 
+// Holds a mapping of cache keys to their dependencies. This is used to invalidate computed keys.
+const dependentCacheKeys = {};
+
 // Keeps a copy of the values of the onyx collection keys as a map for faster lookups
 let onyxCollectionKeyMap = new Map();
 
@@ -191,6 +194,16 @@ const reduceCollectionWithSelector = (collection, selector, withOnyxInstanceStat
         {},
     );
 
+/**
+ * Returns if the key is a computed key.
+ *
+ * @param {Mixed} key
+ * @returns {boolean}
+ */
+function isComputedKey(key) {
+    return typeof key === 'object' && 'compute' in key;
+}
+
 /**
  * Get some data from the store
  *
@@ -311,6 +324,16 @@ function isSafeEvictionKey(testKey) {
     return _.some(evictionAllowList, (key) => isKeyMatch(key, testKey));
 }
 
+/**
+ * Returns a string cache key for a possible computed key.
+ *
+ * @param {Mixed} key
+ * @returns {String}
+ */
+function getCacheKey(key) {
+    return isComputedKey(key) ? key.cacheKey : key;
+}
+
 /**
  * Tries to get a value from the cache. If the value is not present in cache it will return the default value or undefined.
  * If the requested key is a collection, it will return an object with all the collection members.
@@ -320,6 +343,30 @@ function isSafeEvictionKey(testKey) {
  * @returns {Mixed}
  */
 function tryGetCachedValue(key, mapping = {}) {
+    if (isComputedKey(key)) {
+        // Check if we have the value in cache already.
+        let val = cache.getValue(key.cacheKey);
+        if (val !== undefined) {
+            return val;
+        }
+
+        // Check if we can compute the value if all dependencies are in cache.
+        const dependencies = _.mapObject(key.dependencies || {}, (dependencyKey) =>
+            tryGetCachedValue(
+                dependencyKey,
+                // TODO: We could support full mapping here.
+                {key: dependencyKey},
+            ),
+        );
+        if (_.all(dependencies, (dependency) => dependency !== undefined)) {
+            val = key.compute(dependencies);
+            cache.set(key.cacheKey, val);
+            return val;
+        }
+
+        return undefined;
+    }
+
     let val = cache.getValue(key);
 
     if (isCollectionKey(key)) {
@@ -480,6 +527,24 @@ function getCachedCollection(collectionKey) {
     );
 }
 
+function clearComputedCacheForKey(key) {
+    const dependentKeys = dependentCacheKeys[key];
+    if (!dependentKeys) {
+        return;
+    }
+
+    dependentKeys.forEach((dependentKey) => {
+        cache.drop(dependentKey);
+
+        clearComputedCacheForKey(dependentKey);
+    });
+}
+
+function addDependentCacheKey(key, dependentKey) {
+    dependentCacheKeys[key] = dependentCacheKeys[key] || new Set();
+    dependentCacheKeys[key].add(dependentKey);
+}
+
 /**
  * When a collection of keys change, search for any callbacks matching the collection key and trigger those callbacks
  *
@@ -490,6 +555,8 @@ function getCachedCollection(collectionKey) {
  * @param {boolean} [notifyWithOnyxSubscibers=true]
  */
 function keysChanged(collectionKey, partialCollection, notifyRegularSubscibers = true, notifyWithOnyxSubscibers = true) {
+    clearComputedCacheForKey(collectionKey);
+
     // We are iterating over all subscribers similar to keyChanged(). However, we are looking for subscribers who are subscribing to either a collection key or
     // individual collection key member for the collection that is being updated. It is important to note that the collection parameter cane be a PARTIAL collection
     // and does not represent all of the combined keys and values for a collection key. It is just the "new" data that was merged in via mergeCollection().
@@ -667,6 +734,8 @@ function keyChanged(key, data, prevData, canUpdateSubscriber = () => true, notif
         removeLastAccessedKey(key);
     }
 
+    clearComputedCacheForKey(key);
+
     // We are iterating over all subscribers to see if they are interested in the key that has just changed. If the subscriber's  key is a collection key then we will
     // notify them if the key that changed is a collection member. Or if it is a regular key notify them when there is an exact match. Depending on whether the subscriber
     // was connected via withOnyx we will call setState() directly on the withOnyx instance. If it is a regular connection we will pass the data to the provided callback.
@@ -851,7 +920,7 @@ function addKeyToRecentlyAccessedIfNeeded(mapping) {
             throw new Error(`Cannot subscribe to safe eviction key '${mapping.key}' without providing a canEvict value.`);
         }
 
-        addLastAccessedKey(mapping.key);
+        addLastAccessedKey(getCacheKey(mapping.key));
     }
 }
 
@@ -1184,6 +1253,9 @@ const OnyxUtils = {
     prepareKeyValuePairsForStorage,
     applyMerge,
     initializeWithDefaultKeyStates,
+    getCacheKey,
+    isComputedKey,
+    addDependentCacheKey,
 };
 
 export default OnyxUtils;

lib/index.ts

@@ -1,13 +1,14 @@
 import Onyx from './Onyx';
 import type {OnyxUpdate, ConnectOptions} from './Onyx';
-import type {CustomTypeOptions, OnyxCollection, OnyxEntry, NullishDeep, KeyValueMapping, OnyxKey, Selector, WithOnyxInstanceState, OnyxValue} from './types';
+import type {ComputedKey, CustomTypeOptions, OnyxCollection, OnyxEntry, NullishDeep, KeyValueMapping, OnyxKey, Selector, WithOnyxInstanceState, OnyxValue} from './types';
 import type {UseOnyxResult, FetchStatus, ResultMetadata} from './useOnyx';
 import useOnyx from './useOnyx';
 import withOnyx from './withOnyx';
 
 export default Onyx;
 export {withOnyx, useOnyx};
 export type {
+    ComputedKey,
     CustomTypeOptions,
     OnyxCollection,
     OnyxEntry,