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.
177    ///
178    /// The number returned depends on the command, and may not be of
179    /// significance.
180    ///
181    /// This is a shortcut for `Self::parse` and `Self::execute_parsed`.
182    pub fn execute(
183        &self,
184        input: impl Into<StringReader>,
185        source: S,
186    ) -> Result<i32, CommandSyntaxError> {
187        let input = input.into();
188
189        let parse = self.parse(input, source);
190        self.execute_parsed(parse)
191    }
192
193    pub fn add_paths(
194        node: Arc<RwLock<CommandNode<S>>>,
195        result: &mut Vec<Vec<Arc<RwLock<CommandNode<S>>>>>,
196        parents: Vec<Arc<RwLock<CommandNode<S>>>>,
197    ) {
198        let mut current = parents;
199        current.push(node.clone());
200        result.push(current.clone());
201
202        for child in node.read().children.values() {
203            Self::add_paths(child.clone(), result, current.clone());
204        }
205    }
206
207    pub fn get_path(&self, target: CommandNode<S>) -> Vec<String> {
208        let rc_target = Arc::new(RwLock::new(target));
209        let mut nodes: Vec<Vec<Arc<RwLock<CommandNode<S>>>>> = Vec::new();
210        Self::add_paths(self.root.clone(), &mut nodes, vec![]);
211
212        for list in nodes {
213            if *list.last().expect("Nothing in list").read() == *rc_target.read() {
214                let mut result: Vec<String> = Vec::with_capacity(list.len());
215                for node in list {
216                    if !Arc::ptr_eq(&node, &self.root) {
217                        result.push(node.read().name().to_string());
218                    }
219                }
220                return result;
221            }
222        }
223        vec![]
224    }
225
226    pub fn find_node(&self, path: &[&str]) -> Option<Arc<RwLock<CommandNode<S>>>> {
227        let mut node = self.root.clone();
228        for name in path {
229            match node.clone().read().child(name) {
230                Some(child) => {
231                    node = child;
232                }
233                _ => {
234                    return None;
235                }
236            };
237        }
238        Some(node)
239    }
240
241    /// Executes a given pre-parsed command.
242    pub fn execute_parsed(&self, parse: ParseResults<S>) -> Result<i32, CommandSyntaxError> {
243        if parse.reader.can_read() {
244            return Err(if parse.exceptions.len() == 1 {
245                parse.exceptions.values().next().unwrap().clone()
246            } else if parse.context.range.is_empty() {
247                BuiltInError::DispatcherUnknownCommand.create_with_context(&parse.reader)
248            } else {
249                BuiltInError::DispatcherUnknownArgument.create_with_context(&parse.reader)
250            });
251        }
252
253        let command = parse.reader.string();
254        let original = Rc::new(parse.context.build(command));
255        let flat_context = ContextChain::try_flatten(original.clone());
256        let Some(flat_context) = flat_context else {
257            self.consumer.on_command_complete(original, false, 0);
258            return Err(BuiltInError::DispatcherUnknownCommand.create_with_context(&parse.reader));
259        };
260
261        flat_context.execute_all(original.source.clone(), self.consumer.as_ref())
262    }
263
264    pub fn get_all_usage(
265        &self,
266        node: &CommandNode<S>,
267        source: &S,
268        restricted: bool,
269    ) -> Vec<String> {
270        let mut result = vec![];
271        self.get_all_usage_recursive(node, source, &mut result, "", restricted);
272        result
273    }
274
275    fn get_all_usage_recursive(
276        &self,
277        node: &CommandNode<S>,
278        source: &S,
279        result: &mut Vec<String>,
280        prefix: &str,
281        restricted: bool,
282    ) {
283        if restricted && !node.can_use(source) {
284            return;
285        }
286        if node.command.is_some() {
287            result.push(prefix.to_owned());
288        }
289        match &node.redirect {
290            Some(redirect) => {
291                let redirect = if ptr::eq(redirect.data_ptr(), self.root.data_ptr()) {
292                    "...".to_string()
293                } else {
294                    format!("-> {}", redirect.read().usage_text())
295                };
296                if prefix.is_empty() {
297                    result.push(format!("{} {redirect}", node.usage_text()));
298                } else {
299                    result.push(format!("{prefix} {redirect}"));
300                }
301            }
302            _ => {
303                for child in node.children.values() {
304                    let child = child.read();
305                    self.get_all_usage_recursive(
306                        &child,
307                        source,
308                        result,
309                        if prefix.is_empty() {
310                            child.usage_text()
311                        } else {
312                            format!("{prefix} {}", child.usage_text())
313                        }
314                        .as_str(),
315                        restricted,
316                    );
317                }
318            }
319        }
320    }
321
322    /// Gets the possible executable commands from a specified node.
323    ///
324    /// You may use [`Self::root`] as a target to get usage data for the entire
325    /// command tree.
326    pub fn get_smart_usage(
327        &self,
328        node: &CommandNode<S>,
329        source: &S,
330    ) -> Vec<(Arc<RwLock<CommandNode<S>>>, String)> {
331        let mut result = Vec::new();
332
333        let optional = node.command.is_some();
334        for child in node.children.values() {
335            let usage = self.get_smart_usage_recursive(&child.read(), source, optional, false);
336            if let Some(usage) = usage {
337                result.push((child.clone(), usage));
338            }
339        }
340
341        result
342    }
343
344    fn get_smart_usage_recursive(
345        &self,
346        node: &CommandNode<S>,
347        source: &S,
348        optional: bool,
349        deep: bool,
350    ) -> Option<String> {
351        if !node.can_use(source) {
352            return None;
353        }
354
355        let this = if optional {
356            format!("[{}]", node.usage_text())
357        } else {
358            node.usage_text()
359        };
360        let child_optional = node.command.is_some();
361        let open = if child_optional { "[" } else { "(" };
362        let close = if child_optional { "]" } else { ")" };
363
364        if deep {
365            return Some(this);
366        }
367
368        if let Some(redirect) = &node.redirect {
369            let redirect = if ptr::eq(redirect.data_ptr(), self.root.data_ptr()) {
370                "...".to_string()
371            } else {
372                format!("-> {}", redirect.read().usage_text())
373            };
374            return Some(format!("{this} {redirect}"));
375        }
376
377        let children = node
378            .children
379            .values()
380            .filter(|child| child.read().can_use(source))
381            .collect::<Vec<_>>();
382        match children.len().cmp(&1) {
383            Ordering::Less => {}
384            Ordering::Equal => {
385                let usage = self.get_smart_usage_recursive(
386                    &children[0].read(),
387                    source,
388                    child_optional,
389                    child_optional,
390                );
391                if let Some(usage) = usage {
392                    return Some(format!("{this} {usage}"));
393                }
394            }
395            Ordering::Greater => {
396                let mut child_usage = HashSet::new();
397                for child in &children {
398                    let usage =
399                        self.get_smart_usage_recursive(&child.read(), source, child_optional, true);
400                    if let Some(usage) = usage {
401                        child_usage.insert(usage);
402                    }
403                }
404                match child_usage.len().cmp(&1) {
405                    Ordering::Less => {}
406                    Ordering::Equal => {
407                        let usage = child_usage.into_iter().next().unwrap();
408                        let usage = if child_optional {
409                            format!("[{usage}]")
410                        } else {
411                            usage
412                        };
413                        return Some(format!("{this} {usage}"));
414                    }
415                    Ordering::Greater => {
416                        let mut builder = String::new();
417                        builder.push_str(open);
418                        let mut count = 0;
419                        for child in children {
420                            if count > 0 {
421                                builder.push('|');
422                            }
423                            builder.push_str(&child.read().usage_text());
424                            count += 1;
425                        }
426                        if count > 0 {
427                            builder.push_str(close);
428                            return Some(format!("{this} {builder}"));
429                        }
430                    }
431                }
432            }
433        }
434
435        Some(this)
436    }
437
438    pub fn get_completion_suggestions(parse: ParseResults<S>) -> Suggestions {
439        let cursor = parse.reader.total_length();
440        Self::get_completion_suggestions_with_cursor(parse, cursor)
441    }
442
443    pub fn get_completion_suggestions_with_cursor(
444        parse: ParseResults<S>,
445        cursor: usize,
446    ) -> Suggestions {
447        let context = parse.context;
448
449        let node_before_cursor = context.find_suggestion_context(cursor);
450        let parent = node_before_cursor.parent;
451        let start = usize::min(node_before_cursor.start_pos, cursor);
452
453        let full_input = parse.reader.string();
454        let truncated_input = full_input[..cursor].to_string();
455        let truncated_input_lowercase = truncated_input.to_lowercase();
456
457        let mut all_suggestions = Vec::new();
458        for node in parent.read().children.values() {
459            let suggestions = node.read().list_suggestions(
460                context.build(&truncated_input),
461                SuggestionsBuilder::new_with_lowercase(
462                    &truncated_input,
463                    &truncated_input_lowercase,
464                    start,
465                ),
466            );
467            all_suggestions.push(suggestions);
468        }
469
470        Suggestions::merge(full_input, &all_suggestions)
471    }
472}
473
474impl<S> Default for CommandDispatcher<S> {
475    fn default() -> Self {
476        Self::new()
477    }
478}
azalea_brigadier/command_dispatcher.rs

azalea_brigadier/
command_dispatcher.rs
