Implement ComposeDesktopEntryPoint.captureContentToImage (#3155)

Author: m-sasha

compose/ui/ui/api/desktop/ui.api

@@ -136,6 +136,7 @@ public final class androidx/compose/ui/ComposableSingletons$ImageComposeScene_sk
 }
 
 public abstract interface class androidx/compose/ui/ComposeDesktopEntryPoint {
+	public abstract fun captureContentToImage ()Ljava/awt/image/BufferedImage;
 	public abstract fun getSemanticsOwners ()Ljava/util/Collection;
 }
 
@@ -564,6 +565,7 @@ public final class androidx/compose/ui/awt/ComposeDialog : javax/swing/JDialog,
 	public fun addMouseListener (Ljava/awt/event/MouseListener;)V
 	public fun addMouseMotionListener (Ljava/awt/event/MouseMotionListener;)V
 	public fun addMouseWheelListener (Ljava/awt/event/MouseWheelListener;)V
+	public fun captureContentToImage ()Ljava/awt/image/BufferedImage;
 	public fun dispose ()V
 	public final fun getCompositionLocalContext ()Landroidx/compose/runtime/CompositionLocalContext;
 	public fun getPreferredSize ()Ljava/awt/Dimension;
@@ -595,6 +597,7 @@ public final class androidx/compose/ui/awt/ComposePanel : javax/swing/JLayeredPa
 	public fun add (Ljava/awt/Component;)Ljava/awt/Component;
 	public fun addFocusListener (Ljava/awt/event/FocusListener;)V
 	public fun addNotify ()V
+	public fun captureContentToImage ()Ljava/awt/image/BufferedImage;
 	public fun getFocusTraversalKeysEnabled ()Z
 	public fun getMaximumSize ()Ljava/awt/Dimension;
 	public fun getMinimumSize ()Ljava/awt/Dimension;
@@ -635,6 +638,7 @@ public final class androidx/compose/ui/awt/ComposeWindow : javax/swing/JFrame, a
 	public fun addMouseListener (Ljava/awt/event/MouseListener;)V
 	public fun addMouseMotionListener (Ljava/awt/event/MouseMotionListener;)V
 	public fun addMouseWheelListener (Ljava/awt/event/MouseWheelListener;)V
+	public fun captureContentToImage ()Ljava/awt/image/BufferedImage;
 	public fun dispose ()V
 	public final fun getCompositionLocalContext ()Landroidx/compose/runtime/CompositionLocalContext;
 	public final fun getPlacement ()Landroidx/compose/ui/window/WindowPlacement;

compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/ComposeDesktopEntryPoint.kt

@@ -18,6 +18,7 @@ package androidx.compose.ui
 
 import androidx.compose.runtime.tooling.ComposeToolingApi
 import androidx.compose.ui.semantics.SemanticsOwner
+import java.awt.image.BufferedImage
 
 /**
  * The interface for classes that are an entry point for using Compose on the desktop.
@@ -33,4 +34,13 @@ interface ComposeDesktopEntryPoint {
      * changes.
      */
     val semanticsOwners: Collection<SemanticsOwner>
+
+    /**
+     * Captures the content of this entry point into an image.
+     *
+     * Returns `null` if the entry point is not in a state where it has visual content yet.
+     *
+     * May be called only on the event dispatch thread.
+     */
+    fun captureContentToImage(): BufferedImage?
 }

compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeDialog.desktop.kt

@@ -44,6 +44,7 @@ import java.awt.Window
 import java.awt.event.MouseListener
 import java.awt.event.MouseMotionListener
 import java.awt.event.MouseWheelListener
+import java.awt.image.BufferedImage
 import java.util.*
 import javax.swing.JDialog
 import kotlin.coroutines.CoroutineContext
@@ -216,18 +217,6 @@ class ComposeDialog : JDialog, ComposeDesktopEntryPoint {
 
     private val undecoratedWindowResizer = UndecoratedWindowResizer(this)
 
-    /**
-     * Returns the [SemanticsOwner]s corresponding to the roots of the semantics trees in this
-     * [ComposeDialog].
-     *
-     * This is backed by Snapshot state, so reading this property in a restartable function (e.g., a
-     * composable function) will cause the function to restart when the set of semantics owners
-     * changes.
-     */
-    @ComposeToolingApi
-    override val semanticsOwners: Collection<SemanticsOwner>
-        get() = composePanel.semanticsOwners
-
     override fun add(component: Component) = composePanel.add(component)
 
     override fun remove(component: Component) = composePanel.remove(component)
@@ -455,4 +444,28 @@ class ComposeDialog : JDialog, ComposeDesktopEntryPoint {
     internal fun measureContent(constraints: Constraints): IntSize {
         return composePanel.measureContent(constraints)
     }
+
+    /**
+     * Returns the [SemanticsOwner]s corresponding to the roots of the semantics trees in this
+     * [ComposeDialog].
+     *
+     * This is backed by Snapshot state, so reading this property in a restartable function (e.g., a
+     * composable function) will cause the function to restart when the set of semantics owners
+     * changes.
+     */
+    @ComposeToolingApi
+    override val semanticsOwners: Collection<SemanticsOwner>
+        get() = composePanel.semanticsOwners
+
+    /**
+     * Captures the content of this dialog into an image.
+     *
+     * Returns `null` if the dialog has not been made visible yet.
+     *
+     * May be called only on the event dispatching thread.
+     */
+    @ComposeToolingApi
+    override fun captureContentToImage(): BufferedImage? {
+        return composePanel.captureContentToImage()
+    }
 }

compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposePanel.desktop.kt

@@ -46,6 +46,7 @@ import java.awt.FocusTraversalPolicy
 import java.awt.Window
 import java.awt.event.FocusEvent
 import java.awt.event.FocusListener
+import java.awt.image.BufferedImage
 import java.util.*
 import javax.swing.JLayeredPane
 import javax.swing.SwingUtilities
@@ -333,18 +334,6 @@ class ComposePanel @ExperimentalComposeUiApi constructor(
             _composeContainer?.windowContainer = value
         }
 
-    /**
-     * Returns the [SemanticsOwner]s corresponding to the roots of the semantics trees in this
-     * [ComposePanel].
-     *
-     * This is backed by Snapshot state, so reading this property in a restartable function (e.g., a
-     * composable function) will cause the function to restart when the set of semantics owners
-     * changes.
-     */
-    @ComposeToolingApi
-    override val semanticsOwners: Collection<SemanticsOwner>
-        get() = _composeContainer?.semanticsOwners ?: emptyList()
-
     // Needed to preserve binary compatibility
     @Suppress("RedundantOverride")
     override fun add(component: Component): Component = super.add(component)
@@ -520,4 +509,28 @@ class ComposePanel @ExperimentalComposeUiApi constructor(
             field = value
             _composeContainer?.showLayoutBounds = value
         }
+
+    /**
+     * Returns the [SemanticsOwner]s corresponding to the roots of the semantics trees in this
+     * [ComposePanel].
+     *
+     * This is backed by Snapshot state, so reading this property in a restartable function (e.g., a
+     * composable function) will cause the function to restart when the set of semantics owners
+     * changes.
+     */
+    @ComposeToolingApi
+    override val semanticsOwners: Collection<SemanticsOwner>
+        get() = _composeContainer?.semanticsOwners ?: emptyList()
+
+    /**
+     * Captures the content of this panel into an image.
+     *
+     * Returns `null` if the panel has not been made visible yet.
+     *
+     * May be called only on the event dispatching thread.
+     */
+    @ComposeToolingApi
+    override fun captureContentToImage(): BufferedImage? {
+        return _composeContainer?.captureContentToImage()
+    }
 }

compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeWindow.desktop.kt

@@ -41,6 +41,7 @@ import java.awt.GraphicsConfiguration
 import java.awt.event.MouseListener
 import java.awt.event.MouseMotionListener
 import java.awt.event.MouseWheelListener
+import java.awt.image.BufferedImage
 import java.util.*
 import javax.swing.JFrame
 import kotlin.coroutines.CoroutineContext
@@ -86,18 +87,6 @@ class ComposeWindow @ExperimentalComposeUiApi constructor(
     internal val windowContext by composePanel::windowContext
     internal var rootForTestListener by composePanel::rootForTestListener
 
-    /**
-     * Returns the [SemanticsOwner]s corresponding to the roots of the semantics trees in this
-     * [ComposeWindow].
-     *
-     * This is backed by Snapshot state, so reading this property in a restartable function (e.g., a
-     * composable function) will cause the function to restart when the set of semantics owners
-     * changes.
-     */
-    @ComposeToolingApi
-    override val semanticsOwners: Collection<SemanticsOwner>
-        get() = composePanel.semanticsOwners
-
     /**
      * Controls whether mouse-down on an unfocusable element clears focus.
      */
@@ -378,4 +367,28 @@ class ComposeWindow @ExperimentalComposeUiApi constructor(
     internal fun measureContent(constraints: Constraints): IntSize {
         return composePanel.measureContent(constraints)
     }
+
+    /**
+     * Returns the [SemanticsOwner]s corresponding to the roots of the semantics trees in this
+     * [ComposeWindow].
+     *
+     * This is backed by Snapshot state, so reading this property in a restartable function (e.g., a
+     * composable function) will cause the function to restart when the set of semantics owners
+     * changes.
+     */
+    @ComposeToolingApi
+    override val semanticsOwners: Collection<SemanticsOwner>
+        get() = composePanel.semanticsOwners
+
+    /**
+     * Captures the content of this window into an image.
+     *
+     * Returns `null` if the window has not been made visible yet.
+     *
+     * May be called only on the event dispatching thread.
+     */
+    @ComposeToolingApi
+    override fun captureContentToImage(): BufferedImage? {
+        return composePanel.captureContentToImage()
+    }
 }

compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeWindowPanel.desktop.kt

@@ -212,4 +212,6 @@ internal class ComposeWindowPanel(
     fun actualizeSize(size: Dimension, insets: Insets): Dimension {
         return composeContainer.actualizeSize(size, insets)
     }
+
+    fun captureContentToImage() = composeContainer.captureContentToImage()
 }

compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/scene/ComposeContainer.desktop.kt

@@ -36,6 +36,7 @@ import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.util.fastAll
 import androidx.compose.ui.awt.UNSPECIFIED_DIMENSION_VALUE
+import androidx.compose.ui.unit.union
 import androidx.compose.ui.util.fastForEach
 import androidx.compose.ui.util.fastForEachReversed
 import androidx.compose.ui.util.fastRoundToInt
@@ -58,14 +59,14 @@ import java.awt.event.MouseEvent as AwtMouseEvent
 import java.awt.event.WindowEvent
 import java.awt.event.WindowFocusListener
 import java.awt.event.WindowListener
+import java.awt.image.BufferedImage
 import javax.swing.JLayeredPane
 import javax.swing.SwingUtilities
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.math.ceil
 import kotlinx.coroutines.CoroutineExceptionHandler
-import kotlinx.coroutines.Job
 import org.jetbrains.annotations.VisibleForTesting
 import org.jetbrains.skia.Canvas
 import org.jetbrains.skiko.MainUIDispatcher
@@ -621,6 +622,42 @@ internal class ComposeContainer(
         override fun shouldSendMouseEvent(event: AwtMouseEvent): Boolean = noBlockingInputLayers
         override fun shouldSendKeyEvent(event: AwtKeyEvent): Boolean = noBlockingInputLayers
     }
+
+    /**
+     * Captures the content of this container into an image.
+     *
+     * Returns `null` if the window has not been made visible yet.
+     *
+     * This may be called only on the event dispatch thread.
+     */
+    fun captureContentToImage(): BufferedImage? {
+        val mainLayerBounds = mediator.boundsOnScreenPx() ?: return null
+        val layersAndBounds = layers.map {
+            it to it.boundsOnScreenPx()
+        }
+
+        var resultBounds = mainLayerBounds
+        for ((_, bounds) in layersAndBounds) {
+            if (bounds != null) {
+                resultBounds = resultBounds.union(bounds)
+            }
+        }
+
+        val x = resultBounds.left
+        val y = resultBounds.top
+        val width = resultBounds.width
+        val height = resultBounds.height
+
+        val image = BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB)
+        mediator.drawContentInto(image, mainLayerBounds.left - x, mainLayerBounds.top - y)
+        for ((layer, bounds) in layersAndBounds) {
+            if (bounds == null) continue
+            // The offset is from mainLayerBounds because layers draw themselves relative to it
+            layer.drawContentInto(image, mainLayerBounds.left - x, mainLayerBounds.top - y)
+        }
+
+        return image
+    }
 }
 
 /**