Skip to content

Commit bff1419

Browse files
tmiaskotmiasko
committed
Custom MIR: Support cleanup blocks
Cleanup blocks are declared with `bb (cleanup) = { ... }`. `Call` and `Drop` terminators take an additional argument describing the unwind action, which is one of the following: * `Continue()` * `Unreachable()` * `Termiante(reason)`, where reason is `Abi` or `InCleanup` * `Cleanup(block)` Also support resume and terminate terminators: * `Resume()` * `Terminate(reason)`
1 parent 7cc36de commit bff1419

File tree

16 files changed

+282
-51
lines changed

16 files changed

+282
-51
lines changed

‎compiler/rustc_mir_build/src/build/custom/parse.rs

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,10 @@ macro_rules! parse_by_kind {
3030
@call($name:literal, $args:ident) => $call_expr:expr,
3131
)*
3232
$(
33-
$pat:pat => $expr:expr,
33+
@variant($adt:literal, $variant:literal) => $variant_expr:expr,
34+
)*
35+
$(
36+
$pat:pat $(if $guard:expr)? => $expr:expr,
3437
)*
3538
) => {{
3639
let expr_id = $self.preparse($expr_id);
@@ -49,7 +52,13 @@ macro_rules! parse_by_kind {
4952
} => $call_expr,
5053
)*
5154
$(
52-
$pat => $expr,
55+
ExprKind::Adt(box AdtExpr { adt_def, variant_index, .. }) if {
56+
$self.tcx.is_diagnostic_item(rustc_span::Symbol::intern($adt), adt_def.did()) &&
57+
adt_def.variants()[*variant_index].name == rustc_span::Symbol::intern($variant)
58+
} => $variant_expr,
59+
)*
60+
$(
61+
$pat $(if $guard)? => $expr,
5362
)*
5463
#[allow(unreachable_patterns)]
5564
_ => return Err($self.expr_error(expr_id, $expected))
@@ -172,7 +181,8 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
172181
ExprKind::Block { block } => &self.thir[*block].stmts,
173182
);
174183
for (i, block_def) in block_defs.iter().enumerate() {
175-
let block = self.parse_block_def(self.statement_as_expr(*block_def)?)?;
184+
let is_cleanup = self.body.basic_blocks_mut()[BasicBlock::from_usize(i)].is_cleanup;
185+
let block = self.parse_block_def(self.statement_as_expr(*block_def)?, is_cleanup)?;
176186
self.body.basic_blocks_mut()[BasicBlock::from_usize(i)] = block;
177187
}
178188

@@ -181,8 +191,9 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
181191

182192
fn parse_block_decls(&mut self, stmts: impl Iterator<Item = StmtId>) -> PResult<()> {
183193
for stmt in stmts {
184-
let (var, _, _) = self.parse_let_statement(stmt)?;
185-
let data = BasicBlockData::new(None);
194+
let (var, _, _, initializer) = self.parse_let_statement(stmt)?;
195+
let mut data = BasicBlockData::new(None);
196+
data.is_cleanup = initializer.is_some();
186197
let block = self.body.basic_blocks_mut().push(data);
187198
self.block_map.insert(var, block);
188199
}
@@ -195,7 +206,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
195206
self.local_map.insert(ret_var, Local::from_u32(0));
196207

197208
for stmt in stmts {
198-
let (var, ty, span) = self.parse_let_statement(stmt)?;
209+
let (var, ty, span, _) = self.parse_let_statement(stmt)?;
199210
let decl = LocalDecl::new(ty, span);
200211
let local = self.body.local_decls.push(decl);
201212
self.local_map.insert(var, local);
@@ -251,15 +262,17 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
251262
Ok(())
252263
}
253264

254-
fn parse_let_statement(&mut self, stmt_id: StmtId) -> PResult<(LocalVarId, Ty<'tcx>, Span)> {
255-
let pattern = match &self.thir[stmt_id].kind {
256-
StmtKind::Let { pattern, .. } => pattern,
257-
StmtKind::Expr { expr, .. } => {
258-
return Err(self.expr_error(*expr, "let statement"));
265+
fn parse_let_statement(
266+
&mut self,
267+
stmt_id: StmtId,
268+
) -> PResult<(LocalVarId, Ty<'tcx>, Span, Option<ExprId>)> {
269+
match &self.thir[stmt_id].kind {
270+
StmtKind::Let { pattern, initializer, .. } => {
271+
let (var, ty, span) = self.parse_var(pattern)?;
272+
Ok((var, ty, span, *initializer))
259273
}
260-
};
261-
262-
self.parse_var(pattern)
274+
StmtKind::Expr { expr, .. } => Err(self.expr_error(*expr, "let statement")),
275+
}
263276
}
264277

265278
fn parse_var(&mut self, mut pat: &Pat<'tcx>) -> PResult<(LocalVarId, Ty<'tcx>, Span)> {
@@ -281,12 +294,13 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
281294
}
282295
}
283296

284-
fn parse_block_def(&self, expr_id: ExprId) -> PResult<BasicBlockData<'tcx>> {
297+
fn parse_block_def(&self, expr_id: ExprId, is_cleanup: bool) -> PResult<BasicBlockData<'tcx>> {
285298
let block = parse_by_kind!(self, expr_id, _, "basic block",
286299
ExprKind::Block { block } => &self.thir[*block],
287300
);
288301

289302
let mut data = BasicBlockData::new(None);
303+
data.is_cleanup = is_cleanup;
290304
for stmt_id in &*block.stmts {
291305
let stmt = self.statement_as_expr(*stmt_id)?;
292306
let span = self.thir[stmt].span;

‎compiler/rustc_mir_build/src/build/custom/parse/instruction.rs

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,17 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
5252
@call("mir_unreachable", _args) => {
5353
Ok(TerminatorKind::Unreachable)
5454
},
55+
@call("mir_resume", _args) => {
56+
Ok(TerminatorKind::UnwindResume)
57+
},
58+
@call("mir_terminate", args) => {
59+
Ok(TerminatorKind::UnwindTerminate(self.parse_unwind_terminate_reason(args[0])?))
60+
},
5561
@call("mir_drop", args) => {
5662
Ok(TerminatorKind::Drop {
5763
place: self.parse_place(args[0])?,
5864
target: self.parse_block(args[1])?,
59-
unwind: UnwindAction::Continue,
65+
unwind: self.parse_unwind_action(args[2])?,
6066
replace: false,
6167
})
6268
},
@@ -70,6 +76,34 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
7076
)
7177
}
7278

79+
fn parse_unwind_terminate_reason(&self, expr_id: ExprId) -> PResult<UnwindTerminateReason> {
80+
parse_by_kind!(self, expr_id, _, "unwind terminate reason",
81+
@variant("mir_unwind_terminate_reason", "Abi") => {
82+
Ok(UnwindTerminateReason::Abi)
83+
},
84+
@variant("mir_unwind_terminate_reason", "InCleanup") => {
85+
Ok(UnwindTerminateReason::InCleanup)
86+
},
87+
)
88+
}
89+
90+
fn parse_unwind_action(&self, expr_id: ExprId) -> PResult<UnwindAction> {
91+
parse_by_kind!(self, expr_id, _, "unwind action",
92+
@call("mir_continue", _args) => {
93+
Ok(UnwindAction::Continue)
94+
},
95+
@call("mir_unreachable", _args) => {
96+
Ok(UnwindAction::Unreachable)
97+
},
98+
@call("mir_terminate", args) => {
99+
Ok(UnwindAction::Terminate(self.parse_unwind_terminate_reason(args[0])?))
100+
},
101+
@call("mir_cleanup", args) => {
102+
Ok(UnwindAction::Cleanup(self.parse_block(args[0])?))
103+
},
104+
)
105+
}
106+
73107
fn parse_match(&self, arms: &[ArmId], span: Span) -> PResult<SwitchTargets> {
74108
let Some((otherwise, rest)) = arms.split_last() else {
75109
return Err(ParseError {
@@ -113,6 +147,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
113147
);
114148
let destination = self.parse_place(destination)?;
115149
let target = self.parse_block(args[1])?;
150+
let unwind = self.parse_unwind_action(args[2])?;
116151

117152
parse_by_kind!(self, call, _, "function call",
118153
ExprKind::Call { fun, args, from_hir_call, fn_span, .. } => {
@@ -126,7 +161,7 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
126161
args,
127162
destination,
128163
target: Some(target),
129-
unwind: UnwindAction::Continue,
164+
unwind,
130165
call_source: if *from_hir_call { CallSource::Normal } else {
131166
CallSource::OverloadedOperator
132167
},

‎library/core/src/intrinsics/mir.rs

Lines changed: 41 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -110,15 +110,15 @@
110110
//! let popped;
111111
//!
112112
//! {
113-
//! Call(_unused = Vec::push(v, value), pop)
113+
//! Call(_unused = Vec::push(v, value), pop, Continue())
114114
//! }
115115
//!
116116
//! pop = {
117-
//! Call(popped = Vec::pop(v), drop)
117+
//! Call(popped = Vec::pop(v), drop, Continue())
118118
//! }
119119
//!
120120
//! drop = {
121-
//! Drop(popped, ret)
121+
//! Drop(popped, ret, Continue())
122122
//! }
123123
//!
124124
//! ret = {
@@ -238,10 +238,6 @@
238238
//!
239239
//! #### Terminators
240240
//!
241-
//! Custom MIR does not currently support cleanup blocks or non-trivial unwind paths. As such, there
242-
//! are no resume and abort terminators, and terminators that might unwind do not have any way to
243-
//! indicate the unwind block.
244-
//!
245241
//! - [`Goto`], [`Return`], [`Unreachable`] and [`Drop`](Drop()) have associated functions.
246242
//! - `match some_int_operand` becomes a `SwitchInt`. Each arm should be `literal => basic_block`
247243
//! - The exception is the last arm, which must be `_ => basic_block` and corresponds to the
@@ -262,6 +258,18 @@
262258
/// All terminators will have this type as a return type. It helps achieve some type safety.
263259
pub struct BasicBlock;
264260

261+
/// The reason we are terminating the process during unwinding.
262+
#[rustc_diagnostic_item = "mir_unwind_terminate_reason"]
263+
pub enum UnwindTerminateReason {
264+
/// Unwinding is just not possible given the ABI of this function.
265+
Abi,
266+
/// We were already cleaning up for an ongoing unwind, and a *second*, *nested* unwind was
267+
/// triggered by the drop glue.
268+
InCleanup,
269+
}
270+
271+
pub use UnwindTerminateReason::*;
272+
265273
macro_rules! define {
266274
($name:literal, $( #[ $meta:meta ] )* fn $($sig:tt)*) => {
267275
#[rustc_diagnostic_item = $name]
@@ -270,12 +278,15 @@ macro_rules! define {
270278
pub fn $($sig)* { panic!() }
271279
}
272280
}
273-
274281
define!("mir_return", fn Return() -> BasicBlock);
275282
define!("mir_goto", fn Goto(destination: BasicBlock) -> BasicBlock);
276283
define!("mir_unreachable", fn Unreachable() -> BasicBlock);
277-
define!("mir_drop", fn Drop<T>(place: T, goto: BasicBlock));
278-
define!("mir_call", fn Call(call: (), goto: BasicBlock));
284+
define!("mir_continue", fn Continue());
285+
define!("mir_resume", fn Resume());
286+
define!("mir_terminate", fn Terminate(reason: UnwindTerminateReason));
287+
define!("mir_cleanup", fn Cleanup(goto: BasicBlock));
288+
define!("mir_drop", fn Drop<T, U>(place: T, goto: BasicBlock, unwind_action: U));
289+
define!("mir_call", fn Call<U>(call: (), goto: BasicBlock, unwind_action: U));
279290
define!("mir_storage_live", fn StorageLive<T>(local: T));
280291
define!("mir_storage_dead", fn StorageDead<T>(local: T));
281292
define!("mir_deinit", fn Deinit<T>(place: T));
@@ -382,16 +393,15 @@ pub macro mir {
382393
}
383394

384395
$(
385-
$block_name:ident = {
396+
$block_name:ident $(($block_cleanup:ident))? = {
386397
$($block:tt)*
387398
}
388399
)*
389400
) => {{
390401
// First, we declare all basic blocks.
391-
$(
392-
let $block_name: ::core::intrinsics::mir::BasicBlock;
393-
)*
394-
402+
__internal_declare_basic_blocks!($(
403+
$block_name $(($block_cleanup))?
404+
)*);
395405
{
396406
// Now all locals
397407
#[allow(non_snake_case)]
@@ -585,3 +595,19 @@ pub macro __internal_remove_let {
585595
}
586596
},
587597
}
598+
599+
/// Helper macro that declares the basic blocks.
600+
///
601+
/// The cleanup block are distinguished by an initializer.
602+
#[doc(hidden)]
603+
pub macro __internal_declare_basic_blocks {
604+
() => {},
605+
($name:ident (cleanup) $($rest:tt)*) => {
606+
let $name = ::core::intrinsics::mir::BasicBlock;
607+
__internal_declare_basic_blocks!($($rest)*)
608+
},
609+
($name:ident $($rest:tt)*) => {
610+
let $name : ::core::intrinsics::mir::BasicBlock;
611+
__internal_declare_basic_blocks!($($rest)*)
612+
},
613+
}

‎tests/mir-opt/building/custom/terminate.rs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
// compile-flags: --crate-type=lib
2+
// edition:2021
3+
#![feature(custom_mir, core_intrinsics)]
4+
use core::intrinsics::mir::*;
5+
6+
// CHECK-LABEL: fn f()
7+
// CHECK: bb1 (cleanup): {
8+
// CHECK-NEXT: abort(abi);
9+
#[custom_mir(dialect = "runtime", phase = "optimized")]
10+
pub fn f() {
11+
mir!(
12+
{
13+
Return()
14+
}
15+
bb1(cleanup) = {
16+
Terminate(Abi)
17+
}
18+
)
19+
}
20+
21+
// CHECK-LABEL: fn g()
22+
// CHECK: bb1 (cleanup): {
23+
// CHECK-NEXT: abort(cleanup);
24+
#[custom_mir(dialect = "runtime", phase = "optimized")]
25+
pub fn g() {
26+
mir!(
27+
{
28+
Return()
29+
}
30+
bb1(cleanup) = {
31+
Terminate(InCleanup)
32+
}
33+
)
34+
}

‎tests/mir-opt/building/custom/terminators.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ fn ident<T>(t: T) -> T {
1313
fn direct_call(x: i32) -> i32 {
1414
mir!(
1515
{
16-
Call(RET = ident(x), retblock)
16+
Call(RET = ident(x), retblock, Continue())
1717
}
1818

1919
retblock = {
@@ -27,7 +27,7 @@ fn direct_call(x: i32) -> i32 {
2727
fn indirect_call(x: i32, f: fn(i32) -> i32) -> i32 {
2828
mir!(
2929
{
30-
Call(RET = f(x), retblock)
30+
Call(RET = f(x), retblock, Continue())
3131
}
3232

3333
retblock = {
@@ -49,7 +49,7 @@ impl<'a> Drop for WriteOnDrop<'a> {
4949
fn drop_first<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
5050
mir!(
5151
{
52-
Drop(a, retblock)
52+
Drop(a, retblock, Continue())
5353
}
5454

5555
retblock = {
@@ -64,7 +64,7 @@ fn drop_first<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
6464
fn drop_second<'a>(a: WriteOnDrop<'a>, b: WriteOnDrop<'a>) {
6565
mir!(
6666
{
67-
Drop(b, retblock)
67+
Drop(b, retblock, Continue())
6868
}
6969

7070
retblock = {

0 commit comments

Comments
 (0)