[lld/mac] Don't assert when ICFing arm64 code

WordLiteralSection dedupes literals by content.
WordLiteralInputSection::getOffset() used to read a literal at the passed-in
offset and look up this value in the deduping map to find the offset of the
deduped value.

But it's possible that (e.g.) a 16-byte literal's value is accessed 4 bytes in.
To get the offset at that address, we have to get the deduped value at offset 0
and then apply the offset 4 to the result.

(See also WordLiteralSection::finalizeContents() which fills in those maps.)

Only a problem on arm64 because in x86_64 the offset is part of the instruction
instead of a separate ARM64_RELOC_ADDEND relocation. (See bug for more details.)

Fixes PR51999.

Differential Revision: https://reviews.llvm.org/D112584

Author: nico

lld/MachO/InputSection.cpp

@@ -227,14 +227,14 @@ WordLiteralInputSection::WordLiteralInputSection(StringRef segname,
 
 uint64_t WordLiteralInputSection::getOffset(uint64_t off) const {
   auto *osec = cast<WordLiteralSection>(parent);
-  const uint8_t *buf = data.data();
+  const uintptr_t buf = reinterpret_cast<uintptr_t>(data.data());
   switch (sectionType(getFlags())) {
   case S_4BYTE_LITERALS:
-    return osec->getLiteral4Offset(buf + off);
+    return osec->getLiteral4Offset(buf + (off & ~3LLU)) | (off & 3);
   case S_8BYTE_LITERALS:
-    return osec->getLiteral8Offset(buf + off);
+    return osec->getLiteral8Offset(buf + (off & ~7LLU)) | (off & 7);
   case S_16BYTE_LITERALS:
-    return osec->getLiteral16Offset(buf + off);
+    return osec->getLiteral16Offset(buf + (off & ~15LLU)) | (off & 15);
   default:
     llvm_unreachable("invalid literal section type");
   }

lld/MachO/SyntheticSections.h

@@ -562,16 +562,16 @@ class WordLiteralSection final : public SyntheticSection {
            !literal8Map.empty();
   }
 
-  uint64_t getLiteral16Offset(const uint8_t *buf) const {
+  uint64_t getLiteral16Offset(uintptr_t buf) const {
     return literal16Map.at(*reinterpret_cast<const UInt128 *>(buf)) * 16;
   }
 
-  uint64_t getLiteral8Offset(const uint8_t *buf) const {
+  uint64_t getLiteral8Offset(uintptr_t buf) const {
     return literal16Map.size() * 16 +
            literal8Map.at(*reinterpret_cast<const uint64_t *>(buf)) * 8;
   }
 
-  uint64_t getLiteral4Offset(const uint8_t *buf) const {
+  uint64_t getLiteral4Offset(uintptr_t buf) const {
     return literal16Map.size() * 16 + literal8Map.size() * 8 +
            literal4Map.at(*reinterpret_cast<const uint32_t *>(buf)) * 4;
   }

lld/test/MachO/icf-arm64.s

@@ -0,0 +1,109 @@
+# REQUIRES: aarch64
+# RUN: rm -rf %t; split-file %s %t
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin19.0.0 %t/main.s -o %t/main.o
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin19.0.0 %t/f2.s -o %t/f2.o
+# RUN: %lld -arch arm64 -lSystem --icf=all -o %t/main %t/main.o %t/f2.o
+# RUN: llvm-objdump -d --syms --print-imm-hex %t/main | FileCheck %s
+
+# CHECK-LABEL: SYMBOL TABLE:
+# CHECK: [[#%x,F1_REF:]]                 g     F __TEXT,__text _f1
+# CHECK: [[#%x,F1_REF:]]                 g     F __TEXT,__text _f2
+
+# CHECK-LABEL: Disassembly of section __TEXT,__text:
+# CHECK:        <_main>:
+# CHECK: bl 0x[[#%x,F1_REF:]]
+# CHECK: bl 0x[[#%x,F1_REF:]]
+
+#--- main.s
+
+.subsections_via_symbols
+
+.literal16
+.p2align 3
+L_align16:
+.quad 0xffffffffffffffff
+.short 0xaaaa
+.short 0xaaaa
+.space 4, 0xaa
+
+.literal8
+.p2align 3
+L_align8:
+.quad 0xeeeeeeeeeeeeeeee
+
+.literal4
+.p2align 2
+L_align4:
+.short 0xbbbb
+.short 0xbbbb
+
+
+.text
+.p2align 2
+
+.globl _main, _f1, _f2
+
+## Test that loading from __literalN sections at non-literal boundaries
+## doesn't confuse ICF. This function should be folded with the identical
+## _f2 in f2 (which uses literals of the same value in a different isec).
+_f1:
+  adrp x9, L_align16@PAGE + 4
+  add x9, x9, L_align16@PAGEOFF + 4
+  ldr x10, [x9]
+
+  adrp x9, L_align8@PAGE + 4
+  add x9, x9, L_align8@PAGEOFF + 4
+  ldr w11, [x9]
+
+  adrp x9, L_align4@PAGE + 2
+  add x9, x9, L_align4@PAGEOFF + 2
+  ldrh w12, [x9]
+
+  ret
+
+_main:
+  bl _f1
+  bl _f2
+
+#--- f2.s
+
+.subsections_via_symbols
+
+.literal16
+.p2align 3
+L_align16:
+.quad 0xffffffffffffffff
+.short 0xaaaa
+.short 0xaaaa
+.space 4, 170
+
+.literal8
+.p2align 3
+L_align8:
+.quad 0xeeeeeeeeeeeeeeee
+
+.literal4
+.p2align 2
+L_align4:
+.short 0xbbbb
+.short 0xbbbb
+
+.text
+.p2align 2
+
+.globl _f2
+_f2:
+  adrp x9, L_align16@PAGE + 4
+  add x9, x9, L_align16@PAGEOFF + 4
+  ldr x10, [x9]
+
+  adrp x9, L_align8@PAGE + 4
+  add x9, x9, L_align8@PAGEOFF + 4
+  ldr w11, [x9]
+
+  adrp x9, L_align4@PAGE + 2
+  add x9, x9, L_align4@PAGEOFF + 2
+  ldrh w12, [x9]
+
+  ret