Lower the assume intrinsic to a MIR statement

oli-obk · oli-obk · commit 3f0764512024 · 2022-09-06T14:18:32.000Z
diff --git a/compiler/rustc_borrowck/src/dataflow.rs b/compiler/rustc_borrowck/src/dataflow.rs
@@ -392,6 +392,7 @@ impl<'tcx> rustc_mir_dataflow::GenKillAnalysis<'tcx> for Borrows<'_, 'tcx> {
             | mir::StatementKind::AscribeUserType(..)
             | mir::StatementKind::Coverage(..)
             | mir::StatementKind::CopyNonOverlapping(..)
+            | mir::StatementKind::Assume(..)
             | mir::StatementKind::Nop => {}
         }
     }
diff --git a/compiler/rustc_borrowck/src/invalidation.rs b/compiler/rustc_borrowck/src/invalidation.rs
@@ -72,14 +72,14 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                 self.consume_operand(location, dst);
                 self.consume_operand(location, count);
             }
-            StatementKind::Nop
+            // Only relevant for mir typeck
+            StatementKind::AscribeUserType(..)
+            // Doesn't have any language semantics
             | StatementKind::Coverage(..)
-            | StatementKind::AscribeUserType(..)
-            | StatementKind::Retag { .. }
-            | StatementKind::StorageLive(..) => {
-                // `Nop`, `AscribeUserType`, `Retag`, and `StorageLive` are irrelevant
-                // to borrow check.
-            }
+            // Takes a `bool` argument, and has no return value, thus being irrelevant for borrowck
+            | StatementKind::Assume(..)
+            // Does not actually affect borrowck
+            | StatementKind::StorageLive(..) => {}
             StatementKind::StorageDead(local) => {
                 self.access_place(
                     location,
@@ -88,7 +88,10 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                     LocalMutationIsAllowed::Yes,
                 );
             }
-            StatementKind::Deinit(..) | StatementKind::SetDiscriminant { .. } => {
+            StatementKind::Nop
+            | StatementKind::Retag { .. }
+            | StatementKind::Deinit(..)
+            | StatementKind::SetDiscriminant { .. } => {
                 bug!("Statement not allowed in this MIR phase")
             }
         }
diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs
@@ -599,14 +599,14 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                     "Unexpected CopyNonOverlapping, should only appear after lower_intrinsics",
                 )
             }
-            StatementKind::Nop
+            // Only relevant for mir typeck
+            StatementKind::AscribeUserType(..)
+            // Doesn't have any language semantics
             | StatementKind::Coverage(..)
-            | StatementKind::AscribeUserType(..)
-            | StatementKind::Retag { .. }
-            | StatementKind::StorageLive(..) => {
-                // `Nop`, `AscribeUserType`, `Retag`, and `StorageLive` are irrelevant
-                // to borrow check.
-            }
+            // Takes a `bool` argument, and has no return value, thus being irrelevant for borrowck
+            | StatementKind::Assume(..)
+            // Does not actually affect borrowck
+            | StatementKind::StorageLive(..) => {}
             StatementKind::StorageDead(local) => {
                 self.access_place(
                     location,
@@ -616,7 +616,10 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                     flow_state,
                 );
             }
-            StatementKind::Deinit(..) | StatementKind::SetDiscriminant { .. } => {
+            StatementKind::Nop
+            | StatementKind::Retag { .. }
+            | StatementKind::Deinit(..)
+            | StatementKind::SetDiscriminant { .. } => {
                 bug!("Statement not allowed in this MIR phase")
             }
         }
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -1309,6 +1309,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 "Unexpected StatementKind::CopyNonOverlapping, should only appear after lowering_intrinsics",
             ),
             StatementKind::FakeRead(..)
+            | StatementKind::Assume(..)
             | StatementKind::StorageLive(..)
             | StatementKind::StorageDead(..)
             | StatementKind::Retag { .. }
diff --git a/compiler/rustc_codegen_cranelift/src/base.rs b/compiler/rustc_codegen_cranelift/src/base.rs
@@ -791,6 +791,8 @@ fn codegen_stmt<'tcx>(
         | StatementKind::Nop
         | StatementKind::FakeRead(..)
         | StatementKind::Retag { .. }
+        // We ignore `assume` intrinsics, they are only useful for optimizations
+        | StatementKind::Assume(..)
         | StatementKind::AscribeUserType(..) => {}
 
         StatementKind::Coverage { .. } => fx.tcx.sess.fatal("-Zcoverage is unimplemented"),
diff --git a/compiler/rustc_codegen_cranelift/src/constant.rs b/compiler/rustc_codegen_cranelift/src/constant.rs
@@ -540,6 +540,7 @@ pub(crate) fn mir_operand_get_const_val<'tcx>(
                             return None;
                         } // conservative handling
                         StatementKind::Assign(_)
+                        | StatementKind::Assume(_)
                         | StatementKind::FakeRead(_)
                         | StatementKind::SetDiscriminant { .. }
                         | StatementKind::Deinit(_)
diff --git a/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs b/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs
@@ -357,9 +357,6 @@ fn codegen_regular_intrinsic_call<'tcx>(
     let usize_layout = fx.layout_of(fx.tcx.types.usize);
 
     match intrinsic {
-        sym::assume => {
-            intrinsic_args!(fx, args => (_a); intrinsic);
-        }
         sym::likely | sym::unlikely => {
             intrinsic_args!(fx, args => (a); intrinsic);
 
diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs
@@ -77,10 +77,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let result = PlaceRef::new_sized(llresult, fn_abi.ret.layout);
 
         let llval = match name {
-            sym::assume => {
-                bx.assume(args[0].immediate());
-                return;
-            }
             sym::abort => {
                 bx.abort();
                 return;
diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs
@@ -93,6 +93,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty());
                 bx
             }
+            mir::StatementKind::Assume(box ref op) => {
+                let op_val = self.codegen_operand(&mut bx, op);
+                bx.assume(op_val.immediate());
+                bx
+            }
             mir::StatementKind::FakeRead(..)
             | mir::StatementKind::Retag { .. }
             | mir::StatementKind::AscribeUserType(..)
diff --git a/compiler/rustc_const_eval/src/interpret/intrinsics.rs b/compiler/rustc_const_eval/src/interpret/intrinsics.rs
@@ -506,12 +506,6 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 // These just return their argument
                 self.copy_op(&args[0], dest, /*allow_transmute*/ false)?;
             }
-            sym::assume => {
-                let cond = self.read_scalar(&args[0])?.to_bool()?;
-                if !cond {
-                    throw_ub_format!("`assume` intrinsic called with `false`");
-                }
-            }
             sym::raw_eq => {
                 let result = self.raw_eq_intrinsic(&args[0], &args[1])?;
                 self.write_scalar(result, dest)?;
diff --git a/compiler/rustc_const_eval/src/interpret/step.rs b/compiler/rustc_const_eval/src/interpret/step.rs
@@ -122,6 +122,15 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 self.copy_intrinsic(&src, &dst, &count, /* nonoverlapping */ true)?;
             }
 
+            // Call Assume
+            Assume(box op) => {
+                let op = self.eval_operand(op, None)?;
+                let cond = self.read_scalar(&op)?.to_bool()?;
+                if !cond {
+                    throw_ub_format!("`assume` called with `false`");
+                }
+            }
+
             // Statements we do not track.
             AscribeUserType(..) => {}
 
diff --git a/compiler/rustc_const_eval/src/transform/check_consts/check.rs b/compiler/rustc_const_eval/src/transform/check_consts/check.rs
@@ -679,6 +679,7 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::Nop => {}
         }
     }
diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -636,6 +636,15 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     );
                 }
             }
+            StatementKind::Assume(box ref op) => {
+                let ty = op.ty(&self.body.local_decls, self.tcx);
+                if !ty.is_bool() {
+                    self.fail(
+                        location,
+                        format!("`assume` argument must be `bool`, but got: `{}`", ty),
+                    );
+                }
+            }
             StatementKind::CopyNonOverlapping(box rustc_middle::mir::CopyNonOverlapping {
                 ref src,
                 ref dst,
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
@@ -1377,6 +1377,7 @@ impl Debug for Statement<'_> {
             }) => {
                 write!(fmt, "copy_nonoverlapping(src={:?}, dst={:?}, count={:?})", src, dst, count)
             }
+            Assume(box ref cond) => write!(fmt, "assume({:?})", cond),
             Nop => write!(fmt, "nop"),
         }
     }
diff --git a/compiler/rustc_middle/src/mir/spanview.rs b/compiler/rustc_middle/src/mir/spanview.rs
@@ -250,6 +250,7 @@ pub fn statement_kind_name(statement: &Statement<'_>) -> &'static str {
         AscribeUserType(..) => "AscribeUserType",
         Coverage(..) => "Coverage",
         CopyNonOverlapping(..) => "CopyNonOverlapping",
+        Assume(..) => "Assume",
         Nop => "Nop",
     }
 }
diff --git a/compiler/rustc_middle/src/mir/syntax.rs b/compiler/rustc_middle/src/mir/syntax.rs
@@ -342,6 +342,14 @@ pub enum StatementKind<'tcx> {
     /// I vaguely remember Ralf saying somewhere that he thought it should not be.
     CopyNonOverlapping(Box<CopyNonOverlapping<'tcx>>),
 
+    /// Denotes a call to the intrinsic function `assume`.
+    ///
+    /// The operand must be a boolean. Optimizers may use the value of the boolean to backtrack its
+    /// computation to infer information about other variables. So if the boolean came from a
+    /// `x < y` operation, subsequent operations on `x` and `y` could elide various bound checks.
+    /// If the argument is `false`, this operation is equivalent to `TerminatorKind::Unreachable`.
+    Assume(Box<Operand<'tcx>>),
+
     /// No-op. Useful for deleting instructions without affecting statement indices.
     Nop,
 }
diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
@@ -434,6 +434,9 @@ macro_rules! make_mir_visitor {
                       self.visit_operand(dst, location);
                       self.visit_operand(count, location)
                     }
+                    StatementKind::Assume(box ref $($mutability)? val) => {
+                        self.visit_operand(val, location)
+                    }
                     StatementKind::Nop => {}
                 }
             }
diff --git a/compiler/rustc_mir_dataflow/src/impls/liveness.rs b/compiler/rustc_mir_dataflow/src/impls/liveness.rs
@@ -271,6 +271,7 @@ impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'a> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::Nop => None,
         };
         if let Some(destination) = destination {
diff --git a/compiler/rustc_mir_dataflow/src/impls/storage_liveness.rs b/compiler/rustc_mir_dataflow/src/impls/storage_liveness.rs
@@ -143,6 +143,7 @@ impl<'mir, 'tcx> crate::GenKillAnalysis<'tcx> for MaybeRequiresStorage<'mir, 'tc
             | StatementKind::Nop
             | StatementKind::Retag(..)
             | StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::StorageLive(..) => {}
         }
     }
diff --git a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
@@ -331,6 +331,7 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::Nop => {}
         }
     }
diff --git a/compiler/rustc_mir_transform/src/check_unsafety.rs b/compiler/rustc_mir_transform/src/check_unsafety.rs
@@ -105,6 +105,9 @@ impl<'tcx> Visitor<'tcx> for UnsafetyChecker<'_, 'tcx> {
                 // safe (at least as emitted during MIR construction)
             }
 
+            // Move to above list once mir construction uses it.
+            StatementKind::Assume(..) => unreachable!(),
+
             StatementKind::CopyNonOverlapping(..) => unreachable!(),
         }
         self.super_statement(statement, location);
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -826,6 +826,7 @@ pub(super) fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span>
         // Retain spans from all other statements
         StatementKind::FakeRead(box (_, _)) // Not including `ForGuardBinding`
         | StatementKind::CopyNonOverlapping(..)
+        | StatementKind::Assume(..)
         | StatementKind::Assign(_)
         | StatementKind::SetDiscriminant { .. }
         | StatementKind::Deinit(..)
diff --git a/compiler/rustc_mir_transform/src/dead_store_elimination.rs b/compiler/rustc_mir_transform/src/dead_store_elimination.rs
@@ -53,6 +53,7 @@ pub fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, borrowed: &BitS
                 | StatementKind::StorageDead(_)
                 | StatementKind::Coverage(_)
                 | StatementKind::CopyNonOverlapping(_)
+                | StatementKind::Assume(_)
                 | StatementKind::Nop => (),
 
                 StatementKind::FakeRead(_) | StatementKind::AscribeUserType(_, _) => {
diff --git a/compiler/rustc_mir_transform/src/dest_prop.rs b/compiler/rustc_mir_transform/src/dest_prop.rs
@@ -538,6 +538,7 @@ impl<'a> Conflicts<'a> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::Nop => {}
         }
     }
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
@@ -1453,6 +1453,7 @@ impl<'tcx> Visitor<'tcx> for EnsureGeneratorFieldAssignmentsNeverAlias<'_> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::Nop => {}
         }
     }
diff --git a/compiler/rustc_mir_transform/src/lower_intrinsics.rs b/compiler/rustc_mir_transform/src/lower_intrinsics.rs
@@ -62,6 +62,17 @@ impl<'tcx> MirPass<'tcx> for LowerIntrinsics {
                         drop(args);
                         terminator.kind = TerminatorKind::Goto { target };
                     }
+                    sym::assume => {
+                        let target = target.unwrap();
+                        let mut args = args.drain(..);
+                        block.statements.push(Statement {
+                            source_info: terminator.source_info,
+                            kind: StatementKind::Assume(Box::new(args.next().unwrap())),
+                        });
+                        assert_eq!(args.next(), None, "Extra argument for assume intrinsic");
+                        drop(args);
+                        terminator.kind = TerminatorKind::Goto { target };
+                    }
                     sym::wrapping_add | sym::wrapping_sub | sym::wrapping_mul => {
                         if let Some(target) = *target {
                             let lhs;
diff --git a/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs b/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs
@@ -52,6 +52,7 @@ impl RemoveNoopLandingPads {
                 | StatementKind::SetDiscriminant { .. }
                 | StatementKind::Deinit(..)
                 | StatementKind::CopyNonOverlapping(..)
+                | StatementKind::Assume(..)
                 | StatementKind::Retag { .. } => {
                     return false;
                 }
diff --git a/compiler/rustc_mir_transform/src/separate_const_switch.rs b/compiler/rustc_mir_transform/src/separate_const_switch.rs
@@ -250,6 +250,7 @@ fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<
             | StatementKind::Coverage(_)
             | StatementKind::StorageDead(_)
             | StatementKind::CopyNonOverlapping(_)
+            | StatementKind::Assume(_)
             | StatementKind::Nop => {}
         }
     }
@@ -318,6 +319,7 @@ fn find_determining_place<'tcx>(
             | StatementKind::AscribeUserType(_, _)
             | StatementKind::Coverage(_)
             | StatementKind::CopyNonOverlapping(_)
+            | StatementKind::Assume(_)
             | StatementKind::Nop => {}
 
             // If the discriminant is set, it is always set
diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs
@@ -500,6 +500,7 @@ impl<'tcx> Visitor<'tcx> for UsedLocals {
     fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
         match statement.kind {
             StatementKind::CopyNonOverlapping(..)
+            | StatementKind::Assume(..)
             | StatementKind::Retag(..)
             | StatementKind::Coverage(..)
             | StatementKind::FakeRead(..)
diff --git a/src/test/mir-opt/lower_intrinsics.assume.LowerIntrinsics.diff b/src/test/mir-opt/lower_intrinsics.assume.LowerIntrinsics.diff
@@ -0,0 +1,26 @@
+- // MIR for `assume` before LowerIntrinsics
++ // MIR for `assume` after LowerIntrinsics
+  
+  fn assume() -> () {
+      let mut _0: ();                      // return place in scope 0 at $DIR/lower_intrinsics.rs:+0:17: +0:17
+      let _1: ();                          // in scope 0 at $DIR/lower_intrinsics.rs:+2:9: +2:38
+      scope 1 {
+      }
+  
+      bb0: {
+          StorageLive(_1);                 // scope 1 at $DIR/lower_intrinsics.rs:+2:9: +2:38
+-         _1 = std::intrinsics::assume(const true) -> bb1; // scope 1 at $DIR/lower_intrinsics.rs:+2:9: +2:38
+-                                          // mir::Constant
+-                                          // + span: $DIR/lower_intrinsics.rs:97:9: 97:32
+-                                          // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(bool) {std::intrinsics::assume}, val: Value(<ZST>) }
++         assume(const true);              // scope 1 at $DIR/lower_intrinsics.rs:+2:9: +2:38
++         goto -> bb1;                     // scope 1 at $DIR/lower_intrinsics.rs:+2:9: +2:38
+      }
+  
+      bb1: {
+          StorageDead(_1);                 // scope 1 at $DIR/lower_intrinsics.rs:+2:38: +2:39
+          _0 = const ();                   // scope 1 at $DIR/lower_intrinsics.rs:+1:5: +3:6
+          return;                          // scope 0 at $DIR/lower_intrinsics.rs:+4:2: +4:2
+      }
+  }
+  
diff --git a/src/test/mir-opt/lower_intrinsics.f_copy_nonoverlapping.LowerIntrinsics.diff b/src/test/mir-opt/lower_intrinsics.f_copy_nonoverlapping.LowerIntrinsics.diff
diff --git a/src/test/mir-opt/lower_intrinsics.rs b/src/test/mir-opt/lower_intrinsics.rs
diff --git a/src/tools/clippy/clippy_utils/src/qualify_min_const_fn.rs b/src/tools/clippy/clippy_utils/src/qualify_min_const_fn.rs

Original file line number	Diff line number	Diff line change
`@@ -392,6 +392,7 @@ impl<'tcx> rustc_mir_dataflow::GenKillAnalysis<'tcx> for Borrows<'_, 'tcx> {`
`392`	`392`	`\| mir::StatementKind::AscribeUserType(..)`
`393`	`393`	`\| mir::StatementKind::Coverage(..)`
`394`	`394`	`\| mir::StatementKind::CopyNonOverlapping(..)`
	`395`	`+ \| mir::StatementKind::Assume(..)`
`395`	`396`	`\| mir::StatementKind::Nop => {}`
`396`	`397`	`}`
`397`	`398`	`}`