assignment implicit scope

parser/ast.go

@@ -37,11 +37,16 @@ func (p *Program) String() string {
 type AssignStatement struct {
 	Name          Expression // Changed from *Identifier to Expression for member access
 	Value         Expression
+	IsDeclaration bool // true if this is the first assignment in current scope
 }
 
 func (as *AssignStatement) statementNode() {}
 func (as *AssignStatement) String() string {
-	return fmt.Sprintf("%s = %s", as.Name.String(), as.Value.String())
+	prefix := ""
+	if as.IsDeclaration {
+		prefix = "local "
+	}
+	return fmt.Sprintf("%s%s = %s", prefix, as.Name.String(), as.Value.String())
 }
 
 // EchoStatement represents echo output statements

parser/parser.go

@@ -30,6 +30,10 @@ type Parser struct {
 	infixParseFns  map[TokenType]func(Expression) Expression
 
 	errors []ParseError
+
+	// Scope tracking
+	scopes     []map[string]bool // stack of scopes, each tracking declared variables
+	scopeTypes []string          // track what type each scope is: "global", "function", "loop"
 }
 
 // NewParser creates a new parser instance
@@ -37,6 +41,8 @@ func NewParser(lexer *Lexer) *Parser {
 	p := &Parser{
 		lexer:      lexer,
 		errors:     []ParseError{},
+		scopes:     []map[string]bool{make(map[string]bool)}, // start with global scope
+		scopeTypes: []string{"global"},                       // start with global scope type
 	}
 
 	p.prefixParseFns = make(map[TokenType]func() Expression)
@@ -73,6 +79,75 @@ func NewParser(lexer *Lexer) *Parser {
 	return p
 }
 
+// Scope management
+func (p *Parser) enterScope(scopeType string) {
+	p.scopes = append(p.scopes, make(map[string]bool))
+	p.scopeTypes = append(p.scopeTypes, scopeType)
+}
+
+func (p *Parser) exitScope() {
+	if len(p.scopes) > 1 { // never remove global scope
+		p.scopes = p.scopes[:len(p.scopes)-1]
+		p.scopeTypes = p.scopeTypes[:len(p.scopeTypes)-1]
+	}
+}
+
+func (p *Parser) enterFunctionScope() {
+	// Functions create new variable scopes
+	p.enterScope("function")
+}
+
+func (p *Parser) exitFunctionScope() {
+	p.exitScope()
+}
+
+func (p *Parser) enterLoopScope() {
+	// Create temporary scope for loop variables only
+	p.enterScope("loop")
+}
+
+func (p *Parser) exitLoopScope() {
+	// Remove temporary loop scope
+	p.exitScope()
+}
+
+func (p *Parser) enterBlockScope() {
+	// Blocks don't create new variable scopes, just control flow scopes
+	// We don't need to track these for variable declarations
+}
+
+func (p *Parser) exitBlockScope() {
+	// No-op since blocks don't create variable scopes
+}
+
+func (p *Parser) currentVariableScope() map[string]bool {
+	// If we're in a loop scope, declare variables in the parent scope
+	if len(p.scopeTypes) > 1 && p.scopeTypes[len(p.scopeTypes)-1] == "loop" {
+		return p.scopes[len(p.scopes)-2]
+	}
+	// Otherwise use the current scope
+	return p.scopes[len(p.scopes)-1]
+}
+
+func (p *Parser) isVariableDeclared(name string) bool {
+	// Check all scopes from current up to global
+	for i := len(p.scopes) - 1; i >= 0; i-- {
+		if p.scopes[i][name] {
+			return true
+		}
+	}
+	return false
+}
+
+func (p *Parser) declareVariable(name string) {
+	p.currentVariableScope()[name] = true
+}
+
+func (p *Parser) declareLoopVariable(name string) {
+	// Loop variables go in the current loop scope
+	p.scopes[len(p.scopes)-1][name] = true
+}
+
 // registerPrefix registers a prefix parse function
 func (p *Parser) registerPrefix(tokenType TokenType, fn func() Expression) {
 	p.prefixParseFns[tokenType] = fn
@@ -156,10 +231,17 @@ func (p *Parser) parseAssignStatement() *AssignStatement {
 		return nil
 	}
 
-	// Validate assignment target
+	// Validate assignment target and check if it's a declaration
-	switch stmt.Name.(type) {
+	switch name := stmt.Name.(type) {
-	case *Identifier, *DotExpression, *IndexExpression:
+	case *Identifier:
-		// Valid assignment targets
+		// Simple variable assignment - check if it's a declaration
+		stmt.IsDeclaration = !p.isVariableDeclared(name.Value)
+		if stmt.IsDeclaration {
+			p.declareVariable(name.Value)
+		}
+	case *DotExpression, *IndexExpression:
+		// Member access - never a declaration
+		stmt.IsDeclaration = false
 	default:
 		p.addError("invalid assignment target")
 		return nil
@@ -264,8 +346,10 @@ func (p *Parser) parseWhileStatement() *WhileStatement {
 
 	p.nextToken() // move past 'do'
 
-	// Parse loop body
+	// Parse loop body (no new variable scope)
+	p.enterBlockScope()
 	stmt.Body = p.parseBlockStatements(END)
+	p.exitBlockScope()
 
 	if !p.curTokenIs(END) {
 		p.addError("expected 'end' to close while loop")
@@ -349,8 +433,11 @@ func (p *Parser) parseNumericForStatement(variable *Identifier) *ForStatement {
 
 	p.nextToken() // move past 'do'
 
-	// Parse loop body
+	// Create temporary scope for loop variable, assignments in body go to parent scope
+	p.enterLoopScope()
+	p.declareLoopVariable(variable.Value) // loop variable in temporary scope
 	stmt.Body = p.parseBlockStatements(END)
+	p.exitLoopScope() // discard temporary scope with loop variable
 
 	if !p.curTokenIs(END) {
 		p.addError("expected 'end' to close for loop")
@@ -402,8 +489,14 @@ func (p *Parser) parseForInStatement(firstVar *Identifier) *ForInStatement {
 
 	p.nextToken() // move past 'do'
 
-	// Parse loop body
+	// Create temporary scope for loop variables, assignments in body go to parent scope
+	p.enterLoopScope()
+	if stmt.Key != nil {
+		p.declareLoopVariable(stmt.Key.Value) // loop variable in temporary scope
+	}
+	p.declareLoopVariable(stmt.Value.Value) // loop variable in temporary scope
 	stmt.Body = p.parseBlockStatements(END)
+	p.exitLoopScope() // discard temporary scope with loop variables
 
 	if !p.curTokenIs(END) {
 		p.addError("expected 'end' to close for loop")
@@ -432,14 +525,16 @@ func (p *Parser) parseIfStatement() *IfStatement {
 
 	p.nextToken() // move past condition (and optional 'then')
 
-	// Check if we immediately hit END (missing body)
+	// Check if we immediately hit END (empty body should be an error)
 	if p.curTokenIs(END) {
 		p.addError("expected 'end' to close if statement")
 		return nil
 	}
 
-	// Parse if body
+	// Parse if body (no new variable scope)
+	p.enterBlockScope()
 	stmt.Body = p.parseBlockStatements(ELSEIF, ELSE, END)
+	p.exitBlockScope()
 
 	// Parse elseif clauses
 	for p.curTokenIs(ELSEIF) {
@@ -460,14 +555,22 @@ func (p *Parser) parseIfStatement() *IfStatement {
 
 		p.nextToken() // move past condition (and optional 'then')
 
+		// Parse elseif body (no new variable scope)
+		p.enterBlockScope()
 		elseif.Body = p.parseBlockStatements(ELSEIF, ELSE, END)
+		p.exitBlockScope()
+
 		stmt.ElseIfs = append(stmt.ElseIfs, elseif)
 	}
 
 	// Parse else clause
 	if p.curTokenIs(ELSE) {
 		p.nextToken() // move past 'else'
+
+		// Parse else body (no new variable scope)
+		p.enterBlockScope()
 		stmt.Else = p.parseBlockStatements(END)
+		p.exitBlockScope()
 	}
 
 	if !p.curTokenIs(END) {
@@ -659,8 +762,13 @@ func (p *Parser) parseFunctionLiteral() Expression {
 
 	p.nextToken() // move past ')'
 
-	// Parse function body
+	// Enter new function scope and declare parameters
+	p.enterFunctionScope()
+	for _, param := range fn.Parameters {
+		p.declareVariable(param)
+	}
 	fn.Body = p.parseBlockStatements(END)
+	p.exitFunctionScope()
 
 	if !p.curTokenIs(END) {
 		p.addError("expected 'end' to close function")

parser/tests/scope_test.go

@@ -0,0 +1,598 @@
+package parser_test
+
+import (
+	"strings"
+	"testing"
+
+	"git.sharkk.net/Sharkk/Mako/parser"
+)
+
+func TestBasicScopeTracking(t *testing.T) {
+	tests := []struct {
+		input       string
+		assignments []struct {
+			variable      string
+			isDeclaration bool
+		}
+		desc string
+	}{
+		{
+			"x = 5",
+			[]struct {
+				variable      string
+				isDeclaration bool
+			}{
+				{"x", true}, // first assignment to x in global scope
+			},
+			"single global declaration",
+		},
+		{
+			`x = 5
+x = 10`,
+			[]struct {
+				variable      string
+				isDeclaration bool
+			}{
+				{"x", true},  // declaration
+				{"x", false}, // assignment to existing
+			},
+			"global declaration then assignment",
+		},
+		{
+			`x = 5
+y = 10
+x = 15`,
+			[]struct {
+				variable      string
+				isDeclaration bool
+			}{
+				{"x", true},  // declaration
+				{"y", true},  // declaration
+				{"x", false}, // assignment to existing
+			},
+			"multiple variables with reassignment",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			l := parser.NewLexer(tt.input)
+			p := parser.NewParser(l)
+			program := p.ParseProgram()
+			checkParserErrors(t, p)
+
+			assignmentCount := 0
+			for _, stmt := range program.Statements {
+				if assign, ok := stmt.(*parser.AssignStatement); ok {
+					if assignmentCount >= len(tt.assignments) {
+						t.Fatalf("more assignments than expected")
+					}
+
+					expected := tt.assignments[assignmentCount]
+
+					// Check variable name
+					ident, ok := assign.Name.(*parser.Identifier)
+					if !ok {
+						t.Fatalf("expected Identifier, got %T", assign.Name)
+					}
+
+					if ident.Value != expected.variable {
+						t.Errorf("assignment %d: expected variable %s, got %s",
+							assignmentCount, expected.variable, ident.Value)
+					}
+
+					// Check declaration status
+					if assign.IsDeclaration != expected.isDeclaration {
+						t.Errorf("assignment %d (%s): expected IsDeclaration = %t, got %t",
+							assignmentCount, expected.variable, expected.isDeclaration, assign.IsDeclaration)
+					}
+
+					assignmentCount++
+				}
+			}
+
+			if assignmentCount != len(tt.assignments) {
+				t.Errorf("expected %d assignments, found %d", len(tt.assignments), assignmentCount)
+			}
+		})
+	}
+}
+
+func TestBlockScopeTracking(t *testing.T) {
+	input := `x = 5
+if true then
+	y = 10
+	x = 15
+	z = 20
+end
+y = 25
+z = 30`
+
+	l := parser.NewLexer(input)
+	p := parser.NewParser(l)
+	program := p.ParseProgram()
+	checkParserErrors(t, p)
+
+	// With function/global scoping, all variables are declared at global level
+	expectedAssignments := []struct {
+		variable      string
+		isDeclaration bool
+		location      string
+	}{
+		{"x", true, "global"},  // x = 5
+		{"y", true, "global"},  // y = 10 (first in global, even though in if block)
+		{"x", false, "global"}, // x = 15 (exists in global)
+		{"z", true, "global"},  // z = 20 (first in global, even though in if block)
+		{"y", false, "global"}, // y = 25 (already declared in global)
+		{"z", false, "global"}, // z = 30 (already declared in global)
+	}
+
+	assignments := extractAssignments(program)
+
+	if len(assignments) != len(expectedAssignments) {
+		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
+	}
+
+	for i, expected := range expectedAssignments {
+		assign := assignments[i]
+		ident, ok := assign.Name.(*parser.Identifier)
+		if !ok {
+			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Name)
+		}
+
+		if ident.Value != expected.variable {
+			t.Errorf("assignment %d: expected variable %s, got %s",
+				i, expected.variable, ident.Value)
+		}
+
+		if assign.IsDeclaration != expected.isDeclaration {
+			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
+				i, expected.variable, expected.location, expected.isDeclaration, assign.IsDeclaration)
+		}
+	}
+}
+
+func TestNestedBlockScopes(t *testing.T) {
+	input := `x = 1
+if true then
+	y = 2
+	if true then
+		z = 3
+		x = 4
+		y = 5
+	end
+	z = 6
+end`
+
+	l := parser.NewLexer(input)
+	p := parser.NewParser(l)
+	program := p.ParseProgram()
+	checkParserErrors(t, p)
+
+	// With function/global scoping, all variables are declared at global level
+	expectedAssignments := []struct {
+		variable      string
+		isDeclaration bool
+		scope         string
+	}{
+		{"x", true, "global"},  // x = 1
+		{"y", true, "global"},  // y = 2 (first in global)
+		{"z", true, "global"},  // z = 3 (first in global)
+		{"x", false, "global"}, // x = 4 (exists in global)
+		{"y", false, "global"}, // y = 5 (exists in global)
+		{"z", false, "global"}, // z = 6 (exists in global)
+	}
+
+	assignments := extractAssignments(program)
+
+	if len(assignments) != len(expectedAssignments) {
+		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
+	}
+
+	for i, expected := range expectedAssignments {
+		assign := assignments[i]
+		ident, ok := assign.Name.(*parser.Identifier)
+		if !ok {
+			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Name)
+		}
+
+		if ident.Value != expected.variable {
+			t.Errorf("assignment %d: expected variable %s, got %s",
+				i, expected.variable, ident.Value)
+		}
+
+		if assign.IsDeclaration != expected.isDeclaration {
+			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
+				i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
+		}
+	}
+}
+
+func TestFunctionScopeTracking(t *testing.T) {
+	input := `x = 1
+callback = fn(a, b)
+	c = a + b
+	x = 10
+	return c
+end
+c = 20`
+
+	l := parser.NewLexer(input)
+	p := parser.NewParser(l)
+	program := p.ParseProgram()
+	checkParserErrors(t, p)
+
+	// Extract assignments from the program
+	assignments := extractAssignments(program)
+
+	expectedAssignments := []struct {
+		variable      string
+		isDeclaration bool
+		scope         string
+	}{
+		{"x", true, "global"},        // x = 1
+		{"callback", true, "global"}, // callback = fn...
+		{"c", true, "function"},      // c = a + b (first in function scope)
+		{"x", false, "function"},     // x = 10 (exists in global)
+		{"c", true, "global"},        // c = 20 (c was local to function)
+	}
+
+	if len(assignments) != len(expectedAssignments) {
+		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
+	}
+
+	for i, expected := range expectedAssignments {
+		assign := assignments[i]
+		ident, ok := assign.Name.(*parser.Identifier)
+		if !ok {
+			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Name)
+		}
+
+		if ident.Value != expected.variable {
+			t.Errorf("assignment %d: expected variable %s, got %s",
+				i, expected.variable, ident.Value)
+		}
+
+		if assign.IsDeclaration != expected.isDeclaration {
+			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
+				i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
+		}
+	}
+}
+
+func TestLoopScopeTracking(t *testing.T) {
+	tests := []struct {
+		input       string
+		desc        string
+		assignments []struct {
+			variable      string
+			isDeclaration bool
+			scope         string
+		}
+	}{
+		{
+			`x = 1
+for i = 1, 10 do
+	y = i
+	x = 5
+end
+y = 20`,
+			"numeric for loop",
+			[]struct {
+				variable      string
+				isDeclaration bool
+				scope         string
+			}{
+				{"x", true, "global"},  // x = 1
+				{"y", true, "global"},  // y = i (first in global, even though in loop)
+				{"x", false, "global"}, // x = 5 (exists in global)
+				{"y", false, "global"}, // y = 20 (already declared in global)
+			},
+		},
+		{
+			`arr = {1, 2, 3}
+for k, v in arr do
+	sum = sum + v
+	arr = nil
+end
+sum = 0`,
+			"for-in loop",
+			[]struct {
+				variable      string
+				isDeclaration bool
+				scope         string
+			}{
+				{"arr", true, "global"},  // arr = {1, 2, 3}
+				{"sum", true, "global"},  // sum = sum + v (first in global)
+				{"arr", false, "global"}, // arr = nil (exists in global)
+				{"sum", false, "global"}, // sum = 0 (already declared in global)
+			},
+		},
+		{
+			`running = true
+while running do
+	count = count + 1
+	running = false
+end
+count = 0`,
+			"while loop",
+			[]struct {
+				variable      string
+				isDeclaration bool
+				scope         string
+			}{
+				{"running", true, "global"},  // running = true
+				{"count", true, "global"},    // count = count + 1 (first in global)
+				{"running", false, "global"}, // running = false (exists in global)
+				{"count", false, "global"},   // count = 0 (already declared in global)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			l := parser.NewLexer(tt.input)
+			p := parser.NewParser(l)
+			program := p.ParseProgram()
+			checkParserErrors(t, p)
+
+			assignments := extractAssignments(program)
+
+			if len(assignments) != len(tt.assignments) {
+				t.Fatalf("expected %d assignments, got %d", len(tt.assignments), len(assignments))
+			}
+
+			for i, expected := range tt.assignments {
+				assign := assignments[i]
+				ident, ok := assign.Name.(*parser.Identifier)
+				if !ok {
+					t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Name)
+				}
+
+				if ident.Value != expected.variable {
+					t.Errorf("assignment %d: expected variable %s, got %s",
+						i, expected.variable, ident.Value)
+				}
+
+				if assign.IsDeclaration != expected.isDeclaration {
+					t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
+						i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
+				}
+			}
+		})
+	}
+}
+
+func TestMemberAccessNotDeclaration(t *testing.T) {
+	input := `table = {}
+table.key = "value"
+table["index"] = 42
+arr = {1, 2, 3}
+arr[1] = 99`
+
+	l := parser.NewLexer(input)
+	p := parser.NewParser(l)
+	program := p.ParseProgram()
+	checkParserErrors(t, p)
+
+	expectedAssignments := []struct {
+		variable       string
+		isDeclaration  bool
+		isMemberAccess bool
+	}{
+		{"table", true, false}, // table = {} (declaration)
+		{"", false, true},      // table.key = "value" (member access, never declaration)
+		{"", false, true},      // table["index"] = 42 (member access, never declaration)
+		{"arr", true, false},   // arr = {1, 2, 3} (declaration)
+		{"", false, true},      // arr[1] = 99 (member access, never declaration)
+	}
+
+	assignmentCount := 0
+	for _, stmt := range program.Statements {
+		if assign, ok := stmt.(*parser.AssignStatement); ok {
+			if assignmentCount >= len(expectedAssignments) {
+				t.Fatalf("more assignments than expected")
+			}
+
+			expected := expectedAssignments[assignmentCount]
+
+			if expected.isMemberAccess {
+				// Should not be an identifier
+				if _, ok := assign.Name.(*parser.Identifier); ok {
+					t.Errorf("assignment %d: expected member access, got Identifier", assignmentCount)
+				}
+
+				// Member access should never be a declaration
+				if assign.IsDeclaration {
+					t.Errorf("assignment %d: member access should never be declaration", assignmentCount)
+				}
+			} else {
+				// Should be an identifier
+				ident, ok := assign.Name.(*parser.Identifier)
+				if !ok {
+					t.Errorf("assignment %d: expected Identifier, got %T", assignmentCount, assign.Name)
+				} else if ident.Value != expected.variable {
+					t.Errorf("assignment %d: expected variable %s, got %s",
+						assignmentCount, expected.variable, ident.Value)
+				}
+
+				if assign.IsDeclaration != expected.isDeclaration {
+					t.Errorf("assignment %d (%s): expected IsDeclaration = %t, got %t",
+						assignmentCount, expected.variable, expected.isDeclaration, assign.IsDeclaration)
+				}
+			}
+
+			assignmentCount++
+		}
+	}
+}
+
+func TestComplexScopeScenario(t *testing.T) {
+	input := `global_var = "global"
+counter = 0
+
+callback = fn(x)
+	local_var = x * 2
+	global_var = "modified"
+	
+	if local_var > 10 then
+		temp = "high"
+		counter = counter + 1
+	else
+		temp = "low"
+		counter = counter - 1
+	end
+	
+	for i = 1, 3 do
+		temp = temp + i
+		local_var = local_var + i
+	end
+	
+	return local_var
+end
+
+temp = "global_temp"
+local_var = "global_local"`
+
+	l := parser.NewLexer(input)
+	p := parser.NewParser(l)
+	program := p.ParseProgram()
+	checkParserErrors(t, p)
+
+	expectedAssignments := []struct {
+		variable      string
+		isDeclaration bool
+		scope         string
+	}{
+		{"global_var", true, "global"},    // global_var = "global"
+		{"counter", true, "global"},       // counter = 0
+		{"callback", true, "global"},      // callback = fn...
+		{"local_var", true, "function"},   // local_var = x * 2 (first in function)
+		{"global_var", false, "function"}, // global_var = "modified" (exists in global)
+		{"temp", true, "function"},        // temp = "high" (first in function scope)
+		{"counter", false, "function"},    // counter = counter + 1 (exists in global)
+		{"temp", false, "function"},       // temp = "low" (already declared in function)
+		{"counter", false, "function"},    // counter = counter - 1 (exists in global)
+		{"temp", false, "function"},       // temp = temp + i (already declared in function)
+		{"local_var", false, "function"},  // local_var = local_var + i (exists in function)
+		{"temp", true, "global"},          // temp = "global_temp" (temp was local to function)
+		{"local_var", true, "global"},     // local_var = "global_local" (local_var was local to function)
+	}
+
+	assignments := extractAssignments(program)
+
+	if len(assignments) != len(expectedAssignments) {
+		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
+	}
+
+	for i, expected := range expectedAssignments {
+		assign := assignments[i]
+		ident, ok := assign.Name.(*parser.Identifier)
+		if !ok {
+			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Name)
+		}
+
+		if ident.Value != expected.variable {
+			t.Errorf("assignment %d: expected variable %s, got %s",
+				i, expected.variable, ident.Value)
+		}
+
+		if assign.IsDeclaration != expected.isDeclaration {
+			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
+				i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
+		}
+	}
+}
+
+func TestScopeStringRepresentation(t *testing.T) {
+	input := `x = 5
+if true then
+	y = 10
+	x = 15
+end
+y = 20`
+
+	l := parser.NewLexer(input)
+	p := parser.NewParser(l)
+	program := p.ParseProgram()
+	checkParserErrors(t, p)
+
+	// With function/global scoping:
+	// x = 5: declaration in global
+	// y = 10: declaration in global (first occurrence)
+	// x = 15: assignment in global (already exists)
+	// y = 20: assignment in global (already exists)
+	expected := `local x = 5.00
+if true then
+	local y = 10.00
+	x = 15.00
+end
+y = 20.00`
+
+	result := strings.TrimSpace(program.String())
+	if result != expected {
+		t.Errorf("expected:\n%s\n\ngot:\n%s", expected, result)
+	}
+}
+
+// Helper function to extract all assignments from a program recursively
+func extractAssignments(program *parser.Program) []*parser.AssignStatement {
+	var assignments []*parser.AssignStatement
+
+	for _, stmt := range program.Statements {
+		assignments = append(assignments, extractAssignmentsFromStatement(stmt)...)
+	}
+
+	return assignments
+}
+
+func extractAssignmentsFromStatement(stmt parser.Statement) []*parser.AssignStatement {
+	var assignments []*parser.AssignStatement
+
+	switch s := stmt.(type) {
+	case *parser.AssignStatement:
+		assignments = append(assignments, s)
+
+		// Check if the value is a function literal with assignments in body
+		if fn, ok := s.Value.(*parser.FunctionLiteral); ok {
+			for _, bodyStmt := range fn.Body {
+				assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+			}
+		}
+
+	case *parser.IfStatement:
+		// Extract from if body
+		for _, bodyStmt := range s.Body {
+			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+		}
+
+		// Extract from elseif bodies
+		for _, elseif := range s.ElseIfs {
+			for _, bodyStmt := range elseif.Body {
+				assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+			}
+		}
+
+		// Extract from else body
+		for _, bodyStmt := range s.Else {
+			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+		}
+
+	case *parser.WhileStatement:
+		for _, bodyStmt := range s.Body {
+			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+		}
+
+	case *parser.ForStatement:
+		for _, bodyStmt := range s.Body {
+			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+		}
+
+	case *parser.ForInStatement:
+		for _, bodyStmt := range s.Body {
+			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
+		}
+	}
+
+	return assignments
+}