1use std::sync::{
4 Arc,
5 atomic::{AtomicU64, Ordering},
6};
7
8#[cfg(test)]
9use super::CommandContext;
10use super::{
11 BrigadierRuntime, CommandArgumentParser, CommandNodeBuilder, CommandRedirectTarget,
12 CommandRuntime, CommandSyntaxError, CommandSyntaxErrorKind, ContextChain, NodeId, NodeKind,
13 ParseError, ParseResults, ParsedCommandContext, RegistrationError, RegistrationErrorKind,
14 StringRange, StringReader, SuggestionError, Suggestions, SuggestionsBuilder,
15 node::{CommandNode, CommandNodeData, UnregisteredCommandNode},
16};
17
18static NEXT_DISPATCHER_ID: AtomicU64 = AtomicU64::new(1);
19
20pub(crate) struct CommandDispatcher<S, R = BrigadierRuntime>
22where
23 R: CommandRuntime<S>,
24{
25 id: u64,
26 nodes: Vec<CommandNode<S, R>>,
27}
28
29impl<S, R> CommandDispatcher<S, R>
30where
31 R: CommandRuntime<S>,
32{
33 pub(crate) fn new() -> Self {
35 Self {
36 id: NEXT_DISPATCHER_ID.fetch_add(1, Ordering::Relaxed),
37 nodes: vec![CommandNode::root()],
38 }
39 }
40
41 pub(crate) const fn root(&self) -> NodeId {
43 NodeId::new(self.id, 0)
44 }
45
46 pub(crate) fn register(
48 &mut self,
49 builder: CommandNodeBuilder<S, R>,
50 ) -> Result<NodeId, RegistrationError> {
51 let mut node = builder.normalize()?;
52 if node.kind() != NodeKind::Literal {
53 return Err(RegistrationError::new(RegistrationErrorKind::ArgumentRoot));
54 }
55
56 self.validate_redirects(&node)?;
57 let command_root = self
58 .find_child(self.root(), node.name())
59 .unwrap_or_else(|| NodeId::new(self.id, self.nodes.len()));
60 node.resolve_command_root(command_root);
61 self.validate_merge(self.root(), &node)?;
62 Ok(self.apply_merge(self.root(), node))
63 }
64
65 pub(crate) fn parse<'input>(
67 &self,
68 input: &'input str,
69 source: S,
70 ) -> ParseResults<'input, S, R> {
71 self.parse_reader(StringReader::new(input), source)
72 }
73
74 pub(crate) fn parse_reader<'input>(
76 &self,
77 reader: StringReader<'input>,
78 source: S,
79 ) -> ParseResults<'input, S, R> {
80 let context = ParsedCommandContext::new(Arc::new(source), self.root(), reader.cursor());
81 self.parse_nodes(self.root(), reader, context)
82 }
83
84 pub(crate) fn completion_suggestions(
86 &self,
87 parse: &ParseResults<'_, S, R>,
88 ) -> Result<Suggestions, SuggestionError> {
89 let input = parse.reader().input();
90 let cursor = parse.reader().total_length();
91 let Some(context) = parse.context().find_suggestion_context(cursor) else {
92 return Ok(Suggestions::empty());
93 };
94 let Some(children) = self.children(context.parent) else {
95 return Ok(Suggestions::empty());
96 };
97
98 let mut candidate_sets = Vec::with_capacity(children.len());
99 for child_id in children {
100 let child = &self.nodes[child_id.index];
101 if !child.requirement.allows(parse.context().source()) {
104 continue;
105 }
106
107 let mut builder = SuggestionsBuilder::new(input, context.start.min(cursor))?;
108 match &child.data {
109 CommandNodeData::Root => {
110 unreachable!("the command graph never stores a root node as a child")
111 }
112 CommandNodeData::Literal(literal) => {
113 if literal
114 .to_lowercase()
115 .starts_with(builder.remaining_lowercase())
116 {
117 builder.suggest(literal.as_ref());
118 }
119 }
120 CommandNodeData::Argument { argument_type, .. } => {
121 let argument_context = context.context.argument_suggestion_context();
122 argument_type.list_suggestions(&argument_context, &mut builder);
123 }
124 }
125 candidate_sets.push(builder.build()?);
126 }
127
128 Suggestions::merge(input, candidate_sets)
129 }
130
131 pub(crate) fn context_chain(
133 &self,
134 parse: ParseResults<'_, S, R>,
135 ) -> Result<ContextChain<S, R>, CommandSyntaxError> {
136 let (context, reader, mut errors) = parse.into_parts();
137 if reader.can_read() {
138 if errors.len() == 1 {
139 return Err(errors.remove(0).into_error());
140 }
141 let kind = if context.range().is_empty() {
142 CommandSyntaxErrorKind::UnknownCommand
143 } else {
144 CommandSyntaxErrorKind::UnknownArgument
145 };
146 return Err(reader.error(kind));
147 }
148 if context.root().dispatcher != self.id {
149 return Err(reader.error(CommandSyntaxErrorKind::UnknownCommand));
150 }
151
152 let input: Arc<str> = Arc::from(reader.input());
153 let context = context.build(input);
154 ContextChain::try_flatten(context)
155 .ok_or_else(|| reader.error(CommandSyntaxErrorKind::UnknownCommand))
156 }
157
158 pub(crate) fn node(&self, id: NodeId) -> Option<&CommandNode<S, R>> {
160 if id.dispatcher != self.id {
161 return None;
162 }
163 self.nodes.get(id.index)
164 }
165
166 pub(crate) fn children(&self, id: NodeId) -> Option<&[NodeId]> {
168 self.node(id).map(|node| node.children.as_slice())
169 }
170
171 pub(crate) const fn node_count(&self) -> usize {
173 self.nodes.len()
174 }
175
176 fn validate_merge(
177 &self,
178 parent: NodeId,
179 incoming: &UnregisteredCommandNode<S, R>,
180 ) -> Result<(), RegistrationError> {
181 let Some(existing_id) = self.find_child(parent, incoming.name()) else {
182 return Ok(());
183 };
184 let existing = &self.nodes[existing_id.index];
185 existing.validate_compatible(incoming)?;
186 for child in &incoming.children {
187 self.validate_merge(existing_id, child)?;
188 }
189 Ok(())
190 }
191
192 fn validate_redirects(
193 &self,
194 node: &UnregisteredCommandNode<S, R>,
195 ) -> Result<(), RegistrationError> {
196 if let Some(redirect) = &node.redirect
197 && let CommandRedirectTarget::Node(target) = redirect.target
198 && self.node(target).is_none()
199 {
200 return Err(RegistrationError::new(
201 RegistrationErrorKind::InvalidRedirectTarget { target },
202 ));
203 }
204 for child in &node.children {
205 self.validate_redirects(child)?;
206 }
207 Ok(())
208 }
209
210 fn apply_merge(
211 &mut self,
212 parent: NodeId,
213 mut incoming: UnregisteredCommandNode<S, R>,
214 ) -> NodeId {
215 if let Some(existing_id) = self.find_child(parent, incoming.name()) {
216 if incoming.executor.is_some() {
217 self.nodes[existing_id.index].executor = incoming.executor.take();
218 }
219 for child in incoming.children {
220 self.apply_merge(existing_id, child);
221 }
222 return existing_id;
223 }
224
225 let node_id = NodeId::new(self.id, self.nodes.len());
226 let children = incoming.children;
227 self.nodes.push(CommandNode {
228 data: incoming.data,
229 children: Vec::new(),
230 executor: incoming.executor,
231 requirement: incoming.requirement,
232 execution_requirement: incoming.execution_requirement,
233 redirect: incoming.redirect,
234 });
235 self.nodes[parent.index].children.push(node_id);
236 for child in children {
237 self.apply_merge(node_id, child);
238 }
239 node_id
240 }
241
242 fn find_child(&self, parent: NodeId, name: &str) -> Option<NodeId> {
243 let parent = self.node(parent)?;
244 parent.children.iter().copied().find(|child| {
245 self.nodes
246 .get(child.index)
247 .is_some_and(|node| node.name() == name)
248 })
249 }
250
251 fn parse_nodes<'input>(
252 &self,
253 parent: NodeId,
254 original_reader: StringReader<'input>,
255 context_so_far: ParsedCommandContext<S, R>,
256 ) -> ParseResults<'input, S, R> {
257 let mut errors = Vec::new();
258 let mut potentials = Vec::new();
259
260 for child_id in self.relevant_nodes(parent, &original_reader) {
261 let child = &self.nodes[child_id.index];
262 if !child.requirement.allows(context_so_far.source()) {
263 continue;
264 }
265
266 let mut context = context_so_far.branch();
267 let mut reader = original_reader.clone();
268 if let Err(error) = self.parse_node(child_id, &mut reader, &mut context) {
269 errors.push(ParseError::new(child_id, error));
270 continue;
271 }
272 if reader.can_read() && reader.peek() != Some(' ') {
273 errors.push(ParseError::new(
274 child_id,
275 reader.error(CommandSyntaxErrorKind::ExpectedArgumentSeparator),
276 ));
277 continue;
278 }
279
280 let executor = child
281 .executor
282 .as_ref()
283 .filter(|_| child.execution_requirement.allows(context.source()))
284 .map(Arc::clone);
285 context.set_executor(executor);
286 let redirect = child.redirect();
287 let required_remaining = if redirect.is_some() { 1 } else { 2 };
288 if reader.can_read_length(required_remaining) {
289 reader.skip();
290 if let Some(target) = redirect {
291 let child_context = ParsedCommandContext::new(
292 Arc::clone(context.source_arc()),
293 target,
294 reader.cursor(),
295 );
296 let parse = self.parse_nodes(target, reader, child_context);
297 let (child_context, reader, errors) = parse.into_parts();
298 context.set_child(child_context);
299 return ParseResults::new(context, reader, errors);
300 }
301 potentials.push(self.parse_nodes(child_id, reader, context));
302 } else {
303 potentials.push(ParseResults::new(context, reader, Vec::new()));
304 }
305 }
306
307 let mut potentials = potentials.into_iter();
308 let Some(mut best) = potentials.next() else {
309 return ParseResults::new(context_so_far, original_reader, errors);
310 };
311 for potential in potentials {
312 if Self::is_better_parse(&potential, &best) {
313 best = potential;
314 }
315 }
316 best
317 }
318
319 fn parse_node(
320 &self,
321 node_id: NodeId,
322 reader: &mut StringReader<'_>,
323 context: &mut ParsedCommandContext<S, R>,
324 ) -> Result<(), CommandSyntaxError> {
325 let node = &self.nodes[node_id.index];
326 let start = reader.cursor();
327 match &node.data {
328 CommandNodeData::Root => {
329 unreachable!("the command graph never stores a root node as a child")
330 }
331 CommandNodeData::Literal(literal) => {
332 if !reader.try_read_literal(literal) {
333 return Err(
334 reader.error(CommandSyntaxErrorKind::LiteralIncorrect(literal.to_owned()))
335 );
336 }
337 context.with_node(
338 node_id,
339 StringRange::between(start, reader.cursor()),
340 node.redirect.as_ref(),
341 );
342 }
343 CommandNodeData::Argument {
344 name,
345 argument_type,
346 } => {
347 let value = argument_type.parse(reader, context.source())?;
348 let range = StringRange::between(start, reader.cursor());
349 context.with_argument(name, range, value);
350 context.with_node(node_id, range, node.redirect.as_ref());
351 }
352 }
353 Ok(())
354 }
355
356 fn relevant_nodes(&self, parent: NodeId, reader: &StringReader<'_>) -> Vec<NodeId> {
357 let Some(parent) = self.node(parent) else {
358 return Vec::new();
359 };
360 let remaining = reader.remaining();
361 let token = remaining
362 .split_once(' ')
363 .map_or(remaining, |(token, _)| token);
364 let mut arguments = Vec::new();
365
366 for child_id in &parent.children {
367 match &self.nodes[child_id.index].data {
368 CommandNodeData::Literal(literal) if literal.as_ref() == token => {
369 return vec![*child_id];
370 }
371 CommandNodeData::Argument { .. } => arguments.push(*child_id),
372 CommandNodeData::Root | CommandNodeData::Literal(_) => {}
373 }
374 }
375 arguments
376 }
377
378 fn is_better_parse(
379 candidate: &ParseResults<'_, S, R>,
380 current: &ParseResults<'_, S, R>,
381 ) -> bool {
382 if !candidate.reader().can_read() && current.reader().can_read() {
383 return true;
384 }
385 if candidate.reader().can_read() && !current.reader().can_read() {
386 return false;
387 }
388 candidate.errors().is_empty() && !current.errors().is_empty()
389 }
390}
391
392#[cfg(test)]
393impl<S> CommandDispatcher<S, BrigadierRuntime> {
394 pub(super) fn execute_node_for_test(
395 &self,
396 node: NodeId,
397 source: S,
398 ) -> Option<Result<i32, CommandSyntaxError>> {
399 let node = self.node(node)?;
400 if !node.can_execute(&source) {
401 return None;
402 }
403 let executor = node.executor.as_deref()?;
404 Some(executor(&CommandContext::empty(source, self.root())))
405 }
406}
407
408impl<S, R> Default for CommandDispatcher<S, R>
409where
410 R: CommandRuntime<S>,
411{
412 fn default() -> Self {
413 Self::new()
414 }
415}