Add StateMachineNames

Author: lucic71

libcc2rs-macros/src/goto.rs

@@ -2,14 +2,16 @@
 // Distributed under the MIT license that can be found in the LICENSE file.
 
 use proc_macro::TokenStream;
+use proc_macro2::Span;
 use syn::parse::{Parse, ParseStream};
-use syn::{Expr, Lifetime, Token, parse_macro_input};
+use syn::{Block, Expr, ExprBlock, Lifetime, Stmt, parse_macro_input};
 
-use crate::state_machine::{Arm, GotoStateMachine, StateMachine};
+use crate::state_machine::{Arm, GotoStateMachine, StateMachineNames, StateMachine};
 
 pub fn expand(input: TokenStream) -> TokenStream {
     let GotoBlockInput { arms } = parse_macro_input!(input as GotoBlockInput);
     GotoStateMachine {
+        names: StateMachineNames::fresh(),
         arms: arms
             .into_iter()
             .map(|a| Arm {

libcc2rs-macros/src/state_machine.rs

@@ -1,10 +1,13 @@
 // Copyright (c) 2022-present INESC-ID.
 // Distributed under the MIT license that can be found in the LICENSE file.
 
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::{format_ident, quote};
 use syn::visit_mut::{self, VisitMut};
-use syn::{Expr, ExprBreak, ExprContinue, Lifetime, Pat};
+use syn::{Expr, ExprBreak, ExprContinue, Lifetime, Pat, Stmt, parse_quote};
 
 pub struct Arm {
     pub label: String,
@@ -21,12 +24,29 @@ pub trait StateMachine {
     fn emit(self) -> TokenStream2;
 }
 
-fn sm_label() -> Lifetime {
-    Lifetime::new("'__sm", Span::call_site())
+pub(crate) struct StateMachineNames {
+    pub label: Lifetime,
+    pub state: Ident,
+    pub break_flag: Ident,
+    pub cont_flag: Ident,
+}
+
+impl StateMachineNames {
+    pub fn fresh() -> Self {
+        static COUNTER: AtomicU64 = AtomicU64::new(0);
+        let id = COUNTER.fetch_add(1, Ordering::Relaxed);
+        Self {
+            label: Lifetime::new(&format!("'__sm{id}"), Span::call_site()),
+            state: format_ident!("__s{}", id),
+            break_flag: format_ident!("__user_break{}", id),
+            cont_flag: format_ident!("__user_continue{}", id),
+        }
+    }
 }
 
 // Collection of labeled arms that fall-through by default
 pub struct GotoStateMachine {
+    pub names: StateMachineNames,
     pub arms: Vec<Arm>,
 }
 
@@ -87,10 +107,7 @@ impl GotoStateMachine {
 
 impl StateMachine for GotoStateMachine {
     fn emit(self) -> TokenStream2 {
-        let lbl = sm_label();
-        let s = format_ident!("__s");
-        let break_flag = format_ident!("__user_break");
-        let cont_flag = format_ident!("__user_continue");
+        let StateMachineNames { label: lbl, state: s, break_flag, cont_flag } = self.names;
 
         let n = self.arms.len();
         let mut arms_have_break = false;
@@ -186,17 +203,16 @@ impl SwitchStateMachine {
 
 impl StateMachine for SwitchStateMachine {
     fn emit(self) -> TokenStream2 {
-        let lbl = sm_label();
-        let s = format_ident!("__s");
+        let names = StateMachineNames::fresh();
 
-        let user_arms = Self::convert_break_to_switch_exit(&self.goto.arms, &lbl);
-        let dispatch = self.build_dispatch_arm(&user_arms, &lbl, &s);
+        let user_arms = Self::convert_break_to_switch_exit(&self.goto.arms, &names.label);
+        let dispatch = self.build_dispatch_arm(&user_arms, &names.label, &names.state);
 
         let mut arms = Vec::new();
         arms.push(dispatch);
         arms.extend(user_arms);
 
-        GotoStateMachine { arms }.emit()
+        GotoStateMachine { names, arms }.emit()
     }
 }
 

libcc2rs-macros/src/switch.rs

@@ -5,7 +5,9 @@ use proc_macro::TokenStream;
 use syn::parse::{Parse, ParseStream};
 use syn::{Expr, Pat, parse_macro_input};
 
-use crate::state_machine::{Arm, DispatchCase, GotoStateMachine, StateMachine, SwitchStateMachine};
+use crate::state_machine::{
+    Arm, DispatchCase, GotoStateMachine, StateMachineNames, StateMachine, SwitchStateMachine,
+};
 
 pub fn expand(input: TokenStream) -> TokenStream {
     let SwitchInput { condition, arms } = parse_macro_input!(input as SwitchInput);
@@ -24,7 +26,10 @@ pub fn expand(input: TokenStream) -> TokenStream {
         });
     }
     SwitchStateMachine {
-        goto: GotoStateMachine { arms: cfg_arms },
+        goto: GotoStateMachine {
+            names: StateMachineNames::fresh(),
+            arms: cfg_arms,
+        },
         condition,
         cases,
     }