|
| 1 | +// Copyright 2026 Google LLC |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +package ext |
| 16 | + |
| 17 | +import ( |
| 18 | + "github.com/google/cel-go/cel" |
| 19 | + "github.com/google/cel-go/checker" |
| 20 | + "github.com/google/cel-go/common" |
| 21 | + "github.com/google/cel-go/common/types" |
| 22 | + "github.com/google/cel-go/common/types/ref" |
| 23 | + "github.com/google/cel-go/common/types/traits" |
| 24 | + "github.com/google/cel-go/interpreter" |
| 25 | +) |
| 26 | + |
| 27 | +// Maps returns a cel.EnvOption to configure namespaced map functions. |
| 28 | +// |
| 29 | +// CEL has no operator for combining two maps: the `+` operator concatenates |
| 30 | +// strings, bytes, and lists, but is not defined for maps. This library provides |
| 31 | +// map combination as a named function. |
| 32 | +// |
| 33 | +// # Maps.Merge |
| 34 | +// |
| 35 | +// Returns a new map containing the entries of both arguments. When a key is |
| 36 | +// present in both, the value from the second argument wins. Neither input is |
| 37 | +// modified. |
| 38 | +// |
| 39 | +// The merge is shallow: a value that is itself a map is replaced rather than |
| 40 | +// merged recursively. |
| 41 | +// |
| 42 | +// maps.merge(map(K, V), map(K, V)) -> map(K, V) |
| 43 | +// |
| 44 | +// Examples: |
| 45 | +// |
| 46 | +// maps.merge({}, {}) // {} |
| 47 | +// maps.merge({'a': 1}, {'b': 2}) // {'a': 1, 'b': 2} |
| 48 | +// maps.merge({'a': 1}, {'a': 2}) // {'a': 2} |
| 49 | +// maps.merge({'a': {'x': 1}}, {'a': {'y': 2}}) // {'a': {'y': 2}}, values are replaced, not merged |
| 50 | +func Maps(options ...MapsOption) cel.EnvOption { |
| 51 | + l := &mapsLib{} |
| 52 | + for _, o := range options { |
| 53 | + l = o(l) |
| 54 | + } |
| 55 | + return cel.Lib(l) |
| 56 | +} |
| 57 | + |
| 58 | +// MapsOption declares a functional operator for configuring map extensions. |
| 59 | +type MapsOption func(*mapsLib) *mapsLib |
| 60 | + |
| 61 | +// MapsVersion sets the library version for map extensions. |
| 62 | +func MapsVersion(version uint32) MapsOption { |
| 63 | + return func(lib *mapsLib) *mapsLib { |
| 64 | + lib.version = version |
| 65 | + return lib |
| 66 | + } |
| 67 | +} |
| 68 | + |
| 69 | +type mapsLib struct { |
| 70 | + version uint32 |
| 71 | +} |
| 72 | + |
| 73 | +// LibraryName implements the SingletonLibrary interface method. |
| 74 | +func (mapsLib) LibraryName() string { |
| 75 | + return "cel.lib.ext.maps" |
| 76 | +} |
| 77 | + |
| 78 | +// CompileOptions implements the Library interface method. |
| 79 | +func (mapsLib) CompileOptions() []cel.EnvOption { |
| 80 | + mapType := cel.MapType(cel.TypeParamType("K"), cel.TypeParamType("V")) |
| 81 | + return []cel.EnvOption{ |
| 82 | + cel.Function("maps.merge", |
| 83 | + cel.Overload("map_maps_merge_map", []*cel.Type{mapType, mapType}, mapType, |
| 84 | + cel.BinaryBinding(mapsMerge))), |
| 85 | + cel.CostEstimatorOptions( |
| 86 | + checker.OverloadCostEstimate("map_maps_merge_map", estimateMapsMergeCost), |
| 87 | + ), |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +// ProgramOptions implements the Library interface method. |
| 92 | +func (mapsLib) ProgramOptions() []cel.ProgramOption { |
| 93 | + return []cel.ProgramOption{ |
| 94 | + cel.CostTrackerOptions( |
| 95 | + interpreter.OverloadCostTracker("map_maps_merge_map", trackMapsMergeCost), |
| 96 | + ), |
| 97 | + } |
| 98 | +} |
| 99 | + |
| 100 | +// estimateMapsMergeCost charges for visiting every entry of both inputs and for |
| 101 | +// allocating the result map. |
| 102 | +func estimateMapsMergeCost(estimator checker.CostEstimator, _ *checker.AstNode, args []checker.AstNode) *checker.CallEstimate { |
| 103 | + if len(args) != 2 { |
| 104 | + return nil |
| 105 | + } |
| 106 | + lhsSize := estimateSize(estimator, args[0]) |
| 107 | + rhsSize := estimateSize(estimator, args[1]) |
| 108 | + entries := lhsSize.Add(rhsSize) |
| 109 | + cost := entries.MultiplyByCostFactor(1).Add(mapAllocCost).Add(callCostEstimate) |
| 110 | + // The result holds at least as many entries as the larger input, when every |
| 111 | + // key collides, and at most the sum of both, when none do. |
| 112 | + resultSize := rangedSizeEstimate(max(lhsSize.Min, rhsSize.Min), entries.Max) |
| 113 | + return callEstimate(cost, &resultSize) |
| 114 | +} |
| 115 | + |
| 116 | +// trackMapsMergeCost mirrors estimateMapsMergeCost against the actual inputs. |
| 117 | +func trackMapsMergeCost(args []ref.Val, _ ref.Val) *uint64 { |
| 118 | + entries := safeAdd(actualSize(args[0]), actualSize(args[1])) |
| 119 | + cost := safeAdd(callCost, uint64(common.MapCreateBaseCost), entries) |
| 120 | + return &cost |
| 121 | +} |
| 122 | + |
| 123 | +// mapsMerge returns a new map holding the entries of both inputs, with the |
| 124 | +// values of the second input taking precedence on conflicting keys. |
| 125 | +func mapsMerge(lhs, rhs ref.Val) ref.Val { |
| 126 | + first, ok := lhs.(traits.Mapper) |
| 127 | + if !ok { |
| 128 | + return types.MaybeNoSuchOverloadErr(lhs) |
| 129 | + } |
| 130 | + second, ok := rhs.(traits.Mapper) |
| 131 | + if !ok { |
| 132 | + return types.MaybeNoSuchOverloadErr(rhs) |
| 133 | + } |
| 134 | + merged := make(map[ref.Val]ref.Val, actualSize(first)+actualSize(second)) |
| 135 | + if err := copyEntries(first, merged); err != nil { |
| 136 | + return err |
| 137 | + } |
| 138 | + if err := copyEntries(second, merged); err != nil { |
| 139 | + return err |
| 140 | + } |
| 141 | + return types.NewRefValMap(types.DefaultTypeAdapter, merged) |
| 142 | +} |
| 143 | + |
| 144 | +// copyEntries writes every entry of m into dst, overwriting entries whose keys |
| 145 | +// are already present. It returns a non-nil ref.Val only when the map yields an |
| 146 | +// error or unknown value. |
| 147 | +func copyEntries(m traits.Mapper, dst map[ref.Val]ref.Val) ref.Val { |
| 148 | + it := m.Iterator() |
| 149 | + for it.HasNext() == types.True { |
| 150 | + key := it.Next() |
| 151 | + if types.IsUnknownOrError(key) { |
| 152 | + return key |
| 153 | + } |
| 154 | + val, _ := m.Find(key) |
| 155 | + if types.IsUnknownOrError(val) { |
| 156 | + return val |
| 157 | + } |
| 158 | + dst[key] = val |
| 159 | + } |
| 160 | + return nil |
| 161 | +} |
0 commit comments