[IPO/LowerTypesTest] Skip blockaddress when replacing uses

[davide]

Summary

Blockaddresses refer to the function itself, therefore replacing them would cause an assertion in doRAUW.
Fixes https://bugs.llvm.org/show_bug.cgi?id=35201
This was found when trying CFI on a proprietary kernel by Dmitry Mikulin.

[filcab]

Looks good. I’ll ask you to wait a bit and see if @pcc has any objection, though.

Thank you,

Filipe

[pcc]

lib/IR/Value.cpp

464

Can this line just be
if (isa<BlockAddress>(U.getUser()))
?

test/Transforms/LowerTypeTests/blockaddress.ll

17

I think this test case is more complicated than it needs to be. This seems to be enough to trigger the bug:

target triple = "x86_64-unknown-linux"

define void @f1() {
entry:
  %0 = call i1 @llvm.type.test(i8* bitcast (i8* ()* @f2 to i8*), metadata !"_ZTSFvP3bioE")
  ret void
}

declare i1 @llvm.type.test(i8*, metadata)

define i8* @f2() !type !5 {
  br label %b

b:
  ret i8* blockaddress(@f2, %b)
}

!5 = !{i64 0, !"_ZTSFvP3bioE"}

The only important part of the output is the body of f2, so I would only have CHECK lines for that.

[davide]

Address Peter's comments.

lib/IR/Value.cpp

464

yes.

test/Transforms/LowerTypeTests/blockaddress.ll

17

yes, thanks for reducing further.

[pcc]

LGTM

include/llvm/IR/Value.h

public:
 
   /// replaceUsesOutsideBlock - Go through the uses list for this definition and
   /// make each use point to "V" instead of "this" when the use is outside the
   /// block. 'This's use list is expected to have at least one element.
   /// Unlike replaceAllUsesWith this function does not support basic block
   /// values or constant users.
   void replaceUsesOutsideBlock(Value *V, BasicBlock *BB);
 
+  /// replaceUsesExceptBlockAddr - Go through the uses list for this definition
+  /// and make each use point to "V" instead of "this" when the use is outside
+  /// the block. 'This's use list is expected to have at least one element.
+  /// Unlike replaceAllUsesWith this function skips blockaddr uses.
+  void replaceUsesExceptBlockAddr(Value *New);
+
   //----------------------------------------------------------------------
   // Methods for handling the chain of uses of this Value.
   //
   // Materializing a function can introduce new uses, so these methods come in
   // two variants:
   // The methods that start with materialized_ check the uses that are
   // currently known given which functions are materialized. Be very careful
   // when using them since you might not get all uses.

lib/IR/Value.cpp

for (; UI != E;) {
     ++UI;
     auto *Usr = dyn_cast<Instruction>(U.getUser());
     if (Usr && Usr->getParent() == BB)
       continue;
     U.set(New);
   }
 }
 
+void Value::replaceUsesExceptBlockAddr(Value *New) {
+  use_iterator UI = use_begin(), E = use_end();
+  for (; UI != E;) {
+    Use &U = *UI;
+    ++UI;
+
+    if (isa<BlockAddress>(U.getUser()))
+      continue;
+
+    // Must handle Constants specially, we cannot call replaceUsesOfWith on a
+    // constant because they are uniqued.
+    if (auto *C = dyn_cast<Constant>(U.getUser())) {
+      if (!isa<GlobalValue>(C)) {
+        C->handleOperandChange(this, New);
+        continue;
+      }
+    }
+
+    U.set(New);
+  }
+}
+
 namespace {
 // Various metrics for how much to strip off of pointers.
 enum PointerStripKind {
   PSK_ZeroIndices,
   PSK_ZeroIndicesAndAliases,
   PSK_ZeroIndicesAndAliasesAndBarriers,
   PSK_InBoundsConstantIndices,
   PSK_InBounds

lib/Transforms/IPO/LowerTypeTests.cpp

       assert(F->getType()->getAddressSpace() == 0);
 
       GlobalAlias *FAlias = GlobalAlias::create(
           F->getValueType(), 0, F->getLinkage(), "", CombinedGlobalElemPtr, &M);
       FAlias->setVisibility(F->getVisibility());
       FAlias->takeName(F);
       if (FAlias->hasName())
         F->setName(FAlias->getName() + ".cfi");
-      F->replaceAllUsesWith(FAlias);
+      F->replaceUsesExceptBlockAddr(FAlias);
     }
     if (!F->isDeclarationForLinker())
       F->setLinkage(GlobalValue::InternalLinkage);
   }
 
   createJumpTable(JumpTableFn, Functions);
 }
 

test/Transforms/LowerTypeTests/blockaddress.ll

+; RUN: opt -S %s -lowertypetests | FileCheck %s
+
+
+; CHECK: define internal i8* @f2.cfi() !type !0 {
+; CHECK-NEXT:  br label %b
+; CHECK: b:
+; CHECK-NEXT:  ret i8* blockaddress(@f2.cfi, %b)
+; CHECK-NEXT: }
+
+target triple = "x86_64-unknown-linux"
+
+define void @f1() {
+entry:
+  %0 = call i1 @llvm.type.test(i8* bitcast (i8* ()* @f2 to i8*), metadata !"_ZTSFvP3bioE")
+  ret void
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+
+define i8* @f2() !type !5 {
+  br label %b
+
+b:
+  ret i8* blockaddress(@f2, %b)
+}
+
+!5 = !{i64 0, !"_ZTSFvP3bioE"}