VirtualizedList up-to-date state: Enforce correct setState() usage

Summary:
This diff is part of an overall stack, meant to fix incorrect usage of `setState()` in `VirtualizedList`, which triggers new invariant checks added in `VirtualizedList_EXPERIMENTAL`. See the stack summary below for more information on the broader change.

## Diff Summary

This adds a `StateSafePureComponent`, which will override `setState`, along with `this.props` and `this.state`, in order to add runtime enforcement that `this.props` and `this.state` are not accessed during a state update where we should be relying on parameter instead. This should be landed after the fixes for unsafe `this.props`/`this.state` usage.

## Stack Summary
`VirtualizedList`'s component state is a set of cells to render. This state is set via the `setState()` class component API. The main "tick" function `VirtualizedList._updateCellsToRender()` calculates this new state using a combination of the current component state, and instance-local state like maps, measurement caches, etc.

From: https://reactjs.org/docs/state-and-lifecycle.html#state-updates-may-be-asynchronous
 ---
> React may batch multiple setState() calls into a single update for performance. Because this.props and this.state may be updated asynchronously, you should not rely on their values for calculating the next state. For example, this code may fail to update the counter:

```
// Wrong
this.setState({
  counter: this.state.counter + this.props.increment,
});
```
> To fix it, use a second form of setState() that accepts a function rather than an object. That function will receive the previous state as the first argument, and the props at the time the update is applied as the second argument:
```
// Correct
this.setState((state, props) => ({
  counter: state.counter + props.increment
}));
```
 ---
`_updateCellsToRender()` transitively calls many functions which will read directly from `this.props` or `this.state` instead of the value passed by the state updater. This intermittently fires invariant violations, when there is a mismatch.

This diff migrates all usages of `props` and `state` during state update to the values provied in `setState()`. To prevent future mismatch, and to provide better clarity on when it is safe to use `this.props`, `this.state`, I overrode `setState` to fire an invariant violation if it is accessed when it is unsafe to:

{F756963772}

Changelog:
[Internal][Changed] - Enable setState() hooks

Reviewed By: rshest

Differential Revision: D38294338

fbshipit-source-id: d04ff39f68df90adc9f4680887d308d997903675

Libraries/Lists/StateSafePureComponent.js

@@ -0,0 +1,85 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow strict
+ * @format
+ */
+
+import * as React from 'react';
+import invariant from 'invariant';
+
+/**
+ * `setState` is called asynchronously, and should not rely on the value of
+ * `this.props` or `this.state`:
+ * https://reactjs.org/docs/state-and-lifecycle.html#state-updates-may-be-asynchronous
+ *
+ * SafePureComponent adds runtime enforcement, to catch cases where these
+ * variables are read in a state updater function, instead of the ones passed
+ * in.
+ */
+export default class StateSafePureComponent<
+  Props,
+  State: interface {},
+> extends React.PureComponent<Props, State> {
+  _inAsyncStateUpdate = false;
+
+  constructor(props: Props) {
+    super(props);
+    this._installSetStateHooks();
+  }
+
+  setState(
+    partialState: ?($Shape<State> | ((State, Props) => ?$Shape<State>)),
+    callback?: () => mixed,
+  ): void {
+    if (typeof partialState === 'function') {
+      super.setState((state, props) => {
+        this._inAsyncStateUpdate = true;
+        let ret;
+        try {
+          ret = partialState(state, props);
+        } catch (err) {
+          throw err;
+        } finally {
+          this._inAsyncStateUpdate = false;
+        }
+        return ret;
+      }, callback);
+    } else {
+      super.setState(partialState, callback);
+    }
+  }
+
+  _installSetStateHooks() {
+    const that = this;
+    let {props, state} = this;
+
+    Object.defineProperty(this, 'props', {
+      get() {
+        invariant(
+          !that._inAsyncStateUpdate,
+          '"this.props" should not be accessed during state updates',
+        );
+        return props;
+      },
+      set(newProps: Props) {
+        props = newProps;
+      },
+    });
+    Object.defineProperty(this, 'state', {
+      get() {
+        invariant(
+          !that._inAsyncStateUpdate,
+          '"this.state" should not be acceessed during state updates',
+        );
+        return state;
+      },
+      set(newState: State) {
+        state = newState;
+      },
+    });
+  }
+}

Libraries/Lists/VirtualizedList_EXPERIMENTAL.js

@@ -43,6 +43,7 @@ import * as React from 'react';
 
 import {CellRenderMask} from './CellRenderMask';
 import clamp from '../Utilities/clamp';
+import StateSafePureComponent from './StateSafePureComponent';
 
 const RefreshControl = require('../Components/RefreshControl/RefreshControl');
 const ScrollView = require('../Components/ScrollView/ScrollView');
@@ -159,7 +160,7 @@ function findLastWhere<T>(
  * - As an effort to remove defaultProps, use helper functions when referencing certain props
  *
  */
-class VirtualizedList extends React.PureComponent<Props, State> {
+class VirtualizedList extends StateSafePureComponent<Props, State> {
   static contextType: typeof VirtualizedListContext = VirtualizedListContext;
 
   // scrollToEnd may be janky without getItemLayout prop
@@ -2252,4 +2253,5 @@ const styles = StyleSheet.create({
   },
 });
 
+VirtualizedList.displayName = 'VirtualizedList_EXPERIMENTAL';
 module.exports = VirtualizedList;