From a35ac1710d5236172325e90095ab15bd21814256 Mon Sep 17 00:00:00 2001 From: Stephan Renatus Date: Tue, 23 Aug 2022 08:55:30 +0200 Subject: [PATCH] compile: allow opt-out of dependents gathering With `.WithPruneUnused(true)`, the compiler (of the compile package) no longer collects dependents of its entrypoints. The resulting bundle, if used with the wasm target, will no longer be semantically equivalent to the bundle built with the rego target. Since we're unable to have entrypoints for functions, this allows building modules that we couldn't build before. See #5035. Signed-off-by: Stephan Renatus --- compile/compile.go | 63 ++++++++++++++++++++++++++--------------- compile/compile_test.go | 41 +++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 23 deletions(-) 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