Enabling large audio for BLOCK_MODEL

Bug: 298052174
API-Coverage-Bug: 309692716
Change-Id: I22d261f424c1d890d311f593c652d946fcb065fa
diff --git a/core/api/current.txt b/core/api/current.txt
index b59721c..3009730 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -22723,6 +22723,7 @@
   }
 
   public static final class MediaCodec.OutputFrame {
+    method @FlaggedApi("com.android.media.codec.flags.large_audio_frame") @NonNull public java.util.ArrayDeque<android.media.MediaCodec.BufferInfo> getBufferInfos();
     method @NonNull public java.util.Set<java.lang.String> getChangedKeys();
     method public int getFlags();
     method @NonNull public android.media.MediaFormat getFormat();
@@ -22738,6 +22739,7 @@
 
   public final class MediaCodec.QueueRequest {
     method public void queue();
+    method @FlaggedApi("com.android.media.codec.flags.large_audio_frame") @NonNull public android.media.MediaCodec.QueueRequest setBufferInfos(@NonNull java.util.ArrayDeque<android.media.MediaCodec.BufferInfo>);
     method @NonNull public android.media.MediaCodec.QueueRequest setByteBufferParameter(@NonNull String, @NonNull java.nio.ByteBuffer);
     method @NonNull public android.media.MediaCodec.QueueRequest setEncryptedLinearBlock(@NonNull android.media.MediaCodec.LinearBlock, int, int, @NonNull android.media.MediaCodec.CryptoInfo);
     method @NonNull public android.media.MediaCodec.QueueRequest setFlags(int);
diff --git a/media/java/android/media/MediaCodec.java b/media/java/android/media/MediaCodec.java
index 26fb8dd..862ae8d 100644
--- a/media/java/android/media/MediaCodec.java
+++ b/media/java/android/media/MediaCodec.java
@@ -1962,10 +1962,20 @@
                                         index, infos);
                                 break;
                             case BUFFER_MODE_BLOCK:
-                                // TODO
+                                while (mOutputFrames.size() <= index) {
+                                    mOutputFrames.add(null);
+                                }
+                                OutputFrame frame = mOutputFrames.get(index);
+                                if (frame == null) {
+                                    frame = new OutputFrame(index);
+                                    mOutputFrames.set(index, frame);
+                                }
+                                frame.setBufferInfos(infos);
+                                frame.setAccessible(true);
+                                break;
                             default:
                                 throw new IllegalArgumentException(
-                                        "Unrecognized buffer mode: for large frame audio");
+                                        "Unrecognized buffer mode: for large frame output");
                         }
                     }
                     mCallback.onOutputBuffersAvailable(
@@ -3552,6 +3562,26 @@
         }
 
         /**
+         * Sets MediaCodec.BufferInfo objects describing the access units
+         * contained in this queue request. Access units must be laid out
+         * contiguously without gaps and in order.
+         *
+         * @param infos Represents {@link MediaCodec.BufferInfo} objects to mark
+         *              individual access-unit boundaries and the timestamps associated with it.
+         *              The buffer is expected to contain the data in a continuous manner.
+         * @return this object
+         */
+        @FlaggedApi(FLAG_LARGE_AUDIO_FRAME)
+        public @NonNull QueueRequest setBufferInfos(@NonNull ArrayDeque<BufferInfo> infos) {
+            if (!isAccessible()) {
+                throw new IllegalStateException("The request is stale");
+            }
+            mBufferInfos.clear();
+            mBufferInfos.addAll(infos);
+            return this;
+        }
+
+        /**
          * Add an integer parameter.
          * See {@link MediaFormat} for an exhaustive list of supported keys with
          * values of type int, that can also be set with {@link MediaFormat#setInteger}.
@@ -3667,10 +3697,18 @@
                 throw new IllegalStateException("No block is set");
             }
             setAccessible(false);
+            if (mBufferInfos.isEmpty()) {
+                BufferInfo info = new BufferInfo();
+                info.size = mSize;
+                info.offset = mOffset;
+                info.presentationTimeUs = mPresentationTimeUs;
+                info.flags = mFlags;
+                mBufferInfos.add(info);
+            }
             if (mLinearBlock != null) {
                 mCodec.native_queueLinearBlock(
-                        mIndex, mLinearBlock, mOffset, mSize, mCryptoInfo,
-                        mPresentationTimeUs, mFlags,
+                        mIndex, mLinearBlock, mCryptoInfo,
+                        mBufferInfos.toArray(),
                         mTuningKeys, mTuningValues);
             } else if (mHardwareBuffer != null) {
                 mCodec.native_queueHardwareBuffer(
@@ -3688,6 +3726,7 @@
             mHardwareBuffer = null;
             mPresentationTimeUs = 0;
             mFlags = 0;
+            mBufferInfos.clear();
             mTuningKeys.clear();
             mTuningValues.clear();
             return this;
@@ -3711,6 +3750,7 @@
         private HardwareBuffer mHardwareBuffer = null;
         private long mPresentationTimeUs = 0;
         private @BufferFlag int mFlags = 0;
+        private final ArrayDeque<BufferInfo> mBufferInfos = new ArrayDeque<>();
         private final ArrayList<String> mTuningKeys = new ArrayList<>();
         private final ArrayList<Object> mTuningValues = new ArrayList<>();
 
@@ -3720,11 +3760,8 @@
     private native void native_queueLinearBlock(
             int index,
             @NonNull LinearBlock block,
-            int offset,
-            int size,
             @Nullable CryptoInfo cryptoInfo,
-            long presentationTimeUs,
-            int flags,
+            @NonNull Object[] bufferInfos,
             @NonNull ArrayList<String> keys,
             @NonNull ArrayList<Object> values);
 
@@ -4516,6 +4553,22 @@
             return mFlags;
         }
 
+        /*
+         * Returns the BufferInfos associated with this OutputFrame. These BufferInfos
+         * describes the access units present in the OutputFrame. Access units are laid
+         * out contiguously without gaps and in order.
+         */
+        @FlaggedApi(FLAG_LARGE_AUDIO_FRAME)
+        public @NonNull ArrayDeque<BufferInfo> getBufferInfos() {
+            if (mBufferInfos.isEmpty()) {
+                // single BufferInfo could be present.
+                BufferInfo bufferInfo = new BufferInfo();
+                bufferInfo.set(0, 0, mPresentationTimeUs, mFlags);
+                mBufferInfos.add(bufferInfo);
+            }
+            return mBufferInfos;
+        }
+
         /**
          * Returns a read-only {@link MediaFormat} for this frame. The returned
          * object is valid only until the client calls {@link MediaCodec#releaseOutputBuffer}.
@@ -4541,6 +4594,7 @@
             mLinearBlock = null;
             mHardwareBuffer = null;
             mFormat = null;
+            mBufferInfos.clear();
             mChangedKeys.clear();
             mKeySet.clear();
             mLoaded = false;
@@ -4559,6 +4613,11 @@
             mFlags = info.flags;
         }
 
+        void setBufferInfos(ArrayDeque<BufferInfo> infos) {
+            mBufferInfos.clear();
+            mBufferInfos.addAll(infos);
+        }
+
         boolean isLoaded() {
             return mLoaded;
         }
@@ -4573,6 +4632,7 @@
         private long mPresentationTimeUs = 0;
         private @BufferFlag int mFlags = 0;
         private MediaFormat mFormat = null;
+        private final ArrayDeque<BufferInfo> mBufferInfos = new ArrayDeque<>();
         private final ArrayList<String> mChangedKeys = new ArrayList<>();
         private final Set<String> mKeySet = new HashSet<>();
         private boolean mAccessible = false;
diff --git a/media/jni/android_media_MediaCodec.cpp b/media/jni/android_media_MediaCodec.cpp
index 4018747..8cdd59e 100644
--- a/media/jni/android_media_MediaCodec.cpp
+++ b/media/jni/android_media_MediaCodec.cpp
@@ -207,6 +207,7 @@
     jfieldID queueRequestIndexID;
     jfieldID outputFrameLinearBlockID;
     jfieldID outputFrameHardwareBufferID;
+    jfieldID outputFramebufferInfosID;
     jfieldID outputFrameChangedKeysID;
     jfieldID outputFrameFormatID;
     jfieldID bufferInfoFlags;
@@ -458,10 +459,11 @@
 }
 
 status_t JMediaCodec::queueBuffer(
-        size_t index, const std::shared_ptr<C2Buffer> &buffer, int64_t timeUs,
-        uint32_t flags, const sp<AMessage> &tunings, AString *errorDetailMsg) {
+        size_t index, const std::shared_ptr<C2Buffer> &buffer,
+        const sp<RefBase> &infos, const sp<AMessage> &tunings, AString *errorDetailMsg) {
+    sp<BufferInfosWrapper> auInfo((BufferInfosWrapper *)infos.get());
     return mCodec->queueBuffer(
-            index, buffer, timeUs, flags, tunings, errorDetailMsg);
+            index, buffer, auInfo, tunings, errorDetailMsg);
 }
 
 status_t JMediaCodec::queueEncryptedLinearBlock(
@@ -474,13 +476,13 @@
         const uint8_t iv[16],
         CryptoPlugin::Mode mode,
         const CryptoPlugin::Pattern &pattern,
-        int64_t presentationTimeUs,
-        uint32_t flags,
+        const sp<RefBase> &infos,
         const sp<AMessage> &tunings,
         AString *errorDetailMsg) {
+    sp<BufferInfosWrapper> auInfo((BufferInfosWrapper *)infos.get());
     return mCodec->queueEncryptedBuffer(
             index, buffer, offset, subSamples, numSubSamples, key, iv, mode, pattern,
-            presentationTimeUs, flags, tunings, errorDetailMsg);
+            auInfo, tunings, errorDetailMsg);
 }
 
 status_t JMediaCodec::dequeueInputBuffer(size_t *index, int64_t timeoutUs) {
@@ -750,6 +752,42 @@
     return OK;
 }
 
+void maybeSetBufferInfos(JNIEnv *env, jobject &frame, const sp<BufferInfosWrapper> &bufInfos) {
+    if (!bufInfos) {
+        return;
+    }
+    std::vector<AccessUnitInfo> &infos = bufInfos.get()->value;
+    if (infos.empty()) {
+        return;
+    }
+    ScopedLocalRef<jobject> dequeObj{env, env->NewObject(
+            gArrayDequeInfo.clazz, gArrayDequeInfo.ctorId)};
+    jint offset = 0;
+    std::vector<jobject> jObjectInfos;
+    for (int i = 0 ; i < infos.size(); i++) {
+        jobject bufferInfo = env->NewObject(
+                gBufferInfo.clazz, gBufferInfo.ctorId);
+        if (bufferInfo != NULL) {
+            env->CallVoidMethod(bufferInfo, gBufferInfo.setId,
+                    offset,
+                    (jint)(infos)[i].mSize,
+                    (infos)[i].mTimestamp,
+                    (infos)[i].mFlags);
+            (void)env->CallBooleanMethod(
+                    dequeObj.get(), gArrayDequeInfo.addId, bufferInfo);
+            offset += (infos)[i].mSize;
+            jObjectInfos.push_back(bufferInfo);
+        }
+    }
+    env->SetObjectField(
+            frame,
+            gFields.outputFramebufferInfosID,
+            dequeObj.get());
+    for (int i = 0; i < jObjectInfos.size(); i++) {
+        env->DeleteLocalRef(jObjectInfos[i]);
+    }
+}
+
 status_t JMediaCodec::getOutputFrame(
         JNIEnv *env, jobject frame, size_t index) const {
     sp<MediaCodecBuffer> buffer;
@@ -760,6 +798,11 @@
     }
 
     if (buffer->size() > 0) {
+        sp<RefBase> obj;
+        sp<BufferInfosWrapper> bufInfos;
+        if (buffer->meta()->findObject("accessUnitInfo", &obj)) {
+            bufInfos = std::move(((decltype(bufInfos.get()))obj.get()));
+        }
         std::shared_ptr<C2Buffer> c2Buffer = buffer->asC2Buffer();
         if (c2Buffer) {
             switch (c2Buffer->data().type()) {
@@ -775,6 +818,7 @@
                             (jlong)context.release(),
                             true);
                     env->SetObjectField(frame, gFields.outputFrameLinearBlockID, linearBlock.get());
+                    maybeSetBufferInfos(env, frame, bufInfos);
                     break;
                 }
                 case C2BufferData::GRAPHIC: {
@@ -815,6 +859,7 @@
                         (jlong)context.release(),
                         true);
                 env->SetObjectField(frame, gFields.outputFrameLinearBlockID, linearBlock.get());
+                maybeSetBufferInfos(env, frame, bufInfos);
             } else {
                 // No-op.
             }
@@ -2717,8 +2762,7 @@
 
 static void android_media_MediaCodec_native_queueLinearBlock(
         JNIEnv *env, jobject thiz, jint index, jobject bufferObj,
-        jint offset, jint size, jobject cryptoInfoObj,
-        jlong presentationTimeUs, jint flags, jobject keys, jobject values) {
+        jobject cryptoInfoObj, jobjectArray objArray, jobject keys, jobject values) {
     ALOGV("android_media_MediaCodec_native_queueLinearBlock");
 
     sp<JMediaCodec> codec = getMediaCodec(env, thiz);
@@ -2736,7 +2780,24 @@
                 "error occurred while converting tunings from Java to native");
         return;
     }
-
+    jint totalSize;
+    jint initialOffset;
+    std::vector<AccessUnitInfo> infoVec;
+    AString errorDetailMsg;
+    err = extractInfosFromObject(env,
+            &initialOffset,
+            &totalSize,
+            &infoVec,
+            objArray,
+            &errorDetailMsg);
+    if (err != OK) {
+        throwExceptionAsNecessary(
+                env, INVALID_OPERATION, ACTION_CODE_FATAL,
+                codec->getExceptionMessage(errorDetailMsg.c_str()).c_str());
+        return;
+    }
+    sp<BufferInfosWrapper> infos =
+            new BufferInfosWrapper{std::move(infoVec)};
     std::shared_ptr<C2Buffer> buffer;
     sp<hardware::HidlMemory> memory;
     ScopedLocalRef<jobject> lock{env, env->GetObjectField(bufferObj, gLinearBlockInfo.lockId)};
@@ -2745,10 +2806,10 @@
             JMediaCodecLinearBlock *context =
                 (JMediaCodecLinearBlock *)env->GetLongField(bufferObj, gLinearBlockInfo.contextId);
             if (codec->hasCryptoOrDescrambler()) {
-                extractMemoryFromContext(context, offset, size, &memory);
-                offset += context->mHidlMemoryOffset;
+                extractMemoryFromContext(context, initialOffset, totalSize, &memory);
+                initialOffset += context->mHidlMemoryOffset;
             } else {
-                extractBufferFromContext(context, offset, size, &buffer);
+                extractBufferFromContext(context, initialOffset, totalSize, &buffer);
             }
         }
         env->MonitorExit(lock.get());
@@ -2759,7 +2820,6 @@
         return;
     }
 
-    AString errorDetailMsg;
     if (codec->hasCryptoOrDescrambler()) {
         if (!memory) {
             // It means there was an unexpected failure in extractMemoryFromContext above
@@ -2773,7 +2833,7 @@
             return;
         }
         auto cryptoInfo =
-                cryptoInfoObj ? NativeCryptoInfo{env, cryptoInfoObj} : NativeCryptoInfo{size};
+                cryptoInfoObj ? NativeCryptoInfo{env, cryptoInfoObj} : NativeCryptoInfo{totalSize};
         if (env->ExceptionCheck()) {
             // Creation of cryptoInfo failed. Let the exception bubble up.
             return;
@@ -2781,13 +2841,12 @@
         err = codec->queueEncryptedLinearBlock(
                 index,
                 memory,
-                offset,
+                initialOffset,
                 cryptoInfo.mSubSamples, cryptoInfo.mNumSubSamples,
                 (const uint8_t *)cryptoInfo.mKey, (const uint8_t *)cryptoInfo.mIv,
                 cryptoInfo.mMode,
                 cryptoInfo.mPattern,
-                presentationTimeUs,
-                flags,
+                infos,
                 tunings,
                 &errorDetailMsg);
         ALOGI_IF(err != OK, "queueEncryptedLinearBlock returned err = %d", err);
@@ -2804,7 +2863,7 @@
             return;
         }
         err = codec->queueBuffer(
-                index, buffer, presentationTimeUs, flags, tunings, &errorDetailMsg);
+                index, buffer, infos, tunings, &errorDetailMsg);
     }
     throwExceptionAsNecessary(
             env, err, ACTION_CODE_FATAL,
@@ -2862,8 +2921,11 @@
     std::shared_ptr<C2Buffer> buffer = C2Buffer::CreateGraphicBuffer(block->share(
             block->crop(), C2Fence{}));
     AString errorDetailMsg;
+    sp<BufferInfosWrapper> infos =
+        new BufferInfosWrapper{decltype(infos->value)()};
+    infos->value.emplace_back(flags, 0 /*not used*/, presentationTimeUs);
     err = codec->queueBuffer(
-            index, buffer, presentationTimeUs, flags, tunings, &errorDetailMsg);
+            index, buffer, infos, tunings, &errorDetailMsg);
     throwExceptionAsNecessary(
             env, err, ACTION_CODE_FATAL,
             codec->getExceptionMessage(errorDetailMsg.c_str()).c_str());
@@ -3372,6 +3434,10 @@
         env->GetFieldID(clazz.get(), "mLinearBlock", "Landroid/media/MediaCodec$LinearBlock;");
     CHECK(gFields.outputFrameLinearBlockID != NULL);
 
+    gFields.outputFramebufferInfosID =
+        env->GetFieldID(clazz.get(), "mBufferInfos", "Ljava/util/ArrayDeque;");
+    CHECK(gFields.outputFramebufferInfosID != NULL);
+
     gFields.outputFrameHardwareBufferID =
         env->GetFieldID(clazz.get(), "mHardwareBuffer", "Landroid/hardware/HardwareBuffer;");
     CHECK(gFields.outputFrameHardwareBufferID != NULL);
@@ -3891,8 +3957,8 @@
     { "native_closeMediaImage", "(J)V", (void *)android_media_MediaCodec_closeMediaImage },
 
     { "native_queueLinearBlock",
-      "(ILandroid/media/MediaCodec$LinearBlock;IILandroid/media/MediaCodec$CryptoInfo;JI"
-      "Ljava/util/ArrayList;Ljava/util/ArrayList;)V",
+      "(ILandroid/media/MediaCodec$LinearBlock;Landroid/media/MediaCodec$CryptoInfo;"
+      "[Ljava/lang/Object;Ljava/util/ArrayList;Ljava/util/ArrayList;)V",
       (void *)android_media_MediaCodec_native_queueLinearBlock },
 
     { "native_queueHardwareBuffer",
diff --git a/media/jni/android_media_MediaCodec.h b/media/jni/android_media_MediaCodec.h
index 2bc1772..02708ef 100644
--- a/media/jni/android_media_MediaCodec.h
+++ b/media/jni/android_media_MediaCodec.h
@@ -116,7 +116,7 @@
 
     status_t queueBuffer(
             size_t index, const std::shared_ptr<C2Buffer> &buffer,
-            int64_t timeUs, uint32_t flags, const sp<AMessage> &tunings,
+            const sp<RefBase> &infos, const sp<AMessage> &tunings,
             AString *errorDetailMsg);
 
     status_t queueEncryptedLinearBlock(
@@ -129,8 +129,7 @@
             const uint8_t iv[16],
             CryptoPlugin::Mode mode,
             const CryptoPlugin::Pattern &pattern,
-            int64_t presentationTimeUs,
-            uint32_t flags,
+            const sp<RefBase> &infos,
             const sp<AMessage> &tunings,
             AString *errorDetailMsg);