Reduce SurfaceControl feature gap to NDK API

Bug: 200697201
Test: atest android.view.cts.SurfaceControlTest

Change-Id: I1d0cff355714acb5668952b3d7f6fc0c15952523
diff --git a/core/api/current.txt b/core/api/current.txt
index b40a3e5..3b16fd5 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -17695,6 +17695,7 @@
     field public static final int RGB_565 = 4; // 0x4
     field public static final int RGB_888 = 3; // 0x3
     field public static final int S_UI8 = 53; // 0x35
+    field public static final long USAGE_COMPOSER_OVERLAY = 2048L; // 0x800L
     field public static final long USAGE_CPU_READ_OFTEN = 3L; // 0x3L
     field public static final long USAGE_CPU_READ_RARELY = 2L; // 0x2L
     field public static final long USAGE_CPU_WRITE_OFTEN = 48L; // 0x30L
@@ -48825,6 +48826,7 @@
     method @NonNull public android.view.SurfaceControl build();
     method @NonNull public android.view.SurfaceControl.Builder setBufferSize(@IntRange(from=0) int, @IntRange(from=0) int);
     method @NonNull public android.view.SurfaceControl.Builder setFormat(int);
+    method @NonNull public android.view.SurfaceControl.Builder setHidden(boolean);
     method @NonNull public android.view.SurfaceControl.Builder setName(@NonNull String);
     method @NonNull public android.view.SurfaceControl.Builder setOpaque(boolean);
     method @NonNull public android.view.SurfaceControl.Builder setParent(@Nullable android.view.SurfaceControl);
@@ -48839,11 +48841,18 @@
     method @NonNull public android.view.SurfaceControl.Transaction merge(@NonNull android.view.SurfaceControl.Transaction);
     method @NonNull public android.view.SurfaceControl.Transaction reparent(@NonNull android.view.SurfaceControl, @Nullable android.view.SurfaceControl);
     method @NonNull public android.view.SurfaceControl.Transaction setAlpha(@NonNull android.view.SurfaceControl, @FloatRange(from=0.0, to=1.0) float);
+    method @NonNull public android.view.SurfaceControl.Transaction setBuffer(@NonNull android.view.SurfaceControl, @Nullable android.hardware.HardwareBuffer);
     method @NonNull public android.view.SurfaceControl.Transaction setBufferSize(@NonNull android.view.SurfaceControl, @IntRange(from=0) int, @IntRange(from=0) int);
+    method @NonNull public android.view.SurfaceControl.Transaction setBufferTransform(@NonNull android.view.SurfaceControl, int);
+    method @NonNull public android.view.SurfaceControl.Transaction setCrop(@NonNull android.view.SurfaceControl, @Nullable android.graphics.Rect);
+    method @NonNull public android.view.SurfaceControl.Transaction setDamageRegion(@NonNull android.view.SurfaceControl, @Nullable android.graphics.Region);
     method @NonNull public android.view.SurfaceControl.Transaction setFrameRate(@NonNull android.view.SurfaceControl, @FloatRange(from=0.0) float, int);
     method @NonNull public android.view.SurfaceControl.Transaction setFrameRate(@NonNull android.view.SurfaceControl, @FloatRange(from=0.0) float, int, int);
     method @NonNull public android.view.SurfaceControl.Transaction setGeometry(@NonNull android.view.SurfaceControl, @Nullable android.graphics.Rect, @Nullable android.graphics.Rect, int);
     method @NonNull public android.view.SurfaceControl.Transaction setLayer(@NonNull android.view.SurfaceControl, @IntRange(from=java.lang.Integer.MIN_VALUE, to=java.lang.Integer.MAX_VALUE) int);
+    method @NonNull public android.view.SurfaceControl.Transaction setOpaque(@NonNull android.view.SurfaceControl, boolean);
+    method @NonNull public android.view.SurfaceControl.Transaction setPosition(@NonNull android.view.SurfaceControl, float, float);
+    method @NonNull public android.view.SurfaceControl.Transaction setScale(@NonNull android.view.SurfaceControl, float, float);
     method @NonNull public android.view.SurfaceControl.Transaction setVisibility(@NonNull android.view.SurfaceControl, boolean);
     method public void writeToParcel(@NonNull android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.view.SurfaceControl.Transaction> CREATOR;
diff --git a/core/java/android/hardware/HardwareBuffer.java b/core/java/android/hardware/HardwareBuffer.java
index cad30dd..a4a8f31 100644
--- a/core/java/android/hardware/HardwareBuffer.java
+++ b/core/java/android/hardware/HardwareBuffer.java
@@ -25,6 +25,7 @@
 import android.os.Build;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.view.SurfaceControl;
 
 import dalvik.annotation.optimization.CriticalNative;
 import dalvik.annotation.optimization.FastNative;
@@ -129,6 +130,15 @@
     public static final long USAGE_GPU_SAMPLED_IMAGE      = 1 << 8;
     /** Usage: The buffer will be written to by the GPU */
     public static final long USAGE_GPU_COLOR_OUTPUT       = 1 << 9;
+    /**
+     * The buffer will be used as a composer HAL overlay layer.
+     *
+     * This flag is currently only needed when using
+     * {@link android.view.SurfaceControl.Transaction#setBuffer(SurfaceControl, HardwareBuffer)}
+     * to set a buffer. In all other cases, the framework adds this flag
+     * internally to buffers that could be presented in a composer overlay.
+     */
+    public static final long USAGE_COMPOSER_OVERLAY = 1 << 11;
     /** Usage: The buffer must not be used outside of a protected hardware path */
     public static final long USAGE_PROTECTED_CONTENT      = 1 << 14;
     /** Usage: The buffer will be read by a hardware video encoder */
diff --git a/core/java/android/service/wallpaper/WallpaperService.java b/core/java/android/service/wallpaper/WallpaperService.java
index 7b8410b..f9b909f 100644
--- a/core/java/android/service/wallpaper/WallpaperService.java
+++ b/core/java/android/service/wallpaper/WallpaperService.java
@@ -44,7 +44,6 @@
 import android.graphics.BLASTBufferQueue;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
-import android.graphics.GraphicBuffer;
 import android.graphics.Matrix;
 import android.graphics.PixelFormat;
 import android.graphics.Point;
@@ -2041,9 +2040,7 @@
 
             mScreenshotSize.set(mSurfaceSize.x, mSurfaceSize.y);
 
-            GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(hardwareBuffer);
-
-            t.setBuffer(mScreenshotSurfaceControl, graphicBuffer);
+            t.setBuffer(mScreenshotSurfaceControl, hardwareBuffer);
             t.setColorSpace(mScreenshotSurfaceControl, screenshotBuffer.getColorSpace());
             // Place on top everything else.
             t.setLayer(mScreenshotSurfaceControl, Integer.MAX_VALUE);
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index af7d86c..3b52709 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -120,6 +120,8 @@
             long relativeToObject, int zorder);
     private static native void nativeSetPosition(long transactionObj, long nativeObject,
             float x, float y);
+    private static native void nativeSetScale(long transactionObj, long nativeObject,
+            float x, float y);
     private static native void nativeSetSize(long transactionObj, long nativeObject, int w, int h);
     private static native void nativeSetTransparentRegionHint(long transactionObj,
             long nativeObject, Region region);
@@ -202,9 +204,13 @@
     private static native void nativeReparent(long transactionObj, long nativeObject,
             long newParentNativeObject);
     private static native void nativeSetBuffer(long transactionObj, long nativeObject,
-            GraphicBuffer buffer);
+            HardwareBuffer buffer);
+    private static native void nativeSetBufferTransform(long transactionObj, long nativeObject,
+            int transform);
     private static native void nativeSetColorSpace(long transactionObj, long nativeObject,
             int colorSpace);
+    private static native void nativeSetDamageRegion(long transactionObj, long nativeObject,
+            Region region);
 
     private static native void nativeOverrideHdrTypes(IBinder displayToken, int[] modes);
 
@@ -1333,7 +1339,6 @@
          * Set the initial visibility for the SurfaceControl.
          *
          * @param hidden Whether the Surface is initially HIDDEN.
-         * @hide
          */
         @NonNull
         public Builder setHidden(boolean hidden) {
@@ -2840,16 +2845,38 @@
         }
 
         /**
-         * @hide
+         * Sets the SurfaceControl to the specified position relative to the parent
+         * SurfaceControl
+         *
+         * @param sc The SurfaceControl to change position
+         * @param x the X position
+         * @param y the Y position
+         * @return this transaction
          */
-        @UnsupportedAppUsage
-        public Transaction setPosition(SurfaceControl sc, float x, float y) {
+        @NonNull
+        public Transaction setPosition(@NonNull SurfaceControl sc, float x, float y) {
             checkPreconditions(sc);
             nativeSetPosition(mNativeObject, sc.mNativeObject, x, y);
             return this;
         }
 
         /**
+         * Sets the SurfaceControl to the specified scale with (0, 0) as the center point
+         * of the scale.
+         *
+         * @param sc The SurfaceControl to change scale
+         * @param scaleX the X scale
+         * @param scaleY the Y scale
+         * @return this transaction
+         */
+        @NonNull
+        public Transaction setScale(@NonNull SurfaceControl sc, float scaleX, float scaleY) {
+            checkPreconditions(sc);
+            nativeSetScale(mNativeObject, sc.mNativeObject, scaleX, scaleY);
+            return this;
+        }
+
+        /**
          * Set the default buffer size for the SurfaceControl, if there is a
          * {@link Surface} associated with the control, then
          * this will be the default size for buffers dequeued from it.
@@ -3056,7 +3083,9 @@
          * @param sc   SurfaceControl to set crop of.
          * @param crop Bounds of the crop to apply.
          * @hide
+         * @deprecated Use {@link #setCrop(SurfaceControl, Rect)} instead.
          */
+        @Deprecated
         @UnsupportedAppUsage
         public Transaction setWindowCrop(SurfaceControl sc, Rect crop) {
             checkPreconditions(sc);
@@ -3071,6 +3100,28 @@
         }
 
         /**
+         * Bounds the surface and its children to the bounds specified. Size of the surface will be
+         * ignored and only the crop and buffer size will be used to determine the bounds of the
+         * surface. If no crop is specified and the surface has no buffer, the surface bounds is
+         * only constrained by the size of its parent bounds.
+         *
+         * @param sc   SurfaceControl to set crop of.
+         * @param crop Bounds of the crop to apply.
+         * @return this This transaction for chaining
+         */
+        public @NonNull Transaction setCrop(@NonNull SurfaceControl sc, @Nullable Rect crop) {
+            checkPreconditions(sc);
+            if (crop != null) {
+                nativeSetWindowCrop(mNativeObject, sc.mNativeObject,
+                        crop.left, crop.top, crop.right, crop.bottom);
+            } else {
+                nativeSetWindowCrop(mNativeObject, sc.mNativeObject, 0, 0, 0, 0);
+            }
+
+            return this;
+        }
+
+        /**
          * Same as {@link Transaction#setWindowCrop(SurfaceControl, Rect)} but sets the crop rect
          * top left at 0, 0.
          *
@@ -3215,11 +3266,34 @@
         }
 
         /**
-         * Sets the opacity of the surface.  Setting the flag is equivalent to creating the
-         * Surface with the {@link #OPAQUE} flag.
-         * @hide
+         * Indicates whether the surface must be considered opaque, even if its pixel format is
+         * set to translucent. This can be useful if an application needs full RGBA 8888 support
+         * for instance but will still draw every pixel opaque.
+         * <p>
+         * This flag only determines whether opacity will be sampled from the alpha channel.
+         * Plane-alpha from calls to setAlpha() can still result in blended composition
+         * regardless of the opaque setting.
+         *
+         * Combined effects are (assuming a buffer format with an alpha channel):
+         * <ul>
+         * <li>OPAQUE + alpha(1.0) == opaque composition
+         * <li>OPAQUE + alpha(0.x) == blended composition
+         * <li>OPAQUE + alpha(0.0) == no composition
+         * <li>!OPAQUE + alpha(1.0) == blended composition
+         * <li>!OPAQUE + alpha(0.x) == blended composition
+         * <li>!OPAQUE + alpha(0.0) == no composition
+         * </ul>
+         * If the underlying buffer lacks an alpha channel, it is as if setOpaque(true)
+         * were set automatically.
+         *
+         * @see Builder#setOpaque(boolean)
+         *
+         * @param sc The SurfaceControl to update
+         * @param isOpaque true if the buffer's alpha should be ignored, false otherwise
+         * @return this
          */
-        public Transaction setOpaque(SurfaceControl sc, boolean isOpaque) {
+        @NonNull
+        public Transaction setOpaque(@NonNull SurfaceControl sc, boolean isOpaque) {
             checkPreconditions(sc);
             if (isOpaque) {
                 nativeSetFlags(mNativeObject, sc.mNativeObject, SURFACE_OPAQUE, SURFACE_OPAQUE);
@@ -3498,14 +3572,57 @@
          * created as type {@link #FX_SURFACE_BLAST}
          *
          * @hide
+         * @deprecated Use {@link #setBuffer(SurfaceControl, HardwareBuffer)} instead
          */
+        @Deprecated
         public Transaction setBuffer(SurfaceControl sc, GraphicBuffer buffer) {
+            return setBuffer(sc, HardwareBuffer.createFromGraphicBuffer(buffer));
+        }
+
+        /**
+         * Updates the HardwareBuffer displayed for the SurfaceControl.
+         *
+         * Note that the buffer must be allocated with {@link HardwareBuffer#USAGE_COMPOSER_OVERLAY}
+         * as well as {@link HardwareBuffer#USAGE_GPU_SAMPLED_IMAGE} as the surface control might
+         * be composited using either an overlay or using the GPU.
+         *
+         * @param sc The SurfaceControl to update
+         * @param buffer The buffer to be displayed
+         * @return this
+         */
+        public @NonNull Transaction setBuffer(@NonNull SurfaceControl sc,
+                @Nullable HardwareBuffer buffer) {
             checkPreconditions(sc);
             nativeSetBuffer(mNativeObject, sc.mNativeObject, buffer);
             return this;
         }
 
         /**
+         * Sets the buffer transform that should be applied to the current buffer.
+         *
+         * @param sc The SurfaceControl to update
+         * @param transform The transform to apply to the buffer.
+         * @return this
+         */
+        public @NonNull Transaction setBufferTransform(@NonNull SurfaceControl sc,
+                /* TODO: Mark the intdef */ int transform) {
+            checkPreconditions(sc);
+            nativeSetBufferTransform(mNativeObject, sc.mNativeObject, transform);
+            return this;
+        }
+
+        /**
+         * Updates the region for the content on this surface updated in this transaction.
+         *
+         * If unspecified, the complete surface is assumed to be damaged.
+         */
+        public @NonNull Transaction setDamageRegion(@NonNull SurfaceControl sc,
+                @Nullable Region region) {
+            nativeSetDamageRegion(mNativeObject, sc.mNativeObject, region);
+            return this;
+        }
+
+        /**
          * Set the color space for the SurfaceControl. The supported color spaces are SRGB
          * and Display P3, other color spaces will be treated as SRGB. This can only be used for
          * SurfaceControls that were created as type {@link #FX_SURFACE_BLAST}
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index fb5fded..67d0c52 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -597,6 +597,14 @@
     transaction->setPosition(ctrl, x, y);
 }
 
+static void nativeSetScale(JNIEnv* env, jclass clazz, jlong transactionObj, jlong nativeObject,
+                           jfloat xScale, jfloat yScale) {
+    auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
+
+    SurfaceControl* const ctrl = reinterpret_cast<SurfaceControl*>(nativeObject);
+    transaction->setMatrix(ctrl, xScale, 0, 0, yScale);
+}
+
 static void nativeSetGeometry(JNIEnv* env, jclass clazz, jlong transactionObj, jlong nativeObject,
         jobject sourceObj, jobject dstObj, jlong orientation) {
     auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
@@ -620,9 +628,19 @@
                             jobject bufferObject) {
     auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
     SurfaceControl* const ctrl = reinterpret_cast<SurfaceControl*>(nativeObject);
-    sp<GraphicBuffer> buffer(
-            android_graphics_GraphicBuffer_getNativeGraphicsBuffer(env, bufferObject));
-    transaction->setBuffer(ctrl, buffer);
+    sp<GraphicBuffer> graphicBuffer(GraphicBuffer::fromAHardwareBuffer(
+            android_hardware_HardwareBuffer_getNativeHardwareBuffer(env, bufferObject)));
+    transaction->setBuffer(ctrl, graphicBuffer);
+}
+
+static void nativeSetBufferTransform(JNIEnv* env, jclass clazz, jlong transactionObj,
+                                     jlong nativeObject, jint transform) {
+    auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
+    SurfaceControl* const ctrl = reinterpret_cast<SurfaceControl*>(nativeObject);
+    transaction->setTransform(ctrl, transform);
+    bool transformToInverseDisplay = (NATIVE_WINDOW_TRANSFORM_INVERSE_DISPLAY & transform) ==
+            NATIVE_WINDOW_TRANSFORM_INVERSE_DISPLAY;
+    transaction->setTransformToDisplayInverse(ctrl, transformToInverseDisplay);
 }
 
 static void nativeSetColorSpace(JNIEnv* env, jclass clazz, jlong transactionObj, jlong nativeObject,
@@ -740,6 +758,37 @@
     }
 }
 
+static void nativeSetDamageRegion(JNIEnv* env, jclass clazz, jlong transactionObj,
+                                  jlong nativeObject, jobject regionObj) {
+    SurfaceControl* const surfaceControl = reinterpret_cast<SurfaceControl*>(nativeObject);
+    auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
+
+    if (regionObj == nullptr) {
+        transaction->setSurfaceDamageRegion(surfaceControl, Region::INVALID_REGION);
+        return;
+    }
+
+    graphics::RegionIterator iterator(env, regionObj);
+    if (!iterator.isValid()) {
+        transaction->setSurfaceDamageRegion(surfaceControl, Region::INVALID_REGION);
+        return;
+    }
+
+    Region region;
+    while (!iterator.isDone()) {
+        ARect rect = iterator.getRect();
+        region.orSelf(static_cast<const Rect&>(rect));
+        iterator.next();
+    }
+
+    if (region.getBounds().isEmpty()) {
+        transaction->setSurfaceDamageRegion(surfaceControl, Region::INVALID_REGION);
+        return;
+    }
+
+    transaction->setSurfaceDamageRegion(surfaceControl, region);
+}
+
 static void nativeSetAlpha(JNIEnv* env, jclass clazz, jlong transactionObj,
         jlong nativeObject, jfloat alpha) {
     auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
@@ -1905,10 +1954,14 @@
             (void*)nativeSetRelativeLayer },
     {"nativeSetPosition", "(JJFF)V",
             (void*)nativeSetPosition },
+    {"nativeSetScale", "(JJFF)V",
+            (void*)nativeSetScale },
     {"nativeSetSize", "(JJII)V",
             (void*)nativeSetSize },
     {"nativeSetTransparentRegionHint", "(JJLandroid/graphics/Region;)V",
             (void*)nativeSetTransparentRegionHint },
+    { "nativeSetDamageRegion", "(JJLandroid/graphics/Region;)V",
+            (void*)nativeSetDamageRegion },
     {"nativeSetAlpha", "(JJF)V",
             (void*)nativeSetAlpha },
     {"nativeSetColor", "(JJ[F)V",
@@ -2018,8 +2071,9 @@
             (void*)nativeGetDisplayedContentSample },
     {"nativeSetGeometry", "(JJLandroid/graphics/Rect;Landroid/graphics/Rect;J)V",
             (void*)nativeSetGeometry },
-    {"nativeSetBuffer", "(JJLandroid/graphics/GraphicBuffer;)V",
+    {"nativeSetBuffer", "(JJLandroid/hardware/HardwareBuffer;)V",
             (void*)nativeSetBuffer },
+    {"nativeSetBufferTransform", "(JJI)V", (void*) nativeSetBufferTransform},
     {"nativeSetColorSpace", "(JJI)V",
             (void*)nativeSetColorSpace },
     {"nativeSyncInputWindows", "(J)V",
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java
index eea6e3c..c4bd73b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java
@@ -16,7 +16,6 @@
 
 package com.android.wm.shell.common;
 
-import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.view.SurfaceControl;
@@ -63,8 +62,6 @@
             if (buffer == null || buffer.getHardwareBuffer() == null) {
                 return;
             }
-            final GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(
-                buffer.getHardwareBuffer());
             mScreenshot = new SurfaceControl.Builder()
                 .setName("ScreenshotUtils screenshot")
                 .setFormat(PixelFormat.TRANSLUCENT)
@@ -73,7 +70,7 @@
                 .setBLASTLayer()
                 .build();
 
-            mTransaction.setBuffer(mScreenshot, graphicBuffer);
+            mTransaction.setBuffer(mScreenshot, buffer.getHardwareBuffer());
             mTransaction.setColorSpace(mScreenshot, buffer.getColorSpace());
             mTransaction.reparent(mScreenshot, mSurfaceControl);
             mTransaction.setLayer(mScreenshot, mLayer);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
index 6643ca1..4ecc0b6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
@@ -380,9 +380,7 @@
     }
 
     private void drawSizeMatchSnapshot() {
-        GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(
-                mSnapshot.getHardwareBuffer());
-        mTransaction.setBuffer(mSurfaceControl, graphicBuffer)
+        mTransaction.setBuffer(mSurfaceControl, mSnapshot.getHardwareBuffer())
                 .setColorSpace(mSurfaceControl, mSnapshot.getColorSpace())
                 .apply();
     }
@@ -428,20 +426,20 @@
         // Scale the mismatch dimensions to fill the task bounds
         mSnapshotMatrix.setRectToRect(mTmpSnapshotSize, mTmpDstFrame, Matrix.ScaleToFit.FILL);
         mTransaction.setMatrix(childSurfaceControl, mSnapshotMatrix, mTmpFloat9);
-        GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(
-                mSnapshot.getHardwareBuffer());
         mTransaction.setColorSpace(childSurfaceControl, mSnapshot.getColorSpace());
-        mTransaction.setBuffer(childSurfaceControl, graphicBuffer);
+        mTransaction.setBuffer(childSurfaceControl, mSnapshot.getHardwareBuffer());
 
         if (aspectRatioMismatch) {
             GraphicBuffer background = GraphicBuffer.create(mFrame.width(), mFrame.height(),
                     PixelFormat.RGBA_8888,
                     GraphicBuffer.USAGE_HW_TEXTURE | GraphicBuffer.USAGE_HW_COMPOSER
                             | GraphicBuffer.USAGE_SW_WRITE_RARELY);
+            // TODO: Support this on HardwareBuffer
             final Canvas c = background.lockCanvas();
             drawBackgroundAndBars(c, frame);
             background.unlockCanvasAndPost(c);
-            mTransaction.setBuffer(mSurfaceControl, background);
+            mTransaction.setBuffer(mSurfaceControl,
+                    HardwareBuffer.createFromGraphicBuffer(background));
         }
         mTransaction.apply();
         childSurfaceControl.release();
diff --git a/native/android/surface_control.cpp b/native/android/surface_control.cpp
index ceba4d6..f76811e 100644
--- a/native/android/surface_control.cpp
+++ b/native/android/surface_control.cpp
@@ -364,7 +364,7 @@
     sp<SurfaceControl> surfaceControl = ASurfaceControl_to_SurfaceControl(aSurfaceControl);
     Transaction* transaction = ASurfaceTransaction_to_Transaction(aSurfaceTransaction);
 
-    sp<GraphicBuffer> graphic_buffer(reinterpret_cast<GraphicBuffer*>(buffer));
+    sp<GraphicBuffer> graphic_buffer(GraphicBuffer::fromAHardwareBuffer(buffer));
 
     std::optional<sp<Fence>> fence = std::nullopt;
     if (acquire_fence_fd != -1) {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 68eff9a..c5e120e 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -166,7 +166,6 @@
 import android.content.res.Configuration;
 import android.graphics.Bitmap;
 import android.graphics.ColorSpace;
-import android.graphics.GraphicBuffer;
 import android.graphics.Insets;
 import android.graphics.Point;
 import android.graphics.Rect;
@@ -4088,8 +4087,7 @@
         // Make IME snapshot as trusted overlay
         InputMonitor.setTrustedOverlayInputInfo(imeSurface, t, getDisplayId(),
                 "IME-snapshot-surface");
-        GraphicBuffer graphicBuffer = GraphicBuffer.createFromHardwareBuffer(buffer);
-        t.setBuffer(imeSurface, graphicBuffer);
+        t.setBuffer(imeSurface, buffer);
         t.setColorSpace(mSurfaceControl, ColorSpace.get(ColorSpace.Named.SRGB));
         t.setRelativeLayer(imeSurface, activity.getSurfaceControl(), 1);
         t.setPosition(imeSurface, mInputMethodWindow.getDisplayFrame().left,