1 | /* Copyright 2015 The TensorFlow Authors. All Rights Reserved. |
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 | |
16 | #ifndef TENSORFLOW_TSL_PLATFORM_PROTOBUF_H_ |
17 | #define TENSORFLOW_TSL_PLATFORM_PROTOBUF_H_ |
18 | |
19 | #include "tensorflow/tsl/platform/platform.h" |
20 | #include "tensorflow/tsl/platform/types.h" |
21 | |
22 | // Import whatever namespace protobuf comes from into the |
23 | // ::tsl::protobuf namespace. |
24 | // |
25 | // TensorFlow code should use the ::tensorflow::protobuf namespace to |
26 | // refer to all protobuf APIs. |
27 | |
28 | #include "google/protobuf/io/coded_stream.h" // IWYU pragma: export |
29 | #include "google/protobuf/io/tokenizer.h" // IWYU pragma: export |
30 | #include "google/protobuf/io/zero_copy_stream.h" // IWYU pragma: export |
31 | #include "google/protobuf/io/zero_copy_stream_impl_lite.h" // IWYU pragma: export |
32 | #include "google/protobuf/descriptor.pb.h" // IWYU pragma: export |
33 | #include "google/protobuf/arena.h" // IWYU pragma: export |
34 | #include "google/protobuf/descriptor.h" // IWYU pragma: export |
35 | #include "google/protobuf/dynamic_message.h" // IWYU pragma: export |
36 | #include "google/protobuf/map.h" // IWYU pragma: export |
37 | #include "google/protobuf/message.h" // IWYU pragma: export |
38 | #include "google/protobuf/repeated_field.h" // IWYU pragma: export |
39 | #include "google/protobuf/text_format.h" // IWYU pragma: export |
40 | #include "google/protobuf/util/field_comparator.h" // IWYU pragma: export |
41 | #include "google/protobuf/util/json_util.h" // IWYU pragma: export |
42 | #include "google/protobuf/util/message_differencer.h" // IWYU pragma: export |
43 | #include "google/protobuf/util/type_resolver_util.h" // IWYU pragma: export |
44 | |
45 | namespace tsl { |
46 | |
47 | namespace protobuf = ::google::protobuf; |
48 | using protobuf_int64 = ::google::protobuf::int64; |
49 | using protobuf_uint64 = ::google::protobuf::uint64; |
50 | extern const char* kProtobufInt64Typename; |
51 | extern const char* kProtobufUint64Typename; |
52 | |
53 | // Parses a protocol buffer contained in a string in the binary wire format. |
54 | // Returns true on success. Note: Unlike protobuf's builtin ParseFromString, |
55 | // this function has no size restrictions on the total size of the encoded |
56 | // protocol buffer. |
57 | bool ParseProtoUnlimited(protobuf::MessageLite* proto, |
58 | const std::string& serialized); |
59 | bool ParseProtoUnlimited(protobuf::MessageLite* proto, const void* serialized, |
60 | size_t size); |
61 | inline bool ParseProtoUnlimited(protobuf::MessageLite* proto, |
62 | const tstring& serialized) { |
63 | return ParseProtoUnlimited(proto, serialized.data(), serialized.size()); |
64 | } |
65 | |
66 | // Returns the string value for the value of a string or bytes protobuf field. |
67 | inline const std::string& ProtobufStringToString(const std::string& s) { |
68 | return s; |
69 | } |
70 | |
71 | // Set <dest> to <src>. Swapping is allowed, as <src> does not need to be |
72 | // preserved. |
73 | inline void SetProtobufStringSwapAllowed(std::string* src, std::string* dest) { |
74 | *dest = std::move(*src); |
75 | } |
76 | |
77 | #if defined(TENSORFLOW_PROTOBUF_USES_CORD) |
78 | // These versions of ProtobufStringToString and SetProtobufString get used by |
79 | // tools/proto_text's generated code. They have the same name as the versions |
80 | // in tsl/platform/protobuf.h, so the generation code doesn't need to determine |
81 | // if the type is Cord or string at generation time. |
82 | inline std::string ProtobufStringToString(const absl::Cord& s) { |
83 | return std::string(s); |
84 | } |
85 | inline void SetProtobufStringSwapAllowed(std::string* src, absl::Cord* dest) { |
86 | dest->CopyFrom(*src); |
87 | } |
88 | #endif // defined(TENSORFLOW_PROTOBUF_USES_CORD) |
89 | |
90 | inline bool SerializeToTString(const protobuf::MessageLite& proto, |
91 | tstring* output) { |
92 | size_t size = proto.ByteSizeLong(); |
93 | output->resize_uninitialized(size); |
94 | return proto.SerializeWithCachedSizesToArray( |
95 | reinterpret_cast<uint8*>(output->data())); |
96 | } |
97 | |
98 | inline bool ParseFromTString(const tstring& input, |
99 | protobuf::MessageLite* proto) { |
100 | return proto->ParseFromArray(input.data(), static_cast<int>(input.size())); |
101 | } |
102 | |
103 | // Analogue to StringOutputStream for tstring. |
104 | class TStringOutputStream : public protobuf::io::ZeroCopyOutputStream { |
105 | public: |
106 | explicit TStringOutputStream(tstring* target); |
107 | ~TStringOutputStream() override = default; |
108 | |
109 | TStringOutputStream(const TStringOutputStream&) = delete; |
110 | void operator=(const TStringOutputStream&) = delete; |
111 | |
112 | bool Next(void** data, int* size) override; |
113 | void BackUp(int count) override; |
114 | int64_t ByteCount() const override; |
115 | |
116 | private: |
117 | static constexpr int kMinimumSize = 16; |
118 | |
119 | tstring* target_; |
120 | }; |
121 | } // namespace tsl |
122 | |
123 | #endif // TENSORFLOW_TSL_PLATFORM_PROTOBUF_H_ |
124 | |