Internal change

PiperOrigin-RevId: 511820571

Author: lauraharker

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

@@ -502,6 +502,7 @@ public TypedScope processForTesting(Node externsRoot, Node jsRoot) {
     this.process(externsRoot, jsRoot);
 
     compiler.setTypeCheckingHasRun(true);
+    compiler.setTopScope(this.topScope);
 
     return topScope;
   }

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

@@ -116,6 +116,8 @@ public abstract class CompilerTestCase {
   /** Whether to rewrite Closure code before the test is run. */
   private boolean rewriteClosureCode;
 
+  private boolean rewriteModulesAfterTypechecking;
+
   /** Whether to rewrite Closure code before the test is run. */
   private boolean rewriteClosureProvides;
 
@@ -681,6 +683,7 @@ public void setUp() throws Exception {
     this.polymerPass = false;
     this.processCommonJsModules = false;
     this.rewriteClosureCode = false;
+    this.rewriteModulesAfterTypechecking = false;
     this.runTypeCheckAfterProcessing = false;
     this.rewriteEsModulesEnabled = false;
     this.transpileEnabled = false;
@@ -952,6 +955,12 @@ protected final void disableRewriteClosureCode() {
     rewriteClosureCode = false;
   }
 
+  protected final void enableRewriteModulesAfterTypechecking() {
+    checkState(this.setUpRan, "Attempted to configure before running setUp().");
+    enableRewriteClosureCode();
+    this.rewriteModulesAfterTypechecking = true;
+  }
+
   /** Rewrite goog.provides */
   protected final void enableRewriteClosureProvides() {
     checkState(this.setUpRan, "Attempted to configure before running setUp().");
@@ -1447,11 +1456,15 @@ private void testInternal(
           new CheckClosureImports(compiler, compiler.getModuleMetadataMap())
               .process(externsRoot, mainRoot);
           new ClosureRewriteClass(compiler).process(externsRoot, mainRoot);
-          new ClosureRewriteModule(compiler, null, null).process(externsRoot, mainRoot);
           ScopedAliases.builder(compiler).build().process(externsRoot, mainRoot);
           hasCodeChanged = hasCodeChanged || recentChange.hasCodeChanged();
         }
 
+        if (rewriteClosureCode && !rewriteModulesAfterTypechecking && i == 0) {
+          new ClosureRewriteModule(compiler, null, null).process(externsRoot, mainRoot);
+          hasCodeChanged = hasCodeChanged || recentChange.hasCodeChanged();
+        }
+
         // Only run process closure primitives once, if asked.
         if (closurePassEnabled && i == 0) {
           recentChange.reset();
@@ -1495,6 +1508,12 @@ private void testInternal(
           check.processForTesting(externsRoot, mainRoot);
         }
 
+        if (rewriteClosureCode && rewriteModulesAfterTypechecking && i == 0) {
+          new ClosureRewriteModule(compiler, null, compiler.getTopScope())
+              .process(externsRoot, mainRoot);
+          hasCodeChanged = hasCodeChanged || recentChange.hasCodeChanged();
+        }
+
         if (inferConsts && i == 0) {
           new InferConsts(compiler).process(externsRoot, mainRoot);
         }