Skip to content

Commit c4cfc04

Browse files
CEL Dev Teamcopybara-github
CEL Dev Team
authored and
copybara-github
committed
Add a ProtoTypeMask class.
PiperOrigin-RevId: 818745535
1 parent 38f6e94 commit c4cfc04

File tree

4 files changed

+734
-0
lines changed

4 files changed

+734
-0
lines changed

‎checker/internal/BUILD‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,3 +271,38 @@ cc_test(
271271
"@com_google_protobuf//:protobuf",
272272
],
273273
)
274+
275+
cc_library(
276+
name = "proto_type_mask",
277+
srcs = ["proto_type_mask.cc"],
278+
hdrs = ["proto_type_mask.h"],
279+
deps = [
280+
"//common:type",
281+
"//internal:status_macros",
282+
"//util/gtl:iterator_adaptors",
283+
"@com_google_absl//absl/container:flat_hash_map",
284+
"@com_google_absl//absl/memory",
285+
"@com_google_absl//absl/status",
286+
"@com_google_absl//absl/status:statusor",
287+
"@com_google_absl//absl/strings",
288+
"@com_google_absl//absl/strings:str_format",
289+
"@com_google_absl//absl/strings:string_view",
290+
"@com_google_absl//absl/types:optional",
291+
"@com_google_protobuf//:protobuf",
292+
],
293+
)
294+
295+
cc_test(
296+
name = "proto_type_mask_test",
297+
srcs = ["proto_type_mask_test.cc"],
298+
deps = [
299+
":proto_type_mask",
300+
"//internal:testing",
301+
"//internal:testing_descriptor_pool",
302+
"@com_google_absl//absl/container:flat_hash_map",
303+
"@com_google_absl//absl/status",
304+
"@com_google_absl//absl/status:status_matchers",
305+
"@com_google_absl//absl/strings",
306+
"@com_google_absl//absl/strings:string_view",
307+
],
308+
)

‎checker/internal/proto_type_mask.cc‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
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+
#include "checker/internal/proto_type_mask.h"
16+
17+
#include <iterator>
18+
#include <memory>
19+
#include <set>
20+
#include <string>
21+
#include <vector>
22+
23+
#include "absl/container/flat_hash_map.h"
24+
#include "absl/memory/memory.h"
25+
#include "absl/status/status.h"
26+
#include "absl/status/statusor.h"
27+
#include "absl/strings/str_split.h"
28+
#include "absl/strings/string_view.h"
29+
#include "absl/strings/substitute.h"
30+
#include "absl/types/optional.h"
31+
#include "common/type.h"
32+
#include "internal/status_macros.h"
33+
#include "google/protobuf/descriptor.h"
34+
#include "util/gtl/iterator_adaptors.h"
35+
36+
using ::google::protobuf::Descriptor;
37+
using ::google::protobuf::DescriptorPool;
38+
using ::google::protobuf::FieldDescriptor;
39+
using TypeMap = absl::flat_hash_map<std::string, std::set<std::string>>;
40+
41+
absl::StatusOr<std::unique_ptr<ProtoTypeMask>> ProtoTypeMask::Create(
42+
const std::shared_ptr<const DescriptorPool>& descriptor_pool,
43+
const TypeMap& types_and_field_paths) {
44+
if (descriptor_pool == nullptr) {
45+
return absl::InvalidArgumentError(
46+
"ProtoTypeMask descriptor pool cannot be nullptr");
47+
}
48+
CEL_ASSIGN_OR_RETURN(
49+
auto types_and_visible_fields,
50+
ComputeVisibleFieldsMap(descriptor_pool, types_and_field_paths));
51+
std::unique_ptr<ProtoTypeMask> proto_type_mask =
52+
absl::WrapUnique(new ProtoTypeMask(types_and_visible_fields));
53+
return proto_type_mask;
54+
}
55+
56+
bool ProtoTypeMask::FieldIsVisible(absl::string_view type_name,
57+
absl::string_view field_name) {
58+
auto iterator = types_and_visible_fields_.find(type_name);
59+
if (iterator != types_and_visible_fields_.end() &&
60+
!iterator->second.contains(field_name.data())) {
61+
return false;
62+
}
63+
return true;
64+
}
65+
66+
absl::StatusOr<std::set<absl::string_view>> ProtoTypeMask::GetFieldNames(
67+
const std::shared_ptr<const google::protobuf::DescriptorPool>& descriptor_pool,
68+
absl::string_view type_name, const std::set<std::string>& field_paths) {
69+
CEL_ASSIGN_OR_RETURN(const Descriptor* descriptor,
70+
FindMessage(descriptor_pool, type_name));
71+
std::set<absl::string_view> field_numbers;
72+
for (absl::string_view field_path : field_paths) {
73+
std::vector<std::string> field_selection = Split(field_path);
74+
absl::string_view field_name = field_selection.front();
75+
CEL_ASSIGN_OR_RETURN(const FieldDescriptor* field_descriptor,
76+
FindField(descriptor, field_name));
77+
field_numbers.insert(field_descriptor->name());
78+
}
79+
return field_numbers;
80+
}
81+
82+
std::vector<std::string> ProtoTypeMask::Split(absl::string_view field_path) {
83+
return absl::StrSplit(field_path, kPathDelimiter);
84+
}
85+
86+
absl::StatusOr<const Descriptor*> ProtoTypeMask::FindMessage(
87+
const std::shared_ptr<const DescriptorPool>& descriptor_pool,
88+
absl::string_view type_name) {
89+
const Descriptor* descriptor =
90+
descriptor_pool->FindMessageTypeByName(type_name);
91+
if (descriptor == nullptr) {
92+
return absl::InvalidArgumentError(absl::Substitute(
93+
"ProtoTypeMask type not found (type name: '$0')", type_name));
94+
}
95+
return descriptor;
96+
}
97+
98+
absl::StatusOr<const FieldDescriptor*> ProtoTypeMask::FindField(
99+
const Descriptor* descriptor, absl::string_view field_name) {
100+
const FieldDescriptor* field_descriptor =
101+
descriptor->FindFieldByName(field_name);
102+
if (field_descriptor == nullptr) {
103+
return absl::InvalidArgumentError(
104+
absl::Substitute("ProtoTypeMask could not select field from type "
105+
"(type name: '$0', field name: '$1')",
106+
descriptor->full_name(), field_name));
107+
}
108+
return field_descriptor;
109+
}
110+
111+
absl::StatusOr<const Descriptor*> ProtoTypeMask::GetMessage(
112+
const FieldDescriptor* field_descriptor) {
113+
cel::MessageTypeField field(field_descriptor);
114+
cel::Type type = field.GetType();
115+
absl::optional<cel::MessageType> message_type = type.AsMessage();
116+
if (!message_type.has_value()) {
117+
return absl::InvalidArgumentError(
118+
absl::Substitute("ProtoTypeMask subfield is not a message type "
119+
"(field name: '$0', type: '$1')",
120+
field_descriptor->name(), type.name()));
121+
}
122+
return &(*message_type.value());
123+
}
124+
125+
absl::Status ProtoTypeMask::AddAllHiddenFields(
126+
TypeMap& types_and_visible_fields, absl::string_view type_name) {
127+
auto result = types_and_visible_fields.find(type_name);
128+
if (result != types_and_visible_fields.end()) {
129+
if (!result->second.empty()) {
130+
return absl::InvalidArgumentError(absl::Substitute(
131+
"ProtoTypeMask cannot insert all hidden fields when the type has "
132+
"already been inserted with a visible field (type name: '$0')",
133+
type_name));
134+
}
135+
return absl::OkStatus();
136+
}
137+
types_and_visible_fields.insert({type_name.data(), {}});
138+
return absl::OkStatus();
139+
}
140+
141+
absl::Status ProtoTypeMask::AddVisibleField(TypeMap& types_and_visible_fields,
142+
absl::string_view type_name,
143+
absl::string_view field_name) {
144+
auto result = types_and_visible_fields.find(type_name);
145+
if (result != types_and_visible_fields.end()) {
146+
if (result->second.empty()) {
147+
return absl::InvalidArgumentError(
148+
absl::Substitute("ProtoTypeMask cannot insert a visible field when "
149+
"the type has already been "
150+
"inserted with all hidden fields (type name: "
151+
"'$0', field name: '$1')",
152+
type_name, field_name));
153+
}
154+
result->second.insert(field_name.data());
155+
return absl::OkStatus();
156+
}
157+
types_and_visible_fields.insert({type_name.data(), {field_name.data()}});
158+
return absl::OkStatus();
159+
}
160+
161+
absl::StatusOr<TypeMap> ProtoTypeMask::ComputeVisibleFieldsMap(
162+
const std::shared_ptr<const DescriptorPool>& descriptor_pool,
163+
const TypeMap& types_and_field_paths) {
164+
TypeMap types_and_visible_fields;
165+
for (absl::string_view type_name : gtl::key_view(types_and_field_paths)) {
166+
CEL_ASSIGN_OR_RETURN(const Descriptor* descriptor,
167+
FindMessage(descriptor_pool, type_name));
168+
std::set<std::string> field_paths =
169+
types_and_field_paths.find(type_name)->second;
170+
if (field_paths.empty()) {
171+
CEL_RETURN_IF_ERROR(
172+
AddAllHiddenFields(types_and_visible_fields, type_name));
173+
}
174+
for (absl::string_view field_path : field_paths) {
175+
const Descriptor* target_descriptor = descriptor;
176+
std::vector<std::string> field_selection = Split(field_path);
177+
for (auto iterator = field_selection.begin();
178+
iterator != field_selection.end(); ++iterator) {
179+
CEL_ASSIGN_OR_RETURN(const FieldDescriptor* field_descriptor,
180+
FindField(target_descriptor, *iterator));
181+
CEL_RETURN_IF_ERROR(AddVisibleField(types_and_visible_fields,
182+
target_descriptor->full_name(),
183+
*iterator));
184+
if (std::next(iterator) != field_selection.end()) {
185+
CEL_ASSIGN_OR_RETURN(target_descriptor, GetMessage(field_descriptor));
186+
}
187+
}
188+
}
189+
}
190+
return types_and_visible_fields;
191+
}

‎checker/internal/proto_type_mask.h‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
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+
#ifndef THIRD_PARTY_CEL_CPP_CHECKER_PROTO_TYPE_MASK_H_
16+
#define THIRD_PARTY_CEL_CPP_CHECKER_PROTO_TYPE_MASK_H_
17+
18+
#include <memory>
19+
#include <set>
20+
#include <string>
21+
#include <utility>
22+
#include <vector>
23+
24+
#include "absl/container/flat_hash_map.h"
25+
#include "absl/status/status.h"
26+
#include "absl/status/statusor.h"
27+
#include "absl/strings/str_cat.h"
28+
#include "absl/strings/str_format.h"
29+
#include "absl/strings/str_join.h"
30+
#include "absl/strings/string_view.h"
31+
#include "google/protobuf/descriptor.h"
32+
#include "util/gtl/iterator_adaptors.h"
33+
34+
class ProtoTypeMask {
35+
public:
36+
static absl::StatusOr<std::unique_ptr<ProtoTypeMask>> Create(
37+
const std::shared_ptr<const google::protobuf::DescriptorPool>& descriptor_pool,
38+
const absl::flat_hash_map<std::string, std::set<std::string>>&
39+
types_and_field_paths);
40+
41+
absl::flat_hash_map<std::string, std::set<std::string>>
42+
GetTypesAndVisibleFields() {
43+
return types_and_visible_fields_;
44+
}
45+
46+
// Returns true when the field name is visible. A field is visible if:
47+
// 1. The type name is not a key in the map.
48+
// 2. The type name is a key in the map and the field name is in the set of
49+
// field names that are visible for the type.
50+
bool FieldIsVisible(absl::string_view type_name,
51+
absl::string_view field_name);
52+
53+
static absl::StatusOr<std::set<absl::string_view>> GetFieldNames(
54+
const std::shared_ptr<const google::protobuf::DescriptorPool>& descriptor_pool,
55+
absl::string_view type_name, const std::set<std::string>& field_paths);
56+
57+
template <typename Sink>
58+
friend void AbslStringify(Sink& sink, const ProtoTypeMask& proto_type_mask) {
59+
std::string new_line = "\n";
60+
std::string output = absl::StrCat(new_line, "ProtoTypeMask {");
61+
auto types_and_visible_fields = proto_type_mask.types_and_visible_fields_;
62+
for (absl::string_view type_name :
63+
gtl::key_view(types_and_visible_fields)) {
64+
std::set<std::string> visible_fields =
65+
types_and_visible_fields.find(type_name)->second;
66+
absl::StrAppend(
67+
&output, new_line, "{", new_line, " type: ", type_name, new_line,
68+
" fields: ", absl::StrJoin(visible_fields, ", "), new_line, "}");
69+
}
70+
absl::StrAppend(&output, new_line, "}");
71+
absl::Format(&sink, "%v", output);
72+
}
73+
74+
private:
75+
static std::vector<std::string> Split(absl::string_view field_path);
76+
77+
static absl::StatusOr<const google::protobuf::Descriptor*> FindMessage(
78+
const std::shared_ptr<const google::protobuf::DescriptorPool>& descriptor_pool,
79+
absl::string_view type_name);
80+
81+
static absl::StatusOr<const google::protobuf::FieldDescriptor*> FindField(
82+
const google::protobuf::Descriptor* descriptor, absl::string_view field_name);
83+
84+
static absl::StatusOr<const google::protobuf::Descriptor*> GetMessage(
85+
const google::protobuf::FieldDescriptor* field_descriptor);
86+
87+
static absl::Status AddAllHiddenFields(
88+
absl::flat_hash_map<std::string, std::set<std::string>>&
89+
types_and_visible_fields,
90+
absl::string_view type_name);
91+
92+
static absl::Status AddVisibleField(
93+
absl::flat_hash_map<std::string, std::set<std::string>>&
94+
types_and_visible_fields,
95+
absl::string_view type_name, absl::string_view field_name);
96+
97+
static absl::StatusOr<absl::flat_hash_map<std::string, std::set<std::string>>>
98+
ComputeVisibleFieldsMap(
99+
const std::shared_ptr<const google::protobuf::DescriptorPool>& descriptor_pool,
100+
const absl::flat_hash_map<std::string, std::set<std::string>>&
101+
types_and_field_paths);
102+
103+
explicit ProtoTypeMask(absl::flat_hash_map<std::string, std::set<std::string>>
104+
types_and_visible_fields)
105+
: types_and_visible_fields_(std::move(types_and_visible_fields)) {};
106+
107+
static inline constexpr char kPathDelimiter = '.';
108+
109+
// Map of types that have a field mask where the keys are
110+
// fully qualified type names and the values are the set of field names that
111+
// are visible for the type.
112+
absl::flat_hash_map<std::string, std::set<std::string>>
113+
types_and_visible_fields_;
114+
};
115+
116+
#endif // THIRD_PARTY_CEL_CPP_CHECKER_PROTO_TYPE_MASK_H_

0 commit comments

Comments
 (0)