Skip to main content

krabka_protocol_codegen/emit/
wrappers.rs

1//! Generate the `include!` wrapper module bodies under
2//! `crates/protocol/src/{owned,borrowed}/<snake>.rs`.
3
4use proc_macro2::TokenStream;
5use quote::{format_ident, quote};
6
7use crate::{
8    emit::{
9        EmitError, borrowed_quote,
10        common::{banner, validate},
11        tagged_fixture,
12    },
13    ir::{MessageSpec, MessageType},
14    name_conv,
15};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum Flavor {
19    Owned,
20    Borrowed,
21}
22
23impl Flavor {
24    #[must_use]
25    pub fn dir(self) -> &'static str {
26        match self {
27            Self::Owned => "owned",
28            Self::Borrowed => "borrowed",
29        }
30    }
31}
32
33/// Emit a wrapper body for one message + flavor.
34/// # Errors
35/// Returns an error when the generated code is not valid Rust, which is an
36/// emitter bug.
37pub fn emit(
38    spec: &MessageSpec,
39    flavor: Flavor,
40    schemas_version: &str,
41    namespace: Option<&str>,
42) -> Result<String, EmitError> {
43    let type_name = name_conv::type_name(&spec.name);
44    let suffix = match flavor {
45        Flavor::Owned => "owned",
46        Flavor::Borrowed => "borrowed",
47    };
48
49    // 1. Banner string.
50    let mut out = banner(schemas_version);
51
52    // 2. include!(concat!(env!("CARGO_MANIFEST_DIR"), "/generated/{path}.rs"));
53    let path_prefix = namespace.map_or_else(String::new, |ns| format!("{ns}/"));
54    let path = format!("/generated/{path_prefix}{type_name}.{suffix}.rs");
55    let include_tokens: TokenStream = quote! {
56        include!(concat!(env!("CARGO_MANIFEST_DIR"), #path));
57    };
58    validate::<syn::Stmt>(&include_tokens, "include!")?;
59    out.push_str(&include_tokens.to_string());
60    out.push('\n');
61
62    // 3. #[cfg(test)] mod tests { ... } built with quote!.
63    let tests_tokens = match flavor {
64        Flavor::Owned => owned_tests_tokens(&type_name, tagged_fixture::has_tagged_field(spec))?,
65        Flavor::Borrowed => borrowed_tests_tokens(
66            &type_name,
67            borrowed_quote::fallible_structs(spec).contains(&spec.name),
68        )?,
69    };
70    out.push_str(&tests_tokens.to_string());
71
72    Ok(out)
73}
74
75/// Build `#[cfg(test)] mod tests { ... }` for the Owned flavor as a `TokenStream`.
76fn owned_tests_tokens(type_name: &str, has_tagged_field: bool) -> Result<TokenStream, EmitError> {
77    let ty = format_ident!("{type_name}");
78    // `tagged_fixture()` holds a non-default value in every tagged field. It
79    // must round-trip at every version, including the versions where a tag is
80    // out of range and the encoder drops it. At the highest version every tag
81    // is in range, so the bytes must differ from the default message.
82    let tagged_test = has_tagged_field.then(|| {
83        quote! {
84            #[test]
85            fn tagged_fixture_roundtrips_all_versions() {
86                for v in MIN_VERSION..=MAX_VERSION {
87                    roundtrip(&tagged_fixture(), v);
88                    assert!(tagged_fixture_json(v).is_object());
89                }
90                let encode = |msg: &#ty| {
91                    let mut buf = BytesMut::new();
92                    msg.encode(&mut buf, MAX_VERSION).unwrap();
93                    buf
94                };
95                assert!(encode(&tagged_fixture()) != encode(&#ty::default()));
96            }
97        }
98    });
99    let tokens = quote! {
100        #[cfg(test)]
101        mod tests {
102            use assert2::assert;
103            use super::*;
104            use crate::{Decode, Encode};
105            use bytes::BytesMut;
106
107            fn roundtrip(msg: &#ty, v: i16) {
108                let mut buf = BytesMut::new();
109                msg.encode(&mut buf, v).unwrap();
110                assert!(msg.encoded_len(v) == buf.len());
111                let bytes = buf.freeze();
112                let mut cur = &bytes[..];
113                let decoded = #ty::decode(&mut cur, v).unwrap();
114                assert!(cur.is_empty());
115                let mut reencoded = BytesMut::new();
116                decoded.encode(&mut reencoded, v).unwrap();
117                assert!(&reencoded[..] == &bytes[..]);
118                let _ = default_json(v);
119            }
120
121            #[test]
122            fn default_roundtrips_all_versions() {
123                for v in MIN_VERSION..=MAX_VERSION {
124                    roundtrip(&#ty::default(), v);
125                }
126            }
127
128            #[test]
129            fn populated_roundtrips_all_versions() {
130                for v in MIN_VERSION..=MAX_VERSION {
131                    roundtrip(&#ty::populated(v), v);
132                }
133            }
134
135            #tagged_test
136        }
137    };
138    validate::<syn::ItemMod>(&tokens, "owned tests")?;
139    Ok(tokens)
140}
141
142/// Build `#[cfg(test)] mod tests { ... }` for the Borrowed flavor as a `TokenStream`.
143fn borrowed_tests_tokens(type_name: &str, fallible: bool) -> Result<TokenStream, EmitError> {
144    let ty = format_ident!("{type_name}");
145    let to_owned = if fallible {
146        quote!(decoded.to_owned().unwrap())
147    } else {
148        quote!(decoded.to_owned())
149    };
150    let tokens = quote! {
151        #[cfg(test)]
152        mod tests {
153            use assert2::assert;
154            use super::*;
155            use crate::{DecodeBorrow, Encode};
156            use bytes::BytesMut;
157
158            fn check(msg_bytes: &bytes::Bytes, v: i16) {
159                let mut cur: &[u8] = msg_bytes;
160                let decoded = #ty::decode_borrow(&mut cur, v).unwrap();
161                assert!(cur.is_empty());
162                assert!(decoded.encoded_len(v) == msg_bytes.len());
163                let mut reencoded = BytesMut::new();
164                decoded.encode(&mut reencoded, v).unwrap();
165                assert!(&reencoded[..] == &msg_bytes[..]);
166                let owned = #to_owned;
167                let mut owned_buf = BytesMut::new();
168                owned.encode(&mut owned_buf, v).unwrap();
169                assert!(&owned_buf[..] == &msg_bytes[..]);
170            }
171
172            #[test]
173            fn default_roundtrips_all_versions() {
174                for v in MIN_VERSION..=MAX_VERSION {
175                    let msg = #ty::default();
176                    let mut buf = BytesMut::new();
177                    msg.encode(&mut buf, v).unwrap();
178                    check(&buf.freeze(), v);
179                }
180            }
181
182            #[test]
183            fn populated_roundtrips_all_versions() {
184                for v in MIN_VERSION..=MAX_VERSION {
185                    let msg = #ty::populated(v);
186                    let mut buf = BytesMut::new();
187                    msg.encode(&mut buf, v).unwrap();
188                    check(&buf.freeze(), v);
189                }
190            }
191        }
192    };
193    validate::<syn::ItemMod>(&tokens, "borrowed tests")?;
194    Ok(tokens)
195}
196
197/// True if this spec should get a wrapper.
198#[must_use]
199pub fn should_emit_wrapper(spec: &MessageSpec) -> bool {
200    !spec.valid_versions.is_empty()
201        && matches!(
202            spec.message_type,
203            MessageType::Request | MessageType::Response | MessageType::Header | MessageType::Data
204        )
205}