Implement BitMemory{Reader,Writer}

Two simple classes which replace the need to pass
the (BitMemoryRegion, bit_offset) tuple everywhere.

The slightly simplifies the code and it also makes
it possible to optimize those classes in the future.

Test: test-art-host-gtest-stack_map_test
Test: test-art-host-gtest-bit_table_test
Change-Id: I4806c805149a07e1a11b76405ca27960a0012c69
diff --git a/compiler/optimizing/stack_map_stream.cc b/compiler/optimizing/stack_map_stream.cc
index e8b3330..57f47af 100644
--- a/compiler/optimizing/stack_map_stream.cc
+++ b/compiler/optimizing/stack_map_stream.cc
@@ -301,16 +301,16 @@
     }
   }
 
-  size_t bit_offset = 0;
-  stack_maps_.Encode(&out_, &bit_offset);
-  register_masks_.Encode(&out_, &bit_offset);
-  stack_masks_.Encode(&out_, &bit_offset);
-  invoke_infos_.Encode(&out_, &bit_offset);
-  inline_infos_.Encode(&out_, &bit_offset);
-  dex_register_masks_.Encode(&out_, &bit_offset);
-  dex_register_maps_.Encode(&out_, &bit_offset);
-  dex_register_catalog_.Encode(&out_, &bit_offset);
-  EncodeVarintBits(&out_, &bit_offset, num_dex_registers_);
+  BitMemoryWriter<ScopedArenaVector<uint8_t>> out(&out_);
+  stack_maps_.Encode(out);
+  register_masks_.Encode(out);
+  stack_masks_.Encode(out);
+  invoke_infos_.Encode(out);
+  inline_infos_.Encode(out);
+  dex_register_masks_.Encode(out);
+  dex_register_maps_.Encode(out);
+  dex_register_catalog_.Encode(out);
+  EncodeVarintBits(out, num_dex_registers_);
 
   return UnsignedLeb128Size(out_.size()) +  out_.size();
 }
diff --git a/libartbase/base/bit_memory_region.h b/libartbase/base/bit_memory_region.h
index b4764fd..6e491b0 100644
--- a/libartbase/base/bit_memory_region.h
+++ b/libartbase/base/bit_memory_region.h
@@ -99,13 +99,6 @@
     return value & mask;
   }
 
-  // Load bits starting at given `bit_offset`, and advance the `bit_offset`.
-  ALWAYS_INLINE uint32_t LoadBitsAndAdvance(size_t* bit_offset, size_t bit_length) const {
-    uint32_t result = LoadBits(*bit_offset, bit_length);
-    *bit_offset += bit_length;
-    return result;
-  }
-
   // Store `bit_length` bits in `data` starting at given `bit_offset`.
   // The least significant bit is stored in the smallest memory offset.
   ALWAYS_INLINE void StoreBits(size_t bit_offset, uint32_t value, size_t bit_length) {
@@ -132,12 +125,6 @@
     DCHECK_EQ(value, LoadBits(bit_offset, bit_length));
   }
 
-  // Store bits starting at given `bit_offset`, and advance the `bit_offset`.
-  ALWAYS_INLINE void StoreBitsAndAdvance(size_t* bit_offset, uint32_t value, size_t bit_length) {
-    StoreBits(*bit_offset, value, bit_length);
-    *bit_offset += bit_length;
-  }
-
   // Store bits from other bit region.
   ALWAYS_INLINE void StoreBits(size_t bit_offset, const BitMemoryRegion& src, size_t bit_length) {
     DCHECK_LE(bit_offset, bit_size_);
@@ -178,6 +165,62 @@
   size_t bit_size_ = 0;
 };
 
+class BitMemoryReader {
+ public:
+  explicit BitMemoryReader(BitMemoryRegion region, size_t bit_offset = 0)
+      : region_(region), bit_offset_(bit_offset) { }
+
+  size_t GetBitOffset() const { return bit_offset_; }
+
+  ALWAYS_INLINE BitMemoryRegion Skip(size_t bit_length) {
+    BitMemoryRegion result = region_.Subregion(bit_offset_, bit_length);
+    bit_offset_ += bit_length;
+    return result;
+  }
+
+  ALWAYS_INLINE uint32_t ReadBits(size_t bit_length) {
+    uint32_t result = region_.LoadBits(bit_offset_, bit_length);
+    bit_offset_ += bit_length;
+    return result;
+  }
+
+ private:
+  BitMemoryRegion region_;
+  size_t bit_offset_;
+
+  DISALLOW_COPY_AND_ASSIGN(BitMemoryReader);
+};
+
+template<typename Vector>
+class BitMemoryWriter {
+ public:
+  explicit BitMemoryWriter(Vector* out, size_t bit_offset = 0)
+      : out_(out), bit_offset_(bit_offset) { }
+
+  size_t GetBitOffset() const { return bit_offset_; }
+
+  ALWAYS_INLINE BitMemoryRegion Allocate(size_t bit_length) {
+    out_->resize(BitsToBytesRoundUp(bit_offset_ + bit_length));
+    BitMemoryRegion region(MemoryRegion(out_->data(), out_->size()), bit_offset_, bit_length);
+    bit_offset_ += bit_length;
+    return region;
+  }
+
+  ALWAYS_INLINE void WriteBits(uint32_t value, size_t bit_length) {
+    Allocate(bit_length).StoreBits(0, value, bit_length);
+  }
+
+  BitMemoryRegion GetWrittenRegion() const {
+    return BitMemoryRegion(MemoryRegion(out_->data(), out_->size()), 0, bit_offset_);
+  }
+
+ private:
+  Vector* out_;
+  size_t bit_offset_;
+
+  DISALLOW_COPY_AND_ASSIGN(BitMemoryWriter);
+};
+
 }  // namespace art
 
 #endif  // ART_LIBARTBASE_BASE_BIT_MEMORY_REGION_H_
diff --git a/libartbase/base/bit_table.h b/libartbase/base/bit_table.h
index 8036db1..c1619b5 100644
--- a/libartbase/base/bit_table.h
+++ b/libartbase/base/bit_table.h
@@ -39,28 +39,24 @@
 // The first four bits determine the variable length of the encoded integer:
 //   Values 0..11 represent the result as-is, with no further following bits.
 //   Values 12..15 mean the result is in the next 8/16/24/32-bits respectively.
-ALWAYS_INLINE static inline uint32_t DecodeVarintBits(BitMemoryRegion region, size_t* bit_offset) {
-  uint32_t x = region.LoadBitsAndAdvance(bit_offset, kVarintHeaderBits);
+ALWAYS_INLINE static inline uint32_t DecodeVarintBits(BitMemoryReader& reader) {
+  uint32_t x = reader.ReadBits(kVarintHeaderBits);
   if (x > kVarintSmallValue) {
-    x = region.LoadBitsAndAdvance(bit_offset, (x - kVarintSmallValue) * kBitsPerByte);
+    x = reader.ReadBits((x - kVarintSmallValue) * kBitsPerByte);
   }
   return x;
 }
 
 // Store variable-length bit-packed integer from `data` starting at `bit_offset`.
 template<typename Vector>
-ALWAYS_INLINE static inline void EncodeVarintBits(Vector* out, size_t* bit_offset, uint32_t value) {
+ALWAYS_INLINE static inline void EncodeVarintBits(BitMemoryWriter<Vector>& out, uint32_t value) {
   if (value <= kVarintSmallValue) {
-    out->resize(BitsToBytesRoundUp(*bit_offset + kVarintHeaderBits));
-    BitMemoryRegion region(MemoryRegion(out->data(), out->size()));
-    region.StoreBitsAndAdvance(bit_offset, value, kVarintHeaderBits);
+    out.WriteBits(value, kVarintHeaderBits);
   } else {
     uint32_t num_bits = RoundUp(MinimumBitsToStore(value), kBitsPerByte);
-    out->resize(BitsToBytesRoundUp(*bit_offset + kVarintHeaderBits + num_bits));
-    BitMemoryRegion region(MemoryRegion(out->data(), out->size()));
     uint32_t header = kVarintSmallValue + num_bits / kBitsPerByte;
-    region.StoreBitsAndAdvance(bit_offset, header, kVarintHeaderBits);
-    region.StoreBitsAndAdvance(bit_offset, value, num_bits);
+    out.WriteBits(header, kVarintHeaderBits);
+    out.WriteBits(value, num_bits);
   }
 }
 
@@ -74,26 +70,25 @@
   static constexpr uint32_t kValueBias = kNoValue;  // Bias so that -1 is encoded as 0.
 
   BitTableBase() {}
-  BitTableBase(void* data, size_t size, size_t* bit_offset) {
-    Decode(BitMemoryRegion(MemoryRegion(data, size)), bit_offset);
+  explicit BitTableBase(BitMemoryReader& reader) {
+    Decode(reader);
   }
 
-  ALWAYS_INLINE void Decode(BitMemoryRegion region, size_t* bit_offset) {
+  ALWAYS_INLINE void Decode(BitMemoryReader& reader) {
     // Decode row count and column sizes from the table header.
-    size_t initial_bit_offset = *bit_offset;
-    num_rows_ = DecodeVarintBits(region, bit_offset);
+    size_t initial_bit_offset = reader.GetBitOffset();
+    num_rows_ = DecodeVarintBits(reader);
     if (num_rows_ != 0) {
       column_offset_[0] = 0;
       for (uint32_t i = 0; i < kNumColumns; i++) {
-        size_t column_end = column_offset_[i] + DecodeVarintBits(region, bit_offset);
+        size_t column_end = column_offset_[i] + DecodeVarintBits(reader);
         column_offset_[i + 1] = dchecked_integral_cast<uint16_t>(column_end);
       }
     }
-    header_bit_size_ = *bit_offset - initial_bit_offset;
+    header_bit_size_ = reader.GetBitOffset() - initial_bit_offset;
 
     // Record the region which contains the table data and skip past it.
-    table_data_ = region.Subregion(*bit_offset, num_rows_ * NumRowBits());
-    *bit_offset += table_data_.size_in_bits();
+    table_data_ = reader.Skip(num_rows_ * NumRowBits());
   }
 
   ALWAYS_INLINE uint32_t Get(uint32_t row, uint32_t column = 0) const {
@@ -333,25 +328,22 @@
 
   // Encode the stored data into a BitTable.
   template<typename Vector>
-  void Encode(Vector* out, size_t* bit_offset) const {
-    size_t initial_bit_offset = *bit_offset;
+  void Encode(BitMemoryWriter<Vector>& out) const {
+    size_t initial_bit_offset = out.GetBitOffset();
 
     std::array<uint32_t, kNumColumns> column_bits;
     Measure(&column_bits);
-    EncodeVarintBits(out, bit_offset, size());
+    EncodeVarintBits(out, size());
     if (size() != 0) {
       // Write table header.
       for (uint32_t c = 0; c < kNumColumns; c++) {
-        EncodeVarintBits(out, bit_offset, column_bits[c]);
+        EncodeVarintBits(out, column_bits[c]);
       }
 
       // Write table data.
-      uint32_t row_bits = std::accumulate(column_bits.begin(), column_bits.end(), 0u);
-      out->resize(BitsToBytesRoundUp(*bit_offset + row_bits * size()));
-      BitMemoryRegion region(MemoryRegion(out->data(), out->size()));
       for (uint32_t r = 0; r < size(); r++) {
         for (uint32_t c = 0; c < kNumColumns; c++) {
-          region.StoreBitsAndAdvance(bit_offset, rows_[r][c] - kValueBias, column_bits[c]);
+          out.WriteBits(rows_[r][c] - kValueBias, column_bits[c]);
         }
       }
     }
@@ -359,8 +351,8 @@
     // Verify the written data.
     if (kIsDebugBuild) {
       BitTableBase<kNumColumns> table;
-      BitMemoryRegion region(MemoryRegion(out->data(), out->size()));
-      table.Decode(region, &initial_bit_offset);
+      BitMemoryReader reader(out.GetWrittenRegion(), initial_bit_offset);
+      table.Decode(reader);
       DCHECK_EQ(size(), table.NumRows());
       for (uint32_t c = 0; c < kNumColumns; c++) {
         DCHECK_EQ(column_bits[c], table.NumColumnBits(c));
@@ -427,28 +419,26 @@
 
   // Encode the stored data into a BitTable.
   template<typename Vector>
-  void Encode(Vector* out, size_t* bit_offset) const {
-    size_t initial_bit_offset = *bit_offset;
+  void Encode(BitMemoryWriter<Vector>& out) const {
+    size_t initial_bit_offset = out.GetBitOffset();
 
-    EncodeVarintBits(out, bit_offset, size());
+    EncodeVarintBits(out, size());
     if (size() != 0) {
-      EncodeVarintBits(out, bit_offset, max_num_bits_);
+      EncodeVarintBits(out, max_num_bits_);
 
       // Write table data.
-      out->resize(BitsToBytesRoundUp(*bit_offset + max_num_bits_ * size()));
-      BitMemoryRegion region(MemoryRegion(out->data(), out->size()));
       for (MemoryRegion row : rows_) {
         BitMemoryRegion src(row);
-        region.StoreBits(*bit_offset, src, std::min(max_num_bits_, src.size_in_bits()));
-        *bit_offset += max_num_bits_;
+        BitMemoryRegion dst = out.Allocate(max_num_bits_);
+        dst.StoreBits(/* bit_offset */ 0, src, std::min(max_num_bits_, src.size_in_bits()));
       }
     }
 
     // Verify the written data.
     if (kIsDebugBuild) {
       BitTableBase<1> table;
-      BitMemoryRegion region(MemoryRegion(out->data(), out->size()));
-      table.Decode(region, &initial_bit_offset);
+      BitMemoryReader reader(out.GetWrittenRegion(), initial_bit_offset);
+      table.Decode(reader);
       DCHECK_EQ(size(), table.NumRows());
       DCHECK_EQ(max_num_bits_, table.NumColumnBits(0));
       for (uint32_t r = 0; r < size(); r++) {
diff --git a/libartbase/base/bit_table_test.cc b/libartbase/base/bit_table_test.cc
index ee7cb3a..a694148 100644
--- a/libartbase/base/bit_table_test.cc
+++ b/libartbase/base/bit_table_test.cc
@@ -31,13 +31,12 @@
     uint32_t values[] = { 0, 1, 11, 12, 15, 16, 255, 256, ~1u, ~0u };
     for (uint32_t value : values) {
       std::vector<uint8_t> buffer;
-      size_t encode_bit_offset = start_bit_offset;
-      EncodeVarintBits(&buffer, &encode_bit_offset, value);
+      BitMemoryWriter<std::vector<uint8_t>> writer(&buffer, start_bit_offset);
+      EncodeVarintBits(writer, value);
 
-      size_t decode_bit_offset = start_bit_offset;
-      BitMemoryRegion region(MemoryRegion(buffer.data(), buffer.size()));
-      uint32_t result = DecodeVarintBits(region, &decode_bit_offset);
-      EXPECT_EQ(encode_bit_offset, decode_bit_offset);
+      BitMemoryReader reader(writer.GetWrittenRegion(), start_bit_offset);
+      uint32_t result = DecodeVarintBits(reader);
+      EXPECT_EQ(writer.GetBitOffset(), reader.GetBitOffset());
       EXPECT_EQ(value, result);
     }
   }
@@ -49,13 +48,13 @@
   ScopedArenaAllocator allocator(&arena_stack);
 
   std::vector<uint8_t> buffer;
-  size_t encode_bit_offset = 0;
+  BitMemoryWriter<std::vector<uint8_t>> writer(&buffer);
   BitTableBuilderBase<1> builder(&allocator);
-  builder.Encode(&buffer, &encode_bit_offset);
+  builder.Encode(writer);
 
-  size_t decode_bit_offset = 0;
-  BitTableBase<1> table(buffer.data(), buffer.size(), &decode_bit_offset);
-  EXPECT_EQ(encode_bit_offset, decode_bit_offset);
+  BitMemoryReader reader(writer.GetWrittenRegion());
+  BitTableBase<1> table(reader);
+  EXPECT_EQ(writer.GetBitOffset(), reader.GetBitOffset());
   EXPECT_EQ(0u, table.NumRows());
 }
 
@@ -66,17 +65,17 @@
 
   constexpr uint32_t kNoValue = -1;
   std::vector<uint8_t> buffer;
-  size_t encode_bit_offset = 0;
+  BitMemoryWriter<std::vector<uint8_t>> writer(&buffer);
   BitTableBuilderBase<1> builder(&allocator);
   builder.Add({42u});
   builder.Add({kNoValue});
   builder.Add({1000u});
   builder.Add({kNoValue});
-  builder.Encode(&buffer, &encode_bit_offset);
+  builder.Encode(writer);
 
-  size_t decode_bit_offset = 0;
-  BitTableBase<1> table(buffer.data(), buffer.size(), &decode_bit_offset);
-  EXPECT_EQ(encode_bit_offset, decode_bit_offset);
+  BitMemoryReader reader(writer.GetWrittenRegion());
+  BitTableBase<1> table(reader);
+  EXPECT_EQ(writer.GetBitOffset(), reader.GetBitOffset());
   EXPECT_EQ(4u, table.NumRows());
   EXPECT_EQ(42u, table.Get(0));
   EXPECT_EQ(kNoValue, table.Get(1));
@@ -92,14 +91,14 @@
 
   for (size_t start_bit_offset = 0; start_bit_offset <= 32; start_bit_offset++) {
     std::vector<uint8_t> buffer;
-    size_t encode_bit_offset = start_bit_offset;
+    BitMemoryWriter<std::vector<uint8_t>> writer(&buffer, start_bit_offset);
     BitTableBuilderBase<1> builder(&allocator);
     builder.Add({42u});
-    builder.Encode(&buffer, &encode_bit_offset);
+    builder.Encode(writer);
 
-    size_t decode_bit_offset = start_bit_offset;
-    BitTableBase<1> table(buffer.data(), buffer.size(), &decode_bit_offset);
-    EXPECT_EQ(encode_bit_offset, decode_bit_offset) << " start_bit_offset=" << start_bit_offset;
+    BitMemoryReader reader(writer.GetWrittenRegion(), start_bit_offset);
+    BitTableBase<1> table(reader);
+    EXPECT_EQ(writer.GetBitOffset(), reader.GetBitOffset());
     EXPECT_EQ(1u, table.NumRows());
     EXPECT_EQ(42u, table.Get(0));
   }
@@ -112,15 +111,15 @@
 
   constexpr uint32_t kNoValue = -1;
   std::vector<uint8_t> buffer;
-  size_t encode_bit_offset = 0;
+  BitMemoryWriter<std::vector<uint8_t>> writer(&buffer);
   BitTableBuilderBase<4> builder(&allocator);
   builder.Add({42u, kNoValue, 0u, static_cast<uint32_t>(-2)});
   builder.Add({62u, kNoValue, 63u, static_cast<uint32_t>(-3)});
-  builder.Encode(&buffer, &encode_bit_offset);
+  builder.Encode(writer);
 
-  size_t decode_bit_offset = 0;
-  BitTableBase<4> table(buffer.data(), buffer.size(), &decode_bit_offset);
-  EXPECT_EQ(encode_bit_offset, decode_bit_offset);
+  BitMemoryReader reader(writer.GetWrittenRegion());
+  BitTableBase<4> table(reader);
+  EXPECT_EQ(writer.GetBitOffset(), reader.GetBitOffset());
   EXPECT_EQ(2u, table.NumRows());
   EXPECT_EQ(42u, table.Get(0, 0));
   EXPECT_EQ(kNoValue, table.Get(0, 1));
@@ -157,7 +156,7 @@
   ScopedArenaAllocator allocator(&arena_stack);
 
   std::vector<uint8_t> buffer;
-  size_t encode_bit_offset = 0;
+  BitMemoryWriter<std::vector<uint8_t>> writer(&buffer);
   const uint64_t value = 0xDEADBEEF0BADF00Dull;
   BitmapTableBuilder builder(&allocator);
   std::multimap<uint64_t, size_t> indicies;  // bitmap -> row.
@@ -165,12 +164,12 @@
     uint64_t bitmap = value & MaxInt<uint64_t>(bit_length);
     indicies.emplace(bitmap, builder.Dedup(&bitmap, MinimumBitsToStore(bitmap)));
   }
-  builder.Encode(&buffer, &encode_bit_offset);
+  builder.Encode(writer);
   EXPECT_EQ(1 + static_cast<uint32_t>(POPCOUNT(value)), builder.size());
 
-  size_t decode_bit_offset = 0;
-  BitTableBase<1> table(buffer.data(), buffer.size(), &decode_bit_offset);
-  EXPECT_EQ(encode_bit_offset, decode_bit_offset);
+  BitMemoryReader reader(writer.GetWrittenRegion());
+  BitTableBase<1> table(reader);
+  EXPECT_EQ(writer.GetBitOffset(), reader.GetBitOffset());
   for (auto it : indicies) {
     uint64_t expected = it.first;
     BitMemoryRegion actual = table.GetBitMemoryRegion(it.second);
diff --git a/runtime/stack_map.cc b/runtime/stack_map.cc
index f40168b..f9e2d27 100644
--- a/runtime/stack_map.cc
+++ b/runtime/stack_map.cc
@@ -26,6 +26,23 @@
 
 namespace art {
 
+void CodeInfo::Decode(const uint8_t* data) {
+  size_t non_header_size = DecodeUnsignedLeb128(&data);
+  size_ = UnsignedLeb128Size(non_header_size) + non_header_size;
+  MemoryRegion region(const_cast<uint8_t*>(data), non_header_size);
+  BitMemoryReader reader(BitMemoryRegion(region), /* bit_offset */ 0);
+  stack_maps_.Decode(reader);
+  register_masks_.Decode(reader);
+  stack_masks_.Decode(reader);
+  invoke_infos_.Decode(reader);
+  inline_infos_.Decode(reader);
+  dex_register_masks_.Decode(reader);
+  dex_register_maps_.Decode(reader);
+  dex_register_catalog_.Decode(reader);
+  number_of_dex_registers_ = DecodeVarintBits(reader);
+  CHECK_EQ(non_header_size, BitsToBytesRoundUp(reader.GetBitOffset())) << "Invalid CodeInfo";
+}
+
 BitTable<StackMap>::const_iterator CodeInfo::BinarySearchNativePc(uint32_t packed_pc) const {
   return std::partition_point(
       stack_maps_.begin(),
diff --git a/runtime/stack_map.h b/runtime/stack_map.h
index 7aac792..e94f71a 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -445,22 +445,7 @@
                             uint32_t first_dex_register,
                             /*out*/ DexRegisterMap* map) const;
 
-  void Decode(const uint8_t* data) {
-    size_t non_header_size = DecodeUnsignedLeb128(&data);
-    BitMemoryRegion region(MemoryRegion(const_cast<uint8_t*>(data), non_header_size));
-    size_t bit_offset = 0;
-    size_ = UnsignedLeb128Size(non_header_size) + non_header_size;
-    stack_maps_.Decode(region, &bit_offset);
-    register_masks_.Decode(region, &bit_offset);
-    stack_masks_.Decode(region, &bit_offset);
-    invoke_infos_.Decode(region, &bit_offset);
-    inline_infos_.Decode(region, &bit_offset);
-    dex_register_masks_.Decode(region, &bit_offset);
-    dex_register_maps_.Decode(region, &bit_offset);
-    dex_register_catalog_.Decode(region, &bit_offset);
-    number_of_dex_registers_ = DecodeVarintBits(region, &bit_offset);
-    CHECK_EQ(non_header_size, BitsToBytesRoundUp(bit_offset)) << "Invalid CodeInfo";
-  }
+  void Decode(const uint8_t* data);
 
   size_t size_;
   BitTable<StackMap> stack_maps_;