From ebeb5dbbea26ea4bc458cbf4f18563ede13b1d0b Mon Sep 17 00:00:00 2001
From: Nicolas BACQUEY <nicolas.bacquey@tweag.io>
Date: Wed, 4 Dec 2024 12:27:58 +0100
Subject: [PATCH] Add tests for comment extraction

---
 topiary-core/src/comments.rs        | 12 ++--
 topiary-core/src/lib.rs             |  6 +-
 topiary-core/tests/comment_tests.rs | 92 +++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 9 deletions(-)
 create mode 100644 topiary-core/tests/comment_tests.rs

diff --git a/topiary-core/src/comments.rs b/topiary-core/src/comments.rs
index ea5c3999..97a7615e 100644
--- a/topiary-core/src/comments.rs
+++ b/topiary-core/src/comments.rs
@@ -94,13 +94,13 @@ fn find_comments<'a>(
 
 /// The section of code to which a comment refers. We also remember whether the comment
 /// is positioned before or after the section.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum Commented {
     /// The code section is before the comment, as in:
     /// ```
     /// struct Foo {
-    ///     baz: Baz, // this is baz
-    ///     quz: Qux, // this is qux
+    ///     baz: usize, // this is baz
+    ///     quz: usize, // this is qux
     /// }
     /// ```
     CommentedBefore(InputSection),
@@ -108,9 +108,9 @@ pub enum Commented {
     /// ```
     /// struct Foo {
     ///     // let's have a baz
-    ///     baz: Baz,
+    ///     baz: usize,
     ///     // and a qux
-    ///     qux: Qux,
+    ///     qux: usize,
     /// }
     /// ```
     CommentedAfter(InputSection),
@@ -282,7 +282,7 @@ fn find_anchor<'tree>(node: &'tree Node<'tree>, input: &str) -> FormatterResult<
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub struct AnchoredComment {
     pub comment_text: String,
     pub commented: Commented,
diff --git a/topiary-core/src/lib.rs b/topiary-core/src/lib.rs
index f30c95b0..8b8674a7 100644
--- a/topiary-core/src/lib.rs
+++ b/topiary-core/src/lib.rs
@@ -23,13 +23,13 @@ pub use crate::{
 };
 
 mod atom_collection;
-mod comments;
+pub mod comments;
 mod error;
 mod graphviz;
 mod language;
 mod pretty;
-mod tree_sitter;
-mod types;
+pub mod tree_sitter;
+pub mod types;
 
 #[doc(hidden)]
 pub mod test_utils;
diff --git a/topiary-core/tests/comment_tests.rs b/topiary-core/tests/comment_tests.rs
new file mode 100644
index 00000000..876f245a
--- /dev/null
+++ b/topiary-core/tests/comment_tests.rs
@@ -0,0 +1,92 @@
+use topiary_core::{
+    comments::{
+        extract_comments, AnchoredComment,
+        Commented::{CommentedAfter, CommentedBefore},
+        SeparatedInput,
+    },
+    tree_sitter::parse,
+    types::InputSection,
+    Position,
+};
+
+const OCAML_WITH_COMMENTS: &str = r#"(* starting comment *)
+fun (* fun comment *) x (* var comment *) ->
+  (** multi-lined
+    * body comment
+    *)
+  body
+(* final comment *)
+"#;
+
+const OCAML_WITHOUT_COMMENTS: &str = r#"
+fun  x  ->
+  
+  body
+
+"#;
+
+// The section corresponding to `fun` in the curated code
+const FUN_SECTION: InputSection = InputSection {
+    start: Position { row: 2, column: 1 },
+    end: Position { row: 2, column: 4 },
+};
+
+// The section corresponding to `x` in the curated code
+const VAR_SECTION: InputSection = InputSection {
+    start: Position { row: 2, column: 6 },
+    end: Position { row: 2, column: 7 },
+};
+
+// The section corresponding to `body` in the curated code
+const BODY_SECTION: InputSection = InputSection {
+    start: Position { row: 4, column: 3 },
+    end: Position { row: 4, column: 7 },
+};
+
+#[test]
+fn test_extract_comments() {
+    let input = OCAML_WITH_COMMENTS;
+    let ocaml = tree_sitter_ocaml::LANGUAGE_OCAML;
+
+    let (tree, _grammar) = parse(input, &ocaml.into(), false).unwrap();
+
+    let SeparatedInput {
+        input_tree: _,
+        input_string: new_input_string,
+        mut comments,
+    } = extract_comments(&tree, input, &ocaml.into()).unwrap();
+
+    let mut expected_comments: Vec<AnchoredComment> = vec![
+        AnchoredComment {
+            comment_text: "(* starting comment *)".into(),
+            commented: CommentedAfter(FUN_SECTION),
+        },
+        AnchoredComment {
+            comment_text: "(* fun comment *)".into(),
+            commented: CommentedBefore(FUN_SECTION),
+        },
+        AnchoredComment {
+            comment_text: "(* var comment *)".into(),
+            commented: CommentedBefore(VAR_SECTION),
+        },
+        AnchoredComment {
+            comment_text: "(** multi-lined\n    * body comment\n    *)".into(),
+            commented: CommentedAfter(BODY_SECTION),
+        },
+        AnchoredComment {
+            comment_text: "(* final comment *)".into(),
+            commented: CommentedBefore(BODY_SECTION),
+        },
+    ];
+
+    // sort the comments so that we're order-independent
+    comments.sort_by_key(|comment| comment.comment_text.clone());
+    expected_comments.sort_by_key(|comment| comment.comment_text.clone());
+
+    assert_eq!(new_input_string, OCAML_WITHOUT_COMMENTS);
+
+    assert_eq!(comments.len(), 5);
+    for (comment, expected_comment) in comments.iter().zip(expected_comments.iter()) {
+        assert_eq!(comment, expected_comment)
+    }
+}