OnyxCache.ts

import {deepEqual} from 'fast-equals';
import bindAll from 'lodash/bindAll';
import type {ValueOf} from 'type-fest';
import utils from './utils';
import type {CollectionKeyBase, KeyValueMapping, NonUndefined, OnyxCollection, OnyxKey, OnyxValue} from './types';
import OnyxKeys from './OnyxKeys';

/** Frozen object containing all collection members — safe to return by reference */
type CollectionSnapshot = Readonly<NonUndefined<OnyxCollection<KeyValueMapping[OnyxKey]>>>;

/**
 * Stable frozen empty object used as the canonical value for empty collections.
 * Returning the same reference avoids unnecessary re-renders in useSyncExternalStore,
 * which relies on === equality to detect changes.
 */
const FROZEN_EMPTY_COLLECTION: Readonly<NonUndefined<OnyxCollection<KeyValueMapping[OnyxKey]>>> = Object.freeze({});

// Task constants
const TASK = {
    GET: 'get',
    GET_ALL_KEYS: 'getAllKeys',
    CLEAR: 'clear',
} as const;

type CacheTask = ValueOf<typeof TASK> | `${ValueOf<typeof TASK>}:${string}`;

/**
 * In memory cache providing data by reference
 * Encapsulates Onyx cache related functionality
 */
class OnyxCache {
    /** Cache of all the storage keys available in persistent storage */
    private storageKeys: Set<OnyxKey>;

    /** A list of keys where a nullish value has been fetched from storage before, but the key still exists in cache */
    private nullishStorageKeys: Set<OnyxKey>;

    /** A map of cached values */
    private storageMap: Record<OnyxKey, OnyxValue<OnyxKey>>;

    /**
     * Captured pending tasks for already running storage methods
     * Using a map yields better performance on operations such a delete
     */
    private pendingPromises: Map<string, Promise<OnyxValue<OnyxKey> | OnyxKey[]>>;

    /** List of keys that are safe to remove when we reach max storage */
    private evictionAllowList: OnyxKey[] = [];

    /** List of keys that have been directly subscribed to or recently modified from least to most recent */
    private recentlyAccessedKeys = new Set<OnyxKey>();

    /** Frozen collection snapshots for structural sharing */
    private collectionSnapshots: Map<OnyxKey, CollectionSnapshot>;

    /** Collections whose snapshots need rebuilding (lazy — rebuilt on next read) */
    private dirtyCollections: Set<CollectionKeyBase>;

    constructor() {
        this.storageKeys = new Set();
        this.nullishStorageKeys = new Set();
        this.storageMap = {};
        this.pendingPromises = new Map();
        this.collectionSnapshots = new Map();
        this.dirtyCollections = new Set();

        // bind all public methods to prevent problems with `this`
        bindAll(
            this,
            'getAllKeys',
            'get',
            'hasCacheForKey',
            'addKey',
            'addNullishStorageKey',
            'hasNullishStorageKey',
            'clearNullishStorageKeys',
            'set',
            'drop',
            'merge',
            'hasPendingTask',
            'getTaskPromise',
            'captureTask',
            'setAllKeys',
            'setEvictionAllowList',
            'isEvictableKey',
            'removeLastAccessedKey',
            'addLastAccessedKey',
            'addEvictableKeysToRecentlyAccessedList',
            'getKeyForEviction',
            'setCollectionKeys',
            'hasValueChanged',
            'getCollectionData',
        );
    }

    /** Get all the storage keys */
    getAllKeys(): Set<OnyxKey> {
        return this.storageKeys;
    }

    /**
     * Allows to set all the keys at once.
     * This is useful when we are getting
     * all the keys from the storage provider
     * and we want to keep the cache in sync.
     *
     * Previously, we had to call `addKey` in a loop
     * to achieve the same result.
     *
     * @param keys - an array of keys
     */
    setAllKeys(keys: OnyxKey[]) {
        this.storageKeys = new Set(keys);
        for (const key of keys) {
            OnyxKeys.registerMemberKey(key);
        }
    }

    /** Saves a key in the storage keys list
     * Serves to keep the result of `getAllKeys` up to date
     */
    addKey(key: OnyxKey): void {
        this.storageKeys.add(key);
        OnyxKeys.registerMemberKey(key);
    }

    /** Used to set keys that are null/undefined in storage without adding null to the storage map */
    addNullishStorageKey(key: OnyxKey): void {
        this.nullishStorageKeys.add(key);
    }

    /** Used to set keys that are null/undefined in storage without adding null to the storage map */
    hasNullishStorageKey(key: OnyxKey): boolean {
        return this.nullishStorageKeys.has(key);
    }

    /** Used to clear keys that are null/undefined in cache */
    clearNullishStorageKeys(): void {
        this.nullishStorageKeys = new Set();
    }

    /** Check whether cache has data for the given key */
    hasCacheForKey(key: OnyxKey): boolean {
        return this.storageMap[key] !== undefined || this.hasNullishStorageKey(key);
    }

    /** Get a cached value from storage */
    get(key: OnyxKey): OnyxValue<OnyxKey> {
        return this.storageMap[key];
    }

    /**
     * Set's a key value in cache
     * Adds the key to the storage keys list as well
     */
    set(key: OnyxKey, value: OnyxValue<OnyxKey>): OnyxValue<OnyxKey> {
        this.addKey(key);

        // When a key is explicitly set in cache, we can remove it from the list of nullish keys,
        // since it will either be set to a non nullish value or removed from the cache completely.
        this.nullishStorageKeys.delete(key);

        const collectionKey = OnyxKeys.getCollectionKey(key);
        const oldValue = this.storageMap[key];

        if (value === null || value === undefined) {
            delete this.storageMap[key];

            if (collectionKey && oldValue !== undefined) {
                this.dirtyCollections.add(collectionKey);
            }
            return undefined;
        }

        this.storageMap[key] = value;

        if (collectionKey && oldValue !== value) {
            this.dirtyCollections.add(collectionKey);
        }

        return value;
    }

    /** Forget the cached value for the given key */
    drop(key: OnyxKey): void {
        delete this.storageMap[key];

        const collectionKey = OnyxKeys.getCollectionKey(key);
        if (collectionKey) {
            this.dirtyCollections.add(collectionKey);
        }

        // If this is a collection key, clear its snapshot
        if (OnyxKeys.isCollectionKey(key)) {
            this.collectionSnapshots.delete(key);
        }

        this.storageKeys.delete(key);
        OnyxKeys.deregisterMemberKey(key);
    }

    /**
     * Deep merge data to cache, any non existing keys will be created
     * @param data - a map of (cache) key - values
     */
    merge(data: Record<OnyxKey, OnyxValue<OnyxKey>>): void {
        if (typeof data !== 'object' || Array.isArray(data)) {
            throw new Error('data passed to cache.merge() must be an Object of onyx key/value pairs');
        }

        const affectedCollections = new Set<OnyxKey>();

        for (const [key, value] of Object.entries(data)) {
            this.addKey(key);

            const collectionKey = OnyxKeys.getCollectionKey(key);

            if (value === undefined) {
                this.addNullishStorageKey(key);
                // undefined means "no change" — skip storageMap modification
                continue;
            }

            if (value === null) {
                this.addNullishStorageKey(key);
                delete this.storageMap[key];

                if (collectionKey) {
                    affectedCollections.add(collectionKey);
                }
            } else {
                this.nullishStorageKeys.delete(key);

                // Per-key merge instead of spreading the entire storageMap
                const existing = this.storageMap[key];
                const merged = utils.fastMerge(existing, value, {
                    shouldRemoveNestedNulls: true,
                    objectRemovalMode: 'replace',
                }).result;

                // fastMerge is reference-stable: returns the original target when
                // nothing changed, so a simple === check detects no-ops.
                if (merged === existing) {
                    continue;
                }

                this.storageMap[key] = merged;

                if (collectionKey) {
                    affectedCollections.add(collectionKey);
                }
            }
        }

        // Mark affected collections as dirty — snapshots will be lazily rebuilt on next read
        for (const collectionKey of affectedCollections) {
            this.dirtyCollections.add(collectionKey);
        }
    }

    /**
     * Check whether the given task is already running
     * @param taskName - unique name given for the task
     */
    hasPendingTask(taskName: CacheTask): boolean {
        return this.pendingPromises.get(taskName) !== undefined;
    }

    /**
     * Use this method to prevent concurrent calls for the same thing
     * Instead of calling the same task again use the existing promise
     * provided from this function
     * @param taskName - unique name given for the task
     */
    getTaskPromise(taskName: CacheTask): Promise<OnyxValue<OnyxKey> | OnyxKey[]> | undefined {
        return this.pendingPromises.get(taskName);
    }

    /**
     * Capture a promise for a given task so other caller can
     * hook up to the promise if it's still pending
     * @param taskName - unique name for the task
     */
    captureTask(taskName: CacheTask, promise: Promise<OnyxValue<OnyxKey>>): Promise<OnyxValue<OnyxKey>> {
        const returnPromise = promise.finally(() => {
            this.pendingPromises.delete(taskName);
        });

        this.pendingPromises.set(taskName, returnPromise);

        return returnPromise;
    }

    /** Check if the value has changed. Uses reference equality as a fast path, falls back to deep equality. */
    hasValueChanged(key: OnyxKey, value: OnyxValue<OnyxKey>): boolean {
        const currentValue = this.storageMap[key];
        if (currentValue === value) {
            return false;
        }
        return !deepEqual(currentValue, value);
    }

    /**
     * Sets the list of keys that are considered safe for eviction
     * @param keys - Array of OnyxKeys that are safe to evict
     */
    setEvictionAllowList(keys: OnyxKey[]): void {
        this.evictionAllowList = keys;
    }

    /**
     * Checks to see if this key has been flagged as safe for removal.
     * @param testKey - Key to check
     */
    isEvictableKey(testKey: OnyxKey): boolean {
        return this.evictionAllowList.some((key) => OnyxKeys.isKeyMatch(key, testKey));
    }

    /**
     * Remove a key from the recently accessed key list.
     */
    removeLastAccessedKey(key: OnyxKey): void {
        this.recentlyAccessedKeys.delete(key);
    }

    /**
     * Add a key to the list of recently accessed keys. The least
     * recently accessed key should be at the head and the most
     * recently accessed key at the tail.
     */
    addLastAccessedKey(key: OnyxKey, isCollectionKey: boolean): void {
        // Only specific keys belong in this list since we cannot remove an entire collection.
        if (isCollectionKey || !this.isEvictableKey(key)) {
            return;
        }

        this.removeLastAccessedKey(key);
        this.recentlyAccessedKeys.add(key);
    }

    /**
     * Take all the keys that are safe to evict and add them to
     * the recently accessed list when initializing the app. This
     * enables keys that have not recently been accessed to be
     * removed.
     * @param isCollectionKeyFn - Function to determine if a key is a collection key
     * @param getAllKeysFn - Function to get all keys, defaults to Storage.getAllKeys
     */
    addEvictableKeysToRecentlyAccessedList(isCollectionKeyFn: (key: OnyxKey) => boolean, getAllKeysFn: () => Promise<Set<OnyxKey>>): Promise<void> {
        return getAllKeysFn().then((keys: Set<OnyxKey>) => {
            for (const evictableKey of this.evictionAllowList) {
                for (const key of keys) {
                    if (!OnyxKeys.isKeyMatch(evictableKey, key)) {
                        continue;
                    }

                    this.addLastAccessedKey(key, isCollectionKeyFn(key));
                }
            }
        });
    }

    /**
     * Finds the least recently accessed key that can be safely evicted from storage.
     */
    getKeyForEviction(): OnyxKey | undefined {
        // recentlyAccessedKeys is ordered from least to most recently accessed,
        // so the first element is the best candidate for eviction.
        return this.recentlyAccessedKeys.values().next().value;
    }

    /**
     * Set the collection keys for optimized storage
     */
    setCollectionKeys(collectionKeys: Set<OnyxKey>): void {
        OnyxKeys.setCollectionKeys(collectionKeys);

        // Initialize frozen snapshots for collection keys
        for (const collectionKey of collectionKeys) {
            if (!this.collectionSnapshots.has(collectionKey)) {
                this.collectionSnapshots.set(collectionKey, Object.freeze({}));
            }
        }
    }

    /**
     * Rebuilds the frozen collection snapshot from current storageMap references.
     * Uses the indexed collection->members map for O(collectionMembers) instead of O(totalKeys).
     * Returns the previous snapshot reference when all member references are identical,
     * preventing unnecessary re-renders in useSyncExternalStore.
     *
     * @param collectionKey - The collection key to rebuild
     */
    private rebuildCollectionSnapshot(collectionKey: OnyxKey): void {
        const previousSnapshot = this.collectionSnapshots.get(collectionKey);

        const members: NonUndefined<OnyxCollection<KeyValueMapping[OnyxKey]>> = {};
        let hasMemberChanges = false;

        // Use the indexed forward lookup for O(collectionMembers) iteration.
        // Falls back to scanning all storageKeys if the index isn't populated yet.
        const memberKeys = OnyxKeys.getMembersOfCollection(collectionKey);
        const keysToScan = memberKeys ?? this.storageKeys;
        const needsPrefixCheck = !memberKeys;

        for (const key of keysToScan) {
            // When using the fallback path (scanning all storageKeys instead of the indexed
            // forward lookup), skip keys that don't belong to this collection.
            if (needsPrefixCheck && OnyxKeys.getCollectionKey(key) !== collectionKey) {
                continue;
            }
            const val = this.storageMap[key];
            // Skip null/undefined values — they represent deleted or unset keys
            // and should not be included in the frozen collection snapshot.
            if (val !== undefined && val !== null) {
                members[key] = val;

                // Check if this member's reference changed from the old snapshot
                if (!hasMemberChanges && (!previousSnapshot || previousSnapshot[key] !== val)) {
                    hasMemberChanges = true;
                }
            }
        }

        // Check if any members were removed from the previous snapshot.
        // We can't rely on count comparison alone — if one key is removed and another added,
        // the counts match but the snapshot content is different.
        if (!hasMemberChanges && previousSnapshot) {
            // eslint-disable-next-line no-restricted-syntax
            for (const key in previousSnapshot) {
                if (!(key in members)) {
                    hasMemberChanges = true;
                    break;
                }
            }
        }

        // If nothing actually changed, reuse the old snapshot reference.
        // This is critical: useSyncExternalStore uses === to detect changes,
        // so returning the same reference prevents unnecessary re-renders.
        if (!hasMemberChanges && previousSnapshot) {
            return;
        }

        Object.freeze(members);

        this.collectionSnapshots.set(collectionKey, members);
    }

    /**
     * Get all data for a collection key.
     * Returns a frozen snapshot with structural sharing — safe to return by reference.
     * Lazily rebuilds the snapshot if the collection was modified since the last read.
     */
    getCollectionData(collectionKey: OnyxKey): Record<OnyxKey, OnyxValue<OnyxKey>> | undefined {
        if (this.dirtyCollections.has(collectionKey)) {
            this.rebuildCollectionSnapshot(collectionKey);
            this.dirtyCollections.delete(collectionKey);
        }

        const snapshot = this.collectionSnapshots.get(collectionKey);
        if (utils.isEmptyObject(snapshot)) {
            // We check storageKeys.size (not collection-specific keys) to distinguish
            // "init complete, this collection is genuinely empty" from "init not done yet."
            // During init, setAllKeys loads ALL keys at once — so if any key exists,
            // the full storage picture is loaded and an empty collection is truly empty.
            // Returning undefined before init prevents subscribers from seeing a false empty state.
            if (this.storageKeys.size > 0) {
                return FROZEN_EMPTY_COLLECTION;
            }
            return undefined;
        }

        return snapshot;
    }
}

const instance = new OnyxCache();

export default instance;
export {TASK};
export type {CacheTask};