[FuzzMutate] Handle BB without predecessor, avoid insertion after `musttail call`, avoid sinking token type

[HazyFish]

Summary

FuzzMutate didn't consider some corner cases and leads to mutation failure when mutating some modules.
This patch fixes 3 bugs:

• Add null check when encountering basic blocks without predecessor to avoid segmentation fault
• Avoid insertion after musttail call instruction
• Avoid sinking token type

Unit tests are also added.

[HazyFish]

Avoid sinking token type

llvm/lib/FuzzMutate/IRMutator.cpp

auto OpMatchesPred = [&Src](fuzzerop::OpDescriptor &Op) {
     return Op.SourcePreds[0].matches({}, Src);
   };
   auto RS = makeSampler(IB.Rand, make_filter_range(Operations, OpMatchesPred));
   if (RS.isEmpty())
     return std::nullopt;
   return *RS;
 }
 
+static inline iterator_range<BasicBlock::iterator>
+getInsertionRange(BasicBlock &BB) {
+  auto End = BB.getTerminatingMustTailCall() ? std::prev(BB.end()) : BB.end();
+  return make_range(BB.getFirstInsertionPt(), End);
+}
+
 void InjectorIRStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   SmallVector<Instruction *, 32> Insts;
-  for (auto I = BB.getFirstInsertionPt(), E = BB.end(); I != E; ++I)
-    Insts.push_back(&*I);
+  for (Instruction &I : getInsertionRange(BB))
+    Insts.push_back(&I);
   if (Insts.size() < 1)
     return;
 
   // Choose an insertion point for our new instruction.
   size_t IP = uniform<size_t>(IB.Rand, 0, Insts.size() - 1);
 
   auto InstsBefore = ArrayRef(Insts).slice(0, IP);
   auto InstsAfter = ArrayRef(Insts).slice(IP);
void InsertFunctionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   // If nullptr is selected, we will create a new function declaration.
   SmallVector<Function *, 32> Functions({nullptr});
   for (Function &F : M->functions()) {
     Functions.push_back(&F);
   }
 
   auto RS = makeSampler(IB.Rand, Functions);
   Function *F = RS.getSelection();
+  // Some functions accept metadata type or token type as arguments.
+  // We don't call those functions for now.
+  // For example, `@llvm.dbg.declare(metadata, metadata, metadata)`
+  // https://llvm.org/docs/SourceLevelDebugging.html#llvm-dbg-declare
   auto IsUnsupportedTy = [](Type *T) {
     return T->isMetadataTy() || T->isTokenTy();
   };
   if (!F || IsUnsupportedTy(F->getReturnType()) ||
       any_of(F->getFunctionType()->params(), IsUnsupportedTy)) {
     F = IB.createFunctionDeclaration(*M);
   }
 
auto BuilderFunc = [FTy, F, isRetVoid](ArrayRef<Value *> Srcs,
                                          Instruction *Inst) {
     StringRef Name = isRetVoid ? nullptr : "C";
     CallInst *Call = CallInst::Create(FTy, F, Srcs, Name, Inst);
     // Don't return this call inst if it return void as it can't be sinked.
     return isRetVoid ? nullptr : Call;
   };
 
   SmallVector<Instruction *, 32> Insts;
-  for (Instruction &I : make_range(BB.getFirstInsertionPt(), BB.end()))
+  for (Instruction &I : getInsertionRange(BB))
     Insts.push_back(&I);
   if (Insts.size() < 1)
     return;
 
   // Choose an insertion point for our new call instruction.
   uint64_t IP = uniform<uint64_t>(IB.Rand, 0, Insts.size() - 1);
 
   auto InstsBefore = ArrayRef(Insts).slice(0, IP);
void InsertFunctionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   if (Value *Op = BuilderFunc(Srcs, Insts[IP])) {
     // Find a sink and wire up the results of the operation.
     IB.connectToSink(BB, InstsAfter, Op);
   }
 }
 
 void InsertCFGStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   SmallVector<Instruction *, 32> Insts;
-  for (Instruction &I : make_range(BB.getFirstInsertionPt(), BB.end()))
+  for (Instruction &I : getInsertionRange(BB))
     Insts.push_back(&I);
   if (Insts.size() < 1)
     return;
 
   // Choose a point where we split the block.
   uint64_t IP = uniform<uint64_t>(IB.Rand, 0, Insts.size() - 1);
   auto InstsBeforeSplit = ArrayRef(Insts).slice(0, IP);
 
if (!Src) {
       // There is no need to inform IB what previously used values are if we are
       // using `onlyType`
       Src = IB.findOrCreateSource(*Pred, Insts, {}, fuzzerop::onlyType(Ty));
       IncomingValues[Pred] = Src;
     }
     PHI->addIncoming(Src, Pred);
   }
   SmallVector<Instruction *, 32> InstsAfter;
-  for (Instruction &I : make_range(BB.getFirstInsertionPt(), BB.end()))
+  for (Instruction &I : getInsertionRange(BB))
     InstsAfter.push_back(&I);
   IB.connectToSink(BB, InstsAfter, PHI);
 }
 
 void SinkInstructionStrategy::mutate(Function &F, RandomIRBuilder &IB) {
   for (BasicBlock &BB : F) {
     this->mutate(BB, IB);
   }
 }
 void SinkInstructionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   SmallVector<Instruction *, 32> Insts;
-  for (Instruction &I : make_range(BB.getFirstInsertionPt(), BB.end()))
+  for (Instruction &I : getInsertionRange(BB))
     Insts.push_back(&I);
   if (Insts.size() < 1)
     return;
   // Choose an Instruction to mutate.
   uint64_t Idx = uniform<uint64_t>(IB.Rand, 0, Insts.size() - 1);
   Instruction *Inst = Insts[Idx];
   // `Idx + 1` so we don't sink to ourselves.
   auto InstsAfter = ArrayRef(Insts).slice(Idx + 1);
-  LLVMContext &C = BB.getParent()->getParent()->getContext();
-  // Don't sink terminators, void function calls, etc.
-  if (Inst->getType() != Type::getVoidTy(C))
+  Type *Ty = Inst->getType();
+  // Don't sink terminators, void function calls, token, etc.
+  if (!Ty->isVoidTy() && !Ty->isTokenTy())
     // Find a new sink and wire up the results of the operation.
     IB.connectToSink(BB, InstsAfter, Inst);
 }
 
 void ShuffleBlockStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   // A deterministic alternative to SmallPtrSet with the same lookup
   // performance.
   std::map<size_t, Instruction *> AliveInsts;

llvm/lib/FuzzMutate/RandomIRBuilder.cpp

 using namespace llvm;
 using namespace fuzzerop;
 
 /// Return a vector of Blocks that dominates this block, excluding current
 /// block.
 static std::vector<BasicBlock *> getDominators(BasicBlock *BB) {
   std::vector<BasicBlock *> ret;
   DominatorTree DT(*BB->getParent());
-  DomTreeNode *Node = DT[BB]->getIDom();
+  DomTreeNode *Node = DT.getNode(BB);
+  // It's possible that an orphan block is not in the dom tree. In that case we
+  // just return nothing.
+  if (!Node)
+    return ret;
+  Node = Node->getIDom();
   while (Node && Node->getBlock()) {
     ret.push_back(Node->getBlock());
     // Get parent block.
     Node = Node->getIDom();
   }
   return ret;
 }
 
 /// Return a vector of Blocks that is dominated by this block, excluding current
 /// block
 static std::vector<BasicBlock *> getDominatees(BasicBlock *BB) {
   DominatorTree DT(*BB->getParent());
   std::vector<BasicBlock *> ret;
-  for (DomTreeNode *Child : DT[BB]->children())
+  DomTreeNode *Parent = DT.getNode(BB);
+  // It's possible that an orphan block is not in the dom tree. In that case we
+  // just return nothing.
+  if (!Parent)
+    return ret;
+  for (DomTreeNode *Child : Parent->children())
     ret.push_back(Child->getBlock());
   uint64_t Idx = 0;
   while (Idx < ret.size()) {
     DomTreeNode *Node = DT[ret[Idx]];
     Idx++;
     for (DomTreeNode *Child : Node->children())
       ret.push_back(Child->getBlock());
   }

llvm/unittests/FuzzMutate/RandomIRBuilderTest.cpp

for (int i = 0; i < 20; i++) {
     Value *Src = F.getArg(0);
     IB.connectToSink(BB, {I}, Src);
     Value *NewOperand = I->getOperand(0);
     Modified |= (OldOperand != NewOperand);
     ASSERT_FALSE(verifyModule(*M, &errs()));
   }
   ASSERT_FALSE(Modified);
 }
+
+TEST(RandomIRBuilderTest, SrcAndSinkWOrphanBlock) {
+  const char *Source = "\n\
+        define i1 @test(i1 %Bool, i32 %Int, i64 %Long) {   \n\
+        Entry:    \n\
+            %Eq0 = icmp eq i64 %Long, 0 \n\
+            br i1 %Eq0, label %True, label %False \n\
+        True: \n\
+            %Or = or i1 %Bool, %Eq0 \n\
+            ret i1 %Or \n\
+        False: \n\
+            %And = and i1 %Bool, %Eq0 \n\
+            ret i1 %And \n\
+        Orphan_1:  \n\
+            %NotBool = sub i1 1, %Bool \n\
+            ret i1 %NotBool \n\
+        Orphan_2:  \n\
+            %Le42 = icmp sle i32 %Int, 42 \n\
+            ret i1 %Le42 \n\
+        }";
+  LLVMContext Ctx;
+  std::mt19937 mt(Seed);
+  std::uniform_int_distribution<int> RandInt(INT_MIN, INT_MAX);
+  std::array<Type *, 3> IntTys(
+      {Type::getInt64Ty(Ctx), Type::getInt32Ty(Ctx), Type::getInt1Ty(Ctx)});
+  std::vector<Value *> Constants;
+  for (Type *IntTy : IntTys) {
+    for (size_t v : {1, 42}) {
+      Constants.push_back(ConstantInt::get(IntTy, v));
+    }
+  }
+  for (int i = 0; i < 10; i++) {
+    RandomIRBuilder IB(RandInt(mt), IntTys);
+    std::unique_ptr<Module> M = parseAssembly(Source, Ctx);
+    Function &F = *M->getFunction("test");
+    for (BasicBlock &BB : F) {
+      SmallVector<Instruction *, 4> Insts;
+      for (Instruction &I : BB) {
+        Insts.push_back(&I);
+      }
+      for (int j = 0; j < 10; j++) {
+        IB.findOrCreateSource(BB, Insts);
+      }
+      for (Value *V : Constants) {
+        IB.connectToSink(BB, Insts, V);
+      }
+    }
+  }
+}
 } // namespace

llvm/unittests/FuzzMutate/StrategiesTest.cpp

 
 TEST(InjectorIRStrategyTest, LargeInsertion) {
   StringRef Source = "";
   auto Mutator = createInjectorMutator();
   ASSERT_TRUE(Mutator);
   mutateAndVerifyModule(Source, Mutator, 100);
 }
 
+TEST(InjectorIRStrategyTest, InsertWMustTailCall) {
+  StringRef Source = "\n\
+        define i1 @recursive() {    \n\
+        Entry:     \n\
+            %Ret = musttail call i1 @recursive() \n\
+            ret i1 %Ret \n\
+        }";
+  auto Mutator = createInjectorMutator();
+  ASSERT_TRUE(Mutator);
+  mutateAndVerifyModule(Source, Mutator, 100);
+}
+
+TEST(InjectorIRStrategyTest, InsertWTailCall) {
+  StringRef Source = "\n\
+        define i1 @recursive() {    \n\
+        Entry:     \n\
+            %Ret = tail call i1 @recursive() \n\
+            ret i1 %Ret \n\
+        }";
+  auto Mutator = createInjectorMutator();
+  ASSERT_TRUE(Mutator);
+  mutateAndVerifyModule(Source, Mutator, 100);
+}
+
 TEST(InstDeleterIRStrategyTest, EmptyFunction) {
   // Test that we don't crash even if we can't remove from one of the functions.
 
   StringRef Source = ""
                      "define <8 x i32> @func1() {\n"
                      "ret <8 x i32> undef\n"
                      "}\n"
                      "\n"
StringRef Source = "\n\
           %C = and i1 %C2, %C3  \n\
           br i1 %C, label %BB0, label %Exit  \n\
         Exit:  \n\
           ret i32 %I  \n\
       }";
   mutateAndVerifyModule<SinkInstructionStrategy>(Source);
 }
 
+TEST(SinkInstructionStrategy, DoNotSinkTokenType) {
+  StringRef Source = "\n\
+      declare ptr @fake_personality_function() \n\
+      declare token @llvm.experimental.gc.statepoint.p0(i64 immarg %0, i32 immarg %1, ptr %2, i32 immarg %3, i32 immarg %4, ...) \n\
+      define void @test() gc \"statepoint-example\" personality ptr @fake_personality_function { \n\
+      Entry: \n\
+        %token1 = call token (i64, i32, ptr, i32, i32, ...) \
+          @llvm.experimental.gc.statepoint.p0(i64 0, i32 0, ptr elementtype(ptr addrspace(1) ()) undef, i32 0, i32 0, i32 0, i32 0) \n\
+        ret void \n\
+      }";
+  mutateAndVerifyModule<SinkInstructionStrategy>(Source);
+}
+
 static void VerifyBlockShuffle(StringRef Source) {
   LLVMContext Ctx;
   auto Mutator = createMutator<ShuffleBlockStrategy>();
   ASSERT_TRUE(Mutator);
 
   std::unique_ptr<Module> M = parseAssembly(Source.data(), Ctx);
   Function *F = &*M->begin();
   DenseMap<BasicBlock *, int> PreShuffleInstCnt;