[optimizing] Don't record None locations in the stack maps.

- moved environment recording from code generator to stack map stream
- added creation/loading factory methods for the DexRegisterMap (hides
internal details)
- added new tests

Change-Id: Ic8b6d044f0d8255c6759c19a41df332ef37876fe
diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index 863bab2..3168801 100644
--- a/compiler/optimizing/stack_map_stream.h
+++ b/compiler/optimizing/stack_map_stream.h
@@ -20,6 +20,7 @@
 #include "base/bit_vector.h"
 #include "base/value_object.h"
 #include "memory_region.h"
+#include "nodes.h"
 #include "stack_map.h"
 #include "utils/growable_array.h"
 
@@ -32,8 +33,9 @@
 class StackMapStream : public ValueObject {
  public:
   explicit StackMapStream(ArenaAllocator* allocator)
-      : stack_maps_(allocator, 10),
-        dex_register_maps_(allocator, 10 * 4),
+      : allocator_(allocator),
+        stack_maps_(allocator, 10),
+        dex_register_locations_(allocator, 10 * 4),
         inline_infos_(allocator, 2),
         stack_mask_max_(-1),
         number_of_stack_maps_with_inline_info_(0) {}
@@ -52,8 +54,9 @@
     BitVector* sp_mask;
     uint32_t num_dex_registers;
     uint8_t inlining_depth;
-    size_t dex_register_maps_start_index;
+    size_t dex_register_locations_start_index;
     size_t inline_infos_start_index;
+    BitVector* live_dex_registers_mask;
   };
 
   struct InlineInfoEntry {
@@ -65,7 +68,8 @@
                         uint32_t register_mask,
                         BitVector* sp_mask,
                         uint32_t num_dex_registers,
-                        uint8_t inlining_depth) {
+                        uint8_t inlining_depth,
+                        BitVector* live_dex_registers_mask) {
     StackMapEntry entry;
     entry.dex_pc = dex_pc;
     entry.native_pc_offset = native_pc_offset;
@@ -73,8 +77,9 @@
     entry.sp_mask = sp_mask;
     entry.num_dex_registers = num_dex_registers;
     entry.inlining_depth = inlining_depth;
-    entry.dex_register_maps_start_index = dex_register_maps_.Size();
+    entry.dex_register_locations_start_index = dex_register_locations_.Size();
     entry.inline_infos_start_index = inline_infos_.Size();
+    entry.live_dex_registers_mask = live_dex_registers_mask;
     stack_maps_.Add(entry);
 
     if (sp_mask != nullptr) {
@@ -85,11 +90,146 @@
     }
   }
 
-  void AddDexRegisterEntry(DexRegisterLocation::Kind kind, int32_t value) {
-    // Ensure we only use non-compressed location kind at this stage.
-    DCHECK(DexRegisterLocation::IsShortLocationKind(kind))
-        << DexRegisterLocation::PrettyDescriptor(kind);
-    dex_register_maps_.Add(DexRegisterLocation(kind, value));
+  void RecordEnvironment(HEnvironment* environment,
+                         size_t environment_size,
+                         LocationSummary* locations,
+                         uint32_t dex_pc,
+                         uint32_t native_pc,
+                         uint32_t register_mask,
+                         uint32_t inlining_depth) {
+    if (environment == nullptr) {
+      // For stack overflow checks.
+      AddStackMapEntry(dex_pc, native_pc, 0, 0, 0, inlining_depth, nullptr);
+      return;
+    }
+
+    BitVector* live_dex_registers_mask = new (allocator_) ArenaBitVector(allocator_, 0, true);
+
+    AddStackMapEntry(
+        dex_pc, native_pc, register_mask,
+        locations->GetStackMask(), environment_size, inlining_depth, live_dex_registers_mask);
+
+    // Walk over the environment, and record the location of dex registers.
+    for (size_t i = 0; i < environment_size; ++i) {
+      HInstruction* current = environment->GetInstructionAt(i);
+      if (current == nullptr) {
+        // No need to store anything, the `live_dex_registers_mask` will hold the
+        // information that this register is not live.
+        continue;
+      }
+
+      Location location = locations->GetEnvironmentAt(i);
+      switch (location.GetKind()) {
+        case Location::kConstant: {
+          DCHECK_EQ(current, location.GetConstant());
+          if (current->IsLongConstant()) {
+            // TODO: Consider moving setting the bit in AddDexRegisterEntry to avoid
+            // doing it manually here.
+            live_dex_registers_mask->SetBit(i);
+            live_dex_registers_mask->SetBit(i + 1);
+            int64_t value = current->AsLongConstant()->GetValue();
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, Low32Bits(value));
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, High32Bits(value));
+            ++i;
+            DCHECK_LT(i, environment_size);
+          } else if (current->IsDoubleConstant()) {
+            live_dex_registers_mask->SetBit(i);
+            live_dex_registers_mask->SetBit(i + 1);
+            int64_t value = bit_cast<double, int64_t>(current->AsDoubleConstant()->GetValue());
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, Low32Bits(value));
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, High32Bits(value));
+            ++i;
+            DCHECK_LT(i, environment_size);
+          } else if (current->IsIntConstant()) {
+            live_dex_registers_mask->SetBit(i);
+            int32_t value = current->AsIntConstant()->GetValue();
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, value);
+          } else if (current->IsNullConstant()) {
+            live_dex_registers_mask->SetBit(i);
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, 0);
+          } else {
+            DCHECK(current->IsFloatConstant()) << current->DebugName();
+            live_dex_registers_mask->SetBit(i);
+            int32_t value = bit_cast<float, int32_t>(current->AsFloatConstant()->GetValue());
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kConstant, value);
+          }
+          break;
+        }
+
+        case Location::kStackSlot: {
+          live_dex_registers_mask->SetBit(i);
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInStack,
+                              location.GetStackIndex());
+          break;
+        }
+
+        case Location::kDoubleStackSlot: {
+          live_dex_registers_mask->SetBit(i);
+          live_dex_registers_mask->SetBit(i + 1);
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInStack, location.GetStackIndex());
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInStack,
+                              location.GetHighStackIndex(kVRegSize));
+          ++i;
+          DCHECK_LT(i, environment_size);
+          break;
+        }
+
+        case Location::kRegister : {
+          live_dex_registers_mask->SetBit(i);
+          int id = location.reg();
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInRegister, id);
+          if (current->GetType() == Primitive::kPrimLong) {
+            live_dex_registers_mask->SetBit(i + 1);
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kInRegister, id);
+            ++i;
+            DCHECK_LT(i, environment_size);
+          }
+          break;
+        }
+
+        case Location::kFpuRegister : {
+          live_dex_registers_mask->SetBit(i);
+          int id = location.reg();
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInFpuRegister, id);
+          if (current->GetType() == Primitive::kPrimDouble) {
+            live_dex_registers_mask->SetBit(i + 1);
+            AddDexRegisterEntry(DexRegisterLocation::Kind::kInFpuRegister, id);
+            ++i;
+            DCHECK_LT(i, environment_size);
+          }
+          break;
+        }
+
+        case Location::kFpuRegisterPair : {
+          live_dex_registers_mask->SetBit(i);
+          live_dex_registers_mask->SetBit(i + 1);
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInFpuRegister, location.low());
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInFpuRegister, location.high());
+          ++i;
+          DCHECK_LT(i, environment_size);
+          break;
+        }
+
+        case Location::kRegisterPair : {
+          live_dex_registers_mask->SetBit(i);
+          live_dex_registers_mask->SetBit(i + 1);
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInRegister, location.low());
+          AddDexRegisterEntry(DexRegisterLocation::Kind::kInRegister, location.high());
+          ++i;
+          DCHECK_LT(i, environment_size);
+          break;
+        }
+
+        case Location::kInvalid: {
+          // No need to store anything, the `live_dex_registers_mask` will hold the
+          // information that this register is not live.
+          break;
+        }
+
+        default:
+          LOG(FATAL) << "Unexpected kind " << location.GetKind();
+      }
+    }
   }
 
   void AddInlineInfoEntry(uint32_t method_index) {
@@ -118,22 +258,26 @@
   // Compute the size of the Dex register map of `entry`.
   size_t ComputeDexRegisterMapSize(const StackMapEntry& entry) const {
     size_t size = DexRegisterMap::kFixedSize;
-    for (size_t j = 0; j < entry.num_dex_registers; ++j) {
-      DexRegisterLocation dex_register_location =
-          dex_register_maps_.Get(entry.dex_register_maps_start_index + j);
-      size += DexRegisterMap::EntrySize(dex_register_location);
+    // Add the bit mask for the dex register liveness.
+    size += 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);
+        size += DexRegisterMap::EntrySize(dex_register_location);
+        index_in_dex_register_locations++;
+      }
     }
     return size;
   }
 
   // Compute the size of all the Dex register maps.
   size_t ComputeDexRegisterMapsSize() const {
-    size_t size = stack_maps_.Size() * DexRegisterMap::kFixedSize;
-    // The size of each register location depends on the type of
-    // the entry.
-    for (size_t i = 0, e = dex_register_maps_.Size(); i < e; ++i) {
-      DexRegisterLocation entry = dex_register_maps_.Get(i);
-      size += DexRegisterMap::EntrySize(entry);
+    size_t size = 0;
+    for (size_t i = 0; i < stack_maps_.Size(); ++i) {
+      size += ComputeDexRegisterMapSize(stack_maps_.Get(i));
     }
     return size;
   }
@@ -161,7 +305,7 @@
     size_t stack_mask_size = ComputeStackMaskSize();
     uint8_t* memory_start = region.start();
 
-    MemoryRegion dex_register_maps_region = region.Subregion(
+    MemoryRegion dex_register_locations_region = region.Subregion(
       ComputeDexRegisterMapsStart(),
       ComputeDexRegisterMapsSize());
 
@@ -189,7 +333,7 @@
       if (entry.num_dex_registers != 0) {
         // Set the Dex register map.
         MemoryRegion register_region =
-            dex_register_maps_region.Subregion(
+            dex_register_locations_region.Subregion(
                 next_dex_register_map_offset,
                 ComputeDexRegisterMapSize(entry));
         next_dex_register_map_offset += register_region.size();
@@ -198,11 +342,20 @@
 
         // Offset in `dex_register_map` where to store the next register entry.
         size_t offset = DexRegisterMap::kFixedSize;
-        for (size_t j = 0; j < entry.num_dex_registers; ++j) {
-          DexRegisterLocation dex_register_location =
-              dex_register_maps_.Get(entry.dex_register_maps_start_index + j);
-          dex_register_map.SetRegisterInfo(offset, dex_register_location);
-          offset += DexRegisterMap::EntrySize(dex_register_location);
+        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());
@@ -232,12 +385,24 @@
   }
 
  private:
+  void AddDexRegisterEntry(DexRegisterLocation::Kind kind, int32_t value) {
+    // Ensure we only use non-compressed location kind at this stage.
+    DCHECK(DexRegisterLocation::IsShortLocationKind(kind))
+        << DexRegisterLocation::PrettyDescriptor(kind);
+    dex_register_locations_.Add(DexRegisterLocation(kind, value));
+  }
+
+  ArenaAllocator* allocator_;
   GrowableArray<StackMapEntry> stack_maps_;
-  GrowableArray<DexRegisterLocation> dex_register_maps_;
+  GrowableArray<DexRegisterLocation> dex_register_locations_;
   GrowableArray<InlineInfoEntry> inline_infos_;
   int stack_mask_max_;
   size_t number_of_stack_maps_with_inline_info_;
 
+  ART_FRIEND_TEST(StackMapTest, Test1);
+  ART_FRIEND_TEST(StackMapTest, Test2);
+  ART_FRIEND_TEST(StackMapTest, TestNonLiveDexRegisters);
+
   DISALLOW_COPY_AND_ASSIGN(StackMapStream);
 };