From 7c7a4e4edf54604f1caa868be6ff7019061a5d47 Mon Sep 17 00:00:00 2001
From: mlokr <jdoka@crownsterling.io>
Date: Fri, 19 Jul 2024 14:17:45 +0200
Subject: [PATCH] s

---
 hblang/src/lib.rs    | 27 ++++++++++++---------------
 hblang/src/parser.rs | 12 ++++++------
 2 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/hblang/src/lib.rs b/hblang/src/lib.rs
index 553492d..8d0f4c8 100644
--- a/hblang/src/lib.rs
+++ b/hblang/src/lib.rs
@@ -550,30 +550,27 @@ pub fn run_compiler(
 ) -> io::Result<()> {
     let parsed = parse_from_fs(options.extra_threads, root_file)?;
 
-    fn format_to_stdout(ast: parser::Ast) -> std::io::Result<()> {
+    fn format_to(ast: &parser::Ast, out: &mut impl std::io::Write) -> std::io::Result<()> {
         let source = std::fs::read_to_string(&*ast.path)?;
         parser::with_fmt_source(&source, || {
-            for expr in ast.exprs() {
-                use std::io::Write;
-                writeln!(std::io::stdout(), "{expr}")?;
+            for (i, expr) in ast.exprs().iter().enumerate() {
+                write!(out, "{expr}")?;
+                if let Some(expr) = ast.exprs().get(i + 1)
+                    && let Some(rest) = source.get(expr.pos() as usize..)
+                    && parser::insert_needed_semicolon(rest)
+                {
+                    write!(out, ";")?;
+                }
+                writeln!(out)?;
             }
             std::io::Result::Ok(())
         })
     }
 
     fn format_ast(ast: parser::Ast) -> std::io::Result<()> {
-        let source = std::fs::read_to_string(&*ast.path)?;
         let mut output = Vec::new();
-        parser::with_fmt_source(&source, || {
-            for expr in ast.exprs() {
-                use std::io::Write;
-                writeln!(output, "{expr}")?;
-            }
-            std::io::Result::Ok(())
-        })?;
-
+        format_to(&ast, &mut output)?;
         std::fs::write(&*ast.path, output)?;
-
         Ok(())
     }
 
@@ -582,7 +579,7 @@ pub fn run_compiler(
             format_ast(parsed)?;
         }
     } else if options.fmt_current {
-        format_to_stdout(parsed.into_iter().next().unwrap())?;
+        format_to(&parsed.into_iter().next().unwrap(), &mut std::io::stdout())?;
     } else {
         let mut codegen = codegen::Codegen::default();
         codegen.files = parsed;
diff --git a/hblang/src/parser.rs b/hblang/src/parser.rs
index 40b0267..be3a7cc 100644
--- a/hblang/src/parser.rs
+++ b/hblang/src/parser.rs
@@ -942,9 +942,6 @@ impl<'a> std::fmt::Display for Expr<'a> {
                 write!(f, "fn(")?;
                 fmt_list(f, false, "", args, |arg, f| write!(f, "{}: {}", arg.name, arg.ty))?;
                 write!(f, "): {ret} {body}")?;
-                if !matches!(body, Self::Block { .. }) {
-                    write!(f, ";")?;
-                }
                 Ok(())
             }
             Self::Call { func, args, trailing_comma } => {
@@ -966,9 +963,7 @@ impl<'a> std::fmt::Display for Expr<'a> {
                         write!(f, "{stmt}")?;
                         if let Some(expr) = stmts.get(i + 1)
                             && let Some(rest) = source.get(expr.pos() as usize..)
-                            && let kind = lexer::Lexer::new(rest).next().kind
-                            && (kind.precedence().is_some()
-                                || matches!(kind, TokenKind::Struct | TokenKind::Tupl))
+                            && insert_needed_semicolon(rest)
                         {
                             write!(f, ";")?;
                         }
@@ -1004,6 +999,11 @@ impl<'a> std::fmt::Display for Expr<'a> {
     }
 }
 
+pub fn insert_needed_semicolon(source: &str) -> bool {
+    let kind = lexer::Lexer::new(source).next().kind;
+    kind.precedence().is_some() || matches!(kind, TokenKind::Struct | TokenKind::Tupl)
+}
+
 #[repr(C)]
 pub struct AstInner<T: ?Sized> {
     ref_count: AtomicUsize,