implement missing brigadier features and cleanup some more

3 files changed, 230 insertions, 13 deletions

diff --git a/azalea-brigadier/src/context/command_context.rs b/azalea-brigadier/src/context/command_context.rs
index 202d6f21..224f2d63 100644
--- a/azalea-brigadier/src/context/command_context.rs
+++ b/azalea-brigadier/src/context/command_context.rs
@@ -10,16 +10,16 @@ use crate::{
 
 /// A built `CommandContextBuilder`.
 pub struct CommandContext<S> {
-    pub source: Arc<S>,
-    pub input: String,
-    pub arguments: HashMap<String, ParsedArgument>,
-    pub command: Command<S>,
-    pub root_node: Arc<RwLock<CommandNode<S>>>,
-    pub nodes: Vec<ParsedCommandNode<S>>,
-    pub range: StringRange,
-    pub child: Option<Rc<CommandContext<S>>>,
-    pub modifier: Option<Arc<RedirectModifier<S>>>,
-    pub forks: bool,
+    pub(super) 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) range: StringRange,
+    pub(super) child: Option<Rc<CommandContext<S>>>,
+    pub(super) modifier: Option<Arc<RedirectModifier<S>>>,
+    pub(super) forks: bool,
 }
 
 impl<S> Clone for CommandContext<S> {
@@ -59,8 +59,10 @@ impl<S> Debug for CommandContext<S> {
 impl<S> CommandContext<S> {
     pub fn copy_for(&self, source: Arc<S>) -> Self {
         if Arc::ptr_eq(&source, &self.source) {
+            // fast path
             return self.clone();
         }
+
         CommandContext {
             source,
             input: self.input.clone(),
@@ -75,12 +77,56 @@ impl<S> CommandContext<S> {
         }
     }
 
+    pub fn child(&self) -> Option<&CommandContext<S>> {
+        self.child.as_ref().map(|c| c.as_ref())
+    }
+
+    pub fn last_child(&self) -> &CommandContext<S> {
+        let mut result = self;
+        while let Some(child) = result.child() {
+            result = child;
+        }
+        result
+    }
+
+    pub fn command(&self) -> &Command<S> {
+        &self.command
+    }
+
+    pub fn source(&self) -> &Arc<S> {
+        &self.source
+    }
+
+    pub fn argument(&self, name: &str) -> Option<&dyn Any> {
+        let argument = self.arguments.get(name);
+        argument.map(|a| a.result.as_ref())
+    }
+
+    pub fn redirect_modifier(&self) -> Option<&RedirectModifier<S>> {
+        self.modifier.as_ref().map(|m| m.as_ref())
+    }
+
+    pub fn range(&self) -> &StringRange {
+        &self.range
+    }
+
+    pub fn input(&self) -> &str {
+        &self.input
+    }
+
+    pub fn root_node(&self) -> &Arc<RwLock<CommandNode<S>>> {
+        &self.root_node
+    }
+
+    pub fn nodes(&self) -> &[ParsedCommandNode<S>] {
+        &self.nodes
+    }
+
     pub fn has_nodes(&self) -> bool {
         !self.nodes.is_empty()
     }
 
-    pub fn argument(&self, name: &str) -> Option<Arc<dyn Any>> {
-        let argument = self.arguments.get(name);
-        argument.map(|a| a.result.clone())
+    pub fn is_forked(&self) -> bool {
+        self.forks
     }
 }
diff --git a/azalea-brigadier/src/context/context_chain.rs b/azalea-brigadier/src/context/context_chain.rs
new file mode 100644
index 00000000..74fe6e01
--- /dev/null
+++ b/azalea-brigadier/src/context/context_chain.rs
@@ -0,0 +1,169 @@
+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>>>,
+}
+
+impl<S> ContextChain<S> {
+    pub fn new(modifiers: Vec<Rc<CommandContext<S>>>, executable: Rc<CommandContext<S>>) -> Self {
+        if executable.command.is_none() {
+            panic!("Last command in chain must be executable");
+        }
+        Self {
+            modifiers,
+            executable,
+            next_stage_cache: None,
+        }
+    }
+
+    pub fn try_flatten(root_context: Rc<CommandContext<S>>) -> Option<Self> {
+        let mut modifiers = Vec::new();
+        let mut current = root_context;
+        loop {
+            let child = current.child.clone();
+            let Some(child) = child else {
+                // Last entry must be executable command
+                if current.command.is_none() {
+                    return None;
+                }
+
+                return Some(ContextChain::new(modifiers, current));
+            };
+
+            modifiers.push(current);
+            current = child;
+        }
+    }
+
+    pub fn run_modifier(
+        modifier: Rc<CommandContext<S>>,
+        source: Arc<S>,
+        result_consumer: &dyn ResultConsumer<S>,
+        forked_mode: bool,
+    ) -> Result<Vec<Arc<S>>, CommandSyntaxError> {
+        let source_modifier = modifier.redirect_modifier();
+        let Some(source_modifier) = source_modifier else {
+            return Ok(vec![source]);
+        };
+
+        let context_to_use = Rc::new(modifier.copy_for(source));
+        let err = match (source_modifier)(&context_to_use) {
+            Ok(res) => return Ok(res),
+            Err(e) => e,
+        };
+
+        result_consumer.on_command_complete(context_to_use, false, 0);
+        if forked_mode {
+            return Ok(vec![]);
+        }
+        Err(err)
+    }
+
+    pub fn run_executable(
+        &self,
+        executable: Rc<CommandContext<S>>,
+        source: Arc<S>,
+        result_consumer: &dyn ResultConsumer<S>,
+        forked_mode: bool,
+    ) -> Result<i32, 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) };
+            }
+            Err(err) => err,
+        };
+
+        result_consumer.on_command_complete(context_to_use, false, 0);
+        if forked_mode { Ok(0) } else { Err(err) }
+    }
+
+    pub fn execute_all(
+        &self,
+        source: Arc<S>,
+        result_consumer: &(dyn ResultConsumer<S>),
+    ) -> Result<i32, CommandSyntaxError> {
+        if self.modifiers.is_empty() {
+            return self.run_executable(self.executable.clone(), source, result_consumer, false);
+        }
+
+        let mut forked_mode = false;
+        let mut current_sources = vec![source];
+
+        for modifier in &self.modifiers {
+            forked_mode |= modifier.is_forked();
+
+            let mut next_sources = Vec::new();
+            for source_to_run in current_sources {
+                next_sources.extend(Self::run_modifier(
+                    modifier.clone(),
+                    source_to_run.clone(),
+                    result_consumer,
+                    forked_mode,
+                )?);
+            }
+            if next_sources.is_empty() {
+                return Ok(0);
+            }
+            current_sources = next_sources;
+        }
+
+        let mut result = 0;
+        for execution_source in current_sources {
+            result += self.run_executable(
+                self.executable.clone(),
+                execution_source,
+                result_consumer,
+                forked_mode,
+            )?;
+        }
+
+        Ok(result)
+    }
+
+    pub fn stage(&self) -> Stage {
+        if self.modifiers.is_empty() {
+            Stage::Execute
+        } else {
+            Stage::Modify
+        }
+    }
+
+    pub fn top_context(&self) -> Rc<CommandContext<S>> {
+        self.modifiers
+            .first()
+            .cloned()
+            .unwrap_or_else(|| self.executable.clone())
+    }
+
+    pub fn next_stage(&mut self) -> Option<Rc<ContextChain<S>>> {
+        let modifier_count = self.modifiers.len();
+        if modifier_count == 0 {
+            return None;
+        }
+
+        if self.next_stage_cache.is_none() {
+            self.next_stage_cache = Some(Rc::new(ContextChain::new(
+                self.modifiers[1..].to_vec(),
+                self.executable.clone(),
+            )));
+        }
+
+        self.next_stage_cache.clone()
+    }
+}
+
+pub enum Stage {
+    Modify,
+    Execute,
+}
diff --git a/azalea-brigadier/src/context/mod.rs b/azalea-brigadier/src/context/mod.rs
index 28e1a12e..815892bf 100644
--- a/azalea-brigadier/src/context/mod.rs
+++ b/azalea-brigadier/src/context/mod.rs
@@ -1,5 +1,6 @@
 mod command_context;
 mod command_context_builder;
+mod context_chain;
 mod parsed_argument;
 mod parsed_command_node;
 mod string_range;
@@ -7,6 +8,7 @@ pub mod suggestion_context;
 
 pub use command_context::CommandContext;
 pub use command_context_builder::CommandContextBuilder;
+pub use context_chain::ContextChain;
 pub use parsed_argument::ParsedArgument;
 pub use parsed_command_node::ParsedCommandNode;
 pub use string_range::StringRange;