Add compile-time dispatch tree generation for O(len) command routing

3 files changed, 346 insertions, 2 deletions

diff --git a/mingling_macros/src/dispatch_tree_gen.rs b/mingling_macros/src/dispatch_tree_gen.rs
new file mode 100644
index 0000000..840fe15
--- /dev/null
+++ b/mingling_macros/src/dispatch_tree_gen.rs
@@ -0,0 +1,195 @@
+//! Dispatch Tree Generation
+//!
+//! This module generates the dispatch tree code for the `dispatch_tree` feature.
+//! It builds a compact, hardcoded match tree at compile time to achieve O(len)
+//! command dispatch.
+//!
+//! # Algorithm
+//!
+//! For each depth, group nodes by the character at that depth.
+//! - If a group has only one node: emit `starts_with` check for the full name.
+//! - If a group has multiple nodes: emit a `match raw_chars.nth(depth)` arm and recurse.
+//! - At the leaf: call `Dispatcher::begin` on the matched dispatcher.
+//!
+//! Names are matched with a trailing space (e.g. "hello ") to ensure exact boundary.
+
+use proc_macro2::TokenStream;
+use quote::quote;
+use std::collections::BTreeMap;
+
+/// Generate the `get_nodes()` function body for a ProgramCollect impl.
+pub fn gen_get_nodes(entries: &[(String, String, String)]) -> TokenStream {
+    let mut node_entries = Vec::new();
+
+    for (node_name, _disp_type, _entry_name) in entries {
+        let static_name_str = format!("__internal_dispatcher_{}", node_name.replace('.', "_"));
+        let static_ident = syn::Ident::new(&static_name_str, proc_macro2::Span::call_site());
+
+        let node_display_name = node_name.replace('.', " ");
+        let node_display_lit = syn::LitStr::new(&node_display_name, proc_macro2::Span::call_site());
+
+        node_entries.push(quote! {
+            (#node_display_lit.to_string(), & #static_ident)
+        });
+    }
+
+    quote! {
+        fn get_nodes() -> Vec<(String, &'static (dyn ::mingling::Dispatcher<Self::Enum> + Send + Sync))> {
+            vec![
+                #(#node_entries),*
+            ]
+        }
+    }
+}
+
+/// Generate the `dispatch_args_trie()` function body for a ProgramCollect impl.
+///
+/// Builds a hardcoded match tree: at each depth, group nodes by character.
+/// Single-node groups use `starts_with`; multi-node groups recurse with `nth()` match.
+pub fn gen_dispatch_args_trie(entries: &[(String, String, String)]) -> TokenStream {
+    // Prepare (display_name, disp_type) pairs.
+    // display_name = node_name.replace('.', " ")
+    let nodes: Vec<(String, String)> = entries
+        .iter()
+        .map(|(name, disp, _)| (name.replace('.', " "), disp.clone()))
+        .collect();
+
+    let dispatch_body = build_dispatch_body(&nodes, 0);
+
+    quote! {
+        fn dispatch_args_trie(
+            raw: &Vec<String>,
+        ) -> Result<::mingling::AnyOutput<Self::Enum>, ::mingling::error::ProgramInternalExecuteError>
+        {
+            let raw_string = format!("{} ", raw.join(" "));
+            let raw_str = raw_string.as_str();
+            let mut raw_chars = raw_str.chars();
+            #dispatch_body
+        }
+    }
+}
+
+/// Recursively build the trie match body.
+///
+/// `nodes`: slice of (display_name, disp_type) for commands that share the same prefix so far.
+/// `depth`: The character index currently being matched.
+///
+/// Returns a `TokenStream` representing the match block at this depth.
+fn build_dispatch_body(nodes: &[(String, String)], depth: usize) -> TokenStream {
+    if nodes.is_empty() {
+        return quote! {
+            return Ok(Self::build_dispatcher_not_found(raw.clone()));
+        };
+    }
+
+    // Group by character at `depth`
+    // Nodes that end exactly at this depth (name is closed – rare but possible, e.g. "hell")
+    // are collected into `exact_nodes`. All others go into `groups[char]`.
+    let mut groups: BTreeMap<char, Vec<(String, String)>> = BTreeMap::new();
+    let mut exact_nodes: Vec<(String, String)> = Vec::new();
+
+    for (name, disp_type) in nodes {
+        if let Some(ch) = name.chars().nth(depth) {
+            groups
+                .entry(ch)
+                .or_default()
+                .push((name.clone(), disp_type.clone()));
+        } else {
+            exact_nodes.push((name.clone(), disp_type.clone()));
+        }
+    }
+
+    // Build a dispatch arm for a single node via `starts_with`
+    let make_starts_with_arm = |name: &str, disp_type: &str| -> TokenStream {
+        let name_space = format!("{} ", name);
+        let name_lit = syn::LitStr::new(&name_space, proc_macro2::Span::call_site());
+        let disp_ident = syn::Ident::new(disp_type, proc_macro2::Span::call_site());
+        quote! {
+            if let Some(stripped) = raw_str.strip_prefix(#name_lit) {
+                let __cp = <#disp_ident as ::mingling::Dispatcher<Self::Enum>>::begin(
+                    &#disp_ident::default(),
+                    stripped
+                        .split_whitespace()
+                        .map(String::from)
+                        .collect::<Vec<String>>(),
+                );
+                return match __cp {
+                    ::mingling::ChainProcess::Ok(any_output) => Ok(any_output.0),
+                    ::mingling::ChainProcess::Err(chain_process_error) => {
+                        Err(chain_process_error.into())
+                    }
+                };
+            }
+        }
+    };
+
+    // Build match arms
+    let mut arms = Vec::new();
+
+    for (&ch, sub_nodes) in &groups {
+        let ch_char = ch;
+
+        if sub_nodes.len() == 1 {
+            // Only one candidate – use `starts_with` directly at this depth.
+            let (name, disp_type) = &sub_nodes[0];
+            let arm = make_starts_with_arm(name, disp_type);
+            arms.push(quote! {
+                Some(#ch_char) => {
+                    #arm
+                    return Ok(Self::build_dispatcher_not_found(raw.clone()));
+                }
+            });
+        } else {
+            // Multiple candidates – recurse deeper
+            let sub_body = build_dispatch_body(sub_nodes, depth + 1);
+            arms.push(quote! {
+                Some(#ch_char) => {
+                    #sub_body
+                }
+            });
+        }
+    }
+
+    // Exact-match nodes at this depth
+    // These are names that are a prefix of other names (e.g. "hell" when "hello" exists).
+    // They are tried first via `starts_with`, then fall through to `raw_chars.nth(depth)` match.
+    let exact_checks: Vec<TokenStream> = exact_nodes
+        .iter()
+        .map(|(name, disp_type)| make_starts_with_arm(name, disp_type))
+        .collect();
+
+    // Assemble
+    // If there are exact nodes, first check starts_with, then do matching.
+    if !exact_checks.is_empty() && !groups.is_empty() {
+        // Exact nodes + deeper groups: do starts_with checks, then match on nth(depth)
+        let match_body = quote! {
+            match raw_chars.nth(0) {
+                #(#arms)*
+                _ => return Ok(Self::build_dispatcher_not_found(raw.clone())),
+            }
+        };
+        quote! {
+            #(#exact_checks)*
+            #match_body
+        }
+    } else if !exact_checks.is_empty() {
+        // Only exact nodes, no deeper groups
+        quote! {
+            #(#exact_checks)*
+            return Ok(Self::build_dispatcher_not_found(raw.clone()));
+        }
+    } else if arms.is_empty() {
+        // Only fallback (shouldn't happen if nodes is non-empty)
+        quote! {
+            return Ok(Self::build_dispatcher_not_found(raw.clone()));
+        }
+    } else {
+        // Only group arms
+        quote! {
+            match raw_chars.nth(0) {
+                #(#arms)*
+                _ => return Ok(Self::build_dispatcher_not_found(raw.clone())),
+            }
+        }
+    }
+}
diff --git a/mingling_macros/src/dispatcher.rs b/mingling_macros/src/dispatcher.rs
index 3ca5d25..f9dc423 100644
--- a/mingling_macros/src/dispatcher.rs
+++ b/mingling_macros/src/dispatcher.rs
@@ -4,10 +4,14 @@
 //! with automatic implementations of the `Dispatcher` trait.
 
 use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
 use quote::quote;
 use syn::parse::{Parse, ParseStream};
 use syn::{Ident, Result as SynResult, Token};
 
+#[cfg(feature = "dispatch_tree")]
+use crate::COMPILE_TIME_DISPATCHERS;
+
 enum DispatcherChainInput {
     Explicit {
         group_name: Ident,
@@ -96,6 +100,8 @@ pub fn dispatcher(input: TokenStream) -> TokenStream {
 
     let comp_entry = get_comp_entry(&pack);
 
+    let dispatch_tree_entry = get_dispatch_tree_entry(&command_name_str, &command_struct, &pack);
+
     let expanded = {
         let program_ident = if use_default {
             Ident::new("ThisProgram", proc_macro2::Span::call_site())
@@ -110,6 +116,7 @@ pub fn dispatcher(input: TokenStream) -> TokenStream {
             ::mingling::macros::pack!(#program_ident, #pack = Vec<String>);
 
             #comp_entry
+            #dispatch_tree_entry
 
             impl ::mingling::Dispatcher<#program_ident> for #command_struct {
                 fn node(&self) -> ::mingling::Node {
@@ -129,7 +136,7 @@ pub fn dispatcher(input: TokenStream) -> TokenStream {
 }
 
 #[cfg(feature = "comp")]
-fn get_comp_entry(entry_name: &Ident) -> proc_macro2::TokenStream {
+fn get_comp_entry(entry_name: &Ident) -> TokenStream2 {
     let comp_entry = quote! {
         impl ::mingling::CompletionEntry for #entry_name {
             fn get_input(self) -> Vec<String> {
@@ -141,6 +148,84 @@ fn get_comp_entry(entry_name: &Ident) -> proc_macro2::TokenStream {
 }
 
 #[cfg(not(feature = "comp"))]
-fn get_comp_entry(_entry_name: &Ident) -> proc_macro2::TokenStream {
+fn get_comp_entry(_entry_name: &Ident) -> TokenStream2 {
     quote! {}
 }
+
+#[cfg(feature = "dispatch_tree")]
+fn get_dispatch_tree_entry(
+    command_name_str: &str,
+    command_struct: &Ident,
+    entry_name: &Ident,
+) -> TokenStream2 {
+    let node_name_lit = syn::LitStr::new(command_name_str, proc_macro2::Span::call_site());
+    quote! {
+        ::mingling::macros::register_dispatcher!(#node_name_lit, #command_struct, #entry_name);
+    }
+}
+
+#[cfg(not(feature = "dispatch_tree"))]
+fn get_dispatch_tree_entry(
+    _command_name_str: &str,
+    _command_struct: &Ident,
+    _entry_name: &Ident,
+) -> TokenStream2 {
+    quote! {}
+}
+
+#[cfg(feature = "dispatch_tree")]
+/// Input format: ("node.name", DispatcherType, EntryName)
+struct RegisterDispatcherInput {
+    node_name: syn::LitStr,
+    dispatcher_type: Ident,
+    entry_name: Ident,
+}
+
+#[cfg(feature = "dispatch_tree")]
+impl Parse for RegisterDispatcherInput {
+    fn parse(input: ParseStream) -> SynResult<Self> {
+        let node_name = input.parse()?;
+        input.parse::<Token![,]>()?;
+        let dispatcher_type = input.parse()?;
+        input.parse::<Token![,]>()?;
+        let entry_name = input.parse()?;
+        Ok(RegisterDispatcherInput {
+            node_name,
+            dispatcher_type,
+            entry_name,
+        })
+    }
+}
+
+#[cfg(feature = "dispatch_tree")]
+pub fn register_dispatcher(input: TokenStream) -> TokenStream {
+    let RegisterDispatcherInput {
+        node_name,
+        dispatcher_type,
+        entry_name,
+    } = syn::parse_macro_input!(input as RegisterDispatcherInput);
+
+    let node_name_str = node_name.value();
+    let static_name = format!("__internal_dispatcher_{}", node_name_str.replace('.', "_"));
+    let static_ident = Ident::new(&static_name, proc_macro2::Span::call_site());
+
+    // Register node info in the global collection at compile time
+    // Format: "node.name:DispatcherType:EntryName"
+    COMPILE_TIME_DISPATCHERS.lock().unwrap().insert(format!(
+        "{}:{}:{}",
+        node_name_str, dispatcher_type, entry_name
+    ));
+
+    let expanded = quote! {
+        #[doc(hidden)]
+        #[allow(nonstandard_style)]
+        static #static_ident: #dispatcher_type = #dispatcher_type;
+    };
+
+    expanded.into()
+}
+
+#[cfg(not(feature = "dispatch_tree"))]
+pub fn register_dispatcher(_input: TokenStream) -> TokenStream {
+    quote! {}.into()
+}
diff --git a/mingling_macros/src/lib.rs b/mingling_macros/src/lib.rs
index 6d38889..f4aa6cf 100644
--- a/mingling_macros/src/lib.rs
+++ b/mingling_macros/src/lib.rs
@@ -28,6 +28,8 @@ use syn::parse_macro_input;
 mod chain;
 #[cfg(feature = "comp")]
 mod completion;
+#[cfg(feature = "dispatch_tree")]
+mod dispatch_tree_gen;
 mod dispatcher;
 #[cfg(feature = "clap")]
 mod dispatcher_clap;
@@ -50,6 +52,10 @@ pub(crate) static GENERAL_RENDERERS: Lazy<Mutex<BTreeSet<String>>> =
 pub(crate) static COMPLETIONS: Lazy<Mutex<BTreeSet<String>>> =
     Lazy::new(|| Mutex::new(BTreeSet::new()));
 
+#[cfg(feature = "dispatch_tree")]
+pub(crate) static COMPILE_TIME_DISPATCHERS: 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()));
@@ -604,6 +610,24 @@ pub fn register_help(input: TokenStream) -> TokenStream {
     help::register_help(input)
 }
 
+/// Registers a dispatcher at compile time for the `dispatch_tree` feature.
+///
+/// This macro is called internally by [`dispatcher!`](macro.dispatcher.html) when
+/// the `dispatch_tree` feature is enabled. It stores the node name into the global
+/// `COMPILE_TIME_DISPATCHERS` registry and generates a static variable for the
+/// dispatcher instance, which is later used by `gen_program!` to generate the
+/// dispatch tree routing logic.
+///
+/// # Syntax
+///
+/// ```rust,ignore
+/// register_dispatcher!("node.name", DispatcherType, EntryName);
+/// ```
+#[proc_macro]
+pub fn register_dispatcher(input: TokenStream) -> TokenStream {
+    dispatcher::register_dispatcher(input)
+}
+
 /// Declares a help rendering function for an entry type.
 ///
 /// The `#[help]` attribute converts a function into a help provider by:
@@ -1130,6 +1154,45 @@ pub fn program_final_gen(input: TokenStream) -> TokenStream {
     #[cfg(not(feature = "general_renderer"))]
     let general_render = quote! {};
 
+    #[cfg(feature = "dispatch_tree")]
+    let compile_time_dispatchers: Vec<String> = COMPILE_TIME_DISPATCHERS
+        .lock()
+        .unwrap()
+        .clone()
+        .iter()
+        .cloned()
+        .collect();
+
+    #[cfg(feature = "dispatch_tree")]
+    let dispatch_tree_nodes = {
+        let entries: Vec<(String, String, String)> = compile_time_dispatchers
+            .iter()
+            .filter_map(|entry| {
+                let parts: Vec<&str> = entry.split(':').collect();
+                if parts.len() == 3 {
+                    Some((
+                        parts[0].to_string(),
+                        parts[1].to_string(),
+                        parts[2].to_string(),
+                    ))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        let get_nodes_fn = dispatch_tree_gen::gen_get_nodes(&entries);
+        let dispatch_trie_fn = dispatch_tree_gen::gen_dispatch_args_trie(&entries);
+
+        quote! {
+            #get_nodes_fn
+            #dispatch_trie_fn
+        }
+    };
+
+    #[cfg(not(feature = "dispatch_tree"))]
+    let dispatch_tree_nodes = quote! {};
+
     #[cfg(feature = "comp")]
     let completion_tokens: Vec<proc_macro2::TokenStream> = completions
         .iter()
@@ -1215,6 +1278,7 @@ pub fn program_final_gen(input: TokenStream) -> TokenStream {
                     _ => false
                 }
             }
+            #dispatch_tree_nodes
             #general_render
             #comp
         }