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 GOOGLE_PROTOBUF_COMPILER_RUST_CONTEXT_H__
9 #define GOOGLE_PROTOBUF_COMPILER_RUST_CONTEXT_H__
10
11 #include <algorithm>
12 #include <string>
13 #include <vector>
14
15 #include "absl/container/flat_hash_map.h"
16 #include "absl/log/absl_log.h"
17 #include "absl/status/statusor.h"
18 #include "absl/strings/string_view.h"
19 #include "absl/types/span.h"
20 #include "google/protobuf/descriptor.h"
21 #include "google/protobuf/io/printer.h"
22
23 namespace google {
24 namespace protobuf {
25 namespace compiler {
26 namespace rust {
27 // Marks which kernel the Rust codegen should generate code for.
28 enum class Kernel {
29 kUpb,
30 kCpp,
31 };
32
KernelRsName(Kernel kernel)33 inline absl::string_view KernelRsName(Kernel kernel) {
34 switch (kernel) {
35 case Kernel::kUpb:
36 return "upb";
37 case Kernel::kCpp:
38 return "cpp";
39 default:
40 ABSL_LOG(FATAL) << "Unknown kernel type: " << static_cast<int>(kernel);
41 return "";
42 }
43 }
44
45 // Global options for a codegen invocation.
46 struct Options {
47 Kernel kernel;
48 std::string mapping_file_path;
49 bool strip_nonfunctional_codegen = false;
50
51 static absl::StatusOr<Options> Parse(absl::string_view param);
52 };
53
54 class RustGeneratorContext {
55 public:
RustGeneratorContext(const std::vector<const FileDescriptor * > * files_in_current_crate,const absl::flat_hash_map<std::string,std::string> * import_path_to_crate_name)56 explicit RustGeneratorContext(
57 const std::vector<const FileDescriptor*>* files_in_current_crate,
58 const absl::flat_hash_map<std::string, std::string>*
59 import_path_to_crate_name)
60 : files_in_current_crate_(*files_in_current_crate),
61 import_path_to_crate_name_(*import_path_to_crate_name) {}
62
primary_file()63 const FileDescriptor& primary_file() const {
64 return *files_in_current_crate_.front();
65 }
66
is_file_in_current_crate(const FileDescriptor & f)67 bool is_file_in_current_crate(const FileDescriptor& f) const {
68 return std::find(files_in_current_crate_.begin(),
69 files_in_current_crate_.end(),
70 &f) != files_in_current_crate_.end();
71 }
72
73 private:
74 const std::vector<const FileDescriptor*>& files_in_current_crate_;
75 const absl::flat_hash_map<std::string, std::string>&
76 import_path_to_crate_name_;
77
78 friend class Context;
79 };
80
81 // A context for generating a particular kind of definition.
82 class Context {
83 public:
Context(const Options * opts,const RustGeneratorContext * rust_generator_context,io::Printer * printer)84 Context(const Options* opts,
85 const RustGeneratorContext* rust_generator_context,
86 io::Printer* printer)
87 : opts_(opts),
88 rust_generator_context_(rust_generator_context),
89 printer_(printer) {}
90
91 Context(const Context&) = delete;
92 Context& operator=(const Context&) = delete;
93 Context(Context&&) = default;
94 Context& operator=(Context&&) = default;
95
opts()96 const Options& opts() const { return *opts_; }
generator_context()97 const RustGeneratorContext& generator_context() const {
98 return *rust_generator_context_;
99 }
100
is_cpp()101 bool is_cpp() const { return opts_->kernel == Kernel::kCpp; }
is_upb()102 bool is_upb() const { return opts_->kernel == Kernel::kUpb; }
103
104 // NOTE: prefer ctx.Emit() over ctx.printer().Emit();
printer()105 io::Printer& printer() const { return *printer_; }
106
WithPrinter(io::Printer * printer)107 Context WithPrinter(io::Printer* printer) const {
108 return Context(opts_, rust_generator_context_, printer);
109 }
110
111 // Forwards to Emit(), which will likely be called all the time.
112 void Emit(absl::string_view format,
113 io::Printer::SourceLocation loc =
114 io::Printer::SourceLocation::current()) const {
115 printer_->Emit(format, loc);
116 }
117 void Emit(absl::Span<const io::Printer::Sub> vars, absl::string_view format,
118 io::Printer::SourceLocation loc =
119 io::Printer::SourceLocation::current()) const {
120 printer_->Emit(vars, format, loc);
121 }
122
ImportPathToCrateName(absl::string_view import_path)123 absl::string_view ImportPathToCrateName(absl::string_view import_path) const {
124 if (opts_->strip_nonfunctional_codegen) {
125 return "test";
126 }
127 auto it =
128 rust_generator_context_->import_path_to_crate_name_.find(import_path);
129 if (it == rust_generator_context_->import_path_to_crate_name_.end()) {
130 ABSL_LOG(FATAL)
131 << "Path " << import_path
132 << " not found in crate mapping. Crate mapping has "
133 << rust_generator_context_->import_path_to_crate_name_.size()
134 << " entries";
135 }
136 return it->second;
137 }
138
139 private:
140 const Options* opts_;
141 const RustGeneratorContext* rust_generator_context_;
142 io::Printer* printer_;
143 };
144
145 bool IsInCurrentlyGeneratingCrate(Context& ctx, const FileDescriptor& file);
146 bool IsInCurrentlyGeneratingCrate(Context& ctx, const Descriptor& message);
147 bool IsInCurrentlyGeneratingCrate(Context& ctx, const EnumDescriptor& enum_);
148
149 } // namespace rust
150 } // namespace compiler
151 } // namespace protobuf
152 } // namespace google
153
154 #endif // GOOGLE_PROTOBUF_COMPILER_RUST_CONTEXT_H__
155