perf: Reduce template compilation time by ~16% (#206)

* Remove unused rego helpers

Signed-off-by: Will Beason <willbeason@google.com>

* Remove unnecessary re-parse

Signed-off-by: Will Beason <willbeason@google.com>

* Use standard function for parsing hook

Signed-off-by: Will Beason <willbeason@google.com>

* Fix golangci lint errors

Signed-off-by: Will Beason <willbeason@google.com>

* Remove creating CTs from cpu/mem profile

Signed-off-by: Will Beason <willbeason@google.com>

constraint/cmd/rewrite-compatibility/main.go

@@ -66,7 +66,8 @@ func compileSrcs(
 	libs []string,
 	newPkgPrefix string,
 	oldRoot string,
-	newRoot string) error {
+	newRoot string,
+) error {
 	if len(cts) == 0 && len(libs) == 0 {
 		return fmt.Errorf("must specify --ct or --lib or both")
 	}

constraint/pkg/client/client_addtemplate_bench_test.go

@@ -290,6 +290,7 @@ func BenchmarkClient_AddTemplate_Parallel(b *testing.B) {
 					for i := range cts {
 						cts[i] = makeConstraintTemplate(i, tc.module, tc.libs...)
 					}
+					b.ResetTimer()
 
 					for i := 0; i < b.N; i++ {
 						b.StopTimer()

constraint/pkg/client/drivers/local/compilers.go

@@ -116,37 +116,29 @@ func (d *Compilers) list() map[string]map[string]*ast.Compiler {
 	return result
 }
 
-type TargetModule struct {
-	Rego string
-	Libs []string
-}
-
 // parseConstraintTemplate validates the rego in template target by parsing
 // rego modules.
-func parseConstraintTemplate(templ *templates.ConstraintTemplate, externs []string) (map[string]TargetModule, error) {
+func parseConstraintTemplate(templ *templates.ConstraintTemplate, externs []string) (map[string][]*ast.Module, error) {
 	rr, err := regorewriter.New(regorewriter.NewPackagePrefixer(templateLibPrefix), []string{libRoot}, externs)
 	if err != nil {
 		return nil, fmt.Errorf("creating rego rewriter: %w", err)
 	}
 
-	mods := make(map[string]TargetModule)
+	mods := make(map[string][]*ast.Module)
 	for _, target := range templ.Spec.Targets {
 		targetMods, err := parseConstraintTemplateTarget(rr, target)
 		if err != nil {
 			return nil, err
 		}
 
-		mods[target.Target] = TargetModule{
-			Rego: target.Rego,
-			Libs: targetMods,
-		}
+		mods[target.Target] = targetMods
 	}
 
 	return mods, nil
 }
 
-func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, targetSpec templates.Target) ([]string, error) {
-	entryPoint, err := parseModule(targetSpec.Rego)
+func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, targetSpec templates.Target) ([]*ast.Module, error) {
+	entryPoint, err := parseModule(templatePath, targetSpec.Rego)
 	if err != nil {
 		return nil, fmt.Errorf("%w: %v", clienterrors.ErrInvalidConstraintTemplate, err)
 	}
@@ -170,7 +162,14 @@ func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, targetSpec tem
 	rr.AddEntryPointModule(templatePath, entryPoint)
 	for idx, libSrc := range targetSpec.Libs {
 		libPath := fmt.Sprintf(`%s["lib_%d"]`, templateLibPrefix, idx)
-		if err = rr.AddLib(libPath, libSrc); err != nil {
+
+		m, err := parseModule(libPath, libSrc)
+		if err != nil {
+			return nil, fmt.Errorf("%w: %v",
+				clienterrors.ErrInvalidConstraintTemplate, err)
+		}
+
+		if err = rr.AddLib(libPath, m); err != nil {
 			return nil, fmt.Errorf("%w: %v",
 				clienterrors.ErrInvalidConstraintTemplate, err)
 		}
@@ -182,44 +181,28 @@ func parseConstraintTemplateTarget(rr *regorewriter.RegoRewriter, targetSpec tem
 			clienterrors.ErrInvalidConstraintTemplate, err)
 	}
 
-	var mods []string
-	err = sources.ForEachModule(func(m *regorewriter.Module) error {
-		content, err2 := m.Content()
-		if err2 != nil {
-			return err2
-		}
-		mods = append(mods, string(content))
-		return nil
-	})
-
-	if err != nil {
-		return nil, fmt.Errorf("%w: %v",
-			clienterrors.ErrInvalidConstraintTemplate, err)
+	var mods []*ast.Module
+	for _, m := range sources.EntryPoints {
+		mods = append(mods, m.Module)
+	}
+	for _, m := range sources.Libs {
+		mods = append(mods, m.Module)
 	}
 
 	return mods, nil
 }
 
-func compileTemplateTarget(module TargetModule, capabilities *ast.Capabilities, printEnabled bool) (*ast.Compiler, error) {
+func compileTemplateTarget(module []*ast.Module, capabilities *ast.Capabilities, printEnabled bool) (*ast.Compiler, error) {
 	compiler := ast.NewCompiler().
 		WithCapabilities(capabilities).
 		WithEnablePrintStatements(printEnabled)
 
-	modules := make(map[string]*ast.Module)
-
-	builtinModule, err := ast.ParseModule(hookModulePath, hookModule)
-	if err != nil {
-		return nil, fmt.Errorf("%w: %v", clienterrors.ErrParse, err)
-	}
-	modules[hookModulePath] = builtinModule
+	modules := make(map[string]*ast.Module, len(module)+1)
+	modules[hookModulePath] = hookModule
 
-	for i, lib := range module.Libs {
+	for i, lib := range module {
 		libPath := fmt.Sprintf("%s%d", templatePath, i)
-		libModule, err := ast.ParseModule(libPath, lib)
-		if err != nil {
-			return nil, fmt.Errorf("%w: %v", clienterrors.ErrParse, err)
-		}
-		modules[libPath] = libModule
+		modules[libPath] = lib
 	}
 
 	compiler.Compile(modules)
@@ -229,3 +212,18 @@ func compileTemplateTarget(module TargetModule, capabilities *ast.Capabilities,
 
 	return compiler, nil
 }
+
+// parseModule parses the module and also fails empty modules.
+func parseModule(path, rego string) (*ast.Module, error) {
+	module, err := ast.ParseModule(path, rego)
+	if err != nil {
+		return nil, err
+	}
+
+	if module == nil {
+		return nil, fmt.Errorf("%w: module %q is empty",
+			clienterrors.ErrInvalidModule, templatePath)
+	}
+
+	return module, nil
+}

constraint/pkg/client/drivers/local/driver.go

@@ -310,21 +310,6 @@ func (d *Driver) Dump(ctx context.Context) (string, error) {
 	return string(b), nil
 }
 
-// parseModule parses the module and also fails empty modules.
-func parseModule(rego string) (*ast.Module, error) {
-	module, err := ast.ParseModule(templatePath, rego)
-	if err != nil {
-		return nil, err
-	}
-
-	if module == nil {
-		return nil, fmt.Errorf("%w: module %q is empty",
-			clienterrors.ErrInvalidModule, templatePath)
-	}
-
-	return module, nil
-}
-
 // rewriteModulePackage rewrites the module's package path to path.
 func rewriteModulePackage(module *ast.Module) error {
 	pathParts := ast.Ref([]*ast.Term{ast.VarTerm(templatePath)})

constraint/pkg/client/drivers/local/rego.go

@@ -1,5 +1,7 @@
 package local
 
+import "github.com/open-policy-agent/opa/ast"
+
 const (
 	// templatePath is the path the Template's Rego code is stored.
 	// Must match the "data.xxx.violation[r]" path in hookModule.
@@ -16,7 +18,7 @@ const (
 	// This removes boilerplate that would otherwise need to be present in every
 	// Template's Rego code. The violation's response is written to a standard
 	// location we can read from to see if any violations occurred.
-	hookModule = `
+	hookModuleRego = `
 package hooks
 
 # Determine if the object under review violates any passed Constraints.
@@ -45,3 +47,13 @@ violation[response] {
 
 `
 )
+
+var hookModule *ast.Module
+
+func init() {
+	var err error
+	hookModule, err = parseModule(hookModulePath, hookModuleRego)
+	if err != nil {
+		panic(err)
+	}
+}

constraint/pkg/client/drivers/to_result.go

@@ -37,47 +37,42 @@ func ToResults(constraints map[ConstraintKey]*unstructured.Unstructured, resultS
 func ToResult(constraints map[ConstraintKey]*unstructured.Unstructured, r rego.Result) (*types.Result, error) {
 	result := &types.Result{}
 
-	resultMapBinding, found := r.Bindings["result"]
+	resultMap, found, err := unstructured.NestedMap(r.Bindings, "result")
+	if err != nil {
+		return nil, fmt.Errorf("extracting result binding: %v", err)
+	}
+
 	if !found {
 		return nil, errors.New("no binding for result")
 	}
 
-	resultMap, ok := resultMapBinding.(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("result binding was %T but want %T",
-			resultMapBinding, map[string]interface{}{})
+	message, found, err := unstructured.NestedString(resultMap, "msg")
+	if err != nil {
+		return nil, fmt.Errorf("extracting message binding: %v", err)
 	}
 
-	messageBinding, found := resultMap["msg"]
 	if !found {
 		return nil, errors.New("no binding for msg")
 	}
 
-	message, ok := messageBinding.(string)
-	if !ok {
-		return nil, fmt.Errorf("message binding was %T but want %T",
-			messageBinding, "")
-	}
 	result.Msg = message
 
 	result.Metadata = map[string]interface{}{
 		"details": resultMap["details"],
 	}
 
-	keyBinding, found := resultMap["key"]
-	if !found {
-		return nil, errors.New("no binding for Constraint key")
+	keyMap, found, err := unstructured.NestedStringMap(resultMap, "key")
+	if err != nil {
+		return nil, fmt.Errorf("extracting key binding: %v", err)
 	}
 
-	keyMap, ok := keyBinding.(map[string]interface{})
-	if !ok {
-		return nil, fmt.Errorf("key binding was %T but want %T",
-			keyBinding, map[string]interface{}{})
+	if !found {
+		return nil, errors.New("no binding for Constraint key")
 	}
 
 	key := ConstraintKey{
-		Kind: keyMap["kind"].(string),
-		Name: keyMap["name"].(string),
+		Kind: keyMap["kind"],
+		Name: keyMap["name"],
 	}
 
 	constraint := constraints[key]

constraint/pkg/client/e2e_test.go

@@ -207,7 +207,7 @@ func TestClient_Review(t *testing.T) {
 			toReview: handlertest.NewReview("", "foo", "qux"),
 			wantResults: []*types.Result{{
 				Target:            handlertest.TargetName,
-				Msg:               `template0:7: eval_conflict_error: functions must not produce multiple outputs for same inputs`,
+				Msg:               `template:8: eval_conflict_error: functions must not produce multiple outputs for same inputs`,
 				EnforcementAction: constraints.EnforcementActionDeny,
 				Constraint:        cts.MakeConstraint(t, clienttest.KindRuntimeError, "constraint"),
 			}},