8317957: [Lilliput/JDK17] Make C2 LoadNKlassCompactHeader more robust

Reviewed-by: shade

src/hotspot/cpu/aarch64/aarch64.ad

@@ -7485,7 +7485,7 @@ instruct loadNKlass(iRegNNoSp dst, memory4 mem)
   ins_pipe(iload_reg_mem);
 %}
 
-instruct loadNKlassLilliput(iRegNNoSp dst, memory4 mem, rFlagsReg cr)
+instruct loadNKlassCompactHeaders(iRegNNoSp dst, memory4 mem, rFlagsReg cr)
 %{
   match(Set dst (LoadNKlass mem));
   effect(KILL cr);
@@ -7494,19 +7494,7 @@ instruct loadNKlassLilliput(iRegNNoSp dst, memory4 mem, rFlagsReg cr)
   ins_cost(4 * INSN_COST);
   format %{ "ldrw  $dst, $mem\t# compressed class ptr" %}
   ins_encode %{
-    assert($mem$$disp == oopDesc::klass_offset_in_bytes(), "expect correct offset");
-    assert($mem$$index$$Register == noreg, "expect no index");
-    Register dst = $dst$$Register;
-    Register obj = $mem$$base$$Register;
-    C2LoadNKlassStub* stub = new (Compile::current()->comp_arena()) C2LoadNKlassStub(dst);
-    Compile::current()->output()->add_stub(stub);
-    __ ldr(dst, Address(obj, oopDesc::mark_offset_in_bytes()));
-    // NOTE: We can't use tbnz here, because the target is sometimes too far away
-    // and cannot be encoded.
-    __ tst(dst, markWord::monitor_value);
-    __ br(Assembler::NE, stub->entry());
-    __ bind(stub->continuation());
-    __ lsr(dst, dst, markWord::klass_shift);
+    __ load_nklass_compact($dst$$Register, $mem$$base$$Register, $mem$$index$$Register, $mem$$scale, $mem$$disp);
   %}
   ins_pipe(pipe_slow);
 %}

src/hotspot/cpu/aarch64/c2_MacroAssembler_aarch64.cpp

@@ -25,7 +25,10 @@
 #include "precompiled.hpp"
 #include "asm/assembler.hpp"
 #include "asm/assembler.inline.hpp"
+#include "opto/c2_CodeStubs.hpp"
 #include "opto/c2_MacroAssembler.hpp"
+#include "opto/compile.hpp"
+#include "opto/output.hpp"
 #include "opto/intrinsicnode.hpp"
 #include "opto/subnode.hpp"
 #include "runtime/stubRoutines.hpp"
@@ -876,3 +879,30 @@ void C2_MacroAssembler::neon_compare(FloatRegister dst, BasicType bt, FloatRegis
     }
   }
 }
+
+void C2_MacroAssembler::load_nklass_compact(Register dst, Register obj, Register index, int scale, int disp) {
+  C2LoadNKlassStub* stub = new (Compile::current()->comp_arena()) C2LoadNKlassStub(dst);
+  Compile::current()->output()->add_stub(stub);
+
+  // Note: Don't clobber obj anywhere in that method!
+
+  // The incoming address is pointing into obj-start + klass_offset_in_bytes. We need to extract
+  // obj-start, so that we can load from the object's mark-word instead. Usually the address
+  // comes as obj-start in obj and klass_offset_in_bytes in disp. However, sometimes C2
+  // emits code that pre-computes obj-start + klass_offset_in_bytes into a register, and
+  // then passes that register as obj and 0 in disp. The following code extracts the base
+  // and offset to load the mark-word.
+  int offset = oopDesc::mark_offset_in_bytes() + disp - oopDesc::klass_offset_in_bytes();
+  if (index == noreg) {
+    ldr(dst, Address(obj, offset));
+  } else {
+    lea(dst, Address(obj, index, Address::lsl(scale)));
+    ldr(dst, Address(dst, offset));
+  }
+  // NOTE: We can't use tbnz here, because the target is sometimes too far away
+  // and cannot be encoded.
+  tst(dst, markWord::monitor_value);
+  br(Assembler::NE, stub->entry());
+  bind(stub->continuation());
+  lsr(dst, dst, markWord::klass_shift);
+}

src/hotspot/cpu/aarch64/c2_MacroAssembler_aarch64.hpp

@@ -53,4 +53,6 @@
   void neon_compare(FloatRegister dst, BasicType bt, FloatRegister src1,
                     FloatRegister src2, int cond, bool isQ);
 
+  void load_nklass_compact(Register dst, Register obj, Register index, int scale, int disp);
+
 #endif // CPU_AARCH64_C2_MACROASSEMBLER_AARCH64_HPP

src/hotspot/cpu/x86/c2_MacroAssembler_x86.cpp

@@ -4005,3 +4005,25 @@ void C2_MacroAssembler::rearrange_bytes(XMMRegister dst, XMMRegister shuffle, XM
   evshufi64x2(xtmp3, src, src, 0xFF, vlen_enc);
   evpshufb(dst, ktmp, xtmp3, shuffle, true, vlen_enc);
 }
+
+#ifdef _LP64
+void C2_MacroAssembler::load_nklass_compact_c2(Register dst, Register obj, Register index, Address::ScaleFactor scale, int disp) {
+  C2LoadNKlassStub* stub = new (Compile::current()->comp_arena()) C2LoadNKlassStub(dst);
+  Compile::current()->output()->add_stub(stub);
+
+  // Note: Don't clobber obj anywhere in that method!
+
+  // The incoming address is pointing into obj-start + klass_offset_in_bytes. We need to extract
+  // obj-start, so that we can load from the object's mark-word instead. Usually the address
+  // comes as obj-start in obj and klass_offset_in_bytes in disp. However, sometimes C2
+  // emits code that pre-computes obj-start + klass_offset_in_bytes into a register, and
+  // then passes that register as obj and 0 in disp. The following code extracts the base
+  // and offset to load the mark-word.
+  int offset = oopDesc::mark_offset_in_bytes() + disp - oopDesc::klass_offset_in_bytes();
+  movq(dst, Address(obj, index, scale, offset));
+  testb(dst, markWord::monitor_value);
+  jcc(Assembler::notZero, stub->entry());
+  bind(stub->continuation());
+  shrq(dst, markWord::klass_shift);
+}
+#endif

src/hotspot/cpu/x86/c2_MacroAssembler_x86.hpp

@@ -278,4 +278,6 @@
   void rearrange_bytes(XMMRegister dst, XMMRegister shuffle, XMMRegister src, XMMRegister xtmp1,
                        XMMRegister xtmp2, XMMRegister xtmp3, Register rtmp, KRegister ktmp, int vlen_enc);
 
+  void load_nklass_compact_c2(Register dst, Register obj, Register index, Address::ScaleFactor scale, int disp);
+
 #endif // CPU_X86_C2_MACROASSEMBLER_X86_HPP

src/hotspot/cpu/x86/x86_64.ad

@@ -5203,25 +5203,17 @@ instruct loadNKlass(rRegN dst, memory mem)
   ins_pipe(ialu_reg_mem); // XXX
 %}
 
-instruct loadNKlassLilliput(rRegN dst, indOffset8 mem, rFlagsReg cr)
+instruct loadNKlassCompactHeaders(rRegN dst, memory mem, rFlagsReg cr)
 %{
   predicate(UseCompactObjectHeaders);
   match(Set dst (LoadNKlass mem));
   effect(KILL cr);
   ins_cost(125); // XXX
   format %{ "movl    $dst, $mem\t# compressed klass ptr" %}
   ins_encode %{
-    assert($mem$$disp == oopDesc::klass_offset_in_bytes(), "expect correct offset 4, but got: %d", $mem$$disp);
-    assert($mem$$index == 4, "expect no index register: %d", $mem$$index);
-    Register dst = $dst$$Register;
-    Register obj = $mem$$base$$Register;
-    C2LoadNKlassStub* stub = new (Compile::current()->comp_arena()) C2LoadNKlassStub(dst);
-    Compile::current()->output()->add_stub(stub);
-    __ movq(dst, Address(obj, oopDesc::mark_offset_in_bytes()));
-    __ testb(dst, markWord::monitor_value);
-    __ jcc(Assembler::notZero, stub->entry());
-    __ bind(stub->continuation());
-    __ shrq(dst, markWord::klass_shift);
+    Register index = $mem$$index != 4 ? $mem$$index$$Register : noreg;
+    Address::ScaleFactor sf = (index != noreg) ? static_cast<Address::ScaleFactor>($mem$$scale) : Address::no_scale;
+    __ load_nklass_compact_c2($dst$$Register, $mem$$base$$Register, index, sf, $mem$$disp);
   %}
   ins_pipe(pipe_slow); // XXX
 %}