perf: Add ref.CopyNonGround (#7350)

Refs copied for the purpose of modification often don't need
_all_ parts deep-copied, but only those to be modified. It turns
out that in most cases in the context of eval, **no** parts of
the ref needs copying, as those refs are static. This PR adds
a simple new method to refs that allows smarter copying, saving
a million and a half allocations in the `regal lint` benchmark,
and hundreds to thousands of allocations in common policy
evaluation.

**Before**
```
1716967000 ns/op	3125089752 B/op	62201992 allocs/op
```
**After**
```
1679479541 ns/op	3086159368 B/op	60630066 allocs/op
```

Also added an unrelated benchmark of a trivial policy, as I've
been curious to track down costs of doing almost "nothing".

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

v1/ast/compile.go

@@ -1114,7 +1114,7 @@ func (c *Compiler) checkRuleConflicts() {
 		for _, rule := range node.Values {
 			r := rule.(*Rule)
 			ref := r.Ref()
-			name = rw(ref.Copy()).String() // varRewriter operates in-place
+			name = rw(ref.CopyNonGround()).String() // varRewriter operates in-place
 			kinds[r.Head.RuleKind()] = struct{}{}
 			arities[len(r.Head.Args)] = struct{}{}
 			if r.Default {
@@ -1156,7 +1156,7 @@ func (c *Compiler) checkRuleConflicts() {
 			//   data.p.q[r][s] { r := input.r; s := input.s }
 			//   data.p[q].r.s { q := input.q }
 
-			if r.Ref().IsGround() && len(node.Children) > 0 {
+			if ref.IsGround() && len(node.Children) > 0 {
 				conflicts = node.flattenChildren()
 			}
 

v1/ast/term.go

@@ -1017,6 +1017,25 @@ func (ref Ref) Copy() Ref {
 	return termSliceCopy(ref)
 }
 
+// CopyNonGround returns a new ref with deep copies of the non-ground parts and shallow
+// copies of the ground parts. This is a *much* cheaper operation than Copy for operations
+// that only intend to modify (e.g. plug) the non-ground parts. The head element of the ref
+// is always shallow copied.
+func (ref Ref) CopyNonGround() Ref {
+	cpy := make(Ref, len(ref))
+	cpy[0] = ref[0]
+
+	for i := 1; i < len(ref); i++ {
+		if ref[i].Value.IsGround() {
+			cpy[i] = ref[i]
+		} else {
+			cpy[i] = ref[i].Copy()
+		}
+	}
+
+	return cpy
+}
+
 // Equal returns true if ref is equal to other.
 func (ref Ref) Equal(other Value) bool {
 	switch o := other.(type) {
@@ -3063,14 +3082,10 @@ func (c Call) String() string {
 
 func termSliceCopy(a []*Term) []*Term {
 	cpy := make([]*Term, len(a))
-	termSliceCopyTo(a, cpy)
-	return cpy
-}
-
-func termSliceCopyTo(src, dst []*Term) {
-	for i := range src {
-		dst[i] = src[i].Copy()
+	for i := range a {
+		cpy[i] = a[i].Copy()
 	}
+	return cpy
 }
 
 func termSliceEqual(a, b []*Term) bool {

v1/rego/rego_bench_test.go

@@ -154,6 +154,7 @@ func BenchmarkAciTestBuildAndEval(b *testing.B) {
 }
 
 // BenchmarkAciTestOnlyEval-10    12752    92188 ns/op    50005 B/op    1062 allocs/op
+// BenchmarkAciTestOnlyEval-10    13521	   86647 ns/op	  47448 B/op	 967 allocs/op // ref.CopyNonGround
 func BenchmarkAciTestOnlyEval(b *testing.B) {
 	ctx := context.Background()
 
@@ -321,8 +322,8 @@ func BenchmarkObjectIteration(b *testing.B) {
 
 // Comparing the cost of referencing not found data in Go vs. AST storage
 //
-// BenchmarkStoreRefNotFound/inmem-go-10         	    5208	    212288 ns/op	  160609 B/op	    2936 allocs/op
-// BenchmarkStoreRefNotFound/inmem-ast-10        	   13929	     90053 ns/op	   39614 B/op	    1012 allocs/op
+// BenchmarkStoreRefNotFound/inmem-go-10         5208    212288 ns/op    160609 B/op     2936 allocs/op
+// BenchmarkStoreRefNotFound/inmem-ast-10       13929     90053 ns/op     39614 B/op     1012 allocs/op
 func BenchmarkStoreRefNotFound(b *testing.B) {
 	ctx := context.Background()
 
@@ -415,6 +416,34 @@ func BenchmarkStoreRead(b *testing.B) {
 	}
 }
 
+// 233337	      5730 ns/op	    5737 B/op	      93 allocs/op
+// 229280	      5222 ns/op	    5639 B/op	      89 allocs/op // ref.CopyNonGround
+func BenchmarkTrivialPolicy(b *testing.B) {
+	ctx := context.Background()
+	r := New(
+		ParsedQuery(ast.MustParseBody("data.p.r = x")),
+		ParsedModule(ast.MustParseModule(`package p
+		r := 1`)),
+		GenerateJSON(func(*ast.Term, *EvalContext) (any, error) {
+			return nil, nil
+		}),
+	)
+
+	pq, err := r.PrepareForEval(ctx)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+
+	for range b.N {
+		_, err := pq.Eval(ctx)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
 func mustReadFileAsString(b *testing.B, path string) string {
 	b.Helper()
 

v1/topdown/eval.go

@@ -1185,7 +1185,7 @@ func (e *eval) biunifyRef(a, b *ast.Term, b1, b2 *bindings, iter unifyIterator)
 			e:         e,
 			ref:       ref,
 			pos:       1,
-			plugged:   ref.Copy(),
+			plugged:   ref.CopyNonGround(),
 			bindings:  b1,
 			rterm:     b,
 			rbindings: b2,