New J2K: add support of lambda expressions in nullabilityAnalysis

nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/BoundTypeStorage.kt

@@ -16,8 +16,11 @@ import org.jetbrains.kotlin.psi.psiUtil.getNextSiblingIgnoringWhitespace
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCalleeExpressionIfAny
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.isNullable
+import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class BoundTypeStorage(private val analysisAnalysisContext: AnalysisContext, private val printConstraints: Boolean) {
@@ -34,11 +37,12 @@ internal class BoundTypeStorage(private val analysisAnalysisContext: AnalysisCon
                         expression.selectorExpression?.toBoundType(
                             boundTypeFor(expression.receiverExpression)
                         )
+                    is KtLabeledExpression -> expression.baseExpression?.let { boundTypeFor(it) }
+                    is KtLambdaExpression -> lambdaBoundType(expression)
                     else -> expression.getQualifiedExpressionForSelector()?.let { boundTypeFor(it) }
                 } ?: expression.toBoundType(null)
                 ?: LiteralBoundType(expression.isNullable())
 
-
             if (printConstraints) {
                 if (expression.getNextSiblingIgnoringWhitespace() !is PsiComment) {
                     val comment = with(printer) {
@@ -47,9 +51,36 @@ internal class BoundTypeStorage(private val analysisAnalysisContext: AnalysisCon
                     expression.parent.addAfter(comment, expression)
                 }
             }
-            boundType
+            boundType.withForcedNullability(expression.getForcedNullability())
         }
 
+    private fun lambdaBoundType(lambda: KtLambdaExpression): BoundType? {
+        val builtIns = lambda.getType(lambda.analyze())?.builtIns ?: return null
+
+        val descriptor = builtIns.getFunction(lambda.valueParameters.size)
+        val parameterBoundTypes = lambda.valueParameters.map { parameter ->
+            parameter.typeReference?.typeElement?.let { typeElement ->
+                analysisAnalysisContext.typeElementToTypeVariable[typeElement]
+            }?.let {
+                BoundTypeTypeParameter(TypeVariableBoundType(it), Variance.IN_VARIANCE)
+            } ?: return null
+        }
+        val returnBoundType =
+            BoundTypeTypeParameter(
+                TypeVariableBoundType(
+                    analysisAnalysisContext.declarationToTypeVariable[lambda.functionLiteral] ?: return null
+                ),
+                Variance.OUT_VARIANCE
+            )
+        return GenericBoundType(
+            DescriptorClassReference(descriptor),
+            parameterBoundTypes + returnBoundType,
+            forcedNullabilityTo = null,
+            isNull = false
+        )
+    }
+
+
     fun boundTypeForType(
         type: KotlinType,
         contextBoundType: BoundType?,
@@ -103,12 +134,13 @@ internal class BoundTypeStorage(private val analysisAnalysisContext: AnalysisCon
             ?.let { TypeVariableBoundType(it) }
     }
 
-    private fun KtExpression.toBoundType(contextBoundType: BoundType?): BoundType? {
+    private fun KtExpression.toBoundType(contextBoundType: BoundType?): BoundType? =
-        toBoundTypeAsTypeVariable()?.also { return it }
+        run {
-        toBoundTypeAsCallExpression(contextBoundType)?.also { return it }
+            toBoundTypeAsTypeVariable()?.also { return@run it }
-        toBoundTypeAsCastExpression()?.also { return it }
+            toBoundTypeAsCallExpression(contextBoundType)?.also { return@run it }
-        return null
+            toBoundTypeAsCastExpression()?.also { return@run it }
-    }
+            return@run null
+        }?.withForcedNullability(getForcedNullability())
 
     private fun KotlinType.toBoundType(
         contextBoundType: BoundType?,

nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/constraints.kt

@@ -47,6 +47,11 @@ internal data class FunctionTarget(
     override val typeVariable: TypeVariable
 ) : DeclarationTypeVariableOwner
 
+internal data class LambdaTarget(
+    override val target: KtFunctionLiteral,
+    override val typeVariable: TypeVariable
+) : DeclarationTypeVariableOwner
+
 internal data class ParameterTarget(
     override val target: KtParameter,
     override val typeVariable: TypeVariable
@@ -112,7 +117,7 @@ internal class LiteralBoundType(val isNull: Boolean) : BoundType {
 
 
 internal class TypeVariable(
-    val typeElement: KtTypeElement,
+    val typeElement: KtTypeElement?,
     val classReference: ClassReference,
     val typeParameters: List<TypeVariableTypeParameter>,
     var nullability: Nullability

nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt

@@ -9,12 +9,16 @@ import com.intellij.openapi.editor.RangeMarker
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.nj2k.NewJ2kConverterContext
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.nj2k.JKElementInfoLabel
+import org.jetbrains.kotlin.nj2k.NewJ2kConverterContext
 import org.jetbrains.kotlin.nj2k.asLabel
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class ContextCreator(
@@ -25,11 +29,17 @@ internal class ContextCreator(
     private fun KtCallableDeclaration.typeElement(): KtTypeElement? =
         typeReference?.typeElement
 
-    private fun KtElement.asTypeVariableOwner(): TypeVariableOwner? =
+    private fun KtElement.asTypeVariableOwner(): TypeVariableOwner? {
-        when (this) {
+        return when (this) {
             is KtParameter -> typeElement()?.asTypeVariable()?.let { ParameterTarget(this, it) }
             is KtProperty -> typeElement()?.asTypeVariable()?.let { PropertyTarget(this, it) }
             is KtNamedFunction -> typeElement()?.asTypeVariable()?.let { FunctionTarget(this, it) }
+            is KtLambdaExpression -> {
+                val context = analyze()
+                val returnType = getType(context)?.arguments?.lastOrNull()?.type ?: return null
+                val typeVariable = returnType.asTypeVariable() ?: return null
+                LambdaTarget(functionLiteral, typeVariable)
+            }
 
             is KtBinaryExpressionWithTypeRHS -> right?.typeElement
                 ?.takeIf { KtPsiUtil.isUnsafeCast(this) }
@@ -45,6 +55,7 @@ internal class ContextCreator(
                 else null
             else -> null
         }
+    }
 
     fun createContext(analysisScope: AnalysisScope): AnalysisContext {
         val typeVariableOwners = analysisScope.flatMap {
@@ -56,7 +67,8 @@ internal class ContextCreator(
         val typeElementsToTypeVariables =
             typeVariableOwners.flatMap {
                 it.innerTypeVariables() + it.allTypeVariables
-            }.associateBy { it.typeElement }
+            }.filter { it.typeElement != null }
+                .associateBy { it.typeElement!! }
 
         val declarationToTypeVariable = typeVariableOwners.asSequence()
             .mapNotNull { owner ->
@@ -82,6 +94,22 @@ internal class ContextCreator(
             }
         return TypeVariable(this, classReference, typeParameters, nullability)
     }
+
+    private fun KotlinType.asTypeVariable(): TypeVariable? {
+        val classReference = constructor
+            .declarationDescriptor
+            ?.safeAs<ClassDescriptor>()
+            ?.let { DescriptorClassReference(it) }
+            ?: return null
+        val typeParameters =
+            arguments.zip(constructor.parameters) { argument, parameter ->
+                TypeVariableTypeParameterWithTypeParameter(
+                    argument.type.asTypeVariable() ?: return null,
+                    parameter.variance
+                )
+            }
+        return TypeVariable(null, classReference, typeParameters, Nullability.UNKNOWN)
+    }
 }
 
 

nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/nullabilityAnalysis.kt

@@ -22,7 +22,7 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 
 internal fun TypeVariable.changeNullability(toNullable: Boolean) {
-    typeElement.changeNullability(toNullable)
+    typeElement?.changeNullability(toNullable)
 }
 
 internal fun KtTypeElement.changeNullability(toNullable: Boolean) {
@@ -40,6 +40,7 @@ internal fun AnalysisContext.fixTypeVariablesNullability() {
     if (typeElementToTypeVariable.isEmpty()) return
 
     val deepComparator = Comparator<TypeVariable> { o1, o2 ->
+        if (o1.typeElement == null || o2.typeElement == null) return@Comparator -1
         if (o1.typeElement.isAncestor(o2.typeElement)) 1 else -1
     }
     for (typeVariableOwner in typeVariableOwners) {

nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/printing.kt

@@ -70,9 +70,13 @@ internal class Printer(private val analysisContext: AnalysisContext) {
 
 
 private class Namer(analysisContext: AnalysisContext) {
-    val names = analysisContext.typeElementToTypeVariable.values.mapIndexed { index, typeVariable ->
+    val names = run {
-        typeVariable to "T$index"
+        val typeVariables = (analysisContext.typeElementToTypeVariable.values +
-    }.toMap()
+                analysisContext.declarationToTypeVariable.values)
+        typeVariables.mapIndexed { index, typeVariable ->
+            typeVariable to "T$index"
+        }.toMap()
+    }
 
     fun name(typeVariable: TypeVariable): String =
         names.getValue(typeVariable)

nj2k/testData/nullabilityAnalysis/functions.kt

@@ -0,0 +1,19 @@
+fun notNullParameters(f: Function2<Int, Int, String>) {}
+fun nullableParameter(f: Function2<Int, Int, String>) {}
+fun nullableReturnType(f: Function2<Int, Int, String>) {}
+
+
+fun test() {
+    notNullParameters({ i: Int, j: Int ->
+                          if (i < 10 && j > 0) "" else ""
+                      })
+
+    nullableParameter({ i: Int, j: Int ->
+                          if (i == null) "" else ""
+                      })
+
+    nullableReturnType({ i: Int, j: Int ->
+                           if (i < 10) return@nullableReturnType null
+                           return@nullableReturnType "nya"
+                       })
+}

nj2k/testData/nullabilityAnalysis/functions.kt.after

@@ -0,0 +1,19 @@
+fun notNullParameters(f: Function2<Int, Int, String>) {}
+fun nullableParameter(f: Function2<Int?, Int, String>) {}
+fun nullableReturnType(f: Function2<Int, Int, String?>) {}
+
+
+fun test() {
+    notNullParameters({ i: Int, j: Int ->
+                          if (i < 10 && j > 0) "" else ""
+                      })
+
+    nullableParameter({ i: Int?, j: Int ->
+                          if (i == null) "" else ""
+                      })
+
+    nullableReturnType({ i: Int, j: Int ->
+                           if (i < 10) return@nullableReturnType null
+                           return@nullableReturnType "nya"
+                       })
+}

nj2k/tests/org/jetbrains/kotlin/nj2k/NullabilityAnalysisTestGenerated.java

@@ -44,6 +44,11 @@ public class NullabilityAnalysisTestGenerated extends AbstractNullabilityAnalysi
         runTest("nj2k/testData/nullabilityAnalysis/functionTypeParameterNullability.kt");
     }
 
+    @TestMetadata("functions.kt")
+    public void testFunctions() throws Exception {
+        runTest("nj2k/testData/nullabilityAnalysis/functions.kt");
+    }
+
     @TestMetadata("loops.kt")
     public void testLoops() throws Exception {
         runTest("nj2k/testData/nullabilityAnalysis/loops.kt");