Revert "Delete ScreenCapture#captureDisplay and DisplayCaptureArgs."

This reverts commit 8b04d81ffa7c458953f6d940f5007531639f42b7.

Bug: b/301083608
Reason for revert: investigating causes for b/301083608

Change-Id: I24e0387ccd360c22dee7e4a3e35057c986d4fd2b
diff --git a/core/java/android/window/ScreenCapture.java b/core/java/android/window/ScreenCapture.java
index e42193d..95e9e86 100644
--- a/core/java/android/window/ScreenCapture.java
+++ b/core/java/android/window/ScreenCapture.java
@@ -24,6 +24,7 @@
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
 import android.os.Build;
+import android.os.IBinder;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.util.Log;
@@ -34,6 +35,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.function.ObjIntConsumer;
+
 /**
  * Handles display and layer captures for the system.
  *
@@ -43,6 +45,8 @@
     private static final String TAG = "ScreenCapture";
     private static final int SCREENSHOT_WAIT_TIME_S = 4 * Build.HW_TIMEOUT_MULTIPLIER;
 
+    private static native int nativeCaptureDisplay(DisplayCaptureArgs captureArgs,
+            long captureListener);
     private static native int nativeCaptureLayers(LayerCaptureArgs captureArgs,
             long captureListener);
     private static native long nativeCreateScreenCaptureListener(
@@ -52,6 +56,37 @@
     private static native long getNativeListenerFinalizer();
 
     /**
+     * @param captureArgs     Arguments about how to take the screenshot
+     * @param captureListener A listener to receive the screenshot callback
+     * @hide
+     */
+    public static int captureDisplay(@NonNull DisplayCaptureArgs captureArgs,
+            @NonNull ScreenCaptureListener captureListener) {
+        return nativeCaptureDisplay(captureArgs, captureListener.mNativeObject);
+    }
+
+    /**
+     * Captures all the surfaces in a display and returns a {@link ScreenshotHardwareBuffer} with
+     * the content.
+     *
+     * @hide
+     */
+    public static ScreenshotHardwareBuffer captureDisplay(
+            DisplayCaptureArgs captureArgs) {
+        SynchronousScreenCaptureListener syncScreenCapture = createSyncCaptureListener();
+        int status = captureDisplay(captureArgs, syncScreenCapture);
+        if (status != 0) {
+            return null;
+        }
+
+        try {
+            return syncScreenCapture.getBuffer();
+        } catch (Exception e) {
+            return null;
+        }
+    }
+
+    /**
      * Captures a layer and its children and returns a {@link HardwareBuffer} with the content.
      *
      * @param layer      The root layer to capture.
@@ -484,6 +519,92 @@
     }
 
     /**
+     * The arguments class used to make display capture requests.
+     *
+     * @hide
+     * @see #nativeCaptureDisplay(DisplayCaptureArgs, long)
+     */
+    public static class DisplayCaptureArgs extends CaptureArgs {
+        private final IBinder mDisplayToken;
+        private final int mWidth;
+        private final int mHeight;
+        private final boolean mUseIdentityTransform;
+
+        private DisplayCaptureArgs(Builder builder) {
+            super(builder);
+            mDisplayToken = builder.mDisplayToken;
+            mWidth = builder.mWidth;
+            mHeight = builder.mHeight;
+            mUseIdentityTransform = builder.mUseIdentityTransform;
+        }
+
+        /**
+         * The Builder class used to construct {@link DisplayCaptureArgs}
+         */
+        public static class Builder extends CaptureArgs.Builder<Builder> {
+            private IBinder mDisplayToken;
+            private int mWidth;
+            private int mHeight;
+            private boolean mUseIdentityTransform;
+
+            /**
+             * Construct a new {@link LayerCaptureArgs} with the set parameters. The builder
+             * remains valid.
+             */
+            public DisplayCaptureArgs build() {
+                if (mDisplayToken == null) {
+                    throw new IllegalStateException(
+                            "Can't take screenshot with null display token");
+                }
+                return new DisplayCaptureArgs(this);
+            }
+
+            public Builder(IBinder displayToken) {
+                setDisplayToken(displayToken);
+            }
+
+            /**
+             * The display to take the screenshot of.
+             */
+            public Builder setDisplayToken(IBinder displayToken) {
+                mDisplayToken = displayToken;
+                return this;
+            }
+
+            /**
+             * Set the desired size of the returned buffer. The raw screen  will be  scaled down to
+             * this size
+             *
+             * @param width  The desired width of the returned buffer. Caller may pass in 0 if no
+             *               scaling is desired.
+             * @param height The desired height of the returned buffer. Caller may pass in 0 if no
+             *               scaling is desired.
+             */
+            public Builder setSize(int width, int height) {
+                mWidth = width;
+                mHeight = height;
+                return this;
+            }
+
+            /**
+             * Replace the rotation transform of the display with the identity transformation while
+             * taking the screenshot. This ensures the screenshot is taken in the ROTATION_0
+             * orientation. Set this value to false if the screenshot should be taken in the
+             * current screen orientation.
+             */
+            public Builder setUseIdentityTransform(boolean useIdentityTransform) {
+                mUseIdentityTransform = useIdentityTransform;
+                return this;
+            }
+
+            @Override
+            Builder getThis() {
+                return this;
+            }
+        }
+    }
+
+    /**
      * The arguments class used to make layer capture requests.
      *
      * @hide
@@ -561,6 +682,7 @@
 
     /**
      * The object used to receive the results when invoking screen capture requests via
+     * {@link #captureDisplay(DisplayCaptureArgs, ScreenCaptureListener)} or
      * {@link #captureLayers(LayerCaptureArgs, ScreenCaptureListener)}
      *
      * This listener can only be used for a single call to capture content call.
@@ -662,7 +784,8 @@
 
     /**
      * Helper class to synchronously get the {@link ScreenshotHardwareBuffer} when calling
-     * {@link #captureLayers(LayerCaptureArgs, ScreenCaptureListener)}
+     * {@link #captureLayers(LayerCaptureArgs, ScreenCaptureListener)} or
+     * {@link #captureDisplay(DisplayCaptureArgs, ScreenCaptureListener)}
      */
     public abstract static class SynchronousScreenCaptureListener extends ScreenCaptureListener {
         SynchronousScreenCaptureListener(ObjIntConsumer<ScreenshotHardwareBuffer> consumer) {
diff --git a/core/jni/android_window_ScreenCapture.cpp b/core/jni/android_window_ScreenCapture.cpp
index beb8c9b..bdf7eaa 100644
--- a/core/jni/android_window_ScreenCapture.cpp
+++ b/core/jni/android_window_ScreenCapture.cpp
@@ -50,6 +50,13 @@
 } gCaptureArgsClassInfo;
 
 static struct {
+    jfieldID displayToken;
+    jfieldID width;
+    jfieldID height;
+    jfieldID useIdentityTransform;
+} gDisplayCaptureArgsClassInfo;
+
+static struct {
     jfieldID layer;
     jfieldID childrenOnly;
 } gLayerCaptureArgsClassInfo;
@@ -174,6 +181,39 @@
                                  gCaptureArgsClassInfo.hintForSeamlessTransition);
 }
 
+static DisplayCaptureArgs displayCaptureArgsFromObject(JNIEnv* env,
+                                                       jobject displayCaptureArgsObject) {
+    DisplayCaptureArgs captureArgs;
+    getCaptureArgs(env, displayCaptureArgsObject, captureArgs);
+
+    captureArgs.displayToken =
+            ibinderForJavaObject(env,
+                                 env->GetObjectField(displayCaptureArgsObject,
+                                                     gDisplayCaptureArgsClassInfo.displayToken));
+    captureArgs.width =
+            env->GetIntField(displayCaptureArgsObject, gDisplayCaptureArgsClassInfo.width);
+    captureArgs.height =
+            env->GetIntField(displayCaptureArgsObject, gDisplayCaptureArgsClassInfo.height);
+    captureArgs.useIdentityTransform =
+            env->GetBooleanField(displayCaptureArgsObject,
+                                 gDisplayCaptureArgsClassInfo.useIdentityTransform);
+    return captureArgs;
+}
+
+static jint nativeCaptureDisplay(JNIEnv* env, jclass clazz, jobject displayCaptureArgsObject,
+                                 jlong screenCaptureListenerObject) {
+    const DisplayCaptureArgs captureArgs =
+            displayCaptureArgsFromObject(env, displayCaptureArgsObject);
+
+    if (captureArgs.displayToken == nullptr) {
+        return BAD_VALUE;
+    }
+
+    sp<gui::IScreenCaptureListener> captureListener =
+            reinterpret_cast<gui::IScreenCaptureListener*>(screenCaptureListenerObject);
+    return ScreenshotClient::captureDisplay(captureArgs, captureListener);
+}
+
 static jint nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerCaptureArgsObject,
                                 jlong screenCaptureListenerObject) {
     LayerCaptureArgs captureArgs;
@@ -243,6 +283,8 @@
 
 static const JNINativeMethod sScreenCaptureMethods[] = {
         // clang-format off
+    {"nativeCaptureDisplay", "(Landroid/window/ScreenCapture$DisplayCaptureArgs;J)I",
+            (void*)nativeCaptureDisplay },
     {"nativeCaptureLayers",  "(Landroid/window/ScreenCapture$LayerCaptureArgs;J)I",
             (void*)nativeCaptureLayers },
     {"nativeCreateScreenCaptureListener", "(Ljava/util/function/ObjIntConsumer;)J",
@@ -275,6 +317,17 @@
     gCaptureArgsClassInfo.hintForSeamlessTransition =
             GetFieldIDOrDie(env, captureArgsClazz, "mHintForSeamlessTransition", "Z");
 
+    jclass displayCaptureArgsClazz =
+            FindClassOrDie(env, "android/window/ScreenCapture$DisplayCaptureArgs");
+    gDisplayCaptureArgsClassInfo.displayToken =
+            GetFieldIDOrDie(env, displayCaptureArgsClazz, "mDisplayToken", "Landroid/os/IBinder;");
+    gDisplayCaptureArgsClassInfo.width =
+            GetFieldIDOrDie(env, displayCaptureArgsClazz, "mWidth", "I");
+    gDisplayCaptureArgsClassInfo.height =
+            GetFieldIDOrDie(env, displayCaptureArgsClazz, "mHeight", "I");
+    gDisplayCaptureArgsClassInfo.useIdentityTransform =
+            GetFieldIDOrDie(env, displayCaptureArgsClazz, "mUseIdentityTransform", "Z");
+
     jclass layerCaptureArgsClazz =
             FindClassOrDie(env, "android/window/ScreenCapture$LayerCaptureArgs");
     gLayerCaptureArgsClassInfo.layer =