1 use std::fs::{self, OpenOptions}; 2 use std::io::prelude::*; 3 use std::path::{Path, PathBuf}; 4 use walkdir::WalkDir; 5 6 #[derive(Debug)] 7 pub struct CodeGen { 8 inputs: Vec<PathBuf>, 9 output_dir: PathBuf, 10 protoc_path: PathBuf, 11 protoc_gen_upb_path: PathBuf, 12 protoc_gen_upb_minitable_path: PathBuf, 13 includes: Vec<PathBuf>, 14 } 15 16 const VERSION: &str = env!("CARGO_PKG_VERSION"); 17 18 impl CodeGen { new() -> Self19 pub fn new() -> Self { 20 Self { 21 inputs: Vec::new(), 22 output_dir: std::env::current_dir().unwrap().join("src").join("protos"), 23 protoc_path: PathBuf::from("protoc"), 24 protoc_gen_upb_path: PathBuf::from("protoc-gen-upb"), 25 protoc_gen_upb_minitable_path: PathBuf::from("protoc-gen-upb_minitable"), 26 includes: Vec::new(), 27 } 28 } 29 input(&mut self, input: impl AsRef<Path>) -> &mut Self30 pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self { 31 self.inputs.push(input.as_ref().to_owned()); 32 self 33 } 34 inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self35 pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self { 36 self.inputs.extend(inputs.into_iter().map(|input| input.as_ref().to_owned())); 37 self 38 } 39 output_dir(&mut self, output_dir: impl AsRef<Path>) -> &mut Self40 pub fn output_dir(&mut self, output_dir: impl AsRef<Path>) -> &mut Self { 41 self.output_dir = output_dir.as_ref().to_owned(); 42 self 43 } 44 protoc_path(&mut self, protoc_path: impl AsRef<Path>) -> &mut Self45 pub fn protoc_path(&mut self, protoc_path: impl AsRef<Path>) -> &mut Self { 46 self.protoc_path = protoc_path.as_ref().to_owned(); 47 self 48 } 49 protoc_gen_upb_path(&mut self, protoc_gen_upb_path: impl AsRef<Path>) -> &mut Self50 pub fn protoc_gen_upb_path(&mut self, protoc_gen_upb_path: impl AsRef<Path>) -> &mut Self { 51 self.protoc_gen_upb_path = protoc_gen_upb_path.as_ref().to_owned(); 52 self 53 } 54 protoc_gen_upb_minitable_path( &mut self, protoc_gen_upb_minitable_path: impl AsRef<Path>, ) -> &mut Self55 pub fn protoc_gen_upb_minitable_path( 56 &mut self, 57 protoc_gen_upb_minitable_path: impl AsRef<Path>, 58 ) -> &mut Self { 59 self.protoc_gen_upb_minitable_path = protoc_gen_upb_minitable_path.as_ref().to_owned(); 60 self 61 } 62 include(&mut self, include: impl AsRef<Path>) -> &mut Self63 pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self { 64 self.includes.push(include.as_ref().to_owned()); 65 self 66 } 67 includes(&mut self, includes: impl Iterator<Item = impl AsRef<Path>>) -> &mut Self68 pub fn includes(&mut self, includes: impl Iterator<Item = impl AsRef<Path>>) -> &mut Self { 69 self.includes.extend(includes.into_iter().map(|include| include.as_ref().to_owned())); 70 self 71 } 72 73 /// Generate Rust, upb, and upb minitables. Build and link the C code. compile(&self) -> Result<(), String>74 pub fn compile(&self) -> Result<(), String> { 75 let libupb_version = std::env::var("DEP_LIBUPB_VERSION").expect("DEP_LIBUPB_VERSION should have been set, make sure that the Protobuf crate is a dependency"); 76 if VERSION != libupb_version { 77 panic!( 78 "protobuf-codegen version {} does not match protobuf version {}.", 79 VERSION, libupb_version 80 ); 81 } 82 83 let mut cmd = std::process::Command::new(&self.protoc_path); 84 for input in &self.inputs { 85 cmd.arg(input); 86 } 87 if !self.output_dir.exists() { 88 // Attempt to make the directory if it doesn't exist 89 let _ = std::fs::create_dir(&self.output_dir); 90 } 91 cmd.arg(format!("--rust_out={}", self.output_dir.display())) 92 .arg("--rust_opt=experimental-codegen=enabled,kernel=upb") 93 .arg(format!("--plugin=protoc-gen-upb={}", self.protoc_gen_upb_path.display())) 94 .arg(format!( 95 "--plugin=protoc-gen-upb_minitable={}", 96 self.protoc_gen_upb_minitable_path.display() 97 )) 98 .arg(format!("--upb_minitable_out={}", self.output_dir.display())); 99 for include in &self.includes { 100 cmd.arg(format!("--proto_path={}", include.display())); 101 } 102 let output = cmd.output().map_err(|e| format!("failed to run protoc: {}", e))?; 103 println!("{}", std::str::from_utf8(&output.stdout).unwrap()); 104 eprintln!("{}", std::str::from_utf8(&output.stderr).unwrap()); 105 assert!(output.status.success()); 106 107 let mut cc_build = cc::Build::new(); 108 cc_build 109 .include( 110 std::env::var_os("DEP_LIBUPB_INCLUDE") 111 .expect("DEP_LIBUPB_INCLUDE should have been set, make sure that the Protobuf crate is a dependency"), 112 ) 113 .include(self.output_dir.clone()) 114 .flag("-std=c99"); 115 for entry in WalkDir::new(&self.output_dir) { 116 if let Ok(entry) = entry { 117 let path = entry.path(); 118 let file_name = path.file_name().unwrap().to_str().unwrap(); 119 if file_name.ends_with(".upb_minitable.c") { 120 cc_build.file(path); 121 } 122 if file_name.ends_with(".upb.h") { 123 Self::strip_upb_inline(&path); 124 cc_build.file(path.with_extension("c")); 125 } 126 if file_name.ends_with(".pb.rs") { 127 Self::fix_rs_gencode(&path); 128 } 129 } 130 } 131 cc_build.compile(&format!("{}_upb_gen_code", std::env::var("CARGO_PKG_NAME").unwrap())); 132 Ok(()) 133 } 134 135 // Remove UPB_INLINE from the .upb.h file. strip_upb_inline(header: &Path)136 fn strip_upb_inline(header: &Path) { 137 let contents = fs::read_to_string(header).unwrap().replace("UPB_INLINE ", ""); 138 let mut file = 139 OpenOptions::new().write(true).truncate(true).open(header.with_extension("c")).unwrap(); 140 file.write(contents.as_bytes()).unwrap(); 141 } 142 143 // Adjust the generated Rust code to work with the crate structure. fix_rs_gencode(path: &Path)144 fn fix_rs_gencode(path: &Path) { 145 let contents = fs::read_to_string(path) 146 .unwrap() 147 .replace("crate::", "") 148 .replace("protobuf_upb", "protobuf") 149 .replace("::__pb", "__pb") 150 .replace("::__std", "__std"); 151 let mut file = OpenOptions::new().write(true).truncate(true).open(path).unwrap(); 152 file.write(contents.as_bytes()).unwrap(); 153 } 154 } 155