Use immutable map structures

Summary:
Use hashed-array-mapped tries instead of built-in Map for atomValues, to avoid copying the map. This should dramatically increase speed for large numbers of atoms.

In a benchmark I performed, setting a value in a HAMT (just the data structure, not all of Recoil) compared with copying a map and setting the value is:

* 3.7× faster for 10 entries
* 28× faster for 100 entries
* 325× faster for 1,000 entries
* 3,026× faster for 10,000 entries

Reading values is slower:

* 5× slower for 10 entries
* 5× slower for 100 entries
* 6× slower for 1,000 entries
* 8× slower for 10,000 entries

However, in Recoil, we generally only read once per write (since the read value is cached elsewhere until the next write), so this should still be a win overall.

Reviewed By: drarmstr

Differential Revision: D25487986

fbshipit-source-id: cea056cc298520f1b54f9fb12c896086c4ef82c6

package.json

@@ -26,6 +26,9 @@
     "lint": "eslint .",
     "deploy-nightly": "yarn build && node scripts/deploy_nightly_build.js"
   },
+  "dependencies": {
+    "hamt_plus": "1.0.2"
+  },
   "peerDependencies": {
     "react": "^16.13.1 || ^17.0.0"
   },

src/adt/Recoil_PersistentMap.js

@@ -0,0 +1,138 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails oncall+recoil
+ * @flow strict
+ * @format
+ */
+
+'use strict';
+
+import type {HAMTPlusMap} from 'hamt_plus';
+
+const gkx = require('../util/Recoil_gkx');
+const hamt = require('hamt_plus');
+
+export interface PersistentMap<K: string, V> {
+  keys(): Iterable<K>;
+  entries(): Iterable<[K, V]>;
+
+  get(key: K): V | void;
+  has(key: K): boolean;
+  set(key: K, value: V): PersistentMap<K, V>;
+  delete(key: K): PersistentMap<K, V>;
+
+  clone(): PersistentMap<K, V>;
+  toMap(): Map<K, V>;
+}
+
+class BuiltInMap<K: string, V> implements PersistentMap<K, V> {
+  _map: Map<K, V>;
+
+  constructor(existing?: PersistentMap<K, V>) {
+    this._map = new Map(existing?.entries());
+  }
+
+  keys(): Iterable<K> {
+    return this._map.keys();
+  }
+
+  entries(): Iterable<[K, V]> {
+    return this._map.entries();
+  }
+
+  get(k: K): V | void {
+    return this._map.get(k);
+  }
+
+  has(k: K): boolean {
+    return this._map.has(k);
+  }
+
+  set(k: K, v: V): PersistentMap<K, V> {
+    this._map.set(k, v);
+    return this;
+  }
+
+  delete(k: K): PersistentMap<K, V> {
+    this._map.delete(k);
+    return this;
+  }
+
+  clone(): PersistentMap<K, V> {
+    return persistentMap(this);
+  }
+
+  toMap(): Map<K, V> {
+    return new Map(this._map);
+  }
+}
+
+class HashArrayMappedTrieMap<K: string, V> implements PersistentMap<K, V> {
+  // Because hamt.empty is not a function there is no way to introduce type
+  // parameters on it, so empty is typed as HAMTPlusMap<string, mixed>.
+  // flowlint-next-line unclear-type:off
+  _hamt: HAMTPlusMap<K, V> = ((hamt.empty: any).beginMutation(): HAMTPlusMap<
+    K,
+    V,
+  >);
+
+  constructor(existing?: PersistentMap<K, V>) {
+    if (existing instanceof HashArrayMappedTrieMap) {
+      const h = existing._hamt.endMutation();
+      existing._hamt = h.beginMutation();
+      this._hamt = h.beginMutation();
+    } else if (existing) {
+      for (const [k, v] of existing.entries()) {
+        this._hamt.set(k, v);
+      }
+    }
+  }
+
+  keys(): Iterable<K> {
+    return this._hamt.keys();
+  }
+
+  entries(): Iterable<[K, V]> {
+    return this._hamt.entries();
+  }
+
+  get(k: K): V | void {
+    return this._hamt.get(k);
+  }
+
+  has(k: K): boolean {
+    return this._hamt.has(k);
+  }
+
+  set(k: K, v: V): PersistentMap<K, V> {
+    this._hamt.set(k, v);
+    return this;
+  }
+
+  delete(k: K): PersistentMap<K, V> {
+    this._hamt.delete(k);
+    return this;
+  }
+
+  clone(): PersistentMap<K, V> {
+    return persistentMap(this);
+  }
+
+  toMap(): Map<K, V> {
+    return new Map(this._hamt);
+  }
+}
+
+export function persistentMap<K: string, V>(
+  existing?: PersistentMap<K, V>,
+): PersistentMap<K, V> {
+  if (gkx('recoil_hamt_2020')) {
+    return new HashArrayMappedTrieMap(existing);
+  } else {
+    return new BuiltInMap(existing);
+  }
+}

src/core/Recoil_FunctionalCore.js

@@ -14,13 +14,9 @@ import type {Loadable} from '../adt/Recoil_Loadable';
 import type {DependencyMap} from './Recoil_Graph';
 import type {DefaultValue} from './Recoil_Node';
 import type {RecoilValue} from './Recoil_RecoilValue';
-import type {AtomValues, NodeKey, Store, TreeState} from './Recoil_State';
+import type {AtomWrites, NodeKey, Store, TreeState} from './Recoil_State';
 
-const {
-  mapByDeletingFromMap,
-  mapBySettingInMap,
-  setByAddingToSet,
-} = require('../util/Recoil_CopyOnWrite');
+const {setByAddingToSet} = require('../util/Recoil_CopyOnWrite');
 const filterIterable = require('../util/Recoil_filterIterable');
 const mapIterable = require('../util/Recoil_mapIterable');
 const {getNode, getNodeMaybe, recoilValuesForKeys} = require('./Recoil_Node');
@@ -62,12 +58,8 @@ function setUnvalidatedAtomValue_DEPRECATED<T>(
 
   return {
     ...state,
-    atomValues: mapByDeletingFromMap(state.atomValues, key),
-    nonvalidatedAtoms: mapBySettingInMap(
-      state.nonvalidatedAtoms,
-      key,
-      newValue,
-    ),
+    atomValues: state.atomValues.clone().delete(key),
+    nonvalidatedAtoms: state.nonvalidatedAtoms.clone().set(key, newValue),
     dirtyAtoms: setByAddingToSet(state.dirtyAtoms, key),
   };
 }
@@ -80,7 +72,7 @@ function setNodeValue<T>(
   state: TreeState,
   key: NodeKey,
   newValue: T | DefaultValue,
-): [DependencyMap, AtomValues] {
+): [DependencyMap, AtomWrites] {
   const node = getNode(key);
   if (node.set == null) {
     throw new ReadOnlyRecoilValueError(

src/core/Recoil_Node.js

@@ -13,7 +13,7 @@
 import type {Loadable} from '../adt/Recoil_Loadable';
 import type {DependencyMap} from './Recoil_GraphTypes';
 import type {RecoilValue} from './Recoil_RecoilValue';
-import type {AtomValues, NodeKey, Store, TreeState} from './Recoil_State';
+import type {AtomWrites, NodeKey, Store, TreeState} from './Recoil_State';
 
 const expectationViolation = require('../util/Recoil_expectationViolation');
 const mapIterable = require('../util/Recoil_mapIterable');
@@ -70,7 +70,7 @@ export type ReadWriteNodeOptions<T> = $ReadOnly<{
     store: Store,
     state: TreeState,
     newValue: T | DefaultValue,
-  ) => [DependencyMap, AtomValues],
+  ) => [DependencyMap, AtomWrites],
 }>;
 
 type Node<T> = ReadOnlyNodeOptions<T> | ReadWriteNodeOptions<T>;

src/core/Recoil_RecoilRoot.react.js

@@ -21,6 +21,10 @@ const {useContext, useEffect, useMemo, useRef, useState} = require('React');
 // @fb-only: const URI = require('URI');
 
 const Queue = require('../adt/Recoil_Queue');
+const {
+  getNextTreeStateVersion,
+  makeEmptyStoreState,
+} = require('../core/Recoil_State');
 const {mapByDeletingMultipleFromMap} = require('../util/Recoil_CopyOnWrite');
 const expectationViolation = require('../util/Recoil_expectationViolation');
 const nullthrows = require('../util/Recoil_nullthrows');
@@ -37,10 +41,6 @@ const {graph, saveDependencyMapToStore} = require('./Recoil_Graph');
 const {cloneGraph} = require('./Recoil_Graph');
 const {applyAtomValueWrites} = require('./Recoil_RecoilValueInterface');
 const {freshSnapshot} = require('./Recoil_Snapshot');
-const {
-  getNextTreeStateVersion,
-  makeEmptyStoreState,
-} = require('./Recoil_State');
 // @fb-only: const gkx = require('gkx');
 
 type Props = {
@@ -229,10 +229,10 @@ function initialStoreState_DEPRECATED(store, initializeState): StoreState {
 
       saveDependencyMapToStore(depMap, store, state.version);
 
-      const nonvalidatedAtoms = mapByDeletingMultipleFromMap(
-        state.nonvalidatedAtoms,
-        writtenNodes,
-      );
+      const nonvalidatedAtoms = state.nonvalidatedAtoms.clone();
+      for (const n of writtenNodes) {
+        nonvalidatedAtoms.delete(n);
+      }
 
       initial.currentTree = {
         ...state,

src/core/Recoil_RecoilValueInterface.js

@@ -12,10 +12,15 @@
 
 import type {Loadable} from '../adt/Recoil_Loadable';
 import type {ValueOrUpdater} from '../recoil_values/Recoil_selector';
-import type {AtomValues, NodeKey, Store, TreeState} from './Recoil_State';
+import type {
+  AtomValues,
+  AtomWrites,
+  NodeKey,
+  Store,
+  TreeState,
+} from './Recoil_State';
 
 const gkx = require('../util/Recoil_gkx');
-const mapMap = require('../util/Recoil_mapMap');
 const nullthrows = require('../util/Recoil_nullthrows');
 const recoverableViolation = require('../util/Recoil_recoverableViolation');
 const Tracing = require('../util/Recoil_Tracing');
@@ -66,9 +71,9 @@ function getRecoilValueAsLoadable<T>(
 
 function applyAtomValueWrites(
   atomValues: AtomValues,
-  writes: AtomValues,
+  writes: AtomWrites,
 ): AtomValues {
-  const result = mapMap(atomValues, v => v);
+  const result = atomValues.clone();
   writes.forEach((v, k) => {
     if (v.state === 'hasValue' && v.contents instanceof DefaultValue) {
       result.delete(k);
@@ -234,8 +239,8 @@ function batchStart(): () => void {
 function copyTreeState(state) {
   return {
     ...state,
-    atomValues: new Map(state.atomValues),
-    nonvalidatedAtoms: new Map(state.nonvalidatedAtoms),
+    atomValues: state.atomValues.clone(),
+    nonvalidatedAtoms: state.nonvalidatedAtoms.clone(),
     dirtyAtoms: new Set(state.dirtyAtoms),
   };
 }

src/core/Recoil_Snapshot.js

@@ -202,8 +202,8 @@ function cloneStoreState(
           stateID: version,
           transactionMetadata: {...treeState.transactionMetadata},
           dirtyAtoms: new Set(treeState.dirtyAtoms),
-          atomValues: new Map(treeState.atomValues),
-          nonvalidatedAtoms: new Map(treeState.nonvalidatedAtoms),
+          atomValues: treeState.atomValues.clone(),
+          nonvalidatedAtoms: treeState.nonvalidatedAtoms.clone(),
         }
       : treeState,
     nextTree: null,

src/core/Recoil_State.js

@@ -11,14 +11,18 @@
 'use strict';
 
 import type {Loadable} from '../adt/Recoil_Loadable';
+import type {PersistentMap} from '../adt/Recoil_PersistentMap';
 import type {Graph} from './Recoil_GraphTypes';
 import type {ComponentID, NodeKey, StateID, Version} from './Recoil_Keys';
 export type {ComponentID, NodeKey, StateID, Version} from './Recoil_Keys';
 
 const {graph} = require('./Recoil_Graph');
+const {persistentMap} = require('../adt/Recoil_PersistentMap');
 
 // flowlint-next-line unclear-type:off
-export type AtomValues = Map<NodeKey, Loadable<any>>;
+export type AtomValues = PersistentMap<NodeKey, Loadable<any>>;
+// flowlint-next-line unclear-type:off
+export type AtomWrites = Map<NodeKey, Loadable<any>>;
 
 type ComponentCallback = TreeState => void;
 
@@ -39,7 +43,7 @@ export type TreeState = $ReadOnly<{
   // Atoms:
   dirtyAtoms: Set<NodeKey>,
   atomValues: AtomValues,
-  nonvalidatedAtoms: Map<NodeKey, mixed>,
+  nonvalidatedAtoms: PersistentMap<NodeKey, mixed>,
 }>;
 
 // StoreState represents the state of a Recoil context. It is global and mutable.
@@ -128,8 +132,8 @@ function makeEmptyTreeState(): TreeState {
     stateID: version,
     transactionMetadata: {},
     dirtyAtoms: new Set(),
-    atomValues: new Map(),
-    nonvalidatedAtoms: new Map(),
+    atomValues: persistentMap(),
+    nonvalidatedAtoms: persistentMap(),
   };
 }
 

src/hooks/Recoil_Hooks.js

@@ -520,7 +520,7 @@ function useTransactionSubscription(callback: Store => void) {
 function externallyVisibleAtomValuesInState(
   state: TreeState,
 ): Map<NodeKey, mixed> {
-  const atomValues: Map<NodeKey, Loadable<mixed>> = state.atomValues;
+  const atomValues = state.atomValues.toMap();
   const persistedAtomContentsValues = mapMap(
     filterMap(atomValues, (v, k) => {
       const node = getNode(k);
@@ -535,7 +535,10 @@ function externallyVisibleAtomValuesInState(
   );
   // Merge in nonvalidated atoms; we may not have defs for them but they will
   // all have persistence on or they wouldn't be there in the first place.
-  return mergeMaps(state.nonvalidatedAtoms, persistedAtomContentsValues);
+  return mergeMaps(
+    state.nonvalidatedAtoms.toMap(),
+    persistedAtomContentsValues,
+  );
 }
 
 type ExternallyVisibleAtomInfo = {

src/recoil_values/Recoil_atom.js

@@ -66,7 +66,7 @@ import type {
 import type {DependencyMap} from '../core/Recoil_Graph';
 import type {PersistenceInfo, ReadWriteNodeOptions} from '../core/Recoil_Node';
 import type {RecoilState, RecoilValue} from '../core/Recoil_RecoilValue';
-import type {AtomValues, NodeKey, Store, TreeState} from '../core/Recoil_State';
+import type {AtomWrites, NodeKey, Store, TreeState} from '../core/Recoil_State';
 
 // @fb-only: const {scopedAtom} = require('Recoil_ScopedAtom');
 
@@ -409,7 +409,7 @@ function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
     store: Store,
     state: TreeState,
     newValue: T | DefaultValue,
-  ): [DependencyMap, AtomValues] {
+  ): [DependencyMap, AtomWrites] {
     initAtom(store, state, 'set');
 
     // Bail out if we're being set to the existing value, or if we're being

src/recoil_values/Recoil_selector_NEW.js

@@ -63,7 +63,7 @@ import type {
   RecoilValue,
   RecoilValueReadOnly,
 } from '../core/Recoil_RecoilValue';
-import type {AtomValues, NodeKey, Store, TreeState} from '../core/Recoil_State';
+import type {AtomWrites, NodeKey, Store, TreeState} from '../core/Recoil_State';
 
 const {
   loadableWithError,
@@ -910,12 +910,12 @@ function selector<T>(
   }
 
   if (set != null) {
-    function mySet(store, state, newValue): [DependencyMap, AtomValues] {
+    function mySet(store, state, newValue): [DependencyMap, AtomWrites] {
       initSelector(store);
 
       let syncSelectorSetFinished = false;
       const dependencyMap: DependencyMap = new Map();
-      const writes: AtomValues = new Map();
+      const writes: AtomWrites = new Map();
 
       function getRecoilValue<S>({key}: RecoilValue<S>): S {
         if (syncSelectorSetFinished) {