From dcc3e439e202551e5525d00c2fc4c412e275f2a7 Mon Sep 17 00:00:00 2001
From: andrewlewis <andrewlewis@google.com>
Date: Wed, 10 Jul 2024 01:48:05 -0700
Subject: [PATCH] Handle AVI containers with no keyframes in index

If there is no keyframe in the index assume that the first chunk has a
keyframe.

PiperOrigin-RevId: 650915467
---
 .../media3/extractor/avi/AviExtractor.java    |  6 ++--
 .../media3/extractor/avi/ChunkReader.java     | 35 ++++++++++++-------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/avi/AviExtractor.java b/libraries/extractor/src/main/java/androidx/media3/extractor/avi/AviExtractor.java
index 4524fd1101..7fa8133820 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/avi/AviExtractor.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/avi/AviExtractor.java
@@ -411,10 +411,8 @@ public final class AviExtractor implements Extractor {
         // We ignore unknown chunk IDs.
         continue;
       }
-      if ((flags & AVIIF_KEYFRAME) == AVIIF_KEYFRAME) {
-        chunkReader.appendKeyFrameToIndex(offset);
-      }
-      chunkReader.incrementIndexChunkCount();
+      chunkReader.appendIndexChunk(
+          offset, /* isKeyFrame= */ (flags & AVIIF_KEYFRAME) == AVIIF_KEYFRAME);
     }
     for (ChunkReader chunkReader : chunkReaders) {
       chunkReader.compactIndex();
diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/avi/ChunkReader.java b/libraries/extractor/src/main/java/androidx/media3/extractor/avi/ChunkReader.java
index 7ad8197a9b..f87a7817ce 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/avi/ChunkReader.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/avi/ChunkReader.java
@@ -70,17 +70,18 @@ import java.util.Arrays;
 
   private int indexChunkCount;
   private int indexSize;
+  private long firstIndexChunkOffset;
   private long[] keyFrameOffsets;
   private int[] keyFrameIndices;
 
   public ChunkReader(
       int id,
       @C.TrackType int trackType,
-      long durationnUs,
+      long durationUs,
       int streamHeaderChunkCount,
       TrackOutput trackOutput) {
     Assertions.checkArgument(trackType == C.TRACK_TYPE_AUDIO || trackType == C.TRACK_TYPE_VIDEO);
-    this.durationUs = durationnUs;
+    this.durationUs = durationUs;
     this.streamHeaderChunkCount = streamHeaderChunkCount;
     this.trackOutput = trackOutput;
     @ChunkType
@@ -89,18 +90,25 @@ import java.util.Arrays;
     chunkId = getChunkIdFourCc(id, chunkType);
     alternativeChunkId =
         trackType == C.TRACK_TYPE_VIDEO ? getChunkIdFourCc(id, CHUNK_TYPE_VIDEO_UNCOMPRESSED) : -1;
+    firstIndexChunkOffset = C.INDEX_UNSET;
     keyFrameOffsets = new long[INITIAL_INDEX_SIZE];
     keyFrameIndices = new int[INITIAL_INDEX_SIZE];
   }
 
-  public void appendKeyFrameToIndex(long offset) {
-    if (indexSize == keyFrameIndices.length) {
-      keyFrameOffsets = Arrays.copyOf(keyFrameOffsets, keyFrameOffsets.length * 3 / 2);
-      keyFrameIndices = Arrays.copyOf(keyFrameIndices, keyFrameIndices.length * 3 / 2);
+  public void appendIndexChunk(long offset, boolean isKeyFrame) {
+    if (firstIndexChunkOffset == C.INDEX_UNSET) {
+      firstIndexChunkOffset = offset;
     }
-    keyFrameOffsets[indexSize] = offset;
-    keyFrameIndices[indexSize] = indexChunkCount;
-    indexSize++;
+    if (isKeyFrame) {
+      if (indexSize == keyFrameIndices.length) {
+        keyFrameOffsets = Arrays.copyOf(keyFrameOffsets, keyFrameOffsets.length * 3 / 2);
+        keyFrameIndices = Arrays.copyOf(keyFrameIndices, keyFrameIndices.length * 3 / 2);
+      }
+      keyFrameOffsets[indexSize] = offset;
+      keyFrameIndices[indexSize] = indexChunkCount;
+      indexSize++;
+    }
+    indexChunkCount++;
   }
 
   public void advanceCurrentChunk() {
@@ -115,10 +123,6 @@ import java.util.Arrays;
     return getChunkTimestampUs(/* chunkIndex= */ 1);
   }
 
-  public void incrementIndexChunkCount() {
-    indexChunkCount++;
-  }
-
   public void compactIndex() {
     keyFrameOffsets = Arrays.copyOf(keyFrameOffsets, indexSize);
     keyFrameIndices = Arrays.copyOf(keyFrameIndices, indexSize);
@@ -180,6 +184,11 @@ import java.util.Arrays;
   }
 
   public SeekMap.SeekPoints getSeekPoints(long timeUs) {
+    if (indexSize == 0) {
+      // Return the offset of the first chunk as there are no keyframes in the index.
+      return new SeekMap.SeekPoints(
+          new SeekPoint(/* timeUs= */ 0, /* position= */ firstIndexChunkOffset));
+    }
     int targetFrameIndex = (int) (timeUs / getFrameDurationUs());
     int keyFrameIndex =
         Util.binarySearchFloor(