Mako/tests/compiler_test.go

package tests

import (
	"testing"

	assert "git.sharkk.net/Go/Assert"
	"git.sharkk.net/Sharkk/Mako/compiler"
	"git.sharkk.net/Sharkk/Mako/lexer"
	"git.sharkk.net/Sharkk/Mako/parser"
	"git.sharkk.net/Sharkk/Mako/types"
)

func TestCompileConstants(t *testing.T) {
	input := `
	5;
	"hello";
	true;
	false;
	`

	lex := lexer.New(input)
	p := parser.New(lex)
	program := p.ParseProgram()

	bytecode := compiler.Compile(program)

	// Check that we have the right constants
	assert.Equal(t, 4, len(bytecode.Constants))

	// Check constant values
	assert.Equal(t, 5.0, bytecode.Constants[0])
	assert.Equal(t, "hello", bytecode.Constants[1])
	assert.Equal(t, true, bytecode.Constants[2])
	assert.Equal(t, false, bytecode.Constants[3])

	// Check that we have the right instructions
	assert.Equal(t, 10, len(bytecode.Instructions)) // 2 scope + 4 constants + 4 pops

	// Check instruction opcodes - first instruction is enter scope
	assert.Equal(t, types.OpEnterScope, bytecode.Instructions[0].Opcode)

	// Check the constant loading and popping instructions
	for i := 0; i < 4; i++ {
		assert.Equal(t, types.OpConstant, bytecode.Instructions[1+i*2].Opcode)
		assert.Equal(t, i, bytecode.Instructions[1+i*2].Operand)
		assert.Equal(t, types.OpPop, bytecode.Instructions[2+i*2].Opcode)
	}

	// Last instruction is exit scope
	assert.Equal(t, types.OpExitScope, bytecode.Instructions[len(bytecode.Instructions)-1].Opcode)
}

func TestCompileVariableAssignment(t *testing.T) {
	input := `x = 5;`

	lex := lexer.New(input)
	p := parser.New(lex)
	program := p.ParseProgram()

	bytecode := compiler.Compile(program)

	// Constants should be: 5, "x"
	assert.Equal(t, 2, len(bytecode.Constants))
	assert.Equal(t, 5.0, bytecode.Constants[0])
	assert.Equal(t, "x", bytecode.Constants[1])

	// Instructions should be:
	// - OpEnterScope
	// - OpConstant (5)
	// - OpSetGlobal with operand pointing to "x"
	// - OpExitScope
	assert.Equal(t, 4, len(bytecode.Instructions))

	assert.Equal(t, types.OpEnterScope, bytecode.Instructions[0].Opcode)
	assert.Equal(t, types.OpConstant, bytecode.Instructions[1].Opcode)
	assert.Equal(t, 0, bytecode.Instructions[1].Operand)
	assert.Equal(t, types.OpSetGlobal, bytecode.Instructions[2].Opcode)
	assert.Equal(t, 1, bytecode.Instructions[2].Operand)
	assert.Equal(t, types.OpExitScope, bytecode.Instructions[3].Opcode)
}

func TestCompileArithmeticExpressions(t *testing.T) {
	tests := []struct {
		input      string
		constants  []any
		operations []types.Opcode
	}{
		{
			"5 + 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpAdd, types.OpPop},
		},
		{
			"5 - 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpSubtract, types.OpPop},
		},
		{
			"5 * 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpMultiply, types.OpPop},
		},
		{
			"5 / 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpDivide, types.OpPop},
		},
	}

	for _, tt := range tests {
		lex := lexer.New(tt.input)
		p := parser.New(lex)
		program := p.ParseProgram()

		bytecode := compiler.Compile(program)

		// Check constants
		assert.Equal(t, len(tt.constants), len(bytecode.Constants))
		for i, c := range tt.constants {
			assert.Equal(t, c, bytecode.Constants[i])
		}

		// Check operations, skipping the first and last (scope instructions)
		assert.Equal(t, len(tt.operations)+2, len(bytecode.Instructions))
		for i, op := range tt.operations {
			assert.Equal(t, op, bytecode.Instructions[i+1].Opcode)
		}
	}
}

func TestCompileComparisonExpressions(t *testing.T) {
	tests := []struct {
		input      string
		constants  []any
		operations []types.Opcode
	}{
		{
			"5 == 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpEqual, types.OpPop},
		},
		{
			"5 != 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpNotEqual, types.OpPop},
		},
		{
			"5 < 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpLessThan, types.OpPop},
		},
		{
			"5 > 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpGreaterThan, types.OpPop},
		},
		{
			"5 <= 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpLessEqual, types.OpPop},
		},
		{
			"5 >= 10;",
			[]any{5.0, 10.0},
			[]types.Opcode{types.OpConstant, types.OpConstant, types.OpGreaterEqual, types.OpPop},
		},
	}

	for _, tt := range tests {
		lex := lexer.New(tt.input)
		p := parser.New(lex)
		program := p.ParseProgram()

		bytecode := compiler.Compile(program)

		// Check constants
		assert.Equal(t, len(tt.constants), len(bytecode.Constants))
		for i, c := range tt.constants {
			assert.Equal(t, c, bytecode.Constants[i])
		}

		// Check operations, skipping the first and last (scope instructions)
		assert.Equal(t, len(tt.operations)+2, len(bytecode.Instructions))
		for i, op := range tt.operations {
			assert.Equal(t, op, bytecode.Instructions[i+1].Opcode)
		}
	}
}

func TestCompileTable(t *testing.T) {
	input := `table = { name = "John", age = 30 };`

	lex := lexer.New(input)
	p := parser.New(lex)
	program := p.ParseProgram()

	bytecode := compiler.Compile(program)

	// Instead of checking exact order, just verify the constants are present
	assert.Equal(t, 5, len(bytecode.Constants))

	// Check that all expected constants are present
	foundName := false
	foundJohn := false
	foundAge := false
	found30 := false
	foundTable := false

	for _, constant := range bytecode.Constants {
		switch v := constant.(type) {
		case string:
			if v == "name" {
				foundName = true
			} else if v == "John" {
				foundJohn = true
			} else if v == "age" {
				foundAge = true
			} else if v == "table" {
				foundTable = true
			}
		case float64:
			if v == 30.0 {
				found30 = true
			}
		}
	}

	assert.True(t, foundName)
	assert.True(t, foundJohn)
	assert.True(t, foundAge)
	assert.True(t, found30)
	assert.True(t, foundTable)

	// Check opcodes rather than exact operands which depend on constant order
	expectedOpcodes := []types.Opcode{
		types.OpEnterScope,
		types.OpNewTable,  // Create table
		types.OpDup,       // Duplicate to set first property
		types.OpConstant,  // Load key
		types.OpConstant,  // Load value
		types.OpSetIndex,  // Set first property
		types.OpPop,       // Pop result of setindex
		types.OpDup,       // Duplicate for second property
		types.OpConstant,  // Load key
		types.OpConstant,  // Load value
		types.OpSetIndex,  // Set second property
		types.OpPop,       // Pop result of setindex
		types.OpSetGlobal, // Set table variable
		types.OpExitScope,
	}

	assert.Equal(t, len(expectedOpcodes), len(bytecode.Instructions))
	for i, op := range expectedOpcodes {
		assert.Equal(t, op, bytecode.Instructions[i].Opcode)
	}
}

func TestCompileIfStatement(t *testing.T) {
	input := `
	if x < 10 then
		echo "x is less than 10";
	else
		echo "x is not less than 10";
	end
	`

	lex := lexer.New(input)
	p := parser.New(lex)
	program := p.ParseProgram()

	bytecode := compiler.Compile(program)

	// Don't check exact count - just verify key constants exist
	assert.True(t, len(bytecode.Constants) >= 5)

	// Check for key opcodes in the bytecode
	ops := bytecode.Instructions

	// Check for variable lookup
	foundGetGlobal := false
	for _, op := range ops {
		if op.Opcode == types.OpGetGlobal {
			foundGetGlobal = true
			break
		}
	}
	assert.True(t, foundGetGlobal)

	// Check for comparison opcode
	foundLessThan := false
	for _, op := range ops {
		if op.Opcode == types.OpLessThan {
			foundLessThan = true
			break
		}
	}
	assert.True(t, foundLessThan)

	// Check for conditional jump
	foundJumpIfFalse := false
	for _, op := range ops {
		if op.Opcode == types.OpJumpIfFalse {
			foundJumpIfFalse = true
			break
		}
	}
	assert.True(t, foundJumpIfFalse)

	// Check for unconditional jump (for else clause)
	foundJump := false
	for _, op := range ops {
		if op.Opcode == types.OpJump {
			foundJump = true
			break
		}
	}
	assert.True(t, foundJump)

	// Check for echo statements
	foundEcho := false
	echoCount := 0
	for _, op := range ops {
		if op.Opcode == types.OpEcho {
			foundEcho = true
			echoCount++
		}
	}
	assert.True(t, foundEcho)
	assert.Equal(t, 2, echoCount)
}

func TestCompileLogicalOperators(t *testing.T) {
	tests := []struct {
		input            string
		expectedOperator types.Opcode
	}{
		{"not true;", types.OpNot},
	}

	for _, tt := range tests {
		lex := lexer.New(tt.input)
		p := parser.New(lex)
		program := p.ParseProgram()

		bytecode := compiler.Compile(program)

		foundOp := false
		for _, op := range bytecode.Instructions {
			if op.Opcode == tt.expectedOperator {
				foundOp = true
				break
			}
		}
		assert.True(t, foundOp)
	}

	// AND and OR are special as they involve jumps - test separately
	andInput := "true and false;"
	lex := lexer.New(andInput)
	p := parser.New(lex)
	program := p.ParseProgram()
	bytecode := compiler.Compile(program)

	// AND should involve a JumpIfFalse instruction
	foundJumpIfFalse := false
	for _, op := range bytecode.Instructions {
		if op.Opcode == types.OpJumpIfFalse {
			foundJumpIfFalse = true
			break
		}
	}
	assert.True(t, foundJumpIfFalse)

	// OR test
	orInput := "true or false;"
	lex = lexer.New(orInput)
	p = parser.New(lex)
	program = p.ParseProgram()
	bytecode = compiler.Compile(program)

	// OR should involve both Jump and JumpIfFalse instructions
	foundJump := false
	foundJumpIfFalse = false
	for _, op := range bytecode.Instructions {
		if op.Opcode == types.OpJump {
			foundJump = true
		}
		if op.Opcode == types.OpJumpIfFalse {
			foundJumpIfFalse = true
		}
	}
	assert.True(t, foundJump)
	assert.True(t, foundJumpIfFalse)
}