feat(fmt): improve newline insertion when comments are in between leaf spans

swayfmt/src/utils/map/newline.rs

@@ -1,12 +1,6 @@
 use anyhow::Result;
 use ropey::Rope;
-use std::{
-    collections::BTreeMap,
-    fmt::Write,
-    ops::Bound::{Excluded, Included},
-    path::PathBuf,
-    sync::Arc,
-};
+use std::{collections::BTreeMap, fmt::Write, path::PathBuf, sync::Arc};
 use sway_ast::Module;
 
 use crate::{
@@ -17,7 +11,7 @@ use crate::{
 };
 
 /// Represents a series of consecutive newlines
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 struct NewlineSequence {
     sequence_length: usize,
 }
@@ -149,18 +143,65 @@ fn add_newlines(
         .skip(1)
         .zip(formatted_newline_spans.iter().skip(1))
     {
-        if let Some(newline_sequence) = get_newline_sequence_between_spans(
-            previous_unformatted_newline_span,
-            unformatted_newline_span,
-            &newline_map,
-        ) {
-            offset += insert_after_span(
-                previous_formatted_newline_span,
-                newline_sequence,
-                offset,
-                formatted_code,
-                newline_threshold,
-            )?;
+        if previous_unformatted_newline_span.end < unformatted_newline_span.start {
+            let snippet = &unformatted_code
+                [previous_unformatted_newline_span.end..unformatted_newline_span.start];
+
+            let start = previous_unformatted_newline_span.end;
+
+            // Here, we will try to insert newlines that occur before and after comments.
+            // The reason we do this is because comments aren't a part of the AST, and so they aren't
+            // collected as leaf spans - they simply exist between whitespaces.
+            if let Some(start_first_comment) = snippet.find("//") {
+                // Insert newlines that occur before the first comment here
+                if let Some(newline_sequence) = first_newline_sequence_in_span(
+                    &ByteSpan {
+                        start,
+                        end: start + start_first_comment,
+                    },
+                    &newline_map,
+                ) {
+                    let at = previous_formatted_newline_span.end + offset;
+                    offset +=
+                        insert_after_span(at, newline_sequence, formatted_code, newline_threshold)?;
+                }
+
+                // Insert newlines that occur after the last comment here
+                if let Some(start_last_comment) = snippet.rfind("//") {
+                    if start_first_comment != start_last_comment {
+                        if let Some(end_last_comment) = snippet[start_last_comment..].find('\n') {
+                            if let Some(newline_sequence) = first_newline_sequence_in_span(
+                                &ByteSpan {
+                                    start: start + start_last_comment + end_last_comment,
+                                    end: unformatted_newline_span.start,
+                                },
+                                &newline_map,
+                            ) {
+                                let at = previous_formatted_newline_span.end
+                                    + offset
+                                    + start_last_comment
+                                    + end_last_comment;
+                                offset += insert_after_span(
+                                    at,
+                                    newline_sequence,
+                                    formatted_code,
+                                    newline_threshold,
+                                )?;
+                            }
+                        }
+                    }
+                }
+            } else if let Some(newline_sequence) = first_newline_sequence_in_span(
+                &ByteSpan {
+                    start,
+                    end: unformatted_newline_span.start,
+                },
+                &newline_map,
+            ) {
+                let at = previous_formatted_newline_span.end + offset;
+                offset +=
+                    insert_after_span(at, newline_sequence, formatted_code, newline_threshold)?;
+            }
         }
         previous_unformatted_newline_span = unformatted_newline_span;
         previous_formatted_newline_span = formatted_newline_span;
@@ -176,47 +217,47 @@ fn format_newline_sequence(newline_sequence: &NewlineSequence, threshold: usize)
     }
 }
 
-/// Inserts after given span and returns the offset.
+/// Inserts a `NewlineSequence` at position `at` and returns the length of `NewlineSequence` inserted.
+/// The return value is used to calculate the new `at` in a later point.
 fn insert_after_span(
-    from: &ByteSpan,
+    at: usize,
     newline_sequence: NewlineSequence,
-    offset: usize,
     formatted_code: &mut FormattedCode,
     threshold: usize,
 ) -> Result<usize, FormatterError> {
     let mut sequence_string = String::new();
-
     write!(
         sequence_string,
         "{}",
         format_newline_sequence(&newline_sequence, threshold)
     )?;
     let mut src_rope = Rope::from_str(formatted_code);
-    src_rope.insert(from.end + offset, &sequence_string);
+    src_rope.insert(at, &sequence_string);
     formatted_code.clear();
     formatted_code.push_str(&src_rope.to_string());
     Ok(sequence_string.len())
 }
 
-/// Returns a newline sequence between given spans, if found.
-fn get_newline_sequence_between_spans(
-    from: &ByteSpan,
-    to: &ByteSpan,
+/// Returns the first newline sequence contained in a span.
+/// This is inclusive at the start, and exclusive at the end, i.e.
+/// the bounds are [span.start, span.end).
+fn first_newline_sequence_in_span(
+    span: &ByteSpan,
     newline_map: &NewlineMap,
 ) -> Option<NewlineSequence> {
-    if from < to {
-        if let Some((_, newline_sequence)) =
-            newline_map.range((Included(from), Excluded(to))).next()
-        {
-            return Some(newline_sequence.clone());
+    for (range, sequence) in newline_map.iter() {
+        if span.start <= range.start && range.end < span.end {
+            return Some(sequence.clone());
         }
     }
     None
 }
 
 #[cfg(test)]
 mod tests {
-    use super::newline_map_from_src;
+    use crate::utils::map::{byte_span::ByteSpan, newline::first_newline_sequence_in_span};
+
+    use super::{newline_map_from_src, NewlineMap, NewlineSequence};
 
     #[test]
     fn test_newline_map() {
@@ -264,4 +305,29 @@ fn main() {
 
         assert_eq!(newline_sequence_lengths, correct_newline_sequence_lengths);
     }
+
+    #[test]
+    fn test_newline_range_simple() {
+        let mut newline_map = NewlineMap::new();
+        let newline_sequence = NewlineSequence { sequence_length: 2 };
+
+        newline_map.insert(ByteSpan { start: 9, end: 10 }, newline_sequence.clone());
+        assert_eq!(
+            newline_sequence,
+            first_newline_sequence_in_span(&ByteSpan { start: 8, end: 11 }, &newline_map).unwrap()
+        );
+        assert_eq!(
+            newline_sequence,
+            first_newline_sequence_in_span(&ByteSpan { start: 9, end: 11 }, &newline_map).unwrap()
+        );
+        assert!(
+            first_newline_sequence_in_span(&ByteSpan { start: 9, end: 10 }, &newline_map).is_none()
+        );
+        assert!(
+            first_newline_sequence_in_span(&ByteSpan { start: 9, end: 9 }, &newline_map).is_none()
+        );
+        assert!(
+            first_newline_sequence_in_span(&ByteSpan { start: 8, end: 8 }, &newline_map).is_none()
+        );
+    }
 }

swayfmt/tests/mod.rs

@@ -1531,3 +1531,40 @@ use ::utils::numbers::*;
 "#,
     );
 }
+
+#[test]
+fn temporarily_commented_out_fn_with_doc_comments() {
+    check(
+        r#"contract;
+
+abi MyContract {
+    fn test_function() -> bool;
+}
+
+impl MyContract for Contract {
+    /// This is documentation for a commented out function
+    // fn commented_out_function() {
+    //}
+
+    fn test_function() -> bool {
+        true
+    }
+}"#,
+        r#"contract;
+
+abi MyContract {
+    fn test_function() -> bool;
+}
+
+impl MyContract for Contract {
+    /// This is documentation for a commented out function
+    // fn commented_out_function() {
+    //}
+
+    fn test_function() -> bool {
+        true
+    }
+}
+"#,
+    );
+}