From 3be49a5841b22b5eb47f43927b393c2384303ade Mon Sep 17 00:00:00 2001
From: andrewlewis <andrewlewis@google.com>
Date: Tue, 16 Apr 2024 06:19:08 -0700
Subject: [PATCH] Work around incorrect channel count for raw audio

The raw audio decoder's output audio format is stereo when the number of input
channels is (for example) 10 channels. Add a temporary workaround that uses the
input channel count for raw audio. This code should be removed in future when
we bypass the decoder entirely for raw audio.

Tested manually on a WAVE file with 18 audio channels.

PiperOrigin-RevId: 625307624
---
 RELEASENOTES.md                                           | 2 ++
 .../media3/transformer/ExoAssetLoaderBaseRenderer.java    | 8 ++++++++
 2 files changed, 10 insertions(+)

diff --git a/RELEASENOTES.md b/RELEASENOTES.md
index 718bb35eb0..e222a267b2 100644
--- a/RELEASENOTES.md
+++ b/RELEASENOTES.md
@@ -11,6 +11,8 @@
     *   Add `reset` to `BasePreloadManager` to release all the holding sources
         while keep the preload manager instance.
 *   Transformer:
+    *   Work around a decoder bug where the number of audio channels was capped
+        at stereo when handling PCM input.
 *   Track Selection:
 *   Extractors:
 *   Audio:
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java b/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java
index 9d9d99c692..ee8f67e18b 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java
@@ -262,6 +262,14 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
         if (decoderOutputFormat == null) {
           return false;
         }
+
+        // TODO: b/332708880 - The raw decoder caps number of audio channels to stereo. Remove this
+        //  workaround and bypass MediaCodec for raw audio instead.
+        if (checkNotNull(inputFormat.sampleMimeType).equals(MimeTypes.AUDIO_RAW)) {
+          decoderOutputFormat =
+              decoderOutputFormat.buildUpon().setChannelCount(inputFormat.channelCount).build();
+        }
+
         outputFormat = overrideOutputFormat(decoderOutputFormat);
       } else {
         // TODO(b/278259383): Move surface creation out of video sampleConsumer. Init decoder and