Support multi-track in MKV/WebM extractor.

Issue: #514

library/src/androidTest/java/com/google/android/exoplayer/extractor/webm/StreamBuilder.java

@@ -91,41 +91,44 @@ import java.util.List;
     return this;
   }
 
-  public StreamBuilder addVp9Track(int width, int height,
+  public StreamBuilder addVp9Track(byte trackNumber, int width, int height,
       ContentEncodingSettings contentEncodingSettings) {
-    trackEntries.add(createVideoTrackEntry("V_VP9", width, height, contentEncodingSettings, null));
+    trackEntries.add(createVideoTrackEntry(trackNumber, "V_VP9", width, height,
+        contentEncodingSettings, null));
     return this;
   }
 
-  public StreamBuilder addH264Track(int width, int height, byte[] codecPrivate) {
-    trackEntries.add(createVideoTrackEntry("V_MPEG4/ISO/AVC", width, height, null, codecPrivate));
+  public StreamBuilder addH264Track(byte trackNumber, int width, int height, byte[] codecPrivate) {
+    trackEntries.add(createVideoTrackEntry(trackNumber, "V_MPEG4/ISO/AVC", width, height, null,
+        codecPrivate));
     return this;
   }
 
-  public StreamBuilder addOpusTrack(int channelCount, int sampleRate, int codecDelay,
-      int seekPreRoll, byte[] codecPrivate) {
-    trackEntries.add(createAudioTrackEntry("A_OPUS", channelCount, sampleRate, codecPrivate,
-        codecDelay, seekPreRoll, NO_VALUE));
+  public StreamBuilder addOpusTrack(byte trackNumber, int channelCount, int sampleRate,
+      int codecDelay, int seekPreRoll, byte[] codecPrivate) {
+    trackEntries.add(createAudioTrackEntry(trackNumber, "A_OPUS", channelCount, sampleRate,
+        codecPrivate, codecDelay, seekPreRoll, NO_VALUE));
     return this;
   }
 
-  public StreamBuilder addOpusTrack(int channelCount, int sampleRate, int codecDelay,
-      int seekPreRoll, byte[] codecPrivate, int defaultDurationNs) {
-    trackEntries.add(createAudioTrackEntry("A_OPUS", channelCount, sampleRate, codecPrivate,
-        codecDelay, seekPreRoll, defaultDurationNs));
+  public StreamBuilder addOpusTrack(byte trackNumber, int channelCount, int sampleRate,
+      int codecDelay, int seekPreRoll, byte[] codecPrivate, int defaultDurationNs) {
+    trackEntries.add(createAudioTrackEntry(trackNumber, "A_OPUS", channelCount, sampleRate,
+        codecPrivate, codecDelay, seekPreRoll, defaultDurationNs));
     return this;
   }
 
-  public StreamBuilder addVorbisTrack(int channelCount, int sampleRate, byte[] codecPrivate) {
-    trackEntries.add(createAudioTrackEntry("A_VORBIS", channelCount, sampleRate, codecPrivate,
-        NO_VALUE, NO_VALUE, NO_VALUE));
+  public StreamBuilder addVorbisTrack(byte trackNumber, int channelCount, int sampleRate,
+      byte[] codecPrivate) {
+    trackEntries.add(createAudioTrackEntry(trackNumber, "A_VORBIS", channelCount, sampleRate,
+        codecPrivate, NO_VALUE, NO_VALUE, NO_VALUE));
     return this;
   }
 
-  public StreamBuilder addUnsupportedTrack() {
+  public StreamBuilder addUnsupportedTrack(byte trackNumber) {
     trackEntries.add(element(0xAE, // TrackEntry
         element(0x86, "D_WEBVTT/metadata".getBytes()), // CodecID
-        element(0xD7, (byte) 0x03), // TrackNumber
+        element(0xD7, trackNumber), // TrackNumber
         element(0x83, (byte) 0x11))); // TrackType
     return this;
   }
@@ -252,8 +255,8 @@ import java.util.List;
         durationFirst ? timescaleElement : durationElement);
   }
 
-  private static EbmlElement createVideoTrackEntry(String codecId, int pixelWidth, int pixelHeight,
-      ContentEncodingSettings contentEncodingSettings, byte[] codecPrivate) {
+  private static EbmlElement createVideoTrackEntry(byte trackNumber, String codecId, int pixelWidth,
+      int pixelHeight, ContentEncodingSettings contentEncodingSettings, byte[] codecPrivate) {
     byte[] widthBytes = getIntegerBytes(pixelWidth);
     byte[] heightBytes = getIntegerBytes(pixelHeight);
     EbmlElement contentEncodingSettingsElement;
@@ -297,7 +300,7 @@ import java.util.List;
 
     return element(0xAE, // TrackEntry
         element(0x86, codecId.getBytes()), // CodecID
-        element(0xD7, (byte) 0x01), // TrackNumber
+        element(0xD7, trackNumber), // TrackNumber
         element(0x83, (byte) 0x01), // TrackType
         contentEncodingSettingsElement,
         element(0xE0, // Video
@@ -306,13 +309,14 @@ import java.util.List;
         codecPrivateElement);
   }
 
-  private static EbmlElement createAudioTrackEntry(String codecId, int channelCount, int sampleRate,
-      byte[] codecPrivate, int codecDelay, int seekPreRoll, int defaultDurationNs) {
+  private static EbmlElement createAudioTrackEntry(byte trackNumber, String codecId,
+      int channelCount, int sampleRate, byte[] codecPrivate, int codecDelay, int seekPreRoll,
+      int defaultDurationNs) {
     byte channelCountByte = (byte) (channelCount & 0xFF);
     byte[] sampleRateDoubleBytes = getLongBytes(Double.doubleToLongBits(sampleRate));
     return element(0xAE, // TrackEntry
         element(0x86, codecId.getBytes()), // CodecID
-        element(0xD7, (byte) 0x02), // TrackNumber
+        element(0xD7, trackNumber), // TrackNumber
         element(0x83, (byte) 0x02), // TrackType
         // CodecDelay
         codecDelay != NO_VALUE ? element(0x56AA, getIntegerBytes(codecDelay)) : empty(),

library/src/main/java/com/google/android/exoplayer/extractor/webm/WebmExtractor.java

@@ -34,6 +34,7 @@ import com.google.android.exoplayer.util.ParsableByteArray;
 import com.google.android.exoplayer.util.Util;
 
 import android.util.Pair;
+import android.util.SparseArray;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -81,7 +82,6 @@ public final class WebmExtractor implements Extractor {
   private static final int MP3_MAX_INPUT_SIZE = 4096;
   private static final int ENCRYPTION_IV_SIZE = 8;
   private static final int TRACK_TYPE_AUDIO = 2;
-  private static final int TRACK_TYPE_VIDEO = 1;
   private static final int UNKNOWN = -1;
 
   private static final int ID_EBML = 0x1A45DFA3;
@@ -143,6 +143,7 @@ public final class WebmExtractor implements Extractor {
 
   private final EbmlReader reader;
   private final VarintReader varintReader;
+  private final SparseArray<Track> tracks;
 
   // Temporary arrays.
   private final ParsableByteArray nalStartCode;
@@ -158,10 +159,10 @@ public final class WebmExtractor implements Extractor {
   private long durationTimecode = C.UNKNOWN_TIME_US;
   private long durationUs = C.UNKNOWN_TIME_US;
 
+  // The track corresponding to the current TrackEntry element, or null.
   private Track currentTrack;
-  private Track audioTrack;
-  private Track videoTrack;
 
+  // Whether drm init data has been sent to the output.
   private boolean sentDrmInitData;
 
   // Master seek entry related elements.
@@ -208,6 +209,7 @@ public final class WebmExtractor implements Extractor {
     this.reader = reader;
     this.reader.init(new InnerEbmlReaderOutput());
     varintReader = new VarintReader();
+    tracks = new SparseArray<>();
     scratch = new ParsableByteArray(4);
     vorbisNumPageSamples = new ParsableByteArray(ByteBuffer.allocate(4).putInt(-1).array());
     seekEntryIdBytes = new ParsableByteArray(4);
@@ -399,8 +401,7 @@ public final class WebmExtractor implements Extractor {
         if (!sampleSeenReferenceBlock) {
           blockFlags |= C.SAMPLE_FLAG_SYNC;
         }
-        outputSampleMetadata((audioTrack != null && blockTrackNumber == audioTrack.number)
-            ? audioTrack.output : videoTrack.output, blockTimeUs);
+        outputSampleMetadata(tracks.get(blockTrackNumber), blockTimeUs);
         blockState = BLOCK_STATE_START;
         return;
       case ID_CONTENT_ENCODING:
@@ -421,26 +422,16 @@ public final class WebmExtractor implements Extractor {
         }
         return;
       case ID_TRACK_ENTRY:
-        if ((currentTrack.type == TRACK_TYPE_AUDIO && audioTrack != null)
-            || (currentTrack.type == TRACK_TYPE_VIDEO && videoTrack != null)) {
-          // There is more than 1 audio/video track. Ignore everything but the first.
-          currentTrack = null;
-          return;
-        }
-        if (currentTrack.type == TRACK_TYPE_AUDIO && isCodecSupported(currentTrack.codecId)) {
-          audioTrack = currentTrack;
-          audioTrack.initializeOutput(extractorOutput, durationUs);
-        } else if (currentTrack.type == TRACK_TYPE_VIDEO
-            && isCodecSupported(currentTrack.codecId)) {
-          videoTrack = currentTrack;
-          videoTrack.initializeOutput(extractorOutput, durationUs);
+        if (tracks.get(currentTrack.number) == null && isCodecSupported(currentTrack.codecId)) {
+          currentTrack.initializeOutput(extractorOutput, durationUs);
+          tracks.put(currentTrack.number, currentTrack);
         } else {
-          // Unsupported track type. Do nothing.
+          // We've seen this track entry before, or the codec is unsupported. Do nothing.
         }
         currentTrack = null;
         return;
       case ID_TRACKS:
-        if (videoTrack == null && audioTrack == null) {
+        if (tracks.size() == 0) {
           throw new ParserException("No valid tracks were found");
         }
         extractorOutput.endTracks();
@@ -614,16 +605,15 @@ public final class WebmExtractor implements Extractor {
           scratch.reset();
         }
 
-        // Ignore the block if the track number equals neither the audio track nor the video track.
-        if ((audioTrack == null || audioTrack.number != blockTrackNumber)
-            && (videoTrack == null || videoTrack.number != blockTrackNumber)) {
+        Track track = tracks.get(blockTrackNumber);
+
+        // Ignore the block if we don't know about the track to which it belongs.
+        if (track == null) {
           input.skipFully(contentSize - blockTrackNumberLength);
           blockState = BLOCK_STATE_START;
           return;
         }
 
-        Track track = (audioTrack != null && blockTrackNumber == audioTrack.number)
-            ? audioTrack : videoTrack;
         if (blockState == BLOCK_STATE_HEADER) {
           // Read the relative timecode (2 bytes) and flags (1 byte).
           readScratch(input, 3);
@@ -723,7 +713,7 @@ public final class WebmExtractor implements Extractor {
             writeSampleData(input, track, blockLacingSampleSizes[blockLacingSampleIndex]);
             long sampleTimeUs = this.blockTimeUs
                 + (blockLacingSampleIndex * track.defaultSampleDurationNs) / 1000;
-            outputSampleMetadata(track.output, sampleTimeUs);
+            outputSampleMetadata(track, sampleTimeUs);
             blockLacingSampleIndex++;
           }
           blockState = BLOCK_STATE_START;
@@ -739,8 +729,8 @@ public final class WebmExtractor implements Extractor {
     }
   }
 
-  private void outputSampleMetadata(TrackOutput trackOutput, long timeUs) {
-    trackOutput.sampleMetadata(timeUs, blockFlags, sampleBytesWritten, 0, blockEncryptionKeyId);
+  private void outputSampleMetadata(Track track, long timeUs) {
+    track.output.sampleMetadata(timeUs, blockFlags, sampleBytesWritten, 0, blockEncryptionKeyId);
     sampleRead = true;
     resetSample();
   }