KT-33585: Add synchronization between scratch editor and preview

- fix problem with not working shortcuts in the preview window by passing more params into view editor creation

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

@@ -7,15 +7,13 @@ package org.jetbrains.kotlin.idea.scratch.output
 
 import com.intellij.diff.util.DiffUtil
 import com.intellij.openapi.Disposable
- import com.intellij.openapi.application.TransactionGuard
+import com.intellij.openapi.application.TransactionGuard
 import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.FoldRegion
 import com.intellij.openapi.editor.FoldingModel
-import com.intellij.openapi.editor.markup.HighlighterLayer
+import com.intellij.openapi.editor.markup.*
-import com.intellij.openapi.editor.markup.HighlighterTargetArea
-import com.intellij.openapi.editor.markup.MarkupModel
 import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.idea.scratch.ScratchExpression
 import org.jetbrains.kotlin.idea.scratch.ScratchFile
@@ -71,6 +69,13 @@ class PreviewEditorScratchOutputHandler(
 
 private val ScratchExpression.height: Int get() = lineEnd - lineStart + 1
 
+interface ScratchOutputBlock {
+    val sourceExpression: ScratchExpression
+    val lineStart: Int
+    val lineEnd: Int
+    fun addOutput(output: ScratchOutput)
+}
+
 class PreviewOutputBlocksManager(editor: Editor) {
     private val targetDocument: Document = editor.document
     private val foldingModel: FoldingModel = editor.foldingModel
@@ -80,9 +85,9 @@ class PreviewOutputBlocksManager(editor: Editor) {
 
     fun computeSourceToPreviewAlignments(): List<Pair<Int, Int>> = blocks.values.map { it.sourceExpression.lineStart to it.lineStart }
 
-    fun getBlock(expression: ScratchExpression): OutputBlock? = blocks[expression]
+    fun getBlock(expression: ScratchExpression): ScratchOutputBlock? = blocks[expression]
 
-    fun addBlockToTheEnd(expression: ScratchExpression): OutputBlock = OutputBlock(expression).also {
+    fun addBlockToTheEnd(expression: ScratchExpression): ScratchOutputBlock = OutputBlock(expression).also {
         if (blocks.putIfAbsent(expression, it) != null) {
             error("There is already a cell for $expression!")
         }
@@ -97,26 +102,34 @@ class PreviewOutputBlocksManager(editor: Editor) {
         }
     }
 
-    inner class OutputBlock(val sourceExpression: ScratchExpression) {
+    private inner class OutputBlock(override val sourceExpression: ScratchExpression) : ScratchOutputBlock {
         private val outputs: MutableList<ScratchOutput> = mutableListOf()
 
-        var lineStart: Int = computeCellLineStart(sourceExpression)
+        override var lineStart: Int = computeCellLineStart(sourceExpression)
             private set
 
-        val lineEnd: Int get() = lineStart + countNewLines(outputs)
+        override val lineEnd: Int get() = lineStart + countNewLines(outputs)
-        val height: Int get() = lineEnd - lineStart + 1
 
+        val height: Int get() = lineEnd - lineStart + 1
         private var foldRegion: FoldRegion? = null
 
-        fun addOutput(output: ScratchOutput) {
+        override fun addOutput(output: ScratchOutput) {
             printAndSaveOutput(output)
 
+            blocks.lowerEntry(sourceExpression)?.value?.updateFolding()
             blocks.tailMap(sourceExpression).values.forEach {
                 it.recalculatePosition()
                 it.updateFolding()
             }
         }
 
+        /**
+         * We want to make sure that changes in document happen in single edit, because if they are not,
+         * listeners may see inconsistent document, which may cause troubles if they will try to highlight it
+         * in some way. That's why it is important that [insertStringAtLine] does only one insert in the document,
+         * and [output] is inserted into the [outputs] before the edits, so [OutputBlock] can correctly see
+         * all its output expressions and highlight the whole block.
+         */
         private fun printAndSaveOutput(output: ScratchOutput) {
             val beforeAdding = lineEnd
             val currentOutputStartLine = if (outputs.isEmpty()) lineStart else beforeAdding + 1
@@ -129,21 +142,7 @@ class PreviewOutputBlocksManager(editor: Editor) {
                 }
             }
 
-            val insertedTextStart = targetDocument.getLineStartOffset(currentOutputStartLine)
+            markupModel.highlightLines(currentOutputStartLine, lineEnd, getAttributesForOutputType(output.type))
-            val insertedTextEnd = targetDocument.getLineEndOffset(lineEnd)
-            colorRange(insertedTextStart, insertedTextEnd, output.type)
-        }
-
-        private fun colorRange(startOffset: Int, endOffset: Int, outputType: ScratchOutputType) {
-            val textAttributes = getAttributesForOutputType(outputType)
-
-            markupModel.addRangeHighlighter(
-                startOffset,
-                endOffset,
-                HighlighterLayer.SYNTAX,
-                textAttributes,
-                HighlighterTargetArea.EXACT_RANGE
-            )
         }
 
         private fun recalculatePosition() {
@@ -181,6 +180,8 @@ class PreviewOutputBlocksManager(editor: Editor) {
         val compensation = max(differenceBetweenSourceAndOutputHeight, 0)
         return previous.lineEnd + compensation + distanceBetweenSources
     }
+
+    fun getBlockAtLine(line: Int): ScratchOutputBlock? = blocks.values.find { line in it.lineStart..it.lineEnd }
 }
 
 private fun countNewLines(list: List<ScratchOutput>) = list.sumBy { StringUtil.countNewLines(it.text) } + max(list.size - 1, 0)
@@ -188,10 +189,29 @@ private fun countNewLines(list: List<ScratchOutput>) = list.sumBy { StringUtil.c
 private fun Document.getLineContent(lineNumber: Int) =
     DiffUtil.getLinesContent(this, lineNumber, lineNumber + 1).toString()
 
-fun Document.insertStringAtLine(lineNumber: Int, text: String) {
+private fun Document.insertStringAtLine(lineNumber: Int, text: String) {
-    while (DiffUtil.getLineCount(this) <= lineNumber) {
+    val missingNewLines = lineNumber - (DiffUtil.getLineCount(this) - 1)
-        insertString(textLength, "\n")
+    if (missingNewLines > 0) {
+        insertString(textLength, "${"\n".repeat(missingNewLines)}$text")
+    } else {
+        insertString(getLineStartOffset(lineNumber), text)
     }
-
+}
-    insertString(getLineStartOffset(lineNumber), text)
+
+fun MarkupModel.highlightLines(
+    from: Int,
+    to: Int,
+    attributes: TextAttributes,
+    targetArea: HighlighterTargetArea = HighlighterTargetArea.EXACT_RANGE
+): RangeHighlighter {
+    val fromOffset = document.getLineStartOffset(from)
+    val toOffset = document.getLineEndOffset(to)
+
+    return addRangeHighlighter(
+        fromOffset,
+        toOffset,
+        HighlighterLayer.CARET_ROW,
+        attributes,
+        targetArea
+    )
 }

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

@@ -13,7 +13,9 @@ import com.intellij.openapi.actionSystem.*
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorFactory
+import com.intellij.openapi.editor.EditorKind
 import com.intellij.openapi.editor.event.VisibleAreaListener
+import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorPolicy
 import com.intellij.openapi.fileEditor.FileEditorProvider
@@ -28,6 +30,7 @@ import com.intellij.pom.Navigatable
 import com.intellij.psi.PsiManager
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.KotlinBundle
+import org.jetbrains.kotlin.idea.core.util.getLineNumber
 import org.jetbrains.kotlin.idea.scratch.*
 import org.jetbrains.kotlin.idea.scratch.output.*
 import org.jetbrains.kotlin.psi.UserDataProperty
@@ -55,11 +58,13 @@ class KtScratchFileEditorProvider : FileEditorProvider, DumbAware {
 
 class KtScratchFileEditorWithPreview private constructor(
     val scratchFile: ScratchFile,
-    private val sourceTextEditor: TextEditor,
+    sourceTextEditor: TextEditor,
     private val previewTextEditor: TextEditor
-) : TextEditorWithPreview(sourceTextEditor, previewTextEditor), TextEditor {
+) : TextEditorWithPreview(sourceTextEditor, previewTextEditor), TextEditor, ScratchEditorLinesTranslator {
 
-    private val previewOutputManager: PreviewOutputBlocksManager = PreviewOutputBlocksManager(previewTextEditor.editor)
+    private val sourceEditor = sourceTextEditor.editor as EditorEx
+    private val previewEditor = previewTextEditor.editor as EditorEx
+    private val previewOutputManager: PreviewOutputBlocksManager = PreviewOutputBlocksManager(previewEditor)
 
     private val toolWindowHandler: ScratchOutputHandler = requestToolWindowHandler()
     private val inlayScratchOutputHandler = InlayScratchOutputHandler(sourceTextEditor, toolWindowHandler)
@@ -83,14 +88,26 @@ class KtScratchFileEditorWithPreview private constructor(
         scratchFile.replScratchExecutor?.addOutputHandler(commonPreviewOutputHandler)
 
         configureSyncScrollForSourceAndPreview()
+        configureSyncHighlighting(sourceEditor, previewEditor, translator = this)
 
         ScratchFileAutoRunner.addListener(scratchFile.project, sourceTextEditor)
     }
 
-    private fun configureSyncScrollForSourceAndPreview() {
+    override fun previewLineToSourceLines(previewLine: Int): Pair<Int, Int>? {
-        val sourceEditor = sourceTextEditor.editor
+        val expressionUnderCaret = scratchFile.getExpressionAtLine(previewLine) ?: return null
-        val previewEditor = previewTextEditor.editor
+        val outputBlock = previewOutputManager.getBlock(expressionUnderCaret) ?: return null
 
+        return outputBlock.lineStart to outputBlock.lineEnd
+    }
+
+    override fun sourceLineToPreviewLines(sourceLine: Int): Pair<Int, Int>? {
+        val block = previewOutputManager.getBlockAtLine(sourceLine) ?: return null
+        if (!block.sourceExpression.linesInformationIsCorrect()) return null
+
+        return block.sourceExpression.lineStart to block.sourceExpression.lineEnd
+    }
+
+    private fun configureSyncScrollForSourceAndPreview() {
         val scrollable = object : BaseSyncScrollable() {
             override fun processHelper(helper: ScrollHelper) {
                 if (!helper.process(0, 0)) return
@@ -189,7 +206,7 @@ class KtScratchFileEditorWithPreview private constructor(
             val mainEditor = textEditorProvider.createEditor(scratchFile.project, scratchFile.file) as TextEditor
             val editorFactory = EditorFactory.getInstance()
 
-            val viewer = editorFactory.createViewer(editorFactory.createDocument(""))
+            val viewer = editorFactory.createViewer(editorFactory.createDocument(""), scratchFile.project, EditorKind.PREVIEW)
             Disposer.register(mainEditor, Disposable { editorFactory.releaseEditor(viewer) })
 
             val previewEditor = textEditorProvider.getTextEditor(viewer)
@@ -266,3 +283,12 @@ private class LayoutDependantOutputHandler(
             else -> previewOutputHandler
         }
 }
+
+/**
+ * Checks if [ScratchExpression.element] is actually starts at the [ScratchExpression.lineStart]
+ * and ends at the [ScratchExpression.lineEnd].
+ */
+private fun ScratchExpression.linesInformationIsCorrect(): Boolean {
+    if (!element.isValid) return false
+    return element.getLineNumber(start = true) == lineStart && element.getLineNumber(start = false) == lineEnd
+}

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

@@ -0,0 +1,234 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.scratch.ui
+
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.colors.EditorColors
+import com.intellij.openapi.editor.event.CaretEvent
+import com.intellij.openapi.editor.event.CaretListener
+import com.intellij.openapi.editor.event.DocumentEvent
+import com.intellij.openapi.editor.event.DocumentListener
+import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.openapi.editor.ex.FocusChangeListener
+import com.intellij.openapi.editor.markup.HighlighterTargetArea
+import com.intellij.openapi.editor.markup.RangeHighlighter
+import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.psi.PsiDocumentManager
+import org.jetbrains.kotlin.idea.scratch.output.highlightLines
+
+interface ScratchEditorLinesTranslator {
+    fun previewLineToSourceLines(previewLine: Int): Pair<Int, Int>?
+    fun sourceLineToPreviewLines(sourceLine: Int): Pair<Int, Int>?
+}
+
+fun configureSyncHighlighting(sourceEditor: EditorEx, previewEditor: EditorEx, translator: ScratchEditorLinesTranslator) {
+    configureExclusiveCaretRowHighlighting(sourceEditor, previewEditor)
+    configureSourceAndPreviewHighlighting(sourceEditor, previewEditor, translator)
+}
+
+private fun configureSourceAndPreviewHighlighting(
+    sourceEditor: EditorEx,
+    previewEditor: EditorEx,
+    translator: ScratchEditorLinesTranslator
+) {
+    val syncHighlighter = ScratchEditorSyncHighlighter.create(sourceEditor, previewEditor, translator)
+
+    configureHighlightUpdateOnDocumentChange(sourceEditor, previewEditor, syncHighlighter)
+    configureSourceToPreviewHighlighting(sourceEditor, syncHighlighter)
+    configurePreviewToSourceHighlighting(previewEditor, syncHighlighter)
+}
+
+/**
+ * Configures editors such that only one of them have caret row highlighting enabled.
+ */
+private fun configureExclusiveCaretRowHighlighting(sourceEditor: EditorEx, previewEditor: EditorEx) {
+    val exclusiveCaretHighlightingListener = object : FocusChangeListener {
+        override fun focusLost(editor: Editor) {}
+
+        override fun focusGained(editor: Editor) {
+            sourceEditor.settings.isCaretRowShown = false
+            previewEditor.settings.isCaretRowShown = false
+
+            editor.settings.isCaretRowShown = true
+        }
+    }
+
+    sourceEditor.addFocusListener(exclusiveCaretHighlightingListener)
+    previewEditor.addFocusListener(exclusiveCaretHighlightingListener)
+}
+
+/**
+ * When source or preview documents change, we need to update highlighting, because
+ * expression output may become bigger.
+ *
+ * We can do that only when document is fully committed, so [ScratchFile.getExpressions] will return correct expressions
+ * with correct PSIs.
+ */
+private fun configureHighlightUpdateOnDocumentChange(
+    sourceEditor: EditorEx,
+    previewEditor: EditorEx,
+    highlighter: ScratchEditorSyncHighlighter
+) {
+    val updateHighlightOnDocumentChangeListener = object : DocumentListener {
+        override fun documentChanged(event: DocumentEvent) {
+            PsiDocumentManager.getInstance(sourceEditor.project!!).performWhenAllCommitted {
+                highlighter.highlightByCurrentlyFocusedEditor()
+            }
+        }
+    }
+
+    previewEditor.document.addDocumentListener(updateHighlightOnDocumentChangeListener)
+    sourceEditor.document.addDocumentListener(updateHighlightOnDocumentChangeListener)
+}
+
+/**
+ * When caret in [sourceEditor] is moved, highlight is recalculated.
+ *
+ * When focus is switched to the [sourceEditor], highlight is recalculated,
+ * because it is possible to switch focus without changing cursor position,
+ * which would lead to the outdated highlighting.
+ */
+private fun configureSourceToPreviewHighlighting(sourceEditor: EditorEx, highlighter: ScratchEditorSyncHighlighter) {
+    sourceEditor.caretModel.addCaretListener(object : CaretListener {
+        override fun caretPositionChanged(event: CaretEvent) {
+            highlighter.highlightPreviewBySource()
+        }
+    })
+
+    sourceEditor.addFocusListener(object : FocusChangeListener {
+        override fun focusLost(editor: Editor) {}
+
+        override fun focusGained(editor: Editor) {
+            highlighter.highlightPreviewBySource()
+        }
+    })
+}
+
+/**
+ * When caret in [previewEditor] is moved, highlight is recalculated.
+ *
+ * When focus is switched to the [previewEditor], highlight is recalculated,
+ * because it is possible to switch focus without changing cursor position,
+ * which would lead to the outdated highlighting.
+ */
+private fun configurePreviewToSourceHighlighting(previewEditor: EditorEx, highlighter: ScratchEditorSyncHighlighter) {
+    previewEditor.caretModel.addCaretListener(object : CaretListener {
+        override fun caretPositionChanged(event: CaretEvent) {
+            highlighter.highlightSourceByPreview()
+        }
+    })
+
+    previewEditor.addFocusListener(object : FocusChangeListener {
+        override fun focusLost(editor: Editor) {}
+
+        override fun focusGained(editor: Editor) {
+            highlighter.highlightSourceByPreview()
+        }
+    })
+}
+
+private class ScratchEditorsState(private val sourceEditor: EditorEx, private val previewEditor: EditorEx) : FocusChangeListener {
+    private var lastFocusedEditor: Editor = sourceEditor
+
+    enum class FocusedEditor {
+        SOURCE, PREVIEW
+    }
+
+    init {
+        sourceEditor.addFocusListener(this)
+        previewEditor.addFocusListener(this)
+    }
+
+    val sourceEditorCaretLine: Int? get() = sourceEditor.caretModel.allCarets.singleOrNull()?.logicalPosition?.line
+    val previewEditorCaretLine: Int? get() = previewEditor.caretModel.allCarets.singleOrNull()?.logicalPosition?.line
+    val focusedEditor: FocusedEditor get() = if (lastFocusedEditor === sourceEditor) FocusedEditor.SOURCE else FocusedEditor.PREVIEW
+
+    override fun focusLost(editor: Editor) {}
+
+    override fun focusGained(editor: Editor) {
+        lastFocusedEditor = editor
+    }
+}
+
+private class ScratchEditorSyncHighlighter private constructor(
+    private val state: ScratchEditorsState,
+    private val sourceHighlighter: EditorLinesHighlighter,
+    private val previewHighlighter: EditorLinesHighlighter,
+    private val translator: ScratchEditorLinesTranslator
+) {
+    fun highlightSourceByPreview() {
+        clearAllHighlights()
+
+        state.previewEditorCaretLine?.let(::highlightSourceByPreviewLine)
+    }
+
+    fun highlightPreviewBySource() {
+        clearAllHighlights()
+
+        state.sourceEditorCaretLine?.let(::highlightPreviewBySourceLine)
+    }
+
+    fun highlightByCurrentlyFocusedEditor() {
+        when (state.focusedEditor) {
+            ScratchEditorsState.FocusedEditor.SOURCE -> highlightPreviewBySource()
+            ScratchEditorsState.FocusedEditor.PREVIEW -> highlightSourceByPreview()
+        }
+    }
+
+    private fun highlightSourceByPreviewLine(selectedPreviewLine: Int) {
+        val (from, to) = translator.sourceLineToPreviewLines(selectedPreviewLine) ?: return
+
+        sourceHighlighter.highlightLines(from, to)
+    }
+
+    private fun highlightPreviewBySourceLine(selectedSourceLine: Int) {
+        val (from, to) = translator.previewLineToSourceLines(selectedSourceLine) ?: return
+
+        previewHighlighter.highlightLines(from, to)
+    }
+
+    private fun clearAllHighlights() {
+        sourceHighlighter.clearHighlights()
+        previewHighlighter.clearHighlights()
+    }
+
+    companion object {
+        fun create(
+            sourceEditor: EditorEx,
+            previewEditor: EditorEx,
+            translator: ScratchEditorLinesTranslator
+        ): ScratchEditorSyncHighlighter {
+            return ScratchEditorSyncHighlighter(
+                state = ScratchEditorsState(sourceEditor, previewEditor),
+                sourceHighlighter = EditorLinesHighlighter(sourceEditor),
+                previewHighlighter = EditorLinesHighlighter(previewEditor),
+                translator = translator
+            )
+        }
+    }
+}
+
+private class EditorLinesHighlighter(private val targetEditor: Editor) {
+    private var activeHighlight: RangeHighlighter? = null
+
+    fun clearHighlights() {
+        activeHighlight?.let(targetEditor.markupModel::removeHighlighter)
+        activeHighlight = null
+    }
+
+    fun highlightLines(lineStart: Int, lineEnd: Int) {
+        clearHighlights()
+
+        val highlightColor = targetEditor.colorsScheme.getColor(EditorColors.CARET_ROW_COLOR) ?: return
+
+        activeHighlight = targetEditor.markupModel.highlightLines(
+            lineStart,
+            lineEnd,
+            TextAttributes().apply { backgroundColor = highlightColor },
+            HighlighterTargetArea.LINES_IN_RANGE
+        )
+    }
+}