TODO: re-attach comments

topiary-core/src/atom_collection.rs

@@ -8,7 +8,7 @@ use std::{
 use topiary_tree_sitter_facade::Node;
 
 use crate::{
-    tree_sitter::NodeExt, Atom, FormatterError, FormatterResult, ScopeCondition, ScopeInformation,
+    comments::{AnchoredComment, Comment}, tree_sitter::NodeExt, types::InputSection, Atom, FormatterError, FormatterResult, ScopeCondition, ScopeInformation
 };
 
 /// A struct that holds sets of node IDs that have line breaks before or after them.
@@ -37,6 +37,8 @@ pub struct AtomCollection {
     /// something to a node, a new Atom is added to this HashMap.
     /// The key of the hashmap is the identifier of the node.
     append: HashMap<usize, Vec<Atom>>,
+    /// 
+    comments: HashMap<usize, Vec<Comment>>,
     /// A query file can define custom leaf nodes (nodes that Topiary should not
     /// touch during formatting). When such a node is encountered, its id is stored in
     /// this HashSet.
@@ -72,6 +74,7 @@ impl AtomCollection {
             atoms,
             prepend: HashMap::new(),
             append: HashMap::new(),
+            comments: HashMap::new(),
             specified_leaf_nodes: HashSet::new(),
             parent_leaf_nodes: HashMap::new(),
             multi_line_nodes: HashSet::new(),
@@ -87,6 +90,7 @@ impl AtomCollection {
         root: &Node,
         source: &[u8],
         specified_leaf_nodes: HashSet<usize>,
+        comments: Vec<AnchoredComment>,
     ) -> FormatterResult<Self> {
         // Flatten the tree, from the root node, in a depth-first traversal
         let dfs_nodes = dfs_flatten(root);
@@ -100,6 +104,7 @@ impl AtomCollection {
             atoms: Vec::new(),
             prepend: HashMap::new(),
             append: HashMap::new(),
+            comments: HashMap::new(),
             specified_leaf_nodes,
             parent_leaf_nodes: HashMap::new(),
             multi_line_nodes,
@@ -109,7 +114,7 @@ impl AtomCollection {
             counter: 0,
         };
 
-        atoms.collect_leafs_inner(root, source, &Vec::new(), 0)?;
+        atoms.collect_leafs_inner(root, source, &comments, &Vec::new(), 0)?;
 
         Ok(atoms)
     }
@@ -509,6 +514,7 @@ impl AtomCollection {
         &mut self,
         node: &Node,
         source: &[u8],
+        comments: &mut Vec<AnchoredComment>,
         parent_ids: &[usize],
         level: usize,
     ) -> FormatterResult<()> {
@@ -541,9 +547,19 @@ impl AtomCollection {
             });
             // Mark all sub-nodes as having this node as a "leaf parent"
             self.mark_leaf_parent(node, node.id());
+            // Test the node for comments
+            while let Some(comment) = comments.last() {
+                let node_section: InputSection = node.into();
+                while node_section.contains(&comment.into()) {
+                    self.comments.entry(id)
+                        .or_insert(Vec::new())
+                        .push(comment.into());
+                    comments.pop();
+                }
+            }
         } else {
             for child in node.children(&mut node.walk()) {
-                self.collect_leafs_inner(&child, source, &parent_ids, level + 1)?;
+                self.collect_leafs_inner(&child, source, comments, &parent_ids, level + 1)?;
             }
         }
 

topiary-core/src/comments.rs

@@ -127,6 +127,42 @@ impl Diff<InputSection> for Commented {
     }
 }
 
+/// A comment, as part of Topiary's output.
+/// We forget posiiton information here, because the struct
+/// is supposed to be attached to the node it comments.
+#[derive(Debug)]
+pub enum Comment {
+    /// The comment is before the code section, as is:
+    /// ```
+    /// struct Foo {
+    ///     // let's have a baz
+    ///     baz: usize,
+    ///     // and a qux
+    ///     qux: usize,
+    /// }
+    /// ```
+    CommentBefore(String),
+    /// The comment is after the code section, as in:
+    /// The code section is before the comment, as in:
+    /// ```
+    /// struct Foo {
+    ///     baz: usize, // this is baz
+    ///     quz: usize, // this is qux
+    /// }
+    /// ```
+    CommentAfter(String),
+}
+
+impl From<&AnchoredComment> for Comment {
+    fn from(value: &AnchoredComment) -> Self {
+        let content = value.comment_text.clone();
+        match value.commented {
+            Commented::CommentedAfter(_) => Comment::CommentBefore(content),
+            Commented::CommentedBefore(_) => Comment::CommentAfter(content),
+        }
+    }
+}
+
 // TODO: if performance is an issue, use TreeCursor to navigate the tree
 fn next_non_comment_leaf<'tree>(starting_node: Node<'tree>) -> Option<Node<'tree>> {
     let mut node: Node<'tree> = starting_node;
@@ -259,9 +295,9 @@ fn find_anchor<'tree>(node: &'tree Node<'tree>, input: &str) -> FormatterResult<
         })?;
     if prefix.trim_start() == "" {
         if let Some(anchor) = next_non_comment_leaf(node.clone()) {
-            return Ok(Commented::CommentedAfter(anchor.into()));
+            return Ok(Commented::CommentedAfter((&anchor).into()));
         } else if let Some(anchor) = previous_non_comment_leaf(node.clone()) {
-            return Ok(Commented::CommentedBefore(anchor.into()));
+            return Ok(Commented::CommentedBefore((&anchor).into()));
         } else {
             return Err(FormatterError::Internal(
                 format!("Could find no anchor for comment {node:?}",),
@@ -270,9 +306,9 @@ fn find_anchor<'tree>(node: &'tree Node<'tree>, input: &str) -> FormatterResult<
         }
     } else {
         if let Some(anchor) = previous_non_comment_leaf(node.clone()) {
-            return Ok(Commented::CommentedBefore(anchor.into()));
+            return Ok(Commented::CommentedBefore((&anchor).into()));
         } else if let Some(anchor) = next_non_comment_leaf(node.clone()) {
-            return Ok(Commented::CommentedAfter(anchor.into()));
+            return Ok(Commented::CommentedAfter((&anchor).into()));
         } else {
             return Err(FormatterError::Internal(
                 format!("Could find no anchor for comment {node:?}",),
@@ -288,6 +324,15 @@ pub struct AnchoredComment {
     pub commented: Commented,
 }
 
+impl From<&AnchoredComment> for InputSection {
+    fn from(value: &AnchoredComment) -> Self {
+        match value.commented {
+            Commented::CommentedBefore(section) => section,
+            Commented::CommentedAfter(section) => section,
+        }
+    }
+}
+
 pub struct SeparatedInput {
     pub input_tree: Tree,
     pub input_string: String,
@@ -339,7 +384,7 @@ pub fn extract_comments<'a>(
                      comment_text,
                  }|
                  -> FormatterResult<AnchoredComment> {
-                    commented.subtract(comment.clone().into())?;
+                    commented.subtract((&comment).into())?;
                     Ok(AnchoredComment {
                         commented,
                         comment_text: comment_text.to_string(),

topiary-core/src/types.rs

@@ -38,8 +38,14 @@ pub struct InputSection {
     pub end: Position,
 }
 
-impl From<Node<'_>> for InputSection {
-    fn from(value: Node) -> Self {
+impl InputSection {
+    pub fn contains(self: Self, other: &Self) -> bool {
+        self.start <= other.start && other.end <= self.end
+    }
+}
+
+impl From<&Node<'_>> for InputSection {
+    fn from(value: &Node) -> Self {
         InputSection {
             start: value.start_position().into(),
             end: value.end_position().into(),