Refactor emitter and loop compilation to replace positional jumps with labeled jumps, improve label management, and optimize control flow handling. Update integration tests and add utility for safe function execution.

pkg/compiler/compiler.go

@@ -65,7 +65,6 @@ func (c *Compiler) Compile(query string) (program *vm.Program, err error) {
 	p.AddErrorListener(newErrorListener())
 
 	l := NewVisitor(query)
-
 	p.Visit(l)
 
 	if l.Err != nil {

pkg/compiler/internal/core/emitter.go

@@ -1,11 +1,15 @@
 package core
 
 import (
+	"fmt"
 	"github.com/MontFerret/ferret/pkg/vm"
 )
 
 type Emitter struct {
 	instructions []vm.Instruction
+	labels       map[Label]labelDef
+	patches      map[Label][]labelRef
+	nextLabelID  Label
 }
 
 func NewEmitter() *Emitter {
@@ -23,105 +27,44 @@ func (e *Emitter) Size() int {
 	return len(e.instructions)
 }
 
+func (e *Emitter) NewLabel() Label {
+	l := e.nextLabelID
+	e.nextLabelID++
+
+	return l
+}
+
+func (e *Emitter) MarkLabel(label Label) {
+	if e.labels == nil {
+		e.labels = make(map[Label]labelDef)
+	}
+
+	e.labels[label] = labelDef{addr: len(e.instructions)}
+
+	// Back-patch any prior references to this label
+	if refs, ok := e.patches[label]; ok {
+		for _, ref := range refs {
+			e.patchOperand(ref.pos, ref.field, len(e.instructions))
+		}
+
+		delete(e.patches, label)
+	}
+}
+
+func (e *Emitter) LabelPosition(label Label) (int, bool) {
+	def, ok := e.labels[label]
+
+	if !ok {
+		return -1, false
+	}
+
+	return def.addr, true
+}
+
 func (e *Emitter) Position() int {
 	return len(e.instructions) - 1
 }
 
-func (e *Emitter) Patchx(pos int, arg int) {
-	current := e.instructions[pos]
-	e.instructions[pos] = vm.Instruction{
-		Opcode: current.Opcode,
-		Operands: [3]vm.Operand{
-			current.Operands[0],
-			vm.Operand(arg),
-			current.Operands[2],
-		},
-	}
-}
-
-// 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,
-		Operands: [3]vm.Operand{dst, src1, vm.NoopOperand},
-	}
-}
-
-// 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,
-		Operands: [3]vm.Operand{dst, vm.Operand(arg), vm.NoopOperand},
-	}
-}
-
-// 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,
-		Operands: [3]vm.Operand{dst, vm.Operand(arg1), vm.Operand(agr2)},
-	}
-}
-
-// 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[0], seq[len(seq)-1]},
-	}
-}
-
-// PatchInsertAx inserts a new instruction at a specific position in the instructions slice, shifting elements to the right.
-// The inserted instruction includes an opcode and operands, where the third operand is set to a no-op by default.
-func (e *Emitter) PatchInsertAx(pos int, op vm.Opcode, dst vm.Operand, arg int) {
-	// Append a zero value to create space
-	e.instructions = append(e.instructions, vm.Instruction{})
-
-	// Shift elements to the right
-	copy(e.instructions[pos+1:], e.instructions[pos:])
-
-	// Insert the new value
-	e.instructions[pos] = vm.Instruction{
-		Opcode:   op,
-		Operands: [3]vm.Operand{dst, vm.Operand(arg), vm.NoopOperand},
-	}
-}
-
-// PatchInsertAxy inserts an instruction at the specified position in the instruction list, shifting existing elements to the right.
-func (e *Emitter) PatchInsertAxy(pos int, op vm.Opcode, dst vm.Operand, arg1, arg2 int) {
-	// Append a zero value to create space
-	e.instructions = append(e.instructions, vm.Instruction{})
-
-	// Shift elements to the right
-	copy(e.instructions[pos+1:], e.instructions[pos:])
-
-	// Insert the new value
-	e.instructions[pos] = vm.Instruction{
-		Opcode:   op,
-		Operands: [3]vm.Operand{dst, vm.Operand(arg1), vm.Operand(arg2)},
-	}
-}
-
-// PatchJump patches a jump opcode.
-func (e *Emitter) PatchJump(instr int) {
-	e.instructions[instr].Operands[0] = vm.Operand(len(e.instructions) - 1)
-}
-
-// PatchJumpAB patches a jump opcode with a new destination.
-func (e *Emitter) PatchJumpAB(inst int) {
-	e.instructions[inst].Operands[2] = vm.Operand(len(e.instructions) - 1)
-}
-
-// PatchJumpNextAB patches a jump instruction to jump over a current position.
-func (e *Emitter) PatchJumpNextAB(instr int) {
-	e.instructions[instr].Operands[2] = vm.Operand(len(e.instructions))
-}
-
-// PatchJumpNext patches a jump instruction to jump over a current position.
-func (e *Emitter) PatchJumpNext(instr int) {
-	e.instructions[instr].Operands[0] = vm.Operand(len(e.instructions))
-}
-
 // Emit emits an opcode with no arguments.
 func (e *Emitter) Emit(op vm.Opcode) {
 	e.EmitABC(op, 0, 0, 0)
@@ -181,3 +124,137 @@ func (e *Emitter) EmitABC(op vm.Opcode, dest, src1, src2 vm.Operand) {
 		Operands: [3]vm.Operand{dest, src1, src2},
 	})
 }
+
+// SwapAB modifies an instruction at the given position to swap operands and update its operation and destination.
+func (e *Emitter) SwapAB(label Label, op vm.Opcode, dst, src1 vm.Operand) {
+	e.swapInstruction(label, vm.Instruction{
+		Opcode:   op,
+		Operands: [3]vm.Operand{dst, src1, vm.NoopOperand},
+	})
+}
+
+// SwapAx modifies an existing instruction at the specified position with a new opcode, destination, and argument.
+func (e *Emitter) SwapAx(label Label, op vm.Opcode, dst vm.Operand, arg int) {
+	e.swapInstruction(label, vm.Instruction{
+		Opcode:   op,
+		Operands: [3]vm.Operand{dst, vm.Operand(arg), vm.NoopOperand},
+	})
+}
+
+// SwapAxy replaces an instruction at the specified position with a new one using the provided opcode and operands.
+func (e *Emitter) SwapAxy(label Label, op vm.Opcode, dst vm.Operand, arg1, agr2 int) {
+	e.swapInstruction(label, vm.Instruction{
+		Opcode:   op,
+		Operands: [3]vm.Operand{dst, vm.Operand(arg1), vm.Operand(agr2)},
+	})
+}
+
+// SwapAs replaces an instruction at the specified position with a new instruction using the given opcode and operands.
+func (e *Emitter) SwapAs(label Label, op vm.Opcode, dst vm.Operand, seq RegisterSequence) {
+	e.swapInstruction(label, vm.Instruction{
+		Opcode:   op,
+		Operands: [3]vm.Operand{dst, seq[0], seq[len(seq)-1]},
+	})
+}
+
+// InsertAx inserts a new instruction at a specific position in the instructions slice, shifting elements to the right.
+// The inserted instruction includes an opcode and operands, where the third operand is set to a no-op by default.
+func (e *Emitter) InsertAx(label Label, op vm.Opcode, dst vm.Operand, arg int) {
+	e.insertInstruction(label, vm.Instruction{
+		Opcode:   op,
+		Operands: [3]vm.Operand{dst, vm.Operand(arg), vm.NoopOperand},
+	})
+}
+
+// InsertAxy inserts an instruction at the specified position in the instruction list, shifting existing elements to the right.
+func (e *Emitter) InsertAxy(label Label, op vm.Opcode, dst vm.Operand, arg1, arg2 int) {
+	e.insertInstruction(label, vm.Instruction{
+		Opcode:   op,
+		Operands: [3]vm.Operand{dst, vm.Operand(arg1), vm.Operand(arg2)},
+	})
+}
+
+func (e *Emitter) Patchx(label Label, arg int) {
+	pos, ok := e.LabelPosition(label)
+
+	if !ok {
+		panic(fmt.Errorf("label not marked: %d", label))
+	}
+
+	current := e.instructions[pos]
+	e.instructions[pos] = vm.Instruction{
+		Opcode: current.Opcode,
+		Operands: [3]vm.Operand{
+			current.Operands[0],
+			vm.Operand(arg),
+			current.Operands[2],
+		},
+	}
+}
+
+// addLabelRef adds a reference to a label at a specific position and field in the instruction set.
+func (e *Emitter) addLabelRef(pos int, field int, label Label) {
+	if e.labels == nil {
+		e.labels = make(map[Label]labelDef)
+	}
+
+	if def, ok := e.labels[label]; ok {
+		// Already marked → patch immediately
+		e.patchOperand(pos, field, def.addr)
+		return
+	}
+
+	if e.patches == nil {
+		e.patches = make(map[Label][]labelRef)
+	}
+
+	e.patches[label] = append(e.patches[label], labelRef{pos: pos, field: field})
+}
+
+// patchOperand modifies the operand at the specified position and field in the instruction set.
+func (e *Emitter) patchOperand(pos int, field int, val int) {
+	ins := e.instructions[pos]
+	ins.Operands[field] = vm.Operand(val)
+	e.instructions[pos] = ins
+}
+
+// swapInstruction swaps the operands of an instruction at a given position.
+func (e *Emitter) swapInstruction(label Label, ins vm.Instruction) {
+	pos, ok := e.LabelPosition(label)
+
+	if !ok {
+		panic(fmt.Errorf("label not marked: %d", label))
+	}
+
+	e.instructions[pos] = ins
+}
+
+// swapInstruction swaps the operands of an instruction at a given position.
+func (e *Emitter) insertInstruction(label Label, ins vm.Instruction) {
+	pos, ok := e.LabelPosition(label)
+
+	if !ok {
+		panic(fmt.Errorf("label not marked: %d", label))
+	}
+
+	// Insert instruction at position
+	e.instructions = append(e.instructions[:pos],
+		append([]vm.Instruction{ins}, e.instructions[pos:]...)...,
+	)
+
+	// Adjust all subsequent label addresses
+	for l, d := range e.labels {
+		if d.addr >= pos {
+			e.labels[l] = labelDef{addr: d.addr + 1}
+		}
+	}
+
+	// Adjust all patch positions as well
+	for l, refs := range e.patches {
+		for i, ref := range refs {
+			if ref.pos >= pos {
+				e.patches[l][i].pos++
+			}
+		}
+	}
+}

pkg/compiler/internal/core/emitter_helpers.go

@@ -12,8 +12,8 @@ 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) EmitIterNext(iterator vm.Operand, label Label) {
+	e.EmitJumpc(vm.OpIterNext, iterator, label)
 }
 
 func (e *Emitter) EmitIterKey(dst, iterator vm.Operand) {
@@ -24,12 +24,16 @@ 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) EmitIterSkip(state, count vm.Operand, label Label) {
+	e.EmitABx(vm.OpIterSkip, state, count, jumpPlaceholder)
+	pos := len(e.instructions) - 1
+	e.addLabelRef(pos, 2, label)
 }
 
-func (e *Emitter) EmitIterLimit(state, count vm.Operand, jump int) {
-	e.EmitABx(vm.OpIterLimit, state, count, jump)
+func (e *Emitter) EmitIterLimit(state, count vm.Operand, jump Label) {
+	e.EmitABx(vm.OpIterLimit, state, count, jumpPlaceholder)
+	pos := len(e.instructions) - 1
+	e.addLabelRef(pos, 2, jump)
 }
 
 // ─── Value & Memory ──────────────────────────────────────────────────────
@@ -152,40 +156,41 @@ func (e *Emitter) EmitLte(dst, a, b vm.Operand) {
 
 // ─── Control Flow ────────────────────────────────────────────────────────
 
-func (e *Emitter) EmitJump(pos int) int {
-	e.EmitA(vm.OpJump, vm.Operand(pos))
-
-	return len(e.instructions) - 1
+func (e *Emitter) EmitJump(label Label) {
+	e.EmitA(vm.OpJump, vm.Operand(jumpPlaceholder))
+	pos := len(e.instructions) - 1
+	e.addLabelRef(pos, 0, label)
 }
 
 // EmitJumpAB emits a jump opcode with a state and an argument.
-func (e *Emitter) EmitJumpAB(op vm.Opcode, state, cond vm.Operand, pos int) int {
-	e.EmitABC(op, state, cond, vm.Operand(pos))
-
-	return len(e.instructions) - 1
+func (e *Emitter) EmitJumpAB(op vm.Opcode, state, cond vm.Operand, label Label) {
+	e.EmitABC(op, state, cond, vm.Operand(jumpPlaceholder))
+	pos := len(e.instructions) - 1
+	e.addLabelRef(pos, 2, label)
 }
 
 // EmitJumpc emits a conditional jump opcode.
-func (e *Emitter) EmitJumpc(op vm.Opcode, pos int, reg vm.Operand) int {
-	e.EmitAB(op, vm.Operand(pos), reg)
-
-	return len(e.instructions) - 1
+func (e *Emitter) EmitJumpc(op vm.Opcode, reg vm.Operand, label Label) {
+	e.EmitAB(op, vm.Operand(jumpPlaceholder), reg)
+	pos := len(e.instructions) - 1
+	e.addLabelRef(pos, 0, label)
 }
 
-func (e *Emitter) EmitJumpIfFalse(cond vm.Operand, jumpTarget int) int {
-	return e.EmitJumpIf(cond, false, jumpTarget)
+func (e *Emitter) EmitJumpIfFalse(cond vm.Operand, label Label) {
+	e.EmitJumpIf(cond, false, label)
 }
 
-func (e *Emitter) EmitJumpIfTrue(cond vm.Operand, jumpTarget int) int {
-	return e.EmitJumpIf(cond, true, jumpTarget)
+func (e *Emitter) EmitJumpIfTrue(cond vm.Operand, label Label) {
+	e.EmitJumpIf(cond, true, label)
 }
 
-func (e *Emitter) EmitJumpIf(cond vm.Operand, isTrue bool, jumpTarget int) int {
+func (e *Emitter) EmitJumpIf(cond vm.Operand, isTrue bool, label Label) {
 	if isTrue {
-		return e.EmitJumpc(vm.OpJumpIfTrue, jumpTarget, cond)
+		e.EmitJumpc(vm.OpJumpIfTrue, cond, label)
+		return
 	}
 
-	return e.EmitJumpc(vm.OpJumpIfFalse, jumpTarget, cond)
+	e.EmitJumpc(vm.OpJumpIfFalse, cond, label)
 }
 
 func (e *Emitter) EmitReturnValue(val vm.Operand) {

pkg/compiler/internal/core/label.go

@@ -0,0 +1,16 @@
+package core
+
+type (
+	Label int
+
+	labelRef struct {
+		pos   int
+		field int
+	}
+
+	labelDef struct {
+		addr int
+	}
+)
+
+const InvalidLabel Label = -1

pkg/compiler/internal/core/loop.go

@@ -35,9 +35,9 @@ type Loop struct {
 	Distinct bool
 	Allocate bool
 
-	Pos        int
-	Jump       int
-	JumpOffset int
+	Start Label
+	Jump  Label
+	End   Label
 
 	Src      vm.Operand
 	Iterator vm.Operand
@@ -65,20 +65,23 @@ func (l *Loop) DeclareValueVar(name string, st *SymbolTable) {
 }
 
 func (l *Loop) EmitInitialization(alloc *RegisterAllocator, emitter *Emitter) {
+	l.Start = emitter.NewLabel()
+	emitter.MarkLabel(l.Start)
+
 	if l.Allocate {
 		emitter.EmitAb(vm.OpDataSet, l.Dst, l.Distinct)
 	}
 
-	l.Pos = emitter.Position()
-
 	if l.Iterator == vm.NoopOperand {
 		l.Iterator = alloc.Allocate(Temp)
 	}
 
 	emitter.EmitIter(l.Iterator, l.Src)
 
-	// JumpPlaceholder is a placeholder for the exit jump position
-	l.Jump = emitter.EmitJumpc(vm.OpIterNext, JumpPlaceholder, l.Iterator)
+	l.Jump = emitter.NewLabel()
+	l.End = emitter.NewLabel()
+	emitter.MarkLabel(l.Jump)
+	emitter.EmitJumpc(vm.OpIterNext, l.Iterator, l.End)
 
 	if l.canBindVar(l.Value) {
 		l.EmitValue(l.Value, emitter)
@@ -98,34 +101,32 @@ func (l *Loop) EmitKey(dst vm.Operand, emitter *Emitter) {
 }
 
 func (l *Loop) EmitFinalization(emitter *Emitter) {
-	emitter.EmitJump(l.Jump - l.JumpOffset)
+	emitter.EmitJump(l.Jump)
+	emitter.MarkLabel(l.End)
 	emitter.EmitA(vm.OpClose, l.Iterator)
-	emitter.PatchJump(l.Jump)
 }
 
 func (l *Loop) PatchDestinationAx(alloc *RegisterAllocator, emitter *Emitter, op vm.Opcode, arg int) vm.Operand {
 	if l.Allocate {
-		emitter.PatchSwapAx(l.Pos, op, l.Dst, arg)
+		emitter.SwapAx(l.Start, op, l.Dst, arg)
 
 		return l.Dst
 	}
 
 	tmp := alloc.Allocate(Temp)
-	emitter.PatchInsertAx(l.Pos, op, tmp, arg)
-	l.Jump++
+	emitter.InsertAx(l.Start, op, tmp, arg)
 	return tmp
 }
 
 func (l *Loop) PatchDestinationAxy(alloc *RegisterAllocator, emitter *Emitter, op vm.Opcode, arg1, arg2 int) vm.Operand {
 	if l.Allocate {
-		emitter.PatchSwapAxy(l.Pos, op, l.Dst, arg1, arg2)
+		emitter.SwapAxy(l.Start, op, l.Dst, arg1, arg2)
 
 		return l.Dst
 	}
 
 	tmp := alloc.Allocate(Temp)
-	emitter.PatchInsertAxy(l.Pos, op, tmp, arg1, arg2)
-	l.Jump++
+	emitter.InsertAxy(l.Start, op, tmp, arg1, arg2)
 	return tmp
 }
 

pkg/compiler/internal/core/symbols.go

@@ -8,7 +8,7 @@ import (
 )
 
 const (
-	JumpPlaceholder      = -1
+	jumpPlaceholder      = -1
 	UndefinedVariable    = -1
 	IgnorePseudoVariable = "_"
 	PseudoVariable       = "CURRENT"

pkg/compiler/internal/expr.go

@@ -76,22 +76,18 @@ func (ec *ExprCompiler) compileUnary(ctx fql.IUnaryOperatorContext, parent fql.I
 func (ec *ExprCompiler) compileLogicalAnd(ctx fql.IExpressionContext) vm.Operand {
 	dst := ec.ctx.Registers.Allocate(core.Temp)
 
-	// Execute left expression
 	left := ec.Compile(ctx.GetLeft())
-
-	// Execute left expression
 	ec.ctx.Emitter.EmitMove(dst, left)
 
-	// Test if left is false and jump to the end
-	end := ec.ctx.Emitter.EmitJumpIfFalse(dst, core.JumpPlaceholder)
+	end := ec.ctx.Emitter.NewLabel()
+
+	// If left is false, jump to end
+	ec.ctx.Emitter.EmitJumpIfFalse(dst, end)
 
-	// If left is true, execute right expression
 	right := ec.Compile(ctx.GetRight())
-
-	// And move the result to the destination register
 	ec.ctx.Emitter.EmitMove(dst, right)
 
-	ec.ctx.Emitter.PatchJumpNext(end)
+	ec.ctx.Emitter.MarkLabel(end)
 
 	return dst
 }
@@ -100,22 +96,18 @@ func (ec *ExprCompiler) compileLogicalAnd(ctx fql.IExpressionContext) vm.Operand
 func (ec *ExprCompiler) compileLogicalOr(ctx fql.IExpressionContext) vm.Operand {
 	dst := ec.ctx.Registers.Allocate(core.Temp)
 
-	// Execute left expression
 	left := ec.Compile(ctx.GetLeft())
-
-	// Execute left expression
 	ec.ctx.Emitter.EmitMove(dst, left)
 
-	// Test if left is true and jump to the end
-	end := ec.ctx.Emitter.EmitJumpIfTrue(dst, core.JumpPlaceholder)
+	end := ec.ctx.Emitter.NewLabel()
+
+	// If left is true, jump to end
+	ec.ctx.Emitter.EmitJumpIfTrue(dst, end)
 
-	// If left is false, execute right expression
 	right := ec.Compile(ctx.GetRight())
-
-	// And move the result to the destination register
 	ec.ctx.Emitter.EmitMove(dst, right)
 
-	ec.ctx.Emitter.PatchJumpNext(end)
+	ec.ctx.Emitter.MarkLabel(end)
 
 	return dst
 }
@@ -128,8 +120,12 @@ func (ec *ExprCompiler) compileTernary(ctx fql.IExpressionContext) vm.Operand {
 	condReg := ec.Compile(ctx.GetCondition())
 	ec.ctx.Emitter.EmitMove(dst, condReg)
 
-	// Jump to 'false' branch if condition is false
-	otherwise := ec.ctx.Emitter.EmitJumpIfFalse(dst, core.JumpPlaceholder)
+	// Define jump labels
+	elseLabel := ec.ctx.Emitter.NewLabel()
+	endLabel := ec.ctx.Emitter.NewLabel()
+
+	// End to 'false' branch if condition is false
+	ec.ctx.Emitter.EmitJumpIfFalse(dst, elseLabel)
 
 	// True branch
 	if onTrue := ctx.GetOnTrue(); onTrue != nil {
@@ -138,9 +134,10 @@ func (ec *ExprCompiler) compileTernary(ctx fql.IExpressionContext) vm.Operand {
 		ec.ctx.Emitter.EmitMove(dst, trueReg)
 	}
 
-	// Jump over false branch
-	end := ec.ctx.Emitter.EmitJump(core.JumpPlaceholder)
-	ec.ctx.Emitter.PatchJumpNext(otherwise)
+	// End over false branch
+	ec.ctx.Emitter.EmitJump(endLabel)
+	// Mark label for 'else' branch
+	ec.ctx.Emitter.MarkLabel(elseLabel)
 
 	// False branch
 	if onFalse := ctx.GetOnFalse(); onFalse != nil {
@@ -149,7 +146,8 @@ func (ec *ExprCompiler) compileTernary(ctx fql.IExpressionContext) vm.Operand {
 		ec.ctx.Emitter.EmitMove(dst, falseReg)
 	}
 
-	ec.ctx.Emitter.PatchJumpNext(end)
+	// End
+	ec.ctx.Emitter.MarkLabel(endLabel)
 
 	return dst
 }

pkg/compiler/internal/loop.go

@@ -18,6 +18,14 @@ func NewLoopCompiler(ctx *CompilerContext) *LoopCompiler {
 }
 
 func (c *LoopCompiler) Compile(ctx fql.IForExpressionContext) vm.Operand {
+	if ctx.In() != nil {
+		return c.compileForIn(ctx)
+	}
+
+	return c.compileForWhile(ctx)
+}
+
+func (c *LoopCompiler) compileForIn(ctx fql.IForExpressionContext) vm.Operand {
 	returnRuleCtx := c.compileInitialization(ctx)
 
 	// body
@@ -34,6 +42,10 @@ func (c *LoopCompiler) Compile(ctx fql.IForExpressionContext) vm.Operand {
 	return c.compileFinalization(returnRuleCtx)
 }
 
+func (c *LoopCompiler) compileForWhile(ctx fql.IForExpressionContext) vm.Operand {
+	return vm.NoopOperand
+}
+
 func (c *LoopCompiler) compileInitialization(ctx fql.IForExpressionContext) antlr.RuleContext {
 	var distinct bool
 	var returnRuleCtx antlr.RuleContext
@@ -73,7 +85,7 @@ func (c *LoopCompiler) compileInitialization(ctx fql.IForExpressionContext) antl
 				panic("parent loop not found in loop table")
 			}
 
-			c.ctx.Emitter.Patchx(parent.Pos, 1)
+			c.ctx.Emitter.Patchx(parent.Start, 1)
 		}
 	}
 
@@ -196,17 +208,18 @@ func (c *LoopCompiler) compileLimitClauseValue(ctx fql.ILimitClauseValueContext)
 
 func (c *LoopCompiler) compileLimit(src vm.Operand) {
 	state := c.ctx.Registers.Allocate(core.State)
-	c.ctx.Emitter.EmitABx(vm.OpIterLimit, state, src, c.ctx.Loops.Current().Jump)
+	c.ctx.Emitter.EmitIterLimit(state, src, c.ctx.Loops.Current().End)
 }
 
 func (c *LoopCompiler) compileOffset(src vm.Operand) {
 	state := c.ctx.Registers.Allocate(core.State)
-	c.ctx.Emitter.EmitABx(vm.OpIterSkip, state, src, c.ctx.Loops.Current().Jump)
+	c.ctx.Emitter.EmitIterSkip(state, src, c.ctx.Loops.Current().Jump)
 }
 
 func (c *LoopCompiler) compileFilterClause(ctx fql.IFilterClauseContext) {
 	src := c.ctx.ExprCompiler.Compile(ctx.Expression())
-	c.ctx.Emitter.EmitJumpIfFalse(src, c.ctx.Loops.Current().Jump)
+	label := c.ctx.Loops.Current().Jump
+	c.ctx.Emitter.EmitJumpIfFalse(src, label)
 }
 
 func (c *LoopCompiler) compileSortClause(ctx fql.ISortClauseContext) {

pkg/compiler/internal/loop_collect_agg.go

@@ -126,8 +126,11 @@ func (c *LoopCollectCompiler) compileAggregationFuncCall(selectors []fql.ICollec
 	// Check if the number equals to zero
 	c.ctx.Emitter.EmitEq(cond, cond, zero)
 	c.ctx.Registers.Free(zero)
+	elseLabel := c.ctx.Emitter.NewLabel()
+	endLabel := c.ctx.Emitter.NewLabel()
+
 	// We skip the key retrieval and function call of there are no records in the accumulator
-	ifJump := c.ctx.Emitter.EmitJumpIfTrue(cond, core.JumpPlaceholder)
+	c.ctx.Emitter.EmitJumpIfTrue(cond, elseLabel)
 
 	selectorVarRegs := make([]vm.Operand, len(selectors))
 
@@ -166,14 +169,14 @@ func (c *LoopCollectCompiler) compileAggregationFuncCall(selectors []fql.ICollec
 		c.ctx.Registers.Free(result)
 	}
 
-	elseJump := c.ctx.Emitter.EmitJump(core.JumpPlaceholder)
-	c.ctx.Emitter.PatchJumpNext(ifJump)
+	c.ctx.Emitter.EmitJump(endLabel)
+	c.ctx.Emitter.MarkLabel(elseLabel)
 
 	for _, varReg := range selectorVarRegs {
 		c.ctx.Emitter.EmitA(vm.OpLoadNone, varReg)
 	}
 
-	c.ctx.Emitter.PatchJumpNext(elseJump)
+	c.ctx.Emitter.MarkLabel(endLabel)
 	c.ctx.Registers.Free(cond)
 }
 

pkg/compiler/internal/wait.go

@@ -43,24 +43,27 @@ func (wc *WaitCompiler) Compile(ctx fql.IWaitForExpressionContext) vm.Operand {
 	wc.ctx.Emitter.EmitAB(vm.OpStreamIter, streamReg, timeoutReg)
 
 	var valReg vm.Operand
+	start := wc.ctx.Emitter.NewLabel()
+	end := wc.ctx.Emitter.NewLabel()
 
+	wc.ctx.Emitter.MarkLabel(start)
 	// Now we start iterating over the stream
-	jumpToNext := wc.ctx.Emitter.EmitJumpc(vm.OpIterNext, core.JumpPlaceholder, streamReg)
+	wc.ctx.Emitter.EmitIterNext(streamReg, end)
 
 	if filter := ctx.FilterClause(); filter != nil {
 		valReg = wc.ctx.Symbols.DeclareLocal(core.PseudoVariable)
 		wc.ctx.Emitter.EmitAB(vm.OpIterValue, valReg, streamReg)
 
-		wc.ctx.ExprCompiler.Compile(filter.Expression())
+		cond := wc.ctx.ExprCompiler.Compile(filter.Expression())
 
-		wc.ctx.Emitter.EmitJumpc(vm.OpJumpIfFalse, jumpToNext, valReg)
+		wc.ctx.Emitter.EmitJumpIfFalse(cond, start)
 
 		// TODO: Do we need to use timeout here too? We can really get stuck in the loop if no event satisfies the filter
 	}
 
+	wc.ctx.Emitter.MarkLabel(end)
 	// Clean up the stream
 	wc.ctx.Emitter.EmitA(vm.OpClose, streamReg)
-
 	wc.ctx.Symbols.ExitScope()
 
 	return vm.NoopOperand

test/integration/base/try.go

@@ -0,0 +1,30 @@
+package base
+
+import (
+	"fmt"
+	"github.com/pkg/errors"
+)
+
+func Try[T any](f func() T) (T, error) {
+	var v T
+	var e error
+
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				switch x := r.(type) {
+				case error:
+					e = x
+				case string:
+					e = errors.New(x)
+				default:
+					e = errors.New(fmt.Sprintf("unknown error: %v", x))
+				}
+			}
+		}()
+
+		f()
+	}()
+
+	return v, e
+}

test/integration/vm/vm_case.go

@@ -76,6 +76,10 @@ func CaseItems(expression string, expected ...any) UseCase {
 	return NewCase(expression, expected, ShouldHaveSameItems)
 }
 
+func CaseFn(expression string, assertion func(actual any, expected ...any) string) UseCase {
+	return NewCase(expression, nil, assertion)
+}
+
 func SkipCaseItems(expression string, expected ...any) UseCase {
 	return Skip(CaseItems(expression, expected...))
 }

test/integration/vm/vm_for_test.go

@@ -63,9 +63,27 @@ FOR i IN 1..5
 			`FOR i IN ['foo', 'bar', 'qaz'] RETURN i`,
 			[]any{"foo", "bar", "qaz"},
 		),
-		CaseItems(
+		CaseFn(
 			`FOR i IN {a: 'bar', b: 'foo', c: 'qaz'} RETURN i`,
-			[]any{"bar", "foo", "qaz"},
+			func(actual any, expected ...any) string {
+				hashMap := make(map[string]bool)
+				expectedArr := []any{"bar", "foo", "qaz"}
+				actualArr := actual.([]any)
+
+				for _, v := range expectedArr {
+					hashMap[v.(string)] = false
+				}
+
+				for _, v := range actualArr {
+					if _, ok := hashMap[v.(string)]; !ok {
+						return "Unexpected value: " + v.(string)
+					}
+
+					hashMap[v.(string)] = true
+				}
+
+				return ""
+			},
 		),
 		CaseArray(
 			`FOR i, k IN {a: 'foo', b: 'bar', c: 'qaz'} RETURN k`,