[lld-macho] Force-loading should share code path with regular archive loads

[int3]

Summary

This extends D92539: [lld-macho] Support parsing of bitcode within archives to work even when we are loading archive
members via -force_load. I uncovered this issue while trying to
force-load archives containing bitcode -- we were segfaulting.

In addition to fixing the -force_load case, this diff also addresses
the behavior of -ObjC when LTO bitcode is involved -- we need to
force-load those archive members if they contain ObjC categories.

[int3]

minor update

[int3]

update test

[compnerd]

Did I miss the test that would cover the deserialization of bitcode from a -force_load'ed archive?

[int3]

lld/test/MachO/lto-archive.ll

14–15

Did I miss the test that would cover the deserialization of bitcode from a -force_load'ed archive?

It's here

[smeenai]

LGTM

lld/MachO/CMakeLists.txt

add_lld_library(lldMachO2
   Symbols.cpp
   SyntheticSections.cpp
   Target.cpp
   Writer.cpp
 
   LINK_COMPONENTS
   ${LLVM_TARGETS_TO_BUILD}
   BinaryFormat
+  BitReader
   Core
   DebugInfoDWARF
   LTO
   MC
   ObjCARC
   Object
   Option
   Passes

lld/MachO/Driver.h

 std::string createResponseFile(const llvm::opt::InputArgList &args);
 
 // Check for both libfoo.dylib and libfoo.tbd (in that order).
 llvm::Optional<std::string> resolveDylibPath(llvm::StringRef path);
 
 llvm::Optional<DylibFile *> loadDylib(llvm::MemoryBufferRef mbref,
                                       DylibFile *umbrella = nullptr);
 
+llvm::Optional<InputFile *> loadArchiveMember(MemoryBufferRef, uint32_t modTime,
+                                              StringRef archiveName,
+                                              bool objCOnly);
+
 uint32_t getModTime(llvm::StringRef path);
 
 void printArchiveMemberLoad(StringRef reason, const InputFile *);
 
 } // namespace macho
 } // namespace lld
 
 #endif

lld/MachO/Driver.cpp

std::unique_ptr<object::Archive> file = CHECK(
         object::Archive::create(mbref), path + ": failed to parse archive");
 
     if (!file->isEmpty() && !file->hasSymbolTable())
       error(path + ": archive has no index; run ranlib to add one");
 
     if (config->allLoad || forceLoadArchive) {
       if (Optional<MemoryBufferRef> buffer = readFile(path)) {
         for (const ArchiveMember &member : getArchiveMembers(*buffer)) {
-          inputFiles.insert(make<ObjFile>(member.mbref, member.modTime, path));
+          if (Optional<InputFile *> file = loadArchiveMember(
+                  member.mbref, member.modTime, path, /*objCOnly=*/false)) {
+            inputFiles.insert(*file);
             printArchiveMemberLoad(
                 (forceLoadArchive ? "-force_load" : "-all_load"),
                 inputFiles.back());
           }
         }
+      }
     } else if (config->forceLoadObjC) {
       for (const object::Archive::Symbol &sym : file->symbols())
         if (sym.getName().startswith(objc::klass))
           symtab->addUndefined(sym.getName(), /*file=*/nullptr,
                                /*isWeakRef=*/false);
 
       // TODO: no need to look for ObjC sections for a given archive member if
       // we already found that it contains an ObjC symbol. We should also
       // consider creating a LazyObjFile class in order to avoid double-loading
       // these files here and below (as part of the ArchiveFile).
       if (Optional<MemoryBufferRef> buffer = readFile(path)) {
         for (const ArchiveMember &member : getArchiveMembers(*buffer)) {
-          if (hasObjCSection(member.mbref)) {
-            inputFiles.insert(
-                make<ObjFile>(member.mbref, member.modTime, path));
+          if (Optional<InputFile *> file = loadArchiveMember(
+                  member.mbref, member.modTime, path, /*objCOnly=*/true)) {
+            inputFiles.insert(*file);
             printArchiveMemberLoad("-ObjC", inputFiles.back());
           }
         }
       }
     }
 
     newFile = make<ArchiveFile>(std::move(file));
     break;

lld/MachO/DriverUtils.cpp

 //===- DriverUtils.cpp ----------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "Driver.h"
 #include "Config.h"
+#include "Driver.h"
 #include "InputFiles.h"
+#include "ObjC.h"
 
 #include "lld/Common/Args.h"
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/Memory.h"
 #include "lld/Common/Reproduce.h"
 #include "llvm/ADT/CachedHashString.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/LTO/LTO.h"
 #include "llvm/Option/Arg.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Option/Option.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Path.h"
 #include "llvm/TextAPI/MachO/TextAPIReader.h"
 
 using namespace llvm;
Optional<DylibFile *> macho::loadDylib(MemoryBufferRef mbref,
   } else {
     assert(magic == file_magic::macho_dynamically_linked_shared_lib ||
            magic == file_magic::macho_dynamically_linked_shared_lib_stub);
     file = make<DylibFile>(mbref, umbrella);
   }
   return file;
 }
 
+Optional<InputFile *> macho::loadArchiveMember(MemoryBufferRef mb,
+                                               uint32_t modTime,
+                                               StringRef archiveName,
+                                               bool objCOnly) {
+  switch (identify_magic(mb.getBuffer())) {
+  case file_magic::macho_object:
+    if (!objCOnly || hasObjCSection(mb))
+      return make<ObjFile>(mb, modTime, archiveName);
+    return None;
+  case file_magic::bitcode:
+    if (!objCOnly || check(isBitcodeContainingObjCCategory(mb)))
+      return make<BitcodeFile>(mb);
+    return None;
+  default:
+    error(archiveName + ": archive member " + mb.getBufferIdentifier() +
+          " has unhandled file type");
+    return None;
+  }
+}
+
 uint32_t macho::getModTime(StringRef path) {
   fs::file_status stat;
   if (!fs::status(path, stat))
     if (fs::exists(stat))
       return toTimeT(stat.getLastModificationTime());
 
   warn("failed to get modification time of " + path);
   return 0;
 }
 
 void macho::printArchiveMemberLoad(StringRef reason, const InputFile *f) {
   if (config->printEachFile)
     message(toString(f));
   if (config->printWhyLoad)
     message(reason + " forced load of " + toString(f));
 }

lld/MachO/InputFiles.cpp

uint32_t modTime = toTimeT(
                                      "for the member defining symbol " +
                                      toMachOString(sym)));
 
   // `sym` is owned by a LazySym, which will be replace<>() by make<ObjFile>
   // and become invalid after that call. Copy it to the stack so we can refer
   // to it later.
   const object::Archive::Symbol sym_copy = sym;
 
-  InputFile *file;
-  switch (identify_magic(mb.getBuffer())) {
-  case file_magic::macho_object:
-    file = make<ObjFile>(mb, modTime, getName());
-    break;
-  case file_magic::bitcode:
-    file = make<BitcodeFile>(mb);
-    break;
-  default:
-    StringRef bufname =
-        CHECK(c.getName(), toString(this) + ": could not get buffer name");
-    error(toString(this) + ": archive member " + bufname +
-          " has unhandled file type");
-    return;
-  }
-  inputFiles.insert(file);
-
+  if (Optional<InputFile *> file =
+          loadArchiveMember(mb, modTime, getName(), /*objCOnly=*/false)) {
+    inputFiles.insert(*file);
     // ld64 doesn't demangle sym here even with -demangle. Match that, so
     // intentionally no call to toMachOString() here.
-  printArchiveMemberLoad(sym_copy.getName(), file);
+    printArchiveMemberLoad(sym_copy.getName(), *file);
+  }
 }
 
 static macho::Symbol *createBitcodeSymbol(const lto::InputFile::Symbol &objSym,
                                           BitcodeFile &file) {
   StringRef name = saver.save(objSym.getName());
 
   // TODO: support weak references
   if (objSym.isUndefined())

lld/test/MachO/invalid/bad-archive-member.s

 # REQUIRES: x86
 # RUN: rm -rf %t; split-file %s %t
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/foo.s -o %t/foo.o
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/test.s -o %t/test.o
 # RUN: %lld -dylib -lSystem %t/foo.o -o %t/foo.dylib
 # RUN: llvm-ar rcs %t/foo.a %t/foo.dylib
 # RUN: not %lld %t/test.o %t/foo.a -o /dev/null 2>&1 | FileCheck %s -DFILE=%t/foo.a
+# RUN: not %lld %t/test.o -force_load %t/foo.a -o /dev/null 2>&1 | FileCheck %s -DFILE=%t/foo.a
 # CHECK: error: [[FILE]]: archive member foo.dylib has unhandled file type
 
 #--- foo.s
 .globl _foo
 _foo:
   ret
 
 #--- test.s
 .globl _main
 _main:
   callq _foo
   ret

lld/test/MachO/invalid/bad-archive.s

 # REQUIRES: x86
 # RUN: echo "!<arch>" > %t.a
 # RUN: echo "foo" >> %t.a
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %s -o %t.o
 
 # RUN: not %lld %t.o %t.a -o /dev/null 2>&1 | FileCheck -DFILE=%t.a %s
 # RUN: not %lld %t.o -force_load %t.a -o /dev/null 2>&1 | FileCheck -DFILE=%t.a %s
+# RUN: not %lld %t.o -ObjC %t.a -o /dev/null 2>&1 | FileCheck -DFILE=%t.a %s
 # CHECK: error: [[FILE]]: failed to parse archive: truncated or malformed archive (remaining size of archive too small for next archive member header at offset 8)
 
 .global _main
 _main:
   ret

lld/test/MachO/lto-archive.ll

 ; REQUIRES: x86
 ; RUN: rm -rf %t; split-file %s %t
 ; RUN: llvm-as %t/foo.ll -o %t/foo.o
-; RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/test.s -o %t/test.o
+; RUN: llvm-as %t/objc.ll -o %t/objc.o
 ; RUN: llvm-ar rcs %t/foo.a %t/foo.o
-; RUN: %lld -save-temps -lSystem %t/test.o %t/foo.a -o %t/test
-; RUN: llvm-objdump -d --macho --no-show-raw-insn %t/test | FileCheck %s
+; RUN: llvm-ar rcs %t/objc.a %t/objc.o
 
-; CHECK:      _main:
-; CHECK-NEXT: callq _foo
-; CHECK-NEXT: retq
+; RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/main.s -o %t/main.o
+; RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/references-foo.s -o %t/references-foo.o
+
+; RUN: %lld -lSystem %t/references-foo.o %t/foo.a -o %t/test
+; RUN: llvm-objdump --macho --syms %t/test | FileCheck %s --check-prefix=FOO
+
+; RUN: %lld -lSystem -force_load %t/foo.a %t/main.o -o %t/force-load
+; RUN: llvm-objdump --macho --syms %t/force-load | FileCheck %s --check-prefix=FOO
+
+; RUN: %lld -lSystem %t/foo.a %t/main.o -o %t/no-force-load
+; RUN: llvm-objdump --macho --syms %t/no-force-load | FileCheck %s --check-prefix=NO-FOO
+
+; RUN: %lld -lSystem -ObjC -framework CoreFoundation %t/objc.a %t/main.o -o %t/objc
+; RUN: llvm-objdump --macho --syms %t/objc | FileCheck %s --check-prefix=OBJC
+
+; RUN: %lld -lSystem -framework CoreFoundation %t/objc.a %t/main.o -o %t/no-objc
+; RUN: llvm-objdump --macho --syms %t/no-objc | FileCheck %s --check-prefix=NO-OBJC
+
+; FOO: _foo
+
+; NO-FOO-NOT: _foo
+
+; OBJC-DAG: _OBJC_CLASS_$_Foo
+; OBJC-DAG: _OBJC_METACLASS_$_Foo
+
+; NO-OBJC-NOT: _OBJC_CLASS_$_Foo
+; NO-OBJC-NOT: _OBJC_METACLASS_$_Foo
 
 ;--- foo.ll
 
 target triple = "x86_64-apple-macosx10.15.0"
 target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 
 define void @foo() {
   ret void
 }
 
-;--- test.s
+;--- objc.ll
+
+target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.15.0"
+
+%struct._class_t = type { %struct._class_t*, %struct._class_t*, %struct._class_ro_t* }
+%struct._class_ro_t = type { i8* }
+
+@"OBJC_METACLASS_$_NSObject" = external global %struct._class_t
+@OBJC_CLASS_NAME_ = private unnamed_addr constant [4 x i8] c"Foo\00",
+  section "__TEXT,__objc_classname,cstring_literals"
+@"_OBJC_METACLASS_RO_$_Foo" = internal global %struct._class_ro_t {
+    i8* getelementptr inbounds ([4 x i8], [4 x i8]* @OBJC_CLASS_NAME_, i32 0, i32 0)
+  },
+  section "__DATA, __objc_const"
+@"OBJC_METACLASS_$_Foo" = global %struct._class_t {
+    %struct._class_t* @"OBJC_METACLASS_$_NSObject",
+    %struct._class_t* @"OBJC_METACLASS_$_NSObject",
+    %struct._class_ro_t* @"_OBJC_METACLASS_RO_$_Foo"
+  },
+  section "__DATA, __objc_data"
+@"OBJC_CLASS_$_NSObject" = external global %struct._class_t
+@"_OBJC_CLASS_RO_$_Foo" = internal global %struct._class_ro_t {
+    i8* getelementptr inbounds ([4 x i8], [4 x i8]* @OBJC_CLASS_NAME_, i32 0, i32 0)
+  },
+  section "__DATA, __objc_const"
+@"OBJC_CLASS_$_Foo" = global %struct._class_t {
+    %struct._class_t* @"OBJC_METACLASS_$_Foo",
+    %struct._class_t* @"OBJC_CLASS_$_NSObject",
+    %struct._class_ro_t* @"_OBJC_CLASS_RO_$_Foo"
+  },
+  section "__DATA, __objc_data"
+
+;--- main.s
+
+.globl _main
+_main:
+  ret
+
+;--- references-foo.s
 
 .globl _main
 _main:
   callq _foo
   ret