fix: keep the newline char if file ends with that

src/lib.rs

@@ -30,6 +30,9 @@ pub enum TlaError {
 }
 
 pub fn rewrite(input: &str, mode: &Mode, force: bool) -> Result<String, TlaError> {
+    // if the input ends with '\n', we should put the '\n' back to output
+    let end_of_newline = input.chars().last().map_or(false, |x| x == '\n');
+
     let mut parser = Parser::new();
     parser
         .set_language(&tree_sitter_tlaplus::LANGUAGE.into())
@@ -54,6 +57,25 @@ pub fn rewrite(input: &str, mode: &Mode, force: bool) -> Result<String, TlaError
     //println!("{:#?}", tla_lines);
     replace_symbols(&mut tla_lines);
 
+    // push a empty TlaLine,
+    // tla_lines...join("\n") will add a '\n' into end of the output,
+    // the reason why we don't push '\n' into output string directly is
+    // that maybe cause a reallocation.
+    if end_of_newline {
+        // if tla_lines is empty, we should set the text as "\n"
+        // because the join will not work in tla_lines.len < 2.
+        let text = if tla_lines.is_empty() {
+            "\n".to_string()
+        } else {
+            String::new()
+        };
+        tla_lines.push(TlaLine {
+            text,
+            jlists: Vec::new(),
+            symbols: Vec::new(),
+        })
+    }
+
     // Ensure output parse tree is identical to input parse tree
     let output = tla_lines
         .iter()
@@ -432,6 +454,9 @@ fn mark_symbols(tree: &Tree, cursor: &mut QueryCursor, tla_lines: &mut [TlaLine]
 }
 
 fn replace_symbols(tla_lines: &mut [TlaLine]) {
+    if tla_lines.is_empty() {
+        return;
+    }
     for line_number in 0..tla_lines.len() - 1 {
         let (prefix, suffix) = tla_lines.split_at_mut(line_number + 1);
         let line = &mut prefix[line_number];
@@ -850,4 +875,44 @@ op == /\ A
 ===="#,
         );
     }
+
+    // Tests that file ends with newline(or without newline)
+    #[test]
+    fn test_empty_input() {
+        let input = "";
+        let output = rewrite(&input, &Mode::UnicodeToAscii, true);
+        assert_eq!(input, output.unwrap());
+        let output = rewrite(&input, &Mode::AsciiToUnicode, true);
+        assert_eq!(input, output.unwrap());
+    }
+
+    #[test]
+    fn test_single_newline() {
+        let input = "\n";
+        let output = rewrite(&input, &Mode::UnicodeToAscii, true);
+        assert_eq!(input, output.unwrap());
+        let output = rewrite(&input, &Mode::AsciiToUnicode, true);
+        assert_eq!(input, output.unwrap());
+    }
+
+    #[test]
+    fn test_normal_input_without_newline() {
+        run_roundtrip_test(
+            r#"
+---- MODULE Test ----
+op == 1
+===="#,
+        );
+    }
+
+    #[test]
+    fn test_normal_input_with_newline() {
+        run_roundtrip_test(
+            r#"
+---- MODULE Test ----
+op == 1
+====
+"#,
+        );
+    }
 }