pub mod comments;
pub mod fence;
pub mod grpc_attach;
pub mod source;
use std::collections::BTreeMap;
use std::path::PathBuf;
use buffa_descriptor::generated::descriptor::MethodDescriptorProto;
use buffa_descriptor::generated::descriptor::{DescriptorProto, ServiceDescriptorProto};
use switchback_traits::{
CompanionFile, Entity, EntityBody, EntityCategory, EntityId, Group, GroupId, OperationBody,
RefKind, Reference, SchemaBody, ServiceBody, Source, SpecVersion,
};
use crate::category::ProtobufCategory;
use crate::companion::discover_companions;
use crate::descriptor::FileDescriptorProto;
use crate::descriptor_util::split_proto_type_name;
use crate::input::ResolvedInput;
use crate::populate::comments::{CommentIndex, dedent_comment, package_overview};
use crate::populate::fence::{
synthesize_enum_body, synthesize_message_body, synthesize_method_options_body,
synthesize_service_body,
};
use crate::populate::source::SourceCache;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
enum EntityKind {
Message,
Enum,
Service,
Operation,
}
pub struct PopulatedEntity {
pub entity: Entity<ProtobufCategory>,
pub refs: Vec<Reference>,
pub source_file: String,
}
pub struct PopulatedContract {
pub version: SpecVersion,
pub module_id: String,
pub groups: Vec<Group>,
pub entities_by_group: BTreeMap<GroupId, Vec<PopulatedEntity>>,
pub companions: Vec<CompanionFile>,
pub module_root: PathBuf,
}
pub fn populate(resolved: &ResolvedInput) -> switchback_traits::Result<PopulatedContract> {
let by_package = packages_map(&resolved.proto_file, &resolved.file_to_generate);
let mut source = SourceCache::new(resolved.proto_search_paths.clone());
let companions = discover_companions(
&resolved.proto_file,
&resolved.file_to_generate,
&resolved.proto_search_paths,
&resolved.module_root,
)?;
let module_id = default_module_id(&by_package);
let mut groups = Vec::new();
let mut entities_by_group = BTreeMap::new();
for (package, files) in &by_package {
let dir = package_dir(package);
let overview = package_overview(files);
let first_proto = files.first().map(|(n, _)| *n).unwrap_or("");
groups.push(Group {
id: GroupId::from(package.as_str()),
dir,
title: package.clone(),
overview,
source: Some(Source {
file: first_proto.to_string(),
span: None,
}),
entities: Vec::new(),
source_path: resolved.module_root.join(first_proto),
});
let mut entities = Vec::new();
for_each_entity_in_files(files, |kind, name, proto_name, file, mi, si, method_idx| {
let idx = CommentIndex::from_file(file);
match kind {
EntityKind::Message => {
let msg = &file.message_type[mi];
let doc = idx
.leading_message(mi)
.map(dedent_comment)
.filter(|s| !s.is_empty());
let fence_body =
synthesize_message_body(proto_name, &idx, mi, msg, Some(&mut source));
entities.push(PopulatedEntity {
entity: Entity {
id: EntityId::new(
package.as_str(),
ProtobufCategory::Schema.as_str(),
name,
),
category: ProtobufCategory::Schema,
title: name.to_string(),
doc,
source_span: None,
body: EntityBody::Schema(SchemaBody {
fence_language: "protobuf".into(),
fence_body,
payload_format: String::new(),
properties: Vec::new(),
}),
},
refs: message_field_refs(&module_id, msg),
source_file: proto_name.to_string(),
});
}
EntityKind::Enum => {
let en = &file.enum_type[mi];
let doc = idx
.leading_enum(mi)
.map(dedent_comment)
.filter(|s| !s.is_empty());
let fence_body = synthesize_enum_body(proto_name, &idx, mi, en);
entities.push(PopulatedEntity {
entity: Entity {
id: EntityId::new(
package.as_str(),
ProtobufCategory::Schema.as_str(),
name,
),
category: ProtobufCategory::Schema,
title: name.to_string(),
doc,
source_span: None,
body: EntityBody::Schema(SchemaBody {
fence_language: "protobuf".into(),
fence_body,
payload_format: String::new(),
properties: Vec::new(),
}),
},
refs: Vec::new(),
source_file: proto_name.to_string(),
});
}
EntityKind::Service => {
let svc = &file.service[si];
let doc = idx
.leading_service(si)
.map(dedent_comment)
.filter(|s| !s.is_empty());
let fence_body = synthesize_service_body(proto_name, &idx, si, svc);
entities.push(PopulatedEntity {
entity: Entity {
id: EntityId::new(
package.as_str(),
ProtobufCategory::Service.as_str(),
name,
),
category: ProtobufCategory::Service,
title: name.to_string(),
doc,
source_span: None,
body: EntityBody::Service(ServiceBody {
signature: name.to_string(),
fence_language: "protobuf".into(),
fence_body,
}),
},
refs: service_method_refs(&module_id, svc),
source_file: proto_name.to_string(),
});
}
EntityKind::Operation => {
let svc = &file.service[si];
let method = &svc.method[method_idx];
let service_name = svc.name.as_deref().unwrap_or("Service");
let rpc_name = method.name.as_deref().unwrap_or("Rpc");
let op_name = format!("{service_name}.{rpc_name}");
let doc = idx
.leading_method(si, method_idx)
.map(dedent_comment)
.filter(|s| !s.is_empty());
let output_fqn = method
.output_type
.as_deref()
.unwrap_or(".google.protobuf.Empty");
let (signature, protocols, responses, parameters) =
grpc_attach::populate_rpc(method, &module_id, output_fqn);
let fence_body = synthesize_method_options_body(method).unwrap_or_default();
entities.push(PopulatedEntity {
entity: Entity {
id: EntityId::new(
package.as_str(),
ProtobufCategory::Operation.as_str(),
&op_name,
),
category: ProtobufCategory::Operation,
title: rpc_name.to_string(),
doc,
source_span: None,
body: EntityBody::Operation(OperationBody {
signature,
fence_language: "protobuf".into(),
fence_body,
parameters,
responses,
request_body: None,
protocols,
}),
},
refs: operation_refs(&module_id, method),
source_file: proto_name.to_string(),
});
}
}
});
entities_by_group.insert(GroupId::from(package.as_str()), entities);
}
let version = detect_spec_version(&resolved.proto_file, &resolved.file_to_generate);
Ok(PopulatedContract {
version,
module_id,
groups,
entities_by_group,
companions,
module_root: resolved.module_root.clone(),
})
}
fn operation_refs(module_id: &str, method: &MethodDescriptorProto) -> Vec<Reference> {
let mut refs = Vec::new();
for fqn in [method.input_type.as_deref(), method.output_type.as_deref()] {
if let Some(reference) = fqn_type_ref(module_id, fqn) {
refs.push(reference);
}
}
refs
}
fn message_field_refs(module_id: &str, msg: &DescriptorProto) -> Vec<Reference> {
let mut refs = Vec::new();
let mut seen = std::collections::BTreeSet::new();
for field in &msg.field {
let Some(type_name) = field.type_name.as_deref() else {
continue;
};
if let Some(reference) = fqn_type_ref(module_id, Some(type_name)) {
let key = (
reference.target.group.clone(),
reference.target.name.clone(),
);
if seen.insert(key) {
refs.push(reference);
}
}
}
refs
}
fn service_method_refs(module_id: &str, svc: &ServiceDescriptorProto) -> Vec<Reference> {
let mut refs = Vec::new();
let mut seen = std::collections::BTreeSet::new();
for method in &svc.method {
for reference in operation_refs(module_id, method) {
let key = (
reference.target.group.clone(),
reference.target.name.clone(),
);
if seen.insert(key) {
refs.push(reference);
}
}
}
refs
}
fn fqn_type_ref(module_id: &str, fqn: Option<&str>) -> Option<Reference> {
let fqn = fqn?;
let (pkg, msg) = split_proto_type_name(fqn)?;
if pkg.starts_with("google.protobuf") {
return None;
}
Some(Reference {
target: switchback_traits::EntityRef {
module: module_id.to_string(),
group: pkg.to_string(),
category: ProtobufCategory::Schema.as_str().to_string(),
name: msg.to_string(),
},
kind: RefKind::Internal,
})
}
fn default_module_id(by_package: &BTreeMap<String, Vec<(&str, &FileDescriptorProto)>>) -> String {
by_package
.keys()
.next()
.cloned()
.unwrap_or_else(|| "default".into())
}
fn package_dir(package: &str) -> String {
if package.is_empty() {
return "_root".into();
}
package.replace('.', "/")
}
fn detect_spec_version(
proto_file: &[FileDescriptorProto],
file_to_generate: &[String],
) -> SpecVersion {
for name in file_to_generate {
if let Some(file) = proto_file
.iter()
.find(|f| f.name.as_deref() == Some(name.as_str()))
{
if let Some(edition) = &file.edition {
return SpecVersion::from(edition_to_spec_version(*edition));
}
if let Some(syntax) = &file.syntax {
return SpecVersion::from(syntax.as_str());
}
}
}
SpecVersion::from("proto3")
}
fn edition_to_spec_version(
edition: buffa_descriptor::generated::descriptor::Edition,
) -> &'static str {
use buffa_descriptor::generated::descriptor::Edition;
match edition {
Edition::EDITION_PROTO2 => "proto2",
Edition::EDITION_PROTO3 => "proto3",
Edition::EDITION_2023 => "2023",
Edition::EDITION_2024 => "2024",
_ => "proto3",
}
}
pub fn packages_map<'a>(
proto_files: &'a [FileDescriptorProto],
file_to_generate: &'a [String],
) -> BTreeMap<String, Vec<(&'a str, &'a FileDescriptorProto)>> {
let mut by_package: BTreeMap<String, Vec<(&'a str, &'a FileDescriptorProto)>> = BTreeMap::new();
for name in file_to_generate {
let Some(file) = proto_files
.iter()
.find(|f| f.name.as_deref() == Some(name.as_str()))
else {
continue;
};
let pkg = file.package.clone().unwrap_or_default();
by_package
.entry(pkg)
.or_default()
.push((name.as_str(), file));
}
by_package
}
fn for_each_entity_in_files(
files: &[(&str, &FileDescriptorProto)],
mut f: impl FnMut(EntityKind, &str, &str, &FileDescriptorProto, usize, usize, usize),
) {
for (proto_name, file) in files {
for (i, msg) in file.message_type.iter().enumerate() {
if let Some(name) = msg.name.as_deref() {
f(EntityKind::Message, name, proto_name, file, i, 0, 0);
}
}
for (i, en) in file.enum_type.iter().enumerate() {
if let Some(name) = en.name.as_deref() {
f(EntityKind::Enum, name, proto_name, file, i, 0, 0);
}
}
for (si, svc) in file.service.iter().enumerate() {
if let Some(name) = svc.name.as_deref() {
f(EntityKind::Service, name, proto_name, file, 0, si, 0);
for (mi, method) in svc.method.iter().enumerate() {
if method.name.is_some() {
f(EntityKind::Operation, name, proto_name, file, 0, si, mi);
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::descriptor::FileDescriptorProto;
fn file(name: &str, package: &str) -> FileDescriptorProto {
FileDescriptorProto {
name: Some(name.into()),
package: Some(package.into()),
..Default::default()
}
}
#[test]
fn packages_map_preserves_file_to_generate_order_within_package() {
let a = file("pkg/a.proto", "acme.v1");
let b = file("pkg/b.proto", "acme.v1");
let proto_files = vec![a.clone(), b.clone()];
let forward_inputs = ["pkg/a.proto".into(), "pkg/b.proto".into()];
let forward = packages_map(&proto_files, &forward_inputs);
let reverse_inputs = ["pkg/b.proto".into(), "pkg/a.proto".into()];
let reverse = packages_map(&proto_files, &reverse_inputs);
let forward_names: Vec<_> = forward["acme.v1"].iter().map(|(n, _)| *n).collect();
let reverse_names: Vec<_> = reverse["acme.v1"].iter().map(|(n, _)| *n).collect();
assert_eq!(forward_names, ["pkg/a.proto", "pkg/b.proto"]);
assert_eq!(reverse_names, ["pkg/b.proto", "pkg/a.proto"]);
}
}