[UBSan] Diagnose assumption violation (#104741)

This patch extends [D34590](https://reviews.llvm.org/D34590) to check
assumption violations.

---------

Co-authored-by: Vitaly Buka <vitalybuka@google.com>

clang/lib/CodeGen/CGBuiltin.cpp

@@ -1997,8 +1997,8 @@ struct CallObjCArcUse final : EHScopeStack::Cleanup {
 
 Value *CodeGenFunction::EmitCheckedArgForBuiltin(const Expr *E,
                                                  BuiltinCheckKind Kind) {
-  assert((Kind == BCK_CLZPassedZero || Kind == BCK_CTZPassedZero)
-          && "Unsupported builtin check kind");
+  assert((Kind == BCK_CLZPassedZero || Kind == BCK_CTZPassedZero) &&
+         "Unsupported builtin check kind");
 
   Value *ArgValue = EmitScalarExpr(E);
   if (!SanOpts.has(SanitizerKind::Builtin))
@@ -2015,6 +2015,21 @@ Value *CodeGenFunction::EmitCheckedArgForBuiltin(const Expr *E,
   return ArgValue;
 }
 
+Value *CodeGenFunction::EmitCheckedArgForAssume(const Expr *E) {
+  Value *ArgValue = EvaluateExprAsBool(E);
+  if (!SanOpts.has(SanitizerKind::Builtin))
+    return ArgValue;
+
+  SanitizerScope SanScope(this);
+  EmitCheck(
+      std::make_pair(ArgValue, SanitizerKind::Builtin),
+      SanitizerHandler::InvalidBuiltin,
+      {EmitCheckSourceLocation(E->getExprLoc()),
+       llvm::ConstantInt::get(Builder.getInt8Ty(), BCK_AssumePassedFalse)},
+      std::nullopt);
+  return ArgValue;
+}
+
 static Value *EmitAbs(CodeGenFunction &CGF, Value *ArgValue, bool HasNSW) {
   return CGF.Builder.CreateBinaryIntrinsic(
       Intrinsic::abs, ArgValue,
@@ -3428,7 +3443,7 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     if (E->getArg(0)->HasSideEffects(getContext()))
       return RValue::get(nullptr);
 
-    Value *ArgValue = EmitScalarExpr(E->getArg(0));
+    Value *ArgValue = EmitCheckedArgForAssume(E->getArg(0));
     Function *FnAssume = CGM.getIntrinsic(Intrinsic::assume);
     Builder.CreateCall(FnAssume, ArgValue);
     return RValue::get(nullptr);

clang/lib/CodeGen/CGStmt.cpp

@@ -754,7 +754,7 @@ void CodeGenFunction::EmitAttributedStmt(const AttributedStmt &S) {
       const Expr *Assumption = cast<CXXAssumeAttr>(A)->getAssumption();
       if (getLangOpts().CXXAssumptions && Builder.GetInsertBlock() &&
           !Assumption->HasSideEffects(getContext())) {
-        llvm::Value *AssumptionVal = EvaluateExprAsBool(Assumption);
+        llvm::Value *AssumptionVal = EmitCheckedArgForAssume(Assumption);
         Builder.CreateAssumption(AssumptionVal);
       }
     } break;

clang/lib/CodeGen/CodeGenFunction.h

@@ -5084,12 +5084,17 @@ public:
   enum BuiltinCheckKind {
     BCK_CTZPassedZero,
     BCK_CLZPassedZero,
+    BCK_AssumePassedFalse,
   };
 
   /// Emits an argument for a call to a builtin. If the builtin sanitizer is
   /// enabled, a runtime check specified by \p Kind is also emitted.
   llvm::Value *EmitCheckedArgForBuiltin(const Expr *E, BuiltinCheckKind Kind);
 
+  /// Emits an argument for a call to a `__builtin_assume`. If the builtin
+  /// sanitizer is enabled, a runtime check is also emitted.
+  llvm::Value *EmitCheckedArgForAssume(const Expr *E);
+
   /// Emit a description of a type in a format suitable for passing to
   /// a runtime sanitizer handler.
   llvm::Constant *EmitCheckTypeDescriptor(QualType T);

clang/test/CodeGen/ubsan-builtin-checks.c

@@ -51,3 +51,20 @@ void check_clz(int n) {
   // CHECK: call void @__ubsan_handle_invalid_builtin
   __builtin_clzg((unsigned int)n);
 }
+
+// CHECK: define{{.*}} void @check_assume
+void check_assume(int n) {
+  // CHECK: [[TOBOOL:%.*]] = icmp ne i32 [[N:%.*]], 0
+  // CHECK-NEXT: br i1 [[TOBOOL]]
+  //
+  // Handler block:
+  // CHECK: call void @__ubsan_handle_invalid_builtin
+  // CHECK-NEXT: unreachable
+  //
+  // Continuation block:
+  // CHECK: call void @llvm.assume(i1 [[TOBOOL]])
+  __builtin_assume(n);
+
+  // CHECK: call void @__ubsan_handle_invalid_builtin
+  __attribute__((assume(n)));
+}

compiler-rt/lib/ubsan/ubsan_handlers.cpp

@@ -633,9 +633,12 @@ static void handleInvalidBuiltin(InvalidBuiltinData *Data, ReportOptions Opts) {
 
   ScopedReport R(Opts, Loc, ET);
 
-  Diag(Loc, DL_Error, ET,
-       "passing zero to __builtin_%0(), which is not a valid argument")
-      << ((Data->Kind == BCK_CTZPassedZero) ? "ctz" : "clz");
+  if (Data->Kind == BCK_AssumePassedFalse)
+    Diag(Loc, DL_Error, ET, "assumption is violated during execution");
+  else
+    Diag(Loc, DL_Error, ET,
+         "passing zero to __builtin_%0(), which is not a valid argument")
+        << ((Data->Kind == BCK_CTZPassedZero) ? "ctz" : "clz");
 }
 
 void __ubsan::__ubsan_handle_invalid_builtin(InvalidBuiltinData *Data) {

compiler-rt/lib/ubsan/ubsan_handlers.h

@@ -159,6 +159,7 @@ RECOVERABLE(implicit_conversion, ImplicitConversionData *Data, ValueHandle Src,
 enum BuiltinCheckKind : unsigned char {
   BCK_CTZPassedZero,
   BCK_CLZPassedZero,
+  BCK_AssumePassedFalse,
 };
 
 struct InvalidBuiltinData {

compiler-rt/test/ubsan/TestCases/Misc/builtins.cpp

@@ -1,8 +1,8 @@
 // REQUIRES: target={{x86_64.*}}
 //
-// RUN: %clangxx -fsanitize=builtin -w %s -O3 -o %t
+// RUN: %clangxx -fsanitize=builtin -fno-inline -w %s -O3 -o %t
 // RUN: %run %t 2>&1 | FileCheck %s --check-prefix=RECOVER
-// RUN: %clangxx -fsanitize=builtin -fno-sanitize-recover=builtin -w %s -O3 -o %t.abort
+// RUN: %clangxx -fsanitize=builtin -fno-inline -fno-sanitize-recover=builtin -w %s -O3 -o %t.abort
 // RUN: not %run %t.abort 2>&1 | FileCheck %s --check-prefix=ABORT
 
 void check_ctz(int n) {
@@ -28,8 +28,20 @@ void check_clz(int n) {
   __builtin_clzll(n);
 }
 
+void check_assume(int n) {
+  // RECOVER: builtins.cpp:[[@LINE+1]]:20: runtime error: assumption is violated during execution
+  __builtin_assume(n);
+}
+
+void check_assume_attr(int n) {
+  // RECOVER: builtins.cpp:[[@LINE+1]]:25: runtime error: assumption is violated during execution
+  __attribute__((assume(n)));
+}
+
 int main() {
   check_ctz(0);
   check_clz(0);
+  check_assume(0);
+  check_assume_attr(0);
   return 0;
 }