Skip to main content

steel_core/command/brigadier/
dispatcher.rs

1//! Dispatcher-owned command node arena.
2
3use 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
20/// Owns a stable arena of command nodes.
21pub(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    /// Creates an empty dispatcher containing only its root node.
34    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    /// Returns the stable root node ID.
42    pub(crate) const fn root(&self) -> NodeId {
43        NodeId::new(self.id, 0)
44    }
45
46    /// Registers and merges a literal command tree.
47    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    /// Parses `input` into the best Brigadier command branch.
66    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    /// Parses from an existing reader while preserving its current cursor.
75    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    /// Returns completions for the end of a parsed command input.
85    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            // Steel filters internal command completions as an authorization
102            // boundary; Brigadier itself exposes nodes regardless of canUse.
103            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    /// Validates a parse and turns its redirect contexts into executable stages.
132    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    /// Returns a node if the ID belongs to this dispatcher.
159    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    /// Returns a node's children in registration order.
167    pub(crate) fn children(&self, id: NodeId) -> Option<&[NodeId]> {
168        self.node(id).map(|node| node.children.as_slice())
169    }
170
171    /// Returns the number of allocated nodes, including the root.
172    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}