Optimize AvcChunkHandler to use normal ChunkClock if no B Frames.

library/extractor/src/main/java/com/google/android/exoplayer2/extractor/avi/AvcChunkHandler.java

@@ -16,6 +16,7 @@
 package com.google.android.exoplayer2.extractor.avi;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.extractor.ExtractorInput;
@@ -43,28 +44,36 @@ public class AvcChunkHandler extends NalChunkPeeker {
 
   public AvcChunkHandler(int id, @NonNull TrackOutput trackOutput,
       @NonNull ChunkClock clock, Format.Builder formatBuilder) {
-    super(id, trackOutput, new PicCountClock(clock.durationUs, clock.chunks), 16);
+    super(id, trackOutput, clock, 16);
     this.formatBuilder = formatBuilder;
   }
 
-  @NonNull
-  @Override
-  public PicCountClock getClock() {
-    return (PicCountClock) clock;
+  @Nullable
+  @VisibleForTesting
+  PicCountClock getPicCountClock() {
+    if (clock instanceof PicCountClock) {
+      return (PicCountClock)clock;
+    } else {
+      return null;
+    }
   }
 
   @Override
   boolean skip(byte nalType) {
-    return false;
+    if (clock instanceof PicCountClock) {
+      return false;
+    } else {
+      //If the clock is regular clock, skip "normal" frames
+      return nalType >= 0 && nalType <= NAL_TYPE_IDR;
+    }
   }
 
   /**
    * Greatly simplified way to calculate the picOrder
    * Full logic is here
    * https://chromium.googlesource.com/chromium/src/media/+/refs/heads/main/video/h264_poc.cc
-   * @param nalTypeOffset
    */
-  void updatePicCountClock(final int nalTypeOffset) {
+  void updatePicCountClock(final int nalTypeOffset, final PicCountClock picCountClock) {
     final ParsableNalUnitBitArray in = new ParsableNalUnitBitArray(buffer, nalTypeOffset + 1, buffer.length);
     //slide_header()
     in.readUnsignedExpGolombCodedInt(); //first_mb_in_slice
@@ -84,10 +93,10 @@ public class AvcChunkHandler extends NalChunkPeeker {
     if (spsData.picOrderCountType == 0) {
       int picOrderCountLsb = in.readBits(spsData.picOrderCntLsbLength);
       //Log.d("Test", "FrameNum: " + frame + " cnt=" + picOrderCountLsb);
-      getClock().setPicCount(picOrderCountLsb);
+      picCountClock.setPicCount(picOrderCountLsb);
       return;
     } else if (spsData.picOrderCountType == 2) {
-      getClock().setPicCount(frameNum);
+      picCountClock.setPicCount(frameNum);
       return;
     }
     clock.setIndex(clock.getIndex());
@@ -98,11 +107,20 @@ public class AvcChunkHandler extends NalChunkPeeker {
     final int spsStart = nalTypeOffset + 1;
     nalTypeOffset = seekNextNal(input, spsStart);
     spsData = NalUnitUtil.parseSpsNalUnitPayload(buffer, spsStart, pos);
-    if (spsData.picOrderCountType == 0) {
-      getClock().setMaxPicCount(1 << spsData.picOrderCntLsbLength, 2);
-    } else if (spsData.picOrderCountType == 2) {
-      //Plus one because we double the frame number
-      getClock().setMaxPicCount(1 << spsData.frameNumLength, 1);
+    //If we have B Frames, upgrade to PicCountClock
+    final PicCountClock picCountClock;
+    if (spsData.maxNumRefFrames > 1 && !(clock instanceof PicCountClock)) {
+      clock = picCountClock = new PicCountClock(clock.durationUs, clock.chunks);
+    } else {
+      picCountClock = getPicCountClock();
+    }
+    if (picCountClock != null) {
+      if (spsData.picOrderCountType == 0) {
+        picCountClock.setMaxPicCount(1 << spsData.picOrderCntLsbLength, 2);
+      } else if (spsData.picOrderCountType == 2) {
+        //Plus one because we double the frame number
+        picCountClock.setMaxPicCount(1 << spsData.frameNumLength, 1);
+      }
     }
     if (spsData.pixelWidthHeightRatio != pixelWidthHeightRatio) {
       pixelWidthHeightRatio = spsData.pixelWidthHeightRatio;
@@ -121,11 +139,17 @@ public class AvcChunkHandler extends NalChunkPeeker {
         case 2:
         case 3:
         case 4:
-          updatePicCountClock(nalTypeOffset);
+          if (clock instanceof PicCountClock) {
+            updatePicCountClock(nalTypeOffset, (PicCountClock)clock);
+          }
           return;
-        case NAL_TYPE_IDR:
-          getClock().syncIndexes();
+        case NAL_TYPE_IDR: {
+          final PicCountClock picCountClock = getPicCountClock();
+          if (picCountClock != null) {
+            picCountClock.syncIndexes();
+          }
           return;
+        }
         case NAL_TYPE_AUD:
         case NAL_TYPE_SEI:
         case NAL_TYPE_PPS: {

library/extractor/src/main/java/com/google/android/exoplayer2/util/NalUnitUtil.java

@@ -33,6 +33,7 @@ public final class NalUnitUtil {
     public final int constraintsFlagsAndReservedZero2Bits;
     public final int levelIdc;
     public final int seqParameterSetId;
+    public final int maxNumRefFrames;
     public final int width;
     public final int height;
     public final float pixelWidthHeightRatio;
@@ -48,6 +49,7 @@ public final class NalUnitUtil {
         int constraintsFlagsAndReservedZero2Bits,
         int levelIdc,
         int seqParameterSetId,
+        int maxNumRefFrames,
         int width,
         int height,
         float pixelWidthHeightRatio,
@@ -61,6 +63,7 @@ public final class NalUnitUtil {
       this.constraintsFlagsAndReservedZero2Bits = constraintsFlagsAndReservedZero2Bits;
       this.levelIdc = levelIdc;
       this.seqParameterSetId = seqParameterSetId;
+      this.maxNumRefFrames = maxNumRefFrames;
       this.width = width;
       this.height = height;
       this.pixelWidthHeightRatio = pixelWidthHeightRatio;
@@ -367,7 +370,7 @@ public final class NalUnitUtil {
         data.readUnsignedExpGolombCodedInt(); // offset_for_ref_frame[i]
       }
     }
-    data.readUnsignedExpGolombCodedInt(); // max_num_ref_frames
+    int maxNumRefFrames = data.readUnsignedExpGolombCodedInt(); // max_num_ref_frames
     data.skipBit(); // gaps_in_frame_num_value_allowed_flag
 
     int picWidthInMbs = data.readUnsignedExpGolombCodedInt() + 1;
@@ -427,6 +430,7 @@ public final class NalUnitUtil {
         constraintsFlagsAndReservedZero2Bits,
         levelIdc,
         seqParameterSetId,
+        maxNumRefFrames,
         frameWidth,
         frameHeight,
         pixelWidthHeightRatio,

library/extractor/src/test/java/com/google/android/exoplayer2/extractor/avi/AvcChunkPeekerTest.java

@@ -35,8 +35,8 @@ public class AvcChunkPeekerTest {
       setSampleMimeType(MimeTypes.VIDEO_H264).
       setWidth(1280).setHeight(720).setFrameRate(24000f/1001f);
 
-  private static final byte[] P_SLICE = {00,00,00,01,0x41,(byte)0x9A,0x13,0x36,0x21,0x3A,0x5F,
-      (byte)0xFE,(byte)0x9E,0x10,00,00};
+  private static final byte[] P_SLICE = {0,0,0,1,0x41,(byte)0x9A,0x13,0x36,0x21,0x3A,0x5F,
+      (byte)0xFE,(byte)0x9E,0x10,0,0};
 
   private FakeTrackOutput fakeTrackOutput;
   private AvcChunkHandler avcChunkPeeker;
@@ -61,19 +61,20 @@ public class AvcChunkPeekerTest {
   @Test
   public void peek_givenStreamHeader() throws IOException {
     peekStreamHeader();
-    final PicCountClock picCountClock = avcChunkPeeker.getClock();
+    final PicCountClock picCountClock = avcChunkPeeker.getPicCountClock();
+    Assert.assertNotNull(picCountClock);
     Assert.assertEquals(64, picCountClock.getMaxPicCount());
     Assert.assertEquals(0, avcChunkPeeker.getSpsData().picOrderCountType);
     Assert.assertEquals(1.18f, fakeTrackOutput.lastFormat.pixelWidthHeightRatio, 0.01f);
   }
 
   @Test
-  public void peek_givenStreamHeaderAndPSlice() throws IOException {
+  public void newChunk_givenStreamHeaderAndPSlice() throws IOException {
     peekStreamHeader();
-    final PicCountClock picCountClock = avcChunkPeeker.getClock();
+    final PicCountClock picCountClock = avcChunkPeeker.getPicCountClock();
     final FakeExtractorInput input = new FakeExtractorInput.Builder().setData(P_SLICE).build();
 
-    avcChunkPeeker.peek(input, P_SLICE.length);
+    avcChunkPeeker.newChunk(0, P_SLICE.length, input);
 
     Assert.assertEquals(12, picCountClock.getLastPicCount());
   }