azalea_brigadier/
command_dispatcher.rs

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

azalea_brigadier/
command_dispatcher.rs
