diff --git a/checker/internal/type_checker_impl.cc b/checker/internal/type_checker_impl.cc index 6cda11e4f..9adf4f164 100644 --- a/checker/internal/type_checker_impl.cc +++ b/checker/internal/type_checker_impl.cc @@ -807,13 +807,17 @@ void ResolveVisitor::PostVisitComprehensionSubexpression( break; case ComprehensionArg::ITER_RANGE: { Type range_type = GetTypeOrDyn(&comprehension.iter_range()); - Type iter_type = DynType(); + Type iter_type = DynType(); // iter_var for non comprehensions v2. + Type iter_type1 = DynType(); // iter_var for comprehensions v2. + Type iter_type2 = DynType(); // iter_var2 for comprehensions v2. switch (range_type.kind()) { case TypeKind::kList: - iter_type = range_type.GetList().element(); + iter_type1 = IntType(); + iter_type = iter_type2 = range_type.GetList().element(); break; case TypeKind::kMap: - iter_type = range_type.GetMap().key(); + iter_type = iter_type1 = range_type.GetMap().key(); + iter_type2 = range_type.GetMap().value(); break; case TypeKind::kDyn: break; @@ -827,8 +831,15 @@ void ResolveVisitor::PostVisitComprehensionSubexpression( "list, map, or dynamic)"))); break; } - scope.iter_scope->InsertVariableIfAbsent( - MakeVariableDecl(comprehension.iter_var(), iter_type)); + if (comprehension.iter_var2().empty()) { + scope.iter_scope->InsertVariableIfAbsent( + MakeVariableDecl(comprehension.iter_var(), iter_type)); + } else { + scope.iter_scope->InsertVariableIfAbsent( + MakeVariableDecl(comprehension.iter_var(), iter_type1)); + scope.iter_scope->InsertVariableIfAbsent( + MakeVariableDecl(comprehension.iter_var2(), iter_type2)); + } break; } case ComprehensionArg::RESULT: diff --git a/checker/standard_library.cc b/checker/standard_library.cc index dcdda3fb8..18c43a8ad 100644 --- a/checker/standard_library.cc +++ b/checker/standard_library.cc @@ -1058,6 +1058,17 @@ absl::Status AddEnumConstants(TypeCheckerBuilder& builder) { return absl::OkStatus(); } +absl::Status AddComprehensionsV2Functions(TypeCheckerBuilder& builder) { + FunctionDecl map_insert; + map_insert.set_name("@cel.mapInsert"); + CEL_RETURN_IF_ERROR(map_insert.AddOverload( + MakeOverloadDecl("@mapInsert_map_key_value", MapOfAB(), MapOfAB(), + TypeParamA(), TypeParamB()))); + CEL_RETURN_IF_ERROR(map_insert.AddOverload( + MakeOverloadDecl("@mapInsert_map_map", MapOfAB(), MapOfAB(), MapOfAB()))); + return builder.AddFunction(map_insert); +} + absl::Status AddStandardLibraryDecls(TypeCheckerBuilder& builder) { CEL_RETURN_IF_ERROR(AddLogicalOps(builder)); CEL_RETURN_IF_ERROR(AddArithmeticOps(builder)); @@ -1070,7 +1081,7 @@ absl::Status AddStandardLibraryDecls(TypeCheckerBuilder& builder) { CEL_RETURN_IF_ERROR(AddTimeFunctions(builder)); CEL_RETURN_IF_ERROR(AddTypeConstantVariables(builder)); CEL_RETURN_IF_ERROR(AddEnumConstants(builder)); - + CEL_RETURN_IF_ERROR(AddComprehensionsV2Functions(builder)); return absl::OkStatus(); }