1mod expression;
2mod pattern;
3mod sequence;
4mod top_level_item;
5
6use crate::comments::Comments;
7use crate::lindig::{Document, concat, join};
8use syntax::ast::{Attribute, Expression, ImportAlias, Visibility};
9
10pub struct Formatter<'a> {
11 pub(super) comments: Comments<'a>,
12}
13
14impl<'a> Formatter<'a> {
15 pub fn new(comments: Comments<'a>) -> Self {
16 Self { comments }
17 }
18
19 pub fn module(&mut self, top_level_items: &'a [Expression]) -> Document<'a> {
20 let (imports, rest): (Vec<_>, Vec<_>) = top_level_items
21 .iter()
22 .partition(|e| matches!(e, Expression::ModuleImport { .. }));
23
24 let mut docs = Vec::new();
25
26 if !imports.is_empty() {
27 docs.push(self.sort_imports(&imports));
28 }
29
30 let mut prev_end: Option<u32> = None;
31 for (i, item) in rest.iter().enumerate() {
32 let start = Self::item_leading_edge(item);
33
34 let (same_line_trailing, leading, _) = match prev_end {
35 Some(anchor) => self.comments.take_split_by_newline_after(anchor, start),
36 None => (None, self.comments.take_comments_before(start), false),
37 };
38
39 if let Some(t) = same_line_trailing {
40 docs.push(Document::str(" "));
41 docs.push(t);
42 }
43
44 if let Some(comment_doc) = leading {
45 if !docs.is_empty() {
46 docs.push(Document::Newline);
47 docs.push(Document::Newline);
48 }
49 docs.push(comment_doc.force_break());
50 docs.push(Document::Newline);
51 } else if !docs.is_empty() || i > 0 {
52 docs.push(Document::Newline);
53 docs.push(Document::Newline);
54 }
55
56 docs.push(self.definition(item));
57 let span = item.get_span();
58 prev_end = Some(span.byte_offset + span.byte_length);
59 }
60
61 if let Some(comment_doc) = self.comments.take_trailing_comments() {
62 if !docs.is_empty() {
63 docs.push(Document::Newline);
64 docs.push(Document::Newline);
65 }
66 docs.push(comment_doc);
67 }
68
69 if !docs.is_empty() {
70 docs.push(Document::Newline);
71 }
72
73 concat(docs)
74 }
75
76 fn sort_imports(&mut self, imports: &[&'a Expression]) -> Document<'a> {
77 if imports.is_empty() {
78 return Document::Sequence(vec![]);
79 }
80
81 let mut leading_comments: Option<Document<'a>> = None;
82 let mut leading_has_blank_line = false;
83 let mut go_imports: Vec<&'a Expression> = Vec::new();
84 let mut local_imports: Vec<&'a Expression> = Vec::new();
85
86 for (i, import) in imports.iter().enumerate() {
87 let start = import.get_span().byte_offset;
88 let has_blank_line = self.comments.take_empty_lines_before(start);
89
90 let comments = self.comments.take_comments_before(start);
91 if i == 0 && comments.is_some() {
92 leading_comments = comments;
93 leading_has_blank_line = has_blank_line;
94 }
95
96 if let Expression::ModuleImport { name, .. } = import {
97 if name.starts_with("go:") {
98 go_imports.push(import);
99 } else {
100 local_imports.push(import);
101 }
102 }
103 }
104
105 fn import_sort_key(imp: &&Expression) -> (String, String) {
106 if let Expression::ModuleImport { name, alias, .. } = imp {
107 let sort_path = match alias {
108 Some(ImportAlias::Named(a, _)) => a.to_string(),
109 Some(ImportAlias::Blank(_)) => "_".to_string(),
110 None => {
111 let path = name.split_once(':').map(|(_, p)| p).unwrap_or(name);
112 path.to_string()
113 }
114 };
115 (sort_path, name.to_string())
116 } else {
117 (String::new(), String::new())
118 }
119 }
120
121 go_imports.sort_by_key(import_sort_key);
122 local_imports.sort_by_key(import_sort_key);
123
124 let mut group_docs: Vec<Document<'a>> = Vec::new();
125
126 if !go_imports.is_empty() {
127 let docs: Vec<_> = go_imports.iter().map(|imp| self.definition(imp)).collect();
128 group_docs.push(join(docs, Document::Newline));
129 }
130
131 if !local_imports.is_empty() {
132 let docs: Vec<_> = local_imports
133 .iter()
134 .map(|imp| self.definition(imp))
135 .collect();
136 group_docs.push(join(docs, Document::Newline));
137 }
138
139 let imports_doc = join(group_docs, concat([Document::Newline, Document::Newline]));
140
141 match leading_comments {
142 Some(c) => {
143 let separator = if leading_has_blank_line {
144 concat([Document::Newline, Document::Newline])
145 } else {
146 Document::Newline
147 };
148 c.force_break().append(separator).append(imports_doc)
149 }
150 None => imports_doc,
151 }
152 }
153
154 fn definition(&mut self, expression: &'a Expression) -> Document<'a> {
155 let start = expression.get_span().byte_offset;
156 let doc_comments_doc = self.comments.take_doc_comments_before(start);
157
158 let attrs = match expression {
159 Expression::Function { attributes, .. } | Expression::Struct { attributes, .. } => {
160 self.attributes(attributes)
161 }
162 _ => Document::Sequence(vec![]),
163 };
164 let between_attrs_and_keyword = self.comments.take_comments_before(start);
165
166 let (vis, inner) = match expression {
167 Expression::Function {
168 name,
169 generics,
170 params,
171 return_annotation,
172 body,
173 visibility,
174 ..
175 } => (
176 *visibility,
177 self.function(name, generics, params, return_annotation, body),
178 ),
179
180 Expression::Struct {
181 name,
182 generics,
183 fields,
184 kind,
185 visibility,
186 span,
187 ..
188 } => (
189 *visibility,
190 self.struct_definition(name, generics, fields, span, *kind),
191 ),
192
193 Expression::Enum {
194 name,
195 generics,
196 variants,
197 visibility,
198 span,
199 ..
200 } => (
201 *visibility,
202 self.enum_definition(name, generics, variants, span),
203 ),
204
205 Expression::ValueEnum {
206 name,
207 underlying_ty,
208 variants,
209 visibility,
210 span,
211 ..
212 } => (
213 *visibility,
214 self.value_enum_definition(name, underlying_ty.as_ref(), variants, span),
215 ),
216
217 Expression::TypeAlias {
218 name,
219 generics,
220 annotation,
221 visibility,
222 ..
223 } => (*visibility, Self::type_alias(name, generics, annotation)),
224
225 Expression::Interface {
226 name,
227 generics,
228 parents,
229 method_signatures,
230 visibility,
231 span,
232 ..
233 } => (
234 *visibility,
235 self.interface(name, generics, parents, method_signatures, span),
236 ),
237
238 Expression::ImplBlock {
239 annotation,
240 generics,
241 methods,
242 span,
243 ..
244 } => (
245 Visibility::Private,
246 self.impl_block(annotation, generics, methods, span.end()),
247 ),
248
249 Expression::Const {
250 identifier,
251 annotation,
252 expression,
253 visibility,
254 ..
255 } => (
256 *visibility,
257 self.const_definition(identifier, annotation.as_ref(), expression),
258 ),
259
260 Expression::VariableDeclaration {
261 name,
262 annotation,
263 visibility,
264 ..
265 } => (
266 *visibility,
267 Document::str("var ")
268 .append(Document::string(name.to_string()))
269 .append(": ")
270 .append(Self::annotation(annotation)),
271 ),
272
273 Expression::ModuleImport { name, alias, .. } => {
274 let alias_doc = match alias {
275 Some(ImportAlias::Named(a, _)) => Document::string(a.to_string()).append(" "),
276 Some(ImportAlias::Blank(_)) => Document::str("_ "),
277 None => Document::str(""),
278 };
279
280 (
281 Visibility::Private,
282 Document::str("import ")
283 .append(alias_doc)
284 .append("\"")
285 .append(Document::string(name.to_string()))
286 .append("\""),
287 )
288 }
289
290 _ => (Visibility::Private, self.expression(expression)),
291 };
292
293 let vis_inner = match Self::visibility(vis) {
294 Some(pub_doc) => pub_doc.append(inner),
295 None => inner,
296 };
297 let definition_doc = match between_attrs_and_keyword {
298 Some(c) => attrs
299 .append(c.force_break())
300 .append(Document::Newline)
301 .append(vis_inner),
302 None => attrs.append(vis_inner),
303 };
304
305 match doc_comments_doc {
306 Some(doc) => doc.append(Document::Newline).append(definition_doc),
307 None => definition_doc,
308 }
309 }
310
311 fn visibility(vis: Visibility) -> Option<Document<'a>> {
312 match vis {
313 Visibility::Public => Some(Document::str("pub ")),
314 Visibility::Private => None,
315 }
316 }
317
318 fn item_leading_edge(item: &'a Expression) -> u32 {
319 let attrs: &[Attribute] = match item {
320 Expression::Function { attributes, .. } | Expression::Struct { attributes, .. } => {
321 attributes
322 }
323 _ => &[],
324 };
325 attrs
326 .first()
327 .map(|a| a.span.byte_offset)
328 .unwrap_or_else(|| item.get_span().byte_offset)
329 }
330}