feat: add WAITFOR expression support in loop constructs (#942)

pkg/compiler/internal/frontend.go

@@ -55,7 +55,7 @@ func NewCompilationFrontend(session *CompilationSession) *CompilationFrontend {
 	front.Collects.bind(front.Bindings, front.Calls, front.Expressions, front.Recovery, front.TypeFacts)
 	front.Statements.bind(front.Bindings, front.Dispatch, front.Expressions, front.Loops, front.TypeFacts, front.Wait)
 	front.UDFs.bind(front.Calls, front.Expressions, front.TypeFacts, front.Recovery, front.Statements)
-	front.Loops.bind(front.Bindings, front.Collects, front.Dispatch, front.Expressions, front.Literals, front.Recovery, front.Sorts, front.TypeFacts)
+	front.Loops.bind(front.Bindings, front.Collects, front.Dispatch, front.Expressions, front.Literals, front.Recovery, front.Sorts, front.TypeFacts, front.Wait)
 	front.Expressions.bind(front.Bindings, front.Calls, front.Dispatch, front.Literals, front.Loops, front.Recovery, front.TypeFacts, front.Wait)
 
 	return front

pkg/compiler/internal/loop.go

@@ -23,6 +23,7 @@ type (
 		recovery *RecoveryCompiler
 		sorts    *LoopSortCompiler
 		facts    *TypeFacts
+		wait     *WaitCompiler
 	}
 
 	loopOperandKind int
@@ -68,6 +69,7 @@ func (c *LoopCompiler) bind(
 	recovery *RecoveryCompiler,
 	sorts *LoopSortCompiler,
 	facts *TypeFacts,
+	wait *WaitCompiler,
 ) {
 	if c == nil {
 		return
@@ -81,6 +83,7 @@ func (c *LoopCompiler) bind(
 	c.recovery = recovery
 	c.sorts = sorts
 	c.facts = facts
+	c.wait = wait
 }
 
 // Compile processes a FOR expression from the FQL AST and generates the appropriate VM instructions.
@@ -417,6 +420,8 @@ func (c *LoopCompiler) compileForExpressionStatement(ctx fql.IForExpressionState
 	} else if fce := ctx.FunctionCallExpression(); fce != nil {
 		// Handle function calls (e.g., doSomething())
 		_ = c.exprs.CompileFunctionCallExpression(fce)
+	} else if wfe := ctx.WaitForExpression(); wfe != nil {
+		_ = c.wait.Compile(wfe)
 	} else if de := ctx.DispatchExpression(); de != nil {
 		_ = c.dispatch.Compile(de)
 	}

pkg/compiler/internal/name_collision_validator.go

@@ -217,6 +217,8 @@ func (v *NameCollisionValidator) collectForExpressionStatement(scope *nameCollis
 		v.collectCallsInNode(scope, ctx.DeleteStatement())
 	case ctx.FunctionCallExpression() != nil:
 		v.collectCallsInNode(scope, ctx.FunctionCallExpression())
+	case ctx.WaitForExpression() != nil:
+		v.collectCallsInNode(scope, ctx.WaitForExpression())
 	case ctx.DispatchExpression() != nil:
 		v.collectCallsInNode(scope, ctx.DispatchExpression())
 	}

pkg/formatter/internal/statement.go

@@ -478,6 +478,8 @@ func (f *statementFormatter) formatForExpressionBody(ctx *fql.ForExpressionBodyC
 			f.formatDeleteStatement(stmt.DeleteStatement().(*fql.DeleteStatementContext))
 		case stmt.FunctionCallExpression() != nil:
 			f.expression.formatFunctionCallExpression(stmt.FunctionCallExpression().(*fql.FunctionCallExpressionContext))
+		case stmt.WaitForExpression() != nil:
+			f.formatWaitForExpression(stmt.WaitForExpression().(*fql.WaitForExpressionContext))
 		case stmt.DispatchExpression() != nil:
 			f.formatDispatchExpression(stmt.DispatchExpression().(*fql.DispatchExpressionContext))
 		}

pkg/parser/antlr/FqlParser.g4

@@ -382,6 +382,7 @@ forExpressionStatement
     | assignmentStatement
     | deleteStatement
     | functionCallExpression
+    | waitForExpression
     | dispatchExpression
     ;
 

test/integration/compiler/compiler_waitfor_test.go

@@ -99,6 +99,12 @@ func TestWaitforPredicateWhenCompiles(t *testing.T) {
 				TIMEOUT 5ms
 				ON TIMEOUT RETURN NONE
 		`, expectHostFunction("BOOM", 1), "WAITFOR EVENT should compile repeated WHEN host calls and timeout tail"),
+		ProgramCheck(`
+			LET obs = []
+			FOR i IN [1, 2]
+				WAITFOR EVENT "test" IN obs TIMEOUT 5ms ON TIMEOUT RETURN NONE
+				RETURN i
+		`, noCompilerError, "WAITFOR EVENT should compile as a FOR loop body statement"),
 	})
 }
 

test/integration/formatter/formatter_loop_test.go

@@ -27,5 +27,10 @@ FOR WHILE i < 1
 	DISPATCH "click" IN @d
 	RETURN i
 	`, "FOR WHILE i < 1\n    DISPATCH \"click\" IN @d\n    RETURN i"),
+		S(`
+FOR WHILE ready
+WAITFOR EVENT "navigation" IN doc WHEN .type == "match" TIMEOUT 10s
+RETURN ready
+`, "FOR WHILE ready\n    WAITFOR EVENT \"navigation\" IN doc WHEN .type == \"match\" TIMEOUT 10s\n    RETURN ready"),
 	})
 }

test/integration/vm/vm_waitfor_event_test.go

@@ -78,6 +78,15 @@ LET evt = WAITFOR EVENT "test" IN obs WHEN .type != "" WHEN .type == "match"
 RETURN evt.type`, "match", "WAITFOR EVENT repeated filters should return the matched event value").Env(vm.WithParams(map[string]runtime.Value{
 			"obs": matchSecond,
 		})),
+		Array(`LET obs = @obs
+VAR current = 0
+
+FOR WHILE current < 2
+	current += 1
+	WAITFOR EVENT "test" IN obs WHEN .type == "match"
+	RETURN current`, []any{1, 2}, "WAITFOR EVENT should execute as a FOR loop body statement").Env(vm.WithParams(map[string]runtime.Value{
+			"obs": matchFirst,
+		})),
 		S(`LET obs = @obs
 
 LET evt = WAITFOR EVENT "test" IN obs TIMEOUT 1ms ON TIMEOUT RETURN NONE