Consolidate all of ReplaceToggles into a single optimization pass

PiperOrigin-RevId: 560882626

Author: shicks

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

@@ -23,7 +23,6 @@
 import com.google.common.base.Splitter;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.MustBeClosed;
 import com.google.errorprone.annotations.OverridingMethodsMustInvokeSuper;
@@ -713,15 +712,4 @@ public final LogFile createOrReopenIndexedLog(
    * rewriting has not occurred.
    */
   abstract void mergeSyntheticCodeInput();
-
-  /**
-   * Records the mapping of toggle names to ordinals, which is read from a bootstrap JS file by the
-   * first (check) pass of {@link RewriteToggles}.
-   */
-  void setToggleOrdinalMapping(@Nullable ImmutableMap<String, Integer> mapping) {}
-
-  /** Returns the recorded toggle-name-to-ordinal mapping. */
-  @Nullable ImmutableMap<String, Integer> getToggleOrdinalMapping() {
-    return null;
-  }
 }

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

@@ -290,8 +290,6 @@ public SourceFile loadSource(String filename) {
 
   private final Timeline<Node> changeTimeline = new Timeline<>();
 
-  private @Nullable ImmutableMap<String, Integer> toggleOrdinalMapping = null;
-
   /**
    * When mapping symbols from a source map, we must repeatedly combine the path of the original
    * file with the path from the source map to compute the SourceFile of the underlying code. When
@@ -4327,14 +4325,4 @@ private static Node checkNotModule(Node script, String msg, Object... args) {
     checkState(!script.getFirstChild().isModuleBody(), msg, args);
     return script;
   }
-
-  @Override
-  void setToggleOrdinalMapping(@Nullable ImmutableMap<String, Integer> mapping) {
-    this.toggleOrdinalMapping = mapping;
-  }
-
-  @Override
-  @Nullable ImmutableMap<String, Integer> getToggleOrdinalMapping() {
-    return toggleOrdinalMapping;
-  }
 }

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

@@ -351,7 +351,6 @@ protected PassListBuilder getChecks() {
 
     if (options.closurePass) {
       checks.maybeAdd(closurePrimitives);
-      checks.maybeAdd(replaceTogglesCheck);
     }
 
     // It's important that the PolymerPass run *after* the ClosurePrimitives and ChromePass rewrites
@@ -2053,23 +2052,12 @@ public void process(Node externs, Node jsRoot) {
           .setInternalFactory(ConstParamCheck::new)
           .build();
 
-  private final PassFactory replaceTogglesCheck = createReplaceToggles(true);
-  private final PassFactory replaceToggles = createReplaceToggles(false);
-
   /** Replaces goog.toggle calls with toggle lookups. */
-  private final PassFactory createReplaceToggles(boolean check) {
-    return PassFactory.builder()
-        .setName("replaceToggles" + (check ? "Check" : ""))
-        .setInternalFactory(
-            (compiler) ->
-                new CompilerPass() {
-                  @Override
-                  public void process(Node externs, Node root) {
-                    new ReplaceToggles(compiler, check).process(externs, root);
-                  }
-                })
-        .build();
-  }
+  private final PassFactory replaceToggles =
+      PassFactory.builder()
+          .setName("replaceToggles")
+          .setInternalFactory(ReplaceToggles::new)
+          .build();
 
   /** Generates unique ids. */
   private final PassFactory replaceIdGenerators =

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

@@ -16,7 +16,6 @@
 
 package com.google.javascript.jscomp;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.javascript.jscomp.NodeTraversal.AbstractPostOrderCallback;
 import com.google.javascript.rhino.Node;
@@ -26,13 +25,21 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
+import org.jspecify.nullness.Nullable;
 
 /**
  * Replaces calls to id generators with ids.
  *
  * <p>Use this to get unique and short ids.
  */
 class ReplaceToggles implements CompilerPass {
+
+  // NOTE: These diagnostics are only checked in Stage 2 (optimization), since none of this
+  // code should ever be hand-written: the `CLOSURE_TOGGLE_ORDINALS` bootstrap and the calls to
+  // `goog.readToggleInternalDoNotCallDirectly` are all generated automatically by the build
+  // system, so it's unexpected that anyone should ever run into these diagnostics when doing
+  // ordinary development.
+
   static final DiagnosticType INVALID_TOGGLE_PARAMETER =
       DiagnosticType.error(
           "JSC_INVALID_TOGGLE_PARAMETER",
@@ -53,18 +60,17 @@ class ReplaceToggles implements CompilerPass {
   // NOTE: These values are chosen as negative integers because actual toggle ordinals must always
   // be non-negative (at least zero).  Any negative integers would do to distinguish them from real
   // toggle ordinals, but -1 and -2 are the simplest.
-  @VisibleForTesting static final int TRUE_VALUE = -2;
-
-  @VisibleForTesting static final int FALSE_VALUE = -1;
+  private static final int TRUE_VALUE = -2;
+  private static final int FALSE_VALUE = -1;
 
   private final AbstractCompiler compiler;
   private final AstFactory astFactory;
-  private final boolean check;
 
-  ReplaceToggles(AbstractCompiler compiler, boolean check) {
+  private @Nullable ImmutableMap<String, Integer> ordinalMapping = null;
+
+  ReplaceToggles(AbstractCompiler compiler) {
     this.compiler = compiler;
     this.astFactory = compiler.createAstFactory();
-    this.check = check;
   }
 
   @Override
@@ -80,20 +86,21 @@ private class Traversal extends AbstractPostOrderCallback {
 
     @Override
     public void visit(NodeTraversal t, Node n, Node parent) {
-      // Look for `var CLOSURE_TOGGLE_ORDINALS = {...};`.  Note that we only do this in check mode.
-      // It's not our responsibility to delete the unused variable in optimized mode - if all
-      // the calls to readToggle are deleted, then the bootstrap will be unused and deleted, too.
-      if (check
-          && NodeUtil.isNameDeclaration(n)
-          && n.getFirstChild().matchesName(ORDINAL_VAR_NAME)) {
+      // Look for `var CLOSURE_TOGGLE_ORDINALS = {...};`, record the mapping, and then delete the
+      // declaration from the AST since we should no longer need it (and the optimizer won't
+      // delete it in case the global assignment was an intended side effect).
+      if (NodeUtil.isNameDeclaration(n) && n.getFirstChild().matchesName(ORDINAL_VAR_NAME)) {
         Node rhs = n.getFirstFirstChild();
 
         if (rhs == null && n.getToken() == Token.VAR) {
-          // An empty var is a type definition, which is OK.
+          // An empty var is fine; it should get deleted later.
           return;
         } else if (!rhs.isObjectLit()) {
           compiler.report(JSError.make(n, INVALID_ORDINAL_MAPPING, "not an object literal"));
           return;
+        } else if (ordinalMapping != null) {
+          compiler.report(JSError.make(n, INVALID_ORDINAL_MAPPING, "multiple initialized copies"));
+          return;
         }
 
         Map<String, Integer> mapping = new LinkedHashMap<>();
@@ -126,10 +133,10 @@ public void visit(NodeTraversal t, Node n, Node parent) {
           mapping.put(key, intValue);
           ordinals.add(intValue);
         }
-        compiler.setToggleOrdinalMapping(ImmutableMap.copyOf(mapping));
+        ReplaceToggles.this.ordinalMapping = ImmutableMap.copyOf(mapping);
 
         // NOTE: We do not support a simple assignment without `var` since reassignment or later
-        // augmentation is not allowed.
+        // augmentation (i.e. `CLOSURE_TOGGLE_ORDINALS['foo'] = true`) is not allowed.
         return;
       }
 
@@ -147,15 +154,11 @@ public void visit(NodeTraversal t, Node n, Node parent) {
         return;
       }
 
-      ImmutableMap<String, Integer> toggles = compiler.getToggleOrdinalMapping();
-      if (check) {
-        if (toggles != null && !toggles.containsKey(arg.getString())) {
-          compiler.report(JSError.make(n, UNKNOWN_TOGGLE));
-        }
-        return;
+      if (ordinalMapping != null && !ordinalMapping.containsKey(arg.getString())) {
+        compiler.report(JSError.make(n, UNKNOWN_TOGGLE));
       }
 
-      Integer ordinal = toggles != null ? toggles.get(arg.getString()) : null;
+      Integer ordinal = ordinalMapping != null ? ordinalMapping.get(arg.getString()) : null;
       if (ordinal == null || ordinal < 0) {
         // No ordinals given: hard-code `true` if explicitly set as true, or `false` otherwise.
         n.replaceWith(