1 // Protocol Buffers - Google's data interchange format
2 // Copyright 2023 Google LLC. All rights reserved.
3 //
4 // Use of this source code is governed by a BSD-style
5 // license that can be found in the LICENSE file or at
6 // https://developers.google.com/open-source/licenses/bsd
7
8 #ifndef UPB_UPB_GENERATOR_PLUGIN_H_
9 #define UPB_UPB_GENERATOR_PLUGIN_H_
10
11 #include <stdio.h>
12
13 #include <cstring>
14 #include <string>
15 #include <utility>
16 #include <vector>
17
18 #ifdef _WIN32
19 #include <fcntl.h>
20 #include <io.h>
21 #endif
22
23 #include "absl/container/flat_hash_set.h"
24 #include "absl/log/absl_log.h"
25 #include "absl/strings/string_view.h"
26 #include "google/protobuf/compiler/code_generator_lite.h"
27 #include "upb/base/status.hpp"
28 #include "upb/base/string_view.h"
29 #include "upb/mem/arena.h"
30 #include "upb/mem/arena.hpp"
31 #include "upb/reflection/def.hpp"
32 #include "upb/reflection/descriptor_bootstrap.h"
33 #include "upb_generator/plugin_bootstrap.h"
34
35 // Must be last.
36 #include "upb/port/def.inc"
37
38 namespace upb {
39 namespace generator {
40
ParseGeneratorParameter(const absl::string_view text)41 inline std::vector<std::pair<std::string, std::string>> ParseGeneratorParameter(
42 const absl::string_view text) {
43 std::vector<std::pair<std::string, std::string>> ret;
44 google::protobuf::compiler::ParseGeneratorParameter(text, &ret);
45 return ret;
46 }
47
48 class Plugin {
49 public:
Plugin()50 Plugin() { ReadRequest(); }
~Plugin()51 ~Plugin() { WriteResponse(); }
52
parameter()53 absl::string_view parameter() const {
54 return ToStringView(
55 UPB_DESC(compiler_CodeGeneratorRequest_parameter)(request_));
56 }
57
58 template <class T>
GenerateFilesRaw(T && func)59 void GenerateFilesRaw(T&& func) {
60 absl::flat_hash_set<absl::string_view> files_to_generate;
61 size_t size;
62 const upb_StringView* file_to_generate = UPB_DESC(
63 compiler_CodeGeneratorRequest_file_to_generate)(request_, &size);
64 for (size_t i = 0; i < size; i++) {
65 files_to_generate.insert(
66 {file_to_generate[i].data, file_to_generate[i].size});
67 }
68
69 const UPB_DESC(FileDescriptorProto)* const* files =
70 UPB_DESC(compiler_CodeGeneratorRequest_proto_file)(request_, &size);
71 for (size_t i = 0; i < size; i++) {
72 upb::Status status;
73 absl::string_view name =
74 ToStringView(UPB_DESC(FileDescriptorProto_name)(files[i]));
75 func(files[i], files_to_generate.contains(name));
76 }
77 }
78
79 template <class T>
GenerateFiles(T && func)80 void GenerateFiles(T&& func) {
81 GenerateFilesRaw(
82 [this, &func](const UPB_DESC(FileDescriptorProto) * file_proto,
83 bool generate) {
84 upb::Status status;
85 upb::FileDefPtr file = pool_.AddFile(file_proto, &status);
86 if (!file) {
87 absl::string_view name =
88 ToStringView(UPB_DESC(FileDescriptorProto_name)(file_proto));
89 ABSL_LOG(FATAL) << "Couldn't add file " << name
90 << " to DefPool: " << status.error_message();
91 }
92 if (generate) func(file);
93 });
94 }
95
SetError(absl::string_view error)96 void SetError(absl::string_view error) {
97 char* data =
98 static_cast<char*>(upb_Arena_Malloc(arena_.ptr(), error.size()));
99 memcpy(data, error.data(), error.size());
100 UPB_DESC(compiler_CodeGeneratorResponse_set_error)
101 (response_, upb_StringView_FromDataAndSize(data, error.size()));
102 }
103
AddOutputFile(absl::string_view filename,absl::string_view content)104 void AddOutputFile(absl::string_view filename, absl::string_view content) {
105 UPB_DESC(compiler_CodeGeneratorResponse_File)* file = UPB_DESC(
106 compiler_CodeGeneratorResponse_add_file)(response_, arena_.ptr());
107 UPB_DESC(compiler_CodeGeneratorResponse_File_set_name)
108 (file, StringDup(filename));
109 UPB_DESC(compiler_CodeGeneratorResponse_File_set_content)
110 (file, StringDup(content));
111 }
112
113 private:
114 upb::Arena arena_;
115 upb::DefPool pool_;
116 UPB_DESC(compiler_CodeGeneratorRequest) * request_;
117 UPB_DESC(compiler_CodeGeneratorResponse) * response_;
118
ToStringView(upb_StringView sv)119 static absl::string_view ToStringView(upb_StringView sv) {
120 return absl::string_view(sv.data, sv.size);
121 }
122
StringDup(absl::string_view s)123 upb_StringView StringDup(absl::string_view s) {
124 char* data =
125 reinterpret_cast<char*>(upb_Arena_Malloc(arena_.ptr(), s.size()));
126 memcpy(data, s.data(), s.size());
127 return upb_StringView_FromDataAndSize(data, s.size());
128 }
129
ReadAllStdinBinary()130 std::string ReadAllStdinBinary() {
131 std::string data;
132 #ifdef _WIN32
133 _setmode(_fileno(stdin), _O_BINARY);
134 _setmode(_fileno(stdout), _O_BINARY);
135 #endif
136 char buf[4096];
137 while (size_t len = fread(buf, 1, sizeof(buf), stdin)) {
138 data.append(buf, len);
139 }
140 return data;
141 }
142
ReadRequest()143 void ReadRequest() {
144 std::string data = ReadAllStdinBinary();
145 request_ = UPB_DESC(compiler_CodeGeneratorRequest_parse)(
146 data.data(), data.size(), arena_.ptr());
147 if (!request_) {
148 ABSL_LOG(FATAL) << "Failed to parse CodeGeneratorRequest";
149 }
150 response_ = UPB_DESC(compiler_CodeGeneratorResponse_new)(arena_.ptr());
151
152 int features =
153 UPB_DESC(compiler_CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL) |
154 UPB_DESC(compiler_CodeGeneratorResponse_FEATURE_SUPPORTS_EDITIONS);
155 UPB_DESC(compiler_CodeGeneratorResponse_set_supported_features)
156 (response_, features);
157 UPB_DESC(compiler_CodeGeneratorResponse_set_minimum_edition)
158 (response_, UPB_DESC(EDITION_PROTO2));
159 UPB_DESC(compiler_CodeGeneratorResponse_set_maximum_edition)
160 (response_, UPB_DESC(EDITION_2023));
161 }
162
WriteResponse()163 void WriteResponse() {
164 size_t size;
165 char* serialized = UPB_DESC(compiler_CodeGeneratorResponse_serialize)(
166 response_, arena_.ptr(), &size);
167 if (!serialized) {
168 ABSL_LOG(FATAL) << "Failed to serialize CodeGeneratorResponse";
169 }
170
171 if (fwrite(serialized, 1, size, stdout) != size) {
172 ABSL_LOG(FATAL) << "Failed to write response to stdout";
173 }
174 }
175 };
176
177 } // namespace generator
178 } // namespace upb
179
180 #include "upb/port/undef.inc"
181
182 #endif // UPB_UPB_GENERATOR_PLUGIN_H_
183