cargo_zng/
fmt.rs

1use std::{
2    borrow::Cow,
3    fs,
4    io::{self, Read, Write},
5    path::Path,
6    process::Stdio,
7};
8
9use clap::*;
10use once_cell::sync::Lazy;
11use proc_macro2::{Delimiter, TokenStream, TokenTree};
12use rayon::prelude::*;
13use regex::Regex;
14
15use crate::util;
16
17#[derive(Args, Debug, Default)]
18pub struct FmtArgs {
19    /// Only check if files are formatted
20    #[arg(long, action)]
21    check: bool,
22
23    /// Format the crate identified by Cargo.toml
24    #[arg(long)]
25    manifest_path: Option<String>,
26
27    /// Format the workspace crate identified by package name
28    #[arg(short, long)]
29    package: Option<String>,
30
31    /// Format all files matched by glob
32    #[arg(short, long)]
33    files: Option<String>,
34
35    /// Format the stdin to the stdout.
36    #[arg(short, long, action)]
37    stdin: bool,
38}
39
40pub fn run(mut args: FmtArgs) {
41    let (check, action) = if args.check { ("--check", "checking") } else { ("", "formatting") };
42
43    let mut custom_fmt_files = vec![];
44
45    if args.stdin {
46        if args.manifest_path.is_some() || args.package.is_some() || args.files.is_some() {
47            fatal!("stdin can only be used standalone or with --check");
48        }
49
50        let mut code = String::new();
51        if let Err(e) = std::io::stdin().read_to_string(&mut code) {
52            fatal!("stdin read error, {e}");
53        }
54
55        if code.is_empty() {
56            return;
57        }
58
59        if let Some(code) = rustfmt_stdin(&code) {
60            let stream: TokenStream = code.parse().unwrap_or_else(|e| fatal!("cannot parse stdin, {e}"));
61
62            let formatted = fmt_code(&code, stream);
63            if let Err(e) = std::io::stdout().write_all(formatted.as_bytes()) {
64                fatal!("stdout write error, {e}");
65            }
66        }
67    } else if let Some(glob) = args.files {
68        if args.manifest_path.is_some() || args.package.is_some() {
69            fatal!("--files must not be set when crate is set");
70        }
71
72        for file in glob::glob(&glob).unwrap_or_else(|e| fatal!("{e}")) {
73            let file = file.unwrap_or_else(|e| fatal!("{e}"));
74            if let Err(e) = util::cmd("rustfmt", &["--edition", "2021", check, &file.as_os_str().to_string_lossy()], &[]) {
75                fatal!("{e}");
76            }
77            custom_fmt_files.push(file);
78        }
79    } else {
80        if let Some(pkg) = args.package {
81            if args.manifest_path.is_some() {
82                fatal!("expected only one of --package, --manifest-path");
83            }
84            match util::manifest_path_from_package(&pkg) {
85                Some(m) => args.manifest_path = Some(m),
86                None => fatal!("package `{pkg}` not found in workspace"),
87            }
88        }
89        if let Some(path) = args.manifest_path {
90            if let Err(e) = util::cmd("cargo fmt --manifest-path", &[&path, check], &[]) {
91                fatal!("{e}");
92            }
93
94            let files = Path::new(&path)
95                .parent()
96                .unwrap()
97                .join("**/*.rs")
98                .display()
99                .to_string()
100                .replace('\\', "/");
101            for file in glob::glob(&files).unwrap_or_else(|e| fatal!("{e}")) {
102                let file = file.unwrap_or_else(|e| fatal!("{e}"));
103                custom_fmt_files.push(file);
104            }
105        } else {
106            if let Err(e) = util::cmd("cargo fmt", &[check], &[]) {
107                fatal!("{e}");
108            }
109
110            for path in util::workspace_manifest_paths() {
111                let files = path.parent().unwrap().join("**/*.rs").display().to_string().replace('\\', "/");
112                for file in glob::glob(&files).unwrap_or_else(|e| fatal!("{e}")) {
113                    let file = file.unwrap_or_else(|e| fatal!("{e}"));
114                    custom_fmt_files.push(file);
115                }
116            }
117        }
118    }
119
120    custom_fmt_files.par_iter().for_each(|file| {
121        if let Err(e) = custom_fmt(file, args.check) {
122            fatal!("error {action} `{}`, {e}", file.display());
123        }
124    });
125}
126
127fn custom_fmt(rs_file: &Path, check: bool) -> io::Result<()> {
128    let file = fs::read_to_string(rs_file)?;
129
130    // skip UTF-8 BOM
131    let file_code = file.strip_prefix('\u{feff}').unwrap_or(file.as_str());
132    // skip shebang line
133    let file_code = if file_code.starts_with("#!") && !file_code.starts_with("#![") {
134        &file_code[file_code.find('\n').unwrap_or(file_code.len())..]
135    } else {
136        file_code
137    };
138
139    let mut formatted_code = file[..file.len() - file_code.len()].to_owned();
140    formatted_code.reserve(file.len());
141
142    let file_stream: TokenStream = file_code
143        .parse()
144        .unwrap_or_else(|e| fatal!("cannot parse `{}`, {e}", rs_file.display()));
145
146    formatted_code.push_str(&fmt_code(file_code, file_stream));
147
148    if formatted_code != file {
149        if check {
150            fatal!("extended format does not match in file `{}`", rs_file.display());
151        }
152        fs::write(rs_file, formatted_code)?;
153    }
154
155    Ok(())
156}
157
158fn fmt_code(code: &str, stream: TokenStream) -> String {
159    let mut formatted_code = String::new();
160    let mut last_already_fmt_start = 0;
161
162    let mut stream_stack = vec![stream.into_iter()];
163    let next = |stack: &mut Vec<proc_macro2::token_stream::IntoIter>| {
164        while !stack.is_empty() {
165            let tt = stack.last_mut().unwrap().next();
166            if tt.is_some() {
167                return tt;
168            }
169            stack.pop();
170        }
171        None
172    };
173    let mut tail2 = Vec::with_capacity(2);
174
175    let mut skip_next_group = false;
176    while let Some(tt) = next(&mut stream_stack) {
177        match tt {
178            TokenTree::Group(g) => {
179                if tail2.len() == 2
180                    && matches!(g.delimiter(), Delimiter::Brace)
181                    && matches!(&tail2[0], TokenTree::Punct(p) if p.as_char() == '!')
182                    && matches!(&tail2[1], TokenTree::Ident(_))
183                {
184                    // macro! {}
185                    if std::mem::take(&mut skip_next_group) {
186                        continue;
187                    }
188
189                    let bang = tail2[0].span().byte_range().start;
190                    let line_start = code[..bang].rfind('\n').unwrap_or(0);
191                    let base_indent = code[line_start..bang]
192                        .chars()
193                        .skip_while(|&c| c != ' ')
194                        .take_while(|&c| c == ' ')
195                        .count();
196
197                    let group_bytes = g.span().byte_range();
198                    let group_code = &code[group_bytes.clone()];
199
200                    if let Some(formatted) = try_fmt_macro(base_indent, group_code) {
201                        if formatted != group_code {
202                            // changed by custom format
203                            if let Some(stable) = try_fmt_macro(base_indent, &formatted) {
204                                if formatted == stable {
205                                    // change is sable
206                                    let already_fmt = &code[last_already_fmt_start..group_bytes.start];
207                                    formatted_code.push_str(already_fmt);
208                                    formatted_code.push_str(&formatted);
209                                    last_already_fmt_start = group_bytes.end;
210                                }
211                            }
212                        }
213                    }
214                } else if !tail2.is_empty()
215                    && matches!(g.delimiter(), Delimiter::Bracket)
216                    && matches!(&tail2[0], TokenTree::Punct(p) if p.as_char() == '#')
217                {
218                    // #[..]
219                    let mut attr = g.stream().into_iter();
220                    let attr = [attr.next(), attr.next(), attr.next(), attr.next(), attr.next()];
221                    if let [
222                        Some(TokenTree::Ident(i0)),
223                        Some(TokenTree::Punct(p0)),
224                        Some(TokenTree::Punct(p1)),
225                        Some(TokenTree::Ident(i1)),
226                        None,
227                    ] = attr
228                    {
229                        if i0 == "rustfmt" && p0.as_char() == ':' && p1.as_char() == ':' && i1 == "skip" {
230                            // #[rustfmt::skip]
231                            skip_next_group = true;
232                        }
233                    }
234                } else if !std::mem::take(&mut skip_next_group) {
235                    stream_stack.push(g.stream().into_iter());
236                }
237                tail2.clear();
238            }
239            tt => {
240                if tail2.len() == 2 {
241                    tail2.pop();
242                }
243                tail2.insert(0, tt);
244            }
245        }
246    }
247
248    formatted_code.push_str(&code[last_already_fmt_start..]);
249
250    if formatted_code != code {
251        // custom format can cause normal format to change
252        // example: ui_vec![Wgt!{<many properties>}, Wgt!{<same>}]
253        //   Wgt! gets custom formatted onto multiple lines, that causes ui_vec![\n by normal format.
254        formatted_code = rustfmt_stdin_frag(&formatted_code).unwrap_or(formatted_code);
255    }
256
257    formatted_code
258}
259
260fn try_fmt_macro(base_indent: usize, group_code: &str) -> Option<String> {
261    let mut replaced_code = replace_event_args(group_code, false);
262    let is_event_args = matches!(&replaced_code, Cow::Owned(_));
263
264    let mut is_widget = false;
265    if !is_event_args {
266        replaced_code = replace_widget_when(group_code, false);
267        is_widget = matches!(&replaced_code, Cow::Owned(_));
268
269        let tmp = replace_widget_prop(&replaced_code, false);
270        if let Cow::Owned(tmp) = tmp {
271            is_widget = true;
272            replaced_code = Cow::Owned(tmp);
273        }
274    }
275
276    let mut is_expr_var = false;
277    if !is_event_args && !is_widget {
278        replaced_code = replace_expr_var(group_code, false);
279        is_expr_var = matches!(&replaced_code, Cow::Owned(_));
280    }
281
282    let code = rustfmt_stdin_frag(&replaced_code)?;
283
284    let code = if is_event_args {
285        replace_event_args(&code, true)
286    } else if is_widget {
287        let code = replace_widget_when(&code, true);
288        let code = replace_widget_prop(&code, true).into_owned();
289        Cow::Owned(code)
290    } else if is_expr_var {
291        replace_expr_var(&code, true)
292    } else {
293        Cow::Owned(code)
294    };
295
296    let code_stream: TokenStream = code.parse().unwrap_or_else(|e| panic!("{e}\ncode:\n{code}"));
297    let code_tt = code_stream.into_iter().next().unwrap();
298    let code_stream = match code_tt {
299        TokenTree::Group(g) => g.stream(),
300        _ => unreachable!(),
301    };
302    let code = fmt_code(&code, code_stream);
303
304    let mut out = String::new();
305    let mut lb_indent = String::with_capacity(base_indent + 1);
306    for line in code.lines() {
307        if line.is_empty() {
308            if !lb_indent.is_empty() {
309                out.push('\n');
310            }
311        } else {
312            out.push_str(&lb_indent);
313        }
314        out.push_str(line);
315        // "\n    "
316        if lb_indent.is_empty() {
317            lb_indent.push('\n');
318            for _ in 0..base_indent {
319                lb_indent.push(' ');
320            }
321        }
322    }
323    Some(out)
324}
325// replace line with only `..` tokens with:
326//
327// ```
328// // cargo-zng::fmt::dot_dot
329// }
330// impl CargoZngFmt {
331//
332// ```
333fn replace_event_args(code: &str, reverse: bool) -> Cow<str> {
334    static RGX: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*(\.\.)\s*$").unwrap());
335    static MARKER: &str = "// cargo-zng::fmt::dot_dot\n}\nimpl CargoZngFmt {\n";
336    static RGX_REV: Lazy<Regex> =
337        Lazy::new(|| Regex::new(r"(?m)^(\s+)// cargo-zng::fmt::dot_dot\n\s*}\n\s*impl CargoZngFmt\s*\{\n").unwrap());
338
339    if !reverse {
340        RGX.replace_all(code, |caps: &regex::Captures| {
341            format!(
342                "{}{MARKER}{}",
343                &caps[0][..caps.get(1).unwrap().start() - caps.get(0).unwrap().start()],
344                &caps[0][caps.get(1).unwrap().end() - caps.get(0).unwrap().start()..]
345            )
346        })
347    } else {
348        RGX_REV.replace_all(code, "\n$1..\n\n")
349    }
350}
351// replace `prop = 1, 2;` with `prop = (1, 2);`
352// AND replace `prop = { a: 1, b: 2, };` with `prop = __A_ { a: 1, b: 2, }`
353fn replace_widget_prop(code: &str, reverse: bool) -> Cow<str> {
354    static NAMED_RGX: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)\w+\s+=\s+(\{)").unwrap());
355    static NAMED_MARKER: &str = "__A_ ";
356
357    static UNNAMED_RGX: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?ms)\w+\s+=\s+([^\(\{\n\)]+?)(?:;|}$)").unwrap());
358    static UNNAMED_RGX_REV: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?ms)__a_\((.+?)\)").unwrap());
359
360    if !reverse {
361        let named_rpl = NAMED_RGX.replace_all(code, |caps: &regex::Captures| {
362            format!(
363                "{}{NAMED_MARKER} {{",
364                &caps[0][..caps.get(1).unwrap().start() - caps.get(0).unwrap().start()],
365            )
366        });
367        let mut has_unnamed = false;
368        let unnamed_rpl = UNNAMED_RGX.replace_all(&named_rpl, |caps: &regex::Captures| {
369            let cap = caps.get(1).unwrap();
370            let cap_str = cap.as_str().trim();
371            fn more_than_one_expr(code: &str) -> bool {
372                let stream: TokenStream = match code.parse() {
373                    Ok(s) => s,
374                    Err(_e) => {
375                        #[cfg(debug_assertions)]
376                        panic!("{_e}\ncode:\n{code}");
377                        #[cfg(not(debug_assertions))]
378                        return false;
379                    }
380                };
381                for tt in stream {
382                    if let TokenTree::Punct(p) = tt {
383                        if p.as_char() == ',' {
384                            return true;
385                        }
386                    }
387                }
388                false
389            }
390            if cap_str.contains(",") && more_than_one_expr(cap_str) {
391                has_unnamed = true;
392
393                format!(
394                    "{}__a_({cap_str}){}",
395                    &caps[0][..cap.start() - caps.get(0).unwrap().start()],
396                    &caps[0][cap.end() - caps.get(0).unwrap().start()..],
397                )
398            } else {
399                caps.get(0).unwrap().as_str().to_owned()
400            }
401        });
402        if has_unnamed {
403            Cow::Owned(unnamed_rpl.into_owned())
404        } else {
405            named_rpl
406        }
407    } else {
408        let code = UNNAMED_RGX_REV.replace_all(code, |caps: &regex::Captures| {
409            format!(
410                "{}{}{}",
411                &caps[0][..caps.get(1).unwrap().start() - caps.get(0).unwrap().start() - "__a_(".len()],
412                caps.get(1).unwrap().as_str(),
413                &caps[0][caps.get(1).unwrap().end() + ")".len() - caps.get(0).unwrap().start()..]
414            )
415        });
416        Cow::Owned(code.replace(NAMED_MARKER, ""))
417    }
418}
419// replace `when <expr> { <properties> }` with `for cargo_zng_fmt_when in <expr> { <properties> }`
420// AND replace `#expr` with `__P_expr` AND `#{var}` with `__P_!{`
421fn replace_widget_when(code: &str, reverse: bool) -> Cow<str> {
422    static RGX: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?s)\n\s*(when) .+?\{").unwrap());
423    static MARKER: &str = "for cargo_zng_fmt_when in";
424    static POUND_MARKER: &str = "__P_";
425
426    if !reverse {
427        RGX.replace_all(code, |caps: &regex::Captures| {
428            let prefix_spaces = &caps[0][..caps.get(1).unwrap().start() - caps.get(0).unwrap().start()];
429
430            let expr = &caps[0][caps.get(1).unwrap().end() - caps.get(0).unwrap().start()..];
431            let expr = POUND_RGX.replace_all(expr, |caps: &regex::Captures| {
432                let c = &caps[0][caps.get(1).unwrap().end() - caps.get(0).unwrap().start()..];
433                let marker = if c == "{" { POUND_VAR_MARKER } else { POUND_MARKER };
434                format!("{marker}{c}")
435            });
436
437            format!("{prefix_spaces}{MARKER}{expr}")
438        })
439    } else {
440        let code = code.replace(MARKER, "when");
441        let r = POUND_REV_RGX.replace_all(&code, "#").into_owned();
442        Cow::Owned(r)
443    }
444}
445static POUND_RGX: Lazy<Regex> = Lazy::new(|| Regex::new(r"(#)[\w\{]").unwrap());
446static POUND_REV_RGX: Lazy<Regex> = Lazy::new(|| Regex::new(r"__P_!?\s?").unwrap());
447static POUND_VAR_MARKER: &str = "__P_!";
448
449// replace `#{` with `__P_!{`
450fn replace_expr_var(code: &str, reverse: bool) -> Cow<str> {
451    if !reverse {
452        POUND_RGX.replace(code, |caps: &regex::Captures| {
453            let c = &caps[0][caps.get(1).unwrap().end() - caps.get(0).unwrap().start()..];
454            if c == "{" {
455                Cow::Borrowed("__P_!{")
456            } else {
457                Cow::Owned(caps[0].to_owned())
458            }
459        })
460    } else {
461        POUND_REV_RGX.replace(code, "#")
462    }
463}
464
465fn rustfmt_stdin_frag(code: &str) -> Option<String> {
466    let mut s = std::process::Command::new("rustfmt")
467        .arg("--edition")
468        .arg("2021")
469        .stdin(Stdio::piped())
470        .stdout(Stdio::piped())
471        .stderr(Stdio::piped())
472        .spawn()
473        .ok()?;
474    s.stdin.take().unwrap().write_all(format!("fn __try_fmt(){code}").as_bytes()).ok()?;
475    let s = s.wait_with_output().ok()?;
476
477    if s.status.success() {
478        let code = String::from_utf8(s.stdout).ok()?;
479        let code = code.strip_prefix("fn __try_fmt()")?.trim_start().to_owned();
480        Some(code)
481    } else {
482        None
483    }
484}
485
486fn rustfmt_stdin(code: &str) -> Option<String> {
487    let mut s = std::process::Command::new("rustfmt")
488        .arg("--edition")
489        .arg("2021")
490        .stdin(Stdio::piped())
491        .stdout(Stdio::piped())
492        .spawn()
493        .ok()?;
494    s.stdin.take().unwrap().write_all(code.as_bytes()).ok()?;
495    let s = s.wait_with_output().ok()?;
496
497    if s.status.success() {
498        let code = String::from_utf8(s.stdout).ok()?;
499        Some(code)
500    } else {
501        None
502    }
503}