azalea_brigadier/
command_dispatcher.rs

1#![allow(clippy::type_complexity)]
2
3use std::{
4    cmp::Ordering,
5    collections::{HashMap, HashSet},
6    ptr,
7    rc::Rc,
8    sync::Arc,
9};
10
11use parking_lot::RwLock;
12
13use crate::{
14    builder::argument_builder::ArgumentBuilder,
15    context::{CommandContextBuilder, ContextChain},
16    errors::{BuiltInError, CommandResultTrait, CommandSyntaxError},
17    parse_results::ParseResults,
18    result_consumer::{DefaultResultConsumer, ResultConsumer},
19    string_reader::StringReader,
20    suggestion::{Suggestions, SuggestionsBuilder},
21    tree::CommandNode,
22};
23
24/// The root of the command tree. You need to make this to register commands.
25///
26/// ```
27/// # use azalea_brigadier::prelude::*;
28/// # struct CommandSource;
29/// let mut subject = CommandDispatcher::<CommandSource>::new();
30/// ```
31pub struct CommandDispatcher<S, R = i32>
32where
33    Self: Sync + Send,
34{
35    pub root: Arc<RwLock<CommandNode<S, R>>>,
36    consumer: Box<dyn ResultConsumer<S, R> + Send + Sync>,
37}
38
39impl<S, R: CommandResultTrait> CommandDispatcher<S, R> {
40    pub fn new() -> Self {
41        Self {
42            root: Arc::new(RwLock::new(CommandNode::default())),
43            consumer: Box::new(DefaultResultConsumer),
44        }
45    }
46
47    /// Add a new node to the root.
48    ///
49    /// ```
50    /// # use azalea_brigadier::prelude::*;
51    /// # let mut subject = CommandDispatcher::<()>::new();
52    /// subject.register(literal("foo").executes(|_| 42));
53    /// ```
54    pub fn register(&mut self, node: ArgumentBuilder<S, R>) -> Arc<RwLock<CommandNode<S, R>>> {
55        let build = Arc::new(RwLock::new(node.build()));
56        self.root.write().add_child(&build);
57        build
58    }
59
60    pub fn parse(&self, command: StringReader, source: S) -> ParseResults<'_, S, R> {
61        let source = Arc::new(source);
62
63        let context = CommandContextBuilder::new(self, source, self.root.clone(), command.cursor());
64        self.parse_nodes(&self.root, &command, context).unwrap()
65    }
66
67    fn parse_nodes<'a>(
68        &'a self,
69        node: &Arc<RwLock<CommandNode<S, R>>>,
70        original_reader: &StringReader,
71        context_so_far: CommandContextBuilder<'a, S, R>,
72    ) -> Result<ParseResults<'a, S, R>, CommandSyntaxError> {
73        let source = context_so_far.source.clone();
74        #[allow(clippy::mutable_key_type)] // this is fine because we don't mutate the key
75        let mut errors = HashMap::<Rc<CommandNode<S, R>>, CommandSyntaxError>::new();
76        let mut potentials: Vec<ParseResults<S, R>> = vec![];
77        let cursor = original_reader.cursor();
78
79        for child in node.read().get_relevant_nodes(&mut original_reader.clone()) {
80            if !child.read().can_use(&source) {
81                continue;
82            }
83            let mut context = context_so_far.clone();
84            let mut reader = original_reader.clone();
85
86            let parse_with_context_result =
87                child.read().parse_with_context(&mut reader, &mut context);
88            if let Err(ex) = parse_with_context_result {
89                errors.insert(
90                    Rc::new((*child.read()).clone()),
91                    BuiltInError::DispatcherParseException {
92                        message: ex.message(),
93                    }
94                    .create_with_context(&reader),
95                );
96                reader.cursor = cursor;
97                continue;
98            }
99            if reader.can_read() && reader.peek() != ' ' {
100                errors.insert(
101                    Rc::new((*child.read()).clone()),
102                    BuiltInError::DispatcherExpectedArgumentSeparator.create_with_context(&reader),
103                );
104                reader.cursor = cursor;
105                continue;
106            }
107
108            context.with_command(&child.read().command);
109            if reader.can_read_length(if child.read().redirect.is_none() {
110                2
111            } else {
112                1
113            }) {
114                reader.skip();
115                match &child.read().redirect {
116                    Some(redirect) => {
117                        let child_context = CommandContextBuilder::new(
118                            self,
119                            source,
120                            redirect.clone(),
121                            reader.cursor,
122                        );
123                        let parse = self
124                            .parse_nodes(redirect, &reader, child_context)
125                            .expect("Parsing nodes failed");
126                        context.with_child(Rc::new(parse.context));
127                        return Ok(ParseResults {
128                            context,
129                            reader: parse.reader,
130                            exceptions: parse.exceptions,
131                        });
132                    }
133                    _ => {
134                        let parse = self
135                            .parse_nodes(&child, &reader, context)
136                            .expect("Parsing nodes failed");
137                        potentials.push(parse);
138                    }
139                }
140            } else {
141                potentials.push(ParseResults {
142                    context,
143                    reader,
144                    exceptions: HashMap::new(),
145                });
146            }
147        }
148
149        if !potentials.is_empty() {
150            if potentials.len() > 1 {
151                potentials.sort_by(|a, b| {
152                    if !a.reader.can_read() && b.reader.can_read() {
153                        return Ordering::Less;
154                    };
155                    if a.reader.can_read() && !b.reader.can_read() {
156                        return Ordering::Greater;
157                    };
158                    if a.exceptions.is_empty() && !b.exceptions.is_empty() {
159                        return Ordering::Less;
160                    };
161                    if !a.exceptions.is_empty() && b.exceptions.is_empty() {
162                        return Ordering::Greater;
163                    };
164                    Ordering::Equal
165                });
166            }
167            let best_potential = potentials.into_iter().next().unwrap();
168            return Ok(best_potential);
169        }
170
171        Ok(ParseResults {
172            context: context_so_far,
173            reader: original_reader.clone(),
174            exceptions: errors,
175        })
176    }
177
178    /// Parse and execute the command using the given input and context.
179    ///
180    /// The number returned depends on the command, and may not be of
181    /// significance.
182    ///
183    /// This is a shortcut for `Self::parse` and `Self::execute_parsed`.
184    pub fn execute(
185        &self,
186        input: impl Into<StringReader>,
187        source: S,
188    ) -> Result<R, CommandSyntaxError> {
189        let input = input.into();
190
191        let parse = self.parse(input, source);
192        self.execute_parsed(parse)
193    }
194
195    pub fn add_paths(
196        node: Arc<RwLock<CommandNode<S, R>>>,
197        result: &mut Vec<Vec<Arc<RwLock<CommandNode<S, R>>>>>,
198        parents: Vec<Arc<RwLock<CommandNode<S, R>>>>,
199    ) {
200        let mut current = parents;
201        current.push(node.clone());
202        result.push(current.clone());
203
204        for child in node.read().children.values() {
205            Self::add_paths(child.clone(), result, current.clone());
206        }
207    }
208
209    pub fn get_path(&self, target: CommandNode<S, R>) -> Vec<String> {
210        let rc_target = Arc::new(RwLock::new(target));
211        let mut nodes: Vec<Vec<Arc<RwLock<CommandNode<S, R>>>>> = Vec::new();
212        Self::add_paths(self.root.clone(), &mut nodes, vec![]);
213
214        for list in nodes {
215            if *list.last().expect("Nothing in list").read() == *rc_target.read() {
216                let mut result: Vec<String> = Vec::with_capacity(list.len());
217                for node in list {
218                    if !Arc::ptr_eq(&node, &self.root) {
219                        result.push(node.read().name().to_owned());
220                    }
221                }
222                return result;
223            }
224        }
225        vec![]
226    }
227
228    pub fn find_node(&self, path: &[&str]) -> Option<Arc<RwLock<CommandNode<S, R>>>> {
229        let mut node = self.root.clone();
230        for name in path {
231            match node.clone().read().child(name) {
232                Some(child) => {
233                    node = child;
234                }
235                _ => {
236                    return None;
237                }
238            };
239        }
240        Some(node)
241    }
242
243    /// Executes a given pre-parsed command.
244    pub fn execute_parsed(&self, parse: ParseResults<S, R>) -> Result<R, CommandSyntaxError> {
245        if parse.reader.can_read() {
246            return Err(if parse.exceptions.len() == 1 {
247                parse.exceptions.values().next().unwrap().clone()
248            } else if parse.context.range.is_empty() {
249                BuiltInError::DispatcherUnknownCommand.create_with_context(&parse.reader)
250            } else {
251                BuiltInError::DispatcherUnknownArgument.create_with_context(&parse.reader)
252            });
253        }
254
255        let command = parse.reader.string();
256        let original = Rc::new(parse.context.build(command));
257        let flat_context = ContextChain::try_flatten(original.clone());
258        let Some(flat_context) = flat_context else {
259            self.consumer.on_command_complete(original, false, 0);
260            return Err(BuiltInError::DispatcherUnknownCommand.create_with_context(&parse.reader));
261        };
262
263        flat_context.execute_all(original.source.clone(), self.consumer.as_ref())
264    }
265
266    pub fn get_all_usage(
267        &self,
268        node: &CommandNode<S, R>,
269        source: &S,
270        restricted: bool,
271    ) -> Vec<String> {
272        let mut result = vec![];
273        self.get_all_usage_recursive(node, source, &mut result, "", restricted);
274        result
275    }
276
277    fn get_all_usage_recursive(
278        &self,
279        node: &CommandNode<S, R>,
280        source: &S,
281        result: &mut Vec<String>,
282        prefix: &str,
283        restricted: bool,
284    ) {
285        if restricted && !node.can_use(source) {
286            return;
287        }
288        if node.command.is_some() {
289            result.push(prefix.to_owned());
290        }
291        match &node.redirect {
292            Some(redirect) => {
293                let redirect = if ptr::eq(redirect.data_ptr(), self.root.data_ptr()) {
294                    "...".to_owned()
295                } else {
296                    format!("-> {}", redirect.read().usage_text())
297                };
298                if prefix.is_empty() {
299                    result.push(format!("{} {redirect}", node.usage_text()));
300                } else {
301                    result.push(format!("{prefix} {redirect}"));
302                }
303            }
304            _ => {
305                for child in node.children.values() {
306                    let child = child.read();
307                    self.get_all_usage_recursive(
308                        &child,
309                        source,
310                        result,
311                        if prefix.is_empty() {
312                            child.usage_text()
313                        } else {
314                            format!("{prefix} {}", child.usage_text())
315                        }
316                        .as_str(),
317                        restricted,
318                    );
319                }
320            }
321        }
322    }
323
324    /// Gets the possible executable commands from a specified node.
325    ///
326    /// You may use [`Self::root`] as a target to get usage data for the entire
327    /// command tree.
328    pub fn get_smart_usage(
329        &self,
330        node: &CommandNode<S, R>,
331        source: &S,
332    ) -> Vec<(Arc<RwLock<CommandNode<S, R>>>, String)> {
333        let mut result = Vec::new();
334
335        let optional = node.command.is_some();
336        for child in node.children.values() {
337            let usage = self.get_smart_usage_recursive(&child.read(), source, optional, false);
338            if let Some(usage) = usage {
339                result.push((child.clone(), usage));
340            }
341        }
342
343        result
344    }
345
346    fn get_smart_usage_recursive(
347        &self,
348        node: &CommandNode<S, R>,
349        source: &S,
350        optional: bool,
351        deep: bool,
352    ) -> Option<String> {
353        if !node.can_use(source) {
354            return None;
355        }
356
357        let this = if optional {
358            format!("[{}]", node.usage_text())
359        } else {
360            node.usage_text()
361        };
362        let child_optional = node.command.is_some();
363        let open = if child_optional { "[" } else { "(" };
364        let close = if child_optional { "]" } else { ")" };
365
366        if deep {
367            return Some(this);
368        }
369
370        if let Some(redirect) = &node.redirect {
371            let redirect = if ptr::eq(redirect.data_ptr(), self.root.data_ptr()) {
372                "...".to_owned()
373            } else {
374                format!("-> {}", redirect.read().usage_text())
375            };
376            return Some(format!("{this} {redirect}"));
377        }
378
379        let children = node
380            .children
381            .values()
382            .filter(|child| child.read().can_use(source))
383            .collect::<Vec<_>>();
384        match children.len().cmp(&1) {
385            Ordering::Less => {}
386            Ordering::Equal => {
387                let usage = self.get_smart_usage_recursive(
388                    &children[0].read(),
389                    source,
390                    child_optional,
391                    child_optional,
392                );
393                if let Some(usage) = usage {
394                    return Some(format!("{this} {usage}"));
395                }
396            }
397            Ordering::Greater => {
398                let mut child_usage = HashSet::new();
399                for child in &children {
400                    let usage =
401                        self.get_smart_usage_recursive(&child.read(), source, child_optional, true);
402                    if let Some(usage) = usage {
403                        child_usage.insert(usage);
404                    }
405                }
406                match child_usage.len().cmp(&1) {
407                    Ordering::Less => {}
408                    Ordering::Equal => {
409                        let usage = child_usage.into_iter().next().unwrap();
410                        let usage = if child_optional {
411                            format!("[{usage}]")
412                        } else {
413                            usage
414                        };
415                        return Some(format!("{this} {usage}"));
416                    }
417                    Ordering::Greater => {
418                        let mut builder = String::new();
419                        builder.push_str(open);
420                        let mut count = 0;
421                        for child in children {
422                            if count > 0 {
423                                builder.push('|');
424                            }
425                            builder.push_str(&child.read().usage_text());
426                            count += 1;
427                        }
428                        if count > 0 {
429                            builder.push_str(close);
430                            return Some(format!("{this} {builder}"));
431                        }
432                    }
433                }
434            }
435        }
436
437        Some(this)
438    }
439
440    pub fn get_completion_suggestions(parse: ParseResults<S, R>) -> Suggestions {
441        let cursor = parse.reader.total_length();
442        Self::get_completion_suggestions_with_cursor(parse, cursor)
443    }
444
445    pub fn get_completion_suggestions_with_cursor(
446        parse: ParseResults<S, R>,
447        cursor: usize,
448    ) -> Suggestions {
449        let context = parse.context;
450
451        let node_before_cursor = context.find_suggestion_context(cursor);
452        let parent = node_before_cursor.parent;
453        let start = usize::min(node_before_cursor.start_pos, cursor);
454
455        let full_input = parse.reader.string();
456        let truncated_input = full_input[..cursor].to_string();
457        let truncated_input_lowercase = truncated_input.to_lowercase();
458
459        let mut all_suggestions = Vec::new();
460        for node in parent.read().children.values() {
461            let suggestions = node.read().list_suggestions(
462                context.build(&truncated_input),
463                SuggestionsBuilder::new_with_lowercase(
464                    &truncated_input,
465                    &truncated_input_lowercase,
466                    start,
467                ),
468            );
469            all_suggestions.push(suggestions);
470        }
471
472        Suggestions::merge(full_input, &all_suggestions)
473    }
474}
475
476impl<S, R: CommandResultTrait> Default for CommandDispatcher<S, R> {
477    fn default() -> Self {
478        Self::new()
479    }
480}
azalea_brigadier/command_dispatcher.rs

azalea_brigadier/
command_dispatcher.rs
