lib: Extract keys and values from map

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(<map<dyn,dyn>>) -> <list<dyn>>
+//	<map<dyn,dyn>>.keys() -> <list<dyn>>
+//
+// 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(<map<dyn,dyn>>) -> <list<dyn>>
+//	<map<dyn,dyn>>.values() -> <list<dyn>>
+//
+// 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 {

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"
+	]
+}

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
+		}
+	]
+}

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