From 059d5252499bf242f97116aa7a88337ee405b00a Mon Sep 17 00:00:00 2001
From: Anton Efimchuk <anton.efimchuk@jetbrains.com>
Date: Tue, 7 May 2024 13:26:43 +0200
Subject: [PATCH] PY-69893 Fix cell outputs update after clear

GitOrigin-RevId: 42905a11209130d9fde8c8f1af6459550e4ba417
---
 .../notebooks/visualization/ui/EditorCellOutput.kt       | 4 ++--
 .../notebooks/visualization/ui/EditorCellOutputs.kt      | 2 ++
 .../notebooks/visualization/ui/NotebookEditorUiUtil.kt   | 9 +++++++++
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutput.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutput.kt
index ee837f0963d6..bbc6d27cfabb 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutput.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellOutput.kt
@@ -69,8 +69,8 @@ class EditorCellOutput internal constructor(private val editor: EditorEx, privat
 
   fun onViewportChange() {
       val component = component.mainComponent as? NotebookOutputInlayShowable ?: return
-      if (component !is Component) return
-
+      if (component !is JComponent) return
+      validateComponent(component)
       val componentRect = SwingUtilities.convertRectangle(component, component.bounds, editor.scrollPane.viewport.view)
       component.shown = editor.scrollPane.viewport.viewRect.intersects(componentRect)
   }
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 b98c17e8a005..9e8e0cd80c54 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
@@ -98,6 +98,7 @@ class EditorCellOutputs(
       ?: emptyList()
     updateData(outputDataKeys)
     recreateInlayIfNecessary()
+    onViewportChange()
   }
 
   private fun recreateInlayIfNecessary() {
@@ -116,6 +117,7 @@ class EditorCellOutputs(
     }
     else {
       inlay?.let { Disposer.dispose(it) }
+      inlay = null
     }
   }
 
diff --git a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/NotebookEditorUiUtil.kt b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/NotebookEditorUiUtil.kt
index 780e512f88c3..8cbb07423497 100644
--- a/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/NotebookEditorUiUtil.kt
+++ b/notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/NotebookEditorUiUtil.kt
@@ -7,6 +7,7 @@ import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.editor.impl.EditorComponentImpl
 import com.intellij.openapi.editor.impl.EditorEmbeddedComponentManager
 import com.intellij.openapi.editor.impl.EditorEmbeddedComponentManager.Properties.RendererFactory
+import java.awt.Container
 import java.awt.event.ComponentAdapter
 import java.awt.event.ComponentEvent
 import java.awt.event.HierarchyEvent
@@ -112,3 +113,11 @@ fun JComponent.yOffsetFromEditor(editor: Editor): Int? =
   SwingUtilities.convertPoint(this, 0, 0, editor.contentComponent).y
     .takeIf { it >= 0 }
     ?.let { it + insets.top }
+
+fun validateComponent(c: Container) {
+  var validationRoot = c
+  while (!validationRoot.isValidateRoot && validationRoot.parent != null) {
+    validationRoot = validationRoot.parent
+  }
+  validationRoot.validate()
+}