Merge "Move and rewrite bit loading/storing methods."
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_;