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 open-policy-agent#5035.

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

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