azalea-brigadier now allows commands to return a Result

5 files changed, 89 insertions, 66 deletions

diff --git a/azalea-brigadier/src/context/command_context.rs b/azalea-brigadier/src/context/command_context.rs
index a9959895..0139f7e9 100644
--- a/azalea-brigadier/src/context/command_context.rs
+++ b/azalea-brigadier/src/context/command_context.rs
@@ -15,20 +15,20 @@ use crate::{
 };
 
 /// A built `CommandContextBuilder`.
-pub struct CommandContext<S> {
+pub struct CommandContext<S, R = i32> {
     pub source: Arc<S>,
     pub(super) input: String,
     pub(super) arguments: HashMap<String, ParsedArgument>,
-    pub(super) command: Command<S>,
-    pub(super) root_node: Arc<RwLock<CommandNode<S>>>,
-    pub(super) nodes: Vec<ParsedCommandNode<S>>,
+    pub(super) command: Command<S, R>,
+    pub(super) root_node: Arc<RwLock<CommandNode<S, R>>>,
+    pub(super) nodes: Vec<ParsedCommandNode<S, R>>,
     pub(super) range: StringRange,
-    pub(super) child: Option<Rc<CommandContext<S>>>,
-    pub(super) modifier: Option<Arc<RedirectModifier<S>>>,
+    pub(super) child: Option<Rc<CommandContext<S, R>>>,
+    pub(super) modifier: Option<Arc<RedirectModifier<S, R>>>,
     pub(super) forks: bool,
 }
 
-impl<S> Clone for CommandContext<S> {
+impl<S, R> Clone for CommandContext<S, R> {
     fn clone(&self) -> Self {
         Self {
             source: self.source.clone(),
@@ -45,7 +45,7 @@ impl<S> Clone for CommandContext<S> {
     }
 }
 
-impl<S> Debug for CommandContext<S> {
+impl<S, R> Debug for CommandContext<S, R> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("CommandContext")
             // .field("source", &self.source)
@@ -62,7 +62,7 @@ impl<S> Debug for CommandContext<S> {
     }
 }
 
-impl<S> CommandContext<S> {
+impl<S, R> CommandContext<S, R> {
     pub fn copy_for(&self, source: Arc<S>) -> Self {
         if Arc::ptr_eq(&source, &self.source) {
             // fast path
@@ -83,11 +83,11 @@ impl<S> CommandContext<S> {
         }
     }
 
-    pub fn child(&self) -> Option<&CommandContext<S>> {
+    pub fn child(&self) -> Option<&CommandContext<S, R>> {
         self.child.as_ref().map(|c| c.as_ref())
     }
 
-    pub fn last_child(&self) -> &CommandContext<S> {
+    pub fn last_child(&self) -> &CommandContext<S, R> {
         let mut result = self;
         while let Some(child) = result.child() {
             result = child;
@@ -95,7 +95,7 @@ impl<S> CommandContext<S> {
         result
     }
 
-    pub fn command(&self) -> &Command<S> {
+    pub fn command(&self) -> &Command<S, R> {
         &self.command
     }
 
@@ -104,7 +104,7 @@ impl<S> CommandContext<S> {
         argument.map(|a| a.result.as_ref())
     }
 
-    pub fn redirect_modifier(&self) -> Option<&RedirectModifier<S>> {
+    pub fn redirect_modifier(&self) -> Option<&RedirectModifier<S, R>> {
         self.modifier.as_ref().map(|m| m.as_ref())
     }
 
@@ -116,11 +116,11 @@ impl<S> CommandContext<S> {
         &self.input
     }
 
-    pub fn root_node(&self) -> &Arc<RwLock<CommandNode<S>>> {
+    pub fn root_node(&self) -> &Arc<RwLock<CommandNode<S, R>>> {
         &self.root_node
     }
 
-    pub fn nodes(&self) -> &[ParsedCommandNode<S>] {
+    pub fn nodes(&self) -> &[ParsedCommandNode<S, R>] {
         &self.nodes
     }
 
diff --git a/azalea-brigadier/src/context/command_context_builder.rs b/azalea-brigadier/src/context/command_context_builder.rs
index a3819246..e3b09fc1 100644
--- a/azalea-brigadier/src/context/command_context_builder.rs
+++ b/azalea-brigadier/src/context/command_context_builder.rs
@@ -17,20 +17,20 @@ use crate::{
     tree::{Command, CommandNode},
 };
 
-pub struct CommandContextBuilder<'a, S> {
+pub struct CommandContextBuilder<'a, S, R> {
     pub arguments: HashMap<String, ParsedArgument>,
-    pub root: Arc<RwLock<CommandNode<S>>>,
-    pub nodes: Vec<ParsedCommandNode<S>>,
-    pub dispatcher: &'a CommandDispatcher<S>,
+    pub root: Arc<RwLock<CommandNode<S, R>>>,
+    pub nodes: Vec<ParsedCommandNode<S, R>>,
+    pub dispatcher: &'a CommandDispatcher<S, R>,
     pub source: Arc<S>,
-    pub command: Command<S>,
-    pub child: Option<Rc<CommandContextBuilder<'a, S>>>,
+    pub command: Command<S, R>,
+    pub child: Option<Rc<CommandContextBuilder<'a, S, R>>>,
     pub range: StringRange,
-    pub modifier: Option<Arc<RedirectModifier<S>>>,
+    pub modifier: Option<Arc<RedirectModifier<S, R>>>,
     pub forks: bool,
 }
 
-impl<S> Clone for CommandContextBuilder<'_, S> {
+impl<S, R> Clone for CommandContextBuilder<'_, S, R> {
     fn clone(&self) -> Self {
         Self {
             arguments: self.arguments.clone(),
@@ -47,11 +47,11 @@ impl<S> Clone for CommandContextBuilder<'_, S> {
     }
 }
 
-impl<'a, S> CommandContextBuilder<'a, S> {
+impl<'a, S, R> CommandContextBuilder<'a, S, R> {
     pub fn new(
-        dispatcher: &'a CommandDispatcher<S>,
+        dispatcher: &'a CommandDispatcher<S, R>,
         source: Arc<S>,
-        root_node: Arc<RwLock<CommandNode<S>>>,
+        root_node: Arc<RwLock<CommandNode<S, R>>>,
         start: usize,
     ) -> Self {
         Self {
@@ -68,11 +68,11 @@ impl<'a, S> CommandContextBuilder<'a, S> {
         }
     }
 
-    pub fn with_command(&mut self, command: &Command<S>) -> &Self {
+    pub fn with_command(&mut self, command: &Command<S, R>) -> &Self {
         self.command.clone_from(command);
         self
     }
-    pub fn with_child(&mut self, child: Rc<CommandContextBuilder<'a, S>>) -> &Self {
+    pub fn with_child(&mut self, child: Rc<CommandContextBuilder<'a, S, R>>) -> &Self {
         self.child = Some(child);
         self
     }
@@ -80,7 +80,7 @@ impl<'a, S> CommandContextBuilder<'a, S> {
         self.arguments.insert(name.to_owned(), argument);
         self
     }
-    pub fn with_node(&mut self, node: Arc<RwLock<CommandNode<S>>>, range: StringRange) -> &Self {
+    pub fn with_node(&mut self, node: Arc<RwLock<CommandNode<S, R>>>, range: StringRange) -> &Self {
         self.nodes.push(ParsedCommandNode {
             node: node.clone(),
             range,
@@ -91,7 +91,7 @@ impl<'a, S> CommandContextBuilder<'a, S> {
         self
     }
 
-    pub fn build(&self, input: &str) -> CommandContext<S> {
+    pub fn build(&self, input: &str) -> CommandContext<S, R> {
         CommandContext {
             arguments: self.arguments.clone(),
             root_node: self.root.clone(),
@@ -106,7 +106,7 @@ impl<'a, S> CommandContextBuilder<'a, S> {
         }
     }
 
-    pub fn find_suggestion_context(&self, cursor: usize) -> SuggestionContext<S> {
+    pub fn find_suggestion_context(&self, cursor: usize) -> SuggestionContext<S, R> {
         if self.range.start() > cursor {
             panic!("Can't find node before cursor");
         }
@@ -144,7 +144,7 @@ impl<'a, S> CommandContextBuilder<'a, S> {
     }
 }
 
-impl<S> Debug for CommandContextBuilder<'_, S> {
+impl<S, R> Debug for CommandContextBuilder<'_, S, R> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("CommandContextBuilder")
             // .field("arguments", &self.arguments)
diff --git a/azalea-brigadier/src/context/context_chain.rs b/azalea-brigadier/src/context/context_chain.rs
index 8ad5b320..ab66fcf9 100644
--- a/azalea-brigadier/src/context/context_chain.rs
+++ b/azalea-brigadier/src/context/context_chain.rs
@@ -1,16 +1,22 @@
 use std::{rc::Rc, sync::Arc};
 
 use super::CommandContext;
-use crate::{errors::CommandSyntaxError, result_consumer::ResultConsumer};
-
-pub struct ContextChain<S> {
-    modifiers: Vec<Rc<CommandContext<S>>>,
-    executable: Rc<CommandContext<S>>,
-    next_stage_cache: Option<Rc<ContextChain<S>>>,
+use crate::{
+    errors::{CommandResultTrait, CommandSyntaxError},
+    result_consumer::ResultConsumer,
+};
+
+pub struct ContextChain<S, R> {
+    modifiers: Vec<Rc<CommandContext<S, R>>>,
+    executable: Rc<CommandContext<S, R>>,
+    next_stage_cache: Option<Rc<ContextChain<S, R>>>,
 }
 
-impl<S> ContextChain<S> {
-    pub fn new(modifiers: Vec<Rc<CommandContext<S>>>, executable: Rc<CommandContext<S>>) -> Self {
+impl<S, R: CommandResultTrait> ContextChain<S, R> {
+    pub fn new(
+        modifiers: Vec<Rc<CommandContext<S, R>>>,
+        executable: Rc<CommandContext<S, R>>,
+    ) -> Self {
         if executable.command.is_none() {
             panic!("Last command in chain must be executable");
         }
@@ -21,7 +27,7 @@ impl<S> ContextChain<S> {
         }
     }
 
-    pub fn try_flatten(root_context: Rc<CommandContext<S>>) -> Option<Self> {
+    pub fn try_flatten(root_context: Rc<CommandContext<S, R>>) -> Option<Self> {
         let mut modifiers = Vec::new();
         let mut current = root_context;
         loop {
@@ -39,9 +45,9 @@ impl<S> ContextChain<S> {
     }
 
     pub fn run_modifier(
-        modifier: Rc<CommandContext<S>>,
+        modifier: Rc<CommandContext<S, R>>,
         source: Arc<S>,
-        result_consumer: &dyn ResultConsumer<S>,
+        result_consumer: &dyn ResultConsumer<S, R>,
         forked_mode: bool,
     ) -> Result<Vec<Arc<S>>, CommandSyntaxError> {
         let source_modifier = modifier.redirect_modifier();
@@ -64,33 +70,43 @@ impl<S> ContextChain<S> {
 
     pub fn run_executable(
         &self,
-        executable: Rc<CommandContext<S>>,
+        executable: Rc<CommandContext<S, R>>,
         source: Arc<S>,
-        result_consumer: &dyn ResultConsumer<S>,
+        result_consumer: &dyn ResultConsumer<S, R>,
         forked_mode: bool,
-    ) -> Result<i32, CommandSyntaxError> {
+    ) -> Result<R, CommandSyntaxError> {
         let context_to_use = Rc::new(executable.copy_for(source));
         let Some(command) = &executable.command else {
             unimplemented!();
         };
 
-        let err = match (command)(&context_to_use) {
-            Ok(result) => {
-                result_consumer.on_command_complete(context_to_use, true, result);
-                return if forked_mode { Ok(1) } else { Ok(result) };
+        let res = (command)(&context_to_use);
+        let err = match res {
+            Ok(res) => {
+                let Some(res) = res.as_i32() else {
+                    // these are treated as exceptions, so they can bubble up without doing anything
+                    // else
+                    return Ok(res);
+                };
+                result_consumer.on_command_complete(context_to_use, true, res);
+                return if forked_mode {
+                    Ok(R::new(1))
+                } else {
+                    Ok(R::new(res))
+                };
             }
             Err(err) => err,
         };
 
         result_consumer.on_command_complete(context_to_use, false, 0);
-        if forked_mode { Ok(0) } else { Err(err) }
+        if forked_mode { Ok(R::new(0)) } else { Err(err) }
     }
 
     pub fn execute_all(
         &self,
         source: Arc<S>,
-        result_consumer: &dyn ResultConsumer<S>,
-    ) -> Result<i32, CommandSyntaxError> {
+        result_consumer: &dyn ResultConsumer<S, R>,
+    ) -> Result<R, CommandSyntaxError> {
         if self.modifiers.is_empty() {
             return self.run_executable(self.executable.clone(), source, result_consumer, false);
         }
@@ -103,30 +119,37 @@ impl<S> ContextChain<S> {
 
             let mut next_sources = Vec::new();
             for source_to_run in current_sources {
-                next_sources.extend(Self::run_modifier(
+                match Self::run_modifier(
                     modifier.clone(),
                     source_to_run.clone(),
                     result_consumer,
                     forked_mode,
-                )?);
+                ) {
+                    Ok(res) => next_sources.extend(res),
+                    Err(err) => return Err(err),
+                }
             }
             if next_sources.is_empty() {
-                return Ok(0);
+                return Ok(R::new(0));
             }
             current_sources = next_sources;
         }
 
-        let mut result = 0;
+        let mut summed = 0;
         for execution_source in current_sources {
-            result += self.run_executable(
+            let res = self.run_executable(
                 self.executable.clone(),
                 execution_source,
                 result_consumer,
                 forked_mode,
             )?;
+            match res.as_i32() {
+                Some(res) => summed += res,
+                None => return Ok(res),
+            }
         }
 
-        Ok(result)
+        Ok(R::new(summed))
     }
 
     pub fn stage(&self) -> Stage {
@@ -137,14 +160,14 @@ impl<S> ContextChain<S> {
         }
     }
 
-    pub fn top_context(&self) -> Rc<CommandContext<S>> {
+    pub fn top_context(&self) -> Rc<CommandContext<S, R>> {
         self.modifiers
             .first()
             .cloned()
             .unwrap_or_else(|| self.executable.clone())
     }
 
-    pub fn next_stage(&mut self) -> Option<Rc<ContextChain<S>>> {
+    pub fn next_stage(&mut self) -> Option<Rc<ContextChain<S, R>>> {
         let modifier_count = self.modifiers.len();
         if modifier_count == 0 {
             return None;
diff --git a/azalea-brigadier/src/context/parsed_command_node.rs b/azalea-brigadier/src/context/parsed_command_node.rs
index 62da13d8..8dbadb49 100644
--- a/azalea-brigadier/src/context/parsed_command_node.rs
+++ b/azalea-brigadier/src/context/parsed_command_node.rs
@@ -6,12 +6,12 @@ use super::string_range::StringRange;
 use crate::tree::CommandNode;
 
 #[derive(Debug)]
-pub struct ParsedCommandNode<S> {
-    pub node: Arc<RwLock<CommandNode<S>>>,
+pub struct ParsedCommandNode<S, R> {
+    pub node: Arc<RwLock<CommandNode<S, R>>>,
     pub range: StringRange,
 }
 
-impl<S> Clone for ParsedCommandNode<S> {
+impl<S, R> Clone for ParsedCommandNode<S, R> {
     fn clone(&self) -> Self {
         Self {
             node: self.node.clone(),
diff --git a/azalea-brigadier/src/context/suggestion_context.rs b/azalea-brigadier/src/context/suggestion_context.rs
index 58a73fdb..ed3dbfa2 100644
--- a/azalea-brigadier/src/context/suggestion_context.rs
+++ b/azalea-brigadier/src/context/suggestion_context.rs
@@ -5,7 +5,7 @@ use parking_lot::RwLock;
 use crate::tree::CommandNode;
 
 #[derive(Debug)]
-pub struct SuggestionContext<S> {
-    pub parent: Arc<RwLock<CommandNode<S>>>,
+pub struct SuggestionContext<S, R> {
+    pub parent: Arc<RwLock<CommandNode<S, R>>>,
     pub start_pos: usize,
 }