Fix copyright updater in scripts (KT-30133)

#KT-30133 Fixed

idea/src/org/jetbrains/kotlin/idea/copyright/UpdateKotlinCopyright.java

@@ -20,13 +20,17 @@ import com.intellij.openapi.module.Module;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiComment;
-import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.SyntaxTraverser;
+import com.intellij.util.containers.ContainerUtil;
+import com.intellij.util.containers.TreeTraversal;
 import com.maddyhome.idea.copyright.CopyrightProfile;
 import com.maddyhome.idea.copyright.psi.UpdatePsiFileCopyright;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.psi.KtDeclaration;
+
+import java.util.List;
 
 public class UpdateKotlinCopyright extends UpdatePsiFileCopyright {
     UpdateKotlinCopyright(Project project, Module module, VirtualFile root, CopyrightProfile copyrightProfile) {
@@ -35,39 +39,22 @@ public class UpdateKotlinCopyright extends UpdatePsiFileCopyright {
 
     @Override
     protected void scanFile() {
-        PsiRange commentSearchRange = getCommentSearchRange(getFile());
-        PsiElement first = commentSearchRange.first;
-        PsiElement last = commentSearchRange.last;
-
-        if (first != null) {
-            checkComments(first, last, true);
-        }
+        List<PsiComment> comments = getExistentComments(getFile());
+        checkComments(ContainerUtil.getLastItem(comments), true, comments);
     }
 
-    public static class PsiRange {
-        public final @Nullable PsiElement first;
-        public final @Nullable PsiElement last;
-
-        public PsiRange(@Nullable PsiElement first, @Nullable PsiElement second) {
-            this.first = first;
-            this.last = second;
-        }
-    }
-
-    public static @NotNull PsiRange getCommentSearchRange(@NotNull PsiFile psiFile) {
-        PsiElement first = psiFile.getFirstChild();
-        PsiElement last = first;
-        PsiElement next = first;
-        while (next != null) {
-            if (next instanceof PsiComment || next instanceof PsiWhiteSpace || next.getText().isEmpty()) {
-                next = next.getNextSibling();
-            }
-            else {
-                break;
-            }
-            last = next;
-        }
-
-        return new PsiRange(first, last);
+    @NotNull
+    public static List<PsiComment> getExistentComments(@NotNull PsiFile psiFile) {
+        return SyntaxTraverser.psiTraverser(psiFile)
+                .withTraversal(TreeTraversal.LEAVES_DFS)
+                .traverse()
+                .takeWhile(
+                        element ->
+                                (element instanceof PsiComment && !(element.getParent() instanceof KtDeclaration)) ||
+                                element instanceof PsiWhiteSpace ||
+                                element.getText().isEmpty()
+                )
+                .filter(PsiComment.class)
+                .toList();
     }
 }

idea/testData/copyright/Script.kts

@@ -0,0 +1,5 @@
+/* PRESENT */
+
+foo()
+
+// COMMENTS: 1

idea/tests/org/jetbrains/kotlin/copyright/AbstractUpdateKotlinCopyrightTest.kt

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.copyright
 
-import com.intellij.psi.PsiComment
-import com.intellij.psi.PsiElement
 import junit.framework.AssertionFailedError
 import org.jetbrains.kotlin.idea.copyright.UpdateKotlinCopyright
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
@@ -19,60 +17,31 @@ abstract class AbstractUpdateKotlinCopyrightTest : KotlinLightCodeInsightFixture
     fun doTest(path: String) {
         myFixture.configureByFile(path)
 
-        val text = myFixture.file.text.trim()
-        val expectedNumberOfComments = InTextDirectivesUtils.getPrefixedInt(text, "// COMMENTS: ") ?: run {
-            if (!text.isEmpty()) {
+        val fileText = myFixture.file.text.trim()
+        val expectedNumberOfComments = InTextDirectivesUtils.getPrefixedInt(fileText, "// COMMENTS: ") ?: run {
+            if (fileText.isNotEmpty()) {
                 throw AssertionFailedError("Every test should assert number of comments with `COMMENTS` directive")
             } else {
                 0
             }
         }
 
-        var commentsNumber = 0
-        val comments = collectComments(UpdateKotlinCopyright.getCommentSearchRange(myFixture.file))
+        val comments = UpdateKotlinCopyright.getExistentComments(myFixture.file)
         for (comment in comments) {
-            when (comment.text) {
-                "/* PRESENT */" -> {
-                    ++commentsNumber
-                }
+            val commentText = comment.text
+            when (commentText) {
+                "/* PRESENT */" -> {}
                 "/* ABSENT */" -> {
                     throw AssertionFailedError("Unexpected comment found")
                 }
                 else -> {
-                    throw AssertionFailedError("A comment with bad directive found: `$text`")
+                    throw AssertionFailedError("A comment with bad directive found: `$commentText`")
                 }
             }
         }
 
-        Assert.assertEquals("Wrong number of comments found", expectedNumberOfComments, commentsNumber)
+        Assert.assertEquals("Wrong number of comments found", expectedNumberOfComments, comments.size)
     }
 
     override fun getTestDataPath() = File(PluginTestCaseBase.getTestDataPathBase(), "/copyright").path + File.separator
-
-    companion object {
-        private fun collectComments(range: UpdateKotlinCopyright.PsiRange): List<PsiComment> {
-            val comments = mutableListOf<PsiComment>()
-            val first = range.first
-            val last = range.last
-            if (first == null) return listOf()
-            collectComments(first, last, comments)
-            return comments
-        }
-
-        // Copied from UpdatePsiFileCopyright.collectComments()
-        private fun collectComments(first: PsiElement, last: PsiElement?, comments: MutableList<PsiComment>) {
-            if (first === last && first is PsiComment) {
-                comments.add(first)
-                return
-            }
-            var elem: PsiElement? = first
-            while (elem !== last && elem != null) {
-                if (elem is PsiComment) {
-                    comments.add(elem)
-                }
-
-                elem = elem.nextSibling
-            }
-        }
-    }
 }

idea/tests/org/jetbrains/kotlin/copyright/UpdateKotlinCopyrightTestGenerated.java

@@ -44,6 +44,11 @@ public class UpdateKotlinCopyrightTestGenerated extends AbstractUpdateKotlinCopy
         runTest("idea/testData/copyright/NoPackage.kt");
     }
 
+    @TestMetadata("Script.kts")
+    public void testScript() throws Exception {
+        runTest("idea/testData/copyright/Script.kts");
+    }
+
     @TestMetadata("Simple.kt")
     public void testSimple() throws Exception {
         runTest("idea/testData/copyright/Simple.kt");