Implement kro CEL library for maps

pkg/cel/environment.go

@@ -57,6 +57,7 @@ func DefaultEnvironment(options ...EnvOption) (*cel.Env, error) {
 		cel.OptionalTypes(),
 		ext.Encoders(),
 		library.Random(),
+		library.Maps(),
 	}
 
 	opts := &envOptions{}

pkg/cel/library/maps.go

@@ -0,0 +1,106 @@
+package library
+
+import (
+	"math"
+
+	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/common/types"
+	"github.com/google/cel-go/common/types/ref"
+	"github.com/google/cel-go/common/types/traits"
+)
+
+// Maps returns a cel.EnvOption to configure extended functions for map manipulation.
+//
+// # Merge
+//
+// Merges two maps. Keys from the first map take precedence over keys in the second map.
+//
+//	map(string, int).merge(map(string, int)) -> map(string, int)
+//
+// Examples:
+//
+//	{}.merge({}) == {}
+//	{}.merge({'a': 1}) == {'a': 1}
+//	{'a': 1}.merge({}) == {'a': 1}
+//	{'a': 1}.merge({'b': 2}) == {'a': 1, 'b': 2}
+//	{'a': 1}.merge({'a': 2, 'b': 2}) == {'a': 1, 'b': 2}
+func Maps(options ...MapsOption) cel.EnvOption {
+	l := &mapsLib{version: math.MaxUint32}
+	for _, opt := range options {
+		opt(l)
+	}
+	return cel.Lib(l)
+}
+
+type mapsLib struct {
+	version uint32
+}
+
+type MapsOption func(*mapsLib) *mapsLib
+
+// LibraryName implements the cel.SingletonLibrary interface method.
+func (l *mapsLib) LibraryName() string {
+	return "cel.lib.ext.kro.maps"
+}
+
+// CompileOptions implements the cel.Library interface method.
+func (l *mapsLib) CompileOptions() []cel.EnvOption {
+	mapType := cel.MapType(cel.TypeParamType("K"), cel.DynType)
+	// mapDynType := cel.MapType(cel.DynType, cel.DynType)
+	opts := []cel.EnvOption{
+		cel.Function("merge",
+			cel.MemberOverload("map_merge",
+				[]*cel.Type{mapType, mapType},
+				mapType,
+				cel.BinaryBinding(func(arg1, arg2 ref.Val) ref.Val {
+					self, ok := arg1.(traits.Mapper)
+					if !ok {
+						return types.ValOrErr(arg1, "no such overload: %v.merge(%v)", arg1.Type(), arg2.Type())
+					}
+					other, ok := arg2.(traits.Mapper)
+					if !ok {
+						return types.ValOrErr(arg1, "no such overload: %v.merge(%v)", arg1.Type(), arg2.Type())
+					}
+					result, err := merge(self, other)
+					if err != nil {
+						return types.WrapErr(err)
+					}
+					return result
+				}),
+			),
+		),
+	}
+	return opts
+}
+
+// ProgramOptions implements the cel.Library interface method.
+func (l *mapsLib) ProgramOptions() []cel.ProgramOption {
+	return []cel.ProgramOption{}
+}
+
+// merge merges two maps. Keys from the first map take precedence over keys in
+// the second map.
+func merge(self traits.Mapper, other traits.Mapper) (traits.Mapper, error) {
+	var result traits.MutableMapper
+
+	if mapVal, ok := self.Value().(map[ref.Val]ref.Val); ok {
+		result = types.NewMutableMap(types.DefaultTypeAdapter, mapVal)
+	} else {
+		result = types.NewMutableMap(types.DefaultTypeAdapter, nil)
+		for i := self.Iterator(); i.HasNext().(types.Bool); {
+			k := i.Next()
+			v := self.Get(k)
+			result.Insert(k, v)
+		}
+	}
+
+	for i := other.Iterator(); i.HasNext().(types.Bool); {
+		k := i.Next()
+		if result.Contains(k).(types.Bool) {
+			continue
+		}
+		v := other.Get(k)
+		result.Insert(k, v)
+	}
+	return result.ToImmutableMap(), nil
+}

pkg/cel/library/maps_test.go

@@ -0,0 +1,77 @@
+package library
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/google/cel-go/cel"
+)
+
+func TestMaps(t *testing.T) {
+	mapsTests := []struct {
+		expr string
+		err  string
+	}{
+		{expr: `{}.merge({}) == {}`},
+		{expr: `{}.merge({'a': 1}) == {'a': 1}`},
+		{expr: `{'a': 1}.merge({}) == {'a': 1}`},
+		{expr: `{'a': 1}.merge({'b': 2}) == {'a': 1, 'b': 2}`},
+		{expr: `{'a': 1}.merge({'a': 2, 'b': 2}) == {'a': 1, 'b': 2}`},
+		{expr: `{'a': 1}.merge({'b': 2.1}) == {'a': 1, 'b': 2.1}`},
+		{expr: `{'a': 1}.merge({'b': 'foo'}) == {'a': 1, 'b': 'foo'}`},
+		{expr: `{'a': 1}.merge({'b': {'c': 2}}) == {'a': 1, 'b': {'c': 2}}`},
+
+		// {expr: `{}.merge([])`, err: "ERROR: <input>:1:9: found no matching overload for 'merge' applied to 'map(dyn, dyn).(list(dyn))'"},
+	}
+
+	env := testMapsEnv(t)
+	for i, tc := range mapsTests {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			var asts []*cel.Ast
+			pAst, iss := env.Parse(tc.expr)
+			if iss.Err() != nil {
+				t.Fatalf("env.Parse(%v) failed: %v", tc.expr, iss.Err())
+			}
+			asts = append(asts, pAst)
+			cAst, iss := env.Check(pAst)
+			if iss.Err() != nil {
+				t.Fatalf("env.Check(%v) failed: %v", tc.expr, iss.Err())
+			}
+			asts = append(asts, cAst)
+
+			for _, ast := range asts {
+				prg, err := env.Program(ast)
+				if err != nil {
+					t.Fatalf("env.Program() failed: %v", err)
+				}
+				out, _, err := prg.Eval(cel.NoVars())
+				if tc.err != "" {
+					if err == nil {
+						t.Fatalf("got value %v, wanted error %s for expr: %s",
+							out.Value(), tc.err, tc.expr)
+					}
+					if !strings.Contains(err.Error(), tc.err) {
+						t.Errorf("got error %v, wanted error %s for expr: %s", err, tc.err, tc.expr)
+					}
+				} else if err != nil {
+					t.Fatal(err)
+				} else if out.Value() != true {
+					t.Errorf("got %v, wanted true for expr: %s", out.Value(), tc.expr)
+				}
+			}
+		})
+	}
+}
+
+func testMapsEnv(t *testing.T, opts ...cel.EnvOption) *cel.Env {
+	t.Helper()
+	baseOpts := []cel.EnvOption{
+		Maps(),
+	}
+	env, err := cel.NewEnv(append(baseOpts, opts...)...)
+	if err != nil {
+		t.Fatalf("cel.NewEnv(Maps()) failed: %v", err)
+	}
+	return env
+}