Fix tryGetOptimizedFunctionRoot to actually get optimized function root

src/serializer/ResidualHeapSerializer.js

@@ -79,7 +79,7 @@ import { canHoistFunction } from "../react/hoisting.js";
 import { To } from "../singletons.js";
 import { ResidualReactElementSerializer } from "./ResidualReactElementSerializer.js";
 import type { Binding } from "../environment.js";
-import { GlobalEnvironmentRecord, DeclarativeEnvironmentRecord } from "../environment.js";
+import { GlobalEnvironmentRecord, DeclarativeEnvironmentRecord, FunctionEnvironmentRecord } from "../environment.js";
 import type { Referentializer } from "./Referentializer.js";
 import { GeneratorDAG } from "./GeneratorDAG.js";
 import { type Replacement, getReplacement } from "./ResidualFunctionInstantiator.js";
@@ -746,13 +746,34 @@ export class ResidualHeapSerializer {
       }
       let additionalFVEffects = this.additionalFunctionValuesAndEffects;
       if (additionalFVEffects) {
+        // for optimized functions, we should use created objects
         let maybeParentFunctionInfo = additionalFVEffects.get(maybeParentFunction);
         if (maybeParentFunctionInfo && maybeParentFunctionInfo.effects.createdObjects.has(childFunction)) return true;
+      } else {
+        // for other functions, check environment records
+        let env = childFunction.$Environment;
+        while (env.parent !== null) {
+          let envRecord = env.environmentRecord;
+          if (envRecord instanceof FunctionEnvironmentRecord && envRecord.$FunctionObject === maybeParentFunction)
+            return true;
+          env = env.parent;
+        }
       }
     }
     return false;
   }
 
+  // Check if an optimized function defines the given set of functions.
+  definesFunctions(possibleParentFunction: FunctionValue, functions: Set<FunctionValue>): boolean {
+    let additionalFVEffects = this.additionalFunctionValuesAndEffects;
+    invariant(additionalFVEffects);
+    let maybeParentFunctionInfo = additionalFVEffects.get(possibleParentFunction);
+    invariant(maybeParentFunctionInfo);
+    let createdObjects = maybeParentFunctionInfo.effects.createdObjects;
+    for (let func of functions) if (func !== possibleParentFunction && !createdObjects.has(func)) return false;
+    return true;
+  }
+
   // Try and get the root optimized function when passed in an optimized function
   // that may or may not be nested in the tree of said root, or is the root optimized function
   tryGetOptimizedFunctionRoot(val: Value): void | FunctionValue {
@@ -768,28 +789,34 @@ export class ResidualHeapSerializer {
       invariant(s instanceof FunctionValue);
       functionValues.add(s);
     }
-    let additionalFunction;
+    let outermostAdditionalFunctions = new Set();
 
+    // Get the set of optimized functions that may be the root
     for (let functionValue of functionValues) {
       if (this.additionalFunctionGenerators.has(functionValue)) {
-        if (this.isDefinedInsideFunction(functionValue, functionValues)) {
-          continue;
-        }
-        if (additionalFunction !== undefined && additionalFunction !== functionValue) {
-          return undefined;
-        }
-        additionalFunction = functionValue;
+        if (!this.isDefinedInsideFunction(functionValue, functionValues))
+          outermostAdditionalFunctions.add(functionValue);
       } else {
         let f = this.tryGetOptimizedFunctionRoot(functionValue);
-        if (f === undefined) return undefined;
-        if (this.isDefinedInsideFunction(f, functionValues)) {
-          continue;
-        }
-        if (additionalFunction !== undefined && additionalFunction !== f) return undefined;
-        additionalFunction = f;
+        if (f && !this.isDefinedInsideFunction(f, functionValues)) outermostAdditionalFunctions.add(f);
       }
     }
-    return additionalFunction;
+    if (outermostAdditionalFunctions.size === 1) return [...outermostAdditionalFunctions][0];
+
+    let additionalFVEffects = this.additionalFunctionValuesAndEffects;
+    invariant(additionalFVEffects);
+
+    // See if any of the outermost (or any of their parents) are the outermost optimized function
+    let possibleRoots = [...outermostAdditionalFunctions];
+    while (possibleRoots.length > 0) {
+      let possibleRoot = possibleRoots.shift();
+      if (this.definesFunctions(possibleRoot, outermostAdditionalFunctions)) return possibleRoot;
+      let additionalFunctionEffects = additionalFVEffects.get(possibleRoot);
+      invariant(additionalFunctionEffects);
+      let parent = additionalFunctionEffects.parentAdditionalFunction;
+      if (parent) possibleRoots.push(parent);
+    }
+    return undefined;
   }
 
   _getActiveBodyOfGenerator(generator: Generator): void | SerializedBody {

test/serializer/optimized-functions/CommonParentScopeCapture.js

@@ -0,0 +1,22 @@
+(function () {
+  function middle(cond) {
+    if (cond) {
+      var value = 42;
+      function inner1() { return value; }
+      function inner2() { return value; }
+      global.__optimize && __optimize(inner1);
+      global.__optimize && __optimize(inner2);
+      return [inner1, inner2];
+    }
+  }
+  function outer(cond) {
+    return [middle(cond), middle(cond)];
+  }
+  global.outer = outer;
+  global.__optimize && __optimize(outer);
+
+  global.inspect = function() {
+    let [[i1, i2], [i3, i4]] = outer(true);
+    return i1() + i2() + i3() + i4();
+  }
+})();

test/serializer/optimized-functions/Issue2399.js

@@ -0,0 +1,19 @@
+(function () {
+  function outer(cond) {
+    if (cond) {
+      var value = 42;
+      function inner1() { return value; }
+      function inner2() { return value; }
+      global.__optimize && __optimize(inner1);
+      global.__optimize && __optimize(inner2);
+      return [inner1, inner2];
+    }
+  }
+  global.outer = outer;
+  global.__optimize && __optimize(outer);
+
+  global.inspect = function() {
+    let [i1, i2] = outer(true);
+    return i1() + i2();
+  }
+})();

test/serializer/optimized-functions/OptimizedResidualOptimized.js

@@ -0,0 +1,23 @@
+(function () {
+  function middle(cond) {
+    if (cond) {
+      var value = { x: 42 };
+      function inner1() { return value; }
+      function inner2() { return value; }
+      global.__optimize && __optimize(inner1);
+      global.__optimize && __optimize(inner2);
+      return [inner1, inner2];
+    }
+  }
+  function outer(cond) {
+    return [middle(cond), middle(cond)];
+  }
+  global.outer = outer;
+  global.__optimize && __optimize(outer);
+
+  global.inspect = function() {
+    let funcs = [].concat.apply([], outer(true));
+    let objs = funcs.map(f => f());
+    return " " + (objs[0] === objs[1]) + " " + (objs[1] === objs[2]);
+  }
+})();