wit_bindgen_core/
source.rs

1use std::collections::btree_map::Entry;
2use std::collections::BTreeMap;
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        let lines = src.lines().collect::<Vec<_>>();
53        for (i, line) in lines.iter().enumerate() {
54            if !self.continuing_line {
55                if !line.is_empty() {
56                    for _ in 0..self.indent {
57                        self.s.push_str("  ");
58                    }
59                }
60                self.continuing_line = true;
61            }
62
63            let trimmed = line.trim();
64            if trimmed.starts_with("//") {
65                self.in_line_comment = true;
66            }
67
68            if !self.in_line_comment {
69                if trimmed.starts_with('}') && self.s.ends_with("  ") {
70                    self.s.pop();
71                    self.s.pop();
72                }
73            }
74            self.s.push_str(if lines.len() == 1 {
75                line
76            } else {
77                line.trim_start()
78            });
79            if !self.in_line_comment {
80                if trimmed.ends_with('{') {
81                    self.indent += 1;
82                }
83                if trimmed.starts_with('}') {
84                    // Note that a `saturating_sub` is used here to prevent a panic
85                    // here in the case of invalid code being generated in debug
86                    // mode. It's typically easier to debug those issues through
87                    // looking at the source code rather than getting a panic.
88                    self.indent = self.indent.saturating_sub(1);
89                }
90            }
91            if i != lines.len() - 1 || src.ends_with('\n') {
92                self.newline();
93            }
94        }
95    }
96
97    pub fn indent(&mut self, amt: usize) {
98        self.indent += amt;
99    }
100
101    pub fn deindent(&mut self, amt: usize) {
102        self.indent -= amt;
103    }
104
105    /// Set the indentation level, and return the old level.
106    pub fn set_indent(&mut self, amt: usize) -> usize {
107        let old = self.indent;
108        self.indent = amt;
109        old
110    }
111
112    fn newline(&mut self) {
113        self.in_line_comment = false;
114        self.continuing_line = false;
115        self.s.push('\n');
116    }
117
118    pub fn as_mut_string(&mut self) -> &mut String {
119        &mut self.s
120    }
121
122    pub fn as_str(&self) -> &str {
123        &self.s
124    }
125}
126
127impl Write for Source {
128    fn write_str(&mut self, s: &str) -> fmt::Result {
129        self.push_str(s);
130        Ok(())
131    }
132}
133
134impl Deref for Source {
135    type Target = str;
136    fn deref(&self) -> &str {
137        &self.s
138    }
139}
140
141impl From<Source> for String {
142    fn from(s: Source) -> String {
143        s.s
144    }
145}
146
147/// Calls [`write!`] with the passed arguments and unwraps the result.
148///
149/// Useful for writing to things with infallible `Write` implementations like
150/// `Source` and `String`.
151///
152/// [`write!`]: std::write
153#[macro_export]
154macro_rules! uwrite {
155    ($dst:expr, $($arg:tt)*) => {
156        write!($dst, $($arg)*).unwrap()
157    };
158}
159
160/// Calls [`writeln!`] with the passed arguments and unwraps the result.
161///
162/// Useful for writing to things with infallible `Write` implementations like
163/// `Source` and `String`.
164///
165/// [`writeln!`]: std::writeln
166#[macro_export]
167macro_rules! uwriteln {
168    ($dst:expr, $($arg:tt)*) => {
169        writeln!($dst, $($arg)*).unwrap()
170    };
171}
172
173#[cfg(test)]
174mod tests {
175    use super::Source;
176
177    #[test]
178    fn simple_append() {
179        let mut s = Source::default();
180        s.push_str("x");
181        assert_eq!(s.s, "x");
182        s.push_str("y");
183        assert_eq!(s.s, "xy");
184        s.push_str("z ");
185        assert_eq!(s.s, "xyz ");
186        s.push_str(" a ");
187        assert_eq!(s.s, "xyz  a ");
188        s.push_str("\na");
189        assert_eq!(s.s, "xyz  a \na");
190    }
191
192    #[test]
193    fn newline_remap() {
194        let mut s = Source::default();
195        s.push_str("function() {\n");
196        s.push_str("y\n");
197        s.push_str("}\n");
198        assert_eq!(s.s, "function() {\n  y\n}\n");
199    }
200
201    #[test]
202    fn if_else() {
203        let mut s = Source::default();
204        s.push_str("if() {\n");
205        s.push_str("y\n");
206        s.push_str("} else if () {\n");
207        s.push_str("z\n");
208        s.push_str("}\n");
209        assert_eq!(s.s, "if() {\n  y\n} else if () {\n  z\n}\n");
210    }
211
212    #[test]
213    fn trim_ws() {
214        let mut s = Source::default();
215        s.push_str(
216            "function() {
217                x
218        }",
219        );
220        assert_eq!(s.s, "function() {\n  x\n}");
221    }
222}
wit_bindgen_core/source.rs

wit_bindgen_core/
source.rs
