einsum_codegen/codegen/
format.rs

1use std::{
2    io::Write,
3    process::{Command, Stdio},
4};
5
6/// Format generated Rust code using `rustfmt` run as external process.
7pub fn format_block(tt: String) -> String {
8    let tt = format!("fn main() {{ {} }}", tt);
9
10    let mut child = Command::new("rustfmt")
11        .stdin(Stdio::piped())
12        .stdout(Stdio::piped())
13        .spawn()
14        .expect("Failed to spawn rustfmt process");
15
16    // Write input from another thread for avoiding deadlock.
17    // See https://doc.rust-lang.org/std/process/index.html#handling-io
18    let mut stdin = child.stdin.take().expect("Failed to open stdin");
19    std::thread::spawn(move || {
20        stdin
21            .write_all(tt.as_bytes())
22            .expect("Failed to write to stdin");
23    });
24    let output = child
25        .wait_with_output()
26        .expect("Failed to wait output of rustfmt process");
27
28    // non-UTF8 comment should be handled in the tokenize phase,
29    // and not be included in IR.
30    let out = String::from_utf8(output.stdout).expect("rustfmt output contains non-UTF8 input");
31
32    let formatted_lines: Vec<&str> = out
33        .lines()
34        .filter_map(|line| match line {
35            "fn main() {" | "}" => None,
36            _ => line.strip_prefix("    "),
37        })
38        .collect();
39    formatted_lines.join("\n")
40}
einsum_codegen/codegen/format.rs

einsum_codegen/codegen/
format.rs
