Merge "Code cleanup to avoid CompilerDriver abstractions in JIT."
diff --git a/compiler/compiler.h b/compiler/compiler.h
index 8788dc1..3a9ce1b 100644
--- a/compiler/compiler.h
+++ b/compiler/compiler.h
@@ -22,6 +22,10 @@
 
 namespace art {
 
+namespace jit {
+  class JitCodeCache;
+}
+
 class ArtMethod;
 class Backend;
 struct CompilationUnit;
@@ -58,6 +62,13 @@
                                      uint32_t method_idx,
                                      const DexFile& dex_file) const = 0;
 
+  virtual bool JitCompile(Thread* self ATTRIBUTE_UNUSED,
+                          jit::JitCodeCache* code_cache ATTRIBUTE_UNUSED,
+                          ArtMethod* method ATTRIBUTE_UNUSED)
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    return false;
+  }
+
   virtual uintptr_t GetEntryPointOf(ArtMethod* method) const
      SHARED_REQUIRES(Locks::mutator_lock_) = 0;
 
diff --git a/compiler/driver/compiler_driver.cc b/compiler/driver/compiler_driver.cc
index d055b37..527d861 100644
--- a/compiler/driver/compiler_driver.cc
+++ b/compiler/driver/compiler_driver.cc
@@ -559,7 +559,7 @@
     }
   } else if ((access_flags & kAccAbstract) != 0) {
     // Abstract methods don't have code.
-  } else if (Runtime::Current()->IsAotCompiler()) {
+  } else {
     const VerifiedMethod* verified_method =
         driver->GetVerificationResults()->GetVerifiedMethod(method_ref);
     bool compile = compilation_enabled &&
@@ -598,13 +598,6 @@
               ? dex_to_dex_compilation_level
               : optimizer::DexToDexCompilationLevel::kRequired);
     }
-  } else {
-    // This is for the JIT compiler, which has already ensured the class is verified.
-    // We can go straight to compiling.
-    DCHECK(Runtime::Current()->UseJit());
-    compiled_method = driver->GetCompiler()->Compile(code_item, access_flags, invoke_type,
-                                                     class_def_idx, method_idx, class_loader,
-                                                     dex_file, dex_cache);
   }
   if (kTimeCompileMethod) {
     uint64_t duration_ns = NanoTime() - start_ns;
@@ -696,42 +689,6 @@
   self->GetJniEnv()->DeleteGlobalRef(jclass_loader);
 }
 
-CompiledMethod* CompilerDriver::CompileArtMethod(Thread* self, ArtMethod* method) {
-  DCHECK_EQ(method,
-            method->GetInterfaceMethodIfProxy(
-                Runtime::Current()->GetClassLinker()->GetImagePointerSize()));
-  const uint32_t method_idx = method->GetDexMethodIndex();
-  const uint32_t access_flags = method->GetAccessFlags();
-  const InvokeType invoke_type = method->GetInvokeType();
-  StackHandleScope<2> hs(self);
-  Handle<mirror::ClassLoader> class_loader(hs.NewHandle(
-      method->GetDeclaringClass()->GetClassLoader()));
-  Handle<mirror::DexCache> dex_cache(hs.NewHandle(method->GetDexCache()));
-  jobject jclass_loader = class_loader.ToJObject();
-  const DexFile* dex_file = method->GetDexFile();
-  const uint16_t class_def_idx = method->GetClassDefIndex();
-  const DexFile::ClassDef& class_def = dex_file->GetClassDef(class_def_idx);
-  optimizer::DexToDexCompilationLevel dex_to_dex_compilation_level =
-      GetDexToDexCompilationLevel(self, *this, class_loader, *dex_file, class_def);
-  const DexFile::CodeItem* code_item = dex_file->GetCodeItem(method->GetCodeItemOffset());
-  // Go to native so that we don't block GC during compilation.
-  ScopedThreadSuspension sts(self, kNative);
-  CompileMethod(self,
-                this,
-                code_item,
-                access_flags,
-                invoke_type,
-                class_def_idx,
-                method_idx,
-                jclass_loader,
-                *dex_file,
-                dex_to_dex_compilation_level,
-                true,
-                dex_cache);
-  auto* compiled_method = GetCompiledMethod(MethodReference(dex_file, method_idx));
-  return compiled_method;
-}
-
 void CompilerDriver::Resolve(jobject class_loader, const std::vector<const DexFile*>& dex_files,
                              ThreadPool* thread_pool, TimingLogger* timings) {
   for (size_t i = 0; i != dex_files.size(); ++i) {
diff --git a/compiler/driver/compiler_driver.h b/compiler/driver/compiler_driver.h
index 4ed4dc6..07cd077 100644
--- a/compiler/driver/compiler_driver.h
+++ b/compiler/driver/compiler_driver.h
@@ -119,9 +119,6 @@
                   TimingLogger* timings)
       REQUIRES(!Locks::mutator_lock_, !compiled_classes_lock_);
 
-  CompiledMethod* CompileArtMethod(Thread* self, ArtMethod*)
-      SHARED_REQUIRES(Locks::mutator_lock_) REQUIRES(!compiled_methods_lock_) WARN_UNUSED;
-
   // Compile a single Method.
   void CompileOne(Thread* self, ArtMethod* method, TimingLogger* timings)
       SHARED_REQUIRES(Locks::mutator_lock_)
diff --git a/compiler/jit/jit_compiler.cc b/compiler/jit/jit_compiler.cc
index d520208..5f4f472 100644
--- a/compiler/jit/jit_compiler.cc
+++ b/compiler/jit/jit_compiler.cc
@@ -189,13 +189,14 @@
   }
 
   // Do the compilation.
-  CompiledMethod* compiled_method = nullptr;
+  JitCodeCache* const code_cache = runtime->GetJit()->GetCodeCache();
+  bool success = false;
   {
     TimingLogger::ScopedTiming t2("Compiling", &logger);
     // If we get a request to compile a proxy method, we pass the actual Java method
     // of that proxy method, as the compiler does not expect a proxy method.
     ArtMethod* method_to_compile = method->GetInterfaceMethodIfProxy(sizeof(void*));
-    compiled_method = compiler_driver_->CompileArtMethod(self, method_to_compile);
+    success = compiler_driver_->GetCompiler()->JitCompile(self, code_cache, method_to_compile);
   }
 
   // Trim maps to reduce memory usage.
@@ -205,105 +206,14 @@
     runtime->GetArenaPool()->TrimMaps();
   }
 
-  // Check if we failed compiling.
-  if (compiled_method == nullptr) {
-    return false;
-  }
-
   total_time_ += NanoTime() - start_time;
-  bool result = false;
-  const void* code = runtime->GetClassLinker()->GetOatMethodQuickCodeFor(method);
-
-  if (code != nullptr) {
-    // Already have some compiled code, just use this instead of linking.
-    // TODO: Fix recompilation.
-    method->SetEntryPointFromQuickCompiledCode(code);
-    result = true;
-  } else {
-    TimingLogger::ScopedTiming t2("LinkCode", &logger);
-    if (AddToCodeCache(method, compiled_method)) {
-      result = true;
-    }
-  }
-
-  // Remove the compiled method to save memory.
-  compiler_driver_->RemoveCompiledMethod(
-      MethodReference(h_class->GetDexCache()->GetDexFile(), method->GetDexMethodIndex()));
   runtime->GetJit()->AddTimingLogger(logger);
-  return result;
+  return success;
 }
 
 CompilerCallbacks* JitCompiler::GetCompilerCallbacks() const {
   return callbacks_.get();
 }
 
-bool JitCompiler::AddToCodeCache(ArtMethod* method,
-                                 const CompiledMethod* compiled_method) {
-  Runtime* runtime = Runtime::Current();
-  JitCodeCache* const code_cache = runtime->GetJit()->GetCodeCache();
-  auto const quick_code = compiled_method->GetQuickCode();
-  if (quick_code.empty()) {
-    return false;
-  }
-  const auto code_size = quick_code.size();
-  Thread* const self = Thread::Current();
-  auto const mapping_table = compiled_method->GetMappingTable();
-  auto const vmap_table = compiled_method->GetVmapTable();
-  auto const gc_map = compiled_method->GetGcMap();
-  uint8_t* mapping_table_ptr = nullptr;
-  uint8_t* vmap_table_ptr = nullptr;
-  uint8_t* gc_map_ptr = nullptr;
-
-  if (!mapping_table.empty()) {
-    // Write out pre-header stuff.
-    mapping_table_ptr = code_cache->AddDataArray(
-        self, mapping_table.data(), mapping_table.data() + mapping_table.size());
-    if (mapping_table_ptr == nullptr) {
-      return false;  // Out of data cache.
-    }
-  }
-
-  if (!vmap_table.empty()) {
-    vmap_table_ptr = code_cache->AddDataArray(
-        self, vmap_table.data(), vmap_table.data() + vmap_table.size());
-    if (vmap_table_ptr == nullptr) {
-      return false;  // Out of data cache.
-    }
-  }
-
-  if (!gc_map.empty()) {
-    gc_map_ptr = code_cache->AddDataArray(
-        self, gc_map.data(), gc_map.data() + gc_map.size());
-    if (gc_map_ptr == nullptr) {
-      return false;  // Out of data cache.
-    }
-  }
-
-  uint8_t* const code = code_cache->CommitCode(self,
-                                               method,
-                                               mapping_table_ptr,
-                                               vmap_table_ptr,
-                                               gc_map_ptr,
-                                               compiled_method->GetFrameSizeInBytes(),
-                                               compiled_method->GetCoreSpillMask(),
-                                               compiled_method->GetFpSpillMask(),
-                                               compiled_method->GetQuickCode().data(),
-                                               compiled_method->GetQuickCode().size());
-
-  if (code == nullptr) {
-    return false;
-  }
-
-  const size_t thumb_offset = compiled_method->CodeDelta();
-  const uint32_t code_offset = sizeof(OatQuickMethodHeader) + thumb_offset;
-  VLOG(jit)
-      << "JIT added "
-      << PrettyMethod(method) << "@" << method
-      << " ccache_size=" << PrettySize(code_cache->CodeCacheSize()) << ": "
-      << reinterpret_cast<void*>(code + code_offset)
-      << "," << reinterpret_cast<void*>(code + code_offset + code_size);
-  return true;
-}
-
 }  // namespace jit
 }  // namespace art
diff --git a/compiler/optimizing/code_generator.cc b/compiler/optimizing/code_generator.cc
index a1bb5e0..ce92470 100644
--- a/compiler/optimizing/code_generator.cc
+++ b/compiler/optimizing/code_generator.cc
@@ -42,7 +42,7 @@
 
 #include "compiled_method.h"
 #include "dex/verified_method.h"
-#include "driver/dex_compilation_unit.h"
+#include "driver/compiler_driver.h"
 #include "gc_map_builder.h"
 #include "graph_visualizer.h"
 #include "intrinsics.h"
@@ -787,9 +787,10 @@
 }
 
 void CodeGenerator::BuildNativeGCMap(
-    ArenaVector<uint8_t>* data, const DexCompilationUnit& dex_compilation_unit) const {
+    ArenaVector<uint8_t>* data, const CompilerDriver& compiler_driver) const {
   const std::vector<uint8_t>& gc_map_raw =
-      dex_compilation_unit.GetVerifiedMethod()->GetDexGcMap();
+      compiler_driver.GetVerifiedMethod(&GetGraph()->GetDexFile(), GetGraph()->GetMethodIdx())
+          ->GetDexGcMap();
   verifier::DexPcToReferenceMap dex_gc_map(&(gc_map_raw)[0]);
 
   uint32_t max_native_offset = stack_map_stream_.ComputeMaxNativePcOffset();
@@ -911,19 +912,22 @@
   vmap_encoder.PushBackUnsigned(VmapTable::kAdjustedFpMarker);
 }
 
-void CodeGenerator::BuildStackMaps(ArenaVector<uint8_t>* data) {
-  uint32_t size = stack_map_stream_.PrepareForFillIn();
-  data->resize(size);
-  MemoryRegion region(data->data(), size);
+size_t CodeGenerator::ComputeStackMapsSize() {
+  return stack_map_stream_.PrepareForFillIn();
+}
+
+void CodeGenerator::BuildStackMaps(MemoryRegion region) {
   stack_map_stream_.FillIn(region);
 }
 
 void CodeGenerator::RecordNativeDebugInfo(uint32_t dex_pc,
                                           uintptr_t native_pc_begin,
                                           uintptr_t native_pc_end) {
-  if (src_map_ != nullptr && dex_pc != kNoDexPc && native_pc_begin != native_pc_end) {
-    src_map_->push_back(SrcMapElem({static_cast<uint32_t>(native_pc_begin),
-                                    static_cast<int32_t>(dex_pc)}));
+  if (compiler_options_.GetGenerateDebugInfo() &&
+      dex_pc != kNoDexPc &&
+      native_pc_begin != native_pc_end) {
+    src_map_.push_back(SrcMapElem({static_cast<uint32_t>(native_pc_begin),
+                                   static_cast<int32_t>(dex_pc)}));
   }
 }
 
diff --git a/compiler/optimizing/code_generator.h b/compiler/optimizing/code_generator.h
index 47b6f30..a92014d 100644
--- a/compiler/optimizing/code_generator.h
+++ b/compiler/optimizing/code_generator.h
@@ -22,6 +22,7 @@
 #include "base/arena_containers.h"
 #include "base/arena_object.h"
 #include "base/bit_field.h"
+#include "compiled_method.h"
 #include "driver/compiler_options.h"
 #include "globals.h"
 #include "graph_visualizer.h"
@@ -51,13 +52,9 @@
 
 class Assembler;
 class CodeGenerator;
-class DexCompilationUnit;
+class CompilerDriver;
 class LinkerPatch;
 class ParallelMoveResolver;
-class SrcMapElem;
-template <class Alloc>
-class SrcMap;
-using DefaultSrcMap = SrcMap<std::allocator<SrcMapElem>>;
 
 class CodeAllocator {
  public:
@@ -284,13 +281,12 @@
     slow_paths_.push_back(slow_path);
   }
 
-  void SetSrcMap(DefaultSrcMap* src_map) { src_map_ = src_map; }
-
   void BuildMappingTable(ArenaVector<uint8_t>* vector) const;
   void BuildVMapTable(ArenaVector<uint8_t>* vector) const;
   void BuildNativeGCMap(
-      ArenaVector<uint8_t>* vector, const DexCompilationUnit& dex_compilation_unit) const;
-  void BuildStackMaps(ArenaVector<uint8_t>* vector);
+      ArenaVector<uint8_t>* vector, const CompilerDriver& compiler_driver) const;
+  void BuildStackMaps(MemoryRegion region);
+  size_t ComputeStackMapsSize();
 
   bool IsBaseline() const {
     return is_baseline_;
@@ -446,6 +442,10 @@
   // Copy the result of a call into the given target.
   virtual void MoveFromReturnRegister(Location trg, Primitive::Type type) = 0;
 
+  const ArenaVector<SrcMapElem>& GetSrcMappingTable() const {
+    return src_map_;
+  }
+
  protected:
   // Method patch info used for recording locations of required linker patches and
   // target methods. The target method can be used for various purposes, whether for
@@ -488,7 +488,7 @@
         stats_(stats),
         graph_(graph),
         compiler_options_(compiler_options),
-        src_map_(nullptr),
+        src_map_(graph->GetArena()->Adapter(kArenaAllocCodeGenerator)),
         slow_paths_(graph->GetArena()->Adapter(kArenaAllocCodeGenerator)),
         current_block_index_(0),
         is_leaf_(true),
@@ -602,7 +602,7 @@
   const CompilerOptions& compiler_options_;
 
   // Native to dex_pc map used for native debugging/profiling tools.
-  DefaultSrcMap* src_map_;
+  ArenaVector<SrcMapElem> src_map_;
   ArenaVector<SlowPathCode*> slow_paths_;
 
   // The current block index in `block_order_` of the block
diff --git a/compiler/optimizing/optimizing_compiler.cc b/compiler/optimizing/optimizing_compiler.cc
index 8cb2cfc..7e3c5e6 100644
--- a/compiler/optimizing/optimizing_compiler.cc
+++ b/compiler/optimizing/optimizing_compiler.cc
@@ -56,6 +56,7 @@
 #include "inliner.h"
 #include "instruction_simplifier.h"
 #include "intrinsics.h"
+#include "jit/jit_code_cache.h"
 #include "licm.h"
 #include "jni/quick/jni_compiler.h"
 #include "load_store_elimination.h"
@@ -258,15 +259,6 @@
                           const DexFile& dex_file,
                           Handle<mirror::DexCache> dex_cache) const OVERRIDE;
 
-  CompiledMethod* TryCompile(const DexFile::CodeItem* code_item,
-                             uint32_t access_flags,
-                             InvokeType invoke_type,
-                             uint16_t class_def_idx,
-                             uint32_t method_idx,
-                             jobject class_loader,
-                             const DexFile& dex_file,
-                             Handle<mirror::DexCache> dex_cache) const;
-
   CompiledMethod* JniCompile(uint32_t access_flags,
                              uint32_t method_idx,
                              const DexFile& dex_file) const OVERRIDE {
@@ -291,23 +283,45 @@
     }
   }
 
+  bool JitCompile(Thread* self, jit::JitCodeCache* code_cache, ArtMethod* method)
+      OVERRIDE
+      SHARED_REQUIRES(Locks::mutator_lock_);
+
  private:
   // Whether we should run any optimization or register allocation. If false, will
   // just run the code generation after the graph was built.
   const bool run_optimizations_;
 
-  // Optimize and compile `graph`.
-  CompiledMethod* CompileOptimized(HGraph* graph,
-                                   CodeGenerator* codegen,
-                                   CompilerDriver* driver,
-                                   const DexCompilationUnit& dex_compilation_unit,
-                                   PassObserver* pass_observer) const;
+  // Create a 'CompiledMethod' for an optimized graph.
+  CompiledMethod* EmitOptimized(ArenaAllocator* arena,
+                                CodeVectorAllocator* code_allocator,
+                                CodeGenerator* codegen,
+                                CompilerDriver* driver) const;
 
-  // Just compile without doing optimizations.
-  CompiledMethod* CompileBaseline(CodeGenerator* codegen,
-                                  CompilerDriver* driver,
-                                  const DexCompilationUnit& dex_compilation_unit,
-                                  PassObserver* pass_observer) const;
+  // Create a 'CompiledMethod' for a non-optimized graph.
+  CompiledMethod* EmitBaseline(ArenaAllocator* arena,
+                               CodeVectorAllocator* code_allocator,
+                               CodeGenerator* codegen,
+                               CompilerDriver* driver) const;
+
+  // Try compiling a method and return the code generator used for
+  // compiling it.
+  // This method:
+  // 1) Builds the graph. Returns null if it failed to build it.
+  // 2) If `run_optimizations_` is set:
+  //    2.1) Transform the graph to SSA. Returns null if it failed.
+  //    2.2) Run optimizations on the graph, including register allocator.
+  // 3) Generate code with the `code_allocator` provided.
+  CodeGenerator* TryCompile(ArenaAllocator* arena,
+                            CodeVectorAllocator* code_allocator,
+                            const DexFile::CodeItem* code_item,
+                            uint32_t access_flags,
+                            InvokeType invoke_type,
+                            uint16_t class_def_idx,
+                            uint32_t method_idx,
+                            jobject class_loader,
+                            const DexFile& dex_file,
+                            Handle<mirror::DexCache> dex_cache) const;
 
   std::unique_ptr<OptimizingCompilerStats> compilation_stats_;
 
@@ -446,13 +460,32 @@
   }
 }
 
+NO_INLINE  // Avoid increasing caller's frame size by large stack-allocated objects.
+static void AllocateRegisters(HGraph* graph,
+                              CodeGenerator* codegen,
+                              PassObserver* pass_observer) {
+  PrepareForRegisterAllocation(graph).Run();
+  SsaLivenessAnalysis liveness(graph, codegen);
+  {
+    PassScope scope(SsaLivenessAnalysis::kLivenessPassName, pass_observer);
+    liveness.Analyze();
+  }
+  {
+    PassScope scope(RegisterAllocator::kRegisterAllocatorPassName, pass_observer);
+    RegisterAllocator(graph->GetArena(), codegen, liveness).AllocateRegisters();
+  }
+}
+
 static void RunOptimizations(HGraph* graph,
                              CodeGenerator* codegen,
                              CompilerDriver* driver,
                              OptimizingCompilerStats* stats,
                              const DexCompilationUnit& dex_compilation_unit,
-                             PassObserver* pass_observer,
-                             StackHandleScopeCollection* handles) {
+                             PassObserver* pass_observer) {
+  ScopedObjectAccess soa(Thread::Current());
+  StackHandleScopeCollection handles(soa.Self());
+  ScopedThreadSuspension sts(soa.Self(), kNative);
+
   ArenaAllocator* arena = graph->GetArena();
   HDeadCodeElimination* dce1 = new (arena) HDeadCodeElimination(
       graph, stats, HDeadCodeElimination::kInitialDeadCodeEliminationPassName);
@@ -469,7 +502,7 @@
   HInductionVarAnalysis* induction = new (arena) HInductionVarAnalysis(graph);
   BoundsCheckElimination* bce = new (arena) BoundsCheckElimination(graph, induction);
   ReferenceTypePropagation* type_propagation =
-      new (arena) ReferenceTypePropagation(graph, handles);
+      new (arena) ReferenceTypePropagation(graph, &handles);
   HSharpening* sharpening = new (arena) HSharpening(graph, codegen, dex_compilation_unit, driver);
   InstructionSimplifier* simplify2 = new (arena) InstructionSimplifier(
       graph, stats, "instruction_simplifier_after_types");
@@ -492,7 +525,7 @@
 
   RunOptimizations(optimizations1, arraysize(optimizations1), pass_observer);
 
-  MaybeRunInliner(graph, codegen, driver, stats, dex_compilation_unit, pass_observer, handles);
+  MaybeRunInliner(graph, codegen, driver, stats, dex_compilation_unit, pass_observer, &handles);
 
   // TODO: Update passes incompatible with try/catch so we have the same
   //       pipeline for all methods.
@@ -532,6 +565,7 @@
   }
 
   RunArchOptimizations(driver->GetInstructionSet(), graph, stats, pass_observer);
+  AllocateRegisters(graph, codegen, pass_observer);
 }
 
 // The stack map we generate must be 4-byte aligned on ARM. Since existing
@@ -545,22 +579,6 @@
   return ArrayRef<const uint8_t>(vector);
 }
 
-NO_INLINE  // Avoid increasing caller's frame size by large stack-allocated objects.
-static void AllocateRegisters(HGraph* graph,
-                              CodeGenerator* codegen,
-                              PassObserver* pass_observer) {
-  PrepareForRegisterAllocation(graph).Run();
-  SsaLivenessAnalysis liveness(graph, codegen);
-  {
-    PassScope scope(SsaLivenessAnalysis::kLivenessPassName, pass_observer);
-    liveness.Analyze();
-  }
-  {
-    PassScope scope(RegisterAllocator::kRegisterAllocatorPassName, pass_observer);
-    RegisterAllocator(graph->GetArena(), codegen, liveness).AllocateRegisters();
-  }
-}
-
 static ArenaVector<LinkerPatch> EmitAndSortLinkerPatches(CodeGenerator* codegen) {
   ArenaVector<LinkerPatch> linker_patches(codegen->GetGraph()->GetArena()->Adapter());
   codegen->EmitLinkerPatches(&linker_patches);
@@ -574,74 +592,42 @@
   return linker_patches;
 }
 
-CompiledMethod* OptimizingCompiler::CompileOptimized(HGraph* graph,
-                                                     CodeGenerator* codegen,
-                                                     CompilerDriver* compiler_driver,
-                                                     const DexCompilationUnit& dex_compilation_unit,
-                                                     PassObserver* pass_observer) const {
-  ScopedObjectAccess soa(Thread::Current());
-  StackHandleScopeCollection handles(soa.Self());
-  soa.Self()->TransitionFromRunnableToSuspended(kNative);
-  RunOptimizations(graph,
-                   codegen,
-                   compiler_driver,
-                   compilation_stats_.get(),
-                   dex_compilation_unit,
-                   pass_observer,
-                   &handles);
-
-  AllocateRegisters(graph, codegen, pass_observer);
-
-  ArenaAllocator* arena = graph->GetArena();
-  CodeVectorAllocator allocator(arena);
-  DefaultSrcMap src_mapping_table;
-  codegen->SetSrcMap(compiler_driver->GetCompilerOptions().GetGenerateDebugInfo()
-                         ? &src_mapping_table
-                         : nullptr);
-  codegen->CompileOptimized(&allocator);
-
+CompiledMethod* OptimizingCompiler::EmitOptimized(ArenaAllocator* arena,
+                                                  CodeVectorAllocator* code_allocator,
+                                                  CodeGenerator* codegen,
+                                                  CompilerDriver* compiler_driver) const {
   ArenaVector<LinkerPatch> linker_patches = EmitAndSortLinkerPatches(codegen);
-
   ArenaVector<uint8_t> stack_map(arena->Adapter(kArenaAllocStackMaps));
-  codegen->BuildStackMaps(&stack_map);
+  stack_map.resize(codegen->ComputeStackMapsSize());
+  codegen->BuildStackMaps(MemoryRegion(stack_map.data(), stack_map.size()));
 
   MaybeRecordStat(MethodCompilationStat::kCompiledOptimized);
 
   CompiledMethod* compiled_method = CompiledMethod::SwapAllocCompiledMethod(
       compiler_driver,
       codegen->GetInstructionSet(),
-      ArrayRef<const uint8_t>(allocator.GetMemory()),
+      ArrayRef<const uint8_t>(code_allocator->GetMemory()),
       // Follow Quick's behavior and set the frame size to zero if it is
       // considered "empty" (see the definition of
       // art::CodeGenerator::HasEmptyFrame).
       codegen->HasEmptyFrame() ? 0 : codegen->GetFrameSize(),
       codegen->GetCoreSpillMask(),
       codegen->GetFpuSpillMask(),
-      ArrayRef<const SrcMapElem>(src_mapping_table),
+      ArrayRef<const SrcMapElem>(codegen->GetSrcMappingTable()),
       ArrayRef<const uint8_t>(),  // mapping_table.
       ArrayRef<const uint8_t>(stack_map),
       ArrayRef<const uint8_t>(),  // native_gc_map.
       ArrayRef<const uint8_t>(*codegen->GetAssembler()->cfi().data()),
       ArrayRef<const LinkerPatch>(linker_patches));
-  pass_observer->DumpDisassembly();
 
-  soa.Self()->TransitionFromSuspendedToRunnable();
   return compiled_method;
 }
 
-CompiledMethod* OptimizingCompiler::CompileBaseline(
+CompiledMethod* OptimizingCompiler::EmitBaseline(
+    ArenaAllocator* arena,
+    CodeVectorAllocator* code_allocator,
     CodeGenerator* codegen,
-    CompilerDriver* compiler_driver,
-    const DexCompilationUnit& dex_compilation_unit,
-    PassObserver* pass_observer) const {
-  ArenaAllocator* arena = codegen->GetGraph()->GetArena();
-  CodeVectorAllocator allocator(arena);
-  DefaultSrcMap src_mapping_table;
-  codegen->SetSrcMap(compiler_driver->GetCompilerOptions().GetGenerateDebugInfo()
-                         ? &src_mapping_table
-                         : nullptr);
-  codegen->CompileBaseline(&allocator);
-
+    CompilerDriver* compiler_driver) const {
   ArenaVector<LinkerPatch> linker_patches = EmitAndSortLinkerPatches(codegen);
 
   ArenaVector<uint8_t> mapping_table(arena->Adapter(kArenaAllocBaselineMaps));
@@ -649,37 +635,38 @@
   ArenaVector<uint8_t> vmap_table(arena->Adapter(kArenaAllocBaselineMaps));
   codegen->BuildVMapTable(&vmap_table);
   ArenaVector<uint8_t> gc_map(arena->Adapter(kArenaAllocBaselineMaps));
-  codegen->BuildNativeGCMap(&gc_map, dex_compilation_unit);
+  codegen->BuildNativeGCMap(&gc_map, *compiler_driver);
 
   MaybeRecordStat(MethodCompilationStat::kCompiledBaseline);
   CompiledMethod* compiled_method = CompiledMethod::SwapAllocCompiledMethod(
       compiler_driver,
       codegen->GetInstructionSet(),
-      ArrayRef<const uint8_t>(allocator.GetMemory()),
+      ArrayRef<const uint8_t>(code_allocator->GetMemory()),
       // Follow Quick's behavior and set the frame size to zero if it is
       // considered "empty" (see the definition of
       // art::CodeGenerator::HasEmptyFrame).
       codegen->HasEmptyFrame() ? 0 : codegen->GetFrameSize(),
       codegen->GetCoreSpillMask(),
       codegen->GetFpuSpillMask(),
-      ArrayRef<const SrcMapElem>(src_mapping_table),
+      ArrayRef<const SrcMapElem>(codegen->GetSrcMappingTable()),
       AlignVectorSize(mapping_table),
       AlignVectorSize(vmap_table),
       AlignVectorSize(gc_map),
       ArrayRef<const uint8_t>(*codegen->GetAssembler()->cfi().data()),
       ArrayRef<const LinkerPatch>(linker_patches));
-  pass_observer->DumpDisassembly();
   return compiled_method;
 }
 
-CompiledMethod* OptimizingCompiler::TryCompile(const DexFile::CodeItem* code_item,
-                                               uint32_t access_flags,
-                                               InvokeType invoke_type,
-                                               uint16_t class_def_idx,
-                                               uint32_t method_idx,
-                                               jobject class_loader,
-                                               const DexFile& dex_file,
-                                               Handle<mirror::DexCache> dex_cache) const {
+CodeGenerator* OptimizingCompiler::TryCompile(ArenaAllocator* arena,
+                                              CodeVectorAllocator* code_allocator,
+                                              const DexFile::CodeItem* code_item,
+                                              uint32_t access_flags,
+                                              InvokeType invoke_type,
+                                              uint16_t class_def_idx,
+                                              uint32_t method_idx,
+                                              jobject class_loader,
+                                              const DexFile& dex_file,
+                                              Handle<mirror::DexCache> dex_cache) const {
   std::string method_name = PrettyMethod(method_idx, dex_file);
   MaybeRecordStat(MethodCompilationStat::kAttemptCompilation);
   CompilerDriver* compiler_driver = GetCompilerDriver();
@@ -721,13 +708,10 @@
       && compiler_driver->RequiresConstructorBarrier(Thread::Current(),
                                                      dex_compilation_unit.GetDexFile(),
                                                      dex_compilation_unit.GetClassDefIndex());
-  ArenaAllocator arena(Runtime::Current()->GetArenaPool());
-  HGraph* graph = new (&arena) HGraph(
-      &arena, dex_file, method_idx, requires_barrier, compiler_driver->GetInstructionSet(),
+  HGraph* graph = new (arena) HGraph(
+      arena, dex_file, method_idx, requires_barrier, compiler_driver->GetInstructionSet(),
       kInvalidInvokeType, compiler_driver->GetCompilerOptions().GetDebuggable());
 
-  bool shouldOptimize = method_name.find("$opt$reg$") != std::string::npos && run_optimizations_;
-
   std::unique_ptr<CodeGenerator> codegen(
       CodeGenerator::Create(graph,
                             instruction_set,
@@ -779,16 +763,8 @@
     }
   }
 
-  bool can_allocate_registers = RegisterAllocator::CanAllocateRegistersFor(*graph, instruction_set);
-
-  // `run_optimizations_` is set explicitly (either through a compiler filter
-  // or the debuggable flag). If it is set, we can run baseline. Otherwise, we fall back
-  // to Quick.
-  bool can_use_baseline = !run_optimizations_ && builder.CanUseBaselineForStringInit();
-  CompiledMethod* compiled_method = nullptr;
-  if (run_optimizations_ && can_allocate_registers) {
-    VLOG(compiler) << "Optimizing " << method_name;
-
+  VLOG(compiler) << "Optimizing " << method_name;
+  if (run_optimizations_) {
     {
       PassScope scope(SsaBuilder::kSsaBuilderPassName, &pass_observer);
       if (!graph->TryBuildingSsa()) {
@@ -800,37 +776,26 @@
       }
     }
 
-    compiled_method = CompileOptimized(graph,
-                                       codegen.get(),
-                                       compiler_driver,
-                                       dex_compilation_unit,
-                                       &pass_observer);
-  } else if (shouldOptimize && can_allocate_registers) {
-    LOG(FATAL) << "Could not allocate registers in optimizing compiler";
-    UNREACHABLE();
-  } else if (can_use_baseline) {
-    VLOG(compiler) << "Compile baseline " << method_name;
-
-    if (!run_optimizations_) {
-      MaybeRecordStat(MethodCompilationStat::kNotOptimizedDisabled);
-    } else if (!can_allocate_registers) {
-      MaybeRecordStat(MethodCompilationStat::kNotOptimizedRegisterAllocator);
-    }
-
-    compiled_method = CompileBaseline(codegen.get(),
-                                      compiler_driver,
-                                      dex_compilation_unit,
-                                      &pass_observer);
+    RunOptimizations(graph,
+                     codegen.get(),
+                     compiler_driver,
+                     compilation_stats_.get(),
+                     dex_compilation_unit,
+                     &pass_observer);
+    codegen->CompileOptimized(code_allocator);
+  } else {
+    codegen->CompileBaseline(code_allocator);
   }
+  pass_observer.DumpDisassembly();
 
   if (kArenaAllocatorCountAllocations) {
-    if (arena.BytesAllocated() > 4 * MB) {
-      MemStats mem_stats(arena.GetMemStats());
+    if (arena->BytesAllocated() > 4 * MB) {
+      MemStats mem_stats(arena->GetMemStats());
       LOG(INFO) << PrettyMethod(method_idx, dex_file) << " " << Dumpable<MemStats>(mem_stats);
     }
   }
 
-  return compiled_method;
+  return codegen.release();
 }
 
 static bool CanHandleVerificationFailure(const VerifiedMethod* verified_method) {
@@ -852,26 +817,37 @@
                                             Handle<mirror::DexCache> dex_cache) const {
   CompilerDriver* compiler_driver = GetCompilerDriver();
   CompiledMethod* method = nullptr;
-  if (Runtime::Current()->IsAotCompiler()) {
-    const VerifiedMethod* verified_method = compiler_driver->GetVerifiedMethod(&dex_file, method_idx);
-    DCHECK(!verified_method->HasRuntimeThrow());
-    if (compiler_driver->IsMethodVerifiedWithoutFailures(method_idx, class_def_idx, dex_file)
-        || CanHandleVerificationFailure(verified_method)) {
-       method = TryCompile(code_item, access_flags, invoke_type, class_def_idx,
-                           method_idx, jclass_loader, dex_file, dex_cache);
-    } else {
-      if (compiler_driver->GetCompilerOptions().VerifyAtRuntime()) {
-        MaybeRecordStat(MethodCompilationStat::kNotCompiledVerifyAtRuntime);
+  DCHECK(Runtime::Current()->IsAotCompiler());
+  const VerifiedMethod* verified_method = compiler_driver->GetVerifiedMethod(&dex_file, method_idx);
+  DCHECK(!verified_method->HasRuntimeThrow());
+  if (compiler_driver->IsMethodVerifiedWithoutFailures(method_idx, class_def_idx, dex_file)
+      || CanHandleVerificationFailure(verified_method)) {
+    ArenaAllocator arena(Runtime::Current()->GetArenaPool());
+    CodeVectorAllocator code_allocator(&arena);
+    std::unique_ptr<CodeGenerator> codegen(
+        TryCompile(&arena,
+                   &code_allocator,
+                   code_item,
+                   access_flags,
+                   invoke_type,
+                   class_def_idx,
+                   method_idx,
+                   jclass_loader,
+                   dex_file,
+                   dex_cache));
+    if (codegen.get() != nullptr) {
+      if (run_optimizations_) {
+        method = EmitOptimized(&arena, &code_allocator, codegen.get(), compiler_driver);
       } else {
-        MaybeRecordStat(MethodCompilationStat::kNotCompiledClassNotVerified);
+        method = EmitBaseline(&arena, &code_allocator, codegen.get(), compiler_driver);
       }
     }
   } else {
-    // This is for the JIT compiler, which has already ensured the class is verified.
-    // We can go straight to compiling.
-    DCHECK(Runtime::Current()->UseJit());
-    method = TryCompile(code_item, access_flags, invoke_type, class_def_idx,
-                        method_idx, jclass_loader, dex_file, dex_cache);
+    if (compiler_driver->GetCompilerOptions().VerifyAtRuntime()) {
+      MaybeRecordStat(MethodCompilationStat::kNotCompiledVerifyAtRuntime);
+    } else {
+      MaybeRecordStat(MethodCompilationStat::kNotCompiledClassNotVerified);
+    }
   }
 
   if (kIsDebugBuild &&
@@ -896,4 +872,70 @@
   return EndsWith(image, "core.art") || EndsWith(image, "core-optimizing.art");
 }
 
+bool OptimizingCompiler::JitCompile(Thread* self,
+                                    jit::JitCodeCache* code_cache,
+                                    ArtMethod* method) {
+  StackHandleScope<2> hs(self);
+  Handle<mirror::ClassLoader> class_loader(hs.NewHandle(
+      method->GetDeclaringClass()->GetClassLoader()));
+  Handle<mirror::DexCache> dex_cache(hs.NewHandle(method->GetDexCache()));
+
+  jobject jclass_loader = class_loader.ToJObject();
+  const DexFile* dex_file = method->GetDexFile();
+  const uint16_t class_def_idx = method->GetClassDefIndex();
+  const DexFile::CodeItem* code_item = dex_file->GetCodeItem(method->GetCodeItemOffset());
+  const uint32_t method_idx = method->GetDexMethodIndex();
+  const uint32_t access_flags = method->GetAccessFlags();
+  const InvokeType invoke_type = method->GetInvokeType();
+
+  ArenaAllocator arena(Runtime::Current()->GetArenaPool());
+  CodeVectorAllocator code_allocator(&arena);
+  std::unique_ptr<CodeGenerator> codegen;
+  {
+    // Go to native so that we don't block GC during compilation.
+    ScopedThreadSuspension sts(self, kNative);
+
+    DCHECK(run_optimizations_);
+    codegen.reset(
+        TryCompile(&arena,
+                   &code_allocator,
+                   code_item,
+                   access_flags,
+                   invoke_type,
+                   class_def_idx,
+                   method_idx,
+                   jclass_loader,
+                   *dex_file,
+                   dex_cache));
+    if (codegen.get() == nullptr) {
+      return false;
+    }
+  }
+
+  size_t stack_map_size = codegen->ComputeStackMapsSize();
+  uint8_t* stack_map_data = code_cache->ReserveData(self, stack_map_size);
+  if (stack_map_data == nullptr) {
+    return false;
+  }
+  codegen->BuildStackMaps(MemoryRegion(stack_map_data, stack_map_size));
+  const void* code = code_cache->CommitCode(
+      self,
+      method,
+      nullptr,
+      stack_map_data,
+      nullptr,
+      codegen->HasEmptyFrame() ? 0 : codegen->GetFrameSize(),
+      codegen->GetCoreSpillMask(),
+      codegen->GetFpuSpillMask(),
+      code_allocator.GetMemory().data(),
+      code_allocator.GetSize());
+
+  if (code == nullptr) {
+    code_cache->ClearData(self, stack_map_data);
+    return false;
+  }
+
+  return true;
+}
+
 }  // namespace art
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index 4c7cb1e..fbcba1b 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -248,40 +248,49 @@
 
   OatQuickMethodHeader* method_header = nullptr;
   uint8_t* code_ptr = nullptr;
-
-  ScopedThreadSuspension sts(self, kSuspended);
-  MutexLock mu(self, lock_);
-  WaitForPotentialCollectionToComplete(self);
   {
-    ScopedCodeCacheWrite scc(code_map_.get());
-    uint8_t* result = reinterpret_cast<uint8_t*>(
-        mspace_memalign(code_mspace_, alignment, total_size));
-    if (result == nullptr) {
-      return nullptr;
+    ScopedThreadSuspension sts(self, kSuspended);
+    MutexLock mu(self, lock_);
+    WaitForPotentialCollectionToComplete(self);
+    {
+      ScopedCodeCacheWrite scc(code_map_.get());
+      uint8_t* result = reinterpret_cast<uint8_t*>(
+          mspace_memalign(code_mspace_, alignment, total_size));
+      if (result == nullptr) {
+        return nullptr;
+      }
+      code_ptr = result + header_size;
+      DCHECK_ALIGNED_PARAM(reinterpret_cast<uintptr_t>(code_ptr), alignment);
+
+      std::copy(code, code + code_size, code_ptr);
+      method_header = OatQuickMethodHeader::FromCodePointer(code_ptr);
+      new (method_header) OatQuickMethodHeader(
+          (mapping_table == nullptr) ? 0 : code_ptr - mapping_table,
+          (vmap_table == nullptr) ? 0 : code_ptr - vmap_table,
+          (gc_map == nullptr) ? 0 : code_ptr - gc_map,
+          frame_size_in_bytes,
+          core_spill_mask,
+          fp_spill_mask,
+          code_size);
     }
-    code_ptr = result + header_size;
-    DCHECK_ALIGNED_PARAM(reinterpret_cast<uintptr_t>(code_ptr), alignment);
 
-    std::copy(code, code + code_size, code_ptr);
-    method_header = OatQuickMethodHeader::FromCodePointer(code_ptr);
-    new (method_header) OatQuickMethodHeader(
-        (mapping_table == nullptr) ? 0 : code_ptr - mapping_table,
-        (vmap_table == nullptr) ? 0 : code_ptr - vmap_table,
-        (gc_map == nullptr) ? 0 : code_ptr - gc_map,
-        frame_size_in_bytes,
-        core_spill_mask,
-        fp_spill_mask,
-        code_size);
+    __builtin___clear_cache(reinterpret_cast<char*>(code_ptr),
+                            reinterpret_cast<char*>(code_ptr + code_size));
+    method_code_map_.Put(code_ptr, method);
+    // We have checked there was no collection in progress earlier. If we
+    // were, setting the entry point of a method would be unsafe, as the collection
+    // could delete it.
+    DCHECK(!collection_in_progress_);
+    method->SetEntryPointFromQuickCompiledCode(method_header->GetEntryPoint());
   }
+  VLOG(jit)
+      << "JIT added "
+      << PrettyMethod(method) << "@" << method
+      << " ccache_size=" << PrettySize(CodeCacheSize()) << ": "
+      << " dcache_size=" << PrettySize(DataCacheSize()) << ": "
+      << reinterpret_cast<const void*>(method_header->GetEntryPoint()) << ","
+      << reinterpret_cast<const void*>(method_header->GetEntryPoint() + method_header->code_size_);
 
-  __builtin___clear_cache(reinterpret_cast<char*>(code_ptr),
-                          reinterpret_cast<char*>(code_ptr + code_size));
-  method_code_map_.Put(code_ptr, method);
-  // We have checked there was no collection in progress earlier. If we
-  // were, setting the entry point of a method would be unsafe, as the collection
-  // could delete it.
-  DCHECK(!collection_in_progress_);
-  method->SetEntryPointFromQuickCompiledCode(method_header->GetEntryPoint());
   return reinterpret_cast<uint8_t*>(method_header);
 }
 
@@ -304,6 +313,11 @@
   return method_code_map_.size();
 }
 
+void JitCodeCache::ClearData(Thread* self, void* data) {
+  MutexLock mu(self, lock_);
+  mspace_free(data_mspace_, data);
+}
+
 uint8_t* JitCodeCache::ReserveData(Thread* self, size_t size) {
   size = RoundUp(size, sizeof(void*));
   uint8_t* result = nullptr;
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index e10f962..afff657 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -88,6 +88,11 @@
       SHARED_REQUIRES(Locks::mutator_lock_)
       REQUIRES(!lock_);
 
+  // Clear data from the data portion of the code cache.
+  void ClearData(Thread* self, void* data)
+      SHARED_REQUIRES(Locks::mutator_lock_)
+      REQUIRES(!lock_);
+
   // Add a data array of size (end - begin) with the associated contents, returns null if there
   // is no more room.
   uint8_t* AddDataArray(Thread* self, const uint8_t* begin, const uint8_t* end)