cmd/build+compile: allow opt-out of dependents gathering (#5038)

* 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.

Fixes #5035.

* cmd/build: expose new configurable via --prune-unused

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

cmd/build.go

@@ -27,6 +27,7 @@ type buildParams struct {
 	capabilities       *capabilitiesFlag
 	target             *util.EnumFlag
 	bundleMode         bool
+	pruneUnused        bool
 	optimizationLevel  int
 	entrypoints        repeatedStringFlag
 	outputFile         string
@@ -114,6 +115,11 @@ The 'build' command supports targets (specified by -t):
             the input files for each specified entrypoint. The bundle may contain the
             original policy or data files.
 
+    plan    The plan target emits a bundle containing a plan, i.e., an intermediate
+			representation compiled from the input files for each specified entrypoint.
+			This is for further processing, OPA cannot evaluate a "plan bundle" like it
+			can evaluate a wasm or rego bundle.
+
 The -e flag tells the 'build' command which documents will be queried by the software
 asking for policy decisions, so that it can focus optimization efforts and ensure
 that document is not eliminated by the optimizer.
@@ -218,7 +224,8 @@ against OPA v0.22.0:
 	}
 
 	buildCommand.Flags().VarP(buildParams.target, "target", "t", "set the output bundle target type")
-	buildCommand.Flags().BoolVarP(&buildParams.debug, "debug", "", false, "enable debug output")
+	buildCommand.Flags().BoolVar(&buildParams.pruneUnused, "prune-unused", false, "exclude dependents of entrypoints")
+	buildCommand.Flags().BoolVar(&buildParams.debug, "debug", false, "enable debug output")
 	buildCommand.Flags().IntVarP(&buildParams.optimizationLevel, "optimize", "O", 0, "set optimization level")
 	buildCommand.Flags().VarP(&buildParams.entrypoints, "entrypoint", "e", "set slash separated entrypoint path")
 	buildCommand.Flags().VarP(&buildParams.revision, "revision", "r", "set output bundle revision")
@@ -273,6 +280,7 @@ func dobuild(params buildParams, args []string) error {
 		WithCapabilities(capabilities).
 		WithTarget(params.target.String()).
 		WithAsBundle(params.bundleMode).
+		WithPruneUnused(params.pruneUnused).
 		WithOptimizationLevel(params.optimizationLevel).
 		WithOutput(buf).
 		WithEntrypoints(params.entrypoints.v...).

cmd/build_test.go

@@ -220,3 +220,71 @@ func TestBuildVerificationConfigError(t *testing.T) {
 		}
 	})
 }
+
+func TestBuildPlanWithPruneUnused(t *testing.T) {
+
+	files := map[string]string{
+		"test.rego": `
+			package test
+			
+			p[1]
+			
+			f(x) { p[x] }
+		`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+		params := newBuildParams()
+		if err := params.target.Set("plan"); err != nil {
+			t.Fatal(err)
+		}
+		params.pruneUnused = true
+		params.entrypoints.v = []string{"test"}
+		params.outputFile = path.Join(root, "bundle.tar.gz")
+
+		err := dobuild(params, []string{root})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = loader.NewFileLoader().AsBundle(params.outputFile)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Check that manifest is not written given no input manifest and no other flags
+		f, err := os.Open(params.outputFile)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer f.Close()
+
+		gr, err := gzip.NewReader(f)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		tr := tar.NewReader(gr)
+
+		found := false // for plan.json
+
+		for {
+			f, err := tr.Next()
+			if err == io.EOF {
+				break
+			} else if err != nil {
+				t.Fatal(err)
+			}
+			switch {
+			case f.Name == "/plan.json":
+				found = true
+			case f.Name == "/data.json" || strings.HasSuffix(f.Name, "/test.rego"): // expected
+			default:
+				t.Errorf("unexpected file: %s", f.Name)
+			}
+		}
+		if !found {
+			t.Error("plan.json not found")
+		}
+	})
+}

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))
+			}
 		}
 	}
 

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