Add TrackSelectionUtil.getAverageBitrates method

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=219636548

library/core/src/main/java/com/google/android/exoplayer2/source/chunk/BaseMediaChunkIterator.java

@@ -34,10 +34,11 @@ public abstract class BaseMediaChunkIterator implements MediaChunkIterator {
    * @param fromIndex The first available index.
    * @param toIndex The last available index.
    */
+  @SuppressWarnings("method.invocation.invalid")
   public BaseMediaChunkIterator(long fromIndex, long toIndex) {
     this.fromIndex = fromIndex;
     this.toIndex = toIndex;
-    currentIndex = fromIndex - 1;
+    reset();
   }
 
   @Override
@@ -51,6 +52,11 @@ public abstract class BaseMediaChunkIterator implements MediaChunkIterator {
     return !isEnded();
   }
 
+  @Override
+  public void reset() {
+    currentIndex = fromIndex - 1;
+  }
+
   /**
    * Verifies that the iterator points to a valid element.
    *

library/core/src/main/java/com/google/android/exoplayer2/source/chunk/MediaChunkIterator.java

@@ -55,6 +55,11 @@ public interface MediaChunkIterator {
         public long getChunkEndTimeUs() {
           throw new NoSuchElementException();
         }
+
+        @Override
+        public void reset() {
+          // Do nothing.
+        }
       };
 
   /** Returns whether the iteration has reached the end of the available data. */
@@ -93,4 +98,7 @@ public interface MediaChunkIterator {
    *     {@link #next()} or when {@link #isEnded()} is true.
    */
   long getChunkEndTimeUs();
+
+  /** Resets the iterator to the initial position. */
+  void reset();
 }

library/core/src/main/java/com/google/android/exoplayer2/trackselection/TrackSelectionUtil.java

@@ -16,7 +16,9 @@
 package com.google.android.exoplayer2.trackselection;
 
 import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.source.chunk.MediaChunkIterator;
+import com.google.android.exoplayer2.util.Assertions;
 
 /** Track selection related utility methods. */
 public final class TrackSelectionUtil {
@@ -30,7 +32,7 @@ public final class TrackSelectionUtil {
    * @param iterator Iterator for media chunk sequences.
    * @param maxDurationUs Maximum duration of chunks to be included in average bitrate, in
    *     microseconds.
-   * @return Average bitrate for chunks in bits per second, or {@link C#LENGTH_UNSET} if there are
+   * @return Average bitrate for chunks in bits per second, or {@link Format#NO_VALUE} if there are
    *     no chunks or the first chunk length is unknown.
    */
   public static int getAverageBitrate(MediaChunkIterator iterator, long maxDurationUs) {
@@ -51,7 +53,79 @@ public final class TrackSelectionUtil {
       totalLength += chunkLength;
     }
     return totalDurationUs == 0
-        ? C.LENGTH_UNSET
+        ? Format.NO_VALUE
         : (int) (totalLength * C.BITS_PER_BYTE * C.MICROS_PER_SECOND / totalDurationUs);
   }
+
+  /**
+   * Returns average bitrate values for a set of tracks whose upcoming media chunk iterators and
+   * formats are given. If an average bitrate can't be calculated, an estimation is calculated using
+   * average bitrate of another track and the ratio of the bitrate values defined in the formats of
+   * the two tracks.
+   *
+   * @param iterators An array of {@link MediaChunkIterator}s providing information about the
+   *     sequence of upcoming media chunks for each track.
+   * @param formats The track formats.
+   * @param maxDurationUs Maximum duration of chunks to be included in average bitrate values, in
+   *     microseconds.
+   * @return Average bitrate values for the tracks. If for a track, an average bitrate or an
+   *     estimation can't be calculated, {@link Format#NO_VALUE} is set.
+   * @see #getAverageBitrate(MediaChunkIterator, long)
+   */
+  public static int[] getAverageBitrates(
+      MediaChunkIterator[] iterators, Format[] formats, long maxDurationUs) {
+    int trackCount = iterators.length;
+    Assertions.checkArgument(trackCount == formats.length);
+    if (trackCount == 0) {
+      return new int[0];
+    }
+
+    int[] bitrates = new int[trackCount];
+    int[] formatBitrates = new int[trackCount];
+    float[] bitrateRatios = new float[trackCount];
+    boolean needEstimateBitrate = false;
+    boolean canEstimateBitrate = false;
+    for (int i = 0; i < trackCount; i++) {
+      int bitrate = getAverageBitrate(iterators[i], maxDurationUs);
+      if (bitrate != Format.NO_VALUE) {
+        int formatBitrate = formats[i].bitrate;
+        formatBitrates[i] = formatBitrate;
+        if (formatBitrate != Format.NO_VALUE) {
+          bitrateRatios[i] = ((float) bitrate) / formatBitrate;
+          canEstimateBitrate = true;
+        }
+      } else {
+        needEstimateBitrate = true;
+        formatBitrates[i] = Format.NO_VALUE;
+      }
+      bitrates[i] = bitrate;
+    }
+
+    if (needEstimateBitrate && canEstimateBitrate) {
+      for (int i = 0; i < trackCount; i++) {
+        if (bitrates[i] == Format.NO_VALUE) {
+          int formatBitrate = formats[i].bitrate;
+          if (formatBitrate != Format.NO_VALUE) {
+            int closestFormat = findClosestBitrateFormat(formatBitrate, formatBitrates);
+            bitrates[i] = (int) (bitrateRatios[closestFormat] * formatBitrate);
+          }
+        }
+      }
+    }
+    return bitrates;
+  }
+
+  private static int findClosestBitrateFormat(int formatBitrate, int[] formatBitrates) {
+    int closestDistance = Integer.MAX_VALUE;
+    int closestFormat = C.INDEX_UNSET;
+    for (int j = 0; j < formatBitrates.length; j++) {
+      if (formatBitrates[j] != Format.NO_VALUE) {
+        int distance = Math.abs(formatBitrates[j] - formatBitrate);
+        if (distance < closestDistance) {
+          closestFormat = j;
+        }
+      }
+    }
+    return closestFormat;
+  }
 }

library/core/src/test/java/com/google/android/exoplayer2/trackselection/TrackSelectionUtilTest.java

@@ -18,7 +18,9 @@ package com.google.android.exoplayer2.trackselection;
 import static com.google.common.truth.Truth.assertThat;
 
 import android.net.Uri;
+import android.support.annotation.NonNull;
 import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.source.chunk.BaseMediaChunkIterator;
 import com.google.android.exoplayer2.source.chunk.MediaChunkIterator;
 import com.google.android.exoplayer2.upstream.DataSpec;
@@ -33,62 +35,60 @@ public class TrackSelectionUtilTest {
   public static final long MAX_DURATION_US = 30 * C.MICROS_PER_SECOND;
 
   @Test
-  public void getAverageBitrate_emptyIterator_returnsUnsetLength() {
+  public void getAverageBitrate_emptyIterator_returnsNoValue() {
     assertThat(TrackSelectionUtil.getAverageBitrate(MediaChunkIterator.EMPTY, MAX_DURATION_US))
-        .isEqualTo(C.LENGTH_UNSET);
+        .isEqualTo(Format.NO_VALUE);
   }
 
   @Test
   public void getAverageBitrate_oneChunk_returnsChunkBitrate() {
-    long[] chunkTimeBoundariesSec = {0, 5};
+    long[] chunkTimeBoundariesSec = {12, 17};
     long[] chunkLengths = {10};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
-    int expectedAverageBitrate =
-        (int) (chunkLengths[0] * C.BITS_PER_BYTE / chunkTimeBoundariesSec[1]);
-    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US))
-        .isEqualTo(expectedAverageBitrate);
+
+    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US)).isEqualTo(16);
   }
 
   @Test
   public void getAverageBitrate_multipleSameDurationChunks_returnsAverageChunkBitrate() {
     long[] chunkTimeBoundariesSec = {0, 5, 10};
     long[] chunkLengths = {10, 20};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
-    long totalLength = chunkLengths[0] + chunkLengths[1];
-    int expectedAverageBitrate = (int) (totalLength * C.BITS_PER_BYTE / chunkTimeBoundariesSec[2]);
-    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US))
-        .isEqualTo(expectedAverageBitrate);
+
+    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US)).isEqualTo(24);
   }
 
   @Test
   public void getAverageBitrate_multipleDifferentDurationChunks_returnsAverageChunkBitrate() {
     long[] chunkTimeBoundariesSec = {0, 5, 15, 30};
     long[] chunkLengths = {10, 20, 30};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
-    long totalLength = chunkLengths[0] + chunkLengths[1] + chunkLengths[2];
-    int expectedAverageBitrate = (int) (totalLength * C.BITS_PER_BYTE / chunkTimeBoundariesSec[3]);
-    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US))
-        .isEqualTo(expectedAverageBitrate);
+
+    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US)).isEqualTo(16);
   }
 
   @Test
-  public void getAverageBitrate_firstChunkLengthUnset_returnsUnsetLength() {
+  public void getAverageBitrate_firstChunkLengthUnset_returnsNoValue() {
     long[] chunkTimeBoundariesSec = {0, 5, 15, 30};
     long[] chunkLengths = {C.LENGTH_UNSET, 20, 30};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
+
     assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US))
-        .isEqualTo(C.LENGTH_UNSET);
+        .isEqualTo(Format.NO_VALUE);
   }
 
   @Test
   public void getAverageBitrate_secondChunkLengthUnset_returnsFirstChunkBitrate() {
     long[] chunkTimeBoundariesSec = {0, 5, 15, 30};
     long[] chunkLengths = {10, C.LENGTH_UNSET, 30};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
-    int expectedAverageBitrate =
-        (int) (chunkLengths[0] * C.BITS_PER_BYTE / chunkTimeBoundariesSec[1]);
-    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US))
-        .isEqualTo(expectedAverageBitrate);
+
+    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US)).isEqualTo(16);
   }
 
   @Test
@@ -96,22 +96,101 @@ public class TrackSelectionUtilTest {
       getAverageBitrate_chunksExceedingMaxDuration_returnsAverageChunkBitrateUpToMaxDuration() {
     long[] chunkTimeBoundariesSec = {0, 5, 15, 45, 50};
     long[] chunkLengths = {10, 20, 30, 100};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
-    // Just half of the third chunk is in the max duration
-    long totalLength = chunkLengths[0] + chunkLengths[1] + chunkLengths[2] / 2;
-    int expectedAverageBitrate =
-        (int) (totalLength * C.BITS_PER_BYTE * C.MICROS_PER_SECOND / MAX_DURATION_US);
-    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, MAX_DURATION_US))
-        .isEqualTo(expectedAverageBitrate);
+
+    assertThat(TrackSelectionUtil.getAverageBitrate(iterator, 30 * C.MICROS_PER_SECOND))
+        .isEqualTo(12);
   }
 
   @Test
-  public void getAverageBitrate_zeroMaxDuration_returnsUnsetLength() {
+  public void getAverageBitrate_zeroMaxDuration_returnsNoValue() {
     long[] chunkTimeBoundariesSec = {0, 5, 10};
     long[] chunkLengths = {10, 20};
+
     FakeIterator iterator = new FakeIterator(chunkTimeBoundariesSec, chunkLengths);
+
     assertThat(TrackSelectionUtil.getAverageBitrate(iterator, /* maxDurationUs= */ 0))
-        .isEqualTo(C.LENGTH_UNSET);
+        .isEqualTo(Format.NO_VALUE);
+  }
+
+  @Test
+  public void getAverageBitrates_noIterator_returnsEmptyArray() {
+    assertThat(
+            TrackSelectionUtil.getAverageBitrates(
+                new MediaChunkIterator[0], new Format[0], MAX_DURATION_US))
+        .hasLength(0);
+  }
+
+  @Test
+  public void getAverageBitrates_emptyIterator_returnsNoValue() {
+    int[] averageBitrates =
+        TrackSelectionUtil.getAverageBitrates(
+            new MediaChunkIterator[] {MediaChunkIterator.EMPTY},
+            new Format[] {createFormatWithBitrate(10)},
+            MAX_DURATION_US);
+
+    assertThat(averageBitrates).asList().containsExactly(Format.NO_VALUE);
+  }
+
+  @Test
+  public void getAverageBitrates_twoTracks_returnsAverageChunkBitrates() {
+    FakeIterator iterator1 =
+        new FakeIterator(
+            /* chunkTimeBoundariesSec= */ new long[] {0, 10}, /* chunkLengths= */ new long[] {10});
+    FakeIterator iterator2 =
+        new FakeIterator(
+            /* chunkTimeBoundariesSec= */ new long[] {0, 5, 15, 30},
+            /* chunkLengths= */ new long[] {10, 20, 30});
+
+    int[] averageBitrates =
+        TrackSelectionUtil.getAverageBitrates(
+            new MediaChunkIterator[] {iterator1, iterator2},
+            new Format[] {createFormatWithBitrate(10), createFormatWithBitrate(20)},
+            MAX_DURATION_US);
+
+    assertThat(averageBitrates).asList().containsExactly(8, 16).inOrder();
+  }
+
+  @Test
+  public void getAverageBitrates_oneEmptyIteratorOneWithChunks_returnsEstimationForEmpty() {
+    FakeIterator iterator1 =
+        new FakeIterator(
+            /* chunkTimeBoundariesSec= */ new long[] {0, 5}, /* chunkLengths= */ new long[] {10});
+    Format format1 = createFormatWithBitrate(10);
+    MediaChunkIterator iterator2 = MediaChunkIterator.EMPTY;
+    Format format2 = createFormatWithBitrate(20);
+
+    int[] averageBitrates =
+        TrackSelectionUtil.getAverageBitrates(
+            new MediaChunkIterator[] {iterator1, iterator2},
+            new Format[] {format1, format2},
+            MAX_DURATION_US);
+
+    assertThat(averageBitrates).asList().containsExactly(16, 32).inOrder();
+  }
+
+  @Test
+  public void getAverageBitrates_formatWithoutBitrate_returnsNoValueForEmpty() {
+    FakeIterator iterator1 =
+        new FakeIterator(
+            /* chunkTimeBoundariesSec= */ new long[] {0, 5}, /* chunkLengths= */ new long[] {10});
+    Format format1 = createFormatWithBitrate(10);
+    MediaChunkIterator iterator2 = MediaChunkIterator.EMPTY;
+    Format format2 = createFormatWithBitrate(Format.NO_VALUE);
+
+    int[] averageBitrates =
+        TrackSelectionUtil.getAverageBitrates(
+            new MediaChunkIterator[] {iterator1, iterator2},
+            new Format[] {format1, format2},
+            MAX_DURATION_US);
+
+    assertThat(averageBitrates).asList().containsExactly(16, Format.NO_VALUE).inOrder();
+  }
+
+  @NonNull
+  private static Format createFormatWithBitrate(int bitrate) {
+    return Format.createSampleFormat(null, null, null, bitrate, null);
   }
 
   private static final class FakeIterator extends BaseMediaChunkIterator {