1//! Provides the implementation of the `custom_mir` attribute.
2//!
3//! Up until MIR building, this attribute has absolutely no effect. The `mir!` macro is a normal
4//! decl macro that expands like any other, and the code goes through parsing, name resolution and
5//! type checking like all other code. In MIR building we finally detect whether this attribute is
6//! present, and if so we branch off into this module, which implements the attribute by
7//! implementing a custom lowering from THIR to MIR.
8//!
9//! The result of this lowering is returned "normally" from `build_mir`, with the only
10//! notable difference being that the `injected` field in the body is set. Various components of the
11//! MIR pipeline, like borrowck and the pass manager will then consult this field (via
12//! `body.should_skip()`) to skip the parts of the MIR pipeline that precede the MIR phase the user
13//! specified.
14//!
15//! This file defines the general framework for the custom parsing. The parsing for all the
16//! "top-level" constructs can be found in the `parse` submodule, while the parsing for statements,
17//! terminators, and everything below can be found in the `parse::instruction` submodule.
18//!
19
20use rustc_data_structures::fx::FxHashMap;
21use rustc_hir::def_id::DefId;
22use rustc_hir::{Attribute, HirId};
23use rustc_index::{IndexSlice, IndexVec};
24use rustc_middle::mir::*;
25use rustc_middle::span_bug;
26use rustc_middle::thir::*;
27use rustc_middle::ty::{self, Ty, TyCtxt};
28use rustc_span::Span;
29
30mod parse;
31
32pub(super) fn build_custom_mir<'tcx>(
33 tcx: TyCtxt<'tcx>,
34 did: DefId,
35 hir_id: HirId,
36 thir: &Thir<'tcx>,
37 expr: ExprId,
38 params: &IndexSlice<ParamId, Param<'tcx>>,
39 return_ty: Ty<'tcx>,
40 return_ty_span: Span,
41 span: Span,
42 attr: &Attribute,
43) -> Body<'tcx> {
44 let mut body = Body {
45 basic_blocks: BasicBlocks::new(IndexVec::new()),
46 source: MirSource::item(did),
47 phase: MirPhase::Built,
48 source_scopes: IndexVec::new(),
49 coroutine: None,
50 local_decls: IndexVec::new(),
51 user_type_annotations: IndexVec::new(),
52 arg_count: params.len(),
53 spread_arg: None,
54 var_debug_info: Vec::new(),
55 span,
56 required_consts: None,
57 mentioned_items: None,
58 is_polymorphic: false,
59 tainted_by_errors: None,
60 injection_phase: None,
61 pass_count: 0,
62 coverage_info_hi: None,
63 function_coverage_info: None,
64 };
65
66 body.local_decls.push(LocalDecl::new(return_ty, return_ty_span));
67 body.basic_blocks_mut().push(BasicBlockData::new(None, false));
68 body.source_scopes.push(SourceScopeData {
69 span,
70 parent_scope: None,
71 inlined: None,
72 inlined_parent_scope: None,
73 local_data: ClearCrossCrate::Set(SourceScopeLocalData { lint_root: hir_id }),
74 });
75 body.injection_phase = Some(parse_attribute(attr));
76
77 let mut pctxt = ParseCtxt {
78 tcx,
79 typing_env: body.typing_env(tcx),
80 thir,
81 source_scope: OUTERMOST_SOURCE_SCOPE,
82 body: &mut body,
83 local_map: FxHashMap::default(),
84 block_map: FxHashMap::default(),
85 };
86
87 let res: PResult<_> = try {
88 pctxt.parse_args(params)?;
89 pctxt.parse_body(expr)?;
90 };
91 if let Err(err) = res {
92 tcx.dcx().span_fatal(
93 err.span,
94 format!("Could not parse {}, found: {:?}", err.expected, err.item_description),
95 )
96 }
97
98 body
99}
100
101fn parse_attribute(attr: &Attribute) -> MirPhase {
102 let meta_items = attr.meta_item_list().unwrap();
103 let mut dialect: Option<String> = None;
104 let mut phase: Option<String> = None;
105
106 for nested in meta_items {
107 let name = nested.name_or_empty();
108 let value = nested.value_str().unwrap().as_str().to_string();
109 match name.as_str() {
110 "dialect" => {
111 assert!(dialect.is_none());
112 dialect = Some(value);
113 }
114 "phase" => {
115 assert!(phase.is_none());
116 phase = Some(value);
117 }
118 other => {
119 span_bug!(
120 nested.span(),
121 "Unexpected key while parsing custom_mir attribute: '{}'",
122 other
123 );
124 }
125 }
126 }
127
128 let Some(dialect) = dialect else {
129 assert!(phase.is_none());
130 return MirPhase::Built;
131 };
132
133 MirPhase::parse(dialect, phase)
134}
135
136struct ParseCtxt<'a, 'tcx> {
137 tcx: TyCtxt<'tcx>,
138 typing_env: ty::TypingEnv<'tcx>,
139 thir: &'a Thir<'tcx>,
140 source_scope: SourceScope,
141 body: &'a mut Body<'tcx>,
142 local_map: FxHashMap<LocalVarId, Local>,
143 block_map: FxHashMap<LocalVarId, BasicBlock>,
144}
145
146struct ParseError {
147 span: Span,
148 item_description: String,
149 expected: String,
150}
151
152impl<'a, 'tcx> ParseCtxt<'a, 'tcx> {
153 fn expr_error(&self, expr: ExprId, expected: &'static str) -> ParseError {
154 let expr: &{unknown} = &self.thir[expr];
155 ParseError {
156 span: expr.span,
157 item_description: format!("{:?}", expr.kind),
158 expected: expected.to_string(),
159 }
160 }
161
162 fn stmt_error(&self, stmt: StmtId, expected: &'static str) -> ParseError {
163 let stmt: &{unknown} = &self.thir[stmt];
164 let span: Span = match stmt.kind {
165 StmtKind::Expr { expr, .. } => self.thir[expr].span,
166 StmtKind::Let { span: Span, .. } => span,
167 };
168 ParseError {
169 span,
170 item_description: format!("{:?}", stmt.kind),
171 expected: expected.to_string(),
172 }
173 }
174}
175
176type PResult<T> = Result<T, ParseError>;
177