protify/

file.rs

use crate::*;
use hashbrown::HashSet;

/// Struct that represents a protobuf file and its contents.
#[derive(Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "std", derive(Template))]
#[cfg_attr(feature = "std", template(path = "file.proto.j2"))]
#[non_exhaustive]
pub struct ProtoFile {
	pub name: FixedStr,
	pub package: FixedStr,
	pub imports: FileImports,
	pub messages: Vec<MessageSchema>,
	pub enums: Vec<EnumSchema>,
	pub options: Vec<ProtoOption>,
	pub edition: Edition,
	pub services: Vec<Service>,
	pub extensions: Vec<Extension>,
}

impl ProtoFile {
	pub(crate) fn sort_items(&mut self) {
		self.extensions
			.sort_unstable_by_key(|e| e.target.as_str());

		self.messages
			.sort_unstable_by_key(|m| m.name.clone());

		for msg in self.messages.iter_mut() {
			sort_nested(msg);
		}

		self.enums
			.sort_unstable_by_key(|e| e.short_name.clone());
		self.services
			.sort_unstable_by_key(|s| s.name.clone());
	}

	#[doc(hidden)]
	#[must_use]
	pub fn new(name: &'static str, package: &'static str) -> Self {
		Self {
			name: name.into(),
			package: package.into(),
			imports: FileImports::new(name),
			messages: Default::default(),
			enums: Default::default(),
			options: Default::default(),
			edition: Default::default(),
			services: Default::default(),
			extensions: Default::default(),
		}
	}

	#[doc(hidden)]
	#[inline]
	pub fn with_options(&mut self, options: Vec<ProtoOption>) -> &mut Self {
		self.options = options;
		self
	}

	#[doc(hidden)]
	#[inline]
	pub fn with_imports(
		&mut self,
		imports: impl IntoIterator<Item = impl Into<FixedStr>>,
	) -> &mut Self {
		self.imports
			.extend(imports.into_iter().map(|s| s.into()));
		self
	}

	#[doc(hidden)]
	#[inline]
	pub const fn with_edition(&mut self, edition: Edition) -> &mut Self {
		self.edition = edition;
		self
	}

	#[doc(hidden)]
	pub fn with_messages(&mut self, mut messages: Vec<MessageSchema>) -> &mut Self {
		for message in &mut messages {
			message.register_imports(&mut self.imports);
			message.file = self.name.clone();
		}

		self.messages.append(&mut messages);

		self
	}

	#[doc(hidden)]
	#[inline]
	pub fn with_enums(&mut self, mut enums: Vec<EnumSchema>) -> &mut Self {
		for enum_ in &mut enums {
			enum_.file = self.name.clone();
		}

		self.enums.append(&mut enums);

		self
	}

	#[doc(hidden)]
	pub fn with_services(&mut self, mut services: Vec<Service>) -> &mut Self {
		for service in &services {
			for (request, response) in service
				.handlers
				.iter()
				.map(|h| (&h.request, &h.response))
			{
				self.imports.insert_from_path(&request.message);
				self.imports.insert_from_path(&response.message);
			}

			if *service.file != *self.name {
				self.imports.set.insert(service.file.clone());
			}
		}

		self.services.append(&mut services);

		self
	}

	#[doc(hidden)]
	pub fn with_extensions(&mut self, mut extensions: Vec<Extension>) -> &mut Self {
		if !extensions.is_empty() {
			self.imports
				.set
				.insert("google/protobuf/descriptor.proto".into());
		}

		self.extensions.append(&mut extensions);

		self
	}
}

fn sort_nested(message: &mut MessageSchema) {
	message
		.messages
		.sort_unstable_by_key(|m| m.name.clone());
	message
		.enums
		.sort_unstable_by_key(|e| e.name.clone());
}

/// Trait that can generate a [`ProtoFile`].
///
/// Implemented by the unit structs generated by the [`define_proto_file`] macro.
pub trait FileSchema {
	const NAME: &str;
	const PACKAGE: &str;
	const EXTERN_PATH: &str;
	fn file_schema() -> ProtoFile;
}

#[doc(hidden)]
pub struct FileReference {
	pub name: &'static str,
	pub package: &'static str,
	pub extern_path: &'static str,
}

/// The protobuf edition for a file.
#[derive(Default, Debug, PartialEq, Eq, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Edition {
	#[default]
	Proto3,
	E2023,
}

impl Display for Edition {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		match self {
			Self::Proto3 => write!(f, "syntax = \"proto3\""),
			Self::E2023 => write!(f, "edition = \"2023\""),
		}
	}
}

/// HashSet wrapper for a file's imports. Skips insertion if the file is equal to the origin file.
#[derive(PartialEq, Eq, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct FileImports {
	pub set: HashSet<FixedStr>,
	pub file: FixedStr,
	#[cfg_attr(feature = "serde", serde(skip))]
	pub(crate) added_validate_proto: bool,
}

impl Extend<FixedStr> for FileImports {
	fn extend<T: IntoIterator<Item = FixedStr>>(&mut self, iter: T) {
		let iter = iter.into_iter();

		let reserve = if self.set.is_empty() {
			iter.size_hint().0
		} else {
			iter.size_hint().0.div_ceil(2)
		};

		self.set.reserve(reserve);

		for import in iter {
			self.insert_internal(import);
		}
	}
}

impl IntoIterator for FileImports {
	type Item = FixedStr;
	type IntoIter = hashbrown::hash_set::IntoIter<FixedStr>;

	fn into_iter(self) -> Self::IntoIter {
		self.set.into_iter()
	}
}

impl FileImports {
	#[must_use]
	pub(crate) fn new(file: impl Into<FixedStr>) -> Self {
		Self {
			file: file.into(),
			set: HashSet::default(),
			added_validate_proto: false,
		}
	}

	pub(crate) fn insert_validate_proto(&mut self) {
		if !self.added_validate_proto {
			self.set
				.insert("buf/validate/validate.proto".into());
			self.added_validate_proto = true;
		}
	}

	pub(crate) fn insert_internal(&mut self, import: FixedStr) {
		if *import != *self.file {
			if import == "buf/validate/validate.proto" {
				self.insert_validate_proto();
			} else {
				self.set.insert(import);
			}
		}
	}

	pub(crate) fn insert_from_path(&mut self, path: &ProtoPath) {
		if path.file != self.file {
			self.set.insert(path.file.clone());
		}
	}

	#[must_use]
	pub(crate) fn as_sorted_vec(&self) -> Vec<&FixedStr> {
		let mut imports: Vec<&FixedStr> = self.set.iter().collect();

		imports.sort_unstable();

		imports
	}
}