rust-fsm-dsl: refactor for a single pass over transitions

eugene-babichenko · May 12, 2024 · 5649858 · 5649858
1 parent 7abd833
commit 5649858
Showing 1 changed file with 26 additions and 32 deletions.
diff --git a/rust-fsm-dsl/src/lib.rs b/rust-fsm-dsl/src/lib.rs
@@ -42,61 +42,55 @@ pub fn state_machine(tokens: TokenStream) -> TokenStream {
     let fsm_name = input.name;
     let visibility = input.visibility;
 
-    let transitions: Vec<_> = input
-        .transitions
-        .iter()
-        .flat_map(|def| {
-            def.transitions.iter().map(move |transition| Transition {
-                initial_state: &def.initial_state,
-                input_value: &transition.input_value,
-                final_state: &transition.final_state,
-                output: &transition.output,
-            })
+    let transitions = input.transitions.iter().flat_map(|def| {
+        def.transitions.iter().map(move |transition| Transition {
+            initial_state: &def.initial_state,
+            input_value: &transition.input_value,
+            final_state: &transition.final_state,
+            output: &transition.output,
         })
-        .collect();
+    });
 
     let mut states = BTreeSet::new();
     let mut inputs = BTreeSet::new();
     let mut outputs = BTreeSet::new();
+    let mut transition_cases = Vec::new();
+    let mut output_cases = Vec::new();
 
     states.insert(&input.initial_state);
 
-    for transition in transitions.iter() {
-        states.insert(transition.initial_state);
-        states.insert(transition.final_state);
-        inputs.insert(transition.input_value);
-        if let Some(ref output) = transition.output {
-            outputs.insert(output);
-        }
-    }
+    for transition in transitions {
+        let Transition {
+            initial_state,
+            final_state,
+            input_value,
+            output,
+        } = transition;
 
-    let initial_state_name = &input.initial_state;
-
-    let mut transition_cases = vec![];
-    for transition in transitions.iter() {
-        let initial_state = &transition.initial_state;
-        let input_value = &transition.input_value;
-        let final_state = &transition.final_state;
         transition_cases.push(quote! {
             (Self::State::#initial_state, Self::Input::#input_value) => {
                 Some(Self::State::#final_state)
             }
         });
-    }
 
-    let mut output_cases = vec![];
-    for transition in transitions.iter() {
-        if let Some(output_value) = &transition.output {
-            let initial_state = &transition.initial_state;
-            let input_value = &transition.input_value;
+        if let Some(output_value) = output {
             output_cases.push(quote! {
                 (Self::State::#initial_state, Self::Input::#input_value) => {
                     Some(Self::Output::#output_value)
                 }
             });
         }
+
+        states.insert(initial_state);
+        states.insert(final_state);
+        inputs.insert(input_value);
+        if let Some(ref output) = output {
+            outputs.insert(output);
+        }
     }
 
+    let initial_state_name = &input.initial_state;
+
     let (input_type, input_impl) = match input.input_type {
         Some(t) => (quote!(#t), quote!()),
         None => (