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
diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index 9914ef4..5818a37 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 @@
         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 @@
     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 @@
     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 @@
     inline_infos_.Add(entry);
   }
 
-  size_t ComputeNeededSize() const {
+  size_t ComputeNeededSize() {
     size_t size = CodeInfo::kFixedSize
         + ComputeStackMapsSize()
         + ComputeDexRegisterMapsSize()
@@ -118,7 +122,7 @@
     return StackMaskEncodingSize(stack_mask_max_);
   }
 
-  size_t ComputeStackMapsSize() const {
+  size_t ComputeStackMapsSize() {
     return stack_maps_.Size() * StackMap::ComputeStackMapSize(
         ComputeStackMaskSize(),
         ComputeInlineInfoSize(),
@@ -146,10 +150,13 @@
   }
 
   // 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 @@
       + (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 @@
         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 @@
       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 @@
   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 e7075c0..e5a9790 100644
--- a/compiler/optimizing/stack_map_test.cc
+++ b/compiler/optimizing/stack_map_test.cc
@@ -231,4 +231,54 @@
   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 629fc9a..6ec7cc8 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -212,6 +212,14 @@
   // 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_;