Add support for 3-byte headers in VideoToolbox NALU parser.

webrtc/modules/video_coding/codecs/h264/h264_video_toolbox_nalu.cc

 /*
  *  Copyright (c) 2015 The WebRTC project authors. All Rights Reserved.
  *
  *  Use of this source code is governed by a BSD-style license
  *  that can be found in the LICENSE file in the root of the source
  *  tree. An additional intellectual property rights grant can be found
  *  in the file PATENTS.  All contributing project authors may
  *  be found in the AUTHORS file in the root of the source tree.
  *
  */
 
 #include "webrtc/modules/video_coding/codecs/h264/h264_video_toolbox_nalu.h"
 
 #if defined(WEBRTC_VIDEO_TOOLBOX_SUPPORTED)
 
 #include <CoreFoundation/CoreFoundation.h>
 #include <memory>
 #include <vector>
 
 #include "webrtc/base/checks.h"
 #include "webrtc/base/logging.h"
-#include "webrtc/common_video/h264/h264_common.h"
 
 namespace webrtc {
 
-using H264::NaluType;
 using H264::kAud;
 using H264::kSps;
+using H264::NaluIndex;
+using H264::NaluType;
 using H264::ParseNaluType;
 
 const char kAnnexBHeaderBytes[4] = {0, 0, 0, 1};
 const size_t kAvccHeaderByteSize = sizeof(uint32_t);
 
 bool H264CMSampleBufferToAnnexBBuffer(
     CMSampleBufferRef avcc_sample_buffer,
     bool is_keyframe,
     rtc::Buffer* annexb_buffer,
     webrtc::RTPFragmentationHeader** out_header) {
@@ skipping 11 matching lines @@
 
   // Get parameter set information.
   int nalu_header_size = 0;
   size_t param_set_count = 0;
   OSStatus status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(
       description, 0, nullptr, nullptr, &param_set_count, &nalu_header_size);
   if (status != noErr) {
     LOG(LS_ERROR) << "Failed to get parameter set.";
     return false;
   }
-  // TODO(tkchin): handle other potential sizes.
-  RTC_DCHECK_EQ(nalu_header_size, 4);
+  RTC_CHECK_EQ(nalu_header_size, kAvccHeaderByteSize);
   RTC_DCHECK_EQ(param_set_count, 2u);
 
   // Truncate any previous data in the buffer without changing its capacity.
   annexb_buffer->SetSize(0);
 
   size_t nalu_offset = 0;
   std::vector<size_t> frag_offsets;
   std::vector<size_t> frag_lengths;
 
   // Place all parameter sets at the front of buffer.
@@ skipping 61 matching lines @@
     uint32_t* uint32_data_ptr = reinterpret_cast<uint32_t*>(data_ptr);
     uint32_t packet_size = CFSwapInt32BigToHost(*uint32_data_ptr);
     // Update buffer.
     annexb_buffer->AppendData(kAnnexBHeaderBytes, sizeof(kAnnexBHeaderBytes));
     annexb_buffer->AppendData(data_ptr + nalu_header_size, packet_size);
     // Update fragmentation.
     frag_offsets.push_back(nalu_offset + sizeof(kAnnexBHeaderBytes));
     frag_lengths.push_back(packet_size);
     nalu_offset += sizeof(kAnnexBHeaderBytes) + packet_size;
 
-    size_t bytes_written = packet_size + nalu_header_size;
+    size_t bytes_written = packet_size + sizeof(kAnnexBHeaderBytes);
     bytes_remaining -= bytes_written;
     data_ptr += bytes_written;
   }
   RTC_DCHECK_EQ(bytes_remaining, (size_t)0);
 
   std::unique_ptr<webrtc::RTPFragmentationHeader> header;
   header.reset(new webrtc::RTPFragmentationHeader());
   header->VerifyAndAllocateFragmentationHeader(frag_offsets.size());
   RTC_DCHECK_EQ(frag_lengths.size(), frag_offsets.size());
   for (size_t i = 0; i < frag_offsets.size(); ++i) {
@@ skipping 151 matching lines @@
       &description);
   if (status != noErr) {
     LOG(LS_ERROR) << "Failed to create video format description.";
     return nullptr;
   }
   return description;
 }
 
 AnnexBBufferReader::AnnexBBufferReader(const uint8_t* annexb_buffer,
                                        size_t length)
-    : start_(annexb_buffer), offset_(0), next_offset_(0), length_(length) {
+    : start_(annexb_buffer), length_(length) {
   RTC_DCHECK(annexb_buffer);
-  offset_ = FindNextNaluHeader(start_, length_, 0);
-  next_offset_ =
-      FindNextNaluHeader(start_, length_, offset_ + sizeof(kAnnexBHeaderBytes));
+  offsets_ = H264::FindNaluIndices(annexb_buffer, length);
+  offset_ = offsets_.begin();
 }
 
 bool AnnexBBufferReader::ReadNalu(const uint8_t** out_nalu,
                                   size_t* out_length) {
   RTC_DCHECK(out_nalu);
   RTC_DCHECK(out_length);
   *out_nalu = nullptr;
   *out_length = 0;
 
-  size_t data_offset = offset_ + sizeof(kAnnexBHeaderBytes);
-  if (data_offset > length_) {
+  if (offset_ == offsets_.end()) {
     return false;
   }
-  *out_nalu = start_ + data_offset;
-  *out_length = next_offset_ - data_offset;
-  offset_ = next_offset_;
-  next_offset_ =
-      FindNextNaluHeader(start_, length_, offset_ + sizeof(kAnnexBHeaderBytes));
+  *out_nalu = start_ + offset_->payload_start_offset;
+  *out_length = offset_->payload_size;
+  ++offset_;
   return true;
 }
 
 size_t AnnexBBufferReader::BytesRemaining() const {
-  return length_ - offset_;
-}
-
-size_t AnnexBBufferReader::FindNextNaluHeader(const uint8_t* start,
-                                              size_t length,
-                                              size_t offset) const {
-  RTC_DCHECK(start);
-  if (offset + sizeof(kAnnexBHeaderBytes) > length) {
-    return length;
+  if (offset_ == offsets_.end()) {
+    return 0;
   }
-  // NALUs are separated by an 00 00 00 01 header. Scan the byte stream
-  // starting from the offset for the next such sequence.
-  const uint8_t* current = start + offset;
-  // The loop reads sizeof(kAnnexBHeaderBytes) at a time, so stop when there
-  // aren't enough bytes remaining.
-  const uint8_t* const end = start + length - sizeof(kAnnexBHeaderBytes);
-  while (current < end) {
-    if (current[3] > 1) {
-      current += 4;
-    } else if (current[3] == 1 && current[2] == 0 && current[1] == 0 &&
-               current[0] == 0) {
-      return current - start;
-    } else {
-      ++current;
-    }
-  }
-  return length;
+  return length_ - offset_->start_offset;
 }
 
 AvccBufferWriter::AvccBufferWriter(uint8_t* const avcc_buffer, size_t length)
     : start_(avcc_buffer), offset_(0), length_(length) {
   RTC_DCHECK(avcc_buffer);
 }
 
 bool AvccBufferWriter::WriteNalu(const uint8_t* data, size_t data_size) {
   // Check if we can write this length of data.
   if (data_size + kAvccHeaderByteSize > BytesRemaining()) {
     return false;
   }
   // Write length header, which needs to be big endian.
   uint32_t big_endian_length = CFSwapInt32HostToBig(data_size);
   memcpy(start_ + offset_, &big_endian_length, sizeof(big_endian_length));
   offset_ += sizeof(big_endian_length);
   // Write data.
   memcpy(start_ + offset_, data, data_size);
   offset_ += data_size;
   return true;
 }
 
 size_t AvccBufferWriter::BytesRemaining() const {
   return length_ - offset_;
 }
 
 }  // namespace webrtc
 
 #endif  // defined(WEBRTC_VIDEO_TOOLBOX_SUPPORTED)