From edbceda55a9846f95cb3d58b36ddf5e303ab762c Mon Sep 17 00:00:00 2001 From: Roshan Date: Fri, 7 Aug 2026 19:58:46 +0000 Subject: [PATCH] Add a maps extension with maps.merge CEL has no operator for combining two maps: + concatenates strings, bytes, and lists but is not defined for maps, and there is no maps extension. Add maps.merge(a, b), returning a new map with the entries of both and the second argument's values winning on conflicting keys. The merge is shallow, so a value that is itself a map is replaced rather than merged. This is the replace step of the merge semantics discussed in #1240; set-if-absent and recursive merge are left for follow-ups. Includes cost estimation and tracking that scale with the combined size of both inputs, registration in the extension option factory and the repl, and documentation in ext/README.md. --- ext/BUILD.bazel | 2 + ext/README.md | 25 ++++ ext/costs.go | 1 + ext/extension_option_factory.go | 4 + ext/maps.go | 161 +++++++++++++++++++++ ext/maps_test.go | 245 ++++++++++++++++++++++++++++++++ repl/evaluator.go | 1 + 7 files changed, 439 insertions(+) create mode 100644 ext/maps.go create mode 100644 ext/maps_test.go diff --git a/ext/BUILD.bazel b/ext/BUILD.bazel index f362fd97b..9a99a8be8 100644 --- a/ext/BUILD.bazel +++ b/ext/BUILD.bazel @@ -16,6 +16,7 @@ go_library( "formatting_v2.go", "guards.go", "lists.go", + "maps.go", "math.go", "native.go", "network.go", @@ -61,6 +62,7 @@ go_test( "formatting_test.go", "formatting_v2_test.go", "lists_test.go", + "maps_test.go", "math_test.go", "native_test.go", "network_test.go", diff --git a/ext/README.md b/ext/README.md index 6133b5cbf..0b96cb8e5 100644 --- a/ext/README.md +++ b/ext/README.md @@ -427,6 +427,31 @@ Example: proto.hasExt(msg, google.expr.proto2.test.int32_ext) // returns true || false +## Maps + +Extended functions for map manipulation. + +CEL has no operator for combining two maps: the `+` operator concatenates +strings, bytes, and lists, but is not defined for maps. + +### Maps.Merge + +Returns a new map containing the entries of both arguments. When a key is +present in both, the value from the second argument wins. Neither input is +modified. + +The merge is shallow: a value that is itself a map is replaced rather than +merged recursively. + + maps.merge(map(K, V), map(K, V)) -> map(K, V) + +Examples: + + maps.merge({}, {}) // {} + maps.merge({'a': 1}, {'b': 2}) // {'a': 1, 'b': 2} + maps.merge({'a': 1}, {'a': 2}) // {'a': 2} + maps.merge({'a': {'x': 1}}, {'a': {'y': 2}}) // {'a': {'y': 2}} + ## Lists Extended functions for list manipulation. As a general note, all indices are diff --git a/ext/costs.go b/ext/costs.go index d2cf7c757..8b2ef72bc 100644 --- a/ext/costs.go +++ b/ext/costs.go @@ -29,6 +29,7 @@ var ( callCostEstimate = checker.FixedCostEstimate(1) callCost = uint64(1) listAllocCost = checker.FixedCostEstimate(common.ListCreateBaseCost) + mapAllocCost = checker.FixedCostEstimate(common.MapCreateBaseCost) stringCostFactor = common.StringTraversalCostFactor ) diff --git a/ext/extension_option_factory.go b/ext/extension_option_factory.go index e68cf5bc7..512b333e3 100644 --- a/ext/extension_option_factory.go +++ b/ext/extension_option_factory.go @@ -59,6 +59,9 @@ var extFactories = map[string]extensionFactory{ "cel.lib.ext.lists": func(version uint32) cel.EnvOption { return Lists(ListsVersion(version)) }, + "cel.lib.ext.maps": func(version uint32) cel.EnvOption { + return Maps(MapsVersion(version)) + }, "cel.lib.ext.math": func(version uint32) cel.EnvOption { return Math(MathVersion(version)) }, @@ -83,6 +86,7 @@ var extAliases = map[string]string{ "bindings": "cel.lib.ext.cel.bindings", "encoders": "cel.lib.ext.encoders", "lists": "cel.lib.ext.lists", + "maps": "cel.lib.ext.maps", "math": "cel.lib.ext.math", "protos": "cel.lib.ext.protos", "sets": "cel.lib.ext.sets", diff --git a/ext/maps.go b/ext/maps.go new file mode 100644 index 000000000..78d9dd3da --- /dev/null +++ b/ext/maps.go @@ -0,0 +1,161 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ext + +import ( + "github.com/google/cel-go/cel" + "github.com/google/cel-go/checker" + "github.com/google/cel-go/common" + "github.com/google/cel-go/common/types" + "github.com/google/cel-go/common/types/ref" + "github.com/google/cel-go/common/types/traits" + "github.com/google/cel-go/interpreter" +) + +// Maps returns a cel.EnvOption to configure namespaced map functions. +// +// CEL has no operator for combining two maps: the `+` operator concatenates +// strings, bytes, and lists, but is not defined for maps. This library provides +// map combination as a named function. +// +// # Maps.Merge +// +// Returns a new map containing the entries of both arguments. When a key is +// present in both, the value from the second argument wins. Neither input is +// modified. +// +// The merge is shallow: a value that is itself a map is replaced rather than +// merged recursively. +// +// maps.merge(map(K, V), map(K, V)) -> map(K, V) +// +// Examples: +// +// maps.merge({}, {}) // {} +// maps.merge({'a': 1}, {'b': 2}) // {'a': 1, 'b': 2} +// maps.merge({'a': 1}, {'a': 2}) // {'a': 2} +// maps.merge({'a': {'x': 1}}, {'a': {'y': 2}}) // {'a': {'y': 2}}, values are replaced, not merged +func Maps(options ...MapsOption) cel.EnvOption { + l := &mapsLib{} + for _, o := range options { + l = o(l) + } + return cel.Lib(l) +} + +// MapsOption declares a functional operator for configuring map extensions. +type MapsOption func(*mapsLib) *mapsLib + +// MapsVersion sets the library version for map extensions. +func MapsVersion(version uint32) MapsOption { + return func(lib *mapsLib) *mapsLib { + lib.version = version + return lib + } +} + +type mapsLib struct { + version uint32 +} + +// LibraryName implements the SingletonLibrary interface method. +func (mapsLib) LibraryName() string { + return "cel.lib.ext.maps" +} + +// CompileOptions implements the Library interface method. +func (mapsLib) CompileOptions() []cel.EnvOption { + mapType := cel.MapType(cel.TypeParamType("K"), cel.TypeParamType("V")) + return []cel.EnvOption{ + cel.Function("maps.merge", + cel.Overload("map_maps_merge_map", []*cel.Type{mapType, mapType}, mapType, + cel.BinaryBinding(mapsMerge))), + cel.CostEstimatorOptions( + checker.OverloadCostEstimate("map_maps_merge_map", estimateMapsMergeCost), + ), + } +} + +// ProgramOptions implements the Library interface method. +func (mapsLib) ProgramOptions() []cel.ProgramOption { + return []cel.ProgramOption{ + cel.CostTrackerOptions( + interpreter.OverloadCostTracker("map_maps_merge_map", trackMapsMergeCost), + ), + } +} + +// estimateMapsMergeCost charges for visiting every entry of both inputs and for +// allocating the result map. +func estimateMapsMergeCost(estimator checker.CostEstimator, _ *checker.AstNode, args []checker.AstNode) *checker.CallEstimate { + if len(args) != 2 { + return nil + } + lhsSize := estimateSize(estimator, args[0]) + rhsSize := estimateSize(estimator, args[1]) + entries := lhsSize.Add(rhsSize) + cost := entries.MultiplyByCostFactor(1).Add(mapAllocCost).Add(callCostEstimate) + // The result holds at least as many entries as the larger input, when every + // key collides, and at most the sum of both, when none do. + resultSize := rangedSizeEstimate(max(lhsSize.Min, rhsSize.Min), entries.Max) + return callEstimate(cost, &resultSize) +} + +// trackMapsMergeCost mirrors estimateMapsMergeCost against the actual inputs. +func trackMapsMergeCost(args []ref.Val, _ ref.Val) *uint64 { + entries := safeAdd(actualSize(args[0]), actualSize(args[1])) + cost := safeAdd(callCost, uint64(common.MapCreateBaseCost), entries) + return &cost +} + +// mapsMerge returns a new map holding the entries of both inputs, with the +// values of the second input taking precedence on conflicting keys. +func mapsMerge(lhs, rhs ref.Val) ref.Val { + first, ok := lhs.(traits.Mapper) + if !ok { + return types.MaybeNoSuchOverloadErr(lhs) + } + second, ok := rhs.(traits.Mapper) + if !ok { + return types.MaybeNoSuchOverloadErr(rhs) + } + merged := make(map[ref.Val]ref.Val, actualSize(first)+actualSize(second)) + if err := copyEntries(first, merged); err != nil { + return err + } + if err := copyEntries(second, merged); err != nil { + return err + } + return types.NewRefValMap(types.DefaultTypeAdapter, merged) +} + +// copyEntries writes every entry of m into dst, overwriting entries whose keys +// are already present. It returns a non-nil ref.Val only when the map yields an +// error or unknown value. +func copyEntries(m traits.Mapper, dst map[ref.Val]ref.Val) ref.Val { + it := m.Iterator() + for it.HasNext() == types.True { + key := it.Next() + if types.IsUnknownOrError(key) { + return key + } + val, _ := m.Find(key) + if types.IsUnknownOrError(val) { + return val + } + dst[key] = val + } + return nil +} diff --git a/ext/maps_test.go b/ext/maps_test.go new file mode 100644 index 000000000..e48452609 --- /dev/null +++ b/ext/maps_test.go @@ -0,0 +1,245 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ext + +import ( + "testing" + + "github.com/google/cel-go/cel" + "github.com/google/cel-go/checker" + "github.com/google/cel-go/common/ast" + "github.com/google/cel-go/common/types" +) + +func TestMapsMerge(t *testing.T) { + tests := []struct { + name string + expr string + vars []cel.EnvOption + in map[string]any + }{ + { + name: "both empty", + expr: `maps.merge({}, {}) == {}`, + }, + { + name: "disjoint keys", + expr: `maps.merge({'a': 1}, {'b': 2}) == {'a': 1, 'b': 2}`, + }, + { + name: "empty left", + expr: `maps.merge({}, {'b': 2}) == {'b': 2}`, + }, + { + name: "empty right", + expr: `maps.merge({'a': 1}, {}) == {'a': 1}`, + }, + { + name: "conflicting key takes the second value", + expr: `maps.merge({'a': 1}, {'a': 2}) == {'a': 2}`, + }, + { + name: "conflict among other keys", + expr: `maps.merge({'a': 1, 'b': 2}, {'b': 3, 'c': 4}) == {'a': 1, 'b': 3, 'c': 4}`, + }, + { + name: "map values are replaced, not merged", + expr: `maps.merge({'a': {'x': 1}}, {'a': {'y': 2}}) == {'a': {'y': 2}}`, + }, + { + name: "int keys", + expr: `maps.merge({1: 'a'}, {2: 'b'}) == {1: 'a', 2: 'b'}`, + }, + { + name: "inputs are unchanged", + expr: `maps.merge({'a': 1}, {'a': 2}) == {'a': 2} && {'a': 1} == {'a': 1}`, + }, + { + name: "merging a variable", + expr: `maps.merge(x, {'b': 2}) == {'a': 1, 'b': 2}`, + vars: []cel.EnvOption{cel.Variable("x", cel.MapType(cel.StringType, cel.IntType))}, + in: map[string]any{"x": map[string]int64{"a": 1}}, + }, + { + name: "merge is associative when keys are disjoint", + expr: `maps.merge(maps.merge({'a': 1}, {'b': 2}), {'c': 3}) == + maps.merge({'a': 1}, maps.merge({'b': 2}, {'c': 3}))`, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + opts := append([]cel.EnvOption{Maps()}, tc.vars...) + env, err := cel.NewEnv(opts...) + if err != nil { + t.Fatalf("cel.NewEnv() failed: %v", err) + } + ast, iss := env.Compile(tc.expr) + if iss.Err() != nil { + t.Fatalf("env.Compile(%q) failed: %v", tc.expr, iss.Err()) + } + prg, err := env.Program(ast) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + in := tc.in + if in == nil { + in = map[string]any{} + } + out, _, err := prg.Eval(in) + if err != nil { + t.Fatalf("prg.Eval() failed: %v", err) + } + if out != types.True { + t.Errorf("prg.Eval(%q) got %v, wanted true", tc.expr, out) + } + }) + } +} + +func TestMapsMergeTypeChecking(t *testing.T) { + env, err := cel.NewEnv(Maps()) + if err != nil { + t.Fatalf("cel.NewEnv() failed: %v", err) + } + tests := []struct { + name string + expr string + }{ + {name: "list argument", expr: `maps.merge([1], [2])`}, + {name: "string argument", expr: `maps.merge('a', 'b')`}, + {name: "mixed arguments", expr: `maps.merge({'a': 1}, [2])`}, + {name: "single argument", expr: `maps.merge({'a': 1})`}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, iss := env.Compile(tc.expr) + if iss.Err() == nil { + t.Errorf("env.Compile(%q) succeeded, wanted a type-checking error", tc.expr) + } + }) + } +} + +func TestMapsMergeNonMapArgs(t *testing.T) { + env, err := cel.NewEnv(Maps(), + cel.Variable("x", cel.DynType), + cel.Variable("y", cel.DynType)) + if err != nil { + t.Fatalf("cel.NewEnv() failed: %v", err) + } + tests := []struct { + name string + in map[string]any + }{ + {name: "left is not a map", in: map[string]any{"x": []int64{1}, "y": map[string]int64{"a": 1}}}, + {name: "right is not a map", in: map[string]any{"x": map[string]int64{"a": 1}, "y": "b"}}, + {name: "neither is a map", in: map[string]any{"x": 1, "y": 2}}, + } + ast, iss := env.Compile(`maps.merge(x, y)`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + prg, err := env.Program(ast) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, _, err := prg.Eval(tc.in) + if err == nil { + t.Errorf("prg.Eval(%v) succeeded, wanted an error", tc.in) + } + }) + } +} + +func TestMapsMergeCost(t *testing.T) { + env, err := cel.NewEnv(Maps()) + if err != nil { + t.Fatalf("cel.NewEnv() failed: %v", err) + } + ast, iss := env.Compile(`maps.merge({'a': 1, 'b': 2}, {'c': 3})`) + if iss.Err() != nil { + t.Fatalf("env.Compile() failed: %v", iss.Err()) + } + est, err := env.EstimateCost(ast, testCostHintEstimator{}) + if err != nil { + t.Fatalf("env.EstimateCost() failed: %v", err) + } + // The estimate must scale with the inputs rather than collapse to the + // default O(1) bucket a function without an estimator would land in. + if est.Min <= 1 { + t.Errorf("env.EstimateCost() min got %d, wanted a size-dependent estimate", est.Min) + } + prg, err := env.Program(ast, cel.CostTracking(nil)) + if err != nil { + t.Fatalf("env.Program() failed: %v", err) + } + _, det, err := prg.Eval(map[string]any{}) + if err != nil { + t.Fatalf("prg.Eval() failed: %v", err) + } + actual := det.ActualCost() + if actual == nil { + t.Fatal("det.ActualCost() was nil, wanted a tracked cost") + } + if *actual < est.Min || *actual > est.Max { + t.Errorf("det.ActualCost() got %d, outside the estimate [%d, %d]", *actual, est.Min, est.Max) + } +} + +func TestMapsMergeResultSize(t *testing.T) { + // A merge of a 2-entry and a 1-entry map holds 2 entries when every key + // collides and 3 when none do. + est := estimateMapsMergeCost(testCostHintEstimator{}, nil, []checker.AstNode{ + testSizedNode{size: 2}, + testSizedNode{size: 1}, + }) + if est == nil { + t.Fatal("estimateMapsMergeCost() returned nil") + } + if est.ResultSize == nil { + t.Fatal("estimateMapsMergeCost() ResultSize was nil") + } + if est.ResultSize.Min != 2 || est.ResultSize.Max != 3 { + t.Errorf("ResultSize got [%d, %d], wanted [2, 3]", + est.ResultSize.Min, est.ResultSize.Max) + } +} + +// testSizedNode is a checker.AstNode whose size is known, used to pin the +// result-size bounds of the merge estimator. +type testSizedNode struct { + size uint64 +} + +func (n testSizedNode) Path() []string { return nil } + +func (n testSizedNode) Type() *cel.Type { return cel.MapType(cel.StringType, cel.IntType) } + +func (n testSizedNode) Expr() ast.Expr { return nil } + +func (n testSizedNode) ComputedSize() *checker.SizeEstimate { + sz := checker.SizeEstimate{Min: n.size, Max: n.size} + return &sz +} + +func TestMapsVersion(t *testing.T) { + _, err := cel.NewEnv(Maps(MapsVersion(0))) + if err != nil { + t.Fatalf("MapsVersion(0) failed: %v", err) + } +} diff --git a/repl/evaluator.go b/repl/evaluator.go index ecfd72264..514e6d31d 100644 --- a/repl/evaluator.go +++ b/repl/evaluator.go @@ -55,6 +55,7 @@ var ( "encoders": ext.Encoders(), "sets": ext.Sets(), "lists": ext.Lists(), + "maps": ext.Maps(), "two_var_comprehensions": ext.TwoVarComprehensions(), } )