Refactor AbstractPeepholeOptimization::addFeatureToEnclosingScript to not use n.getSourceFile()

This fixes a @closureUnaware crash (as seen in TranspileAndOptimizeClosureUnawareCodeTest).

AbstractPeepholeOptimization::addFeatureToEnclosingScript assumed that "node.getSourceFile()" matched the source file of the containing SCRIPT node. This breaks down if optimizations move code between scripts.

It now also crashed for @closureUnaware code, because we create synthetic file names for each @closureUnaware block within a file, which confuses compiler.getScriptNode.

PiperOrigin-RevId: 858630083

Author: lauraharker

src/com/google/javascript/jscomp/AbstractPeepholeOptimization.java

@@ -17,14 +17,15 @@
 package com.google.javascript.jscomp;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
 
-import com.google.common.collect.LinkedHashMultimap;
+import com.google.common.collect.ImmutableSet;
 import com.google.javascript.jscomp.base.Tri;
 import com.google.javascript.jscomp.parsing.parser.FeatureSet;
 import com.google.javascript.jscomp.parsing.parser.FeatureSet.Feature;
 import com.google.javascript.rhino.Node;
-import com.google.javascript.rhino.StaticSourceFile;
 import java.math.BigInteger;
+import java.util.LinkedHashSet;
 
 /**
  * An abstract class whose implementations run peephole optimizations:
@@ -38,9 +39,14 @@ abstract class AbstractPeepholeOptimization {
   /** Intentionally not exposed to subclasses */
   private AstAnalyzer astAnalyzer;
 
-  /** Features added to scripts during the optimization. */
-  private final LinkedHashMultimap<StaticSourceFile, Feature> newFeatures =
-      LinkedHashMultimap.create();
+  /**
+   * New parser features added in some {@link #optimizeSubtree} call.
+   *
+   * <p>{@link PeepholeOptimizationsPass} uses this to call {@link NodeUtil#addFeaturesToScript} in
+   * closer to O(1) time, as this class doesn't know what script node it's currently in, and would
+   * need to walk the AST.
+   */
+  private final LinkedHashSet<Feature> newFeatures = new LinkedHashSet<>();
 
   /**
    * Given a node to optimize and a traversal, optimize the node. Subclasses
@@ -100,11 +106,11 @@ void beginTraversal(AbstractCompiler compiler) {
    * called again.
    */
   void endTraversal() {
-    for (StaticSourceFile file : newFeatures.keySet()) {
-      Node script = compiler.getScriptNode(file.getName());
-      NodeUtil.addFeaturesToScript(
-          script, FeatureSet.BARE_MINIMUM.with(newFeatures.get(file)), this.compiler);
-    }
+    checkState(
+        this.newFeatures.isEmpty(),
+        "Expected getNewFeatures() to be empty in endTraversal() but found %s for %s",
+        this.newFeatures,
+        this.getClass().getName());
     this.compiler = null;
     astAnalyzer = null;
   }
@@ -244,7 +250,24 @@ protected final void markNewScopesChanged(Node n) {
     NodeUtil.markNewScopesChanged(n, compiler);
   }
 
-  protected final void addFeatureToEnclosingScript(Node n, Feature feature) {
-    newFeatures.put(n.getStaticSourceFile(), feature);
+  /** Calls {@link NodeUtil#addFeatureToScript} with the script currently being visited. */
+  protected final void addFeatureToEnclosingScript(Feature feature) {
+    // NOTE: we could implement this as
+    //  protected final void addFeatureToEnclosingScript(Node node, Feature feature) {
+    //     NodeUtil.addFeatureToScript(NodeUtil.getEnclosingScript(node), ...
+    // This is expected to be behaviorally equivalent.
+    // However, walking the AST to get the enclosing script is O(depth of the AST) per call. With
+    // the  current implementation, we let PeepholeOptimizationsPass to find the script in what's
+    // usually O(1) time, or worst case it only calls getEnclosingScript() once per NodeTraversal
+    // and then caches the result.
+    newFeatures.add(feature);
+  }
+
+  final ImmutableSet<Feature> getNewFeatures() {
+    return ImmutableSet.copyOf(newFeatures);
+  }
+
+  final void clearNewFeatures() {
+    newFeatures.clear();
   }
 }

src/com/google/javascript/jscomp/PeepholeOptimizationsPass.java

@@ -16,8 +16,9 @@
 
 package com.google.javascript.jscomp;
 
+
 import com.google.common.annotations.VisibleForTesting;
-import com.google.javascript.jscomp.NodeTraversal.AbstractPostOrderCallback;
+import com.google.javascript.jscomp.parsing.parser.FeatureSet;
 import com.google.javascript.rhino.Node;
 import java.util.Arrays;
 import java.util.List;
@@ -83,7 +84,7 @@ public void process(Node externs, Node root) {
     endTraversal();
   }
 
-  private class PeepCallback extends AbstractPostOrderCallback {
+  private class PeepCallback extends NodeTraversal.AbstractScopedCallback {
     @Override
     public void visit(NodeTraversal t, Node n, Node parent) {
       Node currentNode = n;
@@ -94,6 +95,29 @@ public void visit(NodeTraversal t, Node n, Node parent) {
         }
       }
     }
+
+    @Override
+    public void enterScope(NodeTraversal t) {}
+
+    @Override
+    public void exitScope(NodeTraversal t) {
+      // Call updateFeatures() here, instead of immediately after
+      // `optim.optimizeSubtree(currentNode)` in visit, to avoid repeating work for
+      // every AST node x every peephole optimization. In practice, profiling shows a small
+      // but measurable improvement for large projects. See comments on cl/856752797.
+      updateFeatures(t);
+    }
+
+    private void updateFeatures(NodeTraversal t) {
+      for (AbstractPeepholeOptimization optim : peepholeOptimizations) {
+        var newFeatures = optim.getNewFeatures();
+        if (!newFeatures.isEmpty()) {
+          NodeUtil.addFeaturesToScript(
+              t.getCurrentScript(), FeatureSet.BARE_MINIMUM.with(newFeatures), compiler);
+        }
+        optim.clearNewFeatures();
+      }
+    }
   }
 
   /** Make sure that all the optimizations have the current compiler so they can report errors. */

src/com/google/javascript/jscomp/PeepholeRemoveDeadCode.java

@@ -1007,7 +1007,7 @@ private void redeclareIfBlockScopedVar(Node decl, Node parent) {
         parent.addChildToBack(var);
       }
     }
-    this.addFeatureToEnclosingScript(parent, Feature.LET_DECLARATIONS);
+    this.addFeatureToEnclosingScript(Feature.LET_DECLARATIONS);
   }
 
   /**

test/com/google/javascript/jscomp/PeepholeOptimizationsPassTest.java

@@ -250,23 +250,23 @@ public void testAddFeatureToEnclosingScript() {
               @Override
               public Node optimizeSubtree(Node node) {
                 if (node.isAdd()) {
-                  this.addFeatureToEnclosingScript(node, Feature.LET_DECLARATIONS);
-                  this.addFeatureToEnclosingScript(node, Feature.LET_DECLARATIONS);
-                  this.addFeatureToEnclosingScript(node, Feature.CLASSES);
+                  this.addFeatureToEnclosingScript(Feature.LET_DECLARATIONS);
+                  this.addFeatureToEnclosingScript(Feature.LET_DECLARATIONS);
+                  this.addFeatureToEnclosingScript(Feature.CLASSES);
                 }
                 return node;
               }
             },
             new AbstractPeepholeOptimization() {
               @Override
               public Node optimizeSubtree(Node node) {
-                if (node.isAdd()) {
-                  this.addFeatureToEnclosingScript(node, Feature.CONST_DECLARATIONS);
+                if (node.isSub()) {
+                  this.addFeatureToEnclosingScript(Feature.CONST_DECLARATIONS);
                 }
                 return node;
               }
             });
-    testSame("(3 + 4);");
+    testSame("(3 + 4); function sub() { return 3 - 4; }");
     Compiler compiler = getLastCompiler();
     Node script = checkNotNull(compiler.getScriptNode("testcode"));
     assertThat(script.getProp(Node.FEATURE_SET))

test/com/google/javascript/jscomp/TranspileAndOptimizeClosureUnawareTest.java

@@ -125,24 +125,26 @@ public void testFoldConstants_multipleClosureUnawareBlocksInFile() {
   public void testFoldConstants_regressionTest_crashInNodeUtilAddFeatureToScript() {
     setLanguageOut(CompilerOptions.LanguageMode.ECMASCRIPT_NEXT);
 
-    // TODO: b/421970620 - fix this crash.
-    assertThrows(
-        RuntimeException.class,
-        () ->
-            testSame(
-                closureUnaware(
-                    """
-                    return function f() {
-                      function referenceX() { return x; }
+    test(
+        closureUnaware(
+            """
+            return function f() {
+              function referenceX() { return x; }
 
-                      return 1;
+              return 1;
 
-                      // Regression test for a crash in PeepholeRemoveDeadCode.
-                      let x = 0;
-                      referenceX();
-                      x++;
-                    }
-                    """)));
+              // Regression test for a crash in PeepholeRemoveDeadCode.
+              let x = 0;
+              referenceX();
+              x++;
+            }
+            """),
+        expectedClosureUnaware(
+            """
+            return function() {
+              return 1;
+            }
+            """));
   }
 
   @Test