perf(topdown): optimize json.remove and .filter (#8193)

Optimize `json.remove` and `json.filter` to reduce memory
allocations and copy overhead.

Changes:

- Pre-allocate result collections in `jsonRemove` using the
  source size as an upper bound.
- Refactor `jsonRemove` for Arrays to accumulate elements in
  a pre-allocated Go slice before creating the AST Array,
  avoiding intermediate allocations.
- Pre-allocate the root object in `pathsToObject` used by both
  built-ins.
- Optimize `ast.Object.Filter` (used by `json.filter`) for
  Arrays to avoid incremental `Append` allocations.

Signed-off-by: Ville Vesilehto <ville@vesilehto.fi>

v1/ast/term.go

@@ -2570,7 +2570,7 @@ func filterObject(o Value, filter Value) (Value, error) {
 	case String, Number, Boolean, Null:
 		return o, nil
 	case *Array:
-		values := NewArray()
+		values := make([]*Term, 0, v.Len())
 		for i := range v.Len() {
 			subFilter := filteredObj.Get(InternedIntegerString(i))
 			if subFilter != nil {
@@ -2578,10 +2578,10 @@ func filterObject(o Value, filter Value) (Value, error) {
 				if err != nil {
 					return nil, err
 				}
-				values = values.Append(NewTerm(filteredValue))
+				values = append(values, NewTerm(filteredValue))
 			}
 		}
-		return values, nil
+		return NewArray(values...), nil
 	case Set:
 		terms := make([]*Term, 0, v.Len())
 		for _, t := range v.Slice() {

v1/topdown/json.go

@@ -65,7 +65,7 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) {
 	case ast.String, ast.Number, ast.Boolean, ast.Null:
 		return a, nil
 	case ast.Object:
-		newObj := ast.NewObject()
+		newObj := ast.NewObjectWithCapacity(aValue.Len())
 		err := aValue.Iter(func(k *ast.Term, v *ast.Term) error {
 			// recurse and add the diff of sub objects as needed
 			diffValue, err := jsonRemove(v, bObj.Get(k))
@@ -80,7 +80,7 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) {
 		}
 		return ast.NewTerm(newObj), nil
 	case ast.Set:
-		newSet := ast.NewSet()
+		newSet := ast.NewSetWithCapacity(aValue.Len())
 		err := aValue.Iter(func(v *ast.Term) error {
 			// recurse and add the diff of sub objects as needed
 			diffValue, err := jsonRemove(v, bObj.Get(v))
@@ -97,7 +97,7 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) {
 	case *ast.Array:
 		// When indexes are removed we shift left to close empty spots in the array
 		// as per the JSON patch spec.
-		newArray := ast.NewArray()
+		newArraySlice := make([]*ast.Term, 0, aValue.Len())
 		for i := range aValue.Len() {
 			v := aValue.Elem(i)
 			// recurse and add the diff of sub objects as needed
@@ -107,10 +107,10 @@ func jsonRemove(a *ast.Term, b *ast.Term) (*ast.Term, error) {
 				return nil, err
 			}
 			if diffValue != nil {
-				newArray = newArray.Append(diffValue)
+				newArraySlice = append(newArraySlice, diffValue)
 			}
 		}
-		return ast.NewTerm(newArray), nil
+		return ast.NewTerm(ast.NewArray(newArraySlice...)), nil
 	default:
 		return nil, fmt.Errorf("invalid value type %T", a)
 	}
@@ -196,7 +196,7 @@ func parsePath(path *ast.Term) (ast.Ref, error) {
 }
 
 func pathsToObject(paths []ast.Ref) ast.Object {
-	root := ast.NewObject()
+	root := ast.NewObjectWithCapacity(len(paths))
 
 	for _, path := range paths {
 		node := root

v1/topdown/json_bench_test.go

@@ -17,6 +17,130 @@ import (
 	inmem "github.com/open-policy-agent/opa/v1/storage/inmem/test"
 )
 
+func BenchmarkJSONRemoveArray(b *testing.B) {
+	sizes := []int{10, 100, 1000, 5000}
+
+	for _, n := range sizes {
+		b.Run(fmt.Sprintf("size=%d", n), func(b *testing.B) {
+			// Create an object wrapping the array: {"a": [0, 1, ...]}
+			terms := make([]*ast.Term, n)
+			for i := range n {
+				terms[i] = ast.IntNumberTerm(i)
+			}
+			arr := ast.NewArray(terms...)
+			obj := ast.NewObject([2]*ast.Term{ast.StringTerm("a"), ast.NewTerm(arr)})
+
+			// Remove something inside the array to force traversal
+			paths := ast.NewSet(ast.StringTerm("a/nonexistent"))
+
+			operands := []*ast.Term{
+				ast.NewTerm(obj),
+				ast.NewTerm(paths),
+			}
+
+			for b.Loop() {
+				if err := builtinJSONRemove(
+					BuiltinContext{
+						Context: context.Background(),
+					},
+					operands,
+					func(*ast.Term) error {
+						return nil
+					},
+				); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkJSONFilterArray(b *testing.B) {
+	sizes := []int{10, 100, 1000, 5000}
+
+	for _, n := range sizes {
+		b.Run(fmt.Sprintf("size=%d", n), func(b *testing.B) {
+			// Create an object with n keys
+			obj := ast.NewObject()
+			pathSlice := make([]*ast.Term, n)
+			for i := range n {
+				k := ast.StringTerm(fmt.Sprintf("k%d", i))
+				obj.Insert(k, ast.IntNumberTerm(i))
+				pathSlice[i] = k
+			}
+			// Filter all keys: json.filter(obj, ["k0", "k1", ...])
+			// This stresses pathsToObject (creating the filter mask)
+			paths := ast.NewSet(pathSlice...)
+
+			operands := []*ast.Term{
+				ast.NewTerm(obj),
+				ast.NewTerm(paths),
+			}
+
+			for b.Loop() {
+				if err := builtinJSONFilter(
+					BuiltinContext{
+						Context: context.Background(),
+					},
+					operands,
+					func(*ast.Term) error {
+						return nil
+					},
+				); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkJSONFilterArrayIndices(b *testing.B) {
+	sizes := []int{10, 100, 1000, 5000}
+
+	for _, n := range sizes {
+		b.Run(fmt.Sprintf("size=%d", n), func(b *testing.B) {
+			// Create an object wrapping an array: {"a": [0, 1, ...]}
+			terms := make([]*ast.Term, n)
+			for i := range n {
+				terms[i] = ast.IntNumberTerm(i)
+			}
+			arr := ast.NewArray(terms...)
+			obj := ast.NewObject([2]*ast.Term{ast.StringTerm("a"), ast.NewTerm(arr)})
+
+			// Filter to keep the first half of the array elements
+			// json.filter(obj, ["a/0", "a/1", ... "a/n/2"])
+			filterSize := n / 2
+			if filterSize == 0 {
+				filterSize = 1
+			}
+			pathSlice := make([]*ast.Term, filterSize)
+			for i := range filterSize {
+				pathSlice[i] = ast.StringTerm(fmt.Sprintf("a/%d", i))
+			}
+			paths := ast.NewSet(pathSlice...)
+
+			operands := []*ast.Term{
+				ast.NewTerm(obj),
+				ast.NewTerm(paths),
+			}
+
+			for b.Loop() {
+				if err := builtinJSONFilter(
+					BuiltinContext{
+						Context: context.Background(),
+					},
+					operands,
+					func(*ast.Term) error {
+						return nil
+					},
+				); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
 func BenchmarkJSONPatchAddShallowScalar(b *testing.B) {
 	ctx := b.Context()