[FuzzMutate] Prevent UB caused by parameter ABI attributes

llvm/lib/FuzzMutate/IRMutator.cpp

@@ -356,6 +356,62 @@ static uint64_t getUniqueCaseValue(SmallSet<uint64_t, 4> &CasesTaken,
   return tmp;
 }
 
+/// Determines whether a function is unsupported by the current mutator's
+/// implementation. The function returns true if any of the following criteria
+/// are met:
+///   * The function accepts metadata or token types as arguments.
+///   * The function has ABI attributes that could cause UB.
+///   * The function uses a non-callable CC that may result in UB.
+static bool isUnsupportedFunction(Function *F) {
+  // 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 (IsUnsupportedTy(F->getReturnType()) ||
+      any_of(F->getFunctionType()->params(), IsUnsupportedTy)) {
+    return true;
+  }
+
+  // ABI attributes must be specified both at the function
+  // declaration/definition and call-site, otherwise the
+  // behavior may be undefined.
+  // We don't call those functions for now to prevent UB from happening.
+  auto IsABIAttribute = [](AttributeSet A) {
+    static const Attribute::AttrKind ABIAttrs[] = {
+        Attribute::StructRet,      Attribute::ByVal,
+        Attribute::InAlloca,       Attribute::InReg,
+        Attribute::StackAlignment, Attribute::SwiftSelf,
+        Attribute::SwiftAsync,     Attribute::SwiftError,
+        Attribute::Preallocated,   Attribute::ByRef,
+        Attribute::ZExt,           Attribute::SExt};
+
+    return std::any_of(
+        std::begin(ABIAttrs), std::end(ABIAttrs),
+        [&](Attribute::AttrKind kind) { return A.hasAttribute(kind); });
+  };
+
+  auto FuncAttrs = F->getAttributes();
+  if (IsABIAttribute(FuncAttrs.getRetAttrs())) {
+    return true;
+  }
+  for (size_t i = 0; i < F->arg_size(); i++) {
+    if (IsABIAttribute(FuncAttrs.getParamAttrs(i))) {
+      return true;
+    }
+  }
+
+  // If it is not satisfied, the IR will be invalid.
+  if (!isCallableCC(F->getCallingConv())) {
+    return true;
+  }
+
+  return false;
+}
+
 void InsertFunctionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   Module *M = BB.getParent()->getParent();
   // If nullptr is selected, we will create a new function declaration.
@@ -366,16 +422,8 @@ void InsertFunctionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
 
   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) ||
-      !isCallableCC(F->getCallingConv())) {
+
+  if (!F || isUnsupportedFunction(F)) {
     F = IB.createFunctionDeclaration(*M);
   }