Scratch: do not store psiFile in ScratchFile

2018-01-25 10:22:06 +03:00
parent c32a7dd852
commit efabf80d37
13 changed files with 119 additions and 114 deletions
@@ -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) {
-
-    override fun getExpressions(): List<ScratchExpression> {
+class KtScratchFile(project: Project, editor: TextEditor) : ScratchFile(project, editor) {
+    override fun getExpressions(psiFile: PsiFile): List<ScratchExpression> {
        // todo multiple expressions at one line
        val doc = PsiDocumentManager.getInstance(psiFile.project).getDocument(psiFile) ?: return emptyList()
        var line = 0
@@ -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? {
-        return (psiFile as? KtFile)?.let { KtScratchFile(psiFile) }
+    override fun createFile(project: Project, editor: TextEditor): ScratchFile? {
+        return KtScratchFile(project, editor)
    }

    override fun createReplExecutor(file: ScratchFile) = KtScratchReplExecutor(file)
@@ -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 fun getExpressions(): List<ScratchExpression>
+abstract class ScratchFile(val project: Project, val editor: TextEditor) {
+    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)
@@ -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)
-        }
-
    }
 }
@@ -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 panel = getScratchPanelFromSelectedEditor(project) ?: return
-        val scratchFile = panel.scratchFile
+        val scratchFile = scratchPanel.scratchFile
+        val psiFile = scratchFile.getPsiFile() ?: return

-        ScratchFileLanguageProvider.get(scratchFile.psiFile.language)?.getOutputHandler()?.clear(scratchFile)
+        ScratchFileLanguageProvider.get(psiFile.language)?.getOutputHandler()?.clear(scratchFile)
    }
 }
@@ -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 = scratchTopPanel.scratchFile
+        val scratchFile = scratchPanel.scratchFile
+        val psiFile = scratchFile.getPsiFile() ?: return

-        val isMakeBeforeRun = scratchTopPanel.isMakeBeforeRun()
-        val isRepl = scratchTopPanel.isRepl()
+        val isMakeBeforeRun = scratchPanel.isMakeBeforeRun()
+        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
            }
@@ -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 {
@@ -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)
-        ScratchToolWindow.clearToolWindow(file.psiFile.project)
+        clearInlays(file.editor)
+        ScratchToolWindow.clearToolWindow(file.project)
    }

    private fun createInlay(file: ScratchFile, line: Int, inlayText: String, outputType: ScratchOutputType) {
        UIUtil.invokeLaterIfNeeded {
-            val project = file.psiFile.project
-            val textEditor = FileEditorManager.getInstance(project).getSelectedEditor(file.psiFile.virtualFile) as? TextEditor ?: return@invokeLaterIfNeeded
-            val editor = textEditor.editor
+            val editor = file.editor.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 {
-        val doc = PsiDocumentManager.getInstance(file.psiFile.project).getDocument(file.psiFile) ?: return -1
+    private fun maxLineLength(file: ScratchFile): Int {
+        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)
-            editors.filterIsInstance<TextEditor>()
-                .flatMap { it.editor.inlayModel.getInlineElementsInRange(0, it.editor.document.textLength) }
+            editor
+                .editor.inlayModel.getInlineElementsInRange(0, editor.editor.document.textLength)
                .filter { it.renderer is ScratchFileRenderer }
                .forEach { Disposer.dispose(it) }
        }
@@ -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}")
@@ -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? {
-    val allTextEditors = getAllTextEditors(project, virtualFile)
-    for (editor in allTextEditors) {
-        if (editor.scratchTopPanel != null) return null
-    }
-    return allTextEditors.firstOrNull()
+fun getEditorWithoutScratchPanel(fileManager: FileEditorManager, virtualFile: VirtualFile): TextEditor? {
+    val editor = fileManager.getSelectedEditor(virtualFile) as? TextEditor
+    if (editor?.scratchTopPanel != null) return null
+    return editor
 }

-fun getEditorWithScratchPanel(project: Project, virtualFile: VirtualFile): Pair<TextEditor, ScratchTopPanel>? {
-    val firstWithPanel = getAllTextEditors(project, virtualFile).firstOrNull { it.scratchTopPanel != null }
-    return firstWithPanel?.let { firstWithPanel to firstWithPanel.scratchTopPanel!! }
-}
-
-fun getScratchPanel(psiFile: PsiFile): ScratchTopPanel? {
-    return getAllTextEditors(psiFile.project, psiFile.virtualFile).mapNotNull { it.scratchTopPanel }.firstOrNull()
+fun getEditorWithScratchPanel(fileManager: FileEditorManager, virtualFile: VirtualFile): Pair<TextEditor, ScratchTopPanel>? {
+    val editor = fileManager.getSelectedEditor(virtualFile) as? TextEditor ?: return null
+    val scratchTopPanel = editor.scratchTopPanel ?: return null
+    return editor to scratchTopPanel
 }

 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? =
-    FileEditorManager.getInstance(project).selectedEditors.asSequence()
-        .filterIsInstance<TextEditor>()
-        .mapNotNull { it.scratchTopPanel }
-        .firstOrNull()
-
-private fun getAllTextEditors(project: Project, virtualFile: VirtualFile) =
-    FileEditorManager.getInstance(project).getAllEditors(virtualFile).filterIsInstance<TextEditor>().asSequence()
+fun TextEditor.getScratchPanel(): ScratchTopPanel? {
+    return scratchTopPanel
+}

 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"))
@@ -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
-            getEditorWithoutScratchPanel(myProject, file)?.let { editor ->
-                editor.addScratchPanel(panel)
-            }
+            val editor = getEditorWithoutScratchPanel(source, file) ?: return
+
+            ScratchTopPanel.createPanel(myProject, file, editor)
        }

-        override fun fileClosed(source: FileEditorManager, file: VirtualFile) {
-            if (!isPluggable(file)) return
+        override fun fileClosed(source: FileEditorManager, file: VirtualFile) {}
+    }

-            getEditorWithScratchPanel(myProject, file)?.let { (editor, panel) ->
-                editor.removeScratchPanel(panel)
-            }
-        }
+    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
    }
 }
@@ -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?
-    get() = getScratchPanel(psiFile)?.takeIf { it.scratchFile == this@scratchTopPanel }
+class ScratchTopPanel private constructor(val scratchFile: ScratchFile) : JPanel(HorizontalLayout(5)), Disposable {
+    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? {
-            val psiFile = PsiManager.getInstance(project).findFile(virtualFile) ?: return null
-            val scratchFile = ScratchFileLanguageProvider.createFile(psiFile) ?: return null
-            return ScratchTopPanel(scratchFile)
+        fun createPanel(project: Project, virtualFile: VirtualFile, editor: TextEditor) {
+            val psiFile = PsiManager.getInstance(project).findFile(virtualFile) ?: return
+            val scratchFile = ScratchFileLanguageProvider.get(psiFile.language)?.createFile(project, editor) ?: return
+            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
        }
    }
 }
@@ -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())
-        launchAction(event, RunScratchAction())
+        val action = RunScratchAction(scratchPanel)
+        val event = getActionEvent(scratchFile, action)
+        launchAction(event, action)

        UIUtil.dispatchAllInvocationEvents()