Add special case for removing whitespace before macro if it has a trailing underscore

2 files changed, 45 insertions, 0 deletions

diff --git a/src/macro_processor/macro_processor.rs b/src/macro_processor/macro_processor.rs
index ad230c0..7839963 100644
--- a/src/macro_processor/macro_processor.rs
+++ b/src/macro_processor/macro_processor.rs
@@ -339,6 +339,20 @@ fn smp_builtin_format_time(
     Ok(format!("{}", dt.format(&args[0])))
 }
 
+fn macro_is_whitespace_deleting(s: &str) -> bool {
+    s.chars().nth(s.len() - 1) == Some('_')
+}
+
+fn macro_name_clean<'a>(macro_name: &'a str) -> &'a str {
+    let mut macro_name = macro_name;
+    if macro_is_whitespace_deleting(macro_name) {
+        let mut macro_chars = macro_name.chars();
+        macro_chars.next_back();
+        macro_name = macro_chars.as_str();
+    }
+    macro_name
+}
+
 /// Types of macros, this is to make it easy to store both functions and strings
 #[derive(Clone)]
 pub enum MacroType {
@@ -498,6 +512,9 @@ impl MacroProcessor {
     /// * `macro_name` - Name of macro to expand if it exists
     /// * `args` - List of arguments parsed along with macro invokation (empty list if no arguments were parsed)
     fn expand_macro(&mut self, macro_name: &str, args: &mut [String]) -> Result<String, SMPError> {
+        // Ignore trailing underscore in macro name, the parser will pop a space in front if
+        // present, but we should ignore it for finding the macro.
+        let macro_name = macro_name_clean(macro_name);
         let Some(macro_body) = self.macros.get(macro_name) else {
             if args.len() == 0 {
                 return Ok(format!("{}", macro_name));
@@ -600,6 +617,12 @@ impl MacroProcessor {
                         parens_level += 1;
                         state = ParserState::InMacroArgs;
                     } else {
+                        if macro_is_whitespace_deleting(&macro_name) {
+                            if output.chars().last() == Some(' ') {
+                                output.pop();
+                            }
+                            macro_name = macro_name_clean(&macro_name).to_string();
+                        }
                         if self.macros.contains_key(&macro_name) {
                             highlight_debug!("\x1b[32m\x1b[7m", input, (macro_name_start -> i));
                         }
@@ -622,6 +645,12 @@ impl MacroProcessor {
                 }
                 ParserState::InMacroArgs => {
                     if (c == ')') && (parens_level == 1) {
+                        if macro_is_whitespace_deleting(&macro_name) {
+                            if output.chars().last() == Some(' ') {
+                                output.pop();
+                            }
+                            macro_name = macro_name_clean(&macro_name).to_string();
+                        }
                         highlight_debug!("\x1b[32m\x1b[7m", input, (macro_name_start -> i));
                         parens_level = 0;
                         macro_args.push(argument.trim().to_string());
@@ -649,6 +678,12 @@ impl MacroProcessor {
 
         // Handle cases where the text ends with a macro without arguments
         if !macro_name.is_empty() {
+            if macro_is_whitespace_deleting(&macro_name) {
+                if output.chars().last() == Some(' ') {
+                    output.pop();
+                }
+                macro_name = macro_name_clean(&macro_name).to_string();
+            }
             output.push_str(&self.expand_macro(&macro_name, &mut [])?);
         }
 
diff --git a/tests/macro_processor.rs b/tests/macro_processor.rs
index cb21df5..c224b54 100644
--- a/tests/macro_processor.rs
+++ b/tests/macro_processor.rs
@@ -257,3 +257,13 @@ fn test_smp_expr_1() {
 ",
     );
 }
+
+#[test]
+fn test_smp_whitespace_deleting_1() {
+    assert_eq!(run_macro_processor("define(a, b) a"), " b",);
+}
+
+#[test]
+fn test_smp_whitespace_deleting_2() {
+    assert_eq!(run_macro_processor("define(a, b) a_"), "b",);
+}