From 161c674ef051b93978b96f6df70f03094144c537 Mon Sep 17 00:00:00 2001
From: Camille GILLOT <gillot.camille@gmail.com>
Date: Sat, 20 Jan 2024 19:01:57 +0000
Subject: [PATCH] Add Assume custom MIR.

---
 .../src/build/custom/parse/instruction.rs     |  4 ++
 compiler/rustc_span/src/symbol.rs             |  1 +
 library/core/src/intrinsics/mir.rs            |  2 +
 .../assume.assume_constant.built.after.mir    | 10 +++++
 .../assume.assume_local.built.after.mir       | 10 +++++
 .../assume.assume_place.built.after.mir       | 10 +++++
 tests/mir-opt/building/custom/assume.rs       | 44 +++++++++++++++++++
 7 files changed, 81 insertions(+)
 create mode 100644 tests/mir-opt/building/custom/assume.assume_constant.built.after.mir
 create mode 100644 tests/mir-opt/building/custom/assume.assume_local.built.after.mir
 create mode 100644 tests/mir-opt/building/custom/assume.assume_place.built.after.mir
 create mode 100644 tests/mir-opt/building/custom/assume.rs

diff --git a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
index 5428333a116..c669d3fd623 100644
--- a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
+++ b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
@@ -20,6 +20,10 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             @call(mir_storage_dead, args) => {
                 Ok(StatementKind::StorageDead(self.parse_local(args[0])?))
             },
+            @call(mir_assume, args) => {
+                let op = self.parse_operand(args[0])?;
+                Ok(StatementKind::Intrinsic(Box::new(NonDivergingIntrinsic::Assume(op))))
+            },
             @call(mir_deinit, args) => {
                 Ok(StatementKind::Deinit(Box::new(self.parse_place(args[0])?)))
             },
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 7b0138d50ba..6d10fdf49d7 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -1027,6 +1027,7 @@ symbols! {
         minnumf32,
         minnumf64,
         mips_target_feature,
+        mir_assume,
         mir_basic_block,
         mir_call,
         mir_cast_transmute,
diff --git a/library/core/src/intrinsics/mir.rs b/library/core/src/intrinsics/mir.rs
index 334e32b26b1..d348e31609d 100644
--- a/library/core/src/intrinsics/mir.rs
+++ b/library/core/src/intrinsics/mir.rs
@@ -357,6 +357,8 @@ define!("mir_unwind_resume",
 
 define!("mir_storage_live", fn StorageLive<T>(local: T));
 define!("mir_storage_dead", fn StorageDead<T>(local: T));
+#[cfg(not(bootstrap))]
+define!("mir_assume", fn Assume(operand: bool));
 define!("mir_deinit", fn Deinit<T>(place: T));
 define!("mir_checked", fn Checked<T>(binop: T) -> (T, bool));
 define!("mir_len", fn Len<T>(place: T) -> usize);
diff --git a/tests/mir-opt/building/custom/assume.assume_constant.built.after.mir b/tests/mir-opt/building/custom/assume.assume_constant.built.after.mir
new file mode 100644
index 00000000000..8e70d0a1e9b
--- /dev/null
+++ b/tests/mir-opt/building/custom/assume.assume_constant.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `assume_constant` after built
+
+fn assume_constant() -> () {
+    let mut _0: ();
+
+    bb0: {
+        assume(const true);
+        return;
+    }
+}
diff --git a/tests/mir-opt/building/custom/assume.assume_local.built.after.mir b/tests/mir-opt/building/custom/assume.assume_local.built.after.mir
new file mode 100644
index 00000000000..7ea1fcd30c2
--- /dev/null
+++ b/tests/mir-opt/building/custom/assume.assume_local.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `assume_local` after built
+
+fn assume_local(_1: bool) -> () {
+    let mut _0: ();
+
+    bb0: {
+        assume(_1);
+        return;
+    }
+}
diff --git a/tests/mir-opt/building/custom/assume.assume_place.built.after.mir b/tests/mir-opt/building/custom/assume.assume_place.built.after.mir
new file mode 100644
index 00000000000..ce914618d3d
--- /dev/null
+++ b/tests/mir-opt/building/custom/assume.assume_place.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `assume_place` after built
+
+fn assume_place(_1: (bool, u8)) -> () {
+    let mut _0: ();
+
+    bb0: {
+        assume((_1.0: bool));
+        return;
+    }
+}
diff --git a/tests/mir-opt/building/custom/assume.rs b/tests/mir-opt/building/custom/assume.rs
new file mode 100644
index 00000000000..a477e12f0e0
--- /dev/null
+++ b/tests/mir-opt/building/custom/assume.rs
@@ -0,0 +1,44 @@
+// skip-filecheck
+#![feature(custom_mir, core_intrinsics)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+
+// EMIT_MIR assume.assume_local.built.after.mir
+#[custom_mir(dialect = "built")]
+fn assume_local(x: bool) {
+    mir!(
+        {
+            Assume(x);
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR assume.assume_place.built.after.mir
+#[custom_mir(dialect = "built")]
+fn assume_place(p: (bool, u8)) {
+    mir!(
+        {
+            Assume(p.0);
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR assume.assume_constant.built.after.mir
+#[custom_mir(dialect = "built")]
+fn assume_constant() {
+    mir!(
+        {
+            Assume(true);
+            Return()
+        }
+    )
+}
+
+fn main() {
+    assume_local(true);
+    assume_place((true, 50));
+    assume_constant();
+}