diff --git a/v1/ast/compile.go b/v1/ast/compile.go index a9455145ef..ae2a728e24 100644 --- a/v1/ast/compile.go +++ b/v1/ast/compile.go @@ -4922,7 +4922,7 @@ func resolveRef(globals map[Var]*usedRef, ignore *declaredVarStack, ref Ref) Ref switch v := x.Value.(type) { case Var: if g, ok := globals[v]; ok && !ignore.Contains(v) { - cpy := g.ref.Copy() + cpy := g.ref.CopyNonGround() for i := range cpy { cpy[i].SetLocation(x.Location) } @@ -5071,7 +5071,7 @@ func resolveRefsInTerm(globals map[Var]*usedRef, ignore *declaredVarStack, term switch v := term.Value.(type) { case Var: if g, ok := globals[v]; ok && !ignore.Contains(v) { - cpy := g.ref.Copy() + cpy := g.ref.CopyNonGround() for i := range cpy { cpy[i].SetLocation(term.Location) } diff --git a/v1/ast/policy.go b/v1/ast/policy.go index 632b5aa6d5..d1c4f17902 100644 --- a/v1/ast/policy.go +++ b/v1/ast/policy.go @@ -989,11 +989,7 @@ func (head *Head) HasDynamicRef() bool { // Copy returns a deep copy of a. func (a Args) Copy() Args { - cpy := make(Args, 0, len(a)) - for _, t := range a { - cpy = append(cpy, t.Copy()) - } - return cpy + return termSliceCopy(a) } func (a Args) String() string { diff --git a/v1/ast/term.go b/v1/ast/term.go index 436b22eb2a..272b412293 100644 --- a/v1/ast/term.go +++ b/v1/ast/term.go @@ -1173,10 +1173,10 @@ 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. +// CopyNonGround returns a new ref with shallow copies of ground parts and deep +// copies of non-ground parts. The head element is always shallow copied. This is +// cheaper than Copy for operations that only modify non-ground parts (e.g. plugging) +// or metadata (e.g. Location). func (ref Ref) CopyNonGround() Ref { cpy := make(Ref, len(ref)) cpy[0] = ref[0] @@ -1729,17 +1729,26 @@ type set struct { // Copy returns a deep copy of s. func (s *set) Copy() Set { + n := len(s.keys) cpy := &set{ hash: s.hash, ground: s.ground, sortGuard: sync.Once{}, - elems: make(map[int]*Term, len(s.elems)), - keys: make([]*Term, 0, len(s.keys)), + elems: make(map[int]*Term, n), + keys: make([]*Term, n), } - for hash := range s.elems { - cpy.elems[hash] = s.elems[hash].Copy() - cpy.keys = append(cpy.keys, cpy.elems[hash]) + if n > 0 { + // Batch-allocate all Term structs in a single contiguous block. + buf := make([]Term, n) + i := 0 + for hash, elem := range s.elems { + buf[i] = *elem + deepCopyTermValue(&buf[i]) + cpy.elems[hash] = &buf[i] + cpy.keys[i] = &buf[i] + i++ + } } return cpy @@ -2364,10 +2373,41 @@ func (obj *object) IsGround() bool { // Copy returns a deep copy of obj. func (obj *object) Copy() Object { - cpy, _ := obj.Map(func(k, v *Term) (*Term, *Term, error) { - return k.Copy(), v.Copy(), nil - }) - cpy.(*object).hash = obj.hash + n := len(obj.keys) + cpy := &object{ + elems: make(map[int]*objectElem, n), + sortGuard: sync.Once{}, + hash: obj.hash, + ground: obj.ground, + } + + if n == 0 { + return cpy + } + + // Batch-allocate all objectElems, keys, and values in contiguous blocks + // (3 allocations instead of 3N). + elems := make([]objectElem, n) + keys := make([]Term, n) + vals := make([]Term, n) + cpy.keys = make([]*objectElem, n) + + for i, srcElem := range obj.keys { + keys[i] = *srcElem.key + deepCopyTermValue(&keys[i]) + vals[i] = *srcElem.value + deepCopyTermValue(&vals[i]) + + elems[i] = objectElem{key: &keys[i], value: &vals[i]} + cpy.keys[i] = &elems[i] + + hash := keys[i].Hash() + if head, ok := cpy.elems[hash]; ok { + elems[i].next = head + } + cpy.elems[hash] = &elems[i] + } + return cpy } @@ -2931,10 +2971,45 @@ func (c Call) String() string { return util.ByteSliceToString(buf) } +// deepCopyTermValue deep copies the Value of term in-place. +// Scalar values (Null, Boolean, Number, String, Var) are already +// copied by struct assignment, so only container types need work. +func deepCopyTermValue(term *Term) { + switch v := term.Value.(type) { + case Null, Boolean, Number, String, Var: + // Already copied by *term = *src struct assignment. + case Ref: + term.Value = v.Copy() + case *Array: + term.Value = v.Copy() + case Set: + term.Value = v.Copy() + case *object: + term.Value = v.Copy() + case *ArrayComprehension: + term.Value = v.Copy() + case *ObjectComprehension: + term.Value = v.Copy() + case *SetComprehension: + term.Value = v.Copy() + case *TemplateString: + term.Value = v.Copy() + case Call: + term.Value = v.Copy() + } +} + func termSliceCopy(a []*Term) []*Term { - cpy := make([]*Term, len(a)) - for i := range a { - cpy[i] = a[i].Copy() + n := len(a) + if n == 0 { + return make([]*Term, 0) + } + // Batch allocate all Term structs in a single contiguous slice (2 allocs + // instead of N+1) using the same pattern as util.NewPtrSlice. + cpy := util.NewPtrSlice[Term](n) + for i := range n { + *cpy[i] = *a[i] // copy Term struct (Value + Location) + deepCopyTermValue(cpy[i]) // deep copy container Values in-place } return cpy } diff --git a/v1/ast/term_bench_test.go b/v1/ast/term_bench_test.go index 6d01bb70bd..f96db317df 100644 --- a/v1/ast/term_bench_test.go +++ b/v1/ast/term_bench_test.go @@ -502,6 +502,67 @@ func BenchmarkObjectCopy(b *testing.B) { } } +func BenchmarkArrayCopy(b *testing.B) { + sizes := []int{5, 50, 500} + for _, n := range sizes { + b.Run(strconv.Itoa(n), func(b *testing.B) { + arr := NewArray() + for i := range n { + arr = arr.Append(IntNumberTerm(i)) + } + b.ResetTimer() + for b.Loop() { + _ = arr.Copy() + } + }) + } +} + +func BenchmarkRefCopy(b *testing.B) { + sizes := []int{5, 10, 20} + for _, n := range sizes { + b.Run(strconv.Itoa(n), func(b *testing.B) { + parts := make([]*Term, n) + parts[0] = VarTerm("data") + for i := range n { + parts[i] = StringTerm("part" + strconv.Itoa(i)) + } + ref := Ref(parts) + b.ResetTimer() + for b.Loop() { + _ = ref.Copy() + } + }) + } +} + +func BenchmarkRefCopyNonGround(b *testing.B) { + b.Run("fully_ground", func(b *testing.B) { + ref := MustParseRef("data.foo.bar.baz.qux") + for b.Loop() { + _ = ref.CopyNonGround() + } + }) + b.Run("mixed", func(b *testing.B) { + ref := MustParseRef("data.foo[x].bar[y]") + for b.Loop() { + _ = ref.CopyNonGround() + } + }) + b.Run("fully_ground/Copy", func(b *testing.B) { + ref := MustParseRef("data.foo.bar.baz.qux") + for b.Loop() { + _ = ref.Copy() + } + }) + b.Run("mixed/Copy", func(b *testing.B) { + ref := MustParseRef("data.foo[x].bar[y]") + for b.Loop() { + _ = ref.Copy() + } + }) +} + func BenchmarkTermHashing(b *testing.B) { sizes := []int{10, 100, 1000} for _, n := range sizes { diff --git a/v1/ast/term_test.go b/v1/ast/term_test.go index ff93964749..47990b7ef5 100644 --- a/v1/ast/term_test.go +++ b/v1/ast/term_test.go @@ -1819,3 +1819,29 @@ func TestAppendNoQuoteExistingBuffer(t *testing.T) { t.Errorf("got %s\nwant %s", got, exp) } } + +func TestRefCopyNonGround(t *testing.T) { + ref := MustParseRef("data.foo[x].bar") + + result := ref.CopyNonGround() + + if &result[0] == &ref[0] { + t.Fatal("CopyNonGround() should return new slice") + } + if !result.Equal(ref) { + t.Fatalf("CopyNonGround() result not equal to original:\ngot: %v\nwant: %v", result, ref) + } + + // Ground elements should be shared (same pointer) + for i := 1; i < len(ref); i++ { + if ref[i].Value.IsGround() { + if result[i] != ref[i] { + t.Errorf("Ground element at index %d should be shared", i) + } + } else { + if result[i] == ref[i] { + t.Errorf("Non-ground element at index %d should be deep copied", i) + } + } + } +} diff --git a/v1/rego/rego_bench_test.go b/v1/rego/rego_bench_test.go index 3eab04e9e4..8475c5cc07 100644 --- a/v1/rego/rego_bench_test.go +++ b/v1/rego/rego_bench_test.go @@ -109,7 +109,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 +// BenchmarkAciTestOnlyEval-10 13521 86647 ns/op 47448 B/op 967 allocs/op // ground-aware Copy // BenchmarkAciTestOnlyEval-12 21007 57551 ns/op 45323 B/op 920 allocs/op func BenchmarkAciTestOnlyEval(b *testing.B) { ctx := b.Context() @@ -357,7 +357,7 @@ func BenchmarkStoreRead(b *testing.B) { } // 5730 ns/op 5737 B/op 93 allocs/op -// 5222 ns/op 5639 B/op 89 allocs/op // ref.CopyNonGround +// 5222 ns/op 5639 B/op 89 allocs/op // ground-aware Copy // 2786 ns/op 5090 B/op 77 allocs/op // Lazy init improvements func BenchmarkTrivialPolicy(b *testing.B) { ctx := b.Context() diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go index 6f93ba530e..6305954ecd 100644 --- a/v1/topdown/eval.go +++ b/v1/topdown/eval.go @@ -1745,7 +1745,7 @@ func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error) // Copy ref here as ref otherwise always escapes to the heap, // whether tracing is enabled or not. - r := ref.Copy() + r := ref.CopyNonGround() e.traceIndex(e.query[e.index], msg.String(), &r) } @@ -3397,7 +3397,7 @@ func (q vcKeyScope) AppendText(buf []byte) ([]byte, error) { // reduce removes vars from the tail of the ref. func (q vcKeyScope) reduce() vcKeyScope { - ref := q.Ref.Copy() + ref := q.Ref.CopyNonGround() var i int for i = len(q.Ref) - 1; i >= 0; i-- { if _, ok := q.Ref[i].Value.(ast.Var); !ok { @@ -4120,7 +4120,7 @@ func (e *eval) comprehensionIndex(term *ast.Term) *ast.ComprehensionIndex { func (e *eval) namespaceRef(ref ast.Ref) ast.Ref { if e.skipSaveNamespace { - return ref.Copy() + return ref.CopyNonGround() } return ref.Insert(e.saveNamespace, 1) }