Refactor logic and add tests

Author: tonihei

RELEASENOTES.md

@@ -49,6 +49,8 @@
     *   MP3: Allow gaps between (and before) ID3 tags at the beginning of MP3
         files ([#811](https://github.com/androidx/media/issues/811),
         [#5718](https://github.com/google/ExoPlayer/issues/5718)).
+    *   Matroska: Add support for DTS-HD detection
+        ([#6225](https://github.com/google/ExoPlayer/issues/6225)).
 *   DataSource:
 *   Audio:
     *   Make `AudioProcessor` instances aware of seeking.

libraries/extractor/src/main/java/androidx/media3/extractor/DtsUtil.java

@@ -27,8 +27,10 @@
 import androidx.media3.common.MimeTypes;
 import androidx.media3.common.ParserException;
 import androidx.media3.common.util.ParsableBitArray;
+import androidx.media3.common.util.ParsableByteArray;
 import androidx.media3.common.util.UnstableApi;
 import androidx.media3.common.util.Util;
+import java.io.IOException;
 import java.lang.annotation.Documented;
 import java.lang.annotation.Retention;
 import java.lang.annotation.Target;
@@ -681,6 +683,39 @@ public static int parseDtsUhdHeaderSize(byte[] headerPrefix) {
         + 1;
   }
 
+  /** Returns whether the sample data at the current {@link ExtractorInput} is a DTS-HD sample. */
+  public static boolean isSampleDtsHd(ExtractorInput input, int sampleSize) throws IOException {
+    // Limit the peek ahead to be up to the sample size plus the sync word of the second frame.
+    int scanLength = sampleSize + 95;
+    ParsableByteArray sampleData = new ParsableByteArray(scanLength);
+    if (!input.peekFully(
+        sampleData.getData(), /* offset= */ 0, scanLength, /* allowEndOfInput= */ true)) {
+      return false;
+    }
+    input.resetPeekPosition();
+
+    while (sampleData.bytesLeft() >= 4) {
+      int word = sampleData.peekInt();
+      if (DtsUtil.getFrameType(word) == DtsUtil.FRAME_TYPE_CORE) {
+        if (sampleData.bytesLeft() < 10) {
+          return false;
+        }
+        byte[] header = new byte[10];
+        sampleData.readBytes(header, /* offset= */ 0, /* length= */ 10);
+        sampleData.setPosition(0);
+        int frameSize = DtsUtil.getDtsFrameSize(header);
+        if (frameSize <= 0 || sampleData.bytesLeft() < frameSize + 4) {
+          return false;
+        }
+        sampleData.skipBytes(frameSize);
+        word = sampleData.readInt();
+        return DtsUtil.getFrameType(word) == DtsUtil.FRAME_TYPE_EXTENSION_SUBSTREAM;
+      }
+      sampleData.skipBytes(4);
+    }
+    return false;
+  }
+
   /**
    * Check if calculated and extracted CRC-16 words match. See ETSI TS 103 491 V1.2.1, Table 6-8.
    */

libraries/extractor/src/main/java/androidx/media3/extractor/mkv/MatroskaExtractor.java

@@ -445,7 +445,7 @@ public static ExtractorsFactory newFactory(SubtitleParser.Factory subtitleParser
   private long durationTimecode = C.TIME_UNSET;
   private long durationUs = C.TIME_UNSET;
   private boolean isWebm;
-  private boolean pendingEndTracks = true;
+  private boolean pendingEndTracks;
 
   // The track corresponding to the current TrackEntry element, or null.
   @Nullable private Track currentTrack;
@@ -557,6 +557,7 @@ public MatroskaExtractor(SubtitleParser.Factory subtitleParserFactory, @Flags in
     encryptionSubsampleData = new ParsableByteArray();
     supplementalData = new ParsableByteArray();
     blockSampleSizes = new int[1];
+    pendingEndTracks = true;
   }
 
   @Override
@@ -797,23 +798,6 @@ protected void startMasterElement(int id, long contentPosition, long contentSize
     }
   }
 
-  /**
-   * Ensures `extractorOutput.endTracks()` gets called only once, and only
-   * if we don't have any pending audio analysis.
-   */
-  private void maybeEndTracks() {
-    if (!pendingEndTracks) {
-      return;
-    }
-    for (int i = 0; i < tracks.size(); i++) {
-      if (tracks.valueAt(i).waitingForDtsAnalysis) {
-        return;
-      }
-    }
-    extractorOutput.endTracks();
-    pendingEndTracks = false;
-  }
-
   /**
    * Called when the end of a master element is encountered.
    *
@@ -1584,51 +1568,14 @@ private int writeSampleData(ExtractorInput input, Track track, int size, boolean
     }
 
     if (track.waitingForDtsAnalysis) {
-      // The format for this DTS track as not been determined yet
-      long remaining = input.getLength() - input.getPosition();
-      // Limit the peek ahead to be up to the max frame size (16383) plus the
-      // sync word of the second frame
-      int scanLength = (int)Math.min(16383 + 95, remaining);
-      byte[] buf = new byte[scanLength];
-
-      input.advancePeekPosition(0);
-      input.peekFully(buf, 0, buf.length);
-      input.resetPeekPosition();
-
-      final ByteBuffer bb = ByteBuffer.wrap(buf);
-      for (int idx = 0; idx + 4 <= buf.length; idx += 4) {
-        int word = bb.getInt(idx);
-
-        if (DtsUtil.getFrameType(word) == DtsUtil.FRAME_TYPE_CORE) {
-          if (idx + 10 > buf.length) {
-            break;
-          }
-
-          bb.mark();
-          bb.position(idx);
-          byte[] header = new byte[10];
-          bb.get(header);
-          bb.reset();
-          int fsize = DtsUtil.getDtsFrameSize(header);
-          if (fsize <= 0 || idx + fsize + 4 > buf.length) {
-            break;
-          }
-
-          word = bb.getInt(idx + fsize);
-
-          if (DtsUtil.getFrameType(word) == DtsUtil.FRAME_TYPE_EXTENSION_SUBSTREAM) {
-            track.formatBuilder.setSampleMimeType(MimeTypes.AUDIO_DTS_HD);
-            track.output.format(track.formatBuilder.build());
-          }
-
-          // After finding a valid DTS core frame we can break the loop, there is no
-          // need to evaluate the rest of the buffer.
-          break;
-        }
+      checkNotNull(track.format);
+      if (DtsUtil.isSampleDtsHd(input, size)) {
+        track.format = track.format.buildUpon().setSampleMimeType(MimeTypes.AUDIO_DTS_HD).build();
       }
+      track.output.format(track.format);
       track.waitingForDtsAnalysis = false;
       maybeEndTracks();
-   }
+    }
 
     TrackOutput output = track.output;
     if (!sampleEncodingHandled) {
@@ -2091,6 +2038,19 @@ private void assertInitialized() {
     checkStateNotNull(extractorOutput);
   }
 
+  private void maybeEndTracks() {
+    if (!pendingEndTracks) {
+      return;
+    }
+    for (int i = 0; i < tracks.size(); i++) {
+      if (tracks.valueAt(i).waitingForDtsAnalysis) {
+        return;
+      }
+    }
+    checkNotNull(extractorOutput).endTracks();
+    pendingEndTracks = false;
+  }
+
   /** Passes events through to the outer {@link MatroskaExtractor}. */
   private final class InnerEbmlProcessor implements EbmlProcessor {
 
@@ -2210,7 +2170,7 @@ protected static final class Track {
 
     // Set when the output is initialized. nalUnitLengthFieldLength is only set for H264/H265.
     public @MonotonicNonNull TrackOutput output;
-    public Format.Builder formatBuilder;
+    public @MonotonicNonNull Format format;
     public int nalUnitLengthFieldLength;
 
     /** Initializes the track with an output. */
@@ -2438,7 +2398,7 @@ public void initializeOutput(ExtractorOutput output, int trackId) throws ParserE
       selectionFlags |= flagForced ? C.SELECTION_FLAG_FORCED : 0;
 
       int type;
-      formatBuilder = new Format.Builder();
+      Format.Builder formatBuilder = new Format.Builder();
       // TODO: Consider reading the name elements of the tracks and, if present, incorporating them
       // into the trackId passed when creating the formats.
       if (MimeTypes.isAudio(mimeType)) {
@@ -2514,7 +2474,7 @@ public void initializeOutput(ExtractorOutput output, int trackId) throws ParserE
         formatBuilder.setLabel(name);
       }
 
-      Format format =
+      format =
           formatBuilder
               .setId(trackId)
               .setContainerMimeType(isWebm ? MimeTypes.VIDEO_WEBM : MimeTypes.VIDEO_MATROSKA)
@@ -2528,7 +2488,9 @@ public void initializeOutput(ExtractorOutput output, int trackId) throws ParserE
               .build();
 
       this.output = output.track(number, type);
-      this.output.format(format);
+      if (!waitingForDtsAnalysis) {
+        this.output.format(format);
+      }
     }
 
     /** Forces any pending sample metadata to be flushed to the output. */

libraries/extractor/src/test/java/androidx/media3/extractor/mkv/MatroskaExtractorTest.java

@@ -211,6 +211,30 @@ public void webmSubsampleEncryptionWithAltrefFrames() throws Exception {
         simulationConfig);
   }
 
+  @Test
+  public void mkvSample_withDts() throws Exception {
+    ExtractorAsserts.assertBehavior(
+        getExtractorFactory(subtitlesParsedDuringExtraction),
+        "media/mkv/sample_with_dts.mkv",
+        simulationConfig);
+  }
+
+  @Test
+  public void mkvSample_withDtsHdMa() throws Exception {
+    ExtractorAsserts.assertBehavior(
+        getExtractorFactory(subtitlesParsedDuringExtraction),
+        "media/mkv/sample_with_dts_hd_ma.mkv",
+        simulationConfig);
+  }
+
+  @Test
+  public void mkvSample_withDtsX() throws Exception {
+    ExtractorAsserts.assertBehavior(
+        getExtractorFactory(subtitlesParsedDuringExtraction),
+        "media/mkv/sample_with_dts_x.mkv",
+        simulationConfig);
+  }
+
   private static ExtractorAsserts.ExtractorFactory getExtractorFactory(
       boolean subtitlesParsedDuringExtraction) {
     SubtitleParser.Factory subtitleParserFactory;