From bd01c6d7827aba841fa5849b9adde7168ddf64da Mon Sep 17 00:00:00 2001
From: Julian Peeters <julianpeeters@gmail.com>
Date: Sun, 24 Dec 2023 14:13:48 -0800
Subject: [PATCH 1/5] Update handling of end-of-line single-line comments

Fix for #321, #605. Solution: distinguish between leading trivia a) previous trailing trivia, b) current leading trivia, c) a footer in case of last statement.
---
 .../mdoc/internal/markdown/Renderer.scala     | 43 ++++++++++++++++---
 .../scala/tests/markdown/DefaultSuite.scala   | 10 ++++-
 .../tests/markdown/VariablePrinterSuite.scala | 40 ++++++++++++++++-
 3 files changed, 84 insertions(+), 9 deletions(-)

diff --git a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
index 059addc17..51edd9c3b 100644
--- a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
+++ b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
@@ -2,6 +2,7 @@ package mdoc.internal.markdown
 
 import java.io.ByteArrayOutputStream
 import java.io.PrintStream
+import java.util.regex.Pattern
 import mdoc.Reporter
 import mdoc.Variable
 import mdoc.document.CompileResult
@@ -90,6 +91,7 @@ object Renderer {
     sb.appendMultiline(string, N)
   }
 
+  // Beneath each binding statement, we insert the evaluated variable, e.g., `x: Int = 1`
   def renderEvaluatedSection(
       doc: EvaluatedDocument,
       section: EvaluatedSection,
@@ -108,15 +110,43 @@ object Renderer {
     section.section.statements.zip(section.source.stats).zipWithIndex.foreach {
       case ((statement, tree), statementIndex) =>
         val pos = tree.pos
-        val leadingStart = stats(statementIndex - 1) match {
+        // for each statement, we need to manage:
+        //   1. the leading trivia: whitespace, newlines
+        //   2. a footer: empty until filled on the last statement in the section
+        //   3. and, internally, the trailing single-line comments of the previous statement
+        val (leading, footer) = stats(statementIndex - 1) match {
           case None =>
-            0
+            (Position.Range(input, 0, pos.start).text, "")
           case Some(previousStatement) =>
-            previousStatement.pos.end
+            // for each statement, we need to:
+            //   1. check the previous statement in order to find where the current statement begins
+            //   2. re-split the input, as a workaround to missing comments in 'stats' and `Position.Range`
+            //     a. use the stored comment-less variables to re-split the input
+            //     b. split again on the newline to collect a comment, the trivia, and maybe a footer
+            val escapedPrevStmt = Pattern.quote(previousStatement.toString())
+            val prevWithTrivia = section.source.pos.text.split(escapedPrevStmt, 2).toList
+            val (prevTrailingSingleLineComment, leadingTrivia, footerTrivia) =
+              if (prevWithTrivia.length == 2) {
+                val a2 = prevWithTrivia(1).split(Pattern.quote("\n"), 2)
+                val prevComment = a2(0)
+                val withTrivia = a2(1).split(Pattern.quote(tree.pos.text), 2)
+                val (trivia, foot) = {
+                  // ensure the last statement includes the footer if present
+                  if ((withTrivia.length == 2) && (statementIndex != (totalStats - 1))) (withTrivia(0), "")
+                  else if (withTrivia.length == 2) (withTrivia(0), withTrivia(1).dropWhile(_ != '\n'))
+                  else (tree.pos.text, "")
+                }
+                (prevComment, trivia, foot)
+              }
+              else ("","","")
+            val lead =
+              // if no trailing single-line comments, then we can use the established `Position.Range` system
+              if (prevWithTrivia.length == 0) Position.Range(input, previousStatement.pos.end, pos.start).text
+              else "\n" + leadingTrivia
+            (lead, footerTrivia)
         }
-        val leadingTrivia = Position.Range(input, leadingStart, pos.start)
         if (!section.mod.isFailOrWarn) {
-          sb.append(leadingTrivia.text)
+          sb.append(leading )
         }
         val endOfLinePosition =
           Position.Range(pos.input, pos.startLine, pos.startColumn, pos.endLine, Int.MaxValue)
@@ -177,8 +207,9 @@ object Renderer {
                 section.mod
               )
               sb.append(printer(variable))
-          }
+          } 
         }
+        sb.append(footer)
     }
     baos.toString.trim
   }
diff --git a/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala b/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala
index 6157960c2..bea487d83 100644
--- a/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala
+++ b/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala
@@ -211,6 +211,8 @@ class DefaultSuite extends BaseMarkdownSuite {
       |val (a, b) = Future.successful(Try(1)) -> 2
       |
       |Future.successful(Try(1))
+      |
+      |// penultimate
       |```
     """.stripMargin,
     """|```scala
@@ -226,6 +228,8 @@ class DefaultSuite extends BaseMarkdownSuite {
        |
        |Future.successful(Try(1))
        |// res1: Future[Try[Int]] = Future(Success(Success(1)))
+       |
+       |// penultimate
        |```
     """.stripMargin
   )
@@ -272,7 +276,8 @@ class DefaultSuite extends BaseMarkdownSuite {
       |
       |/** Docstring */
       |class User()
-      |```
+      |
+      |// ultimate```
     """.stripMargin,
     """|```scala
        |/* Comment 1 */
@@ -290,7 +295,8 @@ class DefaultSuite extends BaseMarkdownSuite {
        |
        |/** Docstring */
        |class User()
-       |```
+       |
+       |// ultimate```
     """.stripMargin
   )
 
diff --git a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
index df0a61c84..00e5093ab 100644
--- a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
+++ b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
@@ -3,6 +3,44 @@ import mdoc.internal.markdown.ReplVariablePrinter
 
 class VariablePrinterSuite extends BaseMarkdownSuite {
 
+  check(
+    "single-line-comment",
+    """
+      |```scala mdoc
+      |import scala.Some // an import statement
+      |val a = Some(1) // a variable
+      |val b = 2 // another variable
+      |```
+    """.stripMargin,
+    """|```scala
+       |import scala.Some // an import statement
+       |val a = Some(1) // a variable
+       |// a: Some[Int] = Some(value = 1)
+       |val b = 2 // another variable
+       |// b: Int = 2
+       |```
+    """.stripMargin,
+    baseSettings
+  )
+
+  check(
+    "single-line-comment:compile-only",
+    """|```scala mdoc:compile-only
+       |// a
+       |val a = 10
+       |val b = 20 // b
+       |val c = 30
+       |```""".stripMargin,
+    """|```scala
+       |// a
+       |val a = 10
+       |val b = 20 // b
+       |val c = 30
+       |```
+    """.stripMargin,
+    baseSettings
+  )
+
   val trailingComment = baseSettings.copy(variablePrinter = { variable =>
     variable.runtimeValue match {
       case n: Int if variable.totalVariablesInStatement == 1 => s" // Number($n)"
@@ -10,7 +48,7 @@ class VariablePrinterSuite extends BaseMarkdownSuite {
     }
   })
 
-  check(
+    check(
     "trailing-comment",
     """
       |```scala mdoc

From ebd446a2dd4a18503a39ef68922c2da80a3b09be Mon Sep 17 00:00:00 2001
From: Julian Peeters <julianpeeters@gmail.com>
Date: Sun, 24 Dec 2023 15:30:29 -0800
Subject: [PATCH 2/5] Formatting

---
 mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala       | 2 +-
 .../src/test/scala/tests/markdown/VariablePrinterSuite.scala    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
index 51edd9c3b..7d3e415f2 100644
--- a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
+++ b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
@@ -207,7 +207,7 @@ object Renderer {
                 section.mod
               )
               sb.append(printer(variable))
-          } 
+          }
         }
         sb.append(footer)
     }
diff --git a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
index 00e5093ab..896e07b51 100644
--- a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
+++ b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
@@ -48,7 +48,7 @@ class VariablePrinterSuite extends BaseMarkdownSuite {
     }
   })
 
-    check(
+  check(
     "trailing-comment",
     """
       |```scala mdoc

From c8a85ccaefe9e1de398b4999a3c535900ec77a81 Mon Sep 17 00:00:00 2001
From: Julian Peeters <julianpeeters@gmail.com>
Date: Sun, 24 Dec 2023 16:10:12 -0800
Subject: [PATCH 3/5] Fix test in CI

Use None instead of Some in order to avoid a Some value, which would otherwise be printed differently accross the test matrix.
---
 .../scala/tests/markdown/VariablePrinterSuite.scala    | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
index 896e07b51..0aa552b71 100644
--- a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
+++ b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
@@ -7,15 +7,15 @@ class VariablePrinterSuite extends BaseMarkdownSuite {
     "single-line-comment",
     """
       |```scala mdoc
-      |import scala.Some // an import statement
-      |val a = Some(1) // a variable
+      |import scala.None // an import statement
+      |val a = None // a variable
       |val b = 2 // another variable
       |```
     """.stripMargin,
     """|```scala
-       |import scala.Some // an import statement
-       |val a = Some(1) // a variable
-       |// a: Some[Int] = Some(value = 1)
+       |import scala.None // an import statement
+       |val a = None // a variable
+       |// a: None.type = None
        |val b = 2 // another variable
        |// b: Int = 2
        |```

From b68b7ba338931f6304b28042930c6e68d58b04fa Mon Sep 17 00:00:00 2001
From: Julian Peeters <julianpeeters@gmail.com>
Date: Fri, 29 Dec 2023 14:34:49 -0800
Subject: [PATCH 4/5] Accept PR suggestion: Why not use positions instead of
 spliting.

This strategy was overlooked, but is preferred since it reuses known machinery, and seems to introduce less complexity for the same behavior.

Co-authored-by: Tomasz Godzik <tgodzik@users.noreply.github.com>
---
 .../mdoc/internal/markdown/Renderer.scala     | 37 +++++++------------
 1 file changed, 14 insertions(+), 23 deletions(-)

diff --git a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
index 7d3e415f2..ccca8c6b2 100644
--- a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
+++ b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
@@ -118,32 +118,23 @@ object Renderer {
           case None =>
             (Position.Range(input, 0, pos.start).text, "")
           case Some(previousStatement) =>
-            // for each statement, we need to:
-            //   1. check the previous statement in order to find where the current statement begins
-            //   2. re-split the input, as a workaround to missing comments in 'stats' and `Position.Range`
-            //     a. use the stored comment-less variables to re-split the input
-            //     b. split again on the newline to collect a comment, the trivia, and maybe a footer
-            val escapedPrevStmt = Pattern.quote(previousStatement.toString())
-            val prevWithTrivia = section.source.pos.text.split(escapedPrevStmt, 2).toList
-            val (prevTrailingSingleLineComment, leadingTrivia, footerTrivia) =
-              if (prevWithTrivia.length == 2) {
-                val a2 = prevWithTrivia(1).split(Pattern.quote("\n"), 2)
-                val prevComment = a2(0)
-                val withTrivia = a2(1).split(Pattern.quote(tree.pos.text), 2)
-                val (trivia, foot) = {
-                  // ensure the last statement includes the footer if present
-                  if ((withTrivia.length == 2) && (statementIndex != (totalStats - 1))) (withTrivia(0), "")
-                  else if (withTrivia.length == 2) (withTrivia(0), withTrivia(1).dropWhile(_ != '\n'))
-                  else (tree.pos.text, "")
-                }
-                (prevComment, trivia, foot)
-              }
-              else ("","","")
+            val betweenStatements =
+              section.source.pos.text.substring(previousStatement.pos.end, pos.start)
+            val Array(prevTrailingSingleLineComment, leadingTrivia) =
+              betweenStatements.split(Pattern.quote("\n"), 2)
+            val footerAtEnd = {
+              if (statementIndex != (totalStats - 1)) ""
+              else section.source.pos.text.substring(pos.end)
+            }
             val lead =
               // if no trailing single-line comments, then we can use the established `Position.Range` system
-              if (prevWithTrivia.length == 0) Position.Range(input, previousStatement.pos.end, pos.start).text
+              if (prevTrailingSingleLineComment.length == 0)
+                Position.Range(input, previousStatement.pos.end, pos.start).text
               else "\n" + leadingTrivia
-            (lead, footerTrivia)
+            val footer = footerAtEnd.split(Pattern.quote("\n")).drop(1)
+            if (footer.nonEmpty)
+              (lead, footer.mkString("\n", "\n", ""))
+            else (lead, "")
         }
         if (!section.mod.isFailOrWarn) {
           sb.append(leading )

From 58c6a0861e2b1bd3a2fa0b234b50185a3064b7da Mon Sep 17 00:00:00 2001
From: Julian Peeters <julianpeeters@gmail.com>
Date: Fri, 29 Dec 2023 18:58:38 -0800
Subject: [PATCH 5/5] Refactor for clarity, simplify test

---
 .../mdoc/internal/markdown/Renderer.scala     | 22 +++++--------------
 .../scala/tests/markdown/DefaultSuite.scala   |  8 +++----
 .../tests/markdown/VariablePrinterSuite.scala | 10 ++++-----
 3 files changed, 14 insertions(+), 26 deletions(-)

diff --git a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
index ccca8c6b2..8fa5bf877 100644
--- a/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
+++ b/mdoc/src/main/scala/mdoc/internal/markdown/Renderer.scala
@@ -2,7 +2,6 @@ package mdoc.internal.markdown
 
 import java.io.ByteArrayOutputStream
 import java.io.PrintStream
-import java.util.regex.Pattern
 import mdoc.Reporter
 import mdoc.Variable
 import mdoc.document.CompileResult
@@ -118,26 +117,15 @@ object Renderer {
           case None =>
             (Position.Range(input, 0, pos.start).text, "")
           case Some(previousStatement) =>
-            val betweenStatements =
-              section.source.pos.text.substring(previousStatement.pos.end, pos.start)
             val Array(prevTrailingSingleLineComment, leadingTrivia) =
-              betweenStatements.split(Pattern.quote("\n"), 2)
-            val footerAtEnd = {
+              section.source.pos.text.substring(previousStatement.pos.end, pos.start).split("\n", 2)
+            val foot =
               if (statementIndex != (totalStats - 1)) ""
-              else section.source.pos.text.substring(pos.end)
-            }
-            val lead =
-              // if no trailing single-line comments, then we can use the established `Position.Range` system
-              if (prevTrailingSingleLineComment.length == 0)
-                Position.Range(input, previousStatement.pos.end, pos.start).text
-              else "\n" + leadingTrivia
-            val footer = footerAtEnd.split(Pattern.quote("\n")).drop(1)
-            if (footer.nonEmpty)
-              (lead, footer.mkString("\n", "\n", ""))
-            else (lead, "")
+              else section.source.pos.text.substring(pos.end).split("\n").drop(1).mkString("\n", "\n", "")
+            ("\n" + leadingTrivia, foot)
         }
         if (!section.mod.isFailOrWarn) {
-          sb.append(leading )
+          sb.append(leading)
         }
         val endOfLinePosition =
           Position.Range(pos.input, pos.startLine, pos.startColumn, pos.endLine, Int.MaxValue)
diff --git a/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala b/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala
index bea487d83..11c4ab29d 100644
--- a/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala
+++ b/tests/unit/src/test/scala/tests/markdown/DefaultSuite.scala
@@ -210,7 +210,7 @@ class DefaultSuite extends BaseMarkdownSuite {
       |Future.successful(Try(1))
       |val (a, b) = Future.successful(Try(1)) -> 2
       |
-      |Future.successful(Try(1))
+      |Future.successful(Try(1)) // last statement
       |
       |// penultimate
       |```
@@ -226,7 +226,7 @@ class DefaultSuite extends BaseMarkdownSuite {
        |// a: Future[Try[Int]] = Future(Success(Success(1)))
        |// b: Int = 2
        |
-       |Future.successful(Try(1))
+       |Future.successful(Try(1)) // last statement
        |// res1: Future[Try[Int]] = Future(Success(Success(1)))
        |
        |// penultimate
@@ -275,7 +275,7 @@ class DefaultSuite extends BaseMarkdownSuite {
       |x + y
       |
       |/** Docstring */
-      |class User()
+      |class User() // Comment 5
       |
       |// ultimate```
     """.stripMargin,
@@ -294,7 +294,7 @@ class DefaultSuite extends BaseMarkdownSuite {
        |// res0: Int = 4
        |
        |/** Docstring */
-       |class User()
+       |class User() // Comment 5
        |
        |// ultimate```
     """.stripMargin
diff --git a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
index 0aa552b71..5f7c9f5eb 100644
--- a/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
+++ b/tests/unit/src/test/scala/tests/markdown/VariablePrinterSuite.scala
@@ -7,15 +7,15 @@ class VariablePrinterSuite extends BaseMarkdownSuite {
     "single-line-comment",
     """
       |```scala mdoc
-      |import scala.None // an import statement
-      |val a = None // a variable
+      |import scala.Int // an import statement
+      |val a: Int = 1 // a variable
       |val b = 2 // another variable
       |```
     """.stripMargin,
     """|```scala
-       |import scala.None // an import statement
-       |val a = None // a variable
-       |// a: None.type = None
+       |import scala.Int // an import statement
+       |val a: Int = 1 // a variable
+       |// a: Int = 1
        |val b = 2 // another variable
        |// b: Int = 2
        |```