From 0f72126c2059d28323f250937395744cee59263a Mon Sep 17 00:00:00 2001
From: claincly <claincly@google.com>
Date: Thu, 7 Mar 2024 07:38:20 -0800
Subject: [PATCH] Add `getLastTimestampUs()` method

PiperOrigin-RevId: 613579674
---
 .../util/ConstantRateTimestampIterator.java   | 31 +++++++++++++------
 .../media3/common/util/TimestampIterator.java | 10 ++++++
 .../ConstantRateTimestampIteratorTest.java    |  8 +++++
 3 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/libraries/common/src/main/java/androidx/media3/common/util/ConstantRateTimestampIterator.java b/libraries/common/src/main/java/androidx/media3/common/util/ConstantRateTimestampIterator.java
index a6a4eea301..d12e891667 100644
--- a/libraries/common/src/main/java/androidx/media3/common/util/ConstantRateTimestampIterator.java
+++ b/libraries/common/src/main/java/androidx/media3/common/util/ConstantRateTimestampIterator.java
@@ -34,8 +34,9 @@ public final class ConstantRateTimestampIterator implements TimestampIterator {
   private final float frameRate;
   private final double framesDurationUs;
   private final long startingTimestampUs;
-  private double currentTimestampUs;
-  private int framesToAdd;
+  private final int totalNumberOfFramesToAdd;
+
+  private int framesAdded;
 
   /**
    * Creates an instance that outputs timestamps from {@code 0}.
@@ -66,27 +67,39 @@ public final class ConstantRateTimestampIterator implements TimestampIterator {
     this.durationUs = durationUs;
     this.frameRate = frameRate;
     this.startingTimestampUs = startingTimestampUs;
-    this.currentTimestampUs = startingTimestampUs;
-    framesToAdd = round(frameRate * (durationUs / (float) C.MICROS_PER_SECOND));
+    this.totalNumberOfFramesToAdd = round(frameRate * (durationUs / (float) C.MICROS_PER_SECOND));
     framesDurationUs = C.MICROS_PER_SECOND / frameRate;
   }
 
   @Override
   public boolean hasNext() {
-    return framesToAdd != 0;
+    return framesAdded < totalNumberOfFramesToAdd;
   }
 
   @Override
   public long next() {
     checkState(hasNext());
-    framesToAdd--;
-    long next = round(currentTimestampUs);
-    currentTimestampUs += framesDurationUs;
-    return next;
+    return getTimestampUsAfter(framesAdded++);
   }
 
   @Override
   public ConstantRateTimestampIterator copyOf() {
     return new ConstantRateTimestampIterator(durationUs, frameRate, startingTimestampUs);
   }
+
+  @Override
+  public long getLastTimestampUs() {
+    if (totalNumberOfFramesToAdd == 0) {
+      return C.TIME_UNSET;
+    }
+    return getTimestampUsAfter(totalNumberOfFramesToAdd - 1);
+  }
+
+  /** Returns the timestamp after {@code numberOfFrames}, in microseconds. */
+  private long getTimestampUsAfter(int numberOfFrames) {
+    long timestampUs = round(startingTimestampUs + framesDurationUs * numberOfFrames);
+    // Check for possible overflow.
+    checkState(timestampUs >= 0);
+    return timestampUs;
+  }
 }
diff --git a/libraries/common/src/main/java/androidx/media3/common/util/TimestampIterator.java b/libraries/common/src/main/java/androidx/media3/common/util/TimestampIterator.java
index 40ccbd875a..9181e8ba4f 100644
--- a/libraries/common/src/main/java/androidx/media3/common/util/TimestampIterator.java
+++ b/libraries/common/src/main/java/androidx/media3/common/util/TimestampIterator.java
@@ -15,6 +15,8 @@
  */
 package androidx.media3.common.util;
 
+import androidx.media3.common.C;
+
 /** A primitive long iterator used for generating sequences of timestamps. */
 @UnstableApi
 public interface TimestampIterator {
@@ -27,4 +29,12 @@ public interface TimestampIterator {
 
   /** Returns fresh copy of the iterator. */
   TimestampIterator copyOf();
+
+  /**
+   * Returns the last (final) timestamp this iterator generates, in microseconds; returns {@link
+   * C#TIME_UNSET} if the last timestamp is unknown, or when no timestamp will be generated.
+   */
+  default long getLastTimestampUs() {
+    return C.TIME_UNSET;
+  }
 }
diff --git a/libraries/common/src/test/java/androidx/media3/common/util/ConstantRateTimestampIteratorTest.java b/libraries/common/src/test/java/androidx/media3/common/util/ConstantRateTimestampIteratorTest.java
index 76d2d30fb8..313604b3f5 100644
--- a/libraries/common/src/test/java/androidx/media3/common/util/ConstantRateTimestampIteratorTest.java
+++ b/libraries/common/src/test/java/androidx/media3/common/util/ConstantRateTimestampIteratorTest.java
@@ -16,6 +16,7 @@
 package androidx.media3.common.util;
 
 import static com.google.common.truth.Truth.assertThat;
+import static java.lang.Math.round;
 
 import androidx.media3.common.C;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -35,6 +36,8 @@ public class ConstantRateTimestampIteratorTest {
 
     assertThat(generateList(constantRateTimestampIterator))
         .containsExactly(0L, C.MICROS_PER_SECOND / 2);
+    assertThat(constantRateTimestampIterator.getLastTimestampUs())
+        .isEqualTo(C.MICROS_PER_SECOND / 2);
   }
 
   @Test
@@ -44,6 +47,8 @@ public class ConstantRateTimestampIteratorTest {
         new ConstantRateTimestampIterator((long) (2.5 * C.MICROS_PER_SECOND), /* frameRate= */ 30);
 
     assertThat(generateList(constantRateTimestampIterator)).hasSize(75);
+    assertThat(constantRateTimestampIterator.getLastTimestampUs())
+        .isEqualTo(round((C.MICROS_PER_SECOND / 30.d) * 74));
   }
 
   @Test
@@ -69,6 +74,7 @@ public class ConstantRateTimestampIteratorTest {
         new ConstantRateTimestampIterator(/* durationUs= */ 1, /* frameRate= */ 2);
 
     assertThat(generateList(constantRateTimestampIterator)).isEmpty();
+    assertThat(constantRateTimestampIterator.getLastTimestampUs()).isEqualTo(C.TIME_UNSET);
   }
 
   @Test
@@ -80,6 +86,8 @@ public class ConstantRateTimestampIteratorTest {
             /* startingTimestampUs= */ 1234);
 
     assertThat(constantRateTimestampIterator.next()).isEqualTo(1234);
+    assertThat(constantRateTimestampIterator.getLastTimestampUs())
+        .isEqualTo(1234 + C.MICROS_PER_SECOND / 2);
   }
 
   @Test