From 989067925127cb331f3959d62c5ecbd63398a972 Mon Sep 17 00:00:00 2001
From: Sirui Mu <msrlancern@gmail.com>
Date: Wed, 4 Dec 2024 21:55:58 +0800
Subject: [PATCH] [CIR][CIRGen] Add CIRGen support for assume statement

This patch adds CIRGen support for C++23 [[assume(expr)]] statement.
---
 clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp   | 11 ++++++++++-
 clang/lib/CIR/CodeGen/CIRGenFunction.h    |  4 ++++
 clang/lib/CIR/CodeGen/CIRGenStmt.cpp      | 10 ++++++++++
 clang/test/CIR/CodeGen/builtin-assume.cpp | 20 ++++++++++++++++++--
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
index 56068ef222e1..6d16b8c2acc1 100644
--- a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
@@ -1110,7 +1110,7 @@ RValue CIRGenFunction::emitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     if (E->getArg(0)->HasSideEffects(getContext()))
       return RValue::get(nullptr);
 
-    mlir::Value argValue = emitScalarExpr(E->getArg(0));
+    mlir::Value argValue = emitCheckedArgForAssume(E->getArg(0));
     builder.create<cir::AssumeOp>(getLoc(E->getExprLoc()), argValue);
     return RValue::get(nullptr);
   }
@@ -2478,6 +2478,15 @@ mlir::Value CIRGenFunction::emitCheckedArgForBuiltin(const Expr *E,
   llvm_unreachable("NYI");
 }
 
+mlir::Value CIRGenFunction::emitCheckedArgForAssume(const Expr *E) {
+  mlir::Value argValue = evaluateExprAsBool(E);
+  if (!SanOpts.has(SanitizerKind::Builtin))
+    return argValue;
+
+  assert(!MissingFeatures::sanitizerBuiltin());
+  llvm_unreachable("NYI");
+}
+
 static mlir::Value emitTargetArchBuiltinExpr(CIRGenFunction *CGF,
                                              unsigned BuiltinID,
                                              const CallExpr *E,
diff --git a/clang/lib/CIR/CodeGen/CIRGenFunction.h b/clang/lib/CIR/CodeGen/CIRGenFunction.h
index 4172ac1e208d..b6b949e47bf2 100644
--- a/clang/lib/CIR/CodeGen/CIRGenFunction.h
+++ b/clang/lib/CIR/CodeGen/CIRGenFunction.h
@@ -1782,6 +1782,10 @@ class CIRGenFunction : public CIRGenTypeCache {
   /// enabled, a runtime check specified by \p Kind is also emitted.
   mlir::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.
+  mlir::Value emitCheckedArgForAssume(const Expr *E);
+
   /// returns true if aggregate type has a volatile member.
   /// TODO(cir): this could be a common AST helper between LLVM / CIR.
   bool hasVolatileMember(QualType T) {
diff --git a/clang/lib/CIR/CodeGen/CIRGenStmt.cpp b/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
index 410c25891bbb..8708eeecb7e5 100644
--- a/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenStmt.cpp
@@ -346,6 +346,16 @@ CIRGenFunction::emitAttributedStmt(const AttributedStmt &S) {
     case attr::AlwaysInline:
     case attr::MustTail:
       llvm_unreachable("NIY attributes");
+    case attr::CXXAssume: {
+      const Expr *assumption = cast<CXXAssumeAttr>(A)->getAssumption();
+      if (getLangOpts().CXXAssumptions && builder.getInsertionBlock() &&
+          !assumption->HasSideEffects(getContext())) {
+        mlir::Value assumptionValue = emitCheckedArgForAssume(assumption);
+        builder.create<cir::AssumeOp>(getLoc(S.getSourceRange()),
+                                      assumptionValue);
+      }
+      break;
+    }
     default:
       break;
     }
diff --git a/clang/test/CIR/CodeGen/builtin-assume.cpp b/clang/test/CIR/CodeGen/builtin-assume.cpp
index 88e8ad11565e..8d7448a2724d 100644
--- a/clang/test/CIR/CodeGen/builtin-assume.cpp
+++ b/clang/test/CIR/CodeGen/builtin-assume.cpp
@@ -1,6 +1,6 @@
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -emit-cir %s -o %t.cir
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++23 -fclangir -emit-cir %s -o %t.cir
 // RUN: FileCheck %s --check-prefix=CIR --input-file=%t.cir
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -emit-llvm %s -o %t.ll
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -std=c++23 -fclangir -emit-llvm %s -o %t.ll
 // RUN: FileCheck %s --check-prefix=LLVM --input-file=%t.ll
 
 int test_assume(int x) {
@@ -19,6 +19,22 @@ int test_assume(int x) {
 //      LLVM: %[[#cond:]] = trunc i8 %{{.+}} to i1
 // LLVM-NEXT: call void @llvm.assume(i1 %[[#cond]])
 
+int test_assume_attr(int x) {
+  [[assume(x > 0)]];
+  return x;
+}
+
+//      CIR: cir.func @_Z16test_assume_attri
+//      CIR:   %[[#x:]] = cir.load %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR-NEXT:   %[[#zero:]] = cir.const #cir.int<0> : !s32i
+// CIR-NEXT:   %[[#cond:]] = cir.cmp(gt, %[[#x]], %[[#zero]]) : !s32i, !cir.bool
+// CIR-NEXT:   cir.assume %[[#cond]] : !cir.bool
+//      CIR: }
+
+//      LLVM: @_Z16test_assume_attri
+//      LLVM: %[[#cond:]] = trunc i8 %{{.+}} to i1
+// LLVM-NEXT: call void @llvm.assume(i1 %[[#cond]])
+
 int test_assume_aligned(int *ptr) {
   int *aligned = (int *)__builtin_assume_aligned(ptr, 8);
   return *aligned;