enhance WAITFOR statement handling

pkg/compiler/internal/wait_event.go

@@ -365,32 +365,20 @@ func (c *WaitCompiler) emitWaitEventCleanupIfReady(state waitEventCompileState,
 
 // CompileWaitForEventName processes the event name expression in a WAITFOR statement.
 func (c *WaitCompiler) CompileWaitForEventName(ctx fql.IWaitForEventNameContext) bytecode.Operand {
-	sl := ctx.StringLiteral()
-	v := ctx.Variable()
-	p := ctx.Param()
-	me := ctx.MemberExpression()
-	fce := ctx.FunctionCall()
-
-	return compileFirstOperand(
-		newOperandBranch(sl != nil, func() bytecode.Operand { return c.literals.CompileStringLiteral(sl) }),
-		newOperandBranch(v != nil, func() bytecode.Operand { return c.exprs.CompileVariable(v) }),
-		newOperandBranch(p != nil, func() bytecode.Operand { return c.exprs.CompileParam(p) }),
-		newOperandBranch(me != nil, func() bytecode.Operand { return c.exprs.CompileMemberExpression(me) }),
-		newOperandBranch(fce != nil, func() bytecode.Operand { return c.exprs.CompileFunctionCall(fce, false) }),
-	)
+	if ctx == nil {
+		return bytecode.NoopOperand
+	}
+
+	return c.exprs.Compile(ctx.Expression())
 }
 
 // CompileWaitForEventSource processes the event source expression in a WAITFOR statement.
 func (c *WaitCompiler) CompileWaitForEventSource(ctx fql.IWaitForEventSourceContext) bytecode.Operand {
-	v := ctx.Variable()
-	me := ctx.MemberExpression()
-	fce := ctx.FunctionCallExpression()
-
-	return compileFirstOperand(
-		newOperandBranch(v != nil, func() bytecode.Operand { return c.exprs.CompileVariable(v) }),
-		newOperandBranch(me != nil, func() bytecode.Operand { return c.exprs.CompileMemberExpression(me) }),
-		newOperandBranch(fce != nil, func() bytecode.Operand { return c.exprs.CompileFunctionCallExpression(fce) }),
-	)
+	if ctx == nil {
+		return bytecode.NoopOperand
+	}
+
+	return c.exprs.Compile(ctx.Expression())
 }
 
 // CompileOptionsClause processes the options clause in a WAITFOR statement.

pkg/formatter/formatter_parentheses_test.go

@@ -2,13 +2,21 @@ package formatter
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"strings"
 	"testing"
 
+	"github.com/antlr4-go/antlr/v4"
+
 	"github.com/MontFerret/ferret/v2/pkg/compiler"
+	"github.com/MontFerret/ferret/v2/pkg/parser"
+	"github.com/MontFerret/ferret/v2/pkg/parser/fql"
+	"github.com/MontFerret/ferret/v2/pkg/runtime"
 	"github.com/MontFerret/ferret/v2/pkg/source"
+	"github.com/MontFerret/ferret/v2/pkg/vm"
 	"github.com/MontFerret/ferret/v2/test/spec"
+	"github.com/MontFerret/ferret/v2/test/spec/mock"
 )
 
 func TestFormatterSimplifiesOnlyRedundantParentheses(t *testing.T) {
@@ -71,6 +79,257 @@ func TestFormatterSimplifiesOnlyRedundantParentheses(t *testing.T) {
 	}
 }
 
+func TestFormatterWaitForEventOperands(t *testing.T) {
+	tests := []struct {
+		name       string
+		input      string
+		want       string
+		wantName   string
+		wantSource string
+	}{
+		{
+			name:       "composed operands",
+			input:      `RETURN WAITFOR EVENT (@kind ?? "message") IN (@source ?? fallback)`,
+			want:       `return waitfor event @kind ?? "message" in @source ?? fallback`,
+			wantName:   `@kind??"message"`,
+			wantSource: `@source??fallback`,
+		},
+		{
+			name:       "membership name removes redundant grouping",
+			input:      `RETURN WAITFOR EVENT (@kind IN @names) IN @source`,
+			want:       `return waitfor event @kind in @names in @source`,
+			wantName:   `@kindin@names`,
+			wantSource: `@source`,
+		},
+		{
+			name:       "membership source retains delimiter boundary",
+			input:      `RETURN WAITFOR EVENT "message" IN (@candidate IN @sources)`,
+			want:       `return waitfor event "message" in (@candidate in @sources)`,
+			wantName:   `"message"`,
+			wantSource: `(@candidatein@sources)`,
+		},
+		{
+			name:       "negated membership source retains delimiter boundary",
+			input:      `RETURN WAITFOR EVENT "message" IN (@candidate NOT IN @sources)`,
+			want:       `return waitfor event "message" in (@candidate not in @sources)`,
+			wantName:   `"message"`,
+			wantSource: `(@candidatenotin@sources)`,
+		},
+		{
+			name:       "function argument membership is bounded",
+			input:      `RETURN WAITFOR EVENT "message" IN (SOURCE(@candidate IN @sources))`,
+			want:       `return waitfor event "message" in SOURCE(@candidate in @sources)`,
+			wantName:   `"message"`,
+			wantSource: `SOURCE(@candidatein@sources)`,
+		},
+		{
+			name:       "array entry membership is bounded",
+			input:      `RETURN WAITFOR EVENT "message" IN ([@candidate IN @sources])`,
+			want:       `return waitfor event "message" in [@candidate in @sources]`,
+			wantName:   `"message"`,
+			wantSource: `[@candidatein@sources]`,
+		},
+		{
+			name:       "object property membership is bounded",
+			input:      `RETURN WAITFOR EVENT "message" IN ({ candidate: @candidate IN @sources })`,
+			want:       `return waitfor event "message" in { candidate: @candidate in @sources }`,
+			wantName:   `"message"`,
+			wantSource: `{candidate:@candidatein@sources}`,
+		},
+		{
+			name:       "ternary true branch membership is bounded",
+			input:      `RETURN WAITFOR EVENT "message" IN (TRUE ? @candidate IN @sources : FALSE)`,
+			want:       `return waitfor event "message" in true ? @candidate in @sources : false`,
+			wantName:   `"message"`,
+			wantSource: `true?@candidatein@sources:false`,
+		},
+		{
+			name:       "precedence grouping bounds nested membership",
+			input:      `RETURN WAITFOR EVENT "message" IN ((@candidate IN @sources) + 1)`,
+			want:       `return waitfor event "message" in (@candidate in @sources) + 1`,
+			wantName:   `"message"`,
+			wantSource: `(@candidatein@sources)+1`,
+		},
+		{
+			name:       "removable inner grouping leaves membership exposed",
+			input:      `RETURN WAITFOR EVENT "message" IN (NOT (@candidate IN @sources))`,
+			want:       `return waitfor event "message" in (not @candidate in @sources)`,
+			wantName:   `"message"`,
+			wantSource: `(not@candidatein@sources)`,
+		},
+		{
+			name:       "ternary condition membership remains exposed",
+			input:      `RETURN WAITFOR EVENT "message" IN (@candidate IN @sources ? TRUE : FALSE)`,
+			want:       `return waitfor event "message" in (@candidate in @sources ? true : false)`,
+			wantName:   `"message"`,
+			wantSource: `(@candidatein@sources?true:false)`,
+		},
+		{
+			name:       "ternary false branch membership remains exposed",
+			input:      `RETURN WAITFOR EVENT "message" IN (TRUE ? FALSE : @candidate IN @sources)`,
+			want:       `return waitfor event "message" in (true ? false : @candidate in @sources)`,
+			wantName:   `"message"`,
+			wantSource: `(true?false:@candidatein@sources)`,
+		},
+		{
+			name:       "membership operands use distinct grouping",
+			input:      `RETURN WAITFOR EVENT (@kind IN @names) IN (@candidate IN @sources)`,
+			want:       `return waitfor event @kind in @names in (@candidate in @sources)`,
+			wantName:   `@kindin@names`,
+			wantSource: `(@candidatein@sources)`,
+		},
+		{
+			name:       "array membership source retains delimiter boundary",
+			input:      `RETURN WAITFOR EVENT "message" IN (@candidate ANY IN @sources)`,
+			want:       `return waitfor event "message" in (@candidate any in @sources)`,
+			wantName:   `"message"`,
+			wantSource: `(@candidateanyin@sources)`,
+		},
+		{
+			name:       "query source",
+			input:      `RETURN WAITFOR EVENT GET_NAME(@id) IN QUERY ONE ".source" IN @registry`,
+			want:       `return waitfor event GET_NAME(@id) in query one ".source" in @registry`,
+			wantName:   `GET_NAME(@id)`,
+			wantSource: `queryone".source"in@registry`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := formatParenthesesStable(t, test.input)
+			if got != test.want {
+				t.Fatalf("formatted output:\n%s\nwant:\n%s", got, test.want)
+			}
+
+			assertWaitForEventOperandBoundaries(t, got, test.wantName, test.wantSource)
+		})
+	}
+}
+
+func assertWaitForEventOperandBoundaries(t *testing.T, input, wantName, wantSource string) {
+	t.Helper()
+
+	p := parser.New(input)
+	program := p.Program()
+	if !p.AtEOF() {
+		t.Fatalf("formatted output did not parse completely:\n%s", input)
+	}
+
+	event := findFirstWaitForEvent(program)
+	if event == nil {
+		t.Fatalf("formatted output has no WAITFOR EVENT expression:\n%s", input)
+	}
+
+	name := event.WaitForEventName().(*fql.WaitForEventNameContext).Expression()
+	if got := name.GetText(); got != wantName {
+		t.Fatalf("reparsed event-name expression = %q, want %q", got, wantName)
+	}
+
+	source := event.WaitForEventSource().(*fql.WaitForEventSourceContext).Expression()
+	if got := source.GetText(); got != wantSource {
+		t.Fatalf("reparsed source expression = %q, want %q", got, wantSource)
+	}
+}
+
+func findFirstWaitForEvent(tree antlr.Tree) *fql.WaitForEventExpressionContext {
+	if tree == nil {
+		return nil
+	}
+
+	if event, ok := tree.(*fql.WaitForEventExpressionContext); ok {
+		return event
+	}
+
+	for i := 0; i < tree.GetChildCount(); i++ {
+		if event := findFirstWaitForEvent(tree.GetChild(i)); event != nil {
+			return event
+		}
+	}
+
+	return nil
+}
+
+func TestFormatterWaitForEventOperandsPreserveExecution(t *testing.T) {
+	tests := []struct {
+		newEnvironment func() []vm.EnvironmentOption
+		name           string
+		input          string
+	}{
+		{
+			name: "composed operands",
+			input: `RETURN WAITFOR EVENT (@eventName ?? "message") IN (@source ?? @fallback)
+WHEN .type == "match"`,
+			newEnvironment: func() []vm.EnvironmentOption {
+				source := mock.NewObservable([]runtime.Value{
+					mock.NewTestEventType("ignored"),
+					mock.NewTestEventType("match"),
+				})
+
+				return []vm.EnvironmentOption{
+					vm.WithParams(map[string]runtime.Value{
+						"eventName": runtime.None,
+						"fallback":  runtime.None,
+						"source":    source,
+					}),
+				}
+			},
+		},
+		{
+			name: "bounded membership source",
+			input: `RETURN WAITFOR EVENT "message" IN (SOURCE(@candidate IN @sources))
+WHEN .type == "match"`,
+			newEnvironment: func() []vm.EnvironmentOption {
+				source := mock.NewObservable([]runtime.Value{
+					mock.NewTestEventType("ignored"),
+					mock.NewTestEventType("match"),
+				})
+
+				return []vm.EnvironmentOption{
+					vm.WithParams(map[string]runtime.Value{
+						"candidate": runtime.NewString("candidate"),
+						"sources":   runtime.NewArrayWith(runtime.NewString("candidate")),
+					}),
+					vm.WithFunction("SOURCE", func(context.Context, ...runtime.Value) (runtime.Value, error) {
+						return source, nil
+					}),
+				}
+			},
+		},
+	}
+
+	for _, test := range tests {
+		formatted := formatParenthesesStable(t, test.input)
+
+		for _, level := range []compiler.OptimizationLevel{compiler.O0, compiler.O1} {
+			t.Run(test.name+"/"+optimizationNameForFormatter(level), func(t *testing.T) {
+				originalProgram, err := spec.Compile(test.input, level)
+				if err != nil {
+					t.Fatalf("compile original: %v", err)
+				}
+
+				formattedProgram, err := spec.Compile(formatted, level)
+				if err != nil {
+					t.Fatalf("compile formatted: %v\n%s", err, formatted)
+				}
+
+				originalResult, err := spec.Run(originalProgram, test.newEnvironment()...)
+				if err != nil {
+					t.Fatalf("run original: %v", err)
+				}
+
+				formattedResult, err := spec.Run(formattedProgram, test.newEnvironment()...)
+				if err != nil {
+					t.Fatalf("run formatted: %v", err)
+				}
+
+				if !bytes.Equal(originalResult, formattedResult) {
+					t.Fatalf("execution changed: original %s, formatted %s\n%s", originalResult, formattedResult, formatted)
+				}
+			})
+		}
+	}
+}
+
 func TestFormatterSimplifiesGroupedForOnlyAtStatementBoundaries(t *testing.T) {
 	tests := []struct {
 		name     string

pkg/formatter/internal/expression.go

@@ -16,6 +16,18 @@ func (f *expressionFormatter) formatExpression(ctx *fql.ExpressionContext) {
 	f.formatExpressionOperand(ctx, expressionOperation{})
 }
 
+// formatWaitForEventSource retains grouping around membership expressions so
+// their IN operator cannot become the WAITFOR EVENT source delimiter.
+func (f *expressionFormatter) formatWaitForEventSource(ctx *fql.ExpressionContext) {
+	if waitForEventSourceNeedsParentheses(ctx) {
+		f.formatExpressionOperand(ctx, expressionOperation{precedence: precedencePrimary})
+
+		return
+	}
+
+	f.formatExpression(ctx)
+}
+
 func (f *expressionFormatter) formatExpressionOperand(ctx *fql.ExpressionContext, outer expressionOperation) {
 	if ctx == nil {
 		return