save and restore ZGC patch addresses for x86 arraycopy stubs in AOT cache

src/hotspot/cpu/x86/gc/z/zBarrierSetAssembler_x86.cpp

@@ -1426,6 +1426,71 @@ void ZBarrierSetAssembler::patch_barriers() {
 #undef __
 #define __ masm->
 
+void ZBarrierSetAssembler::register_reloc_addresses(GrowableArray<address> &entries, int begin, int count) {
+  int formats[] = {
+    ZBarrierRelocationFormatLoadBadAfterTest,
+    ZBarrierRelocationFormatStoreBadAfterTest,
+    ZBarrierRelocationFormatStoreGoodAfterOr
+  };
+  int format_idx = 0;
+  int format = formats[format_idx];
+  for (int i = begin; i < begin + count; i++) {
+    address addr = entries.at(i);
+    // reloc addresses occur in 3 groups terminated with a nullptr
+    if (addr == nullptr) {
+      assert(((uint)format_idx) < (int)(sizeof(formats) / sizeof(formats[0])),
+             "unexpected reloc group count");
+      format = formats[format_idx++];
+    } else {
+      switch(format) {
+      case ZBarrierRelocationFormatLoadBadAfterTest:
+        _load_bad_relocations.append(addr);
+        break;
+      case ZBarrierRelocationFormatStoreBadAfterTest:
+        _store_bad_relocations.append(addr);
+        break;
+      case ZBarrierRelocationFormatStoreGoodAfterOr:
+        _store_good_relocations.append(addr);
+        break;
+      }
+      patch_barrier_relocation(addr, format);
+    }
+  }
+  assert(format_idx == (int)(sizeof(formats) / sizeof(formats[0])),
+         "not enough reloc groups %d - expecting %d", format_idx,  (int)(sizeof(formats) / sizeof(formats[0])));
+}
+
+void ZBarrierSetAssembler::retrieve_reloc_addresses(address start, address end, GrowableArray<address> &entries) {
+  assert(start != nullptr, "start address must not be null");
+  assert(end != nullptr, "start address must not be null");
+  assert(start < end, "stub range must not be empty");
+  for (int i = 0; i < _load_bad_relocations.length(); i++) {
+    address addr = _load_bad_relocations.at(i);
+    assert(addr != nullptr, "load bad reloc address shoudl not be null!");
+    if (start <= addr && addr < end) {
+      entries.append(addr);
+    }
+  }
+  entries.append(nullptr);
+  for (int i = 0; i < _store_bad_relocations.length(); i++) {
+    address addr = _store_bad_relocations.at(i);
+    assert(addr != nullptr, "store bad reloc address shoudl not be null!");
+    if (start <= addr && addr < end) {
+      entries.append(addr);
+    }
+  }
+  entries.append(nullptr);
+  for (int i = 0; i < _store_good_relocations.length(); i++) {
+    address addr = _store_good_relocations.at(i);
+    assert(addr != nullptr, "store good reloc address shoudl not be null!");
+    if (start <= addr && addr < end) {
+      entries.append(addr);
+    }
+  }
+  entries.append(nullptr);
+}
+
+
 void ZBarrierSetAssembler::check_oop(MacroAssembler* masm, Register obj, Register tmp1, Register tmp2, Label& error) {
   // C1 calls verfy_oop in the middle of barriers, before they have been uncolored
   // and after being colored. Therefore, we must deal with colored oops as well.

src/hotspot/cpu/x86/gc/z/zBarrierSetAssembler_x86.hpp

@@ -193,6 +193,10 @@ public:
 
   void patch_barriers();
 
+  void register_reloc_addresses(GrowableArray<address> &entries, int begin, int count);
+
+  void retrieve_reloc_addresses(address start, address end, GrowableArray<address> &entries);
+
   void check_oop(MacroAssembler* masm, Register obj, Register tmp1, Register tmp2, Label& error);
 };
 

src/hotspot/cpu/x86/stubGenerator_x86_64_arraycopy.cpp

@@ -586,7 +586,7 @@ address StubGenerator::generate_disjoint_copy_avx3_masked(StubId stub_id, addres
   if (start != nullptr) {
     assert(entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
-    assert(extras.length() == expected_extra_count,
+    assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
            "unexpected extra addresses count %d", extras.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
@@ -594,6 +594,12 @@ address StubGenerator::generate_disjoint_copy_avx3_masked(StubId stub_id, addres
     if (add_extras) {
       // restore 1 x UMAM {start,end,handler} addresses from extras
       register_unsafe_access_handlers(extras, 0, 1);
+#if INCLUDE_ZGC
+      // register addresses at which ZGC does colour patching
+      if ((UseZGC && is_oop))  {
+        register_reloc_addresses(extras, 3, extras.length());
+      }
+#endif // INCLUDE_ZGC
     }
     return start;
   }
@@ -826,8 +832,14 @@ address StubGenerator::generate_disjoint_copy_avx3_masked(StubId stub_id, addres
   address end = __ pc();
   if (add_extras) {
     retrieve_unsafe_access_handlers(start, end, extras);
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if ((UseZGC && is_oop)) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   }
-  assert(extras.length() == expected_extra_count,
+  assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
          "unexpected extra addresses count %d", extras.length());
 
   // record the stub entry and end plus the no_push entry and any
@@ -961,7 +973,7 @@ address StubGenerator::generate_conjoint_copy_avx3_masked(StubId stub_id, addres
   if (start != nullptr) {
     assert(entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
-    assert(extras.length() == expected_extra_count,
+    assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
            "unexpected extra addresses count %d", extras.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
@@ -969,6 +981,12 @@ address StubGenerator::generate_conjoint_copy_avx3_masked(StubId stub_id, addres
     if (add_extras) {
       // restore 1 x UMAM {start,end,handler} addresses from extras
       register_unsafe_access_handlers(extras, 0, 1);
+#if INCLUDE_ZGC
+      if ((UseZGC && is_oop))  {
+        // register addresses at which ZGC does colour patching
+        register_reloc_addresses(extras, 3, extras.length());
+      }
+#endif // INCLUDE_ZGC
     }
     return start;
   }
@@ -1141,8 +1159,14 @@ address StubGenerator::generate_conjoint_copy_avx3_masked(StubId stub_id, addres
   address end = __ pc();
   if (add_extras) {
     retrieve_unsafe_access_handlers(start, end, extras);
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if ((UseZGC && is_oop)) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   }
-  assert(extras.length() == expected_extra_count,
+  assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
          "unexpected extra addresses count %d", extras.length());
 
   // record the stub entry and end plus the no_push entry and any
@@ -2138,7 +2162,7 @@ address StubGenerator::generate_disjoint_int_oop_copy(StubId stub_id, address* e
   if (start != nullptr) {
     assert(entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
-    assert(extras.length() == expected_extra_count,
+    assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
            "unexpected extra addresses count %d", extras.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
@@ -2146,6 +2170,12 @@ address StubGenerator::generate_disjoint_int_oop_copy(StubId stub_id, address* e
     if (add_extras) {
       // restore 2 UMAM {start,end,handler} addresses from extras
       register_unsafe_access_handlers(extras, 0, 2);
+#if INCLUDE_ZGC
+      // register addresses at which ZGC does colour patching
+      if ((UseZGC && is_oop))  {
+        register_reloc_addresses(extras, 6, extras.length());
+      }
+#endif // INCLUDE_ZGC
     }
     return start;
   }
@@ -2237,8 +2267,14 @@ address StubGenerator::generate_disjoint_int_oop_copy(StubId stub_id, address* e
   address end = __ pc();
   if (add_extras) {
     retrieve_unsafe_access_handlers(start, end, extras);
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if ((UseZGC && is_oop)) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   }
-  assert(extras.length() == expected_extra_count,
+  assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
          "unexpected extra addresses count %d", extras.length());
 
   // record the stub entry and end plus the no_push entry and any
@@ -2306,7 +2342,7 @@ address StubGenerator::generate_conjoint_int_oop_copy(StubId stub_id, address no
   if (start != nullptr) {
     assert(entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
-    assert(extras.length() == expected_extra_count,
+    assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
            "unexpected extra addresses count %d", extras.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
@@ -2314,6 +2350,12 @@ address StubGenerator::generate_conjoint_int_oop_copy(StubId stub_id, address no
     if (add_extras) {
       // restore 2 UMAM {start,end,handler} addresses from extras
       register_unsafe_access_handlers(extras, 0, 2);
+#if INCLUDE_ZGC
+      // register addresses at which ZGC does colour patching
+      if ((UseZGC && is_oop))  {
+        register_reloc_addresses(extras, 6, extras.length());
+      }
+#endif // INCLUDE_ZGC
     }
     return start;
   }
@@ -2409,8 +2451,14 @@ address StubGenerator::generate_conjoint_int_oop_copy(StubId stub_id, address no
   address end = __ pc();
   if (add_extras) {
     retrieve_unsafe_access_handlers(start, end, extras);
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if ((UseZGC && is_oop)) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   }
-  assert(extras.length() == expected_extra_count,
+  assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
          "unexpected extra addresses count %d", extras.length());
 
   // record the stub entry and end plus the no_push entry and any
@@ -2476,7 +2524,7 @@ address StubGenerator::generate_disjoint_long_oop_copy(StubId stub_id, address *
   if (start != nullptr) {
     assert(entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
-    assert(extras.length() == expected_extra_count,
+    assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
            "unexpected extra addresses count %d", extras.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
@@ -2484,6 +2532,12 @@ address StubGenerator::generate_disjoint_long_oop_copy(StubId stub_id, address *
     if (add_extras) {
       // restore 2 UMAM {start,end,handler} addresses from extras
       register_unsafe_access_handlers(extras, 0, 2);
+#if INCLUDE_ZGC
+      // register addresses at which ZGC does colour patching
+      if ((UseZGC && is_oop))  {
+        register_reloc_addresses(extras, 6, extras.length());
+      }
+#endif // INCLUDE_ZGC
     }
     return start;
   }
@@ -2581,8 +2635,14 @@ address StubGenerator::generate_disjoint_long_oop_copy(StubId stub_id, address *
   address end = __ pc();
   if (add_extras) {
     retrieve_unsafe_access_handlers(start, end, extras);
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if ((UseZGC && is_oop)) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   }
-  assert(extras.length() == expected_extra_count,
+  assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
          "unexpected extra addresses count %d", extras.length());
 
   // record the stub entry and end plus the no_push entry and any
@@ -2646,7 +2706,7 @@ address StubGenerator::generate_conjoint_long_oop_copy(StubId stub_id, address n
   if (start != nullptr) {
     assert(entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
-    assert(extras.length() == expected_extra_count,
+    assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
            "unexpected extra addresses count %d", extras.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
@@ -2654,6 +2714,12 @@ address StubGenerator::generate_conjoint_long_oop_copy(StubId stub_id, address n
     if (add_extras) {
       // restore 2 UMAM {start,end,handler} addresses from extras
       register_unsafe_access_handlers(extras, 0, 2);
+#if INCLUDE_ZGC
+      // register addresses at which ZGC does colour patching
+      if ((UseZGC && is_oop))  {
+        register_reloc_addresses(extras, 6, extras.length());
+      }
+#endif // INCLUDE_ZGC
     }
     return start;
   }
@@ -2743,8 +2809,14 @@ address StubGenerator::generate_conjoint_long_oop_copy(StubId stub_id, address n
   address end = __ pc();
   if (add_extras) {
     retrieve_unsafe_access_handlers(start, end, extras);
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if ((UseZGC && is_oop)) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   }
-  assert(extras.length() == expected_extra_count,
+  assert((UseZGC && is_oop) || extras.length() == expected_extra_count,
          "unexpected extra addresses count %d", extras.length());
 
   // record the stub entry and end plus the no_push entry and any
@@ -2807,17 +2879,22 @@ address StubGenerator::generate_checkcast_copy(StubId stub_id, address *entry) {
   }
 
   GrowableArray<address> entries;
+  GrowableArray<address> extras;
   int expected_entry_count = (entry != nullptr ? 2 : 1);
   int entry_count = StubInfo::entry_count(stub_id);
   assert(entry_count == expected_entry_count, "sanity check");
   GrowableArray<address>* entries_ptr = (entry_count == 1 ? nullptr : &entries);
-  address start = load_archive_data(stub_id, entries_ptr);
+  GrowableArray<address>* extras_ptr = (UseZGC ? &extras : nullptr);
+  address start = load_archive_data(stub_id, entries_ptr, extras_ptr);
   if (start != nullptr) {
-    assert(entries.length() == expected_entry_count - 1,
+    assert(UseZGC || entries.length() == expected_entry_count - 1,
            "unexpected extra entry count %d", entries.length());
     if (entry != nullptr) {
       *entry = entries.at(0);
     }
+    if (UseZGC)  {
+      register_reloc_addresses(extras, 0, extras.length());
+    }
     return start;
   }
 
@@ -3010,8 +3087,15 @@ address StubGenerator::generate_checkcast_copy(StubId stub_id, address *entry) {
   __ leave(); // required for proper stackwalking of RuntimeStub frame
   __ ret(0);
 
+  address end = __ pc();
+#if INCLUDE_ZGC
+    // retrieve addresses at which ZGC does colour patching
+    if (UseZGC) {
+      retrieve_reloc_addresses(start, end, extras);
+    }
+#endif // INCLUDE_ZGC
   // record the stub entry and end plus the no_push entry
-  store_archive_data(stub_id, start, __ pc(), entries_ptr);
+    store_archive_data(stub_id, start, end, entries_ptr, extras_ptr);
 
   return start;
 }

src/hotspot/share/gc/z/zBarrierSetAssembler.hpp

@@ -34,6 +34,9 @@ public:
 
   static Address load_bad_mask_from_jni_env(Register env);
   static Address mark_bad_mask_from_jni_env(Register env);
+
+  virtual void register_reloc_addresses(GrowableArray<address> &entries, int begin, int count) { }
+  virtual void retrieve_reloc_addresses(address start, address end, GrowableArray<address> &entries) { }
 };
 
 // Needs to be included after definition of ZBarrierSetAssemblerBase

src/hotspot/share/runtime/stubCodeGenerator.cpp

@@ -31,7 +31,9 @@
 #include "prims/jvmtiExport.hpp"
 #include "runtime/stubCodeGenerator.hpp"
 #include "runtime/stubRoutines.hpp"
-
+#if INCLUDE_ZGC
+#include "gc/z/zBarrierSetAssembler.hpp"
+#endif // INCLUDE_ZGC
 
 // Implementation of StubCodeDesc
 
@@ -139,6 +141,23 @@ void StubCodeGenerator::retrieve_unsafe_access_handlers(address start, address e
   UnsafeMemoryAccess::collect_entries(start, end, entries);
 }
 
+#if INCLUDE_ZGC
+// Helper used to restore ZGC pointer colouring relocation addresses
+// retrieved from the AOT cache.
+void StubCodeGenerator::register_reloc_addresses(GrowableArray<address> &entries, int begin, int count) {
+  ZBarrierSetAssembler *zbs = (ZBarrierSetAssembler*)BarrierSet::barrier_set()->barrier_set_assembler();
+  zbs->register_reloc_addresses(entries, begin, count);
+}
+
+// Helper used to retrieve ranges and handler addresses registered
+// during generation of the stub which spans [start, end) in order to
+// allow them to be saved to an AOT cache.
+void StubCodeGenerator::retrieve_reloc_addresses(address start, address end, GrowableArray<address> &entries) {
+  ZBarrierSetAssembler *zbs = (ZBarrierSetAssembler*)BarrierSet::barrier_set()->barrier_set_assembler();
+  zbs->retrieve_reloc_addresses(start, end, entries);
+}
+#endif // INCLUDE_ZGC
+
 void StubCodeGenerator::stub_prolog(StubCodeDesc* cdesc) {
   // default implementation - do nothing
 }

src/hotspot/share/runtime/stubCodeGenerator.hpp

@@ -118,8 +118,12 @@ class StubCodeGenerator: public StackObj {
   // unsafe handler management
   void register_unsafe_access_handlers(GrowableArray<address> &entries, int begin, int count);
   void retrieve_unsafe_access_handlers(address start, address end, GrowableArray<address> &entries);
+#if INCLUDE_ZGC
+  void register_reloc_addresses(GrowableArray<address> &entries, int begin, int count);
+  void retrieve_reloc_addresses(address start, address end, GrowableArray<address> &entries);
+#endif // INCLUDE_ZGC
 
- public:
+public:
   StubCodeGenerator(CodeBuffer* code, bool print_code = false);
   StubCodeGenerator(CodeBuffer* code, BlobId blob_id, AOTStubData* stub_data = nullptr, bool print_code = false);
   ~StubCodeGenerator();