Make ImageWriter tolerant of image classes implicitly loaded by <clinit>s.

Change-Id: I2cacd47ef63c64c6f48d5a696ec994b63f01f237
diff --git a/src/dex2oat.cc b/src/dex2oat.cc
index 292e41c..0567788 100644
--- a/src/dex2oat.cc
+++ b/src/dex2oat.cc
@@ -162,7 +162,7 @@
 
 
   // Make a list of descriptors for classes to include in the image
-  const std::set<std::string>* GetImageClassDescriptors(const char* image_classes_filename)
+  std::set<std::string>* GetImageClassDescriptors(const char* image_classes_filename)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
     UniquePtr<std::ifstream> image_classes_file(new std::ifstream(image_classes_filename, std::ifstream::in));
     if (image_classes_file.get() == NULL) {
@@ -313,7 +313,7 @@
 
   bool CreateImageFile(const std::string& image_filename,
                        uintptr_t image_base,
-                       const std::set<std::string>* image_classes,
+                       std::set<std::string>* image_classes,
                        const std::string& oat_filename,
                        const std::string& oat_location,
                        const CompilerDriver& compiler)
@@ -934,7 +934,7 @@
   ScopedObjectAccess soa(Thread::Current());
 
   // If --image-classes was specified, calculate the full list of classes to include in the image
-  UniquePtr<const std::set<std::string> > image_classes(NULL);
+  UniquePtr<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) {
diff --git a/src/image_writer.cc b/src/image_writer.cc
index 7e82d6d..fc9aabc 100644
--- a/src/image_writer.cc
+++ b/src/image_writer.cc
@@ -221,17 +221,30 @@
   Runtime* runtime = Runtime::Current();
   ClassLinker* class_linker = runtime->GetClassLinker();
 
+  // Update image_classes_ with classes for objects created by <clinit> methods.
+  Thread* self = Thread::Current();
+  const char* old_cause = self->StartAssertNoThreadSuspension("ImageWriter");
+  Heap* heap = Runtime::Current()->GetHeap();
+  // TODO: Image spaces only?
+  WriterMutexLock mu(self, *Locks::heap_bitmap_lock_);
+  heap->FlushAllocStack();
+  heap->GetLiveBitmap()->Walk(FindClinitImageClassesCallback, this);
+  self->EndAssertNoThreadSuspension(old_cause);
+
+  // Make a list of classes we would like to prune.
   std::set<std::string> non_image_classes;
   NonImageClasses context;
   context.image_writer = this;
   context.non_image_classes = &non_image_classes;
   class_linker->VisitClasses(NonImageClassesVisitor, &context);
 
+  // Remove the undesired classes from the class roots.
   typedef std::set<std::string>::const_iterator ClassIt;  // TODO: C++0x auto
   for (ClassIt it = non_image_classes.begin(), end = non_image_classes.end(); it != end; ++it) {
     class_linker->RemoveClass((*it).c_str(), NULL);
   }
 
+  // Clear references to removed classes from the DexCaches.
   AbstractMethod* resolution_method = runtime->GetResolutionMethod();
   typedef Set::const_iterator CacheIt;  // TODO: C++0x auto
   for (CacheIt it = dex_caches_.begin(), end = dex_caches_.end(); it != end; ++it) {
@@ -258,6 +271,23 @@
   }
 }
 
+void ImageWriter::FindClinitImageClassesCallback(Object* object, void* arg) {
+  DCHECK(object != NULL);
+  DCHECK(arg != NULL);
+  ImageWriter* image_writer = reinterpret_cast<ImageWriter*>(arg);
+  Class* klass = object->GetClass();
+  while (klass->IsArrayClass()) {
+    klass = klass->GetComponentType();
+  }
+  if (klass->IsPrimitive()) {
+    return;
+  }
+  while (!klass->IsObjectClass()) {
+    image_writer->image_classes_->insert(ClassHelper(klass).GetDescriptor());
+    klass = klass->GetSuperClass();
+  }
+}
+
 bool ImageWriter::NonImageClassesVisitor(Class* klass, void* arg) {
   NonImageClasses* context = reinterpret_cast<NonImageClasses*>(arg);
   if (!context->image_writer->IsImageClass(klass)) {
diff --git a/src/image_writer.h b/src/image_writer.h
index 46d134f..0cccf69 100644
--- a/src/image_writer.h
+++ b/src/image_writer.h
@@ -37,7 +37,7 @@
 // Write a Space built during compilation for use during execution.
 class ImageWriter {
  public:
-  explicit ImageWriter(const std::set<std::string>* image_classes)
+  explicit ImageWriter(std::set<std::string>* image_classes)
       : oat_file_(NULL), image_end_(0), image_begin_(NULL), image_classes_(image_classes),
         oat_data_begin_(NULL) {}
 
@@ -110,27 +110,36 @@
     return oat_data_begin_ + offset;
   }
 
+  // Returns true if the class was in the original requested image classes list.
   bool IsImageClass(const mirror::Class* klass) SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
+
+  // Debug aid that list of requested image classes.
   void DumpImageClasses();
 
+  // Preinitializes some otherwise lazy fields (such as Class name) to avoid runtime image dirtying.
   void ComputeLazyFieldsForImageClasses()
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
   static bool ComputeLazyFieldsForClassesVisitor(mirror::Class* klass, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
-  // Wire dex cache resolved strings to strings in the image to avoid runtime resolution
+  // Wire dex cache resolved strings to strings in the image to avoid runtime resolution.
   void ComputeEagerResolvedStrings();
   static void ComputeEagerResolvedStringsCallback(mirror::Object* obj, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  // Remove unwanted classes from various roots.
   void PruneNonImageClasses() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
+  static void FindClinitImageClassesCallback(mirror::Object* object, void* arg)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
   static bool NonImageClassesVisitor(mirror::Class* c, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  // Verify unwanted classes removed.
   void CheckNonImageClassesRemoved();
   static void CheckNonImageClassesRemovedCallback(mirror::Object* obj, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  // Lays out where the image objects will be at runtime.
   void CalculateNewObjectOffsets(size_t oat_loaded_size, size_t oat_data_offset)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
   mirror::ObjectArray<mirror::Object>* CreateImageRoots() const
@@ -138,6 +147,7 @@
   static void CalculateNewObjectOffsetsCallback(mirror::Object* obj, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  // Creates the contiguous image in memory and adjusts pointers.
   void CopyAndFixupObjects();
   static void CopyAndFixupObjectsCallback(mirror::Object* obj, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
@@ -158,32 +168,35 @@
                    bool is_static)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+  // Patches references in OatFile to expect runtime addresses.
   void PatchOatCodeAndMethods(const CompilerDriver& compiler)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
   void SetPatchLocation(const CompilerDriver::PatchInformation* patch, uint32_t value)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
+
+  // Map of Object to where it will be at runtime.
   SafeMap<const mirror::Object*, size_t> offsets_;
 
   // oat file with code for this image
   OatFile* oat_file_;
 
-  // memory mapped for generating the image
+  // Memory mapped for generating the image.
   UniquePtr<MemMap> image_;
 
-  // Offset to the free space in image_
+  // Offset to the free space in image_.
   size_t image_end_;
 
-  // Beginning target image address for the output image
+  // Beginning target image address for the output image.
   byte* image_begin_;
 
   // Set of classes to be include in the image, or NULL for all.
-  const std::set<std::string>* image_classes_;
+  std::set<std::string>* image_classes_;
 
-  // Beginning target oat address for the pointers from the output image to its oat file
+  // Beginning target oat address for the pointers from the output image to its oat file.
   const byte* oat_data_begin_;
 
-  // DexCaches seen while scanning for fixing up CodeAndDirectMethods
+  // DexCaches seen while scanning for fixing up CodeAndDirectMethods.
   typedef std::set<mirror::DexCache*> Set;
   Set dex_caches_;
 };