Mako/compiler/tests/compiler_test.go

package compiler_test

import (
	"testing"

	"git.sharkk.net/Sharkk/Mako/compiler"
	"git.sharkk.net/Sharkk/Mako/parser"
)

// Helper function to compile source code and return chunk
func compileSource(t *testing.T, source string) *compiler.Chunk {
	lexer := parser.NewLexer(source)
	p := parser.NewParser(lexer)
	program := p.ParseProgram()

	if p.HasErrors() {
		t.Fatalf("Parser errors: %v", p.ErrorStrings())
	}

	comp := compiler.NewCompiler()
	chunk, errors := comp.Compile(program)

	if len(errors) > 0 {
		t.Fatalf("Compiler errors: %v", errors)
	}

	return chunk
}

// Helper to check instruction at position
func checkInstruction(t *testing.T, chunk *compiler.Chunk, pos int, expected compiler.Opcode, operands ...uint16) {
	if pos >= len(chunk.Code) {
		t.Fatalf("Position %d out of bounds (code length: %d)", pos, len(chunk.Code))
	}

	op, actualOperands, _ := compiler.DecodeInstruction(chunk.Code, pos)

	if op != expected {
		t.Errorf("Expected opcode %v at position %d, got %v", expected, pos, op)
	}

	if len(actualOperands) != len(operands) {
		t.Errorf("Expected %d operands, got %d", len(operands), len(actualOperands))
		return
	}

	for i, expected := range operands {
		if actualOperands[i] != expected {
			t.Errorf("Expected operand %d to be %d, got %d", i, expected, actualOperands[i])
		}
	}
}

// Test literal compilation with specialized opcodes
func TestNumberLiteral(t *testing.T) {
	chunk := compileSource(t, "echo 42")

	// Should have one constant (42) and load it
	if len(chunk.Constants) != 1 {
		t.Fatalf("Expected 1 constant, got %d", len(chunk.Constants))
	}

	if chunk.Constants[0].Type != compiler.ValueNumber {
		t.Errorf("Expected number constant, got %v", chunk.Constants[0].Type)
	}

	if chunk.Constants[0].Data.(float64) != 42.0 {
		t.Errorf("Expected constant value 42, got %v", chunk.Constants[0].Data)
	}

	// Check bytecode: OpLoadConst 0, OpEcho, OpReturnNil
	checkInstruction(t, chunk, 0, compiler.OpLoadConst, 0)
	checkInstruction(t, chunk, 3, compiler.OpEcho)
	checkInstruction(t, chunk, 4, compiler.OpReturnNil)
}

func TestSpecialNumbers(t *testing.T) {
	tests := []struct {
		source   string
		expected compiler.Opcode
	}{
		{"echo 0", compiler.OpLoadZero},
		{"echo 1", compiler.OpLoadOne},
	}

	for _, test := range tests {
		chunk := compileSource(t, test.source)

		// Should use specialized opcode with no constants
		if len(chunk.Constants) != 0 {
			t.Errorf("Expected 0 constants for %s, got %d", test.source, len(chunk.Constants))
		}

		checkInstruction(t, chunk, 0, test.expected)
		checkInstruction(t, chunk, 1, compiler.OpEcho)
		checkInstruction(t, chunk, 2, compiler.OpReturnNil)
	}
}

func TestStringLiteral(t *testing.T) {
	chunk := compileSource(t, `echo "hello"`)

	if len(chunk.Constants) != 1 {
		t.Fatalf("Expected 1 constant, got %d", len(chunk.Constants))
	}

	if chunk.Constants[0].Type != compiler.ValueString {
		t.Errorf("Expected string constant, got %v", chunk.Constants[0].Type)
	}

	if chunk.Constants[0].Data.(string) != "hello" {
		t.Errorf("Expected constant value 'hello', got %v", chunk.Constants[0].Data)
	}

	checkInstruction(t, chunk, 0, compiler.OpLoadConst, 0)
}

func TestBooleanLiterals(t *testing.T) {
	tests := []struct {
		source   string
		expected compiler.Opcode
	}{
		{"echo true", compiler.OpLoadTrue},
		{"echo false", compiler.OpLoadFalse},
	}

	for _, test := range tests {
		chunk := compileSource(t, test.source)

		// Should use specialized opcode with no constants
		if len(chunk.Constants) != 0 {
			t.Errorf("Expected 0 constants for %s, got %d", test.source, len(chunk.Constants))
		}

		checkInstruction(t, chunk, 0, test.expected)
		checkInstruction(t, chunk, 1, compiler.OpEcho)
		checkInstruction(t, chunk, 2, compiler.OpReturnNil)
	}
}

func TestNilLiteral(t *testing.T) {
	chunk := compileSource(t, "echo nil")

	// Should use specialized opcode with no constants
	if len(chunk.Constants) != 0 {
		t.Errorf("Expected 0 constants, got %d", len(chunk.Constants))
	}

	checkInstruction(t, chunk, 0, compiler.OpLoadNil)
	checkInstruction(t, chunk, 1, compiler.OpEcho)
	checkInstruction(t, chunk, 2, compiler.OpReturnNil)
}

// Test constant folding optimizations
func TestConstantFolding(t *testing.T) {
	// Test simple constants first (these should use specialized opcodes)
	simpleTests := []struct {
		source string
		opcode compiler.Opcode
	}{
		{"echo true", compiler.OpLoadTrue},
		{"echo false", compiler.OpLoadFalse},
		{"echo nil", compiler.OpLoadNil},
		{"echo 0", compiler.OpLoadZero},
		{"echo 1", compiler.OpLoadOne},
	}

	for _, test := range simpleTests {
		chunk := compileSource(t, test.source)
		checkInstruction(t, chunk, 0, test.opcode)
	}

	// Test arithmetic that should be folded (if folding is implemented)
	chunk := compileSource(t, "echo 2 + 3")

	// Check if folding occurred (single constant) or not (two constants + add)
	if len(chunk.Constants) == 1 {
		// Folding worked
		if chunk.Constants[0].Data.(float64) != 5.0 {
			t.Errorf("Expected folded constant 5, got %v", chunk.Constants[0].Data)
		}
	} else if len(chunk.Constants) == 2 {
		// No folding - should have Add instruction
		found := false
		for i := 0; i < len(chunk.Code); i++ {
			op, _, next := compiler.DecodeInstruction(chunk.Code, i)
			if op == compiler.OpAdd {
				found = true
				break
			}
			i = next - 1
		}
		if !found {
			t.Error("Expected OpAdd instruction when folding not implemented")
		}
	} else {
		t.Errorf("Unexpected number of constants: %d", len(chunk.Constants))
	}
}

// Test arithmetic operations (non-foldable)
func TestArithmetic(t *testing.T) {
	// Use variables to prevent constant folding
	chunk := compileSource(t, "x = 1\ny = 2\necho x + y")

	// Find the Add instruction
	found := false
	for i := 0; i < len(chunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpAdd {
			found = true
			break
		}
		i = next - 1
	}
	if !found {
		t.Error("Expected OpAdd instruction")
	}
}

// Test comparison operations
func TestComparison(t *testing.T) {
	tests := []struct {
		source   string
		expected compiler.Opcode
	}{
		{"x = 1\ny = 2\necho x == y", compiler.OpEq},
		{"x = 1\ny = 2\necho x != y", compiler.OpNeq},
		{"x = 1\ny = 2\necho x < y", compiler.OpLt},
		{"x = 1\ny = 2\necho x <= y", compiler.OpLte},
		{"x = 1\ny = 2\necho x > y", compiler.OpGt},
		{"x = 1\ny = 2\necho x >= y", compiler.OpGte},
	}

	for _, test := range tests {
		chunk := compileSource(t, test.source)

		// Find the comparison instruction
		found := false
		for i := 0; i < len(chunk.Code); i++ {
			op, _, next := compiler.DecodeInstruction(chunk.Code, i)
			if op == test.expected {
				found = true
				break
			}
			i = next - 1
		}
		if !found {
			t.Errorf("Expected %v instruction for %s", test.expected, test.source)
		}
	}
}

// Test prefix operations
func TestPrefixOperations(t *testing.T) {
	tests := []struct {
		source   string
		expected compiler.Opcode
	}{
		{"x = 42\necho -x", compiler.OpNeg},
		{"x = true\necho not x", compiler.OpNot},
	}

	for _, test := range tests {
		chunk := compileSource(t, test.source)

		// Find the prefix operation
		found := false
		for i := 0; i < len(chunk.Code); i++ {
			op, _, next := compiler.DecodeInstruction(chunk.Code, i)
			if op == test.expected {
				found = true
				break
			}
			i = next - 1
		}
		if !found {
			t.Errorf("Expected %v instruction for %s", test.expected, test.source)
		}
	}
}

// Test specialized local variable access
func TestSpecializedLocals(t *testing.T) {
	// This test needs to be within a function to have local variables
	chunk := compileSource(t, `
		fn test()
			a = 1
			b = 2
			c = 3
			echo a
			echo b
			echo c
		end
	`)

	// Check that function was compiled
	if len(chunk.Functions) == 0 {
		t.Skip("Function compilation not working")
	}

	funcChunk := &chunk.Functions[0].Chunk

	// Look for specialized local loads in the function
	specializedFound := 0
	for i := 0; i < len(funcChunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(funcChunk.Code, i)
		if op == compiler.OpLoadLocal0 || op == compiler.OpLoadLocal1 || op == compiler.OpLoadLocal2 {
			specializedFound++
		}
		i = next - 1
	}

	if specializedFound == 0 {
		t.Error("Expected specialized local access instructions")
	}
}

// Test variable assignment
func TestGlobalAssignment(t *testing.T) {
	chunk := compileSource(t, "x = 42")

	// Should have: LoadConst 0, StoreGlobal 1, OpReturnNil
	// Constants: [42, "x"]
	if len(chunk.Constants) != 2 {
		t.Fatalf("Expected 2 constants, got %d", len(chunk.Constants))
	}

	// Check that we have the number and variable name
	if chunk.Constants[0].Data.(float64) != 42.0 {
		t.Errorf("Expected first constant to be 42, got %v", chunk.Constants[0].Data)
	}
	if chunk.Constants[1].Data.(string) != "x" {
		t.Errorf("Expected second constant to be 'x', got %v", chunk.Constants[1].Data)
	}

	checkInstruction(t, chunk, 0, compiler.OpLoadConst, 0)   // Load 42
	checkInstruction(t, chunk, 3, compiler.OpStoreGlobal, 1) // Store to "x"
}

func TestZeroAssignment(t *testing.T) {
	chunk := compileSource(t, "x = 0")

	// Should use specialized zero loading
	if len(chunk.Constants) != 1 { // Only "x"
		t.Fatalf("Expected 1 constant, got %d", len(chunk.Constants))
	}

	if chunk.Constants[0].Data.(string) != "x" {
		t.Errorf("Expected constant to be 'x', got %v", chunk.Constants[0].Data)
	}

	checkInstruction(t, chunk, 0, compiler.OpLoadZero)       // Load 0
	checkInstruction(t, chunk, 1, compiler.OpStoreGlobal, 0) // Store to "x"
}

// Test echo statement
func TestEchoStatement(t *testing.T) {
	chunk := compileSource(t, "echo 42")

	// Should have: LoadConst 0, OpEcho, OpReturnNil
	checkInstruction(t, chunk, 0, compiler.OpLoadConst, 0)
	checkInstruction(t, chunk, 3, compiler.OpEcho)
	checkInstruction(t, chunk, 4, compiler.OpReturnNil)
}

// Test if statement
func TestIfStatement(t *testing.T) {
	chunk := compileSource(t, `
		if true then
			echo 1
		end
	`)

	// Should start with: LoadTrue, JumpIfFalse (with offset), Pop
	checkInstruction(t, chunk, 0, compiler.OpLoadTrue) // Load true (specialized)

	// JumpIfFalse has 1 operand (the jump offset)
	op, operands, _ := compiler.DecodeInstruction(chunk.Code, 1)
	if op != compiler.OpJumpIfFalse {
		t.Errorf("Expected OpJumpIfFalse at position 1, got %v", op)
	}
	if len(operands) != 1 {
		t.Errorf("Expected 1 operand for JumpIfFalse, got %d", len(operands))
	}

	checkInstruction(t, chunk, 4, compiler.OpPop) // Pop condition
}

// Test while loop with specialized loop instruction
func TestWhileLoop(t *testing.T) {
	chunk := compileSource(t, `
		while true do
			break
		end
	`)

	// Should have condition evaluation and loop structure
	checkInstruction(t, chunk, 0, compiler.OpLoadTrue) // Load true (specialized)

	// Should have LoopBack instruction instead of regular Jump
	found := false
	for i := 0; i < len(chunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpLoopBack {
			found = true
			break
		}
		i = next - 1
	}
	if !found {
		t.Error("Expected OpLoopBack instruction in while loop")
	}
}

// Test table creation
func TestTableLiteral(t *testing.T) {
	chunk := compileSource(t, "echo {1, 2, 3}")

	// Should start with OpNewTable
	checkInstruction(t, chunk, 0, compiler.OpNewTable)
}

// Test table with key-value pairs
func TestTableWithKeys(t *testing.T) {
	chunk := compileSource(t, `echo {x = 1, y = 2}`)

	checkInstruction(t, chunk, 0, compiler.OpNewTable)
	// Should have subsequent operations to set fields
}

// Test function call optimization
func TestFunctionCall(t *testing.T) {
	chunk := compileSource(t, "print(42)")

	// Should have: LoadGlobal "print", LoadConst 42, Call 1
	found := false
	for i := 0; i < len(chunk.Code)-2; i++ {
		op, operands, _ := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpCall && len(operands) > 0 && operands[0] == 1 {
			found = true
			break
		}
	}
	if !found {
		t.Error("Expected OpCall with 1 argument")
	}
}

// Test optimized local function calls
func TestLocalFunctionCall(t *testing.T) {
	chunk := compileSource(t, `
		fn test()
			f = print
			f(42)
		end
	`)

	if len(chunk.Functions) == 0 {
		t.Skip("Function compilation not working")
	}

	funcChunk := &chunk.Functions[0].Chunk

	// Look for optimized local call
	found := false
	for i := 0; i < len(funcChunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(funcChunk.Code, i)
		if op == compiler.OpCallLocal0 || op == compiler.OpCallLocal1 {
			found = true
			break
		}
		i = next - 1
	}

	if !found {
		t.Log("No optimized local call found (may be expected if function not in slot 0/1)")
	}
}

// Test constant deduplication
func TestConstantDeduplication(t *testing.T) {
	chunk := compileSource(t, "echo 42\necho 42\necho 42")

	// Should only have one constant despite multiple uses
	if len(chunk.Constants) != 1 {
		t.Errorf("Expected 1 constant (deduplicated), got %d", len(chunk.Constants))
	}
}

// Test specialized constant deduplication
func TestSpecializedConstantDeduplication(t *testing.T) {
	chunk := compileSource(t, "echo true\necho true\necho false\necho false")

	// Should have no constants - all use specialized opcodes
	if len(chunk.Constants) != 0 {
		t.Errorf("Expected 0 constants (all specialized), got %d", len(chunk.Constants))
	}

	// Count specialized instructions
	trueCount := 0
	falseCount := 0

	for i := 0; i < len(chunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpLoadTrue {
			trueCount++
		} else if op == compiler.OpLoadFalse {
			falseCount++
		}
		i = next - 1
	}

	if trueCount != 2 {
		t.Errorf("Expected 2 OpLoadTrue instructions, got %d", trueCount)
	}
	if falseCount != 2 {
		t.Errorf("Expected 2 OpLoadFalse instructions, got %d", falseCount)
	}
}

// Test short-circuit evaluation
func TestShortCircuitAnd(t *testing.T) {
	chunk := compileSource(t, "x = 1\ny = 2\necho x and y")

	// Should have conditional jumping for short-circuit
	found := false
	for i := 0; i < len(chunk.Code); i++ {
		op, _, _ := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpJumpIfFalse {
			found = true
			break
		}
	}
	if !found {
		t.Error("Expected JumpIfFalse for short-circuit and")
	}
}

func TestShortCircuitOr(t *testing.T) {
	chunk := compileSource(t, "x = 1\ny = 2\necho x or y")

	// Should have conditional jumping for short-circuit
	foundFalseJump := false
	foundJump := false
	for i := 0; i < len(chunk.Code); i++ {
		op, _, _ := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpJumpIfFalse {
			foundFalseJump = true
		}
		if op == compiler.OpJump {
			foundJump = true
		}
	}
	if !foundFalseJump || !foundJump {
		t.Error("Expected JumpIfFalse and Jump for short-circuit or")
	}
}

// Test increment optimization
func TestIncrementOptimization(t *testing.T) {
	chunk := compileSource(t, `
		fn test()
			x = 5
			y = x + 1
		end
	`)

	if len(chunk.Functions) == 0 {
		t.Skip("Function compilation not working")
	}

	funcChunk := &chunk.Functions[0].Chunk

	// Look for increment optimization (Inc instruction)
	found := false
	for i := 0; i < len(funcChunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(funcChunk.Code, i)
		if op == compiler.OpInc {
			found = true
			break
		}
		i = next - 1
	}

	if !found {
		t.Log("No increment optimization found (pattern may not match exactly)")
	}
}

// Test complex expressions (should prevent some folding)
func TestComplexExpression(t *testing.T) {
	chunk := compileSource(t, "x = 5\necho x + 2 * 3")

	// Should have constants: "x", and numbers for 2*3 (either 2,3 or folded 6)
	if len(chunk.Constants) < 2 {
		t.Errorf("Expected at least 2 constants, got %d", len(chunk.Constants))
	}

	// Check that we have the expected constant values
	hasVarX := false
	hasNumberConstant := false

	for _, constant := range chunk.Constants {
		switch constant.Type {
		case compiler.ValueNumber:
			val := constant.Data.(float64)
			if val == 5 || val == 2 || val == 3 || val == 6 {
				hasNumberConstant = true
			}
		case compiler.ValueString:
			if constant.Data.(string) == "x" {
				hasVarX = true
			}
		}
	}

	if !hasVarX {
		t.Error("Expected variable name 'x'")
	}
	if !hasNumberConstant {
		t.Error("Expected some numeric constant")
	}
}

// Test dead code elimination
func TestDeadCodeElimination(t *testing.T) {
	chunk := compileSource(t, `
		echo 1
		return
		echo 2
	`)

	// Look for NOOP instructions (dead code markers)
	noopCount := 0
	for i := 0; i < len(chunk.Code); i++ {
		op, _, next := compiler.DecodeInstruction(chunk.Code, i)
		if op == compiler.OpNoop {
			noopCount++
		}
		i = next - 1
	}

	if noopCount == 0 {
		t.Log("No dead code elimination detected (may depend on optimization level)")
	}
}