upb_generator/plugin.h

   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
  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:
  50   Plugin() { ReadRequest(); }
  51   ~Plugin() { WriteResponse(); }
  52
  53   absl::string_view parameter() const {
  54     return ToStringView(
  55         UPB_DESC(compiler_CodeGeneratorRequest_parameter)(request_));
  56   }
  57
  58   template <class T>
  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>
  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
  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
 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
 119   static absl::string_view ToStringView(upb_StringView sv) {
 120     return absl::string_view(sv.data, sv.size);
 121   }
 122
 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
 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
 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
 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_