eval: optimize iteration (#7327)

Since we do a lot of this in Rego, I was curious to learn what the cost
of that was, and if it was possible to do better. I identified the two
main costs as these:

- `handleErr` wrapper allocating in each iteration, even though the
  normal condition is to not handle an error
- `biunify` callback function literal leaking to the heap in each
  iteration, with an _high_ cost associated for such a simple one-liner
  function. While it wasn't possible to change that (as it referenced
  the outer scope) it *was* possible to inline the whole biunify flow
  for the array iteration case, as that turned out to be quite simple
  when stepped through with a debugger.

Ideally we'd be able to inline iteation over sets and objects too, but
that's much more complex (see comment in code), so I left that for the
future.

The numbers benchmarked are in the PR, so take a look there. The array
iteration case is *quite* the improvement!

For reference, this amounts to about 2 million allocations less in the
`regal lint bundle` benchmark, and a noticeable improvement in
evaluation time.

Signed-off-by: Anders Eknert <anders@styra.com>

v1/rego/rego_bench_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
 	"testing"
 
 	"github.com/open-policy-agent/opa/internal/runtime"
@@ -175,16 +176,148 @@ func BenchmarkAciTestOnlyEval(b *testing.B) {
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
-
 		res, err := pq.Eval(ctx, EvalParsedInput(input.Value))
 		if err != nil {
 			b.Fatal(err)
 		}
-
 		_ = res
 	}
 }
 
+// BenchmarkArrayIteration-10
+// 15574    77121 ns/op    67249 B/op    1115 allocs/op    // handleErr wrapping, not inlined
+// 33862    35864 ns/op     5768 B/op      93 allocs/op    // handleErr only on error, inlined
+func BenchmarkArrayIteration(b *testing.B) {
+	ctx := context.Background()
+
+	at := make([]*ast.Term, 512)
+	for i := 0; i < 511; i++ {
+		at[i] = ast.StringTerm("a")
+	}
+	at[511] = ast.StringTerm("v")
+
+	input := ast.NewObject(ast.Item(ast.StringTerm("foo"), ast.ArrayTerm(at...)))
+	module := ast.MustParseModule(`package test
+
+	default r := false
+
+	r if input.foo[_] == "v"`)
+
+	r := New(Query("data.test.r = x"), ParsedModule(module))
+
+	pq, err := r.PrepareForEval(ctx)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		res, err := pq.Eval(ctx, EvalParsedInput(input))
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		if res == nil {
+			b.Fatal("expected result")
+		}
+
+		if res[0].Bindings["x"].(bool) != true {
+			b.Fatalf("expected true, got %v", res[0].Bindings["x"])
+		}
+	}
+}
+
+// BenchmarkSetIteration-10
+// 4800    272403 ns/op    80875 B/op    1193 allocs/op    // handleErr wrapping, not inlined
+// 4933	   223234 ns/op	   76772 B/op	  681 allocs/op    // handleErr only on error, not inlined
+func BenchmarkSetIteration(b *testing.B) {
+	ctx := context.Background()
+
+	at := make([]*ast.Term, 512)
+	for i := 0; i < 512; i++ {
+		at[i] = ast.StringTerm(strconv.Itoa(i))
+	}
+
+	input := ast.NewObject(ast.Item(ast.StringTerm("foo"), ast.ArrayTerm(at...)))
+	module := ast.MustParseModule(`package test
+
+	s := {x | x := input.foo[_]}
+
+	default r := false
+
+	r if s[_] == "not found"`)
+
+	r := New(Query("data.test.r = x"), ParsedModule(module))
+
+	pq, err := r.PrepareForEval(ctx)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		res, err := pq.Eval(ctx, EvalParsedInput(input))
+		if err != nil {
+			b.Fatal(err)
+		}
+		if res == nil {
+			b.Fatal("expected result")
+		}
+		if res[0].Bindings["x"].(bool) != false {
+			b.Fatalf("expected false, got %v", res[0].Bindings["x"])
+		}
+	}
+}
+
+// BenchmarkObjectIteration-10
+// 12067    99582 ns/op    72830 B/op   1126 allocs/op    // handleErr wrapping, not inlined
+// 15358    85080 ns/op    27752 B/op    615 allocs/op    // handleErr only on error, not inlined
+func BenchmarkObjectIteration(b *testing.B) {
+	ctx := context.Background()
+
+	at := make([][2]*ast.Term, 512)
+	for i := 0; i < 512; i++ {
+		at[i] = ast.Item(ast.StringTerm(strconv.Itoa(i)), ast.StringTerm(strconv.Itoa(i)))
+	}
+
+	input := ast.NewObject(ast.Item(ast.StringTerm("foo"), ast.ObjectTerm(at...)))
+	module := ast.MustParseModule(`package test
+
+	default r := false
+
+	r if {
+		input.foo[_] == "512"
+	}
+		`)
+
+	r := New(Query("data.test.r = x"), ParsedModule(module))
+
+	pq, err := r.PrepareForEval(ctx)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		res, err := pq.Eval(ctx, EvalParsedInput(input))
+		if err != nil {
+			b.Fatal(err)
+		}
+		if res == nil {
+			b.Fatal("expected result")
+		}
+		if res[0].Bindings["x"].(bool) != false {
+			b.Fatalf("expected false, got %v", res[0].Bindings["x"])
+		}
+	}
+}
+
 func mustReadFileAsString(b *testing.B, path string) string {
 	b.Helper()
 

v1/topdown/eval.go

@@ -3650,28 +3650,55 @@ func (e evalTerm) enumerate(iter unifyIterator) error {
 
 	switch v := e.term.Value.(type) {
 	case *ast.Array:
+		// Note(anders):
+		// For this case (e.g. input.foo[_]), we can avoid the (quite expensive) overhead of a callback
+		// function literal escaping to the heap in each iteration by inlining the biunification logic,
+		// meaning a 10x reduction in both the number of allocations made as well as the memory consumed.
+		// It is possible that such inlining could be done for the set/object cases as well, and that's
+		// worth looking into later, as I imagine set iteration in particular would be an even greater
+		// win across most policies. Those cases are however much more complex, as we need to deal with
+		// any type on either side, not just int/var as is the case here.
 		for i := 0; i < v.Len(); i++ {
-			k := ast.InternedIntNumberTerm(i)
-			if err := handleErr(e.e.biunify(k, e.ref[e.pos], e.bindings, e.bindings, func() error {
-				return e.next(iter, k)
-			})); err != nil {
-				return err
+			a := ast.InternedIntNumberTerm(i)
+			b := e.ref[e.pos]
+
+			if _, ok := b.Value.(ast.Var); ok {
+				if e.e.traceEnabled {
+					e.e.traceUnify(a, b)
+				}
+				var undo undo
+				b, e.bindings = e.bindings.apply(b)
+				e.bindings.bind(b, a, e.bindings, &undo)
+
+				err := e.next(iter, a)
+				undo.Undo()
+				if err != nil {
+					if err := handleErr(err); err != nil {
+						return err
+					}
+				}
 			}
 		}
 	case ast.Object:
 		for _, k := range v.Keys() {
-			if err := handleErr(e.e.biunify(k, e.ref[e.pos], e.termbindings, e.bindings, func() error {
+			err := e.e.biunify(k, e.ref[e.pos], e.termbindings, e.bindings, func() error {
 				return e.next(iter, e.termbindings.Plug(k))
-			})); err != nil {
-				return err
+			})
+			if err != nil {
+				if err := handleErr(err); err != nil {
+					return err
+				}
 			}
 		}
 	case ast.Set:
 		for _, elem := range v.Slice() {
-			if err := handleErr(e.e.biunify(elem, e.ref[e.pos], e.termbindings, e.bindings, func() error {
+			err := e.e.biunify(elem, e.ref[e.pos], e.termbindings, e.bindings, func() error {
 				return e.next(iter, e.termbindings.Plug(elem))
-			})); err != nil {
-				return err
+			})
+			if err != nil {
+				if err := handleErr(err); err != nil {
+					return err
+				}
 			}
 		}
 	}