Share dex register maps between stack maps when possible.

If two stack maps have the same dex register map then one of them will
reference the register map from the other instead of owning an
independent copy.

This saves around 1.5% of space.

Change-Id: Ic2c2c81210c6c45a5c5f650f7ba82a46ff6f45e4

3 files changed, 190 insertions, 37 deletions

diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index 9914ef49c..5818a37a4 100644
--- a/compiler/optimizing/stack_map_stream.h
+++ b/compiler/optimizing/stack_map_stream.h
@@ -17,7 +17,8 @@
 #ifndef ART_COMPILER_OPTIMIZING_STACK_MAP_STREAM_H_
 #define ART_COMPILER_OPTIMIZING_STACK_MAP_STREAM_H_
 
-#include "base/bit_vector.h"
+#include "base/arena_containers.h"
+#include "base/bit_vector-inl.h"
 #include "base/value_object.h"
 #include "memory_region.h"
 #include "nodes.h"
@@ -40,7 +41,8 @@ class StackMapStream : public ValueObject {
         stack_mask_max_(-1),
         dex_pc_max_(0),
         native_pc_offset_max_(0),
-        number_of_stack_maps_with_inline_info_(0) {}
+        number_of_stack_maps_with_inline_info_(0),
+        dex_map_hash_to_stack_map_indices_(std::less<uint32_t>(), allocator->Adapter()) {}
 
   // Compute bytes needed to encode a mask with the given maximum element.
   static uint32_t StackMaskEncodingSize(int max_element) {
@@ -59,6 +61,7 @@ class StackMapStream : public ValueObject {
     size_t dex_register_locations_start_index;
     size_t inline_infos_start_index;
     BitVector* live_dex_registers_mask;
+    uint32_t dex_register_map_hash;
   };
 
   struct InlineInfoEntry {
@@ -80,6 +83,7 @@ class StackMapStream : public ValueObject {
     entry.inlining_depth = inlining_depth;
     entry.dex_register_locations_start_index = dex_register_locations_.Size();
     entry.inline_infos_start_index = inline_infos_.Size();
+    entry.dex_register_map_hash = 0;
     if (num_dex_registers != 0) {
       entry.live_dex_registers_mask =
           new (allocator_) ArenaBitVector(allocator_, num_dex_registers, true);
@@ -105,7 +109,7 @@ class StackMapStream : public ValueObject {
     inline_infos_.Add(entry);
   }
 
-  size_t ComputeNeededSize() const {
+  size_t ComputeNeededSize() {
     size_t size = CodeInfo::kFixedSize
         + ComputeStackMapsSize()
         + ComputeDexRegisterMapsSize()
@@ -118,7 +122,7 @@ class StackMapStream : public ValueObject {
     return StackMaskEncodingSize(stack_mask_max_);
   }
 
-  size_t ComputeStackMapsSize() const {
+  size_t ComputeStackMapsSize() {
     return stack_maps_.Size() * StackMap::ComputeStackMapSize(
         ComputeStackMaskSize(),
         ComputeInlineInfoSize(),
@@ -146,10 +150,13 @@ class StackMapStream : public ValueObject {
   }
 
   // Compute the size of all the Dex register maps.
-  size_t ComputeDexRegisterMapsSize() const {
+  size_t ComputeDexRegisterMapsSize() {
     size_t size = 0;
     for (size_t i = 0; i < stack_maps_.Size(); ++i) {
-      size += ComputeDexRegisterMapSize(stack_maps_.Get(i));
+      if (FindEntryWithTheSameDexMap(i) == kNoSameDexMapFound) {
+        // Entries with the same dex map will have the same offset.
+        size += ComputeDexRegisterMapSize(stack_maps_.Get(i));
+      }
     }
     return size;
   }
@@ -161,11 +168,11 @@ class StackMapStream : public ValueObject {
       + (number_of_stack_maps_with_inline_info_ * InlineInfo::kFixedSize);
   }
 
-  size_t ComputeDexRegisterMapsStart() const {
+  size_t ComputeDexRegisterMapsStart() {
     return CodeInfo::kFixedSize + ComputeStackMapsSize();
   }
 
-  size_t ComputeInlineInfoStart() const {
+  size_t ComputeInlineInfoStart() {
     return ComputeDexRegisterMapsStart() + ComputeDexRegisterMapsSize();
   }
 
@@ -206,38 +213,47 @@ class StackMapStream : public ValueObject {
         stack_map.SetStackMask(code_info, *entry.sp_mask);
       }
 
-      if (entry.num_dex_registers != 0) {
-        // Set the Dex register map.
-        MemoryRegion register_region =
-            dex_register_locations_region.Subregion(
-                next_dex_register_map_offset,
-                ComputeDexRegisterMapSize(entry));
-        next_dex_register_map_offset += register_region.size();
-        DexRegisterMap dex_register_map(register_region);
-        stack_map.SetDexRegisterMapOffset(
+      if (entry.num_dex_registers == 0) {
+        // No dex map available.
+        stack_map.SetDexRegisterMapOffset(code_info, StackMap::kNoDexRegisterMap);
+      } else {
+        // Search for an entry with the same dex map.
+        size_t entry_with_same_map = FindEntryWithTheSameDexMap(i);
+        if (entry_with_same_map != kNoSameDexMapFound) {
+          // If we have a hit reuse the offset.
+          stack_map.SetDexRegisterMapOffset(code_info,
+              code_info.GetStackMapAt(entry_with_same_map).GetDexRegisterMapOffset(code_info));
+        } else {
+          // New dex registers maps should be added to the stack map.
+          MemoryRegion register_region =
+              dex_register_locations_region.Subregion(
+                  next_dex_register_map_offset,
+                  ComputeDexRegisterMapSize(entry));
+          next_dex_register_map_offset += register_region.size();
+          DexRegisterMap dex_register_map(register_region);
+          stack_map.SetDexRegisterMapOffset(
             code_info, register_region.start() - dex_register_locations_region.start());
 
-        // Offset in `dex_register_map` where to store the next register entry.
-        size_t offset = DexRegisterMap::kFixedSize;
-        dex_register_map.SetLiveBitMask(offset,
-                                        entry.num_dex_registers,
-                                        *entry.live_dex_registers_mask);
-        offset += DexRegisterMap::LiveBitMaskSize(entry.num_dex_registers);
-        for (size_t dex_register_number = 0, index_in_dex_register_locations = 0;
-             dex_register_number < entry.num_dex_registers;
-             ++dex_register_number) {
-          if (entry.live_dex_registers_mask->IsBitSet(dex_register_number)) {
-            DexRegisterLocation dex_register_location = dex_register_locations_.Get(
-                entry.dex_register_locations_start_index + index_in_dex_register_locations);
-            dex_register_map.SetRegisterInfo(offset, dex_register_location);
-            offset += DexRegisterMap::EntrySize(dex_register_location);
-            ++index_in_dex_register_locations;
+          // Offset in `dex_register_map` where to store the next register entry.
+          size_t offset = DexRegisterMap::kFixedSize;
+          dex_register_map.SetLiveBitMask(offset,
+                                          entry.num_dex_registers,
+                                          *entry.live_dex_registers_mask);
+          offset += DexRegisterMap::LiveBitMaskSize(entry.num_dex_registers);
+          for (size_t dex_register_number = 0, index_in_dex_register_locations = 0;
+               dex_register_number < entry.num_dex_registers;
+               ++dex_register_number) {
+            if (entry.live_dex_registers_mask->IsBitSet(dex_register_number)) {
+              DexRegisterLocation dex_register_location = dex_register_locations_.Get(
+                  entry.dex_register_locations_start_index + index_in_dex_register_locations);
+              dex_register_map.SetRegisterInfo(offset, dex_register_location);
+              offset += DexRegisterMap::EntrySize(dex_register_location);
+              ++index_in_dex_register_locations;
+            }
           }
+          // Ensure we reached the end of the Dex registers region.
+          DCHECK_EQ(offset, register_region.size());
         }
-        // Ensure we reached the end of the Dex registers region.
-        DCHECK_EQ(offset, register_region.size());
-      } else {
-        stack_map.SetDexRegisterMapOffset(code_info, StackMap::kNoDexRegisterMap);
       }
 
       // Set the inlining info.
@@ -271,11 +287,86 @@ class StackMapStream : public ValueObject {
       DCHECK(DexRegisterLocation::IsShortLocationKind(kind))
           << DexRegisterLocation::PrettyDescriptor(kind);
       dex_register_locations_.Add(DexRegisterLocation(kind, value));
-      stack_maps_.Get(stack_maps_.Size() - 1).live_dex_registers_mask->SetBit(dex_register);
+      StackMapEntry entry = stack_maps_.Get(stack_maps_.Size() - 1);
+      entry.live_dex_registers_mask->SetBit(dex_register);
+      entry.dex_register_map_hash += (1 << dex_register);
+      entry.dex_register_map_hash += static_cast<uint32_t>(value);
+      entry.dex_register_map_hash += static_cast<uint32_t>(kind);
+      stack_maps_.Put(stack_maps_.Size() - 1, entry);
     }
   }
 
  private:
+  // Returns the index of an entry with the same dex register map
+  // or kNoSameDexMapFound if no such entry exists.
+  size_t FindEntryWithTheSameDexMap(size_t entry_index) {
+    StackMapEntry entry = stack_maps_.Get(entry_index);
+    auto entries_it = dex_map_hash_to_stack_map_indices_.find(entry.dex_register_map_hash);
+    if (entries_it == dex_map_hash_to_stack_map_indices_.end()) {
+      // We don't have a perfect hash functions so we need a list to collect all stack maps
+      // which might have the same dex register map.
+      GrowableArray<uint32_t> stack_map_indices(allocator_, 1);
+      stack_map_indices.Add(entry_index);
+      dex_map_hash_to_stack_map_indices_.Put(entry.dex_register_map_hash, stack_map_indices);
+      return kNoSameDexMapFound;
+    }
+
+    // TODO: We don't need to add ourselves to the map if we can guarantee that
+    // FindEntryWithTheSameDexMap is called just once per stack map entry.
+    // A good way to do this is to cache the offset in the stack map entry. This
+    // is easier to do if we add markers when the stack map constructions begins
+    // and when it ends.
+
+    // We might have collisions, so we need to check whether or not we should
+    // add the entry to the map. `needs_to_be_added` keeps track of this.
+    bool needs_to_be_added = true;
+    size_t result = kNoSameDexMapFound;
+    for (size_t i = 0; i < entries_it->second.Size(); i++) {
+      size_t test_entry_index = entries_it->second.Get(i);
+      if (test_entry_index == entry_index) {
+        needs_to_be_added = false;
+      } else if (HaveTheSameDexMaps(stack_maps_.Get(test_entry_index), entry)) {
+        result = test_entry_index;
+        needs_to_be_added = false;
+        break;
+      }
+    }
+    if (needs_to_be_added) {
+      entries_it->second.Add(entry_index);
+    }
+    return result;
+  }
+
+  bool HaveTheSameDexMaps(const StackMapEntry& a, const StackMapEntry& b) const {
+    if (a.live_dex_registers_mask == nullptr && b.live_dex_registers_mask == nullptr) {
+      return true;
+    }
+    if (a.live_dex_registers_mask == nullptr || b.live_dex_registers_mask == nullptr) {
+      return false;
+    }
+    if (a.num_dex_registers != b.num_dex_registers) {
+      return false;
+    }
+
+    int index_in_dex_register_locations = 0;
+    for (uint32_t i = 0; i < a.num_dex_registers; i++) {
+      if (a.live_dex_registers_mask->IsBitSet(i) != b.live_dex_registers_mask->IsBitSet(i)) {
+        return false;
+      }
+      if (a.live_dex_registers_mask->IsBitSet(i)) {
+        DexRegisterLocation a_loc = dex_register_locations_.Get(
+            a.dex_register_locations_start_index + index_in_dex_register_locations);
+        DexRegisterLocation b_loc = dex_register_locations_.Get(
+            b.dex_register_locations_start_index + index_in_dex_register_locations);
+        if (a_loc != b_loc) {
+          return false;
+        }
+        ++index_in_dex_register_locations;
+      }
+    }
+    return true;
+  }
+
   ArenaAllocator* allocator_;
   GrowableArray<StackMapEntry> stack_maps_;
   GrowableArray<DexRegisterLocation> dex_register_locations_;
@@ -285,6 +376,10 @@ class StackMapStream : public ValueObject {
   uint32_t native_pc_offset_max_;
   size_t number_of_stack_maps_with_inline_info_;
 
+  ArenaSafeMap<uint32_t, GrowableArray<uint32_t>> dex_map_hash_to_stack_map_indices_;
+
+  static constexpr uint32_t kNoSameDexMapFound = -1;
+
   ART_FRIEND_TEST(StackMapTest, Test1);
   ART_FRIEND_TEST(StackMapTest, Test2);
   ART_FRIEND_TEST(StackMapTest, TestNonLiveDexRegisters);
diff --git a/compiler/optimizing/stack_map_test.cc b/compiler/optimizing/stack_map_test.cc
index e7075c0ae..e5a979025 100644
--- a/compiler/optimizing/stack_map_test.cc
+++ b/compiler/optimizing/stack_map_test.cc
@@ -231,4 +231,54 @@ TEST(StackMapTest, DexRegisterMapOffsetOverflow) {
   ASSERT_EQ(stack_map.GetDexRegisterMapOffset(code_info), StackMap::kNoDexRegisterMapSmallEncoding);
 }
 
+TEST(StackMapTest, TestShareDexRegisterMap) {
+  ArenaPool pool;
+  ArenaAllocator arena(&pool);
+  StackMapStream stream(&arena);
+
+  ArenaBitVector sp_mask(&arena, 0, false);
+  uint32_t number_of_dex_registers = 2;
+  // First stack map.
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  stream.AddDexRegisterEntry(0, DexRegisterLocation::Kind::kInRegister, 0);
+  stream.AddDexRegisterEntry(1, DexRegisterLocation::Kind::kConstant, -2);
+  // Second stack map, which should share the same dex register map.
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  stream.AddDexRegisterEntry(0, DexRegisterLocation::Kind::kInRegister, 0);
+  stream.AddDexRegisterEntry(1, DexRegisterLocation::Kind::kConstant, -2);
+  // Third stack map (doesn't share the dex register map).
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  stream.AddDexRegisterEntry(0, DexRegisterLocation::Kind::kInRegister, 2);
+  stream.AddDexRegisterEntry(1, DexRegisterLocation::Kind::kConstant, -2);
+
+  size_t size = stream.ComputeNeededSize();
+  void* memory = arena.Alloc(size, kArenaAllocMisc);
+  MemoryRegion region(memory, size);
+  stream.FillIn(region);
+
+  CodeInfo ci(region);
+  // Verify first stack map.
+  StackMap sm0 = ci.GetStackMapAt(0);
+  DexRegisterMap dex_registers0 = ci.GetDexRegisterMapOf(sm0, number_of_dex_registers);
+  ASSERT_EQ(0, dex_registers0.GetMachineRegister(0, number_of_dex_registers));
+  ASSERT_EQ(-2, dex_registers0.GetConstant(1, number_of_dex_registers));
+
+  // Verify second stack map.
+  StackMap sm1 = ci.GetStackMapAt(1);
+  DexRegisterMap dex_registers1 = ci.GetDexRegisterMapOf(sm1, number_of_dex_registers);
+  ASSERT_EQ(0, dex_registers1.GetMachineRegister(0, number_of_dex_registers));
+  ASSERT_EQ(-2, dex_registers1.GetConstant(1, number_of_dex_registers));
+
+  // Verify third stack map.
+  StackMap sm2 = ci.GetStackMapAt(2);
+  DexRegisterMap dex_registers2 = ci.GetDexRegisterMapOf(sm2, number_of_dex_registers);
+  ASSERT_EQ(2, dex_registers2.GetMachineRegister(0, number_of_dex_registers));
+  ASSERT_EQ(-2, dex_registers2.GetConstant(1, number_of_dex_registers));
+
+  // Verify dex register map offsets.
+  ASSERT_EQ(sm0.GetDexRegisterMapOffset(ci), sm1.GetDexRegisterMapOffset(ci));
+  ASSERT_NE(sm0.GetDexRegisterMapOffset(ci), sm2.GetDexRegisterMapOffset(ci));
+  ASSERT_NE(sm1.GetDexRegisterMapOffset(ci), sm2.GetDexRegisterMapOffset(ci));
+}
+
 }  // namespace art
diff --git a/runtime/stack_map.h b/runtime/stack_map.h
index 629fc9a34..6ec7cc852 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -212,6 +212,14 @@ class DexRegisterLocation {
   // Get the actual kind of the location.
   Kind GetInternalKind() const { return kind_; }
 
+  bool operator==(DexRegisterLocation other) const {
+    return kind_ == other.kind_ && value_ == other.value_;
+  }
+
+  bool operator!=(DexRegisterLocation other) const {
+    return !(*this == other);
+  }
+
  private:
   Kind kind_;
   int32_t value_;