diff --git a/v1/compile/compile.go b/v1/compile/compile.go index e51c1f091f..7779f2435a 100644 --- a/v1/compile/compile.go +++ b/v1/compile/compile.go @@ -781,6 +781,27 @@ func findAnnotationsForTerm(term *ast.Term, annotationRefs []*ast.AnnotationsRef return result } +// pruneAnnotationsAndComments filters out annotations and their associated comments based on a predicate. +// It returns the kept annotations and kept comments. +func pruneAnnotationsAndComments( + module *ast.Module, + shouldDiscard func(*ast.Annotations) bool, +) ([]*ast.Annotations, []*ast.Comment) { + keepAnnotations := slices.DeleteFunc(slices.Clone(module.Annotations), shouldDiscard) + + var keepComments []*ast.Comment + for _, comment := range module.Comments { + if slices.ContainsFunc(keepAnnotations, func(a *ast.Annotations) bool { + return comment.Location.Row >= a.Location.Row && + comment.Location.Row <= a.EndLoc().Row + }) { + keepComments = append(keepComments, comment) + } + } + + return keepAnnotations, keepComments +} + // pruneBundleEntrypoints will modify modules in the provided bundle to remove // rules matching the entrypoints along with injecting import statements to // preserve their ability to compile. @@ -816,35 +837,10 @@ func pruneBundleEntrypoints(b *bundle.Bundle, entrypointrefs []*ast.Term) error } } - // Drop any Annotations for rules matching the entrypoint path - var annotations []*ast.Annotations - var prunedAnnotations []*ast.Annotations - for _, annotation := range mf.Parsed.Annotations { - p := annotation.GetTargetPath() - // We prune annotations of dropped rules, but not packages, as the Rego file is always retained - if p.Equal(entrypoint.Value) && !mf.Parsed.Package.Path.Equal(entrypoint.Value) { - prunedAnnotations = append(prunedAnnotations, annotation) - } else { - annotations = append(annotations, annotation) - } - } - - // Drop comments associated with pruned annotations - var comments []*ast.Comment - for _, comment := range mf.Parsed.Comments { - pruned := false - for _, annotation := range prunedAnnotations { - if comment.Location.Row >= annotation.Location.Row && - comment.Location.Row <= annotation.EndLoc().Row { - pruned = true - break - } - } - - if !pruned { - comments = append(comments, comment) - } - } + // Prune annotations and comments for entrypoint rules + annotations, comments := pruneAnnotationsAndComments(mf.Parsed, func(annotation *ast.Annotations) bool { + return annotation.GetTargetPath().Equal(entrypoint.Value) + }) // If any rules or annotations were dropped update the module accordingly if len(rules) != len(mf.Parsed.Rules) || len(comments) != len(mf.Parsed.Comments) { @@ -1215,7 +1211,15 @@ func (*optimizer) merge(a, b []bundle.ModuleFile) []bundle.ModuleFile { } if len(keep) > 0 { + keepAnnotations, keepComments := pruneAnnotationsAndComments(a[i].Parsed, func(annotation *ast.Annotations) bool { + return discarded.Contains(ast.NewTerm(annotation.GetTargetPath())) + }) + a[i].Parsed.Rules = keep + a[i].Parsed.Annotations = keepAnnotations + a[i].Parsed.Comments = keepComments + // Remove the original raw source, we're editing the AST + // directly, so it won't be in sync anymore. a[i].Raw = nil b = append(b, a[i]) } diff --git a/v1/compile/compile_test.go b/v1/compile/compile_test.go index 6b605742c0..5dc9984179 100644 --- a/v1/compile/compile_test.go +++ b/v1/compile/compile_test.go @@ -3531,6 +3531,121 @@ func TestOptimizerOutput(t *testing.T) { } } +func TestCompilerOptimizationMultipleAnnotationEntrypoints(t *testing.T) { + files := map[string]string{ + "test.rego": `package pkg + +# best allow rule ever +# METADATA +# entrypoint: true +allow if { + is_admin + has_permission(input.action) +} + +# nice deny rule +# METADATA +# entrypoint: true +deny if { + input.user == "guest" + input.action == "delete" +} + +# METADATA +# description: this rule should keep its MD +other if input.other + +is_admin if input.user == "admin" +valid_actions := ["read", "write", "delete"] +has_permission(action) if action in valid_actions +`, + } + + tests := []struct { + name string + optimizationLvl int + expectOptimized bool + }{ + { + name: "optimization level 0", + optimizationLvl: 0, + expectOptimized: false, + }, + { + name: "optimization level 1", + optimizationLvl: 1, + expectOptimized: true, + }, + { + name: "optimization level 2", + optimizationLvl: 2, + expectOptimized: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + test.WithTestFS(files, true, func(root string, fsys fs.FS) { + compiler := New(). + WithFS(fsys). + WithPaths(root). + WithOptimizationLevel(tc.optimizationLvl). + WithRegoAnnotationEntrypoints(true) + err := compiler.Build(t.Context()) + if err != nil { + t.Fatal(err) + } + + // Check what the original "test.rego" remnants look like at different optimization levels + var originalModule *bundle.ModuleFile + for i := range compiler.bundle.Modules { + if compiler.bundle.Modules[i].Path == "test.rego" { + originalModule = &compiler.bundle.Modules[i] + break + } + } + + if originalModule == nil { + t.Fatal("expected to find original 'test.rego' in bundle") + } + + // The behavior differs between optimization levels + var expectedOriginal string + switch tc.optimizationLvl { + case 0: // At level 0, all rules remain unchanged + expectedOriginal = files["test.rego"] + case 1: // At level 1, entrypoint rules removed, is_admin and has_permission inlined + expectedOriginal = `package pkg + +# METADATA +# description: this rule should keep its MD +other if input.other + +valid_actions := ["read", "write", "delete"] +` + case 2: // At level 2, entrypoint rules removed, support rules preserved + expectedOriginal = `package pkg + +# METADATA +# description: this rule should keep its MD +other if input.other + +is_admin if input.user == "admin" +valid_actions := ["read", "write", "delete"] +has_permission(action) if action in valid_actions +` + } + + actualOriginal := string(originalModule.Raw) + if actualOriginal != expectedOriginal { + t.Errorf("original module mismatch at level %d:\n\nexpected:\n%s\n\ngot:\n%s", + tc.optimizationLvl, expectedOriginal, actualOriginal) + } + }) + }) + } +} + func TestOptimizerError(t *testing.T) { tests := []struct { note string