staging_protobuf_codegen/
lib.rs

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

staging_protobuf_codegen/
lib.rs
