8364042: UnsafeMemoryAccess will not work with AOT cached code stubs

Reviewed-by: asmehra, kvn

src/hotspot/cpu/aarch64/stubGenerator_aarch64.cpp

@@ -2732,8 +2732,11 @@ class StubGenerator: public StubCodeGenerator {
     address entry_jlong_arraycopy;
     address entry_checkcast_arraycopy;
 
-    address ucm_common_error_exit       =  generate_unsafecopy_common_error_exit();
+    // generate the common exit first so later stubs can rely on it if
-    UnsafeMemoryAccess::set_common_exit_stub_pc(ucm_common_error_exit);
+    // they want an UnsafeMemoryAccess exit non-local to the stub
+    StubRoutines::_unsafecopy_common_exit = generate_unsafecopy_common_error_exit();
+    // register the stub as the default exit with class UnsafeMemoryAccess
+    UnsafeMemoryAccess::set_common_exit_stub_pc(StubRoutines::_unsafecopy_common_exit);
 
     generate_copy_longs(StubId::stubgen_copy_byte_f_id, IN_HEAP | IS_ARRAY, copy_f, r0, r1, r15);
     generate_copy_longs(StubId::stubgen_copy_byte_b_id, IN_HEAP | IS_ARRAY, copy_b, r0, r1, r15);

src/hotspot/cpu/arm/stubGenerator_arm.cpp

@@ -3001,12 +3001,15 @@ class StubGenerator: public StubCodeGenerator {
 
   void generate_arraycopy_stubs() {
 
+    // generate the common exit first so later stubs can rely on it if
+    // they want an UnsafeMemoryAccess exit non-local to the stub
+    StubRoutines::_unsafecopy_common_exit = generate_unsafecopy_common_error_exit();
+    // register the stub as the default exit with class UnsafeMemoryAccess
+    UnsafeMemoryAccess::set_common_exit_stub_pc(StubRoutines::_unsafecopy_common_exit);
+
     // Note:  the disjoint stubs must be generated first, some of
     //        the conjoint stubs use them.
 
-    address ucm_common_error_exit       =  generate_unsafecopy_common_error_exit();
-    UnsafeMemoryAccess::set_common_exit_stub_pc(ucm_common_error_exit);
-
     // these need always status in case they are called from generic_arraycopy
     StubRoutines::_jbyte_disjoint_arraycopy  = generate_primitive_copy(StubId::stubgen_jbyte_disjoint_arraycopy_id);
     StubRoutines::_jshort_disjoint_arraycopy = generate_primitive_copy(StubId::stubgen_jshort_disjoint_arraycopy_id);

src/hotspot/cpu/ppc/stubGenerator_ppc.cpp

@@ -3271,12 +3271,15 @@ class StubGenerator: public StubCodeGenerator {
   }
 
   void generate_arraycopy_stubs() {
+    // generate the common exit first so later stubs can rely on it if
+    // they want an UnsafeMemoryAccess exit non-local to the stub
+    StubRoutines::_unsafecopy_common_exit = generate_unsafecopy_common_error_exit();
+    // register the stub as the default exit with class UnsafeMemoryAccess
+    UnsafeMemoryAccess::set_common_exit_stub_pc(StubRoutines::_unsafecopy_common_exit);
+
     // Note: the disjoint stubs must be generated first, some of
     // the conjoint stubs use them.
 
-    address ucm_common_error_exit       =  generate_unsafecopy_common_error_exit();
-    UnsafeMemoryAccess::set_common_exit_stub_pc(ucm_common_error_exit);
-
     // non-aligned disjoint versions
     StubRoutines::_jbyte_disjoint_arraycopy       = generate_disjoint_byte_copy(StubId::stubgen_jbyte_disjoint_arraycopy_id);
     StubRoutines::_jshort_disjoint_arraycopy      = generate_disjoint_short_copy(StubId::stubgen_jshort_disjoint_arraycopy_id);

src/hotspot/cpu/riscv/stubGenerator_riscv.cpp

@@ -2301,12 +2301,15 @@ class StubGenerator: public StubCodeGenerator {
     address entry_jlong_arraycopy     = nullptr;
     address entry_checkcast_arraycopy = nullptr;
 
+    // generate the common exit first so later stubs can rely on it if
+    // they want an UnsafeMemoryAccess exit non-local to the stub
+    StubRoutines::_unsafecopy_common_exit = generate_unsafecopy_common_error_exit();
+    // register the stub as the default exit with class UnsafeMemoryAccess
+    UnsafeMemoryAccess::set_common_exit_stub_pc(StubRoutines::_unsafecopy_common_exit);
+
     generate_copy_longs(StubId::stubgen_copy_byte_f_id, copy_f, c_rarg0, c_rarg1, t1);
     generate_copy_longs(StubId::stubgen_copy_byte_b_id, copy_b, c_rarg0, c_rarg1, t1);
 
-    address ucm_common_error_exit     = generate_unsafecopy_common_error_exit();
-    UnsafeMemoryAccess::set_common_exit_stub_pc(ucm_common_error_exit);
-
     StubRoutines::riscv::_zero_blocks = generate_zero_blocks();
 
     //*** jbyte

src/hotspot/cpu/s390/stubGenerator_s390.cpp

@@ -1576,12 +1576,14 @@ class StubGenerator: public StubCodeGenerator {
 
   void generate_arraycopy_stubs() {
 
+    // they want an UnsafeMemoryAccess exit non-local to the stub
+    StubRoutines::_unsafecopy_common_exit = generate_unsafecopy_common_error_exit();
+    // register the stub as the default exit with class UnsafeMemoryAccess
+    UnsafeMemoryAccess::set_common_exit_stub_pc(StubRoutines::_unsafecopy_common_exit);
+
     // Note: the disjoint stubs must be generated first, some of
     // the conjoint stubs use them.
 
-    address ucm_common_error_exit       =  generate_unsafecopy_common_error_exit();
-    UnsafeMemoryAccess::set_common_exit_stub_pc(ucm_common_error_exit);
-
     StubRoutines::_jbyte_disjoint_arraycopy      = generate_disjoint_nonoop_copy (StubId::stubgen_jbyte_disjoint_arraycopy_id);
     StubRoutines::_jshort_disjoint_arraycopy     = generate_disjoint_nonoop_copy(StubId::stubgen_jshort_disjoint_arraycopy_id);
     StubRoutines::_jint_disjoint_arraycopy       = generate_disjoint_nonoop_copy  (StubId::stubgen_jint_disjoint_arraycopy_id);

src/hotspot/share/runtime/stubDeclarations.hpp

@@ -883,6 +883,9 @@
   do_stub(final, verify_oop)                                            \
   do_entry(final, verify_oop, verify_oop_subroutine_entry,              \
            verify_oop_subroutine_entry)                                 \
+  do_stub(final, unsafecopy_common)                                     \
+  do_entry(final, unsafecopy_common, unsafecopy_common_exit,            \
+           unsafecopy_common_exit)                                      \
   do_stub(final, jbyte_arraycopy)                                       \
   do_entry_init(final, jbyte_arraycopy, jbyte_arraycopy,                \
                 jbyte_arraycopy, StubRoutines::jbyte_copy)              \

src/hotspot/share/runtime/stubRoutines.cpp

@@ -131,6 +131,35 @@ address UnsafeMemoryAccess::page_error_continue_pc(address pc) {
   return nullptr;
 }
 
+// Used to retrieve mark regions that lie within a generated stub so
+// they can be saved along with the stub and used to reinit the table
+// when the stub is reloaded.
+
+void UnsafeMemoryAccess::collect_entries(address range_start, address range_end, GrowableArray<address>& entries)
+{
+  for (int i = 0; i < _table_length; i++) {
+    UnsafeMemoryAccess& e = _table[i];
+    assert((e._start_pc != nullptr &&
+            e._end_pc != nullptr &&
+            e._error_exit_pc != nullptr),
+           "search for entries found incomplete table entry");
+    if (e._start_pc >= range_start && e._end_pc <= range_end) {
+      assert(((e._error_exit_pc >= range_start &&
+               e._error_exit_pc <= range_end) ||
+              e._error_exit_pc == _common_exit_stub_pc),
+             "unexpected error exit pc");
+      entries.append(e._start_pc);
+      entries.append(e._end_pc);
+      // only return an exit pc when it is within the range of the stub
+      if (e._error_exit_pc != _common_exit_stub_pc) {
+        entries.append(e._error_exit_pc);
+      } else {
+        // an address outside the stub must be the common exit stub address
+        entries.append(nullptr);
+      }
+    }
+  }
+}
 
 static BufferBlob* initialize_stubs(BlobId blob_id,
                                     int code_size, int max_aligned_stubs,

src/hotspot/share/runtime/stubRoutines.hpp

@@ -33,6 +33,7 @@
 #include "runtime/stubCodeGenerator.hpp"
 #include "runtime/stubInfo.hpp"
 #include "runtime/threadWXSetters.inline.hpp"
+#include "utilities/growableArray.hpp"
 #include "utilities/macros.hpp"
 
 // StubRoutines provides entry points to assembly routines used by
@@ -140,6 +141,10 @@ class UnsafeMemoryAccess : public CHeapObj<mtCode> {
   static bool    contains_pc(address pc);
   static address page_error_continue_pc(address pc);
   static void    create_table(int max_size);
+  // Append to entries arrray start, end and exit pcs of all table
+  // entries that identify a sub-interval of range (range_start,
+  // range_end). Append nullptr if the exit pc is not in the range.
+  static void collect_entries(address range_start, address range_end, GrowableArray<address>& entries);
 };
 
 class UnsafeMemoryAccessMark : public StackObj {