From 7db0d91dd9b09c6139294eeecaf6b04afa56d822 Mon Sep 17 00:00:00 2001
From: aquilescanta <aquilescanta@google.com>
Date: Thu, 24 Mar 2022 14:41:28 +0000
Subject: [PATCH] Add maxNumRefFrames to NalUnitUtil

Needed for AVI support.

PiperOrigin-RevId: 436991922
---
 .../main/java/androidx/media3/extractor/NalUnitUtil.java    | 6 +++++-
 .../java/androidx/media3/extractor/NalUnitUtilTest.java     | 1 +
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/NalUnitUtil.java b/libraries/extractor/src/main/java/androidx/media3/extractor/NalUnitUtil.java
index e25b370897..c61d7eaba5 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/NalUnitUtil.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/NalUnitUtil.java
@@ -53,6 +53,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;
@@ -68,6 +69,7 @@ public final class NalUnitUtil {
         int constraintsFlagsAndReservedZero2Bits,
         int levelIdc,
         int seqParameterSetId,
+        int maxNumRefFrames,
         int width,
         int height,
         float pixelWidthHeightRatio,
@@ -81,6 +83,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;
@@ -387,7 +390,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;
@@ -447,6 +450,7 @@ public final class NalUnitUtil {
         constraintsFlagsAndReservedZero2Bits,
         levelIdc,
         seqParameterSetId,
+        maxNumRefFrames,
         frameWidth,
         frameHeight,
         pixelWidthHeightRatio,
diff --git a/libraries/extractor/src/test/java/androidx/media3/extractor/NalUnitUtilTest.java b/libraries/extractor/src/test/java/androidx/media3/extractor/NalUnitUtilTest.java
index ebd2ccc45f..01d7fe15f9 100644
--- a/libraries/extractor/src/test/java/androidx/media3/extractor/NalUnitUtilTest.java
+++ b/libraries/extractor/src/test/java/androidx/media3/extractor/NalUnitUtilTest.java
@@ -126,6 +126,7 @@ public final class NalUnitUtilTest {
   public void parseSpsNalUnit() {
     NalUnitUtil.SpsData data =
         NalUnitUtil.parseSpsNalUnit(SPS_TEST_DATA, SPS_TEST_DATA_OFFSET, SPS_TEST_DATA.length);
+    assertThat(data.maxNumRefFrames).isEqualTo(4);
     assertThat(data.width).isEqualTo(640);
     assertThat(data.height).isEqualTo(360);
     assertThat(data.deltaPicOrderAlwaysZeroFlag).isFalse();