From 7bceb1c7e12449d78ad66ba821b93a91e971368b Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Thu, 29 Aug 2024 00:51:03 +0200
Subject: [PATCH] Trim trailing whitespace in comments (#53)

This patch adds trimming of trailing whitespace inside of comments in
addition to the trimming of trailing whitespace in code. Note that
trailing whitespace inside of multiline is not trimmed since doing so
would change the content of the string.

Closes #50.
---
 README.md        |  6 +++++-
 src/runestone.jl | 24 ++++++++++++++++--------
 test/runtests.jl |  4 ++++
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 8f1943c..c3f8aa0 100644
--- a/README.md
+++ b/README.md
@@ -657,10 +657,14 @@ Braces are consistently used around the right hand side of `where` expressions.
 
 #### Trailing spaces
 
-Trailing spaces are removed. Example:
+Trailing spaces are removed in code and comments (but not inside of multiline strings where
+doing so would change the meaning of the code). Examples:
 ```diff
 -1 + 1 
 +1 + 1
+
+-x = 2 # x is two 
++x = 2 # x is two
 ```
 
 #### Tabs
diff --git a/src/runestone.jl b/src/runestone.jl
index c4a2320..c2e751a 100644
--- a/src/runestone.jl
+++ b/src/runestone.jl
@@ -3,16 +3,24 @@
 # This is the runestone where all the formatting transformations are implemented.
 
 function trim_trailing_whitespace(ctx::Context, node::Node)
-    kind(node) === K"NewlineWs" || return nothing
+    kind(node) in KSet"NewlineWs Comment" || return nothing
     @assert is_leaf(node)
     str = String(read_bytes(ctx, node))
-    str′ = replace(str, r"\h*(\r\n|\r|\n)" => '\n')
-    # If the next sibling is also a NewlineWs we can trim trailing
-    # whitespace from this node too
-    next_kind = next_sibling_kind(ctx)
-    if next_kind === K"NewlineWs"
-        # str′ = replace(str′, r"(\r\n|\r|\n)\h*" => '\n')
-        str′ = replace(str′, r"\n\h*" => '\n')
+    local str′::String
+    if kind(node) === K"NewlineWs"
+        # Strip all whitespace up until the newline while normalizing line endings to \njK:w
+        str′ = replace(str, r"\h*(\r\n|\r|\n)" => '\n')
+        # If the next sibling is also a NewlineWs we can trim trailing
+        # whitespace from this node too
+        next_kind = next_sibling_kind(ctx)
+        if next_kind === K"NewlineWs"
+            # str′ = replace(str′, r"(\r\n|\r|\n)\h*" => '\n')
+            str′ = replace(str′, r"\n\h*" => '\n')
+        end
+    else
+        @assert kind(node) === K"Comment"
+        # Strip trailing spaces and tabs from comments
+        str′ = rstrip(str, (' ', '\t'))
     end
     if str == str′
         return nothing
diff --git a/test/runtests.jl b/test/runtests.jl
index a5563ce..a1a91df 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -70,6 +70,10 @@ end
     # Trailing whitespace just before closing indent token
     @test format_string("begin\n    a = 1 \nend") == "begin\n    a = 1\nend"
     @test format_string("let\n    a = 1 \nend") == "let\n    a = 1\nend"
+    # Trailing whitespace in comments
+    @test format_string("# comment ") == format_string("# comment  ") ==
+        format_string("# comment\t") == format_string("# comment\t\t") ==
+        format_string("# comment \t ") == format_string("# comment\t \t") == "# comment"
 end
 
 @testset "Hex/oct/bin literal integers" begin