From 67acfc67de92d15f3945af9c1b18125187ce0a1c Mon Sep 17 00:00:00 2001
From: aquilescanta <aquilescanta@google.com>
Date: Fri, 25 Mar 2022 15:00:57 +0000
Subject: [PATCH] Support seeking in un-intearleaved tracks in Mp4Extractor

A client can pass the id of the track on which they want to seek.

PiperOrigin-RevId: 437248055
---
 .../media3/extractor/mp4/Mp4Extractor.java    | 38 ++++++++++++++-----
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/Mp4Extractor.java b/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/Mp4Extractor.java
index f56a7dd7c0..4d9c5ec4cc 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/Mp4Extractor.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/mp4/Mp4Extractor.java
@@ -281,6 +281,22 @@ public final class Mp4Extractor implements Extractor, SeekMap {
 
   @Override
   public SeekPoints getSeekPoints(long timeUs) {
+    return getSeekPoints(timeUs, /* trackId= */ C.INDEX_UNSET);
+  }
+
+  // Non-inherited public methods.
+
+  /**
+   * Equivalent to {@link SeekMap#getSeekPoints(long)}, except it adds the {@code trackId}
+   * parameter.
+   *
+   * @param timeUs A seek time in microseconds.
+   * @param trackId The id of the track on which to seek for {@link SeekPoints}. May be {@link
+   *     C#INDEX_UNSET} if the extractor is expected to define the strategy for generating {@link
+   *     SeekPoints}.
+   * @return The corresponding seek points.
+   */
+  public SeekPoints getSeekPoints(long timeUs, int trackId) {
     if (tracks.length == 0) {
       return new SeekPoints(SeekPoint.START);
     }
@@ -290,9 +306,11 @@ public final class Mp4Extractor implements Extractor, SeekMap {
     long secondTimeUs = C.TIME_UNSET;
     long secondOffset = C.POSITION_UNSET;
 
+    // Note that the id matches the index in tracks.
+    int mainTrackIndex = trackId != C.INDEX_UNSET ? trackId : firstVideoTrackIndex;
     // If we have a video track, use it to establish one or two seek points.
-    if (firstVideoTrackIndex != C.INDEX_UNSET) {
-      TrackSampleTable sampleTable = tracks[firstVideoTrackIndex].sampleTable;
+    if (mainTrackIndex != C.INDEX_UNSET) {
+      TrackSampleTable sampleTable = tracks[mainTrackIndex].sampleTable;
       int sampleIndex = getSynchronizationSampleIndex(sampleTable, timeUs);
       if (sampleIndex == C.INDEX_UNSET) {
         return new SeekPoints(SeekPoint.START);
@@ -312,13 +330,15 @@ public final class Mp4Extractor implements Extractor, SeekMap {
       firstOffset = Long.MAX_VALUE;
     }
 
-    // Take into account other tracks.
-    for (int i = 0; i < tracks.length; i++) {
-      if (i != firstVideoTrackIndex) {
-        TrackSampleTable sampleTable = tracks[i].sampleTable;
-        firstOffset = maybeAdjustSeekOffset(sampleTable, firstTimeUs, firstOffset);
-        if (secondTimeUs != C.TIME_UNSET) {
-          secondOffset = maybeAdjustSeekOffset(sampleTable, secondTimeUs, secondOffset);
+    if (trackId == C.INDEX_UNSET) {
+      // Take into account other tracks, but only if the caller has not specified a trackId.
+      for (int i = 0; i < tracks.length; i++) {
+        if (i != firstVideoTrackIndex) {
+          TrackSampleTable sampleTable = tracks[i].sampleTable;
+          firstOffset = maybeAdjustSeekOffset(sampleTable, firstTimeUs, firstOffset);
+          if (secondTimeUs != C.TIME_UNSET) {
+            secondOffset = maybeAdjustSeekOffset(sampleTable, secondTimeUs, secondOffset);
+          }
         }
       }
     }