[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

757

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

757

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.

774

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

[eugenis]

lib/Transforms/IPO/LowerTypeTests.cpp

774

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

773

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

[eugenis]

lib/Transforms/IPO/LowerTypeTests.cpp

773

another good point!

[pcc]

LGTM with nits

lib/Transforms/IPO/LowerTypeTests.cpp

766

Nit: could be GV->getValueType().

1085

Nit: this could be an in-class initializer.

[pcc]

lib/Transforms/IPO/LowerTypeTests.cpp

1085

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

 //===-- LowerTypeTests.cpp - type metadata lowering pass ------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass lowers type metadata and calls to the llvm.type.test intrinsic.
 // See http://llvm.org/docs/TypeMetadata.html for more information.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
 #include "llvm/ADT/EquivalenceClasses.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalObject.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
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 = nullptr;
 
   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);
+  void findGlobalVariableUsersOf(Constant *C,
+                                 SmallSetVector<GlobalVariable *, 8> &Out);
+
 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");
+    // This code is equivalent to relocation application, and should run at the
+    // earliest possible time (i.e. with the highest priority).
+    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->getValueType()));
+}
+
+void LowerTypeTestsModule::findGlobalVariableUsersOf(
+    Constant *C, SmallSetVector<GlobalVariable *, 8> &Out) {
+  for (auto *U : C->users()){
+    if (auto *GV = dyn_cast<GlobalVariable>(U))
+      Out.insert(GV);
+    else if (auto *C2 = dyn_cast<Constant>(U))
+      findGlobalVariableUsersOf(C2, Out);
+  }
+}
+
+// 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.
+  SmallSetVector<GlobalVariable *, 8> GlobalVarUsers;
+  findGlobalVariableUsersOf(F, GlobalVarUsers);
+  for (auto GV : GlobalVarUsers)
+    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();
 }
 
 bool LowerTypeTestsModule::lower() {
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
   if (!TypeTestFunc || TypeTestFunc->use_empty())
     return false;
 
   // Equivalence class set containing type identifiers and the globals that

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
+
+; f + addend
+; CHECK: @x4 = global void ()* null, align 8
+@x4 = global void ()* bitcast (i8* getelementptr (i8, i8* bitcast (void ()* @f to i8*), i64 42) to void ()*), align 8
+
+; aggregate initializer
+; CHECK: @s = global { void ()*, void ()*, i32 } zeroinitializer, align 8
+@s = global { void ()*, void ()*, i32 } { void ()* @f, void ()* @f, i32 42 }, 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-NEXT: entry:
+; CHECK-NEXT: store { void ()*, void ()*, i32 } { void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), i32 42 }, { void ()*, void ()*, i32 }* @s, align 8
+; CHECK-NEXT: store void ()* bitcast (i8* getelementptr (i8, i8* bitcast (void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null) to i8*), i64 42) to void ()*), void ()** @x4, align 8
+; CHECK-NEXT: store void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()** @x3, align 8
+; CHECK-NEXT: store void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()** @x2, align 8
+; CHECK-NEXT: store void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* bitcast ({{.*}}@.cfi.jumptable to void ()*), void ()* null), void ()** @x, align 8
+; CHECK-NEXT: ret void
+; CHECK-NEXT: }
+
+
+!0 = !{i32 0, !"typeid1"}