feat: Propagate rustdoc to generated (super)state enums

Problem: `state_machine!()` users could not document their state and
superstate handler functions with doc comments that would appear on
the generated enum variants. This made the generated code harder to
understand, reduced IDE support, and meant users could not use lints
like `missing_docs`.

Solution: Propagate doc comments through the macro pipeline by
extracting them from the state/superstate handler functions during
analysis, and then attaching the extracted doc comments as attributes
to the generated enum variants. Additionally we added default doc
comments to the generated State and Superstate enums.

Testing: Added trybuild tests that use `#![deny(missing_docs)]` to
validate that the doc comments are propagated.

Author: Owen Phillips

macro/src/analyze.rs

@@ -86,6 +86,8 @@ pub struct State {
     pub context_arg: Option<PatType>,
     /// Whether the function is async or not.
     pub is_async: bool,
+    /// Doc comments from the state handler function.
+    pub doc_comments: Vec<Attribute>,
 }
 
 /// Information regarding a superstate.
@@ -113,6 +115,8 @@ pub struct Superstate {
     pub context_arg: Option<PatType>,
     /// Whether the function is async or not.
     pub is_async: bool,
+    /// Doc comments from the superstate handler function.
+    pub doc_comments: Vec<Attribute>,
 }
 
 /// Information regarding an action.
@@ -540,6 +544,8 @@ pub fn analyze_state(method: &mut ImplItemFn, state_machine: &StateMachine) -> S
         }
     }
 
+    let doc_comments = extract_doc_comments(&method.attrs);
+
     State {
         handler_name,
         superstate,
@@ -552,6 +558,7 @@ pub fn analyze_state(method: &mut ImplItemFn, state_machine: &StateMachine) -> S
         event_arg,
         context_arg,
         is_async,
+        doc_comments,
     }
 }
 
@@ -665,6 +672,8 @@ pub fn analyze_superstate(method: &ImplItemFn, state_machine: &StateMachine) ->
         }
     }
 
+    let doc_comments = extract_doc_comments(&method.attrs);
+
     Superstate {
         handler_name,
         superstate,
@@ -677,6 +686,7 @@ pub fn analyze_superstate(method: &ImplItemFn, state_machine: &StateMachine) ->
         event_arg,
         context_arg,
         is_async,
+        doc_comments,
     }
 }
 
@@ -744,6 +754,15 @@ pub fn get_meta(attrs: &[Attribute], name: &str) -> Vec<Meta> {
         .collect()
 }
 
+/// Extract doc comments from function attributes.
+fn extract_doc_comments(attrs: &[Attribute]) -> Vec<Attribute> {
+    attrs
+        .iter()
+        .filter(|attr| attr.path().is_ident("doc"))
+        .cloned()
+        .collect()
+}
+
 /// Get the ident of the shared storage type.
 pub fn get_shared_storage_path(ty: &Type) -> Path {
     match ty {
@@ -859,6 +878,7 @@ fn valid_state_analyze() {
         }),
         context_arg: None,
         is_async: false,
+        doc_comments: vec![],
     };
 
     let superstate = Superstate {
@@ -877,6 +897,7 @@ fn valid_state_analyze() {
         }),
         context_arg: None,
         is_async: false,
+        doc_comments: vec![],
     };
 
     let entry_action = Action {

macro/src/codegen.rs

@@ -206,6 +206,7 @@ fn codegen_state(ir: &Ir) -> Option<ItemEnum> {
     match &ir.state_machine.state_ident {
         crate::lower::StateIdent::CustomState(_ident) => None,
         crate::lower::StateIdent::StatigState(state_ident) => Some(parse_quote!(
+            /// State enum generated by Statig.
             #[derive(#(#state_derives),*)]
             # visibility enum #state_ident #state_generics {
                 #(#variants),*
@@ -382,6 +383,7 @@ fn codegen_superstate(ir: &Ir) -> ItemEnum {
     let visibility = &ir.state_machine.visibility;
 
     parse_quote!(
+        /// Superstate enum generated by Statig.
         #[derive(#(#superstate_derives),*)]
         #visibility enum #superstate_ident #superstate_generics {
             #(#variants),*

macro/src/lower.rs

@@ -535,7 +535,10 @@ pub fn lower_state(state: &analyze::State, state_machine: &analyze::StateMachine
 
     let handler_inputs: Vec<Ident> = state.inputs.iter().map(fn_arg_to_ident).collect();
 
-    let variant = parse_quote!(#variant_name { #(#variant_fields),* });
+    let mut variant: Variant = parse_quote!(#variant_name { #(#variant_fields),* });
+    // Attach doc comments from the state handler to the enum variant
+    variant.attrs = state.doc_comments.clone();
+
     let pat = parse_quote!(#state_name::#variant_name { #(#pat_fields),*});
     let constructor = parse_quote!(fn #state_handler_name ( #(#constructor_args),* ) -> Self { Self::#variant_name { #(#field_values),*} });
 
@@ -597,7 +600,10 @@ pub fn lower_superstate(
         .collect();
     let handler_inputs: Vec<Ident> = superstate.inputs.iter().map(fn_arg_to_ident).collect();
 
-    let variant = parse_quote!(#superstate_name { #(#variant_fields),* });
+    let mut variant: Variant = parse_quote!(#superstate_name { #(#variant_fields),* });
+    // Attach doc comments from the superstate handler to the enum variant
+    variant.attrs = superstate.doc_comments.clone();
+
     let pat = parse_quote!(#superstate_type::#superstate_name { #(#pat_fields),*});
 
     let handler_call = match &superstate.is_async {
@@ -865,6 +871,7 @@ fn create_analyze_state() -> analyze::State {
             },
         ],
         is_async: false,
+        doc_comments: vec![],
     }
 }
 
@@ -932,6 +939,7 @@ fn create_analyze_superstate() -> analyze::Superstate {
         ],
         is_async: false,
         initial_state: None,
+        doc_comments: vec![],
     }
 }
 

statig/tests/ui/doc_comments.rs

@@ -0,0 +1,60 @@
+#![deny(missing_docs)]
+//! Test doc comment propagation to generated State and Superstate enums
+
+use statig::prelude::*;
+
+/// Machine demonstrating doc comment propagation
+#[derive(Default)]
+pub struct DocMachine {}
+
+/// Test state machine with various doc comment scenarios
+#[state_machine(initial = "State::simple_state()")]
+impl DocMachine {
+    /// Simple state with basic documentation.
+    #[state(superstate = "documented_superstate")]
+    fn simple_state(&mut self) -> Outcome<State> {
+        Transition(State::multi_line_state())
+    }
+
+    /// Multi-line state with comprehensive documentation.
+    ///
+    /// # Purpose
+    /// This state demonstrates multi-line doc comments
+    /// with various formatting including:
+    ///
+    /// - Bullet points
+    /// - Headers
+    /// - Multiple paragraphs
+    ///
+    /// All of this documentation should be preserved
+    /// on the generated State::MultiLineState variant.
+    #[state(superstate = "documented_superstate")]
+    fn multi_line_state(&mut self) -> Outcome<State> {
+        Transition(State::standalone_state())
+    }
+
+    /// Standalone state without superstate grouping.
+    ///
+    /// This tests doc comment propagation for states that
+    /// are not part of any superstate hierarchy.
+    #[state]
+    fn standalone_state(&mut self) -> Outcome<State> {
+        Transition(State::simple_state())
+    }
+
+    /// Documented superstate containing multiple states.
+    ///
+    /// This superstate groups together all the documented states
+    /// and demonstrates that superstate doc comments are also
+    /// properly propagated to the generated Superstate enum.
+    #[superstate]
+    fn documented_superstate(&mut self) -> Outcome<State> {
+        Super
+    }
+}
+
+/// Main function for the test
+fn main() {
+    let machine = DocMachine::default();
+    let _state_machine = machine.uninitialized_state_machine().init();
+}

statig/tests/ui/doc_comments_missing_lint.rs

@@ -0,0 +1,45 @@
+#![deny(missing_docs)]
+//! Critical validation test: This MUST fail to prove our doc comment tests work.
+//!
+//! This test intentionally has missing doc comments on a state handler.
+//! The #![deny(missing_docs)] lint should catch this and cause compilation to fail.
+//!
+//! If this test passes, it means either:
+//! 1. Our doc comment implementation is broken, OR  
+//! 2. The missing_docs lint isn't working
+//!
+//! Either case would invalidate our positive doc comment tests.
+
+use statig::prelude::*;
+
+/// Test machine for missing docs validation
+#[derive(Default)]
+pub struct ValidationMachine {}
+
+/// State machine to validate missing docs detection
+#[state_machine(initial = "State::documented()")]
+impl ValidationMachine {
+    /// This state has documentation.
+    #[state]
+    fn documented(&mut self) -> Outcome<State> {
+        Transition(State::undocumented())
+    }
+
+    // INTENTIONALLY MISSING DOC COMMENT - this should cause compile failure
+    #[state]
+    fn undocumented(&mut self) -> Outcome<State> {
+        Transition(State::documented())
+    }
+
+    /// This superstate has doc comments.
+    #[superstate]
+    fn container(&mut self) -> Outcome<State> {
+        Super
+    }
+}
+
+/// Main function for the test
+fn main() {
+    let machine = ValidationMachine::default();
+    let _state_machine = machine.uninitialized_state_machine().init();
+}

statig/tests/ui/doc_comments_missing_lint.stderr

@@ -0,0 +1,12 @@
+error: missing documentation for a variant
+  --> tests/ui/doc_comments_missing_lint.rs:20:1
+   |
+20 | #[state_machine(initial = "State::documented()")]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |
+note: the lint level is defined here
+  --> tests/ui/doc_comments_missing_lint.rs:1:9
+   |
+1  | #![deny(missing_docs)]
+   |         ^^^^^^^^^^^^
+   = note: this error originates in the attribute macro `state_machine` (in Nightly builds, run with -Z macro-backtrace for more info)

statig/tests/ui_tests.rs

@@ -3,4 +3,8 @@ fn ui() {
     let t = trybuild::TestCases::new();
     t.pass("tests/ui/custom_state.rs");
     t.compile_fail("tests/ui/custom_state_derive_error.rs");
+
+    // Doc comment propagation tests
+    t.pass("tests/ui/doc_comments.rs");
+    t.compile_fail("tests/ui/doc_comments_missing_lint.rs");
 }