[cfi] Fix weak functions handling.

[eugenis]

Summary

When a function pointer is replaced with a jumptable pointer, special
case is needed to preserve the semantics of extern_weak functions.
Since a jumptable entry can not be extern_weak, we emulate that
behaviour by replacing all references to F (the extern_weak function)
with the following expression: F != nullptr ? JumpTablePtr : nullptr.

Extra special care is needed for global initializers, since most (or
probably all) backends can not lower an initializer that includes
this kind of constant expression. Initializers like that are replaced
with a global constructor (i.e. a runtime initializer).

[eugenis]

lib/Transforms/IPO/LowerTypeTests.cpp

747

I wonder if the priority should be lower. Regular constructors have priority of 65535, we definitely want higher than that, but maybe lower than ASan, etc constructors which use 1 and 2 currently.

[pcc]

lib/Transforms/IPO/LowerTypeTests.cpp

747

I think this constructor is equivalent to relocation application so it should run as close to that as possible, i.e. it should have the highest priority.

Please document the reasons for our choice of priority.

764

What if the use is an element of an aggregate that initializes a GlobalVariable?

[eugenis]

lib/Transforms/IPO/LowerTypeTests.cpp

764

Good point. There are also other expressions that are sometimes (always?) allowed in a global initializers, like (char*)f + 42.

[eugenis]

PTAL

[pcc]

lib/Transforms/IPO/LowerTypeTests.cpp

763

Is this correctly handling the case where an initializer contains multiple references to a weak undef?

[eugenis]

lib/Transforms/IPO/LowerTypeTests.cpp

763

another good point!

[pcc]

LGTM with nits

lib/Transforms/IPO/LowerTypeTests.cpp

756

Nit: could be GV->getValueType().

1063

Nit: this could be an in-class initializer.

[pcc]

lib/Transforms/IPO/LowerTypeTests.cpp

1063

I meant that you could write Function *WeakInitializerFn = nullptr; in the class definition

[eugenis]

rebase

[pcc]

still LGTM

[eugenis]

r286636
thanks for the review!

lib/Transforms/IPO/LowerTypeTests.cpp

class LowerTypeTestsModule {
   uint64_t IndirectIndex = 1;
 
   // Mapping from type identifiers to the call sites that test them.
   DenseMap<Metadata *, std::vector<CallInst *>> TypeTestCallSites;
 
   std::vector<ByteArrayInfo> ByteArrayInfos;
 
   Mangler Mang;
+  Function *WeakInitializerFn;
 
   BitSetInfo
   buildBitSet(Metadata *TypeId,
               const DenseMap<GlobalObject *, uint64_t> &GlobalLayout);
   ByteArrayInfo *createByteArray(BitSetInfo &BSI);
   void allocateByteArrays();
   Value *createBitSetTest(IRBuilder<> &B, BitSetInfo &BSI, ByteArrayInfo *&BAI,
                           Value *BitOffset);
void buildBitSetsFromFunctions(ArrayRef<Metadata *> TypeIds,
                                  ArrayRef<Function *> Functions);
   void buildBitSetsFromFunctionsNative(ArrayRef<Metadata *> TypeIds,
                                     ArrayRef<Function *> Functions);
   void buildBitSetsFromFunctionsWASM(ArrayRef<Metadata *> TypeIds,
                                      ArrayRef<Function *> Functions);
   void buildBitSetsFromDisjointSet(ArrayRef<Metadata *> TypeIds,
                                    ArrayRef<GlobalObject *> Globals);
 
+  void replaceWeakDeclarationWithJumpTablePtr(Function *F, Constant *JT);
+  void moveInitializerToModuleConstructor(GlobalVariable *GV);
 public:
   LowerTypeTestsModule(Module &M);
   bool lower();
 };
 
 struct LowerTypeTests : public ModulePass {
   static char ID;
   LowerTypeTests() : ModulePass(ID) {
if (Arch == Triple::x86 || Arch == Triple::x86_64 || Arch == Triple::arm ||
       Arch == Triple::aarch64)
     buildBitSetsFromFunctionsNative(TypeIds, Functions);
   else if (Arch == Triple::wasm32 || Arch == Triple::wasm64)
     buildBitSetsFromFunctionsWASM(TypeIds, Functions);
   else
     report_fatal_error("Unsupported architecture for jump tables");
 }
 
+void LowerTypeTestsModule::moveInitializerToModuleConstructor(
+    GlobalVariable *GV) {
+  if (WeakInitializerFn == nullptr) {
+    WeakInitializerFn = Function::Create(
+        FunctionType::get(Type::getVoidTy(M.getContext()),
+                          /* IsVarArg */ false),
+        GlobalValue::InternalLinkage, "__cfi_global_var_init", &M);
+    BasicBlock *BB =
+        BasicBlock::Create(M.getContext(), "entry", WeakInitializerFn);
+    ReturnInst::Create(M.getContext(), BB);
+    WeakInitializerFn->setSection(
+        ObjectFormat == Triple::MachO
+            ? "__TEXT,__StaticInit,regular,pure_instructions"
+            : ".text.startup");
+    appendToGlobalCtors(M, WeakInitializerFn, /* Priority */ 0);
+  }
+
+  IRBuilder<> IRB(WeakInitializerFn->getEntryBlock().getTerminator());
+  GV->setConstant(false);
+  IRB.CreateAlignedStore(GV->getInitializer(), GV, GV->getAlignment());
+  GV->setInitializer(
+      Constant::getNullValue(GV->getType()->getPointerElementType()));
+}
+
+// Replace all uses of F with (F ? JT : 0).
+void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
+    Function *F, Constant *JT) {
+  // The target expression can not appear in a constant initializer on most
+  // (all?) targets. Switch to a runtime initializer.
+  SmallVector<GlobalVariable *, 8> GlobalVarUses;
+  for (auto *U : F->users())
+    if (auto *GV = dyn_cast<GlobalVariable>(U))
+      GlobalVarUses.push_back(GV);
+  for (auto GV : GlobalVarUses)
+    moveInitializerToModuleConstructor(GV);
+
+  // Can not RAUW F with an expression that uses F. Replace with a temporary
+  // placeholder first.
+  Function *PlaceholderFn =
+      Function::Create(cast<FunctionType>(F->getValueType()),
+                       GlobalValue::ExternalWeakLinkage, "", &M);
+  F->replaceAllUsesWith(PlaceholderFn);
+
+  Constant *Target = ConstantExpr::getSelect(
+      ConstantExpr::getICmp(CmpInst::ICMP_NE, F,
+                            Constant::getNullValue(F->getType())),
+      JT, Constant::getNullValue(F->getType()));
+  PlaceholderFn->replaceAllUsesWith(Target);
+  PlaceholderFn->eraseFromParent();
+}
+
 /// Given a disjoint set of type identifiers and functions, build a jump table
 /// for the functions, build the bit sets and lower the llvm.type.test calls.
 void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
     ArrayRef<Metadata *> TypeIds, ArrayRef<Function *> Functions) {
   // Unlike the global bitset builder, the function bitset builder cannot
   // re-arrange functions in a particular order and base its calculations on the
   // layout of the functions' entry points, as we have no idea how large a
   // particular function will end up being (the size could even depend on what
if (!Functions[I]->hasName())
       Functions[I]->setName("unnamed");
     if (LinkerSubsectionsViaSymbols || Functions[I]->isDeclarationForLinker()) {
       Constant *CombinedGlobalElemPtr = ConstantExpr::getBitCast(
           ConstantExpr::getGetElementPtr(
               JumpTableType, JumpTable,
               ArrayRef<Constant *>{ConstantInt::get(IntPtrTy, 0),
                                    ConstantInt::get(IntPtrTy, I)}),
           Functions[I]->getType());
-      Functions[I]->replaceAllUsesWith(CombinedGlobalElemPtr);
 
-      if (Functions[I]->isWeakForLinker())
+
+      if (Functions[I]->isWeakForLinker()) {
         AsmOS << ".weak " << Functions[I]->getName() << "\n";
+        replaceWeakDeclarationWithJumpTablePtr(Functions[I],
+                                               CombinedGlobalElemPtr);
+      } else {
+        Functions[I]->replaceAllUsesWith(CombinedGlobalElemPtr);
+      }
     } else {
       assert(Functions[I]->getType()->getAddressSpace() == 0);
 
       createJumpTableAlias(AsmOS, Functions[I], JumpTable, I);
 
       Function *DeclAlias =
           Function::Create(cast<FunctionType>(Functions[I]->getValueType()),
                            GlobalValue::ExternalLinkage, "", &M);

 }
 
 /// Lower all type tests in this module.
 LowerTypeTestsModule::LowerTypeTestsModule(Module &M) : M(M) {
   Triple TargetTriple(M.getTargetTriple());
   LinkerSubsectionsViaSymbols = TargetTriple.isMacOSX();
   Arch = TargetTriple.getArch();
   ObjectFormat = TargetTriple.getObjectFormat();
+  WeakInitializerFn = nullptr;
 }
 
 bool LowerTypeTestsModule::lower() {
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
   if (!TypeTestFunc || TypeTestFunc->use_empty())
     return false;
 

test/Transforms/LowerTypeTests/function-weak.ll

+; RUN: opt -S -lowertypetests -mtriple=i686-unknown-linux-gnu < %s | FileCheck %s
+; RUN: opt -S -lowertypetests -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: opt -S -lowertypetests -mtriple=arm-unknown-linux-gnu < %s | FileCheck %s
+; RUN: opt -S -lowertypetests -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK: module asm ".weak f"
+
+; CHECK: @x = global void ()* null, align 8
+@x = global void ()* @f, align 8
+
+; CHECK: @x2 = global void ()* null, align 8
+@x2 = global void ()* @f, align 8
+
+; CHECK: @x3 = internal global void ()* null, align 8
+@x3 = internal constant void ()* @f, align 8
+
+; CHECK: @.cfi.jumptable = external hidden constant [1 x [{{.*}} x i8]]
+; CHECK:  @llvm.global_ctors = appending global {{.*}}{ i32 0, void ()* @__cfi_global_var_init
+
+; CHECK: declare !type !0 extern_weak void @f()
+declare !type !0 extern_weak void @f()
+
+; CHECK: define zeroext i1 @check_f()
+define zeroext i1 @check_f() {
+entry:
+; CHECK: ret i1 icmp ne (void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()* null)
+  ret i1 icmp ne (void ()* @f, void ()* null)
+}
+
+; CHECK: define void @call_f() {
+define void @call_f() {
+entry:
+; CHECK: call void select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null)()
+  call void @f()
+  ret void
+}
+
+declare i1 @llvm.type.test(i8* %ptr, metadata %bitset) nounwind readnone
+
+define i1 @foo(i8* %p) {
+  %x = call i1 @llvm.type.test(i8* %p, metadata !"typeid1")
+  ret i1 %x
+}
+
+; CHECK: define internal void @__cfi_global_var_init() section ".text.startup" {
+; CHECK-DAG: store void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()** @x3, align 8
+; CHECK-DAG: store void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()** @x2, align 8
+; CHECK-DAG: store void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()** @x, align 8
+; CHECK-DAG: ret void
+; CHECK-DAG: }
+
+
+!0 = !{i32 0, !"typeid1"}