codegen_writer/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::cell::RefCell;
4use std::fs::File;
5use std::io;
6use std::io::BufWriter;
7use std::mem;
8
9pub struct Codegen {
10    writer: Box<dyn io::Write>,
11}
12
13impl Codegen {
14    pub fn new(writer: impl io::Write + 'static) -> Self {
15        Self {
16            writer: Box::new(writer),
17        }
18    }
19
20    pub fn create_file(path: &str) -> io::Result<Self> {
21        let file = File::create(path)?;
22        Ok(Self::new(BufWriter::with_capacity(1024 * 1024, file)))
23    }
24}
25
26impl io::Write for Codegen {
27    #[inline]
28    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
29        self.writer.write(buf)
30    }
31
32    #[inline]
33    fn flush(&mut self) -> io::Result<()> {
34        self.writer.flush()
35    }
36}
37
38thread_local! {
39    static CURRENT: RefCell<Option<Codegen>> = RefCell::new(None);
40}
41
42pub fn scoped(g: Codegen, f: impl FnOnce()) -> Codegen {
43    let prev = CURRENT.with(|current| {
44        let mut cur = current.borrow_mut();
45        cur.replace(g)
46    });
47
48    f();
49
50    CURRENT.with(|current| {
51        let mut cur = current.borrow_mut();
52        mem::replace(&mut *cur, prev).unwrap()
53    })
54}
55
56pub fn with<T>(f: impl FnOnce(&mut Codegen) -> T) -> T {
57    CURRENT.with(|current| {
58        let mut cur = current.borrow_mut();
59        let g = cur.as_mut().expect("codegen is not in scope");
60        f(g)
61    })
62}
63
64#[macro_export]
65macro_rules! g {
66    () => {{
67        use ::std::io::Write;
68        $crate::with(|g| writeln!(g)).unwrap();
69    }};
70    ($fmt: literal) => {{
71        use ::std::io::Write;
72        $crate::with(|g| writeln!(g, $fmt)).unwrap();
73    }};
74    ($fmt: literal, $($arg: tt)*) => {{
75        use ::std::io::Write;
76        $crate::with(|g| writeln!(g, $fmt, $($arg)*)).unwrap();
77    }};
78}
79
80#[macro_export]
81macro_rules! glines {
82    [$($line:literal)+] => {{
83        use ::std::io::Write;
84        $crate::with(|g| {
85            $(
86                let line: &str = $line;
87                writeln!(g, "{line}").unwrap();
88            )+
89        });
90    }};
91}