1use super::*;
4use log::warn;
5use syn::{spanned::Spanned as _, ExprBreak, ExprIf, ExprReturn, ExprUnary, Stmt};
6
7use crate::rust_ast::{comment_store, set_span::SetSpan, BytePos, SpanExt};
8
9pub fn structured_cfg(
11 root: &[Structure<Stmt>],
12 comment_store: &mut comment_store::CommentStore,
13 current_block: Box<Expr>,
14 debug_labels: bool,
15 cut_out_trailing_ret: bool,
16) -> TranslationResult<Vec<Stmt>> {
17 let ast: StructuredAST<Box<Expr>, Pat, Label, Stmt> =
18 structured_cfg_help(vec![], &IndexSet::new(), root, &mut IndexSet::new())?;
19
20 let s = StructureState {
21 debug_labels,
22 current_block,
23 };
24 let (mut stmts, _span) = s.to_stmt(ast, comment_store);
25
26 if cut_out_trailing_ret {
29 if let Some(Stmt::Expr(Expr::Return(ExprReturn { expr: None, .. }), _)) = stmts.last() {
30 stmts.pop();
31 }
32 }
33
34 Ok(stmts)
35}
36
37#[derive(Copy, Clone, Debug)]
39pub enum ExitStyle {
40 Continue,
42
43 Break,
45}
46
47pub trait StructuredStatement: Sized {
49 type E;
51
52 type P;
54
55 type L;
57
58 type S;
60
61 fn empty() -> Self;
63
64 fn mk_singleton(stmt: Self::S) -> Self;
66
67 fn mk_append(self, second: Self) -> Self;
69
70 fn mk_goto(to: Self::L) -> Self;
72
73 fn mk_match(
75 cond: Self::E, cases: Vec<(Self::P, Self)>, ) -> Self;
78
79 fn mk_if(cond: Self::E, then: Self, else_: Self) -> Self;
81
82 fn mk_goto_table(
84 cases: Vec<(Self::L, Self)>, then: Self, ) -> Self;
87
88 fn mk_loop(lbl: Option<Self::L>, body: Self) -> Self;
90
91 fn mk_exit(
93 exit_style: ExitStyle, label: Option<Self::L>, ) -> Self;
96
97 fn extend_span(&mut self, span: Span);
98}
99
100#[derive(Debug, Clone)]
101pub struct Spanned<T> {
102 pub node: T,
103 pub span: Span,
104}
105
106pub type StructuredAST<E, P, L, S> = Spanned<StructuredASTKind<E, P, L, S>>;
107
108fn dummy_spanned<T>(inner: T) -> Spanned<T> {
109 Spanned {
110 node: inner,
111 span: Span::dummy(),
112 }
113}
114
115#[allow(missing_docs)]
117#[derive(Debug)]
118pub enum StructuredASTKind<E, P, L, S> {
119 Empty,
120 Singleton(S),
121 Append(
122 Box<StructuredAST<E, P, L, S>>,
123 Box<StructuredAST<E, P, L, S>>,
124 ),
125 Goto(L),
126 Match(E, Vec<(P, StructuredAST<E, P, L, S>)>),
127 If(
128 E,
129 Box<StructuredAST<E, P, L, S>>,
130 Box<StructuredAST<E, P, L, S>>,
131 ),
132 GotoTable(
133 Vec<(L, StructuredAST<E, P, L, S>)>,
134 Box<StructuredAST<E, P, L, S>>,
135 ),
136 Loop(Option<L>, Box<StructuredAST<E, P, L, S>>),
137 Exit(ExitStyle, Option<L>),
138}
139
140impl<E, P, L, S> StructuredStatement for StructuredAST<E, P, L, S> {
141 type E = E;
142 type P = P;
143 type L = L;
144 type S = S;
145
146 fn empty() -> Self {
147 dummy_spanned(StructuredASTKind::Empty)
148 }
149
150 fn mk_singleton(stmt: Self::S) -> Self {
151 dummy_spanned(StructuredASTKind::Singleton(stmt))
152 }
153
154 fn mk_append(self, second: Self) -> Self {
155 dummy_spanned(StructuredASTKind::Append(Box::new(self), Box::new(second)))
156 }
157
158 fn mk_goto(to: Self::L) -> Self {
159 dummy_spanned(StructuredASTKind::Goto(to))
160 }
161
162 fn mk_match(cond: Self::E, cases: Vec<(Self::P, Self)>) -> Self {
163 dummy_spanned(StructuredASTKind::Match(cond, cases))
164 }
165
166 fn mk_if(cond: Self::E, then: Self, else_: Self) -> Self {
167 dummy_spanned(StructuredASTKind::If(cond, Box::new(then), Box::new(else_)))
168 }
169
170 fn mk_goto_table(cases: Vec<(Self::L, Self)>, then: Self) -> Self {
171 dummy_spanned(StructuredASTKind::GotoTable(cases, Box::new(then)))
172 }
173
174 fn mk_loop(lbl: Option<Self::L>, body: Self) -> Self {
175 dummy_spanned(StructuredASTKind::Loop(lbl, Box::new(body)))
176 }
177
178 fn mk_exit(exit_style: ExitStyle, label: Option<Self::L>) -> Self {
179 dummy_spanned(StructuredASTKind::Exit(exit_style, label))
180 }
181
182 fn extend_span(&mut self, span: Span) {
183 if !self.span.is_dummy() {
184 self.span = span_subst_hi(self.span, span).unwrap_or_else(|| {
185 warn!("Could not extend span {:?} to {:?}", self.span, span);
186 self.span
187 });
188 } else {
189 self.span = span;
190 }
191 }
192}
193
194type Exit = (Label, IndexMap<Label, (IndexSet<Label>, ExitStyle)>);
195
196fn structured_cfg_help<S: StructuredStatement<E = Box<Expr>, P = Pat, L = Label, S = Stmt>>(
200 exits: Vec<Exit>,
201 next: &IndexSet<Label>,
202 root: &[Structure<Stmt>],
203 used_loop_labels: &mut IndexSet<Label>,
204) -> TranslationResult<S> {
205 let mut next: &IndexSet<Label> = next;
206 let mut rest: S = S::empty();
207
208 for structure in root.iter().rev() {
209 let mut new_rest: S = S::empty();
210
211 use Structure::*;
212 match structure {
213 Simple {
214 body,
215 terminator,
216 span,
217 ..
218 } => {
219 for s in body.clone() {
220 new_rest = S::mk_append(new_rest, S::mk_singleton(s));
221 }
222 new_rest.extend_span(*span);
223
224 let insert_goto = |to: Label, target: &IndexSet<Label>| -> S {
225 if target.len() == 1 {
226 S::empty()
227 } else {
228 S::mk_goto(to)
229 }
230 };
231
232 let mut branch = |slbl: &StructureLabel<Stmt>| -> TranslationResult<S> {
233 use StructureLabel::*;
234 match slbl {
235 Nested(ref nested) => {
236 structured_cfg_help(exits.clone(), next, nested, used_loop_labels)
237 }
238
239 GoTo(to) | ExitTo(to) if next.contains(to) => {
240 Ok(insert_goto(to.clone(), next))
241 }
242
243 ExitTo(to) => {
244 let mut immediate = true;
245 for (label, local) in &exits {
246 if let Some(&(ref follow, exit_style)) = local.get(to) {
247 let lbl = if immediate {
248 None
249 } else {
250 used_loop_labels.insert(label.clone());
251 Some(label.clone())
252 };
253
254 let mut new_cfg = S::mk_append(
255 insert_goto(to.clone(), follow),
256 S::mk_exit(exit_style, lbl),
257 );
258 new_cfg.extend_span(*span);
259 return Ok(new_cfg);
260 }
261 immediate = false;
262 }
263
264 Err(
265 format_err!("Not a valid exit: {:?} has nothing to exit to", to)
266 .into(),
267 )
268 }
269
270 GoTo(to) => Err(format_err!(
271 "Not a valid exit: {:?} (GoTo isn't falling through to {:?})",
272 to,
273 next
274 )
275 .into()),
276 }
277 };
278
279 new_rest = S::mk_append(
280 new_rest,
281 match terminator {
282 End => S::empty(),
283 Jump(to) => branch(to)?,
284 Branch(c, t, f) => S::mk_if(c.clone(), branch(t)?, branch(f)?),
285 Switch { expr, cases } => {
286 let branched_cases = cases
287 .iter()
288 .map(|(pat, slbl)| Ok((pat.clone(), branch(slbl)?)))
289 .collect::<TranslationResult<_>>()?;
290
291 S::mk_match(expr.clone(), branched_cases)
292 }
293 },
294 );
295 }
296
297 Multiple { branches, then, .. } => {
298 let cases = branches
299 .iter()
300 .map(|(lbl, body)| {
301 let stmts =
302 structured_cfg_help(exits.clone(), next, body, used_loop_labels)?;
303 Ok((lbl.clone(), stmts))
304 })
305 .collect::<TranslationResult<_>>()?;
306
307 let then: S = structured_cfg_help(exits.clone(), next, then, used_loop_labels)?;
308
309 new_rest = S::mk_append(new_rest, S::mk_goto_table(cases, then));
310 }
311
312 Loop { body, entries } => {
313 let label = entries
314 .iter()
315 .next()
316 .ok_or_else(|| format_err!("The loop {:?} has no entry", structure))?;
317
318 let mut these_exits = IndexMap::new();
319 these_exits.extend(
320 entries
321 .iter()
322 .map(|e| (e.clone(), (entries.clone(), ExitStyle::Continue))),
323 );
324 these_exits.extend(
325 next.iter()
326 .map(|e| (e.clone(), (next.clone(), ExitStyle::Break))),
327 );
328
329 let mut exits_new = vec![(label.clone(), these_exits)];
330 exits_new.extend(exits.clone());
331
332 let body = structured_cfg_help(exits_new, entries, body, used_loop_labels)?;
333 let loop_lbl = if used_loop_labels.contains(label) {
334 Some(label.clone())
335 } else {
336 None
337 };
338 new_rest = S::mk_append(new_rest, S::mk_loop(loop_lbl, body));
339 }
340 }
341
342 new_rest = S::mk_append(new_rest, rest);
343
344 rest = new_rest;
345 next = structure.get_entries();
346 }
347
348 Ok(rest)
349}
350
351pub fn has_multiple<Stmt>(root: &[Structure<Stmt>]) -> bool {
354 use Structure::*;
355 root.iter().any(|structure| match structure {
356 Simple { terminator, .. } => {
357 terminator
358 .get_labels()
359 .into_iter()
360 .any(|structure_label| match structure_label {
361 StructureLabel::Nested(nested) => has_multiple(nested),
362 _ => false,
363 })
364 }
365 Multiple { .. } => true,
366 Loop { body, .. } => has_multiple(body),
367 })
368}
369
370struct StructureState {
371 debug_labels: bool,
372 current_block: Box<Expr>,
373}
374
375fn span_subst_lo(span: Span, other: Span) -> Option<Span> {
379 if span.is_dummy() {
380 return Some(other.shrink_to_lo());
381 } else if span.lo() == BytePos(0) {
382 return Some(span.between(other));
383 } else if other.lo() != BytePos(0) && other.lo() != span.lo() {
384 return None;
385 }
386 Some(span)
387}
388
389fn span_subst_hi(span: Span, other: Span) -> Option<Span> {
393 if other.lo() != other.hi() {
394 if span.lo() == span.hi() {
395 return Some(other.between(span));
396 } else if other.hi() != span.hi() {
397 return None;
398 }
399 }
400 Some(span)
401}
402
403impl StructureState {
404 pub fn to_stmt(
405 &self,
406 ast: StructuredAST<Box<Expr>, Pat, Label, Stmt>,
407 comment_store: &mut comment_store::CommentStore,
408 ) -> (Vec<Stmt>, Span) {
409 use crate::cfg::structures::StructuredASTKind::*;
410
411 let span = ast.span;
412
413 let stmt = match ast.node {
414 Empty => return (vec![], ast.span),
415
416 Singleton(mut s) => {
417 let span = s.span().substitute_dummy(ast.span);
418 s.set_span(span);
419 return (vec![s], span);
420 }
421
422 Append(spanned, rhs) if matches!(spanned.node, Empty) => {
423 let lhs_span = spanned.span;
424 let span = ast.span.substitute_dummy(lhs_span);
425 let span = span_subst_lo(span, lhs_span).unwrap_or_else(|| {
426 comment_store.move_comments(lhs_span.lo(), span.lo());
427 span
428 });
429
430 let (mut stmts, stmts_span) = self.to_stmt(*rhs, comment_store);
431 let span = span_subst_hi(span, stmts_span).unwrap_or(span);
432
433 if let Some(stmt) = stmts.first_mut() {
436 stmt.set_span(span_subst_lo(stmt.span(), span).unwrap_or_else(|| {
437 comment_store.move_comments(stmt.span().lo(), span.lo());
438 stmt.span().with_lo(span.lo())
439 }));
440 }
441 if let Some(stmt) = stmts.last_mut() {
442 stmt.set_span(span_subst_hi(stmt.span(), span).unwrap_or_else(|| stmt.span()));
443 }
444 return (stmts, span);
445 }
446
447 Append(lhs, rhs) => {
448 let (mut stmts, lhs_span) = self.to_stmt(*lhs, comment_store);
449 let span = ast.span.substitute_dummy(lhs_span);
450 let span = span_subst_lo(span, lhs_span).unwrap_or_else(|| {
451 comment_store.move_comments(lhs_span.lo(), span.lo());
452 span
453 });
454 let (rhs_stmts, rhs_span) = self.to_stmt(*rhs, comment_store);
455 let span = span_subst_hi(span, rhs_span).unwrap_or(span);
456 stmts.extend(rhs_stmts);
457 if let Some(stmt) = stmts.first_mut() {
460 stmt.set_span(span_subst_lo(stmt.span(), span).unwrap_or_else(|| {
461 comment_store.move_comments(stmt.span().lo(), span.lo());
462 stmt.span().with_lo(span.lo())
463 }));
464 }
465 if let Some(stmt) = stmts.last_mut() {
466 stmt.set_span(span_subst_hi(stmt.span(), span).unwrap_or_else(|| stmt.span()));
467 }
468 return (stmts, span);
469 }
470
471 Goto(to) => {
472 let lbl_expr = if self.debug_labels {
475 to.to_string_expr()
476 } else {
477 to.to_num_expr()
478 };
479 mk().span(span)
480 .semi_stmt(mk().assign_expr(self.current_block.clone(), lbl_expr))
481 }
482
483 Match(cond, cases) => {
484 let arms: Vec<Arm> = cases
487 .into_iter()
488 .map(|(pat, stmts)| -> Arm {
489 let (stmts, span) = self.to_stmt(stmts, comment_store);
490
491 let body = mk().block_expr(mk().span(span).block(stmts));
492 mk().arm(pat, None, body)
493 })
494 .collect();
495
496 let e = mk().match_expr(cond, arms);
497
498 mk().span(span).expr_stmt(e)
499 }
500
501 If(cond, then, els) => {
502 let (then_stmts, then_span) = self.to_stmt(*then, comment_store);
510
511 let (mut els_stmts, els_span) = self.to_stmt(*els, comment_store);
512
513 let mut if_stmt = match (then_stmts.is_empty(), els_stmts.is_empty()) {
514 (true, true) => mk().semi_stmt(cond),
515 (false, true) => {
516 let if_expr =
517 mk().ifte_expr(cond, mk().span(then_span).block(then_stmts), None);
518 mk().expr_stmt(if_expr)
519 }
520 (true, false) => {
521 let negated_cond = not(&cond);
522 let if_expr = mk().ifte_expr(
523 negated_cond,
524 mk().span(els_span).block(els_stmts),
525 None,
526 );
527 mk().expr_stmt(if_expr)
528 }
529 (false, false) => {
530 fn is_expr(kind: &Stmt) -> bool {
531 matches!(kind, Stmt::Expr(Expr::If(..) | Expr::Block(..), None))
532 }
533
534 let is_els_expr = els_stmts.len() == 1 && is_expr(&els_stmts[0]);
538
539 let els_branch = if is_els_expr {
540 let stmt_expr = els_stmts.swap_remove(0);
541 let stmt_expr_span = stmt_expr.span();
542 let mut els_expr = match stmt_expr {
543 Stmt::Expr(e, None) => e,
544 _ => panic!("is_els_expr out of sync"),
545 };
546 els_expr.set_span(stmt_expr_span);
547 Box::new(els_expr)
548 } else {
549 mk().block_expr(mk().span(els_span).block(els_stmts))
550 };
551
552 let if_expr = mk().ifte_expr(
553 cond,
554 mk().span(then_span).block(then_stmts),
555 Some(els_branch),
556 );
557 mk().expr_stmt(if_expr)
558 }
559 };
560
561 if_stmt.set_span(span);
562 if_stmt
563 }
564
565 GotoTable(cases, then) => {
566 let mut arms: Vec<Arm> = cases
569 .into_iter()
570 .map(|(lbl, stmts)| -> Arm {
571 let (stmts, stmts_span) = self.to_stmt(stmts, comment_store);
572
573 let lbl_lit = if self.debug_labels {
574 lbl.to_string_lit()
575 } else {
576 lbl.to_int_lit()
577 };
578 let pat = mk().lit_pat(lbl_lit);
579 let body = mk().block_expr(mk().span(stmts_span).block(stmts));
580 mk().arm(pat, None, body)
581 })
582 .collect();
583
584 let (then, then_span) = self.to_stmt(*then, comment_store);
585
586 arms.push(mk().arm(
587 mk().wild_pat(),
588 None,
589 mk().block_expr(mk().span(then_span).block(then)),
590 ));
591
592 let e = mk().match_expr(self.current_block.clone(), arms);
593
594 mk().span(span).expr_stmt(e)
595 }
596
597 Loop(lbl, body) => {
598 let (body, body_span) = self.to_stmt(*body, comment_store);
604
605 if let Some(stmt @ &Stmt::Expr(ref expr, None)) = body.first() {
607 let stmt_span = stmt.span();
608 let span = if !stmt_span.is_dummy() {
609 stmt_span
610 } else {
611 span
612 };
613 if let syn::Expr::If(ExprIf {
614 cond,
615 then_branch,
616 else_branch: None,
617 ..
618 }) = expr
619 {
620 if let [Stmt::Expr(
621 syn::Expr::Break(ExprBreak {
622 label: None,
623 expr: None,
624 ..
625 }),
626 Some(_),
627 )] = then_branch.stmts.as_slice()
628 {
629 let e = mk().while_expr(
630 not(cond),
631 mk().span(body_span)
632 .block(body.iter().skip(1).cloned().collect()),
633 lbl.map(|l| l.pretty_print()),
634 );
635 return (vec![mk().span(span).expr_stmt(e)], ast.span);
636 }
637 }
638 }
639
640 let e = mk().loop_expr(
641 mk().span(body_span).block(body),
642 lbl.map(|l| l.pretty_print()),
643 );
644
645 mk().span(span).expr_stmt(e)
646 }
647
648 Exit(exit_style, lbl) => {
649 let lbl = lbl.map(|l| l.pretty_print());
652 let e = match exit_style {
653 ExitStyle::Break => mk().break_expr(lbl),
654 ExitStyle::Continue => mk().continue_expr(lbl),
655 };
656
657 mk().span(span).semi_stmt(e)
658 }
659 };
660
661 (vec![stmt], ast.span)
662 }
663}
664
665fn not(bool_expr: &Expr) -> Box<Expr> {
670 use syn::UnOp;
671 match *bool_expr {
672 Expr::Unary(ExprUnary {
673 op: UnOp::Not(_),
674 ref expr,
675 ..
676 }) => Box::new(unparen(expr).clone()),
677 _ => mk().unary_expr(UnOp::Not(Default::default()), Box::new(bool_expr.clone())),
678 }
679}