From f2fd8f5e91824bca24a813be78f74b41a40fd1c2 Mon Sep 17 00:00:00 2001 From: Charles Daniels Date: Tue, 7 Jun 2022 13:59:02 -0400 Subject: [PATCH] built-ins: add object.subset() builtin This implements the new object.subset() builtin. Based on my benchmarking, this offers a 2.77x speedup compared to implementing the same thing in pure Rego, and is also easier to read. Fixes #4358 Signed-off-by: Charles Daniels --- ast/builtins.go | 52 ++- builtin_metadata.json | 26 ++ capabilities.json | 67 ++++ test/cases/testdata/subset/test-subset.yaml | 347 ++++++++++++++++++++ topdown/subset.go | 221 +++++++++++++ topdown/topdown_bench_test.go | 121 +++++++ 6 files changed, 825 insertions(+), 9 deletions(-) create mode 100644 test/cases/testdata/subset/test-subset.yaml create mode 100644 topdown/subset.go diff --git a/ast/builtins.go b/ast/builtins.go index 3bdf93e5fb..f1a8b3d34d 100644 --- a/ast/builtins.go +++ b/ast/builtins.go @@ -161,6 +161,7 @@ var DefaultBuiltins = [...]*Builtin{ ObjectRemove, ObjectFilter, ObjectGet, + ObjectSubset, // JSON Object Manipulation JSONFilter, @@ -1425,18 +1426,36 @@ var JSONPatch = &Builtin{ Categories: objectCat, } -var ObjectGet = &Builtin{ - Name: "object.get", - Description: "Returns value of an object's key if present, otherwise a default. " + - "If the supplied `key` is an `array`, then `object.get` will search through a nested object or array using each key in turn. " + - "For example: `object.get({\"a\": [{ \"b\": true }]}, [\"a\", 0, \"b\"], false)` results in `true`.", +var ObjectSubset = &Builtin{ + Name: "object.subset", + Description: "Determines if an object `sub` is a subset of another object `super`." + + "Object `sub` is a subset of object `super` if and only if every key in `sub` is also in `super`, " + + "**and** for all keys which `sub` and `super` share, they have the same value. " + + "This function works with objects, sets, and arrays. " + + "If both arguments are objects, then the operation is recursive, e.g. " + + "`{\"c\": {\"x\": {10, 15, 20}}` is a subset of `{\"a\": \"b\", \"c\": {\"x\": {10, 15, 20, 25}, \"y\": \"z\"}`. " + + "If both arguments are sets, then this function checks if every element of `sub` is a member of `super`, " + + "but does not attempt to recurse. If both arguments are arrays, " + + "then this function checks if `sub` appears contiguously in order within `super`, " + + "and also does not attempt to recurse.", Decl: types.NewFunction( types.Args( - types.Named("object", types.NewObject(nil, types.NewDynamicProperty(types.A, types.A))).Description("object to get `key` from"), - types.Named("key", types.A).Description("key to lookup in `object`"), - types.Named("default", types.A).Description("default to use when lookup fails"), + types.Named("super", types.NewAny(types.NewObject( + nil, + types.NewDynamicProperty(types.A, types.A), + ), + types.NewSet(types.A), + types.NewArray(nil, types.A), + )).Description("object to test if sub is a subset of"), + types.Named("sub", types.NewAny(types.NewObject( + nil, + types.NewDynamicProperty(types.A, types.A), + ), + types.NewSet(types.A), + types.NewArray(nil, types.A), + )).Description("object to test if super is a superset of"), ), - types.Named("value", types.A).Description("`object[key]` if present, otherwise `default`"), + types.Named("result", types.A).Description("`true` if `sub` is a subset of `super`"), ), } @@ -1513,6 +1532,21 @@ var ObjectFilter = &Builtin{ ), } +var ObjectGet = &Builtin{ + Name: "object.get", + Description: "Returns value of an object's key if present, otherwise a default. " + + "If the supplied `key` is an `array`, then `object.get` will search through a nested object or array using each key in turn. " + + "For example: `object.get({\"a\": [{ \"b\": true }]}, [\"a\", 0, \"b\"], false)` results in `true`.", + Decl: types.NewFunction( + types.Args( + types.Named("object", types.NewObject(nil, types.NewDynamicProperty(types.A, types.A))).Description("object to get `key` from"), + types.Named("key", types.A).Description("key to lookup in `object`"), + types.Named("default", types.A).Description("default to use when lookup fails"), + ), + types.Named("value", types.A).Description("`object[key]` if present, otherwise `default`"), + ), +} + /* * Encoding */ diff --git a/builtin_metadata.json b/builtin_metadata.json index 7e4035f9b7..9ba090d70e 100644 --- a/builtin_metadata.json +++ b/builtin_metadata.json @@ -118,6 +118,7 @@ "object.filter", "object.get", "object.remove", + "object.subset", "object.union", "object.union_n" ], @@ -8394,6 +8395,31 @@ }, "wasm": true }, + "object.subset": { + "args": [ + { + "description": "object to test if sub is a subset of", + "name": "super", + "type": "any\u003carray[any], object[any: any], set[any]\u003e" + }, + { + "description": "object to test if super is a superset of", + "name": "sub", + "type": "any\u003carray[any], object[any: any], set[any]\u003e" + } + ], + "available": [ + "edge" + ], + "description": "Determines if an object `sub` is a subset of another object `super`.Object `sub` is a subset of object `super` if and only if every key in `sub` is also in `super`, **and** for all keys which `sub` and `super` share, they have the same value. This function works with objects, sets, and arrays. If both arguments are objects, then the operation is recursive, e.g. `{\"c\": {\"x\": {10, 15, 20}}` is a subset of `{\"a\": \"b\", \"c\": {\"x\": {10, 15, 20, 25}, \"y\": \"z\"}`. If both arguments are sets, then this function checks if every element of `sub` is a member of `super`, but does not attempt to recurse. If both arguments are arrays, then this function checks if `sub` appears contiguously in order within `super`, and also does not attempt to recurse.", + "introduced": "edge", + "result": { + "description": "`true` if `sub` is a subset of `super`", + "name": "result", + "type": "any" + }, + "wasm": false + }, "object.union": { "args": [ { diff --git a/capabilities.json b/capabilities.json index cf56e6cacd..03c99cffaf 100644 --- a/capabilities.json +++ b/capabilities.json @@ -2671,6 +2671,73 @@ "type": "function" } }, + { + "name": "object.subset", + "decl": { + "args": [ + { + "of": [ + { + "dynamic": { + "type": "any" + }, + "type": "array" + }, + { + "dynamic": { + "key": { + "type": "any" + }, + "value": { + "type": "any" + } + }, + "type": "object" + }, + { + "of": { + "type": "any" + }, + "type": "set" + } + ], + "type": "any" + }, + { + "of": [ + { + "dynamic": { + "type": "any" + }, + "type": "array" + }, + { + "dynamic": { + "key": { + "type": "any" + }, + "value": { + "type": "any" + } + }, + "type": "object" + }, + { + "of": { + "type": "any" + }, + "type": "set" + } + ], + "type": "any" + } + ], + "result": { + "type": "any" + }, + "type": "function" + } + }, { "name": "object.union", "decl": { diff --git a/test/cases/testdata/subset/test-subset.yaml b/test/cases/testdata/subset/test-subset.yaml new file mode 100644 index 0000000000..34e9464ce9 --- /dev/null +++ b/test/cases/testdata/subset/test-subset.yaml @@ -0,0 +1,347 @@ +cases: + - data: + modules: + - | + package test + + A := { + "a": 5, + "b": 7, + "c": 15 + } + + B := { + "a": 5, + } + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/simple object subset 1 + query: data.test.test_result = x + want_result: + - x: true + + + - data: + modules: + - | + package test + + A := { + "a": 5, + "b": 7, + "c": 15 + } + + B := { + "a": 5, + "b": 10, + } + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + not AsubB # B is not a subset of A + not BsubA # A is not a subset of B + } + + note: subset/simple object subset 2 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := { + "a": 5, + "b": 7, + "c": 15, + "nested": { + "x": 10, + "y": 15, + "z": 20 + } + } + + B := { + "a": 5, + "nested": { + "x": 10, + "y": 15, + "z": 20 + } + } + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/nested object subset 1 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := { + "a": 5, + "b": 7, + "c": 15, + "nested": { + "x": 10, + "y": 15, + "z": 20 + } + } + + # The subset operation applies recursively to nested objects + + B := { + "a": 5, + "nested": { + "x": 10, + "y": 15, + } + } + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/nested object subset 2 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := { + "a": 5, + "b": 7, + "c": 15, + "nested": { + "x": 10, + "y": 15, + "z": 20, + "set1": {1,2,3,4}, + "arr1": [6,7,8,9] + } + } + + # The subset operation applies recursively to nested objects + + B := { + "a": 5, + "nested": { + "x": 10, + "y": 15, + "set1": {4, 1}, + "arr1": [6,7] + } + } + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/nested object subset 3 + query: data.test.test_result = x + want_result: + - x: true + + + - data: + modules: + - | + package test + + A := {1, 2, 3} + B := {3, 2} + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/sets 1 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := {1, 2, 3, {"a", "b", "c"}} + B := {3, 2, {"a", "b", "c"}} + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/nested sets 1 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := {1, 2, 3, {"a", "b", "c"}} + B := {3, 2, {"a", "b"}} + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + # Notice that there isn't really a well-defined way to "match" the two + # nested sets to one another, so we B is not a subset of A in this + # case. + + test_result { + not AsubB # B is not a subset of A + not BsubA # A is not a subset of B + } + + note: subset/nested sets 2 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := [1,2,3,4,5,6] + B := [3,4,5] + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + # Notice that there isn't really a well-defined way to "match" the two + # nested sets to one another, so we B is not a subset of A in this + # case. + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/arrays 1 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := [1,2,3,4,5,6] + B := [1,2,3] + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + # Notice that there isn't really a well-defined way to "match" the two + # nested sets to one another, so we B is not a subset of A in this + # case. + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/arrays 2 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := [1,2,3,4,5,6] + B := [4,5,6] + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + # Notice that there isn't really a well-defined way to "match" the two + # nested sets to one another, so we B is not a subset of A in this + # case. + + test_result { + AsubB # B is a subset of A + not BsubA # A is not a subset of B + } + + note: subset/arrays 3 + query: data.test.test_result = x + want_result: + - x: true + + - data: + modules: + - | + package test + + A := [1,2,3,4,5,6] + B := [1,2,3,4,5,6] + + AsubB := object.subset(A, B) + BsubA := object.subset(B, A) + + # Notice that there isn't really a well-defined way to "match" the two + # nested sets to one another, so we B is not a subset of A in this + # case. + + test_result { + AsubB # B is a subset of A + BsubA # A is a subset of B + } + + note: subset/arrays 4 + query: data.test.test_result = x + want_result: + - x: true + diff --git a/topdown/subset.go b/topdown/subset.go new file mode 100644 index 0000000000..9273ac3998 --- /dev/null +++ b/topdown/subset.go @@ -0,0 +1,221 @@ +// Copyright 2022 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/open-policy-agent/opa/topdown/builtins" +) + +func bothObjects(t1, t2 *ast.Term) (bool, ast.Object, ast.Object) { + if (t1 == nil) || (t2 == nil) { + return false, nil, nil + } + + obj1, ok := t1.Value.(ast.Object) + if !ok { + return false, nil, nil + } + + obj2, ok := t2.Value.(ast.Object) + if !ok { + return false, nil, nil + } + + return true, obj1, obj2 +} + +func bothSets(t1, t2 *ast.Term) (bool, ast.Set, ast.Set) { + if (t1 == nil) || (t2 == nil) { + return false, nil, nil + } + + set1, ok := t1.Value.(ast.Set) + if !ok { + return false, nil, nil + } + + set2, ok := t2.Value.(ast.Set) + if !ok { + return false, nil, nil + } + + return true, set1, set2 +} + +func bothArrays(t1, t2 *ast.Term) (bool, *ast.Array, *ast.Array) { + if (t1 == nil) || (t2 == nil) { + return false, nil, nil + } + + array1, ok := t1.Value.(*ast.Array) + if !ok { + return false, nil, nil + } + + array2, ok := t2.Value.(*ast.Array) + if !ok { + return false, nil, nil + } + + return true, array1, array2 +} + +// objectSubset implements the subset operation on a pair of objects. +// +// This function will try to recursively apply the subset operation where it +// can, such as if both super and sub have an object or set as the value +// associated with a key. +func objectSubset(super ast.Object, sub ast.Object) bool { + var superTerm *ast.Term + isSubset := true + + sub.Until(func(key, subTerm *ast.Term) bool { + // This really wants to be a for loop, hence the somewhat + // weird internal structure. However, using Until() in this + // was is a performance optimization, as it avoids performing + // any key hashing on the sub-object. + + superTerm = super.Get(key) + + // subTerm is can't be nil because we got it from Until(), so + // we only need to verify that super is non-nil. + if superTerm == nil { + isSubset = false + return true // break, not a subset + } + + if subTerm.Equal(superTerm) { + return false // continue + } + + // If both of the terms are objects then we want to apply + // the subset operation recursively, otherwise we just compare + // them normally. If only one term is an object, then we + // do a normal comparison which will come up false. + if ok, superObj, subObj := bothObjects(superTerm, subTerm); ok { + if !objectSubset(superObj, subObj) { + isSubset = false + return true // break, not a subset + } + + return false // continue + } + + if ok, superSet, subSet := bothSets(superTerm, subTerm); ok { + if !setSubset(superSet, subSet) { + isSubset = false + return true // break, not a subset + } + + return false // continue + } + + if ok, superArray, subArray := bothArrays(superTerm, subTerm); ok { + if !arraySubset(superArray, subArray) { + isSubset = false + return true // break, not a subset + } + + return false // continue + } + + // We have already checked for exact equality, as well as for + // all of the types of nested subsets we care about, so if we + // get here it means this isn't a subset. + isSubset = false + return true // break, not a subset + }) + + return isSubset +} + +// setSubset implements the subset operation on sets. +// +// Unlike in the object case, this is not recursive, we just compare values +// using ast.Set.Contains() because we have no well defined way to "match up" +// objects that are in different sets. +func setSubset(super ast.Set, sub ast.Set) bool { + isSubset := true + sub.Until(func(t *ast.Term) bool { + if !super.Contains(t) { + isSubset = false + return true + } + return false + }) + + return isSubset +} + +// arraySubset implements the subset operation on arrays. +// +// This is defined to mean that the entire "sub" array must appear in +// the "super" array. For the same rationale as setSubset(), we do not attempt +// to recurse into values. +func arraySubset(super, sub *ast.Array) bool { + // Notice that this is essentially string search. The naive approach + // used here is O(n^2). This should probably be rewritten later to use + // Boyer-Moore or something. + + if sub.Len() > super.Len() { + return false + } + + if sub.Equal(super) { + return true + } + + superCursor := 0 + subCursor := 0 + for { + if subCursor == sub.Len() { + return true + } + + if superCursor == super.Len() { + return false + } + + subElem := sub.Elem(subCursor) + superElem := sub.Elem(superCursor + subCursor) + if superElem == nil { + return false + } + + if superElem.Value.Compare(subElem.Value) == 0 { + subCursor++ + } else { + superCursor++ + subCursor = 0 + } + } +} + +func builtinObjectSubset(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { + superTerm := operands[0] + subTerm := operands[1] + + if ok, superObj, subObj := bothObjects(superTerm, subTerm); ok { + // Both operands are objects. + return iter(ast.BooleanTerm(objectSubset(superObj, subObj))) + } + + if ok, superSet, subSet := bothSets(superTerm, subTerm); ok { + // Both operands are sets. + return iter(ast.BooleanTerm(setSubset(superSet, subSet))) + } + + if ok, superArray, subArray := bothArrays(superTerm, subTerm); ok { + // Both operands are sets. + return iter(ast.BooleanTerm(arraySubset(superArray, subArray))) + } + + return builtins.ErrOperand("both arguments object.subset must be of the same type") +} + +func init() { + RegisterBuiltinFunc(ast.ObjectSubset.Name, builtinObjectSubset) +} diff --git a/topdown/topdown_bench_test.go b/topdown/topdown_bench_test.go index 1d128ac386..4e6272944f 100644 --- a/topdown/topdown_bench_test.go +++ b/topdown/topdown_bench_test.go @@ -729,3 +729,124 @@ func genComprehensionIndexingData(n int) map[string]interface{} { } return map[string]interface{}{"items": items} } + +func BenchmarkObjectSubset(b *testing.B) { + ctx := context.Background() + + sizes := []int{10, 100, 1000, 10000} + + for _, n := range sizes { + b.Run(fmt.Sprint(n), func(b *testing.B) { + all := make(map[string]string) + evens := make(map[string]string) + + for i := 0; i < n; i++ { + all[fmt.Sprint(i)] = fmt.Sprint(i * 2) + if i%2 == 0 { + evens[fmt.Sprint(i)] = fmt.Sprint(i * 2) + } + } + + store := inmem.NewFromObject(map[string]interface{}{"all": all, "evens": evens}) + + module := `package test + main {object.subset(data.all, data.evens)}` + + query := ast.MustParseBody("data.test.main") + compiler := ast.MustCompileModules(map[string]string{ + "test.rego": module, + }) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + + err := storage.Txn(ctx, store, storage.TransactionParams{}, func(txn storage.Transaction) error { + + q := NewQuery(query). + WithCompiler(compiler). + WithStore(store). + WithTransaction(txn) + + _, err := q.Run(ctx) + if err != nil { + return err + } + + return nil + }) + + if err != nil { + b.Fatal(err) + } + } + }) + } +} + +func BenchmarkObjectSubsetSlow(b *testing.B) { + // This benchmarks the suggested means to implement object.subset + // without using the builtin, to give us an idea of whether or not + // the builtin is actually making things any faster. + ctx := context.Background() + + sizes := []int{10, 100, 1000, 10000} + + for _, n := range sizes { + b.Run(fmt.Sprint(n), func(b *testing.B) { + all := make(map[string]string) + evens := make(map[string]string) + + for i := 0; i < n; i++ { + all[fmt.Sprint(i)] = fmt.Sprint(i * 2) + if i%2 == 0 { + evens[fmt.Sprint(i)] = fmt.Sprint(i * 2) + } + } + + store := inmem.NewFromObject(map[string]interface{}{"all": all, "evens": evens}) + + // Code is lifted from here: + // https://github.com/open-policy-agent/opa/issues/4358#issue-1141145857 + + module := `package test + path_matches[match] { + [path, value] := walk(data.evens) + not is_object(value) + + match := object.get(data.all, path, null) == value + } + + main { path_matches == {true} }` + + query := ast.MustParseBody("data.test.main") + compiler := ast.MustCompileModules(map[string]string{ + "test.rego": module, + }) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + + err := storage.Txn(ctx, store, storage.TransactionParams{}, func(txn storage.Transaction) error { + + q := NewQuery(query). + WithCompiler(compiler). + WithStore(store). + WithTransaction(txn) + + _, err := q.Run(ctx) + if err != nil { + return err + } + + return nil + }) + + if err != nil { + b.Fatal(err) + } + } + }) + } +}