From e9a67c74720911c531290c8ebc330e0fb5dc34d9 Mon Sep 17 00:00:00 2001
From: dianqk <dianqk@dianqk.net>
Date: Fri, 2 Jan 2026 22:15:40 +0800
Subject: [PATCH] Propagates assume

---
 .../rustc_mir_transform/src/ssa_range_prop.rs | 16 ++++++
 ...a_range.on_assume.SsaRangePropagation.diff | 56 +++++++++++++++++++
 tests/mir-opt/range/ssa_range.rs              | 10 ++++
 3 files changed, 82 insertions(+)
 create mode 100644 tests/mir-opt/range/ssa_range.on_assume.SsaRangePropagation.diff

diff --git a/compiler/rustc_mir_transform/src/ssa_range_prop.rs b/compiler/rustc_mir_transform/src/ssa_range_prop.rs
index 47e17de0560b..7a8be8efdfd1 100644
--- a/compiler/rustc_mir_transform/src/ssa_range_prop.rs
+++ b/compiler/rustc_mir_transform/src/ssa_range_prop.rs
@@ -140,6 +140,22 @@ impl<'tcx> MutVisitor<'tcx> for RangeSet<'tcx, '_, '_> {
         };
     }
 
+    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
+        self.super_statement(statement, location);
+        match &statement.kind {
+            StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(operand)) => {
+                if let Some(place) = operand.place()
+                    && self.is_ssa(place)
+                {
+                    let successor = location.successor_within_block();
+                    let range = WrappingRange { start: 1, end: 1 };
+                    self.insert_range(place, successor, range);
+                }
+            }
+            _ => {}
+        }
+    }
+
     fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
         self.super_terminator(terminator, location);
         match &terminator.kind {
diff --git a/tests/mir-opt/range/ssa_range.on_assume.SsaRangePropagation.diff b/tests/mir-opt/range/ssa_range.on_assume.SsaRangePropagation.diff
new file mode 100644
index 000000000000..3be6f083604d
--- /dev/null
+++ b/tests/mir-opt/range/ssa_range.on_assume.SsaRangePropagation.diff
@@ -0,0 +1,56 @@
+- // MIR for `on_assume` before SsaRangePropagation
++ // MIR for `on_assume` after SsaRangePropagation
+  
+  fn on_assume(_1: usize, _2: &[u8]) -> u8 {
+      debug i => _1;
+      debug v => _2;
+      let mut _0: u8;
+      let _3: ();
+      let _4: ();
+      let mut _5: bool;
+      let mut _6: usize;
+      let mut _7: usize;
+      let mut _8: &[u8];
+      let _9: usize;
+      let mut _10: usize;
+      let mut _11: bool;
+      scope 1 (inlined core::slice::<impl [u8]>::len) {
+          scope 2 (inlined std::ptr::metadata::<[u8]>) {
+          }
+      }
+  
+      bb0: {
+          StorageLive(_3);
+          StorageLive(_4);
+          nop;
+          StorageLive(_6);
+          _6 = copy _1;
+          nop;
+          StorageLive(_8);
+          _8 = &(*_2);
+          _7 = PtrMetadata(copy _2);
+          StorageDead(_8);
+          _5 = Lt(copy _1, copy _7);
+          nop;
+          StorageDead(_6);
+          assume(copy _5);
+          nop;
+          StorageDead(_4);
+          _3 = const ();
+          StorageDead(_3);
+          StorageLive(_9);
+          _9 = copy _1;
+          _10 = copy _7;
+-         _11 = copy _5;
+-         assert(copy _5, "index out of bounds: the length is {} but the index is {}", copy _7, copy _1) -> [success: bb1, unwind unreachable];
++         _11 = const true;
++         assert(const true, "index out of bounds: the length is {} but the index is {}", copy _7, copy _1) -> [success: bb1, unwind unreachable];
+      }
+  
+      bb1: {
+          _0 = copy (*_2)[_1];
+          StorageDead(_9);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/range/ssa_range.rs b/tests/mir-opt/range/ssa_range.rs
index 964d9b97cf92..c8440a10a408 100644
--- a/tests/mir-opt/range/ssa_range.rs
+++ b/tests/mir-opt/range/ssa_range.rs
@@ -20,6 +20,16 @@ pub fn on_assert(i: usize, v: &[u8]) -> u8 {
     v[i]
 }
 
+// EMIT_MIR ssa_range.on_assume.SsaRangePropagation.diff
+pub fn on_assume(i: usize, v: &[u8]) -> u8 {
+    // CHECK-LABEL: fn on_assume(
+    // CHECK: assert(const true
+    unsafe {
+        std::intrinsics::assume(i < v.len());
+    }
+    v[i]
+}
+
 // EMIT_MIR ssa_range.on_match.SsaRangePropagation.diff
 pub fn on_match(i: u8) -> u8 {
     // CHECK-LABEL: fn on_match(