From 36be62ee8c26ff9f0103158802dfab92356e8995 Mon Sep 17 00:00:00 2001
From: sheenachhabra <sheenachhabra@google.com>
Date: Fri, 10 Jan 2025 13:23:38 -0800
Subject: [PATCH] Skip duplicate chunks in stsc box writing

Earlier muxer wrote a different entry for every chunk but as
per the spec (ISO 14496-12) exact same chunks can be combined
and only the first chunk number can be written.

PiperOrigin-RevId: 714154531
---
 .../java/androidx/media3/muxer/Boxes.java     | 21 ++++++++++++-------
 .../java/androidx/media3/muxer/BoxesTest.java | 13 +++++++++++-
 .../muxerdumps/stsc_box_with_same_chunks.dump |  2 ++
 3 files changed, 28 insertions(+), 8 deletions(-)
 create mode 100644 libraries/test_data/src/test/assets/muxerdumps/stsc_box_with_same_chunks.dump

diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java b/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
index a7e929b957..26f4f64e85 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
@@ -1065,21 +1065,28 @@ import org.checkerframework.checker.nullness.qual.PolyNull;
         ByteBuffer.allocate(writtenChunkSampleCounts.size() * 12 + MAX_FIXED_LEAF_BOX_SIZE);
 
     contents.putInt(0x0); // version and flags
-    contents.putInt(writtenChunkSampleCounts.size()); // entry_count
+    int totalEntryCountIndex = contents.position();
+    contents.putInt(0); // entry_count
 
     int currentChunk = 1;
+    int prevChunkSampleCount = -1;
+    int totalEntryCount = 0;
 
-    // TODO: b/270583563 - Consider optimizing for consecutive chunks having same number of samples.
     for (int i = 0; i < writtenChunkSampleCounts.size(); i++) {
       int samplesInChunk = writtenChunkSampleCounts.get(i);
-      contents.putInt(currentChunk); // first_chunk
-      contents.putInt(samplesInChunk); // samples_per_chunk
-      // sample_description_index: there is only one sample description in each track.
-      contents.putInt(1);
-
+      // For exact same chunks, add only first chunk number.
+      if (samplesInChunk != prevChunkSampleCount) {
+        contents.putInt(currentChunk); // first_chunk
+        contents.putInt(samplesInChunk); // samples_per_chunk
+        // sample_description_index: there is only one sample description in each track.
+        contents.putInt(1);
+        totalEntryCount++;
+        prevChunkSampleCount = samplesInChunk;
+      }
       currentChunk += 1;
     }
 
+    contents.putInt(totalEntryCountIndex, totalEntryCount);
     contents.flip();
     return BoxUtils.wrapIntoBox("stsc", contents);
   }
diff --git a/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java b/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java
index 60e191857e..75a0a12ee0 100644
--- a/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java
+++ b/libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java
@@ -735,7 +735,7 @@ public class BoxesTest {
   }
 
   @Test
-  public void createStscBox_matchesExpected() throws IOException {
+  public void createStscBox_withDifferentChunks_matchesExpected() throws IOException {
     ImmutableList<Integer> chunkSampleCounts = ImmutableList.of(100, 500, 200, 100);
 
     ByteBuffer stscBox = Boxes.stsc(chunkSampleCounts);
@@ -745,6 +745,17 @@ public class BoxesTest {
         context, dumpableBox, MuxerTestUtil.getExpectedDumpFilePath("stsc_box"));
   }
 
+  @Test
+  public void createStscBox_withSameChunks_matchesExpected() throws IOException {
+    ImmutableList<Integer> chunkSampleCounts = ImmutableList.of(100, 100, 100, 100);
+
+    ByteBuffer stscBox = Boxes.stsc(chunkSampleCounts);
+
+    DumpableMp4Box dumpableBox = new DumpableMp4Box(stscBox);
+    DumpFileAsserts.assertOutput(
+        context, dumpableBox, MuxerTestUtil.getExpectedDumpFilePath("stsc_box_with_same_chunks"));
+  }
+
   @Test
   public void createStcoBox_matchesExpected() throws IOException {
     ImmutableList<Long> chunkOffsets = ImmutableList.of(1_000L, 5_000L, 7_000L, 10_000L);
diff --git a/libraries/test_data/src/test/assets/muxerdumps/stsc_box_with_same_chunks.dump b/libraries/test_data/src/test/assets/muxerdumps/stsc_box_with_same_chunks.dump
new file mode 100644
index 0000000000..688ff3b4a5
--- /dev/null
+++ b/libraries/test_data/src/test/assets/muxerdumps/stsc_box_with_same_chunks.dump
@@ -0,0 +1,2 @@
+stsc (28 bytes):
+  Data = length 20, hash 94E19968