feat: add support for WAITFOR VALUE WHEN clause in wait predicates (#939)

pkg/compiler/internal/wait_helpers.go

@@ -102,6 +102,14 @@ func resolveWaitPredicateMode(hasValue, hasExists, hasNot bool) waitForPredicate
 	return waitForPredicateModeBool
 }
 
+func waitPredicateWhenExpression(ctx fql.IWaitForPredicateWhenClauseContext) fql.IExpressionContext {
+	if ctx == nil {
+		return nil
+	}
+
+	return ctx.Expression()
+}
+
 func literalFromExpression(ctx fql.IExpressionContext) fql.ILiteralContext {
 	if ctx == nil {
 		return nil

pkg/compiler/internal/wait_polling.go

@@ -18,6 +18,10 @@ type waitPredicatePollState struct {
 const waitForDefaultEveryMs = 100
 
 func (c *WaitCompiler) tryCompileWaitPredicateFastPath(config waitPredicateCompileConfig) (bytecode.Operand, bool) {
+	if config.whenExpr != nil {
+		return bytecode.NoopOperand, false
+	}
+
 	switch config.mode {
 	case waitForPredicateModeBool:
 		truth, ok := literalTruthinessFromExpression(config.predExpr)
@@ -157,7 +161,16 @@ func (c *WaitCompiler) emitWaitPredicatePollIteration(
 ) bytecode.Operand {
 	valueReg := c.exprs.Compile(config.predExpr)
 	condReg := c.emitWaitPredicateCondition(config.mode, valueReg)
-	c.ctx.Program.Emitter.EmitJumpIfTrue(condReg, successLabel)
+
+	if config.whenExpr == nil {
+		c.ctx.Program.Emitter.EmitJumpIfTrue(condReg, successLabel)
+	} else {
+		retryLabel := c.ctx.Program.Emitter.NewLabel()
+		c.ctx.Program.Emitter.EmitJumpIfFalse(condReg, retryLabel)
+		whenReg := c.emitWaitPredicateWhenCondition(config, valueReg)
+		c.ctx.Program.Emitter.EmitJumpIfTrue(whenReg, successLabel)
+		c.ctx.Program.Emitter.MarkLabel(retryLabel)
+	}
 
 	elapsedReg := c.emitWaitPredicateTimeoutCheck(config.timeoutReg, state.startReg, state.unitReg, timeoutLabel)
 	sleepIntervalReg := c.prepareWaitSleepInterval(config, state.pollReg)
@@ -175,6 +188,19 @@ func (c *WaitCompiler) emitWaitPredicatePollIteration(
 	return valueReg
 }
 
+func (c *WaitCompiler) emitWaitPredicateWhenCondition(config waitPredicateCompileConfig, valueReg bytecode.Operand) bytecode.Operand {
+	if config.whenExpr == nil {
+		return bytecode.NoopOperand
+	}
+
+	c.ctx.Function.Symbols.EnterScope()
+	defer c.ctx.Function.Symbols.ExitScope()
+
+	c.ctx.Function.Symbols.AssignLocal(core.PseudoVariable, core.TypeUnknown, valueReg)
+
+	return c.exprs.CompileWithImplicitCurrent(config.whenExpr)
+}
+
 func (c *WaitCompiler) emitWaitPredicateCondition(mode waitForPredicateMode, valueReg bytecode.Operand) bytecode.Operand {
 	switch mode {
 	case waitForPredicateModeValue, waitForPredicateModeExists:

pkg/compiler/internal/wait_predicate.go

@@ -18,6 +18,7 @@ type (
 
 	waitPredicateCompileConfig struct {
 		predExpr      fql.IExpressionContext
+		whenExpr      fql.IExpressionContext
 		jitterLiteral *float64
 		mode          waitForPredicateMode
 		timeoutReg    bytecode.Operand
@@ -115,6 +116,7 @@ func (c *WaitCompiler) buildWaitPredicateConfig(
 	return waitPredicateCompileConfig{
 		mode:          resolveWaitPredicateMode(predicate.Value() != nil, predicate.Exists() != nil, predicate.Not() != nil),
 		predExpr:      predExpr,
+		whenExpr:      waitPredicateWhenExpression(ctx.WaitForPredicateWhenClause()),
 		timeoutReg:    timeoutReg,
 		everyReg:      everyReg,
 		capEveryReg:   capEveryReg,

pkg/formatter/internal/clause.go

@@ -32,6 +32,19 @@ func (f *clauseFormatter) formatEventFilterClause(ctx *fql.EventFilterClauseCont
 	}
 }
 
+func (f *clauseFormatter) formatWaitForPredicateWhenClause(ctx *fql.WaitForPredicateWhenClauseContext) {
+	if ctx == nil {
+		return
+	}
+
+	f.writeKeyword(keywordWhen)
+	f.p.space()
+
+	if expr := ctx.Expression(); expr != nil {
+		f.expression.formatExpression(expr.(*fql.ExpressionContext))
+	}
+}
+
 func (f *clauseFormatter) formatLimitClause(ctx *fql.LimitClauseContext) {
 	if ctx == nil {
 		return

pkg/formatter/internal/clause_test.go

@@ -35,3 +35,17 @@ func TestClauseFormatter_EventFilterClauseUsesWhen(t *testing.T) {
 		t.Fatalf("unexpected event filter formatting: %q", got)
 	}
 }
+
+func TestClauseFormatter_WaitForPredicateWhenClause(t *testing.T) {
+	input := "WAITFOR VALUE ready WHEN .state == \"ready\""
+	program := parseProgram(t, input)
+	when := mustFirst[*fql.WaitForPredicateWhenClauseContext](t, program)
+
+	var buf bytes.Buffer
+	e := newEngine(source.NewAnonymous(input), &buf, DefaultOptions())
+
+	e.clause.formatWaitForPredicateWhenClause(when)
+	if got := buf.String(); got != "WHEN .state == \"ready\"" {
+		t.Fatalf("unexpected waitfor predicate WHEN formatting: %q", got)
+	}
+}

pkg/formatter/internal/statement.go

@@ -572,6 +572,11 @@ func (f *statementFormatter) formatWaitForPredicateExpression(ctx *fql.WaitForPr
 		f.formatWaitForPredicate(pred.(*fql.WaitForPredicateContext))
 	}
 
+	if when := ctx.WaitForPredicateWhenClause(); when != nil {
+		f.p.space()
+		f.clause.formatWaitForPredicateWhenClause(when.(*fql.WaitForPredicateWhenClauseContext))
+	}
+
 	if timeout := ctx.TimeoutClause(); timeout != nil {
 		f.p.space()
 		f.clause.formatTimeoutClause(timeout.(*fql.TimeoutClauseContext))

pkg/formatter/internal/statement_test.go

@@ -78,6 +78,20 @@ func TestStatementFormatter_WaitForExpressionRecoveryTailCanonicalOrder(t *testi
 	}
 }
 
+func TestStatementFormatter_WaitForExpressionPredicateWhenClause(t *testing.T) {
+	input := `WAITFOR EXISTS rows WHEN LENGTH(.) >= 10 TIMEOUT 1 EVERY 10`
+	program := parseProgram(t, input+"\nRETURN 1")
+	waitExpr := mustFirst[*fql.WaitForExpressionContext](t, program)
+
+	var buf bytes.Buffer
+	e := newEngine(source.NewAnonymous(input), &buf, DefaultOptions())
+
+	e.statement.formatWaitForExpression(waitExpr)
+	if got := buf.String(); got != `WAITFOR EXISTS rows WHEN LENGTH(.) >= 10 TIMEOUT 1 EVERY 10` {
+		t.Fatalf("unexpected waitfor predicate WHEN formatting: %q", got)
+	}
+}
+
 func TestStatementFormatter_WaitForExpressionRetryTailCanonicalOrder(t *testing.T) {
 	input := `WAITFOR VALUE ready TIMEOUT 1 ON ERROR RETRY 3 DELAY 10MS OR RETURN NONE ON TIMEOUT RETURN FALSE`
 	program := parseProgram(t, input+"\nRETURN 1")

pkg/parser/antlr/FqlParser.g4

@@ -405,6 +405,10 @@ eventFilterClause
     : When {p.pushImplicitCurrent()} expression {p.popImplicitCurrent()}
     ;
 
+waitForPredicateWhenClause
+    : When {p.pushImplicitCurrent()} expression {p.popImplicitCurrent()}
+    ;
+
 limitClause
     : Limit limitClauseValue (Comma limitClauseValue)?
     ;
@@ -515,7 +519,7 @@ waitForEventExpression
     ;
 
 waitForPredicateExpression
-    : waitForPredicate (timeoutClause)? (everyClause)? (backoffClause)? (jitterClause)?
+    : waitForPredicate (waitForPredicateWhenClause)? (timeoutClause)? (everyClause)? (backoffClause)? (jitterClause)?
     ;
 
 waitForPredicate

pkg/parser/fql/fqlparser_base_listener.go

@@ -206,6 +206,14 @@ func (s *BaseFqlParserListener) EnterEventFilterClause(ctx *EventFilterClauseCon
 // ExitEventFilterClause is called when production eventFilterClause is exited.
 func (s *BaseFqlParserListener) ExitEventFilterClause(ctx *EventFilterClauseContext) {}
 
+// EnterWaitForPredicateWhenClause is called when production waitForPredicateWhenClause is entered.
+func (s *BaseFqlParserListener) EnterWaitForPredicateWhenClause(ctx *WaitForPredicateWhenClauseContext) {
+}
+
+// ExitWaitForPredicateWhenClause is called when production waitForPredicateWhenClause is exited.
+func (s *BaseFqlParserListener) ExitWaitForPredicateWhenClause(ctx *WaitForPredicateWhenClauseContext) {
+}
+
 // EnterLimitClause is called when production limitClause is entered.
 func (s *BaseFqlParserListener) EnterLimitClause(ctx *LimitClauseContext) {}
 

pkg/parser/fql/fqlparser_base_visitor.go

@@ -131,6 +131,10 @@ func (v *BaseFqlParserVisitor) VisitEventFilterClause(ctx *EventFilterClauseCont
 	return v.VisitChildren(ctx)
 }
 
+func (v *BaseFqlParserVisitor) VisitWaitForPredicateWhenClause(ctx *WaitForPredicateWhenClauseContext) interface{} {
+	return v.VisitChildren(ctx)
+}
+
 func (v *BaseFqlParserVisitor) VisitLimitClause(ctx *LimitClauseContext) interface{} {
 	return v.VisitChildren(ctx)
 }

pkg/parser/fql/fqlparser_listener.go

@@ -100,6 +100,9 @@ type FqlParserListener interface {
 	// EnterEventFilterClause is called when entering the eventFilterClause production.
 	EnterEventFilterClause(c *EventFilterClauseContext)
 
+	// EnterWaitForPredicateWhenClause is called when entering the waitForPredicateWhenClause production.
+	EnterWaitForPredicateWhenClause(c *WaitForPredicateWhenClauseContext)
+
 	// EnterLimitClause is called when entering the limitClause production.
 	EnterLimitClause(c *LimitClauseContext)
 
@@ -583,6 +586,9 @@ type FqlParserListener interface {
 	// ExitEventFilterClause is called when exiting the eventFilterClause production.
 	ExitEventFilterClause(c *EventFilterClauseContext)
 
+	// ExitWaitForPredicateWhenClause is called when exiting the waitForPredicateWhenClause production.
+	ExitWaitForPredicateWhenClause(c *WaitForPredicateWhenClauseContext)
+
 	// ExitLimitClause is called when exiting the limitClause production.
 	ExitLimitClause(c *LimitClauseContext)
 

pkg/parser/fql/fqlparser_visitor.go

@@ -100,6 +100,9 @@ type FqlParserVisitor interface {
 	// Visit a parse tree produced by FqlParser#eventFilterClause.
 	VisitEventFilterClause(ctx *EventFilterClauseContext) interface{}
 
+	// Visit a parse tree produced by FqlParser#waitForPredicateWhenClause.
+	VisitWaitForPredicateWhenClause(ctx *WaitForPredicateWhenClauseContext) interface{}
+
 	// Visit a parse tree produced by FqlParser#limitClause.
 	VisitLimitClause(ctx *LimitClauseContext) interface{}
 

test/integration/compiler/compiler_waitfor_test.go

@@ -3,6 +3,7 @@ package compiler_test
 import (
 	"testing"
 
+	"github.com/MontFerret/ferret/v2/pkg/bytecode"
 	parserd "github.com/MontFerret/ferret/v2/pkg/parser/diagnostics"
 	"github.com/MontFerret/ferret/v2/test/spec"
 	. "github.com/MontFerret/ferret/v2/test/spec/compile"
@@ -64,3 +65,31 @@ func TestWaitforCompilationErrors(t *testing.T) {
 		}, "Out-of-range WAITFOR EVERY float constant should fail compilation"),
 	})
 }
+
+func TestWaitforPredicateWhenCompiles(t *testing.T) {
+	RunSpecs(t, []spec.Spec{
+		ProgramCheck(`
+			RETURN WAITFOR VALUE { state: "ready" }
+				WHEN .state == "ready"
+				TIMEOUT 5ms
+				EVERY 1ms
+				ON TIMEOUT RETURN NONE
+		`, noCompilerError, "WAITFOR VALUE should compile with WHEN and wait tails"),
+		ProgramCheck(`
+			RETURN WAITFOR EXISTS [1, 2, 3]
+				WHEN LENGTH(.) >= 3
+				TIMEOUT 5ms
+				EVERY 1ms
+		`, noCompilerError, "WAITFOR EXISTS should compile with WHEN and wait tails"),
+		ProgramCheck(`
+			RETURN WAITFOR NOT EXISTS []
+				WHEN LENGTH(.) == 0
+				TIMEOUT 5ms
+				EVERY 1ms
+		`, noCompilerError, "WAITFOR NOT EXISTS should compile with WHEN and wait tails"),
+	})
+}
+
+func noCompilerError(*bytecode.Program) error {
+	return nil
+}

test/integration/vm/vm_waitfor_predicate_test.go

@@ -1,8 +1,13 @@
 package vm_test
 
 import (
+	"context"
+	"fmt"
 	"testing"
 
+	"github.com/MontFerret/ferret/v2/pkg/compiler"
+	"github.com/MontFerret/ferret/v2/pkg/runtime"
+	"github.com/MontFerret/ferret/v2/pkg/vm"
 	"github.com/MontFerret/ferret/v2/test/spec"
 	. "github.com/MontFerret/ferret/v2/test/spec/exec"
 )
@@ -89,5 +94,97 @@ func TestWaitforPredicate(t *testing.T) {
 			LET token = WAITFOR VALUE NONE TIMEOUT 20ms EVERY 5ms ON ERROR FAIL
 			RETURN token
 		`, "Explicit FAIL should preserve timeout result semantics"),
+		Object(`
+			LET value = WAITFOR VALUE { ok: true, kind: "candidate" } WHEN .ok
+			RETURN value
+		`, map[string]any{"ok": true, "kind": "candidate"}, "WAITFOR VALUE WHEN should return the candidate value"),
+		S(`
+			LET ok = WAITFOR EXISTS [1, 2, 3] WHEN LENGTH(.) >= 3
+			RETURN ok
+		`, true, "WAITFOR EXISTS WHEN should bind the full candidate array"),
+		S(`
+			LET ok = WAITFOR NOT EXISTS [] WHEN LENGTH(.) == 0
+			RETURN ok
+		`, true, "WAITFOR NOT EXISTS WHEN should bind the not-existing candidate value"),
+		Nil(`
+			LET token = WAITFOR VALUE "ok" WHEN false TIMEOUT 20ms EVERY 5ms ON TIMEOUT RETURN NONE
+			RETURN token
+		`, "WAITFOR VALUE WHEN should honor ON TIMEOUT when the predicate never passes"),
+		S(`
+			LET ok = WAITFOR EXISTS [1] WHEN false TIMEOUT 20ms EVERY 5ms ON TIMEOUT RETURN false
+			RETURN ok
+		`, false, "WAITFOR EXISTS WHEN should honor ON TIMEOUT when the predicate never passes"),
 	})
 }
+
+func TestWaitforPredicateWhenRetriesUntilTrue(t *testing.T) {
+	for _, level := range []compiler.OptimizationLevel{compiler.O0, compiler.O1} {
+		callCount := 0
+
+		RunSpecsWith(
+			t,
+			fmt.Sprintf("VM/O%d", level),
+			compiler.New(compiler.WithOptimizationLevel(level)),
+			[]spec.Spec{
+				S(`
+					LET token = WAITFOR VALUE CANDIDATE() WHEN .state == "ready" TIMEOUT 100ms EVERY 0
+					RETURN token.value
+				`, "ok", "WAITFOR VALUE WHEN should retry until the predicate passes"),
+			},
+			vm.WithFunction("CANDIDATE", func(ctx context.Context, args ...runtime.Value) (runtime.Value, error) {
+				callCount++
+
+				state := "pending"
+				if callCount >= 3 {
+					state = "ready"
+				}
+
+				return runtime.NewObjectWith(map[string]runtime.Value{
+					"state": runtime.NewString(state),
+					"value": runtime.NewString("ok"),
+				}), nil
+			}),
+		)
+
+		if got, want := callCount, 3; got != want {
+			t.Fatalf("unexpected WAITFOR VALUE WHEN candidate call count for O%d: got %d, want %d", level, got, want)
+		}
+	}
+}
+
+func TestWaitforPredicateWhenSkipsPredicateUntilBasePasses(t *testing.T) {
+	for _, level := range []compiler.OptimizationLevel{compiler.O0, compiler.O1} {
+		predicateCalls := 0
+
+		RunSpecsWith(
+			t,
+			fmt.Sprintf("VM/O%d", level),
+			compiler.New(compiler.WithOptimizationLevel(level)),
+			[]spec.Spec{
+				S(`
+					LET ok = WAITFOR EXISTS NONE WHEN PREDICATE(.) TIMEOUT 20ms EVERY 1ms ON TIMEOUT RETURN false
+					RETURN ok
+				`, false, "WAITFOR EXISTS WHEN should not evaluate the predicate before existence passes"),
+			},
+			vm.WithFunction("PREDICATE", func(ctx context.Context, args ...runtime.Value) (runtime.Value, error) {
+				predicateCalls++
+
+				return runtime.True, nil
+			}),
+		)
+
+		if got := predicateCalls; got != 0 {
+			t.Fatalf("WAITFOR EXISTS WHEN evaluated predicate before existence passed for O%d: got %d calls", level, got)
+		}
+	}
+}
+
+func TestWaitforPredicateWhenUsesOperationErrorPolicy(t *testing.T) {
+	RunSpecs(t, []spec.Spec{
+		S(`
+			RETURN WAITFOR VALUE "ok" WHEN FAIL_PREDICATE(.) TIMEOUT 20ms EVERY 0 ON ERROR RETURN "error"
+		`, "error", "WAITFOR VALUE WHEN predicate errors should use the wait error policy"),
+	}, vm.WithFunction("FAIL_PREDICATE", func(ctx context.Context, args ...runtime.Value) (runtime.Value, error) {
+		return runtime.None, fmt.Errorf("predicate failed")
+	}))
+}