From 83d4ecc4a6a879ac36351ee433e4c6c24145aa71 Mon Sep 17 00:00:00 2001
From: Vlad Koshkin <vladislav.koshkin@jetbrains.com>
Date: Sun, 6 Jul 2025 15:57:50 +0200
Subject: [PATCH] [kotlin] K2: refine script use scope handling and improve
 unused symbol detection

#KTIJ-34779

GitOrigin-RevId: 43d8e0a9ee09794e12b67fb72038d5f7a5c0e081
---
 .../k2/navigation/KotlinScriptUseScopeEnlarger.kt | 15 ++++++++++++---
 .../idea/highlighting/K2UnusedSymbolUtil.kt       |  6 ++++--
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/plugins/kotlin/base/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/navigation/KotlinScriptUseScopeEnlarger.kt b/plugins/kotlin/base/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/navigation/KotlinScriptUseScopeEnlarger.kt
index 2ed0bf32f212..b746c94d695f 100644
--- a/plugins/kotlin/base/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/navigation/KotlinScriptUseScopeEnlarger.kt
+++ b/plugins/kotlin/base/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/navigation/KotlinScriptUseScopeEnlarger.kt
@@ -1,8 +1,10 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.core.script.k2.navigation
 
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiElement
-import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.DelegatingGlobalSearchScope
 import com.intellij.psi.search.SearchScope
 import com.intellij.psi.search.UseScopeEnlarger
 import org.jetbrains.kotlin.psi.KtFile
@@ -10,6 +12,13 @@ import org.jetbrains.kotlin.psi.KtFile
 class KotlinScriptUseScopeEnlarger : UseScopeEnlarger() {
     override fun getAdditionalUseScope(element: PsiElement): SearchScope? {
         val file = element.containingFile as? KtFile ?: return null
-        return if (file.isScript()) GlobalSearchScope.allScope(element.project) else null
+        return if (file.isScript()) ProjectFileByExtensionSearchScope(element.project, "kts") else null
     }
-}
\ No newline at end of file
+
+    private class ProjectFileByExtensionSearchScope(project: Project, private val extension: String) :
+        DelegatingGlobalSearchScope(projectScope(project)) {
+        override fun contains(file: VirtualFile): Boolean {
+            return super.contains(file) && file.extension == extension
+        }
+    }
+}
diff --git a/plugins/kotlin/highlighting/highlighting-k2/src/org/jetbrains/kotlin/idea/highlighting/K2UnusedSymbolUtil.kt b/plugins/kotlin/highlighting/highlighting-k2/src/org/jetbrains/kotlin/idea/highlighting/K2UnusedSymbolUtil.kt
index 51ea54b0f953..a4ee94d2801f 100644
--- a/plugins/kotlin/highlighting/highlighting-k2/src/org/jetbrains/kotlin/idea/highlighting/K2UnusedSymbolUtil.kt
+++ b/plugins/kotlin/highlighting/highlighting-k2/src/org/jetbrains/kotlin/idea/highlighting/K2UnusedSymbolUtil.kt
@@ -61,7 +61,6 @@ import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.DataClassResolver
-import org.jetbrains.kotlin.scripting.definitions.isScript
 
 object K2UnusedSymbolUtil {
     private val KOTLIN_ADDITIONAL_ANNOTATIONS: List<String> = listOf("kotlin.test.*", "kotlin.js.JsExport")
@@ -70,7 +69,7 @@ object K2UnusedSymbolUtil {
 
     // Simple PSI-based checks
     fun isApplicableByPsi(declaration: KtNamedDeclaration): Boolean {
-        if (declaration.containingFile.isScript()) return false
+        if (declaration.isScriptTopLevelPublicDeclaration()) return false
         // never mark companion object as unused (there are too many reasons it can be needed for)
         if (declaration is KtObjectDeclaration && declaration.isCompanion()) return false
 
@@ -135,6 +134,9 @@ object K2UnusedSymbolUtil {
         return ownerFunction.containingClass()?.isInterface() == true
     }
 
+    private fun KtNamedDeclaration.isScriptTopLevelPublicDeclaration(): Boolean
+        = parent.parent is KtScript && !hasModifier(KtTokens.PRIVATE_KEYWORD)
+
     context(KaSession)
     fun isHiddenFromResolution(declaration: KtNamedDeclaration): Boolean {
         val anno = declaration.findAnnotation(