Scratch: do not store psiFile in ScratchFile

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

@@ -16,20 +16,21 @@
 
 package org.jetbrains.kotlin.idea.scratch
 
+import com.intellij.openapi.fileEditor.TextEditor
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.util.PsiUtil
 import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
 import org.jetbrains.kotlin.idea.refactoring.getLineNumber
 import org.jetbrains.kotlin.idea.refactoring.getLineStartOffset
 import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtImportDirective
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 
-class KtScratchFile(psiFile: KtFile) : ScratchFile(psiFile) {
+class KtScratchFile(project: Project, editor: TextEditor) : ScratchFile(project, editor) {
-
+    override fun getExpressions(psiFile: PsiFile): List<ScratchExpression> {
-    override fun getExpressions(): List<ScratchExpression> {
         // todo multiple expressions at one line
         val doc = PsiDocumentManager.getInstance(psiFile.project).getDocument(psiFile) ?: return emptyList()
         var line = 0

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

@@ -16,6 +16,8 @@
 
 package org.jetbrains.kotlin.idea.scratch
 
+import com.intellij.openapi.fileEditor.TextEditor
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.idea.scratch.compile.KtCompilingExecutor
 import org.jetbrains.kotlin.idea.scratch.output.InlayScratchOutputHandler
@@ -23,8 +25,8 @@ import org.jetbrains.kotlin.idea.scratch.repl.KtScratchReplExecutor
 import org.jetbrains.kotlin.psi.KtFile
 
 class KtScratchFileLanguageProvider : ScratchFileLanguageProvider() {
-    override fun createFile(psiFile: PsiFile): ScratchFile? {
+    override fun createFile(project: Project, editor: TextEditor): ScratchFile? {
-        return (psiFile as? KtFile)?.let { KtScratchFile(psiFile) }
+        return KtScratchFile(project, editor)
     }
 
     override fun createReplExecutor(file: ScratchFile) = KtScratchReplExecutor(file)

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

@@ -16,11 +16,28 @@
 
 package org.jetbrains.kotlin.idea.scratch
 
+import com.intellij.openapi.fileEditor.TextEditor
+import com.intellij.openapi.module.Module
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 
-abstract class ScratchFile(val psiFile: PsiFile) {
+abstract class ScratchFile(val project: Project, val editor: TextEditor) {
-    abstract fun getExpressions(): List<ScratchExpression>
+    fun getExpressions(): List<ScratchExpression> {
+        val psiFile = getPsiFile() ?: return emptyList()
+        return getExpressions(psiFile)
+    }
+
+    fun getPsiFile(): PsiFile? {
+        return PsiDocumentManager.getInstance(project).getPsiFile(editor.editor.document)
+    }
+
+    fun getModule(): Module? {
+        return editor.getScratchPanel()?.getModule()
+    }
+
+    abstract fun getExpressions(psiFile: PsiFile): List<ScratchExpression>
 }
 
 data class ScratchExpression(val element: PsiElement, val lineStart: Int, val lineEnd: Int = lineStart)

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

@@ -18,13 +18,14 @@ package org.jetbrains.kotlin.idea.scratch
 
 import com.intellij.lang.Language
 import com.intellij.lang.LanguageExtension
+import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.fileTypes.FileType
 import com.intellij.openapi.fileTypes.LanguageFileType
-import com.intellij.psi.PsiFile
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandler
 
 abstract class ScratchFileLanguageProvider {
-    abstract fun createFile(psiFile: PsiFile): ScratchFile?
+    abstract fun createFile(project: Project, editor: TextEditor): ScratchFile?
     abstract fun createReplExecutor(file: ScratchFile): ScratchExecutor?
     abstract fun createCompilingExecutor(file: ScratchFile): ScratchExecutor?
 
@@ -40,10 +41,5 @@ abstract class ScratchFileLanguageProvider {
         fun get(fileType: FileType): ScratchFileLanguageProvider? {
             return (fileType as? LanguageFileType)?.language?.let { get(it) }
         }
-
-        fun createFile(psiFile: PsiFile): ScratchFile? {
-            return get(psiFile.language)?.createFile(psiFile)
-        }
-
     }
 }

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

@@ -21,18 +21,17 @@ import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.scratch.ScratchFileLanguageProvider
-import org.jetbrains.kotlin.idea.scratch.getScratchPanelFromSelectedEditor
+import org.jetbrains.kotlin.idea.scratch.ui.ScratchTopPanel
 
-class ClearScratchAction : AnAction(
+class ClearScratchAction(private val scratchPanel: ScratchTopPanel) : AnAction(
     KotlinBundle.message("scratch.clear.button"),
     KotlinBundle.message("scratch.clear.button"),
     AllIcons.Actions.GC
 ) {
     override fun actionPerformed(e: AnActionEvent) {
-        val project = e.project ?: return
+        val scratchFile = scratchPanel.scratchFile
-        val panel = getScratchPanelFromSelectedEditor(project) ?: return
+        val psiFile = scratchFile.getPsiFile() ?: return
-        val scratchFile = panel.scratchFile
 
-        ScratchFileLanguageProvider.get(scratchFile.psiFile.language)?.getOutputHandler()?.clear(scratchFile)
+        ScratchFileLanguageProvider.get(psiFile.language)?.getOutputHandler()?.clear(scratchFile)
     }
 }

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

@@ -23,11 +23,11 @@ import com.intellij.openapi.compiler.CompilerManager
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.scratch.ScratchFile
 import org.jetbrains.kotlin.idea.scratch.ScratchFileLanguageProvider
-import org.jetbrains.kotlin.idea.scratch.getScratchPanelFromSelectedEditor
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandlerAdapter
 import org.jetbrains.kotlin.idea.scratch.printDebugMessage
+import org.jetbrains.kotlin.idea.scratch.ui.ScratchTopPanel
 
-class RunScratchAction : AnAction(
+class RunScratchAction(private val scratchPanel: ScratchTopPanel) : AnAction(
     KotlinBundle.message("scratch.run.button"),
     KotlinBundle.message("scratch.run.button"),
     AllIcons.Actions.Execute
@@ -35,19 +35,19 @@ class RunScratchAction : AnAction(
     override fun actionPerformed(e: AnActionEvent) {
         val project = e.project ?: return
 
-        val scratchTopPanel = getScratchPanelFromSelectedEditor(project) ?: return
+        val scratchFile = scratchPanel.scratchFile
-        val scratchFile = scratchTopPanel.scratchFile
+        val psiFile = scratchFile.getPsiFile() ?: return
 
-        val isMakeBeforeRun = scratchTopPanel.isMakeBeforeRun()
+        val isMakeBeforeRun = scratchPanel.isMakeBeforeRun()
-        val isRepl = scratchTopPanel.isRepl()
+        val isRepl = scratchPanel.isRepl()
 
-        val provider = ScratchFileLanguageProvider.get(scratchFile.psiFile.language) ?: return
+        val provider = ScratchFileLanguageProvider.get(psiFile.language) ?: return
 
         val handler = provider.getOutputHandler()
 
         org.jetbrains.kotlin.idea.scratch.LOG.printDebugMessage("Run Action: isMakeBeforeRun = $isMakeBeforeRun, isRepl = $isRepl")
 
-        val module = scratchTopPanel.getModule()
+        val module = scratchPanel.getModule()
         if (module == null) {
             handler.error(scratchFile, "Module should be selected")
             handler.onFinish(scratchFile)
@@ -57,7 +57,7 @@ class RunScratchAction : AnAction(
         val runnable = r@ {
             val executor = if (isRepl) provider.createReplExecutor(scratchFile) else provider.createCompilingExecutor(scratchFile)
             if (executor == null) {
-                handler.error(scratchFile, "Couldn't run ${scratchFile.psiFile.name}")
+                handler.error(scratchFile, "Couldn't run ${psiFile.name}")
                 handler.onFinish(scratchFile)
                 return@r
             }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/compile/KtCompilingExecutor.kt

@@ -45,7 +45,6 @@ import org.jetbrains.kotlin.idea.scratch.ScratchFile
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutput
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputType
 import org.jetbrains.kotlin.idea.scratch.printDebugMessage
-import org.jetbrains.kotlin.idea.scratch.ui.scratchTopPanel
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFile
@@ -58,9 +57,10 @@ class KtCompilingExecutor(file: ScratchFile) : ScratchExecutor(file) {
     override fun execute() {
         handlers.forEach { it.onStart(file) }
 
-        val module = file.scratchTopPanel?.getModule() ?: return error("Module should be selected")
+        val module = file.getModule() ?: return error("Module should be selected")
+        val psiFile = file.getPsiFile() ?: return error("Couldn't find psiFile for current editor")
 
-        if (!checkForErrors(file.psiFile as KtFile)) {
+        if (!checkForErrors(psiFile as KtFile)) {
             return error("Compilation Error")
         }
 
@@ -72,7 +72,7 @@ class KtCompilingExecutor(file: ScratchFile) : ScratchExecutor(file) {
                     LOG.printDebugMessage("After processing by KtScratchSourceFileProcessor:\n ${result.code}")
 
                     val modifiedScratchSourceFile =
-                        KtPsiFactory(file.psiFile.project).createFileWithLightClassSupport("tmp.kt", result.code, file.psiFile)
+                        KtPsiFactory(psiFile.project).createFileWithLightClassSupport("tmp.kt", result.code, psiFile)
 
                     try {
                         val tempDir = compileFileToTempDir(modifiedScratchSourceFile) ?: return@invokeLater
@@ -89,7 +89,7 @@ class KtCompilingExecutor(file: ScratchFile) : ScratchExecutor(file) {
                         }
                     } catch (e: Throwable) {
                         LOG.info(result.code, e)
-                        handlers.forEach { it.error(file, e.message ?: "Couldn't compile ${file.psiFile.name}") }
+                        handlers.forEach { it.error(file, e.message ?: "Couldn't compile ${psiFile.name}") }
                     } finally {
                         handlers.forEach { it.onFinish(file) }
                     }
@@ -114,7 +114,7 @@ class KtCompilingExecutor(file: ScratchFile) : ScratchExecutor(file) {
         }
 
         val state = GenerationState.Builder(
-            file.psiFile.project,
+            file.project,
             ClassBuilderFactories.binaries(false),
             resolutionFacade.moduleDescriptor,
             bindingContext,
@@ -178,9 +178,10 @@ class KtCompilingExecutor(file: ScratchFile) : ScratchExecutor(file) {
             val bindingContext = analysisResult.bindingContext
             val diagnostics = bindingContext.diagnostics.filter { it.severity == Severity.ERROR }
             if (diagnostics.isNotEmpty()) {
+                val scratchPsiFile = file.getPsiFile()
                 diagnostics.forEach { diagnostic ->
                     val errorText = DefaultErrorMessages.render(diagnostic)
-                    if (psiFile == file.psiFile) {
+                    if (psiFile == scratchPsiFile) {
                         if (diagnostic.psiElement.containingFile == psiFile) {
                             val scratchExpression = file.findExpression(diagnostic.psiElement)
                             handlers.forEach {

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/output/InlayScratchOutputHandler.kt

@@ -22,12 +22,8 @@ import com.intellij.openapi.editor.EditorCustomElementRenderer
 import com.intellij.openapi.editor.impl.ComplementaryFontsRegistry
 import com.intellij.openapi.editor.impl.FontInfo
 import com.intellij.openapi.editor.markup.TextAttributes
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.TextEditor
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
-import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.psi.PsiDocumentManager
 import com.intellij.ui.Colors
 import com.intellij.ui.JBColor
 import com.intellij.util.ui.UIUtil
@@ -56,7 +52,7 @@ object InlayScratchOutputHandler : ScratchOutputHandler {
     }
 
     override fun error(file: ScratchFile, message: String) {
-        ScratchToolWindow.addMessageToToolWindow(file.psiFile.project, message, ConsoleViewContentType.ERROR_OUTPUT)
+        ScratchToolWindow.addMessageToToolWindow(file.project, message, ConsoleViewContentType.ERROR_OUTPUT)
     }
 
     override fun onFinish(file: ScratchFile) {
@@ -64,15 +60,13 @@ object InlayScratchOutputHandler : ScratchOutputHandler {
     }
 
     override fun clear(file: ScratchFile) {
-        clearInlays(file.psiFile.project, file.psiFile.virtualFile)
+        clearInlays(file.editor)
-        ScratchToolWindow.clearToolWindow(file.psiFile.project)
+        ScratchToolWindow.clearToolWindow(file.project)
     }
 
     private fun createInlay(file: ScratchFile, line: Int, inlayText: String, outputType: ScratchOutputType) {
         UIUtil.invokeLaterIfNeeded {
-            val project = file.psiFile.project
+            val editor = file.editor.editor
-            val textEditor = FileEditorManager.getInstance(project).getSelectedEditor(file.psiFile.virtualFile) as? TextEditor ?: return@invokeLaterIfNeeded
-            val editor = textEditor.editor
 
             val lineStartOffset = editor.document.getLineStartOffset(line)
             val lineEndOffset = editor.document.getLineEndOffset(line)
@@ -96,19 +90,18 @@ object InlayScratchOutputHandler : ScratchOutputHandler {
         }
     }
 
-    fun maxLineLength(file: ScratchFile): Int {
+    private fun maxLineLength(file: ScratchFile): Int {
-        val doc = PsiDocumentManager.getInstance(file.psiFile.project).getDocument(file.psiFile) ?: return -1
+        val doc = file.editor.editor.document
         return (0 until doc.lineCount)
             .map { doc.getLineEndOffset(it) - doc.getLineStartOffset(it) }
             .max()
                 ?: -1
     }
 
-    private fun clearInlays(project: Project, file: VirtualFile) {
+    private fun clearInlays(editor: TextEditor) {
         UIUtil.invokeLaterIfNeeded {
-            val editors = FileEditorManager.getInstance(project).getEditors(file)
+            editor
-            editors.filterIsInstance<TextEditor>()
+                .editor.inlayModel.getInlineElementsInRange(0, editor.editor.document.textLength)
-                .flatMap { it.editor.inlayModel.getInlineElementsInRange(0, it.editor.document.textLength) }
                 .filter { it.renderer is ScratchFileRenderer }
                 .forEach { Disposer.dispose(it) }
         }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/repl/KtScratchReplExecutor.kt

@@ -32,7 +32,6 @@ import org.jetbrains.kotlin.idea.scratch.ScratchFile
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutput
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputType
 import org.jetbrains.kotlin.idea.scratch.printDebugMessage
-import org.jetbrains.kotlin.idea.scratch.ui.scratchTopPanel
 import org.w3c.dom.Element
 import org.xml.sax.InputSource
 import java.io.ByteArrayInputStream
@@ -48,7 +47,7 @@ class KtScratchReplExecutor(file: ScratchFile) : ScratchExecutor(file) {
     override fun execute() {
         handlers.forEach { it.onStart(file) }
 
-        val module = file.scratchTopPanel?.getModule() ?: return error(file, "Module should be selected")
+        val module = file.getModule() ?: return error(file, "Module should be selected")
         val cmdLine = KotlinConsoleKeeper.createCommandLine(module)
 
         LOG.printDebugMessage("Execute REPL: ${cmdLine.commandLineString}")

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

@@ -9,9 +9,9 @@ import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.TextEditor
 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 com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.idea.scratch.ui.ScratchTopPanel
 import org.jetbrains.kotlin.psi.UserDataProperty
 
@@ -20,21 +20,16 @@ internal fun Logger.printDebugMessage(str: String) {
     if (isDebugEnabled) debug("SCRATCH: $str")
 }
 
-fun getEditorWithoutScratchPanel(project: Project, virtualFile: VirtualFile): TextEditor? {
+fun getEditorWithoutScratchPanel(fileManager: FileEditorManager, virtualFile: VirtualFile): TextEditor? {
-    val allTextEditors = getAllTextEditors(project, virtualFile)
+    val editor = fileManager.getSelectedEditor(virtualFile) as? TextEditor
-    for (editor in allTextEditors) {
+    if (editor?.scratchTopPanel != null) return null
-        if (editor.scratchTopPanel != null) return null
+    return editor
-    }
-    return allTextEditors.firstOrNull()
 }
 
-fun getEditorWithScratchPanel(project: Project, virtualFile: VirtualFile): Pair<TextEditor, ScratchTopPanel>? {
+fun getEditorWithScratchPanel(fileManager: FileEditorManager, virtualFile: VirtualFile): Pair<TextEditor, ScratchTopPanel>? {
-    val firstWithPanel = getAllTextEditors(project, virtualFile).firstOrNull { it.scratchTopPanel != null }
+    val editor = fileManager.getSelectedEditor(virtualFile) as? TextEditor ?: return null
-    return firstWithPanel?.let { firstWithPanel to firstWithPanel.scratchTopPanel!! }
+    val scratchTopPanel = editor.scratchTopPanel ?: return null
-}
+    return editor to scratchTopPanel
-
-fun getScratchPanel(psiFile: PsiFile): ScratchTopPanel? {
-    return getAllTextEditors(psiFile.project, psiFile.virtualFile).mapNotNull { it.scratchTopPanel }.firstOrNull()
 }
 
 fun getAllEditorsWithScratchPanel(project: Project): List<Pair<TextEditor, ScratchTopPanel>> =
@@ -43,23 +38,20 @@ fun getAllEditorsWithScratchPanel(project: Project): List<Pair<TextEditor, Scrat
         if (panel != null) it to panel else null
     }
 
-fun getScratchPanelFromSelectedEditor(project: Project): ScratchTopPanel? =
+fun TextEditor.getScratchPanel(): ScratchTopPanel? {
-    FileEditorManager.getInstance(project).selectedEditors.asSequence()
+    return scratchTopPanel
-        .filterIsInstance<TextEditor>()
+}
-        .mapNotNull { it.scratchTopPanel }
-        .firstOrNull()
-
-private fun getAllTextEditors(project: Project, virtualFile: VirtualFile) =
-    FileEditorManager.getInstance(project).getAllEditors(virtualFile).filterIsInstance<TextEditor>().asSequence()
 
 fun TextEditor.addScratchPanel(panel: ScratchTopPanel) {
     scratchTopPanel = panel
-    FileEditorManager.getInstance(panel.scratchFile.psiFile.project).addTopComponent(this, panel)
+    FileEditorManager.getInstance(panel.scratchFile.project).addTopComponent(this, panel)
+
+    Disposer.register(this, panel)
 }
 
 fun TextEditor.removeScratchPanel(panel: ScratchTopPanel) {
     scratchTopPanel = null
-    FileEditorManager.getInstance(panel.scratchFile.psiFile.project).removeTopComponent(this, panel)
+    FileEditorManager.getInstance(panel.scratchFile.project).removeTopComponent(this, panel)
 }
 
 private var TextEditor.scratchTopPanel: ScratchTopPanel? by UserDataProperty<TextEditor, ScratchTopPanel>(Key.create("scratch.panel"))

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

@@ -23,7 +23,10 @@ import com.intellij.openapi.fileEditor.FileEditorManagerListener
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiManager
-import org.jetbrains.kotlin.idea.scratch.*
+import org.jetbrains.kotlin.idea.scratch.ScratchFileLanguageProvider
+import org.jetbrains.kotlin.idea.scratch.getAllEditorsWithScratchPanel
+import org.jetbrains.kotlin.idea.scratch.getEditorWithoutScratchPanel
+import org.jetbrains.kotlin.idea.scratch.removeScratchPanel
 
 class ScratchFileHook(project: Project) : AbstractProjectComponent(project) {
 
@@ -36,28 +39,21 @@ class ScratchFileHook(project: Project) : AbstractProjectComponent(project) {
     }
 
     private inner class ScratchEditorListener : FileEditorManagerListener {
-        private fun isPluggable(file: VirtualFile): Boolean {
-            if (!file.isValid) return false
-            if (!ScratchFileService.isInScratchRoot(file)) return false
-            val psiFile = PsiManager.getInstance(myProject).findFile(file) ?: return false
-            return ScratchFileLanguageProvider.get(psiFile.fileType) != null
-        }
-
         override fun fileOpened(source: FileEditorManager, file: VirtualFile) {
             if (!isPluggable(file)) return
 
-            val panel = ScratchTopPanel.createPanel(myProject, file) ?: return
+            val editor = getEditorWithoutScratchPanel(source, file) ?: return
-            getEditorWithoutScratchPanel(myProject, file)?.let { editor ->
+
-                editor.addScratchPanel(panel)
+            ScratchTopPanel.createPanel(myProject, file, editor)
-            }
         }
 
-        override fun fileClosed(source: FileEditorManager, file: VirtualFile) {
+        override fun fileClosed(source: FileEditorManager, file: VirtualFile) {}
-            if (!isPluggable(file)) return
+    }
 
-            getEditorWithScratchPanel(myProject, file)?.let { (editor, panel) ->
+    private fun isPluggable(file: VirtualFile): Boolean {
-                editor.removeScratchPanel(panel)
+        if (!file.isValid) return false
-            }
+        if (!ScratchFileService.isInScratchRoot(file)) return false
-        }
+        val psiFile = PsiManager.getInstance(myProject).findFile(file) ?: return false
+        return ScratchFileLanguageProvider.get(psiFile.fileType) != null
     }
 }

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

@@ -18,12 +18,15 @@ package org.jetbrains.kotlin.idea.scratch.ui
 
 
 import com.intellij.application.options.ModulesComboBox
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.actionSystem.ActionPlaces
 import com.intellij.openapi.actionSystem.DefaultActionGroup
+import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.ui.components.panels.HorizontalLayout
 import org.jetbrains.annotations.TestOnly
@@ -31,18 +34,20 @@ import org.jetbrains.kotlin.idea.scratch.ScratchFile
 import org.jetbrains.kotlin.idea.scratch.ScratchFileLanguageProvider
 import org.jetbrains.kotlin.idea.scratch.actions.ClearScratchAction
 import org.jetbrains.kotlin.idea.scratch.actions.RunScratchAction
-import org.jetbrains.kotlin.idea.scratch.getScratchPanel
+import org.jetbrains.kotlin.idea.scratch.addScratchPanel
+import org.jetbrains.kotlin.idea.scratch.removeScratchPanel
 import javax.swing.*
 
-val ScratchFile.scratchTopPanel: ScratchTopPanel?
+class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel(HorizontalLayout(5)), Disposable {
-    get() = getScratchPanel(psiFile)?.takeIf { it.scratchFile == this@scratchTopPanel }
+    override fun dispose() {
+        scratchFile.editor.removeScratchPanel(this)
+    }
 
-class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel(HorizontalLayout(5)) {
     companion object {
-        fun createPanel(project: Project, virtualFile: VirtualFile): ScratchTopPanel? {
+        fun createPanel(project: Project, virtualFile: VirtualFile, editor: TextEditor) {
-            val psiFile = PsiManager.getInstance(project).findFile(virtualFile) ?: return null
+            val psiFile = PsiManager.getInstance(project).findFile(virtualFile) ?: return
-            val scratchFile = ScratchFileLanguageProvider.createFile(psiFile) ?: return null
+            val scratchFile = ScratchFileLanguageProvider.get(psiFile.language)?.createFile(project, editor) ?: return
-            return ScratchTopPanel(scratchFile)
+            editor.addScratchPanel(ScratchTopPanel(scratchFile))
         }
     }
 
@@ -53,7 +58,7 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
     init {
         add(createActionsToolbar())
 
-        moduleChooser = createModuleChooser(scratchFile.psiFile.project)
+        moduleChooser = createModuleChooser(scratchFile.project)
         add(JLabel("Use classpath of module"))
         add(moduleChooser)
 
@@ -76,9 +81,13 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
         moduleChooser.selectedModule = module
     }
 
-    fun addModuleListener(f: (Module) -> Unit) {
+    fun addModuleListener(f: (PsiFile, Module) -> Unit) {
         moduleChooser.addActionListener {
-            moduleChooser.selectedModule?.let { f(it) }
+            val selectedModule = moduleChooser.selectedModule
+            val psiFile = scratchFile.getPsiFile()
+            if (selectedModule != null && psiFile != null) {
+                f(psiFile, selectedModule)
+            }
         }
     }
 
@@ -97,9 +106,9 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
 
     private fun createActionsToolbar(): JComponent {
         val toolbarGroup = DefaultActionGroup().apply {
-            add(RunScratchAction())
+            add(RunScratchAction(this@ScratchTopPanel))
             addSeparator()
-            add(ClearScratchAction())
+            add(ClearScratchAction(this@ScratchTopPanel))
         }
 
         return ActionManager.getInstance().createActionToolbar(ActionPlaces.EDITOR_TOOLBAR, toolbarGroup, true).component
@@ -108,7 +117,6 @@ class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel
     private fun createModuleChooser(project: Project): ModulesComboBox {
         return ModulesComboBox().apply {
             fillModules(project)
-            selectedIndex = 0
         }
     }
 }

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

@@ -101,11 +101,12 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
         myFixture.openFileInEditor(scratchFile)
 
         val psiFile = PsiManager.getInstance(project).findFile(scratchFile) ?: error("Couldn't find psi file ${sourceFile.path}")
-        val (editor, scratchPanel) = getEditorWithScratchPanel(project, scratchFile)?: error("Couldn't find scratch panel")
+        val (editor, scratchPanel) = getEditorWithScratchPanel(myManager, scratchFile)?: error("Couldn't find scratch panel")
         scratchPanel.setReplMode(isRepl)
 
-        val event = getActionEvent(scratchFile, RunScratchAction())
+        val action = RunScratchAction(scratchPanel)
-        launchAction(event, RunScratchAction())
+        val event = getActionEvent(scratchFile, action)
+        launchAction(event, action)
 
         UIUtil.dispatchAllInvocationEvents()