allow marking return value from lua as containing macro

src/transpile/lua.rs

@@ -10,8 +10,8 @@ mod enabled {
         base::{self, source_file::SourceElement, Handler},
         syntax::syntax_tree::expression::LuaCode,
         transpile::{
-            error::{LuaRuntimeError, TranspileError, TranspileResult},
-            expression::ComptimeValue,
+            error::{LuaRuntimeError, MismatchedTypes, TranspileError, TranspileResult},
+            expression::{ComptimeValue, ExpectedType},
             Scope, VariableData,
         },
     };
@@ -124,7 +124,7 @@ mod enabled {
                             table.set("path", lua.create_string(path)?)?;
                             Value::Table(table)
                         }
-                        Some(_) => todo!("allow other types"),
+                        Some(_) => todo!("allow other variable types"),
                         None => todo!("throw correct error"),
                     };
                     globals.set(name, value)?;
@@ -155,8 +155,68 @@ mod enabled {
                     handler,
                 ),
                 Value::Boolean(boolean) => Ok(Some(ComptimeValue::Boolean(boolean))),
+                Value::Table(table) => match table.get::<Value>("value") {
+                    Ok(Value::Nil) => {
+                        let err = TranspileError::LuaRuntimeError(LuaRuntimeError {
+                            code_block: self.span(),
+                            error_message: "return table must contain non-nil 'value'".to_string(),
+                        });
+                        handler.receive(err.clone());
+                        Err(err)
+                    }
+                    Ok(value) => {
+                        let value = match self.handle_lua_result(value, handler)? {
+                            Some(ComptimeValue::String(s)) => {
+                                let contains_macro = match table.get::<Value>("contains_macro") {
+                                    Ok(Value::Boolean(boolean)) => Ok(boolean),
+                                    Ok(value) => {
+                                        if let Some(ComptimeValue::Boolean(boolean)) =
+                                            self.handle_lua_result(value, handler)?
+                                        {
+                                            Ok(boolean)
+                                        } else {
+                                            let err =
+                                                TranspileError::MismatchedTypes(MismatchedTypes {
+                                                    expression: self.span(),
+                                                    expected_type: ExpectedType::Boolean,
+                                                });
+                                            handler.receive(err.clone());
+                                            Err(err)
+                                        }
+                                    }
+                                    _ => {
+                                        let err =
+                                            TranspileError::MismatchedTypes(MismatchedTypes {
+                                                expression: self.span(),
+                                                expected_type: ExpectedType::Boolean,
+                                            });
+                                        handler.receive(err.clone());
+                                        Err(err)
+                                    }
+                                }?;
+
+                                if contains_macro {
+                                    Some(ComptimeValue::MacroString(
+                                        s.parse().expect("parsing cannot fail"),
+                                    ))
+                                } else {
+                                    Some(ComptimeValue::String(s))
+                                }
+                            }
+                            value => value,
+                        };
+                        Ok(value)
+                    }
+                    Err(err) => {
+                        let err = TranspileError::LuaRuntimeError(LuaRuntimeError::from_lua_err(
+                            &err,
+                            self.span(),
+                        ));
+                        handler.receive(err.clone());
+                        Err(err)
+                    }
+                },
                 Value::Error(_)
-                | Value::Table(_)
                 | Value::Thread(_)
                 | Value::UserData(_)
                 | Value::LightUserData(_)

src/transpile/transpiler.rs

@@ -334,9 +334,6 @@ impl Transpiler {
 
             let statements = function_data.statements.clone();
 
-            let commands =
-                self.transpile_function(&statements, program_identifier, &function_scope, handler)?;
-
             let modified_name = function_data.annotations.get("deobfuscate").map_or_else(
                 || {
                     let hash_data = program_identifier.to_string() + "\0" + identifier_span.str();
@@ -370,6 +367,16 @@ impl Transpiler {
                 },
             )?;
 
+            let function_location = format!(
+                "{namespace}:{modified_name}",
+                namespace = function_data.namespace
+            );
+
+            function_path.set(function_location.clone()).unwrap();
+
+            let commands =
+                self.transpile_function(&statements, program_identifier, &function_scope, handler)?;
+
             let namespace = self.datapack.namespace_mut(&function_data.namespace);
 
             if namespace.function(&modified_name).is_some() {
@@ -384,19 +391,12 @@ impl Transpiler {
             let function = namespace.function_mut(&modified_name);
             function.get_commands_mut().extend(commands);
 
-            let function_location = format!(
-                "{namespace}:{modified_name}",
-                namespace = function_data.namespace
-            );
-
             if function_data.annotations.contains_key("tick") {
                 self.datapack.add_tick(&function_location);
             }
             if function_data.annotations.contains_key("load") {
                 self.datapack.add_load(&function_location);
             }
-
-            function_path.set(function_location.clone()).unwrap();
         }
 
         let parameters = &function_data.parameters;

src/transpile/util.rs

@@ -1,6 +1,6 @@
 //! Utility methods for transpiling
 
-use std::fmt::Display;
+use std::{fmt::Display, str::FromStr};
 
 use crate::{
     lexical::token::{MacroStringLiteral, MacroStringLiteralPart},
@@ -114,6 +114,61 @@ where
         })
 }
 
+impl FromStr for MacroString {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let pos = s.find("$(");
+        if pos.is_some_and(|pos| s[pos..].contains(')')) {
+            let mut parts = Vec::new();
+            let mut s = s;
+            while let Some(pos) = s.find("$(") {
+                let (before, after) = s.split_at(pos);
+
+                let last_macro_index = after
+                    .char_indices()
+                    .skip(2)
+                    .take_while(|&(_, c)| c.is_ascii_alphanumeric() || c == '_')
+                    .map(|(i, _)| i)
+                    .last();
+
+                match last_macro_index {
+                    Some(last_macro_index) if after[last_macro_index + 1..].starts_with(')') => {
+                        if !before.is_empty() {
+                            parts.push(MacroStringPart::String(before.to_string()));
+                        }
+                        parts.push(MacroStringPart::MacroUsage(
+                            after[2..=last_macro_index].to_string(),
+                        ));
+                        s = &after[last_macro_index + 2..];
+                        if s.is_empty() {
+                            break;
+                        }
+                    }
+                    _ => {
+                        parts.push(MacroStringPart::String(s.to_string()));
+                        s = "";
+                        break;
+                    }
+                }
+            }
+            if !s.is_empty() {
+                parts.push(MacroStringPart::String(s.to_string()));
+            }
+            if parts
+                .iter()
+                .any(|p| matches!(p, MacroStringPart::MacroUsage(_)))
+            {
+                Ok(Self::MacroString(parts))
+            } else {
+                Ok(Self::String(s.to_string()))
+            }
+        } else {
+            Ok(Self::String(s.to_string()))
+        }
+    }
+}
+
 impl<S> From<S> for MacroString
 where
     S: Into<String>,
@@ -172,3 +227,39 @@ impl From<MacroStringLiteral> for MacroString {
         Self::from(&value)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_macro_string() {
+        assert_eq!(
+            MacroString::from_str("Hello, $(world)!").unwrap(),
+            MacroString::MacroString(vec![
+                MacroStringPart::String("Hello, ".to_string()),
+                MacroStringPart::MacroUsage("world".to_string()),
+                MacroStringPart::String("!".to_string())
+            ])
+        );
+        assert_eq!(
+            MacroString::from_str("Hello, $(world)! $(world").unwrap(),
+            MacroString::MacroString(vec![
+                MacroStringPart::String("Hello, ".to_string()),
+                MacroStringPart::MacroUsage("world".to_string()),
+                MacroStringPart::String("! $(world".to_string()),
+            ])
+        );
+        assert_eq!(
+            MacroString::from_str("Hello $(a) from $(b) and $(c)").unwrap(),
+            MacroString::MacroString(vec![
+                MacroStringPart::String("Hello ".to_string()),
+                MacroStringPart::MacroUsage("a".to_string()),
+                MacroStringPart::String(" from ".to_string()),
+                MacroStringPart::MacroUsage("b".to_string()),
+                MacroStringPart::String(" and ".to_string()),
+                MacroStringPart::MacroUsage("c".to_string()),
+            ])
+        );
+    }
+}

src/transpile/variables.rs

@@ -129,6 +129,8 @@ impl<'a> Scope<'a> {
     }
 
     /// Gets the number of times a variable has been shadowed.
+    ///
+    ///
     pub fn get_variable_shadow_count(&self, name: &str) -> usize {
         let count = self
             .shadowed
@@ -138,9 +140,7 @@ impl<'a> Scope<'a> {
             .copied()
             .unwrap_or(0);
         self.parent.as_ref().map_or(count, |parent| {
-            count
-                + parent.get_variable_shadow_count(name)
-                + usize::from(parent.get_variable(name).is_some())
+            count.saturating_sub(1) + parent.get_variable_shadow_count(name)
         })
     }