compile: Add new package that implements bundle compiling and linking

The compiler implements optimization levels that allow callers to
control how aggressively OPA will attempt to optimize the bundle. By
default, optimizations are disabled.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

compile/compile.go

@@ -0,0 +1,608 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+// Package compile implements bundles compilation and linking.
+package compile
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"regexp"
+	"sort"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/bundle"
+	"github.com/open-policy-agent/opa/internal/ref"
+	initload "github.com/open-policy-agent/opa/internal/runtime/init"
+	"github.com/open-policy-agent/opa/loader"
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/storage/inmem"
+)
+
+const (
+	// TargetRego is the default target. The source rego is copied (potentially
+	// rewritten for optimization purpsoes) into the bundle. The target supports
+	// base documents.
+	TargetRego = "rego"
+
+	// TargetWasm is an alternative target that compiles the policy into a wasm
+	// module instead of Rego. The target supports base documents.
+	TargetWasm = "wasm"
+)
+
+var validTargets = map[string]struct{}{
+	TargetRego: struct{}{},
+	TargetWasm: struct{}{},
+}
+
+// Compiler implements bundle compilation and linking.
+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
+	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
+	optimizationLevel int              // how aggressive should optimization be
+	target            string           // target type (wasm, rego, etc.)
+	output            io.Writer        // output stream to write bundle to
+	entrypointrefs    []*ast.Term      // validated entrypoints computed from default decision or manually supplied entrypoints
+	compiler          *ast.Compiler    // rego ast compiler used for semantic checks and rewriting
+	debug             *debugEvents     // debug information produced during build
+}
+
+type debugEvents struct {
+	debug []Debug
+}
+
+func (d *debugEvents) Add(info Debug) {
+	if d != nil {
+		d.debug = append(d.debug, info)
+	}
+}
+
+// Debug contains debugging information produced by the build about optimizations and other operations.
+type Debug struct {
+	Location *ast.Location
+	Message  string
+}
+
+func (d Debug) String() string {
+	return fmt.Sprintf("%v: %v", d.Location, d.Message)
+}
+
+// New returns a new compiler instance that can be invoked.
+func New() *Compiler {
+	return &Compiler{
+		asBundle:          false,
+		optimizationLevel: 0,
+		target:            TargetRego,
+		output:            ioutil.Discard,
+		debug:             &debugEvents{},
+	}
+}
+
+// Debug returns a list of debug events produced by the compiler.
+func (c *Compiler) Debug() []Debug {
+	return c.debug.debug
+}
+
+// WithRevision sets the revision to include in the output bundle manifest.
+func (c *Compiler) WithRevision(r string) *Compiler {
+	c.revision = &r
+	return c
+}
+
+// WithAsBundle sets file loading mode on the compiler.
+func (c *Compiler) WithAsBundle(enabled bool) *Compiler {
+	c.asBundle = enabled
+	return c
+}
+
+// WithEntrypoints sets the policy entrypoints on the compiler. Entrypoints tell the
+// compiler what rules to expect and where optimizations can be targetted. The wasm
+// target requires at least one entrypoint as does optimization.
+func (c *Compiler) WithEntrypoints(e ...string) *Compiler {
+	c.entrypoints = c.entrypoints.Append(e...)
+	return c
+}
+
+// WithOptimizationLevel sets the optimization level on the compiler. By default
+// optimizations are disabled. Higher levels apply more aggressive optimizations
+// but can take longer. Currently only two levels are supported: 0 (disabled) and
+// 1 (enabled).
+func (c *Compiler) WithOptimizationLevel(n int) *Compiler {
+	c.optimizationLevel = n
+	return c
+}
+
+// WithTarget sets the output target type to use.
+func (c *Compiler) WithTarget(t string) *Compiler {
+	c.target = t
+	return c
+}
+
+// WithOutput sets the output stream to write the bundle to.
+func (c *Compiler) WithOutput(w io.Writer) *Compiler {
+	c.output = w
+	return c
+}
+
+// WithPaths adds input filepaths to read policy and data from.
+func (c *Compiler) WithPaths(p ...string) *Compiler {
+	c.paths = append(c.paths, p...)
+	return c
+}
+
+// WithFilter sets the loader filter to use when reading non-bundle input files.
+func (c *Compiler) WithFilter(filter loader.Filter) *Compiler {
+	c.filter = filter
+	return c
+}
+
+// Build compiles and links the input files and outputs a bundle to the writer.
+func (c *Compiler) Build(ctx context.Context) error {
+
+	if err := c.init(); err != nil {
+		return err
+	}
+
+	if err := c.initBundle(); err != nil {
+		return err
+	}
+
+	if err := c.optimize(ctx); err != nil {
+		return err
+	}
+
+	if c.target == TargetWasm {
+		if err := c.compileWasm(ctx); err != nil {
+			return err
+		}
+	}
+
+	if c.revision != nil {
+		c.bundle.Manifest.Revision = *c.revision
+	}
+
+	return bundle.NewWriter(c.output).Write(*c.bundle)
+}
+
+func (c *Compiler) init() error {
+
+	if _, ok := validTargets[c.target]; !ok {
+		return fmt.Errorf("invalid target %q", c.target)
+	}
+
+	for _, e := range c.entrypoints {
+
+		r, err := ref.ParseDataPath(e)
+		if err != nil {
+			return fmt.Errorf("entrypoint %v not valid: use <package>/<rule>", e)
+		}
+
+		if len(r) <= 2 {
+			return fmt.Errorf("entrypoint %v too short: use <package>/<rule>", e)
+		}
+
+		c.entrypointrefs = append(c.entrypointrefs, ast.NewTerm(r))
+	}
+
+	if c.optimizationLevel > 0 && len(c.entrypointrefs) == 0 {
+		return errors.New("bundle optimizations require at least one entrypoint")
+	}
+
+	if c.target == TargetWasm && len(c.entrypointrefs) != 1 {
+		return errors.New("wasm compilation requires exactly one entrypoint")
+	}
+
+	return nil
+}
+
+func (c *Compiler) initBundle() error {
+
+	// TODO(tsandall): the metrics object should passed through here so we that
+	// we can track read and parse times.
+	load, err := initload.LoadPaths(c.paths, c.filter, c.asBundle)
+	if err != nil {
+		return err
+	}
+
+	if c.asBundle {
+		var names []string
+
+		for k := range load.Bundles {
+			names = append(names, k)
+		}
+
+		sort.Strings(names)
+		var bundles []*bundle.Bundle
+
+		for _, k := range names {
+			bundles = append(bundles, load.Bundles[k])
+		}
+
+		result, err := bundle.Merge(bundles)
+		if err != nil {
+			return fmt.Errorf("bundle merge failed: %v", err)
+		}
+
+		c.bundle = result
+		return nil
+	}
+
+	// TODO(tsandall): add support for controlling roots. Either the caller could
+	// supply them or the compiler could infer them based on the packages and data
+	// contents. The latter would require changes to the loader to preserve the
+	// locations where base documents were mounted under data.
+	result := &bundle.Bundle{}
+	result.Manifest.Init()
+	result.Data = load.Files.Documents
+
+	var modules []string
+
+	for k := range load.Files.Modules {
+		modules = append(modules, k)
+	}
+
+	sort.Strings(modules)
+
+	for _, module := range modules {
+		result.Modules = append(result.Modules, bundle.ModuleFile{
+			URL:    load.Files.Modules[module].Name,
+			Path:   load.Files.Modules[module].Name,
+			Parsed: load.Files.Modules[module].Parsed,
+			Raw:    load.Files.Modules[module].Raw,
+		})
+	}
+
+	c.bundle = result
+
+	return nil
+}
+
+func (c *Compiler) optimize(ctx context.Context) error {
+
+	if c.optimizationLevel <= 0 {
+		var err error
+		c.compiler, err = compile(c.bundle)
+		return err
+	}
+
+	o := newOptimizer(c.bundle).
+		WithEntrypoints(c.entrypointrefs).
+		WithDebug(c.debug)
+
+	err := o.Do(ctx)
+	if err != nil {
+		return err
+	}
+
+	c.bundle = o.Bundle()
+
+	return nil
+}
+
+func (c *Compiler) compileWasm(ctx context.Context) error {
+
+	// Lazily compile the modules if needed. If optimizations were run, the
+	// AST compiler will not be set because the default target does not require it.
+	if c.compiler == nil {
+		var err error
+		c.compiler, err = compile(c.bundle)
+		if err != nil {
+			return err
+		}
+	}
+
+	store := inmem.NewFromObject(c.bundle.Data)
+	resultSym := ast.VarTerm(ast.WildcardPrefix + "result")
+
+	cr, err := rego.New(
+		rego.ParsedQuery(ast.NewBody(ast.Equality.Expr(resultSym, c.entrypointrefs[0]))),
+		rego.Compiler(c.compiler),
+		rego.Store(store),
+	).Compile(ctx)
+
+	if err != nil {
+		return err
+	}
+
+	c.bundle.Wasm = cr.Bytes
+
+	return nil
+}
+
+type undefinedEntrypointErr struct {
+	Entrypoint *ast.Term
+}
+
+func (err undefinedEntrypointErr) Error() string {
+	return fmt.Sprintf("undefined entrypoint %v", err.Entrypoint)
+}
+
+type optimizer struct {
+	bundle          *bundle.Bundle
+	compiler        *ast.Compiler
+	entrypoints     []*ast.Term
+	nsprefix        string
+	resultsymprefix string
+	outputprefix    string
+	debug           *debugEvents
+}
+
+func newOptimizer(b *bundle.Bundle) *optimizer {
+	return &optimizer{
+		bundle:          b,
+		nsprefix:        "partial",
+		resultsymprefix: ast.WildcardPrefix,
+		outputprefix:    "optimized",
+	}
+}
+
+func (o *optimizer) WithDebug(debug *debugEvents) *optimizer {
+	o.debug = debug
+	return o
+}
+
+func (o *optimizer) WithEntrypoints(es []*ast.Term) *optimizer {
+	o.entrypoints = es
+	return o
+}
+
+func (o *optimizer) Do(ctx context.Context) error {
+
+	// TODO(tsandall): implement optimization levels. These will just be params on partial evaluation for now.
+	//
+	// Level 1: PE w/ constant folding. Only inline rules that are completely known.
+	// Level 2: L1 except inlining of rules with unknowns.
+	// Level 3: L2 except aggressive inlining using negation and copy propagation optimizations.
+
+	// NOTE(tsandall): if there are multiple entrypoints, copy the bundle because
+	// if any of the optimization steps fail, we do not want to leave the caller's
+	// bundle in a partially modified state.
+	if len(o.entrypoints) > 1 {
+		cpy := o.bundle.Copy()
+		o.bundle = &cpy
+	}
+
+	// initialize other inputs to the optimization process (store, symbols, etc.)
+	data := o.bundle.Data
+	if data == nil {
+		data = map[string]interface{}{}
+	}
+
+	store := inmem.NewFromObject(data)
+	resultsym := ast.VarTerm(o.resultsymprefix + "result")
+	usedFilenames := map[string]int{}
+
+	// NOTE(tsandall): the entrypoints are optimized in order so that the optimization
+	// of entrypoint[1] sees the optimization of entrypoint[0] and so on. This is needed
+	// because otherwise the optimization outputs (e.g., support rules) would have to
+	// merged somehow. Instead of dealing with that, just run the optimizations in the
+	// order the user supplied the entrypoints in.
+	for i, e := range o.entrypoints {
+
+		var err error
+		o.compiler, err = compile(o.bundle)
+		if err != nil {
+			return err
+		}
+
+		r := rego.New(
+			rego.ParsedQuery(ast.NewBody(ast.Equality.Expr(resultsym, e))),
+			rego.PartialNamespace(o.nsprefix),
+			rego.DisableInlining(o.findRequiredDocuments(e)),
+			rego.SkipPartialNamespace(true),
+			rego.Compiler(o.compiler),
+			rego.Store(store),
+		)
+
+		pq, err := r.Partial(ctx)
+		if err != nil {
+			return err
+		}
+
+		// NOTE(tsandall): this might be a bit too strict but in practice it's
+		// unlikely users will want to ignore undefined entrypoints. make this
+		// optional in the future.
+		if len(pq.Queries) == 0 {
+			return undefinedEntrypointErr{Entrypoint: e}
+		}
+
+		if module := o.getSupportForEntrypoint(pq.Queries, e, resultsym); module != nil {
+			pq.Support = append(pq.Support, module)
+		}
+
+		for j, module := range pq.Support {
+			fileName := o.getSupportModuleFilename(usedFilenames, module, i, j)
+			o.bundle.Modules = o.merge(o.bundle.Modules, bundle.ModuleFile{
+				URL:    fileName,
+				Path:   fileName,
+				Parsed: module,
+			})
+		}
+	}
+
+	sort.Slice(o.bundle.Modules, func(i, j int) bool {
+		return o.bundle.Modules[i].URL < o.bundle.Modules[j].URL
+	})
+
+	// NOTE(tsandall): prune out rules and data that are not referenced in the bundle
+	// in the future.
+	o.bundle.Manifest.AddRoot(o.nsprefix)
+	o.bundle.Manifest.Revision = ""
+
+	return nil
+}
+
+func (o *optimizer) Bundle() *bundle.Bundle {
+	return o.bundle
+}
+
+func (o *optimizer) findRequiredDocuments(ref *ast.Term) []string {
+
+	keep := map[string]*ast.Location{}
+	deps := map[*ast.Rule]struct{}{}
+
+	for _, r := range o.compiler.GetRules(ref.Value.(ast.Ref)) {
+		transitiveDependents(o.compiler, r, deps)
+	}
+
+	for rule := range deps {
+		ast.WalkExprs(rule, func(expr *ast.Expr) bool {
+			for _, with := range expr.With {
+				// TODO(tsandall): this should be improved to exclude refs that are
+				// marked as unknown. Since the build command does not allow users to
+				// set unknowns, we can hardcode to assume 'input'.
+				if !with.Target.Value.(ast.Ref).HasPrefix(ast.InputRootRef) {
+					keep[with.Target.String()] = with.Target.Location
+				}
+			}
+			return false
+		})
+	}
+
+	var result []string
+
+	for k := range keep {
+		result = append(result, k)
+	}
+
+	sort.Strings(result)
+
+	for _, k := range result {
+		o.debug.Add(Debug{
+			Location: keep[k],
+			Message:  fmt.Sprintf("disables inlining of %v", k),
+		})
+	}
+
+	return result
+}
+
+func (o *optimizer) getSupportForEntrypoint(queries []ast.Body, e *ast.Term, resultsym *ast.Term) *ast.Module {
+
+	path := e.Value.(ast.Ref)
+	name := ast.Var(path[len(path)-1].Value.(ast.String))
+	module := &ast.Module{Package: &ast.Package{Path: path[:len(path)-1]}}
+
+	for _, query := range queries {
+		// NOTE(tsandall): when the query refers to the original entrypoint, throw it
+		// away since this would create a recursive rule--this occurs if the entrypoint
+		// cannot be partially evaluated.
+		stop := false
+		ast.WalkRefs(query, func(x ast.Ref) bool {
+			if !stop {
+				if x.HasPrefix(path) {
+					stop = true
+				}
+			}
+			return stop
+		})
+		if stop {
+			return nil
+		}
+		module.Rules = append(module.Rules, &ast.Rule{
+			Head:   ast.NewHead(name, nil, resultsym),
+			Body:   query,
+			Module: module,
+		})
+	}
+
+	return module
+}
+
+func (o *optimizer) merge(dst []bundle.ModuleFile, src bundle.ModuleFile) (result []bundle.ModuleFile) {
+
+	// NOTE(tsandlal): replace with trie if this becomes a bottleneck.
+
+	for i := range dst {
+		var keep []*ast.Rule
+		for _, old := range dst[i].Parsed.Rules {
+			var discard bool
+			for _, new := range src.Parsed.Rules {
+				if old.Path().HasPrefix(new.Path()) {
+					discard = true
+					break
+				}
+			}
+			if !discard {
+				keep = append(keep, old)
+			}
+		}
+		if len(keep) > 0 {
+			dst[i].Parsed.Rules = keep
+			dst[i].Raw = nil
+			result = append(result, dst[i])
+		}
+	}
+
+	return append(result, src)
+}
+
+func (o *optimizer) getSupportModuleFilename(used map[string]int, module *ast.Module, entrypointIndex int, supportIndex int) string {
+
+	fileName, err := module.Package.Path.Ptr()
+
+	if err == nil && safePathPattern.MatchString(fileName) {
+		fileName = o.outputprefix + "/" + fileName
+		if c, ok := used[fileName]; ok {
+			fileName += fmt.Sprintf(".%d", c)
+		}
+		used[fileName]++
+		fileName += ".rego"
+		return fileName
+	}
+
+	return fmt.Sprintf("%v/%v/%v/%v.rego", o.outputprefix, o.nsprefix, entrypointIndex, supportIndex)
+}
+
+var safePathPattern = regexp.MustCompile(`^[\w-_/]+$`)
+
+func compile(b *bundle.Bundle) (*ast.Compiler, error) {
+
+	modules := map[string]*ast.Module{}
+
+	for _, mf := range b.Modules {
+		modules[mf.URL] = mf.Parsed
+	}
+
+	c := ast.NewCompiler()
+	c.Compile(modules)
+
+	if c.Failed() {
+		return nil, c.Errors
+	}
+
+	return c, nil
+}
+
+func transitiveDependents(compiler *ast.Compiler, rule *ast.Rule, deps map[*ast.Rule]struct{}) {
+	for x := range compiler.Graph.Dependents(rule) {
+		other := x.(*ast.Rule)
+		deps[other] = struct{}{}
+		transitiveDependents(compiler, other, deps)
+	}
+}
+
+type orderedStringSet []string
+
+func (ss orderedStringSet) Append(s ...string) orderedStringSet {
+	for _, x := range s {
+		var found bool
+		for _, other := range ss {
+			if x == other {
+				found = true
+			}
+		}
+		if !found {
+			ss = append(ss, x)
+		}
+	}
+	return ss
+}

compile/compile_test.go

@@ -0,0 +1,834 @@
+package compile
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"path"
+	"reflect"
+	"sort"
+	"strings"
+	"testing"
+
+	"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/loader"
+	"github.com/open-policy-agent/opa/util"
+	"github.com/open-policy-agent/opa/util/test"
+)
+
+func TestOrderedStringSet(t *testing.T) {
+	var ss orderedStringSet
+	result := ss.Append("a", "b", "b", "a", "e", "c", "e")
+	if !reflect.DeepEqual(result, orderedStringSet{"a", "b", "e", "c"}) {
+		t.Fatal(result)
+	}
+}
+
+func TestCompilerInitErrors(t *testing.T) {
+
+	ctx := context.Background()
+
+	tests := []struct {
+		note string
+		c    *Compiler
+		want error
+	}{
+		{
+			note: "bad target",
+			c:    New().WithTarget("deadbeef"),
+			want: fmt.Errorf("invalid target \"deadbeef\""),
+		},
+		{
+			note: "entrypoint parse error",
+			c:    New().WithEntrypoints("foo%bar"),
+			want: fmt.Errorf("entrypoint foo%%bar not valid: use <package>/<rule>"),
+		},
+		{
+			note: "entrypoint too short error",
+			c:    New().WithEntrypoints("foo"),
+			want: fmt.Errorf("entrypoint foo too short: use <package>/<rule>"),
+		},
+		{
+			note: "optimizations require entrypoint",
+			c:    New().WithOptimizationLevel(1),
+			want: errors.New("bundle optimizations require at least one entrypoint"),
+		},
+		{
+			note: "wasm compilation requires exactly one entrypoint",
+			c:    New().WithTarget("wasm"),
+			want: errors.New("wasm compilation requires exactly one entrypoint"),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			err := tc.c.Build(ctx)
+			if err == nil {
+				t.Fatal("expected error")
+			} else if err.Error() != tc.want.Error() {
+				t.Fatalf("expected %v but got %v", tc.want, err)
+			}
+		})
+	}
+
+}
+
+func TestCompilerLoadError(t *testing.T) {
+
+	test.WithTempFS(nil, func(root string) {
+		err := New().WithPaths(path.Join(root, "does-not-exist")).Build(context.Background())
+		if err == nil {
+			t.Fatal("expected failure")
+		}
+	})
+}
+
+func TestCompilerLoadAsBundleSuccess(t *testing.T) {
+
+	ctx := context.Background()
+
+	files := map[string]string{
+		"b1/.manifest": `{"roots": ["b1"]}`,
+		"b1/test.rego": `
+			package b1.test
+
+			p = 1`,
+		"b1/data.json": `
+			{"b1": {"k": "v"}}`,
+		"b2/.manifest": `{"roots": ["b2"]}`,
+		"b2/data.json": `
+			{"b2": {"k2": "v2"}}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		root1 := path.Join(root, "b1")
+		root2 := path.Join(root, "b2")
+
+		compiler := New().
+			WithPaths(root1, root2).
+			WithAsBundle(true)
+
+		err := compiler.Build(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Verify result is just merger of two bundles.
+		a, err := loader.NewFileLoader().AsBundle(root1)
+		if err != nil {
+			panic(err)
+		}
+
+		b, err := loader.NewFileLoader().AsBundle(root2)
+		if err != nil {
+			panic(err)
+		}
+
+		exp, err := bundle.Merge([]*bundle.Bundle{a, b})
+		if err != nil {
+			panic(err)
+		}
+
+		if !compiler.bundle.Equal(*exp) {
+			t.Fatalf("expected %v but got %v", exp, compiler.bundle)
+		}
+
+		expRoots := []string{"b1", "b2"}
+		expManifest := bundle.Manifest{
+			Roots: &expRoots,
+		}
+
+		if !compiler.bundle.Manifest.Equal(expManifest) {
+			t.Fatalf("expected %v but got %v", compiler.bundle.Manifest, expManifest)
+		}
+	})
+}
+
+func TestCompilerLoadAsBundleMergeError(t *testing.T) {
+
+	ctx := context.Background()
+
+	// Omit manifests (defaulting to '') to trigger a merge error
+	files := map[string]string{
+		"b1/test.rego": `
+			package b1.test
+
+			p = 1`,
+		"b1/data.json": `
+			{"b1": {"k": "v"}}`,
+		"b2/data.json": `
+			{"b2": {"k2": "v2"}}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		root1 := path.Join(root, "b1")
+		root2 := path.Join(root, "b2")
+
+		compiler := New().
+			WithPaths(root1, root2).
+			WithAsBundle(true)
+
+		err := compiler.Build(ctx)
+		if err == nil || err.Error() != "bundle merge failed: manifest has overlapped roots: '' and ''" {
+			t.Fatal(err)
+		}
+	})
+}
+
+func TestCompilerLoadFilesystem(t *testing.T) {
+
+	files := map[string]string{
+		"test.rego": `
+			package b1.test
+
+			p = 1`,
+		"data.json": `
+			{"b1": {"k": "v"}}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().
+			WithPaths(root)
+
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Verify result is just bundle load.
+		exp, err := loader.NewFileLoader().AsBundle(root)
+		if err != nil {
+			panic(err)
+		}
+
+		if !compiler.bundle.Equal(*exp) {
+			t.Fatalf("Expected:\n\n%v\n\nGot:\n\n%v", exp, compiler.bundle)
+		}
+	})
+}
+
+func TestCompilerLoadHonorsFilter(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+			package b1.test
+
+			p = 1`,
+		"data.json": `
+			{"b1": {"k": "v"}}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().
+			WithPaths(root).
+			WithFilter(func(abspath string, _ os.FileInfo, _ int) bool {
+				return strings.HasSuffix(abspath, ".json")
+			})
+
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(compiler.bundle.Data) > 0 {
+			t.Fatal("expected no data to be loaded")
+		}
+	})
+}
+
+func TestCompilerError(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+			package test
+			default p = false
+			p { p }`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().
+			WithPaths(root)
+
+		err := compiler.Build(context.Background())
+		if err == nil {
+			t.Fatal("expected error")
+		}
+
+		astErr, ok := err.(ast.Errors)
+		if !ok || len(astErr) != 1 || astErr[0].Code != ast.RecursionErr {
+			t.Fatal("unexpected error:", err)
+		}
+	})
+}
+
+func TestCompilerOptimization(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+			package test
+			default p = false
+			p { input.x = data.foo }`,
+		"data.json": `
+			{"foo": 1}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().
+			WithPaths(root).
+			WithOptimizationLevel(1).
+			WithEntrypoints("test/p")
+
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		exp := ast.MustParseModule(`
+			package test
+
+			p { input.x = 1}
+			default p = false
+		`)
+
+		if len(compiler.bundle.Modules) != 1 || !compiler.bundle.Modules[0].Parsed.Equal(exp) {
+			t.Fatalf("expected one module but got: %v", compiler.bundle.Modules)
+		}
+	})
+}
+
+func TestCompilerWasmTarget(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `package test
+
+		p = true`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().WithPaths(root).WithTarget("wasm").WithEntrypoints("test/p")
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(compiler.bundle.Wasm) == 0 {
+			t.Fatal("expected to find compiled wasm module")
+		}
+
+	})
+}
+
+func TestCompilerWasmTargetLazyCompile(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `package test
+
+		p { input.x = q }
+		q = "foo"`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().WithPaths(root).WithTarget("wasm").WithEntrypoints("test/p").WithOptimizationLevel(1)
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(compiler.bundle.Wasm) == 0 {
+			t.Fatal("expected to find compiled wasm module")
+		}
+
+		if _, exists := compiler.compiler.Modules["optimized/test.rego"]; !exists {
+			t.Fatal("expected to find optimized module on compiler")
+		}
+	})
+}
+
+func TestCompilerSetRevision(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `package test
+
+		p = true`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().WithPaths(root).WithRevision("deadbeef")
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if compiler.bundle.Manifest.Revision != "deadbeef" {
+			t.Fatal("expected revision to be set but got:", compiler.bundle.Manifest)
+		}
+	})
+}
+
+func TestCompilerOutput(t *testing.T) {
+	// NOTE(tsandall): must use format package here because the compiler formats.
+	files := map[string]string{
+		"test.rego": string(format.MustAst(ast.MustParseModule(`package test
+
+		p { input.x = data.foo }`))),
+		"data.json": `{"foo": 1}`,
+	}
+	test.WithTempFS(files, func(root string) {
+
+		buf := bytes.NewBuffer(nil)
+		compiler := New().WithPaths(root).WithOutput(buf)
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+		result, err := bundle.NewReader(buf).Read()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		exp, err := loader.NewFileLoader().AsBundle(root)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !exp.Equal(result) {
+			t.Fatalf("expected:\n\n%v\n\ngot:\n\n%v", *exp, result)
+		}
+
+		if !exp.Manifest.Equal(result.Manifest) {
+			t.Fatalf("expected:\n\n%v\n\ngot:\n\n%v", exp.Manifest, result.Manifest)
+		}
+
+	})
+}
+
+func TestOptimizerNoops(t *testing.T) {
+	tests := []struct {
+		note        string
+		entrypoints []string
+		modules     map[string]string
+	}{
+		{
+			note:        "recursive result",
+			entrypoints: []string{"data.test.foo"},
+			modules: map[string]string{
+				"test.rego": `
+					package test.foo.bar
+
+					p { input.x = 1 }
+				`,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			o := getOptimizer(tc.modules, "", tc.entrypoints)
+			cpy := o.bundle.Copy()
+			err := o.Do(context.Background())
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !o.bundle.Equal(cpy) {
+				t.Fatalf("Expected no change:\n\n%v\n\nGot:\n\n%v", prettyBundle{cpy}, prettyBundle{*o.bundle})
+			}
+		})
+	}
+}
+
+func TestOptimizerErrors(t *testing.T) {
+	tests := []struct {
+		note        string
+		entrypoints []string
+		modules     map[string]string
+		wantErr     error
+	}{
+		{
+			note:        "undefined entrypoint",
+			entrypoints: []string{"data.test.p"},
+			wantErr:     fmt.Errorf("undefined entrypoint data.test.p"),
+		},
+		{
+			note:        "compile error",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+					p { data.test.p }
+				`,
+			},
+			wantErr: fmt.Errorf("1 error occurred: test.rego:3: rego_recursion_error: rule p is recursive: p -> p"),
+		},
+		{
+			note:        "partial eval error",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+					p { div(1, 0, x) }
+				`,
+			},
+			wantErr: fmt.Errorf("test.rego:3: eval_builtin_error: div: divide by zero"),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			o := getOptimizer(tc.modules, "", tc.entrypoints)
+			cpy := o.bundle.Copy()
+			got := o.Do(context.Background())
+			if got == nil || got.Error() != tc.wantErr.Error() {
+				t.Fatalf("expected error to be %v but got %v", tc.wantErr, got)
+			}
+			if !o.bundle.Equal(cpy) {
+				t.Fatalf("Expected no change:\n\n%v\n\nGot:\n\n%v", prettyBundle{cpy}, prettyBundle{*o.bundle})
+			}
+		})
+	}
+}
+
+func TestOptimizerOutput(t *testing.T) {
+	tests := []struct {
+		note        string
+		entrypoints []string
+		modules     map[string]string
+		data        string
+		roots       []string
+		wantModules map[string]string
+	}{
+		{
+			note:        "rule pruning",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p {
+						q[input.x]
+					}
+
+					q[1]
+					q[2]
+					q[3]
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p = result { 1 = input.x; result = true }
+					p = result { 2 = input.x; result = true }
+					p = result { 3 = input.x; result = true }
+				`,
+				"test.rego": `
+					package test
+
+					q[1]
+					q[2]
+					q[3]
+				`,
+			},
+		},
+		{
+			note:        "support rules",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					default p = false
+
+					p { q[input.x] }
+
+					q[1]
+					q[2]`,
+			},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p = true { 1 = input.x }
+					p = true { 2 = input.x }
+
+					default p = false
+				`,
+				"test.rego": `
+					package test
+
+					q[1]
+					q[2]
+				`,
+			},
+		},
+		{
+			note:        "multiple entrypoints",
+			entrypoints: []string{"data.test.p", "data.test.r"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p {
+						q[input.x]
+					}
+
+					r {
+						q[input.x]
+					}
+
+					q[1]
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p = result { 1 = input.x; result = true }
+				`,
+				"optimized/test.1.rego": `
+					package test
+
+					r = result { 1 = input.x; result = true }
+				`,
+				"test.rego": `
+					package test
+
+					q[1] { true }
+				`,
+			},
+		},
+		{
+			note:        "package pruning",
+			entrypoints: []string{"data.test.foo"},
+			modules: map[string]string{
+				"test.rego": `
+					package test.foo.bar
+
+					p = true
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					foo = result { result = {"bar": {"p": true}} }`,
+			},
+		},
+		{
+			note:        "entrypoint dependent integrity",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p { q[input.x] }
+
+					q[x] {
+						s[x]
+					}
+
+					s[1]
+					s[2]
+
+					t {
+						p
+					}
+
+					r { t with q as {3} }
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/test.1.rego": `
+					package test
+
+					p = result { data.test.q[input.x]; result = true }
+				`,
+				"optimized/test.rego": `
+					package test
+
+					q[1] { true }
+					q[2] { true }
+				`,
+				"test.rego": `
+					package test
+
+					s[1] { true }
+					s[2] { true }
+					t { p }
+					r = true { t with q as {3} }
+				`,
+			},
+		},
+		{
+			note:        "output filename safety",
+			entrypoints: []string{`data.test["foo bar"].p`},
+			modules: map[string]string{
+				"x.rego": `
+					package test["foo bar"]  # package does not match safe pattern so use alt. format
+					p { q[input.x] }
+					q[1]
+					q[2]
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/partial/0/0.rego": `
+					package test["foo bar"]
+					p = result { 1 = input.x; result = true }
+					p = result { 2 = input.x; result = true }
+				`,
+				"x.rego": `
+					package test["foo bar"]
+
+					q[1]
+					q[2]
+				`,
+			},
+		},
+		{
+			note:        "generated package namespace",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p { not q }
+					q { k[input.a]; k[input.b] }  # generate a product that is not inlined
+					k = {1,2,3}
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/partial.rego": `
+					package partial
+
+					__not1_0__ = true { 1 = input.a; 1 = input.b }
+					__not1_0__ = true { 1 = input.a; 2 = input.b }
+					__not1_0__ = true { 1 = input.a; 3 = input.b }
+					__not1_0__ = true { 2 = input.a; 1 = input.b }
+					__not1_0__ = true { 2 = input.a; 2 = input.b }
+					__not1_0__ = true { 2 = input.a; 3 = input.b }
+					__not1_0__ = true { 3 = input.a; 1 = input.b }
+					__not1_0__ = true { 3 = input.a; 2 = input.b }
+					__not1_0__ = true { 3 = input.a; 3 = input.b }
+				`,
+				"optimized/test.rego": `
+					package test
+
+					p = result { not data.partial.__not1_0__; result = true }
+				`,
+				"test.rego": `
+					package test
+
+					q = true { k[input.a]; k[input.b] }
+					k = {1, 2, 3} { true }
+				`,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+
+			o := getOptimizer(tc.modules, tc.data, tc.entrypoints)
+			original := o.bundle.Copy()
+			err := o.Do(context.Background())
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			exp := &bundle.Bundle{
+				Modules: getModuleFiles(tc.wantModules, false),
+				Data:    original.Data, // data is not pruned at all today
+			}
+
+			exp.Manifest.Revision = "" // optimizations must reset the revision.
+
+			if !exp.Equal(*o.bundle) {
+				t.Errorf("Expected:\n\n%v\n\nGot:\n\n%v", prettyBundle{*exp}, prettyBundle{*o.bundle})
+			}
+
+			if !o.bundle.Manifest.Equal(exp.Manifest) {
+				t.Errorf("Expected manifest: %v\n\nGot manifest: %v", exp.Manifest, o.bundle.Manifest)
+			}
+		})
+	}
+}
+
+func getOptimizer(modules map[string]string, data string, entries []string) *optimizer {
+
+	b := &bundle.Bundle{
+		Modules: getModuleFiles(modules, true),
+	}
+
+	if data != "" {
+		b.Data = util.MustUnmarshalJSON([]byte(data)).(map[string]interface{})
+	}
+
+	b.Manifest.Init()
+	b.Manifest.Revision = "DEADBEEF" // ensures that the manifest revision is getting reset
+	entrypoints := make([]*ast.Term, len(entries))
+
+	for i := range entrypoints {
+		entrypoints[i] = ast.MustParseTerm(entries[i])
+	}
+
+	o := newOptimizer(b).
+		WithEntrypoints(entrypoints)
+
+	o.resultsymprefix = ""
+
+	return o
+}
+
+func getModuleFiles(src map[string]string, includeRaw bool) []bundle.ModuleFile {
+
+	var keys []string
+
+	for k := range src {
+		keys = append(keys, k)
+	}
+
+	sort.Strings(keys)
+	var modules []bundle.ModuleFile
+
+	for _, k := range keys {
+		module, err := ast.ParseModule(k, src[k])
+		if err != nil {
+			panic(err)
+		}
+		modules = append(modules, bundle.ModuleFile{
+			Parsed: module,
+			Path:   k,
+			URL:    k,
+		})
+		if includeRaw {
+			modules[len(modules)-1].Raw = []byte(src[k])
+		}
+	}
+
+	return modules
+}
+
+type prettyBundle struct {
+	bundle.Bundle
+}
+
+func (p prettyBundle) String() string {
+
+	buf := []string{fmt.Sprintf("%d module(s) (hiding data):", len(p.Modules)), ""}
+
+	for _, mf := range p.Modules {
+		buf = append(buf, "#")
+		buf = append(buf, fmt.Sprintf("# Module: %q", mf.Path))
+		buf = append(buf, "#")
+		buf = append(buf, mf.Parsed.String())
+		buf = append(buf, "")
+	}
+
+	return strings.Join(buf, "\n")
+}