Recoil Atom Effects (#380)

Summary:
Pull Request resolved: #380

Introduce the concept of effects for Recoil Atoms based on [this proposal](https://fb.quip.com/bauvAeBEX5Ee).  This is a similar concept of using React `useEffect()` for side-effects on components, only for atoms.  It can be challenging to use `useEffect()` with Recoil atoms because the Recoil hierarchy is orthogonal to the React component hierarchy, and we have to worry about lifetime and singleton patterns.  Recoil Atom Effects are an array of callbacks which will be called the first time the atom is used with a `<RecoilRoot>`, so they are always part of a React context.  Atom Effects can be used for:

* Synchronization with browser URI history
* Persistence to local storage
* Bi-directional synchronization with remote asynchronous state
* Maintaining state history
* other side-effects such as "set on get"
* &c.

# Proposed Atom Effects

```
type AtomOptions<T> = {
  key: string,
  default: T | RecoilValue<T> | Promise<T>,
  ...

  effects: $ReadOnlyArray<AtomEffect<T>>,
};

// Effect is called the first time node is used with a <RecoilRoot>
type AtomEffect<T> = ({
  node: RecoilState<T>,
  trigger: 'get' | 'set',

  // Can call either synchronously to initialize value or async to change it later
  setSelf: (T | DefaultValue | (T => T | DefaultValue)) => void,
  resetSelf: () => void,
  getSnapshot: () => Snapshot,

  // Subscribe to events
  // Called when React batch ends, but before global RecoilTransactionObserver
  onSet: ((newValue: T | DefaultValue, oldValue: T | DefaultValue) => void) => void,
}) => void | () => void; // Return optional cleanup handler
```

### Current Diff
This diff introduces a subset of the atom effect functionality:

```
// Effect is called the first time node is used with a <RecoilRoot>
type AtomEffect<T> = ({
  node: RecoilState<T>,
  trigger: 'get' | 'set',

  setSelf: (T | DefaultValue | (T => T | DefaultValue)) => void,
  resetSelf: () => void,
  getSnapshot: () => Snapshot,
}) => void;
```

### Example Usage
Example use of an atom as a local cache of state that can be bi-directionally synchronized with remote async storage.
```
const myAtom = atom({
  key: 'MyAtom',
  default: undefined,

  effects: [({node, setSelf, onSet}) => {
    const storage = getStorage(node.key);

    // set a value for synchronous init that overrides default
    setSelf(storage.get(value));

    // Subscribe to storage updates
    storage.subscribe(value => setSelf(value));

    // Subscribe local changes to update storage
    onSet(value =>
      value === undefined
        ? storage.reset()
        : storage.set(value)
    );

    // Subscribe to cleanup
    onCleanup(() => storage.unsubscribe());
  }],
});
```

Reviewed By: csantos42

Differential Revision: D21919545

fbshipit-source-id: b500602d7ca91972d28790ee3f895517fdc52320

src/core/Recoil_Snapshot.js

@@ -101,6 +101,7 @@ class Snapshot {
 function cloneTreeState(treeState: TreeState): TreeState {
   return {
     transactionMetadata: {...treeState.transactionMetadata},
+    knownAtoms: new Set(treeState.knownAtoms),
     atomValues: new Map(treeState.atomValues),
     nonvalidatedAtoms: new Map(treeState.nonvalidatedAtoms),
     dirtyAtoms: new Set(treeState.dirtyAtoms),

src/core/Recoil_State.js

@@ -23,11 +23,12 @@ type ComponentCallback = TreeState => void;
 export type TreeState = $ReadOnly<{
   // Information about the TreeState itself:
   transactionMetadata: {...},
-  dirtyAtoms: Set<NodeKey>,
 
   // ATOMS
+  knownAtoms: Set<NodeKey>,
   atomValues: AtomValues,
   nonvalidatedAtoms: Map<NodeKey, mixed>,
+  dirtyAtoms: Set<NodeKey>,
 
   // NODE GRAPH
   // Upstream Node dependencies
@@ -84,6 +85,7 @@ export type StoreRef = {
 function makeEmptyTreeState(): TreeState {
   return {
     transactionMetadata: {},
+    knownAtoms: new Set(),
     atomValues: new Map(),
     nonvalidatedAtoms: new Map(),
     dirtyAtoms: new Set(),

src/hooks/__tests__/Recoil_useRecoilCallback-test.js

@@ -229,10 +229,12 @@ describe('useRecoilCallback', () => {
       return null;
     }
 
+    const [Kick, kick] = componentThatReadsAndWritesAtom(kickAtom);
     const c = renderElements(
       <>
         <ReadsAtom atom={myAtom} />
         <Component />
+        <Kick />
       </>,
     );
 
@@ -244,10 +246,11 @@ describe('useRecoilCallback', () => {
       cb('SET');
       cb('UPDATE AGAIN');
     });
+    act(kick);
     expect(c.textContent).toEqual('"UPDATE AGAIN"');
   });
 
-  it('goes to snapshot', async () => {
+  it('goes to snapshot', () => {
     const myAtom = atom({
       key: 'Goto Snapshot From Callback',
       default: 'DEFAULT',
@@ -259,22 +262,29 @@ describe('useRecoilCallback', () => {
         const updatedSnapshot = snapshot.map(({set}) => {
           set(myAtom, 'SET IN SNAPSHOT');
         });
+        expect(updatedSnapshot.getLoadable(myAtom).contents).toEqual(
+          'SET IN SNAPSHOT',
+        );
         gotoSnapshot(updatedSnapshot);
       });
       return null;
     }
 
+    // Something with React and Jest requires this extra kick...
+    const [Kick, kick] = componentThatReadsAndWritesAtom(kickAtom);
+
     const c = renderElements(
       <>
         <ReadsAtom atom={myAtom} />
         <RecoilCallback />
+        <Kick />
       </>,
     );
 
     expect(c.textContent).toEqual('"DEFAULT"');
 
     act(cb);
-    await flushPromisesAndTimers();
+    act(kick);
     expect(c.textContent).toEqual('"SET IN SNAPSHOT"');
   });
 });

src/recoil_values/Recoil_atom.js

@@ -59,7 +59,8 @@
 
 import type {Loadable} from '../adt/Recoil_Loadable';
 import type {RecoilState, RecoilValue} from '../core/Recoil_RecoilValue';
-import type {NodeKey, TreeState} from '../core/Recoil_State';
+import type {Snapshot} from '../core/Recoil_Snapshot';
+import type {NodeKey, Store, TreeState} from '../core/Recoil_State';
 // @fb-only: import type {ScopeRules} from './Recoil_ScopedAtom';
 
 const {loadableWithValue} = require('../adt/Recoil_Loadable');
@@ -69,13 +70,15 @@ const {
   registerNode,
 } = require('../core/Recoil_Node');
 const {isRecoilValue} = require('../core/Recoil_RecoilValue');
+const {cloneSnapshot} = require('../core/Recoil_Snapshot');
 const {
   mapByDeletingFromMap,
   mapBySettingInMap,
   setByAddingToSet,
 } = require('../util/Recoil_CopyOnWrite');
 const deepFreezeValue = require('../util/Recoil_deepFreezeValue');
 const expectationViolation = require('../util/Recoil_expectationViolation');
+const invariant = require('../util/Recoil_invariant');
 const isPromise = require('../util/Recoil_isPromise');
 const nullthrows = require('../util/Recoil_nullthrows');
 // @fb-only: const {scopedAtom} = require('./Recoil_ScopedAtom');
@@ -95,9 +98,23 @@ export type PersistenceSettings<Stored> = $ReadOnly<{
   validator: (mixed, DefaultValue) => Stored | DefaultValue,
 }>;
 
+// Effect is called the first time a node is used with a <RecoilRoot>
+type AtomEffect<T> = ({
+  node: RecoilState<T>,
+  trigger: 'set' | 'get',
+
+  // Call synchronously to initialize value or async to change it later
+  setSelf: (
+    T | DefaultValue | ((T | DefaultValue) => T | DefaultValue),
+  ) => void,
+  resetSelf: () => void,
+  getSnapshot: () => Snapshot,
+}) => void; // TODO Allow returning a cleanup function
+
 export type AtomOptions<T> = $ReadOnly<{
   key: NodeKey,
   default: RecoilValue<T> | Promise<T> | T,
+  effects_UNSTABLE?: $ReadOnlyArray<AtomEffect<T>>,
   persistence_UNSTABLE?: PersistenceSettings<T>,
   // @fb-only: scopeRules_APPEND_ONLY_READ_THE_DOCS?: ScopeRules,
   dangerouslyAllowMutability?: boolean,
@@ -110,86 +127,178 @@ type BaseAtomOptions<T> = $ReadOnly<{
 
 function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
   const {key, persistence_UNSTABLE: persistence} = options;
-  return registerNode({
-    key,
-    options,
 
-    get: (_store, state: TreeState): [TreeState, Loadable<T>] => {
-      if (state.atomValues.has(key)) {
-        // atom value is stored in state
-        return [state, nullthrows(state.atomValues.get(key))];
-      } else if (state.nonvalidatedAtoms.has(key)) {
-        if (persistence == null) {
-          expectationViolation(
-            `Tried to restore a persisted value for atom ${key} but it has no persistence settings.`,
-          );
-          return [state, loadableWithValue(options.default)];
-        }
-        const nonvalidatedValue = state.nonvalidatedAtoms.get(key);
-        const validatedValue: T | DefaultValue = persistence.validator(
-          nonvalidatedValue,
-          DEFAULT_VALUE,
+  function initAtom(
+    store: Store,
+    initState: TreeState,
+    trigger: 'set' | 'get',
+  ): TreeState {
+    if (initState.knownAtoms.has(key)) {
+      return initState;
+    }
+
+    // Run Atom Effects
+    let initValue: T | DefaultValue = DEFAULT_VALUE;
+    if (options.effects_UNSTABLE != null) {
+      let duringInit = true;
+
+      function getSnapshot() {
+        return cloneSnapshot(
+          duringInit
+            ? initState
+            : store.getState().nextTree ?? store.getState().currentTree,
         );
+      }
 
-        return validatedValue instanceof DefaultValue
-          ? [
-              {
-                ...state,
-                nonvalidatedAtoms: mapByDeletingFromMap(
-                  state.nonvalidatedAtoms,
-                  key,
-                ),
-              },
-              loadableWithValue(options.default),
-            ]
-          : [
-              {
-                ...state,
-                atomValues: mapBySettingInMap(
-                  state.atomValues,
-                  key,
-                  loadableWithValue(validatedValue),
-                ),
-                nonvalidatedAtoms: mapByDeletingFromMap(
-                  state.nonvalidatedAtoms,
-                  key,
-                ),
-              },
-              loadableWithValue(validatedValue),
-            ];
-      } else {
+      function setSelf(
+        valueOrUpdater: T | DefaultValue | (T => T | DefaultValue),
+      ) {
+        if (duringInit) {
+          const currentValue: T =
+            initValue instanceof DefaultValue ? options.default : initValue;
+          initValue =
+            typeof valueOrUpdater === 'function'
+              ? // cast to any because we can't restrict type from being a function itself without losing support for opaque types
+                // flowlint-next-line unclear-type:off
+                (valueOrUpdater: any)(currentValue)
+              : valueOrUpdater;
+        } else {
+          store.replaceState(asyncState => {
+            let newState = asyncState;
+            const newValue: T | DefaultValue =
+              typeof valueOrUpdater !== 'function'
+                ? valueOrUpdater
+                : (() => {
+                    const [_, loadable] = myGet(store, asyncState);
+                    invariant(
+                      loadable.state === 'hasValue',
+                      "Recoil doesn't support async Atoms yet",
+                    );
+                    // cast to any because we can't restrict type from being a function itself without losing support for opaque types
+                    // flowlint-next-line unclear-type:off
+                    return (valueOrUpdater: any)(loadable.contents);
+                  })();
+            const [nextState, writtenNodes] = mySet(
+              store,
+              asyncState,
+              newValue,
+            );
+            newState = nextState;
+            store.fireNodeSubscriptions(writtenNodes, 'enqueue');
+            return newState;
+          });
+        }
+      }
+      const resetSelf = () => setSelf(DEFAULT_VALUE);
+
+      for (const effect of options.effects_UNSTABLE ?? []) {
+        effect({node, trigger, setSelf, resetSelf, getSnapshot});
+      }
+
+      duringInit = false;
+    }
+
+    return {
+      ...initState,
+      knownAtoms: setByAddingToSet(initState.knownAtoms, key),
+      atomValues: !(initValue instanceof DefaultValue)
+        ? mapBySettingInMap(
+            initState.atomValues,
+            key,
+            loadableWithValue(initValue),
+          )
+        : initState.atomValues,
+    };
+  }
+
+  function myGet(store: Store, initState: TreeState): [TreeState, Loadable<T>] {
+    const state = initAtom(store, initState, 'get');
+
+    if (state.atomValues.has(key)) {
+      // atom value is stored in state
+      return [state, nullthrows(state.atomValues.get(key))];
+    } else if (state.nonvalidatedAtoms.has(key)) {
+      if (persistence == null) {
+        expectationViolation(
+          `Tried to restore a persisted value for atom ${key} but it has no persistence settings.`,
+        );
         return [state, loadableWithValue(options.default)];
       }
-    },
+      const nonvalidatedValue = state.nonvalidatedAtoms.get(key);
+      const validatedValue: T | DefaultValue = persistence.validator(
+        nonvalidatedValue,
+        DEFAULT_VALUE,
+      );
 
-    set: (
-      _store,
-      state: TreeState,
-      newValue: T | DefaultValue,
-    ): [TreeState, $ReadOnlySet<NodeKey>] => {
-      if (__DEV__) {
-        if (options.dangerouslyAllowMutability !== true) {
-          deepFreezeValue(newValue);
-        }
+      return validatedValue instanceof DefaultValue
+        ? [
+            {
+              ...state,
+              nonvalidatedAtoms: mapByDeletingFromMap(
+                state.nonvalidatedAtoms,
+                key,
+              ),
+            },
+            loadableWithValue(options.default),
+          ]
+        : [
+            {
+              ...state,
+              atomValues: mapBySettingInMap(
+                state.atomValues,
+                key,
+                loadableWithValue(validatedValue),
+              ),
+              nonvalidatedAtoms: mapByDeletingFromMap(
+                state.nonvalidatedAtoms,
+                key,
+              ),
+            },
+            loadableWithValue(validatedValue),
+          ];
+    } else {
+      return [state, loadableWithValue(options.default)];
+    }
+  }
+
+  function mySet(
+    store: Store,
+    initState: TreeState,
+    newValue: T | DefaultValue,
+  ): [TreeState, $ReadOnlySet<NodeKey>] {
+    const state = initAtom(store, initState, 'set');
+
+    if (__DEV__) {
+      if (options.dangerouslyAllowMutability !== true) {
+        deepFreezeValue(newValue);
       }
-      return [
-        {
-          ...state,
-          dirtyAtoms: setByAddingToSet(state.dirtyAtoms, key),
-          atomValues:
-            newValue instanceof DefaultValue
-              ? mapByDeletingFromMap(state.atomValues, key)
-              : mapBySettingInMap(
-                  state.atomValues,
-                  key,
-                  loadableWithValue(newValue),
-                ),
-          nonvalidatedAtoms: mapByDeletingFromMap(state.nonvalidatedAtoms, key),
-        },
-        new Set([key]),
-      ];
-    },
+    }
+
+    return [
+      {
+        ...state,
+        dirtyAtoms: setByAddingToSet(state.dirtyAtoms, key),
+        atomValues:
+          newValue instanceof DefaultValue
+            ? mapByDeletingFromMap(state.atomValues, key)
+            : mapBySettingInMap(
+                state.atomValues,
+                key,
+                loadableWithValue(newValue),
+              ),
+        nonvalidatedAtoms: mapByDeletingFromMap(state.nonvalidatedAtoms, key),
+      },
+      new Set([key]),
+    ];
+  }
+
+  const node = registerNode({
+    key,
+    options,
+    get: myGet,
+    set: mySet,
   });
+  return node;
 }
 
 // prettier-ignore
@@ -240,6 +349,9 @@ function atomWithFallback<T>(
                     DEFAULT_VALUE,
                   ),
           },
+    // TODO Hack for now.
+    // flowlint-next-line unclear-type: off
+    effects_UNSTABLE: (options.effects_UNSTABLE: any),
   });
 
   return selector<T>({