Support overriding module/library/sdk for particular light PsiClass

To support light R classes

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt

@@ -138,6 +138,10 @@ private fun <T> PsiElement.collectInfos(c: ModuleInfoCollector<T>): T {
         return this.processLightElement(c)
     }
 
+    collectModuleInfoByUserData(this).firstOrNull()?.let {
+        return c.onResult(it)
+    }
+
     val containingFile =
         containingFile ?: return c.onFailure("Analyzing element of type ${this::class.java} with no containing file\nText:\n$text")
 
@@ -205,6 +209,8 @@ private inline fun <T> collectInfosByVirtualFile(
     treatAsLibrarySource: Boolean,
     onOccurrence: (IdeaModuleInfo?) -> T
 ): T {
+    collectModuleInfoByUserData(project, virtualFile).map(onOccurrence)
+
     val moduleRelatedModuleInfo = getModuleRelatedModuleInfo(project, virtualFile)
     if (moduleRelatedModuleInfo != null) {
         onOccurrence(moduleRelatedModuleInfo)

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/userDataModuleInfoUtil.kt

@@ -0,0 +1,17 @@
+package org.jetbrains.kotlin.idea.caches.project
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
+
+
+// NO-OP implementation, full implementation only for AS3.3, AS3.4
+
+fun collectModuleInfoByUserData(
+    project: Project,
+    virtualFile: VirtualFile
+): List<IdeaModuleInfo> = emptyList()
+
+fun collectModuleInfoByUserData(
+    psiElement: PsiElement
+): List<IdeaModuleInfo> = emptyList()

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/userDataModuleInfoUtil.kt.as33

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.caches.project
+
+import com.intellij.openapi.module.Module
+import com.intellij.openapi.module.ModuleUtilCore
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.LIBRARY_KEY
+import org.jetbrains.kotlin.idea.MODULE_ROOT_TYPE_KEY
+import org.jetbrains.kotlin.idea.SDK_KEY
+import org.jetbrains.kotlin.idea.caches.project.UserDataModuleContainer.ForPsiElement
+import org.jetbrains.kotlin.idea.caches.project.UserDataModuleContainer.ForVirtualFile
+import org.jetbrains.kotlin.idea.core.getSourceType
+import org.jetbrains.kotlin.utils.addIfNotNull
+
+// This file declares non-exported API for overriding module info with user-data
+
+
+private sealed class UserDataModuleContainer {
+    abstract fun <T> getUserData(key: Key<T>): T?
+    abstract fun getModule(): Module?
+
+    data class ForVirtualFile(val virtualFile: VirtualFile, val project: Project) : UserDataModuleContainer() {
+        override fun <T> getUserData(key: Key<T>): T? = virtualFile.getUserData(key)
+        override fun getModule(): Module? = ModuleUtilCore.findModuleForFile(virtualFile, project)
+    }
+
+    data class ForPsiElement(val psiElement: PsiElement) : UserDataModuleContainer() {
+        override fun <T> getUserData(key: Key<T>): T? {
+            return psiElement.getUserData(key)
+                ?: psiElement.containingFile?.getUserData(key)
+                ?: psiElement.containingFile?.originalFile?.virtualFile?.getUserData(key)
+        }
+
+        override fun getModule(): Module? = ModuleUtilCore.findModuleForPsiElement(psiElement)
+    }
+}
+
+
+private fun collectModuleInfoByUserData(
+    project: Project,
+    container: UserDataModuleContainer
+): List<IdeaModuleInfo> {
+    fun forModule(): ModuleSourceInfo? {
+        val rootType = container.getUserData(MODULE_ROOT_TYPE_KEY) ?: return null
+        val module = container.getModule() ?: return null
+
+        return when (rootType.getSourceType()) {
+            null -> null
+            SourceType.PRODUCTION -> module.productionSourceInfo()
+            SourceType.TEST -> module.testSourceInfo()
+        }
+    }
+
+    val result = mutableListOf<IdeaModuleInfo>()
+    result.addIfNotNull(forModule())
+
+    val library = container.getUserData(LIBRARY_KEY)
+    if (library != null) {
+        result.addAll(createLibraryInfo(project, library))
+    }
+
+    val sdk = container.getUserData(SDK_KEY)
+    if (sdk != null) {
+        result.add(SdkInfo(project, sdk))
+    }
+
+    return result
+}
+
+
+fun collectModuleInfoByUserData(
+    project: Project,
+    virtualFile: VirtualFile
+): List<IdeaModuleInfo> = collectModuleInfoByUserData(project, ForVirtualFile(virtualFile, project))
+
+fun collectModuleInfoByUserData(
+    psiElement: PsiElement
+): List<IdeaModuleInfo> = collectModuleInfoByUserData(psiElement.project, ForPsiElement(psiElement))

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/userDataModuleInfoUtil.kt.as34

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.caches.project
+
+import com.intellij.openapi.module.Module
+import com.intellij.openapi.module.ModuleUtilCore
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.LIBRARY_KEY
+import org.jetbrains.kotlin.idea.MODULE_ROOT_TYPE_KEY
+import org.jetbrains.kotlin.idea.SDK_KEY
+import org.jetbrains.kotlin.idea.caches.project.UserDataModuleContainer.ForPsiElement
+import org.jetbrains.kotlin.idea.caches.project.UserDataModuleContainer.ForVirtualFile
+import org.jetbrains.kotlin.idea.core.getSourceType
+import org.jetbrains.kotlin.utils.addIfNotNull
+
+// This file declares non-exported API for overriding module info with user-data
+
+
+private sealed class UserDataModuleContainer {
+    abstract fun <T> getUserData(key: Key<T>): T?
+    abstract fun getModule(): Module?
+
+    data class ForVirtualFile(val virtualFile: VirtualFile, val project: Project) : UserDataModuleContainer() {
+        override fun <T> getUserData(key: Key<T>): T? = virtualFile.getUserData(key)
+        override fun getModule(): Module? = ModuleUtilCore.findModuleForFile(virtualFile, project)
+    }
+
+    data class ForPsiElement(val psiElement: PsiElement) : UserDataModuleContainer() {
+        override fun <T> getUserData(key: Key<T>): T? {
+            return psiElement.getUserData(key)
+                ?: psiElement.containingFile?.getUserData(key)
+                ?: psiElement.containingFile?.originalFile?.virtualFile?.getUserData(key)
+        }
+
+        override fun getModule(): Module? = ModuleUtilCore.findModuleForPsiElement(psiElement)
+    }
+}
+
+
+private fun collectModuleInfoByUserData(
+    project: Project,
+    container: UserDataModuleContainer
+): List<IdeaModuleInfo> {
+    fun forModule(): ModuleSourceInfo? {
+        val rootType = container.getUserData(MODULE_ROOT_TYPE_KEY) ?: return null
+        val module = container.getModule() ?: return null
+
+        return when (rootType.getSourceType()) {
+            null -> null
+            SourceType.PRODUCTION -> module.productionSourceInfo()
+            SourceType.TEST -> module.testSourceInfo()
+        }
+    }
+
+    val result = mutableListOf<IdeaModuleInfo>()
+    result.addIfNotNull(forModule())
+
+    val library = container.getUserData(LIBRARY_KEY)
+    if (library != null) {
+        result.addAll(createLibraryInfo(project, library))
+    }
+
+    val sdk = container.getUserData(SDK_KEY)
+    if (sdk != null) {
+        result.add(SdkInfo(project, sdk))
+    }
+
+    return result
+}
+
+
+fun collectModuleInfoByUserData(
+    project: Project,
+    virtualFile: VirtualFile
+): List<IdeaModuleInfo> = collectModuleInfoByUserData(project, ForVirtualFile(virtualFile, project))
+
+fun collectModuleInfoByUserData(
+    psiElement: PsiElement
+): List<IdeaModuleInfo> = collectModuleInfoByUserData(psiElement.project, ForPsiElement(psiElement))

idea/idea-analysis/src/org/jetbrains/kotlin/idea/userDataModuleInfo.kt.as33

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea
+
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.libraries.Library
+import com.intellij.openapi.util.Key
+import org.jetbrains.jps.model.module.JpsModuleSourceRootType
+
+// WARNING, this API is used by AS3.3+
+
+
+const val LIBRARY_KEY_NAME = "Kt_Library"
+const val SDK_KEY_NAME = "Kt_Sdk"
+const val MODULE_ROOT_TYPE_KEY_NAME = "Kt_SourceRootType"
+
+@JvmField
+val MODULE_ROOT_TYPE_KEY = getOrCreateKey<JpsModuleSourceRootType<*>>(MODULE_ROOT_TYPE_KEY_NAME)
+
+@JvmField
+val SDK_KEY = getOrCreateKey<Sdk>(SDK_KEY_NAME)
+
+@JvmField
+val LIBRARY_KEY = getOrCreateKey<Library>(LIBRARY_KEY_NAME)
+
+
+inline fun <reified T> getOrCreateKey(name: String): Key<T> {
+    @Suppress("DEPRECATION", "UNCHECKED_CAST")
+    val existingKey = Key.findKeyByName(name) as Key<T>?
+    return existingKey ?: Key.create<T>(name)
+}

idea/idea-analysis/src/org/jetbrains/kotlin/idea/userDataModuleInfo.kt.as34

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea
+
+import com.intellij.openapi.projectRoots.Sdk
+import com.intellij.openapi.roots.libraries.Library
+import com.intellij.openapi.util.Key
+import org.jetbrains.jps.model.module.JpsModuleSourceRootType
+
+// WARNING, this API is used by AS3.3+
+
+
+const val LIBRARY_KEY_NAME = "Kt_Library"
+const val SDK_KEY_NAME = "Kt_Sdk"
+const val MODULE_ROOT_TYPE_KEY_NAME = "Kt_SourceRootType"
+
+@JvmField
+val MODULE_ROOT_TYPE_KEY = getOrCreateKey<JpsModuleSourceRootType<*>>(MODULE_ROOT_TYPE_KEY_NAME)
+
+@JvmField
+val SDK_KEY = getOrCreateKey<Sdk>(SDK_KEY_NAME)
+
+@JvmField
+val LIBRARY_KEY = getOrCreateKey<Library>(LIBRARY_KEY_NAME)
+
+
+inline fun <reified T> getOrCreateKey(name: String): Key<T> {
+    @Suppress("DEPRECATION", "UNCHECKED_CAST")
+    val existingKey = Key.findKeyByName(name) as Key<T>?
+    return existingKey ?: Key.create<T>(name)
+}

idea/idea-core/src/org/jetbrains/kotlin/idea/core/fileIndexUtils.kt

@@ -22,6 +22,19 @@ private val testRootTypes: Set<JpsModuleSourceRootType<*>> = setOf(
     KotlinResourceRootType.TestResource
 )
 
+private val sourceRootTypes = setOf<JpsModuleSourceRootType<*>>(
+    JavaSourceRootType.SOURCE,
+    JavaResourceRootType.RESOURCE,
+    KotlinSourceRootType.Source,
+    KotlinResourceRootType.Resource
+)
+
+fun JpsModuleSourceRootType<*>.getSourceType(): SourceType? = when(this) {
+    in sourceRootTypes -> SourceType.PRODUCTION
+    in testRootTypes -> SourceType.TEST
+    else -> null
+}
+
 fun FileIndex.isInTestSourceContentKotlinAware(fileOrDir: VirtualFile) = isUnderSourceRootOfType(fileOrDir, testRootTypes)
 
 fun FileIndex.getSourceType(fileOrDir: VirtualFile): SourceType? = when {