perf(topdown): pre-alloc in various builtins

This commit introduces ast.NewArrayWithCapacity and leverages it,
along with existing NewSetWithCapacity and NewObjectWithCapacity,
to reduce memory allocations in several hot paths within v1/topdown.

These changes avoid unnecessary map and slice resizing when the size
of the resulting collection is known in advance.

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

v1/ast/term.go

@@ -1413,6 +1413,15 @@ func NewArray(a ...*Term) *Array {
 	return arr
 }
 
+// NewArrayWithCapacity returns a new empty Array with the given capacity pre-allocated.
+func NewArrayWithCapacity(capacity int) *Array {
+	return &Array{
+		elems:  make([]*Term, 0, capacity),
+		hashs:  make([]int, 0, capacity),
+		ground: true,
+	}
+}
+
 // Array represents an array as defined by the language. Arrays are similar to the
 // same types as defined by JSON with the exception that they can contain Vars
 // and References.

v1/topdown/cidr.go

@@ -300,7 +300,7 @@ func builtinNetCIDRMerge(_ BuiltinContext, operands []*ast.Term, iter func(*ast.
 
 	merged := evalNetCIDRMerge(networks)
 
-	result := ast.NewSet()
+	result := ast.NewSetWithCapacity(len(merged))
 	for _, network := range merged {
 		result.Add(ast.StringTerm(network.String()))
 	}

v1/topdown/encoding.go

@@ -287,7 +287,7 @@ func builtinURLQueryDecodeObject(_ BuiltinContext, operands []*ast.Term, iter fu
 		return err
 	}
 
-	queryObject := ast.NewObject()
+	queryObject := ast.NewObjectWithCapacity(len(queryParams))
 	for k, v := range queryParams {
 		paramsArray := make([]*ast.Term, len(v))
 		for i, param := range v {

v1/topdown/eval.go

@@ -1492,20 +1492,20 @@ func (e *eval) amendComprehension(a *ast.Term, b1 *bindings) (*ast.Term, error)
 }
 
 func (e *eval) biunifyComprehensionArray(x *ast.ArrayComprehension, b *ast.Term, b1, b2 *bindings, iter unifyIterator) error {
-	result := ast.NewArray()
+	var elements []*ast.Term
 	child := evalPool.Get()
 
 	e.closure(x.Body, child)
 	defer evalPool.Put(child)
 
 	err := child.Run(func(child *eval) error {
-		result = result.Append(child.bindings.Plug(x.Term))
+		elements = append(elements, child.bindings.Plug(x.Term))
 		return nil
 	})
 	if err != nil {
 		return err
 	}
-	return e.biunify(ast.NewTerm(result), b, b1, b2, iter)
+	return e.biunify(ast.NewTerm(ast.NewArray(elements...)), b, b1, b2, iter)
 }
 
 func (e *eval) biunifyComprehensionSet(x *ast.SetComprehension, b *ast.Term, b1, b2 *bindings, iter unifyIterator) error {

v1/topdown/graphql.go

@@ -146,7 +146,7 @@ func pruneIrrelevantGraphQLASTNodes(value ast.Value) ast.Value {
 	// extant ast type!
 	switch x := value.(type) {
 	case *ast.Array:
-		result := ast.NewArray()
+		result := ast.NewArrayWithCapacity(x.Len())
 		// Iterate over the array's elements, and do the following:
 		// - Drop any Nulls
 		// - Drop any any empty object/array value (after running the pruner)
@@ -173,7 +173,7 @@ func pruneIrrelevantGraphQLASTNodes(value ast.Value) ast.Value {
 		}
 		return result
 	case ast.Object:
-		result := ast.NewObject()
+		result := ast.NewObjectWithCapacity(x.Len())
 		// Iterate over our object's keys, and do the following:
 		// - Drop "Position".
 		// - Drop any key with a Null value.

v1/topdown/jsonschema.go

@@ -110,7 +110,7 @@ func builtinJSONMatchSchema(bctx BuiltinContext, operands []*ast.Term, iter func
 	}
 
 	// In case of validation errors produce Rego array of objects to describe the errors.
-	arr := ast.NewArray()
+	arr := ast.NewArrayWithCapacity(len(result.Errors()))
 	for _, re := range result.Errors() {
 		o := ast.NewObject(
 			[...]*ast.Term{ast.StringTerm("error"), ast.StringTerm(re.String())},

v1/topdown/net.go

@@ -49,7 +49,7 @@ func builtinLookupIPAddr(bctx BuiltinContext, operands []*ast.Term, iter func(*a
 		return err
 	}
 
-	ret := ast.NewSet()
+	ret := ast.NewSetWithCapacity(len(addrs))
 	for _, a := range addrs {
 		ret.Add(ast.StringTerm(a.String()))
 

v1/topdown/object.go

@@ -175,7 +175,7 @@ func builtinObjectKeys(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Te
 func getObjectKeysParam(arrayOrSet ast.Value) (ast.Set, error) {
 	switch v := arrayOrSet.(type) {
 	case *ast.Array:
-		keys := ast.NewSet()
+		keys := ast.NewSetWithCapacity(v.Len())
 		v.Foreach(keys.Add)
 		return keys, nil
 	case ast.Set:

v1/topdown/resolver.go

@@ -99,7 +99,7 @@ func (t *resolverTrie) mktree(e *eval, in resolver.Input) (ast.Value, error) {
 		}
 		return result.Value, nil
 	}
-	obj := ast.NewObject()
+	obj := ast.NewObjectWithCapacity(len(t.children))
 	for k, child := range t.children {
 		v, err := child.mktree(e, resolver.Input{Ref: append(in.Ref, ast.NewTerm(k)), Input: in.Input, Metrics: in.Metrics})
 		if err != nil {