Add --image-classes to dex2oat

Change-Id: Ia88f9d302e0f9cd72be2199ee46d212d99864c67
diff --git a/src/dex2oat.cc b/src/dex2oat.cc
index 231ca9c..0e4c224 100644
--- a/src/dex2oat.cc
+++ b/src/dex2oat.cc
@@ -4,6 +4,8 @@
 #include <stdlib.h>
 #include <sys/file.h>
 
+#include <iostream>
+#include <fstream>
 #include <string>
 #include <vector>
 
@@ -29,13 +31,16 @@
           "      Example: --dex-file=/system/framework/core.jar\n"
           "\n");
   fprintf(stderr,
-          "  --image=<file.art>: specifies the required output image filename.\n"
+          "  --oat=<file.oat>: specifies the required oat filename.\n"
+          "      Example: --oat=/data/art-cache/boot.oat\n"
+          "\n");
+  fprintf(stderr,
+          "  --image=<file.art>: specifies the output image filename.\n"
           "      Example: --image=/data/art-cache/boot.art\n"
           "\n");
-  // TODO: remove this by inferring from --image
   fprintf(stderr,
-          "  --oat=<file.oat>: specifies the required oat filename.\n"
-          "      Example: --image=/data/art-cache/boot.oat\n"
+          "  --image-classes=<classname-file>: specifies classes to include in an image.\n"
+          "      Example: --image=frameworks/base/preloaded-classes\n"
           "\n");
   fprintf(stderr,
           "  --base=<hex-address>: specifies the base address when creating a boot image.\n"
@@ -46,10 +51,6 @@
           "      Example: --boot-image=/data/art-cache/boot.art\n"
           "\n");
   fprintf(stderr,
-          "  --method may be used to limit compilation to a subset of methods.\n"
-          "      Example: --method=Ljava/lang/Object;<init>()V\n"
-          "\n");
-  fprintf(stderr,
           "  --host-prefix may be used to translate host paths to target paths during\n"
           "      cross compilation.\n"
           "      Example: --host-prefix=out/target/product/crespo\n"
@@ -94,33 +95,190 @@
   bool do_unlink_;
 };
 
-// Returns true if dex_files has a dex with the named location.
-bool DexFilesContains(const std::vector<const DexFile*>& dex_files, const std::string& location) {
-  for (size_t i = 0; i < dex_files.size(); ++i) {
-    if (dex_files[i]->GetLocation() == location) {
+class Dex2Oat {
+ public:
+
+  static Dex2Oat* Create(Runtime::Options& options) {
+    UniquePtr<Runtime> runtime(CreateRuntime(options));
+    if (runtime.get() == NULL) {
+      return NULL;
+    }
+    return new Dex2Oat(runtime.release());
+  }
+
+  ~Dex2Oat() {
+    delete runtime_;
+  }
+
+  // Make a list of descriptors for classes to include in the image
+  const std::set<std::string>* GetImageClassDescriptors(const char* image_classes_filename) {
+    UniquePtr<std::ifstream> image_classes_file(new std::ifstream(image_classes_filename, std::ifstream::in));
+    if (image_classes_file.get() == NULL) {
+      LOG(ERROR) << "Failed to open image classes file " << image_classes_filename;
+      return NULL;
+    }
+
+    // Load all the classes specifed in the file
+    ClassLinker* class_linker = runtime_->GetClassLinker();
+    while (image_classes_file->good()) {
+      std::string dot;
+      std::getline(*image_classes_file.get(), dot);
+      if (StringPiece(dot).starts_with("#") || dot.empty()) {
+        continue;
+      }
+      std::string descriptor = DotToDescriptor(dot.c_str());
+      SirtRef<Class> klass(class_linker->FindSystemClass(descriptor));
+      if (klass.get() == NULL) {
+        LOG(WARNING) << "Failed to find class " << descriptor;
+        Thread::Current()->ClearException();
+      }
+    }
+    image_classes_file->close();
+
+    // We walk the roots looking for classes so that we'll pick up the
+    // above classes plus any classes them depend on such super
+    // classes, interfaces, and the required ClassLinker roots.
+    UniquePtr<std::set<std::string> > image_classes(new std::set<std::string>());
+    class_linker->VisitClasses(ClassVisitor, image_classes.get());
+    CHECK_NE(image_classes->size(), 0U);
+    return image_classes.release();
+  }
+
+  bool CreateOatFile(const std::string& boot_image_option,
+                     const std::vector<const char*>& dex_filenames,
+                     const std::string& host_prefix,
+                     File* oat_file,
+                     bool image,
+                     const std::set<std::string>* image_classes) {
+    // SirtRef and ClassLoader creation needs to come after Runtime::Create
+    UniquePtr<SirtRef<ClassLoader> > class_loader(new SirtRef<ClassLoader>(NULL));
+    if (class_loader.get() == NULL) {
+      LOG(ERROR) << "Failed to create SirtRef for class loader";
+      return false;
+    }
+
+    std::vector<const DexFile*> dex_files;
+    if (!boot_image_option.empty()) {
+      ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+      DexFile::OpenDexFiles(dex_filenames, dex_files, host_prefix);
+      std::vector<const DexFile*> class_path_files(dex_files);
+      OpenClassPathFiles(runtime_->GetClassPath(), class_path_files);
+      for (size_t i = 0; i < class_path_files.size(); i++) {
+        class_linker->RegisterDexFile(*class_path_files[i]);
+      }
+      class_loader.get()->reset(PathClassLoader::AllocCompileTime(class_path_files));
+    } else {
+      dex_files = runtime_->GetClassLinker()->GetBootClassPath();
+    }
+
+    Compiler compiler(instruction_set_, image, image_classes);
+    compiler.CompileAll(class_loader->get(), dex_files);
+
+    if (!OatWriter::Create(oat_file, class_loader->get(), compiler)) {
+      LOG(ERROR) << "Failed to create oat file " << oat_file->name();
+      return false;
+    }
+    return true;
+  }
+
+  bool CreateImageFile(const char* image_filename,
+                       uintptr_t image_base,
+                       const std::set<std::string>* image_classes,
+                       const std::string& oat_filename,
+                       const std::string& host_prefix) {
+    // If we have an existing boot image, position new space after its oat file
+    if (Heap::GetSpaces().size() > 1) {
+      Space* last_image_space = Heap::GetSpaces()[Heap::GetSpaces().size()-2];
+      CHECK(last_image_space != NULL);
+      CHECK(last_image_space->IsImageSpace());
+      CHECK(!Heap::GetSpaces()[Heap::GetSpaces().size()-1]->IsImageSpace());
+      byte* oat_limit_addr = last_image_space->GetImageHeader().GetOatLimitAddr();
+      image_base = RoundUp(reinterpret_cast<uintptr_t>(oat_limit_addr), kPageSize);
+    }
+
+    ImageWriter image_writer(image_classes);
+    if (!image_writer.Write(image_filename, image_base, oat_filename, host_prefix)) {
+      LOG(ERROR) << "Failed to create image file " << image_filename;
+      return false;
+    }
+    return true;
+  }
+
+ private:
+
+  Dex2Oat(Runtime* runtime) : runtime_(runtime) {}
+
+  static Runtime* CreateRuntime(Runtime::Options& options) {
+    Runtime* runtime = Runtime::Create(options, false);
+    if (runtime == NULL) {
+      LOG(ERROR) << "Failed to create runtime";
+      return NULL;
+    }
+
+    // if we loaded an existing image, we will reuse values from the image roots.
+    if (!runtime->HasJniDlsymLookupStub()) {
+      runtime->SetJniDlsymLookupStub(Compiler::CreateJniDlysmLookupStub(instruction_set_));
+    }
+    if (!runtime->HasAbstractMethodErrorStubArray()) {
+      runtime->SetAbstractMethodErrorStubArray(Compiler::CreateAbstractMethodErrorStub(instruction_set_));
+    }
+    for (int i = 0; i < Runtime::kLastTrampolineMethodType; i++) {
+      Runtime::TrampolineType type = Runtime::TrampolineType(i);
+      if (!runtime->HasResolutionStubArray(type)) {
+        runtime->SetResolutionStubArray(Compiler::CreateResolutionStub(instruction_set_, type), type);
+      }
+    }
+    for (int i = 0; i < Runtime::kLastCalleeSaveType; i++) {
+      Runtime::CalleeSaveType type = Runtime::CalleeSaveType(i);
+      if (!runtime->HasCalleeSaveMethod(type)) {
+        runtime->SetCalleeSaveMethod(runtime->CreateCalleeSaveMethod(instruction_set_, type), type);
+      }
+    }
+    return runtime;
+  }
+
+  static bool ClassVisitor(Class* klass, void* arg) {
+    std::set<std::string>* image_classes = reinterpret_cast<std::set<std::string>*>(arg);
+    if (klass->IsArrayClass() || klass->IsPrimitive()) {
       return true;
     }
+    image_classes->insert(klass->GetDescriptor()->ToModifiedUtf8());
+    return true;
   }
-  return false;
-}
 
-// Appends to dex_files any elements of class_path that it doesn't already
-// contain. This will open those dex files as necessary.
-void OpenClassPathFiles(const std::string& class_path, std::vector<const DexFile*>& dex_files) {
-  std::vector<std::string> parsed;
-  Split(class_path, ':', parsed);
-  for (size_t i = 0; i < parsed.size(); ++i) {
-    if (DexFilesContains(dex_files, parsed[i])) {
-      continue;
-    }
-    const DexFile* dex_file = DexFile::Open(parsed[i], Runtime::Current()->GetHostPrefix());
-    if (dex_file == NULL) {
-      LOG(WARNING) << "Failed to open dex file " << parsed[i];
-    } else {
-      dex_files.push_back(dex_file);
+  // Appends to dex_files any elements of class_path that it doesn't already
+  // contain. This will open those dex files as necessary.
+  static void OpenClassPathFiles(const std::string& class_path, std::vector<const DexFile*>& dex_files) {
+    std::vector<std::string> parsed;
+    Split(class_path, ':', parsed);
+    for (size_t i = 0; i < parsed.size(); ++i) {
+      if (DexFilesContains(dex_files, parsed[i])) {
+        continue;
+      }
+      const DexFile* dex_file = DexFile::Open(parsed[i], Runtime::Current()->GetHostPrefix());
+      if (dex_file == NULL) {
+        LOG(WARNING) << "Failed to open dex file " << parsed[i];
+      } else {
+        dex_files.push_back(dex_file);
+      }
     }
   }
-}
+
+  // Returns true if dex_files has a dex with the named location.
+  static bool DexFilesContains(const std::vector<const DexFile*>& dex_files, const std::string& location) {
+    for (size_t i = 0; i < dex_files.size(); ++i) {
+      if (dex_files[i]->GetLocation() == location) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  Runtime* runtime_;
+  static const InstructionSet instruction_set_ = kThumb2;
+
+  DISALLOW_IMPLICIT_CONSTRUCTORS(Dex2Oat);
+};
 
 int dex2oat(int argc, char** argv) {
   // Skip over argv[0].
@@ -133,9 +291,9 @@
   }
 
   std::vector<const char*> dex_filenames;
-  std::vector<const char*> method_names;
   std::string oat_filename;
   const char* image_filename = NULL;
+  const char* image_classes_filename = NULL;
   std::string boot_image_option;
   uintptr_t image_base = 0;
   std::string host_prefix;
@@ -148,12 +306,12 @@
     }
     if (option.starts_with("--dex-file=")) {
       dex_filenames.push_back(option.substr(strlen("--dex-file=")).data());
-    } else if (option.starts_with("--method=")) {
-      method_names.push_back(option.substr(strlen("--method=")).data());
     } else if (option.starts_with("--oat=")) {
       oat_filename = option.substr(strlen("--oat=")).data();
     } else if (option.starts_with("--image=")) {
       image_filename = option.substr(strlen("--image=")).data();
+    } else if (option.starts_with("--image-classes=")) {
+      image_classes_filename = option.substr(strlen("--image-classes=")).data();
     } else if (option.starts_with("--base=")) {
       const char* image_base_str = option.substr(strlen("--base=")).data();
       char* end;
@@ -181,18 +339,24 @@
     }
   }
 
-  if (oat_filename == NULL) {
+  if (oat_filename.empty()) {
     fprintf(stderr, "--oat file name not specified\n");
     return EXIT_FAILURE;
   }
 
-  if (image_filename == NULL && boot_image_option.empty()) {
+  bool image = (image_filename != NULL);
+  if (!image && boot_image_option.empty()) {
     fprintf(stderr, "Either --image or --boot-image must be specified\n");
     return EXIT_FAILURE;
   }
 
-  if (dex_filenames.empty()) {
-    fprintf(stderr, "no --dex-file values specified\n");
+  if (image_classes_filename != NULL && !image) {
+    fprintf(stderr, "--image-classes should only be used with --image\n");
+    return EXIT_FAILURE;
+  }
+
+  if (image_classes_filename != NULL && !boot_image_option.empty()) {
+    fprintf(stderr, "--image-classes should not be used with --boot-image\n");
     return EXIT_FAILURE;
   }
 
@@ -220,7 +384,7 @@
   }
 
   // Handles removing the file on failure and unlocking on both failure and success.
-  FileJanitor file_janitor(oat_filename, fd);
+  FileJanitor oat_file_janitor(oat_filename, fd);
 
   // If we won the creation race, block trying to take the lock (since we're going to be doing
   // the work, we need the lock). If we lost the creation race, spin trying to take the lock
@@ -252,7 +416,7 @@
     // TODO: check the creator did a good job by checking the header.
     LOG(INFO) << "Another process finished working on " << oat_filename;
     // Job done.
-    file_janitor.KeepFile();
+    oat_file_janitor.KeepFile();
     return EXIT_SUCCESS;
   }
 
@@ -281,122 +445,47 @@
   for (size_t i = 0; i < runtime_args.size(); i++) {
     options.push_back(std::make_pair(runtime_args[i], reinterpret_cast<void*>(NULL)));
   }
-  UniquePtr<Runtime> runtime(Runtime::Create(options, false));
-  if (runtime.get() == NULL) {
-    LOG(ERROR) << "Could not create runtime";
-    return EXIT_FAILURE;
-  }
-  ClassLinker* class_linker = runtime->GetClassLinker();
 
-  // If we have an existing boot image, position new space after its oat file
-  if (Heap::GetSpaces().size() > 1) {
-    Space* last_image_space = Heap::GetSpaces()[Heap::GetSpaces().size()-2];
-    CHECK(last_image_space != NULL);
-    CHECK(last_image_space->IsImageSpace());
-    CHECK(!Heap::GetSpaces()[Heap::GetSpaces().size()-1]->IsImageSpace());
-    byte* oat_limit_addr = last_image_space->GetImageHeader().GetOatLimitAddr();
-    image_base = RoundUp(reinterpret_cast<uintptr_t>(oat_limit_addr), kPageSize);
-  }
+  UniquePtr<Dex2Oat> dex2oat(Dex2Oat::Create(options));
 
-  // ClassLoader creation needs to come after Runtime::Create
-  SirtRef<ClassLoader> class_loader(NULL);
-  std::vector<const DexFile*> dex_files;
-  if (!boot_image_option.empty()) {
-    DexFile::OpenDexFiles(dex_filenames, dex_files, host_prefix);
-    std::vector<const DexFile*> class_path_files(dex_files);
-    OpenClassPathFiles(runtime->GetClassPath(), class_path_files);
-    for (size_t i = 0; i < class_path_files.size(); i++) {
-      class_linker->RegisterDexFile(*class_path_files[i]);
-    }
-    class_loader.reset(PathClassLoader::AllocCompileTime(class_path_files));
-  } else {
-    dex_files = runtime->GetClassLinker()->GetBootClassPath();
-  }
-
-  // if we loaded an existing image, we will reuse values from the image roots.
-  if (!runtime->HasJniDlsymLookupStub()) {
-    runtime->SetJniDlsymLookupStub(Compiler::CreateJniDlysmLookupStub(kThumb2));
-  }
-  if (!runtime->HasAbstractMethodErrorStubArray()) {
-    runtime->SetAbstractMethodErrorStubArray(Compiler::CreateAbstractMethodErrorStub(kThumb2));
-  }
-  for (int i = 0; i < Runtime::kLastTrampolineMethodType; i++) {
-    Runtime::TrampolineType type = Runtime::TrampolineType(i);
-    if (!runtime->HasResolutionStubArray(type)) {
-      runtime->SetResolutionStubArray(Compiler::CreateResolutionStub(kThumb2, type), type);
-    }
-  }
-  for (int i = 0; i < Runtime::kLastCalleeSaveType; i++) {
-    Runtime::CalleeSaveType type = Runtime::CalleeSaveType(i);
-    if (!runtime->HasCalleeSaveMethod(type)) {
-      runtime->SetCalleeSaveMethod(runtime->CreateCalleeSaveMethod(kThumb2, type), type);
-    }
-  }
-  Compiler compiler(kThumb2, image_filename != NULL);
-  if (method_names.empty()) {
-    compiler.CompileAll(class_loader.get(), dex_files);
-  } else {
-    for (size_t i = 0; i < method_names.size(); i++) {
-      // names are actually class_descriptor + name + signature.
-      // example: Ljava/lang/Object;<init>()V
-      StringPiece method_name = method_names[i];
-      size_t end_of_class_descriptor = method_name.find(';');
-      if (end_of_class_descriptor == method_name.npos) {
-        LOG(ERROR) << "Could not find class descriptor in method " << method_name << "'";
-        return EXIT_FAILURE;
-      }
-      end_of_class_descriptor++;  // want to include ;
-      std::string class_descriptor = method_name.substr(0, end_of_class_descriptor).ToString();
-      size_t end_of_name = method_name.find('(', end_of_class_descriptor);
-      if (end_of_name == method_name.npos) {
-        LOG(ERROR) << "Could not find start of method signature in method '" << method_name << "'";
-        return EXIT_FAILURE;
-      }
-      std::string name = method_name.substr(end_of_class_descriptor,
-                                            end_of_name - end_of_class_descriptor).ToString();
-      std::string signature = method_name.substr(end_of_name).ToString();
-
-      Class* klass = class_linker->FindClass(class_descriptor, class_loader.get());
-      if (klass == NULL) {
-        LOG(ERROR) << "Could not find class for descriptor '" << class_descriptor
-            << "' in method '" << method_name << "'";
-        return EXIT_FAILURE;
-      }
-      Method* method = klass->FindDirectMethod(name, signature);
-      if (method == NULL) {
-          method = klass->FindVirtualMethod(name, signature);
-      }
-      if (method == NULL) {
-        LOG(ERROR) << "Could not find method '" << method_name << "' with signature '"
-            << signature << "' in class '" << class_descriptor << "' for method argument '"
-            << method_name << "'";
-        return EXIT_FAILURE;
-      }
-      compiler.CompileOne(method);
+  // If --image-classes was specified, calculate the full list classes to include in the image
+  UniquePtr<const std::set<std::string> > image_classes(NULL);
+  if (image_classes_filename != NULL) {
+    image_classes.reset(dex2oat->GetImageClassDescriptors(image_classes_filename));
+    if (image_classes.get() == NULL) {
+      LOG(ERROR) << "Failed to create list of image classes from " << image_classes_filename;
+      return EXIT_FAILURE;
     }
   }
 
-  if (!OatWriter::Create(oat_file.get(), class_loader.get(), compiler)) {
-    LOG(ERROR) << "Failed to create oat file " << oat_file->name();
+  if (!dex2oat->CreateOatFile(boot_image_option,
+                              dex_filenames,
+                              host_prefix,
+                              oat_file.get(),
+                              image,
+                              image_classes.get())) {
+    LOG(ERROR) << "Failed to create oat file" << oat_filename;
     return EXIT_FAILURE;
   }
 
-  if (image_filename == NULL) {
-    file_janitor.KeepFile();
-    LOG(INFO) << "Oat file written successfully " << oat_file->name();
+  if (!image) {
+    oat_file_janitor.KeepFile();
+    LOG(INFO) << "Oat file written successfully " << oat_filename;
     return EXIT_SUCCESS;
   }
-  CHECK(compiler.IsImage());
 
-  ImageWriter image_writer;
-  if (!image_writer.Write(image_filename, image_base, oat_file->name(), host_prefix)) {
-    LOG(ERROR) << "Failed to create image file " << image_filename;
+  if (!dex2oat->CreateImageFile(image_filename,
+                                image_base,
+                                image_classes.get(),
+                                oat_filename,
+                                host_prefix)) {
     return EXIT_FAILURE;
   }
 
-  // We wrote the file successfully, and want to keep it.
+  // We wrote the oat file successfully, and want to keep it.
+  oat_file_janitor.KeepFile();
+  LOG(INFO) << "Oat file written successfully " << oat_filename;
   LOG(INFO) << "Image written successfully " << image_filename;
-  file_janitor.KeepFile();
   return EXIT_SUCCESS;
 }