[kotlin] migrate KotlinSuggestedRefactoringAvailability from IdeaModuleInfo to KaModule

KTIJ-31422

GitOrigin-RevId: 96a05811ff024c4a51cd8f68036c137fce34c50b

plugins/kotlin/base/project-structure/src/org/jetbrains/kotlin/idea/base/projectStructure/IDEProjectStructureProvider.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.roots.libraries.Library as OpenapiLibrary
 import com.intellij.openapi.projectRoots.Sdk as OpenapiSdk
 import com.intellij.platform.workspace.jps.entities.LibraryId
 import com.intellij.platform.workspace.jps.entities.ModuleId
+import com.intellij.psi.PsiFile
 import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinProjectStructureProvider
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinProjectStructureProviderBase
@@ -38,6 +39,10 @@ abstract class IDEProjectStructureProvider : KotlinProjectStructureProviderBase(
     abstract fun getOpenapiSdk(module: KaLibraryModule): OpenapiSdk?
 
     abstract fun getContainingKaModules(virtualFile: VirtualFile): List<KaModule>
+
+    abstract fun getForcedKaModule(file: PsiFile): KaModule?
+
+    abstract fun setForcedKaModule(file: PsiFile, kaModule: KaModule?)
 }
 
 

plugins/kotlin/base/project-structure/src/org/jetbrains/kotlin/idea/base/projectStructure/ProjectStructureProviderIdeImpl.kt

@@ -278,6 +278,21 @@ internal class ProjectStructureProviderIdeImpl(private val project: Project) : I
             .mapTo(mutableListOf()) { getKtModuleByModuleInfo(it) }
     }
 
+    override fun getForcedKaModule(file: PsiFile): KaModule? {
+        return file.forcedModuleInfo?.let { getKtModuleByModuleInfo(it) }
+    }
+
+    override fun setForcedKaModule(file: PsiFile, kaModule: KaModule?) {
+        when (kaModule) {
+          null -> {
+              file.forcedModuleInfo = null
+          }
+          is KtModuleByModuleInfoBase -> {
+              file.forcedModuleInfo = kaModule.moduleInfo
+          }
+        }
+    }
+
     companion object {
         // TODO maybe introduce some cache?
         fun getKtModuleByModuleInfo(moduleInfo: ModuleInfo): KaModule {

plugins/kotlin/base/project-structure/src/org/jetbrains/kotlin/idea/base/projectStructure/api.kt

@@ -4,17 +4,19 @@ package org.jetbrains.kotlin.idea.base.projectStructure
 
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.projectRoots.Sdk as OpenapiSdk
-import com.intellij.openapi.roots.libraries.Library as OpenapiLibrary
 import com.intellij.openapi.roots.libraries.Library
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.platform.workspace.jps.entities.LibraryId
 import com.intellij.platform.workspace.jps.entities.ModuleId
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaLibraryModule
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModule
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaModuleProvider
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaSourceModule
+import com.intellij.openapi.projectRoots.Sdk as OpenapiSdk
+import com.intellij.openapi.roots.libraries.Library as OpenapiLibrary
 
 /**
  * Represents kind of [KaSourceModule]
@@ -206,3 +208,18 @@ inline fun <reified M : KaModule> PsiElement.getKaModuleOfType(project: Project,
 */
 fun VirtualFile.getContainingKaModules(project: Project): List<KaModule> =
 project.ideProjectStructureProvider.getContainingKaModules(this)
+
+
+/**
+ * [forcedKaModule] provides a [KaModule] instance for a dummy file. It must not be changed after the first assignment because
+ * [IDEProjectStructureProvider] might cache the module info.
+ */
+var PsiFile.forcedKaModule: KaModule?
+    @ApiStatus.Internal
+    get() {
+        return project.ideProjectStructureProvider.getForcedKaModule(this)
+    }
+    @ApiStatus.Internal
+    set(value) {
+        project.ideProjectStructureProvider.setForcedKaModule(this, value)
+    }

plugins/kotlin/refactorings/kotlin.refactorings.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/suggested/KotlinSuggestedRefactoringAvailability.kt

@@ -21,9 +21,8 @@ import org.jetbrains.kotlin.analysis.api.symbols.KaFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.receiverType
 import org.jetbrains.kotlin.analysis.api.types.KaErrorType
 import org.jetbrains.kotlin.analysis.api.types.KaType
-import org.jetbrains.kotlin.idea.base.projectStructure.forcedModuleInfo
-import org.jetbrains.kotlin.idea.base.projectStructure.moduleInfo
-import org.jetbrains.kotlin.idea.base.util.K1ModeProjectStructureApi
+import org.jetbrains.kotlin.idea.base.projectStructure.forcedKaModule
+import org.jetbrains.kotlin.idea.base.projectStructure.getKaModule
 import org.jetbrains.kotlin.idea.refactoring.isInterfaceClass
 import org.jetbrains.kotlin.idea.refactoring.suggested.KotlinSignatureAdditionalData
 import org.jetbrains.kotlin.idea.refactoring.suggested.KotlinSuggestedRefactoringSupportBase
@@ -92,8 +91,7 @@ class KotlinSuggestedRefactoringAvailability(refactoringSupport: SuggestedRefact
     override fun refineSignaturesWithResolve(state: SuggestedRefactoringState): SuggestedRefactoringState {
         val newDeclaration = state.declaration as? KtCallableDeclaration ?: return state
         val oldDeclaration = state.restoredDeclarationCopy() as? KtCallableDeclaration ?: return state
-        @OptIn(K1ModeProjectStructureApi::class)
-        oldDeclaration.containingKtFile.forcedModuleInfo = newDeclaration.moduleInfo
+        oldDeclaration.containingKtFile.forcedKaModule = newDeclaration.getKaModule(newDeclaration.project, useSiteModule = null)
 
         val descriptorWithOldSignature = allowAnalysisOnEdt { analyzeCopy(oldDeclaration, KaDanglingFileResolutionMode.PREFER_SELF) { signatureTypes(oldDeclaration) } } ?: return state
         val descriptorWithNewSignature = allowAnalysisOnEdt { analyze(newDeclaration) { signatureTypes(newDeclaration) }  } ?: return state