From 6a1b8a69399307871c61d561e35aa4fc15eb6aa8 Mon Sep 17 00:00:00 2001
From: Igor Slobodskov <igor.slobodskov@jetbrains.com>
Date: Mon, 10 Oct 2022 15:12:32 +0200
Subject: [PATCH] [DS-3533] add NotebookIntervalPointer to
 NotebookCellInlayManager

GitOrigin-RevId: f04bb7593b62336a1f90332dd6d82b5f0ad4cde8
---
 .../visualization/NotebookCellInlayManager.kt | 20 +++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt
index 46b56cc2043f..20e62b31c74e 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt
@@ -24,6 +24,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.util.Processor
 import com.intellij.util.SmartList
 import com.intellij.util.asSafely
+import com.intellij.util.containers.SmartHashSet
 import com.intellij.util.ui.update.MergingUpdateQueue
 import com.intellij.util.ui.update.Update
 import org.jetbrains.annotations.TestOnly
@@ -64,6 +65,14 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
     }
   }
 
+  fun update(pointers: Collection<NotebookIntervalPointer>) {
+    updateQueue.queue(UpdateInlaysTask(this, pointers = pointers))
+  }
+
+  fun update(pointer: NotebookIntervalPointer) {
+    updateQueue.queue(UpdateInlaysTask(this, pointers = SmartList(pointer)))
+  }
+
   fun update(interval: NotebookCellLines.Interval) {
     update(interval.lines)
   }
@@ -398,10 +407,16 @@ private object NotebookCellHighlighterRenderer : CustomHighlighterRenderer {
   }
 }
 
-private class UpdateInlaysTask(private val manager: NotebookCellInlayManager, lines: IntRange) : Update(Any()) {
-  private val linesList = SmartList<IntRange>(lines)
+private class UpdateInlaysTask(private val manager: NotebookCellInlayManager,
+                               lines: IntRange? = null,
+                               pointers: Collection<NotebookIntervalPointer>? = null) : Update(Any()) {
+  private val linesList = lines?.let{ SmartList<IntRange>(lines) } ?: SmartList()
+  private val pointersSet = pointers?.let{ SmartHashSet(pointers) } ?: SmartHashSet()
 
   override fun run() {
+    val pointersLines = pointersSet.mapNotNull { it.get()?.lines }.sortedBy { it.first }
+    linesList.mergeAndJoinIntersections(pointersLines)
+
     for (lines in linesList) {
       manager.updateImmediately(lines)
     }
@@ -410,6 +425,7 @@ private class UpdateInlaysTask(private val manager: NotebookCellInlayManager, li
   override fun canEat(update: Update): Boolean {
     update as UpdateInlaysTask
     linesList.mergeAndJoinIntersections(update.linesList)
+    pointersSet.addAll(update.pointersSet)
     return true
   }
 }
\ No newline at end of file