Expose SurfaceControl method to screenshot to GraphicBuffer.

A graphic buffer is most useful, as we can both attach it
to starting windows, and directly use it in Sys-UI. The old
codepath for starting windows/saved surfaces, is co-existing
at the moment, so I don't make large attempts to clean up
the existing screenshot code.

Bug: 31339431
Test: Manual test in combination with other branches
Change-Id: I562fdd5460dbce3201ba090272e8731850780f20
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index b87250e..a12600a 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -17,6 +17,7 @@
 package android.view;
 
 import android.graphics.Bitmap;
+import android.graphics.GraphicBuffer;
 import android.graphics.Rect;
 import android.graphics.Region;
 import android.os.IBinder;
@@ -42,6 +43,9 @@
     private static native Bitmap nativeScreenshot(IBinder displayToken,
             Rect sourceCrop, int width, int height, int minLayer, int maxLayer,
             boolean allLayers, boolean useIdentityTransform, int rotation);
+    private static native GraphicBuffer nativeScreenshotToBuffer(IBinder displayToken,
+            Rect sourceCrop, int width, int height, int minLayer, int maxLayer,
+            boolean allLayers, boolean useIdentityTransform, int rotation);
     private static native void nativeScreenshot(IBinder displayToken, Surface consumer,
             Rect sourceCrop, int width, int height, int minLayer, int maxLayer,
             boolean allLayers, boolean useIdentityTransform);
@@ -828,6 +832,19 @@
     }
 
     /**
+     * Like {@link SurfaceControl#screenshot(Rect, int, int, int, int, boolean, int)}
+     * but returns a GraphicBuffer.
+     */
+    public static GraphicBuffer screenshotToBuffer(Rect sourceCrop, int width, int height,
+            int minLayer, int maxLayer, boolean useIdentityTransform,
+            int rotation) {
+        IBinder displayToken = SurfaceControl.getBuiltInDisplay(
+                SurfaceControl.BUILT_IN_DISPLAY_ID_MAIN);
+        return nativeScreenshotToBuffer(displayToken, sourceCrop, width, height,
+                minLayer, maxLayer, false, useIdentityTransform, rotation);
+    }
+
+    /**
      * Like {@link SurfaceControl#screenshot(int, int, int, int, boolean)} but
      * includes all Surfaces in the screenshot.
      *
diff --git a/core/jni/android/graphics/GraphicBuffer.cpp b/core/jni/android/graphics/GraphicBuffer.cpp
index c61b53e..e661c21 100644
--- a/core/jni/android/graphics/GraphicBuffer.cpp
+++ b/core/jni/android/graphics/GraphicBuffer.cpp
@@ -101,6 +101,13 @@
 // GraphicBuffer lifecycle
 // ----------------------------------------------------------------------------
 
+static jlong android_graphics_GraphicBuffer_wrap(JNIEnv* env, jobject clazz,
+        jlong unwrapped) {
+    sp<GraphicBuffer> b(reinterpret_cast<GraphicBuffer*>(unwrapped));
+    GraphicBufferWrapper* wrapper = new GraphicBufferWrapper(b);
+    return reinterpret_cast<jlong>(wrapper);
+}
+
 static jlong android_graphics_GraphicBuffer_create(JNIEnv* env, jobject clazz,
         jint width, jint height, jint format, jint usage) {
 
@@ -298,7 +305,9 @@
     { "nLockCanvas", "(JLandroid/graphics/Canvas;Landroid/graphics/Rect;)Z",
             (void*) android_graphics_GraphicBuffer_lockCanvas },
     { "nUnlockCanvasAndPost", "(JLandroid/graphics/Canvas;)Z",
-            (void*) android_graphics_GraphicBuffer_unlockCanvasAndPost }
+            (void*) android_graphics_GraphicBuffer_unlockCanvasAndPost },
+    { "nWrapGraphicBuffer", "(J)J",
+            (void*) android_graphics_GraphicBuffer_wrap }
 };
 
 int register_android_graphics_GraphicBuffer(JNIEnv* env) {
@@ -316,4 +325,4 @@
     gRectClassInfo.bottom = GetFieldIDOrDie(env, clazz, "bottom", "I");
 
     return RegisterMethodsOrDie(env, kClassPathName, gMethods, NELEM(gMethods));
-}
\ No newline at end of file
+}
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index ed071cd..a3fef27 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -15,6 +15,7 @@
  */
 
 #define LOG_TAG "SurfaceControl"
+#define LOG_NDEBUG 0
 
 #include "android_os_Parcel.h"
 #include "android_util_Binder.h"
@@ -89,6 +90,11 @@
     jmethodID ctor;
 } gHdrCapabilitiesClassInfo;
 
+static struct {
+    jclass clazz;
+    jmethodID builder;
+} gGraphicBufferClassInfo;
+
 // ----------------------------------------------------------------------------
 
 static jlong nativeCreate(JNIEnv* env, jclass clazz, jobject sessionObj,
@@ -123,6 +129,44 @@
     }
 }
 
+static Rect rectFromObj(JNIEnv* env, jobject rectObj) {
+    int left = env->GetIntField(rectObj, gRectClassInfo.left);
+    int top = env->GetIntField(rectObj, gRectClassInfo.top);
+    int right = env->GetIntField(rectObj, gRectClassInfo.right);
+    int bottom = env->GetIntField(rectObj, gRectClassInfo.bottom);
+    return Rect(left, top, right, bottom);
+}
+
+static jobject nativeScreenshotToBuffer(JNIEnv* env, jclass clazz,
+        jobject displayTokenObj, jobject sourceCropObj, jint width, jint height,
+        jint minLayer, jint maxLayer, bool allLayers, bool useIdentityTransform,
+        int rotation) {
+    sp<IBinder> displayToken = ibinderForJavaObject(env, displayTokenObj);
+    if (displayToken == NULL) {
+        return NULL;
+    }
+    Rect sourceCrop = rectFromObj(env, sourceCropObj);
+    if (allLayers) {
+        minLayer = 0;
+        maxLayer = -1;
+    }
+    sp<GraphicBuffer> buffer;
+    status_t res = ScreenshotClient::captureToBuffer(displayToken,
+            sourceCrop, width, height, minLayer, maxLayer, useIdentityTransform,
+            rotation, &buffer);
+    if (res != NO_ERROR) {
+        return NULL;
+    }
+
+    return env->CallStaticObjectMethod(gGraphicBufferClassInfo.clazz,
+            gGraphicBufferClassInfo.builder,
+            buffer->getWidth(),
+            buffer->getHeight(),
+            buffer->getPixelFormat(),
+            buffer->getUsage(),
+            (void*)buffer.get());
+}
+
 static jobject nativeScreenshotBitmap(JNIEnv* env, jclass clazz,
         jobject displayTokenObj, jobject sourceCropObj, jint width, jint height,
         jint minLayer, jint maxLayer, bool allLayers, bool useIdentityTransform,
@@ -132,11 +176,7 @@
         return NULL;
     }
 
-    int left = env->GetIntField(sourceCropObj, gRectClassInfo.left);
-    int top = env->GetIntField(sourceCropObj, gRectClassInfo.top);
-    int right = env->GetIntField(sourceCropObj, gRectClassInfo.right);
-    int bottom = env->GetIntField(sourceCropObj, gRectClassInfo.bottom);
-    Rect sourceCrop(left, top, right, bottom);
+    Rect sourceCrop = rectFromObj(env, sourceCropObj);
 
     std::unique_ptr<ScreenshotClient> screenshot(new ScreenshotClient());
     status_t res;
@@ -786,7 +826,10 @@
     {"nativeGetHandle", "(J)Landroid/os/IBinder;",
             (void*)nativeGetHandle },
     {"nativeGetTransformToDisplayInverse", "(J)Z",
-            (void*)nativeGetTransformToDisplayInverse },
+     (void*)nativeGetTransformToDisplayInverse },
+    {"nativeScreenshotToBuffer",
+     "(Landroid/os/IBinder;Landroid/graphics/Rect;IIIIZZI)Landroid/graphics/GraphicBuffer;",
+     (void*)nativeScreenshotToBuffer },
 };
 
 int register_android_view_SurfaceControl(JNIEnv* env)
@@ -836,6 +879,11 @@
     gHdrCapabilitiesClassInfo.ctor = GetMethodIDOrDie(env, hdrCapabilitiesClazz, "<init>",
             "([IFFF)V");
 
+    jclass graphicsBufferClazz = FindClassOrDie(env, "android/graphics/GraphicBuffer");
+    gGraphicBufferClassInfo.clazz = MakeGlobalRefOrDie(env, graphicsBufferClazz);
+    gGraphicBufferClassInfo.builder = GetStaticMethodIDOrDie(env, graphicsBufferClazz,
+            "createFromExisting", "(IIIIJ)Landroid/graphics/GraphicBuffer;");
+
     return err;
 }