diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp --- a/llvm/lib/Transforms/Utils/CloneFunction.cpp +++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp @@ -13,6 +13,7 @@ //===----------------------------------------------------------------------===// #include "llvm/ADT/SetVector.h" +#include "llvm/ADT/SmallSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Analysis/DomTreeUpdater.h" #include "llvm/Analysis/InstructionSimplify.h" @@ -22,6 +23,7 @@ #include "llvm/IR/DebugInfo.h" #include "llvm/IR/DerivedTypes.h" #include "llvm/IR/Function.h" +#include "llvm/IR/GlobalVariable.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/LLVMContext.h" @@ -75,6 +77,33 @@ return NewBB; } +static void cloneGlobalVariablesUsingBlockAddress( + SmallSet &GVsUsingBlockAddress, Module *M, + ValueToValueMapTy &VMap) { + for (GlobalVariable *I : GVsUsingBlockAddress) { + GlobalVariable *NewGV = new GlobalVariable( + *M, I->getValueType(), I->isConstant(), I->getLinkage(), + (Constant *)nullptr, I->getName(), I, I->getThreadLocalMode(), + I->getType()->getAddressSpace()); + NewGV->copyAttributesFrom(I); + VMap[I] = NewGV; + + SmallVector, 1> MDs; + I->getAllMetadata(MDs); + for (auto MD : MDs) + NewGV->addMetadata(MD.first, *MapMetadata(MD.second, VMap)); + + assert(I->hasInitializer()); + NewGV->setInitializer(MapValue(I->getInitializer(), VMap)); + + if (const Comdat *SC = I->getComdat()) { + Comdat *DC = M->getOrInsertComdat(SC->getName()); + DC->setSelectionKind(SC->getSelectionKind()); + NewGV->setComdat(DC); + } + } +} + // Clone OldFunc into NewFunc, transforming the old arguments into references to // VMap values. // @@ -134,10 +163,13 @@ // types. Optional DIFinder; + bool CurrentModuleChanges = + Changes < CloneFunctionChangeType::DifferentModule; + // Track the subprogram attachment that needs to be cloned to fine-tune the // mapping within the same module. DISubprogram *SPClonedWithinModule = nullptr; - if (Changes < CloneFunctionChangeType::DifferentModule) { + if (CurrentModuleChanges) { assert((NewFunc->getParent() == nullptr || NewFunc->getParent() == OldFunc->getParent()) && "Expected NewFunc to have the same parent, or no parent"); @@ -162,6 +194,8 @@ } } + SmallSet GVsUsingBlockAddress; + // Loop over all of the basic blocks in the function, cloning them as // appropriate. Note that we save BE this way in order to handle cloning of // recursive functions into themselves. @@ -184,6 +218,16 @@ Constant *OldBBAddr = BlockAddress::get(const_cast(OldFunc), const_cast(&BB)); VMap[OldBBAddr] = BlockAddress::get(NewFunc, CBB); + + if (CurrentModuleChanges) + for (User *U : OldBBAddr->users()) { + while (isa(U) && !isa(U)) { + assert(U->getNumUses() == 1); + U = cast(U)->user_back(); + } + if (auto *GV = dyn_cast(U)) + GVsUsingBlockAddress.insert(GV); + } } // Note return instructions for the caller. @@ -191,6 +235,10 @@ Returns.push_back(RI); } + if (CurrentModuleChanges) + cloneGlobalVariablesUsingBlockAddress(GVsUsingBlockAddress, + NewFunc->getParent(), VMap); + if (Changes < CloneFunctionChangeType::DifferentModule && DIFinder->subprogram_count() > 0) { // Turn on module-level changes, since we need to clone (some of) the diff --git a/llvm/unittests/Transforms/Utils/CloningTest.cpp b/llvm/unittests/Transforms/Utils/CloningTest.cpp --- a/llvm/unittests/Transforms/Utils/CloningTest.cpp +++ b/llvm/unittests/Transforms/Utils/CloningTest.cpp @@ -749,6 +749,45 @@ EXPECT_TRUE(CCI.ContainsDynamicAllocas); } +TEST(CloneFunction, CloneFunctionWithBlockAddressTaken) { + StringRef ImplAssembly = R"( + @ba = internal unnamed_addr constant [1 x ptr] [ptr blockaddress(@foo, %stop)] + + define void @foo() { + entry: + %arrayidx = getelementptr inbounds [1 x ptr], ptr @ba, i64 0, i64 0 + %0 = load ptr, ptr %arrayidx + br label %indirectgoto + + stop: + ret void + + indirectgoto: + %indirect.goto.dest = phi ptr [ %0, %entry ] + indirectbr ptr %indirect.goto.dest, [label %stop] + } + + declare void @bar() + )"; + + LLVMContext Context; + SMDiagnostic Error; + + auto ImplModule = parseAssemblyString(ImplAssembly, Error, Context); + EXPECT_TRUE(ImplModule != nullptr); + auto *ImplFunction = ImplModule->getFunction("foo"); + EXPECT_TRUE(ImplFunction != nullptr); + auto *DeclFunction = ImplModule->getFunction("bar"); + EXPECT_TRUE(DeclFunction != nullptr); + + ValueToValueMapTy VMap; + SmallVector Returns; + CloneFunctionInto(DeclFunction, ImplFunction, VMap, + CloneFunctionChangeType::GlobalChanges, Returns, ""); + + EXPECT_FALSE(verifyModule(*ImplModule, &errs())); +} + TEST(CloneFunction, CloneFunctionWithSubprograms) { // Tests that the debug info is duplicated correctly when a DISubprogram // happens to be one of the operands of the DISubprogram that is being cloned.