petr_fmt/
lib.rs

1pub mod config;
2pub mod constants;
3pub mod ctx;
4#[cfg(test)]
5mod tests;
6
7// potential improvements to the formatter:
8// don't use Strings in the Lines struct, use a Vec<String> and join, separated by single spaces
9use std::{
10    fs,
11    io::{self, Write},
12    path::PathBuf,
13    rc::Rc,
14};
15
16pub use config::{FormatterConfig, FormatterConfigBuilder};
17use constants::{CLOSE_COMMENT_STR, INDENTATION_CHARACTER, OPEN_COMMENT_STR};
18pub use ctx::FormatterContext;
19use petr_ast::*;
20use petr_parse::Parser;
21use petr_utils::{render_error, PrettyPrint, SpannedItem};
22
23pub fn format_sources(
24    sources: Vec<(PathBuf, String)>,
25    config: FormatterConfig,
26) -> io::Result<()> {
27    let longest_source_name = sources.iter().map(|(path, _)| path.display().to_string().len()).max().unwrap_or(0);
28
29    let distance_to_check = longest_source_name + 5;
30
31    for (source_name, source) in sources {
32        let num_dots_to_display = distance_to_check - source_name.display().to_string().len();
33        print!("formatting {}...", source_name.display());
34        let string_source_name = source_name.to_string_lossy();
35        let parser = Parser::new(vec![(string_source_name.clone(), source.clone())]);
36        let (ast, errs, interner, source_map) = parser.into_result();
37        print!("{}", ".".repeat(num_dots_to_display));
38        if !errs.is_empty() {
39            errs.into_iter().for_each(|err| eprintln!("{:?}", render_error(&source_map, err)));
40            panic!("fmt failed: code didn't parse");
41        }
42        let mut ctx = FormatterContext::from_interner(interner).with_config(Default::default());
43        let formatted_content = ast.line_length_aware_format(&mut ctx).render();
44        // Create a new formatter context
45        print!("...");
46
47        if config.backup() {
48            let backup_path = format!("{string_source_name}.bak");
49            fs::write(backup_path, &source)?;
50        }
51
52        // Write the formatted content back to the file
53        let mut file = fs::File::create(source_name)?;
54        file.write_all(formatted_content.as_bytes())?;
55        println!("✅");
56    }
57    Ok(())
58}
59
60impl<T> Formattable for Commented<T>
61where
62    T: Formattable,
63{
64    fn format(
65        &self,
66        ctx: &mut FormatterContext,
67    ) -> FormattedLines {
68        let comments = self.comments();
69        let mut lines = Vec::new();
70        // if we are joining comments, join their contents
71        if ctx.config.join_comments() && !comments.is_empty() {
72            let mut buf = String::from(OPEN_COMMENT_STR);
73            for c in comments.iter().take(comments.len() - 1) {
74                // if this is not the first line, add enough space to offset the open comment syntax
75                if !lines.is_empty() {
76                    buf.push_str(&" ".repeat(OPEN_COMMENT_STR.len()));
77                }
78                buf.push_str(&c.content);
79                lines.push(ctx.new_line(buf));
80                buf = Default::default();
81            }
82            if !lines.is_empty() {
83                buf.push_str(&" ".repeat(OPEN_COMMENT_STR.len()));
84            }
85            buf.push_str(&comments.last().expect("invariant: is_empty() checked above").content);
86            buf.push_str(CLOSE_COMMENT_STR);
87            lines.push(ctx.new_line(buf));
88        } else {
89            for comment in comments {
90                lines.push(ctx.new_line(comment.pretty_print(&ctx.interner, ctx.indentation())));
91            }
92        }
93        for _ in 0..ctx.config.newlines_between_comment_and_item() {
94            lines.push(ctx.new_line(""));
95        }
96        let mut formatted_inner = self.item().format(ctx).lines;
97        lines.append(&mut formatted_inner);
98        FormattedLines::new(lines)
99    }
100}
101
102impl Formattable for FunctionDeclaration {
103    fn format(
104        &self,
105        ctx: &mut FormatterContext,
106    ) -> FormattedLines {
107        let mut lines: Vec<Line> = Vec::new();
108        let mut buf: String = if self.visibility == Visibility::Exported { "export fn " } else { "fn " }.to_string();
109
110        buf.push_str(&ctx.interner.get(self.name.id));
111
112        buf.push('(');
113        if ctx.config.put_fn_params_on_new_lines() && !self.parameters.is_empty() {
114            lines.push(ctx.new_line(buf));
115            buf = Default::default();
116        }
117        // parameter contents are indented by one
118        ctx.indented(|ctx| {
119            for (ix, param) in self.parameters.iter().enumerate() {
120                let mut param = (param).format(ctx).into_single_line().content.to_string();
121                let is_last = ix == self.parameters.len() - 1;
122
123                // if this is not the last parameter OR we are putting parameters on new lines, add a comma
124                if !is_last || ctx.config.put_fn_params_on_new_lines() {
125                    param.push(',');
126                }
127                // if we are putting params on a new line, push a new line
128                if ctx.config.put_fn_params_on_new_lines() {
129                    lines.push(ctx.new_line(param));
130                } else {
131                    buf.push_str(&format!("{param} "));
132                }
133            }
134        });
135        buf.push_str(") returns ");
136
137        buf.push_str(&self.return_type.pretty_print(&ctx.interner, ctx.indentation()));
138
139        lines.push(ctx.new_line(buf));
140
141        let mut body = ctx.indented(|ctx| self.body.format(ctx));
142
143        if ctx.config.put_fn_body_on_new_line() {
144            lines.append(&mut body.lines);
145        } else {
146            let first_line_of_body = body.lines.remove(0);
147            lines.last_mut().expect("invariant").join_with_line(first_line_of_body, " ");
148            lines.append(&mut body.lines);
149        }
150
151        FormattedLines::new(lines)
152    }
153}
154
155impl Formattable for FunctionParameter {
156    fn format(
157        &self,
158        ctx: &mut FormatterContext,
159    ) -> FormattedLines {
160        let mut buf = String::new();
161        buf.push_str(&ctx.interner.get(self.name.id));
162
163        let ty_in = if ctx.config.use_set_notation_for_types() { "∈" } else { "in" };
164
165        buf.push_str(&format!(" {ty_in} "));
166        buf.push_str(&self.ty.pretty_print(&ctx.interner, ctx.indentation()));
167
168        FormattedLines::new(vec![ctx.new_line(buf)])
169    }
170}
171
172impl Formattable for Expression {
173    fn format(
174        &self,
175        ctx: &mut FormatterContext,
176    ) -> FormattedLines {
177        match self {
178            Expression::Operator(op) => {
179                let mut buf = op.op.item().as_str().to_string();
180                buf.push(' ');
181                let (mut lhs, mut rhs) = ctx.indented(|ctx| {
182                    let lhs = op.lhs.item().format(ctx).lines;
183                    let rhs = op.rhs.item().format(ctx).lines;
184                    (lhs, rhs)
185                });
186                if lhs.len() == 1 && rhs.len() == 1 {
187                    buf.push_str(&lhs[0].content);
188                    buf.push(' ');
189                    buf.push_str(&rhs[0].content);
190                    FormattedLines::new(vec![ctx.new_line(buf)])
191                } else {
192                    let mut lines = Vec::new();
193                    buf.push_str(" (");
194                    lines.push(ctx.new_line(buf));
195                    lines.append(&mut lhs);
196                    lines.append(&mut rhs);
197                    FormattedLines::new(lines)
198                }
199            },
200            Expression::Literal(lit) => FormattedLines::new(vec![ctx.new_line(lit.to_string())]),
201            Expression::Variable(var) => {
202                let ident_as_string = ctx.interner.get(var.id);
203                FormattedLines::new(vec![ctx.new_line(ident_as_string)])
204            },
205            Expression::List(list) => list.format(ctx),
206            Expression::TypeConstructor(..) => unreachable!("this is only constructed after binding, which the formatter doesn't do"),
207            Expression::FunctionCall(f) => f.format(ctx),
208            Expression::IntrinsicCall(i) => i.format(ctx),
209            Expression::Binding(binding) => binding.format(ctx),
210        }
211    }
212}
213
214impl Formattable for ExpressionWithBindings {
215    fn format(
216        &self,
217        ctx: &mut FormatterContext,
218    ) -> FormattedLines {
219        let mut lines = Vec::new();
220        // First, format the bindings, if any.
221        if !self.bindings.is_empty() {
222            for (ix, binding) in self.bindings.iter().enumerate() {
223                let is_first = ix == 0;
224                let is_last = ix == self.bindings.len() - 1;
225                let mut buf = if is_first { "let ".to_string() } else { "    ".to_string() };
226
227                let name = ctx.interner.get(binding.name.id);
228                buf.push_str(&name);
229                buf.push_str(" = ");
230                let expr_lines = ctx.indented(|ctx| binding.val.format(ctx).lines);
231
232                if expr_lines.len() == 1 {
233                    buf.push_str(&expr_lines[0].content);
234                    lines.push(ctx.new_line(buf));
235                } else {
236                    // extend buf with first line
237                    // then add the rest of the lines
238                    buf.push_str(&expr_lines[0].content);
239                    lines.push(ctx.new_line(buf));
240                    lines.append(&mut expr_lines[1..].to_vec());
241                }
242                // add comma to the end of the last line
243                if !is_last || ctx.config.put_trailing_semis_on_let_bindings() {
244                    let last_line = lines.last().expect("invariant");
245                    let last_line_indentation = last_line.indentation;
246                    let mut last_line_content = last_line.content.to_string();
247                    last_line_content.push(';');
248                    *(lines.last_mut().expect("invariant")) = Line {
249                        content:     Rc::from(last_line_content),
250                        indentation: last_line_indentation,
251                    };
252                }
253            }
254        }
255
256        // Then, format the expression itself.
257        let expr_lines = self.expression.format(ctx).lines;
258        lines.append(&mut expr_lines.to_vec());
259
260        FormattedLines::new(lines)
261    }
262}
263
264impl Formattable for IntrinsicCall {
265    fn format(
266        &self,
267        ctx: &mut FormatterContext,
268    ) -> FormattedLines {
269        {
270            let mut buf = String::new();
271            let mut lines = vec![];
272
273            buf.push_str(&format!("@{}", self.intrinsic));
274            buf.push('(');
275
276            if ctx.config.put_fn_args_on_new_lines() {
277                lines.push(ctx.new_line(buf));
278                buf = Default::default();
279            }
280
281            ctx.indented(|ctx| {
282                for (ix, arg) in self.args.iter().enumerate() {
283                    let mut arg = (arg).format(ctx).into_single_line().content.to_string();
284                    let is_last = ix == self.args.len() - 1;
285
286                    if !is_last || ctx.config.put_fn_args_on_new_lines() {
287                        arg.push(',');
288                    }
289
290                    if !ctx.config.put_fn_args_on_new_lines() && !is_last {
291                        arg.push(' ');
292                    }
293                    if ctx.config.put_fn_args_on_new_lines() {
294                        lines.push(ctx.new_line(arg));
295                        buf = Default::default();
296                    } else {
297                        buf.push_str(&arg);
298                    }
299                }
300            });
301
302            buf.push(')');
303
304            lines.push(ctx.new_line(buf));
305            FormattedLines::new(lines)
306        }
307    }
308}
309
310impl Formattable for FunctionCall {
311    fn format(
312        &self,
313        ctx: &mut FormatterContext,
314    ) -> FormattedLines {
315        // function calls look like this: ~foo bar, baz
316        // format as such
317        let mut buf = String::new();
318        let mut lines = vec![];
319
320        buf.push('~');
321        buf.push_str(&ctx.interner.get_path(&self.func_name).join("."));
322        if self.args_were_parenthesized {
323            buf.push('(');
324        } else if !ctx.config.put_fn_args_on_new_lines() && !self.args.is_empty() {
325            buf.push(' ');
326        }
327
328        if ctx.config.put_fn_args_on_new_lines() {
329            lines.push(ctx.new_line(buf));
330            buf = Default::default();
331        }
332
333        ctx.indented(|ctx| {
334            for (ix, arg) in self.args.iter().enumerate() {
335                let mut arg = (arg).format(ctx).into_single_line().content.to_string();
336                let is_last = ix == self.args.len() - 1;
337
338                if !is_last || ctx.config.put_fn_args_on_new_lines() {
339                    arg.push(',');
340                }
341
342                if !ctx.config.put_fn_args_on_new_lines() && !is_last {
343                    arg.push(' ');
344                }
345                if ctx.config.put_fn_args_on_new_lines() {
346                    lines.push(ctx.new_line(arg));
347                    buf = Default::default();
348                } else {
349                    buf.push_str(&arg);
350                }
351            }
352        });
353        if self.args_were_parenthesized {
354            buf.push(')');
355        }
356
357        lines.push(ctx.new_line(buf));
358        FormattedLines::new(lines)
359    }
360}
361
362impl Formattable for Ast {
363    fn format(
364        &self,
365        ctx: &mut FormatterContext,
366    ) -> FormattedLines {
367        // realistically this will only get called on individual files, so there's only ever one module,
368        // since there's no inline module syntax.
369        let mut lines = Vec::new();
370        for (ix, item) in self.modules.iter().enumerate() {
371            lines.append(&mut item.format(ctx).lines);
372            if ix != self.modules.len() - 1 {
373                for _ in 0..ctx.config.newlines_between_items() {
374                    lines.push(ctx.new_line(""));
375                }
376            }
377        }
378        FormattedLines::new(lines)
379    }
380}
381
382impl Formattable for Module {
383    fn format(
384        &self,
385        ctx: &mut FormatterContext,
386    ) -> FormattedLines {
387        let mut lines = Vec::new();
388        for (ix, item) in self.nodes.iter().enumerate() {
389            lines.append(&mut item.format(ctx).lines);
390            if ix != self.nodes.len() - 1 {
391                for _ in 0..ctx.config.newlines_between_items() {
392                    lines.push(ctx.new_line(""));
393                }
394            }
395        }
396        FormattedLines::new(lines)
397    }
398}
399
400impl Formattable for AstNode {
401    fn format(
402        &self,
403        ctx: &mut FormatterContext,
404    ) -> FormattedLines {
405        match self {
406            AstNode::FunctionDeclaration(fd) => fd.format(ctx),
407            AstNode::TypeDeclaration(ty) => ty.format(ctx),
408            AstNode::ImportStatement(_) => todo!(),
409        }
410    }
411}
412
413impl Formattable for TypeDeclaration {
414    fn format(
415        &self,
416        ctx: &mut FormatterContext,
417    ) -> FormattedLines {
418        let mut lines = Vec::new();
419        let mut buf: String = if self.visibility == Visibility::Exported { "Type " } else { "type " }.to_string();
420        buf.push_str(&ctx.interner.get(self.name.id));
421        let mut variants = self.variants.iter();
422        if let Some(first_variant) = variants.next() {
423            buf.push_str(" = ");
424            let first_variant = first_variant.format(ctx).into_single_line().content;
425            buf.push_str(&first_variant);
426        } else {
427            // this is a variantless struct
428            buf.push(';');
429            return FormattedLines::new(vec![ctx.new_line(buf)]);
430        }
431        let len_to_eq = if ctx.config.put_variants_on_new_lines() {
432            buf.find('=').expect("invariant")
433        } else {
434            0
435        };
436        if ctx.config.put_variants_on_new_lines() {
437            lines.push(ctx.new_line(buf));
438            buf = Default::default();
439        }
440        ctx.indent_by(len_to_eq, |ctx| {
441            // format variants 2..n
442            for variant in variants {
443                if ctx.config.put_variants_on_new_lines() && !buf.is_empty() {
444                    lines.push(ctx.new_line(buf));
445                    buf = Default::default();
446                }
447                if !ctx.config.put_variants_on_new_lines() {
448                    buf.push(' ');
449                }
450                buf.push_str("| ");
451                let variant = variant.format(ctx).into_single_line().content;
452                buf.push_str(&variant);
453            }
454            lines.push(ctx.new_line(buf));
455        });
456        FormattedLines::new(lines)
457    }
458}
459
460impl Formattable for TypeVariant {
461    fn format(
462        &self,
463        ctx: &mut FormatterContext,
464    ) -> FormattedLines {
465        let name = ctx.interner.get(self.name.id);
466        let mut buf = name.to_string();
467        if !self.fields.is_empty() {
468            buf.push(' ');
469        }
470        let mut fields_buf = Vec::with_capacity(self.fields.len());
471        for field in &*self.fields {
472            fields_buf.push(field.format(ctx).into_single_line().content.to_string());
473        }
474        buf.push_str(&fields_buf.join(" "));
475        FormattedLines::new(vec![ctx.new_line(buf)])
476    }
477}
478
479// TODO: would be nice to format types and type fields as such
480// type Ptr =
481//      Unsized
482//          address 'int
483//    | Sized
484//          address 'int
485//          size 'int
486
487impl Formattable for TypeField {
488    fn format(
489        &self,
490        ctx: &mut FormatterContext,
491    ) -> FormattedLines {
492        let name = ctx.interner.get(self.name.id);
493        let mut buf = name.to_string();
494        buf.push(' ');
495        buf.push_str(&self.ty.pretty_print(&ctx.interner, ctx.indentation()));
496        FormattedLines::new(vec![ctx.new_line(buf)])
497    }
498}
499
500impl Formattable for Ty {
501    fn format(
502        &self,
503        ctx: &mut FormatterContext,
504    ) -> FormattedLines {
505        let name = match self {
506            Ty::Bool => "bool".to_string(),
507            Ty::Int => "int".to_string(),
508            Ty::String => "string".to_string(),
509            Ty::Unit => "unit".to_string(),
510            Ty::Named(name) => ctx.interner.get(name.id).to_string(),
511        };
512        FormattedLines::new(vec![ctx.new_line(format!("'{name}"))])
513    }
514}
515
516impl Formattable for List {
517    fn format(
518        &self,
519        ctx: &mut FormatterContext,
520    ) -> FormattedLines {
521        let mut lines = Vec::new();
522        let items = self.elements.iter();
523        let mut item_buf = vec![];
524        ctx.indented(|ctx| {
525            for item in items {
526                let element = item.format(ctx).into_single_line();
527                item_buf.push(element);
528            }
529        });
530
531        if ctx.config.put_list_elements_on_new_lines() {
532            lines.push(ctx.new_line("["));
533            for mut line in item_buf {
534                line.content = Rc::from(format!("{},", line.content).as_str());
535                lines.push(line);
536            }
537            lines.push(ctx.new_line("]"));
538        } else {
539            let text = item_buf.iter().map(|item| item.content.trim()).collect::<Vec<_>>().join(", ");
540            lines.push(ctx.new_line(format!("[{}]", text)));
541        }
542
543        FormattedLines::new(lines)
544    }
545}
546
547impl<T: Formattable> Formattable for SpannedItem<T> {
548    fn format(
549        &self,
550        ctx: &mut FormatterContext,
551    ) -> FormattedLines {
552        self.item().format(ctx)
553    }
554}
555
556// TODO: methods like "continue on current line" are going to be useful
557// this can also hold line length context
558// Instead of pushing/appending, we should be using custom joins with "continues from previous line" logic
559#[derive(Debug)]
560pub struct FormattedLines {
561    lines: Vec<Line>,
562}
563
564impl FormattedLines {
565    pub fn new(lines: Vec<Line>) -> Self {
566        Self { lines }
567    }
568
569    pub fn max_length(&self) -> usize {
570        self.lines
571            .iter()
572            .map(|line| line.content.len() + INDENTATION_CHARACTER.repeat(line.indentation).len())
573            .max()
574            .unwrap_or(0)
575    }
576
577    pub fn render(&self) -> String {
578        let mut buf = String::new();
579        for Line { indentation, content } in &self.lines {
580            // don't print indentation if the line is empty
581            // it would just be trailing whitespace
582            if content.trim().is_empty() {
583                buf.push('\n');
584                continue;
585            }
586
587            // the line has non-whitespace content, so we indent and print it
588            buf.push_str(&format!("{}{}\n", INDENTATION_CHARACTER.repeat(*indentation), content));
589        }
590        buf
591    }
592
593    /// Forces a multi-line `FormattedLines` into a single line.
594    fn into_single_line(self) -> Line {
595        let Some(indentation) = self.lines.first().map(|x| x.indentation) else {
596            return Line {
597                indentation: 0,
598                content:     Rc::from(""),
599            };
600        };
601        let content = self.lines.into_iter().map(|line| line.content).collect::<Vec<_>>().join(" ");
602        Line {
603            indentation,
604            content: Rc::from(content),
605        }
606    }
607}
608
609#[derive(Debug, Clone)]
610pub struct Line {
611    indentation: usize,
612    content:     Rc<str>,
613}
614
615impl Line {
616    pub fn join_with_line(
617        &mut self,
618        other: Line,
619        join_str: &str,
620    ) {
621        self.content = Rc::from(format!("{}{join_str}{}", self.content, other.content).as_str());
622    }
623}
624
625pub trait Formattable {
626    fn format(
627        &self,
628        ctx: &mut FormatterContext,
629    ) -> FormattedLines;
630    /// the below is the actual entry point to the interface, which attempts to reformat the item
631    /// with more things broken up across lines if the line length exceeds the limit
632    fn line_length_aware_format(
633        &self,
634        ctx: &mut FormatterContext,
635    ) -> FormattedLines {
636        let base_result = self.format(ctx);
637        // try the first configs, then the later ones
638        // these are in order of preference
639        let configs = vec![
640            ctx.config.as_builder().put_fn_params_on_new_lines(true).build(),
641            ctx.config
642                .as_builder()
643                .put_fn_params_on_new_lines(true)
644                .put_fn_args_on_new_lines(true)
645                .build(),
646            ctx.config
647                .as_builder()
648                .put_fn_params_on_new_lines(true)
649                .put_fn_body_on_new_line(true)
650                .build(),
651            ctx.config.as_builder().put_list_elements_on_new_lines(true).build(),
652            ctx.config
653                .as_builder()
654                .put_fn_params_on_new_lines(true)
655                .put_fn_body_on_new_line(true)
656                .put_variants_on_new_lines(true)
657                .build(),
658            ctx.config
659                .as_builder()
660                .put_fn_params_on_new_lines(true)
661                .put_list_elements_on_new_lines(true)
662                .put_variants_on_new_lines(true)
663                .put_list_elements_on_new_lines(true)
664                .build(),
665        ];
666
667        for config in &configs {
668            let result = self.try_config(ctx, *config);
669            if result.max_length() < ctx.config.max_line_length() {
670                return result;
671            }
672        }
673
674        // if none of the above were good enough, pick the shortest one
675        vec![base_result]
676            .into_iter()
677            .chain(configs.into_iter().map(|config| self.try_config(ctx, config)))
678            .min_by_key(|fl| fl.max_length())
679            .unwrap()
680    }
681    fn try_config(
682        &self,
683        ctx: &mut FormatterContext,
684        config: FormatterConfig,
685    ) -> FormattedLines {
686        ctx.with_new_config(config, |ctx| self.format(ctx))
687    }
688}
petr_fmt/lib.rs

petr_fmt/
lib.rs
