Rework how we handle state and calling and returning

main.go

@@ -40,6 +40,7 @@ func main() {
 		fmt.Fprintf(os.Stderr, "   or direct assignment: %s \"<value>(\\d+)</value>\" \"=0\" data.xml\n", os.Args[0])
 		fmt.Fprintf(os.Stderr, "\nNote: v1, v2, etc. are used to refer to capture groups.\n")
 		fmt.Fprintf(os.Stderr, "      If expression starts with an operator like *, /, +, -, =, etc., v1 is automatically prepended\n")
+		fmt.Fprintf(os.Stderr, "      You can use any valid Lua code, including if statements, loops, etc.\n")
 	}
 
 	flag.Parse()
@@ -54,8 +55,8 @@ func main() {
 	luaExpr := args[1]
 	files := args[2:]
 
-	// Generate the Lua script
+	// Prepare the Lua expression
-	luaScript := buildLuaScript(luaExpr)
+	luaExpr = buildLuaScript(luaExpr)
 
 	if strings.Contains(regexPattern, "!num") {
 		regexPattern = strings.ReplaceAll(regexPattern, "!num", "(\\d*\\.?\\d+)")
@@ -75,7 +76,7 @@ func main() {
 
 	// Process each file
 	for _, file := range files {
-		err := processFile(file, pattern, luaScript)
+		err := processFile(file, pattern, luaExpr)
 		if err != nil {
 			Error.Printf("Error processing file %s: %v", file, err)
 		}
@@ -84,65 +85,32 @@ func main() {
 
 // buildLuaScript creates a complete Lua script from the expression
 func buildLuaScript(luaExpr string) string {
-	log.Printf("Building Lua script from expression: %s", luaExpr)
+	// Check if the expression starts with an operator that needs special handling
-
+	if strings.HasPrefix(luaExpr, "*") {
-	// Check if the expression needs v1 to be prepended
+		luaExpr = "v1 = v1" + luaExpr
-	if strings.HasPrefix(luaExpr, "*") || strings.HasPrefix(luaExpr, "/") ||
+	} else if strings.HasPrefix(luaExpr, "/") {
-		strings.HasPrefix(luaExpr, "+") || strings.HasPrefix(luaExpr, "-") ||
+		luaExpr = "v1 = v1" + luaExpr
-		strings.HasPrefix(luaExpr, "^") || strings.HasPrefix(luaExpr, "%") {
+	} else if strings.HasPrefix(luaExpr, "+") {
-		luaExpr = "v1" + luaExpr
+		luaExpr = "v1 = v1" + luaExpr
-		log.Printf("Prepended 'v1' to expression, new expression: %s", luaExpr)
+	} else if strings.HasPrefix(luaExpr, "-") {
+		luaExpr = "v1 = v1" + luaExpr
+	} else if strings.HasPrefix(luaExpr, "^") {
+		luaExpr = "v1 = v1" + luaExpr
+	} else if strings.HasPrefix(luaExpr, "%") {
+		luaExpr = "v1 = v1" + luaExpr
 	} else if strings.HasPrefix(luaExpr, "=") {
 		// Handle direct assignment with = operator
 		luaExpr = "v1 " + luaExpr
-		log.Printf("Handled direct assignment, new expression: %s", luaExpr)
 	}
 
-	// Replace shorthand v1, v2, etc. with their direct variable names (no need for array notation)
+	// Replace shorthand v1, v2, etc. with their direct variable names
-	// Note: we're keeping this regex in case the user writes v[1] style, but we'll convert it to v1
 	shorthandRegex := regexp.MustCompile(`\bv\[(\d+)\]\b`)
 	luaExpr = shorthandRegex.ReplaceAllString(luaExpr, "v$1")
-	log.Printf("Converted shorthand variables, new expression: %s", luaExpr)
 
-	// Add custom script header with helper functions
+	return luaExpr
-	scriptHeader := `
--- Custom Lua helpers for math operations
-function min(a, b) return math.min(a, b) end
-function max(a, b) return math.max(a, b) end
-function round(x) return math.floor(x + 0.5) end
-function floor(x) return math.floor(x) end
-function ceil(x) return math.ceil(x) end
-`
-
-	// Generate the full script
-	var fullScript string
-	if strings.Contains(luaExpr, "v1 =") || strings.Contains(luaExpr, "v2 =") ||
-		strings.Contains(luaExpr, "v3 =") || strings.Contains(luaExpr, "return") {
-		// Already has assignments, use as is
-		fullScript = fmt.Sprintf(`%s
-function process(...)
-    %s
-    return v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12
-end
-`, scriptHeader, luaExpr)
-		log.Println("Generated Lua script with assignments.")
-	} else {
-		// Simple expression, assign to v1
-		fullScript = fmt.Sprintf(`%s
-function process(...)
-    local result = %s
-    v1 = result
-    return v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12
-end
-`, scriptHeader, luaExpr)
-		log.Println("Generated Lua script for simple expression.")
-	}
-
-	log.Printf("Final Lua script: %s", fullScript)
-	return fullScript
 }
 
-func processFile(filename string, pattern *regexp.Regexp, luaScript string) error {
+func processFile(filename string, pattern *regexp.Regexp, luaExpr string) error {
 	fullPath := filepath.Join(".", filename)
 	log.Printf("Processing file: %s", fullPath)
 
@@ -154,7 +122,7 @@ func processFile(filename string, pattern *regexp.Regexp, luaScript string) erro
 	log.Printf("Successfully read file: %s, content length: %d bytes", fullPath, len(content))
 
 	fileContent := string(content)
-	result, err := process(fileContent, pattern, luaScript)
+	result, err := process(fileContent, pattern, luaExpr)
 	if err != nil {
 		Error.Printf("Error processing content of file %s: %v", fullPath, err)
 		return err
@@ -171,7 +139,7 @@ func processFile(filename string, pattern *regexp.Regexp, luaScript string) erro
 	return nil
 }
 
-func process(data string, pattern *regexp.Regexp, luaScript string) (string, error) {
+func process(data string, pattern *regexp.Regexp, luaExpr string) (string, error) {
 	L := lua.NewState()
 	defer L.Close()
 
@@ -182,22 +150,26 @@ func process(data string, pattern *regexp.Regexp, luaScript string) (string, err
 		Error.Printf("Error loading Lua math library: %v", err)
 		return data, fmt.Errorf("error loading Lua math library: %v", err)
 	}
-	log.Println("Lua math library loaded successfully.")
 
-	// Load the Lua script
+	// Initialize helper functions
-	log.Printf("Executing Lua script: %s", luaScript)
+	helperScript := `
-	if err := L.DoString(luaScript); err != nil {
+-- Custom Lua helpers for math operations
-		Error.Printf("Error in Lua script: %v", err)
+function min(a, b) return math.min(a, b) end
-		return data, fmt.Errorf("error in Lua script: %v", err)
+function max(a, b) return math.max(a, b) end
+function round(x) return math.floor(x + 0.5) end
+function floor(x) return math.floor(x) end
+function ceil(x) return math.ceil(x) end
+`
+	if err := L.DoString(helperScript); err != nil {
+		Error.Printf("Error loading helper functions: %v", err)
+		return data, fmt.Errorf("error loading helper functions: %v", err)
 	}
-	log.Printf("Lua script executed successfully. Script: %s", luaScript)
 
 	// Process all regex matches
 	result := pattern.ReplaceAllStringFunc(data, func(match string) string {
 		captures := pattern.FindStringSubmatch(match)
 		if len(captures) <= 1 {
 			// No capture groups, return unchanged
-			log.Printf("No capture groups found for match: %s", match)
 			return match
 		}
 
@@ -207,41 +179,29 @@ func process(data string, pattern *regexp.Regexp, luaScript string) (string, err
 			floatVal, err := strconv.ParseFloat(capture, 64)
 			if err == nil {
 				L.SetGlobal(fmt.Sprintf("v%d", i+1), lua.LNumber(floatVal))
-				log.Printf("Set v%d to %f (capture: %s)", i+1, floatVal, capture)
 			} else {
 				L.SetGlobal(fmt.Sprintf("v%d", i+1), lua.LString(capture))
-				log.Printf("Set v%d to %s (capture: %s)", i+1, capture, capture)
 			}
 		}
 
-		// Call the Lua function
+		// Execute the user's Lua code
-		L.Push(L.GetGlobal("process"))
+		if err := L.DoString(luaExpr); err != nil {
-		if err := L.PCall(0, 12, nil); err != nil { // We're returning up to 12 values
 			Error.Printf("Error executing Lua script: %v", err)
 			return match // Return unchanged on error
 		}
-		log.Println("Lua function 'process' called successfully.")
 
-		// Get the return values (up to 12)
+		// Get the values of v1-v12 after code execution
 		returnValues := make([]lua.LValue, 12)
 		for i := 0; i < 12; i++ {
-			// Get the value from the stack (if exists)
+			varName := fmt.Sprintf("v%d", i+1)
-			if L.GetTop() >= i+1 {
+			returnValues[i] = L.GetGlobal(varName)
-				returnValues[i] = L.Get(-12 + i)
-				log.Printf("Return value v%d: %v", i+1, returnValues[i])
-			} else {
-				returnValues[i] = lua.LNil
 		}
-		}
-		// Pop all return values
-		L.Pop(L.GetTop())
 
 		// Replace each capture group with its new value (if available)
 		result := match
 		for i := 0; i < len(captures)-1 && i < 12; i++ {
 			if returnValues[i] == lua.LNil {
-				log.Printf("Return value v%d is nil, skipping replacement.", i+1)
+				continue // Skip if nil
-				continue // Skip if nil return
 			}
 
 			oldVal := captures[i+1]
@@ -258,7 +218,6 @@ func process(data string, pattern *regexp.Regexp, luaScript string) (string, err
 
 			// If we have a value, replace it
 			if newVal != "" {
-				log.Printf("Replacing '%s' with '%s'", oldVal, newVal)
 				result = strings.Replace(result, oldVal, newVal, 1)
 			}
 		}
@@ -266,6 +225,5 @@ func process(data string, pattern *regexp.Regexp, luaScript string) (string, err
 		return result
 	})
 
-	log.Println("Processing completed successfully.")
 	return result, nil
 }

main_test.go

@@ -65,7 +65,7 @@ func TestShorthandNotation(t *testing.T) {
 	`
 
 	regex := regexp.MustCompile(`(?s)<value>(\d+)</value>`)
-	luaExpr := `v1 * 1.5`
+	luaExpr := `v1 = v1 * 1.5` // Use direct assignment syntax
 	luaScript := buildLuaScript(luaExpr)
 
 	modifiedContent, err := process(fileContents, regex, luaScript)
@@ -97,7 +97,7 @@ func TestArrayNotation(t *testing.T) {
 	`
 
 	regex := regexp.MustCompile(`(?s)<value>(\d+)</value>`)
-	luaExpr := `v1 * 1.5`
+	luaExpr := `v1 = v1 * 1.5` // Use direct assignment syntax
 	luaScript := buildLuaScript(luaExpr)
 
 	modifiedContent, err := process(fileContents, regex, luaScript)
@@ -172,7 +172,7 @@ func TestMultipleCaptureGroups(t *testing.T) {
 
 	// Use (?s) flag to match across multiple lines
 	regex := regexp.MustCompile(`(?s)<value>(\d+)</value>.*?<multiplier>(\d+)</multiplier>`)
-	luaExpr := `v1 * v2`
+	luaExpr := `v1 = v1 * v2` // Use direct assignment syntax
 	luaScript := buildLuaScript(luaExpr)
 
 	// Verify the regex matches before processing
@@ -256,7 +256,7 @@ func TestDecimalValues(t *testing.T) {
 	`
 
 	regex := regexp.MustCompile(`(?s)<value>([0-9.]+)</value>.*?<multiplier>([0-9.]+)</multiplier>`)
-	luaExpr := `v1 * v2`
+	luaExpr := `v1 = v1 * v2` // Use direct assignment syntax
 	luaScript := buildLuaScript(luaExpr)
 
 	modifiedContent, err := process(fileContents, regex, luaScript)
@@ -288,7 +288,7 @@ func TestLuaMathFunctions(t *testing.T) {
 	`
 
 	regex := regexp.MustCompile(`(?s)<value>(\d+)</value>`)
-	luaExpr := `math.sqrt(v1)`
+	luaExpr := `v1 = math.sqrt(v1)` // Use direct assignment syntax
 	luaScript := buildLuaScript(luaExpr)
 
 	modifiedContent, err := process(fileContents, regex, luaScript)
@@ -361,7 +361,7 @@ func TestProcessingSampleFiles(t *testing.T) {
 			name:         "Complex file - multiply values by multiplier and divide by divider",
 			fileContent:  string(complexFile),
 			regexPattern: `(?s)<value>(\d+)</value>.*?<multiplier>(\d+)</multiplier>.*?<divider>(\d+)</divider>`,
-			luaExpr:      `v1 * v2 / v3`,
+			luaExpr:      `v1 = v1 * v2 / v3`,
 			expectedFunc: func(content string) string {
 				// Replace values manually for verification
 				r := regexp.MustCompile(`(?s)<item>\s*<value>150</value>\s*<multiplier>2</multiplier>\s*<divider>4</divider>\s*</item>`)
@@ -377,7 +377,7 @@ func TestProcessingSampleFiles(t *testing.T) {
 			name:         "Test data - simple multiplication",
 			fileContent:  string(testDataFile),
 			regexPattern: `(?s)<test id="simple">.*?<value>(\d+)</value>.*?</test>`,
-			luaExpr:      `*1.5`,
+			luaExpr:      `v1 = v1 * 1.5`,
 			expectedFunc: func(content string) string {
 				r := regexp.MustCompile(`<value>100</value>`)
 				return r.ReplaceAllString(content, "<value>150</value>")
@@ -387,7 +387,7 @@ func TestProcessingSampleFiles(t *testing.T) {
 			name:         "Test data - multiple capture groups",
 			fileContent:  string(testDataFile),
 			regexPattern: `(?s)<test id="multi">.*?<value>(\d+)</value>.*?<multiplier>(\d+)</multiplier>.*?<divider>(\d+)</divider>.*?</test>`,
-			luaExpr:      `v1 * v2 / v3`,
+			luaExpr:      `v1 = v1 * v2 / v3`,
 			expectedFunc: func(content string) string {
 				r := regexp.MustCompile(`<value>50</value>`)
 				return r.ReplaceAllString(content, "<value>75</value>")
@@ -397,7 +397,7 @@ func TestProcessingSampleFiles(t *testing.T) {
 			name:         "Test data - decimal values",
 			fileContent:  string(testDataFile),
 			regexPattern: `(?s)<test id="decimal">.*?<value>([0-9.]+)</value>.*?<multiplier>([0-9.]+)</multiplier>.*?</test>`,
-			luaExpr:      `v1 * v2`,
+			luaExpr:      `v1 = v1 * v2`,
 			expectedFunc: func(content string) string {
 				r := regexp.MustCompile(`<value>10.5</value>`)
 				return r.ReplaceAllString(content, "<value>26.25</value>")
@@ -455,7 +455,7 @@ func TestFileOperations(t *testing.T) {
 
 		// Configure test
 		regexPattern := `(?s)<value>(\d+)</value>.*?<multiplier>(\d+)</multiplier>.*?<divider>(\d+)</divider>`
-		luaExpr := `v1 * v2 / v3`
+		luaExpr := `v1 = v1 * v2 / v3` // Use direct assignment
 
 		// Execute test
 		regex := regexp.MustCompile(regexPattern)
@@ -492,7 +492,7 @@ func TestFileOperations(t *testing.T) {
 
 		// Configure test for simple value
 		regexPattern := `(?s)<test id="simple">.*?<value>(\d+)</value>.*?</test>`
-		luaExpr := `*1.5`
+		luaExpr := `v1 = v1 * 1.5` // Use direct assignment
 
 		// Execute test
 		regex := regexp.MustCompile(regexPattern)
@@ -527,7 +527,7 @@ func TestFileOperations(t *testing.T) {
 
 		// Configure test for decimal values
 		regexPattern := `(?s)<test id="decimal">.*?<value>([0-9.]+)</value>.*?<multiplier>([0-9.]+)</multiplier>.*?</test>`
-		luaExpr := `v1 * v2`
+		luaExpr := `v1 = v1 * v2` // Use direct assignment
 
 		// Execute test
 		regex := regexp.MustCompile(regexPattern)
@@ -569,7 +569,7 @@ func TestHigherVariableIndices(t *testing.T) {
 	// Test using v3, v4, v5 in the expression
 	t.Run("Using v3-v5 variables", func(t *testing.T) {
 		regex := regexp.MustCompile(`(?s)<value1>(\d+)</value1>.*?<value2>(\d+)</value2>.*?<value3>(\d+)</value3>.*?<value4>(\d+)</value4>.*?<value5>(\d+)</value5>`)
-		luaExpr := `v1 + v2 * v3 / v4 - v5`
+		luaExpr := `v1 = v1 + v2 * v3 / v4 - v5`
 		luaScript := buildLuaScript(luaExpr)
 
 		// Expected: 10 + 20 * 30 / 40 - 50 = 10 + 15 - 50 = -25
@@ -589,7 +589,7 @@ func TestHigherVariableIndices(t *testing.T) {
 	// For double digit indexes, we need to capture it as the second variable (v2)
 	t.Run("Using v11 variable", func(t *testing.T) {
 		regex := regexp.MustCompile(`(?s)<value1>(\d+)</value1>.*?<value11>(\d+)</value11>`)
-		luaExpr := `v1 * v2`
+		luaExpr := `v1 = v1 * v2`
 		luaScript := buildLuaScript(luaExpr)
 
 		// Expected: 10 * 110 = 1100
@@ -609,7 +609,7 @@ func TestHigherVariableIndices(t *testing.T) {
 	t.Run("Using v0 variable", func(t *testing.T) {
 		// For this test, we'll capture the tag content and manipulate it
 		regex := regexp.MustCompile(`(?s)<value1>(\d+)</value1>`)
-		luaExpr := `tonumber(v1) * 2`
+		luaExpr := `v1 = tonumber(v1) * 2`
 		luaScript := buildLuaScript(luaExpr)
 
 		// This should double the value
@@ -624,3 +624,88 @@ func TestHigherVariableIndices(t *testing.T) {
 		}
 	})
 }
+
+func TestMultiStatementExpression(t *testing.T) {
+	fileContents := `
+		<config>
+			<item>
+				<value1>100</value1>
+				<value2>200</value2>
+			</item>
+		</config>
+	`
+	expected := `
+		<config>
+			<item>
+				<value1>0</value1>
+				<value2>0</value2>
+			</item>
+		</config>
+	`
+
+	regex := regexp.MustCompile(`(?s)<value1>(\d+)</value1>.*?<value2>(\d+)</value2>`)
+	luaExpr := `v1=0 v2=0` // Multiple statements without semicolons
+	luaScript := buildLuaScript(luaExpr)
+
+	t.Logf("Generated Lua script: %s", luaScript)
+
+	modifiedContent, err := process(fileContents, regex, luaScript)
+	if err != nil {
+		t.Fatalf("Error processing file: %v", err)
+	}
+
+	normalizedModified := normalizeWhitespace(modifiedContent)
+	normalizedExpected := normalizeWhitespace(expected)
+	if normalizedModified != normalizedExpected {
+		t.Fatalf("Expected modified content to be %q, but got %q", normalizedExpected, normalizedModified)
+	}
+}
+
+func TestComplexLuaScripts(t *testing.T) {
+	fileContents := `
+		<config>
+			<item>
+				<value1>100</value1>
+				<value2>200</value2>
+				<value3>50</value3>
+			</item>
+		</config>
+	`
+	expected := `
+		<config>
+			<item>
+				<value1>300</value1>
+				<value2>0</value2>
+				<value3>150</value3>
+			</item>
+		</config>
+	`
+
+	regex := regexp.MustCompile(`(?s)<value1>(\d+)</value1>.*?<value2>(\d+)</value2>.*?<value3>(\d+)</value3>`)
+	luaExpr := `
+local sum = v1 + v2
+if sum > 250 then
+    v1 = sum
+    v2 = 0
+    v3 = v3 * 3
+else
+    v1 = 0
+    v2 = sum
+    v3 = v3 * 2
+end
+`
+	luaScript := buildLuaScript(luaExpr)
+
+	t.Logf("Generated Lua script: %s", luaScript)
+
+	modifiedContent, err := process(fileContents, regex, luaScript)
+	if err != nil {
+		t.Fatalf("Error processing file: %v", err)
+	}
+
+	normalizedModified := normalizeWhitespace(modifiedContent)
+	normalizedExpected := normalizeWhitespace(expected)
+	if normalizedModified != normalizedExpected {
+		t.Fatalf("Expected modified content to be %q, but got %q", normalizedExpected, normalizedModified)
+	}
+}