Support most constant kinds in custom mir

6 files changed, 170 insertions, 78 deletions

diff --git a/compiler/rustc_mir_build/src/build/custom/parse.rs b/compiler/rustc_mir_build/src/build/custom/parse.rs
index 52cb0a4826d..b2c5aead430 100644
--- a/compiler/rustc_mir_build/src/build/custom/parse.rs
+++ b/compiler/rustc_mir_build/src/build/custom/parse.rs
@@ -23,6 +23,7 @@ macro_rules! parse_by_kind {
     (
         $self:ident,
         $expr_id:expr,
+        $expr_name:pat,
         $expected:literal,
         $(
             @call($name:literal, $args:ident) => $call_expr:expr,
@@ -33,6 +34,8 @@ macro_rules! parse_by_kind {
     ) => {{
         let expr_id = $self.preparse($expr_id);
         let expr = &$self.thir[expr_id];
+        debug!("Trying to parse {:?} as {}", expr.kind, $expected);
+        let $expr_name = expr;
         match &expr.kind {
             $(
                 ExprKind::Call { ty, fun: _, args: $args, .. } if {
@@ -137,10 +140,10 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     /// This allows us to easily parse the basic blocks declarations, local declarations, and
     /// basic block definitions in order.
     pub fn parse_body(&mut self, expr_id: ExprId) -> PResult<()> {
-        let body = parse_by_kind!(self, expr_id, "whole body",
+        let body = parse_by_kind!(self, expr_id, _, "whole body",
             ExprKind::Block { block } => self.thir[*block].expr.unwrap(),
         );
-        let (block_decls, rest) = parse_by_kind!(self, body, "body with block decls",
+        let (block_decls, rest) = parse_by_kind!(self, body, _, "body with block decls",
             ExprKind::Block { block } => {
                 let block = &self.thir[*block];
                 (&block.stmts, block.expr.unwrap())
@@ -148,7 +151,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         );
         self.parse_block_decls(block_decls.iter().copied())?;
 
-        let (local_decls, rest) = parse_by_kind!(self, rest, "body with local decls",
+        let (local_decls, rest) = parse_by_kind!(self, rest, _, "body with local decls",
             ExprKind::Block { block } => {
                 let block = &self.thir[*block];
                 (&block.stmts, block.expr.unwrap())
@@ -156,7 +159,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
         );
         self.parse_local_decls(local_decls.iter().copied())?;
 
-        let block_defs = parse_by_kind!(self, rest, "body with block defs",
+        let block_defs = parse_by_kind!(self, rest, _, "body with block defs",
             ExprKind::Block { block } => &self.thir[*block].stmts,
         );
         for (i, block_def) in block_defs.iter().enumerate() {
@@ -223,7 +226,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     fn parse_block_def(&self, expr_id: ExprId) -> PResult<BasicBlockData<'tcx>> {
-        let block = parse_by_kind!(self, expr_id, "basic block",
+        let block = parse_by_kind!(self, expr_id, _, "basic block",
             ExprKind::Block { block } => &self.thir[*block],
         );
 
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 6d6176584f5..cd809807a12 100644
--- a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
+++ b/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs
@@ -4,7 +4,7 @@ use super::{parse_by_kind, PResult, ParseCtxt};
 
 impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     pub fn parse_statement(&self, expr_id: ExprId) -> PResult<StatementKind<'tcx>> {
-        parse_by_kind!(self, expr_id, "statement",
+        parse_by_kind!(self, expr_id, _, "statement",
             @call("mir_retag", args) => {
                 Ok(StatementKind::Retag(RetagKind::Default, Box::new(self.parse_place(args[0])?)))
             },
@@ -20,7 +20,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     pub fn parse_terminator(&self, expr_id: ExprId) -> PResult<TerminatorKind<'tcx>> {
-        parse_by_kind!(self, expr_id, "terminator",
+        parse_by_kind!(self, expr_id, _, "terminator",
             @call("mir_return", _args) => {
                 Ok(TerminatorKind::Return)
             },
@@ -31,7 +31,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     fn parse_rvalue(&self, expr_id: ExprId) -> PResult<Rvalue<'tcx>> {
-        parse_by_kind!(self, expr_id, "rvalue",
+        parse_by_kind!(self, expr_id, _, "rvalue",
             ExprKind::Borrow { borrow_kind, arg } => Ok(
                 Rvalue::Ref(self.tcx.lifetimes.re_erased, *borrow_kind, self.parse_place(*arg)?)
             ),
@@ -43,14 +43,24 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     fn parse_operand(&self, expr_id: ExprId) -> PResult<Operand<'tcx>> {
-        parse_by_kind!(self, expr_id, "operand",
+        parse_by_kind!(self, expr_id, expr, "operand",
             @call("mir_move", args) => self.parse_place(args[0]).map(Operand::Move),
+            ExprKind::Literal { .. }
+            | ExprKind::NamedConst { .. }
+            | ExprKind::NonHirLiteral { .. }
+            | ExprKind::ZstLiteral { .. }
+            | ExprKind::ConstParam { .. }
+            | ExprKind::ConstBlock { .. } => {
+                Ok(Operand::Constant(Box::new(
+                    crate::build::expr::as_constant::as_constant_inner(expr, |_| None, self.tcx)
+                )))
+            },
             _ => self.parse_place(expr_id).map(Operand::Copy),
         )
     }
 
     fn parse_place(&self, expr_id: ExprId) -> PResult<Place<'tcx>> {
-        parse_by_kind!(self, expr_id, "place",
+        parse_by_kind!(self, expr_id, _, "place",
             ExprKind::Deref { arg } => Ok(
                 self.parse_place(*arg)?.project_deeper(&[PlaceElem::Deref], self.tcx)
             ),
@@ -59,13 +69,13 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     fn parse_local(&self, expr_id: ExprId) -> PResult<Local> {
-        parse_by_kind!(self, expr_id, "local",
+        parse_by_kind!(self, expr_id, _, "local",
             ExprKind::VarRef { id } => Ok(self.local_map[id]),
         )
     }
 
     fn parse_block(&self, expr_id: ExprId) -> PResult<BasicBlock> {
-        parse_by_kind!(self, expr_id, "basic block",
+        parse_by_kind!(self, expr_id, _, "basic block",
             ExprKind::VarRef { id } => Ok(self.block_map[id]),
         )
     }
diff --git a/compiler/rustc_mir_build/src/build/expr/as_constant.rs b/compiler/rustc_mir_build/src/build/expr/as_constant.rs
index 32c0207cb68..717c6231574 100644
--- a/compiler/rustc_mir_build/src/build/expr/as_constant.rs
+++ b/compiler/rustc_mir_build/src/build/expr/as_constant.rs
@@ -8,7 +8,9 @@ use rustc_middle::mir::interpret::{
 };
 use rustc_middle::mir::*;
 use rustc_middle::thir::*;
-use rustc_middle::ty::{self, CanonicalUserTypeAnnotation, TyCtxt};
+use rustc_middle::ty::{
+    self, CanonicalUserType, CanonicalUserTypeAnnotation, TyCtxt, UserTypeAnnotationIndex,
+};
 use rustc_span::DUMMY_SP;
 use rustc_target::abi::Size;
 
@@ -19,84 +21,87 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         let this = self;
         let tcx = this.tcx;
         let Expr { ty, temp_lifetime: _, span, ref kind } = *expr;
-        match *kind {
+        match kind {
             ExprKind::Scope { region_scope: _, lint_level: _, value } => {
-                this.as_constant(&this.thir[value])
-            }
-            ExprKind::Literal { lit, neg } => {
-                let literal =
-                    match lit_to_mir_constant(tcx, LitToConstInput { lit: &lit.node, ty, neg }) {
-                        Ok(c) => c,
-                        Err(LitToConstError::Reported(guar)) => {
-                            ConstantKind::Ty(tcx.const_error_with_guaranteed(ty, guar))
-                        }
-                        Err(LitToConstError::TypeError) => {
-                            bug!("encountered type error in `lit_to_mir_constant")
-                        }
-                    };
-
-                Constant { span, user_ty: None, literal }
+                this.as_constant(&this.thir[*value])
             }
-            ExprKind::NonHirLiteral { lit, ref user_ty } => {
-                let user_ty = user_ty.as_ref().map(|user_ty| {
-                    this.canonical_user_type_annotations.push(CanonicalUserTypeAnnotation {
+            _ => as_constant_inner(
+                expr,
+                |user_ty| {
+                    Some(this.canonical_user_type_annotations.push(CanonicalUserTypeAnnotation {
                         span,
                         user_ty: user_ty.clone(),
                         inferred_ty: ty,
-                    })
-                });
-                let literal = ConstantKind::Val(ConstValue::Scalar(Scalar::Int(lit)), ty);
+                    }))
+                },
+                tcx,
+            ),
+        }
+    }
+}
 
-                Constant { span, user_ty: user_ty, literal }
-            }
-            ExprKind::ZstLiteral { ref user_ty } => {
-                let user_ty = user_ty.as_ref().map(|user_ty| {
-                    this.canonical_user_type_annotations.push(CanonicalUserTypeAnnotation {
-                        span,
-                        user_ty: user_ty.clone(),
-                        inferred_ty: ty,
-                    })
-                });
-                let literal = ConstantKind::Val(ConstValue::ZeroSized, ty);
+pub fn as_constant_inner<'tcx>(
+    expr: &Expr<'tcx>,
+    push_cuta: impl FnMut(&Box<CanonicalUserType<'tcx>>) -> Option<UserTypeAnnotationIndex>,
+    tcx: TyCtxt<'tcx>,
+) -> Constant<'tcx> {
+    let Expr { ty, temp_lifetime: _, span, ref kind } = *expr;
+    match *kind {
+        ExprKind::Literal { lit, neg } => {
+            let literal =
+                match lit_to_mir_constant(tcx, LitToConstInput { lit: &lit.node, ty, neg }) {
+                    Ok(c) => c,
+                    Err(LitToConstError::Reported(guar)) => {
+                        ConstantKind::Ty(tcx.const_error_with_guaranteed(ty, guar))
+                    }
+                    Err(LitToConstError::TypeError) => {
+                        bug!("encountered type error in `lit_to_mir_constant")
+                    }
+                };
 
-                Constant { span, user_ty: user_ty, literal }
-            }
-            ExprKind::NamedConst { def_id, substs, ref user_ty } => {
-                let user_ty = user_ty.as_ref().map(|user_ty| {
-                    this.canonical_user_type_annotations.push(CanonicalUserTypeAnnotation {
-                        span,
-                        user_ty: user_ty.clone(),
-                        inferred_ty: ty,
-                    })
-                });
+            Constant { span, user_ty: None, literal }
+        }
+        ExprKind::NonHirLiteral { lit, ref user_ty } => {
+            let user_ty = user_ty.as_ref().map(push_cuta).flatten();
 
-                let uneval =
-                    mir::UnevaluatedConst::new(ty::WithOptConstParam::unknown(def_id), substs);
-                let literal = ConstantKind::Unevaluated(uneval, ty);
+            let literal = ConstantKind::Val(ConstValue::Scalar(Scalar::Int(lit)), ty);
 
-                Constant { user_ty, span, literal }
-            }
-            ExprKind::ConstParam { param, def_id: _ } => {
-                let const_param = tcx.mk_const(param, expr.ty);
-                let literal = ConstantKind::Ty(const_param);
+            Constant { span, user_ty: user_ty, literal }
+        }
+        ExprKind::ZstLiteral { ref user_ty } => {
+            let user_ty = user_ty.as_ref().map(push_cuta).flatten();
 
-                Constant { user_ty: None, span, literal }
-            }
-            ExprKind::ConstBlock { did: def_id, substs } => {
-                let uneval =
-                    mir::UnevaluatedConst::new(ty::WithOptConstParam::unknown(def_id), substs);
-                let literal = ConstantKind::Unevaluated(uneval, ty);
+            let literal = ConstantKind::Val(ConstValue::ZeroSized, ty);
 
-                Constant { user_ty: None, span, literal }
-            }
-            ExprKind::StaticRef { alloc_id, ty, .. } => {
-                let const_val = ConstValue::Scalar(Scalar::from_pointer(alloc_id.into(), &tcx));
-                let literal = ConstantKind::Val(const_val, ty);
+            Constant { span, user_ty: user_ty, literal }
+        }
+        ExprKind::NamedConst { def_id, substs, ref user_ty } => {
+            let user_ty = user_ty.as_ref().map(push_cuta).flatten();
 
-                Constant { span, user_ty: None, literal }
-            }
-            _ => span_bug!(span, "expression is not a valid constant {:?}", kind),
+            let uneval = mir::UnevaluatedConst::new(ty::WithOptConstParam::unknown(def_id), substs);
+            let literal = ConstantKind::Unevaluated(uneval, ty);
+
+            Constant { user_ty, span, literal }
+        }
+        ExprKind::ConstParam { param, def_id: _ } => {
+            let const_param = tcx.mk_const(ty::ConstKind::Param(param), expr.ty);
+            let literal = ConstantKind::Ty(const_param);
+
+            Constant { user_ty: None, span, literal }
+        }
+        ExprKind::ConstBlock { did: def_id, substs } => {
+            let uneval = mir::UnevaluatedConst::new(ty::WithOptConstParam::unknown(def_id), substs);
+            let literal = ConstantKind::Unevaluated(uneval, ty);
+
+            Constant { user_ty: None, span, literal }
+        }
+        ExprKind::StaticRef { alloc_id, ty, .. } => {
+            let const_val = ConstValue::Scalar(Scalar::from_pointer(alloc_id.into(), &tcx));
+            let literal = ConstantKind::Val(const_val, ty);
+
+            Constant { span, user_ty: None, literal }
         }
+        _ => span_bug!(span, "expression is not a valid constant {:?}", kind),
     }
 }
 
diff --git a/library/core/src/intrinsics/mir.rs b/library/core/src/intrinsics/mir.rs
index 6e3cf974119..18011aa5e98 100644
--- a/library/core/src/intrinsics/mir.rs
+++ b/library/core/src/intrinsics/mir.rs
@@ -151,6 +151,13 @@ pub macro __internal_extract_let {
         let $var $(: $ty)?;
         ::core::intrinsics::mir::__internal_extract_let!($($rest)*);
     },
+    // Due to #86730, we have to handle const blocks separately
+    (
+        let $var:ident $(: $ty:ty)? = const $block:block; $($rest:tt)*
+    ) => {
+        let $var $(: $ty)?;
+        ::core::intrinsics::mir::__internal_extract_let!($($rest)*);
+    },
     // Otherwise, output nothing
     (
         $stmt:stmt; $($rest:tt)*
@@ -218,6 +225,28 @@ pub macro __internal_remove_let {
             }
         }
     )},
+    // Due to #86730 , we have to handle const blocks separately
+    (
+        {
+            {
+                $($already_parsed:tt)*
+            }
+            {
+                let $var:ident $(: $ty:ty)? = const $block:block;
+                $($rest:tt)*
+            }
+        }
+    ) => { ::core::intrinsics::mir::__internal_remove_let!(
+        {
+            {
+                $($already_parsed)*
+                $var = const $block;
+            }
+            {
+                $($rest)*
+            }
+        }
+    )},
     // Otherwise, keep going
     (
         {
diff --git a/src/test/mir-opt/building/custom/consts.consts.built.after.mir b/src/test/mir-opt/building/custom/consts.consts.built.after.mir
new file mode 100644
index 00000000000..e384cdeb465
--- /dev/null
+++ b/src/test/mir-opt/building/custom/consts.consts.built.after.mir
@@ -0,0 +1,22 @@
+// MIR for `consts` after built
+
+fn consts() -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/consts.rs:10:27: 10:27
+    let mut _1: u8;                      // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+    let mut _2: i8;                      // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+    let mut _3: u32;                     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+    let mut _4: i32;                     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+    let mut _5: fn() {consts::<10>};     // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+
+    bb0: {
+        _1 = const 5_u8;                 // scope 0 at $DIR/consts.rs:+0:1: +0:26
+        _2 = const _;                    // scope 0 at $DIR/consts.rs:+0:1: +0:26
+        _3 = const C;                    // scope 0 at $DIR/consts.rs:+0:1: +0:26
+        _4 = const _;                    // scope 0 at $DIR/consts.rs:+0:1: +0:26
+        _5 = consts::<10>;               // scope 0 at $DIR/consts.rs:+0:1: +0:26
+                                         // mir::Constant
+                                         // + span: $DIR/consts.rs:16:18: 16:30
+                                         // + literal: Const { ty: fn() {consts::<10>}, val: Value(<ZST>) }
+        return;                          // scope 0 at $DIR/consts.rs:+0:1: +0:26
+    }
+}
diff --git a/src/test/mir-opt/building/custom/consts.rs b/src/test/mir-opt/building/custom/consts.rs
new file mode 100644
index 00000000000..98b087f1e58
--- /dev/null
+++ b/src/test/mir-opt/building/custom/consts.rs
@@ -0,0 +1,23 @@
+#![feature(custom_mir, core_intrinsics, inline_const)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+
+const D: i32 = 5;
+
+// EMIT_MIR consts.consts.built.after.mir
+#[custom_mir(dialect = "built")]
+fn consts<const C: u32>() {
+    mir!({
+        let _a = 5_u8;
+        let _b = const { 5_i8 };
+        let _c = C;
+        let _d = D;
+        let _e = consts::<10>;
+        Return()
+    })
+}
+
+fn main() {
+    consts::<5>();
+}