Patch to fix cloning block addresses correctly in CloneModule

[Hatsunespica]

Summary

Current CloneModule failed to clone block addresses, which leads to crashes in optimizing cloned modules.
https://github.com/llvm/llvm-project/issues/55991
As a result, we should create empty basic blocks inside functions referred by block addresses, and fill them with instructions later.

llvm/include/llvm/Transforms/Utils/Cloning.h

 /// names, specify it as the optional third parameter.
 ///
 /// If you would like the basic block to be auto-inserted into the end of a
 /// function, you can specify it as the optional fourth parameter.
 ///
 /// If you would like to collect additional information about the cloned
 /// function, you can specify a ClonedCodeInfo object with the optional fifth
 /// parameter.
+///
+/// If NewBB is not empty, all instructions will be inserted into newBB in
+/// the case of creating a basic block first and cloning it later.
 BasicBlock *CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
                             const Twine &NameSuffix = "", Function *F = nullptr,
                             ClonedCodeInfo *CodeInfo = nullptr,
-                            DebugInfoFinder *DIFinder = nullptr);
+                            DebugInfoFinder *DIFinder = nullptr,
+                            BasicBlock *NewBB = nullptr);
 
 /// Return a copy of the specified function and add it to that
 /// function's module.  Also, any references specified in the VMap are changed
 /// to refer to their mapped value instead of the original one.  If any of the
 /// arguments to the function are in the VMap, the arguments are deleted from
 /// the resultant function.  The VMap is updated to include mappings from all of
 /// the instructions and basicblocks in the function from their old to new
 /// values.  The final argument captures information about the cloned code if

 /// CloneFunctionMetadataInto() and expecting / updating callers to call it
 /// first when / how it's needed.
 void CloneFunctionInto(Function *NewFunc, const Function *OldFunc,
                        ValueToValueMapTy &VMap, CloneFunctionChangeType Changes,
                        SmallVectorImpl<ReturnInst *> &Returns,
                        const char *NameSuffix = "",
                        ClonedCodeInfo *CodeInfo = nullptr,
                        ValueMapTypeRemapper *TypeMapper = nullptr,
-                       ValueMaterializer *Materializer = nullptr);
+                       ValueMaterializer *Materializer = nullptr,
+                       bool shouldCreateBB = true);
 
 void CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
                                const Instruction *StartingInst,
                                ValueToValueMapTy &VMap, bool ModuleLevelChanges,
                                SmallVectorImpl<ReturnInst *> &Returns,
                                const char *NameSuffix = "",
                                ClonedCodeInfo *CodeInfo = nullptr);
 

llvm/lib/Transforms/Utils/CloneFunction.cpp

 using namespace llvm;
 
 #define DEBUG_TYPE "clone-function"
 
 /// See comments in Cloning.h.
 BasicBlock *llvm::CloneBasicBlock(const BasicBlock *BB, ValueToValueMapTy &VMap,
                                   const Twine &NameSuffix, Function *F,
                                   ClonedCodeInfo *CodeInfo,
-                                  DebugInfoFinder *DIFinder) {
-  BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "", F);
+                                  DebugInfoFinder *DIFinder,
+                                  BasicBlock *NewBB) {
+  if (NewBB == nullptr)
+    NewBB = BasicBlock::Create(BB->getContext(), "", F);
   if (BB->hasName())
     NewBB->setName(BB->getName() + NameSuffix);
 
   bool hasCalls = false, hasDynamicAllocas = false;
   Module *TheModule = F ? F->getParent() : nullptr;
 
   // Loop over all instructions, and copy them over.
   for (const Instruction &I : *BB) {

 // VMap values.
 //
 void llvm::CloneFunctionInto(Function *NewFunc, const Function *OldFunc,
                              ValueToValueMapTy &VMap,
                              CloneFunctionChangeType Changes,
                              SmallVectorImpl<ReturnInst *> &Returns,
                              const char *NameSuffix, ClonedCodeInfo *CodeInfo,
                              ValueMapTypeRemapper *TypeMapper,
-                             ValueMaterializer *Materializer) {
+                             ValueMaterializer *Materializer,
+                             bool shouldCreateBB) {
   assert(NameSuffix && "NameSuffix cannot be null!");
 
 #ifndef NDEBUG
   for (const Argument &I : OldFunc->args())
     assert(VMap.count(&I) && "No mapping from source argument specified!");
 #endif
 
   bool ModuleLevelChanges = Changes > CloneFunctionChangeType::LocalChangesOnly;
if (Changes == CloneFunctionChangeType::DifferentModule) {
       // Need to find all the compile units.
       DIFinder.emplace();
     }
   }
 
   // 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.
+  auto NewBBIt = NewFunc->begin();
   for (const BasicBlock &BB : *OldFunc) {
 
+    BasicBlock *CBB = nullptr;
     // Create a new basic block and copy instructions into it!
-    BasicBlock *CBB = CloneBasicBlock(&BB, VMap, NameSuffix, NewFunc, CodeInfo,
-                                      DIFinder ? &*DIFinder : nullptr);
+    if (shouldCreateBB) {
+      CBB = CloneBasicBlock(&BB, VMap, NameSuffix, NewFunc, CodeInfo,
+    } else {
+      CBB = CloneBasicBlock(&BB, VMap, NameSuffix, NewFunc, CodeInfo,
+                            DIFinder ? &*DIFinder : nullptr, &*NewBBIt);
+      ++NewBBIt;
+    }
 
     // Add basic block mapping.
     VMap[&BB] = CBB;
 
     // It is only legal to clone a function if a block address within that
     // function is never referenced outside of the function.  Given that, we
     // want to map block addresses from the old function to block addresses in
     // the clone. (This is different from the generic ValueMapper

llvm/lib/Transforms/Utils/CloneModule.cpp

 //===- CloneModule.cpp - Clone an entire module ---------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the CloneModule interface which makes a copy of an
 // entire module.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 using namespace llvm;
 
 namespace llvm {
 class Constant;
std::unique_ptr<Module> llvm::CloneModule(
 
   // Loop over the functions in the module, making external functions as before
   for (const Function &I : M) {
     Function *NF =
         Function::Create(cast<FunctionType>(I.getValueType()), I.getLinkage(),
                          I.getAddressSpace(), I.getName(), New.get());
     NF->copyAttributesFrom(&I);
     VMap[&I] = NF;
+
+    // We should create basic blocks in the function in case we have to
+    // clone them because BlockAddress might reference to them.
+    if (!I.isDeclaration() && ShouldCloneDefinition(&I)) {
+      for (const BasicBlock &BB : I) {
+        BasicBlock *CBB = BasicBlock::Create(BB.getContext(), "", NF);
+
+        VMap[&BB] = CBB;
+        if (BB.hasAddressTaken()) {
+          Constant *OldBBAddr = BlockAddress::get(
+              const_cast<Function *>(&I), const_cast<BasicBlock *>(&BB));
+          VMap[OldBBAddr] = BlockAddress::get(NF, CBB);
+        }
+      }
+    }
   }
 
   // Loop over the aliases in the module
   for (const GlobalAlias &I : M.aliases()) {
     if (!ShouldCloneDefinition(&I)) {
       // An alias cannot act as an external reference, so we need to create
       // either a function or a global variable depending on the value type.
       // FIXME: Once pointee types are gone we can probably pick one or the
for (const Function &I : M) {
     Function::arg_iterator DestI = F->arg_begin();
     for (const Argument &J : I.args()) {
       DestI->setName(J.getName());
       VMap[&J] = &*DestI++;
     }
 
     SmallVector<ReturnInst *, 8> Returns; // Ignore returns cloned.
     CloneFunctionInto(F, &I, VMap, CloneFunctionChangeType::ClonedModule,
-                      Returns);
+                      Returns, "", nullptr, nullptr, nullptr, false);
 
     if (I.hasPersonalityFn())
       F->setPersonalityFn(MapValue(I.getPersonalityFn(), VMap));
 
     copyComdat(F, &I);
   }
 
   // And aliases

llvm/unittests/Transforms/Utils/CloningTest.cpp

 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
 
 class CloneInstruction : public ::testing::Test {
TEST_F(CloneModule, Comdat) {
   auto *CD = NewGV->getComdat();
   ASSERT_NE(nullptr, CD);
   EXPECT_EQ("comdat", CD->getName());
   EXPECT_EQ(Comdat::ExactMatch, CD->getSelectionKind());
 
   Function *NewF = NewM->getFunction("f");
   EXPECT_EQ(CD, NewF->getComdat());
 }
+
+static void optimizeModule(llvm::Module &M) {
+  llvm::LoopAnalysisManager LAM;
+  llvm::FunctionAnalysisManager FAM;
+  llvm::CGSCCAnalysisManager CGAM;
+  llvm::ModuleAnalysisManager MAM;
+  llvm::PassBuilder PB;
+  llvm::ModulePassManager MPM;
+  PB.registerModuleAnalyses(MAM);
+  PB.registerCGSCCAnalyses(CGAM);
+  PB.registerFunctionAnalyses(FAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  MPM = PB.buildPerModuleDefaultPipeline(OptimizationLevel::O2);
+  MPM.run(M, MAM);
+}
+
+TEST(CloneModule2, CloneCrash) {
+  StringRef ImplAssembly = R"(
[e-mail] = constant [2 x i8*] [i8* blockaddress(@loop, %bb0), i8* blockaddress(@loop, %bb1)]
[e-mail] = constant i8* blockaddress(@loop, %bb1)
+
+define void @loop(i64* %p) {
+entry:
+  br label %bb0
+bb0:
+  %target = load i8*, i8** @loop.target, align 8
+  indirectbr i8* %target, [label %bb0, label %bb1]
+bb1:
+  ret void
+}
+    )";
+  LLVMContext Context;
+  SMDiagnostic Error;
+  auto ImplModule = parseAssemblyString(ImplAssembly, Error, Context);
+  EXPECT_TRUE(ImplModule != nullptr);
+  EXPECT_FALSE(verifyModule(*ImplModule));
+  llvm::ValueToValueMapTy vmap;
+  auto NewModule = llvm::CloneModule(*ImplModule, vmap);
+  EXPECT_FALSE(llvm::verifyModule(*NewModule, &llvm::errs()));
+  NewModule->print(llvm::errs(), nullptr);
+  optimizeModule(*NewModule);
+  NewModule->print(llvm::errs(), nullptr);
+}
 }