Mako/compiler/state.go

package compiler

import "fmt"

// Constants for compiler limits
const (
	MaxLocals    = 256   // Maximum local variables per function
	MaxUpvalues  = 256   // Maximum upvalues per function
	MaxConstants = 65536 // Maximum constants per chunk
)

// CompilerState holds state during compilation
type CompilerState struct {
	Chunk         *Chunk         // Current chunk being compiled
	Constants     map[string]int // Constant pool index mapping
	Functions     []Function     // Compiled functions
	Structs       []Struct       // Compiled structs
	Locals        []Local        // Local variable stack
	Upvalues      []UpvalueRef   // Upvalue definitions
	ScopeDepth    int            // Current scope nesting level
	FunctionType  FunctionType   // Type of function being compiled
	BreakJumps    []int          // Break jump addresses for loops
	ContinueJumps []int          // Continue jump addresses for loops
	LoopStart     int            // Start of current loop for continue
	LoopDepth     int            // Current loop nesting depth
	parent        *CompilerState // Parent compiler state for nested functions
	CurrentLine   int            // Current source line being compiled
}

// Local represents a local variable during compilation
type Local struct {
	Name       string // Variable name
	Depth      int    // Scope depth where declared
	IsCaptured bool   // Whether variable is captured by closure
	Slot       int    // Stack slot index
}

// UpvalueRef represents an upvalue reference during compilation
type UpvalueRef struct {
	Index   uint8 // Index in enclosing function's locals or upvalues
	IsLocal bool  // True if captures local, false if captures upvalue
}

// FunctionType represents the type of function being compiled
type FunctionType uint8

const (
	FunctionTypeScript   FunctionType = iota // Top-level script
	FunctionTypeFunction                     // Regular function
	FunctionTypeMethod                       // Struct method
)

// CompileError represents a compilation error with location information
type CompileError struct {
	Message string
	Line    int
	Column  int
}

func (ce CompileError) Error() string {
	return fmt.Sprintf("Compile error at line %d, column %d: %s", ce.Line, ce.Column, ce.Message)
}

// NewCompilerState creates a new compiler state for compilation
func NewCompilerState(functionType FunctionType) *CompilerState {
	return &CompilerState{
		Chunk:         NewChunk(),
		Constants:     make(map[string]int),
		Functions:     make([]Function, 0),
		Structs:       make([]Struct, 0),
		Locals:        make([]Local, 0, MaxLocals),
		Upvalues:      make([]UpvalueRef, 0, MaxUpvalues),
		ScopeDepth:    0,
		FunctionType:  functionType,
		BreakJumps:    make([]int, 0),
		ContinueJumps: make([]int, 0),
		LoopStart:     -1,
		LoopDepth:     0,
		parent:        nil,
	}
}

// NewChunk creates a new bytecode chunk
func NewChunk() *Chunk {
	return &Chunk{
		Code:      make([]uint8, 0, 256),
		Constants: make([]Value, 0, 64),
		Lines:     make([]int, 0, 256),
		Functions: make([]Function, 0),
		Structs:   make([]Struct, 0),
	}
}

// Scope management methods
func (cs *CompilerState) BeginScope() {
	cs.ScopeDepth++
}

func (cs *CompilerState) EndScope() {
	cs.ScopeDepth--

	// Remove locals that go out of scope
	for len(cs.Locals) > 0 && cs.Locals[len(cs.Locals)-1].Depth > cs.ScopeDepth {
		local := cs.Locals[len(cs.Locals)-1]
		if local.IsCaptured {
			// Emit close upvalue instruction
			cs.EmitByte(uint8(OpCloseUpvalue))
		} else {
			// Emit pop instruction
			cs.EmitByte(uint8(OpPop))
		}
		cs.Locals = cs.Locals[:len(cs.Locals)-1]
	}
}

// Local variable management
func (cs *CompilerState) AddLocal(name string) error {
	if len(cs.Locals) >= MaxLocals {
		return CompileError{
			Message: "too many local variables in function",
		}
	}

	local := Local{
		Name:       name,
		Depth:      -1, // Mark as uninitialized
		IsCaptured: false,
		Slot:       len(cs.Locals),
	}

	cs.Locals = append(cs.Locals, local)
	return nil
}

func (cs *CompilerState) MarkInitialized() {
	if len(cs.Locals) > 0 {
		cs.Locals[len(cs.Locals)-1].Depth = cs.ScopeDepth
	}
}

func (cs *CompilerState) ResolveLocal(name string) int {
	for i := len(cs.Locals) - 1; i >= 0; i-- {
		local := &cs.Locals[i]
		if local.Name == name {
			if local.Depth == -1 {
				// Variable used before initialization
				return -2
			}
			return i
		}
	}
	return -1
}

// Upvalue management
func (cs *CompilerState) AddUpvalue(index uint8, isLocal bool) int {
	upvalueCount := len(cs.Upvalues)

	// Check if upvalue already exists
	for i := range upvalueCount {
		upvalue := &cs.Upvalues[i]
		if upvalue.Index == index && upvalue.IsLocal == isLocal {
			return i
		}
	}

	if upvalueCount >= MaxUpvalues {
		return -1 // Too many upvalues
	}

	cs.Upvalues = append(cs.Upvalues, UpvalueRef{
		Index:   index,
		IsLocal: isLocal,
	})

	return upvalueCount
}

// Constant pool management
func (cs *CompilerState) AddConstant(value Value) int {
	// Check if constant already exists to avoid duplicates
	key := cs.valueKey(value)
	if index, exists := cs.Constants[key]; exists {
		return index
	}

	if len(cs.Chunk.Constants) >= MaxConstants {
		return -1 // Too many constants
	}

	index := len(cs.Chunk.Constants)
	cs.Chunk.Constants = append(cs.Chunk.Constants, value)
	cs.Constants[key] = index
	return index
}

// Generate unique key for value in constant pool
func (cs *CompilerState) valueKey(value Value) string {
	switch value.Type {
	case ValueNil:
		return "nil"
	case ValueBool:
		if value.Data.(bool) {
			return "bool:true"
		}
		return "bool:false"
	case ValueNumber:
		return fmt.Sprintf("number:%g", value.Data.(float64))
	case ValueString:
		return fmt.Sprintf("string:%s", value.Data.(string))
	default:
		// For complex types, use memory address as fallback
		return fmt.Sprintf("%T:%p", value.Data, value.Data)
	}
}

// Bytecode emission methods
func (cs *CompilerState) EmitByte(byte uint8) {
	cs.Chunk.Code = append(cs.Chunk.Code, byte)
	cs.Chunk.Lines = append(cs.Chunk.Lines, cs.CurrentLine)
}

func (cs *CompilerState) EmitBytes(bytes ...uint8) {
	for _, b := range bytes {
		cs.EmitByte(b)
	}
}

func (cs *CompilerState) EmitInstruction(op Opcode, operands ...uint16) {
	bytes := EncodeInstruction(op, operands...)
	cs.EmitBytes(bytes...)
}

func (cs *CompilerState) EmitJump(op Opcode) int {
	cs.EmitByte(uint8(op))
	cs.EmitByte(0xFF)             // Placeholder
	cs.EmitByte(0xFF)             // Placeholder
	return len(cs.Chunk.Code) - 2 // Return offset of jump address
}

func (cs *CompilerState) PatchJump(offset int) {
	// Calculate jump distance
	jump := len(cs.Chunk.Code) - offset - 2

	if jump > 65535 {
		// Jump too large - would need long jump instruction
		return
	}

	cs.Chunk.Code[offset] = uint8(jump & 0xFF)
	cs.Chunk.Code[offset+1] = uint8((jump >> 8) & 0xFF)
}

// Loop management
func (cs *CompilerState) EnterLoop() {
	cs.LoopStart = len(cs.Chunk.Code)
	cs.LoopDepth++
}

func (cs *CompilerState) ExitLoop() {
	cs.LoopDepth--
	if cs.LoopDepth == 0 {
		cs.LoopStart = -1
	}

	// Patch break jumps
	for _, jumpOffset := range cs.BreakJumps {
		cs.PatchJump(jumpOffset)
	}
	cs.BreakJumps = cs.BreakJumps[:0]

	// Patch continue jumps
	for _, jumpOffset := range cs.ContinueJumps {
		jump := cs.LoopStart - jumpOffset - 2
		if jump < 65535 {
			cs.Chunk.Code[jumpOffset] = uint8(jump & 0xFF)
			cs.Chunk.Code[jumpOffset+1] = uint8((jump >> 8) & 0xFF)
		}
	}
	cs.ContinueJumps = cs.ContinueJumps[:0]
}

func (cs *CompilerState) EmitBreak() {
	jumpOffset := cs.EmitJump(OpJump)
	cs.BreakJumps = append(cs.BreakJumps, jumpOffset)
}

func (cs *CompilerState) EmitContinue() {
	if cs.LoopStart != -1 {
		jumpOffset := cs.EmitJump(OpJump)
		cs.ContinueJumps = append(cs.ContinueJumps, jumpOffset)
	}
}

func (cs *CompilerState) SetLine(line int) {
	cs.CurrentLine = line
}