Enable new type inference for UAST tests as for IDE tests

compiler/tests-common/tests/org/jetbrains/kotlin/checkers/CompilerTestLanguageVersionSettings.kt

@@ -28,13 +28,13 @@ data class CompilerTestLanguageVersionSettings(
         private val initialLanguageFeatures: Map<LanguageFeature, LanguageFeature.State>,
         override val apiVersion: ApiVersion,
         override val languageVersion: LanguageVersion,
-        private val analysisFlags: Map<AnalysisFlag<*>, Any?> = emptyMap()
+        val analysisFlags: Map<AnalysisFlag<*>, Any?> = emptyMap()
 ) : LanguageVersionSettings {
-    private val languageFeatures = specificFeaturesForTests() + initialLanguageFeatures
-    private val delegate = LanguageVersionSettingsImpl(languageVersion, apiVersion, emptyMap(), languageFeatures)
+    val extraLanguageFeatures = specificFeaturesForTests() + initialLanguageFeatures
+    private val delegate = LanguageVersionSettingsImpl(languageVersion, apiVersion, emptyMap(), extraLanguageFeatures)
 
     override fun getFeatureSupport(feature: LanguageFeature): LanguageFeature.State =
-            languageFeatures[feature] ?: delegate.getFeatureSupport(feature)
+            extraLanguageFeatures[feature] ?: delegate.getFeatureSupport(feature)
 
     override fun isPreRelease(): Boolean = KotlinCompilerVersion.isPreRelease()
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/NewInferenceForIDEAnalysisComponent.kt

@@ -10,7 +10,7 @@ import com.intellij.openapi.project.Project
 
 object NewInferenceForIDEAnalysisComponent {
     private const val inferenceOption = "kotlin.use.new.inference.for.ide.analysis"
-    private const val defaultState = true
+    const val defaultState = true
 
     @JvmStatic
     fun setEnabled(project: Project, state: Boolean) {

plugins/uast-kotlin/tests/AbstractKotlinUastTest.kt

@@ -8,6 +8,7 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.util.io.URLUtil
+import org.jetbrains.kotlin.checkers.CompilerTestLanguageVersionSettings
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoot
 import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
@@ -16,8 +17,8 @@ import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.compiler.NoScopeRecordCliBindingTrace
 import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
-import org.jetbrains.kotlin.config.CommonConfigurationKeys
-import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.*
+import org.jetbrains.kotlin.idea.project.NewInferenceForIDEAnalysisComponent
 import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.script.loadScriptingPlugin
@@ -51,6 +52,8 @@ abstract class AbstractKotlinUastTest : AbstractUastTest() {
 
         initializeKotlinEnvironment()
 
+        enableNewTypeInferenceIfNeeded()
+
         val trace = NoScopeRecordCliBindingTrace()
 
         val environment = kotlinCoreEnvironment!!
@@ -66,6 +69,28 @@ abstract class AbstractKotlinUastTest : AbstractUastTest() {
         return vfs.findFileByPath(testFile.canonicalPath)!!
     }
 
+    private fun enableNewTypeInferenceIfNeeded() {
+        if (!NewInferenceForIDEAnalysisComponent.defaultState) return
+
+        val currentLanguageVersionSettings = compilerConfiguration.languageVersionSettings
+        if (currentLanguageVersionSettings.supportsFeature(LanguageFeature.NewInference)) return
+
+        val extraLanguageFeatures = mutableMapOf<LanguageFeature, LanguageFeature.State>()
+        val extraAnalysisFlags = mutableMapOf<AnalysisFlag<*>, Any?>()
+
+        if (currentLanguageVersionSettings is CompilerTestLanguageVersionSettings) {
+            extraLanguageFeatures += currentLanguageVersionSettings.extraLanguageFeatures
+            extraAnalysisFlags += currentLanguageVersionSettings.analysisFlags
+        }
+
+        compilerConfiguration.languageVersionSettings = CompilerTestLanguageVersionSettings(
+            extraLanguageFeatures + (LanguageFeature.NewInference to LanguageFeature.State.ENABLED),
+            currentLanguageVersionSettings.apiVersion,
+            currentLanguageVersionSettings.languageVersion,
+            extraAnalysisFlags
+        )
+    }
+
     private fun initializeKotlinEnvironment() {
         val area = Extensions.getRootArea()
         area.getExtensionPoint(UastLanguagePlugin.extensionPointName)