diff --git a/src/class_linker.cc b/src/class_linker.cc
index e26677d..10e2e5e 100644
--- a/src/class_linker.cc
+++ b/src/class_linker.cc
@@ -570,7 +570,7 @@
 
   Heap* heap = Runtime::Current()->GetHeap();
   std::string boot_image_option_string("--boot-image=");
-  boot_image_option_string += heap->GetSpaces()[0]->AsImageSpace()->GetImageFilename();
+  boot_image_option_string += heap->GetImageSpace()->GetImageFilename();
   const char* boot_image_option = boot_image_option_string.c_str();
 
   std::string dex_file_option_string("--dex-file=");
@@ -875,40 +875,33 @@
   CHECK(!init_done_);
 
   Heap* heap = Runtime::Current()->GetHeap();
-  const std::vector<Space*>& spaces = heap->GetSpaces();
-  for (size_t i = 0; i < spaces.size(); ++i) {
-    if (spaces[i]->IsImageSpace()) {
-      ImageSpace* space = spaces[i]->AsImageSpace();
-      OatFile* oat_file = OpenOat(space);
-      CHECK(oat_file != NULL) << "Failed to open oat file for image";
-      Object* dex_caches_object = space->GetImageHeader().GetImageRoot(ImageHeader::kDexCaches);
-      ObjectArray<DexCache>* dex_caches = dex_caches_object->AsObjectArray<DexCache>();
+  ImageSpace* space = heap->GetImageSpace();
+  OatFile* oat_file = OpenOat(space);
+  CHECK(oat_file != NULL) << "Failed to open oat file for image";
+  Object* dex_caches_object = space->GetImageHeader().GetImageRoot(ImageHeader::kDexCaches);
+  ObjectArray<DexCache>* dex_caches = dex_caches_object->AsObjectArray<DexCache>();
 
-      if (i == 0) {
-        // Special case of setting up the String class early so that we can test arbitrary objects
-        // as being Strings or not
-        Class* java_lang_String = space->GetImageHeader().GetImageRoot(ImageHeader::kClassRoots)
-            ->AsObjectArray<Class>()->Get(kJavaLangString);
-        String::SetClass(java_lang_String);
-      }
+  // Special case of setting up the String class early so that we can test arbitrary objects
+  // as being Strings or not
+  Class* java_lang_String = space->GetImageHeader().GetImageRoot(ImageHeader::kClassRoots)
+      ->AsObjectArray<Class>()->Get(kJavaLangString);
+  String::SetClass(java_lang_String);
 
-      CHECK_EQ(oat_file->GetOatHeader().GetDexFileCount(),
-               static_cast<uint32_t>(dex_caches->GetLength()));
-      for (int i = 0; i < dex_caches->GetLength(); i++) {
-        SirtRef<DexCache> dex_cache(dex_caches->Get(i));
-        const std::string& dex_file_location(dex_cache->GetLocation()->ToModifiedUtf8());
-        const OatFile::OatDexFile* oat_dex_file = oat_file->GetOatDexFile(dex_file_location);
-        const DexFile* dex_file = oat_dex_file->OpenDexFile();
-        if (dex_file == NULL) {
-          LOG(FATAL) << "Failed to open dex file " << dex_file_location
-                     << " from within oat file " << oat_file->GetLocation();
-        }
-
-        CHECK_EQ(dex_file->GetLocationChecksum(), oat_dex_file->GetDexFileLocationChecksum());
-
-        AppendToBootClassPath(*dex_file, dex_cache);
-      }
+  CHECK_EQ(oat_file->GetOatHeader().GetDexFileCount(),
+           static_cast<uint32_t>(dex_caches->GetLength()));
+  for (int i = 0; i < dex_caches->GetLength(); i++) {
+    SirtRef<DexCache> dex_cache(dex_caches->Get(i));
+    const std::string& dex_file_location(dex_cache->GetLocation()->ToModifiedUtf8());
+    const OatFile::OatDexFile* oat_dex_file = oat_file->GetOatDexFile(dex_file_location);
+    const DexFile* dex_file = oat_dex_file->OpenDexFile();
+    if (dex_file == NULL) {
+      LOG(FATAL) << "Failed to open dex file " << dex_file_location
+                 << " from within oat file " << oat_file->GetLocation();
     }
+
+    CHECK_EQ(dex_file->GetLocationChecksum(), oat_dex_file->GetDexFileLocationChecksum());
+
+    AppendToBootClassPath(*dex_file, dex_cache);
   }
 
   HeapBitmap* heap_bitmap = heap->GetLiveBits();
@@ -919,7 +912,7 @@
 
   // reinit class_roots_
   Object* class_roots_object =
-      spaces[0]->AsImageSpace()->GetImageHeader().GetImageRoot(ImageHeader::kClassRoots);
+      heap->GetImageSpace()->GetImageHeader().GetImageRoot(ImageHeader::kClassRoots);
   class_roots_ = class_roots_object->AsObjectArray<Class>();
 
   // reinit array_iftable_ from any array class instance, they should be ==
diff --git a/src/dex2oat.cc b/src/dex2oat.cc
index debbabf..1f02a44 100644
--- a/src/dex2oat.cc
+++ b/src/dex2oat.cc
@@ -245,23 +245,6 @@
                        const std::set<std::string>* image_classes,
                        const std::string& oat_filename,
                        const std::string& oat_location) {
-    // If we have an existing boot image, position new space after its oat file
-    Heap* heap = Runtime::Current()->GetHeap();
-    if (heap->GetSpaces().size() > 1) {
-      ImageSpace* last_image_space = NULL;
-      const std::vector<Space*>& spaces = heap->GetSpaces();
-      for (size_t i=0; i < spaces.size(); i++) {
-        if (spaces[i]->IsImageSpace()) {
-          last_image_space = spaces[i]->AsImageSpace();
-        }
-      }
-      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().GetOatEnd();
-      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, oat_location)) {
       LOG(ERROR) << "Failed to create image file " << image_filename;
diff --git a/src/heap.cc b/src/heap.cc
index f4f2996..189bb91 100644
--- a/src/heap.cc
+++ b/src/heap.cc
@@ -132,6 +132,7 @@
 Heap::Heap(size_t initial_size, size_t growth_limit, size_t capacity,
            const std::string& original_image_file_name)
     : lock_(NULL),
+      image_space_(NULL),
       alloc_space_(NULL),
       mark_bitmap_(NULL),
       live_bitmap_(NULL),
@@ -163,35 +164,34 @@
   byte* requested_begin = NULL;
   std::string image_file_name(original_image_file_name);
   if (!image_file_name.empty()) {
-    ImageSpace* space = NULL;
     if (OS::FileExists(image_file_name.c_str())) {
       // If the /system file exists, it should be up-to-date, don't try to generate
-      space = Space::CreateImageSpace(image_file_name);
+      image_space_ = Space::CreateImageSpace(image_file_name);
     } else {
       // If the /system file didn't exist, we need to use one from the art-cache.
       // If the cache file exists, try to open, but if it fails, regenerate.
       // If it does not exist, generate.
       image_file_name = GetArtCacheFilenameOrDie(image_file_name);
       if (OS::FileExists(image_file_name.c_str())) {
-        space = Space::CreateImageSpace(image_file_name);
+        image_space_ = Space::CreateImageSpace(image_file_name);
       }
-      if (space == NULL) {
+      if (image_space_ == NULL) {
         if (!GenerateImage(image_file_name)) {
           LOG(FATAL) << "Failed to generate image: " << image_file_name;
         }
-        space = Space::CreateImageSpace(image_file_name);
+        image_space_ = Space::CreateImageSpace(image_file_name);
       }
     }
-    if (space == NULL) {
+    if (image_space_ == NULL) {
       LOG(FATAL) << "Failed to create space from " << image_file_name;
     }
 
-    AddSpace(space);
-    UpdateFirstAndLastSpace(&first_space, &last_space, space);
+    AddSpace(image_space_);
+    UpdateFirstAndLastSpace(&first_space, &last_space, image_space_);
     // Oat files referenced by image files immediately follow them in memory, ensure alloc space
     // isn't going to get in the middle
-    byte* oat_end_addr = space->GetImageHeader().GetOatEnd();
-    CHECK(oat_end_addr > space->End());
+    byte* oat_end_addr = image_space_->GetImageHeader().GetOatEnd();
+    CHECK(oat_end_addr > image_space_->End());
     if (oat_end_addr > requested_begin) {
       requested_begin = reinterpret_cast<byte*>(RoundUp(reinterpret_cast<uintptr_t>(oat_end_addr),
                                                         kPageSize));
diff --git a/src/heap.h b/src/heap.h
index ccc2d23..374fa84 100644
--- a/src/heap.h
+++ b/src/heap.h
@@ -33,6 +33,7 @@
 
 class AllocSpace;
 class Class;
+class ImageSpace;
 class Object;
 class Space;
 class Thread;
@@ -208,7 +209,13 @@
   size_t GetBytesAllocated() { return num_bytes_allocated_; }
   size_t GetObjectsAllocated() { return num_objects_allocated_; }
 
+  ImageSpace* GetImageSpace() {
+    CHECK(image_space_ != NULL);
+    return image_space_;
+  }
+
   AllocSpace* GetAllocSpace() {
+    CHECK(alloc_space_ != NULL);
     return alloc_space_;
   }
 
@@ -242,6 +249,8 @@
 
   std::vector<Space*> spaces_;
 
+  ImageSpace* image_space_;
+
   // default Space for allocations
   AllocSpace* alloc_space_;
 
diff --git a/src/image_test.cc b/src/image_test.cc
index e886e49..fcea695 100644
--- a/src/image_test.cc
+++ b/src/image_test.cc
@@ -64,6 +64,7 @@
     Space* space = heap->GetSpaces()[0];
     ASSERT_FALSE(space->IsImageSpace());
     ASSERT_TRUE(space != NULL);
+    ASSERT_EQ(space, heap->GetAllocSpace());
     ASSERT_GE(sizeof(image_header) + space->Size(), static_cast<size_t>(file->Length()));
   }
 
@@ -88,9 +89,13 @@
   Heap* heap = Runtime::Current()->GetHeap();
   ASSERT_EQ(2U, heap->GetSpaces().size());
   ASSERT_TRUE(heap->GetSpaces()[0]->IsImageSpace());
+  ASSERT_FALSE(heap->GetSpaces()[0]->IsAllocSpace());
   ASSERT_FALSE(heap->GetSpaces()[1]->IsImageSpace());
+  ASSERT_TRUE(heap->GetSpaces()[1]->IsAllocSpace());
+  ASSERT_TRUE(heap->GetImageSpace() != NULL);
+  ASSERT_TRUE(heap->GetAllocSpace() != NULL);
 
-  ImageSpace* image_space = heap->GetSpaces()[0]->AsImageSpace();
+  ImageSpace* image_space = heap->GetImageSpace();
   byte* image_begin = image_space->Begin();
   byte* image_end = image_space->End();
   CHECK_EQ(requested_image_base, reinterpret_cast<uintptr_t>(image_begin));
diff --git a/src/image_writer.cc b/src/image_writer.cc
index 4579c0c..cc37e5b 100644
--- a/src/image_writer.cc
+++ b/src/image_writer.cc
@@ -49,11 +49,7 @@
   image_begin_ = reinterpret_cast<byte*>(image_begin);
 
   Heap* heap = Runtime::Current()->GetHeap();
-  const std::vector<Space*>& spaces = heap->GetSpaces();
-  // currently just write the last space, assuming it is the space that was being used for allocation
-  CHECK_GE(spaces.size(), 1U);
-  source_space_ = spaces[spaces.size()-1];
-  CHECK(!source_space_->IsImageSpace());
+  source_space_ = heap->GetAllocSpace();
 
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
   const std::vector<DexCache*>& all_dex_caches = class_linker->GetDexCaches();
diff --git a/src/oatdump.cc b/src/oatdump.cc
index 9641390..295c74f 100644
--- a/src/oatdump.cc
+++ b/src/oatdump.cc
@@ -1154,7 +1154,7 @@
   }
 
   Heap* heap = Runtime::Current()->GetHeap();
-  ImageSpace* image_space = heap->GetSpaces()[heap->GetSpaces().size()-2]->AsImageSpace();
+  ImageSpace* image_space = heap->GetImageSpace();
   CHECK(image_space != NULL);
   const ImageHeader& image_header = image_space->GetImageHeader();
   if (!image_header.IsValid()) {