protoc-gen-prost-crate 0.4.1

Protocol Buffers compiler plugin powered by Prost!
Documentation
use std::rc::Rc;

use once_cell::sync::Lazy;
use prost_build::Module;
use prost_types::compiler::code_generator_response::File;
use protoc_gen_prost::{Generator, ModuleRequestSet, Result};

use crate::PackageLimiter;

pub(crate) struct IncludeFileGenerator<'a> {
    filename: &'a str,
    limiter: Rc<PackageLimiter>,
}

impl<'a> IncludeFileGenerator<'a> {
    pub(crate) fn new(filename: &'a str, limiter: Rc<PackageLimiter>) -> Self {
        Self { filename, limiter }
    }
}

impl<'a> Generator for IncludeFileGenerator<'a> {
    fn generate(&mut self, module_request_set: &ModuleRequestSet) -> Result {
        let mut context = CodeGenContext::new();

        let _: () = module_request_set
            .requests()
            .filter_map(|(module, request)| {
                let filename = request.output_filename()?;
                self.limiter
                    .is_allowed(request.proto_package_name())
                    .then(|| {
                        context.move_to(module, request.proto_package_name());
                        context.push_include(filename);
                        context.push_insertion_point(request.proto_package_name());
                    })
            })
            .collect();

        let content = context.finish();

        let file = File {
            name: Some(self.filename.to_string()),
            content: Some(content),
            ..File::default()
        };

        Ok(vec![file])
    }
}

static ROOT_MODULE: Lazy<Module> = Lazy::new(|| Module::from_parts([] as [String; 0]));

#[must_use]
#[derive(Debug)]
struct CodeGenContext<'a> {
    last: &'a Module,
    indent: String,
    buf: String,
}

const INDENT: &str = "    ";

impl<'a> CodeGenContext<'a> {
    fn new() -> Self {
        // Reserve an initial 16 kiB
        let mut buf = String::with_capacity(16_384);
        buf.push_str("// @generated\n");

        let indent = String::with_capacity(INDENT.len() * 8);

        Self {
            last: &*ROOT_MODULE,
            indent,
            buf,
        }
    }

    fn indent(&mut self) {
        self.indent.push_str(INDENT);
    }

    fn dedent(&mut self) {
        self.indent.truncate(
            self.indent
                .len()
                .checked_sub(INDENT.len())
                .expect("indent underflow"),
        );
    }

    fn push_indent(&mut self) {
        self.buf.push_str(&self.indent);
    }

    fn move_to(&mut self, next: &'a Module, package: &str) {
        let (down, prefix) = difference(self.last, next);

        for _ in 0..down {
            self.close_module();
        }

        let take = next.len() - prefix - 1;

        for module_name in next.parts().skip(prefix).take(take) {
            self.open_module(module_name);
        }

        self.push_attribute_insertion_point(package);
        self.open_module(next.parts().last().unwrap());

        self.last = next;
    }

    fn finish(mut self) -> String {
        while !self.indent.is_empty() {
            self.close_module()
        }
        self.buf
    }

    fn push_include(&mut self, filename: &str) {
        self.push_indent();
        self.buf.push_str("include!(\"");
        self.buf.push_str(filename);
        self.buf.push_str("\");\n");
    }

    fn push_insertion_point(&mut self, package_name: &str) {
        self.push_indent();
        self.buf.push_str("// @@protoc_insertion_point(");
        self.buf.push_str(package_name);
        self.buf.push_str(")\n");
    }

    fn push_attribute_insertion_point(&mut self, package_name: &str) {
        self.push_indent();
        self.buf.push_str("// @@protoc_insertion_point(attribute:");
        self.buf.push_str(package_name);
        self.buf.push_str(")\n");
    }

    fn close_module(&mut self) {
        self.dedent();
        self.push_indent();
        self.buf.push_str("}\n");
    }

    fn open_module(&mut self, module_name: &str) {
        self.push_indent();
        self.buf.push_str("pub mod ");
        self.buf.push_str(module_name);
        self.buf.push_str(" {\n");
        self.indent();
    }
}

fn difference(left: &Module, right: &Module) -> (usize, usize) {
    let mut left_parts = left.parts();
    let mut right_parts = right.parts();

    let mut prefix = 0;

    loop {
        match (left_parts.next(), right_parts.next()) {
            (Some(left), Some(right)) if left == right => prefix += 1,
            (Some(_), Some(_)) => return (left_parts.count() + 1, prefix),
            (Some(_), None) => return (left_parts.count() + 1, prefix),
            (None, Some(_)) => return (0, prefix),
            (None, None) => return (0, prefix),
        }
    }
}