protoc_gen_prost_crate/generator/
include_file.rs

1use std::rc::Rc;
2
3use once_cell::sync::Lazy;
4use prost_build::Module;
5use prost_types::compiler::code_generator_response::File;
6use protoc_gen_prost::{Generator, ModuleRequestSet, Result};
7
8use crate::PackageLimiter;
9
10pub(crate) struct IncludeFileGenerator<'a> {
11    filename: &'a str,
12    limiter: Rc<PackageLimiter>,
13}
14
15impl<'a> IncludeFileGenerator<'a> {
16    pub(crate) fn new(filename: &'a str, limiter: Rc<PackageLimiter>) -> Self {
17        Self { filename, limiter }
18    }
19}
20
21impl<'a> Generator for IncludeFileGenerator<'a> {
22    fn generate(&mut self, module_request_set: &ModuleRequestSet) -> Result {
23        let mut context = CodeGenContext::new();
24
25        let _: () = module_request_set
26            .requests()
27            .filter_map(|(module, request)| {
28                let filename = request.output_filename()?;
29                self.limiter
30                    .is_allowed(request.proto_package_name())
31                    .then(|| {
32                        context.move_to(module, request.proto_package_name());
33                        context.push_include(filename);
34                        context.push_insertion_point(request.proto_package_name());
35                    })
36            })
37            .collect();
38
39        let content = context.finish();
40
41        let file = File {
42            name: Some(self.filename.to_string()),
43            content: Some(content),
44            ..File::default()
45        };
46
47        Ok(vec![file])
48    }
49}
50
51static ROOT_MODULE: Lazy<Module> = Lazy::new(|| Module::from_parts([] as [String; 0]));
52
53#[must_use]
54#[derive(Debug)]
55struct CodeGenContext<'a> {
56    last: &'a Module,
57    indent: String,
58    buf: String,
59}
60
61const INDENT: &str = "    ";
62
63impl<'a> CodeGenContext<'a> {
64    fn new() -> Self {
65        // Reserve an initial 16 kiB
66        let mut buf = String::with_capacity(16_384);
67        buf.push_str("// @generated\n");
68
69        let indent = String::with_capacity(INDENT.len() * 8);
70
71        Self {
72            last: &*ROOT_MODULE,
73            indent,
74            buf,
75        }
76    }
77
78    fn indent(&mut self) {
79        self.indent.push_str(INDENT);
80    }
81
82    fn dedent(&mut self) {
83        self.indent.truncate(
84            self.indent
85                .len()
86                .checked_sub(INDENT.len())
87                .expect("indent underflow"),
88        );
89    }
90
91    fn push_indent(&mut self) {
92        self.buf.push_str(&self.indent);
93    }
94
95    fn move_to(&mut self, next: &'a Module, package: &str) {
96        let (down, prefix) = difference(self.last, next);
97
98        for _ in 0..down {
99            self.close_module();
100        }
101
102        let take = next.len() - prefix - 1;
103
104        for module_name in next.parts().skip(prefix).take(take) {
105            self.open_module(module_name);
106        }
107
108        self.push_attribute_insertion_point(package);
109        self.open_module(next.parts().last().unwrap());
110
111        self.last = next;
112    }
113
114    fn finish(mut self) -> String {
115        while !self.indent.is_empty() {
116            self.close_module()
117        }
118        self.buf
119    }
120
121    fn push_include(&mut self, filename: &str) {
122        self.push_indent();
123        self.buf.push_str("include!(\"");
124        self.buf.push_str(filename);
125        self.buf.push_str("\");\n");
126    }
127
128    fn push_insertion_point(&mut self, package_name: &str) {
129        self.push_indent();
130        self.buf.push_str("// @@protoc_insertion_point(");
131        self.buf.push_str(package_name);
132        self.buf.push_str(")\n");
133    }
134
135    fn push_attribute_insertion_point(&mut self, package_name: &str) {
136        self.push_indent();
137        self.buf.push_str("// @@protoc_insertion_point(attribute:");
138        self.buf.push_str(package_name);
139        self.buf.push_str(")\n");
140    }
141
142    fn close_module(&mut self) {
143        self.dedent();
144        self.push_indent();
145        self.buf.push_str("}\n");
146    }
147
148    fn open_module(&mut self, module_name: &str) {
149        self.push_indent();
150        self.buf.push_str("pub mod ");
151        self.buf.push_str(module_name);
152        self.buf.push_str(" {\n");
153        self.indent();
154    }
155}
156
157fn difference(left: &Module, right: &Module) -> (usize, usize) {
158    let mut left_parts = left.parts();
159    let mut right_parts = right.parts();
160
161    let mut prefix = 0;
162
163    loop {
164        match (left_parts.next(), right_parts.next()) {
165            (Some(left), Some(right)) if left == right => prefix += 1,
166            (Some(_), Some(_)) => return (left_parts.count() + 1, prefix),
167            (Some(_), None) => return (left_parts.count() + 1, prefix),
168            (None, Some(_)) => return (0, prefix),
169            (None, None) => return (0, prefix),
170        }
171    }
172}
protoc_gen_prost_crate/generator/include_file.rs

protoc_gen_prost_crate/generator/
include_file.rs
