Export filename extension from KotlinScriptDefinition, register...

extensions after loading definitions

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -20,10 +20,7 @@ import com.intellij.codeInsight.ContainerProvider
 import com.intellij.codeInsight.ExternalAnnotationsManager
 import com.intellij.codeInsight.InferredAnnotationsManager
 import com.intellij.codeInsight.runner.JavaMainMethodProvider
-import com.intellij.core.CoreApplicationEnvironment
-import com.intellij.core.CoreJavaFileManager
-import com.intellij.core.JavaCoreApplicationEnvironment
-import com.intellij.core.JavaCoreProjectEnvironment
+import com.intellij.core.*
 import com.intellij.ide.highlighter.JavaFileType
 import com.intellij.ide.plugins.PluginManagerCore
 import com.intellij.lang.MetaLanguage
@@ -36,6 +33,7 @@ import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.extensions.ExtensionsArea
 import com.intellij.openapi.fileTypes.FileTypeExtensionPoint
+import com.intellij.openapi.fileTypes.FileTypeRegistry
 import com.intellij.openapi.fileTypes.PlainTextFileType
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
@@ -220,9 +218,6 @@ class KotlinCoreEnvironment private constructor(
             extension.updateConfiguration(configuration)
         }
 
-        sourceFiles += createKtFiles(project)
-        sourceFiles.sortBy { it.virtualFile.path }
-
         // If not disabled explicitly, we should always support at least the standard script definition
         if (!configuration.getBoolean(JVMConfigurationKeys.DISABLE_STANDARD_SCRIPT_DEFINITION) &&
             StandardScriptDefinition !in configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS)
@@ -235,6 +230,20 @@ class KotlinCoreEnvironment private constructor(
             scriptDefinitionProvider.setScriptDefinitionsSources(configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES))
             scriptDefinitionProvider.setScriptDefinitions(configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS))
 
+            // Register new file extensions
+            val fileTypeRegistry = FileTypeRegistry.getInstance() as CoreFileTypeRegistry
+
+            scriptDefinitionProvider.getKnownFilenameExtensions().filter {
+                fileTypeRegistry.getFileTypeByExtension(it) != KotlinFileType.INSTANCE
+            }.forEach {
+                fileTypeRegistry.registerFileType(KotlinFileType.INSTANCE, it)
+            }
+        }
+
+        sourceFiles += createKtFiles(project)
+        sourceFiles.sortBy { it.virtualFile.path }
+
+        if (scriptDefinitionProvider != null) {
             ScriptDependenciesProvider.getInstance(project).let { importsProvider ->
                 configuration.addJvmClasspathRoots(
                     sourceFiles.mapNotNull(importsProvider::getScriptDependencies)

compiler/psi/src/org/jetbrains/kotlin/script/KotlinScriptDefinition.kt

@@ -45,6 +45,9 @@ open class KotlinScriptDefinition(open val template: KClass<out Any>) : UserData
     open fun getScriptName(script: KtScript): Name =
         NameUtils.getScriptNameForFile(script.containingKtFile.name)
 
+    open val fileExtension: String
+        get() = "kts"
+
     open val dependencyResolver: DependenciesResolver get() = DependenciesResolver.NoDependencies
 
     open val acceptedAnnotations: List<KClass<out Annotation>> get() = emptyList()

compiler/psi/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionProvider.kt

@@ -34,6 +34,8 @@ interface ScriptDefinitionProvider {
     fun isScript(fileName: String): Boolean
     fun getDefaultScriptDefinition(): KotlinScriptDefinition
 
+    fun getKnownFilenameExtensions(): Sequence<String>
+
     companion object {
         fun getInstance(project: Project): ScriptDefinitionProvider =
             ServiceManager.getService(project, ScriptDefinitionProvider::class.java)
@@ -94,6 +96,10 @@ abstract class LazyScriptDefinitionProvider : ScriptDefinitionProvider {
 
     override fun isScript(fileName: String) = findScriptDefinition(fileName) != null
 
+    override fun getKnownFilenameExtensions(): Sequence<String> = lock.read {
+        cachedDefinitions?.map { it.fileExtension } ?: emptySequence()
+    }
+
     companion object {
         // TODO: find a common place for storing kotlin-related extensions and reuse values from it everywhere
         protected val nonScriptFilenameSuffixes = arrayOf(".${KotlinFileType.EXTENSION}", ".${JavaFileType.DEFAULT_EXTENSION}")

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt

@@ -41,7 +41,6 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.script.*
 import org.jetbrains.kotlin.script.util.scriptCompilationClasspathFromContextOrStlib
 import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPI
-import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPIBase
 import org.jetbrains.kotlin.utils.PathUtil
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.flattenTo
@@ -136,18 +135,17 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
         assert(lock.isWriteLocked) { "updateDefinitions should only be called under the write lock" }
         definitions = definitionsByContributor.values.flattenTo(mutableListOf()).asSequence()
 
-        // Register new file extensions
         val fileTypeManager = FileTypeManager.getInstance()
-        val extensions = definitions?.mapNotNull { definition ->
-            (definition as? KotlinScriptDefinitionAdapterFromNewAPIBase)
-                ?.scriptFileExtensionWithDot?.removePrefix(".")
-                ?.takeIf { fileTypeManager.getFileTypeByExtension(it) != KotlinFileType.INSTANCE }
-        }?.toList()
 
-        if (extensions?.isNotEmpty() == true) {
+        val newExtensions = getKnownFilenameExtensions().filter {
+            fileTypeManager.getFileTypeByExtension(it) != KotlinFileType.INSTANCE
+        }.toList()
+
+        if (newExtensions.any()) {
+            // Register new file extensions
             ApplicationManager.getApplication().invokeLater {
                 runWriteAction {
-                    extensions.forEach {
+                    newExtensions.forEach {
                         fileTypeManager.associateExtension(KotlinFileType.INSTANCE, it)
                     }
                 }

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/KotlinScriptDefinitionAdapterFromNewAPI.kt

@@ -29,8 +29,6 @@ abstract class KotlinScriptDefinitionAdapterFromNewAPIBase : KotlinScriptDefinit
 
     protected abstract val hostConfiguration: ScriptingHostConfiguration
 
-    abstract val scriptFileExtensionWithDot: String
-
     open val baseClass: KClass<*> by lazy(LazyThreadSafetyMode.PUBLICATION) {
         getScriptingClass(scriptCompilationConfiguration.getOrError(ScriptCompilationConfiguration.baseClass))
     }
@@ -43,11 +41,11 @@ abstract class KotlinScriptDefinitionAdapterFromNewAPIBase : KotlinScriptDefinit
     override val fileType: LanguageFileType = KotlinFileType.INSTANCE
 
     override fun isScript(fileName: String): Boolean =
-        fileName.endsWith(scriptFileExtensionWithDot)
+        fileName.endsWith(".$fileExtension")
 
     override fun getScriptName(script: KtScript): Name {
         val fileBasedName = NameUtils.getScriptNameForFile(script.containingKtFile.name)
-        return Name.identifier(fileBasedName.identifier.removeSuffix(scriptFileExtensionWithDot))
+        return Name.identifier(fileBasedName.identifier.removeSuffix(".$fileExtension"))
     }
 
     override val annotationsForSamWithReceivers: List<String>
@@ -105,6 +103,6 @@ class KotlinScriptDefinitionAdapterFromNewAPI(
 
     override val name: String get() = scriptCompilationConfiguration[ScriptCompilationConfiguration.displayName] ?: super.name
 
-    override val scriptFileExtensionWithDot =
-        "." + (scriptCompilationConfiguration[ScriptCompilationConfiguration.fileExtension] ?: "kts")
+    override val fileExtension: String
+        get() = scriptCompilationConfiguration[ScriptCompilationConfiguration.fileExtension] ?: super.fileExtension
 }

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/LazyScriptDefinitionFromDiscoveredClass.kt

@@ -64,14 +64,12 @@ class LazyScriptDefinitionFromDiscoveredClass internal constructor(
         }
     }
 
-    override val scriptFileExtensionWithDot: String by lazy(LazyThreadSafetyMode.PUBLICATION) {
-        val extFromAnn = annotationsFromAsm.find { it.name == KotlinScript::class.simpleName }?.args
+    override val fileExtension: String by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        annotationsFromAsm.find { it.name == KotlinScript::class.simpleName }?.args
             ?.find { it.name == "extension" }?.value
-        val ext = extFromAnn
             ?: scriptCompilationConfiguration.let {
-                it[ScriptCompilationConfiguration.fileExtension] ?: "kts"
+                it[ScriptCompilationConfiguration.fileExtension] ?: super.fileExtension
             }
-        ".$ext"
     }
 
     override val name: String by lazy(LazyThreadSafetyMode.PUBLICATION) {