flatbuffers-tonic-build 0.1.0

gRPC with flatbuffers code generator.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use std::path::Path;

use flatbuffers_util::reflect::MessageType;
use heck::ToSnakeCase;
use proc_macro2::TokenStream;
use quote::quote;

pub(crate) fn compile_flatbuffers_tonic_only(
    path: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    let schema = flatbuffers_util::reflect::compile_reflection_schema(path);
    let schema_ref = schema.get_ref();
    let services = flatbuffers_util::reflect::get_services_from_schema(&schema_ref);
    let types = flatbuffers_util::reflect::collect_in_out_types(&services);

    // println!("Processing service: {services:?}");
    // println!("Processing types: {types:?}");

    // generate types alias file
    let out_dir = std::env::var("OUT_DIR").unwrap();
    let out_dir = Path::new(&out_dir);
    let package = types[0].namespace.as_ref().expect("no package found");
    let tonic_services = get_tonic_services(&services);
    tonic_prost_build::manual::Builder::new().compile(&tonic_services);

    let content = compile_flatbuffers_tonic_file(package, &types, &services);
    write_tonic_flatbuffers_file(out_dir, &content, package)?;
    Ok(())
}

fn get_tonic_services(
    services: &Vec<flatbuffers_util::reflect::Service>,
) -> Vec<tonic_prost_build::manual::Service> {
    let mut tonic_services = Vec::new();
    for service in services {
        let namespace = service.namespace.as_ref().expect("no namespace found");
        let mut tonic_service = tonic_prost_build::manual::Service::builder()
            .package(namespace)
            .name(&service.name);
        for method in &service.methods {
            let method_name = get_method_name(method);
            let tonic_method = tonic_prost_build::manual::Method::builder()
                .name(&method_name)
                .route_name(&method.name)
                .input_type(format!("super::Owned{}", method.request_type()))
                .output_type(format!("super::Owned{}", method.response_type()))
                .codec_path("flatbuffers_tonic::FlatBuffersCodec")
                .build();
            tonic_service = tonic_service.method(tonic_method);
        }
        tonic_services.push(tonic_service.build());
    }
    tonic_services
}

fn get_method_name(method: &flatbuffers_util::reflect::Method) -> String {
    // remove the prefix if any, and convert to snake_case
    let name = method.name.rsplit('.').next().unwrap().to_string();
    name.to_snake_case()
}

pub fn compile_types_alias(package: &str, types: &Vec<MessageType>) -> TokenStream {
    let mut content = TokenStream::new();

    for t in types {
        let wrapper_type = quote::format_ident!("Owned{}", t.fb_type);
        let rs_type = quote::format_ident!("{}", t.fb_type);
        let rs_type_mod = quote::format_ident!("{}", package);
        // add definition
        content.extend(quote! {
            pub struct #wrapper_type(pub flatbuffers_tonic::OwnedFB<#rs_type_mod::#rs_type<'static>>);
            impl #wrapper_type {
                pub fn get_ref<'a>(&'a self) -> #rs_type_mod::#rs_type<'a> {
                    self.0.get_ref()
                }
            }
        });
        // add impl
        content.extend(quote! {
            impl flatbuffers_tonic::OwnedFBCodecable for #wrapper_type {
                fn new_boxed(buf: Box<[u8]>) -> Result<Self, flatbuffers::InvalidFlatbuffer> {
                    let owned = flatbuffers_tonic::OwnedFB::<#rs_type_mod::#rs_type<'static>>::new_boxed(buf)?;
                    Ok(#wrapper_type(owned))
                }

                fn get_slice(&self) -> &[u8] {
                    self.0.get_slice()
                }
            }
        });
    }
    content
}

pub fn compile_flatbuffers_tonic_file(
    package: &str,
    types: &Vec<MessageType>,
    services: &Vec<flatbuffers_util::reflect::Service>,
) -> TokenStream {
    let mut content = TokenStream::new();

    let types_content = compile_types_alias(package, types);

    content.extend(quote! {
        /// This file is generated by flatbuffers-tonic-build
        /// Do not edit this file manually
    });
    // add flatbuffers includes
    content.extend(quote! {
        // flatbuffers generated code has warnings.
        include!(concat!(env!("OUT_DIR"), "/flatbuffers/", "mod.rs"));
    });

    // add include for flatbuffers and tonic files
    for service in services {
        let namespace = service.namespace.as_ref().expect("no namespace found");
        let service_name = &service.name;
        content.extend(quote! {
                include!(concat!(env!("OUT_DIR"), "/", #namespace, ".", #service_name, ".rs"));
        });
    }
    // add types content
    content.extend(types_content);
    content
}

fn write_tonic_flatbuffers_file(
    out_dir: &Path,
    content: &TokenStream,
    package: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    use std::fs::File;
    use std::io::Write;

    let dest_path = out_dir.join(format!("flatbuffers_tonic.{package}.rs"));
    let mut f = File::create(&dest_path)?;

    // Parse TokenStream to syn::File and pretty-print
    let syntax_tree: syn::File = syn::parse2(content.clone()).unwrap();
    let formatted = prettyplease::unparse(&syntax_tree);
    f.write_all(formatted.as_bytes())?;
    Ok(())
}