1 files changed, 55 insertions, 19 deletions

diff --git a/compiler/optimizing/optimizing_unit_test.h b/compiler/optimizing/optimizing_unit_test.h
index 33f1a4affe..f31ad828eb 100644
--- a/compiler/optimizing/optimizing_unit_test.h
+++ b/compiler/optimizing/optimizing_unit_test.h
@@ -17,6 +17,7 @@
 #ifndef ART_COMPILER_OPTIMIZING_OPTIMIZING_UNIT_TEST_H_
 #define ART_COMPILER_OPTIMIZING_OPTIMIZING_UNIT_TEST_H_
 
+#include "base/scoped_arena_allocator.h"
 #include "builder.h"
 #include "common_compiler_test.h"
 #include "dex_file.h"
@@ -78,30 +79,65 @@ void RemoveSuspendChecks(HGraph* graph) {
   }
 }
 
-inline HGraph* CreateGraph(ArenaAllocator* allocator) {
-  return new (allocator) HGraph(
-      allocator,
-      *reinterpret_cast<DexFile*>(allocator->Alloc(sizeof(DexFile))),
+class ArenaPoolAndAllocator {
+ public:
+  ArenaPoolAndAllocator() : pool_(), allocator_(&pool_), arena_stack_(&pool_) { }
+
+  ArenaAllocator* GetAllocator() { return &allocator_; }
+  ArenaStack* GetArenaStack() { return &arena_stack_; }
+
+ private:
+  ArenaPool pool_;
+  ArenaAllocator allocator_;
+  ArenaStack arena_stack_;
+};
+
+inline HGraph* CreateGraph(ArenaPoolAndAllocator* pool_and_allocator) {
+  return new (pool_and_allocator->GetAllocator()) HGraph(
+      pool_and_allocator->GetAllocator(),
+      pool_and_allocator->GetArenaStack(),
+      *reinterpret_cast<DexFile*>(pool_and_allocator->GetAllocator()->Alloc(sizeof(DexFile))),
       /*method_idx*/-1,
       kRuntimeISA);
 }
 
-// Create a control-flow graph from Dex instructions.
-inline HGraph* CreateCFG(ArenaAllocator* allocator,
-                         const uint16_t* data,
-                         DataType::Type return_type = DataType::Type::kInt32) {
-  const DexFile::CodeItem* item =
-    reinterpret_cast<const DexFile::CodeItem*>(data);
-  HGraph* graph = CreateGraph(allocator);
-
-  {
-    ScopedObjectAccess soa(Thread::Current());
-    VariableSizedHandleScope handles(soa.Self());
-    HGraphBuilder builder(graph, *item, &handles, return_type);
-    bool graph_built = (builder.BuildGraph() == kAnalysisSuccess);
-    return graph_built ? graph : nullptr;
+class OptimizingUnitTest : public CommonCompilerTest {
+ protected:
+  OptimizingUnitTest() : pool_and_allocator_(new ArenaPoolAndAllocator()) { }
+
+  ArenaAllocator* GetAllocator() { return pool_and_allocator_->GetAllocator(); }
+  ArenaStack* GetArenaStack() { return pool_and_allocator_->GetArenaStack(); }
+
+  void ResetPoolAndAllocator() {
+    pool_and_allocator_.reset(new ArenaPoolAndAllocator());
+    handles_.reset();  // When getting rid of the old HGraph, we can also reset handles_.
   }
-}
+
+  HGraph* CreateGraph() {
+    return art::CreateGraph(pool_and_allocator_.get());
+  }
+
+  // Create a control-flow graph from Dex instructions.
+  HGraph* CreateCFG(const uint16_t* data, DataType::Type return_type = DataType::Type::kInt32) {
+    const DexFile::CodeItem* item =
+      reinterpret_cast<const DexFile::CodeItem*>(data);
+    HGraph* graph = CreateGraph();
+
+    {
+      ScopedObjectAccess soa(Thread::Current());
+      if (handles_ == nullptr) {
+        handles_.reset(new VariableSizedHandleScope(soa.Self()));
+      }
+      HGraphBuilder builder(graph, *item, handles_.get(), return_type);
+      bool graph_built = (builder.BuildGraph() == kAnalysisSuccess);
+      return graph_built ? graph : nullptr;
+    }
+  }
+
+ private:
+  std::unique_ptr<ArenaPoolAndAllocator> pool_and_allocator_;
+  std::unique_ptr<VariableSizedHandleScope> handles_;
+};
 
 // Naive string diff data type.
 typedef std::list<std::pair<std::string, std::string>> diff_t;