Preserve exit targets in nested control-flow macros

Signed-off-by: Perfloop Agent <agent@perfloop.ai>
Signed-off-by: Tomás Senart <ts@perfloop.ai>

Author: perfloop-agent

libcc2rs-macros/src/state_machine.rs

@@ -189,10 +189,7 @@ impl GotoRewriter<'_> {
 
     fn recurse_into_inner_goto_block(&mut self, mac: &mut syn::Macro) -> bool {
         if mac.path.is_ident("switch") || mac.path.is_ident("goto_block") {
-            if let Ok(mut inner) = syn::parse2::<Expr>(mac.tokens.clone()) {
-                self.visit_expr_mut(&mut inner);
-                mac.tokens = quote!(#inner);
-            }
+            visit_control_flow_macro(self, mac, true);
             return true;
         }
         false
@@ -284,12 +281,47 @@ impl StateMachine for SwitchStateMachine {
     }
 }
 
+// Macro tokens are opaque to syn's normal traversal. A nested goto block is
+// transparent to exits; a nested switch owns breaks in its arms, but not in
+// its condition or guards. Neither macro owns continues.
+fn visit_control_flow_macro(
+    visitor: &mut impl VisitMut,
+    mac: &mut syn::Macro,
+    visit_switch_arms: bool,
+) {
+    let Some(name) = mac.path.segments.last().map(|s| &s.ident) else {
+        return;
+    };
+    if name == "goto_block"
+        && let Ok(mut body) = syn::parse2::<Expr>(mac.tokens.clone())
+    {
+        visitor.visit_expr_mut(&mut body);
+        mac.tokens = quote!(#body);
+    } else if name == "switch"
+        && let Ok(mut body) = syn::parse2::<syn::ExprMatch>(mac.tokens.clone())
+    {
+        visitor.visit_expr_mut(&mut body.expr);
+        for arm in &mut body.arms {
+            if let Some((_, guard)) = &mut arm.guard {
+                visitor.visit_expr_mut(guard);
+            }
+            if visit_switch_arms {
+                visitor.visit_expr_mut(&mut arm.body);
+            }
+        }
+        mac.tokens = quote!(#body);
+    }
+}
+
 // Rewrite break into break '__sm
 struct ExitSwitchRewriter {
     label: Lifetime,
 }
 
 impl VisitMut for ExitSwitchRewriter {
+    fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
+        visit_control_flow_macro(self, mac, false);
+    }
     fn visit_expr_break_mut(&mut self, node: &mut ExprBreak) {
         if node.label.is_none() {
             node.label = Some(self.label.clone());
@@ -333,6 +365,9 @@ impl PropagateRewriter {
 }
 
 impl VisitMut for PropagateRewriter {
+    fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
+        visit_control_flow_macro(self, mac, matches!(self.kind, ControlKind::Continue));
+    }
     fn visit_expr_mut(&mut self, expr: &mut Expr) {
         let hit = match self.kind {
             ControlKind::Break => {

tests/unit/goto_nested_loop_control.c

@@ -0,0 +1,46 @@
+#include <assert.h>
+
+int main(void) {
+  int visits = 0;
+  int tail = 0;
+  for (int i = 0; i < 3; i++) {
+  outer:
+    visits++;
+    {
+    inner:
+      if (i == 0) {
+        break;
+      }
+    }
+    tail++;
+  }
+  assert(visits == 1);
+  assert(tail == 0);
+
+  visits = 0;
+  tail = 0;
+  for (int i = 0; i < 3; i++) {
+    switch (i) {
+    case 0: {
+    switch_outer:
+      visits++;
+      {
+      switch_inner:
+        if (i == 0) {
+          break;
+        }
+      }
+      tail++;
+    }
+    case 1:
+      tail += 10;
+      break;
+    default:
+      tail += 100;
+      break;
+    }
+  }
+  assert(visits == 1);
+  assert(tail == 110);
+  return 0;
+}

tests/unit/out/refcount/goto_nested_loop_control.rs

@@ -0,0 +1,80 @@
+extern crate libcc2rs;
+use libcc2rs::*;
+use std::cell::RefCell;
+use std::collections::BTreeMap;
+use std::io::prelude::*;
+use std::io::{Read, Seek, Write};
+use std::os::fd::AsFd;
+use std::rc::{Rc, Weak};
+pub fn main() {
+    __cpp2rust_init_globals();
+    std::process::exit(main_0());
+}
+fn main_0() -> i32 {
+    let mut visits: i32 = 0;
+    let mut tail: i32 = 0;
+    let mut i: i32 = 0;
+    'loop_: while (((i < 3) as i32) != 0) {
+        goto_block!({
+            '__entry: {}
+            'outer: {
+                visits.postfix_inc();
+                {
+                    goto_block!({
+                        '__entry: {}
+                        'inner: {
+                            if (((i == 0) as i32) != 0) {
+                                break;
+                            }
+                        }
+                    });
+                }
+                tail.postfix_inc();
+            }
+        });
+        i.postfix_inc();
+    }
+    assert!((((visits == 1) as i32) != 0));
+    assert!((((tail == 0) as i32) != 0));
+    visits = 0;
+    tail = 0;
+    let mut i: i32 = 0;
+    'loop_: while (((i < 3) as i32) != 0) {
+        switch!(match i {
+            __v if __v == 0 => {
+                {
+                    goto_block!({
+                        '__entry: {}
+                        'switch_outer: {
+                            visits.postfix_inc();
+                            {
+                                goto_block!({
+                                    '__entry: {}
+                                    'switch_inner: {
+                                        if (((i == 0) as i32) != 0) {
+                                            break;
+                                        }
+                                    }
+                                });
+                            }
+                            tail.postfix_inc();
+                        }
+                    });
+                }
+            }
+            __v if __v == 1 => {
+                tail += 10;
+                break;
+            }
+            _ => {
+                tail += 100;
+                break;
+            }
+        });
+        i.postfix_inc();
+    }
+    assert!((((visits == 1) as i32) != 0));
+    assert!((((tail == 110) as i32) != 0));
+    return 0;
+}
+pub fn __cpp2rust_init_globals() {}

tests/unit/out/unsafe/goto_nested_loop_control.rs

@@ -0,0 +1,82 @@
+extern crate libc;
+use libc::*;
+extern crate libcc2rs;
+use libcc2rs::*;
+use std::collections::BTreeMap;
+use std::io::{Read, Seek, Write};
+use std::os::fd::{AsFd, FromRawFd, IntoRawFd};
+use std::rc::Rc;
+pub fn main() {
+    unsafe {
+        __cpp2rust_init_globals();
+        std::process::exit(main_0() as i32);
+    }
+}
+unsafe fn main_0() -> i32 {
+    let mut visits: i32 = 0;
+    let mut tail: i32 = 0;
+    let mut i: i32 = 0;
+    'loop_: while ((((i) < (3)) as i32) != 0) {
+        goto_block!({
+            '__entry: {}
+            'outer: {
+                visits.postfix_inc();
+                {
+                    goto_block!({
+                        '__entry: {}
+                        'inner: {
+                            if ((((i) == (0)) as i32) != 0) {
+                                break;
+                            }
+                        }
+                    });
+                }
+                tail.postfix_inc();
+            }
+        });
+        i.postfix_inc();
+    }
+    assert!(((((visits) == (1)) as i32) != 0));
+    assert!(((((tail) == (0)) as i32) != 0));
+    visits = 0;
+    tail = 0;
+    let mut i: i32 = 0;
+    'loop_: while ((((i) < (3)) as i32) != 0) {
+        switch!(match i {
+            __v if __v == 0 => {
+                {
+                    goto_block!({
+                        '__entry: {}
+                        'switch_outer: {
+                            visits.postfix_inc();
+                            {
+                                goto_block!({
+                                    '__entry: {}
+                                    'switch_inner: {
+                                        if ((((i) == (0)) as i32) != 0) {
+                                            break;
+                                        }
+                                    }
+                                });
+                            }
+                            tail.postfix_inc();
+                        }
+                    });
+                }
+            }
+            __v if __v == 1 => {
+                tail += 10;
+                break;
+            }
+            _ => {
+                tail += 100;
+                break;
+            }
+        });
+        i.postfix_inc();
+    }
+    assert!(((((visits) == (1)) as i32) != 0));
+    assert!(((((tail) == (110)) as i32) != 0));
+    return 0;
+}
+pub unsafe fn __cpp2rust_init_globals() {}