feat: Propagate rustdoc to generated (super)state enums

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");
 }