KT-31295: Enable support for .ws.kts files as for scratch files

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt

@@ -807,9 +807,13 @@ fun main(args: Array<String>) {
         }
 
         testClass<AbstractScratchRunActionTest> {
-            model("scratch", extension = "kts", testMethod = "doCompilingTest", testClassName = "Compiling", recursive = false)
-            model("scratch", extension = "kts", testMethod = "doReplTest", testClassName = "Repl", recursive = false)
-            model("scratch/multiFile", extension = null, testMethod = "doMultiFileTest", recursive = false)
+            model("scratch", extension = "kts", testMethod = "doScratchCompilingTest", testClassName = "ScratchCompiling", recursive = false)
+            model("scratch", extension = "kts", testMethod = "doScratchReplTest", testClassName = "ScratchRepl", recursive = false)
+            model("scratch/multiFile", extension = null, testMethod = "doScratchMultiFileTest", testClassName = "ScratchMultiFile", recursive = false)
+
+            model("worksheet", extension = "ws.kts", testMethod = "doWorksheetCompilingTest", testClassName = "WorksheetCompiling", recursive = false)
+            model("worksheet", extension = "ws.kts", testMethod = "doWorksheetReplTest", testClassName = "WorksheetRepl", recursive = false)
+            model("worksheet/multiFile", extension = null, testMethod = "doWorksheetMultiFileTest", testClassName = "WorksheetMultiFile", recursive = false)
         }
 
         testClass<AbstractScratchLineMarkersTest> {

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/actions/ScratchRunLineMarkerContributor.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.idea.core.util.getLineCount
 import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.scratch.ScratchExpression
 import org.jetbrains.kotlin.idea.scratch.getEditorWithScratchPanel
+import org.jetbrains.kotlin.idea.scratch.isKotlinWorksheet
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
@@ -27,7 +28,7 @@ class ScratchRunLineMarkerContributor : RunLineMarkerContributor() {
         val ktFile = element.containingFile as? KtFile
         if (ktFile?.isScript() != true) return null
         val file = ktFile.virtualFile
-        if (ScratchFileService.getInstance().getRootType(file) !is ScratchRootType) return null
+        if (!(ScratchFileService.getInstance().getRootType(file) is ScratchRootType || file.isKotlinWorksheet)) return null
 
         val declaration = element.getStrictParentOfType<KtNamedDeclaration>()
         if (declaration != null && declaration !is KtParameter && declaration.nameIdentifier == element) {

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/scratchUtils.kt

@@ -13,6 +13,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.idea.actions.KOTLIN_WORKSHEET_EXTENSION
 import org.jetbrains.kotlin.idea.scratch.ui.ScratchPanelListener
 import org.jetbrains.kotlin.idea.scratch.ui.ScratchTopPanel
 import org.jetbrains.kotlin.idea.syncPublisherWithDisposeCheck
@@ -23,6 +24,8 @@ internal fun Logger.printDebugMessage(str: String) {
     if (isDebugEnabled) debug("SCRATCH: $str")
 }
 
+val VirtualFile.isKotlinWorksheet: Boolean get() = name.endsWith(".$KOTLIN_WORKSHEET_EXTENSION")
+
 fun getEditorWithoutScratchPanel(fileManager: FileEditorManager, virtualFile: VirtualFile): TextEditor? {
     val editor = fileManager.getSelectedEditor(virtualFile) as? TextEditor
     if (editor?.scratchTopPanel != null) return null

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/ScratchFileHook.kt

@@ -52,7 +52,7 @@ class ScratchFileHook(val project: Project) : ProjectComponent {
 
     private fun isPluggable(file: VirtualFile): Boolean {
         if (!file.isValid) return false
-        if (ScratchFileService.getInstance().getRootType(file) !is ScratchRootType) return false
+        if (!(ScratchFileService.getInstance().getRootType(file) is ScratchRootType || file.isKotlinWorksheet)) return false
         val psiFile = PsiManager.getInstance(project).findFile(file) ?: return false
         return ScratchFileLanguageProvider.get(psiFile.fileType) != null
     }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/ScratchTopPanel.kt

@@ -107,7 +107,7 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
         add(JLabel("Use classpath of module"))
         add(moduleChooser)
 
-        isMakeBeforeRunCheckbox = JCheckBox("Make before Run")
+        isMakeBeforeRunCheckbox = JCheckBox("Make module before Run")
         add(isMakeBeforeRunCheckbox)
         isMakeBeforeRunCheckbox.addItemListener {
             scratchFile.saveOptions {

idea/testData/worksheet/multiFile/inlineFunScriptRuntime/inlineFunScriptRuntime.ws.comp.after

@@ -0,0 +1,5 @@
+// REPL_MODE: false
+
+import inlineFun.*
+
+foo { 1 + 3 }         // RESULT: 4

idea/testData/worksheet/multiFile/inlineFunScriptRuntime/inlineFunScriptRuntime.ws.kts

@@ -0,0 +1,5 @@
+// REPL_MODE: ~REPL_MODE~
+
+import inlineFun.*
+
+foo { 1 + 3 }

idea/testData/worksheet/multiFile/inlineFunScriptRuntime/inlineFunScriptRuntime.ws.repl.after

@@ -0,0 +1,5 @@
+// REPL_MODE: true
+
+import inlineFun.*
+
+foo { 1 + 3 }         // RESULT: res1: kotlin.Int = 4

idea/testData/worksheet/multiFile/inlineFunScriptRuntime/myFun.kt

@@ -0,0 +1,3 @@
+package inlineFun
+
+inline fun foo(f: () -> Int): Int = f()

idea/testData/worksheet/multiFile/javaDepScriptRuntime/javaDepScriptRuntime.ws.comp.after

@@ -0,0 +1,5 @@
+// REPL_MODE: false
+
+import myTest.MyJavaClass
+
+MyJavaClass().test()         // RESULT: 1

idea/testData/worksheet/multiFile/javaDepScriptRuntime/javaDepScriptRuntime.ws.kts

@@ -0,0 +1,5 @@
+// REPL_MODE: ~REPL_MODE~
+
+import myTest.MyJavaClass
+
+MyJavaClass().test()

idea/testData/worksheet/multiFile/javaDepScriptRuntime/javaDepScriptRuntime.ws.repl.after

@@ -0,0 +1,5 @@
+// REPL_MODE: true
+
+import myTest.MyJavaClass
+
+MyJavaClass().test()         // RESULT: res1: kotlin.Int = 1

idea/testData/worksheet/multiFile/javaDepScriptRuntime/myTest/MyJavaClass.java

@@ -0,0 +1,7 @@
+package myTest;
+
+public class MyJavaClass {
+    public int test() {
+        return 1;
+    }
+}

idea/testData/worksheet/simpleScriptRuntime.ws.comp.after

@@ -0,0 +1,4 @@
+// REPL_MODE: false
+
+val a = 1    // RESULT: val a: Int
+a            // RESULT: 1

idea/testData/worksheet/simpleScriptRuntime.ws.kts

@@ -0,0 +1,4 @@
+// REPL_MODE: ~REPL_MODE~
+
+val a = 1
+a

idea/testData/worksheet/simpleScriptRuntime.ws.repl.after

@@ -0,0 +1,4 @@
+// REPL_MODE: true
+
+val a = 1
+a            // RESULT: res1: kotlin.Int = 1

idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt

@@ -22,6 +22,8 @@ import com.intellij.testFramework.TestActionEvent
 import com.intellij.util.ui.UIUtil
 import org.jetbrains.kotlin.codegen.forTestCompile.ForTestCompileRuntime
 import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.idea.actions.KOTLIN_SCRIPT_EXTENSION
+import org.jetbrains.kotlin.idea.actions.KOTLIN_WORKSHEET_EXTENSION
 import org.jetbrains.kotlin.idea.core.script.ScriptDependenciesManager
 import org.jetbrains.kotlin.idea.highlighter.KotlinHighlightingUtil
 import org.jetbrains.kotlin.idea.scratch.actions.ClearScratchAction
@@ -42,15 +44,34 @@ import org.junit.Assert
 import java.io.File
 
 abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
-    fun doReplTest(fileName: String) {
-        doTest(fileName, true)
+
+    fun doWorksheetReplTest(fileName: String) {
+        doTest(fileName = fileName, isRepl = true, isWorksheet = true)
     }
 
-    fun doCompilingTest(fileName: String) {
-        doTest(fileName, false)
+    fun doScratchReplTest(fileName: String) {
+        doTest(fileName = fileName, isRepl = true, isWorksheet = false)
     }
 
-    fun doMultiFileTest(dirName: String) {
+    fun doWorksheetCompilingTest(fileName: String) {
+        doTest(fileName = fileName, isRepl = false, isWorksheet = true)
+    }
+
+    fun doScratchCompilingTest(fileName: String) {
+        doTest(fileName = fileName, isRepl = false, isWorksheet = false)
+    }
+
+    fun doWorksheetMultiFileTest(dirName: String) {
+        doMultiFileTest(dirName, isWorksheet = true)
+    }
+
+    fun doScratchMultiFileTest(dirName: String) {
+        doMultiFileTest(dirName, isWorksheet = false)
+    }
+
+    private fun doMultiFileTest(dirName: String, isWorksheet: Boolean) {
+        val mainFileExtension = if (isWorksheet) KOTLIN_WORKSHEET_EXTENSION else KOTLIN_SCRIPT_EXTENSION
+
         val javaFiles = arrayListOf<File>()
         val kotlinFiles = arrayListOf<File>()
         val baseDir = File(testDataPath, dirName)
@@ -75,23 +96,27 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
 
         PsiTestUtil.setCompilerOutputPath(myFixture.module, outputDir.path, false)
 
-        val mainFileName = "$dirName/${getTestName(true)}.kts"
-        doCompilingTest(mainFileName)
+        val mainFileName = "$dirName/${getTestName(true)}.$mainFileExtension"
+        doTest(mainFileName, isRepl = false, isWorksheet = isWorksheet)
 
         launchAction(ClearScratchAction())
 
-        doReplTest(mainFileName)
+        doTest(mainFileName, isRepl = true, isWorksheet = isWorksheet)
 
         ModuleRootModificationUtil.updateModel(myFixture.module) { model ->
             model.getModuleExtension(CompilerModuleExtension::class.java).inheritCompilerOutputPath(true)
         }
     }
 
-    fun doTest(fileName: String, isRepl: Boolean) {
+    fun doTest(fileName: String, isRepl: Boolean, isWorksheet: Boolean) {
         val sourceFile = File(testDataPath, fileName)
         val fileText = sourceFile.readText().inlinePropertiesValues(isRepl)
 
-        configureScratchByText(sourceFile.name, fileText)
+        if (isWorksheet) {
+            configureWorksheetByText(sourceFile.name, fileText)
+        } else {
+            configureScratchByText(sourceFile.name, fileText)
+        }
 
         if (!KotlinHighlightingUtil.shouldHighlight(myFixture.file)) error("Highlighting for scratch file is switched off")
 
@@ -159,6 +184,20 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
         return scratchPanel
     }
 
+    protected fun configureWorksheetByText(name: String, text: String): ScratchTopPanel {
+        val worksheetFile = myFixture.configureByText(name, text).virtualFile
+
+        ScriptDependenciesManager.updateScriptDependenciesSynchronously(worksheetFile, project)
+
+        val (_, scratchPanel) = getEditorWithScratchPanel(myManager, myFixture.file.virtualFile)
+            ?: error("Couldn't find scratch panel")
+
+        configureOptions(scratchPanel, text, myFixture.module)
+
+        return scratchPanel
+    }
+
+
     protected fun launchScratch() {
         val action = RunScratchAction()
         launchAction(action)
@@ -204,6 +243,7 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
         return when {
             testName.endsWith("WithKotlinTest") -> INSTANCE_WITH_KOTLIN_TEST
             testName.endsWith("NoRuntime") -> INSTANCE_WITHOUT_RUNTIME
+            testName.endsWith("ScriptRuntime") -> INSTANCE_WITH_SCRIPT_RUNTIME
             else -> KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE_FULL_JDK
         }
     }
@@ -246,6 +286,15 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
             override fun getSdk() = PluginTestCaseBase.fullJdk()
         }
 
+        private val INSTANCE_WITH_SCRIPT_RUNTIME = object : KotlinWithJdkAndRuntimeLightProjectDescriptor(
+            arrayListOf(
+                ForTestCompileRuntime.runtimeJarForTests(),
+                ForTestCompileRuntime.scriptRuntimeJarForTests()
+            )
+        ) {
+            override fun getSdk() = PluginTestCaseBase.fullJdk()
+        }
+
         fun configureOptions(
             scratchPanel: ScratchTopPanel,
             fileText: String,

idea/tests/org/jetbrains/kotlin/idea/scratch/ScratchRunActionTestGenerated.java

@@ -22,12 +22,12 @@ public class ScratchRunActionTestGenerated extends AbstractScratchRunActionTest
     @TestMetadata("idea/testData/scratch")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
-    public static class Compiling extends AbstractScratchRunActionTest {
+    public static class ScratchCompiling extends AbstractScratchRunActionTest {
         private void runTest(String testDataFilePath) throws Exception {
-            KotlinTestUtils.runTest(this::doCompilingTest, TargetBackend.ANY, testDataFilePath);
+            KotlinTestUtils.runTest(this::doScratchCompilingTest, TargetBackend.ANY, testDataFilePath);
         }
 
-        public void testAllFilesPresentInCompiling() throws Exception {
+        public void testAllFilesPresentInScratchCompiling() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/scratch"), Pattern.compile("^(.+)\\.kts$"), TargetBackend.ANY, false);
         }
 
@@ -110,12 +110,12 @@ public class ScratchRunActionTestGenerated extends AbstractScratchRunActionTest
     @TestMetadata("idea/testData/scratch")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
-    public static class Repl extends AbstractScratchRunActionTest {
+    public static class ScratchRepl extends AbstractScratchRunActionTest {
         private void runTest(String testDataFilePath) throws Exception {
-            KotlinTestUtils.runTest(this::doReplTest, TargetBackend.ANY, testDataFilePath);
+            KotlinTestUtils.runTest(this::doScratchReplTest, TargetBackend.ANY, testDataFilePath);
         }
 
-        public void testAllFilesPresentInRepl() throws Exception {
+        public void testAllFilesPresentInScratchRepl() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/scratch"), Pattern.compile("^(.+)\\.kts$"), TargetBackend.ANY, false);
         }
 
@@ -198,12 +198,12 @@ public class ScratchRunActionTestGenerated extends AbstractScratchRunActionTest
     @TestMetadata("idea/testData/scratch/multiFile")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
-    public static class MultiFile extends AbstractScratchRunActionTest {
+    public static class ScratchMultiFile extends AbstractScratchRunActionTest {
         private void runTest(String testDataFilePath) throws Exception {
-            KotlinTestUtils.runTest(this::doMultiFileTest, TargetBackend.ANY, testDataFilePath);
+            KotlinTestUtils.runTest(this::doScratchMultiFileTest, TargetBackend.ANY, testDataFilePath);
         }
 
-        public void testAllFilesPresentInMultiFile() throws Exception {
+        public void testAllFilesPresentInScratchMultiFile() throws Exception {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/scratch/multiFile"), Pattern.compile("^([^\\.]+)$"), TargetBackend.ANY, false);
         }
 
@@ -217,4 +217,63 @@ public class ScratchRunActionTestGenerated extends AbstractScratchRunActionTest
             runTest("idea/testData/scratch/multiFile/javaDep/");
         }
     }
+
+    @TestMetadata("idea/testData/worksheet")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class WorksheetCompiling extends AbstractScratchRunActionTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doWorksheetCompilingTest, TargetBackend.ANY, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInWorksheetCompiling() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/worksheet"), Pattern.compile("^(.+)\\.ws.kts$"), TargetBackend.ANY, false);
+        }
+
+        @TestMetadata("simpleScriptRuntime.ws.kts")
+        public void testSimpleScriptRuntime() throws Exception {
+            runTest("idea/testData/worksheet/simpleScriptRuntime.ws.kts");
+        }
+    }
+
+    @TestMetadata("idea/testData/worksheet")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class WorksheetRepl extends AbstractScratchRunActionTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doWorksheetReplTest, TargetBackend.ANY, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInWorksheetRepl() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/worksheet"), Pattern.compile("^(.+)\\.ws.kts$"), TargetBackend.ANY, false);
+        }
+
+        @TestMetadata("simpleScriptRuntime.ws.kts")
+        public void testSimpleScriptRuntime() throws Exception {
+            runTest("idea/testData/worksheet/simpleScriptRuntime.ws.kts");
+        }
+    }
+
+    @TestMetadata("idea/testData/worksheet/multiFile")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class WorksheetMultiFile extends AbstractScratchRunActionTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doWorksheetMultiFileTest, TargetBackend.ANY, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInWorksheetMultiFile() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/worksheet/multiFile"), Pattern.compile("^([^\\.]+)$"), TargetBackend.ANY, false);
+        }
+
+        @TestMetadata("inlineFunScriptRuntime")
+        public void testInlineFunScriptRuntime() throws Exception {
+            runTest("idea/testData/worksheet/multiFile/inlineFunScriptRuntime/");
+        }
+
+        @TestMetadata("javaDepScriptRuntime")
+        public void testJavaDepScriptRuntime() throws Exception {
+            runTest("idea/testData/worksheet/multiFile/javaDepScriptRuntime/");
+        }
+    }
 }