Simplify StackMapStream encoding.

Test: test-art-host-gtest-stack_map_test
Change-Id: Ife021d03e4e486043ec609f9af8673ace7bde497
diff --git a/compiler/exception_test.cc b/compiler/exception_test.cc
index b0e0337..fd17364 100644
--- a/compiler/exception_test.cc
+++ b/compiler/exception_test.cc
@@ -81,7 +81,9 @@
     stack_maps.BeginStackMapEntry(kDexPc, native_pc_offset);
     stack_maps.EndStackMapEntry();
     stack_maps.EndMethod();
-    const size_t stack_maps_size = stack_maps.PrepareForFillIn();
+    ScopedArenaVector<uint8_t> stack_map = stack_maps.Encode();
+
+    const size_t stack_maps_size = stack_map.size();
     const size_t header_size = sizeof(OatQuickMethodHeader);
     const size_t code_alignment = GetInstructionSetAlignment(kRuntimeISA);
 
@@ -90,9 +92,8 @@
     uint8_t* code_ptr =
       AlignUp(&fake_header_code_and_maps_[stack_maps_size + header_size], code_alignment);
 
-    MemoryRegion stack_maps_region(&fake_header_code_and_maps_[0], stack_maps_size);
-    stack_maps.FillInCodeInfo(stack_maps_region);
-    OatQuickMethodHeader method_header(code_ptr - stack_maps_region.begin(), code_size);
+    memcpy(&fake_header_code_and_maps_[0], stack_map.data(), stack_maps_size);
+    OatQuickMethodHeader method_header(code_ptr - fake_header_code_and_maps_.data(), code_size);
     static_assert(std::is_trivially_copyable<OatQuickMethodHeader>::value, "Cannot use memcpy");
     memcpy(code_ptr - header_size, &method_header, header_size);
     memcpy(code_ptr, fake_code_.data(), fake_code_.size());
diff --git a/compiler/optimizing/code_generator.cc b/compiler/optimizing/code_generator.cc
index f6a104b..21ecedd 100644
--- a/compiler/optimizing/code_generator.cc
+++ b/compiler/optimizing/code_generator.cc
@@ -961,12 +961,6 @@
 
 CodeGenerator::~CodeGenerator() {}
 
-void CodeGenerator::ComputeStackMapSize(size_t* stack_map_size) {
-  DCHECK(stack_map_size != nullptr);
-  StackMapStream* stack_map_stream = GetStackMapStream();
-  *stack_map_size = stack_map_stream->PrepareForFillIn();
-}
-
 size_t CodeGenerator::GetNumberOfJitRoots() const {
   DCHECK(code_generation_data_ != nullptr);
   return code_generation_data_->GetNumberOfJitRoots();
@@ -1033,13 +1027,12 @@
   }
 }
 
-void CodeGenerator::BuildStackMaps(MemoryRegion stack_map_region,
-                                   const DexFile::CodeItem* code_item_for_osr_check) {
-  StackMapStream* stack_map_stream = GetStackMapStream();
-  stack_map_stream->FillInCodeInfo(stack_map_region);
-  if (kIsDebugBuild && code_item_for_osr_check != nullptr) {
-    CheckLoopEntriesCanBeUsedForOsr(*graph_, CodeInfo(stack_map_region), *code_item_for_osr_check);
+ScopedArenaVector<uint8_t> CodeGenerator::BuildStackMaps(const DexFile::CodeItem* code_item) {
+  ScopedArenaVector<uint8_t> stack_map = GetStackMapStream()->Encode();
+  if (kIsDebugBuild && code_item != nullptr) {
+    CheckLoopEntriesCanBeUsedForOsr(*graph_, CodeInfo(stack_map.data()), *code_item);
   }
+  return stack_map;
 }
 
 void CodeGenerator::RecordPcInfo(HInstruction* instruction,
diff --git a/compiler/optimizing/code_generator.h b/compiler/optimizing/code_generator.h
index 3d58d29..e77d621 100644
--- a/compiler/optimizing/code_generator.h
+++ b/compiler/optimizing/code_generator.h
@@ -350,9 +350,7 @@
 
   void AddSlowPath(SlowPathCode* slow_path);
 
-  void BuildStackMaps(MemoryRegion stack_map_region,
-                      const DexFile::CodeItem* code_item_for_osr_check);
-  void ComputeStackMapSize(size_t* stack_map_size);
+  ScopedArenaVector<uint8_t> BuildStackMaps(const DexFile::CodeItem* code_item_for_osr_check);
   size_t GetNumberOfJitRoots() const;
 
   // Fills the `literals` array with literals collected during code generation.
diff --git a/compiler/optimizing/optimizing_compiler.cc b/compiler/optimizing/optimizing_compiler.cc
index b2733ee..c40cbcf 100644
--- a/compiler/optimizing/optimizing_compiler.cc
+++ b/compiler/optimizing/optimizing_compiler.cc
@@ -711,15 +711,7 @@
                                          CodeGenerator* codegen,
                                          const DexFile::CodeItem* code_item_for_osr_check) const {
   ArenaVector<linker::LinkerPatch> linker_patches = EmitAndSortLinkerPatches(codegen);
-  ArenaVector<uint8_t> stack_map(allocator->Adapter(kArenaAllocStackMaps));
-  ArenaVector<uint8_t> method_info(allocator->Adapter(kArenaAllocStackMaps));
-  size_t stack_map_size = 0;
-  size_t method_info_size = 0;
-  codegen->ComputeStackMapSize(&stack_map_size);
-  stack_map.resize(stack_map_size);
-  method_info.resize(method_info_size);
-  codegen->BuildStackMaps(MemoryRegion(stack_map.data(), stack_map.size()),
-                          code_item_for_osr_check);
+  ScopedArenaVector<uint8_t> stack_map = codegen->BuildStackMaps(code_item_for_osr_check);
 
   CompiledMethod* compiled_method = CompiledMethod::SwapAllocCompiledMethod(
       GetCompilerDriver(),
@@ -1097,22 +1089,19 @@
   return compiled_method;
 }
 
-static void CreateJniStackMap(ArenaStack* arena_stack,
-                              const JniCompiledMethod& jni_compiled_method,
-                              /* out */ ArenaVector<uint8_t>* stack_map) {
-  ScopedArenaAllocator allocator(arena_stack);
+static ScopedArenaVector<uint8_t> CreateJniStackMap(ScopedArenaAllocator* allocator,
+                                                    const JniCompiledMethod& jni_compiled_method) {
   // StackMapStream is quite large, so allocate it using the ScopedArenaAllocator
   // to stay clear of the frame size limit.
   std::unique_ptr<StackMapStream> stack_map_stream(
-      new (&allocator) StackMapStream(&allocator, jni_compiled_method.GetInstructionSet()));
+      new (allocator) StackMapStream(allocator, jni_compiled_method.GetInstructionSet()));
   stack_map_stream->BeginMethod(
       jni_compiled_method.GetFrameSize(),
       jni_compiled_method.GetCoreSpillMask(),
       jni_compiled_method.GetFpSpillMask(),
       /* num_dex_registers */ 0);
   stack_map_stream->EndMethod();
-  stack_map->resize(stack_map_stream->PrepareForFillIn());
-  stack_map_stream->FillInCodeInfo(MemoryRegion(stack_map->data(), stack_map->size()));
+  return stack_map_stream->Encode();
 }
 
 CompiledMethod* OptimizingCompiler::JniCompile(uint32_t access_flags,
@@ -1166,8 +1155,9 @@
       compiler_options, access_flags, method_idx, dex_file);
   MaybeRecordStat(compilation_stats_.get(), MethodCompilationStat::kCompiledNativeStub);
 
-  ArenaVector<uint8_t> stack_map(allocator.Adapter(kArenaAllocStackMaps));
-  CreateJniStackMap(&arena_stack, jni_compiled_method, &stack_map);
+  ScopedArenaAllocator stack_map_allocator(&arena_stack);  // Will hold the stack map.
+  ScopedArenaVector<uint8_t> stack_map = CreateJniStackMap(&stack_map_allocator,
+                                                           jni_compiled_method);
   return CompiledMethod::SwapAllocCompiledMethod(
       GetCompilerDriver(),
       jni_compiled_method.GetInstructionSet(),
@@ -1232,11 +1222,11 @@
     ScopedNullHandle<mirror::ObjectArray<mirror::Object>> roots;
     ArenaSet<ArtMethod*, std::less<ArtMethod*>> cha_single_implementation_list(
         allocator.Adapter(kArenaAllocCHA));
-    ArenaVector<uint8_t> stack_map(allocator.Adapter(kArenaAllocStackMaps));
     ArenaStack arena_stack(runtime->GetJitArenaPool());
     // StackMapStream is large and it does not fit into this frame, so we need helper method.
-    // TODO: Try to avoid the extra memory copy that results from this.
-    CreateJniStackMap(&arena_stack, jni_compiled_method, &stack_map);
+    ScopedArenaAllocator stack_map_allocator(&arena_stack);  // Will hold the stack map.
+    ScopedArenaVector<uint8_t> stack_map = CreateJniStackMap(&stack_map_allocator,
+                                                             jni_compiled_method);
     uint8_t* stack_map_data = nullptr;
     uint8_t* roots_data = nullptr;
     uint32_t data_size = code_cache->ReserveData(self,
@@ -1329,8 +1319,7 @@
     }
   }
 
-  size_t stack_map_size = 0;
-  codegen->ComputeStackMapSize(&stack_map_size);
+  ScopedArenaVector<uint8_t> stack_map = codegen->BuildStackMaps(code_item);
   size_t number_of_roots = codegen->GetNumberOfJitRoots();
   // We allocate an object array to ensure the JIT roots that we will collect in EmitJitRoots
   // will be visible by the GC between EmitLiterals and CommitCode. Once CommitCode is
@@ -1348,7 +1337,7 @@
   uint8_t* stack_map_data = nullptr;
   uint8_t* roots_data = nullptr;
   uint32_t data_size = code_cache->ReserveData(self,
-                                               stack_map_size,
+                                               stack_map.size(),
                                                number_of_roots,
                                                method,
                                                &stack_map_data,
@@ -1357,7 +1346,7 @@
     MaybeRecordStat(compilation_stats_.get(), MethodCompilationStat::kJitOutOfMemoryForCommit);
     return false;
   }
-  codegen->BuildStackMaps(MemoryRegion(stack_map_data, stack_map_size), code_item);
+  memcpy(stack_map_data, stack_map.data(), stack_map.size());
   codegen->EmitJitRoots(code_allocator.GetData(), roots, roots_data);
 
   const void* code = code_cache->CommitCode(
@@ -1398,7 +1387,7 @@
     info.code_address = code_address;
     info.code_size = code_allocator.GetMemory().size();
     info.frame_size_in_bytes = method_header->GetFrameSizeInBytes();
-    info.code_info = stack_map_size == 0 ? nullptr : stack_map_data;
+    info.code_info = stack_map.size() == 0 ? nullptr : stack_map_data;
     info.cfi = ArrayRef<const uint8_t>(*codegen->GetAssembler()->cfi().data());
     GenerateJitDebugInfo(method, info);
   }
diff --git a/compiler/optimizing/stack_map_stream.cc b/compiler/optimizing/stack_map_stream.cc
index b26b53f..8858cbe 100644
--- a/compiler/optimizing/stack_map_stream.cc
+++ b/compiler/optimizing/stack_map_stream.cc
@@ -57,6 +57,15 @@
 void StackMapStream::EndMethod() {
   DCHECK(in_method_) << "Mismatched Begin/End calls";
   in_method_ = false;
+
+  // Read the stack masks now. The compiler might have updated them.
+  for (size_t i = 0; i < lazy_stack_masks_.size(); i++) {
+    BitVector* stack_mask = lazy_stack_masks_[i];
+    if (stack_mask != nullptr && stack_mask->GetNumberOfBits() != 0) {
+      stack_maps_[i][StackMap::kStackMaskIndex] =
+          stack_masks_.Dedup(stack_mask->GetRawStorage(), stack_mask->GetNumberOfBits());
+    }
+  }
 }
 
 void StackMapStream::BeginStackMapEntry(uint32_t dex_pc,
@@ -281,19 +290,12 @@
   bit_table.Encode(out);
 }
 
-size_t StackMapStream::PrepareForFillIn() {
-  DCHECK_EQ(out_.size(), 0u);
+ScopedArenaVector<uint8_t> StackMapStream::Encode() {
+  DCHECK(in_stack_map_ == false) << "Mismatched Begin/End calls";
+  DCHECK(in_inline_info_ == false) << "Mismatched Begin/End calls";
 
-  // Read the stack masks now. The compiler might have updated them.
-  for (size_t i = 0; i < lazy_stack_masks_.size(); i++) {
-    BitVector* stack_mask = lazy_stack_masks_[i];
-    if (stack_mask != nullptr && stack_mask->GetNumberOfBits() != 0) {
-      stack_maps_[i][StackMap::kStackMaskIndex] =
-        stack_masks_.Dedup(stack_mask->GetRawStorage(), stack_mask->GetNumberOfBits());
-    }
-  }
-
-  BitMemoryWriter<ScopedArenaVector<uint8_t>> out(&out_);
+  ScopedArenaVector<uint8_t> buffer(allocator_->Adapter(kArenaAllocStackMapStream));
+  BitMemoryWriter<ScopedArenaVector<uint8_t>> out(&buffer);
   EncodeVarintBits(out, packed_frame_size_);
   EncodeVarintBits(out, core_spill_mask_);
   EncodeVarintBits(out, fp_spill_mask_);
@@ -307,20 +309,9 @@
   EncodeTable(out, dex_register_maps_);
   EncodeTable(out, dex_register_catalog_);
 
-  return out_.size();
-}
-
-void StackMapStream::FillInCodeInfo(MemoryRegion region) {
-  DCHECK(in_stack_map_ == false) << "Mismatched Begin/End calls";
-  DCHECK(in_inline_info_ == false) << "Mismatched Begin/End calls";
-  DCHECK_NE(0u, out_.size()) << "PrepareForFillIn not called before FillIn";
-  DCHECK_EQ(region.size(), out_.size());
-
-  region.CopyFromVector(0, out_);
-
   // Verify that we can load the CodeInfo and check some essentials.
-  CodeInfo code_info(region);
-  CHECK_EQ(code_info.Size(), out_.size());
+  CodeInfo code_info(buffer.data());
+  CHECK_EQ(code_info.Size(), buffer.size());
   CHECK_EQ(code_info.GetNumberOfStackMaps(), stack_maps_.size());
 
   // Verify all written data (usually only in debug builds).
@@ -329,6 +320,8 @@
       dcheck(code_info);
     }
   }
+
+  return buffer;
 }
 
 }  // namespace art
diff --git a/compiler/optimizing/stack_map_stream.h b/compiler/optimizing/stack_map_stream.h
index cd04ff0..df11709 100644
--- a/compiler/optimizing/stack_map_stream.h
+++ b/compiler/optimizing/stack_map_stream.h
@@ -37,7 +37,8 @@
 class StackMapStream : public DeletableArenaObject<kArenaAllocStackMapStream> {
  public:
   explicit StackMapStream(ScopedArenaAllocator* allocator, InstructionSet instruction_set)
-      : instruction_set_(instruction_set),
+      : allocator_(allocator),
+        instruction_set_(instruction_set),
         stack_maps_(allocator),
         inline_infos_(allocator),
         method_infos_(allocator),
@@ -46,7 +47,6 @@
         dex_register_masks_(allocator),
         dex_register_maps_(allocator),
         dex_register_catalog_(allocator),
-        out_(allocator->Adapter(kArenaAllocStackMapStream)),
         lazy_stack_masks_(allocator->Adapter(kArenaAllocStackMapStream)),
         current_stack_map_(),
         current_inline_infos_(allocator->Adapter(kArenaAllocStackMapStream)),
@@ -88,16 +88,16 @@
   uint32_t GetStackMapNativePcOffset(size_t i);
   void SetStackMapNativePcOffset(size_t i, uint32_t native_pc_offset);
 
-  // Prepares the stream to fill in a memory region. Must be called before FillIn.
-  // Returns the size (in bytes) needed to store this stream.
-  size_t PrepareForFillIn();
-  void FillInCodeInfo(MemoryRegion region);
+  // Encode all stack map data.
+  // The returned vector is allocated using the allocator passed to the StackMapStream.
+  ScopedArenaVector<uint8_t> Encode();
 
  private:
   static constexpr uint32_t kNoValue = -1;
 
   void CreateDexRegisterMap();
 
+  ScopedArenaAllocator* allocator_;
   const InstructionSet instruction_set_;
   uint32_t packed_frame_size_ = 0;
   uint32_t core_spill_mask_ = 0;
@@ -111,7 +111,6 @@
   BitmapTableBuilder dex_register_masks_;
   BitTableBuilder<MaskInfo> dex_register_maps_;
   BitTableBuilder<DexRegisterInfo> dex_register_catalog_;
-  ScopedArenaVector<uint8_t> out_;
 
   ScopedArenaVector<BitVector*> lazy_stack_masks_;
 
diff --git a/compiler/optimizing/stack_map_test.cc b/compiler/optimizing/stack_map_test.cc
index 16a9216..a281bb3 100644
--- a/compiler/optimizing/stack_map_test.cc
+++ b/compiler/optimizing/stack_map_test.cc
@@ -62,12 +62,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo code_info(region);
+  CodeInfo code_info(memory.data());
   ASSERT_EQ(1u, code_info.GetNumberOfStackMaps());
 
   uint32_t number_of_catalog_entries = code_info.GetNumberOfLocationCatalogEntries();
@@ -151,12 +148,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo code_info(region);
+  CodeInfo code_info(memory.data());
   ASSERT_EQ(4u, code_info.GetNumberOfStackMaps());
 
   uint32_t number_of_catalog_entries = code_info.GetNumberOfLocationCatalogEntries();
@@ -324,12 +318,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo code_info(region);
+  CodeInfo code_info(memory.data());
   ASSERT_EQ(1u, code_info.GetNumberOfStackMaps());
 
   uint32_t number_of_catalog_entries = code_info.GetNumberOfLocationCatalogEntries();
@@ -382,12 +373,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo code_info(region);
+  CodeInfo code_info(memory.data());
   ASSERT_EQ(1u, code_info.GetNumberOfStackMaps());
 
   uint32_t number_of_catalog_entries = code_info.GetNumberOfLocationCatalogEntries();
@@ -444,12 +432,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo ci(region);
+  CodeInfo ci(memory.data());
 
   // Verify first stack map.
   StackMap sm0 = ci.GetStackMapAt(0);
@@ -495,12 +480,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo code_info(region);
+  CodeInfo code_info(memory.data());
   ASSERT_EQ(2u, code_info.GetNumberOfStackMaps());
 
   uint32_t number_of_catalog_entries = code_info.GetNumberOfLocationCatalogEntries();
@@ -597,12 +579,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo ci(region);
+  CodeInfo ci(memory.data());
 
   {
     // Verify first stack map.
@@ -744,12 +723,9 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  size_t size = stream.PrepareForFillIn();
-  void* memory = allocator.Alloc(size, kArenaAllocMisc);
-  MemoryRegion region(memory, size);
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
-  CodeInfo code_info(region);
+  CodeInfo code_info(memory.data());
   ASSERT_EQ(2u, code_info.GetNumberOfStackMaps());
 
   StackMap stack_map1 = code_info.GetStackMapForNativePcOffset(4 * kPcAlign);
@@ -771,9 +747,7 @@
   stream.EndStackMapEntry();
 
   stream.EndMethod();
-  std::vector<uint8_t> memory(stream.PrepareForFillIn());
-  MemoryRegion region(memory.data(), memory.size());
-  stream.FillInCodeInfo(region);
+  ScopedArenaVector<uint8_t> memory = stream.Encode();
 
   std::vector<uint8_t> out;
   CodeInfo::DedupeMap dedupe_map;
diff --git a/runtime/stack_map.h b/runtime/stack_map.h
index 0b1c4f9..acde3e3 100644
--- a/runtime/stack_map.h
+++ b/runtime/stack_map.h
@@ -280,10 +280,6 @@
     Decode(reinterpret_cast<const uint8_t*>(data), flags);
   }
 
-  explicit CodeInfo(MemoryRegion region) : CodeInfo(region.begin()) {
-    DCHECK_EQ(Size(), region.size());
-  }
-
   explicit CodeInfo(const OatQuickMethodHeader* header, DecodeFlags flags = DecodeFlags::Default);
 
   size_t Size() const {