Render View into Bitmap using hardware acceleration (1/2)

This CLs switches the screenshot tests to render using hardware
rendering rather than software rendering. The reason for this is that
some features are not supported in software rendering (like clipping to
an outline).

Even though hardware rendering are not meant to be 100% deterministic
(e.g. shadows and ripples), I'd like us to try to still go for
pixel-perfect matching of screenshots whenever possible, especially
given that we don't really care about testing things like
shadows/elevation. In the future, we might either use software rendering
or more lenient matchers in case we want to test some UIs that are
impossible to make deterministic.

Because the AndroidX View.captureToBitmap() API unfortunately does not
work for dialogs (see b/195673633), I had to fork ViewCapture.kt and
WindowCapture.kt to ensure that we use the correct window we are sending
over to PixelCopy for the hardware rendering.

Bug: 230832101
Test: atest SystemUIGoogleScreenshotTests
Change-Id: I8cb6398c0c446b754d5c1af27296a18d53ce738e
diff --git a/packages/SystemUI/screenshot/Android.bp b/packages/SystemUI/screenshot/Android.bp
index 601e92f..f449398 100644
--- a/packages/SystemUI/screenshot/Android.bp
+++ b/packages/SystemUI/screenshot/Android.bp
@@ -38,6 +38,7 @@
         "androidx.test.espresso.core",
         "androidx.appcompat_appcompat",
         "platform-screenshot-diff-core",
+        "guava",
     ],
 
     kotlincflags: ["-Xjvm-default=all"],
diff --git a/packages/SystemUI/screenshot/res/values/themes.xml b/packages/SystemUI/screenshot/res/values/themes.xml
index 40e50bb..a7f8a26 100644
--- a/packages/SystemUI/screenshot/res/values/themes.xml
+++ b/packages/SystemUI/screenshot/res/values/themes.xml
@@ -19,6 +19,12 @@
         <item name="android:windowActionBar">false</item>
         <item name="android:windowNoTitle">true</item>
 
+        <!-- We make the status and navigation bars transparent so that the screenshotted content is
+             not clipped by the status bar height when drawn into the Bitmap (which is what happens
+             given that we draw the view into the Bitmap using hardware acceleration). -->
+        <item name="android:statusBarColor">@android:color/transparent</item>
+        <item name="android:navigationBarColor">@android:color/transparent</item>
+
         <!-- Make sure that device specific cutouts don't impact the outcome of screenshot tests -->
         <item name="android:windowLayoutInDisplayCutoutMode">shortEdges</item>
     </style>
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
index 3d26cda..a4a70a4 100644
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/Bitmap.kt
@@ -24,6 +24,8 @@
 import platform.test.screenshot.matchers.PixelPerfectMatcher
 
 /** Draw this [View] into a [Bitmap]. */
+// TODO(b/195673633): Remove this once Compose screenshot tests use hardware rendering for their
+// tests.
 fun View.drawIntoBitmap(): Bitmap {
     val bitmap =
         Bitmap.createBitmap(
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt
new file mode 100644
index 0000000..c609e6f
--- /dev/null
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewCapture.kt
@@ -0,0 +1,180 @@
+package com.android.systemui.testing.screenshot
+
+import android.app.Activity
+import android.content.Context
+import android.content.ContextWrapper
+import android.graphics.Bitmap
+import android.graphics.Canvas
+import android.graphics.Rect
+import android.os.Build
+import android.os.Handler
+import android.os.Looper
+import android.util.Log
+import android.view.PixelCopy
+import android.view.SurfaceView
+import android.view.View
+import android.view.ViewTreeObserver
+import android.view.Window
+import androidx.annotation.RequiresApi
+import androidx.concurrent.futures.ResolvableFuture
+import androidx.test.annotation.ExperimentalTestApi
+import androidx.test.core.internal.os.HandlerExecutor
+import androidx.test.platform.graphics.HardwareRendererCompat
+import com.google.common.util.concurrent.ListenableFuture
+
+/*
+ * This file was forked from androidx/test/core/view/ViewCapture.kt to add [Window] parameter to
+ * [View.captureToBitmap].
+ * TODO(b/195673633): Remove this fork and use the AndroidX version instead.
+ */
+
+/**
+ * Asynchronously captures an image of the underlying view into a [Bitmap].
+ *
+ * For devices below [Build.VERSION_CODES#O] (or if the view's window cannot be determined), the
+ * image is obtained using [View#draw]. Otherwise, [PixelCopy] is used.
+ *
+ * This method will also enable [HardwareRendererCompat#setDrawingEnabled(boolean)] if required.
+ *
+ * This API is primarily intended for use in lower layer libraries or frameworks. For test authors,
+ * its recommended to use espresso or compose's captureToImage.
+ *
+ * This API is currently experimental and subject to change or removal.
+ */
+@ExperimentalTestApi
+@RequiresApi(Build.VERSION_CODES.JELLY_BEAN)
+fun View.captureToBitmap(window: Window? = null): ListenableFuture<Bitmap> {
+    val bitmapFuture: ResolvableFuture<Bitmap> = ResolvableFuture.create()
+    val mainExecutor = HandlerExecutor(Handler(Looper.getMainLooper()))
+
+    // disable drawing again if necessary once work is complete
+    if (!HardwareRendererCompat.isDrawingEnabled()) {
+        HardwareRendererCompat.setDrawingEnabled(true)
+        bitmapFuture.addListener({ HardwareRendererCompat.setDrawingEnabled(false) }, mainExecutor)
+    }
+
+    mainExecutor.execute {
+        val forceRedrawFuture = forceRedraw()
+        forceRedrawFuture.addListener({ generateBitmap(bitmapFuture, window) }, mainExecutor)
+    }
+
+    return bitmapFuture
+}
+
+/**
+ * Trigger a redraw of the given view.
+ *
+ * Should only be called on UI thread.
+ *
+ * @return a [ListenableFuture] that will be complete once ui drawing is complete
+ */
+// NoClassDefFoundError occurs on API 15
+@RequiresApi(Build.VERSION_CODES.JELLY_BEAN)
+// @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
+@ExperimentalTestApi
+fun View.forceRedraw(): ListenableFuture<Void> {
+    val future: ResolvableFuture<Void> = ResolvableFuture.create()
+
+    if (Build.VERSION.SDK_INT >= 29 && isHardwareAccelerated) {
+        viewTreeObserver.registerFrameCommitCallback() { future.set(null) }
+    } else {
+        viewTreeObserver.addOnDrawListener(
+            object : ViewTreeObserver.OnDrawListener {
+                var handled = false
+                override fun onDraw() {
+                    if (!handled) {
+                        handled = true
+                        future.set(null)
+                        // cannot remove on draw listener inside of onDraw
+                        Handler(Looper.getMainLooper()).post {
+                            viewTreeObserver.removeOnDrawListener(this)
+                        }
+                    }
+                }
+            }
+        )
+    }
+    invalidate()
+    return future
+}
+
+private fun View.generateBitmap(
+    bitmapFuture: ResolvableFuture<Bitmap>,
+    window: Window? = null,
+) {
+    if (bitmapFuture.isCancelled) {
+        return
+    }
+    val destBitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
+    when {
+        Build.VERSION.SDK_INT < 26 -> generateBitmapFromDraw(destBitmap, bitmapFuture)
+        this is SurfaceView -> generateBitmapFromSurfaceViewPixelCopy(destBitmap, bitmapFuture)
+        else -> {
+            val window = window ?: getActivity()?.window
+            if (window != null) {
+                generateBitmapFromPixelCopy(window, destBitmap, bitmapFuture)
+            } else {
+                Log.i(
+                    "View.captureToImage",
+                    "Could not find window for view. Falling back to View#draw instead of PixelCopy"
+                )
+                generateBitmapFromDraw(destBitmap, bitmapFuture)
+            }
+        }
+    }
+}
+
+@SuppressWarnings("NewApi")
+private fun SurfaceView.generateBitmapFromSurfaceViewPixelCopy(
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    val onCopyFinished =
+        PixelCopy.OnPixelCopyFinishedListener { result ->
+            if (result == PixelCopy.SUCCESS) {
+                bitmapFuture.set(destBitmap)
+            } else {
+                bitmapFuture.setException(
+                    RuntimeException(String.format("PixelCopy failed: %d", result))
+                )
+            }
+        }
+    PixelCopy.request(this, null, destBitmap, onCopyFinished, handler)
+}
+
+internal fun View.generateBitmapFromDraw(
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    destBitmap.density = resources.displayMetrics.densityDpi
+    computeScroll()
+    val canvas = Canvas(destBitmap)
+    canvas.translate((-scrollX).toFloat(), (-scrollY).toFloat())
+    draw(canvas)
+    bitmapFuture.set(destBitmap)
+}
+
+private fun View.getActivity(): Activity? {
+    fun Context.getActivity(): Activity? {
+        return when (this) {
+            is Activity -> this
+            is ContextWrapper -> this.baseContext.getActivity()
+            else -> null
+        }
+    }
+    return context.getActivity()
+}
+
+private fun View.generateBitmapFromPixelCopy(
+    window: Window,
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    val locationInWindow = intArrayOf(0, 0)
+    getLocationInWindow(locationInWindow)
+    val x = locationInWindow[0]
+    val y = locationInWindow[1]
+    val boundsInWindow = Rect(x, y, x + width, y + height)
+
+    return window.generateBitmapFromPixelCopy(boundsInWindow, destBitmap, bitmapFuture)
+}
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt
index 3209c8b..60130e1 100644
--- a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/ViewScreenshotTestRule.kt
@@ -18,10 +18,22 @@
 
 import android.app.Activity
 import android.app.Dialog
+import android.graphics.Bitmap
+import android.graphics.HardwareRenderer
+import android.os.Looper
 import android.view.View
 import android.view.ViewGroup
 import android.view.ViewGroup.LayoutParams
+import android.view.ViewGroup.LayoutParams.MATCH_PARENT
+import android.view.ViewGroup.LayoutParams.WRAP_CONTENT
+import android.view.Window
+import androidx.activity.ComponentActivity
+import androidx.test.espresso.Espresso
 import androidx.test.ext.junit.rules.ActivityScenarioRule
+import com.google.common.util.concurrent.FutureCallback
+import com.google.common.util.concurrent.Futures
+import kotlin.coroutines.suspendCoroutine
+import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertEquals
 import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
@@ -59,29 +71,39 @@
      */
     fun screenshotTest(
         goldenIdentifier: String,
-        layoutParams: LayoutParams =
-            LayoutParams(LayoutParams.WRAP_CONTENT, LayoutParams.WRAP_CONTENT),
-        viewProvider: (Activity) -> View,
+        mode: Mode = Mode.WrapContent,
+        viewProvider: (ComponentActivity) -> View,
     ) {
         activityRule.scenario.onActivity { activity ->
             // Make sure that the activity draws full screen and fits the whole display instead of
             // the system bars.
-            activity.window.setDecorFitsSystemWindows(false)
-            activity.setContentView(viewProvider(activity), layoutParams)
+            val window = activity.window
+            window.setDecorFitsSystemWindows(false)
+
+            // Set the content.
+            activity.setContentView(viewProvider(activity), mode.layoutParams)
+
+            // Elevation/shadows is not deterministic when doing hardware rendering, so we disable
+            // it for any view in the hierarchy.
+            window.decorView.removeElevationRecursively()
         }
 
         // We call onActivity again because it will make sure that our Activity is done measuring,
         // laying out and drawing its content (that we set in the previous onActivity lambda).
+        var contentView: View? = null
         activityRule.scenario.onActivity { activity ->
             // Check that the content is what we expected.
             val content = activity.requireViewById<ViewGroup>(android.R.id.content)
             assertEquals(1, content.childCount)
-            screenshotRule.assertBitmapAgainstGolden(
-                content.getChildAt(0).drawIntoBitmap(),
-                goldenIdentifier,
-                matcher
-            )
+            contentView = content.getChildAt(0)
         }
+
+        val bitmap = contentView?.toBitmap() ?: error("contentView is null")
+        screenshotRule.assertBitmapAgainstGolden(
+            bitmap,
+            goldenIdentifier,
+            matcher,
+        )
     }
 
     /**
@@ -104,25 +126,78 @@
                     create()
                     window.setWindowAnimations(0)
 
+                    // Elevation/shadows is not deterministic when doing hardware rendering, so we
+                    // disable it for any view in the hierarchy.
+                    window.decorView.removeElevationRecursively()
+
                     // Show the dialog.
                     show()
                 }
         }
 
-        // We call onActivity again because it will make sure that our Dialog is done measuring,
-        // laying out and drawing its content (that we set in the previous onActivity lambda).
-        activityRule.scenario.onActivity {
-            // Check that the content is what we expected.
-            val dialog = dialog ?: error("dialog is null")
-            try {
-                screenshotRule.assertBitmapAgainstGolden(
-                    dialog.window.decorView.drawIntoBitmap(),
-                    goldenIdentifier,
-                    matcher,
+        try {
+            val bitmap = dialog?.toBitmap() ?: error("dialog is null")
+            screenshotRule.assertBitmapAgainstGolden(
+                bitmap,
+                goldenIdentifier,
+                matcher,
+            )
+        } finally {
+            dialog?.dismiss()
+        }
+    }
+
+    private fun View.removeElevationRecursively() {
+        this.elevation = 0f
+
+        if (this is ViewGroup) {
+            repeat(childCount) { i -> getChildAt(i).removeElevationRecursively() }
+        }
+    }
+
+    private fun Dialog.toBitmap(): Bitmap {
+        val window = window
+        return window.decorView.toBitmap(window)
+    }
+
+    private fun View.toBitmap(window: Window? = null): Bitmap {
+        if (Looper.getMainLooper() == Looper.myLooper()) {
+            error("toBitmap() can't be called from the main thread")
+        }
+
+        if (!HardwareRenderer.isDrawingEnabled()) {
+            error("Hardware rendering is not enabled")
+        }
+
+        // Make sure we are idle.
+        Espresso.onIdle()
+
+        val mainExecutor = context.mainExecutor
+        return runBlocking {
+            suspendCoroutine { continuation ->
+                Futures.addCallback(
+                    captureToBitmap(window),
+                    object : FutureCallback<Bitmap> {
+                        override fun onSuccess(result: Bitmap?) {
+                            continuation.resumeWith(Result.success(result!!))
+                        }
+
+                        override fun onFailure(t: Throwable) {
+                            continuation.resumeWith(Result.failure(t))
+                        }
+                    },
+                    // We know that we are not on the main thread, so we can block the current
+                    // thread and wait for the result in the main thread.
+                    mainExecutor,
                 )
-            } finally {
-                dialog.dismiss()
             }
         }
     }
+
+    enum class Mode(val layoutParams: LayoutParams) {
+        WrapContent(LayoutParams(WRAP_CONTENT, WRAP_CONTENT)),
+        MatchSize(LayoutParams(MATCH_PARENT, MATCH_PARENT)),
+        MatchWidth(LayoutParams(MATCH_PARENT, WRAP_CONTENT)),
+        MatchHeight(LayoutParams(WRAP_CONTENT, MATCH_PARENT)),
+    }
 }
diff --git a/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/WindowCapture.kt b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/WindowCapture.kt
new file mode 100644
index 0000000..d34f46b
--- /dev/null
+++ b/packages/SystemUI/screenshot/src/com/android/systemui/testing/screenshot/WindowCapture.kt
@@ -0,0 +1,37 @@
+package com.android.systemui.testing.screenshot
+
+import android.graphics.Bitmap
+import android.graphics.Rect
+import android.os.Handler
+import android.os.Looper
+import android.view.PixelCopy
+import android.view.Window
+import androidx.concurrent.futures.ResolvableFuture
+
+/*
+ * This file was forked from androidx/test/core/view/WindowCapture.kt.
+ * TODO(b/195673633): Remove this fork and use the AndroidX version instead.
+ */
+fun Window.generateBitmapFromPixelCopy(
+    boundsInWindow: Rect? = null,
+    destBitmap: Bitmap,
+    bitmapFuture: ResolvableFuture<Bitmap>
+) {
+    val onCopyFinished =
+        PixelCopy.OnPixelCopyFinishedListener { result ->
+            if (result == PixelCopy.SUCCESS) {
+                bitmapFuture.set(destBitmap)
+            } else {
+                bitmapFuture.setException(
+                    RuntimeException(String.format("PixelCopy failed: %d", result))
+                )
+            }
+        }
+    PixelCopy.request(
+        this,
+        boundsInWindow,
+        destBitmap,
+        onCopyFinished,
+        Handler(Looper.getMainLooper())
+    )
+}