[lld][MachO] Tail merge strings

lld/MachO/Config.h

@@ -223,6 +223,7 @@ struct Configuration {
   bool warnThinArchiveMissingMembers;
   bool disableVerify;
   bool separateCstringLiteralSections;
+  bool tailMergeStrings;
 
   bool callGraphProfileSort = false;
   llvm::StringRef printSymbolOrder;

lld/MachO/Driver.cpp

@@ -1986,6 +1986,8 @@ bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
   config->separateCstringLiteralSections =
       args.hasFlag(OPT_separate_cstring_literal_sections,
                    OPT_no_separate_cstring_literal_sections, false);
+  config->tailMergeStrings =
+      args.hasFlag(OPT_tail_merge_strings, OPT_no_tail_merge_strings, false);
 
   auto IncompatWithCGSort = [&](StringRef firstArgStr) {
     // Throw an error only if --call-graph-profile-sort is explicitly specified

lld/MachO/Options.td

@@ -1091,6 +1091,10 @@ defm separate_cstring_literal_sections
          "Emit all cstring literals into the __cstring section. As a special "
          "case, the __objc_methname section will still be emitted. (default)">,
       Group<grp_rare>;
+defm tail_merge_strings
+    : BB<"tail-merge-strings", "Enable string tail merging",
+         "Disable string tail merging to improve link-time performance">,
+      Group<grp_rare>;
 
 def grp_deprecated : OptionGroup<"deprecated">, HelpText<"DEPRECATED">;
 

lld/MachO/SyntheticSections.cpp

@@ -1746,6 +1746,8 @@ void CStringSection::finalizeContents() {
 void DeduplicatedCStringSection::finalizeContents() {
   // Find the largest alignment required for each string.
   DenseMap<CachedHashStringRef, Align> strToAlignment;
+  // Used for tail merging only
+  std::vector<CachedHashStringRef> deduplicatedStrs;
   for (const CStringInputSection *isec : inputs) {
     for (const auto &[i, piece] : llvm::enumerate(isec->pieces)) {
       if (!piece.live)
@@ -1754,27 +1756,79 @@ void DeduplicatedCStringSection::finalizeContents() {
       assert(isec->align != 0);
       auto align = getStringPieceAlignment(isec, piece);
       auto [it, wasInserted] = strToAlignment.try_emplace(s, align);
+      if (config->tailMergeStrings && wasInserted)
+        deduplicatedStrs.push_back(s);
       if (!wasInserted && it->second < align)
         it->second = align;
     }
   }
 
+  // Like lexigraphical sort, except we read strings in reverse and take the
+  // longest string first
+  // TODO: We could improve performance by implementing our own sort that avoids
+  // comparing characters we know to be the same. See
+  // StringTableBuilder::multikeySort() for details
+  llvm::sort(deduplicatedStrs, [](const auto &left, const auto &right) {
+    for (const auto &[leftChar, rightChar] :
+         llvm::zip(llvm::reverse(left.val()), llvm::reverse(right.val()))) {
+      if (leftChar == rightChar)
+        continue;
+      return leftChar < rightChar;
+    }
+    return left.size() > right.size();
+  });
+  std::optional<CachedHashStringRef> mergeCandidate;
+  DenseMap<CachedHashStringRef, std::pair<CachedHashStringRef, uint64_t>>
+      tailMergeMap;
+  for (auto &s : deduplicatedStrs) {
+    if (!mergeCandidate || !mergeCandidate->val().ends_with(s.val())) {
+      mergeCandidate = s;
+      continue;
+    }
+    uint64_t tailMergeOffset = mergeCandidate->size() - s.size();
+    // TODO: If the tail offset is incompatible with this string's alignment, we
+    // might be able to find another superstring with a compatible tail offset.
+    // The difficulty is how to do this efficiently
+    const auto &align = strToAlignment.at(s);
+    if (!isAligned(align, tailMergeOffset))
+      continue;
+    auto &mergeCandidateAlign = strToAlignment[*mergeCandidate];
+    if (align > mergeCandidateAlign)
+      mergeCandidateAlign = align;
+    tailMergeMap.try_emplace(s, *mergeCandidate, tailMergeOffset);
+  }
+
   // Sort the strings for performance and compression size win, and then
   // assign an offset for each string and save it to the corresponding
   // StringPieces for easy access.
   for (auto &[isec, i] : priorityBuilder.buildCStringPriorities(inputs)) {
     auto &piece = isec->pieces[i];
     auto s = isec->getCachedHashStringRef(i);
+    // Any string can be tail merged with itself with an offset of zero
+    uint64_t tailMergeOffset = 0;
+    auto mergeIt =
+        config->tailMergeStrings ? tailMergeMap.find(s) : tailMergeMap.end();
+    if (mergeIt != tailMergeMap.end()) {
+      auto &[superString, offset] = mergeIt->second;
+      // s can be tail merged with superString. Do not layout s. Instead layout
+      // superString if we haven't already
+      assert(superString.val().ends_with(s.val()));
+      s = superString;
+      tailMergeOffset = offset;
+    }
     auto [it, wasInserted] = stringOffsetMap.try_emplace(s, /*placeholder*/ 0);
     if (wasInserted) {
       // Avoid computing the offset until we are sure we will need to
       uint64_t offset = alignTo(size, strToAlignment.at(s));
       it->second = offset;
       size = offset + s.size() + 1; // account for null terminator
     }
-    // If the string was already in stringOffsetMap, it is a duplicate and we
-    // only need to assign the offset.
-    piece.outSecOff = it->second;
+    piece.outSecOff = it->second + tailMergeOffset;
+    if (mergeIt != tailMergeMap.end()) {
+      auto &tailMergedString = mergeIt->first;
+      stringOffsetMap[tailMergedString] = piece.outSecOff;
+      assert(isAligned(strToAlignment.at(tailMergedString), piece.outSecOff));
+    }
   }
   for (CStringInputSection *isec : inputs)
     isec->isFinal = true;

lld/docs/ReleaseNotes.rst

@@ -46,6 +46,8 @@ MachO Improvements
 
 * ``--separate-cstring-literal-sections`` emits cstring literal sections into sections defined by their section name.
   (`#158720 <https://github.com/llvm/llvm-project/pull/158720>`_)
+* ``--tail-merge-strings`` enables tail merging of cstring literals.
+  (`#161262 <https://github.com/llvm/llvm-project/pull/161262>`_)
 
 WebAssembly Improvements
 ------------------------

lld/test/MachO/cstring-tailmerge-objc.s

@@ -0,0 +1,144 @@
+; REQUIRES: aarch64
+; RUN: rm -rf %t && split-file %s %t
+
+; Test that ObjC method names are tail merged and
+; ObjCSelRefsHelper::makeSelRef() still works correctly
+
+; RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/a.s -o %t/a.o
+; RUN: %lld -dylib -arch arm64 --tail-merge-strings %t/a.o -o %t/a
+; RUN: llvm-objdump --macho --section="__TEXT,__objc_methname" %t/a | FileCheck %s --implicit-check-not=error
+
+; RUN: %lld -dylib -arch arm64 --no-tail-merge-strings %t/a.o -o %t/nomerge
+; RUN: llvm-objdump --macho --section="__TEXT,__objc_methname" %t/nomerge | FileCheck %s --check-prefixes=CHECK,NOMERGE --implicit-check-not=error
+
+; CHECK: withBar:error:
+; NOMERGE: error:
+
+;--- a.mm
+__attribute__((objc_root_class))
+@interface Foo
+- (void)withBar:(int)bar error:(int)error;
+- (void)error:(int)error;
+@end
+
+@implementation Foo
+- (void)withBar:(int)bar error:(int)error {}
+- (void)error:(int)error {}
+@end
+
+void *_objc_empty_cache;
+void *_objc_empty_vtable;
+;--- gen
+clang -Oz -target arm64-apple-darwin a.mm -S -o -
+;--- a.s
+	.build_version macos, 11, 0
+	.section	__TEXT,__text,regular,pure_instructions
+	.p2align	2                               ; -- Begin function -[Foo withBar:error:]
+"-[Foo withBar:error:]":                ; @"\01-[Foo withBar:error:]"
+	.cfi_startproc
+; %bb.0:
+	ret
+	.cfi_endproc
+                                        ; -- End function
+	.p2align	2                               ; -- Begin function -[Foo error:]
+"-[Foo error:]":                        ; @"\01-[Foo error:]"
+	.cfi_startproc
+; %bb.0:
+	ret
+	.cfi_endproc
+                                        ; -- End function
+	.globl	__objc_empty_vtable             ; @_objc_empty_vtable
+.zerofill __DATA,__common,__objc_empty_vtable,8,3
+	.section	__DATA,__objc_data
+	.globl	_OBJC_CLASS_$_Foo               ; @"OBJC_CLASS_$_Foo"
+	.p2align	3, 0x0
+_OBJC_CLASS_$_Foo:
+	.quad	_OBJC_METACLASS_$_Foo
+	.quad	0
+	.quad	__objc_empty_cache
+	.quad	__objc_empty_vtable
+	.quad	__OBJC_CLASS_RO_$_Foo
+
+	.globl	_OBJC_METACLASS_$_Foo           ; @"OBJC_METACLASS_$_Foo"
+	.p2align	3, 0x0
+_OBJC_METACLASS_$_Foo:
+	.quad	_OBJC_METACLASS_$_Foo
+	.quad	_OBJC_CLASS_$_Foo
+	.quad	__objc_empty_cache
+	.quad	__objc_empty_vtable
+	.quad	__OBJC_METACLASS_RO_$_Foo
+
+	.section	__TEXT,__objc_classname,cstring_literals
+l_OBJC_CLASS_NAME_:                     ; @OBJC_CLASS_NAME_
+	.asciz	"Foo"
+
+	.section	__DATA,__objc_const
+	.p2align	3, 0x0                          ; @"_OBJC_METACLASS_RO_$_Foo"
+__OBJC_METACLASS_RO_$_Foo:
+	.long	3                               ; 0x3
+	.long	40                              ; 0x28
+	.long	40                              ; 0x28
+	.space	4
+	.quad	0
+	.quad	l_OBJC_CLASS_NAME_
+	.quad	0
+	.quad	0
+	.quad	0
+	.quad	0
+	.quad	0
+
+	.section	__TEXT,__objc_methname,cstring_literals
+l_OBJC_METH_VAR_NAME_:                  ; @OBJC_METH_VAR_NAME_
+	.asciz	"withBar:error:"
+
+	.section	__TEXT,__objc_methtype,cstring_literals
+l_OBJC_METH_VAR_TYPE_:                  ; @OBJC_METH_VAR_TYPE_
+	.asciz	"v24@0:8i16i20"
+
+	.section	__TEXT,__objc_methname,cstring_literals
+l_OBJC_METH_VAR_NAME_.1:                ; @OBJC_METH_VAR_NAME_.1
+	.asciz	"error:"
+
+	.section	__TEXT,__objc_methtype,cstring_literals
+l_OBJC_METH_VAR_TYPE_.2:                ; @OBJC_METH_VAR_TYPE_.2
+	.asciz	"v20@0:8i16"
+
+	.section	__DATA,__objc_const
+	.p2align	3, 0x0                          ; @"_OBJC_$_INSTANCE_METHODS_Foo"
+__OBJC_$_INSTANCE_METHODS_Foo:
+	.long	24                              ; 0x18
+	.long	2                               ; 0x2
+	.quad	l_OBJC_METH_VAR_NAME_
+	.quad	l_OBJC_METH_VAR_TYPE_
+	.quad	"-[Foo withBar:error:]"
+	.quad	l_OBJC_METH_VAR_NAME_.1
+	.quad	l_OBJC_METH_VAR_TYPE_.2
+	.quad	"-[Foo error:]"
+
+	.p2align	3, 0x0                          ; @"_OBJC_CLASS_RO_$_Foo"
+__OBJC_CLASS_RO_$_Foo:
+	.long	2                               ; 0x2
+	.long	0                               ; 0x0
+	.long	0                               ; 0x0
+	.space	4
+	.quad	0
+	.quad	l_OBJC_CLASS_NAME_
+	.quad	__OBJC_$_INSTANCE_METHODS_Foo
+	.quad	0
+	.quad	0
+	.quad	0
+	.quad	0
+
+	.globl	__objc_empty_cache              ; @_objc_empty_cache
+.zerofill __DATA,__common,__objc_empty_cache,8,3
+	.section	__DATA,__objc_classlist,regular,no_dead_strip
+	.p2align	3, 0x0                          ; @"OBJC_LABEL_CLASS_$"
+l_OBJC_LABEL_CLASS_$:
+	.quad	_OBJC_CLASS_$_Foo
+
+	.section	__DATA,__objc_imageinfo,regular,no_dead_strip
+L_OBJC_IMAGE_INFO:
+	.long	0
+	.long	64
+
+.subsections_via_symbols

lld/test/MachO/cstring-tailmerge.s

@@ -0,0 +1,85 @@
+; REQUIRES: aarch64
+; RUN: rm -rf %t && split-file %s %t
+
+; RUN: sed "s/<ALIGN>/0/g" %t/align.s.template > %t/align-1.s
+; RUN: sed "s/<ALIGN>/1/g" %t/align.s.template > %t/align-2.s
+; RUN: sed "s/<ALIGN>/2/g" %t/align.s.template > %t/align-4.s
+
+; RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/first.s -o %t/first.o
+; RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/align-1.s -o %t/align-1.o
+; RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/align-2.s -o %t/align-2.o
+; RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/align-4.s -o %t/align-4.o
+
+; RUN: %lld -dylib -arch arm64 --tail-merge-strings %t/first.o %t/align-1.o -o %t/align-1
+; RUN: llvm-objdump --macho --section="__TEXT,__cstring" --syms %t/align-1 | FileCheck %s --check-prefixes=CHECK,ALIGN1
+
+; RUN: %lld -dylib -arch arm64 --tail-merge-strings %t/first.o %t/align-2.o -o %t/align-2
+; RUN: llvm-objdump --macho --section="__TEXT,__cstring" --syms %t/align-2 | FileCheck %s --check-prefixes=CHECK,ALIGN2
+
+; RUN: %lld -dylib -arch arm64 --tail-merge-strings %t/first.o %t/align-4.o -o %t/align-4
+; RUN: llvm-objdump --macho --section="__TEXT,__cstring" --syms %t/align-4 | FileCheck %s --check-prefixes=CHECK,ALIGN4
+
+; CHECK: Contents of (__TEXT,__cstring) section
+; CHECK: [[#%.16x,START:]] get awkward offset{{$}}
+
+; ALIGN1: [[#%.16x,START+19]] myotherlongstr{{$}}
+; ALIGN1: [[#%.16x,START+19+15]] otherstr{{$}}
+
+; ALIGN2: [[#%.16x,START+20]] myotherlongstr{{$}}
+; ALIGN2: [[#%.16x,START+20+16]] longstr{{$}}
+; ALIGN2: [[#%.16x,START+20+16+8]] otherstr{{$}}
+; ALIGN2: [[#%.16x,START+20+16+8+10]] str{{$}}
+
+; ALIGN4: [[#%.16x,START+20]] myotherlongstr{{$}}
+; ALIGN4: [[#%.16x,START+20+16]] otherlongstr{{$}}
+; ALIGN4: [[#%.16x,START+20+16+16]] longstr{{$}}
+; ALIGN4: [[#%.16x,START+20+16+16+8]] otherstr{{$}}
+; ALIGN4: [[#%.16x,START+20+16+16+8+12]] str{{$}}
+
+; CHECK: SYMBOL TABLE:
+
+; ALIGN1: [[#%.16x,START+19]] l     O __TEXT,__cstring _myotherlongstr
+; ALIGN1: [[#%.16x,START+21]] l     O __TEXT,__cstring _otherlongstr
+; ALIGN1: [[#%.16x,START+26]] l     O __TEXT,__cstring _longstr
+; ALIGN1: [[#%.16x,START+34]] l     O __TEXT,__cstring _otherstr
+; ALIGN1: [[#%.16x,START+39]] l     O __TEXT,__cstring _str
+
+; ALIGN2: [[#%.16x,START+20]] l     O __TEXT,__cstring _myotherlongstr
+; ALIGN2: [[#%.16x,START+20+2]] l     O __TEXT,__cstring _otherlongstr
+; ALIGN2: [[#%.16x,START+20+16]] l     O __TEXT,__cstring _longstr
+; ALIGN2: [[#%.16x,START+20+16+8]] l     O __TEXT,__cstring _otherstr
+; ALIGN2: [[#%.16x,START+20+16+8+10]] l     O __TEXT,__cstring _str
+
+; ALIGN4: [[#%.16x,START+20]] l     O __TEXT,__cstring _myotherlongstr
+; ALIGN4: [[#%.16x,START+20+16]] l     O __TEXT,__cstring _otherlongstr
+; ALIGN4: [[#%.16x,START+20+16+16]] l     O __TEXT,__cstring _longstr
+; ALIGN4: [[#%.16x,START+20+16+16+8]] l     O __TEXT,__cstring _otherstr
+; ALIGN4: [[#%.16x,START+20+16+16+8+12]] l     O __TEXT,__cstring _str
+
+;--- first.s
+.cstring
+.p2align 2
+.asciz "get awkward offset"  ; length = 19
+
+;--- align.s.template
+.cstring
+
+.p2align <ALIGN>
+  _myotherlongstr:
+.asciz "myotherlongstr"      ; length = 15
+
+.p2align <ALIGN>
+  _otherlongstr:
+.asciz   "otherlongstr"      ; length = 13, tail offset = 2
+
+.p2align <ALIGN>
+  _longstr:
+.asciz        "longstr"      ; length = 8, tail offset = 7
+
+.p2align <ALIGN>
+  _otherstr:
+.asciz       "otherstr"      ; length = 9
+
+.p2align <ALIGN>
+  _str:
+.asciz            "str"      ; length = 4, tail offset = 5