azalea_brigadier/
command_dispatcher.rs

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

azalea_brigadier/
command_dispatcher.rs
