NI: Optimize some potential hot places

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ResolutionParts.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.types.model.typeConstructor
 import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.types.typeUtil.makeNullable
+import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal object CheckVisibility : ResolutionPart() {
@@ -314,27 +315,28 @@ internal object CollectionTypeVariableUsagesInfo : ResolutionPart() {
         dependentTypeParametersSeen: List<Pair<TypeConstructorMarker, KotlinTypeMarker?>> = listOf()
     ): List<Pair<TypeConstructorMarker, KotlinTypeMarker?>> {
         val context = asConstraintSystemCompleterContext()
-        val dependentTypeParameters = getBuilder().currentStorage().notFixedTypeVariables.mapNotNull { (typeConstructor, constraints) ->
+        val dependentTypeParameters = getBuilder().currentStorage().notFixedTypeVariables.asSequence()
-            val upperBounds = constraints.constraints.filter {
+            .flatMap { (typeConstructor, constraints) ->
-                it.position.from is DeclaredUpperBoundConstraintPosition && it.kind == ConstraintKind.UPPER
+                val upperBounds = constraints.constraints.filter {
-            }
+                    it.position.from is DeclaredUpperBoundConstraintPosition && it.kind == ConstraintKind.UPPER
+                }
 
-            upperBounds.mapNotNull { constraint ->
+                upperBounds.mapNotNull { constraint ->
-                if (constraint.type.typeConstructor(context) != variable) {
+                    if (constraint.type.typeConstructor(context) != variable) {
-                    val suitableUpperBound = upperBounds.find { upperBound ->
+                        val suitableUpperBound = upperBounds.find { upperBound ->
-                        with(context) { upperBound.type.contains { it.typeConstructor() == variable } }
+                            with(context) { upperBound.type.contains { it.typeConstructor() == variable } }
-                    }?.type
+                        }?.type
 
-                    if (suitableUpperBound != null) typeConstructor to suitableUpperBound else null
+                        if (suitableUpperBound != null) typeConstructor to suitableUpperBound else null
-                } else typeConstructor to null
+                    } else typeConstructor to null
-            }
+                }
-        }.flatten().filter { it !in dependentTypeParametersSeen && it.first != variable }
+            }.filter { it !in dependentTypeParametersSeen && it.first != variable }.toList()
 
-        return dependentTypeParameters + dependentTypeParameters.mapNotNull { (typeConstructor, _) ->
+        return dependentTypeParameters + dependentTypeParameters.flatMapTo(SmartList()) { (typeConstructor, _) ->
             if (typeConstructor != variable) {
                 getDependentTypeParameters(typeConstructor, dependentTypeParameters + dependentTypeParametersSeen)
-            } else null
+            } else emptyList()
-        }.flatten()
+        }
     }
 
     private fun NewConstraintSystem.isContainedInInvariantOrContravariantPositionsAmongUpperBound(

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/ConstraintIncorporator.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.resolve.calls.components.CreateFreshVariablesSubstit
 import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.model.*
+import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
@@ -98,7 +99,7 @@ class ConstraintIncorporator(
 
         for (otherTypeVariable in otherInMyConstraint) {
             // to avoid ConcurrentModificationException
-            val otherConstraints = ArrayList(this.getConstraintsForVariable(otherTypeVariable))
+            val otherConstraints = SmartList(this.getConstraintsForVariable(otherTypeVariable))
             for (otherConstraint in otherConstraints) {
                 generateNewConstraint(typeVariable, constraint, otherTypeVariable, otherConstraint)
             }
@@ -230,7 +231,7 @@ class ConstraintIncorporator(
             )
         ) return
 
-        val derivedFrom = (baseConstraint.derivedFrom + otherConstraint.derivedFrom).toMutableSet()
+        val derivedFrom = SmartSet.create(baseConstraint.derivedFrom).also { it.addAll(otherConstraint.derivedFrom) }
         if (otherVariable in derivedFrom) return
 
         derivedFrom.add(otherVariable)
@@ -273,7 +274,7 @@ class ConstraintIncorporator(
     }
 
     fun Context.getNestedTypeVariables(type: KotlinTypeMarker): List<TypeVariableMarker> =
-        getNestedArguments(type).mapNotNull { getTypeVariable(it.getType().typeConstructor()) }
+        getNestedArguments(type).mapNotNullTo(SmartList()) { getTypeVariable(it.getType().typeConstructor()) }
 
 
     private fun KotlinTypeMarker.substitute(c: Context, typeVariable: TypeVariableMarker, value: KotlinTypeMarker): KotlinTypeMarker {
@@ -288,7 +289,7 @@ class ConstraintIncorporator(
 }
 
 private fun TypeSystemInferenceExtensionContext.getNestedArguments(type: KotlinTypeMarker): List<TypeArgumentMarker> {
-    val result = ArrayList<TypeArgumentMarker>()
+    val result = SmartList<TypeArgumentMarker>()
     val stack = ArrayDeque<TypeArgumentMarker>()
 
     when (type) {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/PostponedArgumentInputTypesResolver.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.utils.SmartSet
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class PostponedArgumentInputTypesResolver(
@@ -21,7 +22,7 @@ class PostponedArgumentInputTypesResolver(
 ) {
     interface Context : KotlinConstraintSystemCompleter.Context
 
-    data class ParameterTypesInfo(
+    private class ParameterTypesInfo(
         val parametersFromDeclaration: List<UnwrappedType?>?,
         val parametersFromDeclarationOfRelatedLambdas: Set<List<UnwrappedType?>>?,
         val parametersFromConstraints: Set<List<TypeWithKind>>?,
@@ -48,11 +49,11 @@ class PostponedArgumentInputTypesResolver(
         val dependentVariables =
             variableDependencyProvider.getShallowlyDependentVariables(typeVariableTypeConstructor).orEmpty() + typeVariableTypeConstructor
 
-        return dependentVariables.mapNotNull { type ->
+        return dependentVariables.flatMap { type ->
-            val constraints = notFixedTypeVariables[type]?.constraints ?: return@mapNotNull null
+            val constraints = notFixedTypeVariables[type]?.constraints ?: return@flatMap emptyList()
             val constraintsWithFunctionalType = constraints.filter { (it.type as? KotlinType)?.isBuiltinFunctionalTypeOrSubtype == true }
             constraintsWithFunctionalType.extractFunctionalTypes()
-        }.flatten()
+        }
     }
 
     private fun extractParameterTypesFromDeclaration(atom: ResolutionAtom) =
@@ -81,32 +82,42 @@ class PostponedArgumentInputTypesResolver(
         val parameterTypesFromDeclaration =
             if (argument is LambdaWithTypeVariableAsExpectedTypeAtom) argument.parameterTypesFromDeclaration else null
 
-        val parameterTypesFromConstraints = functionalTypesFromConstraints?.map { typeWithKind ->
+        val parameterTypesFromConstraints = functionalTypesFromConstraints?.mapTo(SmartSet.create()) { typeWithKind ->
             typeWithKind.type.getPureArgumentsForFunctionalTypeOrSubtype().map {
                 // We should use opposite kind as lambda's parameters are contravariant
                 TypeWithKind(it, typeWithKind.direction.opposite())
             }
-        }?.toSet()
+        }
 
         // An extension function flag can only come from a declaration of anonymous function: `select({ this + it }, fun Int.(x: Int) = 10)`
         val (parameterTypesFromDeclarationOfRelatedLambdas, isThereExtensionFunctionAmongRelatedLambdas) =
             getDeclaredParametersFromRelatedLambdas(argument, postponedArguments, variableDependencyProvider)
 
         val annotationsFromConstraints = functionalTypesFromConstraints?.run {
-            Annotations.create(map { it.type.annotations }.flatten())
+            Annotations.create(flatMap { it.type.annotations })
         } ?: Annotations.EMPTY
 
         val annotations = if (isThereExtensionFunctionAmongRelatedLambdas) {
             annotationsFromConstraints.withExtensionFunctionAnnotation(expectedType.builtIns)
         } else annotationsFromConstraints
 
+        var isSuspend = false
+        var isNullable = false
+        if (!functionalTypesFromConstraints.isNullOrEmpty()) {
+            isNullable = true
+            for (funType in functionalTypesFromConstraints) {
+                if (!isSuspend && funType.type.isSuspendFunctionTypeOrSubtype) isSuspend = true
+                if (isNullable && !funType.type.isMarkedNullable) isNullable = false
+                if (isSuspend && !isNullable) break
+            }
+        }
         return ParameterTypesInfo(
             parameterTypesFromDeclaration,
             parameterTypesFromDeclarationOfRelatedLambdas,
             parameterTypesFromConstraints,
             annotations,
-            isSuspend = !functionalTypesFromConstraints.isNullOrEmpty() && functionalTypesFromConstraints.any { it.type.isSuspendFunctionTypeOrSubtype },
+            isSuspend = isSuspend,
-            isNullable = !functionalTypesFromConstraints.isNullOrEmpty() && functionalTypesFromConstraints.all { it.type.isMarkedNullable }
+            isNullable = isNullable
         )
     }
 
@@ -116,24 +127,28 @@ class PostponedArgumentInputTypesResolver(
         dependencyProvider: TypeVariableDependencyInformationProvider
     ): Pair<Set<List<UnwrappedType?>>?, Boolean> {
         val parameterTypesFromDeclarationOfRelatedLambdas = postponedArguments
-            .filterIsInstance<LambdaWithTypeVariableAsExpectedTypeAtom>()
-            .filter { it.parameterTypesFromDeclaration != null && it != argument }
             .mapNotNull { anotherArgument ->
-                val argumentExpectedTypeConstructor = argument.expectedType?.typeConstructor() ?: return@mapNotNull null
+                when {
-                val anotherArgumentExpectedTypeConstructor = anotherArgument.expectedType.typeConstructor()
+                    anotherArgument !is LambdaWithTypeVariableAsExpectedTypeAtom -> null
-                val areTypeVariablesRelated = dependencyProvider.areVariablesDependentShallowly(
+                    anotherArgument.parameterTypesFromDeclaration == null || anotherArgument == argument -> null
-                    argumentExpectedTypeConstructor, anotherArgumentExpectedTypeConstructor
+                    else -> {
-                )
+                        val argumentExpectedTypeConstructor = argument.expectedType?.typeConstructor() ?: return@mapNotNull null
-                val anotherAtom = anotherArgument.atom
+                        val anotherArgumentExpectedTypeConstructor = anotherArgument.expectedType.typeConstructor()
-                val isAnonymousExtensionFunction = anotherAtom is FunctionExpression && anotherAtom.receiverType != null
+                        val areTypeVariablesRelated = dependencyProvider.areVariablesDependentShallowly(
-                val parameterTypesFromDeclarationOfRelatedLambda = anotherArgument.parameterTypesFromDeclaration
+                            argumentExpectedTypeConstructor, anotherArgumentExpectedTypeConstructor
+                        )
+                        val anotherAtom = anotherArgument.atom
+                        val isAnonymousExtensionFunction = anotherAtom is FunctionExpression && anotherAtom.receiverType != null
+                        val parameterTypesFromDeclarationOfRelatedLambda = anotherArgument.parameterTypesFromDeclaration
 
-                if (areTypeVariablesRelated && parameterTypesFromDeclarationOfRelatedLambda != null) {
+                        if (areTypeVariablesRelated && parameterTypesFromDeclarationOfRelatedLambda != null) {
-                    parameterTypesFromDeclarationOfRelatedLambda to isAnonymousExtensionFunction
+                            parameterTypesFromDeclarationOfRelatedLambda to isAnonymousExtensionFunction
-                } else null
+                        } else null
+                    }
+                }
             }
 
-        return parameterTypesFromDeclarationOfRelatedLambdas.run { map { it.first }.toSet() to any { it.second } }
+        return parameterTypesFromDeclarationOfRelatedLambdas.run { mapTo(SmartSet.create()) { it.first } to any { it.second } }
     }
 
     private fun Context.createTypeVariableForReturnType(argument: PostponedAtomWithRevisableExpectedType): NewTypeVariable {
@@ -186,7 +201,8 @@ class PostponedArgumentInputTypesResolver(
         return allGroupedParameterTypes.mapIndexed { index, types ->
             val parameterTypeVariable = createTypeVariableForParameterType(argument, index)
 
-            for (typeWithKind in types.filterNotNull()) {
+            for (typeWithKind in types) {
+                if (typeWithKind == null) continue
                 when (typeWithKind.direction) {
                     ConstraintKind.EQUALITY -> csBuilder.addEqualityConstraint(
                         parameterTypeVariable.defaultType, typeWithKind.type, ArgumentConstraintPosition(atom)
@@ -327,29 +343,23 @@ class PostponedArgumentInputTypesResolver(
         dependencyProvider: TypeVariableDependencyInformationProvider
     ): Boolean {
         // We can collect parameter types from declaration in any mode, they can't change during completion.
-        val postponedArgumentsToCollectTypesFromDeclaredParameters = postponedArguments
+        for (argument in postponedArguments) {
-            .filterIsInstance<LambdaWithTypeVariableAsExpectedTypeAtom>()
+            if (argument !is LambdaWithTypeVariableAsExpectedTypeAtom) continue
-            .filter { it.parameterTypesFromDeclaration == null }
+            if (argument.parameterTypesFromDeclaration != null) continue
-
-        for (argument in postponedArgumentsToCollectTypesFromDeclaredParameters) {
             argument.parameterTypesFromDeclaration = extractParameterTypesFromDeclaration(argument.atom)
         }
 
-        /*
+        return postponedArguments.any { argument ->
-         * We can build new functional expected types in partial mode only for anonymous functions,
+            /*
-         * because more exact type can't appear from constraints in full mode (anonymous functions have fully explicit declaration).
+             * We can build new functional expected types in partial mode only for anonymous functions,
-         * It can be so for lambdas: for instance, an extension function type can appear in full mode (it may not be known in partial mode).
+             * because more exact type can't appear from constraints in full mode (anonymous functions have fully explicit declaration).
-         *
+             * It can be so for lambdas: for instance, an extension function type can appear in full mode (it may not be known in partial mode).
-         * TODO: investigate why we can't do it for anonymous functions in full mode always (see `diagnostics/tests/resolve/resolveWithSpecifiedFunctionLiteralWithId.kt`)
+             *
-         */
+             * TODO: investigate why we can't do it for anonymous functions in full mode always (see `diagnostics/tests/resolve/resolveWithSpecifiedFunctionLiteralWithId.kt`)
-        val postponedArgumentsToCollectParameterTypesAndBuildNewExpectedType =
+             */
-            if (completionMode == KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.PARTIAL) {
+            if (completionMode == KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.PARTIAL && !isAnonymousFunction(argument))
-                postponedArguments.filter(::isAnonymousFunction)
+                return@any false
-            } else {
+            if (argument.revisedExpectedType != null) return@any false
-                postponedArguments
-            }
-
-        return postponedArgumentsToCollectParameterTypesAndBuildNewExpectedType.filter { it.revisedExpectedType == null }.any { argument ->
             val parameterTypesInfo =
                 c.extractParameterTypesInfo(argument, postponedArguments, dependencyProvider) ?: return@any false
             val newExpectedType =
@@ -393,26 +403,28 @@ class PostponedArgumentInputTypesResolver(
                 listOf(typeConstructor) + relatedVariables.filterIsInstance<TypeVariableTypeConstructor>()
             }
             type.arguments.isNotEmpty() -> {
-                type.arguments.map { getAllDeeplyRelatedTypeVariables(it.type, variableDependencyProvider) }.flatten()
+                type.arguments.flatMap { getAllDeeplyRelatedTypeVariables(it.type, variableDependencyProvider) }
             }
-            else -> listOf()
+            else -> emptyList()
         }
     }
 
-    private fun getDeclaredParametersConsideringExtensionFunctionsPresence(parameterTypesInfo: ParameterTypesInfo): List<UnwrappedType?>? {
+    private fun getDeclaredParametersConsideringExtensionFunctionsPresence(parameterTypesInfo: ParameterTypesInfo): List<UnwrappedType?>? =
-        val (parametersFromDeclaration, _, parametersFromConstraints, annotations) = parameterTypesInfo
+        with (parameterTypesInfo) {
 
-        if (parametersFromConstraints.isNullOrEmpty() || parametersFromDeclaration.isNullOrEmpty())
+            if (parametersFromConstraints.isNullOrEmpty() || parametersFromDeclaration.isNullOrEmpty())
-            return parametersFromDeclaration
+                parametersFromDeclaration
+            else {
+                val oneLessParameterInDeclarationThanInConstraints =
+                    parametersFromConstraints.first().size == parametersFromDeclaration.size + 1
 
-        val oneLessParameterInDeclarationThanInConstraints = parametersFromConstraints.first().size == parametersFromDeclaration.size + 1
+                if (oneLessParameterInDeclarationThanInConstraints && annotations.hasExtensionFunctionAnnotation()) {
-
+                    listOf(null) + parametersFromDeclaration
-        return if (oneLessParameterInDeclarationThanInConstraints && annotations.hasExtensionFunctionAnnotation()) {
+                } else {
-            listOf(null) + parametersFromDeclaration
+                    parametersFromDeclaration
-        } else {
+                }
-            parametersFromDeclaration
+            }
         }
-    }
 
     fun fixNextReadyVariableForParameterTypeIfNeeded(
         c: Context,
@@ -443,7 +455,7 @@ class PostponedArgumentInputTypesResolver(
         dependencyProvider: TypeVariableDependencyInformationProvider
     ): Boolean {
         val relatedVariables = type.getPureArgumentsForFunctionalTypeOrSubtype()
-            .map { getAllDeeplyRelatedTypeVariables(it, dependencyProvider) }.flatten()
+            .flatMap { getAllDeeplyRelatedTypeVariables(it, dependencyProvider) }
         val variableForFixation = variableFixationFinder.findFirstVariableForFixation(
             this, relatedVariables, postponedArguments, KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL, topLevelType
         )

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/model/MutableConstraintStorage.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 import org.jetbrains.kotlin.types.model.TypeVariableMarker
 import org.jetbrains.kotlin.types.typeUtil.unCapture
+import org.jetbrains.kotlin.utils.SmartList
 
 
 class MutableVariableWithConstraints private constructor(
@@ -36,12 +37,15 @@ class MutableVariableWithConstraints private constructor(
     // see @OnlyInputTypes annotation
     val projectedInputCallTypes: Collection<UnwrappedType>
         get() = mutableConstraints
-            .filter { it.position.from is OnlyInputTypeConstraintPosition || it.inputTypePositionBeforeIncorporation != null }
+            .mapNotNullTo(SmartList()) {
-            .map { (it.type as KotlinType).unCapture().unwrap() }
+                if (it.position.from is OnlyInputTypeConstraintPosition || it.inputTypePositionBeforeIncorporation != null)
+                    (it.type as KotlinType).unCapture().unwrap()
+                else null
+            }
 
-    private val mutableConstraints = if (constraints == null) ArrayList() else ArrayList(constraints)
+    private val mutableConstraints = if (constraints == null) SmartList() else SmartList(constraints)
 
-    private var simplifiedConstraints: ArrayList<Constraint>? = mutableConstraints
+    private var simplifiedConstraints: SmartList<Constraint>? = mutableConstraints
 
     // return new actual constraint, if this constraint is new
     fun addConstraint(constraint: Constraint): Constraint? {
@@ -110,13 +114,13 @@ class MutableVariableWithConstraints private constructor(
         }
     }
 
-    private fun ArrayList<Constraint>.simplifyConstraints(): ArrayList<Constraint> {
+    private fun SmartList<Constraint>.simplifyConstraints(): SmartList<Constraint> {
         val equalityConstraints =
             filter { it.kind == ConstraintKind.EQUALITY }
                 .groupBy { it.typeHashCode }
         return when {
             equalityConstraints.isEmpty() -> this
-            else -> filterTo(ArrayList()) { isUsefulConstraint(it, equalityConstraints) }
+            else -> filterTo(SmartList()) { isUsefulConstraint(it, equalityConstraints) }
         }
     }
 
@@ -134,10 +138,10 @@ class MutableVariableWithConstraints private constructor(
 internal class MutableConstraintStorage : ConstraintStorage {
     override val allTypeVariables: MutableMap<TypeConstructorMarker, TypeVariableMarker> = LinkedHashMap()
     override val notFixedTypeVariables: MutableMap<TypeConstructorMarker, MutableVariableWithConstraints> = LinkedHashMap()
-    override val initialConstraints: MutableList<InitialConstraint> = ArrayList()
+    override val initialConstraints: MutableList<InitialConstraint> = SmartList()
     override var maxTypeDepthFromInitialConstraints: Int = 1
-    override val errors: MutableList<KotlinCallDiagnostic> = ArrayList()
+    override val errors: MutableList<KotlinCallDiagnostic> = SmartList()
     override val hasContradiction: Boolean get() = errors.any { !it.candidateApplicability.isSuccess }
     override val fixedTypeVariables: MutableMap<TypeConstructorMarker, KotlinTypeMarker> = LinkedHashMap()
-    override val postponedTypeVariables: ArrayList<TypeVariableMarker> = ArrayList()
+    override val postponedTypeVariables: MutableList<TypeVariableMarker> = SmartList()
 }