Fix for ReactElement temporals in React branches (#2597)

Summary:
Release notes: none

This fixes a React reconciliation bug where lazy ReactElement temporals were not correctly being put into the correct branches. Previously, `wrapReactElementInBranchOrReturnValue` was only called on the "result" of an `applyBranchedLogicValue` where it should have been applied after each ReactElement creation given that `applyBranchedLogicValue` is recursive.
Pull Request resolved: https://github.com/facebook/prepack/pull/2597

Differential Revision: D10298262

Pulled By: trueadm

fbshipit-source-id: 4868ab7f2514734913a626e4e7e7e3259e4199b5

src/react/branching.js

@@ -180,14 +180,14 @@ export function getValueWithBranchingLogicApplied(
   searchAndFlagMismatchingNonHostTypes(parentX, parentY, 0);
 
   if (needsKeys) {
-    return applyBranchedLogicValue(realm, value);
+    return applyBranchedLogicValue(realm, value, false);
   }
   return value;
 }
 
 // When we apply branching logic, it means to add keys to all ReactElement nodes
 // we encounter, thus returning new ReactElements with the keys on them
-function applyBranchedLogicValue(realm: Realm, value: Value): Value {
+function applyBranchedLogicValue(realm: Realm, value: Value, inBranch: boolean): Value {
   if (
     value instanceof StringValue ||
     value instanceof NumberValue ||
@@ -197,9 +197,13 @@ function applyBranchedLogicValue(realm: Realm, value: Value): Value {
   ) {
     // terminal values
   } else if (value instanceof ObjectValue && isReactElement(value)) {
-    return addKeyToReactElement(realm, value);
+    if (inBranch) {
+      return wrapReactElementInBranchOrReturnValue(realm, addKeyToReactElement(realm, value));
+    } else {
+      return addKeyToReactElement(realm, value);
+    }
   } else if (value instanceof ArrayValue) {
-    let newArray = mapArrayValue(realm, value, elementValue => applyBranchedLogicValue(realm, elementValue));
+    let newArray = mapArrayValue(realm, value, elementValue => applyBranchedLogicValue(realm, elementValue, inBranch));
     newArray.makeFinal();
     return newArray;
   } else if (value instanceof AbstractValue && value.kind === "conditional") {
@@ -210,14 +214,14 @@ function applyBranchedLogicValue(realm: Realm, value: Value): Value {
       condValue,
       () => {
         return realm.evaluateForEffects(
-          () => wrapReactElementInBranchOrReturnValue(realm, applyBranchedLogicValue(realm, consequentVal)),
+          () => applyBranchedLogicValue(realm, consequentVal, true),
           null,
           "applyBranchedLogicValue consequent"
         );
       },
       () => {
         return realm.evaluateForEffects(
-          () => wrapReactElementInBranchOrReturnValue(realm, applyBranchedLogicValue(realm, alternateVal)),
+          () => applyBranchedLogicValue(realm, alternateVal, true),
           null,
           "applyBranchedLogicValue alternate"
         );

test/react/Reconciliation-test.js

@@ -48,6 +48,14 @@ it("Key nesting 8", () => {
   runTest(__dirname + "/Reconciliation/key-nesting-8.js");
 });
 
+it("Key nesting 9", () => {
+  runTest(__dirname + "/Reconciliation/key-nesting-9.js", {
+    expectedCreateElementCalls:
+      /* original 3 reactElements for 6 test cases */ 18 +
+      /* prepacked: one removed by inlining, but we have 6 test cases */ 12,
+  });
+});
+
 it("Key change", () => {
   runTest(__dirname + "/Reconciliation/key-change.js");
 });

test/react/Reconciliation/key-nesting-9.js

@@ -0,0 +1,39 @@
+var React = require("react");
+
+function Foo(props) {
+  return [<div ref={props.deepProperty.callback} key="0" />];
+}
+
+function Bar(props) {
+  return [<div ref={props.deepProperty.callback} key="0" />];
+}
+
+function App(props) {
+  return props.foo ? <Foo {...props} /> : <Bar {...props} />;
+}
+
+if (this.__optimizeReactComponentTree) {
+  __optimizeReactComponentTree(App);
+}
+
+App.getTrials = function(renderer, Root) {
+  let counter = 0;
+  let nodes = [];
+  function callback(node) {
+    nodes.push(node);
+    counter++;
+  }
+  renderer.update(<Root deepProperty={{ callback: callback }} foo={true} />);
+  renderer.update(<Root deepProperty={{ callback: callback }} foo={true} />);
+  renderer.update(<Root deepProperty={{ callback: callback }} foo={false} />);
+  renderer.update(<Root deepProperty={{ callback: callback }} foo={false} />);
+  renderer.update(<Root deepProperty={{ callback: callback }} foo={true} />);
+  renderer.update(<Root deepProperty={{ callback: callback }} foo={true} />);
+  let results = [];
+  results.push(["ensure ref is called on every change", counter]);
+  results.push(["ensure refs are cleared", nodes.map(Boolean)]);
+  results.push(["ensure refs are different", new Set(nodes).size]);
+  return results;
+};
+
+module.exports = App;

test/react/__snapshots__/Reconciliation-test.js.snap

@@ -2700,6 +2700,126 @@ ReactStatistics {
 }
 `;
 
+exports[`Key nesting 9: (JSX => JSX) 1`] = `
+ReactStatistics {
+  "componentsEvaluated": 3,
+  "evaluatedRootNodes": Array [
+    Object {
+      "children": Array [
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Foo",
+          "status": "INLINED",
+        },
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Bar",
+          "status": "INLINED",
+        },
+      ],
+      "message": "",
+      "name": "App",
+      "status": "ROOT",
+    },
+  ],
+  "inlinedComponents": 2,
+  "optimizedNestedClosures": 0,
+  "optimizedTrees": 1,
+}
+`;
+
+exports[`Key nesting 9: (JSX => createElement) 1`] = `
+ReactStatistics {
+  "componentsEvaluated": 3,
+  "evaluatedRootNodes": Array [
+    Object {
+      "children": Array [
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Foo",
+          "status": "INLINED",
+        },
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Bar",
+          "status": "INLINED",
+        },
+      ],
+      "message": "",
+      "name": "App",
+      "status": "ROOT",
+    },
+  ],
+  "inlinedComponents": 2,
+  "optimizedNestedClosures": 0,
+  "optimizedTrees": 1,
+}
+`;
+
+exports[`Key nesting 9: (createElement => JSX) 1`] = `
+ReactStatistics {
+  "componentsEvaluated": 3,
+  "evaluatedRootNodes": Array [
+    Object {
+      "children": Array [
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Foo",
+          "status": "INLINED",
+        },
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Bar",
+          "status": "INLINED",
+        },
+      ],
+      "message": "",
+      "name": "App",
+      "status": "ROOT",
+    },
+  ],
+  "inlinedComponents": 2,
+  "optimizedNestedClosures": 0,
+  "optimizedTrees": 1,
+}
+`;
+
+exports[`Key nesting 9: (createElement => createElement) 1`] = `
+ReactStatistics {
+  "componentsEvaluated": 3,
+  "evaluatedRootNodes": Array [
+    Object {
+      "children": Array [
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Foo",
+          "status": "INLINED",
+        },
+        Object {
+          "children": Array [],
+          "message": "",
+          "name": "Bar",
+          "status": "INLINED",
+        },
+      ],
+      "message": "",
+      "name": "App",
+      "status": "ROOT",
+    },
+  ],
+  "inlinedComponents": 2,
+  "optimizedNestedClosures": 0,
+  "optimizedTrees": 1,
+}
+`;
+
 exports[`Key nesting: (JSX => JSX) 1`] = `
 ReactStatistics {
   "componentsEvaluated": 6,