1 files changed, 63 insertions, 28 deletions

diff --git a/compiler/optimizing/sharpening.cc b/compiler/optimizing/sharpening.cc
index 8637db13ad..b8471e3721 100644
--- a/compiler/optimizing/sharpening.cc
+++ b/compiler/optimizing/sharpening.cc
@@ -19,6 +19,7 @@
 #include "art_method-inl.h"
 #include "base/casts.h"
 #include "base/enums.h"
+#include "base/logging.h"
 #include "class_linker.h"
 #include "code_generator.h"
 #include "driver/compiler_options.h"
@@ -26,29 +27,29 @@
 #include "gc/heap.h"
 #include "gc/space/image_space.h"
 #include "handle_scope-inl.h"
+#include "jit/jit.h"
 #include "mirror/dex_cache.h"
 #include "mirror/string.h"
 #include "nodes.h"
 #include "runtime.h"
 #include "scoped_thread_state_change-inl.h"
-#include "utils/dex_cache_arrays_layout-inl.h"
 
 namespace art {
 
 static bool IsInBootImage(ArtMethod* method) {
-  const std::vector<gc::space::ImageSpace*>& image_spaces =
-      Runtime::Current()->GetHeap()->GetBootImageSpaces();
-  for (gc::space::ImageSpace* image_space : image_spaces) {
-    const ImageSection& method_section = image_space->GetImageHeader().GetMethodsSection();
-    if (method_section.Contains(reinterpret_cast<uint8_t*>(method) - image_space->Begin())) {
-      return true;
-    }
-  }
-  return false;
+  gc::Heap* heap = Runtime::Current()->GetHeap();
+  DCHECK_EQ(heap->IsBootImageAddress(method),
+            std::any_of(heap->GetBootImageSpaces().begin(),
+                        heap->GetBootImageSpaces().end(),
+                        [=](gc::space::ImageSpace* space) REQUIRES_SHARED(Locks::mutator_lock_) {
+                          return space->GetImageHeader().GetMethodsSection().Contains(
+                              reinterpret_cast<uint8_t*>(method) - space->Begin());
+                        }));
+  return heap->IsBootImageAddress(method);
 }
 
 static bool BootImageAOTCanEmbedMethod(ArtMethod* method, const CompilerOptions& compiler_options) {
-  DCHECK(compiler_options.IsBootImage());
+  DCHECK(compiler_options.IsBootImage() || compiler_options.IsBootImageExtension());
   ScopedObjectAccess soa(Thread::Current());
   ObjPtr<mirror::Class> klass = method->GetDeclaringClass();
   DCHECK(klass != nullptr);
@@ -86,10 +87,13 @@ HInvokeStaticOrDirect::DispatchInfo HSharpening::SharpenInvokeStaticOrDirect(
     // Recursive call.
     method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kRecursive;
     code_ptr_location = HInvokeStaticOrDirect::CodePtrLocation::kCallSelf;
-  } else if (compiler_options.IsBootImage()) {
+  } else if (compiler_options.IsBootImage() || compiler_options.IsBootImageExtension()) {
     if (!compiler_options.GetCompilePic()) {
       // Test configuration, do not sharpen.
       method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kRuntimeCall;
+    } else if (IsInBootImage(callee)) {
+      DCHECK(compiler_options.IsBootImageExtension());
+      method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kBootImageRelRo;
     } else if (BootImageAOTCanEmbedMethod(callee, compiler_options)) {
       method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kBootImageLinkTimePcRelative;
     } else {
@@ -98,11 +102,18 @@ HInvokeStaticOrDirect::DispatchInfo HSharpening::SharpenInvokeStaticOrDirect(
     }
     code_ptr_location = HInvokeStaticOrDirect::CodePtrLocation::kCallArtMethod;
   } else if (Runtime::Current()->UseJitCompilation()) {
-    // JIT or on-device AOT compilation referencing a boot image method.
-    // Use the method address directly.
-    method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kJitDirectAddress;
-    method_load_data = reinterpret_cast<uintptr_t>(callee);
-    code_ptr_location = HInvokeStaticOrDirect::CodePtrLocation::kCallArtMethod;
+    ScopedObjectAccess soa(Thread::Current());
+    if (Runtime::Current()->GetJit()->CanEncodeMethod(
+            callee,
+            codegen->GetGraph()->IsCompilingForSharedJitCode())) {
+      method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kJitDirectAddress;
+      method_load_data = reinterpret_cast<uintptr_t>(callee);
+      code_ptr_location = HInvokeStaticOrDirect::CodePtrLocation::kCallArtMethod;
+    } else {
+      // Do not sharpen.
+      method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kRuntimeCall;
+      code_ptr_location = HInvokeStaticOrDirect::CodePtrLocation::kCallArtMethod;
+    }
   } else if (IsInBootImage(callee)) {
     // Use PC-relative access to the .data.bimg.rel.ro methods array.
     method_load_kind = HInvokeStaticOrDirect::MethodLoadKind::kBootImageRelRo;
@@ -152,19 +163,22 @@ HLoadClass::LoadKind HSharpening::ComputeLoadClassKind(
     HLoadClass::LoadKind desired_load_kind = HLoadClass::LoadKind::kInvalid;
     Runtime* runtime = Runtime::Current();
     const CompilerOptions& compiler_options = codegen->GetCompilerOptions();
-    if (compiler_options.IsBootImage()) {
-      // Compiling boot image. Check if the class is a boot image class.
+    if (compiler_options.IsBootImage() || compiler_options.IsBootImageExtension()) {
+      // Compiling boot image or boot image extension. Check if the class is a boot image class.
       DCHECK(!runtime->UseJitCompilation());
       if (!compiler_options.GetCompilePic()) {
         // Test configuration, do not sharpen.
         desired_load_kind = HLoadClass::LoadKind::kRuntimeCall;
+      } else if (klass != nullptr && runtime->GetHeap()->ObjectIsInBootImageSpace(klass.Get())) {
+        DCHECK(compiler_options.IsBootImageExtension());
+        is_in_boot_image = true;
+        desired_load_kind = HLoadClass::LoadKind::kBootImageRelRo;
       } else if ((klass != nullptr) &&
                  compiler_options.IsImageClass(dex_file.StringByTypeIdx(type_index))) {
         is_in_boot_image = true;
         desired_load_kind = HLoadClass::LoadKind::kBootImageLinkTimePcRelative;
       } else {
         // Not a boot image class.
-        DCHECK(ContainsElement(compiler_options.GetDexFilesForOatFile(), &dex_file));
         desired_load_kind = HLoadClass::LoadKind::kBssEntry;
       }
     } else {
@@ -175,7 +189,16 @@ HLoadClass::LoadKind HSharpening::ComputeLoadClassKind(
         if (is_in_boot_image) {
           desired_load_kind = HLoadClass::LoadKind::kJitBootImageAddress;
         } else if (klass != nullptr) {
-          desired_load_kind = HLoadClass::LoadKind::kJitTableAddress;
+          if (runtime->GetJit()->CanEncodeClass(
+                  klass.Get(),
+                  codegen->GetGraph()->IsCompilingForSharedJitCode())) {
+            desired_load_kind = HLoadClass::LoadKind::kJitTableAddress;
+          } else {
+            // Shared JIT code cannot encode a literal that the GC can move.
+            VLOG(jit) << "Unable to encode in shared region class literal: "
+                      << klass->PrettyClass();
+            desired_load_kind = HLoadClass::LoadKind::kRuntimeCall;
+          }
         } else {
           // Class not loaded yet. This happens when the dex code requesting
           // this `HLoadClass` hasn't been executed in the interpreter.
@@ -299,12 +322,11 @@ void HSharpening::ProcessLoadString(
     ObjPtr<mirror::String> string = nullptr;
 
     const CompilerOptions& compiler_options = codegen->GetCompilerOptions();
-    if (compiler_options.IsBootImage()) {
-      // Compiling boot image. Resolve the string and allocate it if needed, to ensure
-      // the string will be added to the boot image.
+    if (compiler_options.IsBootImage() || compiler_options.IsBootImageExtension()) {
+      // Compiling boot image or boot image extension. Resolve the string and allocate it
+      // if needed, to ensure the string will be added to the boot image.
       DCHECK(!runtime->UseJitCompilation());
       if (compiler_options.GetCompilePic()) {
-        DCHECK(ContainsElement(compiler_options.GetDexFilesForOatFile(), &dex_file));
         if (compiler_options.IsForceDeterminism()) {
           // Strings for methods we're compiling should be pre-resolved but Strings in inlined
           // methods may not be if these inlined methods are not in the boot image profile.
@@ -319,7 +341,12 @@ void HSharpening::ProcessLoadString(
           CHECK(string != nullptr);
         }
         if (string != nullptr) {
-          desired_load_kind = HLoadString::LoadKind::kBootImageLinkTimePcRelative;
+          if (runtime->GetHeap()->ObjectIsInBootImageSpace(string)) {
+            DCHECK(compiler_options.IsBootImageExtension());
+            desired_load_kind = HLoadString::LoadKind::kBootImageRelRo;
+          } else {
+            desired_load_kind = HLoadString::LoadKind::kBootImageLinkTimePcRelative;
+          }
         } else {
           desired_load_kind = HLoadString::LoadKind::kBssEntry;
         }
@@ -331,10 +358,18 @@ void HSharpening::ProcessLoadString(
       DCHECK(!codegen->GetCompilerOptions().GetCompilePic());
       string = class_linker->LookupString(string_index, dex_cache.Get());
       if (string != nullptr) {
-        if (runtime->GetHeap()->ObjectIsInBootImageSpace(string)) {
+        gc::Heap* heap = runtime->GetHeap();
+        if (heap->ObjectIsInBootImageSpace(string)) {
           desired_load_kind = HLoadString::LoadKind::kJitBootImageAddress;
-        } else {
+        } else if (runtime->GetJit()->CanEncodeString(
+                  string,
+                  codegen->GetGraph()->IsCompilingForSharedJitCode())) {
           desired_load_kind = HLoadString::LoadKind::kJitTableAddress;
+        } else {
+          // Shared JIT code cannot encode a literal that the GC can move.
+          VLOG(jit) << "Unable to encode in shared region string literal: "
+                    << string->ToModifiedUtf8();
+          desired_load_kind = HLoadString::LoadKind::kRuntimeCall;
         }
       } else {
         desired_load_kind = HLoadString::LoadKind::kRuntimeCall;