2 files changed, 56 insertions, 0 deletions

diff --git a/compiler/driver/compiler_driver.cc b/compiler/driver/compiler_driver.cc
index 70cbb01569..273bd50b9e 100644
--- a/compiler/driver/compiler_driver.cc
+++ b/compiler/driver/compiler_driver.cc
@@ -872,6 +872,14 @@ void CompilerDriver::PreCompile(jobject class_loader,
                                 TimingLogger* timings) {
   CheckThreadPools();
 
+  if (kUseBitstringTypeCheck &&
+      !compiler_options_->IsBootImage() &&
+      compiler_options_->IsAotCompilationEnabled()) {
+    RecordBootImageClassesWithAssignedBitstring();
+    VLOG(compiler) << "RecordBootImageClassesWithAssignedBitstring: "
+        << GetMemoryUsageString(false);
+  }
+
   LoadImageClasses(timings);
   VLOG(compiler) << "LoadImageClasses: " << GetMemoryUsageString(false);
 
@@ -940,6 +948,43 @@ void CompilerDriver::PreCompile(jobject class_loader,
   }
 }
 
+void CompilerDriver::RecordBootImageClassesWithAssignedBitstring() {
+  if (boot_image_classes_with_assigned_bitstring_ != nullptr) {
+    return;  // Already recorded. (Happens because of class unloading between dex files.)
+  }
+
+  class Visitor : public ClassVisitor {
+   public:
+    explicit Visitor(std::unordered_set<mirror::Class*>* recorded_classes)
+        : recorded_classes_(recorded_classes) {}
+
+    bool operator()(ObjPtr<mirror::Class> klass) OVERRIDE
+        REQUIRES(Locks::subtype_check_lock_) REQUIRES_SHARED(Locks::mutator_lock_) {
+      DCHECK(klass != nullptr);
+      SubtypeCheckInfo::State state = SubtypeCheck<ObjPtr<mirror::Class>>::GetState(klass);
+      if (state == SubtypeCheckInfo::kAssigned) {
+        recorded_classes_->insert(klass.Ptr());
+      }
+      return true;
+    }
+
+   private:
+    std::unordered_set<mirror::Class*>* const recorded_classes_;
+  };
+
+  boot_image_classes_with_assigned_bitstring_.reset(new std::unordered_set<mirror::Class*>());
+  Visitor visitor(boot_image_classes_with_assigned_bitstring_.get());
+  ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  ScopedObjectAccess soa(Thread::Current());
+  MutexLock subtype_check_lock(soa.Self(), *Locks::subtype_check_lock_);
+  class_linker->VisitClasses(&visitor);
+}
+
+bool CompilerDriver::IsBootImageClassWithAssignedBitstring(ObjPtr<mirror::Class> klass) {
+  DCHECK(boot_image_classes_with_assigned_bitstring_ != nullptr);
+  return boot_image_classes_with_assigned_bitstring_->count(klass.Ptr()) != 0u;
+}
+
 bool CompilerDriver::IsImageClass(const char* descriptor) const {
   if (image_classes_ != nullptr) {
     // If we have a set of image classes, use those.
diff --git a/compiler/driver/compiler_driver.h b/compiler/driver/compiler_driver.h
index 4b5916d572..30042751b3 100644
--- a/compiler/driver/compiler_driver.h
+++ b/compiler/driver/compiler_driver.h
@@ -385,12 +385,17 @@ class CompilerDriver {
     return dex_to_dex_compiler_;
   }
 
+  bool IsBootImageClassWithAssignedBitstring(ObjPtr<mirror::Class> klass)
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
  private:
   void PreCompile(jobject class_loader,
                   const std::vector<const DexFile*>& dex_files,
                   TimingLogger* timings)
       REQUIRES(!Locks::mutator_lock_);
 
+  void RecordBootImageClassesWithAssignedBitstring() REQUIRES(!Locks::mutator_lock_);
+
   void LoadImageClasses(TimingLogger* timings) REQUIRES(!Locks::mutator_lock_);
 
   // Attempt to resolve all type, methods, fields, and strings
@@ -513,6 +518,12 @@ class CompilerDriver {
   // This option may be restricted to the boot image, depending on a flag in the implementation.
   std::unique_ptr<std::unordered_set<std::string>> methods_to_compile_;
 
+  // For AOT app compilation, we keep the set of boot image classes with assigned type check
+  // bitstring. We need to retrieve this set before we initialize app image classes as the
+  // initialization can cause more boot image bitstrings to be assigned.
+  // Note that boot image classes are non-moveable, so it's OK to keep raw pointers.
+  std::unique_ptr<std::unordered_set<mirror::Class*>> boot_image_classes_with_assigned_bitstring_;
+
   std::atomic<uint32_t> number_of_soft_verifier_failures_;
   bool had_hard_verifier_failure_;