Skip to content

Commit aaf917e

Browse files
CEL Dev Teamcopybara-github
CEL Dev Team
authored andcommitted
Adding mapInsert internal runtime function
PiperOrigin-RevId: 747580855
1 parent 5b4b234 commit aaf917e

File tree

5 files changed

+378
-2
lines changed

5 files changed

+378
-2
lines changed

extensions/src/main/java/dev/cel/extensions/BUILD.bazel

+20
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ java_library(
1717
],
1818
deps = [
1919
":bindings",
20+
":comprehensions",
2021
":encoders",
2122
":lists",
2223
":math",
@@ -171,3 +172,22 @@ java_library(
171172
"@maven//:com_google_guava_guava",
172173
],
173174
)
175+
176+
java_library(
177+
name = "comprehensions",
178+
srcs = ["CelComprehensions.java"],
179+
tags = [
180+
],
181+
deps = [
182+
"//checker:checker_builder",
183+
"//common:compiler_common",
184+
"//common/ast",
185+
"//common/types",
186+
"//compiler:compiler_builder",
187+
"//parser:macro",
188+
"//parser:parser_builder",
189+
"//runtime",
190+
"//runtime:function_binding",
191+
"@maven//:com_google_guava_guava",
192+
],
193+
)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
// Copyright 2025 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+
// https://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 dev.cel.extensions;
16+
17+
import com.google.common.collect.ImmutableList;
18+
import com.google.common.collect.ImmutableSet;
19+
import com.google.common.collect.Maps;
20+
import dev.cel.checker.CelCheckerBuilder;
21+
import dev.cel.common.CelFunctionDecl;
22+
import dev.cel.common.CelIssue;
23+
import dev.cel.common.CelOverloadDecl;
24+
import dev.cel.common.ast.CelExpr;
25+
import dev.cel.common.ast.CelExpr.ExprKind.Kind;
26+
import dev.cel.common.types.MapType;
27+
import dev.cel.common.types.TypeParamType;
28+
import dev.cel.compiler.CelCompilerLibrary;
29+
import dev.cel.parser.CelMacro;
30+
import dev.cel.parser.CelMacroExprFactory;
31+
import dev.cel.parser.CelParserBuilder;
32+
import dev.cel.runtime.CelFunctionBinding;
33+
import dev.cel.runtime.CelRuntimeBuilder;
34+
import dev.cel.runtime.CelRuntimeLibrary;
35+
import java.util.List;
36+
import java.util.Map;
37+
import java.util.Optional;
38+
import java.util.Set;
39+
40+
/** */
41+
final class CelComprehensions implements CelCompilerLibrary, CelRuntimeLibrary {
42+
43+
private static final TypeParamType TYPE_PARAM_A = TypeParamType.create("A");
44+
private static final TypeParamType TYPE_PARAM_B = TypeParamType.create("B");
45+
private static final MapType MAP_OF_AB = MapType.create(TYPE_PARAM_A, TYPE_PARAM_B);
46+
private static final String CEL_NAMESPACE = "cel";
47+
private static final String MAP_INSERT_FUNCTION = "cel.@mapInsert";
48+
private static final String MAP_INSERT_OVERLOAD_MAP_MAP = "@mapInsert_map_map";
49+
private static final String MAP_INSERT_OVERLOAD_KEY_VALUE = "@mapInsert_map_key_value";
50+
51+
public enum Function {
52+
MAP_INSERT(
53+
CelFunctionDecl.newFunctionDeclaration(
54+
MAP_INSERT_FUNCTION,
55+
CelOverloadDecl.newGlobalOverload(
56+
MAP_INSERT_OVERLOAD_MAP_MAP, "map insertion", MAP_OF_AB, MAP_OF_AB, MAP_OF_AB),
57+
CelOverloadDecl.newGlobalOverload(
58+
MAP_INSERT_OVERLOAD_KEY_VALUE,
59+
"map insertion",
60+
MAP_OF_AB,
61+
MAP_OF_AB,
62+
TYPE_PARAM_A,
63+
TYPE_PARAM_B)),
64+
CelFunctionBinding.from(
65+
MAP_INSERT_OVERLOAD_MAP_MAP, Map.class, Map.class, CelComprehensions::mapInsert),
66+
CelFunctionBinding.from(
67+
MAP_INSERT_OVERLOAD_KEY_VALUE,
68+
ImmutableList.of(Map.class, Object.class, Object.class),
69+
CelComprehensions::mapInsert));
70+
71+
private final CelFunctionDecl functionDecl;
72+
private final ImmutableSet<CelFunctionBinding> functionBindings;
73+
74+
String getFunction() {
75+
return functionDecl.name();
76+
}
77+
78+
Function(CelFunctionDecl functionDecl, CelFunctionBinding... functionBindings) {
79+
this.functionDecl = functionDecl;
80+
this.functionBindings = ImmutableSet.copyOf(functionBindings);
81+
}
82+
}
83+
84+
private final ImmutableSet<Function> functions;
85+
86+
CelComprehensions() {
87+
this.functions = ImmutableSet.copyOf(Function.values());
88+
}
89+
90+
CelComprehensions(Set<Function> functions) {
91+
this.functions = ImmutableSet.copyOf(functions);
92+
}
93+
94+
@Override
95+
public void setParserOptions(CelParserBuilder parserBuilder) {
96+
parserBuilder.addMacros(
97+
CelMacro.newReceiverVarArgMacro("mapInsert", CelComprehensions::expandMapInsertMacro));
98+
}
99+
100+
@Override
101+
public void setCheckerOptions(CelCheckerBuilder checkerBuilder) {
102+
functions.forEach(function -> checkerBuilder.addFunctionDeclarations(function.functionDecl));
103+
}
104+
105+
@Override
106+
public void setRuntimeOptions(CelRuntimeBuilder runtimeBuilder) {
107+
functions.forEach(function -> runtimeBuilder.addFunctionBindings(function.functionBindings));
108+
}
109+
110+
private static Map<Object, Object> mapInsert(Map<?, ?> first, Map<?, ?> second) {
111+
// TODO: return a mutable map instead of an actual copy.
112+
Map<Object, Object> result = Maps.newHashMapWithExpectedSize(first.size() + second.size());
113+
result.putAll(first);
114+
for (Map.Entry<?, ?> entry : second.entrySet()) {
115+
if (result.containsKey(entry.getKey())) {
116+
throw new IllegalArgumentException(
117+
String.format("insert failed: key '%s' already exists", entry.getKey()));
118+
}
119+
result.put(entry.getKey(), entry.getValue());
120+
}
121+
return result;
122+
}
123+
124+
private static Map<Object, Object> mapInsert(Object[] args) {
125+
Map<?, ?> map = (Map<?, ?>) args[0];
126+
Object key = args[1];
127+
Object value = args[2];
128+
// TODO: return a mutable map instead of an actual copy.
129+
if (map.containsKey(key)) {
130+
throw new IllegalArgumentException(
131+
String.format("insert failed: key '%s' already exists", key));
132+
}
133+
Map<Object, Object> result = Maps.newHashMapWithExpectedSize(map.size() + 1);
134+
result.putAll(map);
135+
result.put(key, value);
136+
return result;
137+
}
138+
139+
private static Optional<CelExpr> expandMapInsertMacro(
140+
CelMacroExprFactory exprFactory, CelExpr target, ImmutableList<CelExpr> arguments) {
141+
if (!isTargetInNamespace(target)) {
142+
// Return empty to indicate that we're not interested in expanding this macro, and
143+
// that the parser should default to a function call on the receiver.
144+
return Optional.empty();
145+
}
146+
147+
switch (arguments.size()) {
148+
case 2:
149+
Optional<CelExpr> invalidArg =
150+
checkInvalidArgument(exprFactory, MAP_INSERT_OVERLOAD_MAP_MAP, arguments);
151+
if (invalidArg.isPresent()) {
152+
return invalidArg;
153+
}
154+
155+
return Optional.of(exprFactory.newGlobalCall(MAP_INSERT_FUNCTION, arguments));
156+
case 3:
157+
invalidArg = checkInvalidArgument(exprFactory, MAP_INSERT_OVERLOAD_KEY_VALUE, arguments);
158+
if (invalidArg.isPresent()) {
159+
return invalidArg;
160+
}
161+
162+
return Optional.of(exprFactory.newGlobalCall(MAP_INSERT_FUNCTION, arguments));
163+
default:
164+
return newError(
165+
exprFactory,
166+
"cel.mapInsert() arguments must be either two maps or a map and a key-value pair",
167+
target);
168+
}
169+
}
170+
171+
private static boolean isTargetInNamespace(CelExpr target) {
172+
return target.exprKind().getKind().equals(Kind.IDENT)
173+
&& target.ident().name().equals(CEL_NAMESPACE);
174+
}
175+
176+
private static Optional<CelExpr> checkInvalidArgument(
177+
CelMacroExprFactory exprFactory, String functionName, List<CelExpr> arguments) {
178+
179+
if (functionName.equals(MAP_INSERT_OVERLOAD_MAP_MAP)) {
180+
for (CelExpr arg : arguments) {
181+
if (arg.exprKind().getKind() != Kind.MAP) {
182+
return newError(
183+
exprFactory, String.format("Invalid argument '%s': must be a map", arg), arg);
184+
}
185+
}
186+
}
187+
if (functionName.equals(MAP_INSERT_OVERLOAD_KEY_VALUE)) {
188+
if (arguments.get(0).exprKind().getKind() != Kind.MAP) {
189+
return newError(
190+
exprFactory,
191+
String.format("Invalid argument '%s': must be a map", arguments.get(0)),
192+
arguments.get(0));
193+
}
194+
if (arguments.get(1).exprKind().getKind() != Kind.CONSTANT) {
195+
return newError(
196+
exprFactory,
197+
String.format("'%s' is an invalid Key", arguments.get(1)),
198+
arguments.get(1));
199+
}
200+
}
201+
202+
return Optional.empty();
203+
}
204+
205+
private static Optional<CelExpr> newError(
206+
CelMacroExprFactory exprFactory, String errorMessage, CelExpr argument) {
207+
return Optional.of(
208+
exprFactory.reportError(
209+
CelIssue.formatError(exprFactory.getSourceLocation(argument), errorMessage)));
210+
}
211+
}

extensions/src/main/java/dev/cel/extensions/CelExtensions.java

+7-1
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,10 @@ public static CelListsExtensions lists(Set<CelListsExtensions.Function> function
258258
return new CelListsExtensions(functions);
259259
}
260260

261+
public static CelComprehensions comprehensions() {
262+
return new CelComprehensions();
263+
}
264+
261265
/**
262266
* Retrieves all function names used by every extension libraries.
263267
*
@@ -276,7 +280,9 @@ public static ImmutableSet<String> getAllFunctionNames() {
276280
stream(CelEncoderExtensions.Function.values())
277281
.map(CelEncoderExtensions.Function::getFunction),
278282
stream(CelListsExtensions.Function.values())
279-
.map(CelListsExtensions.Function::getFunction))
283+
.map(CelListsExtensions.Function::getFunction),
284+
stream(CelComprehensions.Function.values())
285+
.map(CelComprehensions.Function::getFunction))
280286
.collect(toImmutableSet());
281287
}
282288

0 commit comments

Comments
 (0)