Skip to main content

olai_codegen/
postprocess.rs

1//! Rewrites of the buffa model output, whose JSON choices are fixed in the buf plugin:
2//!
3//! - **Field names.** [`JsonFieldNames::Proto`] serializes the snake_case proto names instead of
4//!   lowerCamelCase (both are still read). An explicit `json_name` is kept.
5//! - **64-bit integers.** `int64_module` points buffa's string helpers at the `json_int64`
6//!   module, which writes numbers.
7//! - **Free-form JSON.** `json_value_module` makes `google.protobuf.Value`/`Struct`/`ListValue`
8//!   fields write whole numbers as integers (`1`, not `1.0`).
9//! - **Field options** ([`field_options`]): explicit `json_name`s.
10//!
11//! The rewrites are idempotent and skip the files trestle writes itself.
12
13use std::collections::BTreeMap;
14use std::path::{Path, PathBuf};
15use std::sync::LazyLock;
16
17use regex::Regex;
18
19use crate::error::Result;
20use crate::utils::last_segment;
21use crate::utils::strings::lower_camel;
22
23/// The entries of the models directory that trestle writes itself.
24const TRESTLE_FILES: &[&str] = &[
25    "mod.rs",
26    "labels.rs",
27    "pyo3_impls.rs",
28    crate::codegen::EXTRACTORS_SUBDIR,
29];
30
31/// How the JSON field names of one proto package are serialized.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
33pub enum JsonFieldNames {
34    /// lowerCamelCase field names, buffa's proto3-JSON default.
35    #[default]
36    Json,
37    /// snake_case (proto) field names.
38    Proto,
39}
40
41/// The rewrites to apply to the buffa model output.
42#[derive(Debug, Clone, Default)]
43pub struct ModelRewrites {
44    /// Field-name style of every package not in [`packages`](Self::packages).
45    pub field_names: JsonFieldNames,
46    /// Field-name style per proto package (`my.pkg.v1`), overriding
47    /// [`field_names`](Self::field_names). A package file is `<package>.rs`.
48    pub packages: BTreeMap<String, JsonFieldNames>,
49    /// Path of the `json_int64` helper module (e.g. `crate::models::json_int64`); when set,
50    /// buffa's 64-bit integer helpers are pointed at it.
51    pub int64_module: Option<String>,
52    /// Path of the `json_value` helper module (e.g. `crate::models::json_value`); when set, the
53    /// free-form JSON fields serialize through it.
54    pub json_value_module: Option<String>,
55    /// The fields with options that change their JSON, by proto package; see [`field_options`].
56    pub fields: BTreeMap<String, PackageFields>,
57}
58
59/// The fields of one proto package with JSON options, keyed `Message.field`.
60#[derive(Debug, Clone, Default, PartialEq, Eq)]
61pub struct PackageFields {
62    /// Explicit `json_name`s, by field.
63    pub json_names: BTreeMap<String, String>,
64}
65
66/// The fields of `metadata` with JSON options, by package, for [`ModelRewrites::fields`].
67pub fn field_options(
68    metadata: &crate::parsing::CodeGenMetadata,
69) -> BTreeMap<String, PackageFields> {
70    let mut packages: BTreeMap<String, PackageFields> = BTreeMap::new();
71    for message in metadata.messages.values() {
72        let name = message.name.trim_start_matches('.');
73        let package = name
74            .split('.')
75            .take_while(|s| s.starts_with(|c: char| c.is_ascii_lowercase()))
76            .collect::<Vec<_>>()
77            .join(".");
78        let simple = last_segment(name);
79        for field in &message.fields {
80            let key = format!("{simple}.{}", field.name);
81            if let Some(json) = &field.json_name {
82                packages
83                    .entry(package.clone())
84                    .or_default()
85                    .json_names
86                    .insert(key, json.clone());
87            }
88        }
89    }
90    packages
91}
92
93/// The entry of `by_package` for the model file `file_name`: `<package>.rs` (the buf plugin) or
94/// `<package>.<proto file>.rs` (buffa-build).
95fn package_entry<'a, T>(by_package: &'a BTreeMap<String, T>, file_name: &str) -> Option<&'a T> {
96    let stem = file_name.strip_suffix(".rs")?;
97    let parent = stem.rsplit_once('.').map(|(package, _)| package);
98    by_package
99        .get(stem)
100        .or_else(|| parent.and_then(|package| by_package.get(package)))
101}
102
103impl ModelRewrites {
104    /// Whether any rewrite is enabled.
105    pub fn is_empty(&self) -> bool {
106        self.field_names == JsonFieldNames::Json
107            && self.packages.values().all(|n| *n == JsonFieldNames::Json)
108            && self.int64_module.is_none()
109            && self.json_value_module.is_none()
110            && self.fields.is_empty()
111    }
112
113    /// The field-name style of the model file `file_name` (see [`package_entry`]).
114    fn field_names_for(&self, file_name: &str) -> JsonFieldNames {
115        package_entry(&self.packages, file_name)
116            .copied()
117            .unwrap_or(self.field_names)
118    }
119
120    /// The [`PackageFields`] of the model file `file_name` (see [`package_entry`]).
121    fn fields_for(&self, file_name: &str) -> Option<&PackageFields> {
122        package_entry(&self.fields, file_name)
123    }
124
125    /// Rewrite the source of the model file `file_name`.
126    pub fn rewrite(&self, file_name: &str, src: &str) -> String {
127        let mut out = src.to_string();
128        if self.field_names_for(file_name) == JsonFieldNames::Proto {
129            out = swap_serde_renames(&out);
130            out = snakeify_serialize_entries(&out);
131        }
132        if let Some(fields) = self.fields_for(file_name) {
133            out = apply_field_options(&out, fields);
134        }
135        if let Some(module) = &self.int64_module {
136            out = int64_helpers_to(&out, module);
137            out = int64_containers_to(&out, module);
138        }
139        if let Some(module) = &self.json_value_module {
140            out = json_value_serializers(&out, module);
141        }
142        out
143    }
144
145    /// Rewrite the model files under `dir` in place. Returns the number of files changed.
146    pub fn rewrite_dir(&self, dir: &Path) -> Result<usize> {
147        if self.is_empty() || !dir.exists() {
148            return Ok(0);
149        }
150        let mut changed = 0;
151        for path in model_files(dir)? {
152            let src = std::fs::read_to_string(&path)?;
153            let file_name = path
154                .file_name()
155                .and_then(|n| n.to_str())
156                .unwrap_or_default();
157            let rewritten = self.rewrite(file_name, &src);
158            if rewritten != src {
159                std::fs::write(&path, rewritten)?;
160                changed += 1;
161            }
162        }
163        Ok(changed)
164    }
165}
166
167/// Every `.rs` file under `dir`, sorted, without the [`TRESTLE_FILES`].
168fn model_files(dir: &Path) -> Result<Vec<PathBuf>> {
169    let mut out = Vec::new();
170    let mut stack = vec![dir.to_path_buf()];
171    while let Some(d) = stack.pop() {
172        for entry in std::fs::read_dir(&d)? {
173            let path = entry?.path();
174            let trestle_file = d == dir
175                && path
176                    .file_name()
177                    .and_then(|n| n.to_str())
178                    .is_some_and(|n| TRESTLE_FILES.contains(&n));
179            if trestle_file {
180                continue;
181            }
182            if path.is_dir() {
183                stack.push(path);
184            } else if path.extension().is_some_and(|e| e == "rs") {
185                out.push(path);
186            }
187        }
188    }
189    out.sort();
190    Ok(out)
191}
192
193/// Swap `rename`/`alias` in each `#[serde(...)]` whose `rename` is the default proto3-JSON name
194/// of its `alias`, so serialization uses the proto name. An explicit `json_name` is kept.
195fn swap_serde_renames(src: &str) -> String {
196    // One `#[serde( ... )]` attribute; the body never contains `)]`.
197    static SERDE_ATTR: LazyLock<Regex> =
198        LazyLock::new(|| Regex::new(r"#\[serde\((?P<body>[^)]*)\)\]").expect("valid regex"));
199    static RENAME: LazyLock<Regex> =
200        LazyLock::new(|| Regex::new(r#"rename = "(?P<v>[^"]*)""#).expect("valid regex"));
201    static ALIAS: LazyLock<Regex> =
202        LazyLock::new(|| Regex::new(r#"alias = "(?P<v>[^"]*)""#).expect("valid regex"));
203
204    SERDE_ATTR
205        .replace_all(src, |caps: &regex::Captures| {
206            let whole = &caps[0];
207            let body = &caps["body"];
208            let (Some(rn), Some(al)) = (RENAME.captures(body), ALIAS.captures(body)) else {
209                return whole.to_string();
210            };
211            let rename_val = rn["v"].to_string();
212            let alias_val = al["v"].to_string();
213            // Swap only the default JSON name; this also keeps a second run a no-op.
214            if rename_val == alias_val || rename_val != lower_camel(&alias_val) {
215                return whole.to_string();
216            }
217            let body = ALIAS.replace(body, format!(r#"alias = "{rename_val}""#));
218            let body = RENAME.replace(&body, format!(r#"rename = "{alias_val}""#));
219            format!("#[serde({body})]")
220        })
221        .into_owned()
222}
223
224/// Convert the default `serialize_entry("camelCase", …)` names of the `*View` `Serialize` impls
225/// to snake_case.
226fn snakeify_serialize_entries(src: &str) -> String {
227    static ENTRY: LazyLock<Regex> =
228        LazyLock::new(|| Regex::new(r#"serialize_entry\("(?P<name>[^"]*)""#).expect("valid regex"));
229    ENTRY
230        .replace_all(src, |caps: &regex::Captures| {
231            let name = &caps["name"];
232            let snake = camel_to_snake(name);
233            let is_default_json_name =
234                name.starts_with(|c: char| c.is_ascii_lowercase()) && lower_camel(&snake) == name;
235            if is_default_json_name {
236                format!(r#"serialize_entry("{snake}""#)
237            } else {
238                caps[0].to_string()
239            }
240        })
241        .into_owned()
242}
243
244/// Point buffa's `int64`, `opt_int64`, `uint64` and `opt_uint64` helpers at `module`.
245fn int64_helpers_to(src: &str, module: &str) -> String {
246    static HELPER: LazyLock<Regex> = LazyLock::new(|| {
247        Regex::new(r"::buffa::json_helpers::(?P<name>(?:opt_)?u?int64)\b").expect("valid regex")
248    });
249    HELPER
250        .replace_all(src, |caps: &regex::Captures| {
251            format!("{module}::{}", &caps["name"])
252        })
253        .into_owned()
254}
255
256/// Apply the field options: an explicit `json_name` becomes the `rename` (restoring what
257/// [`swap_serde_renames`] may have swapped).
258fn apply_field_options(src: &str, fields: &PackageFields) -> String {
259    static STRUCT: LazyLock<Regex> =
260        LazyLock::new(|| Regex::new(r"^\s*pub struct (\w+)").expect("valid regex"));
261    static FIELD: LazyLock<Regex> =
262        LazyLock::new(|| Regex::new(r"^\s*pub (?:r#)?(\w+):").expect("valid regex"));
263    static RENAME: LazyLock<Regex> =
264        LazyLock::new(|| Regex::new(r#"rename = "[^"]*""#).expect("valid regex"));
265    static ALIAS: LazyLock<Regex> =
266        LazyLock::new(|| Regex::new(r#"alias = "[^"]*""#).expect("valid regex"));
267
268    let lines: Vec<&str> = src.lines().collect();
269    let mut out: Vec<String> = Vec::with_capacity(lines.len());
270    let mut message = None;
271    for (i, line) in lines.iter().enumerate() {
272        if let Some(caps) = STRUCT.captures(line) {
273            message = Some(caps[1].to_string());
274        } else if let (Some(message), Some(caps)) = (&message, FIELD.captures(line)) {
275            let name = &caps[1];
276            let key = format!("{message}.{name}");
277            if let Some(json) = fields.json_names.get(&key) {
278                let attrs = attribute_lines(&lines, i);
279                let text = lines[attrs.clone()].join("\n");
280                let text = RENAME
281                    .replace(&text, format!(r#"rename = "{json}""#))
282                    .into_owned();
283                let text = ALIAS
284                    .replace(&text, format!(r#"alias = "{name}""#))
285                    .into_owned();
286                out.truncate(out.len() - attrs.len());
287                out.extend(text.lines().map(str::to_string));
288            }
289        }
290        out.push(line.to_string());
291    }
292    join_lines_like(&out, src)
293}
294
295/// `lines` joined by newlines, with a trailing newline when `src` has one.
296fn join_lines_like(lines: &[String], src: &str) -> String {
297    let mut out = lines.join("\n");
298    if src.ends_with('\n') {
299        out.push('\n');
300    }
301    out
302}
303
304/// Point repeated `int64`/`uint64` fields and maps of them at the `repeated_int64` and
305/// `map_int64` helpers of `module`, instead of buffa's `proto_seq`/`proto_map`.
306fn int64_containers_to(src: &str, module: &str) -> String {
307    let lines: Vec<&str> = src.lines().collect();
308    let mut out: Vec<String> = lines.iter().map(|l| l.to_string()).collect();
309    for i in 0..lines.len() {
310        let trimmed = lines[i].trim_start();
311        if !trimmed.starts_with("pub ") || !trimmed.contains(':') && !trimmed.ends_with('<') {
312            continue;
313        }
314        let (decl, _) = declaration(&lines, i);
315        let int64 = |ty: &str| ty == "i64" || ty == "u64";
316        let (from, to) = if decl
317            .split_once("Vec<")
318            .is_some_and(|(_, rest)| int64(rest.split('>').next().unwrap_or_default()))
319        {
320            ("::buffa::json_helpers::proto_seq", "repeated_int64")
321        } else if decl
322            .split_once("HashMap<")
323            .and_then(|(_, rest)| rest.rsplit_once('>'))
324            .and_then(|(args, _)| args.rsplit_once(','))
325            .is_some_and(|(_, value)| int64(value.trim()))
326        {
327            ("::buffa::json_helpers::proto_map", "map_int64")
328        } else {
329            continue;
330        };
331        for j in attribute_lines(&lines, i) {
332            out[j] = out[j].replace(from, &format!("{module}::{to}"));
333        }
334    }
335    join_lines_like(&out, src)
336}
337
338/// The field declaration at line `start`, whitespace collapsed (`pub a: Vec<i64>,`), and the
339/// line it ends on.
340fn declaration(lines: &[&str], start: usize) -> (String, usize) {
341    let mut end = start;
342    let mut decl = String::new();
343    loop {
344        decl.push_str(lines[end].trim());
345        decl.push(' ');
346        let depth = decl.matches('<').count() as isize - decl.matches('>').count() as isize;
347        if (depth <= 0 && lines[end].trim_end().ends_with(',')) || end + 1 == lines.len() {
348            break;
349        }
350        end += 1;
351    }
352    let decl = decl
353        .replace(", >", ">")
354        .replace("< ", "<")
355        .replace(" >", ">");
356    (decl, end)
357}
358
359/// Give each free-form JSON field (singular, repeated or map value) a `serialize_with` of
360/// `module`, unless it has one already.
361fn json_value_serializers(src: &str, module: &str) -> String {
362    const JSON: &str = "::buffa_types::google::protobuf::";
363    let lines: Vec<&str> = src.lines().collect();
364    let mut out = String::with_capacity(src.len() + 256);
365    let mut i = 0;
366    while i < lines.len() {
367        let line = lines[i];
368        let trimmed = line.trim_start();
369        let is_field = trimmed.starts_with("pub ");
370        let json_field_on_one_line = is_field && trimmed.contains(": ") && line.contains(JSON);
371        // A type that rustfmt broke onto the next lines.
372        let field_type_continues = is_field && trimmed.ends_with('<');
373        if json_field_on_one_line || field_type_continues {
374            let (decl, end) = declaration(&lines, i);
375            let kind = if !decl.contains(JSON) {
376                None
377            } else if decl.contains(": ::buffa::MessageField<") {
378                Some("message_field")
379            } else if decl.contains(": ::buffa::alloc::vec::Vec<") {
380                Some("repeated")
381            } else if decl.contains("HashMap<") {
382                Some("map")
383            } else {
384                None
385            };
386            let attrs = field_attributes(&lines, i);
387            if let Some(kind) = kind
388                && !has_serde_key(&attrs, "serialize_with")
389                && !has_serde_key(&attrs, "with")
390            {
391                let indent = &line[..line.len() - trimmed.len()];
392                out.push_str(&format!(
393                    "{indent}#[serde(serialize_with = \"{module}::{kind}\")]\n"
394                ));
395            }
396            for l in &lines[i..=end] {
397                out.push_str(l);
398                out.push('\n');
399            }
400            i = end + 1;
401            continue;
402        }
403        out.push_str(line);
404        out.push('\n');
405        i += 1;
406    }
407    if !src.ends_with('\n') {
408        out.pop();
409    }
410    out
411}
412
413/// Whether the serde attributes `attrs` set `key` as a whole word (`deserialize_with` does not
414/// set `with`).
415fn has_serde_key(attrs: &str, key: &str) -> bool {
416    attrs.match_indices(key).any(|(at, _)| {
417        let before = attrs[..at].chars().next_back();
418        let after = attrs[at + key.len()..].trim_start();
419        !before.is_some_and(|c| c.is_ascii_alphanumeric() || c == '_') && after.starts_with('=')
420    })
421}
422
423/// The attributes above the field declaration at line `field`, as one string.
424fn field_attributes(lines: &[&str], field: usize) -> String {
425    attribute_lines(lines, field)
426        .map(|j| lines[j].trim())
427        .collect()
428}
429
430/// The lines of the attributes and doc comments above the field declaration at line `field`.
431fn attribute_lines(lines: &[&str], field: usize) -> std::ops::Range<usize> {
432    let mut j = field;
433    while j > 0 {
434        let prev = lines[j - 1].trim();
435        if prev.starts_with("//") || prev.starts_with("#[") && prev.ends_with(']') {
436            j -= 1;
437        } else if prev.ends_with(")]") {
438            // A multi-line attribute: back up to its `#[`.
439            let mut k = j - 1;
440            while k > 0 && !lines[k].trim_start().starts_with("#[") {
441                k -= 1;
442            }
443            j = k;
444        } else {
445            break;
446        }
447    }
448    j..field
449}
450
451/// Invert the proto3 JSON mapping: each uppercase letter becomes `_` + its lowercase.
452fn camel_to_snake(name: &str) -> String {
453    let mut out = String::with_capacity(name.len() + 4);
454    for ch in name.chars() {
455        if ch.is_ascii_uppercase() {
456            out.push('_');
457            out.push(ch.to_ascii_lowercase());
458        } else {
459            out.push(ch);
460        }
461    }
462    out
463}
464
465/// Put blank lines between the items of a pretty-printed Rust file (prettyplease prints none;
466/// rustfmt keeps them but never adds them): before an item that follows a `}` or a `use` group
467/// at the same indentation, and after the file's inner attributes. Idempotent.
468pub fn space_items(src: &str) -> String {
469    let lines: Vec<&str> = src.lines().collect();
470    let mut out = String::with_capacity(src.len() + src.len() / 16);
471    for (i, line) in lines.iter().enumerate() {
472        out.push_str(line);
473        out.push('\n');
474        if lines
475            .get(i + 1)
476            .is_some_and(|next| blank_between(line, next))
477        {
478            out.push('\n');
479        }
480    }
481    out
482}
483
484/// Whether a blank line goes between the lines `line` and `next`.
485fn blank_between(line: &str, next: &str) -> bool {
486    let indent = |l: &str| l.len() - l.trim_start().len();
487    if line.trim().is_empty() || next.trim().is_empty() || indent(line) != indent(next) {
488        return false;
489    }
490    let (line, next) = (line.trim(), next.trim());
491    let is_use = |l: &str| {
492        ["use ", "pub use ", "pub(crate) use ", "pub(super) use "]
493            .iter()
494            .any(|p| l.starts_with(p))
495    };
496    let is_inner = |l: &str| l.starts_with("//!") || l.starts_with("#![");
497    if is_inner(line) {
498        return !is_inner(next);
499    }
500    if !starts_item(next) {
501        return false;
502    }
503    match line {
504        "}" => true,
505        // A multi-line `use` ends with `};`: the next `use` stays in the group.
506        "};" => !is_use(next),
507        _ => is_use(line) && line.ends_with(';') && !is_use(next),
508    }
509}
510
511/// Whether `line` (trimmed) starts an item, or the attributes or docs of one.
512fn starts_item(line: &str) -> bool {
513    const STARTS: &[&str] = &[
514        "#[",
515        "///",
516        "/**",
517        "pub ",
518        "pub(",
519        "fn ",
520        "async fn ",
521        "const fn ",
522        "unsafe ",
523        "extern ",
524        "impl ",
525        "impl<",
526        "struct ",
527        "enum ",
528        "trait ",
529        "mod ",
530        "type ",
531        "const ",
532        "static ",
533        "use ",
534        "macro_rules!",
535        "include!",
536    ];
537    STARTS.iter().any(|p| line.starts_with(p))
538}
539
540#[cfg(test)]
541mod tests {
542    use super::*;
543
544    fn proto() -> ModelRewrites {
545        ModelRewrites {
546            field_names: JsonFieldNames::Proto,
547            ..Default::default()
548        }
549    }
550
551    #[test]
552    fn camel_to_snake_conversions() {
553        assert_eq!(camel_to_snake("displayName"), "display_name");
554        assert_eq!(camel_to_snake("nextPageToken"), "next_page_token");
555        assert_eq!(camel_to_snake("name"), "name");
556    }
557
558    #[test]
559    fn snakeifies_view_serialize_entry_literal() {
560        let src = r#"                __map.serialize_entry("displayName", self.display_name)?;"#;
561        let out = proto().rewrite("demo.v1.rs", src);
562        assert!(out.contains(r#"serialize_entry("display_name", self.display_name)"#));
563    }
564
565    #[test]
566    fn view_serialize_entry_explicit_names_untouched() {
567        for src in [
568            r#"__map.serialize_entry("name", self.name)?;"#,
569            r#"__map.serialize_entry("metadata-location", self.metadata_location)?;"#,
570            r#"__map.serialize_entry("Resources", self.resources)?;"#,
571        ] {
572            assert_eq!(proto().rewrite("demo.v1.rs", src), src);
573        }
574    }
575
576    #[test]
577    fn swaps_multiword_rename_alias_pair() {
578        let src = r#"    #[serde(
579        rename = "storageRoot",
580        alias = "storage_root",
581        skip_serializing_if = "::core::option::Option::is_none"
582    )]
583    pub storage_root: Option<String>,"#;
584        let out = swap_serde_renames(src);
585        assert!(out.contains(r#"rename = "storage_root""#));
586        assert!(out.contains(r#"alias = "storageRoot""#));
587        assert!(out.contains(r#"skip_serializing_if = "::core::option::Option::is_none""#));
588    }
589
590    #[test]
591    fn single_word_rename_only_is_untouched() {
592        let src = r#"    #[serde(rename = "name", skip_serializing_if = "Option::is_none")]"#;
593        assert_eq!(swap_serde_renames(src), src);
594    }
595
596    #[test]
597    fn explicit_json_name_is_kept() {
598        // `json_name = "metadata-location"`: not the default JSON name of `metadata_location`.
599        let src = r#"#[serde(rename = "metadata-location", alias = "metadata_location")]"#;
600        assert_eq!(swap_serde_renames(src), src);
601        let src = r#"#[serde(rename = "Resources", alias = "resources")]"#;
602        assert_eq!(swap_serde_renames(src), src);
603    }
604
605    #[test]
606    fn enum_and_scalar_with_helpers_preserved() {
607        let src = r#"    #[serde(
608        rename = "catalogType",
609        alias = "catalog_type",
610        with = "::buffa::json_helpers::opt_enum",
611        skip_serializing_if = "::core::option::Option::is_none"
612    )]"#;
613        let out = swap_serde_renames(src);
614        assert!(out.contains(r#"rename = "catalog_type""#));
615        assert!(out.contains(r#"alias = "catalogType""#));
616        assert!(out.contains(r#"with = "::buffa::json_helpers::opt_enum""#));
617    }
618
619    #[test]
620    fn swap_is_idempotent_on_already_swapped_output() {
621        let camel = r#"    #[serde(rename = "storageRoot", alias = "storage_root", skip_serializing_if = "x")]"#;
622        let once = swap_serde_renames(camel);
623        assert!(once.contains(r#"rename = "storage_root", alias = "storageRoot""#));
624        assert_eq!(swap_serde_renames(&once), once);
625    }
626
627    #[test]
628    fn non_field_serde_attr_is_untouched() {
629        assert_eq!(swap_serde_renames("#[serde(default)]"), "#[serde(default)]");
630    }
631
632    #[test]
633    fn package_override_keeps_camel_case() {
634        let mut rewrites = proto();
635        rewrites
636            .packages
637            .insert("demo.scim.v2".into(), JsonFieldNames::Json);
638        let src = r#"#[serde(rename = "userName", alias = "user_name")]"#;
639        assert_eq!(rewrites.rewrite("demo.scim.v2.rs", src), src);
640        assert_eq!(rewrites.rewrite("demo.scim.v2.users.rs", src), src);
641        assert!(
642            rewrites
643                .rewrite("demo.v1.rs", src)
644                .contains(r#"rename = "user_name""#)
645        );
646    }
647
648    #[test]
649    fn field_options_set_json_names() {
650        let src = r#"pub struct Spec {
651    #[serde(rename = "spec_id", alias = "specId")]
652    pub spec_id: i32,
653}
654pub struct Other {
655    #[serde(rename = "spec_id", alias = "specId")]
656    pub spec_id: i32,
657}
658"#;
659        let fields = PackageFields {
660            json_names: [("Spec.spec_id".to_string(), "specId".to_string())].into(),
661        };
662        let out = apply_field_options(src, &fields);
663        assert!(
664            out.contains(r#"rename = "specId", alias = "spec_id""#),
665            "{out}"
666        );
667        // Only `Spec.spec_id` is renamed.
668        assert_eq!(out.matches(r#"rename = "specId""#).count(), 1, "{out}");
669    }
670
671    #[test]
672    fn int64_containers_write_numbers() {
673        let src = r#"    #[serde(rename = "sizes", with = "::buffa::json_helpers::proto_seq")]
674    pub sizes: ::buffa::alloc::vec::Vec<i64>,
675    #[serde(
676        rename = "counts",
677        with = "::buffa::json_helpers::proto_map"
678    )]
679    pub counts: ::buffa::__private::HashMap<
680        ::buffa::alloc::string::String,
681        i64,
682    >,
683    #[serde(rename = "ratios", with = "::buffa::json_helpers::proto_seq")]
684    pub ratios: ::buffa::alloc::vec::Vec<f64>,
685"#;
686        let out = int64_containers_to(src, "crate::models::json_int64");
687        assert!(out.contains(r#"with = "crate::models::json_int64::repeated_int64""#));
688        assert!(out.contains(r#"with = "crate::models::json_int64::map_int64""#));
689        assert_eq!(
690            out.matches("::buffa::json_helpers::proto_seq").count(),
691            1,
692            "{out}"
693        );
694        assert!(out.ends_with(",\n"));
695    }
696
697    #[test]
698    fn int64_helpers_point_at_module() {
699        let rewrites = ModelRewrites {
700            int64_module: Some("crate::models::json_int64".into()),
701            ..Default::default()
702        };
703        let src = r#"with = "::buffa::json_helpers::opt_int64",
704            ::buffa::json_helpers::int64::serialize(&self.0, __s)
705            with = "::buffa::json_helpers::uint64",
706            with = "::buffa::json_helpers::int32","#;
707        let out = rewrites.rewrite("demo.v1.rs", src);
708        assert!(out.contains(r#"with = "crate::models::json_int64::opt_int64""#));
709        assert!(out.contains("crate::models::json_int64::int64::serialize(&self.0, __s)"));
710        assert!(out.contains(r#"with = "crate::models::json_int64::uint64""#));
711        assert!(out.contains(r#"with = "::buffa::json_helpers::int32""#));
712        // Idempotent.
713        assert_eq!(rewrites.rewrite("demo.v1.rs", &out), out);
714    }
715
716    #[test]
717    fn rewrite_dir_leaves_trestle_files_alone() {
718        let rewrites = ModelRewrites {
719            int64_module: Some("crate::models::json_int64".into()),
720            ..Default::default()
721        };
722        let dir = tempfile::tempdir().unwrap();
723        let src = r#"with = "::buffa::json_helpers::int64","#;
724        std::fs::create_dir(dir.path().join("extractors")).unwrap();
725        for file in ["demo.v1.rs", "mod.rs", "extractors/mod.rs"] {
726            std::fs::write(dir.path().join(file), src).unwrap();
727        }
728        assert_eq!(rewrites.rewrite_dir(dir.path()).unwrap(), 1);
729        let read = |file: &str| std::fs::read_to_string(dir.path().join(file)).unwrap();
730        assert!(read("demo.v1.rs").contains("crate::models::json_int64::int64"));
731        assert_eq!(read("mod.rs"), src);
732        assert_eq!(read("extractors/mod.rs"), src);
733    }
734
735    #[test]
736    fn space_items_separates_items_and_functions() {
737        let src = "//! Docs.
738use a::B;
739use c::D;
740/// A thing.
741pub struct S {
742    x: u8,
743}
744impl S {
745    fn a(&self) {
746        if true {
747            f();
748        }
749        g();
750    }
751    #[inline]
752    fn b(&self) -> u8 {
753        match self.x {
754            0 => {
755                1
756            }
757            _ => 2,
758        }
759    }
760}
761";
762        let spaced = space_items(src);
763        assert_eq!(
764            spaced,
765            "//! Docs.
766
767use a::B;
768use c::D;
769
770/// A thing.
771pub struct S {
772    x: u8,
773}
774
775impl S {
776    fn a(&self) {
777        if true {
778            f();
779        }
780        g();
781    }
782
783    #[inline]
784    fn b(&self) -> u8 {
785        match self.x {
786            0 => {
787                1
788            }
789            _ => 2,
790        }
791    }
792}
793"
794        );
795        assert_eq!(space_items(&spaced), spaced);
796    }
797
798    #[test]
799    fn json_value_fields_serialize_through_the_module() {
800        let rewrites = ModelRewrites {
801            json_value_module: Some("crate::models::json_value".into()),
802            ..Default::default()
803        };
804        let src = r#"    #[serde(rename = "attributes")]
805    pub attributes: ::buffa::MessageField<::buffa_types::google::protobuf::Struct>,
806    #[serde(rename = "notes", deserialize_with = "::buffa::json_helpers::null_as_default")]
807    pub notes: ::buffa::alloc::vec::Vec<::buffa_types::google::protobuf::Value>,
808    #[serde(rename = "rows")]
809    pub rows: ::buffa::__private::HashMap<
810        ::buffa::alloc::string::String,
811        ::buffa_types::google::protobuf::ListValue,
812    >,
813    #[serde(rename = "name")]
814    pub name: ::buffa::alloc::string::String,
815"#;
816        let out = rewrites.rewrite("demo.v1.rs", src);
817        for kind in ["message_field", "repeated", "map"] {
818            assert!(
819                out.contains(&format!(
820                    "#[serde(serialize_with = \"crate::models::json_value::{kind}\")]"
821                )),
822                "{out}"
823            );
824        }
825        assert_eq!(
826            out.matches("serialize_with = \"crate::models").count(),
827            3,
828            "{out}"
829        );
830        assert!(out.contains("    pub rows: ::buffa::__private::HashMap<\n"));
831    }
832}