krabka_protocol_codegen/emit/
differential_table.rs

1//! Emit a Rust source file that provides:
2//!   - `CASES: &[Case]` — a static table of `(message, api_key, version, kind)` cases
3//!   - `default_json_for(name)` — JSON the oracle should accept for default fixture
4//!   - `encode_default(name, version)` — Rust-encoded bytes from `Default::default()`
5//!   - `TAGGED_CASES`, `encode_tagged_fixture(name, version)` and
6//!     `tagged_fixture_json_for(name, version)` — the same sweep over every
7//!     message that has a tagged field, with each tagged field set
8//!
9//! `crates/protocol/tests/differential_all.rs` consumes the file with
10//! `include!`.
11//!
12//! This module emits with `proc_macro2` and `quote!` instead of `writeln!`.
13//! See `emit::api_key_enum_quote` for the overall pattern.
14
15use proc_macro2::{Literal, TokenStream};
16use quote::{format_ident, quote};
17
18use crate::{
19    emit::{EmitError, common::validate, tagged_fixture},
20    ir::{MessageSpec, MessageType},
21    name_conv,
22};
23
24/// Maximum version this emitter ever iterates to. Schemas with open-ended
25/// `validVersions`, such as `0+`, would produce `i16::MAX` iterations, which is
26/// wrong. In practice all active Kafka 4.2 schemas have bounded ranges, but
27/// this cap is a defence against enormous output.
28const MAX_VERSION_CAP: i16 = 127;
29
30/// # Errors
31/// Returns an error when the generated table is not valid Rust, which is an
32/// emitter bug.
33pub fn emit(specs: &[MessageSpec], schemas_version: &str) -> Result<String, EmitError> {
34    let banner = format!(
35        "// AUTO-GENERATED by krabka-protocol-codegen against {schemas_version}. Do not edit.\n\n"
36    );
37
38    let tokens = build_body(specs);
39
40    // Validate at generation time; rustfmt (applied by the caller) does the
41    // actual formatting. See `crate::fmt`.
42    validate::<syn::File>(&tokens, "differential table")?;
43
44    Ok(format!("{banner}{tokens}"))
45}
46
47fn build_body(specs: &[MessageSpec]) -> TokenStream {
48    let mut ts = TokenStream::new();
49    ts.extend(build_preamble());
50    ts.extend(build_cases_table(specs));
51    ts.extend(build_encode_default(specs));
52    ts.extend(build_default_json_for(specs));
53    ts.extend(build_tagged_fixture(specs));
54    ts.extend(build_roundtrip(specs));
55    ts.extend(build_header_versions(specs));
56    ts.extend(build_strip_frame_header());
57    ts
58}
59
60fn build_preamble() -> TokenStream {
61    quote! {
62        use std::{error::Error, fmt};
63
64        use bytes::BytesMut;
65        use krabka_protocol::{Encode, ProtocolError};
66
67        #[derive(Debug, Clone, Copy)]
68        pub struct Case {
69            pub name: &'static str,
70            pub api_key: i16,
71            pub version: i16,
72            pub kind: Kind,
73        }
74
75        #[derive(Debug, Clone, Copy)]
76        pub enum Kind { Request, Response, RequestHeader, ResponseHeader }
77
78        #[doc = " The owned and the borrowed encoding of one message."]
79        pub type EncodedPair = (Vec<u8>, Vec<u8>);
80
81        #[doc = " Why a table lookup, encode or round trip failed."]
82        #[derive(Debug)]
83        pub enum TableError {
84            #[doc = " `name` identifies no message the table covers."]
85            UnknownMessage(String),
86            #[doc = " The codec rejected the message."]
87            Protocol(ProtocolError),
88            #[doc = " A decode left bytes behind."]
89            TrailingBytes { name: String, remaining: usize },
90        }
91
92        impl fmt::Display for TableError {
93            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94                match self {
95                    Self::UnknownMessage(name) => write!(f, "unknown message {name}"),
96                    Self::Protocol(e) => write!(f, "{e}"),
97                    Self::TrailingBytes { name, remaining } => {
98                        write!(f, "{name}: {remaining} bytes left after decode")
99                    }
100                }
101            }
102        }
103
104        impl Error for TableError {
105            fn source(&self) -> Option<&(dyn Error + 'static)> {
106                match self {
107                    Self::Protocol(e) => Some(e),
108                    Self::UnknownMessage(_) | Self::TrailingBytes { .. } => None,
109                }
110            }
111        }
112
113        impl From<ProtocolError> for TableError {
114            fn from(e: ProtocolError) -> Self {
115                Self::Protocol(e)
116            }
117        }
118
119        fn encode_with<T: Encode>(message: &T, version: i16) -> Result<Vec<u8>, TableError> {
120            let mut buf = BytesMut::new();
121            message.encode(&mut buf, version)?;
122            Ok(buf.to_vec())
123        }
124
125        #[doc = " Fails when a decode of `name` left bytes in `rest`."]
126        fn expect_consumed(name: &str, rest: &[u8]) -> Result<(), TableError> {
127            if rest.is_empty() {
128                Ok(())
129            } else {
130                Err(TableError::TrailingBytes { name: name.to_owned(), remaining: rest.len() })
131            }
132        }
133    }
134}
135
136fn build_cases_table(specs: &[MessageSpec]) -> TokenStream {
137    let rows: Vec<TokenStream> = specs
138        .iter()
139        .flat_map(|s| {
140            if s.valid_versions.is_empty() || s.internal {
141                return Vec::new();
142            }
143            let kind_tokens: TokenStream = match s.message_type {
144                MessageType::Request => quote!(Kind::Request),
145                MessageType::Response => quote!(Kind::Response),
146                MessageType::Header => match s.name.as_str() {
147                    "RequestHeader" => quote!(Kind::RequestHeader),
148                    "ResponseHeader" => quote!(Kind::ResponseHeader),
149                    _ => return Vec::new(), // Unknown header type — skip
150                },
151                MessageType::Data => return Vec::new(),
152            };
153            let api_key = Literal::i16_unsuffixed(s.api_key.unwrap_or(0));
154            let name_lit = Literal::string(&s.name);
155            let max_v = s.valid_versions.max.min(MAX_VERSION_CAP);
156            (s.valid_versions.min..=max_v)
157                .map(|v| {
158                    let v_lit = Literal::i16_unsuffixed(v);
159                    let kind = kind_tokens.clone();
160                    quote! {
161                        Case { name: #name_lit, api_key: #api_key, version: #v_lit, kind: #kind },
162                    }
163                })
164                .collect()
165        })
166        .collect();
167
168    quote! {
169        pub const CASES: &[Case] = &[
170            #(#rows)*
171        ];
172    }
173}
174
175/// The schemas whose owned and borrowed codecs the table dispatches to.
176fn table_specs(specs: &[MessageSpec]) -> impl Iterator<Item = &MessageSpec> {
177    specs.iter().filter(|s| {
178        !s.valid_versions.is_empty()
179            && !s.internal
180            && matches!(
181                s.message_type,
182                MessageType::Request | MessageType::Response | MessageType::Header
183            )
184    })
185}
186
187/// Chunk `arms` into private helpers named `<prefix>_<n>`, each a `match` on
188/// `name` that returns `Ok(None)` for a name outside its chunk, plus the
189/// public `dispatcher` that tries each helper and fails with
190/// `TableError::UnknownMessage`.
191///
192/// `params` and `args` are the helpers' extra parameters and arguments after
193/// `name`, and `ty` is the value one arm produces.
194fn dispatch(
195    arms: &[TokenStream],
196    chunk: usize,
197    prefix: &str,
198    signature: (&TokenStream, &TokenStream, &TokenStream),
199    public: &TokenStream,
200) -> TokenStream {
201    let (params, args, ty) = signature;
202    let helpers: Vec<_> = arms
203        .chunks(chunk)
204        .enumerate()
205        .map(|(index, chunk)| {
206            let helper = format_ident!("{prefix}_{index}");
207            quote! {
208                fn #helper(name: &str, #params) -> Result<Option<#ty>, TableError> {
209                    Ok(Some(match name {
210                        #(#chunk)*
211                        _ => return Ok(None),
212                    }))
213                }
214            }
215        })
216        .collect();
217    let tries = (0..helpers.len()).map(|index| {
218        let helper = format_ident!("{prefix}_{index}");
219        quote! {
220            if let Some(value) = #helper(name, #args)? {
221                return Ok(value);
222            }
223        }
224    });
225    quote! {
226        #(#helpers)*
227        #public {
228            #(#tries)*
229            Err(TableError::UnknownMessage(name.to_owned()))
230        }
231    }
232}
233
234fn build_encode_default(specs: &[MessageSpec]) -> TokenStream {
235    let arms: Vec<TokenStream> = table_specs(specs)
236        .map(|s| {
237            let snake_ident = format_ident!("{}", name_conv::module_name(&s.name));
238            let type_ident = format_ident!("{}", name_conv::type_name(&s.name));
239            let name_lit = Literal::string(&s.name);
240            quote! {
241                #name_lit => encode_with(
242                    &krabka_protocol::owned::#snake_ident::#type_ident::default(),
243                    version,
244                )?,
245            }
246        })
247        .collect();
248    dispatch(
249        &arms,
250        8,
251        "encode_default",
252        (&quote!(version: i16), &quote!(version), &quote!(Vec<u8>)),
253        &quote! {
254            #[doc = " The default value of message `name`, encoded at `version`."]
255            #[doc = ""]
256            #[doc = " # Errors"]
257            #[doc = ""]
258            #[doc = " Fails when `name` is unknown or the encode fails."]
259            pub fn encode_default(name: &str, version: i16) -> Result<Vec<u8>, TableError>
260        },
261    )
262}
263
264fn build_default_json_for(specs: &[MessageSpec]) -> TokenStream {
265    let arms: Vec<TokenStream> = table_specs(specs)
266        .map(|s| {
267            let snake_ident = format_ident!("{}", name_conv::module_name(&s.name));
268            let name_lit = Literal::string(&s.name);
269            quote! {
270                #name_lit => krabka_protocol::owned::#snake_ident::default_json(version),
271            }
272        })
273        .collect();
274    dispatch(
275        &arms,
276        60,
277        "default_json_for",
278        (
279            &quote!(version: i16),
280            &quote!(version),
281            &quote!(::serde_json::Value),
282        ),
283        &quote! {
284            #[doc = " The JSON the oracle reads as the default value of message `name`."]
285            #[doc = ""]
286            #[doc = " # Errors"]
287            #[doc = ""]
288            #[doc = " Fails when `name` is unknown."]
289            pub fn default_json_for(name: &str, version: i16) -> Result<::serde_json::Value, TableError>
290        },
291    )
292}
293
294/// The tagged-field sweep. It covers requests, responses and data records,
295/// because the metadata records carry tagged fields too.
296fn build_tagged_fixture(specs: &[MessageSpec]) -> TokenStream {
297    let tagged: Vec<&MessageSpec> = specs
298        .iter()
299        .filter(|s| {
300            !s.valid_versions.is_empty() && !s.internal && tagged_fixture::has_tagged_field(s)
301        })
302        .collect();
303
304    let rows = tagged.iter().flat_map(|s| {
305        let name_lit = Literal::string(&s.name);
306        let max_v = s.valid_versions.max.min(MAX_VERSION_CAP);
307        (s.valid_versions.min..=max_v).map(move |v| {
308            let v_lit = Literal::i16_unsuffixed(v);
309            quote! { TaggedCase { name: #name_lit, version: #v_lit }, }
310        })
311    });
312
313    let encode_arms: Vec<TokenStream> = tagged
314        .iter()
315        .map(|s| {
316            let snake_ident = format_ident!("{}", name_conv::module_name(&s.name));
317            let type_ident = format_ident!("{}", name_conv::type_name(&s.name));
318            let name_lit = Literal::string(&s.name);
319            quote! {
320                #name_lit => {
321                    let owned = krabka_protocol::owned::#snake_ident::tagged_fixture();
322                    let max = krabka_protocol::owned::#snake_ident::MAX_VERSION;
323                    let full = encode_with(&owned, max)?;
324                    let mut cur = full.as_slice();
325                    let borrowed =
326                        <krabka_protocol::borrowed::#snake_ident::#type_ident as krabka_protocol::DecodeBorrow>::decode_borrow(&mut cur, max)?;
327                    expect_consumed(name, cur)?;
328                    (encode_with(&owned, version)?, encode_with(&borrowed, version)?)
329                }
330            }
331        })
332        .collect();
333    let encode = dispatch(
334        &encode_arms,
335        6,
336        "encode_tagged_fixture",
337        (
338            &quote!(version: i16),
339            &quote!(version),
340            &quote!(EncodedPair),
341        ),
342        &quote! {
343            #[doc = " The owned `tagged_fixture()` encoded at `version`, and the borrowed flavor of"]
344            #[doc = " the same message encoded at `version`."]
345            #[doc = ""]
346            #[doc = " The borrowed message is decoded from the owned bytes at the highest version,"]
347            #[doc = " where every tag is in range."]
348            #[doc = ""]
349            #[doc = " # Errors"]
350            #[doc = ""]
351            #[doc = " Fails when `name` has no tagged field, or an encode or decode fails."]
352            pub fn encode_tagged_fixture(name: &str, version: i16) -> Result<EncodedPair, TableError>
353        },
354    );
355
356    let json_arms = tagged.iter().map(|s| {
357        let snake_ident = format_ident!("{}", name_conv::module_name(&s.name));
358        let name_lit = Literal::string(&s.name);
359        quote! {
360            #name_lit => Ok(krabka_protocol::owned::#snake_ident::tagged_fixture_json(version)),
361        }
362    });
363
364    quote! {
365        #[derive(Debug, Clone, Copy)]
366        pub struct TaggedCase {
367            pub name: &'static str,
368            pub version: i16,
369        }
370
371        pub const TAGGED_CASES: &[TaggedCase] = &[
372            #(#rows)*
373        ];
374
375        #encode
376
377        #[doc = " The JSON form of the tagged fixture of message `name` at `version`."]
378        #[doc = ""]
379        #[doc = " # Errors"]
380        #[doc = ""]
381        #[doc = " Fails when `name` has no tagged field."]
382        pub fn tagged_fixture_json_for(name: &str, version: i16) -> Result<::serde_json::Value, TableError> {
383            match name {
384                #(#json_arms)*
385                _ => Err(TableError::UnknownMessage(name.to_owned())),
386            }
387        }
388    }
389}
390
391fn build_roundtrip(specs: &[MessageSpec]) -> TokenStream {
392    let arms: Vec<TokenStream> = table_specs(specs)
393        .map(|s| {
394            let snake_ident = format_ident!("{}", name_conv::module_name(&s.name));
395            let type_ident = format_ident!("{}", name_conv::type_name(&s.name));
396            let name_lit = Literal::string(&s.name);
397            quote! {
398                #name_lit => {
399                    let mut cur = bytes;
400                    let msg = <krabka_protocol::owned::#snake_ident::#type_ident as krabka_protocol::Decode>::decode(&mut cur, version)?;
401                    expect_consumed(name, cur)?;
402                    encode_with(&msg, version)?
403                }
404            }
405        })
406        .collect();
407    dispatch(
408        &arms,
409        6,
410        "roundtrip",
411        (
412            &quote!(version: i16, bytes: &[u8]),
413            &quote!(version, bytes),
414            &quote!(Vec<u8>),
415        ),
416        &quote! {
417            #[doc = " Decode `bytes` as message `name` at `version` and encode it again."]
418            #[doc = ""]
419            #[doc = " # Errors"]
420            #[doc = ""]
421            #[doc = " Fails when `name` is unknown, the decode or encode fails, or the decode"]
422            #[doc = " leaves bytes behind."]
423            pub fn roundtrip(name: &str, version: i16, bytes: &[u8]) -> Result<Vec<u8>, TableError>
424        },
425    )
426}
427
428fn build_header_versions(specs: &[MessageSpec]) -> TokenStream {
429    let flex_entries = |message_type: MessageType| -> Vec<TokenStream> {
430        specs
431            .iter()
432            .filter(|s| {
433                !s.valid_versions.is_empty()
434                    && !s.internal
435                    && s.message_type == message_type
436                    && s.name != "ApiVersionsResponse"
437            })
438            .map(|s| {
439                let fm_lit = Literal::i16_unsuffixed(super::owned::flex_min(s));
440                let name_lit = Literal::string(&s.name);
441                quote! { (#name_lit, #fm_lit), }
442            })
443            .collect()
444    };
445    let req_entries = flex_entries(MessageType::Request);
446    let resp_entries = flex_entries(MessageType::Response);
447
448    quote! {
449        const REQUEST_FLEX_MIN: &[(&str, i16)] = &[#(#req_entries)*];
450        const RESPONSE_FLEX_MIN: &[(&str, i16)] = &[#(#resp_entries)*];
451
452        fn header_flex_min(table: &[(&str, i16)], name: &str) -> Result<i16, TableError> {
453            table
454                .iter()
455                .find_map(|(candidate, minimum)| (*candidate == name).then_some(*minimum))
456                .ok_or_else(|| TableError::UnknownMessage(name.to_owned()))
457        }
458
459        #[doc = " The request header version that request `name` uses at `version`."]
460        #[doc = ""]
461        #[doc = " # Errors"]
462        #[doc = ""]
463        #[doc = " Fails when `name` is not a known request."]
464        pub fn request_header_version(name: &str, version: i16) -> Result<i16, TableError> {
465            if name == "ControlledShutdownRequest" && version == 0 {
466                return Ok(0);
467            }
468            let minimum = header_flex_min(REQUEST_FLEX_MIN, name)?;
469            Ok(if minimum == i16::MIN || version >= minimum { 2 } else { 1 })
470        }
471
472        #[doc = " The response header version that response `name` uses at `version`."]
473        #[doc = ""]
474        #[doc = " # Errors"]
475        #[doc = ""]
476        #[doc = " Fails when `name` is not a known response."]
477        pub fn response_header_version(name: &str, version: i16) -> Result<i16, TableError> {
478            if name == "ApiVersionsResponse" {
479                return Ok(0);
480            }
481            let minimum = header_flex_min(RESPONSE_FLEX_MIN, name)?;
482            Ok(i16::from(minimum != i16::MIN && version >= minimum))
483        }
484    }
485}
486
487fn build_strip_frame_header() -> TokenStream {
488    quote! {
489        #[doc = " Decode and discard the request/response header from a full frame body,"]
490        #[doc = " returning the remaining message bytes."]
491        #[doc = ""]
492        #[doc = " # Errors"]
493        #[doc = ""]
494        #[doc = " Fails when the frame header is malformed or `name` is unknown."]
495        pub fn strip_frame_header(
496            name: &str,
497            version: i16,
498            is_request: bool,
499            frame: &[u8],
500        ) -> Result<Vec<u8>, TableError> {
501            let mut cur = frame;
502            if is_request {
503                let hv = request_header_version(name, version)?;
504                <krabka_protocol::owned::request_header::RequestHeader as krabka_protocol::Decode>::decode(&mut cur, hv)?;
505            } else {
506                let hv = response_header_version(name, version)?;
507                <krabka_protocol::owned::response_header::ResponseHeader as krabka_protocol::Decode>::decode(&mut cur, hv)?;
508            }
509            Ok(cur.to_vec())
510        }
511    }
512}
krabka_protocol_codegen/emit/differential_table.rs

krabka_protocol_codegen/emit/
differential_table.rs
