From d198695615ec3854ad9e256a183a9241551df6c8 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Mon, 30 May 2016 18:04:43 -0700 Subject: [PATCH 01/11] Extend Body and Expr with Hash/IsGround This will allow us to treat the Body as term which will be needed for comprehensions. --- ast/policy.go | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/ast/policy.go b/ast/policy.go index e373efc230..6f9dfaff90 100644 --- a/ast/policy.go +++ b/ast/policy.go @@ -219,6 +219,25 @@ func (body Body) Equal(other Body) bool { return true } +// Hash returns the hash code for the Body. +func (body Body) Hash() int { + s := 0 + for _, e := range body { + s += e.Hash() + } + return s +} + +// IsGround returns true if all of the expressions in the Body are ground. +func (body Body) IsGround() bool { + for _, e := range body { + if !e.IsGround() { + return false + } + } + return true +} + func (body Body) String() string { var buf []string for _, v := range body { @@ -264,6 +283,23 @@ func (expr *Expr) Equal(other *Expr) bool { return false } +// Hash returns the hash code of the Expr. +func (expr *Expr) Hash() int { + s := 0 + switch ts := expr.Terms.(type) { + case []*Term: + for _, t := range ts { + s += t.Value.Hash() + } + case *Term: + s += ts.Value.Hash() + } + if expr.Negated { + s++ + } + return s +} + // IsEquality returns true if this is an equality expression. func (expr *Expr) IsEquality() bool { terms, ok := expr.Terms.([]*Term) @@ -276,6 +312,21 @@ func (expr *Expr) IsEquality() bool { return terms[0].Equal(VarTerm("=")) } +// IsGround returns true if all of the expression terms are ground. +func (expr *Expr) IsGround() bool { + switch ts := expr.Terms.(type) { + case []*Term: + for _, t := range ts[1:] { + if !t.IsGround() { + return false + } + } + case *Term: + return ts.IsGround() + } + return true +} + // OutputVars returns the set of variables that would be bound by // evaluating this expression in isolation. func (expr *Expr) OutputVars() VarSet { From 7a096fabb64955908b7f4dc5c2cd2a70ef609b87 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Tue, 31 May 2016 13:41:59 -0700 Subject: [PATCH 02/11] Refactor Term unmarshalling This will make it easier to introduce array comprehensions. --- ast/policy.go | 31 ++++------- ast/policy_test.go | 7 +-- ast/term.go | 128 ++++++++++++++++++--------------------------- ast/term_test.go | 46 ++++++---------- 4 files changed, 82 insertions(+), 130 deletions(-) diff --git a/ast/policy.go b/ast/policy.go index 6f9dfaff90..75204e20b7 100644 --- a/ast/policy.go +++ b/ast/policy.go @@ -401,15 +401,12 @@ func (expr *Expr) UnmarshalJSON(bs []byte) error { return err } - n, ok := v["Negated"] - if !ok { - expr.Negated = false - } else { - b, ok := n.(bool) - if !ok { - return unmarshalError(n, "bool") + if x, ok := v["Negated"]; ok { + if b, ok := x.(bool); ok { + expr.Negated = b + } else { + return fmt.Errorf("ast: unable to unmarshal Negated field with type: %T (expected true or false)", v["Negated"]) } - expr.Negated = b } switch ts := v["Terms"].(type) { @@ -420,21 +417,13 @@ func (expr *Expr) UnmarshalJSON(bs []byte) error { } expr.Terms = &Term{Value: v} case []interface{}: - buf := []*Term{} - for _, v := range ts { - e, ok := v.(map[string]interface{}) - if !ok { - return unmarshalError(v, "map[string]interface{}") - } - v, err := unmarshalValue(e) - if err != nil { - return err - } - buf = append(buf, &Term{Value: v}) + terms, err := unmarshalTermSlice(ts) + if err != nil { + return err } - expr.Terms = buf + expr.Terms = terms default: - return unmarshalError(v["Terms"], "Term or []Term") + return fmt.Errorf(`ast: unable to unmarshal Terms field with type: %T (expected {"Value": ..., "Type": ...} or [{"Value": ..., "Type": ...}, ...])`, v["Terms"]) } return nil } diff --git a/ast/policy_test.go b/ast/policy_test.go index f14ece4df4..3751b3e3d2 100644 --- a/ast/policy_test.go +++ b/ast/policy_test.go @@ -6,6 +6,7 @@ package ast import ( "encoding/json" + "fmt" "reflect" "testing" ) @@ -187,7 +188,7 @@ func TestExprBadJSON(t *testing.T) { } ` - exp := unmarshalError(100.0, "bool") + exp := fmt.Errorf("ast: unable to unmarshal Negated field with type: float64 (expected true or false)") assert(js, exp) js = ` @@ -197,7 +198,7 @@ func TestExprBadJSON(t *testing.T) { ] } ` - exp = unmarshalError("foo", "map[string]interface{}") + exp = fmt.Errorf("ast: unable to unmarshal term") assert(js, exp) js = ` @@ -205,7 +206,7 @@ func TestExprBadJSON(t *testing.T) { "Terms": "bad value" } ` - exp = unmarshalError("bad value", "Term or []Term") + exp = fmt.Errorf(`ast: unable to unmarshal Terms field with type: string (expected {"Value": ..., "Type": ...} or [{"Value": ..., "Type": ...}, ...])`) assert(js, exp) } diff --git a/ast/term.go b/ast/term.go index a731882c21..82ae3ac891 100644 --- a/ast/term.go +++ b/ast/term.go @@ -631,104 +631,80 @@ func termSliceIsGround(a []*Term) bool { return true } -func unmarshalError(v interface{}, e string) error { - return fmt.Errorf("ast: cannot unmarshal %T into Go value of type %v", v, e) -} +// TODO(tsandall): The unmarshalling errors in these functions are not +// helpful for callers because they do not identify the source of the +// unmarshalling error. Because OPA doesn't accept JSON describing ASTs +// from callers, this is acceptable (for now). If that changes in the future, +// the error messages should be revisited. The current approach focuses +// on the happy path and treats all errors the same. If better error +// reporting is needed, the error paths will need to be fleshed out. -func unmarshalTermSlice(d map[string]interface{}) ([]*Term, error) { - s, ok := d["Value"].([]interface{}) - if !ok { - return nil, unmarshalError(d["Value"], "[]interface{}") - } +func unmarshalTermSlice(s []interface{}) ([]*Term, error) { buf := []*Term{} - for _, i := range s { - m, ok := i.(map[string]interface{}) - if !ok { - return nil, unmarshalError(i, "map[string]interface{}") + for _, x := range s { + if m, ok := x.(map[string]interface{}); ok { + if v, err := unmarshalValue(m); err == nil { + buf = append(buf, &Term{Value: v}) + continue + } } - v, err := unmarshalValue(m) - if err != nil { - return nil, err - } - buf = append(buf, &Term{Value: v}) + return nil, fmt.Errorf("ast: unable to unmarshal term") } return buf, nil } +func unmarshalTermSliceValue(d map[string]interface{}) ([]*Term, error) { + if s, ok := d["Value"].([]interface{}); ok { + return unmarshalTermSlice(s) + } + return nil, fmt.Errorf(`ast: unable to unmarshal term (expected {"Value": [...], "Type": ...} where type is one of: array, reference)`) +} + func unmarshalValue(d map[string]interface{}) (Value, error) { + v := d["Value"] switch d["Type"] { case "null": return Null{}, nil case "boolean": - b, ok := d["Value"].(bool) - if !ok { - return nil, unmarshalError(d["Value"], "bool") + if b, ok := v.(bool); ok { + return Boolean(b), nil } - return Boolean(b), nil case "number": - f, ok := d["Value"].(float64) - if !ok { - return nil, unmarshalError(d["Value"], "float64") + if n, ok := v.(float64); ok { + return Number(n), nil } - return Number(f), nil case "string": - s, ok := d["Value"].(string) - if !ok { - return nil, unmarshalError(d["Value"], "string") + if s, ok := v.(string); ok { + return String(s), nil } - return String(s), nil - case "ref": - s, err := unmarshalTermSlice(d) - if err != nil { - return nil, err - } - return Ref(s), nil case "var": - s, ok := d["Value"].(string) - if !ok { - return nil, unmarshalError(d["Value"], "ast.Var") + if s, ok := v.(string); ok { + return Var(s), nil + } + case "ref": + if s, err := unmarshalTermSliceValue(d); err == nil { + return Ref(s), nil } - return Var(s), nil case "array": - s, err := unmarshalTermSlice(d) - if err != nil { - return nil, err + if s, err := unmarshalTermSliceValue(d); err == nil { + return Array(s), nil } - return Array(s), nil case "object": - buf := Object{} - s, ok := d["Value"].([]interface{}) - if !ok { - return nil, unmarshalError(d["Value"], "[]interface{}") + if s, ok := v.([]interface{}); ok { + buf := Object{} + for _, x := range s { + if i, ok := x.([]interface{}); ok && len(i) == 2 { + p, err := unmarshalTermSlice(i) + if err == nil { + buf = append(buf, Item(p[0], p[1])) + continue + } + } + goto unmarshal_error + } + return buf, nil } - for _, i := range s { - p, ok := i.([]interface{}) - if !ok { - return nil, unmarshalError(i, "[]interface{}") - } - if len(p) != 2 { - return nil, unmarshalError(p, "[2]interface{}") - } - km, ok := p[0].(map[string]interface{}) - if !ok { - return nil, unmarshalError(p[0], "map[string]interface{}") - } - k, err := unmarshalValue(km) - if err != nil { - return nil, err - } - vm, ok := p[1].(map[string]interface{}) - if !ok { - return nil, unmarshalError(p[1], "map[string]interface{}") - } - v, err := unmarshalValue(vm) - if err != nil { - return nil, err - } - buf = append(buf, [2]*Term{&Term{Value: k}, &Term{Value: v}}) - } - return buf, nil - default: - return nil, fmt.Errorf("ast: cannot unmarshal Term with Type %v", d["Type"]) } +unmarshal_error: + return nil, fmt.Errorf("ast: unable to unmarshal term") } diff --git a/ast/term_test.go b/ast/term_test.go index b55aa4d146..e855438c60 100644 --- a/ast/term_test.go +++ b/ast/term_test.go @@ -109,38 +109,24 @@ func TestQuery(t *testing.T) { func TestTermBadJSON(t *testing.T) { - assert := func(js string, exp error) { - term := Term{} - err := json.Unmarshal([]byte(js), &term) - if !reflect.DeepEqual(exp, err) { - t.Errorf("Expected %v but got: %v", exp, err) - } + input := `{ + "Value": [[ + {"Value": [{"Value": "a", "Type": "var"}, {"Value": "x", "Type": "string"}], "Type": "ref"}, + {"Value": [{"Value": "x", "Type": "var"}], "Type": "array"} + ], [ + {"Value": 100, "Type": "array"}, + {"Value": "foo", "Type": "string"} + ]], + "Type": "object" + }` + + term := Term{} + err := json.Unmarshal([]byte(input), &term) + expected := fmt.Errorf("ast: unable to unmarshal term") + if !reflect.DeepEqual(expected, err) { + t.Errorf("Expected %v but got: %v", expected, err) } - castTests := []struct { - input string - val interface{} - expected string - }{ - {`{"Value": null, "Type": "boolean"}`, nil, "bool"}, - {`{"Value": false, "Type": "number"}`, false, "float64"}, - {`{"Value": 100, "Type": "string"}`, 100.0, "string"}, - {`{"Value": "hello", "Type": "number"}`, "hello", "float64"}, - {`{"Value": 100, "Type": "var"}`, 100.0, "ast.Var"}, - {`{"Value": "abc", "Type": "ref"}`, "abc", "[]interface{}"}, - {`{"Value": ["abc"], "Type": "ref"}`, "abc", "map[string]interface{}"}, - {`{"Value": "abc", "Type": "array"}`, "abc", "[]interface{}"}, - {`{"Value": ["abc"], "Type": "array"}`, "abc", "map[string]interface{}"}, - {`{"Value": "abc", "Type": "object"}`, "abc", "[]interface{}"}, - {`{"Value": ["abc"], "Type": "object"}`, "abc", "[]interface{}"}, - {`{"Value": [["abc"]], "Type": "object"}`, []interface{}{}, "[2]interface{}"}, - {`{"Value": [["abc", "abc"]], "Type": "object"}`, "abc", "map[string]interface{}"}, - {`{"Value": [[{"Value": "abc", "Type": "string"}, "abc"]], "Type": "object"}`, "abc", "map[string]interface{}"}, - } - - for _, tc := range castTests { - assert(tc.input, unmarshalError(tc.val, tc.expected)) - } } func TestTermEqual(t *testing.T) { From d25f2a4c304c7fa7dfc3e247e9a8acf397bd1146 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Tue, 31 May 2016 14:26:13 -0700 Subject: [PATCH 03/11] Add array comprehension parsing Also, move wildcard mangling into post processing step in parser extensions. Wildcards needs to be mangled after Parse() because otherwise the generated variable names will reset when handling closures. --- ast/parser_ext.go | 13 +++++ ast/parser_test.go | 67 +++++++++++++++++++++++- ast/policy.go | 27 +--------- ast/policy_test.go | 1 + ast/rego.peg | 13 +++-- ast/term.go | 124 +++++++++++++++++++++++++++++++++++++++++++-- ast/term_test.go | 6 ++- ast/visit.go | 3 ++ ast/visit_test.go | 25 +++++++-- 9 files changed, 241 insertions(+), 38 deletions(-) diff --git a/ast/parser_ext.go b/ast/parser_ext.go index 4980acaf01..4077ab2e20 100644 --- a/ast/parser_ext.go +++ b/ast/parser_ext.go @@ -166,6 +166,7 @@ func ParseStatements(input string) ([]interface{}, error) { return nil, err } stmts := parsed.([]interface{}) + postProcess(stmts) return stmts, err } @@ -266,6 +267,18 @@ func parseModule(stmts []interface{}) (*Module, error) { return mod, nil } +func postProcess(stmts []interface{}) { + mangleWildcards(stmts) +} + +func mangleWildcards(stmts []interface{}) { + + mangler := &wildcardMangler{} + for _, stmt := range stmts { + Walk(mangler, stmt) + } +} + type wildcardMangler struct { c int } diff --git a/ast/parser_test.go b/ast/parser_test.go index cb4bc11c00..5ca0882255 100644 --- a/ast/parser_test.go +++ b/ast/parser_test.go @@ -140,6 +140,47 @@ func TestCompositesWithRefs(t *testing.T) { assertParseOneTerm(t, "ref values", "[{8: a[i].b, f: c[0][\"d\"].e[j]}]", ArrayTerm(ObjectTerm(Item(NumberTerm(8), ref1), Item(VarTerm("f"), ref2)))) } +func TestArrayComprehensions(t *testing.T) { + + input := `[ + {"x": [a[i] | xs = [{"a": ["baz", j]} | q[p], p.a != "bar", j = "foo"], + xs[j].a[k] = "foo"]} + ]` + + expected := ArrayTerm( + ObjectTerm(Item( + StringTerm("x"), + ArrayComprehensionTerm( + RefTerm(VarTerm("a"), VarTerm("i")), + Body{ + NewBuiltinExpr( + VarTerm("="), + VarTerm("xs"), + ArrayComprehensionTerm( + ObjectTerm(Item(StringTerm("a"), ArrayTerm(StringTerm("baz"), VarTerm("j")))), + Body{ + &Expr{ + Terms: RefTerm(VarTerm("q"), VarTerm("p")), + }, + NewBuiltinExpr(VarTerm("!="), RefTerm(VarTerm("p"), StringTerm("a")), StringTerm("bar")), + NewBuiltinExpr(VarTerm("="), VarTerm("j"), StringTerm("foo")), + }, + ), + ), + NewBuiltinExpr( + VarTerm("="), + RefTerm(VarTerm("xs"), VarTerm("j"), StringTerm("a"), VarTerm("k")), + StringTerm("foo"), + ), + }, + ), + )), + ) + + assertParseOneTerm(t, "nested", input, expected) + +} + func TestInfixExpr(t *testing.T) { assertParseOneExpr(t, "scalars 1", "true = false", NewBuiltinExpr(VarTerm("="), BooleanTerm(true), BooleanTerm(false))) assertParseOneExpr(t, "scalars 2", "3.14 = null", NewBuiltinExpr(VarTerm("="), NumberTerm(3.14), NullTerm())) @@ -295,6 +336,10 @@ func TestComments(t *testing.T) { :- m = [1,2, 3], a = m[i] + + r[x] :- x = [ a | # inside comprehension + a = z[i], + b[i].a = a ] ` assertParseModule(t, "module comments", testModule, &Module{ @@ -307,6 +352,7 @@ func TestComments(t *testing.T) { Rules: []*Rule{ MustParseStatement("p[x] = y :- y = \"foo\", x = \"bar\", x != y, q[x]").(*Rule), MustParseStatement("q[a] :- m = [1,2,3], a = m[i]").(*Rule), + MustParseStatement("r[x] :- x = [a | a = z[i], b[i].a = a]").(*Rule), }, }) } @@ -416,6 +462,25 @@ func TestWildcards(t *testing.T) { ), }, }) + + assertParseOneExpr(t, "comprehension", "_ = [x | a = a[_]]", &Expr{ + Terms: []*Term{ + VarTerm("="), + VarTerm("$0"), + ArrayComprehensionTerm( + VarTerm("x"), + Body{ + &Expr{ + Terms: []*Term{ + VarTerm("="), + VarTerm("a"), + RefTerm(VarTerm("a"), VarTerm("$1")), + }, + }, + }, + ), + }, + }) } func assertParse(t *testing.T, msg string, input string, correct func([]interface{})) { @@ -485,7 +550,7 @@ func assertParseOneExpr(t *testing.T, msg string, input string, correct *Expr) { } expr := body[0] if !expr.Equal(correct) { - t.Errorf("Error on test %s: expressions not equal: %v (parsed), %v (correct)", msg, expr, correct) + t.Errorf("Error on test %s: expressions not equal:\n%v (parsed)\n%v (correct)", msg, expr, correct) } }) } diff --git a/ast/policy.go b/ast/policy.go index 75204e20b7..5506a20a82 100644 --- a/ast/policy.go +++ b/ast/policy.go @@ -400,32 +400,7 @@ func (expr *Expr) UnmarshalJSON(bs []byte) error { if err := json.Unmarshal(bs, &v); err != nil { return err } - - if x, ok := v["Negated"]; ok { - if b, ok := x.(bool); ok { - expr.Negated = b - } else { - return fmt.Errorf("ast: unable to unmarshal Negated field with type: %T (expected true or false)", v["Negated"]) - } - } - - switch ts := v["Terms"].(type) { - case map[string]interface{}: - v, err := unmarshalValue(ts) - if err != nil { - return err - } - expr.Terms = &Term{Value: v} - case []interface{}: - terms, err := unmarshalTermSlice(ts) - if err != nil { - return err - } - expr.Terms = terms - default: - return fmt.Errorf(`ast: unable to unmarshal Terms field with type: %T (expected {"Value": ..., "Type": ...} or [{"Value": ..., "Type": ...}, ...])`, v["Terms"]) - } - return nil + return unmarshalExpr(expr, v) } // Vars returns a VarSet containing all of the variables in the expression. diff --git a/ast/policy_test.go b/ast/policy_test.go index 3751b3e3d2..a8661ff83a 100644 --- a/ast/policy_test.go +++ b/ast/policy_test.go @@ -19,6 +19,7 @@ func TestModuleJSONRoundTrip(t *testing.T) { p = [1,2,{"foo":3}] :- r[x] = 1, not q[x] r[y] = v :- i[1] = y, v = i[2] q[x] :- a=[true,false,null,{"x":[1,2,3]}], a[i] = x + t = true :- xs = [{"x": a[i].a} | a[i].n = "bob", b[x]] `) bs, err := json.Marshal(mod) diff --git a/ast/rego.peg b/ast/rego.peg index bae81804e9..2d6c778f3f 100644 --- a/ast/rego.peg +++ b/ast/rego.peg @@ -90,7 +90,6 @@ Import <- "import" ws path:(Ref / Var) alias:(ws "as" ws Var)? { return imp, nil } -// TODO(tsandall): update to handle underscore variables Rule <- name:Var key:( _ "[" _ Term _ "]" _ )? value:( _ "=" _ Term )? body:( _ ":-" _ Body) { rule := &Rule{} @@ -130,8 +129,6 @@ Body <- head:Expr tail:( _ "," _ Expr)* { expr := s.([]interface{})[3].(*Expr) buf = append(buf, expr) } - mangler := &wildcardMangler{} - Walk(mangler, buf) return buf, nil } @@ -169,10 +166,18 @@ PrefixExpr <- op:Var "(" _ head:Term? tail:( _ "," _ Term )* _ ")" { return buf, nil } -Term <- val:( Composite / Scalar / Ref / Var ) { +Term <- val:( Comprehension / Composite / Scalar / Ref / Var ) { return val, nil } +Comprehension <- ArrayComprehension + +ArrayComprehension <- "[" _ term:Term _ "|" _ body:Body _ "]" { + ac := ArrayComprehensionTerm(term.(*Term), body.(Body)) + ac.Location = currentLocation(c) + return ac, nil +} + Composite <- Object / Array Scalar <- Number / String / Bool / Null diff --git a/ast/term.go b/ast/term.go index 82ae3ac891..b98841d82b 100644 --- a/ast/term.go +++ b/ast/term.go @@ -34,6 +34,7 @@ func NewLocation(text []byte, file string, row int, col int) *Location { // - Object, Array // - Variables // - References +// - Array Comprehensions // type Value interface { // Equal returns true if this value equals the other value. @@ -70,6 +71,11 @@ func (term *Term) Equal(other *Term) bool { return term.Value.Equal(other.Value) } +// Hash returns the hash code of the Term's value. +func (term *Term) Hash() int { + return term.Value.Hash() +} + // IsGround returns true if this terms' Value is ground. func (term *Term) IsGround() bool { return term.Value.IsGround() @@ -97,6 +103,8 @@ func (term *Term) MarshalJSON() ([]byte, error) { typ = "array" case Object: typ = "object" + case *ArrayComprehension: + typ = "array-comprehension" } d := map[string]interface{}{ "Type": typ, @@ -582,6 +590,48 @@ func (obj Object) queryRec(ref Ref, keys map[Var]Value, iter QueryIterator) erro } } +// ArrayComprehension represents an array comprehension as defined in the language. +type ArrayComprehension struct { + Term *Term + Body Body +} + +// ArrayComprehensionTerm creates a new Term with an ArrayComprehension value. +func ArrayComprehensionTerm(term *Term, body Body) *Term { + return &Term{ + Value: &ArrayComprehension{ + Term: term, + Body: body, + }, + } +} + +// Equal returns true if this array comprehension is syntactically equal to another. +func (ac *ArrayComprehension) Equal(other Value) bool { + if ac == other { + return true + } + o, ok := other.(*ArrayComprehension) + if !ok { + return false + } + return o.Term.Equal(ac.Term) && o.Body.Equal(ac.Body) +} + +// Hash returns the hash code of the Value. +func (ac *ArrayComprehension) Hash() int { + return ac.Term.Hash() + ac.Body.Hash() +} + +// IsGround returns true if the Term and Body are ground. +func (ac *ArrayComprehension) IsGround() bool { + return ac.Term.IsGround() && ac.Body.IsGround() +} + +func (ac *ArrayComprehension) String() string { + return "[" + ac.Term.String() + " | " + ac.Body.String() + "]" +} + func queryRec(v Value, ref Ref, tail Ref, keys map[Var]Value, iter QueryIterator, skipScalar bool) error { if len(tail) == 0 { if err := iter(keys, v); err != nil { @@ -631,7 +681,7 @@ func termSliceIsGround(a []*Term) bool { return true } -// TODO(tsandall): The unmarshalling errors in these functions are not +// NOTE(tsandall): The unmarshalling errors in these functions are not // helpful for callers because they do not identify the source of the // unmarshalling error. Because OPA doesn't accept JSON describing ASTs // from callers, this is acceptable (for now). If that changes in the future, @@ -639,12 +689,64 @@ func termSliceIsGround(a []*Term) bool { // on the happy path and treats all errors the same. If better error // reporting is needed, the error paths will need to be fleshed out. +func unmarshalBody(b []interface{}) (Body, error) { + buf := Body{} + for _, e := range b { + if m, ok := e.(map[string]interface{}); ok { + expr := &Expr{} + if err := unmarshalExpr(expr, m); err == nil { + buf = append(buf, expr) + continue + } + } + goto unmarshal_error + } + return buf, nil +unmarshal_error: + return nil, fmt.Errorf("ast: unable to unmarshal body") +} + +func unmarshalExpr(expr *Expr, v map[string]interface{}) error { + if x, ok := v["Negated"]; ok { + if b, ok := x.(bool); ok { + expr.Negated = b + } else { + return fmt.Errorf("ast: unable to unmarshal Negated field with type: %T (expected true or false)", v["Negated"]) + } + } + switch ts := v["Terms"].(type) { + case map[string]interface{}: + t, err := unmarshalTerm(ts) + if err != nil { + return err + } + expr.Terms = t + case []interface{}: + terms, err := unmarshalTermSlice(ts) + if err != nil { + return err + } + expr.Terms = terms + default: + return fmt.Errorf(`ast: unable to unmarshal Terms field with type: %T (expected {"Value": ..., "Type": ...} or [{"Value": ..., "Type": ...}, ...])`, v["Terms"]) + } + return nil +} + +func unmarshalTerm(m map[string]interface{}) (*Term, error) { + v, err := unmarshalValue(m) + if err != nil { + return nil, err + } + return &Term{Value: v}, nil +} + func unmarshalTermSlice(s []interface{}) ([]*Term, error) { buf := []*Term{} for _, x := range s { if m, ok := x.(map[string]interface{}); ok { - if v, err := unmarshalValue(m); err == nil { - buf = append(buf, &Term{Value: v}) + if t, err := unmarshalTerm(m); err == nil { + buf = append(buf, t) continue } } @@ -704,6 +806,22 @@ func unmarshalValue(d map[string]interface{}) (Value, error) { } return buf, nil } + case "array-comprehension": + if m, ok := v.(map[string]interface{}); ok { + if t, ok := m["Term"].(map[string]interface{}); ok { + if term, err := unmarshalTerm(t); err == nil { + if b, ok := m["Body"].([]interface{}); ok { + if body, err := unmarshalBody(b); err == nil { + buf := &ArrayComprehension{ + Term: term, + Body: body, + } + return buf, nil + } + } + } + } + } } unmarshal_error: return nil, fmt.Errorf("ast: unable to unmarshal term") diff --git a/ast/term_test.go b/ast/term_test.go index e855438c60..68693fd99c 100644 --- a/ast/term_test.go +++ b/ast/term_test.go @@ -141,6 +141,7 @@ func TestTermEqual(t *testing.T) { assertTermEqual(t, ArrayTerm(NumberTerm(1), NumberTerm(2), NumberTerm(3)), ArrayTerm(NumberTerm(1), NumberTerm(2), NumberTerm(3))) assertTermEqual(t, VarTerm("foo"), VarTerm("foo")) assertTermEqual(t, RefTerm(VarTerm("foo"), VarTerm("i"), NumberTerm(2)), RefTerm(VarTerm("foo"), VarTerm("i"), NumberTerm(2))) + assertTermEqual(t, ArrayComprehensionTerm(VarTerm("x"), Body{&Expr{Terms: RefTerm(VarTerm("a"), VarTerm("i"))}}), ArrayComprehensionTerm(VarTerm("x"), Body{&Expr{Terms: RefTerm(VarTerm("a"), VarTerm("i"))}})) assertTermNotEqual(t, NullTerm(), BooleanTerm(true)) assertTermNotEqual(t, BooleanTerm(true), BooleanTerm(false)) assertTermNotEqual(t, NumberTerm(5), NumberTerm(7)) @@ -153,6 +154,7 @@ func TestTermEqual(t *testing.T) { assertTermNotEqual(t, ArrayTerm(NumberTerm(1), NumberTerm(2), NumberTerm(3)), ArrayTerm(NumberTerm(1), NumberTerm(2), NumberTerm(4))) assertTermNotEqual(t, VarTerm("foo"), VarTerm("bar")) assertTermNotEqual(t, RefTerm(VarTerm("foo"), VarTerm("i"), NumberTerm(2)), RefTerm(VarTerm("foo"), StringTerm("i"), NumberTerm(2))) + assertTermNotEqual(t, ArrayComprehensionTerm(VarTerm("x"), Body{&Expr{Terms: RefTerm(VarTerm("a"), VarTerm("j"))}}), ArrayComprehensionTerm(VarTerm("x"), Body{&Expr{Terms: RefTerm(VarTerm("a"), VarTerm("i"))}})) } func TestHash(t *testing.T) { @@ -164,7 +166,8 @@ func TestHash(t *testing.T) { ], "e": { 100: a[i].b - } + }, + "k": [ "foo" | true ] } ` @@ -195,6 +198,7 @@ func TestTermString(t *testing.T) { assertToString(t, ArrayTerm().Value, "[]") assertToString(t, ObjectTerm().Value, "{}") assertToString(t, ArrayTerm(ObjectTerm(Item(VarTerm("foo"), ArrayTerm(RefTerm(VarTerm("bar"), VarTerm("i"))))), StringTerm("foo"), BooleanTerm(true), NullTerm(), NumberTerm(42.1)).Value, "[{foo: [bar[i]]}, \"foo\", true, null, 42.1]") + assertToString(t, ArrayComprehensionTerm(ArrayTerm(VarTerm("x")), Body{&Expr{Terms: RefTerm(VarTerm("a"), VarTerm("i"))}}).Value, "[[x] | a[i]]") } func TestRefUnderlying(t *testing.T) { diff --git a/ast/visit.go b/ast/visit.go index 3251660862..f249348c04 100644 --- a/ast/visit.go +++ b/ast/visit.go @@ -72,5 +72,8 @@ func Walk(v Visitor, x interface{}) { for _, t := range x { Walk(w, t.Value) } + case *ArrayComprehension: + Walk(w, x.Term) + Walk(w, x.Body) } } diff --git a/ast/visit_test.go b/ast/visit_test.go index d295692242..bbdf87cbcb 100644 --- a/ast/visit_test.go +++ b/ast/visit_test.go @@ -20,7 +20,10 @@ func TestVisitor(t *testing.T) { rule := MustParseModule(` package a.b import x.y as z - t[x] = y :- p[x] = {"foo": [y,2,{"bar": 3}]}, not q[x] + t[x] = y :- + p[x] = {"foo": [y,2,{"bar": 3}]}, + not q[x], + y = [ [x,z] | x = "x", z = "z" ] `) vis := &testVis{} Walk(vis, rule) @@ -59,9 +62,25 @@ func TestVisitor(t *testing.T) { ref2 q x + expr3 + = + y + compr + array + x + z + body + expr4 + = + x + "x" + expr5 + = + z + "z" */ - if len(vis.elems) != 33 { - t.Errorf("Expected exactly 33 elements in AST but got %d: %v", len(vis.elems), vis.elems) + if len(vis.elems) != 49 { + t.Errorf("Expected exactly 49 elements in AST but got %d: %v", len(vis.elems), vis.elems) } } From fe37cc03a82437da52b03c78024e19c4a54b064e Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Wed, 8 Jun 2016 17:06:38 -0700 Subject: [PATCH 04/11] Improve coverage of ground-ness checks --- ast/policy_test.go | 6 ++++++ ast/term_test.go | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/ast/policy_test.go b/ast/policy_test.go index a8661ff83a..40b8a95f64 100644 --- a/ast/policy_test.go +++ b/ast/policy_test.go @@ -133,6 +133,12 @@ func TestExprEquals(t *testing.T) { assertExprNotEqual(t, expr20, expr23) } +func TestBodyIsGround(t *testing.T) { + if MustParseBody(`a.b[0] = 1, a = [1,2,x]`).IsGround() { + t.Errorf("Expected body to be non-ground") + } +} + func TestExprOutputVars(t *testing.T) { body := MustParseBody(`{"a": [{x: y}, b[z]]} = c[i], [{"a": d[j][k]}] != xs`) one := body[0] diff --git a/ast/term_test.go b/ast/term_test.go index 68693fd99c..3c0cc0dd4f 100644 --- a/ast/term_test.go +++ b/ast/term_test.go @@ -181,6 +181,42 @@ func TestHash(t *testing.T) { } } +func TestTermIsGround(t *testing.T) { + + tests := []struct { + note string + term string + expected bool + }{ + {"null", "null", true}, + {"string", `"foo"`, true}, + {"number", "42.1", true}, + {"boolean", "false", true}, + {"var", "x", false}, + {"ref ground", "a.b[0]", true}, + {"ref non-ground", "a.b[i].x", false}, + {"array ground", "[1,2,3]", true}, + {"array non-ground", "[1,2,x]", false}, + {"object ground", `{"a": 1}`, true}, + {"object non-ground key", `{"x": 1, y: 2}`, false}, + {"object non-ground value", `{"x": 1, "y": y}`, false}, + {"array compr ground", `["a" | true]`, true}, + {"array compr non-ground", `[x | x = a[i]]`, false}, + } + + for i, tc := range tests { + term := MustParseTerm(tc.term) + if term.IsGround() != tc.expected { + expected := "ground" + if !tc.expected { + expected = "non-ground" + } + t.Errorf("Expected term %v to be %s (test case %d: %v)", term, expected, i, tc.note) + } + } + +} + func TestTermString(t *testing.T) { assertToString(t, Null{}, "null") assertToString(t, Boolean(true), "true") From 0decd227ab93e51ca1c92bb03c5396a7b72cc97b Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Thu, 9 Jun 2016 13:34:25 -0700 Subject: [PATCH 05/11] Add safety check on closures/array comprehensions This change set does away with the old way of determining which variables are outputs. Equality is now handled with special care. Outputs that would make a variable safe by depending on another unsafe variable are no longer included. As a result, the occurs check in the topdown implementation is no longer needed. This change was introduced to handle odd cases involving comprehensions, e.g., x = y, x = [ y | y = 1 ]. In this case, without exluding unsafe vars, the query would evaluate with x/[1]. This would violate the semantics, because in the comprehension y/1. Also, fix bug in reordering whereby potentially unsafe expressions were added to the reordered body multiple times. This occurred because the expression would be added once when the preceeding expression made it safe and then again once the outer loop got it. With the fix, we reprocess the body each time an expression is added to the reordered body. --- ast/builtins.go | 35 +++------ ast/builtins_test.go | 21 +---- ast/compile.go | 170 +++++++++++++++++++++++++++++++++------- ast/compile_test.go | 157 +++++++++++++++++++++++++++++++------ ast/policy.go | 160 ++++++++++++++++++++++++++----------- ast/policy_test.go | 39 +++++++-- ast/term.go | 18 +++++ ast/unify.go | 165 ++++++++++++++++++++++++++++++++++++++ ast/unify_test.go | 78 ++++++++++++++++++ ast/varset.go | 22 ++++++ ast/visit.go | 40 ++++++++++ topdown/topdown.go | 26 ------ topdown/topdown_test.go | 6 -- 13 files changed, 761 insertions(+), 176 deletions(-) create mode 100644 ast/unify.go create mode 100644 ast/unify_test.go diff --git a/ast/builtins.go b/ast/builtins.go index 52febabea2..24d8bb5e64 100644 --- a/ast/builtins.go +++ b/ast/builtins.go @@ -29,10 +29,10 @@ var BuiltinMap map[Var]*Builtin // Equality represents the "=" operator. var Equality = &Builtin{ - Name: Var("="), - Alias: Var("eq"), - NumArgs: 2, - RecTargetPos: []int{0, 1}, + Name: Var("="), + Alias: Var("eq"), + NumArgs: 2, + TargetPos: []int{0, 1}, } // GreaterThan represents the ">" comparison operator. @@ -73,11 +73,10 @@ var NotEqual = &Builtin{ // Builtin represents a built-in function supported by OPA. Every // built-in function is uniquely identified by a name. type Builtin struct { - Name Var - Alias Var - NumArgs int - TargetPos []int - RecTargetPos []int + Name Var + Alias Var + NumArgs int + TargetPos []int } // GetPrintableName returns a printable name for the builtin. @@ -91,26 +90,14 @@ func (b *Builtin) GetPrintableName() string { return b.Name.String() } -// Unifies returns true if a term in the given position will unify -// non-recursively or recursively. -func (b *Builtin) Unifies(i int) bool { +// IsTargetPos returns true if a variable in the i-th position will be +// bound when the expression is evaluated. +func (b *Builtin) IsTargetPos(i int) bool { for _, x := range b.TargetPos { if x == i { return true } } - return b.UnifiesRecursively(i) -} - -// UnifiesRecursively returns true if a term in the given position will -// unify recursively, i.e., variables embedded inside a collection type -// will unify. -func (b *Builtin) UnifiesRecursively(i int) bool { - for _, x := range b.RecTargetPos { - if x == i { - return true - } - } return false } diff --git a/ast/builtins_test.go b/ast/builtins_test.go index 0967636c41..74ebd8bb18 100644 --- a/ast/builtins_test.go +++ b/ast/builtins_test.go @@ -10,26 +10,11 @@ import ( ) func TestUnifies(t *testing.T) { - b := &Builtin{Name: Var("dummy"), NumArgs: 4, RecTargetPos: []int{2, 3}, TargetPos: []int{1}} - expected := []int{1, 2, 3} + b := &Builtin{Name: Var("dummy"), NumArgs: 4, TargetPos: []int{1, 3}} + expected := []int{1, 3} result := []int{} for i := 0; i < 4; i++ { - if b.Unifies(i) { - result = append(result, i) - } - } - if !reflect.DeepEqual(expected, result) { - t.Errorf("Expected %v but got: %v", expected, result) - } -} - -func TestUnifiesRecursively(t *testing.T) { - - b := &Builtin{Name: Var("dummy"), NumArgs: 4, RecTargetPos: []int{2, 3}, TargetPos: []int{1}} - expected := []int{2, 3} - result := []int{} - for i := 0; i < 4; i++ { - if b.UnifiesRecursively(i) { + if b.IsTargetPos(i) { result = append(result, i) } } diff --git a/ast/compile.go b/ast/compile.go index 4222fba88a..6888ce1f0d 100644 --- a/ast/compile.go +++ b/ast/compile.go @@ -204,7 +204,7 @@ func (c *Compiler) checkSafetyHead() { for _, m := range c.Modules { for _, r := range m.Rules { headVars := r.HeadVars() - bodyVars := r.Body.Vars() + bodyVars := r.Body.Vars(true) for headVar := range headVars { if _, ok := bodyVars[headVar]; !ok { c.err("unsafe variable from head of %v: %v", r.Name, headVar) @@ -533,6 +533,19 @@ func (vs unsafeVars) Add(e *Expr, v Var) { } } +func (vs unsafeVars) Set(e *Expr, s VarSet) { + vs[e] = s +} + +func (vs unsafeVars) Update(o unsafeVars) { + for k, v := range o { + if _, ok := vs[k]; !ok { + vs[k] = VarSet{} + } + vs[k].Update(v) + } +} + func (vs unsafeVars) Vars() VarSet { r := VarSet{} for _, s := range vs { @@ -604,50 +617,149 @@ func findRulesRec(node *ModuleTreeNode, ref Ref) []*Rule { // contains a mapping of expressions to unsafe variables in those expressions. func reorderBodyForSafety(globals VarSet, body Body) (Body, unsafeVars) { + body, unsafe := reorderBodyForClosures(globals, body) + if len(unsafe) != 0 { + return nil, unsafe + } + reordered := Body{} - unsafe := unsafeVars{} + safe := VarSet{} for _, e := range body { - for v := range e.Vars() { - if !globals.Contains(v) { + for v := range e.Vars(true) { + if globals.Contains(v) { + safe.Add(v) + } else { unsafe.Add(e, v) } } } - safe := VarSet{} + for { + n := len(reordered) - for _, e := range body { + for _, e := range body { + if reordered.Contains(e) { + continue + } - safe.Update(e.OutputVars()) + safe.Update(e.OutputVars(safe)) - for v := range unsafe[e] { - if safe.Contains(v) { - delete(unsafe[e], v) + for v := range unsafe[e] { + if safe.Contains(v) { + delete(unsafe[e], v) + } + } + + if len(unsafe[e]) == 0 { + delete(unsafe, e) + reordered = append(reordered, e) } } - if len(unsafe[e]) == 0 { - reordered = append(reordered, e) - delete(unsafe, e) + if len(reordered) == n { + break + } + } - // Check if other expressions in the body are considered safe - // now. If they are considered safe now, they can be added - // to the end of the re-ordered body. - for _, e := range body { - if reordered.Contains(e) { - continue - } - for v := range unsafe[e] { - if safe.Contains(v) { - delete(unsafe[e], v) - } - } - if len(unsafe[e]) == 0 { - reordered = append(reordered, e) - delete(unsafe, e) - } + // Recursively visit closures and perform the safety checks on them. + // Update the globals at each expression to include the variables that could + // be closed over. + g := globals.Copy() + for i, e := range reordered { + if i > 0 { + g.Update(reordered[i-1].Vars(true)) + } + vis := &bodySafetyVisitor{ + current: e, + globals: g, + unsafe: unsafe, + } + Walk(vis, e) + } + + return reordered, unsafe +} + +type bodySafetyVisitor struct { + current *Expr + globals VarSet + unsafe unsafeVars +} + +func (vis *bodySafetyVisitor) Visit(x interface{}) Visitor { + switch x := x.(type) { + case *Expr: + cpy := *vis + cpy.current = x + return &cpy + case *ArrayComprehension: + vis.checkArrayComprehensionSafety(x) + return nil + } + return vis +} + +func (vis *bodySafetyVisitor) checkArrayComprehensionSafety(ac *ArrayComprehension) { + // Check term for safety. This is analagous to the rule head safety check. + tv := ac.Term.Vars() + bv := ac.Body.Vars(true) + bv.Update(vis.globals) + uv := tv.Diff(bv) + for v := range uv { + vis.unsafe.Add(vis.current, v) + } + + // Check body for safety, reordering as necessary. + r, u := reorderBodyForSafety(vis.globals, ac.Body) + if len(u) == 0 { + ac.Body = r + } else { + vis.unsafe.Update(u) + } +} + +// reorderBodyForClosures returns a copy of the body ordered such that +// expressions (such as array comprehensions) that close over variables are ordered +// after other expressions that contain the same variable in an output position. +func reorderBodyForClosures(globals VarSet, body Body) (Body, unsafeVars) { + + reordered := Body{} + unsafe := unsafeVars{} + + for { + n := len(reordered) + + for _, e := range body { + if reordered.Contains(e) { + continue } + + // Collect vars that are contained in closures within this + // expression. + vs := VarSet{} + WalkClosures(e, func(x interface{}) bool { + vis := &varVisitor{vars: vs} + Walk(vis, x) + return true + }) + + // Compute vars that are closed over from the body but not yet + // contained in the output position of an expression in the reordered + // body. These vars are considered unsafe. + cv := vs.Intersect(body.Vars(true)).Diff(globals) + uv := cv.Diff(reordered.OutputVars(globals)) + + if len(uv) == 0 { + reordered = append(reordered, e) + delete(unsafe, e) + } else { + unsafe.Set(e, uv) + } + } + + if len(reordered) == n { + break } } diff --git a/ast/compile_test.go b/ast/compile_test.go index aa3f7a4d55..5da555d0f2 100644 --- a/ast/compile_test.go +++ b/ast/compile_test.go @@ -8,6 +8,7 @@ import ( "fmt" "reflect" "sort" + "strings" "testing" ) @@ -178,31 +179,111 @@ func TestCompilerCheckSafetyHead(t *testing.T) { } func TestCompilerCheckSafetyBodyReordering(t *testing.T) { - c := NewCompiler() - c.Modules = getCompilerTestModules() - c.Modules["newMod"] = MustParseModule(` - package a.b - needsReorder = true :- a[i] = x, a = [1,2,3,4] - needsReorderNegated = true :- a = [true, false], b = [true, false], not a[i], b[i] - `) - compileStages(c, "", "checkSafetyBody") + tests := []struct { + note string + body string + expected interface{} + }{ + // trivial cases + {"noop", "x = 1, x != 0", "x = 1, x != 0"}, + {"var/ref", "a[i] = x, a = [1,2,3,4]", "a = [1,2,3,4], a[i] = x"}, + {"negation", + "a = [true, false], b = [true, false], not a[i], b[i]", + "a = [true, false], b = [true, false], b[i], not a[i]"}, + {"built-in", "x != 0, count([1,2,3], x)", "count([1,2,3], x), x != 0"}, + {"var/var 1", "x = y, z = 1, y = z", "z = 1, y = z, x = y"}, + {"var/var 2", "x = y, 1 = z, z = y", "1 = z, z = y, x = y"}, + {"var/var 3", "x != 0, y = x, y = 1", "y = 1, y = x, x != 0"}, - assertNotFailed(t, c) - - expected1 := MustParseBody(`a = [1,2,3,4], a[i] = x`) - - reordered1 := c.Modules["newMod"].Rules[0].Body - - if !expected1.Equal(reordered1) { - t.Errorf("Expected body to be re-ordered and equal to %v but got: %v", expected1, reordered1) + // comprehensions + {"array compr/var", "x != 0, [y | y = 1] = x", "[y | y = 1] = x, x != 0"}, + {"array compr/array", "[1] != [x], [y | y = 1] = [x]", "[y | y = 1] = [x], [1] != [x]"}, } - expected2 := MustParseBody(`a = [true, false], b = [true, false], b[i], not a[i]`) + for i, tc := range tests { + c := NewCompiler() + c.Modules = map[string]*Module{ + "mod": MustParseModule( + fmt.Sprintf(`package test + p :- %s`, tc.body)), + } - reordered2 := c.Modules["newMod"].Rules[1].Body + compileStages(c, "", "checkSafetyBody") + switch exp := tc.expected.(type) { + case string: + if c.Failed() { + t.Errorf("%v (#%d): Unexpected compilation error: %v", tc.note, i, c.FlattenErrors()) + return + } + e := MustParseBody(exp) + if !e.Equal(c.Modules["mod"].Rules[0].Body) { + t.Errorf("%v (#%d): Expected body to be ordered and equal to %v but got: %v", tc.note, i, e, c.Modules["mod"].Rules[0].Body) + } + case error: + if len(c.Errors) > 0 { + if !reflect.DeepEqual(c.Errors[0], exp) { + t.Errorf("%v (#%d): Expected compiler error %v but got: %v", tc.note, i, exp, c.Errors[0]) + } + } else { + t.Errorf("%v (#%d): Expected compiler error but got: %v", tc.note, i, c.Modules["mod"].Rules[0]) + } + } + } +} - if !expected2.Equal(reordered2) { - t.Errorf("Expected body to be re-ordered and equal to %v but got: %v", expected2, reordered2) +func TestCompilerCheckSafetyBodyReorderingClosures(t *testing.T) { + c := NewCompiler() + c.Modules = map[string]*Module{ + "mod": MustParseModule( + ` + package compr + + import data.b + import data.c + + p :- v = [null | true], # leave untouched + xs = [x | a[i] = x, a = [y | y != 1, y = c[j]]], # close over 'i' and 'j', 2-level reorder + xs[j] > 0, + b[i] = j + + # test that reordering is not performed when closing over different globals, e.g., + # built-ins, data, imports. + q :- _ = [x | x = b[i]], + _ = b[j], + + _ = [x | x = true, x != false], + true != false, + + _ = [x | data.foo[_] = x], + data.foo[_] = _ + `), + } + + compileStages(c, "", "checkSafetyBody") + assertNotFailed(t, c) + + result1 := c.Modules["mod"].Rules[0].Body + expected1 := MustParseBody(` + v = [null | true], + b[i] = j, + xs = [x | a = [y | y = c[j], y != 1], a[i] = x], + xs[j] > 0 + `) + if !result1.Equal(expected1) { + t.Errorf("Expected reordered body to be equal to:\n%v\nBut got:\n%v", expected1, result1) + } + + result2 := c.Modules["mod"].Rules[1].Body + expected2 := MustParseBody(` + _ = [x | x = b[i]], + _ = b[j], + _ = [x | x = true, x != false], + true != false, + _ = [x | data.foo[_] = x], + data.foo[_] = _ + `) + if !result2.Equal(expected2) { + t.Errorf("Expected pre-ordered body to equal:\n%v\nBut got:\n%v", expected2, result2) } } @@ -240,7 +321,10 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) { unboundNegated3[x] = true :- a = [1,2,3,4], b = [1,2,3,4], not a[i] = x, not b[j] = x # i and j would be unbound even though they are in embedded references - unboundNegated4 = true :- a = [{"foo": ["bar", "baz"]}], not a[0].foo = [a[0].foo[i], a[0].foo[j]] + unboundNegated4 = true :- a = [{"foo": ["bar", "baz"]}], not a[0].foo = [a[0].foo[i], a[0].foo[j]] + + # x would be unbound as input to count + unsafeBuiltin :- count([1,2,x], x) # i and x would be bound in the last expression so the third expression is safe negatedSafe = true :- a = [1,2,3,4], b = [1,2,3,4], not a[i] = x, b[i] = x @@ -248,6 +332,17 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) { # x would be unbound because it does not appear in the target position of any expression unboundNoTarget = true :- x > 0, x <= 3, x != 2 + unboundArrayComprBody1 :- _ = [x | x = data.a[_], y > 1] + unboundArrayComprBody2 :- _ = [x | x = a[_], a = [y | y = data.a[_], z > 1]] + unboundArrayComprBody3 :- _ = [v | v = [x | x = data.a[_]], x > 1] + unboundArrayComprTerm1 :- _ = [u | true] + unboundArrayComprTerm2 :- _ = [v | v = [w | w != 0]] + unboundArrayComprTerm3 :- _ = [x[i] | x = []] + unboundArrayComprMixed1 :- _ = [x | y = [a | a = z[i]]] + + unsafeClosure1 :- x = [x | x = 1] + unsafeClosure2 :- x = y, x = [y | y = 1] + negatedImport1 = true :- not foo negatedImport2 = true :- not bar negatedImport3 = true :- not baz @@ -262,11 +357,29 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) { fmt.Errorf("unsafe variables in unboundNegated2: [i x]"), fmt.Errorf("unsafe variables in unboundNegated3: [i j x]"), fmt.Errorf("unsafe variables in unboundNegated4: [i j]"), + fmt.Errorf("unsafe variables in unsafeBuiltin: [x]"), fmt.Errorf("unsafe variables in unboundNoTarget: [x]"), + fmt.Errorf("unsafe variables in unboundArrayComprBody1: [y]"), + fmt.Errorf("unsafe variables in unboundArrayComprBody2: [z]"), + fmt.Errorf("unsafe variables in unboundArrayComprBody3: [x]"), + fmt.Errorf("unsafe variables in unboundArrayComprTerm1: [u]"), + fmt.Errorf("unsafe variables in unboundArrayComprTerm2: [w]"), + fmt.Errorf("unsafe variables in unboundArrayComprTerm3: [i]"), + fmt.Errorf("unsafe variables in unboundArrayComprMixed1: [x z]"), + fmt.Errorf("unsafe variables in unsafeClosure1: [x]"), + fmt.Errorf("unsafe variables in unsafeClosure2: [y]"), } if !reflect.DeepEqual(expected, c.Errors) { - t.Errorf("Expected %v but got:%v", expected, c.Errors) + e := []string{} + for _, x := range expected { + e = append(e, x.Error()) + } + r := []string{} + for _, x := range c.Errors { + r = append(r, x.Error()) + } + t.Errorf("Expected:\n%v\nBut got:\n%v", strings.Join(e, "\n"), strings.Join(r, "\n")) } } diff --git a/ast/policy.go b/ast/policy.go index 5506a20a82..9f199886a9 100644 --- a/ast/policy.go +++ b/ast/policy.go @@ -238,6 +238,16 @@ func (body Body) IsGround() bool { return true } +// OutputVars returns a VarSet containing the variables that would be bound by evaluating +// the body. +func (body Body) OutputVars(safe VarSet) VarSet { + o := safe.Copy() + for _, e := range body { + o.Update(e.OutputVars(o)) + } + return o.Diff(safe) +} + func (body Body) String() string { var buf []string for _, v := range body { @@ -246,10 +256,13 @@ func (body Body) String() string { return strings.Join(buf, ", ") } -// Vars returns map where keys represent all of the variables found in the -// body. The values of the map are ignored. -func (body Body) Vars() VarSet { - vis := &varVisitor{vars: VarSet{}} +// Vars returns a VarSet containing all of the variables in the body. If skipClosures is true, +// variables contained inside closures within the body will be ignored. +func (body Body) Vars(skipClosures bool) VarSet { + vis := &varVisitor{ + vars: VarSet{}, + skipClosures: skipClosures, + } Walk(vis, body) return vis.vars } @@ -327,46 +340,24 @@ func (expr *Expr) IsGround() bool { return true } -// OutputVars returns the set of variables that would be bound by -// evaluating this expression in isolation. -func (expr *Expr) OutputVars() VarSet { - - result := VarSet{} - if expr.Negated { - return result - } - - vis := &varVisitor{ - skipRefHead: true, - skipObjectKeys: true, - vars: VarSet{}, - } - - switch ts := expr.Terms.(type) { - case *Term: - if r, ok := ts.Value.(Ref); ok { - Walk(vis, r) - } - case []*Term: - b := BuiltinMap[ts[0].Value.(Var)] - for i, t := range ts[1:] { - switch v := t.Value.(type) { - case Object, Array: - if b.UnifiesRecursively(i) { - Walk(vis, v) +// OutputVars returns a VarSet containing variables that would be bound by evaluating +// this expression. +func (expr *Expr) OutputVars(safe VarSet) VarSet { + if !expr.Negated { + switch terms := expr.Terms.(type) { + case *Term: + return expr.outputVarsRefs() + case []*Term: + name := terms[0].Value.(Var) + if b := BuiltinMap[name]; b != nil { + if b.Name.Equal(Equality.Name) { + return expr.outputVarsEquality(safe) } - case Var: - if b.Unifies(i) { - result.Add(v) - } - case Ref: - Walk(vis, v) + return expr.outputVarsBuiltins(b, safe) } } } - - result.Update(vis.vars) - return result + return VarSet{} } func (expr *Expr) String() string { @@ -404,12 +395,73 @@ func (expr *Expr) UnmarshalJSON(bs []byte) error { } // Vars returns a VarSet containing all of the variables in the expression. -func (expr *Expr) Vars() VarSet { - vis := &varVisitor{vars: VarSet{}} +// If skipClosures is true then variables contained inside closures within this +// expression will not be included in the VarSet. +func (expr *Expr) Vars(skipClosures bool) VarSet { + vis := &varVisitor{ + skipClosures: skipClosures, + vars: VarSet{}, + } Walk(vis, expr) return vis.vars } +func (expr *Expr) outputVarsBuiltins(b *Builtin, safe VarSet) VarSet { + + o := expr.outputVarsRefs() + terms := expr.Terms.([]*Term) + + // Check that all input terms are ground or safe. + for i, t := range terms[1:] { + if b.IsTargetPos(i) { + continue + } + if t.Value.IsGround() { + continue + } + vis := &varVisitor{ + skipClosures: true, + skipObjectKeys: true, + skipRefHead: true, + skipBuiltinNames: true, + vars: VarSet{}, + } + Walk(vis, t) + unsafe := vis.vars.Diff(o).Diff(safe) + if len(unsafe) > 0 { + return VarSet{} + } + } + + // Add vars in target positions to result. + for i, t := range terms[1:] { + if v, ok := t.Value.(Var); ok { + if b.IsTargetPos(i) { + o.Add(v) + } + } + } + + return o +} + +func (expr *Expr) outputVarsEquality(safe VarSet) VarSet { + ts := expr.Terms.([]*Term) + o := expr.outputVarsRefs() + o.Update(safe) + o.Update(Unify(o, ts[1], ts[2])) + return o.Diff(safe) +} + +func (expr *Expr) outputVarsRefs() VarSet { + o := VarSet{} + WalkRefs(expr, func(r Ref) bool { + o.Update(r.OutputVars()) + return false + }) + return o +} + // NewBuiltinExpr creates a new Expr object with the supplied terms. // The builtin operator must be the first term. func NewBuiltinExpr(terms ...*Term) *Expr { @@ -417,9 +469,11 @@ func NewBuiltinExpr(terms ...*Term) *Expr { } type varVisitor struct { - skipRefHead bool - skipObjectKeys bool - vars VarSet + skipRefHead bool + skipObjectKeys bool + skipClosures bool + skipBuiltinNames bool + vars VarSet } func (vis *varVisitor) Visit(v interface{}) Visitor { @@ -439,6 +493,22 @@ func (vis *varVisitor) Visit(v interface{}) Visitor { return nil } } + if vis.skipClosures { + switch v.(type) { + case *ArrayComprehension: + return nil + } + } + if vis.skipBuiltinNames { + if v, ok := v.(*Expr); ok { + if ts, ok := v.Terms.([]*Term); ok { + for _, t := range ts[1:] { + Walk(vis, t) + } + return nil + } + } + } if v, ok := v.(Var); ok { vis.vars.Add(v) } diff --git a/ast/policy_test.go b/ast/policy_test.go index 40b8a95f64..daa823b4f8 100644 --- a/ast/policy_test.go +++ b/ast/policy_test.go @@ -140,12 +140,39 @@ func TestBodyIsGround(t *testing.T) { } func TestExprOutputVars(t *testing.T) { - body := MustParseBody(`{"a": [{x: y}, b[z]]} = c[i], [{"a": d[j][k]}] != xs`) - one := body[0] - vars := one.OutputVars() - expected := NewVarSet(Var("y"), Var("z"), Var("i")) - if !reflect.DeepEqual(expected, vars) { - t.Errorf("Expected output vars %v from %v but got: %v", expected, one, vars) + + tests := []struct { + note string + expr string + safe string + expected string + }{ + {"ref 1", "a[i].b[j]", "[a]", "[i, j]"}, + {"ref 2", "[1,2,a[i]]", "[a]", "[i]"}, + {"simple unify", `{"a": [{x: y}, b[z]]} = c[i]`, "[b, c]", "[y, z, i]"}, + {"built-in", "count([], x)", "[]", "[x]"}, + } + + for i, tc := range tests { + + expr := MustParseBody(tc.expr)[0] + safe := VarSet{} + for _, x := range MustParseTerm(tc.safe).Value.(Array) { + safe.Add(x.Value.(Var)) + } + + result := expr.OutputVars(safe) + + expected := VarSet{} + for _, x := range MustParseTerm(tc.expected).Value.(Array) { + expected.Add(x.Value.(Var)) + } + + missing := expected.Diff(result) + extra := result.Diff(expected) + if len(missing) != 0 || len(extra) != 0 { + t.Errorf("%s (%d): Missing output vars: %v, extra output vars: %v", tc.note, i, missing, extra) + } } } diff --git a/ast/term.go b/ast/term.go index b98841d82b..63043a216c 100644 --- a/ast/term.go +++ b/ast/term.go @@ -132,6 +132,13 @@ func (term *Term) UnmarshalJSON(bs []byte) error { return nil } +// Vars returns a VarSet with variables contained in this term. +func (term *Term) Vars() VarSet { + vis := &varVisitor{vars: VarSet{}} + Walk(vis, term) + return vis.vars +} + // Null represents the null value defined by JSON. type Null struct{} @@ -415,6 +422,17 @@ func (ref Ref) Underlying() ([]interface{}, error) { return r, nil } +// OutputVars returns a VarSet containing variables that would be bound by evaluating +// this expression in isolation. +func (ref Ref) OutputVars() VarSet { + vis := &varVisitor{ + vars: VarSet{}, + skipRefHead: true, + } + Walk(vis, ref) + return vis.vars +} + // QueryIterator defines the interface for querying AST documents with references. type QueryIterator func(map[Var]Value, Value) error diff --git a/ast/unify.go b/ast/unify.go new file mode 100644 index 0000000000..9b6b48235b --- /dev/null +++ b/ast/unify.go @@ -0,0 +1,165 @@ +// 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 + +// Unify returns a set of variables that will be unified when the equality expression defined by +// terms a and b is evaluated. The unifier assumes that variables in the VarSet safe are already +// unified. +func Unify(safe VarSet, a *Term, b *Term) VarSet { + u := &unifier{ + safe: safe, + unified: VarSet{}, + unknown: map[Var]VarSet{}, + } + u.unify(a, b) + return u.unified +} + +type unifier struct { + safe VarSet + unified VarSet + unknown map[Var]VarSet +} + +func (u *unifier) isSafe(x Var) bool { + return u.safe.Contains(x) || u.unified.Contains(x) +} + +func (u *unifier) unify(a *Term, b *Term) { + + switch a := a.Value.(type) { + + case Var: + switch b := b.Value.(type) { + case Var: + if u.isSafe(b) { + u.markSafe(a) + } else if u.isSafe(a) { + u.markSafe(b) + } else { + u.markUnknown(a, b) + u.markUnknown(b, a) + } + case Array, Object: + u.unifyAll(a, b) + default: + u.markSafe(a) + } + + case Ref: + switch b := b.Value.(type) { + case Var: + u.markSafe(b) + case Array, Object: + u.markAllSafe(b, a) + } + + case *ArrayComprehension: + switch b := b.Value.(type) { + case Var: + u.markSafe(b) + case Array: + u.markAllSafe(b, a) + } + + case Array: + switch b := b.Value.(type) { + case Var: + u.unifyAll(b, a) + case Ref, *ArrayComprehension: + u.markAllSafe(a, b) + case Array: + if len(a) == len(b) { + for i := range a { + u.unify(a[i], b[i]) + } + } + } + + case Object: + switch b := b.Value.(type) { + case Var: + u.unifyAll(b, a) + case Ref: + u.markAllSafe(a, b) + case Object: + if len(a) == len(b) { + for i := range a { + u.unify(a[i][1], b[i][1]) + } + } + } + + default: + switch b := b.Value.(type) { + case Var: + u.markSafe(b) + } + } +} + +func (u *unifier) markAllSafe(x Value, y Value) { + vis := u.varVisitor() + Walk(vis, x) + for v := range vis.vars { + u.markSafe(v) + } +} + +func (u *unifier) markSafe(x Var) { + u.unified.Add(x) + + // Add dependencies of 'x' to safe set + vs := u.unknown[x] + delete(u.unknown, x) + for v := range vs { + u.markSafe(v) + } + + // Add dependants of 'x' to safe set if they have no more + // dependencies. + for v, deps := range u.unknown { + if deps.Contains(x) { + delete(deps, x) + if len(deps) == 0 { + u.markSafe(v) + } + } + } +} + +func (u *unifier) markUnknown(a, b Var) { + if _, ok := u.unknown[a]; !ok { + u.unknown[a] = NewVarSet() + } + u.unknown[a].Add(b) +} + +func (u *unifier) unifyAll(a Var, b Value) { + if u.isSafe(a) { + u.markAllSafe(b, a) + } else { + vis := u.varVisitor() + Walk(vis, b) + unsafe := vis.vars.Diff(u.safe).Diff(u.unified) + if len(unsafe) == 0 { + u.markSafe(a) + } else { + for v := range unsafe { + u.markUnknown(a, v) + } + } + } +} + +func (u *unifier) varVisitor() *varVisitor { + return &varVisitor{ + skipRefHead: true, + skipObjectKeys: true, + skipClosures: true, + skipBuiltinNames: true, + vars: VarSet{}, + } +} diff --git a/ast/unify_test.go b/ast/unify_test.go new file mode 100644 index 0000000000..9ee34ab94c --- /dev/null +++ b/ast/unify_test.go @@ -0,0 +1,78 @@ +// 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 "testing" + +func TestUnify(t *testing.T) { + + tests := []struct { + note string + expr string + safe string + expected string + }{ + // collection cases + {"array/ref", "[1,2,x] = a[_]", "[a]", "[x]"}, + {"array/ref (reversed)", "a[_] = [1,2,x]", "[a]", "[x]"}, + {"array/var", "[1,2,x] = y", "[x]", "[y]"}, + {"array/var (reversed)", "y = [1,2,x]", "[x]", "[y]"}, + {"array/var-2", "[1,2,x] = y", "[y]", "[x]"}, + {"array/var-2 (reversed)", "y = [1,2,x]", "[y]", "[x]"}, + {"array/uneven", "[1,2,x] = [y,x]", "[]", "[]"}, + {"array/uneven-2", "[1,2,x] = [y,x]", "[x]", "[]"}, + {"object/ref", `{"x": x} = a[_]`, "[a]", "[x]"}, + {"object/ref (reversed)", `a[_] = {"x": x}`, "[a]", "[x]"}, + {"object/var", `{"x": 1, "y": x} = y`, "[x]", "[y]"}, + {"object/var (reversed)", `y = {"x": 1, "y": x}`, "[x]", "[y]"}, + {"object/var-2", `{"x": 1, "y": x} = y`, "[y]", "[x]"}, + {"object/var-3", `{"x": 1, "y": x} = y`, "[]", "[]"}, + {"object/uneven", `{"x": x, "y": 1} = {"x": y}`, "[]", "[]"}, + {"object/uneven", `{"x": x, "y": 1} = {"x": y}`, "[x]", "[]"}, + + // transitive cases + {"trans/redundant", "[x, x] = [x, 0]", "[]", "[x]"}, + {"trans/simple", "[x, 1] = [y, y]", "[]", "[y, x]"}, + {"trans/array", "[x, y] = [y, [z, a]]", "[x]", "[a, y, z]"}, + {"trans/object", `[x, y] = [y, {"a":a,"z":z}]`, "[x]", "[a, y, z]"}, + {"trans/ref", "[x, y, [x, y, i]] = [1, a[i], z]", "[a, i]", "[x, y, z]"}, + {"trans/lazy", "[x, z, 2] = [1, [y, x], y]", "[]", "[x, y, z]"}, + {"trans/redundant-nested", "[x, z, z] = [1, [y, x], [2, 1]]", "[]", "[x, y, z]"}, + {"trans/bidirectional", "[x, z, y] = [[z,y], [1,y], 2]", "[]", "[x, y, z]"}, + {"trans/occurs", "[x, z, y] = [[y,z], [y, 1], [2, x]]", "[]", "[]"}, + } + + for i, tc := range tests { + + expr := MustParseBody(tc.expr)[0] + safe := VarSet{} + for _, x := range MustParseTerm(tc.safe).Value.(Array) { + safe.Add(x.Value.(Var)) + } + + terms := expr.Terms.([]*Term) + if !terms[0].Value.Equal(Equality.Name) { + panic(terms) + } + + a, b := terms[1], terms[2] + unified := Unify(safe, a, b) + result := VarSet{} + for k := range unified { + result.Add(k) + } + + expected := VarSet{} + for _, x := range MustParseTerm(tc.expected).Value.(Array) { + expected.Add(x.Value.(Var)) + } + + missing := expected.Diff(result) + extra := result.Diff(expected) + if len(missing) != 0 || len(extra) != 0 { + t.Errorf("%s (%d): Missing vars: %v, extra vars: %v", tc.note, i, missing, extra) + } + } +} diff --git a/ast/varset.go b/ast/varset.go index f9f7ecb291..176c9ea2e2 100644 --- a/ast/varset.go +++ b/ast/varset.go @@ -41,6 +41,28 @@ func (s VarSet) Copy() VarSet { return cpy } +// Diff returns a VarSet containing variables in s that are not in vs. +func (s VarSet) Diff(vs VarSet) VarSet { + r := VarSet{} + for v := range s { + if !vs.Contains(v) { + r.Add(v) + } + } + return r +} + +// Intersect returns a VarSet containing variables in s that are in vs. +func (s VarSet) Intersect(vs VarSet) VarSet { + r := VarSet{} + for v := range s { + if vs.Contains(v) { + r.Add(v) + } + } + return r +} + // Update merges the other VarSet into this VarSet. func (s VarSet) Update(vs VarSet) { for v := range vs { diff --git a/ast/visit.go b/ast/visit.go index f249348c04..8cfda36c1c 100644 --- a/ast/visit.go +++ b/ast/visit.go @@ -77,3 +77,43 @@ func Walk(v Visitor, x interface{}) { Walk(w, x.Body) } } + +// WalkClosures calls the function f on all closures under x. If the function f +// returns true, AST nodes under the last node will not be visited. +func WalkClosures(x interface{}, f func(interface{}) bool) { + vis := &GenericVisitor{func(x interface{}) bool { + switch x.(type) { + case *ArrayComprehension: + return f(x) + } + return false + }} + Walk(vis, x) +} + +// WalkRefs calls the function f on all references under x. If the function f +// returns true, AST nodes under the last node will not be visited. +func WalkRefs(x interface{}, f func(Ref) bool) { + vis := &GenericVisitor{func(x interface{}) bool { + if r, ok := x.(Ref); ok { + return f(r) + } + return false + }} + Walk(vis, x) +} + +// GenericVisitor implements the Visitor interface to provide +// a utility to walk over AST nodes using a closure. If the closure +// returns true, the visitor will not walk over AST nodes under x. +type GenericVisitor struct { + f func(x interface{}) bool +} + +// Visit calls the function f on the GenericVisitor. +func (vis *GenericVisitor) Visit(x interface{}) Visitor { + if vis.f(x) { + return nil + } + return vis +} diff --git a/topdown/topdown.go b/topdown/topdown.go index d50677f325..5e680f454c 100644 --- a/topdown/topdown.go +++ b/topdown/topdown.go @@ -78,15 +78,6 @@ func (ctx *Context) BindVar(variable ast.Var, value ast.Value) *Context { if variable.Equal(value) { return ctx } - occurs := walkValue(value, func(other ast.Value) bool { - if variable.Equal(other) { - return true - } - return false - }) - if occurs { - return nil - } cpy := *ctx cpy.Locals = storage.NewBindings() @@ -389,23 +380,6 @@ func dereferenceVar(v ast.Var, ctx *Context) (interface{}, error) { func evalContext(ctx *Context, iter Iterator) error { if ctx.Index >= len(ctx.Query) { - - // Check if the bindings contain values that are non-ground. E.g., - // suppose the query's final expression is "x = y" and "x" and "y" - // do not appear elsewhere in the query. In this case, "x" and "y" - // will be bound to each other; they will not be ground and so - // the proof should not be considered successful. - isNonGround := ctx.Locals.Iter(func(k, v ast.Value) bool { - if !v.IsGround() { - return true - } - return false - }) - - if isNonGround { - return nil - } - ctx.traceFinish() return iter(ctx) } diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index 7dbb099736..cdd8e9720e 100644 --- a/topdown/topdown_test.go +++ b/topdown/topdown_test.go @@ -234,7 +234,6 @@ func TestTopDownCompleteDoc(t *testing.T) { {`object/nested composites: {"a": [1], "b": [2], "c": [3]}`, `p = {"a": [1], "b": [2], "c": [3]} :- true`, `{"a": [1], "b": [2], "c": [3]}`}, - {"var/var", "p = true :- x = y", ""}, } data := loadSmallTestData() @@ -258,7 +257,6 @@ func TestTopDownPartialSetDoc(t *testing.T) { {"nested composites", "p[x] :- f[i] = x", `[{"xs": [1.0], "ys": [2.0]}, {"xs": [2.0], "ys": [3.0]}]`}, {"deep ref/heterogeneous", "p[x] :- c[i][j][k] = x", `[null, 3.14159, true, false, true, false, "foo"]`}, {"composite var value", "p[x] :- x = [i, a[i]]", "[[0,1],[1,2],[2,3],[3,4]]"}, - {"var/var", "p[x] :- x = y", "[]"}, } data := loadSmallTestData() @@ -278,8 +276,6 @@ func TestTopDownPartialObjectDoc(t *testing.T) { {"composites", "p[k] = v :- d[k] = v", `{"e": ["bar", "baz"]}`}, {"non-string key", "p[k] = v :- a[k] = v", fmt.Errorf("illegal object key type float64: 0")}, {"body/join var", "p[k] = v :- a[i] = v, g[k][i] = v", `{"a": 1, "b": 2, "c": 4}`}, - {"var/var key", "p[k] = v :- v = 1, k = x", "{}"}, - {"var/var val", `p[k] = v :- k = "x", v = x`, "{}"}, } data := loadSmallTestData() @@ -310,8 +306,6 @@ func TestTopDownEqExpr(t *testing.T) { {"undefined: array deep var 2", "p = true :- [[1,x],[3,4]] = [[1,2],[x,4]]", ""}, {"undefined: array uneven", `p = true :- [true, false, "foo", "deadbeef"] = c[i][j]`, ""}, {"undefined: object uneven", `p = true :- {"a": 1, "b": 2} = {"a": 1}`, ""}, - {"undefined: occurs 1", "p = true :- [y,x] = [[x],y]", ""}, - {"undefined: occurs 2", "p = true :- [y,x] = [{\"a\": x}, y]", ""}, // ground terms {"ground: bool", `p = true :- true = true`, "true"}, From 1e628e33dfa4adf3b087f4dde0b5c3774a5cc2b5 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Fri, 10 Jun 2016 14:26:51 -0700 Subject: [PATCH 06/11] Resolve references within array comprehensions --- ast/compile.go | 49 ++++++++++++----- ast/compile_test.go | 128 +++++++++++++++++++++++++++++--------------- 2 files changed, 120 insertions(+), 57 deletions(-) diff --git a/ast/compile.go b/ast/compile.go index 6888ce1f0d..08dcc217b8 100644 --- a/ast/compile.go +++ b/ast/compile.go @@ -231,16 +231,7 @@ func (c *Compiler) err(f string, a ...interface{}) { func (c *Compiler) resolveAllRefs() { for _, mod := range c.Modules { for _, rule := range mod.Rules { - for _, expr := range rule.Body { - switch ts := expr.Terms.(type) { - case *Term: - expr.Terms = c.resolveRefs(c.Globals[mod], ts) - case []*Term: - for i, t := range ts { - ts[i] = c.resolveRefs(c.Globals[mod], t) - } - } - } + rule.Body = c.resolveRefsInBody(c.Globals[mod], rule.Body) } for i := range mod.Imports { mod.Imports[i].Alias = Var("") @@ -287,7 +278,30 @@ func (c *Compiler) resolveRef(globals map[Var]Value, ref Ref) Ref { return fqn } -func (c *Compiler) resolveRefs(globals map[Var]Value, term *Term) *Term { +func (c *Compiler) resolveRefsInBody(globals map[Var]Value, body Body) Body { + r := Body{} + for _, expr := range body { + r = append(r, c.resolveRefsInExpr(globals, expr)) + } + return r +} + +func (c *Compiler) resolveRefsInExpr(globals map[Var]Value, expr *Expr) *Expr { + cpy := *expr + switch ts := expr.Terms.(type) { + case *Term: + cpy.Terms = c.resolveRefsInTerm(globals, ts) + case []*Term: + buf := []*Term{} + for _, t := range ts { + buf = append(buf, c.resolveRefsInTerm(globals, t)) + } + cpy.Terms = buf + } + return &cpy +} + +func (c *Compiler) resolveRefsInTerm(globals map[Var]Value, term *Term) *Term { switch v := term.Value.(type) { case Var: if r, ok := globals[v]; ok { @@ -304,8 +318,8 @@ func (c *Compiler) resolveRefs(globals map[Var]Value, term *Term) *Term { case Object: o := Object{} for _, i := range v { - k := c.resolveRefs(globals, i[0]) - v := c.resolveRefs(globals, i[1]) + k := c.resolveRefsInTerm(globals, i[0]) + v := c.resolveRefsInTerm(globals, i[1]) o = append(o, Item(k, v)) } cpy := *term @@ -314,12 +328,19 @@ func (c *Compiler) resolveRefs(globals map[Var]Value, term *Term) *Term { case Array: a := Array{} for _, e := range v { - x := c.resolveRefs(globals, e) + x := c.resolveRefsInTerm(globals, e) a = append(a, x) } cpy := *term cpy.Value = a return &cpy + case *ArrayComprehension: + ac := &ArrayComprehension{} + ac.Term = c.resolveRefsInTerm(globals, v.Term) + ac.Body = c.resolveRefsInBody(globals, v.Body) + cpy := *term + cpy.Value = ac + return &cpy default: return term } diff --git a/ast/compile_test.go b/ast/compile_test.go index 5da555d0f2..65ef2ef1c1 100644 --- a/ast/compile_test.go +++ b/ast/compile_test.go @@ -17,7 +17,7 @@ func TestModuleTree(t *testing.T) { mods := getCompilerTestModules() tree := NewModuleTree(mods) - if tree.Size() != 4 { + if tree.Size() != 5 { t.Errorf("Expected size of 4 in module tree but got: %v", tree.Size()) } @@ -120,47 +120,6 @@ func TestCompilerSetGlobals(t *testing.T) { bar: data.bar}`) } -func TestCompilerResolveAllRefs(t *testing.T) { - c := NewCompiler() - c.Modules = getCompilerTestModules() - compileStages(c, "", "resolveAllRefs") - - assertNotFailed(t, c) - - mod1 := c.Modules["mod1"] - p := mod1.Rules[0] - expr1 := p.Body[0] - term := expr1.Terms.(*Term) - e := MustParseTerm("data.a.b.c.q[x]") - if !term.Equal(e) { - t.Errorf("Wrong term (global in same module): expected %v but got: %v", e, term) - } - - expr2 := p.Body[1] - term = expr2.Terms.(*Term) - e = MustParseTerm("data.a.b.c.r[x]") - if !term.Equal(e) { - t.Errorf("Wrong term (global in same package/diff module): expected %v but got: %v", e, term) - } - - mod2 := c.Modules["mod2"] - r := mod2.Rules[0] - expr3 := r.Body[1] - term = expr3.Terms.([]*Term)[1] - e = MustParseTerm("data.x.y.p") - if !term.Equal(e) { - t.Errorf("Wrong term (var import): expected %v but got: %v", e, term) - } - - mod3 := c.Modules["mod3"] - expr4 := mod3.Rules[0].Body[0] - term = expr4.Terms.([]*Term)[2] - e = MustParseTerm("{x.secret: [x.keyid]}") - if !term.Equal(e) { - t.Errorf("Wrong term (nested refs): expected %v but got: %v", e, term) - } -} - func TestCompilerCheckSafetyHead(t *testing.T) { c := NewCompiler() c.Modules = getCompilerTestModules() @@ -384,6 +343,68 @@ func TestCompilerCheckSafetyBodyErrors(t *testing.T) { } +func TestCompilerResolveAllRefs(t *testing.T) { + c := NewCompiler() + c.Modules = getCompilerTestModules() + compileStages(c, "", "resolveAllRefs") + + assertNotFailed(t, c) + + mod1 := c.Modules["mod1"] + p := mod1.Rules[0] + expr1 := p.Body[0] + term := expr1.Terms.(*Term) + e := MustParseTerm("data.a.b.c.q[x]") + if !term.Equal(e) { + t.Errorf("Wrong term (global in same module): expected %v but got: %v", e, term) + } + + expr2 := p.Body[1] + term = expr2.Terms.(*Term) + e = MustParseTerm("data.a.b.c.r[x]") + if !term.Equal(e) { + t.Errorf("Wrong term (global in same package/diff module): expected %v but got: %v", e, term) + } + + mod2 := c.Modules["mod2"] + r := mod2.Rules[0] + expr3 := r.Body[1] + term = expr3.Terms.([]*Term)[1] + e = MustParseTerm("data.x.y.p") + if !term.Equal(e) { + t.Errorf("Wrong term (var import): expected %v but got: %v", e, term) + } + + mod3 := c.Modules["mod3"] + expr4 := mod3.Rules[0].Body[0] + term = expr4.Terms.([]*Term)[2] + e = MustParseTerm("{x.secret: [x.keyid]}") + if !term.Equal(e) { + t.Errorf("Wrong term (nested refs): expected %v but got: %v", e, term) + } + + // Array comprehensions. + mod5 := c.Modules["mod5"] + + ac := func(r *Rule) *ArrayComprehension { + return r.Body[0].Terms.(*Term).Value.(*ArrayComprehension) + } + + acTerm1 := ac(mod5.Rules[0]) + assertTermEqual(t, acTerm1.Term, MustParseTerm("x.a")) + acTerm2 := ac(mod5.Rules[1]) + assertTermEqual(t, acTerm2.Term, MustParseTerm("a.b.c.q.a")) + acTerm3 := ac(mod5.Rules[2]) + assertTermEqual(t, acTerm3.Body[0].Terms.([]*Term)[1], MustParseTerm("x.a")) + acTerm4 := ac(mod5.Rules[3]) + assertTermEqual(t, acTerm4.Body[0].Terms.([]*Term)[1], MustParseTerm("a.b.c.q[i]")) + acTerm5 := ac(mod5.Rules[4]) + assertTermEqual(t, acTerm5.Body[0].Terms.([]*Term)[2].Value.(*ArrayComprehension).Term, MustParseTerm("x.a")) + acTerm6 := ac(mod5.Rules[5]) + assertTermEqual(t, acTerm6.Body[0].Terms.([]*Term)[2].Value.(*ArrayComprehension).Body[0].Terms.([]*Term)[1], MustParseTerm("a.b.c.q[i]")) + +} + func TestCompilerSetRuleGraph(t *testing.T) { c := NewCompiler() c.Modules = getCompilerTestModules() @@ -438,6 +459,11 @@ func TestCompilerCheckRecursion(t *testing.T) { import data.rec3.p q[x] = y :- p[x] = y `), + "newMod6": MustParseModule(` + package rec5 + acp[x] :- acq[x] + acq[x] :- a = [x | acp[x]], a[i] = x + `), } compileStages(c, "", "checkRecursion") @@ -451,6 +477,8 @@ func TestCompilerCheckRecursion(t *testing.T) { fmt.Errorf("recursion found in e: e, a, b, c, e"), fmt.Errorf("recursion found in p: p, q, p"), fmt.Errorf("recursion found in q: q, p, q"), + fmt.Errorf("recursion found in acq: acq, acp, acq"), + fmt.Errorf("recursion found in acp: acp, acq, acp"), } if len(c.Errors) != len(expected) { @@ -630,5 +658,19 @@ func getCompilerTestModules() map[string]*Module { package a.b.empty `) - return map[string]*Module{"mod2": mod2, "mod3": mod3, "mod1": mod1, "mod4": mod4} + mod5 := MustParseModule(` + package a.b.compr + + import x as y + import a.b.c.q + + p :- [y.a | true] + r :- [q.a | true] + s :- [true | y.a = 0] + t :- [true | q[i] = 1] + u :- [true | _ = [y.a | true]] + v :- [true | _ = [ true | q[i] = 1]] + `) + + return map[string]*Module{"mod2": mod2, "mod3": mod3, "mod1": mod1, "mod4": mod4, "mod5": mod5} } From 77653b10c359c3933ed2ddcc8c5207cccd47d1a2 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Fri, 10 Jun 2016 14:47:59 -0700 Subject: [PATCH 07/11] Add support for comprehensions in topdown --- topdown/topdown.go | 109 +++++++++++++++++++++++++++------------- topdown/topdown_test.go | 24 +++++++++ 2 files changed, 99 insertions(+), 34 deletions(-) diff --git a/topdown/topdown.go b/topdown/topdown.go index 5e680f454c..ea15c99d81 100644 --- a/topdown/topdown.go +++ b/topdown/topdown.go @@ -50,11 +50,11 @@ func (ctx *Context) Binding(k ast.Value) ast.Value { return nil } -// BindRef returns a new Context with bindings that map the reference to the value. -func (ctx *Context) BindRef(ref ast.Ref, value ast.Value) *Context { +// BindValue returns a new Context with bindings that map the key to the value. +func (ctx *Context) BindValue(key ast.Value, value ast.Value) *Context { cpy := *ctx cpy.Locals = ctx.Locals.Copy() - cpy.Locals.Put(ref, value) + cpy.Locals.Put(key, value) return &cpy } @@ -97,10 +97,10 @@ func (ctx *Context) BindVar(variable ast.Var, value ast.Value) *Context { return &cpy } -// Child returns a new context to evaluate a rule that was referenced by this context. -func (ctx *Context) Child(rule *ast.Rule, locals *storage.Bindings) *Context { +// Child returns a new context to evaluate a query that was referenced by this context. +func (ctx *Context) Child(query ast.Body, locals *storage.Bindings) *Context { cpy := *ctx - cpy.Query = rule.Body + cpy.Query = query cpy.Locals = locals cpy.Previous = ctx cpy.Index = 0 @@ -602,7 +602,7 @@ func evalRefRuleCompleteDoc(ctx *Context, ref ast.Ref, suffix ast.Ref, rules []* for _, rule := range rules { bindings := storage.NewBindings() - child := ctx.Child(rule, bindings) + child := ctx.Child(rule.Body, bindings) isTrue := false err := Eval(child, func(child *Context) error { @@ -647,7 +647,7 @@ func evalRefRulePartialObjectDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule * // NOTE: if at some point multiple variables are supported here, it may be // cleaner to generalize this (instead of having two separate branches). if !key.IsGround() { - child := ctx.Child(rule, storage.NewBindings()) + child := ctx.Child(rule.Body, storage.NewBindings()) return Eval(child, func(child *Context) error { key := child.Binding(rule.Key.Value) if key == nil { @@ -664,7 +664,7 @@ func evalRefRulePartialObjectDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule * bindings := storage.NewBindings() bindings.Put(rule.Key.Value, key) - child := ctx.Child(rule, bindings) + child := ctx.Child(rule.Body, bindings) return Eval(child, func(child *Context) error { value := child.Binding(rule.Value.Value) @@ -687,7 +687,7 @@ func evalRefRulePartialObjectDocFull(ctx *Context, ref ast.Ref, rules []*ast.Rul for _, rule := range rules { bindings := storage.NewBindings() - child := ctx.Child(rule, bindings) + child := ctx.Child(rule.Body, bindings) err := Eval(child, func(child *Context) error { key := child.Binding(rule.Key.Value) @@ -705,7 +705,7 @@ func evalRefRulePartialObjectDocFull(ctx *Context, ref ast.Ref, rules []*ast.Rul } } - ctx = ctx.BindRef(ref, result) + ctx = ctx.BindValue(ref, result) return iter(ctx) } @@ -725,7 +725,7 @@ func evalRefRulePartialSetDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule *ast key := plugValue(suffix[0].Value, ctx) if !key.IsGround() { - child := ctx.Child(rule, storage.NewBindings()) + child := ctx.Child(rule.Body, storage.NewBindings()) return Eval(child, func(child *Context) error { value := child.Binding(rule.Key.Value) if value == nil { @@ -737,18 +737,18 @@ func evalRefRulePartialSetDoc(ctx *Context, ref ast.Ref, path ast.Ref, rule *ast // "p = true :- q[x]", we say that "p" should be defined if "q" // is defined for some value "x". ctx = ctx.BindVar(key.(ast.Var), value) - ctx = ctx.BindRef(ref[:len(path)+1], ast.Boolean(true)) + ctx = ctx.BindValue(ref[:len(path)+1], ast.Boolean(true)) return iter(ctx) }) } bindings := storage.NewBindings() bindings.Put(rule.Key.Value, key) - child := ctx.Child(rule, bindings) + child := ctx.Child(rule.Body, bindings) return Eval(child, func(child *Context) error { // See comment above for explanation of why the reference is bound to true. - ctx = ctx.BindRef(ref[:len(path)+1], ast.Boolean(true)) + ctx = ctx.BindValue(ref[:len(path)+1], ast.Boolean(true)) return iter(ctx) }) @@ -771,7 +771,7 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val binding = append(binding, result...) binding = append(binding, suffix...) return evalRefRec(ctx, result, suffix, func(ctx *Context) error { - ctx = ctx.BindRef(ref, plugValue(binding, ctx)) + ctx = ctx.BindValue(ref, plugValue(binding, ctx)) return iter(ctx) }) @@ -782,14 +782,14 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val pluggedSuffix = append(pluggedSuffix, plugTerm(t, ctx)) } return result.Query(pluggedSuffix, func(keys map[ast.Var]ast.Value, value ast.Value) error { - ctx = ctx.BindRef(ref, value) + ctx = ctx.BindValue(ref, value) for k, v := range keys { ctx = ctx.BindVar(k, v) } return iter(ctx) }) } - ctx = ctx.BindRef(ref, result) + ctx = ctx.BindValue(ref, result) return iter(ctx) case ast.Object: @@ -799,14 +799,14 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val pluggedSuffix = append(pluggedSuffix, plugTerm(t, ctx)) } return result.Query(pluggedSuffix, func(keys map[ast.Var]ast.Value, value ast.Value) error { - ctx = ctx.BindRef(ref, value) + ctx = ctx.BindValue(ref, value) for k, v := range keys { ctx = ctx.BindVar(k, v) } return iter(ctx) }) } - ctx = ctx.BindRef(ref, result) + ctx = ctx.BindValue(ref, result) return iter(ctx) default: @@ -814,18 +814,12 @@ func evalRefRuleResult(ctx *Context, ref ast.Ref, suffix ast.Ref, result ast.Val // This is not defined because it attempts to dereference a scalar. return nil } - ctx = ctx.BindRef(ref, result) + ctx = ctx.BindValue(ref, result) return iter(ctx) } } -// evalTerms is used to get bindings for variables in individual terms. -// -// Before an expression is evaluated, this function is called to find bindings -// for variables used in references inside the expression. Finding bindings for -// variables used in references involves iterating collections in storage or -// evaluating rules identified by the references. In either case, this function -// will invoke the iterator with each set of bindings that should be evaluated. +// TODO(tsandall): func evalTerms(ctx *Context, iter Iterator) error { expr := ctx.Current() @@ -876,6 +870,25 @@ func evalTerms(ctx *Context, iter Iterator) error { return evalTermsRec(ctx, iter, ts) } +func evalTermsComprehension(ctx *Context, comp ast.Value, iter Iterator) error { + switch comp := comp.(type) { + case *ast.ArrayComprehension: + r := ast.Array{} + c := ctx.Child(comp.Body, ctx.Locals) + err := Eval(c, func(c *Context) error { + r = append(r, plugTerm(comp.Term, c)) + return nil + }) + if err != nil { + return err + } + ctx = ctx.BindValue(comp, r) + return iter(ctx) + default: + panic(fmt.Sprintf("illegal argument: %v %v", ctx, comp)) + } +} + func evalTermsIndexed(ctx *Context, iter Iterator, indexed ast.Ref, nonIndexed *ast.Term) error { iterateIndex := func(ctx *Context) error { @@ -927,6 +940,10 @@ func evalTermsRec(ctx *Context, iter Iterator, ts []*ast.Term) error { return evalTermsRecObject(ctx, head, 0, func(ctx *Context) error { return evalTermsRec(ctx, iter, tail) }) + case *ast.ArrayComprehension: + return evalTermsComprehension(ctx, head, func(ctx *Context) error { + return evalTermsRec(ctx, iter, tail) + }) default: return evalTermsRec(ctx, iter, tail) } @@ -949,6 +966,10 @@ func evalTermsRecArray(ctx *Context, arr ast.Array, idx int, iter Iterator) erro return evalTermsRecObject(ctx, v, 0, func(ctx *Context) error { return evalTermsRecArray(ctx, arr, idx+1, iter) }) + case *ast.ArrayComprehension: + return evalTermsComprehension(ctx, v, func(ctx *Context) error { + return evalTermsRecArray(ctx, arr, idx+1, iter) + }) default: return evalTermsRecArray(ctx, arr, idx+1, iter) } @@ -974,6 +995,10 @@ func evalTermsRecObject(ctx *Context, obj ast.Object, idx int, iter Iterator) er return evalTermsRecObject(ctx, v, 0, func(ctx *Context) error { return evalTermsRecObject(ctx, obj, idx+1, iter) }) + case *ast.ArrayComprehension: + return evalTermsComprehension(ctx, v, func(ctx *Context) error { + return evalTermsRecObject(ctx, obj, idx+1, iter) + }) default: return evalTermsRecObject(ctx, obj, idx+1, iter) } @@ -992,6 +1017,10 @@ func evalTermsRecObject(ctx *Context, obj ast.Object, idx int, iter Iterator) er return evalTermsRecObject(ctx, v, 0, func(ctx *Context) error { return evalTermsRecObject(ctx, obj, idx+1, iter) }) + case *ast.ArrayComprehension: + return evalTermsComprehension(ctx, v, func(ctx *Context) error { + return evalTermsRecObject(ctx, obj, idx+1, iter) + }) default: return evalTermsRecObject(ctx, obj, idx+1, iter) } @@ -1139,6 +1168,11 @@ func plugTerm(term *ast.Term, ctx *Context) *ast.Term { plugged.Value = plugValue(v, ctx) return &plugged + case *ast.ArrayComprehension: + plugged := *term + plugged.Value = plugValue(v, ctx) + return &plugged + default: if !term.IsGround() { panic("unreachable") @@ -1151,15 +1185,22 @@ func plugValue(v ast.Value, ctx *Context) ast.Value { switch v := v.(type) { case ast.Var: - binding := ctx.Binding(v) - if binding == nil { + b := ctx.Binding(v) + if b == nil { return v } - return binding.(ast.Value) + return b + + case *ast.ArrayComprehension: + b := ctx.Binding(v) + if b == nil { + return v + } + return b case ast.Ref: - if binding := ctx.Binding(v); binding != nil { - return binding.(ast.Value) + if b := ctx.Binding(v); b != nil { + return b } if v.IsGround() { return v @@ -1189,7 +1230,7 @@ func plugValue(v ast.Value, ctx *Context) ast.Value { default: if !v.IsGround() { - panic("unreachable") + panic(fmt.Sprintf("illegal value: %v %v", ctx, v)) } return v } diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index cdd8e9720e..13fe7cef48 100644 --- a/topdown/topdown_test.go +++ b/topdown/topdown_test.go @@ -554,6 +554,30 @@ func TestTopDownNegation(t *testing.T) { } } +func TestTopDownAggregates(t *testing.T) { + + tests := []struct { + note string + rules []string + expected interface{} + }{ + {"simple", []string{"p[i] :- xs = [x | x = a[_]], xs[i] > 1"}, "[1,2,3]"}, + {"nested", []string{"p[i] :- ys = [y | y = x[_], x = [z | z = a[_]]], ys[i] > 1"}, "[1,2,3]"}, + {"embedded array", []string{"p[i] :- xs = [[x | x = a[_]]], xs[0][i] > 1"}, "[1,2,3]"}, + {"embedded object", []string{`p[i] :- xs = {"a": [x | x = a[_]]}, xs["a"][i] > 1`}, "[1,2,3]"}, + {"recursive", []string{"p :- y = 1, x = y, x = [y | y = 1]"}, ""}, + // TODO(tsandall): semantics? + {"recursive", []string{"p :- x = y, x = [y | y = 1]"}, "true"}, + {"recursive", []string{"p :- x = [x | x = 1]"}, "true"}, + } + + data := loadSmallTestData() + + for i, tc := range tests { + runTopDownTestCase(t, data, i, tc.note, tc.rules, tc.expected) + } +} + func TestTopDownEmbeddedVirtualDoc(t *testing.T) { mods := compileModules([]string{ From a593a8a581939a054d50530f7ab3f548f73f60c3 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Thu, 26 May 2016 17:15:54 -0700 Subject: [PATCH 08/11] Add built-ins for basic aggregation support --- ast/builtins.go | 53 ++++++++++++++++++++++++++++++++ topdown/aggregates.go | 68 +++++++++++++++++++++++++++++++++++++++++ topdown/arithmetic.go | 38 +++++++++++++++++++++++ topdown/builtins.go | 4 +++ topdown/casts.go | 55 +++++++++++++++++++++++++++++++++ topdown/topdown_test.go | 61 +++++++++++++++++++++++++++++++++--- 6 files changed, 274 insertions(+), 5 deletions(-) create mode 100644 topdown/aggregates.go create mode 100644 topdown/arithmetic.go create mode 100644 topdown/casts.go diff --git a/ast/builtins.go b/ast/builtins.go index 24d8bb5e64..8df45e985a 100644 --- a/ast/builtins.go +++ b/ast/builtins.go @@ -21,12 +21,19 @@ func RegisterBuiltin(b *Builtin) { var DefaultBuiltins = [...]*Builtin{ Equality, GreaterThan, GreaterThanEq, LessThan, LessThanEq, NotEqual, + Plus, + Count, Sum, + ToNumber, } // BuiltinMap provides a convenient mapping of built-in names to // built-in definitions. var BuiltinMap map[Var]*Builtin +/** + * Unification + */ + // Equality represents the "=" operator. var Equality = &Builtin{ Name: Var("="), @@ -35,6 +42,10 @@ var Equality = &Builtin{ TargetPos: []int{0, 1}, } +/** + * Comparisons + */ + // GreaterThan represents the ">" comparison operator. var GreaterThan = &Builtin{ Name: Var(">"), @@ -70,6 +81,48 @@ var NotEqual = &Builtin{ NumArgs: 2, } +/** + * Arithmetic + */ + +// Plus adds two numbers together. +var Plus = &Builtin{ + Name: Var("plus"), + NumArgs: 3, + TargetPos: []int{2}, +} + +/** + * Aggregates + */ + +// Count takes a collection and counts the number of elements in it. +var Count = &Builtin{ + Name: Var("count"), + NumArgs: 2, + TargetPos: []int{1}, +} + +// Sum takes an array of numbers and sums them. +var Sum = &Builtin{ + Name: Var("sum"), + NumArgs: 2, + TargetPos: []int{1}, +} + +/** + * Casting + */ + +// ToNumber takes a string, bool, or number value and converts it to a number. +// Strings are converted to numbers using strconv.Atoi. +// Boolean false is converted to 0 and boolean true is converted to 1. +var ToNumber = &Builtin{ + Name: Var("to_number"), + NumArgs: 2, + TargetPos: []int{1}, +} + // Builtin represents a built-in function supported by OPA. Every // built-in function is uniquely identified by a name. type Builtin struct { diff --git a/topdown/aggregates.go b/topdown/aggregates.go new file mode 100644 index 0000000000..8de8df78f1 --- /dev/null +++ b/topdown/aggregates.go @@ -0,0 +1,68 @@ +// 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 ( + "fmt" + + "github.com/open-policy-agent/opa/ast" + "github.com/pkg/errors" +) + +func evalCount(ctx *Context, expr *ast.Expr, iter Iterator) error { + ops := expr.Terms.([]*ast.Term) + src, dst := ops[1].Value, ops[2].Value + s, err := ValueToInterface(src, ctx) + if err != nil { + return errors.Wrapf(err, "count") + } + + var count ast.Number + + switch s := s.(type) { + case []interface{}: + count = ast.Number(len(s)) + case map[string]interface{}: + count = ast.Number(len(s)) + default: + return fmt.Errorf("count: source must be a collection: %v", src) + } + + switch dst := dst.(type) { + case ast.Var: + ctx = ctx.BindVar(dst, count) + return iter(ctx) + default: + if dst.Equal(count) { + return iter(ctx) + } + return nil + } +} + +func evalSum(ctx *Context, expr *ast.Expr, iter Iterator) error { + ops := expr.Terms.([]*ast.Term) + src, dst := ops[1].Value, ops[2].Value + s, err := ValueToSlice(src, ctx) + if err != nil { + return errors.Wrapf(err, "sum") + } + + sum := ast.Number(0) + for _, x := range s { + sum += ast.Number(x.(float64)) + } + + switch dst := dst.(type) { + case ast.Var: + ctx = ctx.BindVar(dst, sum) + return iter(ctx) + default: + if dst.Equal(sum) { + return iter(ctx) + } + return nil + } +} diff --git a/topdown/arithmetic.go b/topdown/arithmetic.go new file mode 100644 index 0000000000..53d5b394e9 --- /dev/null +++ b/topdown/arithmetic.go @@ -0,0 +1,38 @@ +// 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/ast" + "github.com/pkg/errors" +) + +func evalPlus(ctx *Context, expr *ast.Expr, iter Iterator) error { + ops := expr.Terms.([]*ast.Term) + + a, err := ValueToFloat64(ops[1].Value, ctx) + if err != nil { + return errors.Wrapf(err, "plus") + } + + b, err := ValueToFloat64(ops[2].Value, ctx) + if err != nil { + return errors.Wrapf(err, "plus") + } + + c := ops[3].Value + r := ast.Number(a + b) + + switch c := c.(type) { + case ast.Var: + ctx = ctx.BindVar(c, r) + return iter(ctx) + default: + if r.Equal(c) { + return iter(ctx) + } + return nil + } +} diff --git a/topdown/builtins.go b/topdown/builtins.go index b0b602f993..187e96959a 100644 --- a/topdown/builtins.go +++ b/topdown/builtins.go @@ -31,6 +31,10 @@ var defaultBuiltinFuncs = map[ast.Var]BuiltinFunc{ ast.LessThan.Name: evalIneq(compareLessThan), ast.LessThanEq.Name: evalIneq(compareLessThanEq), ast.NotEqual.Name: evalIneq(compareNotEq), + ast.Plus.Name: evalPlus, + ast.Count.Name: evalCount, + ast.Sum.Name: evalSum, + ast.ToNumber.Name: evalToNumber, } func init() { diff --git a/topdown/casts.go b/topdown/casts.go new file mode 100644 index 0000000000..6dd111572a --- /dev/null +++ b/topdown/casts.go @@ -0,0 +1,55 @@ +// 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 ( + "fmt" + "strconv" + + "github.com/open-policy-agent/opa/ast" + "github.com/pkg/errors" +) + +func evalToNumber(ctx *Context, expr *ast.Expr, iter Iterator) error { + ops := expr.Terms.([]*ast.Term) + a, b := ops[1].Value, ops[2].Value + + x, err := ValueToInterface(a, ctx) + if err != nil { + return fmt.Errorf("to_number") + } + + var n ast.Number + + switch x := x.(type) { + case string: + f, err := strconv.ParseFloat(string(x), 64) + if err != nil { + return errors.Wrapf(err, "to_number") + } + n = ast.Number(f) + case float64: + n = ast.Number(x) + case bool: + if x { + n = ast.Number(1) + } else { + n = ast.Number(0) + } + default: + return fmt.Errorf("to_number: source must be a string, boolean, or number: %T", a) + } + + switch b := b.(type) { + case ast.Var: + ctx = ctx.BindVar(b, n) + return iter(ctx) + default: + if n.Equal(b) { + return iter(ctx) + } + return nil + } +} diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index 13fe7cef48..e3dd862d89 100644 --- a/topdown/topdown_test.go +++ b/topdown/topdown_test.go @@ -554,7 +554,7 @@ func TestTopDownNegation(t *testing.T) { } } -func TestTopDownAggregates(t *testing.T) { +func TestTopDownComprehensions(t *testing.T) { tests := []struct { note string @@ -565,10 +565,61 @@ func TestTopDownAggregates(t *testing.T) { {"nested", []string{"p[i] :- ys = [y | y = x[_], x = [z | z = a[_]]], ys[i] > 1"}, "[1,2,3]"}, {"embedded array", []string{"p[i] :- xs = [[x | x = a[_]]], xs[0][i] > 1"}, "[1,2,3]"}, {"embedded object", []string{`p[i] :- xs = {"a": [x | x = a[_]]}, xs["a"][i] > 1`}, "[1,2,3]"}, - {"recursive", []string{"p :- y = 1, x = y, x = [y | y = 1]"}, ""}, - // TODO(tsandall): semantics? - {"recursive", []string{"p :- x = y, x = [y | y = 1]"}, "true"}, - {"recursive", []string{"p :- x = [x | x = 1]"}, "true"}, + {"closure", []string{"p[x] :- y = 1, x = [y | y = 1]"}, "[[1]]"}, + } + + data := loadSmallTestData() + + for i, tc := range tests { + runTopDownTestCase(t, data, i, tc.note, tc.rules, tc.expected) + } +} + +func TestTopDownAggregates(t *testing.T) { + + tests := []struct { + note string + rules []string + expected interface{} + }{ + {"count", []string{"p[x] :- count(a, x)"}, "[4]"}, + {"count virtual", []string{"p[x] :- count([y | q[y]], x)", "q[x] :- x = a[_]"}, "[4]"}, + {"count keys", []string{"p[x] :- count(b, x)"}, "[2]"}, + {"count keys virtual", []string{"p[x] :- count([k | q[k] = _], x)", "q[k] = v :- b[k] = v"}, "[2]"}, + {"sum", []string{"p[x] :- sum([1,2,3,4], x)"}, "[10]"}, + {"sum virtual", []string{"p[x] :- sum([y | q[y]], x)", "q[x] :- a[_] = x"}, "[10]"}, + } + + data := loadSmallTestData() + + for i, tc := range tests { + runTopDownTestCase(t, data, i, tc.note, tc.rules, tc.expected) + } +} + +func TestTopDownArithmetic(t *testing.T) { + tests := []struct { + note string + rules []string + expected interface{} + }{ + {"plus", []string{"p[y] :- a[i] = x, plus(i, x, y)"}, "[1,3,5,7]"}, + } + + data := loadSmallTestData() + + for i, tc := range tests { + runTopDownTestCase(t, data, i, tc.note, tc.rules, tc.expected) + } +} + +func TestTopDownCasts(t *testing.T) { + tests := []struct { + note string + rules []string + expected interface{} + }{ + {"to_number", []string{`p[x] :- to_number("-42.0", y), to_number(false, z), x = [y, z]`}, "[[-42.0, 0]]"}, } data := loadSmallTestData() From fa97fdd10abad40a76c8f36754d1758e10737ddb Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Mon, 13 Jun 2016 12:03:22 -0700 Subject: [PATCH 09/11] Add comprehensions to language reference --- docs/docs/lang.md | 68 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) diff --git a/docs/docs/lang.md b/docs/docs/lang.md index 8d40a05645..50365599cd 100644 --- a/docs/docs/lang.md +++ b/docs/docs/lang.md @@ -470,6 +470,71 @@ The result: +-------+ ``` +## Comprehensions + +Comprehensions provide a concise way of building [Composite Values](#composite-values) from sub-queries. + +Like [Rules](#rules), comprehensions consist of a head and a body. The body of a comprehension can be understood in exactly the same way as the body of a rule, that is, one or more expressions that must all be true in order for the overall body to be true. When the body evaluates to true, the head of the comprehension is evaluated to produce an element in the result. + +The body of a comprehension is able to refer to variables defined in the outer body. For example: + +``` +> region = "west", names = [name | sites[i].region = region, sites[i].name = name] ++-----------------+--------+ +| NAMES | REGION | ++-----------------+--------+ +| ["smoke","dev"] | "west" | ++-----------------+--------+ +``` + +In the above query, the second expression contains an [Array Comprehension](#array-comprehension) that refers to the "region" variable. The region variable will be bound in the outer body. + +> When a comprehension refers to a variable in an outer body, OPA will reorder expressions in the outer body so that variables referred to in the comprehension are bound by the time the comprehension is evaluated. + +Comprehensions are similar to the same constructs found in other languages like Python. For example, we could write the above comprehension in Python as follows: + +```python +# Python equivalent of Rego comprehension shown above. +names = [site.name for site in sites if site.region = "west"] +``` + +Comprehensions are often used to group elements by some key. A common use case for comprehensions is to assist in computing aggregate values (e.g., the number of containers running on a host). + +### Array Comprehensions + +Array Comprehensions build array values out of sub-queries. Array Comprehensions have the form: + +``` +[ | ] +``` + +For example, the following rule defines an object where the keys are application names and the values are hostnames of servers where the application is deployed. The hostnames of servers are represented as an array. + +```rego +app_to_hostnames[app_name] = hostnames :- + apps[_] = app, + app_name = app.name, + hostnames = [hostname | name = app.servers[_], + sites[_].servers[_] = s, + s.name = name, + hostname = s.hostname] +``` + +The result: + +``` +> app_to_hostnames[app] = hostnames ++-----------+-----------------------------------------------------+ +| APP | HOSTNAMES | ++-----------+-----------------------------------------------------+ +| "web" | ["hydrogen","helium","berylium","boron","nitrogen"] | +| "mysql" | ["lithium","carbon"] | +| "mongodb" | ["oxygen"] | ++-----------+-----------------------------------------------------+ +``` + +In the future, Rego will support Set and Object comprehensions. + ## Rules Rules define the content of [Virtual Documents](/docs/arch.html#data-model) in @@ -864,7 +929,8 @@ literal = expr | "not" expr expr = term | expr-builtin | expr-infix expr-builtin = var "(" [ term { , term } ] ")" expr-infix = term bool-operator term -term = ref | var | scalar | array | object +term = ref | var | scalar | array | object | array-compr +array-compr = "[" term "|" rule-body "]" bool-operator = "=" | "!=" | "<" | ">" | ">=" | "<=" ref = var { ref-arg } ref-arg = ref-arg-dot | ref-arg-brack From 1f94f392313608ea175eca0c81852afc413abf56 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Tue, 14 Jun 2016 09:14:04 -0700 Subject: [PATCH 10/11] Add other simple arithmetic built-ins --- ast/builtins.go | 30 ++++++++++++++- topdown/arithmetic.go | 84 +++++++++++++++++++++++++++++++++-------- topdown/builtins.go | 6 ++- topdown/topdown_test.go | 4 ++ 4 files changed, 107 insertions(+), 17 deletions(-) diff --git a/ast/builtins.go b/ast/builtins.go index 8df45e985a..13a1174818 100644 --- a/ast/builtins.go +++ b/ast/builtins.go @@ -21,7 +21,7 @@ func RegisterBuiltin(b *Builtin) { var DefaultBuiltins = [...]*Builtin{ Equality, GreaterThan, GreaterThanEq, LessThan, LessThanEq, NotEqual, - Plus, + Plus, Minus, Multiply, Divide, Round, Count, Sum, ToNumber, } @@ -92,6 +92,34 @@ var Plus = &Builtin{ TargetPos: []int{2}, } +// Minus subtracts the second number from the first number. +var Minus = &Builtin{ + Name: Var("minus"), + NumArgs: 3, + TargetPos: []int{2}, +} + +// Multiply multiplies two numbers together. +var Multiply = &Builtin{ + Name: Var("mul"), + NumArgs: 3, + TargetPos: []int{2}, +} + +// Divide divides the first number by the second number. +var Divide = &Builtin{ + Name: Var("div"), + NumArgs: 3, + TargetPos: []int{2}, +} + +// Round rounds the number up to the nearest integer. +var Round = &Builtin{ + Name: Var("round"), + NumArgs: 2, + TargetPos: []int{1}, +} + /** * Aggregates */ diff --git a/topdown/arithmetic.go b/topdown/arithmetic.go index 53d5b394e9..40baa67be4 100644 --- a/topdown/arithmetic.go +++ b/topdown/arithmetic.go @@ -5,34 +5,88 @@ package topdown import ( + "fmt" + "math" + "github.com/open-policy-agent/opa/ast" "github.com/pkg/errors" ) -func evalPlus(ctx *Context, expr *ast.Expr, iter Iterator) error { - ops := expr.Terms.([]*ast.Term) +type arithmeticFunc func(a, b float64) (ast.Number, error) +func arithPlus(a, b float64) (ast.Number, error) { + return ast.Number(a + b), nil +} + +func arithMinus(a, b float64) (ast.Number, error) { + return ast.Number(a - b), nil +} + +func arithMultiply(a, b float64) (ast.Number, error) { + return ast.Number(a * b), nil +} + +func arithDivide(a, b float64) (ast.Number, error) { + if b == 0 { + return 0, fmt.Errorf("divide: by zero") + } + return ast.Number(a / b), nil +} + +func arithRound(a float64) (ast.Number, error) { + return ast.Number(math.Floor(a + 0.5)), nil +} + +func evalRound(ctx *Context, expr *ast.Expr, iter Iterator) error { + ops := expr.Terms.([]*ast.Term) a, err := ValueToFloat64(ops[1].Value, ctx) if err != nil { - return errors.Wrapf(err, "plus") + return errors.Wrapf(err, "round") } - - b, err := ValueToFloat64(ops[2].Value, ctx) - if err != nil { - return errors.Wrapf(err, "plus") - } - - c := ops[3].Value - r := ast.Number(a + b) - - switch c := c.(type) { + r := ast.Number(math.Floor(a + 0.5)) + b := ops[2].Value + switch b := b.(type) { case ast.Var: - ctx = ctx.BindVar(c, r) + ctx = ctx.BindVar(b, r) return iter(ctx) default: - if r.Equal(c) { + if b.Equal(r) { return iter(ctx) } return nil } } + +func evalArithmetic(f arithmeticFunc) BuiltinFunc { + return func(ctx *Context, expr *ast.Expr, iter Iterator) error { + ops := expr.Terms.([]*ast.Term) + + a, err := ValueToFloat64(ops[1].Value, ctx) + if err != nil { + return errors.Wrapf(err, "arithemtic") + } + + b, err := ValueToFloat64(ops[2].Value, ctx) + if err != nil { + return errors.Wrapf(err, "arithemtic") + } + + c, err := f(a, b) + if err != nil { + return err + } + + cv := ops[3].Value + + switch cv := cv.(type) { + case ast.Var: + ctx = ctx.BindVar(cv, c) + return iter(ctx) + default: + if cv.Equal(c) { + return iter(ctx) + } + return nil + } + } +} diff --git a/topdown/builtins.go b/topdown/builtins.go index 187e96959a..561acd117c 100644 --- a/topdown/builtins.go +++ b/topdown/builtins.go @@ -31,7 +31,11 @@ var defaultBuiltinFuncs = map[ast.Var]BuiltinFunc{ ast.LessThan.Name: evalIneq(compareLessThan), ast.LessThanEq.Name: evalIneq(compareLessThanEq), ast.NotEqual.Name: evalIneq(compareNotEq), - ast.Plus.Name: evalPlus, + ast.Plus.Name: evalArithmetic(arithPlus), + ast.Minus.Name: evalArithmetic(arithMinus), + ast.Multiply.Name: evalArithmetic(arithMultiply), + ast.Divide.Name: evalArithmetic(arithDivide), + ast.Round.Name: evalRound, ast.Count.Name: evalCount, ast.Sum.Name: evalSum, ast.ToNumber.Name: evalToNumber, diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index e3dd862d89..0fec7a6a7e 100644 --- a/topdown/topdown_test.go +++ b/topdown/topdown_test.go @@ -604,6 +604,10 @@ func TestTopDownArithmetic(t *testing.T) { expected interface{} }{ {"plus", []string{"p[y] :- a[i] = x, plus(i, x, y)"}, "[1,3,5,7]"}, + {"minus", []string{"p[y] :- a[i] = x, minus(i, x, y)"}, "[-1,-1,-1,-1]"}, + {"multiply", []string{"p[y] :- a[i] = x, mul(i, x, y)"}, "[0,2,6,12]"}, + {"divide+round", []string{"p[z] :- a[i] = x, div(i, x, y), round(y, z)"}, "[0,1,1,1]"}, + {"divide+error", []string{"p[y] :- a[i] = x, div(x, i, y)"}, fmt.Errorf("divide: by zero")}, } data := loadSmallTestData() From 34b1c8628b79b93450f6af3021c66a82152f7109 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Thu, 16 Jun 2016 11:59:59 -0700 Subject: [PATCH 11/11] Evaluate single, non-boolean term expressions Previously evaluation would stop with an error if the term was not a boolean. Now, we continue if the term is defined and not false. --- topdown/topdown.go | 10 ++++------ topdown/topdown_test.go | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/topdown/topdown.go b/topdown/topdown.go index ea15c99d81..c81d75738a 100644 --- a/topdown/topdown.go +++ b/topdown/topdown.go @@ -439,15 +439,13 @@ func evalExpr(ctx *Context, iter Iterator) error { return iter(ctx) }) case *ast.Term: - switch tv := tt.Value.(type) { - case ast.Boolean: - if tv.Equal(ast.Boolean(true)) { + v := tt.Value + if !v.Equal(ast.Boolean(false)) { + if v.IsGround() { return iter(ctx) } - return nil - default: - return fmt.Errorf("illegal implicit cast: %v", tv) } + return nil default: panic(fmt.Sprintf("illegal argument: %v", tt)) } diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index 0fec7a6a7e..2eff5af2fb 100644 --- a/topdown/topdown_test.go +++ b/topdown/topdown_test.go @@ -285,6 +285,38 @@ func TestTopDownPartialObjectDoc(t *testing.T) { } } +func TestTopDownEvalTermExpr(t *testing.T) { + + tests := []struct { + note string + rule string + expected string + }{ + {"true", "p :- true", "true"}, + {"false", "p :- false", ""}, + {"number non-zero", "p :- -3.14", "true"}, + {"number zero", "p :- null", "true"}, + {"null", "p :- null", "true"}, + {"string non-empty", `p :- "abc"`, "true"}, + {"string empty", `p :- ""`, "true"}, + {"array non-empty", "p :- [1,2,3]", "true"}, + {"array empty", "p :- []", "true"}, + {"object non-empty", `p :- {"a": 1}`, "true"}, + {"object empty", `p :- {}`, "true"}, + {"ref", "p :- a[i]", "true"}, + {"ref undefined", "p :- data.deadbeef[i]", ""}, + {"array comprehension", "p :- [x | x = 1]", "true"}, + {"array comprehension empty", "p :- [x | x = 1, x = 2]", "true"}, + {"arbitrary position", "p :- a[i] = x, x, i", "true"}, + } + + data := loadSmallTestData() + + for i, tc := range tests { + runTopDownTestCase(t, data, i, tc.note, []string{tc.rule}, tc.expected) + } +} + func TestTopDownEqExpr(t *testing.T) { tests := []struct {