Fix app image class pruning in `CompilerDriver`.

Fix app image class pruning when loading image classes and
after verification and initialization. This is required to
keep `CompilerOptions::IsImageClass()` consistent with the
actual content of the app image produced by `ImageWriter`.
Flag a bug for `--initialize-app-image-classes=true` where
the inconsistency still remains.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing --speed-profile
Bug: 38313278
Change-Id: I196d97c69bccb255fb3190a6b636596805ab55bd

11 files changed, 176 insertions, 20 deletions

diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index c2c4fb27a1..9f2d263de3 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -1966,6 +1966,9 @@ class Dex2Oat final {
         }
       }
     }
+    if (IsAppImage()) {
+      AotClassLinker::SetAppImageDexFiles(&compiler_options_->GetDexFilesForOatFile());
+    }
 
     // Register dex caches and key them to the class loader so that they only unload when the
     // class loader unloads.
diff --git a/dex2oat/driver/compiler_driver.cc b/dex2oat/driver/compiler_driver.cc
index 9ee7e64cad..f7f86ed4fc 100644
--- a/dex2oat/driver/compiler_driver.cc
+++ b/dex2oat/driver/compiler_driver.cc
@@ -889,6 +889,12 @@ void CompilerDriver::PreCompile(jobject class_loader,
                                << "Please check the log.";
       _exit(1);
     }
+
+    if (GetCompilerOptions().IsAppImage() && had_hard_verifier_failure_) {
+      // Prune erroneous classes and classes that depend on them.
+      UpdateImageClasses(timings, image_classes);
+      VLOG(compiler) << "verify/UpdateImageClasses: " << GetMemoryUsageString(false);
+    }
   }
 
   if (GetCompilerOptions().IsGeneratingImage()) {
@@ -911,8 +917,10 @@ void CompilerDriver::PreCompile(jobject class_loader,
       visitor.FillAllIMTAndConflictTables();
     }
 
-    UpdateImageClasses(timings, image_classes);
-    VLOG(compiler) << "UpdateImageClasses: " << GetMemoryUsageString(false);
+    if (GetCompilerOptions().IsBootImage() || GetCompilerOptions().IsBootImageExtension()) {
+      UpdateImageClasses(timings, image_classes);
+      VLOG(compiler) << "UpdateImageClasses: " << GetMemoryUsageString(false);
+    }
 
     if (kBitstringSubtypeCheckEnabled &&
         GetCompilerOptions().IsForceDeterminism() && GetCompilerOptions().IsBootImage()) {
@@ -1380,7 +1388,8 @@ class ClinitImageUpdate {
     bool operator()(ObjPtr<mirror::Class> klass) override REQUIRES_SHARED(Locks::mutator_lock_) {
       bool resolved = klass->IsResolved();
       DCHECK(resolved || klass->IsErroneousUnresolved());
-      bool can_include_in_image = LIKELY(resolved) && CanIncludeInCurrentImage(klass);
+      bool can_include_in_image =
+          LIKELY(resolved) && LIKELY(!klass->IsErroneous()) && CanIncludeInCurrentImage(klass);
       std::string temp;
       std::string_view descriptor(klass->GetDescriptor(&temp));
       auto it = data_->image_class_descriptors_->find(descriptor);
@@ -1451,17 +1460,16 @@ class ClinitImageUpdate {
 
 void CompilerDriver::UpdateImageClasses(TimingLogger* timings,
                                         /*inout*/ HashSet<std::string>* image_classes) {
-  if (GetCompilerOptions().IsBootImage() || GetCompilerOptions().IsBootImageExtension()) {
-    TimingLogger::ScopedTiming t("UpdateImageClasses", timings);
+  DCHECK(GetCompilerOptions().IsGeneratingImage());
+  TimingLogger::ScopedTiming t("UpdateImageClasses", timings);
 
-    // Suspend all threads.
-    ScopedSuspendAll ssa(__FUNCTION__);
+  // Suspend all threads.
+  ScopedSuspendAll ssa(__FUNCTION__);
 
-    ClinitImageUpdate update(image_classes, Thread::Current());
+  ClinitImageUpdate update(image_classes, Thread::Current());
 
-    // Do the marking.
-    update.Walk();
-  }
+  // Do the marking.
+  update.Walk();
 }
 
 void CompilerDriver::ProcessedInstanceField(bool resolved) {
@@ -2039,7 +2047,8 @@ class VerifyClassVisitor : public CompilationVisitor {
                                                klass,
                                                log_level_);
 
-      if (klass->IsErroneous()) {
+      DCHECK_EQ(klass->IsErroneous(), failure_kind == verifier::FailureKind::kHardFailure);
+      if (failure_kind == verifier::FailureKind::kHardFailure) {
         // ClassLinker::VerifyClass throws, which isn't useful in the compiler.
         CHECK(soa.Self()->IsExceptionPending());
         soa.Self()->ClearException();
diff --git a/dex2oat/linker/image_writer.cc b/dex2oat/linker/image_writer.cc
index 63ac71fc12..abc3354421 100644
--- a/dex2oat/linker/image_writer.cc
+++ b/dex2oat/linker/image_writer.cc
@@ -1101,7 +1101,17 @@ bool ImageWriter::KeepClass(ObjPtr<mirror::Class> klass) {
     // the boot image spaces since these may have already been loaded at
     // run time when this image is loaded. Keep classes in the boot image
     // spaces we're compiling against since we don't want to re-resolve these.
-    return !PruneImageClass(klass);
+    // FIXME: Update image classes in the `CompilerOptions` after initializing classes
+    // with `--initialize-app-image-classes=true`. This experimental flag can currently
+    // cause an inconsistency between `CompilerOptions::IsImageClass()` and what actually
+    // ends up in the app image as seen in the run-test `660-clinit` where the class
+    // `ObjectRef` is considered an app image class during compilation but in the end
+    // it's pruned here. This inconsistency should be fixed if we want to properly
+    // initialize app image classes. b/38313278
+    bool keep = !PruneImageClass(klass);
+    CHECK_IMPLIES(!compiler_options_.InitializeAppImageClasses(), keep)
+        << klass->PrettyDescriptor();
+    return keep;
   }
   return true;
 }
diff --git a/runtime/oat/aot_class_linker.cc b/runtime/oat/aot_class_linker.cc
index 6d6887595c..09fb92452d 100644
--- a/runtime/oat/aot_class_linker.cc
+++ b/runtime/oat/aot_class_linker.cc
@@ -16,6 +16,7 @@
 
 #include "aot_class_linker.h"
 
+#include "base/stl_util.h"
 #include "class_status.h"
 #include "compiler_callbacks.h"
 #include "dex/class_reference.h"
@@ -131,6 +132,12 @@ verifier::FailureKind AotClassLinker::PerformClassVerification(
   return ClassLinker::PerformClassVerification(self, verifier_deps, klass, log_level, error_msg);
 }
 
+static const std::vector<const DexFile*>* gAppImageDexFiles = nullptr;
+
+void AotClassLinker::SetAppImageDexFiles(const std::vector<const DexFile*>* app_image_dex_files) {
+  gAppImageDexFiles = app_image_dex_files;
+}
+
 bool AotClassLinker::CanReferenceInBootImageExtensionOrAppImage(
     ObjPtr<mirror::Class> klass, gc::Heap* heap) {
   // Do not allow referencing a class or instance of a class defined in a dex file
@@ -160,8 +167,25 @@ bool AotClassLinker::CanReferenceInBootImageExtensionOrAppImage(
     }
   }
 
+  auto can_reference_dex_cache = [&](ObjPtr<mirror::DexCache> dex_cache)
+      REQUIRES_SHARED(Locks::mutator_lock_) {
+    // We cannot reference a boot image dex cache for classes
+    // that were not themselves in the boot image.
+    if (heap->ObjectIsInBootImageSpace(dex_cache)) {
+      return false;
+    }
+    // App image compilation can pull in dex files from parent or library class loaders.
+    // Classes from such dex files cannot be included or referenced in the current app image
+    // to avoid conflicts with classes in the parent or library class loader's app image.
+    if (gAppImageDexFiles != nullptr &&
+        !ContainsElement(*gAppImageDexFiles, dex_cache->GetDexFile())) {
+      return false;
+    }
+    return true;
+  };
+
   // Check the class itself.
-  if (heap->ObjectIsInBootImageSpace(klass->GetDexCache())) {
+  if (!can_reference_dex_cache(klass->GetDexCache())) {
     return false;
   }
 
@@ -169,7 +193,7 @@ bool AotClassLinker::CanReferenceInBootImageExtensionOrAppImage(
   ObjPtr<mirror::Class> superclass = klass->GetSuperClass();
   while (!heap->ObjectIsInBootImageSpace(superclass)) {
     DCHECK(superclass != nullptr);  // Cannot skip Object which is in the primary boot image.
-    if (heap->ObjectIsInBootImageSpace(superclass->GetDexCache())) {
+    if (!can_reference_dex_cache(superclass->GetDexCache())) {
       return false;
     }
     superclass = superclass->GetSuperClass();
@@ -181,7 +205,7 @@ bool AotClassLinker::CanReferenceInBootImageExtensionOrAppImage(
     ObjPtr<mirror::Class> interface = if_table->GetInterface(i);
     DCHECK(interface != nullptr);
     if (!heap->ObjectIsInBootImageSpace(interface) &&
-        heap->ObjectIsInBootImageSpace(interface->GetDexCache())) {
+        !can_reference_dex_cache(interface->GetDexCache())) {
       return false;
     }
   }
@@ -194,7 +218,7 @@ bool AotClassLinker::CanReferenceInBootImageExtensionOrAppImage(
       for (auto& m : k->GetVirtualMethods(pointer_size)) {
         ObjPtr<mirror::Class> declaring_class = m.GetDeclaringClass();
         CHECK(heap->ObjectIsInBootImageSpace(declaring_class) ||
-              !heap->ObjectIsInBootImageSpace(declaring_class->GetDexCache()));
+              can_reference_dex_cache(declaring_class->GetDexCache()));
       }
       k = k->GetSuperClass();
     }
diff --git a/runtime/oat/aot_class_linker.h b/runtime/oat/aot_class_linker.h
index 4396a50efd..18689bed7f 100644
--- a/runtime/oat/aot_class_linker.h
+++ b/runtime/oat/aot_class_linker.h
@@ -35,6 +35,8 @@ class AotClassLinker : public ClassLinker {
   explicit AotClassLinker(InternTable *intern_table);
   ~AotClassLinker();
 
+  EXPORT static void SetAppImageDexFiles(const std::vector<const DexFile*>* app_image_dex_files);
+
   EXPORT static bool CanReferenceInBootImageExtensionOrAppImage(
       ObjPtr<mirror::Class> klass, gc::Heap* heap) REQUIRES_SHARED(Locks::mutator_lock_);
 
diff --git a/test/732-checker-app-image/expected-stdout.txt b/test/732-checker-app-image/expected-stdout.txt
index a8deeec6de..cb6b036407 100644
--- a/test/732-checker-app-image/expected-stdout.txt
+++ b/test/732-checker-app-image/expected-stdout.txt
@@ -1,2 +1,5 @@
 AppImageClass
 NonAppImageClass
+SecondaryAppImageClass
+SecondaryNonAppImageClass
+SecondaryClassWithUnmetDependency
diff --git a/test/732-checker-app-image/profile b/test/732-checker-app-image/profile
index 924c319453..a35f977d7c 100644
--- a/test/732-checker-app-image/profile
+++ b/test/732-checker-app-image/profile
@@ -5,3 +5,8 @@ SHLMain;->$noinline$getNonAppImageClass()Ljava/lang/Class;
 SHLMain;->$noinline$callAppImageClassNop()V
 SHLMain;->$noinline$callAppImageClassWithClinitNop()V
 SHLMain;->$noinline$callNonAppImageClassNop()V
+LSecondaryAppImageClass;
+LSecondaryClassWithUnmetDependency;
+SHLSecondary;->$noinline$getSecondaryAppImageClass()Ljava/lang/Class;
+SHLSecondary;->$noinline$getSecondaryNonAppImageClass()Ljava/lang/Class;
+SHLSecondary;->$noinline$getSecondaryClassWithUnmetDependency()Ljava/lang/Class;
diff --git a/test/732-checker-app-image/run.py b/test/732-checker-app-image/run.py
index 3dc165b8e8..16658852ad 100644
--- a/test/732-checker-app-image/run.py
+++ b/test/732-checker-app-image/run.py
@@ -16,5 +16,12 @@
 
 
 def run(ctx, args):
-  # We need a profile to tell dex2oat to include classes in the final app image
-  ctx.default_run(args, profile=True)
+  if args.jvm:
+    ctx.default_run(args)
+  else:
+    ctx.default_run(
+        args,
+        # We need a profile to tell dex2oat to include classes in the final app image
+        profile=True,
+        # Append graphs for checker tests (we run dex2oat twice) with `--dump-cfg-append`.
+        Xcompiler_option=["--dump-cfg-append"])
diff --git a/test/732-checker-app-image/src-ex/Secondary.java b/test/732-checker-app-image/src-ex/Secondary.java
new file mode 100644
index 0000000000..52354856d6
--- /dev/null
+++ b/test/732-checker-app-image/src-ex/Secondary.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public class Secondary {
+    public static void main() {
+        System.out.println($noinline$getSecondaryAppImageClass().getName());
+        System.out.println($noinline$getSecondaryNonAppImageClass().getName());
+        System.out.println($noinline$getSecondaryClassWithUnmetDependency().getName());
+    }
+
+    /// CHECK-START: java.lang.Class Secondary.$noinline$getSecondaryAppImageClass() builder (after)
+    /// CHECK:            LoadClass load_kind:BssEntry in_image:true
+    public static Class<?> $noinline$getSecondaryAppImageClass() {
+        return SecondaryAppImageClass.class;
+    }
+
+    /// CHECK-START: java.lang.Class Secondary.$noinline$getSecondaryNonAppImageClass() builder (after)
+    /// CHECK:            LoadClass load_kind:BssEntry in_image:false
+    public static Class<?> $noinline$getSecondaryNonAppImageClass() {
+        return SecondaryNonAppImageClass.class;
+    }
+
+    /// CHECK-START: java.lang.Class Secondary.$noinline$getSecondaryClassWithUnmetDependency() builder (after)
+    /// CHECK:            LoadClass load_kind:BssEntry in_image:false
+    public static Class<?> $noinline$getSecondaryClassWithUnmetDependency() {
+        return SecondaryClassWithUnmetDependency.class;
+    }
+}
+
+class SecondaryAppImageClass {  // Included in the profile.
+}
+
+class SecondaryNonAppImageClass {  // Not included in the profile.
+}
+
+// Included in the profile but with interface defined in parent class loader
+// and therefore unsuitable for inclusion in the app image.
+class SecondaryClassWithUnmetDependency implements PrimaryInterface {}
diff --git a/test/732-checker-app-image/src/Main.java b/test/732-checker-app-image/src/Main.java
index 200708a20b..ac8d2254c9 100644
--- a/test/732-checker-app-image/src/Main.java
+++ b/test/732-checker-app-image/src/Main.java
@@ -14,14 +14,39 @@
  * limitations under the License.
  */
 
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
+
 public class Main {
-    public static void main(String args[]) {
+    public static String TEST_NAME = "732-checker-app-image";
+
+    public static ClassLoader getSecondaryClassLoader() throws Exception {
+        String location = System.getenv("DEX_LOCATION");
+        try {
+            Class<?> class_loader_class = Class.forName("dalvik.system.PathClassLoader");
+            Constructor<?> ctor =
+                    class_loader_class.getConstructor(String.class, ClassLoader.class);
+            return (ClassLoader) ctor.newInstance(location + "/" + TEST_NAME + "-ex.jar",
+                                                  Main.class.getClassLoader());
+        } catch (ClassNotFoundException e) {
+            // Running on RI. Use URLClassLoader.
+            return new java.net.URLClassLoader(
+                    new java.net.URL[] { new java.net.URL("file://" + location + "/classes-ex/") });
+        }
+    }
+
+    public static void main(String args[]) throws Exception {
         System.out.println($noinline$getAppImageClass().getName());
         System.out.println($noinline$getNonAppImageClass().getName());
 
         $noinline$callAppImageClassNop();
         $noinline$callAppImageClassWithClinitNop();
         $noinline$callNonAppImageClassNop();
+
+        ClassLoader secondaryLoader = getSecondaryClassLoader();
+        Class<?> secondaryClass = Class.forName("Secondary", true, secondaryLoader);
+        Method secondaryMain = secondaryClass.getMethod("main");
+        secondaryMain.invoke(null);
     }
 
     /// CHECK-START: java.lang.Class Main.$noinline$getAppImageClass() builder (after)
diff --git a/test/732-checker-app-image/src/PrimaryInterface.java b/test/732-checker-app-image/src/PrimaryInterface.java
new file mode 100644
index 0000000000..4b2e3f4a7e
--- /dev/null
+++ b/test/732-checker-app-image/src/PrimaryInterface.java
@@ -0,0 +1,17 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public interface PrimaryInterface {}