From f97d785e7a45dc8ed42246e810ad6b3336a8ea81 Mon Sep 17 00:00:00 2001
From: Yevhenii Babichenko <eugene.babichenko@protonmail.com>
Date: Mon, 13 May 2024 02:23:40 +0300
Subject: [PATCH] optional Mermaid diagrams

---
 CHANGELOG.md            |  1 +
 README.md               |  8 ++++++++
 rust-fsm-dsl/Cargo.toml |  3 +++
 rust-fsm-dsl/src/lib.rs | 32 ++++++++++++++++++++++++++++++++
 rust-fsm/Cargo.toml     |  2 ++
 rust-fsm/src/lib.rs     |  3 +++
 6 files changed, 49 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc1fa4d..9a2d822 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@ adheres to [Semantic Versioning][semver].
 * A type alias `StateMachine` for `rust_fsm::StateMachine<Impl>` is now
   generated inside the said module.
 * Supplying ones own enums for state, input and output in the proc-macro (#10).
+* An optional possibility to generate Mermaid diagrams.
 
 ## [0.6.2] - 2024-05-11
 ### Changed
diff --git a/README.md b/README.md
index 52035d4..2785b4b 100644
--- a/README.md
+++ b/README.md
@@ -179,6 +179,13 @@ state_machine! {
 }
 ```
 
+#### Diagrams
+
+`state_machine` macro can document your state machines with diagrams. This is
+controlled by the `diagram` feature, which is non-default. The diagrams are
+generated in the [Mermaid][mermaid] format. This feature includes the Mermaid
+script into the documentation page.
+
 ### Without DSL
 
 The `state_machine` macro has limited capabilities (for example, a state
@@ -197,3 +204,4 @@ You can see an example of the Circuit Breaker state machine in the
 [docs-link]: https://docs.rs/rust-fsm
 [crate-badge]: https://img.shields.io/crates/v/rust-fsm.svg
 [crate-link]: https://crates.io/crates/rust-fsm
+[mermaid]: https://mermaid.js.org/
diff --git a/rust-fsm-dsl/Cargo.toml b/rust-fsm-dsl/Cargo.toml
index 82f3b9f..8129905 100644
--- a/rust-fsm-dsl/Cargo.toml
+++ b/rust-fsm-dsl/Cargo.toml
@@ -15,6 +15,9 @@ edition = "2018"
 [lib]
 proc-macro = true
 
+[features]
+diagram = []
+
 [dependencies]
 proc-macro2 = "1"
 syn = "2"
diff --git a/rust-fsm-dsl/src/lib.rs b/rust-fsm-dsl/src/lib.rs
index 5b13e69..ab29e68 100644
--- a/rust-fsm-dsl/src/lib.rs
+++ b/rust-fsm-dsl/src/lib.rs
@@ -57,6 +57,12 @@ pub fn state_machine(tokens: TokenStream) -> TokenStream {
     let mut transition_cases = Vec::new();
     let mut output_cases = Vec::new();
 
+    #[cfg(feature = "diagram")]
+    let mut mermaid_diagram = format!(
+        "///```mermaid\n///stateDiagram-v2\n///    [*] --> {}\n",
+        input.initial_state
+    );
+
     states.insert(&input.initial_state);
 
     for transition in transitions {
@@ -67,6 +73,11 @@ pub fn state_machine(tokens: TokenStream) -> TokenStream {
             output,
         } = transition;
 
+        #[cfg(feature = "diagram")]
+        mermaid_diagram.push_str(&format!(
+            "///    {initial_state} --> {final_state}: {input_value}"
+        ));
+
         transition_cases.push(quote! {
             (Self::State::#initial_state, Self::Input::#input_value) => {
                 Some(Self::State::#final_state)
@@ -79,8 +90,14 @@ pub fn state_machine(tokens: TokenStream) -> TokenStream {
                     Some(Self::Output::#output_value)
                 }
             });
+
+            #[cfg(feature = "diagram")]
+            mermaid_diagram.push_str(&format!(" [{output_value}]"));
         }
 
+        #[cfg(feature = "diagram")]
+        mermaid_diagram.push('\n');
+
         states.insert(initial_state);
         states.insert(final_state);
         inputs.insert(input_value);
@@ -89,6 +106,11 @@ pub fn state_machine(tokens: TokenStream) -> TokenStream {
         }
     }
 
+    #[cfg(feature = "diagram")]
+    mermaid_diagram.push_str("///```");
+    #[cfg(feature = "diagram")]
+    let mermaid_diagram: proc_macro2::TokenStream = mermaid_diagram.parse().unwrap();
+
     let initial_state_name = &input.initial_state;
 
     let (input_type, input_impl) = match input.input_type {
@@ -139,7 +161,17 @@ pub fn state_machine(tokens: TokenStream) -> TokenStream {
         }
     };
 
+    #[cfg(feature = "diagram")]
+    let diagram = quote! {
+        #[cfg_attr(doc, rust_fsm::aquamarine)]
+        #mermaid_diagram
+    };
+
+    #[cfg(not(feature = "diagram"))]
+    let diagram = quote!();
+
     let output = quote! {
+        #diagram
         #visibility mod #fsm_name {
             #attrs
             pub struct Impl;
diff --git a/rust-fsm/Cargo.toml b/rust-fsm/Cargo.toml
index 34a100d..1a6962e 100644
--- a/rust-fsm/Cargo.toml
+++ b/rust-fsm/Cargo.toml
@@ -16,8 +16,10 @@ edition = "2018"
 default = ["std", "dsl"]
 std = []
 dsl = ["rust-fsm-dsl"]
+diagram = ["aquamarine", "rust-fsm-dsl/diagram"]
 
 [dependencies]
+aquamarine = { version = "*", optional = true }
 rust-fsm-dsl = { path = "../rust-fsm-dsl", version = "0.6.2", optional = true }
 
 [profile.dev]
diff --git a/rust-fsm/src/lib.rs b/rust-fsm/src/lib.rs
index 95a5ea9..2decd67 100644
--- a/rust-fsm/src/lib.rs
+++ b/rust-fsm/src/lib.rs
@@ -8,6 +8,9 @@ use std::error::Error;
 #[cfg(feature = "dsl")]
 pub use rust_fsm_dsl::state_machine;
 
+#[cfg(feature = "diagram")]
+pub use aquamarine::aquamarine;
+
 /// This trait is designed to describe any possible deterministic finite state
 /// machine/transducer. This is just a formal definition that may be
 /// inconvenient to be used in practical programming, but it is used throughout