mirror of
https://github.com/open-policy-agent/opa.git
synced 2026-08-12 19:32:48 -06:00
Ville Vesilehto
7f82c0b3b9
Add NewSetWithCapacity to ast package, exposing pre-allocation capability for external packages. Use it in builtinSetUnion to pre-allocate the result set by first counting total elements across all input sets, avoiding incremental reallocations as the result grows. Signed-off-by: Ville Vesilehto <ville@vesilehto.fi>
107 lines
2.6 KiB
Go
107 lines
2.6 KiB
Go
// Copyright 2016 The OPA Authors. All rights reserved.
|
|
// Use of this source code is governed by an Apache2
|
|
// license that can be found in the LICENSE file.
|
|
|
|
package topdown
|
|
|
|
import (
|
|
"github.com/open-policy-agent/opa/v1/ast"
|
|
"github.com/open-policy-agent/opa/v1/topdown/builtins"
|
|
)
|
|
|
|
// Deprecated: deprecated in v0.4.2 in favour of minus/infix "-" operation.
|
|
func builtinSetDiff(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
|
|
|
|
s1, err := builtins.SetOperand(operands[0].Value, 1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
s2, err := builtins.SetOperand(operands[1].Value, 2)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return iter(ast.NewTerm(s1.Diff(s2)))
|
|
}
|
|
|
|
// builtinSetIntersection returns the intersection of the given input sets
|
|
func builtinSetIntersection(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
|
|
|
|
inputSet, err := builtins.SetOperand(operands[0].Value, 1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// empty input set
|
|
if inputSet.Len() == 0 {
|
|
return iter(ast.InternedEmptySet)
|
|
}
|
|
|
|
var result ast.Set
|
|
|
|
err = inputSet.Iter(func(x *ast.Term) error {
|
|
n, err := builtins.SetOperand(x.Value, 1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if result == nil {
|
|
result = n
|
|
} else {
|
|
result = result.Intersect(n)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return iter(ast.NewTerm(result))
|
|
}
|
|
|
|
// builtinSetUnion returns the union of the given input sets
|
|
func builtinSetUnion(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
|
|
// The set union logic here is manually inlined on purpose. By lifting
|
|
// this logic up a level and not doing pairwise set unions, we avoid
|
|
// many heap allocations. We also pre-allocate the result set by first
|
|
// counting total elements across all input sets.
|
|
inputSet, err := builtins.SetOperand(operands[0].Value, 1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// First pass: count total elements for pre-allocation
|
|
totalSize := 0
|
|
err = inputSet.Iter(func(x *ast.Term) error {
|
|
item, err := builtins.SetOperand(x.Value, 1)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
totalSize += item.Len()
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Pre-allocate result set with estimated capacity
|
|
result := ast.NewSetWithCapacity(totalSize)
|
|
|
|
err = inputSet.Iter(func(x *ast.Term) error {
|
|
item, _ := builtins.SetOperand(x.Value, 1) // error checked above
|
|
item.Foreach(result.Add)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return iter(ast.NewTerm(result))
|
|
}
|
|
|
|
func init() {
|
|
RegisterBuiltinFunc(ast.SetDiff.Name, builtinSetDiff)
|
|
RegisterBuiltinFunc(ast.Intersection.Name, builtinSetIntersection)
|
|
RegisterBuiltinFunc(ast.Union.Name, builtinSetUnion)
|
|
}
|