diff --git a/compile/compile.go b/compile/compile.go index 1aebe4c62e..dabe3314e2 100644 --- a/compile/compile.go +++ b/compile/compile.go @@ -61,6 +61,7 @@ type Compiler struct { bundle *bundle.Bundle // the bundle that the compiler operates on revision *string // the revision to set on the output bundle asBundle bool // whether to assume bundle layout on file loading or not + pruneUnused bool // whether to extend the entrypoint set for semantic equivalence of built bundles filter loader.Filter // filter to apply to file loader paths []string // file paths to load. TODO(tsandall): add support for supplying readers for embedded users. entrypoints orderedStringSet // policy entrypoints required for optimization and certain targets @@ -100,6 +101,20 @@ func (c *Compiler) WithAsBundle(enabled bool) *Compiler { return c } +// WithPruneUnused will make rules be ignored that are defined on the same +// package as the entrypoint, but that are not in the entrypoint set. +// +// Notably this includes functions (they can't be entrypoints) and causes +// the built bundle to no longer be semantically equivalent to the bundle built +// without wasm. +// +// This affects the 'wasm' and 'plan' targets only. It has no effect on +// building 'rego' bundles, i.e., "ordinary bundles". +func (c *Compiler) WithPruneUnused(enabled bool) *Compiler { + c.pruneUnused = enabled + return c +} + // WithEntrypoints sets the policy entrypoints on the compiler. Entrypoints tell the // compiler what rules to expect and where optimizations can be targeted. The wasm // target requires at least one entrypoint as does optimization. @@ -189,7 +204,7 @@ func (c *Compiler) WithCapabilities(capabilities *ast.Capabilities) *Compiler { return c } -//WithMetadata sets the additional data to be included in .manifest +// WithMetadata sets the additional data to be included in .manifest func (c *Compiler) WithMetadata(metadata *map[string]interface{}) *Compiler { c.metadata = metadata return c @@ -415,32 +430,34 @@ func (c *Compiler) compilePlan(ctx context.Context) error { } } - // Find transitive dependents of entrypoints and add them to the set to compile. - // - // NOTE(tsandall): We compile entrypoints because the evaluator does not support - // evaluation of wasm-compiled rules when 'with' statements are in-scope. Compiling - // out the dependents avoids the need to support that case for now. - deps := map[*ast.Rule]struct{}{} - for i := range c.entrypointrefs { - transitiveDocumentDependents(c.compiler, c.entrypointrefs[i], deps) - } + if !c.pruneUnused { + // Find transitive dependents of entrypoints and add them to the set to compile. + // + // NOTE(tsandall): We compile entrypoints because the evaluator does not support + // evaluation of wasm-compiled rules when 'with' statements are in-scope. Compiling + // out the dependents avoids the need to support that case for now. + deps := map[*ast.Rule]struct{}{} + for i := range c.entrypointrefs { + transitiveDocumentDependents(c.compiler, c.entrypointrefs[i], deps) + } - extras := ast.NewSet() - for rule := range deps { - extras.Add(ast.NewTerm(rule.Path())) - } + extras := ast.NewSet() + for rule := range deps { + extras.Add(ast.NewTerm(rule.Path())) + } - sorted := extras.Sorted() + sorted := extras.Sorted() - for i := 0; i < sorted.Len(); i++ { - p, err := sorted.Elem(i).Value.(ast.Ref).Ptr() - if err != nil { - return err - } + for i := 0; i < sorted.Len(); i++ { + p, err := sorted.Elem(i).Value.(ast.Ref).Ptr() + if err != nil { + return err + } - if !c.entrypoints.Contains(p) { - c.entrypoints = append(c.entrypoints, p) - c.entrypointrefs = append(c.entrypointrefs, sorted.Elem(i)) + if !c.entrypoints.Contains(p) { + c.entrypoints = append(c.entrypoints, p) + c.entrypointrefs = append(c.entrypointrefs, sorted.Elem(i)) + } } } diff --git a/compile/compile_test.go b/compile/compile_test.go index 55a30ef54b..b1387035eb 100644 --- a/compile/compile_test.go +++ b/compile/compile_test.go @@ -3,6 +3,7 @@ package compile import ( "bytes" "context" + "encoding/json" "errors" "fmt" "os" @@ -15,6 +16,7 @@ import ( "github.com/open-policy-agent/opa/ast" "github.com/open-policy-agent/opa/bundle" "github.com/open-policy-agent/opa/format" + "github.com/open-policy-agent/opa/internal/ir" "github.com/open-policy-agent/opa/internal/ref" "github.com/open-policy-agent/opa/loader" "github.com/open-policy-agent/opa/util" @@ -740,6 +742,45 @@ func TestCompilerPlanTarget(t *testing.T) { }) } +func TestCompilerPlanTargetPruneUnused(t *testing.T) { + files := map[string]string{ + "test.rego": `package test + p[1] + f(x) { p[x] }`, + } + + test.WithTempFS(files, func(root string) { + + compiler := New(). + WithPaths(root). + WithTarget("plan"). + WithEntrypoints("test"). + WithPruneUnused(true) + err := compiler.Build(context.Background()) + if err != nil { + t.Fatal(err) + } + + if len(compiler.bundle.PlanModules) == 0 { + t.Fatal("expected to find compiled plan module") + } + + plan := compiler.bundle.PlanModules[0].Raw + var policy ir.Policy + + if err := json.Unmarshal(plan, &policy); err != nil { + t.Fatal(err) + } + if exp, act := 1, len(policy.Funcs.Funcs); act != exp { + t.Fatalf("expected %d funcs, got %d", exp, act) + } + f := policy.Funcs.Funcs[0] + if exp, act := "g0.data.test.p", f.Name; act != exp { + t.Fatalf("expected func named %v, got %v", exp, act) + } + }) +} + func TestCompilerSetRevision(t *testing.T) { files := map[string]string{ "test.rego": `package test