From 1a46affb45c5e3c2e8645cde04a8690583d8b912 Mon Sep 17 00:00:00 2001
From: Davide Italiano <davide@freebsd.org>
Date: Tue, 7 Nov 2017 00:09:25 +0000
Subject: [PATCH] [IPO/LowerTypesTest] Skip blockaddress(es) when replacing
 uses.

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.

Differential Revision:  https://reviews.llvm.org/D39695

llvm-svn: 317527
---
 llvm/include/llvm/IR/Value.h                       |  6 +++++
 llvm/lib/IR/Value.cpp                              | 22 ++++++++++++++++++
 llvm/lib/Transforms/IPO/LowerTypeTests.cpp         |  2 +-
 .../test/Transforms/LowerTypeTests/blockaddress.ll | 27 ++++++++++++++++++++++
 4 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LowerTypeTests/blockaddress.ll

diff --git a/llvm/include/llvm/IR/Value.h b/llvm/include/llvm/IR/Value.h
index 9e49149..f50f017 100644
--- a/llvm/include/llvm/IR/Value.h
+++ b/llvm/include/llvm/IR/Value.h
@@ -299,6 +299,12 @@ public:
   /// 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.
   //
diff --git a/llvm/lib/IR/Value.cpp b/llvm/lib/IR/Value.cpp
index 51a7d42..5df0c6d 100644
--- a/llvm/lib/IR/Value.cpp
+++ b/llvm/lib/IR/Value.cpp
@@ -454,6 +454,28 @@ void Value::replaceUsesOutsideBlock(Value *New, BasicBlock *BB) {
   }
 }
 
+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 {
diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 9fa5ed9..6cef866 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -1401,7 +1401,7 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
       FAlias->takeName(F);
       if (FAlias->hasName())
         F->setName(FAlias->getName() + ".cfi");
-      F->replaceAllUsesWith(FAlias);
+      F->replaceUsesExceptBlockAddr(FAlias);
     }
     if (!F->isDeclarationForLinker())
       F->setLinkage(GlobalValue::InternalLinkage);
diff --git a/llvm/test/Transforms/LowerTypeTests/blockaddress.ll b/llvm/test/Transforms/LowerTypeTests/blockaddress.ll
new file mode 100644
index 0000000..ecc4814
--- /dev/null
+++ b/llvm/test/Transforms/LowerTypeTests/blockaddress.ll
@@ -0,0 +1,27 @@
+; 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"}
-- 
2.7.4