Skip to main content

wit_bindgen_core/
source.rs

1use std::collections::BTreeMap;
2use std::collections::btree_map::Entry;
3use std::fmt::{self, Write};
4use std::ops::Deref;
5
6#[derive(Default)]
7pub struct Files {
8    files: BTreeMap<String, Vec<u8>>,
9}
10
11impl Files {
12    pub fn push(&mut self, name: &str, contents: &[u8]) {
13        match self.files.entry(name.to_owned()) {
14            Entry::Vacant(entry) => {
15                entry.insert(contents.to_owned());
16            }
17            Entry::Occupied(ref mut entry) => {
18                entry.get_mut().extend_from_slice(contents);
19            }
20        }
21    }
22
23    pub fn get_size(&mut self, name: &str) -> Option<usize> {
24        self.files.get(name).map(|data| data.len())
25    }
26
27    pub fn remove(&mut self, name: &str) -> Option<Vec<u8>> {
28        self.files.remove(name)
29    }
30
31    pub fn iter(&self) -> impl Iterator<Item = (&'_ str, &'_ [u8])> {
32        self.files.iter().map(|p| (p.0.as_str(), p.1.as_slice()))
33    }
34}
35
36#[derive(Default)]
37pub struct Source {
38    s: String,
39    indent: usize,
40    in_line_comment: bool,
41    continuing_line: bool,
42}
43
44impl Source {
45    pub fn append_src(&mut self, src: &Source) {
46        self.s.push_str(&src.s);
47        self.indent += src.indent;
48        self.in_line_comment = src.in_line_comment;
49    }
50
51    pub fn push_str(&mut self, src: &str) {
52        self.push_str_impl(src, true);
53    }
54
55    /// Appends literal text with the current indentation without interpreting
56    /// braces or line comments as source syntax.
57    pub fn push_str_literal(&mut self, src: &str) {
58        self.push_str_impl(src, false);
59    }
60
61    fn push_str_impl(&mut self, src: &str, interpret_syntax: bool) {
62        let lines = src.lines().collect::<Vec<_>>();
63        for (i, line) in lines.iter().enumerate() {
64            if !self.continuing_line {
65                if !line.is_empty() {
66                    for _ in 0..self.indent {
67                        self.s.push_str("  ");
68                    }
69                }
70                self.continuing_line = true;
71            }
72
73            let trimmed = line.trim();
74            if interpret_syntax && trimmed.starts_with("//") {
75                self.in_line_comment = true;
76            }
77
78            if interpret_syntax && !self.in_line_comment {
79                if trimmed.starts_with('}') && self.s.ends_with("  ") {
80                    self.s.pop();
81                    self.s.pop();
82                }
83            }
84            self.s.push_str(if lines.len() == 1 {
85                line
86            } else {
87                line.trim_start()
88            });
89            if interpret_syntax && !self.in_line_comment {
90                if trimmed.ends_with('{') {
91                    self.indent += 1;
92                }
93                if trimmed.starts_with('}') {
94                    // Note that a `saturating_sub` is used here to prevent a panic
95                    // here in the case of invalid code being generated in debug
96                    // mode. It's typically easier to debug those issues through
97                    // looking at the source code rather than getting a panic.
98                    self.indent = self.indent.saturating_sub(1);
99                }
100            }
101            if i != lines.len() - 1 || src.ends_with('\n') {
102                self.newline();
103            }
104        }
105    }
106
107    pub fn indent(&mut self, amt: usize) {
108        self.indent += amt;
109    }
110
111    pub fn deindent(&mut self, amt: usize) {
112        self.indent -= amt;
113    }
114
115    /// Set the indentation level, and return the old level.
116    pub fn set_indent(&mut self, amt: usize) -> usize {
117        let old = self.indent;
118        self.indent = amt;
119        old
120    }
121
122    fn newline(&mut self) {
123        self.in_line_comment = false;
124        self.continuing_line = false;
125        self.s.push('\n');
126    }
127
128    pub fn as_mut_string(&mut self) -> &mut String {
129        &mut self.s
130    }
131
132    pub fn as_str(&self) -> &str {
133        &self.s
134    }
135}
136
137impl Write for Source {
138    fn write_str(&mut self, s: &str) -> fmt::Result {
139        self.push_str(s);
140        Ok(())
141    }
142}
143
144impl Deref for Source {
145    type Target = str;
146    fn deref(&self) -> &str {
147        &self.s
148    }
149}
150
151impl From<Source> for String {
152    fn from(s: Source) -> String {
153        s.s
154    }
155}
156
157/// Calls [`write!`] with the passed arguments and unwraps the result.
158///
159/// Useful for writing to things with infallible `Write` implementations like
160/// `Source` and `String`.
161///
162/// [`write!`]: std::write
163#[macro_export]
164macro_rules! uwrite {
165    ($dst:expr, $($arg:tt)*) => {
166        write!($dst, $($arg)*).unwrap()
167    };
168}
169
170/// Calls [`writeln!`] with the passed arguments and unwraps the result.
171///
172/// Useful for writing to things with infallible `Write` implementations like
173/// `Source` and `String`.
174///
175/// [`writeln!`]: std::writeln
176#[macro_export]
177macro_rules! uwriteln {
178    ($dst:expr, $($arg:tt)*) => {
179        writeln!($dst, $($arg)*).unwrap()
180    };
181}
182
183#[cfg(test)]
184mod tests {
185    use super::Source;
186
187    #[test]
188    fn simple_append() {
189        let mut s = Source::default();
190        s.push_str("x");
191        assert_eq!(s.s, "x");
192        s.push_str("y");
193        assert_eq!(s.s, "xy");
194        s.push_str("z ");
195        assert_eq!(s.s, "xyz ");
196        s.push_str(" a ");
197        assert_eq!(s.s, "xyz  a ");
198        s.push_str("\na");
199        assert_eq!(s.s, "xyz  a \na");
200    }
201
202    #[test]
203    fn newline_remap() {
204        let mut s = Source::default();
205        s.push_str("function() {\n");
206        s.push_str("y\n");
207        s.push_str("}\n");
208        assert_eq!(s.s, "function() {\n  y\n}\n");
209    }
210
211    #[test]
212    fn if_else() {
213        let mut s = Source::default();
214        s.push_str("if() {\n");
215        s.push_str("y\n");
216        s.push_str("} else if () {\n");
217        s.push_str("z\n");
218        s.push_str("}\n");
219        assert_eq!(s.s, "if() {\n  y\n} else if () {\n  z\n}\n");
220    }
221
222    #[test]
223    fn trim_ws() {
224        let mut s = Source::default();
225        s.push_str(
226            "function() {
227                x
228        }",
229        );
230        assert_eq!(s.s, "function() {\n  x\n}");
231    }
232
233    #[test]
234    fn literal_text_does_not_change_indentation() {
235        let mut s = Source::default();
236        s.indent(1);
237        s.push_str_literal("}\n{");
238        s.deindent(1);
239        assert_eq!(s.s, "  }\n  {");
240    }
241}