Replace trackIndex with TrackToken in Transformer

This is to eventually replace Transformer/Muxer.java with Muxer/Muxer.java

PiperOrigin-RevId: 627027254

libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerPauseResumeTest.java

@@ -23,6 +23,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import android.content.Context;
+import androidx.annotation.Nullable;
 import androidx.media3.common.C;
 import androidx.media3.common.Effect;
 import androidx.media3.common.Format;
@@ -33,6 +34,7 @@ import androidx.media3.common.audio.AudioProcessor;
 import androidx.media3.common.audio.SonicAudioProcessor;
 import androidx.media3.common.util.Util;
 import androidx.media3.effect.RgbFilter;
+import androidx.media3.muxer.Muxer.TrackToken;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
 import com.google.common.base.Ascii;
@@ -430,28 +432,27 @@ public class TransformerPauseResumeTest {
     private final FrameBlockingMuxer.Listener listener;
 
     private boolean notifiedListener;
-    private int videoTrackIndex;
+    @Nullable private TrackToken videoTrackToken;
 
     private FrameBlockingMuxer(Muxer wrappedMuxer, FrameBlockingMuxer.Listener listener) {
       this.wrappedMuxer = wrappedMuxer;
       this.listener = listener;
-      videoTrackIndex = C.INDEX_UNSET;
     }
 
     @Override
-    public int addTrack(Format format) throws MuxerException {
+    public TrackToken addTrack(Format format) throws MuxerException {
-      int trackIndex = wrappedMuxer.addTrack(format);
+      TrackToken trackToken = wrappedMuxer.addTrack(format);
       if (MimeTypes.isVideo(format.sampleMimeType)) {
-        videoTrackIndex = trackIndex;
+        videoTrackToken = trackToken;
       }
-      return trackIndex;
+      return trackToken;
     }
 
     @Override
     public void writeSampleData(
-        int trackIndex, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
+        TrackToken trackToken, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
         throws MuxerException {
-      if (trackIndex == videoTrackIndex
+      if (trackToken == videoTrackToken
           && presentationTimeUs >= DEFAULT_PRESENTATION_TIME_US_TO_BLOCK_FRAME) {
         if (!notifiedListener) {
           listener.onFrameBlocked();
@@ -459,7 +460,7 @@ public class TransformerPauseResumeTest {
         }
         return;
       }
-      wrappedMuxer.writeSampleData(trackIndex, data, presentationTimeUs, flags);
+      wrappedMuxer.writeSampleData(trackToken, data, presentationTimeUs, flags);
     }
 
     @Override

libraries/transformer/src/main/java/androidx/media3/transformer/DefaultMuxer.java

@@ -19,6 +19,7 @@ import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.Metadata;
 import androidx.media3.common.util.UnstableApi;
+import androidx.media3.muxer.Muxer.TrackToken;
 import com.google.common.collect.ImmutableList;
 import java.nio.ByteBuffer;
 
@@ -64,15 +65,15 @@ public final class DefaultMuxer implements Muxer {
   }
 
   @Override
-  public int addTrack(Format format) throws MuxerException {
+  public TrackToken addTrack(Format format) throws MuxerException {
     return muxer.addTrack(format);
   }
 
   @Override
   public void writeSampleData(
-      int trackIndex, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
+      TrackToken trackToken, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
       throws MuxerException {
-    muxer.writeSampleData(trackIndex, data, presentationTimeUs, flags);
+    muxer.writeSampleData(trackToken, data, presentationTimeUs, flags);
   }
 
   @Override

libraries/transformer/src/main/java/androidx/media3/transformer/FrameworkMuxer.java

@@ -24,7 +24,7 @@ import android.annotation.SuppressLint;
 import android.media.MediaCodec;
 import android.media.MediaFormat;
 import android.media.MediaMuxer;
-import android.util.SparseLongArray;
+import androidx.annotation.Nullable;
 import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.Metadata;
@@ -32,10 +32,13 @@ import androidx.media3.common.MimeTypes;
 import androidx.media3.common.util.MediaFormatUtil;
 import androidx.media3.common.util.Util;
 import androidx.media3.container.Mp4LocationData;
+import androidx.media3.muxer.Muxer.TrackToken;
 import com.google.common.collect.ImmutableList;
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
 
 /** {@link Muxer} implementation that uses a {@link MediaMuxer}. */
 /* package */ final class FrameworkMuxer implements Muxer {
@@ -78,10 +81,10 @@ import java.nio.ByteBuffer;
   private final MediaMuxer mediaMuxer;
   private final long videoDurationUs;
   private final MediaCodec.BufferInfo bufferInfo;
-  private final SparseLongArray trackIndexToLastPresentationTimeUs;
+  private final Map<TrackToken, Long> trackTokenToLastPresentationTimeUs;
-  private final SparseLongArray trackIndexToPresentationTimeOffsetUs;
+  private final Map<TrackToken, Long> trackTokenToPresentationTimeOffsetUs;
 
-  private int videoTrackIndex;
+  @Nullable private TrackToken videoTrackToken;
 
   private boolean isStarted;
   private boolean isReleased;
@@ -90,13 +93,12 @@ import java.nio.ByteBuffer;
     this.mediaMuxer = mediaMuxer;
     this.videoDurationUs = Util.msToUs(videoDurationMs);
     bufferInfo = new MediaCodec.BufferInfo();
-    trackIndexToLastPresentationTimeUs = new SparseLongArray();
+    trackTokenToLastPresentationTimeUs = new HashMap<>();
-    trackIndexToPresentationTimeOffsetUs = new SparseLongArray();
+    trackTokenToPresentationTimeOffsetUs = new HashMap<>();
-    videoTrackIndex = C.INDEX_UNSET;
   }
 
   @Override
-  public int addTrack(Format format) throws MuxerException {
+  public TrackToken addTrack(Format format) throws MuxerException {
     String sampleMimeType = checkNotNull(format.sampleMimeType);
     MediaFormat mediaFormat;
     boolean isVideo = MimeTypes.isVideo(sampleMimeType);
@@ -122,27 +124,27 @@ import java.nio.ByteBuffer;
       throw new MuxerException("Failed to add track with format=" + format, e);
     }
 
+    TrackToken trackToken = new TrackTokenImpl(trackIndex);
     if (isVideo) {
-      videoTrackIndex = trackIndex;
+      videoTrackToken = trackToken;
     }
 
-    return trackIndex;
+    return trackToken;
   }
 
   @Override
   public void writeSampleData(
-      int trackIndex, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
+      TrackToken trackToken, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
       throws MuxerException {
 
     if (videoDurationUs != C.TIME_UNSET
-        && trackIndex == videoTrackIndex
+        && trackToken == videoTrackToken
         && presentationTimeUs > videoDurationUs) {
       return;
     }
-
     if (!isStarted) {
       if (Util.SDK_INT < 30 && presentationTimeUs < 0) {
-        trackIndexToPresentationTimeOffsetUs.put(trackIndex, -presentationTimeUs);
+        trackTokenToPresentationTimeOffsetUs.put(trackToken, -presentationTimeUs);
       }
       startMuxer();
     }
@@ -150,11 +152,13 @@ import java.nio.ByteBuffer;
     int offset = data.position();
     int size = data.limit() - offset;
 
-    long presentationTimeOffsetUs = trackIndexToPresentationTimeOffsetUs.get(trackIndex);
+    long presentationTimeOffsetUs =
+        trackTokenToPresentationTimeOffsetUs.getOrDefault(trackToken, 0L);
     presentationTimeUs += presentationTimeOffsetUs;
 
     bufferInfo.set(offset, size, presentationTimeUs, TransformerUtil.getMediaCodecFlags(flags));
-    long lastSamplePresentationTimeUs = trackIndexToLastPresentationTimeUs.get(trackIndex);
+    long lastSamplePresentationTimeUs =
+        trackTokenToLastPresentationTimeUs.getOrDefault(trackToken, 0L);
     // writeSampleData blocks on old API versions, so check here to avoid calling the method.
     checkState(
         Util.SDK_INT > 24 || presentationTimeUs >= lastSamplePresentationTimeUs,
@@ -163,7 +167,7 @@ import java.nio.ByteBuffer;
             + " < "
             + lastSamplePresentationTimeUs
             + ") unsupported on this API version");
-    trackIndexToLastPresentationTimeUs.put(trackIndex, presentationTimeUs);
+    trackTokenToLastPresentationTimeUs.put(trackToken, presentationTimeUs);
 
     checkState(
         presentationTimeOffsetUs == 0 || presentationTimeUs >= lastSamplePresentationTimeUs,
@@ -174,15 +178,11 @@ import java.nio.ByteBuffer;
             + ") unsupported when using negative PTS workaround");
 
     try {
-      mediaMuxer.writeSampleData(trackIndex, data, bufferInfo);
+      checkState(trackToken instanceof TrackTokenImpl);
+      mediaMuxer.writeSampleData(((TrackTokenImpl) trackToken).trackIndex, data, bufferInfo);
     } catch (RuntimeException e) {
       throw new MuxerException(
-          "Failed to write sample for trackIndex="
+          "Failed to write sample for presentationTimeUs=" + presentationTimeUs + ", size=" + size,
-              + trackIndex
-              + ", presentationTimeUs="
-              + presentationTimeUs
-              + ", size="
-              + size,
           e);
     }
   }
@@ -207,9 +207,9 @@ import java.nio.ByteBuffer;
       startMuxer();
     }
 
-    if (videoDurationUs != C.TIME_UNSET && videoTrackIndex != C.INDEX_UNSET) {
+    if (videoDurationUs != C.TIME_UNSET && videoTrackToken != null) {
       writeSampleData(
-          videoTrackIndex,
+          videoTrackToken,
           ByteBuffer.allocateDirect(0),
           videoDurationUs,
           C.BUFFER_FLAG_END_OF_STREAM);
@@ -278,4 +278,12 @@ import java.nio.ByteBuffer;
     }
     return supportedMimeTypes.build();
   }
+
+  private static class TrackTokenImpl implements TrackToken {
+    public final int trackIndex;
+
+    public TrackTokenImpl(int trackIndex) {
+      this.trackIndex = trackIndex;
+    }
+  }
 }

libraries/transformer/src/main/java/androidx/media3/transformer/InAppMuxer.java

@@ -32,9 +32,7 @@ import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.LinkedHashSet;
-import java.util.List;
 import java.util.Set;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -161,7 +159,6 @@ public final class InAppMuxer implements Muxer {
 
   private final androidx.media3.muxer.Muxer muxer;
   private final @Nullable MetadataProvider metadataProvider;
-  private final List<TrackToken> trackTokenList;
   private final BufferInfo bufferInfo;
   private final Set<Metadata.Entry> metadataEntries;
 
@@ -169,26 +166,22 @@ public final class InAppMuxer implements Muxer {
       androidx.media3.muxer.Muxer muxer, @Nullable MetadataProvider metadataProvider) {
     this.muxer = muxer;
     this.metadataProvider = metadataProvider;
-    trackTokenList = new ArrayList<>();
     bufferInfo = new BufferInfo();
     metadataEntries = new LinkedHashSet<>();
   }
 
   @Override
-  public int addTrack(Format format) {
+  public TrackToken addTrack(Format format) {
     TrackToken trackToken = muxer.addTrack(format);
-    trackTokenList.add(trackToken);
-
     if (MimeTypes.isVideo(format.sampleMimeType)) {
       muxer.addMetadata(new Mp4OrientationData(format.rotationDegrees));
     }
-
+    return trackToken;
-    return trackTokenList.size() - 1;
   }
 
   @Override
   public void writeSampleData(
-      int trackIndex, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
+      TrackToken trackToken, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
       throws MuxerException {
 
     int size = data.remaining();
@@ -196,15 +189,10 @@ public final class InAppMuxer implements Muxer {
         data.position(), size, presentationTimeUs, TransformerUtil.getMediaCodecFlags(flags));
 
     try {
-      muxer.writeSampleData(trackTokenList.get(trackIndex), data, bufferInfo);
+      muxer.writeSampleData(trackToken, data, bufferInfo);
     } catch (IOException e) {
       throw new MuxerException(
-          "Failed to write sample for trackIndex="
+          "Failed to write sample for presentationTimeUs=" + presentationTimeUs + ", size=" + size,
-              + trackIndex
-              + ", presentationTimeUs="
-              + presentationTimeUs
-              + ", size="
-              + size,
           e);
     }
   }

libraries/transformer/src/main/java/androidx/media3/transformer/Muxer.java

@@ -20,6 +20,7 @@ import androidx.media3.common.Format;
 import androidx.media3.common.Metadata;
 import androidx.media3.common.MimeTypes;
 import androidx.media3.common.util.UnstableApi;
+import androidx.media3.muxer.Muxer.TrackToken;
 import com.google.common.collect.ImmutableList;
 import java.nio.ByteBuffer;
 
@@ -71,15 +72,17 @@ public interface Muxer {
    * Adds a track with the specified format.
    *
    * @param format The {@link Format} of the track.
-   * @return The index for this track, which should be passed to {@link #writeSampleData}.
+   * @return The {@link TrackToken} for this track, which should be passed to {@link
+   *     #writeSampleData}.
    * @throws MuxerException If the muxer encounters a problem while adding the track.
    */
-  int addTrack(Format format) throws MuxerException;
+  TrackToken addTrack(Format format) throws MuxerException;
 
   /**
    * Writes the specified sample.
    *
-   * @param trackIndex The index of the track, previously returned by {@link #addTrack(Format)}.
+   * @param trackToken The {@link TrackToken} of the track, previously returned by {@link
+   *     #addTrack(Format)}.
    * @param data A buffer containing the sample data to write to the container.
    * @param presentationTimeUs The presentation time of the sample in microseconds.
    * @param flags The {@link C.BufferFlags} associated with the data. Only {@link
@@ -87,7 +90,7 @@ public interface Muxer {
    * @throws MuxerException If the muxer fails to write the sample.
    */
   void writeSampleData(
-      int trackIndex, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
+      TrackToken trackToken, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
       throws MuxerException;
 
   /** Adds {@linkplain Metadata.Entry metadata} about the output file. */

libraries/transformer/src/main/java/androidx/media3/transformer/MuxerWrapper.java

@@ -40,6 +40,7 @@ import androidx.media3.common.MimeTypes;
 import androidx.media3.common.util.Util;
 import androidx.media3.container.NalUnitUtil;
 import androidx.media3.effect.DebugTraceUtil;
+import androidx.media3.muxer.Muxer.TrackToken;
 import com.google.common.collect.ImmutableList;
 import java.io.File;
 import java.lang.annotation.Documented;
@@ -552,7 +553,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     resetAbortTimer();
     checkStateNotNull(muxer);
     muxer.writeSampleData(
-        trackInfo.index, data, presentationTimeUs, isKeyFrame ? C.BUFFER_FLAG_KEY_FRAME : 0);
+        trackInfo.trackToken, data, presentationTimeUs, isKeyFrame ? C.BUFFER_FLAG_KEY_FRAME : 0);
     if (trackType == C.TRACK_TYPE_VIDEO) {
       DebugTraceUtil.logEvent(DebugTraceUtil.EVENT_MUXER_WRITE_SAMPLE_VIDEO, presentationTimeUs);
     } else if (trackType == C.TRACK_TYPE_AUDIO) {
@@ -737,15 +738,15 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
   private static final class TrackInfo {
     public final Format format;
-    public final int index;
+    public final TrackToken trackToken;
 
     public long bytesWritten;
     public int sampleCount;
     public long timeUs;
 
-    public TrackInfo(Format format, int index) {
+    public TrackInfo(Format format, TrackToken trackToken) {
       this.format = format;
-      this.index = index;
+      this.trackToken = trackToken;
     }
 
     /**

libraries/transformer/src/test/java/androidx/media3/transformer/CapturingMuxer.java

@@ -25,6 +25,7 @@ import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.Metadata;
 import androidx.media3.common.util.Util;
+import androidx.media3.muxer.Muxer.TrackToken;
 import androidx.media3.test.utils.DumpableFormat;
 import androidx.media3.test.utils.Dumper;
 import androidx.media3.test.utils.Dumper.Dumpable;
@@ -85,7 +86,7 @@ public final class CapturingMuxer implements Muxer, Dumpable {
   private final boolean handleAudioAsPcm;
   private final SparseArray<DumpableFormat> dumpableFormatByTrackType;
   private final SparseArray<DumpableStream> dumpableStreamByTrackType;
-  private final Map<Integer, Integer> trackIndexToType;
+  private final Map<TrackToken, Integer> trackTokenToType;
   private final ArrayList<Metadata.Entry> metadataList;
   private boolean released;
 
@@ -95,18 +96,18 @@ public final class CapturingMuxer implements Muxer, Dumpable {
     this.handleAudioAsPcm = handleAudioAsPcm;
     dumpableFormatByTrackType = new SparseArray<>();
     dumpableStreamByTrackType = new SparseArray<>();
-    trackIndexToType = new HashMap<>();
+    trackTokenToType = new HashMap<>();
     metadataList = new ArrayList<>();
   }
 
   // Muxer implementation.
 
   @Override
-  public int addTrack(Format format) throws MuxerException {
+  public TrackToken addTrack(Format format) throws MuxerException {
-    int trackIndex = wrappedMuxer.addTrack(format);
+    TrackToken trackToken = wrappedMuxer.addTrack(format);
     @C.TrackType int trackType = getProcessedTrackType(format.sampleMimeType);
 
-    trackIndexToType.put(trackIndex, trackType);
+    trackTokenToType.put(trackToken, trackType);
 
     dumpableFormatByTrackType.append(
         trackType, new DumpableFormat(format, /* tag= */ Util.getTrackTypeString(trackType)));
@@ -117,19 +118,19 @@ public final class CapturingMuxer implements Muxer, Dumpable {
             ? new DumpablePcmAudioStream(trackType)
             : new DumpableStream(trackType));
 
-    return trackIndex;
+    return trackToken;
   }
 
   @Override
   public void writeSampleData(
-      int trackIndex, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
+      TrackToken trackToken, ByteBuffer data, long presentationTimeUs, @C.BufferFlags int flags)
       throws MuxerException {
-    @C.TrackType int trackType = checkNotNull(trackIndexToType.get(trackIndex));
+    @C.TrackType int trackType = checkNotNull(trackTokenToType.get(trackToken));
     dumpableStreamByTrackType
         .get(trackType)
         .addSample(
             data, (flags & C.BUFFER_FLAG_KEY_FRAME) == C.BUFFER_FLAG_KEY_FRAME, presentationTimeUs);
-    wrappedMuxer.writeSampleData(trackIndex, data, presentationTimeUs, flags);
+    wrappedMuxer.writeSampleData(trackToken, data, presentationTimeUs, flags);
   }
 
   @Override