Eliminate memory copy of H264 data through H264 reader.

I think this is the limit of how far we should be pushing complexity
v.s. efficiency. It's a little complicated to understand, but probably
worth it since the H264 bitstream is the majority of the data.

Issue: #278

library/src/main/java/com/google/android/exoplayer/hls/parser/AdtsReader.java

@@ -85,7 +85,7 @@ import java.util.Collections;
           break;
         case STATE_READING_SAMPLE:
           int bytesToRead = Math.min(data.bytesLeft(), sampleSize - bytesRead);
-          appendSampleData(data, bytesToRead);
+          appendData(data, bytesToRead);
           bytesRead += bytesToRead;
           if (bytesRead == sampleSize) {
             commitSample(true);

library/src/main/java/com/google/android/exoplayer/hls/parser/H264Reader.java

@@ -18,11 +18,13 @@ package com.google.android.exoplayer.hls.parser;
 import com.google.android.exoplayer.MediaFormat;
 import com.google.android.exoplayer.mp4.Mp4Util;
 import com.google.android.exoplayer.upstream.BufferPool;
+import com.google.android.exoplayer.util.Assertions;
 import com.google.android.exoplayer.util.MimeTypes;
 import com.google.android.exoplayer.util.ParsableBitArray;
 import com.google.android.exoplayer.util.ParsableByteArray;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -37,54 +39,72 @@ import java.util.List;
   private static final int NAL_UNIT_TYPE_AUD = 9;
 
   private final SeiReader seiReader;
-  private final ParsableByteArray pendingSampleWrapper;
+  private final boolean[] prefixFlags;
+  private final NalUnitTargetBuffer sps;
+  private final NalUnitTargetBuffer pps;
+  private final NalUnitTargetBuffer sei;
 
-  // TODO: Ideally we wouldn't need to have a copy step through a byte array here.
-  private byte[] pendingSampleData;
-  private int pendingSampleSize;
-  private long pendingSampleTimeUs;
+  private boolean isKeyframe;
 
   public H264Reader(BufferPool bufferPool, SeiReader seiReader) {
     super(bufferPool);
     this.seiReader = seiReader;
-    this.pendingSampleData = new byte[1024];
-    this.pendingSampleWrapper = new ParsableByteArray();
+    prefixFlags = new boolean[3];
+    sps = new NalUnitTargetBuffer(NAL_UNIT_TYPE_SPS, 128);
+    pps = new NalUnitTargetBuffer(NAL_UNIT_TYPE_PPS, 128);
+    sei = new NalUnitTargetBuffer(NAL_UNIT_TYPE_SEI, 128);
   }
 
   @Override
   public void consume(ParsableByteArray data, long pesTimeUs, boolean startOfPacket) {
     while (data.bytesLeft() > 0) {
-      boolean sampleFinished = readToNextAudUnit(data, pesTimeUs);
-      if (!sampleFinished) {
-        continue;
-      }
+      int offset = data.getPosition();
+      int limit = data.limit();
+      byte[] dataArray = data.data;
+
+      // Append the data to the buffer.
+      appendData(data, data.bytesLeft());
+
+      // Scan the appended data, processing NAL units as they are encountered
+      while (offset < limit) {
+        int nextNalUnitOffset = Mp4Util.findNalUnit(dataArray, offset, limit, prefixFlags);
+        if (nextNalUnitOffset < limit) {
+          // We've seen the start of a NAL unit.
 
-      // Scan the sample to find relevant NAL units.
-      int position = 0;
-      int idrNalUnitPosition = Integer.MAX_VALUE;
-      while (position < pendingSampleSize) {
-        position = Mp4Util.findNalUnit(pendingSampleData, position, pendingSampleSize);
-        if (position < pendingSampleSize) {
-          int type = Mp4Util.getNalUnitType(pendingSampleData, position);
-          if (type == NAL_UNIT_TYPE_IDR) {
-            idrNalUnitPosition = position;
-          } else if (type == NAL_UNIT_TYPE_SEI) {
-            seiReader.read(pendingSampleData, position, pendingSampleTimeUs);
+          // This is the length to the start of the unit. It may be negative if the NAL unit
+          // actually started in previously consumed data.
+          int lengthToNalUnit = nextNalUnitOffset - offset;
+          if (lengthToNalUnit > 0) {
+            feedNalUnitTargetBuffersData(dataArray, offset, nextNalUnitOffset);
           }
-          position += 4;
-        }
-      }
 
-      // Determine whether the sample is a keyframe.
-      boolean isKeyframe = pendingSampleSize > idrNalUnitPosition;
-      if (!hasMediaFormat() && isKeyframe) {
-        parseMediaFormat(pendingSampleData, pendingSampleSize);
-      }
+          int nalUnitType = Mp4Util.getNalUnitType(dataArray, nextNalUnitOffset);
+          int nalUnitOffsetInData = nextNalUnitOffset - limit;
+          if (nalUnitType == NAL_UNIT_TYPE_AUD) {
+            if (writingSample()) {
+              if (isKeyframe && !hasMediaFormat() && sps.isCompleted() && pps.isCompleted()) {
+                parseMediaFormat(sps, pps);
+              }
+              commitSample(isKeyframe, nalUnitOffsetInData);
+            }
+            startSample(pesTimeUs, nalUnitOffsetInData);
+            isKeyframe = false;
+          } else if (nalUnitType == NAL_UNIT_TYPE_IDR) {
+            isKeyframe = true;
+          }
 
-      // Commit the sample to the queue.
-      pendingSampleWrapper.reset(pendingSampleData, pendingSampleSize);
-      appendSampleData(pendingSampleWrapper, pendingSampleSize);
-      commitSample(isKeyframe);
+          // If the length to the start of the unit is negative then we wrote too many bytes to the
+          // NAL buffers. Discard the excess bytes when notifying that the unit has ended.
+          feedNalUnitTargetEnd(pesTimeUs, lengthToNalUnit < 0 ? -lengthToNalUnit : 0);
+          // Notify the start of the next NAL unit.
+          feedNalUnitTargetBuffersStart(nalUnitType);
+          // Continue scanning the data.
+          offset = nextNalUnitOffset + 4;
+        } else {
+          feedNalUnitTargetBuffersData(dataArray, offset, limit);
+          offset = limit;
+        }
+      }
     }
   }
 
@@ -93,71 +113,41 @@ import java.util.List;
     // Do nothing.
   }
 
-  /**
-   * Reads data up to (but not including) the start of the next AUD unit.
-   *
-   * @param data The data to consume.
-   * @param pesTimeUs The corresponding time.
-   * @return True if the current sample is now complete. False otherwise.
-   */
-  private boolean readToNextAudUnit(ParsableByteArray data, long pesTimeUs) {
-    int pesOffset = data.getPosition();
-    int pesLimit = data.limit();
-
-    // TODO: We probably need to handle the case where the AUD start code was split across the
-    // previous and current data buffers.
-    int audOffset = Mp4Util.findNalUnit(data.data, pesOffset, pesLimit, NAL_UNIT_TYPE_AUD);
-    int bytesToNextAud = audOffset - pesOffset;
-    if (bytesToNextAud == 0) {
-      if (!writingSample()) {
-        startSample(pesTimeUs);
-        pendingSampleSize = 0;
-        pendingSampleTimeUs = pesTimeUs;
-        appendToSample(data, 4);
-        return false;
-      } else {
-        return true;
-      }
-    } else if (writingSample()) {
-      appendToSample(data, bytesToNextAud);
-      return data.bytesLeft() > 0;
-    } else {
-      data.skip(bytesToNextAud);
-      return false;
+  private void feedNalUnitTargetBuffersStart(int nalUnitType) {
+    if (!hasMediaFormat()) {
+      sps.startNalUnit(nalUnitType);
+      pps.startNalUnit(nalUnitType);
     }
+    sei.startNalUnit(nalUnitType);
   }
 
-  private void appendToSample(ParsableByteArray data, int length) {
-    int requiredSize = pendingSampleSize + length;
-    if (pendingSampleData.length < requiredSize) {
-      byte[] newPendingSampleData = new byte[(requiredSize * 3) / 2];
-      System.arraycopy(pendingSampleData, 0, newPendingSampleData, 0, pendingSampleSize);
-      pendingSampleData = newPendingSampleData;
+  private void feedNalUnitTargetBuffersData(byte[] dataArray, int offset, int limit) {
+    if (!hasMediaFormat()) {
+      sps.appendToNalUnit(dataArray, offset, limit);
+      pps.appendToNalUnit(dataArray, offset, limit);
     }
-    data.readBytes(pendingSampleData, pendingSampleSize, length);
-    pendingSampleSize += length;
+    sei.appendToNalUnit(dataArray, offset, limit);
   }
 
-  private void parseMediaFormat(byte[] sampleData, int sampleSize) {
-    // Locate the SPS and PPS units.
-    int spsOffset = Mp4Util.findNalUnit(sampleData, 0, sampleSize, NAL_UNIT_TYPE_SPS);
-    int ppsOffset = Mp4Util.findNalUnit(sampleData, 0, sampleSize, NAL_UNIT_TYPE_PPS);
-    if (spsOffset == sampleSize || ppsOffset == sampleSize) {
-      return;
+  private void feedNalUnitTargetEnd(long pesTimeUs, int discardPadding) {
+    sps.endNalUnit(discardPadding);
+    pps.endNalUnit(discardPadding);
+    if (sei.endNalUnit(discardPadding)) {
+      seiReader.read(sei.nalData, 0, pesTimeUs);
     }
-    // Determine the length of the units, and copy them to build the initialization data.
-    int spsLength = Mp4Util.findNalUnit(sampleData, spsOffset + 3, sampleSize) - spsOffset;
-    int ppsLength = Mp4Util.findNalUnit(sampleData, ppsOffset + 3, sampleSize) - ppsOffset;
-    byte[] spsData = new byte[spsLength];
-    byte[] ppsData = new byte[ppsLength];
-    System.arraycopy(sampleData, spsOffset, spsData, 0, spsLength);
-    System.arraycopy(sampleData, ppsOffset, ppsData, 0, ppsLength);
+  }
+
+  private void parseMediaFormat(NalUnitTargetBuffer sps, NalUnitTargetBuffer pps) {
+    byte[] spsData = new byte[sps.nalLength];
+    byte[] ppsData = new byte[pps.nalLength];
+    System.arraycopy(sps.nalData, 0, spsData, 0, sps.nalLength);
+    System.arraycopy(pps.nalData, 0, ppsData, 0, pps.nalLength);
     List<byte[]> initializationData = new ArrayList<byte[]>();
     initializationData.add(spsData);
     initializationData.add(ppsData);
 
     // Unescape and then parse the SPS unit.
-    byte[] unescapedSps = unescapeStream(spsData, 0, spsLength);
+    byte[] unescapedSps = unescapeStream(spsData, 0, spsData.length);
     ParsableBitArray bitArray = new ParsableBitArray(unescapedSps);
     bitArray.skipBits(32); // NAL header
     int profileIdc = bitArray.readBits(8);
@@ -293,4 +283,83 @@ import java.util.List;
     return limit;
   }
 
+  /**
+   * A buffer that fills itself with data corresponding to a specific NAL unit, as it is
+   * encountered in the stream.
+   */
+  private static final class NalUnitTargetBuffer {
+
+    private final int targetType;
+
+    private boolean isFilling;
+    private boolean isCompleted;
+
+    public byte[] nalData;
+    public int nalLength;
+
+    public NalUnitTargetBuffer(int targetType, int initialCapacity) {
+      this.targetType = targetType;
+      // Initialize data, writing the known NAL prefix into the first four bytes.
+      nalData = new byte[4 + initialCapacity];
+      nalData[2] = 1;
+      nalData[3] = (byte) targetType;
+    }
+
+    public boolean isCompleted() {
+      return isCompleted;
+    }
+
+    /**
+     * Invoked to indicate that a NAL unit has started.
+     *
+     * @param type The type of the NAL unit.
+     */
+    public void startNalUnit(int type) {
+      Assertions.checkState(!isFilling);
+      isFilling = type == targetType;
+      if (isFilling) {
+        // Length is initially the length of the NAL prefix.
+        nalLength = 4;
+        isCompleted = false;
+      }
+    }
+
+    /**
+     * Invoked to pass stream data. The data passed should not include 4 byte NAL unit prefixes.
+     *
+     * @param data Holds the data being passed.
+     * @param offset The offset of the data in {@code data}.
+     * @param limit The limit (exclusive) of the data in {@code data}.
+     */
+    public void appendToNalUnit(byte[] data, int offset, int limit) {
+      if (!isFilling) {
+        return;
+      }
+      int readLength = limit - offset;
+      if (nalData.length < nalLength + readLength) {
+        nalData = Arrays.copyOf(nalData, (nalLength + readLength) * 2);
+      }
+      System.arraycopy(data, offset, nalData, nalLength, readLength);
+      nalLength += readLength;
+    }
+
+    /**
+     * Invoked to indicate that a NAL unit has ended.
+     *
+     * @param discardPadding The number of excess bytes that were passed to
+     *     {@link #appendData(byte[], int, int)}, which should be discarded.
+     * @return True if the ended NAL unit is of the target type. False otherwise.
+     */
+    public boolean endNalUnit(int discardPadding) {
+      if (!isFilling) {
+        return false;
+      }
+      nalLength -= discardPadding;
+      isFilling = false;
+      isCompleted = true;
+      return true;
+    }
+
+  }
+
 }

library/src/main/java/com/google/android/exoplayer/hls/parser/Id3Reader.java

@@ -35,7 +35,7 @@ import com.google.android.exoplayer.util.ParsableByteArray;
       startSample(pesTimeUs);
     }
     if (writingSample()) {
-      appendSampleData(data, data.bytesLeft());
+      appendData(data, data.bytesLeft());
     }
   }
 

library/src/main/java/com/google/android/exoplayer/hls/parser/RollingSampleBuffer.java

@@ -17,6 +17,7 @@ package com.google.android.exoplayer.hls.parser;
 
 import com.google.android.exoplayer.SampleHolder;
 import com.google.android.exoplayer.upstream.BufferPool;
+import com.google.android.exoplayer.util.Assertions;
 import com.google.android.exoplayer.util.ParsableByteArray;
 
 import android.annotation.SuppressLint;
@@ -44,7 +45,6 @@ import java.util.concurrent.ConcurrentLinkedQueue;
   private long totalBytesWritten;
   private byte[] lastFragment;
   private int lastFragmentOffset;
-  private int pendingSampleSize;
   private long pendingSampleTimeUs;
   private long pendingSampleOffset;
 
@@ -141,23 +141,25 @@ import java.util.concurrent.ConcurrentLinkedQueue;
   // Called by the loading thread.
 
   /**
-   * Starts writing the next sample.
+   * Indicates the start point for the next sample.
    *
    * @param sampleTimeUs The sample timestamp.
+   * @param offset The offset of the sample's data, relative to the total number of bytes written
+   *     to the buffer. Must be negative or zero.
    */
-  public void startSample(long sampleTimeUs) {
+  public void startSample(long sampleTimeUs, int offset) {
+    Assertions.checkState(offset <= 0);
     pendingSampleTimeUs = sampleTimeUs;
-    pendingSampleOffset = totalBytesWritten;
-    pendingSampleSize = 0;
+    pendingSampleOffset = totalBytesWritten + offset;
   }
 
   /**
-   * Appends data to the sample currently being written.
+   * Appends data to the rolling buffer.
    *
    * @param buffer A buffer containing the data to append.
    * @param length The length of the data to append.
    */
-  public void appendSampleData(ParsableByteArray buffer, int length) {
+  public void appendData(ParsableByteArray buffer, int length) {
     int remainingWriteLength = length;
     while (remainingWriteLength > 0) {
       if (dataQueue.isEmpty() || lastFragmentOffset == fragmentLength) {
@@ -171,17 +173,20 @@ import java.util.concurrent.ConcurrentLinkedQueue;
       remainingWriteLength -= thisWriteLength;
     }
     totalBytesWritten += length;
-    pendingSampleSize += length;
   }
 
   /**
-   * Commits the sample currently being written, making it available for consumption.
+   * Indicates the end point for the current sample, making it available for consumption.
    *
    * @param isKeyframe True if the sample being committed is a keyframe. False otherwise.
+   * @param offset The offset of the first byte after the end of the sample's data, relative to
+   *     the total number of bytes written to the buffer. Must be negative or zero.
    */
   @SuppressLint("InlinedApi")
-  public void commitSample(boolean isKeyframe) {
-    infoQueue.commitSample(pendingSampleTimeUs, pendingSampleOffset, pendingSampleSize,
+  public void commitSample(boolean isKeyframe, int offset) {
+    Assertions.checkState(offset <= 0);
+    int sampleSize = (int) (totalBytesWritten + offset - pendingSampleOffset);
+    infoQueue.commitSample(pendingSampleTimeUs, pendingSampleOffset, sampleSize,
         isKeyframe ? MediaExtractor.SAMPLE_FLAG_SYNC : 0);
   }
 

library/src/main/java/com/google/android/exoplayer/hls/parser/SampleQueue.java

@@ -181,17 +181,25 @@ import android.media.MediaExtractor;
   }
 
   protected void startSample(long sampleTimeUs) {
+    startSample(sampleTimeUs, 0);
+  }
+
+  protected void startSample(long sampleTimeUs, int offset) {
     writingSample = true;
     largestParsedTimestampUs = Math.max(largestParsedTimestampUs, sampleTimeUs);
-    rollingBuffer.startSample(sampleTimeUs);
+    rollingBuffer.startSample(sampleTimeUs, offset);
   }
 
-  protected void appendSampleData(ParsableByteArray buffer, int size) {
-    rollingBuffer.appendSampleData(buffer, size);
+  protected void appendData(ParsableByteArray buffer, int length) {
+    rollingBuffer.appendData(buffer, length);
   }
 
   protected void commitSample(boolean isKeyframe) {
-    rollingBuffer.commitSample(isKeyframe);
+    commitSample(isKeyframe, 0);
+  }
+
+  protected void commitSample(boolean isKeyframe, int offset) {
+    rollingBuffer.commitSample(isKeyframe, offset);
     writingSample = false;
   }
 

library/src/main/java/com/google/android/exoplayer/hls/parser/SeiReader.java

@@ -42,7 +42,7 @@ import com.google.android.exoplayer.util.ParsableByteArray;
     int ccDataSize = Eia608Parser.parseHeader(seiBuffer);
     if (ccDataSize > 0) {
       startSample(pesTimeUs);
-      appendSampleData(seiBuffer, ccDataSize);
+      appendData(seiBuffer, ccDataSize);
       commitSample(true);
     }
   }

library/src/main/java/com/google/android/exoplayer/mp4/Mp4Util.java

@@ -15,6 +15,7 @@
  */
 package com.google.android.exoplayer.mp4;
 
+import com.google.android.exoplayer.util.Assertions;
 import com.google.android.exoplayer.util.CodecSpecificDataUtil;
 import com.google.android.exoplayer.util.ParsableByteArray;
 
@@ -100,30 +101,66 @@ public final class Mp4Util {
   }
 
   /**
-   * Like {@link #findNalUnit(byte[], int, int, int)} with {@code type == -1}.
+   * Finds the first NAL unit in {@code data}.
+   * <p>
+   * For a NAL unit to be found, its first four bytes must be contained within the part of the
+   * array being searched.
    *
    * @param data The data to search.
    * @param startOffset The offset (inclusive) in the data to start the search.
    * @param endOffset The offset (exclusive) in the data to end the search.
+   * @param type The type of the NAL unit to search for, or -1 for any NAL unit.
    * @return The offset of the NAL unit, or {@code endOffset} if a NAL unit was not found.
    */
-  public static int findNalUnit(byte[] data, int startOffset, int endOffset) {
-    return findNalUnit(data, startOffset, endOffset, -1);
+  public static int findNalUnit(byte[] data, int startOffset, int endOffset, int type) {
+    return findNalUnit(data, startOffset, endOffset, type, null);
   }
 
   /**
-   * Finds the first NAL unit in {@code data}.
+   * Like {@link #findNalUnit(byte[], int, int, int)}, but supports finding of NAL units across
+   * array boundaries.
    * <p>
-   * For a NAL unit to be found, its first four bytes must be contained within the part of the
-   * array being searched.
+   * To use this method, pass the same {@code prefixFlags} parameter to successive calls where the
+   * data passed represents a contiguous stream. The state maintained in this parameter allows the
+   * detection of NAL units where the NAL unit prefix spans array boundaries.
+   * <p>
+   * Note that when using {@code prefixFlags} the return value may be 3, 2 or 1 less than
+   * {@code startOffset}, to indicate a NAL unit starting 3, 2 or 1 bytes before the first byte in
+   * the current array.
    *
    * @param data The data to search.
    * @param startOffset The offset (inclusive) in the data to start the search.
    * @param endOffset The offset (exclusive) in the data to end the search.
    * @param type The type of the NAL unit to search for, or -1 for any NAL unit.
+   * @param prefixFlags A boolean array whose first three elements are used to store the state
+   *     required to detect NAL units where the NAL unit prefix spans array boundaries. The array
+   *     must be at least 3 elements long.
    * @return The offset of the NAL unit, or {@code endOffset} if a NAL unit was not found.
    */
-  public static int findNalUnit(byte[] data, int startOffset, int endOffset, int type) {
+  public static int findNalUnit(byte[] data, int startOffset, int endOffset, int type,
+      boolean[] prefixFlags) {
+    int length = endOffset - startOffset;
+
+    Assertions.checkState(length >= 0);
+    if (length == 0) {
+      return endOffset;
+    }
+
+    if (prefixFlags != null) {
+      if (prefixFlags[0] && matchesType(data, startOffset, type)) {
+        clearPrefixFlags(prefixFlags);
+        return startOffset - 3;
+      } else if (length > 1 && prefixFlags[1] && data[startOffset] == 1
+          && matchesType(data, startOffset + 1, type)) {
+        clearPrefixFlags(prefixFlags);
+        return startOffset - 2;
+      } else if (length > 2 && prefixFlags[2] && data[startOffset] == 0
+          && data[startOffset + 1] == 1 && matchesType(data, startOffset + 2, type)) {
+        clearPrefixFlags(prefixFlags);
+        return startOffset - 1;
+      }
+    }
+
     int limit = endOffset - 2;
     // We're looking for the NAL unit start code prefix 0x000001, followed by a byte that matches
     // the specified type. The value of i tracks the index of the third byte in the four bytes
@@ -132,8 +169,8 @@ public final class Mp4Util {
       if ((data[i] & 0xFE) != 0) {
         // There isn't a NAL prefix here, or at the next two positions. Do nothing and let the
         // loop advance the index by three.
-      } else if ((data[i - 2] == 0 && data[i - 1] == 0 && data[i] == 1)
-          && (type == -1 || (type == (data[i + 1] & 0x1F)))) {
+      } else if (data[i - 2] == 0 && data[i - 1] == 0 && data[i] == 1
+          && matchesType(data, i + 1, type)) {
         return i - 2;
       } else {
         // There isn't a NAL prefix here, but there might be at the next position. We should
@@ -141,18 +178,77 @@ public final class Mp4Util {
         i -= 2;
       }
     }
+
+    if (prefixFlags != null) {
+      // True if the last three bytes in the data seen so far are {0,0,1}.
+      prefixFlags[0] = length > 2
+          ? (data[endOffset - 3] == 0 && data[endOffset - 2] == 0 && data[endOffset - 1] == 1)
+          : length == 2 ? (prefixFlags[2] && data[endOffset - 2] == 0 && data[endOffset - 1] == 1)
+          : (prefixFlags[1] && data[endOffset - 1] == 1);
+      // True if the last three bytes in the data seen so far are {0,0}.
+      prefixFlags[1] = length > 1 ? data[endOffset - 2] == 0 && data[endOffset - 1] == 0
+          : prefixFlags[2] && data[endOffset - 1] == 0;
+      // True if the last three bytes in the data seen so far are {0}.
+      prefixFlags[2] = data[endOffset - 1] == 0;
+    }
+
     return endOffset;
   }
 
   /**
+   * Like {@link #findNalUnit(byte[], int, int, int)} with {@code type == -1}.
+   *
+   * @param data The data to search.
+   * @param startOffset The offset (inclusive) in the data to start the search.
+   * @param endOffset The offset (exclusive) in the data to end the search.
+   * @return The offset of the NAL unit, or {@code endOffset} if a NAL unit was not found.
+   */
+  public static int findNalUnit(byte[] data, int startOffset, int endOffset) {
+    return findNalUnit(data, startOffset, endOffset, null);
+  }
+
+  /**
+   * Like {@link #findNalUnit(byte[], int, int, int, boolean[])} with {@code type == -1}.
+   *
+   * @param data The data to search.
+   * @param startOffset The offset (inclusive) in the data to start the search.
+   * @param endOffset The offset (exclusive) in the data to end the search.
+   * @param prefixFlags A boolean array of length at least 3.
+   * @return The offset of the NAL unit, or {@code endOffset} if a NAL unit was not found.
+   */
+  public static int findNalUnit(byte[] data, int startOffset, int endOffset,
+      boolean[] prefixFlags) {
+    return findNalUnit(data, startOffset, endOffset, -1, prefixFlags);
+  }
+
+  /**
    * Gets the type of the NAL unit in {@code data} that starts at {@code offset}.
    *
    * @param data The data to search.
-   * @param offset The start offset of a NAL unit.
+   * @param offset The start offset of a NAL unit. Must lie between {@code -3} (inclusive) and
+   *     {@code data.length - 3} (exclusive).
    * @return The type of the unit.
    */
   public static int getNalUnitType(byte[] data, int offset) {
     return data[offset + 3] & 0x1F;
   }
 
+  /**
+   * Clears prefix flags, as used by {@link #findNalUnit(byte[], int, int, int, boolean[])}.
+   *
+   * @param prefixFlags The flags to clear.
+   */
+  private static void clearPrefixFlags(boolean[] prefixFlags) {
+    prefixFlags[0] = false;
+    prefixFlags[1] = false;
+    prefixFlags[2] = false;
+  }
+
+  /**
+   * Returns true if the type at {@code offset} is equal to {@code type}, or if {@code type == -1}.
+   */
+  private static boolean matchesType(byte[] data, int offset, int type) {
+    return type == -1 || (data[offset] & 0x1F) == type;
+  }
+
 }