Hoist function assignment to the top of the inner-most *block* scope when rewriting a function declaration as a function expression for eager browser JS compilation parenthesization.

PiperOrigin-RevId: 499470585

Author: Closure Team

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

@@ -1600,7 +1600,7 @@ private static boolean isFirstChild(Node n) {
 
   private void addArrowFunction(Node n, Node first, Node last, Context context) {
     checkState(first.getString().isEmpty(), first);
-    boolean funcNeedsParens = arrowFunctionNeedsParens(n);
+    boolean funcNeedsParens = arrowFunctionNeedsParens(n) || n.getMarkForParenthesize();
     if (funcNeedsParens) {
       add("(");
     }
@@ -1630,8 +1630,7 @@ private void addArrowFunction(Node n, Node first, Node last, Context context) {
   }
 
   private void addFunction(Node n, Node first, Node last, Context context) {
-    boolean funcNeedsParens =
-        (context == Context.START_OF_EXPR || n.getBooleanProp(Node.MARK_FOR_PARENTHESIZE));
+    boolean funcNeedsParens = (context == Context.START_OF_EXPR || n.getMarkForParenthesize());
     if (funcNeedsParens) {
       add("(");
     }

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

@@ -16,12 +16,22 @@
 
 package com.google.javascript.jscomp;
 
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ListMultimap;
 import com.google.javascript.jscomp.diagnostic.LogFile;
 import com.google.javascript.rhino.IR;
 import com.google.javascript.rhino.Node;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Deque;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.jspecify.nullness.Nullable;
 
 /**
  * Marks all functions of specified chunks for eager parsing by adding the node property
@@ -38,12 +48,14 @@
  * for eager compile for each specified chunk.
  */
 public final class ParenthesizeFunctionsInChunks implements CompilerPass {
-  final AbstractCompiler compiler;
-  final Set<String> parenthesizeFunctionsInChunks;
-
-  // Map for recording diagnostic information about what was marked for eager compilation.
-  final Map<String, Long> chunkToEagerCompileFnCount = new HashMap<>();
+  private final AbstractCompiler compiler;
+  private final Set<String> parenthesizeFunctionsInChunks;
 
+  /**
+   * @param compiler An abstract compiler.
+   * @param parenthesizeFunctionsInChunks The set of chunk names in which to parenthesize top level
+   *     functions.
+   */
   public ParenthesizeFunctionsInChunks(
       AbstractCompiler compiler, Set<String> parenthesizeFunctionsInChunks) {
     this.compiler = compiler;
@@ -52,52 +64,155 @@ public ParenthesizeFunctionsInChunks(
 
   @Override
   public void process(Node externs, Node root) {
-    NodeTraversal.traverse(compiler, root, new Traversal());
-    for (String key : chunkToEagerCompileFnCount.keySet()) {
-      try (LogFile log = compiler.createOrReopenLog(getClass(), "eager_compile_chunks.log")) {
-        log.log("%s: %d fn's marked for eager compile", key, chunkToEagerCompileFnCount.get(key));
+    Traversal traversal = new Traversal(parenthesizeFunctionsInChunks);
+    NodeTraversal.traverse(compiler, root, traversal);
+
+    Map<String, Long> chunkToEagerCompileFnCounts = traversal.getChunkToEagerCompileFnCounts();
+
+    try (LogFile log = compiler.createOrReopenLog(getClass(), "eager_compile_chunks.log")) {
+      for (Map.Entry<String, Long> entry : chunkToEagerCompileFnCounts.entrySet()) {
+        log.log("%s: %d fn's marked for eager compile", entry.getKey(), entry.getValue());
       }
     }
   }
 
-  private class Traversal implements NodeTraversal.Callback {
+  private static class Traversal implements NodeTraversal.Callback {
+    private final Set<String> parenthesizeFunctionsInChunks;
+
+    // The stack of nested block scopes for the node we're currently visiting.
+    private final Deque<Node> nestedBlockScopes = new ArrayDeque<>();
+
+    // A multimap relating a scope node to any children nodes which should be hosted into its scope.
+    private final ListMultimap<Node, Node> hoistNodesToScope = ArrayListMultimap.create();
+
+    // Map for recording diagnostic information about what was marked for eager compilation.
+    private final Map<String, Long> chunkToEagerCompileFnCounts = new HashMap<>();
+
+    public Traversal(Set<String> parenthesizeFunctionsInChunks) {
+      this.parenthesizeFunctionsInChunks = parenthesizeFunctionsInChunks;
+    }
+
+    @Override
+    public boolean shouldTraverse(NodeTraversal t, Node n, Node parent) {
+      if (!shouldParenthesizeTree(t, n)) {
+        return false;
+      }
+
+      if (parent != null && parent.isFunction()) {
+        return false; // Don't visit the contents of any functions.
+      }
+
+      if (NodeUtil.isStatementBlock(n)) {
+        nestedBlockScopes.push(n); // Enter block scope.
+      }
+
+      return true;
+    }
+
+    @Override
+    public void visit(NodeTraversal t, Node n, Node parent) {
+      if (isInnerMostBlockScope(n)) {
+        hoistChildrenToTopOfScope(n);
+        nestedBlockScopes.pop(); // Exit block scope.
+      } else if (NodeUtil.isFunctionExpression(n)) {
+        n.setMarkForParenthesize(true);
+        incrementEagerlyCompiledFunctionCount(t);
+      } else if (NodeUtil.isFunctionDeclaration(n)) {
+        n.setMarkForParenthesize(true);
+        addChildForHoistToScope(functionDeclarationToFunctionExpression(t, n));
+        incrementEagerlyCompiledFunctionCount(t);
+      } else {
+        // Ex: Method function definitions.
+        // Do nothing.
+      }
+    }
+
+    public Map<String, Long> getChunkToEagerCompileFnCounts() {
+      checkState(
+          nestedBlockScopes.isEmpty(), "Expected empty scope stack. Got: %s", nestedBlockScopes);
+      checkState(
+          hoistNodesToScope.isEmpty(), "Expected empty hoist map. Got: %s", hoistNodesToScope);
+      return chunkToEagerCompileFnCounts;
+    }
+
+    /**
+     * Whether we should parenthesize functions in the given tree of the traversal. If there is no
+     * chunk, then we are compiling a single-chunk output, so we will just parenthesize all
+     * top-level functions.
+     */
+    private boolean shouldParenthesizeTree(NodeTraversal t, Node n) {
+      if (!n.isScript()) {
+        return true;
+      }
 
-    Traversal() {}
+      String chunkName = getChunkName(t);
+      return chunkName == null || parenthesizeFunctionsInChunks.contains(chunkName);
+    }
+
+    private void incrementEagerlyCompiledFunctionCount(NodeTraversal t) {
+      chunkToEagerCompileFnCounts.merge(getChunkName(t), 1L, (oldValue, value) -> oldValue + 1);
+    }
 
-    private void rewriteAsFunctionExpression(Node n) {
+    /** Gets the chunk name of the current traveral or null if it doesn't belong to a chunk. */
+    private @Nullable String getChunkName(NodeTraversal t) {
+      JSChunk chunk = t.getChunk();
+      return (chunk != null ? chunk.getName() : null);
+    }
+
+    /** Whether this node is the inner-most block scope. */
+    private boolean isInnerMostBlockScope(Node n) {
+      return !nestedBlockScopes.isEmpty() && nestedBlockScopes.peek() == n;
+    }
+
+    /**
+     * Converts the given function declaration into a function expression suitable for wrapping in
+     * parenthesis. The function expression is assigned to a `var` declaration so that it is
+     * semantically hoisted to the inner-most function scope like with function declarations.
+     */
+    private Node functionDeclarationToFunctionExpression(NodeTraversal t, Node n) {
+      AbstractCompiler compiler = t.getCompiler();
       Node nameNode = n.getFirstChild();
       Node name = IR.name(nameNode.getString()).srcref(nameNode);
       Node var = IR.var(name).srcref(n);
-      // Add the VAR, remove the FUNCTION
-      n.replaceWith(var);
-      // readd the function, but remove the function name, to guard against
+      // read the function, but remove the function name, to guard against
       // functions that re-assign to themselves, which will end up causing a
       // recursive loop.
-      n.getFirstChild().setString("");
+      nameNode.setString("");
+      compiler.reportChangeToEnclosingScope(n.getLastChild());
+      // Add the VAR, remove the FUNCTION
+      n.replaceWith(var);
+      compiler.reportChangeToEnclosingScope(var);
       name.addChildToFront(n);
+      return var;
     }
 
-    @Override
-    public boolean shouldTraverse(NodeTraversal t, Node n, Node parent) {
-      if (n.isScript() || n.isRoot() || n.isModuleBody()) {
-        return true;
-      }
-      if (n.isFunction()) {
-        String chunkName = t.getChunk().getName();
-        chunkToEagerCompileFnCount.putIfAbsent(chunkName, 0L);
-        if (parenthesizeFunctionsInChunks.contains(chunkName)) {
-          n.putBooleanProp(Node.MARK_FOR_PARENTHESIZE, true);
-          if (!NodeUtil.isFunctionExpression(n)) {
-            rewriteAsFunctionExpression(n);
-          }
-          chunkToEagerCompileFnCount.put(chunkName, chunkToEagerCompileFnCount.get(chunkName) + 1);
-        }
-        return false;
+    /**
+     * Marks a node to be moved to the top of its parent *block* scope. This is useful for hoisting
+     * function expression assignments to emulate how function declarations are assigned.
+     *
+     * <pre>
+     * Function declarations have the following semantics:
+     * 1. Hoists the function variable declaration to the top of the inner-most *function* scope.
+     * 2. Hoists the function assignment to the top of the inner-most *block* scope.
+     * </pre>
+     */
+    private void addChildForHoistToScope(Node node) {
+      if (nestedBlockScopes.isEmpty()) { // This node is already at the root scope.
+        return;
       }
-      return true;
+
+      Node innerMostBlockScope = nestedBlockScopes.peek();
+      hoistNodesToScope.put(innerMostBlockScope, node.detach());
     }
 
-    @Override
-    public void visit(NodeTraversal t, Node n, Node parent) {}
+    /** Hoists any marked nodes to the beginning of this scope. */
+    private void hoistChildrenToTopOfScope(Node scope) {
+      List<Node> nodes = new ArrayList<>(hoistNodesToScope.removeAll(scope));
+      Collections.reverse(nodes); // Maintain node order when pushing as first sibling.
+
+      for (Node node : nodes) {
+        scope.addChildToFront(node);
+      }
+    }
   }
 }

src/com/google/javascript/rhino/Node.java

@@ -253,17 +253,28 @@ public final Node setTrailingNonJSDocComment(NonJSDocComment comment) {
     return this;
   }
 
-  /** Sets whether this node is inside parentheses. */
+  /** Sets whether this node was inside original source-level parentheses. */
   public final void setIsParenthesized(boolean b) {
     checkState(IR.mayBeExpression(this));
     putBooleanProp(Prop.IS_PARENTHESIZED, b);
   }
 
-  /** Check whether node was inside parentheses. */
+  /** Check whether node was inside original source-level parentheses. */
   public final boolean getIsParenthesized() {
     return getBooleanProp(Prop.IS_PARENTHESIZED);
   }
 
+  /** Sets whether this node is should be parenthesized in output. */
+  public final void setMarkForParenthesize(boolean value) {
+    checkState(IR.mayBeExpression(this));
+    putBooleanProp(Prop.MARK_FOR_PARENTHESIZE, value);
+  }
+
+  /** Check whether node should be parenthesized in output. */
+  public final boolean getMarkForParenthesize() {
+    return getBooleanProp(Prop.MARK_FOR_PARENTHESIZE);
+  }
+
   // TODO(sdh): Get rid of these by using accessor methods instead.
   // These export instances of a private type, which is awkward but a side effect is that it
   // prevents anyone from introducing problemmatic uses of the general-purpose accessors.
@@ -297,7 +308,6 @@ public final boolean getIsParenthesized() {
   public static final Prop MODULE_EXPORT = Prop.MODULE_EXPORT;
   public static final Prop IS_SHORTHAND_PROPERTY = Prop.IS_SHORTHAND_PROPERTY;
   public static final Prop ES6_MODULE = Prop.ES6_MODULE;
-  public static final Prop MARK_FOR_PARENTHESIZE = Prop.MARK_FOR_PARENTHESIZE;
 
   private static final class NumberNode extends Node {
 
@@ -1002,6 +1012,7 @@ public void validateProperties(Consumer<String> violationMessageConsumer) {
       // this method was created to cover the checks previously done there.
       switch (prop) {
         case IS_PARENTHESIZED:
+        case MARK_FOR_PARENTHESIZE:
           if (!IR.mayBeExpression(this)) {
             violationMessageConsumer.accept("non-expression is parenthesized");
           }

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

@@ -93,7 +93,7 @@ public void setUp() throws Exception {
   }
 
   @Test
-  public void testParenthesizedProperty() {
+  public void testIsParenthesizedProperty() {
     // Since we're building the AST by hand, there won't be any types on it.
     typeInfoValidationMode = TypeInfoValidation.NONE;
 
@@ -112,6 +112,25 @@ public void testParenthesizedProperty() {
     expectInvalid(objNode, Check.EXPRESSION, "non-expression is parenthesized");
   }
 
+  @Test
+  public void testMarkForParenthesizeProperty() {
+    // Since we're building the AST by hand, there won't be any types on it.
+    typeInfoValidationMode = TypeInfoValidation.NONE;
+
+    Node n = IR.string("a");
+    n.setMarkForParenthesize(true);
+    setTestSourceLocationForTree(n);
+    expectValid(n, Check.EXPRESSION);
+
+    n.setToken(Token.STRING_KEY); // A string key cannot be parenthesized
+    // We have to put the STRING_KEY into an object and give it a child, so we have an
+    // expression to validate that is valid other than the bad parenthesized property.
+    Node objNode = IR.objectlit(n);
+    n.addChildToFront(IR.number(0));
+    objNode.srcrefTree(n);
+    expectInvalid(objNode, Check.EXPRESSION, "non-expression is parenthesized");
+  }
+
   @Test
   public void testValidGetPropAndGetElem() {
     valid(