Add completion system with shell context and dispatcher integration

6 files changed, 155 insertions, 122 deletions

diff --git a/mingling_macros/src/README.txt b/mingling_macros/src/README.txt
index a882bd0..e653795 100644
--- a/mingling_macros/src/README.txt
+++ b/mingling_macros/src/README.txt
@@ -1,12 +1,7 @@
-*Listen up*, this is a messy macro written by AI.
-It works, but it looks like expired spaghetti and smells like freshly laid shit.
+接下来的优化计划
 
-But the general idea is mostly right. If I'm in a good mood, this code will definitely need a refactor.
+1. 将全局注册包装成 registry.rs
+2. 声明类似 register_chain_match ... 等函数
+...
 
-----------------------------------------------------
-Chinese original version:
-
-*听着*，这是一坨由 AI 写的乱七八糟的宏
-它能用，但是看起来就像过期的意大利面，闻起来就像刚拉出来的新鲜的屎
-
-不过思路大抵是对的，如果心情好这坨代码肯定要重构
+还在思考
diff --git a/mingling_macros/src/chain.rs b/mingling_macros/src/chain.rs
index 924931e..fbf6a6b 100644
--- a/mingling_macros/src/chain.rs
+++ b/mingling_macros/src/chain.rs
@@ -176,13 +176,8 @@ pub fn chain_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
     };
 
     // Record the chain mapping
-    let chain_entry = quote! {
-        #struct_name => #previous_type,
-    };
-
-    let chain_exist_entry = quote! {
-        Self::#previous_type => true,
-    };
+    let chain_entry = build_chain_arm(&struct_name, &previous_type);
+    let chain_exist_entry = build_chain_exist_arm(&previous_type);
 
     let mut chains = crate::CHAINS.lock().unwrap();
     let mut chain_exist = crate::CHAINS_EXIST.lock().unwrap();
@@ -192,17 +187,23 @@ pub fn chain_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
     let chain_exist_entry = chain_exist_entry.to_string();
     let previous_type_str = previous_type.to_token_stream().to_string();
 
-    if !chains.contains(&chain_entry) {
-        chains.push(chain_entry);
-    }
+    chains.insert(chain_entry);
+    chain_exist.insert(chain_exist_entry);
+    packed_types.insert(previous_type_str);
 
-    if !chain_exist.contains(&chain_exist_entry) {
-        chain_exist.push(chain_exist_entry);
-    }
+    expanded.into()
+}
 
-    if !packed_types.contains(&previous_type_str) {
-        packed_types.push(previous_type_str);
+/// Builds a match arm for chain mapping
+pub fn build_chain_arm(struct_name: &Ident, previous_type: &TypePath) -> proc_macro2::TokenStream {
+    quote! {
+        #struct_name => #previous_type,
     }
+}
 
-    expanded.into()
+/// Builds a match arm for chain existence check
+pub fn build_chain_exist_arm(previous_type: &TypePath) -> proc_macro2::TokenStream {
+    quote! {
+        Self::#previous_type => true,
+    }
 }
diff --git a/mingling_macros/src/completion.rs b/mingling_macros/src/completion.rs
index 23b509f..334affd 100644
--- a/mingling_macros/src/completion.rs
+++ b/mingling_macros/src/completion.rs
@@ -93,9 +93,7 @@ pub fn completion_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
 
     let mut completions = crate::COMPLETIONS.lock().unwrap();
     let completion_str = completion_entry.to_string();
-    if !completions.contains(&completion_str) {
-        completions.push(completion_str);
-    }
+    completions.insert(completion_str);
 
     expanded.into()
 }
diff --git a/mingling_macros/src/lib.rs b/mingling_macros/src/lib.rs
index 7291f0e..f97c458 100644
--- a/mingling_macros/src/lib.rs
+++ b/mingling_macros/src/lib.rs
@@ -3,9 +3,12 @@
 //! This crate provides procedural macros for the Mingling framework.
 //! Macros are implemented in separate modules and re-exported here.
 
+use once_cell::sync::Lazy;
 use proc_macro::TokenStream;
 use proc_macro2::Ident;
 use quote::quote;
+use std::collections::BTreeSet;
+use std::sync::Mutex;
 use syn::parse_macro_input;
 
 mod chain;
@@ -21,21 +24,23 @@ mod renderer;
 #[cfg(feature = "comp")]
 mod suggest;
 
-use once_cell::sync::Lazy;
-use std::sync::Mutex;
-
 // Global variables
 #[cfg(feature = "general_renderer")]
-pub(crate) static GENERAL_RENDERERS: Lazy<Mutex<Vec<String>>> =
-    Lazy::new(|| Mutex::new(Vec::new()));
+pub(crate) static GENERAL_RENDERERS: Lazy<Mutex<BTreeSet<String>>> =
+    Lazy::new(|| Mutex::new(BTreeSet::new()));
 #[cfg(feature = "comp")]
-pub(crate) static COMPLETIONS: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
-
-pub(crate) static PACKED_TYPES: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
-pub(crate) static CHAINS: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
-pub(crate) static RENDERERS: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
-pub(crate) static CHAINS_EXIST: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
-pub(crate) static RENDERERS_EXIST: Lazy<Mutex<Vec<String>>> = Lazy::new(|| Mutex::new(Vec::new()));
+pub(crate) static COMPLETIONS: Lazy<Mutex<BTreeSet<String>>> =
+    Lazy::new(|| Mutex::new(BTreeSet::new()));
+
+pub(crate) static PACKED_TYPES: Lazy<Mutex<BTreeSet<String>>> =
+    Lazy::new(|| Mutex::new(BTreeSet::new()));
+pub(crate) static CHAINS: Lazy<Mutex<BTreeSet<String>>> = Lazy::new(|| Mutex::new(BTreeSet::new()));
+pub(crate) static RENDERERS: Lazy<Mutex<BTreeSet<String>>> =
+    Lazy::new(|| Mutex::new(BTreeSet::new()));
+pub(crate) static CHAINS_EXIST: Lazy<Mutex<BTreeSet<String>>> =
+    Lazy::new(|| Mutex::new(BTreeSet::new()));
+pub(crate) static RENDERERS_EXIST: Lazy<Mutex<BTreeSet<String>>> =
+    Lazy::new(|| Mutex::new(BTreeSet::new()));
 
 #[proc_macro]
 pub fn node(input: TokenStream) -> TokenStream {
@@ -101,24 +106,69 @@ pub fn derive_groupped_serialize(input: TokenStream) -> TokenStream {
 
 #[proc_macro]
 pub fn gen_program(input: TokenStream) -> TokenStream {
-    let name = if input.is_empty() {
-        Ident::new("ThisProgram", proc_macro2::Span::call_site())
-    } else {
-        parse_macro_input!(input as Ident)
+    let name = read_name(&input);
+
+    #[cfg(feature = "comp")]
+    let out = TokenStream::from(quote! {
+        ::mingling::macros::program_gen_completion!(#name);
+        ::mingling::macros::program_final_gen!(#name);
+    });
+    #[cfg(not(feature = "comp"))]
+    let out = TokenStream::from(quote! {
+        ::mingling::macros::program_final_gen!(#name);
+    });
+
+    out
+}
+
+#[proc_macro]
+#[cfg(feature = "comp")]
+pub fn program_gen_completion(input: TokenStream) -> TokenStream {
+    let name = read_name(&input);
+
+    let comp_dispatcher = quote! {
+        #[allow(unused)]
+        use __completion_gen::*;
+        pub mod __completion_gen {
+            use super::*;
+            use mingling::marker::NextProcess;
+            ::mingling::macros::dispatcher!(#name, "__comp", CompletionDispatcher => CompletionContext);
+            ::mingling::macros::pack!(
+                #name,
+                CompletionSuggest = (::mingling::ShellContext, ::mingling::Suggest)
+            );
+
+            #[::mingling::macros::chain(#name)]
+            pub async fn __exec_completion(prev: CompletionContext) -> NextProcess {
+                let read_ctx = ::mingling::ShellContext::try_from(prev.inner);
+                match read_ctx {
+                    Ok(ctx) => {
+                        let suggest = ::mingling::CompletionHelper::exec_completion::<#name>(&ctx);
+                        CompletionSuggest::new((ctx, suggest)).to_render()
+                    }
+                    Err(_) => std::process::exit(1),
+                }
+            }
+
+            #[::mingling::macros::renderer(#name)]
+            pub fn __render_completion(prev: CompletionSuggest) {
+                let (ctx, suggest) = prev.inner;
+                ::mingling::CompletionHelper::render_suggest::<#name>(ctx, suggest);
+            }
+        }
     };
 
-    let mut packed_types = PACKED_TYPES.lock().unwrap().clone();
-    packed_types.push("DispatcherNotFound".to_string());
-    packed_types.push("RendererNotFound".to_string());
+    TokenStream::from(comp_dispatcher)
+}
 
-    #[cfg(feature = "comp")]
-    {
-        packed_types.push("CompletionContext".to_string());
-        packed_types.push("CompletionSuggest".to_string());
-    }
+#[proc_macro]
+pub fn program_final_gen(input: TokenStream) -> TokenStream {
+    let name = read_name(&input);
+
+    let mut packed_types = PACKED_TYPES.lock().unwrap().clone();
+    packed_types.insert("DispatcherNotFound".to_string());
+    packed_types.insert("RendererNotFound".to_string());
 
-    packed_types.sort();
-    packed_types.dedup();
     let renderers = RENDERERS.lock().unwrap().clone();
     let chains = CHAINS.lock().unwrap().clone();
     let renderer_exist = RENDERERS_EXIST.lock().unwrap().clone();
@@ -178,36 +228,6 @@ pub fn gen_program(input: TokenStream) -> TokenStream {
     let general_render = quote! {};
 
     #[cfg(feature = "comp")]
-    let comp_dispatcher = quote! {
-        ::mingling::macros::dispatcher!(#name, "__comp", CompletionDispatcher => CompletionContext);
-        ::mingling::macros::pack!(
-            #name,
-            CompletionSuggest = (::mingling::ShellContext, ::mingling::Suggest)
-        );
-
-        #[::mingling::macros::chain]
-        async fn __completion(prev: CompletionContext) -> NextProcess {
-            let read_ctx = ::mingling::ShellContext::try_from(prev.inner);
-            match read_ctx {
-                Ok(ctx) => {
-                    let suggest = ::mingling::CompletionHelper::exec_completion::<#name>(&ctx);
-                    CompletionSuggest::new((ctx, suggest)).to_render()
-                }
-                Err(_) => std::process::exit(1),
-            }
-        }
-
-        #[::mingling::macros::renderer]
-        fn __render_completion(prev: CompletionSuggest) {
-            let (ctx, suggest) = prev.inner;
-            ::mingling::CompletionHelper::render_suggest::<#name>(ctx, suggest);
-        }
-    };
-
-    #[cfg(not(feature = "comp"))]
-    let comp_dispatcher = quote! {};
-
-    #[cfg(feature = "comp")]
     let completion_tokens: Vec<proc_macro2::TokenStream> = completions
         .iter()
         .map(|s| syn::parse_str::<proc_macro2::TokenStream>(s).unwrap())
@@ -238,8 +258,6 @@ pub fn gen_program(input: TokenStream) -> TokenStream {
             #(#packed_types),*
         }
 
-        #comp_dispatcher
-
         impl ::std::fmt::Display for #name {
             fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                 match self {
@@ -299,7 +317,7 @@ pub fn __register_chain(input: TokenStream) -> TokenStream {
     let chain_entry = parse_macro_input!(input as syn::LitStr);
     let entry_str = chain_entry.value();
 
-    CHAINS.lock().unwrap().push(entry_str);
+    CHAINS.lock().unwrap().insert(entry_str);
 
     TokenStream::new()
 }
@@ -314,7 +332,7 @@ pub fn __register_renderer(input: TokenStream) -> TokenStream {
     let renderer_entry = parse_macro_input!(input as syn::LitStr);
     let entry_str = renderer_entry.value();
 
-    RENDERERS.lock().unwrap().push(entry_str);
+    RENDERERS.lock().unwrap().insert(entry_str);
 
     TokenStream::new()
 }
@@ -324,3 +342,11 @@ pub fn __register_renderer(input: TokenStream) -> TokenStream {
 pub fn suggest(input: TokenStream) -> TokenStream {
     suggest::suggest(input)
 }
+
+fn read_name(input: &TokenStream) -> Ident {
+    if input.is_empty() {
+        Ident::new("ThisProgram", proc_macro2::Span::call_site())
+    } else {
+        syn::parse(input.clone()).unwrap()
+    }
+}
diff --git a/mingling_macros/src/node.rs b/mingling_macros/src/node.rs
index 8cb88d3..be47374 100644
--- a/mingling_macros/src/node.rs
+++ b/mingling_macros/src/node.rs
@@ -39,7 +39,13 @@ pub fn node(input: TokenStream) -> TokenStream {
     // Split the path by dots
     let parts: Vec<String> = path_str
         .split('.')
-        .map(|s| kebab_case!(s).to_string())
+        .map(|s| {
+            if s.starts_with('_') {
+                s.to_string()
+            } else {
+                kebab_case!(s).to_string()
+            }
+        })
         .collect();
 
     // Build the expression starting from Node::default()
diff --git a/mingling_macros/src/renderer.rs b/mingling_macros/src/renderer.rs
index f700cbf..61c67b2 100644
--- a/mingling_macros/src/renderer.rs
+++ b/mingling_macros/src/renderer.rs
@@ -101,25 +101,10 @@ pub fn renderer_attr(item: TokenStream) -> TokenStream {
     let struct_name = syn::Ident::new(&pascal_case_name, fn_name.span());
 
     // Register the renderer in the global list
-    let renderer_entry = quote! {
-        #struct_name => #previous_type,
-    };
-
-    let renderer_exist_entry = quote! {
-        Self::#previous_type => true,
-    };
-
+    let renderer_entry = build_renderer_entry(&struct_name, &previous_type);
+    let renderer_exist_entry = build_renderer_exist_entry(&previous_type);
     #[cfg(feature = "general_renderer")]
-    let general_renderer_entry = quote! {
-        Self::#previous_type => {
-            // SAFETY: Only types that match will enter this branch for forced conversion,
-            // and `AnyOutput::new` ensures the type implements serde::Serialize
-            let raw = unsafe { any.restore::<#previous_type>().unwrap_unchecked() };
-            let mut r = ::mingling::RenderResult::default();
-            ::mingling::GeneralRenderer::render(&raw, setting, &mut r)?;
-            Ok(r)
-        }
-    };
+    let general_renderer_entry = build_general_renderer_entry(&previous_type);
 
     let mut renderers = crate::RENDERERS.lock().unwrap();
     let mut renderer_exist = crate::RENDERERS_EXIST.lock().unwrap();
@@ -135,22 +120,12 @@ pub fn renderer_attr(item: TokenStream) -> TokenStream {
     #[cfg(feature = "general_renderer")]
     let general_renderer_entry_str = general_renderer_entry.to_string();
 
-    if !renderers.contains(&renderer_entry_str) {
-        renderers.push(renderer_entry_str);
-    }
-
-    if !renderer_exist.contains(&renderer_exist_entry_str) {
-        renderer_exist.push(renderer_exist_entry_str);
-    }
-
-    if !packed_types.contains(&previous_type_str) {
-        packed_types.push(previous_type_str);
-    }
+    renderers.insert(renderer_entry_str);
+    renderer_exist.insert(renderer_exist_entry_str);
+    packed_types.insert(previous_type_str);
 
     #[cfg(feature = "general_renderer")]
-    if !general_renderers.contains(&general_renderer_entry_str) {
-        general_renderers.push(general_renderer_entry_str);
-    }
+    general_renderers.insert(general_renderer_entry_str);
 
     // Generate the struct and implementation
     // We need to create a wrapper function that adds the r parameter
@@ -185,3 +160,35 @@ pub fn renderer_attr(item: TokenStream) -> TokenStream {
 
     expanded.into()
 }
+
+/// Builds the renderer entry for the global renderers list
+pub fn build_renderer_entry(
+    struct_name: &syn::Ident,
+    previous_type: &TypePath,
+) -> proc_macro2::TokenStream {
+    quote! {
+        #struct_name => #previous_type,
+    }
+}
+
+/// Builds the renderer existence check entry
+pub fn build_renderer_exist_entry(previous_type: &TypePath) -> proc_macro2::TokenStream {
+    quote! {
+        Self::#previous_type => true,
+    }
+}
+
+/// Builds the general renderer entry
+#[cfg(feature = "general_renderer")]
+pub fn build_general_renderer_entry(previous_type: &TypePath) -> proc_macro2::TokenStream {
+    quote! {
+        Self::#previous_type => {
+            // SAFETY: Only types that match will enter this branch for forced conversion,
+            // and `AnyOutput::new` ensures the type implements serde::Serialize
+            let raw = unsafe { any.restore::<#previous_type>().unwrap_unchecked() };
+            let mut r = ::mingling::RenderResult::default();
+            ::mingling::GeneralRenderer::render(&raw, setting, &mut r)?;
+            Ok(r)
+        }
+    }
+}