Move and rewrite bit loading/storing methods.

Move bit loading/storing methods to BitMemoryRegion and rewrite them.

Enforce natural alignment of the data pointer in BitMemoryRegion.
This probably would not be reasonable in MemoryRegion, but it is
fine here since the BitMemoryRegion already has extra bit offset.
The alignment makes it possible to simplify and optimize the methods.

This makes the stackmap reading code 33% faster.
(measured by the time needed to verify all stackmap fields on ARM).

Test: m -j40 test-art-host-gtest
Change-Id: I00be8052969a6056b262df4cd2066ffd86043196
diff --git a/compiler/optimizing/stack_map_stream.cc b/compiler/optimizing/stack_map_stream.cc
index 7010e3f..bf7c554 100644
--- a/compiler/optimizing/stack_map_stream.cc
+++ b/compiler/optimizing/stack_map_stream.cc
@@ -427,10 +427,11 @@
   if (stack_mask_bits > 0) {
     size_t stack_mask_bytes = RoundUp(stack_mask_bits, kBitsPerByte) / kBitsPerByte;
     for (size_t i = 0; i < encoding.stack_mask.num_entries; ++i) {
-      MemoryRegion source(&stack_masks_[i * stack_mask_bytes], stack_mask_bytes);
-      BitMemoryRegion stack_mask = code_info.GetStackMask(i, encoding);
-      for (size_t bit_index = 0; bit_index < stack_mask_bits; ++bit_index) {
-        stack_mask.StoreBit(bit_index, source.LoadBit(bit_index));
+      BitMemoryRegion src(MemoryRegion(&stack_masks_[i * stack_mask_bytes], stack_mask_bytes));
+      BitMemoryRegion dst = code_info.GetStackMask(i, encoding);
+      for (size_t bit_index = 0; bit_index < stack_mask_bits; bit_index += BitSizeOf<uint32_t>()) {
+        size_t num_bits = std::min<size_t>(stack_mask_bits - bit_index, BitSizeOf<uint32_t>());
+        dst.StoreBits(bit_index, src.LoadBits(bit_index, num_bits), num_bits);
       }
     }
   }
@@ -600,8 +601,9 @@
   for (StackMapEntry& stack_map : stack_maps_) {
     size_t index = dedup.size();
     MemoryRegion stack_mask(stack_masks_.data() + index * byte_entry_size, byte_entry_size);
+    BitMemoryRegion stack_mask_bits(stack_mask);
     for (size_t i = 0; i < entry_size_in_bits; i++) {
-      stack_mask.StoreBit(i, stack_map.sp_mask != nullptr && stack_map.sp_mask->IsBitSet(i));
+      stack_mask_bits.StoreBit(i, stack_map.sp_mask != nullptr && stack_map.sp_mask->IsBitSet(i));
     }
     stack_map.stack_mask_index = dedup.emplace(stack_mask, index).first->second;
   }
diff --git a/libartbase/Android.bp b/libartbase/Android.bp
index c8a06ed..692c97f 100644
--- a/libartbase/Android.bp
+++ b/libartbase/Android.bp
@@ -147,6 +147,7 @@
         "arch/instruction_set_test.cc",
         "base/arena_allocator_test.cc",
         "base/bit_field_test.cc",
+        "base/bit_memory_region_test.cc",
         "base/bit_string_test.cc",
         "base/bit_struct_test.cc",
         "base/bit_utils_test.cc",
diff --git a/libartbase/base/bit_memory_region.h b/libartbase/base/bit_memory_region.h
index f3926bc..dd16957 100644
--- a/libartbase/base/bit_memory_region.h
+++ b/libartbase/base/bit_memory_region.h
@@ -19,6 +19,9 @@
 
 #include "memory_region.h"
 
+#include "bit_utils.h"
+#include "memory_tool.h"
+
 namespace art {
 
 // Bit memory region is a bit offset subregion of a normal memoryregion. This is useful for
@@ -26,46 +29,113 @@
 class BitMemoryRegion FINAL : public ValueObject {
  public:
   BitMemoryRegion() = default;
-  ALWAYS_INLINE BitMemoryRegion(MemoryRegion region, size_t bit_offset, size_t bit_size) {
-    bit_start_ = bit_offset % kBitsPerByte;
-    const size_t start = bit_offset / kBitsPerByte;
-    const size_t end = (bit_offset + bit_size + kBitsPerByte - 1) / kBitsPerByte;
-    region_ = region.Subregion(start, end - start);
+  ALWAYS_INLINE explicit BitMemoryRegion(MemoryRegion region)
+    : data_(reinterpret_cast<uintptr_t*>(AlignDown(region.pointer(), sizeof(uintptr_t)))),
+      bit_start_(8 * (reinterpret_cast<uintptr_t>(region.pointer()) % sizeof(uintptr_t))),
+      bit_size_(region.size_in_bits()) {
+  }
+  ALWAYS_INLINE BitMemoryRegion(MemoryRegion region, size_t bit_offset, size_t bit_length)
+    : BitMemoryRegion(region) {
+    DCHECK_LE(bit_offset, bit_size_);
+    DCHECK_LE(bit_length, bit_size_ - bit_offset);
+    bit_start_ += bit_offset;
+    bit_size_ = bit_length;
   }
 
-  void* pointer() const { return region_.pointer(); }
-  size_t size() const { return region_.size(); }
-  size_t BitOffset() const { return bit_start_; }
+  ALWAYS_INLINE bool IsValid() const { return data_ != nullptr; }
+
   size_t size_in_bits() const {
-    return region_.size_in_bits();
+    return bit_size_;
   }
 
-  ALWAYS_INLINE BitMemoryRegion Subregion(size_t bit_offset, size_t bit_size) const {
-    return BitMemoryRegion(region_, bit_start_ + bit_offset, bit_size);
+  ALWAYS_INLINE BitMemoryRegion Subregion(size_t bit_offset, size_t bit_length) const {
+    DCHECK_LE(bit_offset, bit_size_);
+    DCHECK_LE(bit_length, bit_size_ - bit_offset);
+    BitMemoryRegion result = *this;
+    result.bit_start_ += bit_offset;
+    result.bit_size_ = bit_length;
+    return result;
   }
 
   // Load a single bit in the region. The bit at offset 0 is the least
   // significant bit in the first byte.
+  ATTRIBUTE_NO_SANITIZE_ADDRESS  // We might touch extra bytes due to the alignment.
   ALWAYS_INLINE bool LoadBit(uintptr_t bit_offset) const {
-    return region_.LoadBit(bit_offset + bit_start_);
+    DCHECK_LT(bit_offset, bit_size_);
+    size_t index = (bit_start_ + bit_offset) / kBitsPerIntPtrT;
+    size_t shift = (bit_start_ + bit_offset) % kBitsPerIntPtrT;
+    return ((data_[index] >> shift) & 1) != 0;
   }
 
   ALWAYS_INLINE void StoreBit(uintptr_t bit_offset, bool value) const {
-    region_.StoreBit(bit_offset + bit_start_, value);
+    DCHECK_LT(bit_offset, bit_size_);
+    uint8_t* data = reinterpret_cast<uint8_t*>(data_);
+    size_t index = (bit_start_ + bit_offset) / kBitsPerByte;
+    size_t shift = (bit_start_ + bit_offset) % kBitsPerByte;
+    data[index] &= ~(1 << shift);  // Clear bit.
+    data[index] |= (value ? 1 : 0) << shift;  // Set bit.
+    DCHECK_EQ(value, LoadBit(bit_offset));
   }
 
-  ALWAYS_INLINE uint32_t LoadBits(uintptr_t bit_offset, size_t length) const {
-    return region_.LoadBits(bit_offset + bit_start_, length);
+  // Load `bit_length` bits from `data` starting at given `bit_offset`.
+  // The least significant bit is stored in the smallest memory offset.
+  ATTRIBUTE_NO_SANITIZE_ADDRESS  // We might touch extra bytes due to the alignment.
+  ALWAYS_INLINE uint32_t LoadBits(size_t bit_offset, size_t bit_length) const {
+    DCHECK(IsAligned<sizeof(uintptr_t)>(data_));
+    DCHECK_LE(bit_offset, bit_size_);
+    DCHECK_LE(bit_length, bit_size_ - bit_offset);
+    DCHECK_LE(bit_length, BitSizeOf<uint32_t>());
+    if (bit_length == 0) {
+      return 0;
+    }
+    uintptr_t mask = std::numeric_limits<uintptr_t>::max() >> (kBitsPerIntPtrT - bit_length);
+    size_t index = (bit_start_ + bit_offset) / kBitsPerIntPtrT;
+    size_t shift = (bit_start_ + bit_offset) % kBitsPerIntPtrT;
+    uintptr_t value = data_[index] >> shift;
+    size_t finished_bits = kBitsPerIntPtrT - shift;
+    if (finished_bits < bit_length) {
+      value |= data_[index + 1] << finished_bits;
+    }
+    return value & mask;
   }
 
-  // Store at a bit offset from inside the bit memory region.
-  ALWAYS_INLINE void StoreBits(uintptr_t bit_offset, uint32_t value, size_t length) {
-    region_.StoreBits(bit_offset + bit_start_, value, length);
+  // 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) {
+    DCHECK_LE(bit_offset, bit_size_);
+    DCHECK_LE(bit_length, bit_size_ - bit_offset);
+    DCHECK_LE(bit_length, BitSizeOf<uint32_t>());
+    DCHECK_LE(value, MaxInt<uint32_t>(bit_length));
+    if (bit_length == 0) {
+      return;
+    }
+    // Write data byte by byte to avoid races with other threads
+    // on bytes that do not overlap with this region.
+    uint8_t* data = reinterpret_cast<uint8_t*>(data_);
+    uint32_t mask = std::numeric_limits<uint32_t>::max() >> (BitSizeOf<uint32_t>() - bit_length);
+    size_t index = (bit_start_ + bit_offset) / kBitsPerByte;
+    size_t shift = (bit_start_ + bit_offset) % kBitsPerByte;
+    data[index] &= ~(mask << shift);  // Clear bits.
+    data[index] |= (value << shift);  // Set bits.
+    size_t finished_bits = kBitsPerByte - shift;
+    for (int i = 1; finished_bits < bit_length; i++, finished_bits += kBitsPerByte) {
+      data[index + i] &= ~(mask >> finished_bits);  // Clear bits.
+      data[index + i] |= (value >> finished_bits);  // Set bits.
+    }
+    DCHECK_EQ(value, LoadBits(bit_offset, bit_length));
+  }
+
+  ALWAYS_INLINE bool Equals(const BitMemoryRegion& other) const {
+    return data_ == other.data_ &&
+           bit_start_ == other.bit_start_ &&
+           bit_size_ == other.bit_size_;
   }
 
  private:
-  MemoryRegion region_;
+  // The data pointer must be naturally aligned. This makes loading code faster.
+  uintptr_t* data_ = nullptr;
   size_t bit_start_ = 0;
+  size_t bit_size_ = 0;
 };
 
 }  // namespace art
diff --git a/libartbase/base/bit_memory_region_test.cc b/libartbase/base/bit_memory_region_test.cc
new file mode 100644
index 0000000..b754698
--- /dev/null
+++ b/libartbase/base/bit_memory_region_test.cc
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "bit_memory_region.h"
+
+#include "gtest/gtest.h"
+
+namespace art {
+
+static void CheckBits(uint8_t* data,
+                      size_t size,
+                      uint32_t init,
+                      size_t offset,
+                      size_t length,
+                      uint32_t value) {
+  for (size_t i = 0; i < size * kBitsPerByte; i++) {
+    uint8_t expected = (offset <= i && i < offset + length) ? value >> (i - offset) : init;
+    uint8_t actual = data[i / kBitsPerByte] >> (i % kBitsPerByte);
+    EXPECT_EQ(expected & 1, actual & 1);
+  }
+}
+
+TEST(BitMemoryRegion, TestBit) {
+  uint8_t data[sizeof(uint32_t) * 2];
+  for (size_t bit_offset = 0; bit_offset < 2 * sizeof(uint32_t) * kBitsPerByte; ++bit_offset) {
+    for (uint32_t initial_value = 0; initial_value <= 1; initial_value++) {
+      for (uint32_t value = 0; value <= 1; value++) {
+        // Check Store and Load with bit_offset set on the region.
+        std::fill_n(data, sizeof(data), initial_value * 0xFF);
+        BitMemoryRegion bmr1(MemoryRegion(&data, sizeof(data)), bit_offset, 1);
+        bmr1.StoreBit(0, value);
+        EXPECT_EQ(bmr1.LoadBit(0), value);
+        CheckBits(data, sizeof(data), initial_value, bit_offset, 1, value);
+        // Check Store and Load with bit_offset set on the methods.
+        std::fill_n(data, sizeof(data), initial_value * 0xFF);
+        BitMemoryRegion bmr2(MemoryRegion(&data, sizeof(data)));
+        bmr2.StoreBit(bit_offset, value);
+        EXPECT_EQ(bmr2.LoadBit(bit_offset), value);
+        CheckBits(data, sizeof(data), initial_value, bit_offset, 1, value);
+      }
+    }
+  }
+}
+
+TEST(BitMemoryRegion, TestBits) {
+  uint8_t data[sizeof(uint32_t) * 4];
+  for (size_t bit_offset = 0; bit_offset < 3 * sizeof(uint32_t) * kBitsPerByte; ++bit_offset) {
+    uint32_t mask = 0;
+    for (size_t bit_length = 0; bit_length < sizeof(uint32_t) * kBitsPerByte; ++bit_length) {
+      const uint32_t value = 0xDEADBEEF & mask;
+      for (uint32_t initial_value = 0; initial_value <= 1; initial_value++) {
+        // Check Store and Load with bit_offset set on the region.
+        std::fill_n(data, sizeof(data), initial_value * 0xFF);
+        BitMemoryRegion bmr1(MemoryRegion(&data, sizeof(data)), bit_offset, bit_length);
+        bmr1.StoreBits(0, value, bit_length);
+        EXPECT_EQ(bmr1.LoadBits(0, bit_length), value);
+        CheckBits(data, sizeof(data), initial_value, bit_offset, bit_length, value);
+        // Check Store and Load with bit_offset set on the methods.
+        std::fill_n(data, sizeof(data), initial_value * 0xFF);
+        BitMemoryRegion bmr2(MemoryRegion(&data, sizeof(data)));
+        bmr2.StoreBits(bit_offset, value, bit_length);
+        EXPECT_EQ(bmr2.LoadBits(bit_offset, bit_length), value);
+        CheckBits(data, sizeof(data), initial_value, bit_offset, bit_length, value);
+      }
+      mask = (mask << 1) | 1;
+    }
+  }
+}
+
+}  // namespace art
diff --git a/libartbase/base/bit_utils.h b/libartbase/base/bit_utils.h
index 04f0e85..58cc78c 100644
--- a/libartbase/base/bit_utils.h
+++ b/libartbase/base/bit_utils.h
@@ -22,6 +22,7 @@
 
 #include <android-base/logging.h>
 
+#include "globals.h"
 #include "stl_util_identity.h"
 
 namespace art {
@@ -499,6 +500,10 @@
   return bitfield_unsigned;
 }
 
+inline static constexpr size_t BitsToBytesRoundUp(size_t num_bits) {
+  return RoundUp(num_bits, kBitsPerByte) / kBitsPerByte;
+}
+
 }  // namespace art
 
 #endif  // ART_LIBARTBASE_BASE_BIT_UTILS_H_
diff --git a/libartbase/base/memory_region.cc b/libartbase/base/memory_region.cc
index 862ff73..d207872 100644
--- a/libartbase/base/memory_region.cc
+++ b/libartbase/base/memory_region.cc
@@ -29,36 +29,4 @@
   memmove(reinterpret_cast<void*>(begin() + offset), from.pointer(), from.size());
 }
 
-void MemoryRegion::StoreBits(uintptr_t bit_offset, uint32_t value, size_t length) {
-  DCHECK_LE(value, MaxInt<uint32_t>(length));
-  DCHECK_LE(length, BitSizeOf<uint32_t>());
-  DCHECK_LE(bit_offset + length, size_in_bits());
-  if (length == 0) {
-    return;
-  }
-  // Bits are stored in this order {7 6 5 4 3 2 1 0}.
-  // How many remaining bits in current byte is (bit_offset % kBitsPerByte) + 1.
-  uint8_t* out = ComputeInternalPointer<uint8_t>(bit_offset >> kBitsPerByteLog2);
-  size_t orig_len = length;
-  uint32_t orig_value = value;
-  uintptr_t bit_remainder = bit_offset % kBitsPerByte;
-  while (true) {
-    const uintptr_t remaining_bits = kBitsPerByte - bit_remainder;
-    if (length <= remaining_bits) {
-      // Length is smaller than all of remainder bits.
-      size_t mask = ((1 << length) - 1) << bit_remainder;
-      *out = (*out & ~mask) | (value << bit_remainder);
-      break;
-    }
-    // Copy remaining bits in current byte.
-    size_t value_mask = (1 << remaining_bits) - 1;
-    *out = (*out & ~(value_mask << bit_remainder)) | ((value & value_mask) << bit_remainder);
-    value >>= remaining_bits;
-    bit_remainder = 0;
-    length -= remaining_bits;
-    ++out;
-  }
-  DCHECK_EQ(LoadBits(bit_offset, orig_len), orig_value) << bit_offset << " " << orig_len;
-}
-
 }  // namespace art
diff --git a/libartbase/base/memory_region.h b/libartbase/base/memory_region.h
index 3d00f5b..2060329 100644
--- a/libartbase/base/memory_region.h
+++ b/libartbase/base/memory_region.h
@@ -109,67 +109,6 @@
     return ComputeInternalPointer<T>(offset);
   }
 
-  // Load a single bit in the region. The bit at offset 0 is the least
-  // significant bit in the first byte.
-  ALWAYS_INLINE bool LoadBit(uintptr_t bit_offset) const {
-    uint8_t bit_mask;
-    uint8_t byte = *ComputeBitPointer(bit_offset, &bit_mask);
-    return byte & bit_mask;
-  }
-
-  ALWAYS_INLINE void StoreBit(uintptr_t bit_offset, bool value) const {
-    uint8_t bit_mask;
-    uint8_t* byte = ComputeBitPointer(bit_offset, &bit_mask);
-    if (value) {
-      *byte |= bit_mask;
-    } else {
-      *byte &= ~bit_mask;
-    }
-  }
-
-  // Load `length` bits from the region starting at bit offset `bit_offset`.
-  // The bit at the smallest offset is the least significant bit in the
-  // loaded value.  `length` must not be larger than the number of bits
-  // contained in the return value (32).
-  ALWAYS_INLINE uint32_t LoadBits(uintptr_t bit_offset, size_t length) const {
-    DCHECK_LE(length, BitSizeOf<uint32_t>());
-    DCHECK_LE(bit_offset + length, size_in_bits());
-    if (UNLIKELY(length == 0)) {
-      // Do not touch any memory if the range is empty.
-      return 0;
-    }
-    const uint8_t* address = begin() + bit_offset / kBitsPerByte;
-    const uint32_t shift = bit_offset & (kBitsPerByte - 1);
-    // Load the value (reading only the strictly needed bytes).
-    const uint32_t load_bit_count = shift + length;
-    uint32_t value = address[0] >> shift;
-    if (load_bit_count > 8) {
-      value |= static_cast<uint32_t>(address[1]) << (8 - shift);
-      if (load_bit_count > 16) {
-        value |= static_cast<uint32_t>(address[2]) << (16 - shift);
-        if (load_bit_count > 24) {
-          value |= static_cast<uint32_t>(address[3]) << (24 - shift);
-          if (load_bit_count > 32) {
-            value |= static_cast<uint32_t>(address[4]) << (32 - shift);
-          }
-        }
-      }
-    }
-    // Clear unwanted most significant bits.
-    uint32_t clear_bit_count = BitSizeOf(value) - length;
-    value = (value << clear_bit_count) >> clear_bit_count;
-    for (size_t i = 0; i < length; ++i) {
-      DCHECK_EQ((value >> i) & 1, LoadBit(bit_offset + i));
-    }
-    return value;
-  }
-
-  // Store `value` on `length` bits in the region starting at bit offset
-  // `bit_offset`.  The bit at the smallest offset is the least significant
-  // bit of the stored `value`.  `value` must not be larger than `length`
-  // bits.
-  void StoreBits(uintptr_t bit_offset, uint32_t value, size_t length);
-
   void CopyFrom(size_t offset, const MemoryRegion& from) const;
 
   template<class Vector>
diff --git a/libartbase/base/memory_region_test.cc b/libartbase/base/memory_region_test.cc
index e3aead4..72e03a4 100644
--- a/libartbase/base/memory_region_test.cc
+++ b/libartbase/base/memory_region_test.cc
@@ -18,8 +18,6 @@
 
 #include "gtest/gtest.h"
 
-#include "bit_memory_region.h"
-
 namespace art {
 
 TEST(MemoryRegion, LoadUnaligned) {
@@ -57,35 +55,4 @@
   }
 }
 
-TEST(MemoryRegion, TestBits) {
-  const size_t n = 8;
-  uint8_t data[n] = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };
-  MemoryRegion region(&data, n);
-  uint32_t value = 0xDEADBEEF;
-  // Try various offsets and lengths.
-  for (size_t bit_offset = 0; bit_offset < 2 * kBitsPerByte; ++bit_offset) {
-    for (size_t length = 0; length < 2 * kBitsPerByte; ++length) {
-      const uint32_t length_mask = (1 << length) - 1;
-      uint32_t masked_value = value & length_mask;
-      BitMemoryRegion bmr(region, bit_offset, length);
-      region.StoreBits(bit_offset, masked_value, length);
-      EXPECT_EQ(region.LoadBits(bit_offset, length), masked_value);
-      EXPECT_EQ(bmr.LoadBits(0, length), masked_value);
-      // Check adjacent bits to make sure they were not incorrectly cleared.
-      EXPECT_EQ(region.LoadBits(0, bit_offset), (1u << bit_offset) - 1);
-      EXPECT_EQ(region.LoadBits(bit_offset + length, length), length_mask);
-      region.StoreBits(bit_offset, length_mask, length);
-      // Store with bit memory region.
-      bmr.StoreBits(0, masked_value, length);
-      EXPECT_EQ(bmr.LoadBits(0, length), masked_value);
-      // Check adjacent bits to make sure they were not incorrectly cleared.
-      EXPECT_EQ(region.LoadBits(0, bit_offset), (1u << bit_offset) - 1);
-      EXPECT_EQ(region.LoadBits(bit_offset + length, length), length_mask);
-      region.StoreBits(bit_offset, length_mask, length);
-      // Flip the value to try different edge bit combinations.
-      value = ~value;
-    }
-  }
-}
-
 }  // namespace art
diff --git a/runtime/method_info.h b/runtime/method_info.h
index b00ddc6..6f74678 100644
--- a/runtime/method_info.h
+++ b/runtime/method_info.h
@@ -21,7 +21,7 @@
 
 #include "base/leb128.h"
 #include "base/macros.h"
-#include "base/memory_region.h"
+#include "base/bit_memory_region.h"
 
 namespace art {
 
@@ -35,8 +35,8 @@
   explicit MethodInfo(const uint8_t* ptr) {
     if (ptr != nullptr) {
       num_method_indices_ = DecodeUnsignedLeb128(&ptr);
-      region_ = MemoryRegion(const_cast<uint8_t*>(ptr),
-                             num_method_indices_ * sizeof(MethodIndexType));
+      region_ = BitMemoryRegion(
+          MemoryRegion(const_cast<uint8_t*>(ptr), num_method_indices_ * sizeof(MethodIndexType)));
     }
   }
 
@@ -44,7 +44,7 @@
   MethodInfo(uint8_t* ptr, size_t num_method_indices) : num_method_indices_(num_method_indices) {
     DCHECK(ptr != nullptr);
     ptr = EncodeUnsignedLeb128(ptr, num_method_indices_);
-    region_ = MemoryRegion(ptr, num_method_indices_ * sizeof(MethodIndexType));
+    region_ = BitMemoryRegion(MemoryRegion(ptr, num_method_indices_ * sizeof(MethodIndexType)));
   }
 
   static size_t ComputeSize(size_t num_method_indices) {
@@ -71,7 +71,7 @@
 
  private:
   size_t num_method_indices_ = 0u;
-  MemoryRegion region_;
+  BitMemoryRegion region_;
 };
 
 }  // namespace art
diff --git a/runtime/stack_map.h b/runtime/stack_map.h
index 3839764..7d3d56a 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -452,7 +452,7 @@
   explicit DexRegisterMap(MemoryRegion region) : region_(region) {}
   DexRegisterMap() {}
 
-  bool IsValid() const { return region_.pointer() != nullptr; }
+  bool IsValid() const { return region_.IsValid(); }
 
   // Get the surface kind of Dex register `dex_register_number`.
   DexRegisterLocation::Kind GetLocationKind(uint16_t dex_register_number,
@@ -627,7 +627,7 @@
 
   // Return the size of the DexRegisterMap object, in bytes.
   size_t Size() const {
-    return region_.size();
+    return BitsToBytesRoundUp(region_.size_in_bits());
   }
 
   void Dump(VariableIndentationOutputStream* vios,
@@ -650,7 +650,7 @@
 
   static constexpr int kFixedSize = 0;
 
-  MemoryRegion region_;
+  BitMemoryRegion region_;
 
   friend class CodeInfo;
   friend class StackMapStream;
@@ -678,7 +678,7 @@
 
   template <typename Region>
   ALWAYS_INLINE void Store(Region region, int32_t value) const {
-    region.StoreBits(start_offset_, value - min_value_, BitSize());
+    region.StoreBits(start_offset_, static_cast<uint32_t>(value - min_value_), BitSize());
     DCHECK_EQ(Load(region), value);
   }
 
@@ -805,7 +805,7 @@
   StackMap() {}
   explicit StackMap(BitMemoryRegion region) : region_(region) {}
 
-  ALWAYS_INLINE bool IsValid() const { return region_.pointer() != nullptr; }
+  ALWAYS_INLINE bool IsValid() const { return region_.IsValid(); }
 
   ALWAYS_INLINE uint32_t GetDexPc(const StackMapEncoding& encoding) const {
     return encoding.GetDexPcEncoding().Load(region_);
@@ -868,9 +868,7 @@
   }
 
   ALWAYS_INLINE bool Equals(const StackMap& other) const {
-    return region_.pointer() == other.region_.pointer() &&
-           region_.size() == other.region_.size() &&
-           region_.BitOffset() == other.region_.BitOffset();
+    return region_.Equals(other.region_);
   }
 
   void Dump(VariableIndentationOutputStream* vios,
@@ -1257,7 +1255,7 @@
     return method_info.GetMethodIndex(GetMethodIndexIdx(encoding));
   }
 
-  bool IsValid() const { return region_.pointer() != nullptr; }
+  bool IsValid() const { return region_.IsValid(); }
 
  private:
   BitMemoryRegion region_;