Refactor and improve mod-union tables.

Allow support for adding more mod union tables, reduces the amount
of baked in logic. Adds support for updating mod union table references
from compaction (not for ReferenceCache table yet).

Change-Id: I1beeda00839ed86ef0e853beff5ce10d0ab2b9d1
diff --git a/runtime/gc/accounting/mod_union_table-inl.h b/runtime/gc/accounting/mod_union_table-inl.h
index 29450c1..fb425df 100644
--- a/runtime/gc/accounting/mod_union_table-inl.h
+++ b/runtime/gc/accounting/mod_union_table-inl.h
@@ -28,9 +28,11 @@
 // A mod-union table to record image references to the Zygote and alloc space.
 class ModUnionTableToZygoteAllocspace : public ModUnionTableReferenceCache {
  public:
-  explicit ModUnionTableToZygoteAllocspace(Heap* heap) : ModUnionTableReferenceCache(heap) {}
+  explicit ModUnionTableToZygoteAllocspace(const std::string& name, Heap* heap,
+                                           space::ContinuousSpace* space)
+      : ModUnionTableReferenceCache(name, heap, space) {}
 
-  bool AddReference(const mirror::Object* /* obj */, const mirror::Object* ref) {
+  bool AddReference(const mirror::Object* /* obj */, const mirror::Object* ref) ALWAYS_INLINE {
     const std::vector<space::ContinuousSpace*>& spaces = GetHeap()->GetContinuousSpaces();
     typedef std::vector<space::ContinuousSpace*>::const_iterator It;
     for (It it = spaces.begin(); it != spaces.end(); ++it) {
@@ -47,16 +49,18 @@
 // A mod-union table to record Zygote references to the alloc space.
 class ModUnionTableToAllocspace : public ModUnionTableReferenceCache {
  public:
-  explicit ModUnionTableToAllocspace(Heap* heap) : ModUnionTableReferenceCache(heap) {}
+  explicit ModUnionTableToAllocspace(const std::string& name, Heap* heap,
+                                     space::ContinuousSpace* space)
+      : ModUnionTableReferenceCache(name, heap, space) {}
 
-  bool AddReference(const mirror::Object* /* obj */, const mirror::Object* ref) {
+  bool AddReference(const mirror::Object* /* obj */, const mirror::Object* ref) ALWAYS_INLINE {
     const std::vector<space::ContinuousSpace*>& spaces = GetHeap()->GetContinuousSpaces();
     typedef std::vector<space::ContinuousSpace*>::const_iterator It;
     for (It it = spaces.begin(); it != spaces.end(); ++it) {
       space::ContinuousSpace* space = *it;
       if (space->Contains(ref)) {
         // The allocation space is always considered for collection whereas the Zygote space is
-        //
+        // only considered for full GC.
         return space->GetGcRetentionPolicy() == space::kGcRetentionPolicyAlwaysCollect;
       }
     }
diff --git a/runtime/gc/accounting/mod_union_table.cc b/runtime/gc/accounting/mod_union_table.cc
index 4865219..7cbe94d 100644
--- a/runtime/gc/accounting/mod_union_table.cc
+++ b/runtime/gc/accounting/mod_union_table.cc
@@ -19,6 +19,7 @@
 #include "base/stl_util.h"
 #include "card_table-inl.h"
 #include "heap_bitmap.h"
+#include "gc/collector/mark_sweep.h"
 #include "gc/collector/mark_sweep-inl.h"
 #include "gc/heap.h"
 #include "gc/space/space.h"
@@ -67,60 +68,87 @@
   std::vector<byte*>* const cleared_cards_;
 };
 
-class ModUnionScanImageRootVisitor {
+class ModUnionUpdateObjectReferencesVisitor {
  public:
-  explicit ModUnionScanImageRootVisitor(collector::MarkSweep* const mark_sweep)
-      : mark_sweep_(mark_sweep) {}
+  ModUnionUpdateObjectReferencesVisitor(RootVisitor visitor, void* arg)
+    : visitor_(visitor),
+      arg_(arg) {
+  }
 
-  void operator()(const Object* root) const
-      EXCLUSIVE_LOCKS_REQUIRED(Locks::heap_bitmap_lock_)
-      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
-    DCHECK(root != NULL);
-    mark_sweep_->ScanRoot(root);
+  // Extra parameters are required since we use this same visitor signature for checking objects.
+  void operator()(Object* obj, Object* ref, const MemberOffset& offset,
+                  bool /* is_static */) const SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    // Only add the reference if it is non null and fits our criteria.
+    if (ref != nullptr) {
+      Object* new_ref = visitor_(ref, arg_);
+      if (new_ref != ref) {
+        obj->SetFieldObject(offset, ref, false, true);
+      }
+    }
   }
 
  private:
-  collector::MarkSweep* const mark_sweep_;
+  RootVisitor* visitor_;
+  void* arg_;
 };
 
-void ModUnionTableReferenceCache::ClearCards(space::ContinuousSpace* space) {
+class ModUnionScanImageRootVisitor {
+ public:
+  ModUnionScanImageRootVisitor(RootVisitor visitor, void* arg)
+      : visitor_(visitor), arg_(arg) {}
+
+  void operator()(Object* root) const
+      EXCLUSIVE_LOCKS_REQUIRED(Locks::heap_bitmap_lock_)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    DCHECK(root != NULL);
+    ModUnionUpdateObjectReferencesVisitor ref_visitor(visitor_, arg_);
+    collector::MarkSweep::VisitObjectReferences(root, ref_visitor, true);
+  }
+
+ private:
+  RootVisitor* visitor_;
+  void* arg_;
+};
+
+void ModUnionTableReferenceCache::ClearCards() {
   CardTable* card_table = GetHeap()->GetCardTable();
   ModUnionClearCardSetVisitor visitor(&cleared_cards_);
   // Clear dirty cards in the this space and update the corresponding mod-union bits.
-  card_table->ModifyCardsAtomic(space->Begin(), space->End(), AgeCardVisitor(), visitor);
+  card_table->ModifyCardsAtomic(space_->Begin(), space_->End(), AgeCardVisitor(), visitor);
 }
 
 class AddToReferenceArrayVisitor {
  public:
   explicit AddToReferenceArrayVisitor(ModUnionTableReferenceCache* mod_union_table,
-                                      std::vector<const Object*>* references)
+                                      std::vector<Object**>* references)
     : mod_union_table_(mod_union_table),
       references_(references) {
   }
 
   // Extra parameters are required since we use this same visitor signature for checking objects.
-  void operator()(const Object* obj, const Object* ref, const MemberOffset& /* offset */,
+  void operator()(Object* obj, Object* ref, const MemberOffset& offset,
                   bool /* is_static */) const {
     // Only add the reference if it is non null and fits our criteria.
-    if (ref != NULL && mod_union_table_->AddReference(obj, ref)) {
-      references_->push_back(ref);
+    if (ref != nullptr && mod_union_table_->AddReference(obj, ref)) {
+      // Push the adddress of the reference.
+      references_->push_back(obj->GetFieldObjectAddr(offset));
     }
   }
 
  private:
   ModUnionTableReferenceCache* const mod_union_table_;
-  std::vector<const Object*>* const references_;
+  std::vector<Object**>* const references_;
 };
 
 class ModUnionReferenceVisitor {
  public:
   explicit ModUnionReferenceVisitor(ModUnionTableReferenceCache* const mod_union_table,
-                                    std::vector<const Object*>* references)
+                                    std::vector<Object**>* references)
     : mod_union_table_(mod_union_table),
       references_(references) {
   }
 
-  void operator()(const Object* obj) const
+  void operator()(Object* obj) const
       SHARED_LOCKS_REQUIRED(Locks::heap_bitmap_lock_, Locks::mutator_lock_) {
     DCHECK(obj != NULL);
     // We don't have an early exit since we use the visitor pattern, an early
@@ -130,7 +158,7 @@
   }
  private:
   ModUnionTableReferenceCache* const mod_union_table_;
-  std::vector<const Object*>* const references_;
+  std::vector<Object**>* const references_;
 };
 
 class CheckReferenceVisitor {
@@ -143,8 +171,8 @@
 
   // Extra parameters are required since we use this same visitor signature for checking objects.
   // TODO: Fixme when anotatalysis works with visitors.
-  void operator()(const Object* obj, const Object* ref, const MemberOffset& /* offset */,
-                  bool /* is_static */) const
+  void operator()(const Object* obj, const Object* ref,
+                  const MemberOffset& /* offset */, bool /* is_static */) const
       SHARED_LOCKS_REQUIRED(Locks::heap_bitmap_lock_, Locks::mutator_lock_) {
     Heap* heap = mod_union_table_->GetHeap();
     if (ref != NULL && mod_union_table_->AddReference(obj, ref) &&
@@ -174,7 +202,7 @@
       : mod_union_table_(mod_union_table), references_(references) {
   }
 
-  void operator()(const Object* obj) const NO_THREAD_SAFETY_ANALYSIS {
+  void operator()(Object* obj) const NO_THREAD_SAFETY_ANALYSIS {
     Locks::heap_bitmap_lock_->AssertSharedHeld(Thread::Current());
     DCHECK(obj != NULL);
     CheckReferenceVisitor visitor(mod_union_table_, references_);
@@ -188,26 +216,25 @@
 
 void ModUnionTableReferenceCache::Verify() {
   // Start by checking that everything in the mod union table is marked.
-  Heap* heap = GetHeap();
-  for (const std::pair<const byte*, std::vector<const Object*> >& it : references_) {
-    for (const Object* ref : it.second) {
-      CHECK(heap->IsLiveObjectLocked(ref));
+  for (const auto& ref_pair : references_) {
+    for (Object** ref : ref_pair.second) {
+      CHECK(heap_->IsLiveObjectLocked(*ref));
     }
   }
 
   // Check the references of each clean card which is also in the mod union table.
-  CardTable* card_table = heap->GetCardTable();
-  for (const std::pair<const byte*, std::vector<const Object*> > & it : references_) {
-    const byte* card = it.first;
+  CardTable* card_table = heap_->GetCardTable();
+  SpaceBitmap* live_bitmap = space_->GetLiveBitmap();
+  for (const auto& ref_pair : references_) {
+    const byte* card = ref_pair.first;
     if (*card == CardTable::kCardClean) {
-      std::set<const Object*> reference_set(it.second.begin(), it.second.end());
+      std::set<const Object*> reference_set;
+      for (Object** obj_ptr : ref_pair.second) {
+        reference_set.insert(*obj_ptr);
+      }
       ModUnionCheckReferences visitor(this, reference_set);
       uintptr_t start = reinterpret_cast<uintptr_t>(card_table->AddrFromCard(card));
-      uintptr_t end = start + CardTable::kCardSize;
-      auto* space = heap->FindContinuousSpaceFromObject(reinterpret_cast<Object*>(start), false);
-      DCHECK(space != nullptr);
-      SpaceBitmap* live_bitmap = space->GetLiveBitmap();
-      live_bitmap->VisitMarkedRange(start, end, visitor);
+      live_bitmap->VisitMarkedRange(start, start + CardTable::kCardSize, visitor);
     }
   }
 }
@@ -221,24 +248,24 @@
     os << reinterpret_cast<void*>(start) << "-" << reinterpret_cast<void*>(end) << ",";
   }
   os << "]\nModUnionTable references: [";
-  for (const std::pair<const byte*, std::vector<const Object*> >& it : references_) {
-    const byte* card_addr = it.first;
+  for (const auto& ref_pair : references_) {
+    const byte* card_addr = ref_pair.first;
     uintptr_t start = reinterpret_cast<uintptr_t>(card_table->AddrFromCard(card_addr));
     uintptr_t end = start + CardTable::kCardSize;
     os << reinterpret_cast<void*>(start) << "-" << reinterpret_cast<void*>(end) << "->{";
-    for (const mirror::Object* ref : it.second) {
-      os << reinterpret_cast<const void*>(ref) << ",";
+    for (Object** ref : ref_pair.second) {
+      os << reinterpret_cast<const void*>(*ref) << ",";
     }
     os << "},";
   }
 }
 
-void ModUnionTableReferenceCache::Update() {
+void ModUnionTableReferenceCache::UpdateAndMarkReferences(RootVisitor visitor, void* arg) {
   Heap* heap = GetHeap();
   CardTable* card_table = heap->GetCardTable();
 
-  std::vector<const Object*> cards_references;
-  ModUnionReferenceVisitor visitor(this, &cards_references);
+  std::vector<Object**> cards_references;
+  ModUnionReferenceVisitor add_visitor(this, &cards_references);
 
   for (const auto& card : cleared_cards_) {
     // Clear and re-compute alloc space references associated with this card.
@@ -248,7 +275,7 @@
     auto* space = heap->FindContinuousSpaceFromObject(reinterpret_cast<Object*>(start), false);
     DCHECK(space != nullptr);
     SpaceBitmap* live_bitmap = space->GetLiveBitmap();
-    live_bitmap->VisitMarkedRange(start, end, visitor);
+    live_bitmap->VisitMarkedRange(start, end, add_visitor);
 
     // Update the corresponding references for the card.
     auto found = references_.find(card);
@@ -263,46 +290,41 @@
     }
   }
   cleared_cards_.clear();
-}
-
-void ModUnionTableReferenceCache::MarkReferences(collector::MarkSweep* mark_sweep) {
   size_t count = 0;
-
   for (const auto& ref : references_) {
-    for (const auto& obj : ref.second) {
-      mark_sweep->MarkRoot(obj);
-      ++count;
+    for (const auto& obj_ptr : ref.second) {
+      Object* obj = *obj_ptr;
+      if (obj != nullptr) {
+        Object* new_obj = visitor(obj, arg);
+        // Avoid dirtying pages in the image unless necessary.
+        if (new_obj != obj) {
+          *obj_ptr = new_obj;
+        }
+      }
     }
+    count += ref.second.size();
   }
   if (VLOG_IS_ON(heap)) {
     VLOG(gc) << "Marked " << count << " references in mod union table";
   }
 }
 
-void ModUnionTableCardCache::ClearCards(space::ContinuousSpace* space) {
+void ModUnionTableCardCache::ClearCards() {
   CardTable* card_table = GetHeap()->GetCardTable();
   ModUnionClearCardSetVisitor visitor(&cleared_cards_);
   // Clear dirty cards in the this space and update the corresponding mod-union bits.
-  card_table->ModifyCardsAtomic(space->Begin(), space->End(), AgeCardVisitor(), visitor);
+  card_table->ModifyCardsAtomic(space_->Begin(), space_->End(), AgeCardVisitor(), visitor);
 }
 
 // Mark all references to the alloc space(s).
-void ModUnionTableCardCache::MarkReferences(collector::MarkSweep* mark_sweep) {
+void ModUnionTableCardCache::UpdateAndMarkReferences(RootVisitor visitor, void* arg) {
   CardTable* card_table = heap_->GetCardTable();
-  ModUnionScanImageRootVisitor visitor(mark_sweep);
-  space::ContinuousSpace* space = nullptr;
-  SpaceBitmap* bitmap = nullptr;
+  ModUnionScanImageRootVisitor scan_visitor(visitor, arg);
+  SpaceBitmap* bitmap = space_->GetLiveBitmap();
   for (const byte* card_addr : cleared_cards_) {
-    auto start = reinterpret_cast<uintptr_t>(card_table->AddrFromCard(card_addr));
-    auto end = start + CardTable::kCardSize;
-    auto obj_start = reinterpret_cast<Object*>(start);
-    if (UNLIKELY(space == nullptr || !space->Contains(obj_start))) {
-      space = heap_->FindContinuousSpaceFromObject(obj_start, false);
-      DCHECK(space != nullptr);
-      bitmap = space->GetLiveBitmap();
-      DCHECK(bitmap != nullptr);
-    }
-    bitmap->VisitMarkedRange(start, end, visitor);
+    uintptr_t start = reinterpret_cast<uintptr_t>(card_table->AddrFromCard(card_addr));
+    DCHECK(space_->HasAddress(reinterpret_cast<Object*>(start)));
+    bitmap->VisitMarkedRange(start, start + CardTable::kCardSize, scan_visitor);
   }
 }
 
diff --git a/runtime/gc/accounting/mod_union_table.h b/runtime/gc/accounting/mod_union_table.h
index eb7a754..d874c60 100644
--- a/runtime/gc/accounting/mod_union_table.h
+++ b/runtime/gc/accounting/mod_union_table.h
@@ -19,6 +19,7 @@
 
 #include "gc_allocator.h"
 #include "globals.h"
+#include "root_visitor.h"
 #include "safe_map.h"
 
 #include <set>
@@ -52,21 +53,23 @@
  public:
   typedef std::set<byte*, std::less<byte*>, GCAllocator<byte*> > CardSet;
 
-  explicit ModUnionTable(Heap* heap) : heap_(heap) {}
+  explicit ModUnionTable(const std::string& name, Heap* heap, space::ContinuousSpace* space)
+      : name_(name),
+        heap_(heap),
+        space_(space) {
+  }
 
   virtual ~ModUnionTable() {}
 
   // Clear cards which map to a memory range of a space. This doesn't immediately update the
   // mod-union table, as updating the mod-union table may have an associated cost, such as
   // determining references to track.
-  virtual void ClearCards(space::ContinuousSpace* space) = 0;
+  virtual void ClearCards() = 0;
 
   // Update the mod-union table using data stored by ClearCards. There may be multiple ClearCards
-  // before a call to update, for example, back-to-back sticky GCs.
-  virtual void Update() = 0;
-
-  // Mark the bitmaps for all references which are stored in the mod-union table.
-  virtual void MarkReferences(collector::MarkSweep* mark_sweep) = 0;
+  // before a call to update, for example, back-to-back sticky GCs. Also mark references to other
+  // spaces which are stored in the mod-union table.
+  virtual void UpdateAndMarkReferences(RootVisitor visitor, void* arg) = 0;
 
   // Verification, sanity checks that we don't have clean cards which conflict with out cached data
   // for said cards. Exclusive lock is required since verify sometimes uses
@@ -75,31 +78,35 @@
   virtual void Verify() EXCLUSIVE_LOCKS_REQUIRED(Locks::heap_bitmap_lock_) = 0;
 
   virtual void Dump(std::ostream& os) = 0;
-
+  space::ContinuousSpace* GetSpace() {
+    return space_;
+  }
   Heap* GetHeap() const {
     return heap_;
   }
+  const std::string& GetName() const {
+    return name_;
+  }
 
  protected:
+  const std::string name_;
   Heap* const heap_;
+  space::ContinuousSpace* const space_;
 };
 
 // Reference caching implementation. Caches references pointing to alloc space(s) for each card.
 class ModUnionTableReferenceCache : public ModUnionTable {
  public:
-  explicit ModUnionTableReferenceCache(Heap* heap) : ModUnionTable(heap) {}
+  explicit ModUnionTableReferenceCache(const std::string& name, Heap* heap,
+                                       space::ContinuousSpace* space)
+      : ModUnionTable(name, heap, space) {}
   virtual ~ModUnionTableReferenceCache() {}
 
   // Clear and store cards for a space.
-  void ClearCards(space::ContinuousSpace* space);
+  void ClearCards();
 
-  // Update table based on cleared cards.
-  void Update()
-      EXCLUSIVE_LOCKS_REQUIRED(Locks::heap_bitmap_lock_)
-      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
-
-  // Mark all references to the alloc space(s).
-  void MarkReferences(collector::MarkSweep* mark_sweep)
+  // Update table based on cleared cards and mark all references to the other spaces.
+  void UpdateAndMarkReferences(RootVisitor visitor, void* arg)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
       EXCLUSIVE_LOCKS_REQUIRED(Locks::heap_bitmap_lock_);
 
@@ -117,24 +124,22 @@
   ModUnionTable::CardSet cleared_cards_;
 
   // Maps from dirty cards to their corresponding alloc space references.
-  SafeMap<const byte*, std::vector<const mirror::Object*>, std::less<const byte*>,
-    GCAllocator<std::pair<const byte*, std::vector<const mirror::Object*> > > > references_;
+  SafeMap<const byte*, std::vector<mirror::Object**>, std::less<const byte*>,
+    GCAllocator<std::pair<const byte*, std::vector<mirror::Object**> > > > references_;
 };
 
 // Card caching implementation. Keeps track of which cards we cleared and only this information.
 class ModUnionTableCardCache : public ModUnionTable {
  public:
-  explicit ModUnionTableCardCache(Heap* heap) : ModUnionTable(heap) {}
+  explicit ModUnionTableCardCache(const std::string& name, Heap* heap, space::ContinuousSpace* space)
+      : ModUnionTable(name, heap, space) {}
   virtual ~ModUnionTableCardCache() {}
 
   // Clear and store cards for a space.
-  void ClearCards(space::ContinuousSpace* space);
-
-  // Nothing to update as all dirty cards were placed into cleared cards during clearing.
-  void Update() {}
+  void ClearCards();
 
   // Mark all references to the alloc space(s).
-  void MarkReferences(collector::MarkSweep* mark_sweep)
+  void UpdateAndMarkReferences(RootVisitor visitor, void* arg)
       EXCLUSIVE_LOCKS_REQUIRED(Locks::heap_bitmap_lock_)
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
diff --git a/runtime/gc/accounting/space_bitmap.h b/runtime/gc/accounting/space_bitmap.h
index f975692..4cf8872 100644
--- a/runtime/gc/accounting/space_bitmap.h
+++ b/runtime/gc/accounting/space_bitmap.h
@@ -247,8 +247,8 @@
 
   template <typename Visitor>
   void Visit(const Visitor& visitor) NO_THREAD_SAFETY_ANALYSIS {
-    for (Objects::iterator it = contained_.begin(); it != contained_.end(); ++it) {
-      visitor(*it);
+    for (const mirror::Object* obj : contained_) {
+      visitor(const_cast<mirror::Object*>(obj));
     }
   }