Ensure we clone temporalAlias values when we clone ReactElements (#2112)

Summary:
Release notes: none

When working on https://github.com/facebook/prepack/pull/2107, I noticed that we weren't cloning `temporalAlias` values, which is the wrong approach. We instead were re-using them on cloned ReactElements. The current approach worked fine with concrete ReactElement values, but when we introduce temporal ReactElement values that have dependencies on a shared `temporalAlias`, we run into very complex problems.

This PR makes it so we can clone `temporalAlias` values on `ObjectValues` (used for snapshotting).  This only used for React code – so it should not affect the other parts of Prepack. I've tested this on our internal bundle and our tests and they all pass, we create more code – but this is a *better* approach. We can remove the bloated code at a later point, but this aims to keep our approach consistent.
Closes https://github.com/facebook/prepack/pull/2112

Differential Revision: D8379812

Pulled By: trueadm

fbshipit-source-id: 04bf133be7a97cec3389a99584be1c042371ee07

src/intrinsics/ecma262/Object.js

@@ -148,10 +148,11 @@ export default function(realm: Realm): NativeFunctionValue {
         to.makeSimple();
 
         // Tell serializer that it may add properties to to only after temporalTo has been emitted
+        let temporalArgs = [ObjectAssign, to, ...delayedSources];
         let temporalTo = AbstractValue.createTemporalFromBuildFunction(
           realm,
           ObjectValue,
-          [ObjectAssign, to, ...delayedSources],
+          temporalArgs,
           ([methodNode, targetNode, ...sourceNodes]: Array<BabelNodeExpression>) => {
             return t.callExpression(methodNode, [targetNode, ...sourceNodes]);
           },
@@ -165,6 +166,10 @@ export default function(realm: Realm): NativeFunctionValue {
           temporalTo.values = new ValuesDomain(to);
         }
         to.temporalAlias = temporalTo;
+        // Store the args for the temporal so we can easily clone
+        // and reconstruct the temporal at another point, rather than
+        // mutate the existing temporal
+        realm.temporalAliasArgs.set(temporalTo, temporalArgs);
       }
       return to;
     });

src/react/elements.js

@@ -172,10 +172,13 @@ function createPropsObject(
         }
         let defaultPropsHelper = realm.react.defaultPropsHelper;
         invariant(defaultPropsHelper !== undefined);
+        let snapshot = props.getSnapshot();
+        props.temporalAlias = snapshot;
+        let temporalArgs = [defaultPropsHelper, snapshot, defaultProps];
         let temporalTo = AbstractValue.createTemporalFromBuildFunction(
           realm,
           ObjectValue,
-          [defaultPropsHelper, props.getSnapshot(), defaultProps],
+          temporalArgs,
           ([methodNode, ..._args]) => {
             return t.callExpression(methodNode, ((_args: any): Array<any>));
           },
@@ -189,6 +192,10 @@ function createPropsObject(
           temporalTo.values = new ValuesDomain(props);
         }
         props.temporalAlias = temporalTo;
+        // Store the args for the temporal so we can easily clone
+        // and reconstruct the temporal at another point, rather than
+        // mutate the existing temporal
+        realm.temporalAliasArgs.set(temporalTo, temporalArgs);
       }
     }
   } else {

src/react/utils.js

@@ -12,7 +12,7 @@
 import { Realm, Effects } from "../realm.js";
 import { ValuesDomain } from "../domains/index.js";
 import { AbruptCompletion, PossiblyNormalCompletion } from "../completions.js";
-import type { BabelNode, BabelNodeJSXIdentifier } from "babel-types";
+import type { BabelNode, BabelNodeJSXIdentifier, BabelNodeExpression } from "babel-types";
 import { parseExpression } from "babylon";
 import {
   AbstractObjectValue,
@@ -931,6 +931,38 @@ export function createInternalReactElement(
   return obj;
 }
 
+function applyClonedTemporalAlias(realm: Realm, props: ObjectValue, clonedProps: ObjectValue): void {
+  let temporalAlias = props.temporalAlias;
+  invariant(temporalAlias !== undefined);
+  if (temporalAlias.kind === "conditional") {
+    // Leave in for now, we should deal with this later, but there might
+    // be a better option.
+    invariant(false, "TODO applyClonedTemporalAlias conditional");
+  }
+  let temporalArgs = realm.temporalAliasArgs.get(temporalAlias);
+  invariant(temporalArgs !== undefined);
+  // replace the original props with the cloned one
+  let newTemporalArgs = temporalArgs.map(arg => (arg === props ? clonedProps : arg));
+
+  let temporalTo = AbstractValue.createTemporalFromBuildFunction(
+    realm,
+    ObjectValue,
+    newTemporalArgs,
+    ([methodNode, targetNode, ...sourceNodes]: Array<BabelNodeExpression>) => {
+      return t.callExpression(methodNode, [targetNode, ...sourceNodes]);
+    },
+    { skipInvariant: true }
+  );
+  invariant(temporalTo instanceof AbstractObjectValue);
+  invariant(clonedProps instanceof ObjectValue);
+  temporalTo.values = new ValuesDomain(clonedProps);
+  clonedProps.temporalAlias = temporalTo;
+  // Store the args for the temporal so we can easily clone
+  // and reconstruct the temporal at another point, rather than
+  // mutate the existing temporal
+  realm.temporalAliasArgs.set(temporalTo, newTemporalArgs);
+}
+
 export function cloneProps(
   realm: Realm,
   props: ObjectValue,
@@ -959,7 +991,7 @@ export function cloneProps(
     flagPropsWithNoPartialKeyOrRef(realm, clonedProps);
   }
   if (props.temporalAlias !== undefined) {
-    clonedProps.temporalAlias = props.temporalAlias;
+    applyClonedTemporalAlias(realm, props, clonedProps);
   }
   clonedProps.makeFinal();
   return clonedProps;
@@ -1003,7 +1035,9 @@ export function applyObjectAssignConfigsForReactElement(realm: Realm, to: Object
                 Properties.Set(realm, to, propName, Get(realm, source, propName), true);
               }
             }
-            delayedSources.push(source.getSnapshot());
+            let snapshot = source.getSnapshot();
+            delayedSources.push(snapshot);
+            source.temporalAlias = snapshot;
           } else {
             // if we are dealing with an abstract object or one that is partial, then
             // we don't try and copy its properties over as there's no guarantee they are
@@ -1012,7 +1046,9 @@ export function applyObjectAssignConfigsForReactElement(realm: Realm, to: Object
               // Make it implicit again since it is getting delayed into an Object.assign call.
               delayedSources.push(source.args[0]);
             } else {
-              delayedSources.push(source.getSnapshot());
+              let snapshot = source.getSnapshot();
+              delayedSources.push(snapshot);
+              source.temporalAlias = snapshot;
             }
             // if to has properties, we better remove them because after the temporal call to Object.assign we don't know their values anymore
             if (to.hasStringOrSymbolProperties()) {
@@ -1024,10 +1060,11 @@ export function applyObjectAssignConfigsForReactElement(realm: Realm, to: Object
         // prepare our temporal Object.assign fallback
         to.makePartial();
         to.makeSimple();
+        let temporalArgs = [objAssign, to, ...delayedSources];
         let temporalTo = AbstractValue.createTemporalFromBuildFunction(
           realm,
           ObjectValue,
-          [objAssign, to, ...delayedSources],
+          temporalArgs,
           ([methodNode, ..._args]) => {
             return t.callExpression(methodNode, ((_args: any): Array<any>));
           },
@@ -1036,6 +1073,10 @@ export function applyObjectAssignConfigsForReactElement(realm: Realm, to: Object
         invariant(temporalTo instanceof AbstractObjectValue);
         temporalTo.values = new ValuesDomain(to);
         to.temporalAlias = temporalTo;
+        // Store the args for the temporal so we can easily clone
+        // and reconstruct the temporal at another point, rather than
+        // mutate the existing temporal
+        realm.temporalAliasArgs.set(temporalTo, temporalArgs);
         return;
       } else {
         throw error;

src/realm.js

@@ -244,6 +244,7 @@ export class Realm {
     this.evaluators = (Object.create(null): any);
     this.partialEvaluators = (Object.create(null): any);
     this.$GlobalEnv = ((undefined: any): LexicalEnvironment);
+    this.temporalAliasArgs = new WeakMap();
 
     this.react = {
       abstractHints: new WeakMap(),
@@ -328,6 +329,12 @@ export class Realm {
   $GlobalEnv: LexicalEnvironment;
   intrinsics: Intrinsics;
 
+  // temporalAliasArgs is used to map a temporal abstract object value
+  // to its respective temporal args used to originally create the temporal.
+  // This is used to "clone" immutable objects where they have a dependency
+  // on a temporal alias (for example, Object.assign) when used with snapshotting
+  temporalAliasArgs: WeakMap<AbstractObjectValue | ObjectValue, Array<Value>>;
+
   react: {
     // reactHints are generated to help improve the effeciency of the React reconciler when
     // operating on a tree of React components. We can use reactHint to mark AbstractValues

src/serializer/ResidualReactElementSerializer.js

@@ -14,7 +14,7 @@ import { ResidualHeapSerializer } from "./ResidualHeapSerializer.js";
 import { canHoistReactElement } from "../react/hoisting.js";
 import * as t from "babel-types";
 import type { BabelNode, BabelNodeExpression } from "babel-types";
-import { AbstractObjectValue, AbstractValue, ObjectValue, SymbolValue, Value } from "../values/index.js";
+import { AbstractValue, ObjectValue, SymbolValue, Value } from "../values/index.js";
 import { convertExpressionToJSXIdentifier, convertKeyValueToJSXAttribute } from "../react/jsx.js";
 import { Logger } from "../utils/logger.js";
 import invariant from "../invariant.js";
@@ -283,12 +283,7 @@ export class ResidualReactElementSerializer {
       visitAbstractOrPartialProps: (propsValue: AbstractValue | ObjectValue) => {
         let reactElementSpread = this._createReactElementAttribute();
         this._serializeNowOrAfterWaitingForDependencies(propsValue, reactElement, () => {
-          let expr;
-          if (propsValue.temporalAlias instanceof AbstractObjectValue) {
-            expr = this.residualHeapSerializer.serializeValue(propsValue.temporalAlias);
-          } else {
-            expr = this.residualHeapSerializer.serializeValue(propsValue);
-          }
+          let expr = this.residualHeapSerializer.serializeValue(propsValue);
           reactElementSpread.expr = expr;
           reactElementSpread.type = "SPREAD";
         });

src/serializer/ResidualReactElementVisitor.js

@@ -10,7 +10,7 @@
 /* @flow strict-local */
 
 import { Realm } from "../realm.js";
-import { AbstractObjectValue, AbstractValue, ObjectValue, SymbolValue, Value } from "../values/index.js";
+import { AbstractValue, ObjectValue, SymbolValue, Value } from "../values/index.js";
 import { ResidualHeapVisitor } from "./ResidualHeapVisitor.js";
 import { determineIfReactElementCanBeHoisted } from "../react/hoisting.js";
 import { traverseReactElement } from "../react/elements.js";
@@ -52,11 +52,7 @@ export class ResidualReactElementVisitor {
         this.residualHeapVisitor.visitValue(refValue);
       },
       visitAbstractOrPartialProps: (propsValue: AbstractValue | ObjectValue) => {
-        if (propsValue.temporalAlias instanceof AbstractObjectValue) {
-          this.residualHeapVisitor.visitValue(propsValue.temporalAlias);
-        } else {
-          this.residualHeapVisitor.visitValue(propsValue);
-        }
+        this.residualHeapVisitor.visitValue(propsValue);
       },
       visitConcreteProps: (propsValue: ObjectValue) => {
         for (let [propName, binding] of propsValue.properties) {