From 61f28360b22fe4ec221dd8764681c998095c3e23 Mon Sep 17 00:00:00 2001
From: Kun Li <kun@moderne.io>
Date: Tue, 7 Nov 2023 19:55:07 -0800
Subject: [PATCH] more tunning

---
 .../internal/KotlinTreeParserVisitor.java     | 25 ++++++++++++++-----
 .../kotlin/tree/CompilationUnitTest.java      |  6 ++---
 .../kotlin/tree/VariableDeclarationTest.java  |  2 +-
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
index bc09b6d73..23b0df66f 100644
--- a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
+++ b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -2139,7 +2139,7 @@ public J visitClassBody(KtClassBody classBody, ExecutionContext data) {
                 continue;
             }
             Statement statement = convertToStatement(d.accept(this, data));
-            list.add(maybeFollowingSemicolon(statement, d));
+            list.add(maybeSemicolon(statement, d));
         }
 
         return new J.Block(
@@ -3337,7 +3337,7 @@ private J.Block convertToBlock(KtExpression ktExpression, ExecutionContext data)
     }
 
     private <J2 extends J> JRightPadded<J2> maybeSemicolon(J2 j, KtElement element) {
-        PsiElement maybeSemicolon = element.getLastChild();
+        PsiElement maybeSemicolon = findLastNotSpaceChild(element);
         boolean hasSemicolon = maybeSemicolon instanceof LeafPsiElement && ((LeafPsiElement) maybeSemicolon).getElementType() == KtTokens.SEMICOLON;
 
         if (hasSemicolon) {
@@ -3346,7 +3346,7 @@ private <J2 extends J> JRightPadded<J2> maybeSemicolon(J2 j, KtElement element)
 
         maybeSemicolon = findFirstNonSpaceNextSibling(element);
         if (maybeSemicolon instanceof LeafPsiElement && ((LeafPsiElement) maybeSemicolon).getElementType() == KtTokens.SEMICOLON) {
-            return new JRightPadded<>(j, prefix(maybeSemicolon), Markers.EMPTY.add(new Semicolon(randomId())));
+            return new JRightPadded<>(j, deepPrefix(maybeSemicolon), Markers.EMPTY.add(new Semicolon(randomId())));
         }
 
         return padRight(j, Space.EMPTY);
@@ -3502,7 +3502,7 @@ private Space endFix(@Nullable PsiElement element, @Nullable Set<PsiElement> con
             return Space.EMPTY;
         }
 
-        Space end = endFix(element.getLastChild());
+        Space end = endFix(findLastNotSpaceChild(element));
         PsiElement lastSpace = findLastSpaceChild(element);
         if (consumedSpaces != null && consumedSpaces.contains(element)) {
             return end;
@@ -3898,13 +3898,26 @@ private PsiElement findFirstNotSpaceChild(PsiElement parent) {
     }
 
     @Nullable
-    private PsiElement findLastSpaceChild(@Nullable PsiElement parent) {
-        PsiElement ret = null;
+    private PsiElement findLastNotSpaceChild(@Nullable PsiElement parent) {
+        if (parent == null) {
+            return null;
+        }
+
+        for (PsiElement child = parent.getLastChild(); child != null; child = child.getPrevSibling()) {
+            if (!isSpace(child.getNode())) {
+                return child;
+            }
+        }
+        return null;
+    }
 
+    @Nullable
+    private PsiElement findLastSpaceChild(@Nullable PsiElement parent) {
         if (parent == null) {
             return null;
         }
 
+        PsiElement ret = null;
         for (PsiElement child = parent.getLastChild(); child != null; child = child.getPrevSibling()) {
             if (isSpace(child.getNode())) {
                 ret = child;
diff --git a/src/test/java/org/openrewrite/kotlin/tree/CompilationUnitTest.java b/src/test/java/org/openrewrite/kotlin/tree/CompilationUnitTest.java
index 6459761a8..8672405f3 100644
--- a/src/test/java/org/openrewrite/kotlin/tree/CompilationUnitTest.java
+++ b/src/test/java/org/openrewrite/kotlin/tree/CompilationUnitTest.java
@@ -55,12 +55,12 @@ void packageAndComments() {
         rewriteRun(
           kotlin(
             """
-              /* Comment */
+              /* C0 */
               package a
               import java.util.List
               
-              class A
-              // comment
+              class A /*C1*/
+              // C2
               """,
             SourceSpec::noTrim
           )
diff --git a/src/test/java/org/openrewrite/kotlin/tree/VariableDeclarationTest.java b/src/test/java/org/openrewrite/kotlin/tree/VariableDeclarationTest.java
index 37b3dd63d..2b13960f2 100644
--- a/src/test/java/org/openrewrite/kotlin/tree/VariableDeclarationTest.java
+++ b/src/test/java/org/openrewrite/kotlin/tree/VariableDeclarationTest.java
@@ -668,7 +668,7 @@ void trailingSemicolon() {
           kotlin(
             """
               class Test {
-                  var t = 1;
+                  var t = 1 /*C1*/  ; /*C2*/
                   fun method() {}
               }
               """