Курнат Андрей
132 lines
2.5 KiB
Go
132 lines
2.5 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 ast
|
|
|
|
import (
|
|
"fmt"
|
|
"slices"
|
|
|
|
"github.com/open-policy-agent/opa/v1/util"
|
|
)
|
|
|
|
// VarSet represents a set of variables.
|
|
type VarSet map[Var]struct{ *Location }
|
|
|
|
// NewVarSet returns a new VarSet containing the specified variables.
|
|
func NewVarSet(vs ...Var) VarSet {
|
|
s := make(VarSet, len(vs))
|
|
for _, v := range vs {
|
|
s.Add(v)
|
|
}
|
|
return s
|
|
}
|
|
|
|
// NewVarSet returns a new VarSet containing the specified variables.
|
|
func NewVarSetOfSize(size int) VarSet {
|
|
return make(VarSet, size)
|
|
}
|
|
|
|
// Add updates the set to include the variable "v".
|
|
func (s VarSet) Add(v Var) {
|
|
if _, ok := s[v]; !ok {
|
|
s[v] = struct{ *Location }{}
|
|
}
|
|
}
|
|
|
|
func (s VarSet) AddLocation(v Var, l *Location) {
|
|
if entry, ok := s[v]; ok {
|
|
entry.Location = l
|
|
s[v] = entry
|
|
}
|
|
}
|
|
|
|
// Contains returns true if the set contains the variable "v".
|
|
func (s VarSet) Contains(v Var) bool {
|
|
_, ok := s[v]
|
|
return ok
|
|
}
|
|
|
|
// Copy returns a shallow copy of the VarSet.
|
|
func (s VarSet) Copy() VarSet {
|
|
cpy := NewVarSetOfSize(len(s))
|
|
for v := range s {
|
|
cpy.Add(v)
|
|
}
|
|
return cpy
|
|
}
|
|
|
|
// Diff returns a VarSet containing variables in s that are not in vs.
|
|
func (s VarSet) Diff(vs VarSet) VarSet {
|
|
r := NewVarSetOfSize(s.DiffCount(vs))
|
|
for v := range s {
|
|
if !vs.Contains(v) {
|
|
r.Add(v)
|
|
r.AddLocation(v, s[v].Location)
|
|
}
|
|
}
|
|
return r
|
|
}
|
|
|
|
// DiffCount returns the number of variables in s that are not in vs.
|
|
func (s VarSet) DiffCount(vs VarSet) (i int) {
|
|
for v := range s {
|
|
if !vs.Contains(v) {
|
|
i++
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
// Equal returns true if s contains exactly the same elements as vs.
|
|
func (s VarSet) Equal(vs VarSet) bool {
|
|
if len(s) != len(vs) {
|
|
return false
|
|
}
|
|
for v := range s {
|
|
if !vs.Contains(v) {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
// Intersect returns a VarSet containing variables in s that are in vs.
|
|
func (s VarSet) Intersect(vs VarSet) VarSet {
|
|
i := 0
|
|
for v := range s {
|
|
if vs.Contains(v) {
|
|
i++
|
|
}
|
|
}
|
|
r := NewVarSetOfSize(i)
|
|
for v := range s {
|
|
if vs.Contains(v) {
|
|
r.Add(v)
|
|
}
|
|
}
|
|
return r
|
|
}
|
|
|
|
// Sorted returns a new sorted slice of vars from s.
|
|
func (s VarSet) Sorted() []Var {
|
|
sorted := make([]Var, 0, len(s))
|
|
for v := range s {
|
|
sorted = append(sorted, v)
|
|
}
|
|
slices.SortFunc(sorted, VarCompare)
|
|
return sorted
|
|
}
|
|
|
|
// Update merges the other VarSet into this VarSet.
|
|
func (s VarSet) Update(vs VarSet) {
|
|
for v := range vs {
|
|
s.Add(v)
|
|
}
|
|
}
|
|
|
|
func (s VarSet) String() string {
|
|
return fmt.Sprintf("%v", util.KeysSorted(s))
|
|
}
|