protoc_gen_prost_crate/generator/
features.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    rc::Rc,
4};
5
6use prost_types::compiler::code_generator_response::File;
7use protoc_gen_prost::{Generator, ModuleRequest, ModuleRequestSet, Result};
8
9use crate::PackageLimiter;
10
11pub(crate) struct FeaturesGenerator<'a> {
12    include_filename: &'a str,
13    package_separator: &'a str,
14    limiter: Rc<PackageLimiter>,
15}
16
17impl<'a> FeaturesGenerator<'a> {
18    pub(crate) fn new(
19        include_filename: &'a str,
20        package_separator: &'a str,
21        limiter: Rc<PackageLimiter>,
22    ) -> Self {
23        Self {
24            include_filename,
25            package_separator,
26            limiter,
27        }
28    }
29}
30
31impl<'a> Generator for FeaturesGenerator<'a> {
32    fn generate(&mut self, module_request_set: &ModuleRequestSet) -> Result {
33        let mut files = Vec::new();
34
35        let dep_tree = module_request_set
36            .requests()
37            .map(|(_, v)| v)
38            .collect::<PackageDependencies>();
39
40        let mut buf = String::new();
41
42        if !dep_tree.dependency_graph.is_empty() {
43            buf.push_str(
44                "# @@protoc_deletion_point(features)\n# This section is automatically generated \
45                 by protoc-gen-prost-crate.\n# Changes in this area may be lost on regeneration.\n",
46            );
47            buf.push_str("proto_full = [");
48            for feature in dep_tree.dependency_graph.keys() {
49                if !self.limiter.is_allowed(feature) {
50                    continue;
51                }
52
53                let feature_name = feature.replace('.', self.package_separator);
54                buf.push('"');
55                buf.push_str(&feature_name);
56                buf.push_str("\",");
57
58                files.push(File {
59                    name: Some(self.include_filename.to_owned()),
60                    content: Some(format!("#[cfg(feature = \"{feature_name}\")]\n")),
61                    insertion_point: Some(format!("attribute:{feature}")),
62                    ..File::default()
63                });
64            }
65            if let Some('[') = buf.pop() {
66                buf.push('[');
67            }
68            buf.push_str("]\n");
69
70            for (feature, dependencies) in dep_tree.dependency_graph {
71                if !self.limiter.is_allowed(feature) {
72                    continue;
73                }
74
75                buf.push('"');
76                buf.push_str(&feature.replace('.', self.package_separator));
77                buf.push('"');
78                if dependencies.is_empty() {
79                    buf.push_str(" = []\n");
80                } else {
81                    buf.push_str(" = [");
82                    for feature in dependencies {
83                        if !self.limiter.is_allowed(feature) {
84                            continue;
85                        }
86
87                        buf.push('"');
88                        buf.push_str(&feature.replace('.', self.package_separator));
89                        buf.push_str("\",");
90                    }
91                    if let Some('[') = buf.pop() {
92                        buf.push('[');
93                    }
94                    buf.push_str("]\n");
95                }
96            }
97        }
98
99        files.push(File {
100            name: Some("Cargo.toml".to_string()),
101            content: Some(buf),
102            insertion_point: Some("features".to_string()),
103            ..File::default()
104        });
105
106        Ok(files)
107    }
108}
109
110struct PackageDependencies<'a> {
111    pub dependency_graph: BTreeMap<&'a str, BTreeSet<&'a str>>,
112}
113
114impl<'a> FromIterator<&'a ModuleRequest> for PackageDependencies<'a> {
115    fn from_iter<T: IntoIterator<Item = &'a ModuleRequest>>(iter: T) -> Self {
116        let requests = iter.into_iter().collect::<Vec<&ModuleRequest>>();
117        let mut features = requests
118            .iter()
119            .filter(|r| r.output_filename().is_some() && !r.proto_package_name().is_empty())
120            .map(|r| (r.proto_package_name(), BTreeSet::new()))
121            .collect::<BTreeMap<&str, BTreeSet<&str>>>();
122
123        let mut depend_on_type = |current_package: &'a str, depends_on_type: &str| {
124            // Only deal with fully-qualified paths
125            if depends_on_type.starts_with('.') {
126                // Search in reverse, relying on the fact that, for a given prefix, the more
127                // specific package names will be after the higher-level
128                // And don't make the package depend on itself
129                if let Some(&key) = features
130                    .keys()
131                    .rev()
132                    .filter(|&&key| key != current_package)
133                    .find(|&&key| depends_on_type[1..].starts_with(&format!("{key}.")))
134                {
135                    features.entry(current_package).or_default().insert(key);
136                }
137            }
138        };
139
140        for file in requests.into_iter().flat_map(|r| r.files()) {
141            for message in &file.message_type {
142                for field in &message.field {
143                    depend_on_type(file.package(), field.type_name());
144                }
145            }
146
147            for service in &file.service {
148                for method in &service.method {
149                    depend_on_type(file.package(), method.input_type());
150                    depend_on_type(file.package(), method.output_type());
151                }
152            }
153        }
154
155        Self {
156            dependency_graph: features,
157        }
158    }
159}
protoc_gen_prost_crate/generator/features.rs

protoc_gen_prost_crate/generator/
features.rs
