JS: fix generation of main call and test run

js/js.translator/src/org/jetbrains/kotlin/js/translate/general/Translation.java

@@ -274,6 +274,20 @@ public final class Translation {
             merger.addFragment(staticContext.getFragment());
         }
 
+        JsProgramFragment testFragment = mayBeGenerateTests(files, config, bindingTrace, moduleDescriptor);
+        fragments.add(testFragment);
+        merger.addFragment(testFragment);
+        rootFunction.getParameters().add(new JsParameter(internalModuleName));
+
+        if (mainCallParameters.shouldBeGenerated()) {
+            JsProgramFragment mainCallFragment = generateCallToMain(
+                    bindingTrace, config, moduleDescriptor, files, mainCallParameters.arguments());
+            if (mainCallFragment != null) {
+                fragments.add(mainCallFragment);
+                merger.addFragment(mainCallFragment);
+            }
+        }
+
         merger.merge();
 
         JsBlock rootBlock = rootFunction.getBody();
@@ -285,9 +299,6 @@ public final class Translation {
             defineModule(program, statements, config.getModuleId());
         }
 
-        //mayBeGenerateTests(files, rootBlock, context);
-        rootFunction.getParameters().add(new JsParameter(internalModuleName));
-
         // Invoke function passing modules as arguments
         // This should help minifier tool to recognize references to these modules as local variables and make them shorter.
         List<JsImportedModule> importedModuleList = merger.getImportedModules();
@@ -296,15 +307,6 @@ public final class Translation {
             rootFunction.getParameters().add(new JsParameter(importedModule.getInternalName()));
         }
 
-        /*
-        if (mainCallParameters.shouldBeGenerated()) {
-            JsStatement statement = generateCallToMain(context, files, mainCallParameters.arguments());
-            if (statement != null) {
-                statements.add(statement);
-            }
-        }
-        */
-
         statements.add(new JsReturn(internalModuleName.makeRef()));
 
         JsBlock block = program.getGlobalBlock();
@@ -355,20 +357,27 @@ public final class Translation {
         }
     }
 
-    private static void mayBeGenerateTests(
+    @NotNull
-            @NotNull Collection<KtFile> files, @NotNull JsBlock rootBlock, @NotNull TranslationContext context
+    private static JsProgramFragment mayBeGenerateTests(
+            @NotNull Collection<KtFile> files, @NotNull JsConfig config, @NotNull BindingTrace trace,
+            @NotNull ModuleDescriptor moduleDescriptor
     ) {
-        JSTester tester = new QUnitTester();
+        StaticContext staticContext = new StaticContext(trace, config, moduleDescriptor);
-        tester.initialize(context, rootBlock);
+        TranslationContext context = TranslationContext.rootContext(staticContext);
+        JSTester tester = new QUnitTester(context);
         JSTestGenerator.generateTestCalls(context, files, tester);
-        tester.deinitialize();
+
+        return staticContext.getFragment();
     }
 
     //TODO: determine whether should throw exception
     @Nullable
-    private static JsStatement generateCallToMain(
+    private static JsProgramFragment generateCallToMain(
-            @NotNull TranslationContext context, @NotNull Collection<KtFile> files, @NotNull List<String> arguments
+            @NotNull BindingTrace trace, @NotNull JsConfig config, @NotNull ModuleDescriptor moduleDescriptor,
+            @NotNull Collection<KtFile> files, @NotNull List<String> arguments
     ) {
+        StaticContext staticContext = new StaticContext(trace, config, moduleDescriptor);
+        TranslationContext context = TranslationContext.rootContext(staticContext);
         MainFunctionDetector mainFunctionDetector = new MainFunctionDetector(context.bindingContext());
         KtNamedFunction mainFunction = mainFunctionDetector.getMainFunction(files);
         if (mainFunction == null) {
@@ -376,6 +385,8 @@ public final class Translation {
         }
         FunctionDescriptor functionDescriptor = getFunctionDescriptor(context.bindingContext(), mainFunction);
         JsArrayLiteral argument = new JsArrayLiteral(toStringLiteralList(arguments, context.program()));
-        return CallTranslator.INSTANCE.buildCall(context, functionDescriptor, Collections.singletonList(argument), null).makeStmt();
+        JsExpression call = CallTranslator.INSTANCE.buildCall(context, functionDescriptor, Collections.singletonList(argument), null);
+        context.addTopLevelStatement(call.makeStmt());
+        return staticContext.getFragment();
     }
 }

js/js.translator/src/org/jetbrains/kotlin/js/translate/test/CommonUnitTester.java

@@ -18,15 +18,19 @@ package org.jetbrains.kotlin.js.translate.test;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.js.backend.ast.*;
+import org.jetbrains.kotlin.js.translate.context.TranslationContext;
 
 public abstract class CommonUnitTester extends JSTester {
+    public CommonUnitTester(@NotNull TranslationContext context) {
+        super(context);
+    }
 
     @Override
     public void constructTestMethodInvocation(@NotNull JsExpression functionToTestCall,
             @NotNull JsStringLiteral testName) {
         JsFunction functionToTest = new JsFunction(getContext().scope(), "test function");
         functionToTest.setBody(new JsBlock(functionToTestCall.makeStmt()));
-        getBlock().getStatements().add(new JsInvocation(getTestMethodRef(), testName, functionToTest).makeStmt());
+        getContext().addTopLevelStatement(new JsInvocation(getTestMethodRef(), testName, functionToTest).makeStmt());
     }
 
     @NotNull

js/js.translator/src/org/jetbrains/kotlin/js/translate/test/JSTester.java

@@ -16,47 +16,24 @@
 
 package org.jetbrains.kotlin.js.translate.test;
 
-import org.jetbrains.kotlin.js.backend.ast.JsBlock;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.js.backend.ast.JsExpression;
 import org.jetbrains.kotlin.js.backend.ast.JsStringLiteral;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.js.translate.context.TranslationContext;
 
 public abstract class JSTester {
 
-    @Nullable
+    @NotNull
-    private JsBlock block;
+    private final TranslationContext context;
 
-    @Nullable
+    public JSTester(@NotNull TranslationContext context) {
-    private TranslationContext context;
+        this.context = context;
-
-    public JSTester() {
-        this.block = null;
-        this.context = null;
     }
 
     public abstract void constructTestMethodInvocation(@NotNull JsExpression call, @NotNull JsStringLiteral name);
 
-    @NotNull
-    protected JsBlock getBlock() {
-        assert block != null : "Call initialize before using tester.";
-        return block;
-    }
-
     @NotNull
     protected TranslationContext getContext() {
-        assert context != null : "Call initialize before using tester.";
         return context;
     }
-
-    public void initialize(@NotNull TranslationContext context, @NotNull JsBlock block) {
-        this.block = block;
-        this.context = context;
-    }
-
-    public void deinitialize() {
-        this.block = null;
-        this.context = null;
-    }
 }

js/js.translator/src/org/jetbrains/kotlin/js/translate/test/QUnitTester.java

@@ -19,8 +19,13 @@ package org.jetbrains.kotlin.js.translate.test;
 import org.jetbrains.kotlin.js.backend.ast.JsExpression;
 import org.jetbrains.kotlin.js.backend.ast.JsNameRef;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.js.translate.context.TranslationContext;
 
 public final class QUnitTester extends CommonUnitTester {
+    public QUnitTester(@NotNull TranslationContext context) {
+        super(context);
+    }
+
     @NotNull
     private static final JsNameRef TEST_FUN_REF = new JsNameRef("test", "QUnit");