use std::path::{Path, PathBuf};
use std::process::Command;
use buffa::Message;
use buffa_codegen::generated::descriptor::FileDescriptorSet;
#[doc(inline)]
pub use buffa_codegen::CodeGenConfig;
#[doc(inline)]
pub use buffa_codegen::ReflectMode;
#[doc(inline)]
pub use buffa_codegen::StringRepr;
#[derive(Debug, Clone, Default)]
enum DescriptorSource {
#[default]
Protoc,
Buf,
Precompiled(PathBuf),
}
pub struct Config {
files: Vec<PathBuf>,
includes: Vec<PathBuf>,
out_dir: Option<PathBuf>,
codegen_config: CodeGenConfig,
descriptor_source: DescriptorSource,
include_file: Option<String>,
}
impl Config {
pub fn new() -> Self {
Self {
files: Vec::new(),
includes: Vec::new(),
out_dir: None,
codegen_config: CodeGenConfig::default(),
descriptor_source: DescriptorSource::default(),
include_file: None,
}
}
#[must_use]
pub fn files(mut self, files: &[impl AsRef<Path>]) -> Self {
self.files
.extend(files.iter().map(|f| f.as_ref().to_path_buf()));
self
}
#[must_use]
pub fn includes(mut self, includes: &[impl AsRef<Path>]) -> Self {
self.includes
.extend(includes.iter().map(|i| i.as_ref().to_path_buf()));
self
}
#[must_use]
pub fn out_dir(mut self, dir: impl Into<PathBuf>) -> Self {
self.out_dir = Some(dir.into());
self
}
#[must_use]
pub fn generate_views(mut self, enabled: bool) -> Self {
self.codegen_config.generate_views = enabled;
self
}
#[must_use]
pub fn generate_json(mut self, enabled: bool) -> Self {
self.codegen_config.generate_json = enabled;
self
}
#[must_use]
pub fn generate_text(mut self, enabled: bool) -> Self {
self.codegen_config.generate_text = enabled;
self
}
#[must_use]
pub fn generate_arbitrary(mut self, enabled: bool) -> Self {
self.codegen_config.generate_arbitrary = enabled;
self
}
#[must_use]
pub fn gate_impls_on_crate_features(mut self, enabled: bool) -> Self {
self.codegen_config.gate_impls_on_crate_features = enabled;
self
}
#[doc(hidden)]
#[must_use]
pub fn gate_reflect_on_crate_feature(mut self, enabled: bool) -> Self {
self.codegen_config.gate_reflect_on_crate_feature = enabled;
self
}
#[must_use]
pub fn generate_with_setters(mut self, enabled: bool) -> Self {
self.codegen_config.generate_with_setters = enabled;
self
}
#[must_use]
pub fn generate_reflection(mut self, enabled: bool) -> Self {
let mode = if enabled {
ReflectMode::VTable
} else {
ReflectMode::Off
};
mode.apply(&mut self.codegen_config);
self
}
#[must_use]
pub fn reflect_mode(mut self, mode: ReflectMode) -> Self {
mode.apply(&mut self.codegen_config);
self
}
#[must_use]
pub fn idiomatic_enum_aliases(mut self, enabled: bool) -> Self {
self.codegen_config.idiomatic_enum_aliases = enabled;
self
}
#[must_use]
pub fn preserve_unknown_fields(mut self, enabled: bool) -> Self {
self.codegen_config.preserve_unknown_fields = enabled;
self
}
#[must_use]
pub fn strict_utf8_mapping(mut self, enabled: bool) -> Self {
self.codegen_config.strict_utf8_mapping = enabled;
self
}
#[must_use]
pub fn allow_message_set(mut self, enabled: bool) -> Self {
self.codegen_config.allow_message_set = enabled;
self
}
#[must_use]
pub fn extern_path(
mut self,
proto_path: impl Into<String>,
rust_path: impl Into<String>,
) -> Self {
let mut proto_path = proto_path.into();
if !proto_path.starts_with('.') {
proto_path.insert(0, '.');
}
self.codegen_config
.extern_paths
.push((proto_path, rust_path.into()));
self
}
#[must_use]
pub fn use_bytes_type_in(mut self, paths: &[impl AsRef<str>]) -> Self {
self.codegen_config
.bytes_fields
.extend(paths.iter().map(|p| p.as_ref().to_string()));
self
}
#[must_use]
pub fn use_bytes_type(mut self) -> Self {
self.codegen_config.bytes_fields.push(".".to_string());
self
}
#[must_use]
pub fn string_type_in(mut self, repr: StringRepr, paths: &[impl AsRef<str>]) -> Self {
self.codegen_config
.string_fields
.extend(paths.iter().map(|p| (p.as_ref().to_string(), repr)));
self
}
#[must_use]
pub fn string_type(mut self, repr: StringRepr) -> Self {
self.codegen_config
.string_fields
.push((".".to_string(), repr));
self
}
#[must_use]
pub fn type_attribute(mut self, path: impl Into<String>, attribute: impl Into<String>) -> Self {
self.codegen_config
.type_attributes
.push((normalize_attr_path(path.into()), attribute.into()));
self
}
#[must_use]
pub fn field_attribute(
mut self,
path: impl Into<String>,
attribute: impl Into<String>,
) -> Self {
self.codegen_config
.field_attributes
.push((normalize_attr_path(path.into()), attribute.into()));
self
}
#[must_use]
pub fn message_attribute(
mut self,
path: impl Into<String>,
attribute: impl Into<String>,
) -> Self {
self.codegen_config
.message_attributes
.push((normalize_attr_path(path.into()), attribute.into()));
self
}
#[must_use]
pub fn enum_attribute(mut self, path: impl Into<String>, attribute: impl Into<String>) -> Self {
self.codegen_config
.enum_attributes
.push((normalize_attr_path(path.into()), attribute.into()));
self
}
#[must_use]
pub fn use_buf(mut self) -> Self {
self.descriptor_source = DescriptorSource::Buf;
self
}
#[must_use]
pub fn descriptor_set(mut self, path: impl Into<PathBuf>) -> Self {
self.descriptor_source = DescriptorSource::Precompiled(path.into());
self
}
#[must_use]
pub fn include_file(mut self, name: impl Into<String>) -> Self {
self.include_file = Some(name.into());
self
}
pub fn compile(self) -> Result<(), Box<dyn std::error::Error>> {
let relative_includes = self.out_dir.is_some();
let out_dir = self
.out_dir
.or_else(|| std::env::var("OUT_DIR").ok().map(PathBuf::from))
.ok_or("OUT_DIR not set and no out_dir configured")?;
let descriptor_bytes = match &self.descriptor_source {
DescriptorSource::Protoc => invoke_protoc(&self.files, &self.includes)?,
DescriptorSource::Buf => invoke_buf()?,
DescriptorSource::Precompiled(path) => std::fs::read(path).map_err(|e| {
format!("failed to read descriptor set '{}': {}", path.display(), e)
})?,
};
let fds = FileDescriptorSet::decode_from_slice(&descriptor_bytes)
.map_err(|e| format!("failed to decode FileDescriptorSet: {}", e))?;
let files_to_generate: Vec<String> = if matches!(
self.descriptor_source,
DescriptorSource::Precompiled(_) | DescriptorSource::Buf
) {
self.files
.iter()
.filter_map(|f| f.to_str().map(str::to_string))
.collect()
} else {
self.files
.iter()
.map(|f| proto_relative_name(f, &self.includes))
.filter(|s| !s.is_empty())
.collect()
};
let (generated, warnings) = buffa_codegen::generate_with_diagnostics(
&fds.file,
&files_to_generate,
&self.codegen_config,
)?;
for warning in warnings {
println!("cargo:warning=buffa: {warning}");
}
let mut output_entries: Vec<(String, String)> = Vec::new();
for file in generated {
let path = out_dir.join(&file.name);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
write_if_changed(&path, file.content.as_bytes())?;
if file.kind == buffa_codegen::GeneratedFileKind::PackageMod {
output_entries.push((file.name, file.package));
}
}
if let Some(ref include_name) = self.include_file {
let include_content = generate_include_file(&output_entries, relative_includes);
let include_path = out_dir.join(include_name);
write_if_changed(&include_path, include_content.as_bytes())?;
}
match self.descriptor_source {
DescriptorSource::Buf => emit_buf_rerun_if_changed(),
DescriptorSource::Protoc => {
println!("cargo:rerun-if-env-changed=PROTOC");
for proto_file in &self.files {
println!("cargo:rerun-if-changed={}", proto_file.display());
}
}
DescriptorSource::Precompiled(ref path) => {
println!("cargo:rerun-if-changed={}", path.display());
}
}
Ok(())
}
}
impl Default for Config {
fn default() -> Self {
Self::new()
}
}
fn normalize_attr_path(mut path: String) -> String {
if !path.starts_with('.') {
path.insert(0, '.');
}
if path.len() > 1 {
while path.ends_with('.') {
path.pop();
}
}
path
}
fn write_if_changed(path: &Path, content: &[u8]) -> std::io::Result<()> {
if let Ok(existing) = std::fs::read(path) {
if existing == content {
return Ok(());
}
}
std::fs::write(path, content)
}
fn invoke_protoc(
files: &[PathBuf],
includes: &[PathBuf],
) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
let protoc = std::env::var("PROTOC").unwrap_or_else(|_| "protoc".to_string());
let descriptor_file =
tempfile::NamedTempFile::new().map_err(|e| format!("failed to create temp file: {}", e))?;
let descriptor_path = descriptor_file.path().to_path_buf();
let mut cmd = Command::new(&protoc);
cmd.arg("--include_imports");
cmd.arg("--include_source_info");
cmd.arg(format!(
"--descriptor_set_out={}",
descriptor_path.display()
));
for include in includes {
cmd.arg(format!("--proto_path={}", include.display()));
}
for file in files {
cmd.arg(file.as_os_str());
}
let output = cmd
.output()
.map_err(|e| format!("failed to run protoc ({}): {}", protoc, e))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
return Err(format!("protoc failed: {}", stderr).into());
}
let bytes = std::fs::read(&descriptor_path)
.map_err(|e| format!("failed to read descriptor set: {}", e))?;
Ok(bytes)
}
fn invoke_buf() -> Result<Vec<u8>, Box<dyn std::error::Error>> {
let output = Command::new("buf")
.arg("build")
.arg("--as-file-descriptor-set")
.arg("-o")
.arg("-")
.output()
.map_err(|e| format!("failed to run buf (is it installed and on PATH?): {e}"))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
return Err(
format!("buf build failed (is buf.yaml present at crate root?): {stderr}").into(),
);
}
Ok(output.stdout)
}
fn emit_buf_rerun_if_changed() {
println!("cargo:rerun-if-changed=buf.yaml");
if Path::new("buf.lock").exists() {
println!("cargo:rerun-if-changed=buf.lock");
}
match Command::new("buf").arg("ls-files").output() {
Ok(out) if out.status.success() => {
for line in String::from_utf8_lossy(&out.stdout).lines() {
let path = line.trim();
if !path.is_empty() {
println!("cargo:rerun-if-changed={path}");
}
}
}
_ => {
}
}
}
fn proto_relative_name(file: &Path, includes: &[PathBuf]) -> String {
let mut best: Option<&Path> = None;
for include in includes {
if let Ok(rel) = file.strip_prefix(include) {
match best {
Some(prev) if prev.as_os_str().len() <= rel.as_os_str().len() => {}
_ => best = Some(rel),
}
}
}
best.unwrap_or(file).to_str().unwrap_or("").to_string()
}
fn generate_include_file(entries: &[(String, String)], relative: bool) -> String {
let mode = if relative {
buffa_codegen::IncludeMode::Relative("")
} else {
buffa_codegen::IncludeMode::OutDir
};
buffa_codegen::generate_module_tree(entries, mode, false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn proto_relative_name_strips_include() {
let got = proto_relative_name(
Path::new("proto/my/service.proto"),
&[PathBuf::from("proto/")],
);
assert_eq!(got, "my/service.proto");
}
#[test]
fn proto_relative_name_longest_prefix_wins() {
let got = proto_relative_name(
Path::new("proto/vendor/ext.proto"),
&[PathBuf::from("proto/"), PathBuf::from("proto/vendor/")],
);
assert_eq!(got, "ext.proto");
let got = proto_relative_name(
Path::new("proto/vendor/ext.proto"),
&[PathBuf::from("proto/vendor/"), PathBuf::from("proto/")],
);
assert_eq!(got, "ext.proto");
}
#[test]
fn proto_relative_name_no_match_returns_full_path() {
let got = proto_relative_name(Path::new("my/pkg/service.proto"), &[]);
assert_eq!(got, "my/pkg/service.proto");
}
#[test]
fn proto_relative_name_no_match_with_unrelated_includes() {
let got = proto_relative_name(
Path::new("src/my.proto"),
&[PathBuf::from("other/"), PathBuf::from("third/")],
);
assert_eq!(got, "src/my.proto");
}
#[test]
fn include_file_out_dir_mode_uses_env_var() {
let entries = vec![
("foo.bar.rs".to_string(), "foo".to_string()),
("root.rs".to_string(), String::new()),
];
let out = generate_include_file(&entries, false);
assert!(
out.contains(r#"include!(concat!(env!("OUT_DIR"), "/foo.bar.rs"));"#),
"nested-package file should use env!(OUT_DIR): {out}"
);
assert!(
out.contains(r#"include!(concat!(env!("OUT_DIR"), "/root.rs"));"#),
"empty-package file should use env!(OUT_DIR): {out}"
);
assert!(!out.contains(r#"include!("foo.bar.rs")"#));
}
#[test]
fn include_file_relative_mode_uses_sibling_paths() {
let entries = vec![
("foo.bar.rs".to_string(), "foo".to_string()),
("root.rs".to_string(), String::new()),
];
let out = generate_include_file(&entries, true);
assert!(
out.contains(r#"include!("foo.bar.rs");"#),
"nested-package file should use relative path: {out}"
);
assert!(
out.contains(r#"include!("root.rs");"#),
"empty-package file should use relative path: {out}"
);
assert!(
!out.contains("OUT_DIR"),
"relative mode must not reference OUT_DIR: {out}"
);
}
#[test]
fn include_file_relative_mode_nested_packages() {
let entries = vec![
("a.b.one.rs".to_string(), "a.b".to_string()),
("a.b.two.rs".to_string(), "a.b".to_string()),
];
let out = generate_include_file(&entries, true);
let indent = " "; assert!(
out.contains(&format!(r#"{indent}include!("a.b.one.rs");"#)),
"first file at depth 2: {out}"
);
assert!(
out.contains(&format!(r#"{indent}include!("a.b.two.rs");"#)),
"second file at depth 2: {out}"
);
assert_eq!(
out.matches("pub mod b {").count(),
1,
"both files share one `mod b`: {out}"
);
assert!(!out.contains("OUT_DIR"));
}
#[test]
fn write_if_changed_creates_new_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("new.rs");
write_if_changed(&path, b"hello").unwrap();
assert_eq!(std::fs::read(&path).unwrap(), b"hello");
}
#[test]
fn write_if_changed_skips_identical_content() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("same.rs");
std::fs::write(&path, b"content").unwrap();
let mtime_before = std::fs::metadata(&path).unwrap().modified().unwrap();
std::thread::sleep(std::time::Duration::from_millis(50));
write_if_changed(&path, b"content").unwrap();
let mtime_after = std::fs::metadata(&path).unwrap().modified().unwrap();
assert_eq!(mtime_before, mtime_after);
}
#[test]
fn write_if_changed_overwrites_different_content() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("changed.rs");
std::fs::write(&path, b"old").unwrap();
write_if_changed(&path, b"new").unwrap();
assert_eq!(std::fs::read(&path).unwrap(), b"new");
}
#[test]
fn normalize_attr_path_prepends_leading_dot() {
assert_eq!(normalize_attr_path("my.pkg".into()), ".my.pkg");
}
#[test]
fn normalize_attr_path_preserves_leading_dot() {
assert_eq!(normalize_attr_path(".my.pkg".into()), ".my.pkg");
}
#[test]
fn normalize_attr_path_trims_trailing_dot() {
assert_eq!(normalize_attr_path("my.pkg.".into()), ".my.pkg");
assert_eq!(normalize_attr_path(".my.pkg.".into()), ".my.pkg");
assert_eq!(normalize_attr_path(".my.pkg...".into()), ".my.pkg");
}
#[test]
fn normalize_attr_path_preserves_catchall() {
assert_eq!(normalize_attr_path(".".into()), ".");
assert_eq!(normalize_attr_path("".into()), ".");
}
#[test]
fn type_attribute_forwards_normalized_path() {
let cfg = Config::new().type_attribute("my.pkg.", "#[derive(Foo)]");
assert_eq!(
cfg.codegen_config.type_attributes,
vec![(".my.pkg".to_string(), "#[derive(Foo)]".to_string())]
);
}
#[test]
fn field_attribute_forwards_normalized_path() {
let cfg = Config::new().field_attribute("pkg.Msg.f", "#[serde(skip)]");
assert_eq!(
cfg.codegen_config.field_attributes,
vec![(".pkg.Msg.f".to_string(), "#[serde(skip)]".to_string())]
);
}
#[test]
fn message_attribute_forwards_normalized_path() {
let cfg = Config::new().message_attribute(".", "#[serde(default)]");
assert_eq!(
cfg.codegen_config.message_attributes,
vec![(".".to_string(), "#[serde(default)]".to_string())]
);
}
#[test]
fn enum_attribute_forwards_normalized_path() {
let cfg = Config::new().enum_attribute("my.pkg.", "#[derive(strum::EnumIter)]");
assert_eq!(
cfg.codegen_config.enum_attributes,
vec![(
".my.pkg".to_string(),
"#[derive(strum::EnumIter)]".to_string(),
)]
);
assert!(cfg.codegen_config.type_attributes.is_empty());
assert!(cfg.codegen_config.message_attributes.is_empty());
assert!(cfg.codegen_config.field_attributes.is_empty());
}
#[test]
fn attribute_calls_accumulate_in_insertion_order() {
let cfg = Config::new()
.type_attribute(".", "#[derive(A)]")
.type_attribute(".pkg.M", "#[derive(B)]")
.type_attribute(".", "#[derive(C)]");
let paths: Vec<_> = cfg
.codegen_config
.type_attributes
.iter()
.map(|(_, a)| a.as_str())
.collect();
assert_eq!(paths, vec!["#[derive(A)]", "#[derive(B)]", "#[derive(C)]"]);
}
}