KT-32366: Add sync scroll for source and preview editor

- ^KT-32366 Fixed
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/output/PreviewEditorScratchOutputHandler.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/output/PreviewEditorScratchOutputHandler.kt
index 7351e64..5bcdafd 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/output/PreviewEditorScratchOutputHandler.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/output/PreviewEditorScratchOutputHandler.kt
@@ -33,12 +33,17 @@
 ) : ScratchOutputHandler {
     private val previewOutputBlocksManager: PreviewOutputBlocksManager = PreviewOutputBlocksManager(previewTextEditor.editor)
 
+    /**
+     * Returns pairs of line numbers which should be on the same visual positions during scrolling.
+     */
+    val sourceToPreviewAlignments: Sequence<Pair<Int, Int>> get() = previewOutputBlocksManager.alignments
+
     override fun onStart(file: ScratchFile) {
         toolwindowHandler.onStart(file)
     }
 
     override fun handle(file: ScratchFile, expression: ScratchExpression, output: ScratchOutput) {
-        printToPreviewEditor(file, expression, output)
+        printToPreviewEditor(expression, output)
     }
 
     override fun error(file: ScratchFile, message: String) {
@@ -57,7 +62,7 @@
         clearPreviewEditor()
     }
 
-    private fun printToPreviewEditor(file: ScratchFile, expression: ScratchExpression, output: ScratchOutput) {
+    private fun printToPreviewEditor(expression: ScratchExpression, output: ScratchOutput) {
         TransactionGuard.submitTransaction(previewTextEditor, Runnable {
             val targetCell = previewOutputBlocksManager.getBlock(expression) ?: previewOutputBlocksManager.addBlockToTheEnd(expression)
             targetCell.addOutput(output)
@@ -84,7 +89,7 @@
 
     val blocks: NavigableMap<ScratchExpression, OutputBlock> = TreeMap(Comparator.comparingInt { it.lineStart })
 
-    val alignments: List<Pair<Int, Int>> get() = blocks.values.map { it.sourceExpression.lineStart to it.lineStart }
+    val alignments: Sequence<Pair<Int, Int>> get() = blocks.values.asSequence().map { it.sourceExpression.lineStart to it.lineStart }
 
     fun getBlock(expression: ScratchExpression): OutputBlock? = blocks[expression]
 
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt
index 9a89deb..47f1156 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt
@@ -6,11 +6,14 @@
 package org.jetbrains.kotlin.idea.scratch.ui
 
 import com.intellij.codeInsight.daemon.DaemonCodeAnalyzer
+import com.intellij.diff.tools.util.BaseSyncScrollable
+import com.intellij.diff.tools.util.SyncScrollSupport.TwosideSyncScrollSupport
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.actionSystem.ActionToolbar
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorFactory
+import com.intellij.openapi.editor.event.VisibleAreaListener
 import com.intellij.openapi.fileEditor.FileEditor
 import com.intellij.openapi.fileEditor.FileEditorPolicy
 import com.intellij.openapi.fileEditor.FileEditorProvider
@@ -51,7 +54,7 @@
 
 class KtScratchFileEditorWithPreview private constructor(
     val scratchFile: ScratchFile,
-    sourceTextEditor: TextEditor,
+    private val sourceTextEditor: TextEditor,
     private val previewTextEditor: TextEditor
 ) : TextEditorWithPreview(sourceTextEditor, previewTextEditor), TextEditor {
 
@@ -73,9 +76,39 @@
         scratchFile.compilingScratchExecutor?.addOutputHandler(commonPreviewOutputHandler)
         scratchFile.replScratchExecutor?.addOutputHandler(commonPreviewOutputHandler)
 
+        configureSyncScrollForSourceAndPreview()
+
         ScratchFileAutoRunner.addListener(scratchFile.project, sourceTextEditor)
     }
 
+    private fun configureSyncScrollForSourceAndPreview() {
+        val sourceEditor = sourceTextEditor.editor
+        val previewEditor = previewTextEditor.editor
+
+        val scrollable = object : BaseSyncScrollable() {
+            override fun processHelper(helper: ScrollHelper) {
+                if (!helper.process(0, 0)) return
+
+                val alignments = previewEditorScratchOutputHandler.sourceToPreviewAlignments
+
+                for ((fromSource, fromPreview) in alignments) {
+                    if (!helper.process(fromSource, fromPreview)) return
+                    if (!helper.process(fromSource, fromPreview)) return
+                }
+
+                helper.process(sourceEditor.document.lineCount, previewEditor.document.lineCount)
+            }
+
+            override fun isSyncScrollEnabled(): Boolean = true
+        }
+
+        val scrollSupport = TwosideSyncScrollSupport(listOf(sourceEditor, previewEditor), scrollable)
+        val listener = VisibleAreaListener { e -> scrollSupport.visibleAreaChanged(e) }
+
+        sourceEditor.scrollingModel.addVisibleAreaListener(listener)
+        previewEditor.scrollingModel.addVisibleAreaListener(listener)
+    }
+
     override fun dispose() {
         scratchFile.replScratchExecutor?.stop()
         scratchFile.compilingScratchExecutor?.stop()