From b94c7d08c13210d16725b774f2eda8b1b808e2ab Mon Sep 17 00:00:00 2001
From: tofunmi <tofunmi@google.com>
Date: Wed, 24 Jan 2024 04:41:52 -0800
Subject: [PATCH] MuxerWrapper rotation degree fix

Allow setAdditionalRotationDegrees to be called with same rotation after tracks added. This is needed for processes that mux files partially trim optimization so they don't error out after hitting the check state

Manually tested to ensure trim optimization succeeds, automated test added here as well

PiperOrigin-RevId: 601081778
---
 .../transformer/TransformerEndToEndTest.java  | 13 +++++--
 .../media3/transformer/MuxerWrapper.java      |  2 +-
 .../media3/transformer/MuxerWrapperTest.java  | 35 +++++++++++++++++++
 3 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java
index 4af37d478a..971efef128 100644
--- a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java
+++ b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java
@@ -65,6 +65,7 @@ import androidx.media3.effect.DefaultVideoFrameProcessor;
 import androidx.media3.effect.FrameCache;
 import androidx.media3.effect.Presentation;
 import androidx.media3.effect.RgbFilter;
+import androidx.media3.effect.ScaleAndRotateTransformation;
 import androidx.media3.effect.TimestampWrapper;
 import androidx.media3.exoplayer.audio.TeeAudioProcessor;
 import androidx.test.core.app.ApplicationProvider;
@@ -587,8 +588,9 @@ public class TransformerEndToEndTest {
   }
 
   @Test
-  public void clippedMedia_trimOptimizationEnabled_audioRemoved_completesWithOptimizationApplied()
-      throws Exception {
+  public void
+      clippedMedia_trimOptimizationEnabled_audioRemovedAndRotated_completesWithOptimizationApplied()
+          throws Exception {
     String testId = "clippedMedia_trimOptimizationEnabled_completesWithOptimizationApplied";
     if (!isRunningOnEmulator() || Util.SDK_INT != 33) {
       // The trim optimization is only guaranteed to work on emulator for this (emulator-transcoded)
@@ -607,8 +609,13 @@ public class TransformerEndToEndTest {
                     .setEndPositionMs(2500)
                     .build())
             .build();
+    Effects effects =
+        new Effects(
+            /* audioProcessors= */ ImmutableList.of(),
+            ImmutableList.of(
+                new ScaleAndRotateTransformation.Builder().setRotationDegrees(90).build()));
     EditedMediaItem editedMediaItem =
-        new EditedMediaItem.Builder(mediaItem).setRemoveAudio(true).build();
+        new EditedMediaItem.Builder(mediaItem).setRemoveAudio(true).setEffects(effects).build();
 
     ExportTestResult result =
         new TransformerAndroidTestRunner.Builder(context, transformer)
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java b/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java
index 578b3065fe..aaa8324c25 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java
@@ -174,7 +174,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
    */
   public void setAdditionalRotationDegrees(int additionalRotationDegrees) {
     checkState(
-        trackTypeToInfo.size() == 0,
+        trackTypeToInfo.size() == 0 || this.additionalRotationDegrees == additionalRotationDegrees,
         "The additional rotation cannot be changed after adding track formats.");
     this.additionalRotationDegrees = additionalRotationDegrees;
   }
diff --git a/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java b/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java
index e89289b0bd..87a6a935c3 100644
--- a/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java
+++ b/libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java
@@ -72,6 +72,41 @@ public class MuxerWrapperTest {
     }
   }
 
+  @Test
+  public void setAdditionalRotationDegrees_sameRotationSetAfterTracksAdded_doesNotThrow()
+      throws Exception {
+    muxerWrapper =
+        new MuxerWrapper(
+            temporaryFolder.newFile().getPath(),
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_DEFAULT,
+            /* dropSamplesBeforeFirstVideoSample= */ false);
+    muxerWrapper.setAdditionalRotationDegrees(90);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.setAdditionalRotationDegrees(180);
+    muxerWrapper.addTrackFormat(FAKE_AUDIO_TRACK_FORMAT);
+    muxerWrapper.setAdditionalRotationDegrees(180);
+  }
+
+  @Test
+  public void setAdditionalRotationDegrees_differentRotationSetAfterTracksAdded_throws()
+      throws Exception {
+    muxerWrapper =
+        new MuxerWrapper(
+            temporaryFolder.newFile().getPath(),
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_DEFAULT,
+            /* dropSamplesBeforeFirstVideoSample= */ false);
+    muxerWrapper.setAdditionalRotationDegrees(90);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.setAdditionalRotationDegrees(180);
+    muxerWrapper.addTrackFormat(FAKE_AUDIO_TRACK_FORMAT);
+
+    assertThrows(IllegalStateException.class, () -> muxerWrapper.setAdditionalRotationDegrees(90));
+  }
+
   @Test
   public void changeToAppendMode_afterDefaultMode_throws() throws Exception {
     muxerWrapper =