feat(anodized-fmt): handle arrays formatting in specs by refactoring comment extraction and formatting

Author: chammika

crates/anodized-fmt/src/collect_comments.rs

@@ -1,211 +1,144 @@
 use std::collections::HashMap;
 
 use crop::Rope;
-use proc_macro2::{LineColumn, Span, TokenStream};
+use proc_macro2::LineColumn;
 
-/// Extract comments and empty lines from the gaps between tokens.
+/// Extract standalone comments (comments on their own line, no code before them).
 ///
-/// This function traverses a TokenStream and examines the text between consecutive
-/// tokens. When tokens appear on different lines, it extracts any comments (lines
-/// starting with //) and records empty lines.
+/// Uses a scan_gap-like approach but only extracts comments that are on their own line.
+/// This avoids issues with inline comments and comments inside nested structures
+/// which would be lost during prettyplease formatting.
 ///
 /// Returns a HashMap mapping line numbers (0-indexed) to optional comment text:
-/// - Some(comment_text) for lines with comments
+/// - Some(comment_text) for lines with standalone comments
 /// - None for empty lines (preserves vertical spacing)
-pub(crate) fn extract_whitespace_and_comments(
-    source: &Rope,
-    tokens: TokenStream,
+pub(crate) fn extract_standalone_comments(
+    source: &str,
+    start_line: usize,
 ) -> HashMap<usize, Option<String>> {
-    let mut whitespace_and_comments = HashMap::new();
-    let mut last_span: Option<Span> = None;
-
-    traverse_token_stream(tokens, &mut |span: Span| {
-        if let Some(last_span) = last_span
-            && last_span.end().line != span.start().line
-        {
-            let text = get_text_between_spans(source, last_span.end(), span.start());
-            for (idx, line) in text.lines().enumerate() {
-                let comment = line
-                    .to_string()
-                    .split_once("//")
-                    .map(|(_, txt)| txt)
-                    .map(str::trim)
-                    .map(ToOwned::to_owned);
-
-                let line_index = last_span.end().line - 1 + idx;
-
-                // Skip empty lines at token boundaries, but keep comments
-                if comment.is_none()
-                    && (line_index == last_span.end().line - 1
-                        || line_index == span.start().line - 1)
-                {
-                    continue;
-                }
-
-                whitespace_and_comments.insert(line_index, comment);
+    let mut result = HashMap::new();
+
+    for (idx, line) in source.lines().enumerate() {
+        let line_index = start_line + idx;
+        let trimmed = line.trim();
+
+        // Check if line is a standalone comment (only whitespace before //)
+        if let Some(comment_text) = trimmed.strip_prefix("//") {
+            let comment_text = comment_text.trim().to_string();
+            if !comment_text.is_empty() {
+                result.insert(line_index, Some(comment_text));
+            } else {
+                // Empty comment line
+                result.insert(line_index, None);
             }
+        } else if trimmed.is_empty() {
+            // Empty line (no comment, no code)
+            result.insert(line_index, None);
         }
-        last_span = Some(span);
-    });
+    }
 
-    whitespace_and_comments
+    result
 }
 
-/// Recursively traverse a TokenStream, calling the callback for each token's span.
-fn traverse_token_stream(tokens: TokenStream, cb: &mut impl FnMut(Span)) {
-    for token in tokens {
-        match token {
-            proc_macro2::TokenTree::Group(group) => {
-                cb(group.span_open());
-                traverse_token_stream(group.stream(), cb);
-                cb(group.span_close());
+/// Check if source text has inline comments (comments after code on the same line)
+/// or comments inside block/paren delimiters.
+///
+/// Comments inside `[]` brackets are allowed — those are handled by the formatter's
+/// comment-aware array formatting. Comments inside `{}` or `()` are problematic
+/// because prettyplease strips them when reformatting expressions.
+pub(crate) fn has_problematic_comments(source: &str) -> bool {
+    for line in source.lines() {
+        if let Some((before, _after)) = line.split_once("//") {
+            let before_trimmed = before.trim();
+            if !before_trimmed.is_empty() && !before_trimmed.starts_with("//") {
+                return true;
             }
-            _ => cb(token.span()),
         }
     }
-}
 
-/// Extract text from the source Rope between two line/column positions.
-fn get_text_between_spans(rope: &Rope, start: LineColumn, end: LineColumn) -> String {
-    let start_byte = line_column_to_byte(rope, start);
-    let end_byte = line_column_to_byte(rope, end);
+    let mut block_depth: i32 = 0;
 
-    rope.byte_slice(start_byte..end_byte).to_string()
-}
+    for line in source.lines() {
+        let trimmed = line.trim();
 
-/// Convert a LineColumn position to a byte offset in the Rope.
-pub fn line_column_to_byte(source: &Rope, point: LineColumn) -> usize {
-    let line_byte = source.byte_of_line(point.line - 1);
-    let line = source.line(point.line - 1);
-    let char_byte: usize = line.chars().take(point.column).map(|c| c.len_utf8()).sum();
-    line_byte + char_byte
-}
-
-/// Check if a TokenStream contains comments inside nested structures.
-pub fn has_nested_structure_comments(source: &Rope, tokens: TokenStream) -> bool {
-    has_nested_structure_comments_inner(source, tokens, 0)
-}
+        let code_part = trimmed
+            .split_once("//")
+            .map_or(trimmed, |(before, _)| before);
 
-// Recursively check for comments inside nested groups, tracking depth to avoid false positives.
-fn has_nested_structure_comments_inner(source: &Rope, tokens: TokenStream, _depth: usize) -> bool {
-    for token in tokens {
-        if let proc_macro2::TokenTree::Group(group) = token {
-            // Check for comments inside any nested structure (arrays, blocks, structs, etc.)
-            // We need to include the opening delimiter to check for comments after the opening brace/bracket
-            let mut group_comments = HashMap::new();
-            let opening_span = group.span_open();
-
-            // Check gaps between tokens inside the group, including after the opening delimiter
-            let mut last_span: Option<Span> = Some(opening_span);
-            traverse_token_stream(group.stream(), &mut |span: Span| {
-                if let Some(last) = last_span
-                    && last.end().line != span.start().line
-                {
-                    let text = get_text_between_spans(source, last.end(), span.start());
-                    for (idx, line) in text.lines().enumerate() {
-                        let comment = line
-                            .to_string()
-                            .split_once("//")
-                            .map(|(_, txt)| txt)
-                            .map(str::trim)
-                            .map(ToOwned::to_owned);
-
-                        let line_index = last.end().line - 1 + idx;
-
-                        if comment.is_some() {
-                            group_comments.insert(line_index, comment);
-                        }
-                    }
-                }
-                last_span = Some(span);
-            });
-
-            if !group_comments.is_empty() {
-                return true;
+        for ch in code_part.chars() {
+            match ch {
+                '{' => block_depth += 1,
+                '}' => block_depth = block_depth.saturating_sub(1),
+                _ => {}
             }
+        }
 
-            // Recursively check nested groups at increased depth
-            if has_nested_structure_comments_inner(source, group.stream(), _depth + 1) {
-                return true;
-            }
+        if block_depth > 0 && trimmed.starts_with("//") {
+            return true;
         }
     }
 
     false
 }
 
+/// Convert a LineColumn position to a byte offset in the Rope.
+pub fn line_column_to_byte(source: &Rope, point: LineColumn) -> usize {
+    let line_byte = source.byte_of_line(point.line - 1);
+    let line = source.line(point.line - 1);
+    let char_byte: usize = line.chars().take(point.column).map(|c| c.len_utf8()).sum();
+    line_byte + char_byte
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_extract_simple_comment() {
-        let source_text = r#"
-            requires: x > 0,
-            // This is a comment
-            ensures: *output > 0
-            "#;
-        let rope = Rope::from(source_text);
-        let tokens: TokenStream = source_text.parse().unwrap();
-
-        let comments = extract_whitespace_and_comments(&rope, tokens);
-
-        // Should find the comment on line 2 (0-indexed)
-        assert!(comments.contains_key(&2));
+    fn test_extract_standalone_comments() {
+        let source = "requires: x > 0,\n// This is a comment\nensures: *output > 0";
+        let comments = extract_standalone_comments(source, 0);
+
         assert_eq!(
-            comments.get(&2),
+            comments.get(&1),
             Some(&Some("This is a comment".to_string()))
         );
+        assert!(!comments.contains_key(&0)); // Code line, not a comment
+        assert!(!comments.contains_key(&2)); // Code line, not a comment
     }
 
     #[test]
-    fn test_no_comments() {
-        let source_text = r#"requires: x > 0, ensures: *output > 0"#;
-        let rope = Rope::from(source_text);
-        let tokens: TokenStream = source_text.parse().unwrap();
-
-        let comments = extract_whitespace_and_comments(&rope, tokens);
-
-        // Should be empty - no gaps between lines
-        assert!(comments.is_empty());
+    fn test_extract_multiple_comments_and_empty_lines() {
+        let source = "// First comment\n\n// Second comment\nrequires: x > 0";
+        let comments = extract_standalone_comments(source, 5);
+
+        assert_eq!(comments.get(&5), Some(&Some("First comment".to_string())));
+        assert_eq!(comments.get(&6), Some(&None)); // Empty line
+        assert_eq!(comments.get(&7), Some(&Some("Second comment".to_string())));
+        assert!(!comments.contains_key(&8)); // Code line
     }
 
     #[test]
-    fn test_multiple_comments() {
-        let source_text = r#"
-            requires: x > 0,
-            // First comment
-            // Second comment
-            ensures: *output > 0,
-            // Third comment
-            binds: result
-            "#;
-        let rope = Rope::from(source_text);
-        let tokens: TokenStream = source_text.parse().unwrap();
-
-        let comments = extract_whitespace_and_comments(&rope, tokens);
-
-        // Should find all three comments
-        assert!(comments.len() == 3);
-        assert_eq!(comments.get(&2), Some(&Some("First comment".to_string())));
-        assert_eq!(comments.get(&3), Some(&Some("Second comment".to_string())));
-        assert_eq!(comments.get(&5), Some(&Some("Third comment".to_string())));
+    fn test_has_problematic_comments_inline() {
+        let source = "requires: x > 0, // inline comment\nensures: *output > 0";
+        assert!(has_problematic_comments(source));
     }
 
     #[test]
-    fn test_empty_lines() {
-        let source_text = r#"
-            requires: x > 0,
-
-            ensures: *output > 0
-            "#;
-        let rope = Rope::from(source_text);
-        let tokens: TokenStream = source_text.parse().unwrap();
+    fn test_has_problematic_comments_in_block() {
+        let source = "requires: {\n    // comment in block\n    x > 0\n}";
+        assert!(has_problematic_comments(source));
+    }
 
-        let comments = extract_whitespace_and_comments(&rope, tokens);
+    #[test]
+    fn test_has_no_problematic_comments_in_array() {
+        let source = "requires: [\n    // comment in array\n    x > 0,\n    y > 0\n]";
+        assert!(!has_problematic_comments(source));
+    }
 
-        // Should record the empty line as None
-        let has_empty_line = comments.values().any(|v| v.is_none());
-        assert!(has_empty_line);
+    #[test]
+    fn test_has_no_problematic_comments_standalone() {
+        let source =
+            "// standalone comment\nrequires: x > 0,\n// another standalone\nensures: *output > 0";
+        assert!(!has_problematic_comments(source));
     }
 }