KT-32366: Get replace editor with virtual file in ScratchFile

- remove `getOutputHandler` method from `KtScratchFileLanguageProvider`
- now `InlayScratchOutputHandler` is a class, it has reference to the target editor and is (temporarily) attached to the editor

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

@@ -16,8 +16,8 @@
 
 package org.jetbrains.kotlin.idea.scratch
 
-import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
@@ -33,7 +33,7 @@ import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.AnalyzingUtils
 
-class KtScratchFile(project: Project, editor: TextEditor) : ScratchFile(project, editor) {
+class KtScratchFile(project: Project, file: VirtualFile) : ScratchFile(project, file) {
     override fun getExpressions(psiFile: PsiFile): List<ScratchExpression> {
         // todo multiple expressions at one line
         val doc = PsiDocumentManager.getInstance(psiFile.project).getLastCommittedDocument(psiFile) ?: return emptyList()

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

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

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

@@ -17,16 +17,16 @@
 package org.jetbrains.kotlin.idea.scratch
 
 import com.intellij.openapi.application.runReadAction
-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.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.util.messages.Topic
+import org.jetbrains.kotlin.idea.core.util.toPsiFile
 import org.jetbrains.kotlin.idea.scratch.ui.scratchFileOptions
 
-abstract class ScratchFile(val project: Project, val editor: TextEditor) {
+abstract class ScratchFile(val project: Project, private val file: VirtualFile) {
     var replScratchExecutor: SequentialScratchExecutor? = null
     var compilingScratchExecutor: ScratchExecutor? = null
 
@@ -39,7 +39,7 @@ abstract class ScratchFile(val project: Project, val editor: TextEditor) {
     }
 
     fun getPsiFile(): PsiFile? = runReadAction {
-        PsiDocumentManager.getInstance(project).getPsiFile(editor.editor.document)
+        file.toPsiFile(project)
     }
 
     fun setModule(value: Module?) {

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

@@ -18,18 +18,17 @@ 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.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.kotlin.idea.scratch.actions.ScratchCompilationSupport
-import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandler
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandlerAdapter
 import org.jetbrains.kotlin.idea.syncPublisherWithDisposeCheck
 
 abstract class ScratchFileLanguageProvider {
-    fun newScratchFile(project: Project, editor: TextEditor): ScratchFile? {
+    fun newScratchFile(project: Project, file: VirtualFile): ScratchFile? {
-        val scratchFile = createFile(project, editor) ?: return null
+        val scratchFile = createFile(project, file) ?: return null
 
         scratchFile.replScratchExecutor = createReplExecutor(scratchFile)
         scratchFile.compilingScratchExecutor = createCompilingExecutor(scratchFile)
@@ -43,7 +42,6 @@ abstract class ScratchFileLanguageProvider {
     }
 
     private fun ScratchExecutor.addOutputHandlers() {
-        addOutputHandler(getOutputHandler())
         addOutputHandler(object : ScratchOutputHandlerAdapter() {
             override fun onStart(file: ScratchFile) {
                 ScratchCompilationSupport.start(file, this@addOutputHandlers)
@@ -55,12 +53,10 @@ abstract class ScratchFileLanguageProvider {
         })
     }
 
-    protected abstract fun createFile(project: Project, editor: TextEditor): ScratchFile?
+    protected abstract fun createFile(project: Project, file: VirtualFile): ScratchFile?
     protected abstract fun createReplExecutor(file: ScratchFile): SequentialScratchExecutor?
     protected abstract fun createCompilingExecutor(file: ScratchFile): ScratchExecutor?
 
-    abstract fun getOutputHandler(): ScratchOutputHandler
-
     companion object {
         private val EXTENSION = LanguageExtension<ScratchFileLanguageProvider>("org.jetbrains.kotlin.scratchFileLanguageProvider")
 

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

@@ -18,9 +18,11 @@ package org.jetbrains.kotlin.idea.scratch.actions
 
 import com.intellij.icons.AllIcons
 import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.fileEditor.impl.text.TextEditorProvider
 import org.jetbrains.kotlin.idea.KotlinBundle
-import org.jetbrains.kotlin.idea.scratch.ScratchFileLanguageProvider
+import org.jetbrains.kotlin.idea.scratch.attachedOutputHandler
-import org.jetbrains.kotlin.idea.scratch.getScratchFileFromSelectedEditor
+import org.jetbrains.kotlin.idea.scratch.getScratchFile
 
 class ClearScratchAction : ScratchAction(
     KotlinBundle.message("scratch.clear.button"),
@@ -28,9 +30,11 @@ class ClearScratchAction : ScratchAction(
 ) {
     override fun actionPerformed(e: AnActionEvent) {
         val project = e.project ?: return
-        val scratchFile = getScratchFileFromSelectedEditor(project) ?: return
-        val psiFile = scratchFile.getPsiFile() ?: return
 
-        ScratchFileLanguageProvider.get(psiFile.language)?.getOutputHandler()?.clear(scratchFile)
+        val selectedEditor = FileEditorManager.getInstance(project).selectedTextEditor ?: return
+        val textEditor = TextEditorProvider.getInstance().getTextEditor(selectedEditor)
+        val scratchFile = textEditor.getScratchFile() ?: return
+
+        textEditor.attachedOutputHandler?.clear(scratchFile)
     }
 }

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

@@ -23,13 +23,15 @@ import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.idea.scratch.ScratchExpression
 import org.jetbrains.kotlin.idea.scratch.ScratchFile
 
-object InlayScratchOutputHandler : ScratchOutputHandler {
+class InlayScratchOutputHandler(private val textEditor: TextEditor) : ScratchOutputHandler {
-    private const val maxLineLength = 120
+    private val maxLineLength = 120
-    private const val maxInsertOffset = 60
+    private val maxInsertOffset = 60
-    private const val minSpaceCount = 4
+    private val minSpaceCount = 4
+
+    private val toolwindowHandler = getToolwindowHandler(textEditor)
 
     override fun onStart(file: ScratchFile) {
-        getToolwindowHandler().onStart(file)
+        toolwindowHandler.onStart(file)
     }
 
     override fun handle(file: ScratchFile, expression: ScratchExpression, output: ScratchOutput) {
@@ -38,26 +40,26 @@ object InlayScratchOutputHandler : ScratchOutputHandler {
         createInlay(file, expression, output)
 
         if (output.type == ScratchOutputType.ERROR) {
-            getToolwindowHandler().handle(file, expression, output)
+            toolwindowHandler.handle(file, expression, output)
         }
     }
 
     override fun error(file: ScratchFile, message: String) {
-        getToolwindowHandler().error(file, message)
+        toolwindowHandler.error(file, message)
     }
 
     override fun onFinish(file: ScratchFile) {
-        getToolwindowHandler().onFinish(file)
+        toolwindowHandler.onFinish(file)
     }
 
     override fun clear(file: ScratchFile) {
-        clearInlays(file.editor)
+        clearInlays(textEditor)
-        getToolwindowHandler().clear(file)
+        toolwindowHandler.clear(file)
     }
 
     private fun createInlay(file: ScratchFile, expression: ScratchExpression, output: ScratchOutput) {
         TransactionGuard.submitTransaction(file.project, Runnable {
-            val editor = file.editor.editor
+            val editor = textEditor.editor
             val line = expression.lineStart
 
             val lineStartOffset = editor.document.getLineStartOffset(line)
@@ -94,12 +96,12 @@ object InlayScratchOutputHandler : ScratchOutputHandler {
 
     private fun printToToolWindow(file: ScratchFile, expression: ScratchExpression, output: ScratchOutput) {
         if (output.type != ScratchOutputType.ERROR) {
-            getToolwindowHandler().handle(file, expression, output)
+            toolwindowHandler.handle(file, expression, output)
         }
     }
 
     private fun maxLineLength(file: ScratchFile): Int {
-        val doc = file.editor.editor.document
+        val doc = textEditor.editor.document
         return file.getExpressions()
             .flatMap { it.lineStart..it.lineEnd }
             .map { doc.getLineEndOffset(it) - doc.getLineStartOffset(it) }

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

@@ -38,15 +38,15 @@ import org.jetbrains.kotlin.idea.scratch.ScratchFile
 import org.jetbrains.kotlin.psi.KtPsiFactory
 
 
-fun getToolwindowHandler(): ScratchOutputHandler {
+fun getToolwindowHandler(parentDisposable: Disposable): ScratchOutputHandler {
     return if (ApplicationManager.getApplication().isUnitTestMode) {
         TestOutputHandler
     } else {
-        ToolWindowScratchOutputHandler
+        ToolWindowScratchOutputHandler(parentDisposable)
     }
 }
 
-private object ToolWindowScratchOutputHandler : ScratchOutputHandlerAdapter() {
+private class ToolWindowScratchOutputHandler(private val parentDisposable: Disposable) : ScratchOutputHandlerAdapter() {
 
     override fun handle(file: ScratchFile, expression: ScratchExpression, output: ScratchOutput) {
         printToConsole(file) {
@@ -134,7 +134,7 @@ private object ToolWindowScratchOutputHandler : ScratchOutputHandlerAdapter() {
         val window = toolWindowManager.getToolWindow(ScratchToolWindowFactory.ID)
         ScratchToolWindowFactory().createToolWindowContent(project, window)
 
-        Disposer.register(file.editor, Disposable {
+        Disposer.register(parentDisposable, Disposable {
             window.setAvailable(false, null)
         })
 

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

@@ -17,6 +17,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.actions.KOTLIN_WORKSHEET_EXTENSION
+import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandler
 import org.jetbrains.kotlin.idea.scratch.ui.ScratchTopPanel
 import org.jetbrains.kotlin.psi.UserDataProperty
 
@@ -69,4 +70,12 @@ fun TextEditor.addScratchPanel(panel: ScratchTopPanel) {
     Disposer.register(this, panel)
 }
 
+fun TextEditor.attachOutputHandler(handler: ScratchOutputHandler) {
+    outputHandler = handler
+}
+
+val TextEditor.attachedOutputHandler: ScratchOutputHandler? get() = outputHandler
+
 private var TextEditor.scratchTopPanel: ScratchTopPanel? by UserDataProperty<TextEditor, ScratchTopPanel>(Key.create("scratch.panel"))
+
+private var TextEditor.outputHandler: ScratchOutputHandler? by UserDataProperty<TextEditor, ScratchOutputHandler>(Key.create("scratch.output.handler"))

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

@@ -21,7 +21,6 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.ProjectComponent
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.FileEditorManagerListener
-import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiManager
@@ -41,9 +40,14 @@ class ScratchFileHook(val project: Project) : ProjectComponent {
 
             val editor = getEditorWithoutScratchFile(source, file) ?: return
 
-            val scratchFile = createScratchFile(project, file, editor) ?: return
+            val scratchFile = createScratchFile(project, file) ?: return
-            val panel = ScratchTopPanel(scratchFile)
+
-            editor.addScratchPanel(panel)
+            val inlayOutputHandler = InlayScratchOutputHandler(editor)
+            scratchFile.compilingScratchExecutor?.addOutputHandler(inlayOutputHandler)
+            scratchFile.replScratchExecutor?.addOutputHandler(inlayOutputHandler)
+            editor.attachOutputHandler(inlayOutputHandler)
+
+            editor.addScratchPanel(ScratchTopPanel(scratchFile))
 
             ScratchFileAutoRunner.addListener(project, editor)
         }
@@ -59,9 +63,9 @@ class ScratchFileHook(val project: Project) : ProjectComponent {
     }
 }
 
-private fun createScratchFile(project: Project, file: VirtualFile, editor: TextEditor): ScratchFile? {
+private fun createScratchFile(project: Project, file: VirtualFile): ScratchFile? {
     val psiFile = PsiManager.getInstance(project).findFile(file) ?: return null
-    val scratchFile = ScratchFileLanguageProvider.get(psiFile.language)?.newScratchFile(project, editor) ?: return null
+    val scratchFile = ScratchFileLanguageProvider.get(psiFile.language)?.newScratchFile(project, file) ?: return null
     setupReplRestartingOutputHandler(project, scratchFile)
 
     return scratchFile