From 05d99a3f9ef01ecbe2ffe714c91e7abbfaa4290c Mon Sep 17 00:00:00 2001
From: Anton Efimchuk <anton.efimchuk@jetbrains.com>
Date: Thu, 1 Aug 2024 11:16:35 +0200
Subject: [PATCH] PY-74339 [Jupyter] Do not report not cell produced inlays in
 offset check

GitOrigin-RevId: 4c85483e4410f1b7aa6ec5c2f1e88fc881ab6697
---
 .../notebooks/visualization/NotebookCellInlayManager.kt  | 6 +++++-
 .../ui/ControllerEditorCellViewComponent.kt              | 5 +++++
 .../notebooks/visualization/ui/EditorCellOutputs.kt      | 4 ++++
 .../plugins/notebooks/visualization/ui/EditorCellView.kt | 5 +++++
 .../visualization/ui/EditorCellViewComponent.kt          | 9 +++++++++
 5 files changed, 28 insertions(+), 1 deletion(-)

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 13c1f1349779..45c097db8709 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt
@@ -323,8 +323,12 @@ class NotebookCellInlayManager private constructor(
         add(editor.document.getLineEndOffset(cell.interval.lines.last))
       }
     }
-    val wronglyPlacedInlays = editor.inlayModel.getBlockElementsInRange(0, editor.document.textLength)
+
+    val wronglyPlacedInlays = _cells.asSequence()
+      .mapNotNull { it.view }
+      .flatMap { it.getInlays() }
       .filter { it.offset !in inlaysOffsets }
+      .toSet()
     if (wronglyPlacedInlays.isNotEmpty()) {
       thisLogger().error("Expected offsets: $inlaysOffsets. Wrongly placed offsets: ${wronglyPlacedInlays.map { it.offset }} of inlays $wronglyPlacedInlays, for file = '${editor.virtualFile?.name}'")
     }
diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/ControllerEditorCellViewComponent.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/ControllerEditorCellViewComponent.kt
index 0bcbb65988d0..f479137f7990 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/ControllerEditorCellViewComponent.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/ControllerEditorCellViewComponent.kt
@@ -3,6 +3,7 @@ package org.jetbrains.plugins.notebooks.visualization.ui
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.FoldRegion
+import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.editor.ex.FoldingModelEx
 import com.intellij.openapi.util.Disposer
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayController
@@ -66,4 +67,8 @@ class ControllerEditorCellViewComponent(
 
   private fun createFoldRegion(foldingModel: FoldingModelEx, regionToFold: IntRange): FoldRegion? =
     foldingModel.createFoldRegion(regionToFold.first, regionToFold.last, "", null, true)
+
+  override fun doGetInlays(): Sequence<Inlay<*>> {
+    return sequenceOf(controller.inlay)
+  }
 }
\ No newline at end of file
diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutputs.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutputs.kt
index dd45fbad2309..8ab6f3df6dd2 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutputs.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutputs.kt
@@ -288,6 +288,10 @@ class EditorCellOutputs(
       }
     }
   }
+
+  override fun doGetInlays(): Sequence<Inlay<*>> {
+    return inlay?.let { sequenceOf(it) } ?: emptySequence()
+  }
 }
 
 private var JComponent.outputComponentFactory: NotebookOutputComponentFactory<*, *>? by SwingClientProperty("outputComponentFactory")
\ No newline at end of file
diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellView.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellView.kt
index 8097e757107b..aece793b4b53 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellView.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellView.kt
@@ -10,6 +10,7 @@ import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorKind
 import com.intellij.openapi.editor.FoldRegion
+import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.editor.VisualPosition
 import com.intellij.openapi.editor.ex.RangeHighlighterEx
 import com.intellij.openapi.editor.impl.EditorImpl
@@ -456,6 +457,10 @@ class EditorCellView(
       }
   }
 
+  override fun doGetInlays(): Sequence<Inlay<*>> {
+    return controllers.map { it.inlay }.asSequence()
+  }
+
   companion object {
     val WAS_FOLDED_IN_RENDERED_STATE_KEY = Key<Boolean>("jupyter.markdown.folding.was.rendered")
   }
diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellViewComponent.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellViewComponent.kt
index 963d43c059f1..e132c75faac6 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellViewComponent.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellViewComponent.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.plugins.notebooks.visualization.ui
 
+import com.intellij.openapi.editor.Inlay
 import org.jetbrains.plugins.notebooks.visualization.UpdateContext
 import java.awt.Rectangle
 
@@ -39,4 +40,12 @@ abstract class EditorCellViewComponent {
       it.updateCellFolding(updateContext)
     }
   }
+
+  fun getInlays(): Sequence<Inlay<*>> {
+    return doGetInlays() + children.asSequence().flatMap { it.getInlays() }
+  }
+
+  open fun doGetInlays(): Sequence<Inlay<*>> {
+    return emptySequence()
+  }
 }
\ No newline at end of file