[IR] Do not store Function inside BlockAddress

Currently BlockAddresses store both the Function and the BasicBlock
they reference, and the BlockAddress is part of the use list of
both the Function and BasicBlock.

This is quite awkward, because this is not really a use of the
function itself (and walks of function uses generally skip block
addresses for that reason). This also has weird implications
on function RAUW (as that will replace the function in block
addresses in a way that generally doesn't make sense), and
causes other peculiar issues, like the ability to have multiple
block addresses for one block (with different functions).

Instead, I believe it makes more sense to specify only the basic
block and let the function be implied by the BB parent. This
does mean that we may have block addresses without a function
(if the BB is not inserted), but this should only happen during
IR construction.

Author: nikic

clang/lib/CodeGen/CodeGenFunction.cpp

@@ -2286,7 +2286,7 @@ llvm::BlockAddress *CodeGenFunction::GetAddrOfLabel(const LabelDecl *L) {
 
   // Make sure the indirect branch includes all of the address-taken blocks.
   IndirectBranch->addDestination(BB);
-  return llvm::BlockAddress::get(CurFn, BB);
+  return llvm::BlockAddress::get(CurFn->getType(), BB);
 }
 
 llvm::BasicBlock *CodeGenFunction::GetIndirectGotoBlock() {

llvm/include/llvm/IR/Constants.h

@@ -893,9 +893,9 @@ class ConstantTargetNone final : public ConstantData {
 class BlockAddress final : public Constant {
   friend class Constant;
 
-  constexpr static IntrusiveOperandsAllocMarker AllocMarker{2};
+  constexpr static IntrusiveOperandsAllocMarker AllocMarker{1};
 
-  BlockAddress(Function *F, BasicBlock *BB);
+  BlockAddress(Type *Ty, BasicBlock *BB);
 
   void *operator new(size_t S) { return User::operator new(S, AllocMarker); }
 
@@ -912,6 +912,11 @@ class BlockAddress final : public Constant {
   /// block must be embedded into a function.
   static BlockAddress *get(BasicBlock *BB);
 
+  /// Return a BlockAddress for the specified basic block, which may not be
+  /// part of a function. The specified type must match the type of the function
+  /// the block will be inserted into.
+  static BlockAddress *get(Type *Ty, BasicBlock *BB);
+
   /// Lookup an existing \c BlockAddress constant for the given BasicBlock.
   ///
   /// \returns 0 if \c !BB->hasAddressTaken(), otherwise the \c BlockAddress.
@@ -920,8 +925,8 @@ class BlockAddress final : public Constant {
   /// Transparently provide more efficient getOperand methods.
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Value);
 
-  Function *getFunction() const { return (Function *)Op<0>().get(); }
-  BasicBlock *getBasicBlock() const { return (BasicBlock *)Op<1>().get(); }
+  BasicBlock *getBasicBlock() const { return (BasicBlock *)Op<0>().get(); }
+  Function *getFunction() const { return getBasicBlock()->getParent(); }
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const Value *V) {
@@ -931,7 +936,7 @@ class BlockAddress final : public Constant {
 
 template <>
 struct OperandTraits<BlockAddress>
-    : public FixedNumOperandTraits<BlockAddress, 2> {};
+    : public FixedNumOperandTraits<BlockAddress, 1> {};
 
 DEFINE_TRANSPARENT_OPERAND_ACCESSORS(BlockAddress, Value)
 

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -1649,7 +1649,7 @@ Expected<Value *> BitcodeReader::materializeValue(unsigned StartValID,
               FwdBBs[BBID] = BasicBlock::Create(Context);
             BB = FwdBBs[BBID];
           }
-          C = BlockAddress::get(Fn, BB);
+          C = BlockAddress::get(Fn->getType(), BB);
           break;
         }
         case BitcodeConstant::ConstantStructOpcode: {

llvm/lib/CodeGen/AsmPrinter/WinCFGuard.cpp

@@ -49,10 +49,6 @@ static bool isPossibleIndirectCallTarget(const Function *F) {
     const Value *FnOrCast = Users.pop_back_val();
     for (const Use &U : FnOrCast->uses()) {
       const User *FnUser = U.getUser();
-      if (isa<BlockAddress>(FnUser)) {
-        // Block addresses are illegal to call.
-        continue;
-      }
       if (const auto *Call = dyn_cast<CallBase>(FnUser)) {
         if ((!Call->isCallee(&U) || U.get() != F) &&
             !Call->getFunction()->getName().ends_with("$exit_thunk")) {

llvm/lib/IR/Constants.cpp

@@ -1891,77 +1891,61 @@ void PoisonValue::destroyConstantImpl() {
   getContext().pImpl->PVConstants.erase(getType());
 }
 
+BlockAddress *BlockAddress::get(Type *Ty, BasicBlock *BB) {
+  BlockAddress *&BA = BB->getContext().pImpl->BlockAddresses[BB];
+  if (!BA)
+    BA = new BlockAddress(Ty, BB);
+  return BA;
+}
+
 BlockAddress *BlockAddress::get(BasicBlock *BB) {
   assert(BB->getParent() && "Block must have a parent");
-  return get(BB->getParent(), BB);
+  return get(BB->getParent()->getType(), BB);
 }
 
 BlockAddress *BlockAddress::get(Function *F, BasicBlock *BB) {
-  BlockAddress *&BA =
-    F->getContext().pImpl->BlockAddresses[std::make_pair(F, BB)];
-  if (!BA)
-    BA = new BlockAddress(F, BB);
-
-  assert(BA->getFunction() == F && "Basic block moved between functions");
-  return BA;
+  assert(BB->getParent() == F && "Block not part of specified function");
+  return get(BB->getParent()->getType(), BB);
 }
 
-BlockAddress::BlockAddress(Function *F, BasicBlock *BB)
-    : Constant(PointerType::get(F->getContext(), F->getAddressSpace()),
-               Value::BlockAddressVal, AllocMarker) {
-  setOperand(0, F);
-  setOperand(1, BB);
+BlockAddress::BlockAddress(Type *Ty, BasicBlock *BB)
+    : Constant(Ty, Value::BlockAddressVal, AllocMarker) {
+  setOperand(0, BB);
   BB->AdjustBlockAddressRefCount(1);
 }
 
 BlockAddress *BlockAddress::lookup(const BasicBlock *BB) {
   if (!BB->hasAddressTaken())
     return nullptr;
 
-  const Function *F = BB->getParent();
-  assert(F && "Block must have a parent");
-  BlockAddress *BA =
-      F->getContext().pImpl->BlockAddresses.lookup(std::make_pair(F, BB));
+  BlockAddress *BA = BB->getContext().pImpl->BlockAddresses.lookup(BB);
   assert(BA && "Refcount and block address map disagree!");
   return BA;
 }
 
 /// Remove the constant from the constant table.
 void BlockAddress::destroyConstantImpl() {
-  getFunction()->getType()->getContext().pImpl
-    ->BlockAddresses.erase(std::make_pair(getFunction(), getBasicBlock()));
+  getType()->getContext().pImpl->BlockAddresses.erase(getBasicBlock());
   getBasicBlock()->AdjustBlockAddressRefCount(-1);
 }
 
 Value *BlockAddress::handleOperandChangeImpl(Value *From, Value *To) {
-  // This could be replacing either the Basic Block or the Function.  In either
-  // case, we have to remove the map entry.
-  Function *NewF = getFunction();
-  BasicBlock *NewBB = getBasicBlock();
-
-  if (From == NewF)
-    NewF = cast<Function>(To->stripPointerCasts());
-  else {
-    assert(From == NewBB && "From does not match any operand");
-    NewBB = cast<BasicBlock>(To);
-  }
+  assert(From == getBasicBlock());
+  BasicBlock *NewBB = cast<BasicBlock>(To);
 
   // See if the 'new' entry already exists, if not, just update this in place
   // and return early.
-  BlockAddress *&NewBA =
-    getContext().pImpl->BlockAddresses[std::make_pair(NewF, NewBB)];
+  BlockAddress *&NewBA = getContext().pImpl->BlockAddresses[NewBB];
   if (NewBA)
     return NewBA;
 
   getBasicBlock()->AdjustBlockAddressRefCount(-1);
 
   // Remove the old entry, this can't cause the map to rehash (just a
   // tombstone will get added).
-  getContext().pImpl->BlockAddresses.erase(std::make_pair(getFunction(),
-                                                          getBasicBlock()));
+  getContext().pImpl->BlockAddresses.erase(getBasicBlock());
   NewBA = this;
-  setOperand(0, NewF);
-  setOperand(1, NewBB);
+  setOperand(0, NewBB);
   getBasicBlock()->AdjustBlockAddressRefCount(1);
 
   // If we just want to keep the existing value, then return null.

llvm/lib/IR/Function.cpp

@@ -966,9 +966,6 @@ bool Function::hasAddressTaken(const User **PutOffender,
                                bool IgnoreCastedDirectCall) const {
   for (const Use &U : uses()) {
     const User *FU = U.getUser();
-    if (isa<BlockAddress>(FU))
-      continue;
-
     if (IgnoreCallbackUses) {
       AbstractCallSite ACS(&U);
       if (ACS && ACS.isCallbackCall())
@@ -1033,12 +1030,7 @@ bool Function::isDefTriviallyDead() const {
       !hasAvailableExternallyLinkage())
     return false;
 
-  // Check if the function is used by anything other than a blockaddress.
-  for (const User *U : users())
-    if (!isa<BlockAddress>(U))
-      return false;
-
-  return true;
+  return use_empty();
 }
 
 /// callsFunctionThatReturnsTwice - Return true if the function has a call to

llvm/lib/IR/LLVMContextImpl.h

@@ -1688,8 +1688,7 @@ class LLVMContextImpl {
 
   StringMap<std::unique_ptr<ConstantDataSequential>> CDSConstants;
 
-  DenseMap<std::pair<const Function *, const BasicBlock *>, BlockAddress *>
-      BlockAddresses;
+  DenseMap<const BasicBlock *, BlockAddress *> BlockAddresses;
 
   DenseMap<const GlobalValue *, DSOLocalEquivalent *> DSOLocalEquivalents;
 

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -1936,9 +1936,6 @@ bool Attributor::checkForAllCallSites(function_ref<bool(AbstractCallSite)> Pred,
       LLVM_DEBUG(dbgs() << "[Attributor] Function " << Fn.getName()
                         << " has non call site use " << *U.get() << " in "
                         << *U.getUser() << "\n");
-      // BlockAddress users are allowed.
-      if (isa<BlockAddress>(U.getUser()))
-        continue;
       return false;
     }
 
@@ -3061,14 +3058,6 @@ ChangeStatus Attributor::rewriteFunctionSignatures(
     // function empty.
     NewFn->splice(NewFn->begin(), OldFn);
 
-    // Fixup block addresses to reference new function.
-    SmallVector<BlockAddress *, 8u> BlockAddresses;
-    for (User *U : OldFn->users())
-      if (auto *BA = dyn_cast<BlockAddress>(U))
-        BlockAddresses.push_back(BA);
-    for (auto *BA : BlockAddresses)
-      BA->replaceAllUsesWith(BlockAddress::get(NewFn, BA->getBasicBlock()));
-
     // Set of all "call-like" instructions that invoke the old function mapped
     // to their new replacements.
     SmallVector<std::pair<CallBase *, CallBase *>, 8> CallSitePairs;

llvm/lib/Transforms/IPO/GlobalOpt.cpp

@@ -1678,11 +1678,8 @@ processGlobal(GlobalValue &GV,
 /// Walk all of the direct calls of the specified function, changing them to
 /// FastCC.
 static void ChangeCalleesToFastCall(Function *F) {
-  for (User *U : F->users()) {
-    if (isa<BlockAddress>(U))
-      continue;
+  for (User *U : F->users())
     cast<CallBase>(U)->setCallingConv(CallingConv::Fast);
-  }
 }
 
 static AttributeList StripAttr(LLVMContext &C, AttributeList Attrs,
@@ -1696,8 +1693,6 @@ static AttributeList StripAttr(LLVMContext &C, AttributeList Attrs,
 static void RemoveAttribute(Function *F, Attribute::AttrKind A) {
   F->setAttributes(StripAttr(F->getContext(), F->getAttributes(), A));
   for (User *U : F->users()) {
-    if (isa<BlockAddress>(U))
-      continue;
     CallBase *CB = cast<CallBase>(U);
     CB->setAttributes(StripAttr(F->getContext(), CB->getAttributes(), A));
   }
@@ -1722,8 +1717,6 @@ static bool hasChangeableCCImpl(Function *F) {
   // Can't change CC of the function that either has musttail calls, or is a
   // musttail callee itself
   for (User *U : F->users()) {
-    if (isa<BlockAddress>(U))
-      continue;
     CallInst* CI = dyn_cast<CallInst>(U);
     if (!CI)
       continue;
@@ -1772,9 +1765,6 @@ isValidCandidateForColdCC(Function &F,
     return false;
 
   for (User *U : F.users()) {
-    if (isa<BlockAddress>(U))
-      continue;
-
     CallBase &CB = cast<CallBase>(*U);
     Function *CallerFunc = CB.getParent()->getParent();
     BlockFrequencyInfo &CallerBFI = GetBFI(*CallerFunc);
@@ -1787,11 +1777,8 @@ isValidCandidateForColdCC(Function &F,
 }
 
 static void changeCallSitesToColdCC(Function *F) {
-  for (User *U : F->users()) {
-    if (isa<BlockAddress>(U))
-      continue;
+  for (User *U : F->users())
     cast<CallBase>(U)->setCallingConv(CallingConv::Cold);
-  }
 }
 
 // This function iterates over all the call instructions in the input Function
@@ -1832,12 +1819,7 @@ hasOnlyColdCalls(Function &F,
 
 static bool hasMustTailCallers(Function *F) {
   for (User *U : F->users()) {
-    CallBase *CB = dyn_cast<CallBase>(U);
-    if (!CB) {
-      assert(isa<BlockAddress>(U) &&
-             "Expected either CallBase or BlockAddress");
-      continue;
-    }
+    CallBase *CB = cast<CallBase>(U);
     if (CB->isMustTailCall())
       return true;
   }

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

@@ -1920,9 +1920,9 @@ void LowerTypeTestsModule::replaceCfiUses(Function *Old, Value *New,
                                           bool IsJumpTableCanonical) {
   SmallSetVector<Constant *, 4> Constants;
   for (Use &U : llvm::make_early_inc_range(Old->uses())) {
-    // Skip block addresses and no_cfi values, which refer to the function
-    // body instead of the jump table.
-    if (isa<BlockAddress, NoCFIValue>(U.getUser()))
+    // Skip no_cfi values, which refer to the function body instead of the jump
+    // table.
+    if (isa<NoCFIValue>(U.getUser()))
       continue;
 
     // Skip direct calls to externally defined or non-dso_local functions.