Additional aggregation tests showcasing multiple aggregation rules with nested match behavior

Author: TristonianJones

policy/BUILD.bazel

@@ -72,8 +72,9 @@ go_test(
         "//test:go_default_library",
         "//common/debug:go_default_library",
         "//common/types:go_default_library",
-        "//interpreter:go_default_library",
         "//common/types/ref:go_default_library",
+        "//common/types/traits:go_default_library",
+        "//interpreter:go_default_library",
         "//test/proto3pb:go_default_library",
         "@in_yaml_go_yaml_v3//:go_default_library",
         "@com_github_google_go_cmp//cmp:go_default_library",

policy/compiler_test.go

@@ -46,52 +46,6 @@ func TestCompile(t *testing.T) {
 	}
 }
 
-func TestRuleComposerError(t *testing.T) {
-	env, err := cel.NewEnv()
-	if err != nil {
-		t.Fatalf("NewEnv() failed: %v", err)
-	}
-	_, err = NewRuleComposer(env, ExpressionUnnestHeight(-1))
-	if err == nil || !strings.Contains(err.Error(), "invalid unnest") {
-		t.Errorf("NewRuleComposer() got %v, wanted 'invalid unnest'", err)
-	}
-}
-
-func TestRuleComposerUnnest(t *testing.T) {
-	for _, tst := range composerUnnestTests {
-		tc := tst
-		t.Run(tc.name, func(t *testing.T) {
-			r := newRunner(tc.name, tc.expr, []ParserOption{})
-			env, rule, iss := r.compileRule(t)
-			if iss.Err() != nil {
-				t.Fatalf("CompileRule() failed: %v", iss.Err())
-			}
-			rc, err := NewRuleComposer(env, tc.composerOpts...)
-			if err != nil {
-				t.Fatalf("NewRuleComposer() failed: %v", err)
-			}
-			ast, iss := rc.Compose(rule)
-			if iss.Err() != nil {
-				t.Fatalf("Compose(rule) failed: %v", iss.Err())
-			}
-			policy := parsePolicy(t, tc.name, []ParserOption{})
-			verifySourceInfoCoverage(t, policy, ast)
-			unparsed, err := cel.AstToString(ast)
-			if err != nil {
-				t.Fatalf("cel.AstToString() failed: %v", err)
-			}
-			if normalize(unparsed) != normalize(tc.composed) {
-				t.Errorf("cel.AstToString() got %s, wanted %s", unparsed, tc.composed)
-			}
-			if !ast.OutputType().IsEquivalentType(tc.outputType) {
-				t.Errorf("ast.OutputType() got %v, wanted %v", ast.OutputType(), tc.outputType)
-			}
-			r.setup(t, env, ast)
-			r.run(t)
-		})
-	}
-}
-
 func TestCompileError(t *testing.T) {
 	for _, tst := range policyErrorTests {
 		policy := parsePolicy(t, tst.name, []ParserOption{})
@@ -290,6 +244,36 @@ func BenchmarkCompile(b *testing.B) {
 	}
 }
 
+func parsePolicySource(t testing.TB, name string, policySource string, parseOpts ...ParserOption) *Policy {
+	t.Helper()
+	p := StringSource(policySource, name)
+	parser, err := NewParser(parseOpts...)
+	if err != nil {
+		t.Fatalf("NewParser() failed: %v", err)
+	}
+	policy, iss := parser.Parse(p)
+	if iss.Err() != nil {
+		t.Fatalf("parser.Parse() failed: %v", iss.Err())
+	}
+	return policy
+}
+
+func parseAndCompilePolicy(t testing.TB, name string, policySource string, envOpts []cel.EnvOption, compilerOpts []CompilerOption) (*cel.Env, *cel.Ast, *cel.Issues) {
+	t.Helper()
+	policy := parsePolicySource(t, name, policySource)
+	envOpts = append([]cel.EnvOption{
+		cel.OptionalTypes(),
+		cel.EnableMacroCallTracking(),
+		ext.Bindings(),
+	}, envOpts...)
+	env, err := cel.NewEnv(envOpts...)
+	if err != nil {
+		t.Fatalf("cel.NewEnv() failed: %v", err)
+	}
+	ast, iss := Compile(env, policy, compilerOpts...)
+	return env, ast, iss
+}
+
 func newRunner(name, expr string, parseOpts []ParserOption, opts ...cel.EnvOption) *runner {
 	return &runner{
 		name:      name,
@@ -772,27 +756,7 @@ rule:
 	for _, tst := range tests {
 		tc := tst
 		t.Run(tc.name, func(t *testing.T) {
-			p := StringSource(tc.policy, "<input>")
-			parser, err := NewParser()
-			if err != nil {
-				t.Fatalf("NewParser() failed: %v", err)
-			}
-			policy, iss := parser.Parse(p)
-			if iss.Err() != nil {
-				t.Fatalf("parser.Parse() failed: %v", iss.Err())
-			}
-
-			envOpts := append([]cel.EnvOption{
-				cel.OptionalTypes(),
-				cel.EnableMacroCallTracking(),
-				ext.Bindings(),
-			}, tc.envOpts...)
-			env, err := cel.NewEnv(envOpts...)
-			if err != nil {
-				t.Fatalf("cel.NewEnv() failed: %v", err)
-			}
-
-			ast, iss := Compile(env, policy)
+			env, ast, iss := parseAndCompilePolicy(t, tc.name, tc.policy, tc.envOpts, nil)
 			if tc.wantErr != "" {
 				if iss.Err() == nil {
 					t.Fatalf("Compile() succeeded, wanted error %q", tc.wantErr)
@@ -844,15 +808,7 @@ rule:
   aggregate:
     - condition: 'true'
       emit: "'foo'"`
-	p := StringSource(policySource, "<input>")
-	parser, err := NewParser()
-	if err != nil {
-		t.Fatalf("NewParser() failed: %v", err)
-	}
-	policy, iss := parser.Parse(p)
-	if iss.Err() != nil {
-		t.Fatalf("parser.Parse() failed: %v", iss.Err())
-	}
+	policy := parsePolicySource(t, "aggregate_semantic", policySource)
 	env, err := cel.NewEnv()
 	if err != nil {
 		t.Fatalf("cel.NewEnv() failed: %v", err)
@@ -872,20 +828,7 @@ rule:
   aggregate:
     - condition: 'false'
       emit: "'foo'"`
-	p := StringSource(policySource, "<input>")
-	parser, err := NewParser()
-	if err != nil {
-		t.Fatalf("NewParser() failed: %v", err)
-	}
-	policy, iss := parser.Parse(p)
-	if iss.Err() != nil {
-		t.Fatalf("parser.Parse() failed: %v", iss.Err())
-	}
-	env, err := cel.NewEnv()
-	if err != nil {
-		t.Fatalf("cel.NewEnv() failed: %v", err)
-	}
-	_, iss = Compile(env, policy)
+	_, _, iss := parseAndCompilePolicy(t, "condition_always_false", policySource, nil, nil)
 	if iss.Err() == nil {
 		t.Fatalf("Compile() succeeded, wanted error")
 	}

policy/composer.go

@@ -195,7 +195,10 @@ func (opt *ruleComposerImpl) optimizeRule(ctx *cel.OptimizerContext, r *Compiled
 	matches := r.Matches()
 	matchCount := len(matches)
 	var output compositionStep = nil
-	if !isAggregate && r.HasOptionalOutput() {
+	// If the rule is non-aggregate and has an optional output, the last result in the ternary should return
+	// `optional.none`. This output is implicit and created here to reflect the desired
+	// last possible output of this type of rule.
+	if !returnList && r.HasOptionalOutput() {
 		output = newOptionalCompositionStep(ctx, ctx.NewLiteral(types.True), ctx.NewCall("optional.none"))
 	}
 	// Build the rule subgraph.
@@ -205,6 +208,10 @@ func (opt *ruleComposerImpl) optimizeRule(ctx *cel.OptimizerContext, r *Compiled
 
 		var currentStep compositionStep
 		if m.Output() != nil {
+			// If the output is non-nil, then it is considered a non-optional output since
+			// it is explicitly stated. If the rule itself is optional, then the base case value
+			// of output being optional.none() will convert the non-optional value to an optional
+			// one.
 			out := ctx.CopyASTAndMetadata(m.Output().Expr().NativeRep())
 			if returnList {
 				outList := ctx.NewList([]ast.Expr{out}, []int32{})
@@ -213,9 +220,19 @@ func (opt *ruleComposerImpl) optimizeRule(ctx *cel.OptimizerContext, r *Compiled
 				currentStep = newNonOptionalCompositionStep(ctx, cond, out)
 			}
 		} else if m.NestedRule() != nil {
+			// If the match has a nested rule, then compute the rule and whether it has
+			// an optional return value.
+			//
+			// Semantics for nesting:
+			// - With optional values (nestedHasOptional = true): The step is treated as optional.
+			//   If the nested rule yields optional.none, composition allows fall-through to
+			//   subsequent match cases.
+			// - Without optional values (nestedHasOptional = false): The step is treated as non-optional.
+			//   A matching result produces a concrete value that short-circuits further match evaluation,
+			//   though it may be wrapped into optional.of(...) if the outer rule produces optional output.
 			child := m.NestedRule()
 			nestedRule := opt.optimizeRule(ctx, child, returnList)
-			nestedHasOptional := child.HasOptionalOutput()
+			nestedHasOptional := !returnList && child.HasOptionalOutput()
 			if nestedHasOptional {
 				currentStep = newOptionalCompositionStep(ctx, cond, nestedRule)
 			} else {
@@ -522,12 +539,17 @@ func (s nonOptionalCompositionStep) combine(step compositionStep) compositionSte
 	if !s.isConditional() {
 		return s
 	}
+	stepExpr := step.expr()
+	if step.isConditional() {
+		emptyList := ctx.NewList([]ast.Expr{}, []int32{})
+		stepExpr = ctx.NewCall(operators.Conditional, step.condition(), step.expr(), emptyList)
+	}
 	return newNonOptionalCompositionStep(ctx,
 		trueCondition,
 		ctx.NewCall(operators.Conditional,
 			s.condition(),
 			s.expr(),
-			step.expr()))
+			stepExpr))
 }
 
 // newOptionalCompositionStep returns an output step with an optional policy output.
@@ -615,9 +637,7 @@ func isOptionalNone(e ast.Expr) bool {
 
 func removeIneligibleSubExprs(e ast.NavigableExpr, unnestMap map[int64]bool) {
 	for _, id := range comprehensionSubExprIDs(e) {
-		if _, found := unnestMap[id]; found {
-			delete(unnestMap, id)
-		}
+		delete(unnestMap, id)
 	}
 }