hongming
484
vendor/github.com/open-policy-agent/opa/topdown/copypropagation/copypropagation.go
generated
vendored
Normal file
484
vendor/github.com/open-policy-agent/opa/topdown/copypropagation/copypropagation.go
generated
vendored
Normal file
@@ -0,0 +1,484 @@
|
||||
// Copyright 2018 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 copypropagation
|
||||
|
||||
import (
|
||||
"sort"
|
||||
|
||||
"github.com/open-policy-agent/opa/ast"
|
||||
)
|
||||
|
||||
// CopyPropagator implements a simple copy propagation optimization to remove
|
||||
// intermediate variables in partial evaluation results.
|
||||
//
|
||||
// For example, given the query: input.x > 1 where 'input' is unknown, the
|
||||
// compiled query would become input.x = a; a > 1 which would remain in the
|
||||
// partial evaluation result. The CopyPropagator will remove the variable
|
||||
// assignment so that partial evaluation simply outputs input.x > 1.
|
||||
//
|
||||
// In many cases, copy propagation can remove all variables from the result of
|
||||
// partial evaluation which simplifies evaluation for non-OPA consumers.
|
||||
//
|
||||
// In some cases, copy propagation cannot remove all variables. If the output of
|
||||
// a built-in call is subsequently used as a ref head, the output variable must
|
||||
// be kept. For example. sort(input, x); x[0] == 1. In this case, copy
|
||||
// propagation cannot replace x[0] == 1 with sort(input, x)[0] == 1 as this is
|
||||
// not legal.
|
||||
type CopyPropagator struct {
|
||||
livevars ast.VarSet // vars that must be preserved in the resulting query
|
||||
sorted []ast.Var // sorted copy of vars to ensure deterministic result
|
||||
ensureNonEmptyBody bool
|
||||
}
|
||||
|
||||
// New returns a new CopyPropagator that optimizes queries while preserving vars
|
||||
// in the livevars set.
|
||||
func New(livevars ast.VarSet) *CopyPropagator {
|
||||
|
||||
sorted := make([]ast.Var, 0, len(livevars))
|
||||
for v := range livevars {
|
||||
sorted = append(sorted, v)
|
||||
}
|
||||
|
||||
sort.Slice(sorted, func(i, j int) bool {
|
||||
return sorted[i].Compare(sorted[j]) < 0
|
||||
})
|
||||
|
||||
return &CopyPropagator{livevars: livevars, sorted: sorted}
|
||||
}
|
||||
|
||||
// WithEnsureNonEmptyBody configures p to ensure that results are always non-empty.
|
||||
func (p *CopyPropagator) WithEnsureNonEmptyBody(yes bool) *CopyPropagator {
|
||||
p.ensureNonEmptyBody = yes
|
||||
return p
|
||||
}
|
||||
|
||||
// Apply executes the copy propagation optimization and returns a new query.
|
||||
func (p *CopyPropagator) Apply(query ast.Body) (result ast.Body) {
|
||||
|
||||
uf, ok := makeDisjointSets(p.livevars, query)
|
||||
if !ok {
|
||||
return query
|
||||
}
|
||||
|
||||
// Compute set of vars that appear in the head of refs in the query. If a var
|
||||
// is dereferenced, we cannot plug it with a constant value so the constant on
|
||||
// the union-find root must be unset (e.g., [1][0] is not legal.)
|
||||
headvars := ast.NewVarSet()
|
||||
ast.WalkRefs(query, func(x ast.Ref) bool {
|
||||
if v, ok := x[0].Value.(ast.Var); ok {
|
||||
if root, ok := uf.Find(v); ok {
|
||||
root.constant = nil
|
||||
headvars.Add(root.key)
|
||||
} else {
|
||||
headvars.Add(v)
|
||||
}
|
||||
}
|
||||
return false
|
||||
})
|
||||
|
||||
bindings := map[ast.Var]*binding{}
|
||||
|
||||
for _, expr := range query {
|
||||
|
||||
pctx := &plugContext{
|
||||
bindings: bindings,
|
||||
uf: uf,
|
||||
negated: expr.Negated,
|
||||
headvars: headvars,
|
||||
}
|
||||
|
||||
if expr, keep := p.plugBindings(pctx, expr); keep {
|
||||
if p.updateBindings(pctx, expr) {
|
||||
result.Append(expr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Run post-processing step on the query to ensure that all live vars are bound
|
||||
// in the result. The plugging that happens above substitutes all vars in the
|
||||
// same set with the root.
|
||||
//
|
||||
// This step should run before the next step to prevent unnecessary bindings
|
||||
// from being added to the result. For example:
|
||||
//
|
||||
// - Given the following result: <empty>
|
||||
// - Given the following bindings: x/input.x and y/input
|
||||
// - Given the following liveset: {x}
|
||||
//
|
||||
// If this step were to run AFTER the following step, the output would be:
|
||||
//
|
||||
// x = input.x; y = input
|
||||
//
|
||||
// Even though y = input is not required.
|
||||
for _, v := range p.sorted {
|
||||
if root, ok := uf.Find(v); ok {
|
||||
if root.constant != nil {
|
||||
result.Append(ast.Equality.Expr(ast.NewTerm(v), root.constant))
|
||||
} else if b, ok := bindings[root.key]; ok {
|
||||
result.Append(ast.Equality.Expr(ast.NewTerm(v), ast.NewTerm(b.v)))
|
||||
} else if root.key != v {
|
||||
result.Append(ast.Equality.Expr(ast.NewTerm(v), ast.NewTerm(root.key)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Run post-processing step on query to ensure that all killed exprs are
|
||||
// accounted for. If an expr is killed but the binding is never used, the query
|
||||
// must still include the expr. For example, given the query 'input.x = a' and
|
||||
// an empty livevar set, the result must include the ref input.x otherwise the
|
||||
// query could be satisfied without input.x being defined. When exprs are
|
||||
// killed we initialize the binding counter to zero and then increment it each
|
||||
// time the binding is substituted. if the binding was never substituted it
|
||||
// means the binding value must be added back into the query.
|
||||
for _, b := range sortbindings(bindings) {
|
||||
if !b.containedIn(result) {
|
||||
result.Append(ast.Equality.Expr(ast.NewTerm(b.k), ast.NewTerm(b.v)))
|
||||
}
|
||||
}
|
||||
|
||||
if p.ensureNonEmptyBody && len(result) == 0 {
|
||||
result = append(result, ast.NewExpr(ast.BooleanTerm(true)))
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// plugBindings applies the binding list and union-find to x. This process
|
||||
// removes as many variables as possible.
|
||||
func (p *CopyPropagator) plugBindings(pctx *plugContext, expr *ast.Expr) (*ast.Expr, bool) {
|
||||
|
||||
// Kill single term expressions that are in the binding list. They will be
|
||||
// re-added during post-processing if needed.
|
||||
if term, ok := expr.Terms.(*ast.Term); ok {
|
||||
if v, ok := term.Value.(ast.Var); ok {
|
||||
if root, ok := pctx.uf.Find(v); ok {
|
||||
if _, ok := pctx.bindings[root.key]; ok {
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
xform := bindingPlugTransform{
|
||||
pctx: pctx,
|
||||
}
|
||||
|
||||
// Deep copy the expression as it may be mutated during the transform and
|
||||
// the caller running copy propagation may have references to the
|
||||
// expression. Note, the transform does not contain any error paths and
|
||||
// should never return a non-expression value for the root so consider
|
||||
// errors unreachable.
|
||||
x, err := ast.Transform(xform, expr.Copy())
|
||||
|
||||
if expr, ok := x.(*ast.Expr); !ok || err != nil {
|
||||
panic("unreachable")
|
||||
} else {
|
||||
return expr, true
|
||||
}
|
||||
}
|
||||
|
||||
type bindingPlugTransform struct {
|
||||
pctx *plugContext
|
||||
}
|
||||
|
||||
func (t bindingPlugTransform) Transform(x interface{}) (interface{}, error) {
|
||||
switch x := x.(type) {
|
||||
case ast.Var:
|
||||
return t.plugBindingsVar(t.pctx, x), nil
|
||||
case ast.Ref:
|
||||
return t.plugBindingsRef(t.pctx, x), nil
|
||||
default:
|
||||
return x, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (t bindingPlugTransform) plugBindingsVar(pctx *plugContext, v ast.Var) (result ast.Value) {
|
||||
|
||||
result = v
|
||||
|
||||
// Apply union-find to remove redundant variables from input.
|
||||
if root, ok := pctx.uf.Find(v); ok {
|
||||
result = root.Value()
|
||||
}
|
||||
|
||||
// Apply binding list to substitute remaining vars.
|
||||
if v, ok := result.(ast.Var); ok {
|
||||
if b, ok := pctx.bindings[v]; ok {
|
||||
if !pctx.negated || b.v.IsGround() {
|
||||
result = b.v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (t bindingPlugTransform) plugBindingsRef(pctx *plugContext, v ast.Ref) ast.Ref {
|
||||
|
||||
// Apply union-find to remove redundant variables from input.
|
||||
if root, ok := pctx.uf.Find(v[0].Value.(ast.Var)); ok {
|
||||
v[0].Value = root.Value()
|
||||
}
|
||||
|
||||
result := v
|
||||
|
||||
// Refs require special handling. If the head of the ref was killed, then
|
||||
// the rest of the ref must be concatenated with the new base.
|
||||
//
|
||||
// Invariant: ref heads can only be replaced by refs (not calls).
|
||||
if b, ok := pctx.bindings[v[0].Value.(ast.Var)]; ok {
|
||||
if !pctx.negated || b.v.IsGround() {
|
||||
result = b.v.(ast.Ref).Concat(v[1:])
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// updateBindings returns false if the expression can be killed. If the
|
||||
// expression is killed, the binding list is updated to map a var to value.
|
||||
func (p *CopyPropagator) updateBindings(pctx *plugContext, expr *ast.Expr) bool {
|
||||
if pctx.negated || len(expr.With) > 0 {
|
||||
return true
|
||||
}
|
||||
if expr.IsEquality() {
|
||||
a, b := expr.Operand(0), expr.Operand(1)
|
||||
if a.Equal(b) {
|
||||
return false
|
||||
}
|
||||
k, v, keep := p.updateBindingsEq(a, b)
|
||||
if !keep {
|
||||
if v != nil {
|
||||
pctx.bindings[k] = newbinding(k, v)
|
||||
}
|
||||
return false
|
||||
}
|
||||
} else if expr.IsCall() {
|
||||
terms := expr.Terms.([]*ast.Term)
|
||||
output := terms[len(terms)-1]
|
||||
if k, ok := output.Value.(ast.Var); ok && !p.livevars.Contains(k) && !pctx.headvars.Contains(k) {
|
||||
pctx.bindings[k] = newbinding(k, ast.CallTerm(terms[:len(terms)-1]...).Value)
|
||||
return false
|
||||
}
|
||||
}
|
||||
return !isNoop(expr)
|
||||
}
|
||||
|
||||
func (p *CopyPropagator) updateBindingsEq(a, b *ast.Term) (ast.Var, ast.Value, bool) {
|
||||
k, v, keep := p.updateBindingsEqAsymmetric(a, b)
|
||||
if !keep {
|
||||
return k, v, keep
|
||||
}
|
||||
return p.updateBindingsEqAsymmetric(b, a)
|
||||
}
|
||||
|
||||
func (p *CopyPropagator) updateBindingsEqAsymmetric(a, b *ast.Term) (ast.Var, ast.Value, bool) {
|
||||
k, ok := a.Value.(ast.Var)
|
||||
if !ok || p.livevars.Contains(k) {
|
||||
return "", nil, true
|
||||
}
|
||||
|
||||
switch b.Value.(type) {
|
||||
case ast.Ref, ast.Call:
|
||||
return k, b.Value, false
|
||||
}
|
||||
|
||||
return "", nil, true
|
||||
}
|
||||
|
||||
type plugContext struct {
|
||||
bindings map[ast.Var]*binding
|
||||
uf *unionFind
|
||||
headvars ast.VarSet
|
||||
negated bool
|
||||
}
|
||||
|
||||
type binding struct {
|
||||
k ast.Var
|
||||
v ast.Value
|
||||
}
|
||||
|
||||
func newbinding(k ast.Var, v ast.Value) *binding {
|
||||
return &binding{k: k, v: v}
|
||||
}
|
||||
|
||||
func (b *binding) containedIn(query ast.Body) bool {
|
||||
var stop bool
|
||||
switch v := b.v.(type) {
|
||||
case ast.Ref:
|
||||
ast.WalkRefs(query, func(other ast.Ref) bool {
|
||||
if stop || other.HasPrefix(v) {
|
||||
stop = true
|
||||
return stop
|
||||
}
|
||||
return false
|
||||
})
|
||||
default:
|
||||
ast.WalkTerms(query, func(other *ast.Term) bool {
|
||||
if stop || other.Value.Compare(v) == 0 {
|
||||
stop = true
|
||||
return stop
|
||||
}
|
||||
return false
|
||||
})
|
||||
}
|
||||
return stop
|
||||
}
|
||||
|
||||
func sortbindings(bindings map[ast.Var]*binding) []*binding {
|
||||
sorted := make([]*binding, 0, len(bindings))
|
||||
for _, b := range bindings {
|
||||
sorted = append(sorted, b)
|
||||
}
|
||||
sort.Slice(sorted, func(i, j int) bool {
|
||||
return sorted[i].k.Compare(sorted[j].k) < 0
|
||||
})
|
||||
return sorted
|
||||
}
|
||||
|
||||
type unionFind struct {
|
||||
roots map[ast.Var]*unionFindRoot
|
||||
parents map[ast.Var]ast.Var
|
||||
rank rankFunc
|
||||
}
|
||||
|
||||
// makeDisjointSets builds the union-find structure for the query. The structure
|
||||
// is built by processing all of the equality exprs in the query. Sets represent
|
||||
// vars that must be equal to each other. In addition to vars, each set can have
|
||||
// at most one constant. If the query contains expressions that cannot be
|
||||
// satisfied (e.g., because a set has multiple constants) this function returns
|
||||
// false.
|
||||
func makeDisjointSets(livevars ast.VarSet, query ast.Body) (*unionFind, bool) {
|
||||
uf := newUnionFind(func(r1, r2 *unionFindRoot) (*unionFindRoot, *unionFindRoot) {
|
||||
if livevars.Contains(r1.key) {
|
||||
return r1, r2
|
||||
}
|
||||
return r2, r1
|
||||
})
|
||||
for _, expr := range query {
|
||||
if expr.IsEquality() && !expr.Negated && len(expr.With) == 0 {
|
||||
a, b := expr.Operand(0), expr.Operand(1)
|
||||
varA, ok1 := a.Value.(ast.Var)
|
||||
varB, ok2 := b.Value.(ast.Var)
|
||||
if ok1 && ok2 {
|
||||
if _, ok := uf.Merge(varA, varB); !ok {
|
||||
return nil, false
|
||||
}
|
||||
} else if ok1 && ast.IsConstant(b.Value) {
|
||||
root := uf.MakeSet(varA)
|
||||
if root.constant != nil && !root.constant.Equal(b) {
|
||||
return nil, false
|
||||
}
|
||||
root.constant = b
|
||||
} else if ok2 && ast.IsConstant(a.Value) {
|
||||
root := uf.MakeSet(varB)
|
||||
if root.constant != nil && !root.constant.Equal(a) {
|
||||
return nil, false
|
||||
}
|
||||
root.constant = a
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return uf, true
|
||||
}
|
||||
|
||||
type rankFunc func(*unionFindRoot, *unionFindRoot) (*unionFindRoot, *unionFindRoot)
|
||||
|
||||
func newUnionFind(rank rankFunc) *unionFind {
|
||||
return &unionFind{
|
||||
roots: map[ast.Var]*unionFindRoot{},
|
||||
parents: map[ast.Var]ast.Var{},
|
||||
rank: rank,
|
||||
}
|
||||
}
|
||||
|
||||
func (uf *unionFind) MakeSet(v ast.Var) *unionFindRoot {
|
||||
|
||||
root, ok := uf.Find(v)
|
||||
if ok {
|
||||
return root
|
||||
}
|
||||
|
||||
root = newUnionFindRoot(v)
|
||||
uf.parents[v] = v
|
||||
uf.roots[v] = root
|
||||
return uf.roots[v]
|
||||
}
|
||||
|
||||
func (uf *unionFind) Find(v ast.Var) (*unionFindRoot, bool) {
|
||||
|
||||
parent, ok := uf.parents[v]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
if parent == v {
|
||||
return uf.roots[v], true
|
||||
}
|
||||
|
||||
return uf.Find(parent)
|
||||
}
|
||||
|
||||
func (uf *unionFind) Merge(a, b ast.Var) (*unionFindRoot, bool) {
|
||||
|
||||
r1 := uf.MakeSet(a)
|
||||
r2 := uf.MakeSet(b)
|
||||
|
||||
if r1 != r2 {
|
||||
|
||||
r1, r2 = uf.rank(r1, r2)
|
||||
|
||||
uf.parents[r2.key] = r1.key
|
||||
delete(uf.roots, r2.key)
|
||||
|
||||
// Sets can have at most one constant value associated with them. When
|
||||
// unioning, we must preserve this invariant. If a set has two constants,
|
||||
// there will be no way to prove the query.
|
||||
if r1.constant != nil && r2.constant != nil && !r1.constant.Equal(r2.constant) {
|
||||
return nil, false
|
||||
} else if r1.constant == nil {
|
||||
r1.constant = r2.constant
|
||||
}
|
||||
}
|
||||
|
||||
return r1, true
|
||||
}
|
||||
|
||||
type unionFindRoot struct {
|
||||
key ast.Var
|
||||
constant *ast.Term
|
||||
}
|
||||
|
||||
func newUnionFindRoot(key ast.Var) *unionFindRoot {
|
||||
return &unionFindRoot{
|
||||
key: key,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *unionFindRoot) Value() ast.Value {
|
||||
if r.constant != nil {
|
||||
return r.constant.Value
|
||||
}
|
||||
return r.key
|
||||
}
|
||||
|
||||
func isNoop(expr *ast.Expr) bool {
|
||||
|
||||
if !expr.IsCall() {
|
||||
term := expr.Terms.(*ast.Term)
|
||||
if !ast.IsConstant(term.Value) {
|
||||
return false
|
||||
}
|
||||
return !ast.Boolean(false).Equal(term.Value)
|
||||
}
|
||||
|
||||
// A==A can be ignored
|
||||
if expr.Operator().Equal(ast.Equal.Ref()) {
|
||||
return expr.Operand(0).Equal(expr.Operand(1))
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
Reference in New Issue
Block a user