protobuf_codegen/
lib.rs

1use std::fs::File;
2use std::io::Write;
3use std::path::{Path, PathBuf};
4
5#[derive(Debug, Clone)]
6pub struct Dependency {
7    pub crate_name: String,
8    pub proto_import_paths: Vec<PathBuf>,
9    pub c_include_paths: Vec<PathBuf>,
10    pub proto_files: Vec<String>,
11}
12
13#[derive(Debug)]
14pub struct CodeGen {
15    inputs: Vec<PathBuf>,
16    output_dir: PathBuf,
17    includes: Vec<PathBuf>,
18    dependencies: Vec<Dependency>,
19}
20
21const VERSION: &str = env!("CARGO_PKG_VERSION");
22
23fn missing_protoc_error_message() -> String {
24    format!(
25        "
26Please make sure you have protoc available in your PATH. You can build it \
27from source as follows: \
28git clone https://github.com/protocolbuffers/protobuf.git; \
29cd protobuf; \
30git checkout rust-prerelease-{}; \
31cmake . -Dprotobuf_FORCE_FETCH_DEPENDENCIES=ON; \
32cmake --build . --parallel 12",
33        VERSION
34    )
35}
36
37// Given the output of "protoc --version", returns a shortened version string
38// suitable for comparing against the protobuf crate version.
39//
40// The output of protoc --version looks something like "libprotoc XX.Y",
41// optionally followed by "-dev" or "-rcN". We want to strip the "-dev" suffix
42// if present and return something like "30.0" or "30.0-rc1".
43fn protoc_version(protoc_output: &str) -> String {
44    let mut s = protoc_output.strip_prefix("libprotoc ").unwrap().trim().to_string();
45    let first_dash = s.find("-dev");
46    if let Some(i) = first_dash {
47        s.truncate(i);
48    }
49    s
50}
51
52// Given a crate version string, returns just the part of it suitable for
53// comparing against the protoc version. The crate version is of the form
54// "X.Y.Z" with an optional suffix starting with a dash. We want to drop the
55// major version ("X.") and only keep the suffix if it starts with "-rc".
56fn expected_protoc_version(cargo_version: &str) -> String {
57    let mut s = cargo_version.replace("-rc.", "-rc");
58    let is_release_candidate = s.find("-rc") != None;
59    if !is_release_candidate {
60        if let Some(i) = s.find('-') {
61            s.truncate(i);
62        }
63    }
64    let mut v: Vec<&str> = s.split('.').collect();
65    assert_eq!(v.len(), 3);
66    v.remove(0);
67    v.join(".")
68}
69
70impl CodeGen {
71    pub fn new() -> Self {
72        Self {
73            inputs: Vec::new(),
74            output_dir: PathBuf::from(std::env::var("OUT_DIR").unwrap()).join("protobuf_generated"),
75            includes: Vec::new(),
76            dependencies: Vec::new(),
77        }
78    }
79
80    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
81        self.inputs.push(input.as_ref().to_owned());
82        self
83    }
84
85    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
86        self.inputs.extend(inputs.into_iter().map(|input| input.as_ref().to_owned()));
87        self
88    }
89
90    pub fn output_dir(&mut self, output_dir: impl AsRef<Path>) -> &mut Self {
91        self.output_dir = output_dir.as_ref().to_owned();
92        self
93    }
94
95    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
96        self.includes.push(include.as_ref().to_owned());
97        self
98    }
99
100    pub fn includes(&mut self, includes: impl Iterator<Item = impl AsRef<Path>>) -> &mut Self {
101        self.includes.extend(includes.into_iter().map(|include| include.as_ref().to_owned()));
102        self
103    }
104
105    pub fn dependency(&mut self, deps: Vec<Dependency>) -> &mut Self {
106        self.dependencies.extend(deps);
107        self
108    }
109
110    fn expected_generated_rs_files(&self) -> Vec<PathBuf> {
111        self.inputs
112            .iter()
113            .map(|input| {
114                let mut input = input.clone();
115                assert!(input.set_extension("u.pb.rs"));
116                self.output_dir.join(input)
117            })
118            .collect()
119    }
120
121    fn expected_generated_c_files(&self) -> Vec<PathBuf> {
122        self.inputs
123            .iter()
124            .map(|input| {
125                let mut input = input.clone();
126                assert!(input.set_extension("upb_minitable.c"));
127                self.output_dir.join(input)
128            })
129            .collect()
130    }
131
132    fn generate_crate_mapping_file(&self) -> PathBuf {
133        let crate_mapping_path = self.output_dir.join("crate_mapping.txt");
134        let mut file = File::create(crate_mapping_path.clone()).unwrap();
135        for dep in &self.dependencies {
136            file.write_all(format!("{}\n", dep.crate_name).as_bytes()).unwrap();
137            file.write_all(format!("{}\n", dep.proto_files.len()).as_bytes()).unwrap();
138            for f in &dep.proto_files {
139                file.write_all(format!("{}\n", f).as_bytes()).unwrap();
140            }
141        }
142        crate_mapping_path
143    }
144
145    pub fn generate_and_compile(&self) -> Result<(), String> {
146        let upb_version = std::env::var("DEP_UPB_VERSION").expect("DEP_UPB_VERSION should have been set, make sure that the Protobuf crate is a dependency");
147        if VERSION != upb_version {
148            panic!(
149                "protobuf-codegen version {} does not match protobuf version {}.",
150                VERSION, upb_version
151            );
152        }
153
154        let mut version_cmd = std::process::Command::new("protoc");
155        let output = version_cmd.arg("--version").output().map_err(|e| {
156            format!("failed to run protoc --version: {} {}", e, missing_protoc_error_message())
157        })?;
158
159        let protoc_version = protoc_version(&String::from_utf8(output.stdout).unwrap());
160        let expected_protoc_version = expected_protoc_version(VERSION);
161        if protoc_version != expected_protoc_version {
162            panic!(
163                "Expected protoc version {} but found {}",
164                expected_protoc_version, protoc_version
165            );
166        }
167
168        let mut cmd = std::process::Command::new("protoc");
169        for input in &self.inputs {
170            cmd.arg(input);
171        }
172        if !self.output_dir.exists() {
173            // Attempt to make the directory if it doesn't exist
174            let _ = std::fs::create_dir(&self.output_dir);
175        }
176
177        for include in &self.includes {
178            println!("cargo:rerun-if-changed={}", include.display());
179        }
180        for dep in &self.dependencies {
181            for path in &dep.proto_import_paths {
182                println!("cargo:rerun-if-changed={}", path.display());
183            }
184        }
185
186        let crate_mapping_path = self.generate_crate_mapping_file();
187
188        cmd.arg(format!("--rust_out={}", self.output_dir.display()))
189            .arg("--rust_opt=experimental-codegen=enabled,kernel=upb")
190            .arg(format!("--upb_minitable_out={}", self.output_dir.display()));
191        for include in &self.includes {
192            cmd.arg(format!("--proto_path={}", include.display()));
193        }
194        for dep in &self.dependencies {
195            for path in &dep.proto_import_paths {
196                cmd.arg(format!("--proto_path={}", path.display()));
197            }
198        }
199        cmd.arg(format!("--rust_opt=crate_mapping={}", crate_mapping_path.display()));
200        let output = cmd.output().map_err(|e| format!("failed to run protoc: {}", e))?;
201        println!("{}", std::str::from_utf8(&output.stdout).unwrap());
202        eprintln!("{}", std::str::from_utf8(&output.stderr).unwrap());
203        assert!(output.status.success());
204        self.compile_only()
205    }
206
207    /// Builds and links the C code.
208    pub fn compile_only(&self) -> Result<(), String> {
209        let mut cc_build = cc::Build::new();
210        cc_build
211            .include(
212                std::env::var_os("DEP_UPB_INCLUDE")
213                    .expect("DEP_UPB_INCLUDE should have been set, make sure that the Protobuf crate is a dependency"),
214            )
215            .include(self.output_dir.clone())
216            .flag("-std=c99");
217
218        for dep in &self.dependencies {
219            for path in &dep.c_include_paths {
220                cc_build.include(path);
221            }
222        }
223
224        for path in &self.expected_generated_rs_files() {
225            if !path.exists() {
226                return Err(format!("expected generated file {} does not exist", path.display()));
227            }
228            println!("cargo:rerun-if-changed={}", path.display());
229        }
230        for path in &self.expected_generated_c_files() {
231            if !path.exists() {
232                return Err(format!("expected generated file {} does not exist", path.display()));
233            }
234            println!("cargo:rerun-if-changed={}", path.display());
235            cc_build.file(path);
236        }
237        cc_build.compile(&format!("{}_upb_gen_code", std::env::var("CARGO_PKG_NAME").unwrap()));
238        Ok(())
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245    use googletest::prelude::*;
246
247    #[gtest]
248    fn test_protoc_version() {
249        assert_that!(protoc_version("libprotoc 30.0"), eq("30.0"));
250        assert_that!(protoc_version("libprotoc 30.0\n"), eq("30.0"));
251        assert_that!(protoc_version("libprotoc 30.0-dev"), eq("30.0"));
252        assert_that!(protoc_version("libprotoc 30.0-rc1"), eq("30.0-rc1"));
253    }
254
255    #[googletest::test]
256    fn test_expected_protoc_version() {
257        assert_that!(expected_protoc_version("4.30.0"), eq("30.0"));
258        assert_that!(expected_protoc_version("4.30.0-alpha"), eq("30.0"));
259        assert_that!(expected_protoc_version("4.30.0-beta"), eq("30.0"));
260        assert_that!(expected_protoc_version("4.30.0-pre"), eq("30.0"));
261        assert_that!(expected_protoc_version("4.30.0-rc.1"), eq("30.0-rc1"));
262    }
263}
protobuf_codegen/lib.rs

protobuf_codegen/
lib.rs
