special handling for return command

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - support for commands using macros
 - support for registering scoreboards (automatic creation and deletion)
+- "return" command with special handling in groups and conditionals
 
 ### Changed
 - use "return" command for conditionals instead of data storage when using supported pack format

Cargo.toml

@@ -25,10 +25,10 @@ getset = "0.1.5"
 serde = { version = "1.0.219", optional = true, features = ["derive"] }
 serde_json = "1.0.140"
 tracing = "0.1.41"
-zip = { version = "2.6.1", default-features = false, features = ["deflate", "time"], optional = true }
+zip = { version = "4.0.0", default-features = false, features = ["deflate", "time"], optional = true }
 
 [dev-dependencies]
-tempfile = "3.19.1"
+tempfile = "3.20.0"
 
 [package.metadata.docs.rs]
 all-features = true

src/datapack/command/execute/conditional.rs

@@ -28,10 +28,8 @@ pub fn compile_if_cond(
     global_state: &MutCompilerState,
     function_state: &FunctionCompilerState,
 ) -> Vec<CompiledCommand> {
-    // TODO: special handling for return command
-
     if options.pack_format < 20 {
-        compile_pre_20_format(
+        compile_using_data_storage(
             cond,
             then,
             el,
@@ -56,7 +54,7 @@ pub fn compile_if_cond(
 }
 
 #[expect(clippy::too_many_lines, clippy::too_many_arguments)]
-fn compile_pre_20_format(
+fn compile_using_data_storage(
     cond: &Condition,
     then: &Execute,
     el: Option<&Execute>,
@@ -218,31 +216,45 @@ fn compile_since_20_format(
 
     // if the conditions have multiple parts joined by a disjunction or an else part, commands need to be grouped
     if el.is_some() || str_cond.len() > 1 {
-        let group_cmds = handle_return_group_case_since_20(
+        // change to compilation using data storage when the conditional contains a return
-            str_cond,
+        if !then.contains_return() && !el.is_some_and(super::Execute::contains_return) {
-            then,
+            let group_cmds = handle_return_group_case_since_20(
-            el,
+                str_cond,
-            prefix,
+                then,
-            options,
+                el,
-            global_state,
+                prefix,
-            function_state,
+                options,
-        );
+                global_state,
-        let group = Command::Group(group_cmds);
+                function_state,
-        let cmds = group.compile(options, global_state, function_state);
+            );
-        if contains_macros {
+            let group = Command::Group(group_cmds);
-            cmds.into_iter()
+            let cmds = group.compile(options, global_state, function_state);
-                .map(|cmd| cmd.or_contains_macros(true))
+            if contains_macros {
-                .collect()
+                cmds.into_iter()
+                    .map(|cmd| cmd.or_contains_macros(true))
+                    .collect()
+            } else {
+                cmds
+            }
         } else {
-            cmds
+            compile_using_data_storage(
+                cond,
+                then,
+                el,
+                prefix,
+                prefix_contains_macros,
+                options,
+                global_state,
+                function_state,
+            )
         }
     } else if then_count > 1 {
-        let then_cmd = match then.clone() {
+        let then_cmds = match then.clone() {
             Execute::Run(cmd) => vec![*cmd],
             Execute::Runs(cmds) => cmds,
             ex => vec![Command::Execute(ex)],
         };
-        let group_cmd = Command::Group(then_cmd);
+        let group_cmd = Command::Group(then_cmds);
         let then_cmd = if group_cmd.forbid_prefix() {
             group_cmd
         } else {

src/datapack/command/execute/mod.rs

@@ -296,6 +296,29 @@ impl Execute {
             Self::Runs(cmds) => cmds.iter().flat_map(|cmd| cmd.get_macros()).collect(),
         }
     }
+
+    /// Check if the execute command contains a return command
+    pub(crate) fn contains_return(&self) -> bool {
+        match self {
+            Self::Align(_, next)
+            | Self::Anchored(_, next)
+            | Self::As(_, next)
+            | Self::At(_, next)
+            | Self::AsAt(_, next)
+            | Self::Facing(_, next)
+            | Self::In(_, next)
+            | Self::On(_, next)
+            | Self::Positioned(_, next)
+            | Self::Rotated(_, next)
+            | Self::Store(_, next)
+            | Self::Summon(_, next) => next.contains_return(),
+            Self::If(_, then, el) => {
+                then.contains_return() || el.as_deref().is_some_and(Self::contains_return)
+            }
+            Self::Run(cmd) => cmd.contains_return(),
+            Self::Runs(cmds) => cmds.iter().any(super::Command::contains_return),
+        }
+    }
 }
 
 impl From<Execute> for Command {

src/datapack/command/mod.rs

@@ -68,23 +68,7 @@ impl Command {
             Self::Comment(comment) => {
                 vec![CompiledCommand::new("#".to_string() + comment).with_forbid_prefix(true)]
             }
-            Self::Return(return_cmd) => match return_cmd {
+            Self::Return(return_cmd) => return_cmd.compile(options, global_state, function_state),
-                ReturnCommand::Value(value) => {
-                    vec![CompiledCommand::new(format!("return {}", value.compile()))]
-                }
-                ReturnCommand::Command(cmd) => {
-                    let compiled_cmd = Self::Group(vec![*cmd.clone()]).compile(
-                        options,
-                        global_state,
-                        function_state,
-                    );
-                    let compiled_cmd = compiled_cmd
-                        .into_iter()
-                        .next()
-                        .expect("group will always return exactly one command");
-                    vec![compiled_cmd.apply_prefix("return run ")]
-                }
-            },
             Self::Concat(a, b) => {
                 let a = a.compile(options, global_state, function_state);
                 let b = b.compile(options, global_state, function_state);
@@ -195,6 +179,20 @@ impl Command {
             Self::Concat(a, _) => a.forbid_prefix(),
         }
     }
+
+    // Check whether the command contains a return command.
+    #[must_use]
+    pub fn contains_return(&self) -> bool {
+        match self {
+            Self::Comment(_) | Self::Debug(_) => false,
+            Self::Return(_) => true,
+            Self::Concat(a, b) => a.contains_return() || b.contains_return(),
+            Self::Execute(exec) => exec.contains_return(),
+            Self::Raw(cmd) => cmd.starts_with("return "),
+            Self::UsesMacro(m) => m.compile().starts_with("return "),
+            Self::Group(g) => g.iter().any(Self::contains_return),
+        }
+    }
 }
 
 impl From<&str> for Command {
@@ -250,19 +248,11 @@ fn compile_group(
         0 => Vec::new(),
         1 => commands[0].compile(options, global_state, function_state),
         _ => {
-            let uid = function_state.request_uid();
             let pass_macros = group_contains_macro(commands);
+            let contains_return = commands.iter().any(Command::contains_return);
 
             // calculate a hashed path for the function in the `sb` subfolder
-            let function_path = {
+            let function_path = generate_group_function_path(function_state);
-                let function_path = function_state.path();
-                let function_path = function_path.strip_prefix("sb/").unwrap_or(function_path);
-
-                let pre_hash_path = function_path.to_owned() + ":" + &uid.to_string();
-                let hash = md5::hash(pre_hash_path).to_hex_lowercase();
-
-                "sb/".to_string() + function_path + "/" + &hash[..16]
-            };
 
             let namespace = function_state.namespace();
 
@@ -273,6 +263,66 @@ fn compile_group(
 
             let mut function_invocation = format!("function {namespace}:{function_path}");
 
+            let additional_return_cmds = if contains_return {
+                let full_path = format!("{namespace}:{function_path}");
+                let return_data_path = md5::hash(&full_path).to_hex_lowercase();
+
+                let pre_cmds = Command::Raw(format!(
+                    "data remove storage shulkerbox:return {return_data_path}"
+                ))
+                .compile(options, global_state, function_state)
+                .into_iter()
+                .map(|c| c.with_forbid_prefix(true))
+                .collect::<Vec<_>>();
+                let post_condition = Condition::Atom(
+                    format!("data storage shulkerbox:return {return_data_path}").into(),
+                );
+
+                let post_cmd_store = global_state
+                    .read()
+                    .unwrap()
+                    .functions_with_special_return
+                    .get(&format!(
+                        "{}:{}",
+                        function_state.namespace(),
+                        function_state.path()
+                    ))
+                    .cloned().map(|parent_return_data_path| {
+                        Command::Execute(Execute::If(
+                        post_condition.clone(),
+                        Box::new(Execute::Run(Box::new(Command::Raw(format!(
+                            "data modify storage shulkerbox:return {parent_return_data_path} set from storage shulkerbox:return {return_data_path}"
+                        ))))),
+                        None,
+                    ))
+                    });
+
+                let post_cmd_return = Command::Execute(Execute::If(
+                    post_condition,
+                    Box::new(Execute::Run(Box::new(Command::Raw(format!(
+                        "return run data get storage shulkerbox:return {return_data_path}"
+                    ))))),
+                    None,
+                ));
+
+                let post_cmds = post_cmd_store
+                    .into_iter()
+                    .chain(std::iter::once(post_cmd_return))
+                    .flat_map(|cmd| cmd.compile(options, global_state, function_state))
+                    .map(|c| c.with_forbid_prefix(true))
+                    .collect::<Vec<_>>();
+
+                global_state
+                    .write()
+                    .unwrap()
+                    .functions_with_special_return
+                    .insert(full_path, return_data_path);
+
+                Some((pre_cmds, post_cmds))
+            } else {
+                None
+            };
+
             if pass_macros {
                 // WARNING: this seems to be the only way to pass macros to the function called.
                 // Because everything is passed as a string, it looses one "level" of escaping per pass.
@@ -284,11 +334,30 @@ fn compile_group(
                 function_invocation.push_str(&format!(" {{{macros_block}}}"));
             }
 
-            vec![CompiledCommand::new(function_invocation).with_contains_macros(pass_macros)]
+            if let Some((mut pre_cmds, post_cmds)) = additional_return_cmds {
+                pre_cmds.push(
+                    CompiledCommand::new(function_invocation).with_contains_macros(pass_macros),
+                );
+                pre_cmds.extend(post_cmds);
+                pre_cmds
+            } else {
+                vec![CompiledCommand::new(function_invocation).with_contains_macros(pass_macros)]
+            }
         }
     }
 }
 
+fn generate_group_function_path(function_state: &FunctionCompilerState) -> String {
+    let uid = function_state.request_uid();
+    let function_path = function_state.path();
+    let function_path = function_path.strip_prefix("sb/").unwrap_or(function_path);
+
+    let pre_hash_path = function_path.to_owned() + ":" + &uid.to_string();
+    let hash = md5::hash(pre_hash_path).to_hex_lowercase();
+
+    "sb/".to_string() + function_path + "/" + &hash[..16]
+}
+
 fn group_contains_macro(commands: &[Command]) -> bool {
     commands.iter().any(Command::contains_macro)
 }
@@ -301,6 +370,66 @@ fn group_get_macros(commands: &[Command]) -> HashSet<&str> {
     macros
 }
 
+impl ReturnCommand {
+    pub fn compile(
+        &self,
+        options: &CompileOptions,
+        global_state: &MutCompilerState,
+        function_state: &FunctionCompilerState,
+    ) -> Vec<CompiledCommand> {
+        let return_data_path = global_state
+            .read()
+            .unwrap()
+            .functions_with_special_return
+            .get(&format!(
+                "{}:{}",
+                function_state.namespace(),
+                function_state.path()
+            ))
+            .cloned();
+        match (self, return_data_path) {
+            (Self::Value(value), None) => {
+                vec![CompiledCommand::new(format!("return {}", value.compile()))]
+            }
+            (Self::Value(value), Some(data_path)) => {
+                let value = value.compile();
+                let store_cmd = CompiledCommand::new(format!(
+                    "data modify storage shulkerbox:return {data_path} set value {value}",
+                ));
+                let return_cmd = CompiledCommand::new(format!("return {value}"));
+                vec![store_cmd, return_cmd]
+            }
+            (Self::Command(cmd), None) => {
+                let compiled_cmd = Command::Group(vec![*cmd.clone()]).compile(
+                    options,
+                    global_state,
+                    function_state,
+                );
+                let compiled_cmd = compiled_cmd
+                    .into_iter()
+                    .next()
+                    .expect("group will always return exactly one command");
+                vec![compiled_cmd.apply_prefix("return run ")]
+            }
+            (Self::Command(cmd), Some(data_path)) => {
+                let compiled_cmd = Command::Execute(Execute::Store(
+                    format!("result storage shulkerbox:return {data_path} int 1.0").into(),
+                    Box::new(Execute::Run(Box::new(Command::Group(vec![*cmd.clone()])))),
+                ))
+                .compile(options, global_state, function_state);
+                let compiled_cmd = compiled_cmd
+                    .into_iter()
+                    .next()
+                    .expect("group will always return exactly one command");
+                let return_cmd = CompiledCommand::new(format!(
+                    "return run data get storage shulkerbox:return {data_path} value"
+                ));
+                vec![compiled_cmd, return_cmd]
+            }
+        }
+    }
+}
+
 #[allow(clippy::too_many_lines)]
 fn validate_raw_cmd(cmd: &str, pack_formats: &RangeInclusive<u8>) -> bool {
     static CMD_FORMATS: OnceLock<HashMap<&str, RangeInclusive<u8>>> = OnceLock::new();
@@ -426,7 +555,7 @@ fn validate_raw_cmd(cmd: &str, pack_formats: &RangeInclusive<u8>) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Mutex;
+    use std::sync::RwLock;
 
     use crate::util::compile::CompilerState;
 
@@ -438,7 +567,7 @@ mod tests {
         let command_b = Command::raw("say foo bar");
 
         let options = &CompileOptions::default();
-        let global_state = &Mutex::new(CompilerState::default());
+        let global_state = &RwLock::new(CompilerState::default());
         let function_state = &FunctionCompilerState::default();
 
         assert_eq!(
@@ -458,7 +587,7 @@ mod tests {
         let comment = Command::Comment("this is a comment".to_string());
 
         let options = &CompileOptions::default();
-        let global_state = &Mutex::new(CompilerState::default());
+        let global_state = &RwLock::new(CompilerState::default());
         let function_state = &FunctionCompilerState::default();
 
         assert_eq!(

src/datapack/function.rs

@@ -98,7 +98,7 @@ mod tests {
     use super::*;
 
     use crate::util::compile::CompilerState;
-    use std::sync::Mutex;
+    use std::sync::RwLock;
 
     #[test]
     fn test_function() {
@@ -111,7 +111,7 @@ mod tests {
         assert_eq!(function.get_commands().len(), 1);
 
         let options = &CompileOptions::default();
-        let global_state = &Mutex::new(CompilerState::default());
+        let global_state = &RwLock::new(CompilerState::default());
         let function_state = &FunctionCompilerState::default();
 
         let compiled = function.compile(options, global_state, function_state);

src/datapack/mod.rs

@@ -8,7 +8,7 @@ pub use command::{Command, Condition, Execute, ReturnCommand};
 pub use function::Function;
 pub use namespace::Namespace;
 
-use std::{borrow::Cow, collections::BTreeMap, ops::RangeInclusive, sync::Mutex};
+use std::{borrow::Cow, collections::BTreeMap, ops::RangeInclusive, sync::RwLock};
 
 use crate::{
     util::compile::{CompileOptions, CompilerState, MutCompilerState},
@@ -158,7 +158,7 @@ impl Datapack {
             ..options.clone()
         };
 
-        let compiler_state = Mutex::new(CompilerState::default());
+        let compiler_state = RwLock::new(CompilerState::default());
 
         let mut root_folder = self.custom_files.clone();
         let mcmeta = generate_mcmeta(self, &options, &compiler_state);
@@ -290,7 +290,7 @@ mod tests {
     #[test]
     fn test_generate_mcmeta() {
         let dp = &Datapack::new("main", Datapack::LATEST_FORMAT).with_description("foo");
-        let state = Mutex::new(CompilerState::default());
+        let state = RwLock::new(CompilerState::default());
         let mcmeta = generate_mcmeta(dp, &CompileOptions::default(), &state);
 
         let json = if let VFile::Text(text) = mcmeta {

src/util/compile.rs

@@ -1,6 +1,11 @@
 //! Compile options for the compiler.
 
-use std::{fmt::Display, ops::Deref, sync::Mutex};
+use std::{
+    collections::HashMap,
+    fmt::Display,
+    ops::Deref,
+    sync::{Mutex, RwLock},
+};
 
 use getset::Getters;
 
@@ -52,9 +57,12 @@ impl Default for CompileOptions {
 #[allow(missing_copy_implementations)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Debug, Clone, Default)]
-pub struct CompilerState {}
+pub struct CompilerState {
+    /// Functions and their return value data path.
+    pub(crate) functions_with_special_return: HashMap<String, String>,
+}
 /// Mutex for the compiler state.
-pub type MutCompilerState = Mutex<CompilerState>;
+pub type MutCompilerState = RwLock<CompilerState>;
 
 /// State of the compiler for each function that can change during compilation.
 #[derive(Debug, Getters, Default)]