diff --git a/lib/collections.go b/lib/collections.go index 8fea491..5780984 100644 --- a/lib/collections.go +++ b/lib/collections.go @@ -19,6 +19,7 @@ package lib import ( "reflect" + "sort" "strings" "github.com/google/cel-go/cel" @@ -236,6 +237,30 @@ import ( // // zip(["a", "b"], [1, 2]) // return {"a":1, "b":2} // ["a", "b"].zip([1, 2]) // return {"a":1, "b":2} +// +// # Keys +// +// Returns a list of keys from a map: +// +// keys(>) -> > +// >.keys() -> > +// +// Examples: +// +// keys({"a":1, "b":2}) // return ["a", "b"] +// {1:"a", 2:"b"}.keys() // return [1, 2] +// +// # Values +// +// Returns a list of values from a map: +// +// values(>) -> > +// >.values() -> > +// +// Examples: +// +// values({"a":1, "b":2}) // return [1, 2] +// {1:"a", 2:"b"}.values() // return ["a", "b"] func Collections() cel.EnvOption { return cel.Lib(collectionsLib{}) } @@ -379,6 +404,34 @@ func (collectionsLib) CompileOptions() []cel.EnvOption { []string{"K", "V"}, ), ), + decls.NewFunction("keys", + decls.NewParameterizedInstanceOverload( + "map_keys", + []*expr.Type{mapKV}, + listK, + []string{"K"}, + ), + decls.NewParameterizedOverload( + "keys_map", + []*expr.Type{mapKV}, + listK, + []string{"K"}, + ), + ), + decls.NewFunction("values", + decls.NewParameterizedInstanceOverload( + "map_values", + []*expr.Type{mapKV}, + listV, + []string{"V"}, + ), + decls.NewParameterizedOverload( + "values_map", + []*expr.Type{mapKV}, + listV, + []string{"V"}, + ), + ), ), } } @@ -481,6 +534,26 @@ func (collectionsLib) ProgramOptions() []cel.ProgramOption { Binary: zipLists, }, ), + cel.Functions( + &functions.Overload{ + Operator: "map_keys", + Unary: mapKeys, + }, + &functions.Overload{ + Operator: "keys_map", + Unary: mapKeys, + }, + ), + cel.Functions( + &functions.Overload{ + Operator: "map_values", + Unary: mapValues, + }, + &functions.Overload{ + Operator: "values_map", + Unary: mapValues, + }, + ), } } @@ -949,6 +1022,79 @@ func zipLists(arg0, arg1 ref.Val) ref.Val { return types.NewRefValMap(types.DefaultTypeAdapter, m) } +func mapKeys(val ref.Val) ref.Val { + mapK, ok := val.(traits.Mapper) + if !ok { + return types.ValOrErr(mapK, "no such overload") + } + n, ok := mapK.Size().(types.Int) + if !ok { + return types.NewErr("unable to get size of map") + } + keys := make([]ref.Val, 0, n) + if mapK.Size() != types.IntZero { + canSort := true + it := mapK.Iterator() + for it.HasNext() == types.True { + k := it.Next() + keys = append(keys, k) + _, ok := k.(traits.Comparer) + if !ok { + canSort = false + } + } + if canSort { + sort.Slice(keys, func(i, j int) bool { + return keys[i].(traits.Comparer).Compare(keys[j]) == types.Int(-1) + }) + } + } + return types.NewRefValList(types.DefaultTypeAdapter, keys) +} + +func mapValues(val ref.Val) ref.Val { + mapK, ok := val.(traits.Mapper) + if !ok { + return types.ValOrErr(mapK, "no such overload") + } + n, ok := mapK.Size().(types.Int) + if !ok { + return types.NewErr("unable to get size of map") + } + values := make([]ref.Val, 0, n) + type valComparer interface { + ref.Val + traits.Comparer + } + type kv struct { + Key valComparer + Value ref.Val + } + if mapK.Size() != types.IntZero { + canSort := true + it := mapK.Iterator() + ss := make([]kv, 0, n) + for it.HasNext() == types.True { + k := it.Next() + v := mapK.Get(k) + ck, ok := k.(valComparer) + if !ok { + canSort = false + } + ss = append(ss, kv{ck, v}) + } + if canSort { + sort.Slice(ss, func(i, j int) bool { + return ss[i].Key.Compare(ss[j].Key) == types.Int(-1) + }) + } + for _, kv := range ss { + values = append(values, kv.Value) + } + } + return types.NewRefValList(types.DefaultTypeAdapter, values) +} + func makeAs(eh parser.ExprHelper, target ast.Expr, args []ast.Expr) (ast.Expr, *common.Error) { ident := args[0] if ident.Kind() != ast.IdentKind { diff --git a/testdata/keys.txt b/testdata/keys.txt new file mode 100644 index 0000000..dcd40df --- /dev/null +++ b/testdata/keys.txt @@ -0,0 +1,20 @@ +mito -use collections,try src.cel +! stderr . +cmp stdout want.txt + +-- src.cel -- +{ + "good_instance": {"b": [2, 3], "a": {1: "aa"}}.keys(), + "good_function": keys({1: "a", 2: "b"}), +} +-- want.txt -- +{ + "good_function": [ + 1, + 2 + ], + "good_instance": [ + "a", + "b" + ] +} diff --git a/testdata/values.txt b/testdata/values.txt new file mode 100644 index 0000000..50b9fc5 --- /dev/null +++ b/testdata/values.txt @@ -0,0 +1,30 @@ +mito -use collections,try src.cel +! stderr . +cmp stdout want.txt + +-- src.cel -- +{ + "good_instance": {"c": {"cc": 4, "cd": 5}, "a": ["aa", "ab"], "b": [2, 3]}.values(), + "good_function": values({2: "b", 1: "a"}), +} +-- want.txt -- +{ + "good_function": [ + "a", + "b" + ], + "good_instance": [ + [ + "aa", + "ab" + ], + [ + 2, + 3 + ], + { + "cc": 4, + "cd": 5 + } + ] +} diff --git a/testdata/zip.txt b/testdata/zip.txt index 20b14e0..be7de9a 100644 --- a/testdata/zip.txt +++ b/testdata/zip.txt @@ -6,6 +6,7 @@ cmp stdout want.txt { "good_instance": ["a", "b"].zip([1, 2]), "good_function": zip(["a", "b"], [1, 2]), + "good_function_2": zip(keys({"a": 1, "b": 2}), values({"a": 1, "b": 2})), "bad_instance": try(["a", "b"].zip([1, 2, 3])), "bad_function": try(zip(["a", "b"], [1, 2, 3])), } @@ -17,6 +18,10 @@ cmp stdout want.txt "a": 1, "b": 2 }, + "good_function_2": { + "a": 1, + "b": 2 + }, "good_instance": { "a": 1, "b": 2