Bug 1409664 - P20. Add BitWriter class. r=bryce

Summary:

Depends on D1633

Tags: #secure-revision

Bug #: 1409664

Differential Revision: https://phabricator.services.mozilla.com/D1678

dom/media/BitReader.cpp

@@ -157,4 +157,50 @@ BitReader::FillReservoir()
   mReservoir <<= 32 - mNumBitsLeft;
 }
 
+/* static */ uint32_t
+BitReader::GetBitLength(const mozilla::MediaByteBuffer* aNAL)
+{
+  size_t size = aNAL->Length();
+
+  while (size > 0 && aNAL->ElementAt(size - 1) == 0) {
+    size--;
+  }
+
+  if (!size) {
+    return 0;
+  }
+
+  if (size > UINT32_MAX / 8) {
+    // We can't represent it, we'll use as much as we can.
+    return UINT32_MAX;
+  }
+
+  uint8_t v = aNAL->ElementAt(size - 1);
+  size *= 8;
+
+  // Remove the stop bit and following trailing zeros.
+  if (v) {
+    // Count the consecutive zero bits (trailing) on the right by binary search.
+    // Adapted from Matt Whitlock algorithm to only bother with 8 bits integers.
+    uint32_t c;
+    if (v & 1) {
+      // Special case for odd v (assumed to happen half of the time).
+      c = 0;
+    } else {
+      c = 1;
+      if ((v & 0xf) == 0) {
+        v >>= 4;
+        c += 4;
+      }
+      if ((v & 0x3) == 0) {
+        v >>= 2;
+        c += 2;
+      }
+      c -= v & 0x1;
+    }
+    size -= c + 1;
+  }
+  return size;
+}
+
 } // namespace mozilla

dom/media/BitReader.h

@@ -13,12 +13,12 @@ namespace mozilla
 class BitReader
 {
 public:
-  explicit BitReader(const mozilla::MediaByteBuffer* aBuffer);
-  BitReader(const mozilla::MediaByteBuffer* aBuffer, size_t aBits);
+  explicit BitReader(const MediaByteBuffer* aBuffer);
+  BitReader(const MediaByteBuffer* aBuffer, size_t aBits);
   BitReader(const uint8_t* aBuffer, size_t aBits);
   ~BitReader();
   uint32_t ReadBits(size_t aNum);
-  uint32_t ReadBit() { return ReadBits(1); }
+  bool ReadBit() { return ReadBits(1) != 0; }
   uint32_t ReadU32() { return ReadBits(32); }
   uint64_t ReadU64();
 
@@ -35,6 +35,9 @@ public:
   // Return the number of bits left.
   size_t BitsLeft() const;
 
+  // Return RBSP bit length.
+  static uint32_t GetBitLength(const MediaByteBuffer* aNAL);
+
 private:
   void FillReservoir();
   const uint8_t* mData;

dom/media/BitWriter.cpp

@@ -0,0 +1,96 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "BitWriter.h"
+#include "MediaData.h"
+#include "mozilla/MathAlgorithms.h"
+
+namespace mozilla
+{
+
+constexpr uint8_t golombLen[256] = {
+  1,  3,  3,  5,  5,  5,  5,  7,  7,  7,  7,  7,  7,  7,  7,  9,  9,  9,  9,
+  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  11, 11, 11, 11, 11, 11, 11,
+  11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
+  11, 11, 11, 11, 11, 11, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13,
+  13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13,
+  13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13,
+  13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
+  15, 15, 15, 15, 15, 15, 15, 15, 17,
+};
+
+BitWriter::BitWriter(MediaByteBuffer* aBuffer)
+  : mBuffer(aBuffer)
+{
+}
+
+BitWriter::~BitWriter() {}
+
+void
+BitWriter::WriteBits(uint64_t aValue, size_t aBits)
+{
+  MOZ_ASSERT(aBits <= sizeof(uint64_t) * 8);
+
+  while (aBits) {
+    if (mBitIndex == 0) {
+      mBuffer->AppendElement(0);
+    }
+
+    const uint8_t clearMask = ~(~0 << (8 - mBitIndex));
+    uint8_t mask = 0;
+
+    if (mBitIndex + aBits > 8) {
+      // Not enough bits in the current byte to write all the bits
+      // required, we'll process what we can and continue with the left over.
+      const uint8_t leftOverBits = mBitIndex + aBits - 8;
+      const uint64_t leftOver = aValue & (~uint64_t(0) >> (8 - mBitIndex));
+      mask = aValue >> leftOverBits;
+
+      mBitIndex = 8;
+      aValue = leftOver;
+      aBits = leftOverBits;
+    } else {
+      const uint8_t offset = 8 - mBitIndex - aBits;
+      mask = aValue << offset;
+
+      mBitIndex += aBits;
+      aBits = 0;
+    }
+
+    mBuffer->ElementAt(mPosition) |= mask & clearMask;
+
+    if (mBitIndex == 8) {
+      mPosition++;
+      mBitIndex = 0;
+    }
+  }
+}
+
+void
+BitWriter::WriteUE(uint32_t aValue)
+{
+  MOZ_ASSERT(aValue <= (UINT32_MAX - 1));
+
+  if (aValue < 256) {
+    WriteBits(aValue + 1, golombLen[aValue]);
+  } else {
+    const uint32_t e = FloorLog2(aValue + 1);
+    WriteBits(aValue + 1, e * 2 + 1);
+  }
+}
+
+void
+BitWriter::CloseWithRbspTrailing()
+{
+  WriteBit(true);
+  WriteBits(0, (8 - mBitIndex) & 7);
+}
+
+} // namespace mozilla

dom/media/BitWriter.h

@@ -0,0 +1,43 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef BIT_WRITER_H_
+#define BIT_WRITER_H_
+
+#include "mozilla/RefPtr.h"
+
+namespace mozilla
+{
+
+class MediaByteBuffer;
+
+class BitWriter
+{
+public:
+  explicit BitWriter(MediaByteBuffer* aBuffer);
+  virtual ~BitWriter();
+  void WriteBits(uint64_t aValue, size_t aBits);
+  void WriteBit(bool aValue) { WriteBits(aValue, 1); }
+  void WriteU8(uint8_t aValue) { WriteBits(aValue, 8); }
+  void WriteU32(uint32_t aValue) { WriteBits(aValue, 32); }
+  void WriteU64(uint64_t aValue) { WriteBits(aValue, 64); }
+
+  // Write unsigned integer into Exp-Golomb-coded. 2^32-2 at most
+  void WriteUE(uint32_t aValue);
+
+  // Write RBSP trailing bits.
+  void CloseWithRbspTrailing();
+
+  // Return the number of bits written so far;
+  size_t BitCount() const { return mPosition * 8 + mBitIndex; }
+
+private:
+  RefPtr<MediaByteBuffer> mBuffer;
+  size_t mPosition = 0;
+  uint8_t mBitIndex = 0;
+};
+
+} // namespace mozilla
+
+#endif // BIT_WRITER_H_

dom/media/gtest/TestBitWriter.cpp

@@ -0,0 +1,50 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "gtest/gtest.h"
+#include "BitReader.h"
+#include "BitWriter.h"
+
+using namespace mozilla;
+
+TEST(BitWriter, BitWriter)
+{
+  RefPtr<MediaByteBuffer> test = new MediaByteBuffer();
+  BitWriter b(test);
+  b.WriteBit(false);
+  b.WriteBits(~1ULL, 1); // ensure that extra bits don't modify byte buffer.
+  b.WriteBits(3, 1);
+  b.WriteUE(1280 / 16 - 1);
+  b.WriteUE(720 / 16 - 1);
+  b.WriteUE(1280);
+  b.WriteUE(720);
+  b.WriteBit(true);
+  b.WriteBit(false);
+  b.WriteBit(true);
+  b.WriteU8(7);
+  b.WriteU32(16356);
+  b.WriteU64(116356);
+  b.WriteBits(~(0ULL) & ~1ULL, 16);
+  const uint32_t length = b.BitCount();
+  b.CloseWithRbspTrailing();
+
+  BitReader c(test);
+
+  EXPECT_EQ(c.ReadBit(), false);
+  EXPECT_EQ(c.ReadBit(), false);
+  EXPECT_EQ(c.ReadBit(), true);
+  EXPECT_EQ(c.ReadUE(), 1280u / 16 - 1);
+  EXPECT_EQ(c.ReadUE(), 720u / 16 - 1);
+  EXPECT_EQ(c.ReadUE(), 1280u);
+  EXPECT_EQ(c.ReadUE(), 720u);
+  EXPECT_EQ(c.ReadBit(), true);
+  EXPECT_EQ(c.ReadBit(), false);
+  EXPECT_EQ(c.ReadBit(), true);
+  EXPECT_EQ(c.ReadBits(8), 7u);
+  EXPECT_EQ(c.ReadU32(), 16356u);
+  EXPECT_EQ(c.ReadU64(), 116356u);
+  EXPECT_EQ(c.ReadBits(16), 0xfffeu);
+  EXPECT_EQ(length, BitReader::GetBitLength(test));
+}

dom/media/gtest/moz.build

@@ -12,6 +12,7 @@ UNIFIED_SOURCES += [
     'TestAudioPacketizer.cpp',
     'TestAudioSegment.cpp',
     'TestAudioTrackEncoder.cpp',
+    'TestBitWriter.cpp',
     'TestBlankVideoDataCreator.cpp',
     'TestCDMStorage.cpp',
     'TestDataMutex.cpp',

dom/media/moz.build

@@ -105,6 +105,7 @@ EXPORTS += [
     'BackgroundVideoDecodingPermissionObserver.h',
     'Benchmark.h',
     'BitReader.h',
+    'BitWriter.h',
     'BufferMediaResource.h',
     'BufferReader.h',
     'ByteWriter.h',
@@ -224,6 +225,7 @@ UNIFIED_SOURCES += [
     'BaseMediaResource.cpp',
     'Benchmark.cpp',
     'BitReader.cpp',
+    'BitWriter.cpp',
     'CanvasCaptureMediaStream.cpp',
     'ChannelMediaDecoder.cpp',
     'ChannelMediaResource.cpp',

dom/media/platforms/agnostic/bytestreams/H264.cpp

@@ -120,52 +120,6 @@ scaling_list(BitReader& aBr, uint8_t (&aScalingList)[N], const uint8_t (&aDefaul
   detail::scaling_list(aBr, aScalingList, N, aDefaultList, nullptr);
 }
 
-static uint32_t
-GetBitLength(const mozilla::MediaByteBuffer* aNAL)
-{
-  size_t size = aNAL->Length();
-
-  while (size > 0 && aNAL->ElementAt(size - 1) == 0) {
-    size--;
-  }
-
-  if (!size) {
-    return 0;
-  }
-
-  if (size > UINT32_MAX / 8) {
-    // We can't represent it, we'll use as much as we can.
-    return UINT32_MAX;
-  }
-
-  uint8_t v = aNAL->ElementAt(size - 1);
-  size *= 8;
-
-  // Remove the stop bit and following trailing zeros.
-  if (v) {
-    // Count the consecutive zero bits (trailing) on the right by binary search.
-    // Adapted from Matt Whitlock algorithm to only bother with 8 bits integers.
-    uint32_t c;
-    if (v & 1) {
-      // Special case for odd v (assumed to happen half of the time).
-      c = 0;
-    } else {
-      c = 1;
-      if ((v & 0xf) == 0) {
-        v >>= 4;
-        c += 4;
-      }
-      if ((v & 0x3) == 0) {
-        v >>= 2;
-        c += 2;
-      }
-      c -= v & 0x1;
-    }
-    size -= c + 1;
-  }
-  return size;
-}
-
 SPSData::SPSData()
 {
   PodZero(this);
@@ -205,7 +159,7 @@ public:
     }
     mDecodedNAL = H264::DecodeNALUnit(aPtr, aLength);
     if (mDecodedNAL) {
-      mLength = GetBitLength(mDecodedNAL);
+      mLength = BitReader::GetBitLength(mDecodedNAL);
     }
   }
 
@@ -405,7 +359,7 @@ H264::DecodeSPS(const mozilla::MediaByteBuffer* aSPS, SPSData& aDest)
   if (!aSPS) {
     return false;
   }
-  BitReader br(aSPS, GetBitLength(aSPS));
+  BitReader br(aSPS, BitReader::GetBitLength(aSPS));
 
   aDest.profile_idc = br.ReadBits(8);
   aDest.constraint_set0_flag = br.ReadBit();