compile: Add support for shallow inlining

This commit plumbs the new partial evaluation mode into the compiler
and defines the new optimization levels (0=off, 1=shallow inlining,
2=aggressive inlining).

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

compile/compile.go

@@ -115,8 +115,7 @@ func (c *Compiler) WithEntrypoints(e ...string) *Compiler {
 
 // 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).
+// but can take longer.
 func (c *Compiler) WithOptimizationLevel(n int) *Compiler {
 	c.optimizationLevel = n
 	return c
@@ -277,7 +276,8 @@ func (c *Compiler) optimize(ctx context.Context) error {
 
 	o := newOptimizer(c.bundle).
 		WithEntrypoints(c.entrypointrefs).
-		WithDebug(c.debug)
+		WithDebug(c.debug).
+		WithShallowInlining(c.optimizationLevel <= 1)
 
 	err := o.Do(ctx)
 	if err != nil {
@@ -334,6 +334,7 @@ type optimizer struct {
 	nsprefix        string
 	resultsymprefix string
 	outputprefix    string
+	shallow         bool
 	debug           *debugEvents
 }
 
@@ -356,13 +357,12 @@ func (o *optimizer) WithEntrypoints(es []*ast.Term) *optimizer {
 	return o
 }
 
-func (o *optimizer) Do(ctx context.Context) error {
+func (o *optimizer) WithShallowInlining(yes bool) *optimizer {
+	o.shallow = yes
+	return o
+}
 
-	// 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.
+func (o *optimizer) Do(ctx context.Context) error {
 
 	// 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
@@ -399,6 +399,7 @@ func (o *optimizer) Do(ctx context.Context) error {
 			rego.ParsedQuery(ast.NewBody(ast.Equality.Expr(resultsym, e))),
 			rego.PartialNamespace(o.nsprefix),
 			rego.DisableInlining(o.findRequiredDocuments(e)),
+			rego.ShallowInlining(o.shallow),
 			rego.SkipPartialNamespace(true),
 			rego.Compiler(o.compiler),
 			rego.Store(store),

compile/compile_test.go

@@ -268,12 +268,14 @@ func TestCompilerError(t *testing.T) {
 	})
 }
 
-func TestCompilerOptimization(t *testing.T) {
+func TestCompilerOptimizationL1(t *testing.T) {
+
 	files := map[string]string{
 		"test.rego": `
 			package test
 			default p = false
-			p { input.x = data.foo }`,
+			p { q }
+			q { input.x = data.foo }`,
 		"data.json": `
 			{"foo": 1}`,
 	}
@@ -290,15 +292,82 @@ func TestCompilerOptimization(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		exp := ast.MustParseModule(`
+		optimizedExp := ast.MustParseModule(`
 			package test
 
-			p { input.x = 1}
+			q { input.x = 1 }
+			p { data.test.q = X; X }
 			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)
+		// NOTE(tsandall): PE generates vars with wildcard prefix. Instead of
+		// constructing the AST manually, just rewrite to the expected value
+		// here. If this becomes a common pattern, we could refactor (e.g.,
+		// allow caller to control var prefix, split into a reusable function,
+		// etc.)
+		ast.TransformVars(optimizedExp, func(x ast.Var) (ast.Value, error) {
+			if x == ast.Var("X") {
+				return ast.Var("$_term_1_01"), nil
+			}
+			return x, nil
+		})
+
+		if len(compiler.bundle.Modules) != 1 {
+			t.Fatalf("expected 1 module but got: %v", compiler.bundle.Modules)
+		}
+
+		if !compiler.bundle.Modules[0].Parsed.Equal(optimizedExp) {
+			t.Fatalf("expected optimized module to be:\n\n%v\n\ngot:\n\n%v", optimizedExp, compiler.bundle.Modules[0])
+		}
+	})
+}
+
+func TestCompilerOptimizationL2(t *testing.T) {
+
+	files := map[string]string{
+		"test.rego": `
+			package test
+			default p = false
+			p { q }
+			q { input.x = data.foo }`,
+		"data.json": `
+			{"foo": 1}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().
+			WithPaths(root).
+			WithOptimizationLevel(2).
+			WithEntrypoints("test/p")
+
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		prunedExp := ast.MustParseModule(`
+			package test
+
+			q { input.x = data.foo }`)
+
+		optimizedExp := ast.MustParseModule(`
+			package test
+
+			p { input.x = 1 }
+			default p = false
+		`)
+
+		if len(compiler.bundle.Modules) != 2 {
+			t.Fatalf("expected two modules but got: %v", compiler.bundle.Modules)
+		}
+
+		if !compiler.bundle.Modules[0].Parsed.Equal(prunedExp) {
+			t.Fatalf("expected pruned module to be:\n\n%v\n\ngot:\n\n%v", prunedExp, compiler.bundle.Modules[0])
+		}
+
+		if !compiler.bundle.Modules[1].Parsed.Equal(optimizedExp) {
+			t.Fatalf("expected optimized module to be:\n\n%v\n\ngot:\n\n%v", optimizedExp, compiler.bundle.Modules[1])
 		}
 	})
 }