Merge "[RenderScript] Add API to map Allocation mallocptr to Java ByteBuffer"

3 files changed, 107 insertions, 5 deletions

diff --git a/rs/java/android/renderscript/Allocation.java b/rs/java/android/renderscript/Allocation.java
index a4876b92fadf..a71ba6364ac8 100644
--- a/rs/java/android/renderscript/Allocation.java
+++ b/rs/java/android/renderscript/Allocation.java
@@ -16,14 +16,16 @@
 
 package android.renderscript;
 
+import java.nio.ByteBuffer;
 import java.util.HashMap;
+
 import android.content.res.Resources;
 import android.graphics.Bitmap;
 import android.graphics.BitmapFactory;
-import android.view.Surface;
-import android.util.Log;
 import android.graphics.Canvas;
 import android.os.Trace;
+import android.util.Log;
+import android.view.Surface;
 
 /**
  * <p> This class provides the primary method through which data is passed to
@@ -78,6 +80,8 @@ public class Allocation extends BaseObj {
     OnBufferAvailableListener mBufferNotifier;
 
     private Surface mGetSurfaceSurface = null;
+    private ByteBuffer mByteBuffer = null;
+    private long mByteBufferStride = -1;
 
     private Element.DataType validateObjectIsPrimitiveArray(Object d, boolean checkType) {
         final Class c = d.getClass();
@@ -2050,6 +2054,59 @@ public class Allocation extends BaseObj {
     }
 
     /**
+     * @hide
+     * Gets or creates a ByteBuffer that contains the raw data of the current Allocation.
+     * If the Allocation is created with USAGE_IO_INPUT, the returned ByteBuffer
+     * would contain the up-to-date data as READ ONLY.
+     * For a 2D or 3D Allocation, the raw data maybe padded so that each row of
+     * the Allocation has certain alignment. The size of each row including padding,
+     * called stride, can be queried using the {@link #getStride()} method.
+     *
+     * Note: Operating on the ByteBuffer of a destroyed Allocation will triger errors.
+     *
+     * @return ByteBuffer The ByteBuffer associated with raw data pointer of the Allocation.
+     */
+    public ByteBuffer getByteBuffer() {
+        // Create a new ByteBuffer if it is not initialized or using IO_INPUT.
+        if (mType.hasFaces()) {
+            throw new RSInvalidStateException("Cubemap is not supported for getByteBuffer().");
+        }
+        if (mType.getYuv() == android.graphics.ImageFormat.NV21 ||
+            mType.getYuv() == android.graphics.ImageFormat.YV12 ||
+            mType.getYuv() == android.graphics.ImageFormat.YUV_420_888 ) {
+            throw new RSInvalidStateException("YUV format is not supported for getByteBuffer().");
+        }
+        if (mByteBuffer == null || (mUsage & USAGE_IO_INPUT) != 0) {
+            int xBytesSize = mType.getX() * mType.getElement().getBytesSize();
+            long[] stride = new long[1];
+            mByteBuffer = mRS.nAllocationGetByteBuffer(getID(mRS), stride, xBytesSize, mType.getY(), mType.getZ());
+            mByteBufferStride = stride[0];
+        }
+        if ((mUsage & USAGE_IO_INPUT) != 0) {
+            return mByteBuffer.asReadOnlyBuffer();
+        }
+        return mByteBuffer;
+    }
+
+    /**
+     * @hide
+     * Gets the stride of the Allocation.
+     * For a 2D or 3D Allocation, the raw data maybe padded so that each row of
+     * the Allocation has certain alignment. The size of each row including such
+     * padding is called stride.
+     *
+     * @return the stride. For 1D Allocation, the stride will be the number of
+     *         bytes of this Allocation. For 2D and 3D Allocations, the stride
+     *         will be the stride in X dimension measuring in bytes.
+     */
+    public long getStride() {
+        if (mByteBufferStride == -1) {
+            getByteBuffer();
+        }
+        return mByteBufferStride;
+    }
+
+    /**
      * Returns the handle to a raw buffer that is being managed by the screen
      * compositor. This operation is only valid for Allocations with {@link
      * #USAGE_IO_INPUT}.
diff --git a/rs/java/android/renderscript/RenderScript.java b/rs/java/android/renderscript/RenderScript.java
index 3668ddd62303..4788223b4357 100644
--- a/rs/java/android/renderscript/RenderScript.java
+++ b/rs/java/android/renderscript/RenderScript.java
@@ -18,17 +18,18 @@ package android.renderscript;
 
 import java.io.File;
 import java.lang.reflect.Method;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import android.content.Context;
 import android.content.res.AssetManager;
 import android.graphics.Bitmap;
 import android.graphics.SurfaceTexture;
-import android.util.Log;
-import android.view.Surface;
 import android.os.SystemProperties;
 import android.os.Trace;
-import java.util.ArrayList;
+import android.util.Log;
+import android.view.Surface;
 
 // TODO: Clean up the whitespace that separates methods in this class.
 
@@ -489,6 +490,13 @@ public class RenderScript {
         validate();
         rsnAllocationSyncAll(mContext, alloc, src);
     }
+
+    native ByteBuffer rsnAllocationGetByteBuffer(long con, long alloc, long[] stride, int xBytesSize, int dimY, int dimZ);
+    synchronized ByteBuffer nAllocationGetByteBuffer(long alloc, long[] stride, int xBytesSize, int dimY, int dimZ) {
+        validate();
+        return rsnAllocationGetByteBuffer(mContext, alloc, stride, xBytesSize, dimY, dimZ);
+    }
+
     native Surface rsnAllocationGetSurface(long con, long alloc);
     synchronized Surface nAllocationGetSurface(long alloc) {
         validate();
diff --git a/rs/jni/android_renderscript_RenderScript.cpp b/rs/jni/android_renderscript_RenderScript.cpp
index 39540709b16e..398d89b20b1e 100644
--- a/rs/jni/android_renderscript_RenderScript.cpp
+++ b/rs/jni/android_renderscript_RenderScript.cpp
@@ -2752,7 +2752,43 @@ nSystemGetPointerSize(JNIEnv *_env, jobject _this) {
     return (jint)sizeof(void*);
 }
 
+static jobject
+nAllocationGetByteBuffer(JNIEnv *_env, jobject _this, jlong con, jlong alloc,
+                        jlongArray strideArr, jint xBytesSize,
+                        jint dimY, jint dimZ) {
+    if (kLogApi) {
+        ALOGD("nAllocationGetByteBuffer, con(%p), alloc(%p)", (RsContext)con, (RsAllocation)alloc);
+    }
+
+    jlong *jStridePtr = _env->GetLongArrayElements(strideArr, nullptr);
+    if (jStridePtr == nullptr) {
+        ALOGE("Failed to get Java array elements: strideArr");
+        return 0;
+    }
 
+    size_t strideIn = xBytesSize;
+    void* ptr = nullptr;
+    if (alloc != 0) {
+        ptr = rsAllocationGetPointer((RsContext)con, (RsAllocation)alloc, 0,
+                                     RS_ALLOCATION_CUBEMAP_FACE_POSITIVE_X, 0, 0,
+                                     &strideIn, sizeof(size_t));
+    }
+
+    jobject byteBuffer = nullptr;
+    if (ptr != nullptr) {
+        size_t bufferSize = strideIn;
+        jStridePtr[0] = strideIn;
+        if (dimY > 0) {
+            bufferSize *= dimY;
+        }
+        if (dimZ > 0) {
+            bufferSize *= dimZ;
+        }
+        byteBuffer = _env->NewDirectByteBuffer(ptr, (jlong) bufferSize);
+    }
+    _env->ReleaseLongArrayElements(strideArr, jStridePtr, 0);
+    return byteBuffer;
+}
 // ---------------------------------------------------------------------------
 
 
@@ -2909,6 +2945,7 @@ static const JNINativeMethod methods[] = {
 {"rsnMeshGetIndices",                "(JJ[J[II)V",                            (void*)nMeshGetIndices },
 
 {"rsnSystemGetPointerSize",          "()I",                                   (void*)nSystemGetPointerSize },
+{"rsnAllocationGetByteBuffer",       "(JJ[JIII)Ljava/nio/ByteBuffer;",        (void*)nAllocationGetByteBuffer },
 };
 
 static int registerFuncs(JNIEnv *_env)