Improve inlining of React.createContext (#2098)

Summary:
Release notes: none

Whilst working on adding React Native mocks, I ran into cases where context should be inlining but it wasn't. It now is inlining far better, with more test coverage. Furthermore, we now have a new config flag to pass to `__optimizeReactComponentTree`, in the case of `isRoot` to state that the tree is a root component tree (rather than a a branch of another tree).
Closes https://github.com/facebook/prepack/pull/2098

Differential Revision: D8348381

Pulled By: trueadm

fbshipit-source-id: 5e01bd77437e8bc3d1f22ff47d668897152203a0

scripts/test-react.js

@@ -784,7 +784,7 @@ function runTestSuite(outputJsx, shouldTranspileSource) {
         await runTest(directory, "react-context3.js");
       });
 
-      it("React Context 4", async () => {
+      it.only("React Context 4", async () => {
         await runTest(directory, "react-context4.js");
       });
 
@@ -795,6 +795,26 @@ function runTestSuite(outputJsx, shouldTranspileSource) {
       it("React Context 6", async () => {
         await runTest(directory, "react-context6.js");
       });
+
+      it("React Context 7", async () => {
+        await runTest(directory, "react-context7.js");
+      });
+
+      it("React Context from root tree", async () => {
+        await runTest(directory, "react-root-context.js");
+      });
+
+      it("React Context from root tree 2", async () => {
+        await runTest(directory, "react-root-context2.js");
+      });
+
+      it("React Context from root tree 3", async () => {
+        await runTest(directory, "react-root-context3.js");
+      });
+
+      it("React Context from root tree 4", async () => {
+        await runTest(directory, "react-root-context4.js");
+      });
     });
 
     describe("First render only", () => {
@@ -852,6 +872,10 @@ function runTestSuite(outputJsx, shouldTranspileSource) {
         await runTest(directory, "react-context5.js");
       });
 
+      it("React Context 6", async () => {
+        await runTest(directory, "react-context6.js");
+      });
+
       it.skip("Replace this in callbacks", async () => {
         await runTest(directory, "replace-this-in-callbacks.js");
       });

src/intrinsics/fb-www/react-mocks.js

@@ -559,7 +559,6 @@ export function createMockReact(realm: Realm, reactRequireName: string): ObjectV
       Properties.Set(realm, providerObject, "context", consumer, true);
 
       Properties.Set(realm, consumerObject, "Provider", provider, true);
-
       return consumer;
     }
   );

src/react/experimental-server-rendering/rendering.js

@@ -475,7 +475,7 @@ export function renderToString(
   staticMarkup: boolean
 ): StringValue | AbstractValue {
   let reactStatistics = new ReactStatistics();
-  let reconciler = new Reconciler(realm, { firstRenderOnly: true }, reactStatistics);
+  let reconciler = new Reconciler(realm, { firstRenderOnly: true, isRoot: true }, reactStatistics);
   let typeValue = getProperty(realm, reactElement, "type");
   let propsValue = getProperty(realm, reactElement, "props");
   let evaluatedRootNode = createReactEvaluatedNode("ROOT", getComponentName(realm, typeValue));

src/react/reconcilation.js

@@ -139,10 +139,10 @@ export class Reconciler {
     this.realm = realm;
     this.statistics = statistics;
     this.logger = logger;
+    this.componentTreeConfig = componentTreeConfig;
     this.componentTreeState = this._createComponentTreeState();
     this.alreadyEvaluatedRootNodes = new Map();
     this.alreadyEvaluatedNestedClosures = new Set();
-    this.componentTreeConfig = componentTreeConfig;
     this.nestedOptimizedClosures = [];
     this.branchedComponentTrees = [];
   }
@@ -420,13 +420,16 @@ export class Reconciler {
     let lastValueProp = getProperty(this.realm, contextConsumer, "currentValue");
     this._incremementReferenceForContextNode(contextConsumer);
 
+    let valueProp;
     // if we have a value prop, set it
     if (propsValue instanceof ObjectValue || propsValue instanceof AbstractObjectValue) {
-      let valueProp = Get(this.realm, propsValue, "value");
+      valueProp = Get(this.realm, propsValue, "value");
       setContextCurrentValue(contextConsumer, valueProp);
     }
-    if (this.componentTreeConfig.firstRenderOnly) {
-      if (propsValue instanceof ObjectValue) {
+    if (propsValue instanceof ObjectValue) {
+      // if the value is abstract, we need to keep the render prop as unless
+      // we are in firstRenderOnly mode, where we can just inline the abstract value
+      if (!(valueProp instanceof AbstractValue) || this.componentTreeConfig.firstRenderOnly) {
         let resolvedReactElement = this._resolveReactElementHostChildren(
           componentType,
           reactElement,
@@ -481,7 +484,10 @@ export class Reconciler {
     this.componentTreeState.contextNodeReferences.set(contextNode, references);
   }
 
-  _hasReferenceForContextNode(contextNode: ObjectValue | AbstractObjectValue): boolean {
+  _isContextValueKnown(contextNode: ObjectValue | AbstractObjectValue): boolean {
+    if (this.componentTreeConfig.isRoot) {
+      return true;
+    }
     if (this.componentTreeState.contextNodeReferences.has(contextNode)) {
       let references = this.componentTreeState.contextNodeReferences.get(contextNode);
       if (!references) {
@@ -496,6 +502,7 @@ export class Reconciler {
     componentType: Value,
     reactElement: ObjectValue,
     context: ObjectValue | AbstractObjectValue,
+    branchStatus: BranchStatusEnum,
     evaluatedNode: ReactEvaluatedNode
   ): Value | void {
     let typeValue = getProperty(this.realm, reactElement, "type");
@@ -510,18 +517,20 @@ export class Reconciler {
 
         this._findReactComponentTrees(propsValue, evaluatedChildNode, "NORMAL_FUNCTIONS");
         if (renderProp instanceof ECMAScriptSourceFunctionValue) {
-          if (this.componentTreeConfig.firstRenderOnly) {
-            if (typeValue instanceof ObjectValue || typeValue instanceof AbstractObjectValue) {
-              // make sure this context is in our tree
-              if (this._hasReferenceForContextNode(typeValue)) {
-                let valueProp = Get(this.realm, typeValue, "currentValue");
+          if (typeValue instanceof ObjectValue || typeValue instanceof AbstractObjectValue) {
+            // make sure this context is in our tree
+            if (this._isContextValueKnown(typeValue)) {
+              let valueProp = Get(this.realm, typeValue, "currentValue");
+              // if the value is abstract, we need to keep the render prop as unless
+              // we are in firstRenderOnly mode, where we can just inline the abstract value
+              if (!(valueProp instanceof AbstractValue) || this.componentTreeConfig.firstRenderOnly) {
                 let result = getValueFromFunctionCall(this.realm, renderProp, this.realm.intrinsics.undefined, [
                   valueProp,
                 ]);
                 this.statistics.inlinedComponents++;
                 this.statistics.componentsEvaluated++;
                 evaluatedChildNode.status = "INLINED";
-                return result;
+                return this._resolveDeeply(componentType, result, context, branchStatus, evaluatedNode);
               }
             }
           }
@@ -1152,7 +1161,13 @@ export class Reconciler {
           );
         }
         case "CONTEXT_CONSUMER": {
-          result = this._resolveContextConsumerComponent(componentType, reactElement, context, evaluatedNode);
+          result = this._resolveContextConsumerComponent(
+            componentType,
+            reactElement,
+            context,
+            branchStatus,
+            evaluatedNode
+          );
           break;
         }
         case "FORWARD_REF": {

src/react/utils.js

@@ -760,6 +760,7 @@ export function convertConfigObjectToReactComponentTreeConfig(
 ): ReactComponentTreeConfig {
   // defaults
   let firstRenderOnly = false;
+  let isRoot = false;
 
   if (!(config instanceof UndefinedValue)) {
     for (let [key] of config.properties) {
@@ -771,6 +772,8 @@ export function convertConfigObjectToReactComponentTreeConfig(
         if (typeof value === "boolean") {
           if (key === "firstRenderOnly") {
             firstRenderOnly = value;
+          } else if (key === "isRoot") {
+            isRoot = value;
           }
         }
       } else {
@@ -787,6 +790,7 @@ export function convertConfigObjectToReactComponentTreeConfig(
   }
   return {
     firstRenderOnly,
+    isRoot,
   };
 }
 

src/types.js

@@ -347,6 +347,7 @@ export type ReactHint = {| firstRenderValue: Value, object: ObjectValue, propert
 
 export type ReactComponentTreeConfig = {
   firstRenderOnly: boolean,
+  isRoot: boolean,
 };
 
 export type DebugServerType = {

test/react/first-render-only/react-context6.js

@@ -0,0 +1,45 @@
+var React = require('React');
+// the JSX transform converts to React, so we need to add it back in
+this['React'] = React;
+
+var { Provider, Consumer } = React.createContext(null);
+
+function Child2(props) {
+  return <span>{props.title}</span>;
+}
+
+function Child(props) {
+  return (
+    <div>
+      <Consumer>
+        {value => {
+          return <span><Child2 title={value} /></span>
+        }}
+      </Consumer>
+    </div>
+  )
+}
+
+function App(props) {
+  return (
+    <div>
+      <Provider value="b">
+        <Child />
+      </Provider>
+      <Child />
+    </div>
+  );
+}
+
+App.getTrials = function(renderer, Root) {
+  renderer.update(<Root />);
+  return [['render props context', renderer.toJSON()]];
+};
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App, {
+    firstRenderOnly: true,
+  });
+}
+
+module.exports = App;

test/react/render-props/react-context.js

@@ -2,14 +2,14 @@ var React = require('React');
 // the JSX transform converts to React, so we need to add it back in
 this['React'] = React;
 
-var { Provider, Consumer } = React.createContext(null);
+var { Provider, Consumer } = React.createContext("bar");
 
 function Child(props) {
   return (
     <div>
       <Consumer>
         {context => {
-          return <span>123</span>
+          return <span>{context}</span>
         }}
       </Consumer>
     </div>
@@ -18,15 +18,19 @@ function Child(props) {
 
 function App(props) {
   return (
-    <Provider>
+    <Provider value={"foo"}>
       <Child />
     </Provider>
   );
 }
 
 App.getTrials = function(renderer, Root) {
+  let results = [];
   renderer.update(<Root />);
-  return [['render props context', renderer.toJSON()]];
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
 };
 
 if (this.__optimizeReactComponentTree) {

test/react/render-props/react-context2.js

@@ -2,7 +2,7 @@ var React = require('React');
 // the JSX transform converts to React, so we need to add it back in
 this['React'] = React;
 
-var { Provider, Consumer } = React.createContext(null);
+var { Provider, Consumer } = React.createContext("foo");
 
 function Child(props) {
   return (
@@ -25,8 +25,12 @@ function App(props) {
 }
 
 App.getTrials = function(renderer, Root) {
+  let results = [];
   renderer.update(<Root />);
-  return [['render props context', renderer.toJSON()]];
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
 };
 
 if (this.__optimizeReactComponentTree) {

test/react/render-props/react-context3.js

@@ -22,12 +22,20 @@ function App(props) {
 App.Ctx = Ctx;
 
 App.getTrials = function(renderer, Root) {
+  let results = [];
   renderer.update((
     <Root.Ctx.Provider value={5}>
       <Root />
     </Root.Ctx.Provider>
   ));
-  return [['render props context', renderer.toJSON()]];
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update((
+    <Root.Ctx.Provider value={5}>
+      <Root />
+    </Root.Ctx.Provider>
+  ));
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
 };
 
 if (this.__optimizeReactComponentTree) {

test/react/render-props/react-context4.js

@@ -28,8 +28,12 @@ function App(props) {
 }
 
 App.getTrials = function(renderer, Root) {
+  let results = [];
   renderer.update(<Root />);
-  return [['render props context', renderer.toJSON()]];
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
 };
 
 if (this.__optimizeReactComponentTree) {

test/react/render-props/react-context5.js

@@ -28,8 +28,12 @@ function App(props) {
 }
 
 App.getTrials = function(renderer, Root) {
+  let results = [];
   renderer.update(<Root />);
-  return [['render props context', renderer.toJSON()]];
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
 };
 
 if (this.__optimizeReactComponentTree) {

test/react/render-props/react-context6.js

@@ -31,8 +31,12 @@ function App(props) {
 }
 
 App.getTrials = function(renderer, Root) {
+  let results = [];
   renderer.update(<Root />);
-  return [['render props context', renderer.toJSON()]];
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
 };
 
 if (this.__optimizeReactComponentTree) {

test/react/render-props/react-context7.js

@@ -0,0 +1,40 @@
+var React = require('React');
+// the JSX transform converts to React, so we need to add it back in
+this['React'] = React;
+
+var { Provider, Consumer } = React.createContext(null);
+
+function Child(props) {
+  var renderProp = function(value) {
+    return <span>{value}</span>
+  }
+
+  return (
+    <div>
+      <Consumer>{renderProp}</Consumer>
+    </div>
+  )
+}
+
+function App(props) {
+  return (
+    <Provider value={props.dynamicValue}>
+      <Child />
+    </Provider>
+  );
+}
+
+App.getTrials = function(renderer, Root) {
+  let results = [];
+  renderer.update(<Root dynamicValue={5} />);
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root dynamicValue={7} />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
+};
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App);
+}
+
+module.exports = App;

test/react/render-props/react-root-context.js

@@ -0,0 +1,51 @@
+var React = require('React');
+// the JSX transform converts to React, so we need to add it back in
+this['React'] = React;
+
+var { Provider, Consumer } = React.createContext("bar");
+
+function Child(props) {
+  var x = function(context) {
+    var click = function () {
+      return x;
+    }
+
+    return <span onClick={click}>{props.x}</span>
+  }
+
+  return (
+    <div>
+      <Consumer>
+        {x}
+      </Consumer>
+    </div>
+  )
+}
+
+function App(props) {
+  return (
+    <div>
+      <Provider value={"foo"}>
+        <Child />
+      </Provider>
+      <Child />
+    </div>
+  );
+}
+
+App.getTrials = function(renderer, Root) {
+  let results = [];
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
+};
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App, {
+    isRoot: true,
+  });
+}
+
+module.exports = App;

test/react/render-props/react-root-context2.js

@@ -0,0 +1,45 @@
+var React = require('React');
+// the JSX transform converts to React, so we need to add it back in
+this['React'] = React;
+
+var { Provider, Consumer } = React.createContext("bar");
+
+function Child(props) {
+  return (
+    <div>
+      <Consumer>
+        {value => {
+          return <span>{value}</span>
+        }}
+      </Consumer>
+    </div>
+  )
+}
+
+function App(props) {
+  return (
+    <Provider value="a">
+      <Provider value="b">
+        <Child />
+      </Provider>
+      <Child />
+    </Provider>
+  );
+}
+
+App.getTrials = function(renderer, Root) {
+  let results = [];
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
+};
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App, {
+    isRoot: true,
+  });
+}
+
+module.exports = App;

test/react/render-props/react-root-context3.js

@@ -0,0 +1,47 @@
+var React = require('React');
+// the JSX transform converts to React, so we need to add it back in
+this['React'] = React;
+
+var { Provider, Consumer } = React.createContext("bar");
+
+function Child(props) {
+  return (
+    <div>
+      <Consumer>
+        {value => {
+          return <span>{value}</span>
+        }}
+      </Consumer>
+    </div>
+  )
+}
+
+var x = (
+  <Provider value="a">
+    <Provider value="b">
+      <Child />
+    </Provider>
+    <Child />
+  </Provider>
+);
+
+function App(props) {
+  return x;
+}
+
+App.getTrials = function(renderer, Root) {
+  let results = [];
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
+};
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App, {
+    isRoot: true,
+  });
+}
+
+module.exports = App;

test/react/render-props/react-root-context4.js

@@ -0,0 +1,42 @@
+var React = require('React');
+// the JSX transform converts to React, so we need to add it back in
+this['React'] = React;
+
+var { Provider, Consumer } = React.createContext(null);
+
+function Child(props) {
+  var renderProp = function(value) {
+    return <span>{value}</span>
+  }
+
+  return (
+    <div>
+      <Consumer>{renderProp}</Consumer>
+    </div>
+  )
+}
+
+function App(props) {
+  return (
+    <Provider value={props.dynamicValue}>
+      <Child />
+    </Provider>
+  );
+}
+
+App.getTrials = function(renderer, Root) {
+  let results = [];
+  renderer.update(<Root dynamicValue={5} />);
+  results.push(['render props context', renderer.toJSON()]);
+  renderer.update(<Root dynamicValue={7} />);
+  results.push(['render props context', renderer.toJSON()]);
+  return results;
+};
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App, {
+    isRoot: true,
+  });
+}
+
+module.exports = App;