diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp --- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp +++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp @@ -19,6 +19,7 @@ #include "llvm/ADT/None.h" #include "llvm/ADT/Optional.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SetVector.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" @@ -3365,7 +3366,7 @@ bool NeedsMetadataAttachment = F.hasMetadata(); DILocation *LastDL = nullptr; - SmallPtrSet BlockAddressUsers; + SmallSetVector BlockAddressUsers; // Finally, emit all the instructions, in order. for (const BasicBlock &BB : F) { @@ -3401,11 +3402,24 @@ } if (BlockAddress *BA = BlockAddress::lookup(&BB)) { - for (User *U : BA->users()) { - if (auto *I = dyn_cast(U)) { - Function *P = I->getParent()->getParent(); - if (P != &F) - BlockAddressUsers.insert(P); + SmallVector BlockAddressUsersStack { BA }; + SmallPtrSet BlockAddressUsersVisited { BA }; + + while (!BlockAddressUsersStack.empty()) { + Value *V = BlockAddressUsersStack.pop_back_val(); + + for (User *U : V->users()) { + if ((isa(U) || isa(U)) && + !BlockAddressUsersVisited.contains(U)) { + BlockAddressUsersStack.push_back(U); + BlockAddressUsersVisited.insert(U); + } + + if (auto *I = dyn_cast(U)) { + Function *P = I->getParent()->getParent(); + if (P != &F) + BlockAddressUsers.insert(P); + } } } } diff --git a/llvm/test/Bitcode/blockaddress-aggregate-users.ll b/llvm/test/Bitcode/blockaddress-aggregate-users.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Bitcode/blockaddress-aggregate-users.ll @@ -0,0 +1,40 @@ +; RUN: llvm-as %s -o %t.bc +; RUN: llvm-bcanalyzer -dump %t.bc | FileCheck %s +; RUN: llvm-dis %t.bc + +; There's a curious case where blockaddress constants may refer to functions +; outside of the function they're used in. There's a special bitcode function +; code, FUNC_CODE_BLOCKADDR_USERS, used to signify that this is the case. + +; The intent of this test is two-fold: +; 1. Ensure we produce BLOCKADDR_USERS bitcode function code on the first fn, +; @repro, since @fun and @fun2 both refer to @repro via blockaddress +; constants. +; 2. Ensure we can round-trip serializing+desearlizing such case. + +; CHECK: