From d9a9d0b008bd29a13f9811aefa674bd6119d3cbd Mon Sep 17 00:00:00 2001
From: Rob Dixon <rob@mindmage.ai>
Date: Tue, 19 Nov 2024 17:33:22 +0000
Subject: [PATCH] Apply whitespace rules for LStrip and Trim to comment blocks

---
 .../com/hubspot/jinjava/tree/TreeParser.java  | 19 +++++++++++--------
 .../jinjava/tree/parse/TokenScanner.java      | 17 +++++++++++------
 .../hubspot/jinjava/tree/TreeParserTest.java  | 12 ++++++++++++
 .../tree/parse/TokenWhitespaceTest.java       | 10 ++++++++++
 .../tokenizer/whitespace-comment-tags.jinja   |  8 ++++++++
 5 files changed, 52 insertions(+), 14 deletions(-)
 create mode 100644 src/test/resources/parse/tokenizer/whitespace-comment-tags.jinja

diff --git a/src/main/java/com/hubspot/jinjava/tree/TreeParser.java b/src/main/java/com/hubspot/jinjava/tree/TreeParser.java
index 8bdc29653..9904e7c7a 100644
--- a/src/main/java/com/hubspot/jinjava/tree/TreeParser.java
+++ b/src/main/java/com/hubspot/jinjava/tree/TreeParser.java
@@ -170,14 +170,17 @@ private Node getLastSibling() {
 
   private Node text(TextToken textToken) {
     if (interpreter.getConfig().isLstripBlocks()) {
-      if (scanner.hasNext() && scanner.peek().getType() == symbols.getTag()) {
-        textToken =
-          new TextToken(
-            StringUtils.stripEnd(textToken.getImage(), "\t "),
-            textToken.getLineNumber(),
-            textToken.getStartPosition(),
-            symbols
-          );
+      if (scanner.hasNext()) {
+        final int nextTokenType = scanner.peek().getType();
+        if (nextTokenType == symbols.getTag() || nextTokenType == symbols.getNote()) {
+          textToken =
+            new TextToken(
+              StringUtils.stripEnd(textToken.getImage(), "\t "),
+              textToken.getLineNumber(),
+              textToken.getStartPosition(),
+              symbols
+            );
+        }
       }
     }
 
diff --git a/src/main/java/com/hubspot/jinjava/tree/parse/TokenScanner.java b/src/main/java/com/hubspot/jinjava/tree/parse/TokenScanner.java
index 54671d644..c336969d3 100644
--- a/src/main/java/com/hubspot/jinjava/tree/parse/TokenScanner.java
+++ b/src/main/java/com/hubspot/jinjava/tree/parse/TokenScanner.java
@@ -253,13 +253,18 @@ private Token newToken(int kind) {
       lastStart - lastNewlinePos + 1
     );
 
-    if (t instanceof TagToken) {
-      if (config.isTrimBlocks() && currPost < length && is[currPost] == '\n') {
-        lastNewlinePos = currPost;
-        ++currPost;
-        ++tokenStart;
-      }
+    if (
+      (t instanceof TagToken || t instanceof NoteToken) &&
+      config.isTrimBlocks() &&
+      currPost < length &&
+      is[currPost] == '\n'
+    ) {
+      lastNewlinePos = currPost;
+      ++currPost;
+      ++tokenStart;
+    }
 
+    if (t instanceof TagToken) {
       TagToken tt = (TagToken) t;
       if ("raw".equals(tt.getTagName())) {
         inRaw = 1;
diff --git a/src/test/java/com/hubspot/jinjava/tree/TreeParserTest.java b/src/test/java/com/hubspot/jinjava/tree/TreeParserTest.java
index d1522b23b..c54a33735 100644
--- a/src/test/java/com/hubspot/jinjava/tree/TreeParserTest.java
+++ b/src/test/java/com/hubspot/jinjava/tree/TreeParserTest.java
@@ -147,6 +147,18 @@ public void trimAndLstripBlocks() {
       .isEqualTo("<div>\n" + "        yay\n" + "</div>\n");
   }
 
+  @Test
+  public void trimAndLstripCommentBlocks() {
+    interpreter =
+      new Jinjava(
+        JinjavaConfig.newBuilder().withLstripBlocks(true).withTrimBlocks(true).build()
+      )
+        .newInterpreter();
+
+    assertThat(interpreter.render(parse("parse/tokenizer/whitespace-comment-tags.jinja")))
+      .isEqualTo("<div>\n" + "        yay\n" + "        whoop\n" + "</div>\n");
+  }
+
   @Test
   public void itWarnsAgainstMissingStartTags() {
     String expression = "{% if true %} foo {% endif %} {% endif %}";
diff --git a/src/test/java/com/hubspot/jinjava/tree/parse/TokenWhitespaceTest.java b/src/test/java/com/hubspot/jinjava/tree/parse/TokenWhitespaceTest.java
index ee79b1606..ab7243e74 100644
--- a/src/test/java/com/hubspot/jinjava/tree/parse/TokenWhitespaceTest.java
+++ b/src/test/java/com/hubspot/jinjava/tree/parse/TokenWhitespaceTest.java
@@ -21,6 +21,16 @@ public void trimBlocksTrimsAfterTag() {
     assertThat(tokens.get(2).getImage()).isEqualTo("        yay\n    ");
   }
 
+  @Test
+  public void trimBlocksTrimsAfterCommentTag() {
+    List<Token> tokens = scanTokens(
+      "parse/tokenizer/whitespace-comment-tags.jinja",
+      trimBlocksConfig()
+    );
+    assertThat(tokens.get(2).getImage()).isEqualTo("        yay\n    ");
+    assertThat(tokens.get(4).getImage()).isEqualTo("        whoop\n</div>\n");
+  }
+
   private List<Token> scanTokens(String srcPath, JinjavaConfig config) {
     try {
       return Lists.newArrayList(
diff --git a/src/test/resources/parse/tokenizer/whitespace-comment-tags.jinja b/src/test/resources/parse/tokenizer/whitespace-comment-tags.jinja
new file mode 100644
index 000000000..6cb5a61bb
--- /dev/null
+++ b/src/test/resources/parse/tokenizer/whitespace-comment-tags.jinja
@@ -0,0 +1,8 @@
+<div>
+    {# a comment #}
+        yay
+    {# another comment
+       This time, over multiple lines
+     #}
+        whoop
+</div>