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}