Support comments on attributes (#5850)

Signed-off-by: Nick Cameron <nrc@ncameron.org>

rust/kcl-lib/src/parsing/ast/types/mod.rs

@@ -53,6 +53,12 @@ pub struct Node<T> {
     pub module_id: ModuleId,
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub outer_attrs: NodeList<Annotation>,
+    // Some comments are kept here, some are kept in NonCodeMeta, and some are ignored. See how each
+    // node is parsed to check for certain. In any case, only comments which are strongly associated
+    // with an item are kept here.
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub pre_comments: Vec<String>,
+    pub comment_start: usize,
 }
 
 impl<T: JsonSchema> schemars::JsonSchema for Node<T> {
@@ -85,6 +91,20 @@ impl<T> Node<T> {
             end,
             module_id,
             outer_attrs: Vec::new(),
+            pre_comments: Vec::new(),
+            comment_start: start,
+        }
+    }
+
+    pub fn new_node(start: usize, end: usize, module_id: ModuleId, inner: T) -> Self {
+        Self {
+            inner,
+            start,
+            end,
+            module_id,
+            outer_attrs: Vec::new(),
+            pre_comments: Vec::new(),
+            comment_start: start,
         }
     }
 
@@ -95,6 +115,8 @@ impl<T> Node<T> {
             end: 0,
             module_id: ModuleId::default(),
             outer_attrs: Vec::new(),
+            pre_comments: Vec::new(),
+            comment_start: 0,
         }
     }
 
@@ -105,6 +127,8 @@ impl<T> Node<T> {
             end,
             module_id,
             outer_attrs: Vec::new(),
+            pre_comments: Vec::new(),
+            comment_start: start,
         })
     }
 
@@ -133,8 +157,15 @@ impl<T> Node<T> {
             end: self.end,
             module_id: self.module_id,
             outer_attrs: self.outer_attrs,
+            pre_comments: self.pre_comments,
+            comment_start: self.comment_start,
         }
     }
+
+    pub fn set_comments(&mut self, comments: Vec<String>, start: usize) {
+        self.pre_comments = comments;
+        self.comment_start = start;
+    }
 }
 
 impl<T> Deref for Node<T> {
@@ -373,6 +404,26 @@ impl Program {
         if self.non_code_meta.in_comment(pos) {
             return true;
         }
+
+        for item in &self.body {
+            let r = item.comment_range();
+            eprintln!("item {r:?}");
+            if pos >= r.0 && pos < r.1 {
+                return true;
+            }
+            if pos < r.0 {
+                break;
+            }
+        }
+        for n in &self.inner_attrs {
+            if pos >= n.comment_start && pos < n.start {
+                return true;
+            }
+            if pos < n.comment_start {
+                break;
+            }
+        }
+
         let item = self.get_body_item_for_position(pos);
 
         // Recurse over the item.
@@ -660,6 +711,36 @@ impl BodyItem {
             BodyItem::ReturnStatement(node) => &mut node.outer_attrs,
         }
     }
+
+    pub(crate) fn set_comments(&mut self, comments: Vec<String>, start: usize) {
+        match self {
+            BodyItem::ImportStatement(node) => node.set_comments(comments, start),
+            BodyItem::ExpressionStatement(node) => node.set_comments(comments, start),
+            BodyItem::VariableDeclaration(node) => node.set_comments(comments, start),
+            BodyItem::TypeDeclaration(node) => node.set_comments(comments, start),
+            BodyItem::ReturnStatement(node) => node.set_comments(comments, start),
+        }
+    }
+
+    pub(crate) fn get_comments(&self) -> &[String] {
+        match self {
+            BodyItem::ImportStatement(node) => &node.pre_comments,
+            BodyItem::ExpressionStatement(node) => &node.pre_comments,
+            BodyItem::VariableDeclaration(node) => &node.pre_comments,
+            BodyItem::TypeDeclaration(node) => &node.pre_comments,
+            BodyItem::ReturnStatement(node) => &node.pre_comments,
+        }
+    }
+
+    pub(crate) fn comment_range(&self) -> (usize, usize) {
+        match self {
+            BodyItem::ImportStatement(node) => (node.comment_start, node.start),
+            BodyItem::ExpressionStatement(node) => (node.comment_start, node.start),
+            BodyItem::VariableDeclaration(node) => (node.comment_start, node.start),
+            BodyItem::TypeDeclaration(node) => (node.comment_start, node.start),
+            BodyItem::ReturnStatement(node) => (node.comment_start, node.start),
+        }
+    }
 }
 
 impl From<BodyItem> for SourceRange {
@@ -1171,6 +1252,23 @@ pub enum CommentStyle {
     Block,
 }
 
+impl CommentStyle {
+    pub fn render_comment(&self, comment: &str) -> String {
+        match self {
+            CommentStyle::Line => {
+                let comment = comment.trim();
+                let mut result = "//".to_owned();
+                if !comment.is_empty() && !comment.starts_with('/') {
+                    result.push(' ');
+                }
+                result.push_str(comment);
+                result
+            }
+            CommentStyle::Block => format!("/* {comment} */"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
 #[ts(export)]
 #[serde(tag = "type", rename_all = "camelCase")]
@@ -1186,7 +1284,7 @@ pub enum NonCodeValue {
     },
     /// A block comment.
     /// An example of this is the following:
-    /// ```python,no_run
+    /// ```no_run
     /// /* This is a
     ///     block comment */
     /// 1 + 1
@@ -3890,7 +3988,6 @@ startSketchOn('XY')"#;
             formatted,
             r#"@settings(defaultLengthUnit = mm)
 
-
 startSketchOn('XY')
 "#
         );
@@ -3925,6 +4022,7 @@ startSketchOn('XY')
         assert_eq!(
             formatted,
             r#"@settings(defaultLengthUnit = mm)
+
 startSketchOn('XY')
 "#
         );