Refactor register allocation and symbol management; introduce ConstantsPool for deduplication

pkg/compiler/internal/allocator.go

@@ -1,177 +1,157 @@
 package internal
 
 import (
+	"fmt"
+
 	"github.com/MontFerret/ferret/pkg/vm"
 )
 
-type (
-	// RegisterType represents the type of a register
-	RegisterType int
-
-	// RegisterStatus tracks register usage
-	RegisterStatus struct {
-		IsAllocated bool
-		Type        RegisterType // Type of variable stored
-	}
-
-	RegisterLifetime struct {
-		Start int // Instruction number where register becomes live
-		End   int // Instruction number where register dies
-	}
-
-	RegisterSequence struct {
-		Registers []vm.Operand
-	}
-
-	// RegisterAllocator manages register allocation
-	RegisterAllocator struct {
-		registers    map[vm.Operand]*RegisterStatus
-		nextRegister vm.Operand
-		currentInstr int
-	}
-)
+type RegisterType int
 
 const (
-	Temp   RegisterType = iota // Short-lived intermediate results
-	Var                        // Local variables
-	State                      // FOR loop state
-	Result                     // FOR loop result
+	Temp RegisterType = iota
+	Var
+	State
+	Result
 )
 
-func NewRegisterSequence(registers ...vm.Operand) *RegisterSequence {
-	return &RegisterSequence{
-		Registers: registers,
-	}
+type RegisterSequence []vm.Operand
+
+type registerInfo struct {
+	typ       RegisterType
+	allocated bool
+}
+
+type RegisterAllocator struct {
+	next     vm.Operand
+	freelist map[RegisterType][]vm.Operand
+	all      map[vm.Operand]*registerInfo
 }
 
 func NewRegisterAllocator() *RegisterAllocator {
 	return &RegisterAllocator{
-		registers:    make(map[vm.Operand]*RegisterStatus),
-		nextRegister: vm.NoopOperand + 1, // we start at 1 to avoid NoopOperand
+		next:     vm.NoopOperand + 1,
+		freelist: make(map[RegisterType][]vm.Operand),
+		all:      make(map[vm.Operand]*registerInfo),
 	}
 }
 
-func (ra *RegisterAllocator) Size() int {
-	return int(ra.nextRegister)
-}
-
-// Allocate assigns a register based on variable type
-func (ra *RegisterAllocator) Allocate(regType RegisterType) vm.Operand {
-	// Try to find a free register first
-	reg, found := ra.findFreeRegister()
-
-	if found {
+func (ra *RegisterAllocator) Allocate(typ RegisterType) vm.Operand {
+	// Reuse from freelist
+	if regs, ok := ra.freelist[typ]; ok && len(regs) > 0 {
+		last := len(regs) - 1
+		reg := regs[last]
+		ra.freelist[typ] = regs[:last]
+		ra.all[reg].allocated = true
 		return reg
 	}
 
-	// If no free registers, create a new one
-	newReg := ra.nextRegister
-	ra.nextRegister++
+	// New register
+	reg := ra.next
+	ra.next++
 
-	// Initialize register status
-	ra.registers[newReg] = &RegisterStatus{
-		IsAllocated: true,
-		Type:        regType,
+	ra.all[reg] = &registerInfo{
+		typ:       typ,
+		allocated: true,
 	}
 
-	return newReg
+	return reg
 }
 
-// Free marks a register as available
 func (ra *RegisterAllocator) Free(reg vm.Operand) {
-	//if status, exists := ra.registers[reg]; exists {
-	//status.IsAllocated = false
-	//status.Lifetime.End = ra.currentInstr
+	info, ok := ra.all[reg]
+	if !ok || !info.allocated {
+		return // double-free or unknown
+	}
 
-	//// Clean up interference graph
-	//delete(ra.usageGraph, reg)
-	//
-	//for _, edges := range ra.usageGraph {
-	//	delete(edges, reg)
-	//}
-	//}
+	info.allocated = false
+	ra.freelist[info.typ] = append(ra.freelist[info.typ], reg)
 }
 
-// findFreeRegister looks for an unused register
-func (ra *RegisterAllocator) findFreeRegister() (vm.Operand, bool) {
-	// First, try to find a completely free register
-	for reg, status := range ra.registers {
-		if !status.IsAllocated {
-			return reg, true
+func (ra *RegisterAllocator) AllocateSequence(count int) RegisterSequence {
+	// Try to find a contiguous free block
+	start, found := ra.findContiguousFree(count)
+
+	if found {
+		seq := make(RegisterSequence, count)
+		for i := 0; i < count; i++ {
+			reg := start + vm.Operand(i)
+			info, ok := ra.all[reg]
+			if !ok {
+				info = &registerInfo{}
+				ra.all[reg] = info
+			}
+			info.typ = Temp
+			info.allocated = true
+			seq[i] = reg
+		}
+		return seq
+	}
+
+	// Otherwise, allocate new block
+	start = ra.next
+	seq := make(RegisterSequence, count)
+
+	for i := 0; i < count; i++ {
+		reg := start + vm.Operand(i)
+		ra.all[reg] = &registerInfo{
+			typ:       Temp,
+			allocated: true,
+		}
+		seq[i] = reg
+	}
+
+	ra.next += vm.Operand(count)
+	return seq
+}
+
+func (ra *RegisterAllocator) FreeSequence(seq RegisterSequence) {
+	for _, reg := range seq {
+		ra.Free(reg)
+	}
+}
+
+func (ra *RegisterAllocator) findContiguousFree(count int) (vm.Operand, bool) {
+	if count <= 0 {
+		return 0, false
+	}
+
+	limit := ra.next
+	for start := vm.NoopOperand + 1; start+vm.Operand(count) <= limit; start++ {
+		ok := true
+		for i := 0; i < count; i++ {
+			reg := start + vm.Operand(i)
+			info := ra.all[reg]
+			if info == nil || info.allocated {
+				ok = false
+				break
+			}
+		}
+		if ok {
+			return start, true
 		}
 	}
 
 	return 0, false
 }
 
-// AllocateSequence allocates a sequence of registers for function arguments or similar uses
-func (ra *RegisterAllocator) AllocateSequence(count int) *RegisterSequence {
-	sequence := &RegisterSequence{
-		Registers: make([]vm.Operand, count),
-	}
+func (ra *RegisterAllocator) Size() int {
+	return int(ra.next)
+}
 
-	// First pass: try to find contiguous free registers
-	startReg, found := ra.findContiguousRegisters(count)
-
-	if found {
-		// Use contiguous block
-		for i := 0; i < count; i++ {
-			reg := startReg + vm.Operand(i)
-			sequence.Registers[i] = reg
-
-			// Initialize or update register status
-			ra.registers[reg] = &RegisterStatus{
-				IsAllocated: true,
-				Type:        Temp,
+func (ra *RegisterAllocator) DebugView() string {
+	out := ""
+	for reg := vm.NoopOperand + 1; reg < ra.next; reg++ {
+		info := ra.all[reg]
+		status := "FREE"
+		typ := "?"
+		if info != nil {
+			if info.allocated {
+				status = "USED"
 			}
+			typ = fmt.Sprintf("%v", info.typ)
 		}
-	} else {
-		// Allocate registers individually if contiguous block not available
-		for i := 0; i < count; i++ {
-			reg := ra.Allocate(Temp)
-			sequence.Registers[i] = reg
-		}
-	}
-
-	return sequence
-}
-
-// findContiguousRegisters attempts to find a block of consecutive free registers
-func (ra *RegisterAllocator) findContiguousRegisters(count int) (vm.Operand, bool) {
-	if count <= 0 {
-		return 0, false
-	}
-
-	// First, try to find a contiguous block in existing registers
-	maxReg := ra.nextRegister
-	for start := vm.NoopOperand + 1; start < maxReg; start++ {
-		if ra.isContiguousBlockFree(start, count) {
-			return start, true
-		}
-	}
-
-	// If no existing contiguous block found, allocate new block
-	startReg := ra.nextRegister
-	ra.nextRegister += vm.Operand(count)
-
-	return startReg, true
-}
-
-// isContiguousBlockFree checks if a block of registers is available
-func (ra *RegisterAllocator) isContiguousBlockFree(start vm.Operand, count int) bool {
-	for i := 0; i < count; i++ {
-		reg := start + vm.Operand(i)
-
-		if status := ra.registers[reg]; status != nil && status.IsAllocated {
-			return false
-		}
-	}
-	return true
-}
-
-// FreeSequence frees all registers in a sequence
-func (ra *RegisterAllocator) FreeSequence(seq *RegisterSequence) {
-	for _, reg := range seq.Registers {
-		ra.Free(reg)
+		out += fmt.Sprintf("R%d: %s [%s]\n", reg, status, typ)
 	}
+	return out
 }

pkg/compiler/internal/constants.go

@@ -0,0 +1,62 @@
+package internal
+
+import (
+	"github.com/MontFerret/ferret/pkg/runtime"
+	"github.com/MontFerret/ferret/pkg/vm"
+	"strconv"
+)
+
+// ConstantsPool stores and deduplicates constants
+type ConstantsPool struct {
+	values []runtime.Value
+	index  map[uint64]int
+}
+
+func NewConstantsPool() *ConstantsPool {
+	return &ConstantsPool{
+		values: make([]runtime.Value, 0),
+		index:  make(map[uint64]int),
+	}
+}
+
+func (cp *ConstantsPool) Add(val runtime.Value) vm.Operand {
+	var hash uint64
+	isNone := val == runtime.None
+
+	if runtime.IsScalar(val) {
+		hash = val.Hash()
+	}
+
+	if hash > 0 || isNone {
+		if idx, ok := cp.index[hash]; ok {
+			return vm.NewConstantOperand(idx)
+		}
+	}
+
+	cp.values = append(cp.values, val)
+	idx := len(cp.values) - 1
+
+	if hash > 0 || isNone {
+		cp.index[hash] = idx
+	}
+
+	return vm.NewConstantOperand(idx)
+}
+
+func (cp *ConstantsPool) Get(addr vm.Operand) runtime.Value {
+	if !addr.IsConstant() {
+		panic(runtime.Error(ErrInvalidOperandType, strconv.Itoa(int(addr))))
+	}
+
+	idx := addr.Constant()
+
+	if idx < 0 || idx >= len(cp.values) {
+		panic(runtime.Error(ErrConstantNotFound, strconv.Itoa(idx)))
+	}
+
+	return cp.values[idx]
+}
+
+func (cp *ConstantsPool) All() []runtime.Value {
+	return cp.values
+}

pkg/compiler/internal/emitter.go

@@ -18,6 +18,7 @@ func (e *Emitter) Bytecode() []vm.Instruction {
 	return e.instructions
 }
 
+// Size returns the number of instructions currently stored in the Emitter.
 func (e *Emitter) Size() int {
 	return len(e.instructions)
 }
@@ -43,6 +44,7 @@ func (e *Emitter) EmitJumpc(op vm.Opcode, pos int, reg vm.Operand) int {
 	return len(e.instructions) - 1
 }
 
+// PatchSwapAB modifies an instruction at the given position to swap operands and update its operation and destination.
 func (e *Emitter) PatchSwapAB(pos int, op vm.Opcode, dst, src1 vm.Operand) {
 	e.instructions[pos] = vm.Instruction{
 		Opcode:   op,
@@ -50,6 +52,7 @@ func (e *Emitter) PatchSwapAB(pos int, op vm.Opcode, dst, src1 vm.Operand) {
 	}
 }
 
+// PatchSwapAx modifies an existing instruction at the specified position with a new opcode, destination, and argument.
 func (e *Emitter) PatchSwapAx(pos int, op vm.Opcode, dst vm.Operand, arg int) {
 	e.instructions[pos] = vm.Instruction{
 		Opcode:   op,
@@ -57,6 +60,7 @@ func (e *Emitter) PatchSwapAx(pos int, op vm.Opcode, dst vm.Operand, arg int) {
 	}
 }
 
+// PatchSwapAxy replaces an instruction at the specified position with a new one using the provided opcode and operands.
 func (e *Emitter) PatchSwapAxy(pos int, op vm.Opcode, dst vm.Operand, arg1, agr2 int) {
 	e.instructions[pos] = vm.Instruction{
 		Opcode:   op,
@@ -64,10 +68,11 @@ func (e *Emitter) PatchSwapAxy(pos int, op vm.Opcode, dst vm.Operand, arg1, agr2
 	}
 }
 
-func (e *Emitter) PatchSwapAs(pos int, op vm.Opcode, dst vm.Operand, seq *RegisterSequence) {
+// PatchSwapAs replaces an instruction at the specified position with a new instruction using the given opcode and operands.
+func (e *Emitter) PatchSwapAs(pos int, op vm.Opcode, dst vm.Operand, seq RegisterSequence) {
 	e.instructions[pos] = vm.Instruction{
 		Opcode:   op,
-		Operands: [3]vm.Operand{dst, seq.Registers[0], seq.Registers[len(seq.Registers)-1]},
+		Operands: [3]vm.Operand{dst, seq[0], seq[len(seq)-1]},
 	}
 }
 
@@ -123,10 +128,10 @@ func (e *Emitter) EmitAx(op vm.Opcode, dest vm.Operand, arg int) {
 }
 
 // EmitAs emits an opcode with a destination register and a sequence of registers.
-func (e *Emitter) EmitAs(op vm.Opcode, dest vm.Operand, seq *RegisterSequence) {
+func (e *Emitter) EmitAs(op vm.Opcode, dest vm.Operand, seq RegisterSequence) {
 	if seq != nil {
-		src1 := seq.Registers[0]
-		src2 := seq.Registers[len(seq.Registers)-1]
+		src1 := seq[0]
+		src2 := seq[len(seq)-1]
 		e.EmitABC(op, dest, src1, src2)
 	} else {
 		e.EmitA(op, dest)

pkg/compiler/internal/emitter_helpers.go

@@ -0,0 +1,148 @@
+package internal
+
+import (
+	"github.com/MontFerret/ferret/pkg/runtime"
+	"github.com/MontFerret/ferret/pkg/vm"
+)
+
+// Emitter Helpers - Common opcode shortcut methods
+
+// ─── Loop & Iterator ──────────────────────────────────────────────────────
+
+func (e *Emitter) EmitIter(dst, src vm.Operand) {
+	e.EmitAB(vm.OpIter, dst, src)
+}
+
+func (e *Emitter) EmitIterNext(jumpTarget int, iterator vm.Operand) int {
+	return e.EmitJumpc(vm.OpIterNext, jumpTarget, iterator)
+}
+
+func (e *Emitter) EmitIterKey(dst, iterator vm.Operand) {
+	e.EmitAB(vm.OpIterKey, dst, iterator)
+}
+
+func (e *Emitter) EmitIterValue(dst, iterator vm.Operand) {
+	e.EmitAB(vm.OpIterValue, dst, iterator)
+}
+
+func (e *Emitter) EmitIterSkip(state, count vm.Operand, jump int) {
+	e.EmitABx(vm.OpIterSkip, state, count, jump)
+}
+
+func (e *Emitter) EmitIterLimit(state, count vm.Operand, jump int) {
+	e.EmitABx(vm.OpIterLimit, state, count, jump)
+}
+
+// ─── Value & Memory ──────────────────────────────────────────────────────
+
+func (e *Emitter) EmitMove(dst, src vm.Operand) {
+	e.EmitAB(vm.OpMove, dst, src)
+}
+
+func (e *Emitter) EmitPush(dst, src vm.Operand) {
+	e.EmitAB(vm.OpPush, dst, src)
+}
+
+func (e *Emitter) EmitPushKV(dst, key, val vm.Operand) {
+	e.EmitABC(vm.OpPushKV, dst, key, val)
+}
+
+func (e *Emitter) EmitClose(reg vm.Operand) {
+	e.EmitA(vm.OpClose, reg)
+}
+
+func (e *Emitter) EmitLoadConst(dst vm.Operand, val runtime.Value, symbols *SymbolTable) {
+	e.EmitAB(vm.OpLoadConst, dst, symbols.AddConstant(val))
+}
+
+// ─── Data Structures ──────────────────────────────────────────────────────
+
+func (e *Emitter) EmitList(dst vm.Operand, seq RegisterSequence) {
+	e.EmitAs(vm.OpList, dst, seq)
+}
+
+func (e *Emitter) EmitMap(dst vm.Operand, seq RegisterSequence) {
+	e.EmitAs(vm.OpMap, dst, seq)
+}
+
+func (e *Emitter) EmitRange(dst, start, end vm.Operand) {
+	e.EmitABC(vm.OpRange, dst, start, end)
+}
+
+func (e *Emitter) EmitLoadIndex(dst, arr, idx vm.Operand) {
+	e.EmitABC(vm.OpLoadIndex, dst, arr, idx)
+}
+
+func (e *Emitter) EmitLoadProperty(dst, obj, prop vm.Operand) {
+	e.EmitABC(vm.OpLoadProperty, dst, obj, prop)
+}
+
+func (e *Emitter) EmitLoadPropertyOptional(dst, obj, prop vm.Operand) {
+	e.EmitABC(vm.OpLoadPropertyOptional, dst, obj, prop)
+}
+
+// ─── Arithmetic and Logical ──────────────────────────────────────────────
+
+func (e *Emitter) EmitAdd(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpAdd, dst, a, b)
+}
+
+func (e *Emitter) EmitSub(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpSub, dst, a, b)
+}
+
+func (e *Emitter) EmitMul(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpMulti, dst, a, b)
+}
+
+func (e *Emitter) EmitDiv(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpDiv, dst, a, b)
+}
+
+func (e *Emitter) EmitMod(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpMod, dst, a, b)
+}
+
+func (e *Emitter) EmitEq(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpEq, dst, a, b)
+}
+
+func (e *Emitter) EmitNeq(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpNeq, dst, a, b)
+}
+
+func (e *Emitter) EmitGt(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpGt, dst, a, b)
+}
+
+func (e *Emitter) EmitLt(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpLt, dst, a, b)
+}
+
+func (e *Emitter) EmitGte(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpGte, dst, a, b)
+}
+
+func (e *Emitter) EmitLte(dst, a, b vm.Operand) {
+	e.EmitABC(vm.OpLte, dst, a, b)
+}
+
+// ─── Control Flow ────────────────────────────────────────────────────────
+
+func (e *Emitter) EmitJumpIf(cond vm.Operand, isTrue bool, jumpTarget int) int {
+	if isTrue {
+		return e.EmitJumpc(vm.OpJumpIfTrue, jumpTarget, cond)
+	}
+
+	return e.EmitJumpc(vm.OpJumpIfFalse, jumpTarget, cond)
+}
+
+func (e *Emitter) EmitReturnValue(val vm.Operand) {
+	if val.IsConstant() {
+		e.EmitAB(vm.OpLoadGlobal, vm.NoopOperand, val)
+	} else {
+		e.EmitAB(vm.OpMove, vm.NoopOperand, val)
+	}
+
+	e.Emit(vm.OpReturn)
+}

pkg/compiler/internal/symbols.go

@@ -1,55 +1,62 @@
 package internal
 
 import (
-	"strconv"
-
+	"fmt"
 	"github.com/MontFerret/ferret/pkg/runtime"
 	"github.com/MontFerret/ferret/pkg/vm"
 )
 
-type (
-	Variable struct {
-		Name     string
-		Register vm.Operand
-		Depth    int
-	}
+type SymbolKind int
 
-	SymbolTable struct {
-		registers      *RegisterAllocator
-		constants      []runtime.Value
-		constantsIndex map[uint64]int
-		params         map[string]string
-		globals        map[string]vm.Operand
-		locals         []*Variable
-		scope          int
-	}
+const (
+	SymbolVar SymbolKind = iota
+	SymbolConst
+	SymbolParam
+	SymbolGlobal
 )
 
+type ValueType int
+
+const (
+	TypeUnknown ValueType = iota
+	TypeInt
+	TypeFloat
+	TypeString
+	TypeBool
+	TypeList
+	TypeMap
+	TypeAny
+)
+
+type Variable struct {
+	Name     string
+	Kind     SymbolKind
+	Register vm.Operand
+	Depth    int
+	Type     ValueType
+}
+
+type SymbolTable struct {
+	registers *RegisterAllocator
+	constants *ConstantsPool
+
+	params  map[string]string
+	globals map[string]vm.Operand
+	locals  []*Variable
+
+	scope int
+}
+
 func NewSymbolTable(registers *RegisterAllocator) *SymbolTable {
 	return &SymbolTable{
-		registers:      registers,
-		constants:      make([]runtime.Value, 0),
-		constantsIndex: make(map[uint64]int),
-		params:         make(map[string]string),
-		globals:        make(map[string]vm.Operand),
-		locals:         make([]*Variable, 0),
+		registers: registers,
+		constants: NewConstantsPool(),
+		params:    make(map[string]string),
+		globals:   make(map[string]vm.Operand),
+		locals:    make([]*Variable, 0),
 	}
 }
 
-func (st *SymbolTable) Params() []string {
-	params := make([]string, 0, len(st.params))
-
-	for _, name := range st.params {
-		params = append(params, name)
-	}
-
-	return params
-}
-
-func (st *SymbolTable) Constants() []runtime.Value {
-	return st.constants
-}
-
 func (st *SymbolTable) Scope() int {
 	return st.scope
 }
@@ -58,142 +65,104 @@ func (st *SymbolTable) EnterScope() {
 	st.scope++
 }
 
-func (st *SymbolTable) AddParam(name string) vm.Operand {
-	st.params[name] = name
-
-	return st.AddConstant(runtime.NewString(name))
-}
-
-// AddConstant adds a constant to the constants pool and returns its index.
-// If the constant is a scalar, it will be deduplicated.
-// If the constant is not a scalar, it will be added to the pool without deduplication.
-func (st *SymbolTable) AddConstant(constant runtime.Value) vm.Operand {
-	var hash uint64
-	isNone := constant == runtime.None
-
-	if runtime.IsScalar(constant) {
-		hash = constant.Hash()
-	}
-
-	if hash > 0 || isNone {
-		if p, ok := st.constantsIndex[hash]; ok {
-			return vm.NewConstantOperand(p)
-		}
-	}
-
-	st.constants = append(st.constants, constant)
-	p := len(st.constants) - 1
-
-	if hash > 0 || isNone {
-		st.constantsIndex[hash] = p
-	}
-
-	// We flip the sign to indicate that this is a constant index, not a register.
-	return vm.NewConstantOperand(p)
-}
-
-// Constant returns a constant by its index.
-func (st *SymbolTable) Constant(addr vm.Operand) runtime.Value {
-	if !addr.IsConstant() {
-		panic(runtime.Error(ErrInvalidOperandType, strconv.Itoa(int(addr))))
-	}
-
-	index := addr.Constant()
-
-	if index < 0 || index >= len(st.constants) {
-		panic(runtime.Error(ErrConstantNotFound, strconv.Itoa(index)))
-	}
-
-	return st.constants[index]
-}
-
-func (st *SymbolTable) DefineVariable(name string) vm.Operand {
-	if st.scope == 0 {
-		// Check for duplicate global variable names.
-		_, ok := st.globals[name]
-
-		if ok {
-			panic(runtime.Error(ErrVariableNotUnique, name))
-		}
-
-		op := st.AddConstant(runtime.NewString(name))
-		// Define global variable.
-		st.globals[name] = op
-
-		return op
-	}
-
-	register := st.registers.Allocate(Var)
-
-	if st.scope == 0 {
-		panic("cannot define scoped variable in global scope")
-	}
-
-	st.locals = append(st.locals, &Variable{
-		Name:     name,
-		Depth:    st.scope,
-		Register: register,
-	})
-
-	return register
-}
-
-func (st *SymbolTable) DefineVariableInScope(name string, scope int) vm.Operand {
-	if scope == 0 {
-		panic("cannot define scoped variable in global scope")
-	}
-
-	if scope > st.scope {
-		panic("cannot define variable in a scope that is deeper than the current scope")
-	}
-
-	register := st.registers.Allocate(Var)
-
-	st.locals = append(st.locals, &Variable{
-		Name:     name,
-		Depth:    scope,
-		Register: register,
-	})
-
-	return register
-}
-
-func (st *SymbolTable) Variable(name string) vm.Operand {
-	for i := len(st.locals) - 1; i >= 0; i-- {
-		variable := st.locals[i]
-		if variable.Name == name {
-			return vm.NewRegisterOperand(int(variable.Register))
-		}
-	}
-
-	op, ok := st.globals[name]
-
-	if !ok {
-		panic(runtime.Error(ErrVariableNotFound, name))
-	}
-
-	return op
-}
-
-// GlobalVariable returns a global variable by its name.
-func (st *SymbolTable) GlobalVariable(name string) (vm.Operand, bool) {
-	op, ok := st.globals[name]
-
-	return op, ok
-}
-
 func (st *SymbolTable) ExitScope() {
 	st.scope--
 
-	// Pop all local variables from the stack within the closed scope.
 	for len(st.locals) > 0 && st.locals[len(st.locals)-1].Depth > st.scope {
-		popped := st.locals[len(st.locals)-1:]
-
-		// Free the register.
-		for _, v := range popped {
-			st.registers.Free(v.Register)
-		}
-
+		popped := st.locals[len(st.locals)-1]
+		st.registers.Free(popped.Register)
 		st.locals = st.locals[:len(st.locals)-1]
 	}
 }
+
+func (st *SymbolTable) DeclareLocal(name string) vm.Operand {
+	return st.DeclareLocalTyped(name, TypeUnknown)
+}
+
+func (st *SymbolTable) DeclareLocalTyped(name string, typ ValueType) vm.Operand {
+	reg := st.registers.Allocate(Var)
+
+	st.locals = append(st.locals, &Variable{
+		Name:     name,
+		Kind:     SymbolVar,
+		Register: reg,
+		Depth:    st.scope,
+		Type:     typ,
+	})
+
+	return reg
+}
+
+func (st *SymbolTable) DeclareGlobal(name string) vm.Operand {
+	if _, exists := st.globals[name]; exists {
+		panic(runtime.Error(ErrVariableNotUnique, name))
+	}
+
+	op := st.constants.Add(runtime.NewString(name))
+	st.globals[name] = op
+	return op
+}
+
+func (st *SymbolTable) BindParam(name string) vm.Operand {
+	st.params[name] = name
+	return st.constants.Add(runtime.NewString(name))
+}
+
+func (st *SymbolTable) Constants() []runtime.Value {
+	return st.constants.All()
+}
+
+func (st *SymbolTable) AddConstant(val runtime.Value) vm.Operand {
+	return st.constants.Add(val)
+}
+
+func (st *SymbolTable) Constant(addr vm.Operand) runtime.Value {
+	return st.constants.Get(addr)
+}
+
+func (st *SymbolTable) Resolve(name string) (vm.Operand, SymbolKind, bool) {
+	for i := len(st.locals) - 1; i >= 0; i-- {
+		v := st.locals[i]
+		if v.Name == name {
+			return vm.NewRegisterOperand(int(v.Register)), v.Kind, true
+		}
+	}
+
+	if reg, ok := st.globals[name]; ok {
+		return reg, SymbolGlobal, true
+	}
+
+	return vm.NoopOperand, SymbolVar, false
+}
+
+func (st *SymbolTable) Lookup(name string) (*Variable, bool) {
+	for i := len(st.locals) - 1; i >= 0; i-- {
+		if st.locals[i].Name == name {
+			return st.locals[i], true
+		}
+	}
+
+	return nil, false
+}
+
+func (st *SymbolTable) Params() []string {
+	out := make([]string, 0, len(st.params))
+	for k := range st.params {
+		out = append(out, k)
+	}
+	return out
+}
+
+func (st *SymbolTable) DebugSymbols() []string {
+	var out []string
+
+	for _, v := range st.locals {
+		out = append(out, fmt.Sprintf("[local] %s -> R%d (%v)", v.Name, v.Register, v.Type))
+	}
+
+	for k, r := range st.globals {
+		out = append(out, fmt.Sprintf("[global] %s -> R%d", k, r))
+	}
+
+	return out
+}

pkg/compiler/internal/visitor.go

@@ -144,7 +144,7 @@ func (v *Visitor) visitWaitForEventExpression(ctx *fql.WaitForExpressionContext)
 	jumpToNext := v.Emitter.EmitJumpc(vm.OpIterNext, jumpPlaceholder, streamReg)
 
 	if filter := ctx.FilterClause(); filter != nil {
-		valReg = v.Symbols.DefineVariable(pseudoVariable)
+		valReg = v.Symbols.DeclareLocal(pseudoVariable)
 		v.Emitter.EmitAB(vm.OpIterValue, valReg, streamReg)
 
 		filter.Expression().Accept(v)
@@ -256,14 +256,14 @@ func (v *Visitor) VisitForExpression(ctx *fql.ForExpressionContext) interface{}
 		if val := ctx.GetValueVariable(); val != nil {
 			if txt := val.GetText(); txt != "" && txt != ignorePseudoVariable {
 				loop.ValueName = txt
-				loop.Value = v.Symbols.DefineVariable(txt)
+				loop.Value = v.Symbols.DeclareLocal(txt)
 			}
 		}
 
 		if ctr := ctx.GetCounterVariable(); ctr != nil {
 			if txt := ctr.GetText(); txt != "" && txt != ignorePseudoVariable {
 				loop.KeyName = txt
-				loop.Key = v.Symbols.DefineVariable(txt)
+				loop.Key = v.Symbols.DeclareLocal(txt)
 			}
 		}
 	} else {
@@ -467,7 +467,7 @@ func (v *Visitor) VisitCollectClause(ctx *fql.CollectClauseContext) interface{}
 		if projectionVariableName != "" {
 			// Now we need to expand group variables from the dataset
 			v.emitIterKey(loop, kvValReg)
-			v.emitIterValue(loop, v.Symbols.DefineVariable(projectionVariableName))
+			v.emitIterValue(loop, v.Symbols.DeclareLocal(projectionVariableName))
 		} else {
 			v.emitIterKey(loop, kvKeyReg)
 			v.emitIterValue(loop, kvValReg)
@@ -527,11 +527,11 @@ func (v *Visitor) emitCollectAggregator(c fql.ICollectAggregatorContext, parentL
 	}
 
 	// Store upper scope for aggregators
-	mainScope := v.Symbols.Scope()
+	//mainScope := v.Symbols.Scope()
 	// Nested scope for aggregators
 	v.Symbols.EnterScope()
 
-	aggrIterVal := v.Symbols.DefineVariable(loop.ValueName)
+	aggrIterVal := v.Symbols.DeclareLocal(loop.ValueName)
 	v.Emitter.EmitAB(vm.OpIterValue, aggrIterVal, loop.Iterator)
 
 	// Now we add value selectors to the accumulators
@@ -577,11 +577,11 @@ func (v *Visitor) emitCollectAggregator(c fql.ICollectAggregatorContext, parentL
 
 			// We execute the function call with the accumulator as an argument
 			accum := accums[i]
-			result := v.emitFunctionCall(fcx.FunctionCall(), fcx.ErrorOperator() != nil, NewRegisterSequence(accum))
+			result := v.emitFunctionCall(fcx.FunctionCall(), fcx.ErrorOperator() != nil, RegisterSequence{accum})
 
 			// We define the variable for the selector result in the upper scope
 			// Since this temporary scope is only for aggregators and will be closed after the aggregation
-			varReg := v.Symbols.DefineVariableInScope(selectorVarName, mainScope)
+			varReg := v.Symbols.DeclareLocal(selectorVarName)
 			v.Emitter.EmitAB(vm.OpMove, varReg, result)
 			v.Registers.Free(result)
 		}
@@ -623,11 +623,11 @@ func (v *Visitor) emitCollectAggregator(c fql.ICollectAggregatorContext, parentL
 			value := v.Registers.Allocate(Temp)
 			v.Emitter.EmitABC(vm.OpLoadKey, value, aggregator, key)
 
-			result := v.emitFunctionCall(fcx.FunctionCall(), fcx.ErrorOperator() != nil, NewRegisterSequence(value))
+			result := v.emitFunctionCall(fcx.FunctionCall(), fcx.ErrorOperator() != nil, RegisterSequence{value})
 
 			// We define the variable for the selector result in the upper scope
 			// Since this temporary scope is only for aggregators and will be closed after the aggregation
-			varReg := v.Symbols.DefineVariableInScope(selectorVarName, mainScope)
+			varReg := v.Symbols.DeclareLocal(selectorVarName)
 			v.Emitter.EmitAB(vm.OpMove, varReg, result)
 			v.Registers.Free(result)
 			v.Registers.Free(value)
@@ -660,7 +660,7 @@ func (v *Visitor) emitCollectGroupKeySelectors(selectors []fql.ICollectSelectorC
 
 		for i, selector := range selectors {
 			reg := selector.Accept(v).(vm.Operand)
-			v.Emitter.EmitAB(vm.OpMove, selectorRegs.Registers[i], reg)
+			v.Emitter.EmitAB(vm.OpMove, selectorRegs[i], reg)
 			// Free the register after moving its value to the sequence register
 			v.Registers.Free(reg)
 		}
@@ -683,7 +683,7 @@ func (v *Visitor) emitCollectGroupKeySelectorVariables(selectors []fql.ICollectS
 			name := selector.Identifier().GetText()
 
 			if variables[i] == vm.NoopOperand {
-				variables[i] = v.Symbols.DefineVariable(name)
+				variables[i] = v.Symbols.DeclareLocal(name)
 			}
 
 			reg := kvValReg
@@ -703,7 +703,7 @@ func (v *Visitor) emitCollectGroupKeySelectorVariables(selectors []fql.ICollectS
 		// Get the variable name
 		name := selectors[0].Identifier().GetText()
 		// Define a variable for each selector
-		varReg := v.Symbols.DefineVariable(name)
+		varReg := v.Symbols.DeclareLocal(name)
 
 		reg := kvValReg
 
@@ -722,8 +722,8 @@ func (v *Visitor) emitCollectDefaultGroupProjection(loop *Loop, kvValReg vm.Oper
 
 		// TODO: Review this. It's quite a questionable ArrangoDB feature of wrapping group items by a nested object
 		// We will keep it for now for backward compatibility.
-		v.loadConstantTo(runtime.String(loop.ValueName), seq.Registers[0]) // Map key
-		v.Emitter.EmitAB(vm.OpMove, seq.Registers[1], kvValReg)            // Map value
+		v.loadConstantTo(runtime.String(loop.ValueName), seq[0]) // Map key
+		v.Emitter.EmitAB(vm.OpMove, seq[1], kvValReg)            // Map value
 		v.Emitter.EmitAs(vm.OpMap, kvValReg, seq)
 
 		v.Registers.FreeSequence(seq)
@@ -733,8 +733,15 @@ func (v *Visitor) emitCollectDefaultGroupProjection(loop *Loop, kvValReg vm.Oper
 
 		for i, j := 0, 0; i < len(variables); i, j = i+1, j+2 {
 			varName := variables[i].GetText()
-			v.loadConstantTo(runtime.String(varName), seq.Registers[j])
-			v.Emitter.EmitAB(vm.OpMove, seq.Registers[j+1], v.Symbols.Variable(varName))
+			v.loadConstantTo(runtime.String(varName), seq[j])
+
+			variable, _, found := v.Symbols.Resolve(varName)
+
+			if !found {
+				panic("variable not found: " + varName)
+			}
+
+			v.Emitter.EmitAB(vm.OpMove, seq[j+1], variable)
 		}
 
 		v.Emitter.EmitAs(vm.OpMap, kvValReg, seq)
@@ -785,8 +792,8 @@ func (v *Visitor) VisitSortClause(ctx *fql.SortClauseContext) interface{} {
 
 		for i, clause := range clauses {
 			clauseReg := clause.Accept(v).(vm.Operand)
-			v.Emitter.EmitAB(vm.OpMove, keyRegs.Registers[i], clauseReg)
-			clausesRegs[i] = keyRegs.Registers[i]
+			v.Emitter.EmitAB(vm.OpMove, keyRegs[i], clauseReg)
+			clausesRegs[i] = keyRegs[i]
 			directions[i] = v.sortDirection(clause.SortDirection())
 			// TODO: Free Registers
 		}
@@ -983,7 +990,7 @@ func (v *Visitor) VisitRangeOperand(ctx *fql.RangeOperandContext) interface{} {
 func (v *Visitor) VisitParam(ctx *fql.ParamContext) interface{} {
 	name := ctx.Identifier().GetText()
 	reg := v.Registers.Allocate(Temp)
-	v.Emitter.EmitAB(vm.OpLoadParam, reg, v.Symbols.AddParam(name))
+	v.Emitter.EmitAB(vm.OpLoadParam, reg, v.Symbols.BindParam(name))
 
 	return reg
 }
@@ -1000,15 +1007,18 @@ func (v *Visitor) VisitVariableDeclaration(ctx *fql.VariableDeclarationContext)
 	src := ctx.Expression().Accept(v).(vm.Operand)
 
 	if name != ignorePseudoVariable {
-		dest := v.Symbols.DefineVariable(name)
+		var dest vm.Operand
 
 		if src.IsConstant() {
+			dest = v.Symbols.DeclareGlobal(name)
 			tmp := v.Registers.Allocate(Temp)
 			v.Emitter.EmitAB(vm.OpLoadConst, tmp, src)
 			v.Emitter.EmitAB(vm.OpStoreGlobal, dest, tmp)
 		} else if v.Symbols.Scope() == 0 {
+			dest = v.Symbols.DeclareGlobal(name)
 			v.Emitter.EmitAB(vm.OpStoreGlobal, dest, src)
 		} else {
+			dest = v.Symbols.DeclareLocal(name)
 			v.Emitter.EmitAB(vm.OpMove, dest, src)
 		}
 
@@ -1020,7 +1030,7 @@ func (v *Visitor) VisitVariableDeclaration(ctx *fql.VariableDeclarationContext)
 
 func (v *Visitor) VisitVariable(ctx *fql.VariableContext) interface{} {
 	// Just return the register / constant index
-	op := v.Symbols.Variable(ctx.GetText())
+	op, _, _ := v.Symbols.Resolve(ctx.GetText())
 
 	if op.IsRegister() {
 		return op
@@ -1051,7 +1061,7 @@ func (v *Visitor) VisitArrayLiteral(ctx *fql.ArrayLiteralContext) interface{} {
 				srcReg := exp.Accept(v).(vm.Operand)
 
 				// TODO: Figure out how to remove OpMove and use Registers returned from each expression
-				v.Emitter.EmitAB(vm.OpMove, seq.Registers[i], srcReg)
+				v.Emitter.EmitAB(vm.OpMove, seq[i], srcReg)
 
 				// Free source register if temporary
 				if srcReg.IsRegister() {
@@ -1106,8 +1116,8 @@ func (v *Visitor) VisitObjectLiteral(ctx *fql.ObjectLiteralContext) interface{}
 
 		regIndex := i * 2
 
-		v.Emitter.EmitAB(vm.OpMove, seq.Registers[regIndex], propOp)
-		v.Emitter.EmitAB(vm.OpMove, seq.Registers[regIndex+1], valOp)
+		v.Emitter.EmitAB(vm.OpMove, seq[regIndex], propOp)
+		v.Emitter.EmitAB(vm.OpMove, seq[regIndex+1], valOp)
 
 		// Free source register if temporary
 		if propOp.IsRegister() {
@@ -1535,7 +1545,7 @@ func (v *Visitor) VisitExpressionAtom(ctx *fql.ExpressionAtomContext) interface{
 
 func (v *Visitor) visitFunctionCall(ctx *fql.FunctionCallContext, protected bool) interface{} {
 	var size int
-	var seq *RegisterSequence
+	var seq RegisterSequence
 
 	if list := ctx.ArgumentList(); list != nil {
 		// Get all array element expressions
@@ -1552,7 +1562,7 @@ func (v *Visitor) visitFunctionCall(ctx *fql.FunctionCallContext, protected bool
 				srcReg := exp.Accept(v).(vm.Operand)
 
 				// TODO: Figure out how to remove OpMove and use Registers returned from each expression
-				v.Emitter.EmitAB(vm.OpMove, seq.Registers[i], srcReg)
+				v.Emitter.EmitAB(vm.OpMove, seq[i], srcReg)
 
 				// Free source register if temporary
 				if srcReg.IsRegister() {
@@ -1565,38 +1575,38 @@ func (v *Visitor) visitFunctionCall(ctx *fql.FunctionCallContext, protected bool
 	return v.emitFunctionCall(ctx, protected, seq)
 }
 
-func (v *Visitor) emitFunctionCall(ctx fql.IFunctionCallContext, protected bool, seq *RegisterSequence) vm.Operand {
+func (v *Visitor) emitFunctionCall(ctx fql.IFunctionCallContext, protected bool, seq RegisterSequence) vm.Operand {
 	name := v.functionName(ctx)
 
 	switch name {
 	case runtimeLength:
 		dst := v.Registers.Allocate(Temp)
 
-		if seq == nil || len(seq.Registers) != 1 {
+		if seq == nil || len(seq) != 1 {
 			panic(runtime.Error(runtime.ErrInvalidArgument, runtimeLength+": expected 1 argument"))
 		}
 
-		v.Emitter.EmitAB(vm.OpLength, dst, seq.Registers[0])
+		v.Emitter.EmitAB(vm.OpLength, dst, seq[0])
 
 		return dst
 	case runtimeTypename:
 		dst := v.Registers.Allocate(Temp)
 
-		if seq == nil || len(seq.Registers) != 1 {
+		if seq == nil || len(seq) != 1 {
 			panic(runtime.Error(runtime.ErrInvalidArgument, runtimeTypename+": expected 1 argument"))
 		}
 
-		v.Emitter.EmitAB(vm.OpType, dst, seq.Registers[0])
+		v.Emitter.EmitAB(vm.OpType, dst, seq[0])
 
 		return dst
 	case runtimeWait:
-		if seq == nil || len(seq.Registers) != 1 {
+		if len(seq) != 1 {
 			panic(runtime.Error(runtime.ErrInvalidArgument, runtimeWait+": expected 1 argument"))
 		}
 
-		v.Emitter.EmitA(vm.OpSleep, seq.Registers[0])
+		v.Emitter.EmitA(vm.OpSleep, seq[0])
 
-		return seq.Registers[0]
+		return seq[0]
 	default:
 		nameAndDest := v.loadConstant(v.functionName(ctx))
 

test/integration/vm/vm_func_test.go

@@ -23,7 +23,7 @@ func TestFunctionCall(t *testing.T) {
 		SkipCase("LET duration = 10 WAIT(duration) RETURN 1", 1),
 		CaseNil("RETURN (FALSE OR T::FAIL())?"),
 		CaseNil("RETURN T::FAIL()?"),
-		SkipCaseArray(`FOR i IN [1, 2, 3, 4]
+		CaseArray(`FOR i IN [1, 2, 3, 4]
 				LET duration = 10
 		
 				WAIT(duration)