Some space handling improvement

openrewrite · Nov 7, 2023 · 17c3324 · 17c3324
1 parent da67534
commit 17c3324
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 11 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -1976,11 +1976,8 @@ private J visitClass0(KtClass klass, ExecutionContext data) {
         JContainer<TypeTree> implementings = null;
         Markers markers = Markers.EMPTY;
         J.MethodDeclaration primaryConstructor;
-        PsiElement prefixConsumed = findFirstPrefixSpace(findFirstNotSpaceChild(klass));
         Set<PsiElement> prefixConsumedSet = new HashSet<>();
-        if (prefixConsumed != null) {
-            prefixConsumedSet.add(prefixConsumed);
-        }
+        addIfNotNull(prefixConsumedSet, getFirstSpaceChildOrNull(klass));
 
         List<J.Modifier> modifiers = mapModifiers(klass.getModifierList(), leadingAnnotations, lastAnnotations, data);
 
@@ -2420,9 +2417,9 @@ private J visitNamedFunction0(KtNamedFunction function, ExecutionContext data) {
         List<J.Modifier> modifiers = mapModifiers(function.getModifierList(), leadingAnnotations, lastAnnotations, data);
         J.TypeParameters typeParameters = null;
         TypeTree returnTypeExpression = null;
-        PsiElement prefixConsumed = findFirstPrefixSpace(findFirstNotSpaceChild(function));
+
         Set<PsiElement> prefixConsumedSet = new HashSet<>();
-        prefixConsumedSet.add(prefixConsumed);
+        addIfNotNull(prefixConsumedSet, getFirstSpaceChildOrNull(function));
 
         if (function.getTypeParameterList() != null) {
             typeParameters = new J.TypeParameters(
@@ -2614,6 +2611,7 @@ public J visitObjectDeclaration(KtObjectDeclaration declaration, ExecutionContex
         Markers markers = Markers.EMPTY;
         List<J.Modifier> modifiers = new ArrayList<>();
         JContainer<TypeTree> implementings = null;
+        Set<PsiElement> consumedSpaces = new HashSet<>();
 
         List<J.Annotation> leadingAnnotations = new ArrayList<>();
         List<J.Annotation> lastAnnotations = new ArrayList<>();
@@ -2659,6 +2657,8 @@ public J visitObjectDeclaration(KtObjectDeclaration declaration, ExecutionContex
                     .withMarkers(Markers.EMPTY.addIfAbsent(new Implicit(randomId())));
         }
 
+        addIfNotNull(consumedSpaces, getFirstSpaceChildOrNull(declaration));
+
         return new J.ClassDeclaration(
                 randomId(),
                 endFixPrefixAndInfix(declaration),
@@ -2667,7 +2667,7 @@ public J visitObjectDeclaration(KtObjectDeclaration declaration, ExecutionContex
                 modifiers,
                 new J.ClassDeclaration.Kind(
                         randomId(),
-                        prefix(declaration.getObjectKeyword()),
+                        prefix(declaration.getObjectKeyword(), consumedSpaces),
                         Markers.EMPTY,
                         lastAnnotations,
                         J.ClassDeclaration.Kind.Type.Class
@@ -2763,9 +2763,8 @@ public J visitProperty(KtProperty property, ExecutionContext data) {
                 JContainer.build(prefix(property.getTypeParameterList()), mapTypeParameters(property, data), Markers.EMPTY) : null;
         K.TypeConstraints typeConstraints = null;
         boolean isSetterFirst = false;
-        PsiElement prefixConsumed = findFirstPrefixSpace(findFirstNotSpaceChild(property));
         Set<PsiElement> prefixConsumedSet = new HashSet<>();
-        prefixConsumedSet.add(prefixConsumed);
+        addIfNotNull(prefixConsumedSet, getFirstSpaceChildOrNull(property));
 
         modifiers.add(new J.Modifier(
                 Tree.randomId(),
@@ -3432,6 +3431,22 @@ private PsiElement findFirstPrefixSpace(@Nullable PsiElement element) {
         return pre;
     }
 
+    @Nullable
+    private PsiElement getFirstSpaceChildOrNull(@Nullable PsiElement element) {
+        if (element == null) {
+            return null;
+        }
+
+        Iterator<PsiElement> iterator = PsiUtilsKt.getAllChildren(element).iterator();
+        PsiElement first = iterator.next();
+
+        if (first != null) {
+           return isSpace(first.getNode()) ?  first : null;
+        }
+
+        return null;
+    }
+
     @Nullable
     private PsiElement findFirstNonSpacePrevSibling(@Nullable PsiElement element) {
         PsiElement maybeSpace = findFirstPrefixSpace(element);
@@ -3894,6 +3909,12 @@ private PsiElement findLastSpaceChild(@Nullable PsiElement parent) {
         return ret;
     }
 
+    private static <E> void addIfNotNull(Set<E> set, @Nullable E element) {
+        if (element != null) {
+            set.add(element);
+        }
+    }
+
     @Nullable
     private PsiElement findTrailingComma(PsiElement element) {
         PsiElement nextSibling = element.getNextSibling();

diff --git a/src/test/java/org/openrewrite/kotlin/tree/CommentTest.java b/src/test/java/org/openrewrite/kotlin/tree/CommentTest.java
@@ -188,7 +188,7 @@ fun method(){} // END
                   /**
                    * Comment
                    */
-                  companion object AssertValidQueryCompanion {
+                  companion  /*C2*/  object AssertValidQueryCompanion {
                   }
               }
               """

diff --git a/src/test/java/org/openrewrite/kotlin/tree/MethodInvocationTest.java b/src/test/java/org/openrewrite/kotlin/tree/MethodInvocationTest.java
@@ -61,7 +61,7 @@ fun method(): Any = calleeMethod()
                   fun method2(): Any = CALLEE_FIELD
               }
               
-              object Callee {
+              /*42*/ object Callee {
                   const val CALLEE_FIELD = ""
                   fun calleeMethod(): Unit = Unit
               }
-Original file line number
+Diff line change
@@ Expand Up / @@ -188,7 +188,7 @@ fun method(){} // END @@
                       /**
                        * Comment
                        */
-                      companion object AssertValidQueryCompanion {
+                      companion  /*C2*/  object AssertValidQueryCompanion {
                       }
                   }
                   """
@@ Expand Down @@