Fix image writer for generation of app images

Fixed IsInBootImage and IsInBootOatFile. Also re-enabled generation
of app images for tests.

Bug: 26317072
Bug: 22858531
Change-Id: I19cbe7a8143b5ff7582eb27dfd9d2d8121c67b3c
diff --git a/compiler/image_writer.cc b/compiler/image_writer.cc
index 3134297..da15bc9 100644
--- a/compiler/image_writer.cc
+++ b/compiler/image_writer.cc
@@ -76,23 +76,35 @@
 
 // Return true if an object is already in an image space.
 bool ImageWriter::IsInBootImage(const void* obj) const {
+  gc::Heap* const heap = Runtime::Current()->GetHeap();
   if (!compile_app_image_) {
-    DCHECK(boot_image_space_ == nullptr);
+    DCHECK(heap->GetBootImageSpaces().empty());
     return false;
   }
-  const uint8_t* image_begin = boot_image_space_->Begin();
-  // Real image end including ArtMethods and ArtField sections.
-  const uint8_t* image_end = image_begin + boot_image_space_->GetImageHeader().GetImageSize();
-  return image_begin <= obj && obj < image_end;
+  for (gc::space::ImageSpace* boot_image_space : heap->GetBootImageSpaces()) {
+    const uint8_t* image_begin = boot_image_space->Begin();
+    // Real image end including ArtMethods and ArtField sections.
+    const uint8_t* image_end = image_begin + boot_image_space->GetImageHeader().GetImageSize();
+    if (image_begin <= obj && obj < image_end) {
+      return true;
+    }
+  }
+  return false;
 }
 
 bool ImageWriter::IsInBootOatFile(const void* ptr) const {
+  gc::Heap* const heap = Runtime::Current()->GetHeap();
   if (!compile_app_image_) {
-    DCHECK(boot_image_space_ == nullptr);
+    DCHECK(heap->GetBootImageSpaces().empty());
     return false;
   }
-  const ImageHeader& image_header = boot_image_space_->GetImageHeader();
-  return image_header.GetOatFileBegin() <= ptr && ptr < image_header.GetOatFileEnd();
+  for (gc::space::ImageSpace* boot_image_space : heap->GetBootImageSpaces()) {
+    const ImageHeader& image_header = boot_image_space->GetImageHeader();
+    if (image_header.GetOatFileBegin() <= ptr && ptr < image_header.GetOatFileEnd()) {
+      return true;
+    }
+  }
+  return false;
 }
 
 static void CheckNoDexObjectsCallback(Object* obj, void* arg ATTRIBUTE_UNUSED)
@@ -109,14 +121,6 @@
 bool ImageWriter::PrepareImageAddressSpace() {
   target_ptr_size_ = InstructionSetPointerSize(compiler_driver_.GetInstructionSet());
   gc::Heap* const heap = Runtime::Current()->GetHeap();
-  // Cache boot image space.
-    for (gc::space::ContinuousSpace* space : heap->GetContinuousSpaces()) {
-      if (space->IsImageSpace()) {
-        CHECK(compile_app_image_);
-        CHECK(boot_image_space_ == nullptr) << "Multiple image spaces";
-        boot_image_space_ = space->AsImageSpace();
-      }
-    }
   {
     ScopedObjectAccess soa(Thread::Current());
     PruneNonImageClasses();  // Remove junk
@@ -1973,7 +1977,7 @@
   copy->SetDeclaringClass(GetImageAddress(orig->GetDeclaringClassUnchecked()));
 
   const char* oat_filename;
-  if (orig->IsRuntimeMethod()) {
+  if (orig->IsRuntimeMethod() || compile_app_image_) {
     oat_filename = default_oat_filename_;
   } else {
     auto it = dex_file_oat_filename_map_.find(orig->GetDexFile());
@@ -2192,7 +2196,6 @@
       oat_file_(nullptr),
       compile_pic_(compile_pic),
       compile_app_image_(compile_app_image),
-      boot_image_space_(nullptr),
       target_ptr_size_(InstructionSetPointerSize(compiler_driver_.GetInstructionSet())),
       image_method_array_(ImageHeader::kImageMethodsCount),
       dirty_methods_(0u),
diff --git a/compiler/image_writer.h b/compiler/image_writer.h
index f491a5a..ad69038 100644
--- a/compiler/image_writer.h
+++ b/compiler/image_writer.h
@@ -470,9 +470,6 @@
   const bool compile_pic_;
   const bool compile_app_image_;
 
-  // Cache the boot image space in this class for faster lookups.
-  gc::space::ImageSpace* boot_image_space_;
-
   // Size of pointers on the target architecture.
   size_t target_ptr_size_;
 
diff --git a/test/etc/run-test-jar b/test/etc/run-test-jar
index dacb7b9..e004b6c 100755
--- a/test/etc/run-test-jar
+++ b/test/etc/run-test-jar
@@ -361,9 +361,7 @@
 dex2oat_cmdline="true"
 mkdir_cmdline="mkdir -p ${DEX_LOCATION}/dalvik-cache/$ISA"
 
-# TODO: allow app-image to work with multi-image. b/26317072
-app_image=""
-# app_image="--app-image-file=$DEX_LOCATION/dalvik-cache/$ISA/$(echo $DEX_LOCATION/$TEST_NAME.jar/classes.art | cut -d/ -f 2- | sed "s:/:@:g")"
+app_image="--app-image-file=$DEX_LOCATION/dalvik-cache/$ISA/$(echo $DEX_LOCATION/$TEST_NAME.jar/classes.art | cut -d/ -f 2- | sed "s:/:@:g")"
 
 if [ "$PREBUILD" = "y" ]; then
   dex2oat_cmdline="$INVOKE_WITH $ANDROID_ROOT/bin/dex2oatd \