From 430fafded635ef5e6a7dc1503223c241f2e51a87 Mon Sep 17 00:00:00 2001
From: Googler <noreply@google.com>
Date: Mon, 22 Apr 2024 08:56:19 -0700
Subject: [PATCH] Boxes: Update STTS duration calculation.

Update the function convertPresentationTimestampsToDurationsVu
to return a list of duration in decoding order,used in the
creation of STTS boxes.

PiperOrigin-RevId: 627052898
---
 .../java/androidx/media3/muxer/Boxes.java     | 20 +++++++++++++++++--
 .../java/androidx/media3/muxer/BoxesTest.java | 16 +++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java b/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
index d87860b31c..94d31c0748 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
@@ -629,15 +629,31 @@ import java.util.Locale;
       long firstSamplePresentationTimeUs,
       int videoUnitTimescale,
       @Mp4Muxer.LastFrameDurationBehavior int lastDurationBehavior) {
+    List<Long> presentationTimestampsUs = new ArrayList<>(samplesInfo.size());
     List<Long> durationsVu = new ArrayList<>(samplesInfo.size());
 
     if (samplesInfo.isEmpty()) {
       return durationsVu;
     }
 
+    boolean hasBframe = false;
+    long lastSamplePresentationTimeUs = 0L;
+    for (int sampleId = 0; sampleId < samplesInfo.size(); sampleId++) {
+      long currentSamplePresentationTimeUs = samplesInfo.get(sampleId).presentationTimeUs;
+      presentationTimestampsUs.add(currentSamplePresentationTimeUs);
+      if (currentSamplePresentationTimeUs < lastSamplePresentationTimeUs) {
+        hasBframe = true;
+      }
+      lastSamplePresentationTimeUs = currentSamplePresentationTimeUs;
+    }
+
+    if (hasBframe) {
+      Collections.sort(presentationTimestampsUs);
+    }
+
     long currentSampleTimeUs = firstSamplePresentationTimeUs;
-    for (int nextSampleId = 1; nextSampleId < samplesInfo.size(); nextSampleId++) {
-      long nextSampleTimeUs = samplesInfo.get(nextSampleId).presentationTimeUs;
+    for (int nextSampleId = 1; nextSampleId < presentationTimestampsUs.size(); nextSampleId++) {
+      long nextSampleTimeUs = presentationTimestampsUs.get(nextSampleId);
       // TODO: b/316158030 - First calculate the duration and then convert us to vu to avoid
       //  rounding error.
       long currentSampleDurationVu =
diff --git a/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java b/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java
index 6ce162ae64..b93dc7fc52 100644
--- a/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java
+++ b/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java
@@ -461,6 +461,22 @@ public class BoxesTest {
     assertThat(durationsVu).containsExactly(3_000L, 5_000L, 5_000L);
   }
 
+  @Test
+  public void
+      convertPresentationTimestampsToDurationsVu_withOutOfOrderSampleTimestamps_returnsExpectedDurations() {
+    List<MediaCodec.BufferInfo> sampleBufferInfos =
+        createBufferInfoListWithSamplePresentationTimestamps(0L, 10_000L, 1_000L, 2_000L, 11_000L);
+
+    List<Long> durationsVu =
+        Boxes.convertPresentationTimestampsToDurationsVu(
+            sampleBufferInfos,
+            /* firstSamplePresentationTimeUs= */ 0L,
+            VU_TIMEBASE,
+            LAST_FRAME_DURATION_BEHAVIOR_INSERT_SHORT_FRAME);
+
+    assertThat(durationsVu).containsExactly(100L, 100L, 800L, 100L, 0L);
+  }
+
   @Test
   public void createSttsBox_withSingleSampleDuration_matchesExpected() throws IOException {
     ImmutableList<Long> sampleDurations = ImmutableList.of(500L);