diff --git a/crates/pumpkin/src/command/context/command_context.rs b/crates/pumpkin/src/command/context/command_context.rs index 1f8f6fe9b..c28de6aa1 100644 --- a/crates/pumpkin/src/command/context/command_context.rs +++ b/crates/pumpkin/src/command/context/command_context.rs @@ -216,17 +216,17 @@ impl<'a> ContextChain<'a> { #[must_use] pub fn try_flatten(root: &CommandContext<'a>) -> Option { let mut modifiers = Vec::new(); - let mut current = root; - + let mut current = Arc::new(root.clone()); loop { if let Some(child) = current.get_child() { - modifiers.push(child.clone()); - current = child; + let temp_child = child.clone(); + modifiers.push(current); + current = temp_child; } else { return current .command .is_some() - .then(|| Self::new(modifiers, Arc::new(current.clone()))); + .then(|| Self::new(modifiers, current)); } } } @@ -547,8 +547,6 @@ impl<'a> CommandContextBuilder<'a> { #[cfg(test)] mod test { - use std::sync::Arc; - use crate::command::argument_builder::{ArgumentBuilder, CommandArgumentBuilder}; use crate::command::context::command_context::{ CommandContext, CommandContextBuilder, ContextChain, ParsedArgument, Stage, @@ -558,7 +556,11 @@ mod test { use crate::command::errors::command_syntax_error::CommandSyntaxError; use crate::command::node::dispatcher::{CommandDispatcher, EmptyResultConsumer}; use crate::command::node::tree::ROOT_NODE_ID; - use crate::command::node::{CommandExecutor, CommandExecutorResult, Redirection}; + use crate::command::node::{ + CommandExecutor, CommandExecutorResult, RedirectModifier, Redirection, + }; + use pumpkin_util::math::vector3::Vector3; + use std::sync::Arc; struct TenExecutor; impl CommandExecutor for TenExecutor { @@ -710,4 +712,49 @@ mod test { assert!(ContextChain::try_flatten(&top_context).is_none()); } + + struct CustomExecutor; + impl CommandExecutor for CustomExecutor { + fn execute<'a>(&'a self, context: &'a CommandContext) -> CommandExecutorResult<'a> { + Box::pin(async move { + let source = &context.source; + assert_eq!(source.position, Vector3::new(0f64, 10f64, 0f64)); + Ok(1) + }) + } + } + + #[tokio::test] + async fn multi_stage_modifier_execution() { + let mut dispatcher = CommandDispatcher::new(); + dispatcher.register( + CommandArgumentBuilder::new("foo", "A test command").executes(CustomExecutor), + ); + dispatcher.register( + CommandArgumentBuilder::new("bar", "Another test command").redirect_with_modifier( + Redirection::Root, + RedirectModifier::Custom(Arc::new(|context| { + Box::pin(async move { + let mut new_source = context.source.as_ref().clone(); + new_source.position = Vector3::new(0f64, 10f64, 0f64); + Ok(vec![Arc::new(new_source)]) + }) + })), + ), + ); + let source = Arc::new(CommandSource::dummy()); + let result = dispatcher.parse_input("bar foo", &source).await; + let top_context = result.context.build("bar foo"); + let chain = ContextChain::try_flatten(&top_context) + .expect("The context should have properly flattened, as it has a command to execute"); + assert_eq!(chain.get_top_context().source.position, Vector3::default()); + let chain2 = chain + .next_stage() + .expect("There should have been the next stage"); + assert!(chain2.next_stage().is_none()); + let res = chain + .execute_all(&source, dispatcher.consumer.as_ref()) + .await; + assert!(res.is_ok_and(|val| val == 1)); + } }