New J2K: calculate for loop variable nullability

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

@@ -78,6 +78,29 @@ data class OneTimeProcessingGroup(val processings: List<Processing>) : Processin
     constructor(vararg processings: Processing) : this(processings.toList())
 }
 
+private abstract class CheckableProcessing<E : PsiElement>(private val classTag: KClass<E>) : NewJ2kPostProcessing {
+    protected open fun check(element: E): Boolean =
+        check(element, null)
+
+    protected open fun check(element: E, settings: ConverterSettings?): Boolean =
+        check(element)
+
+    protected abstract fun action(element: E)
+
+    override fun createAction(element: PsiElement, diagnostics: Diagnostics, settings: ConverterSettings?): (() -> Unit)? {
+        if (!element::class.isSubclassOf(classTag)) return null
+        @Suppress("UNCHECKED_CAST")
+        if (!check(element as E, settings)) return null
+        return {
+            if (element.isValid && check(element, settings)) {
+                action(element)
+            }
+        }
+    }
+
+    override val writeActionNeeded: Boolean = true
+}
+
 object NewJ2KPostProcessingRegistrar {
 
     private fun Processing.processings(): Sequence<NewJ2kPostProcessing> =
@@ -308,7 +331,8 @@ object NewJ2KPostProcessingRegistrar {
                 if (context.diagnostics.forElement(element).any { it.factory == Errors.UNNECESSARY_NOT_NULL_ASSERTION }) {
                     exclExclExpr.replace(baseExpression)
                 }
-            }
+            },
+            RemoveForExpressionLoopParameterTypeProcessing()
         )
     )
 
@@ -790,6 +814,17 @@ object NewJ2KPostProcessingRegistrar {
         override val writeActionNeeded: Boolean = true
     }
 
+    private class RemoveForExpressionLoopParameterTypeProcessing : CheckableProcessing<KtForExpression>(KtForExpression::class) {
+        override fun check(element: KtForExpression, settings: ConverterSettings?): Boolean =
+            element.loopParameter?.typeReference?.typeElement != null
+                    && settings?.specifyLocalVariableTypeByDefault != true
+
+        override fun action(element: KtForExpression) {
+            element.loopParameter?.typeReference = null
+        }
+
+    }
+
     private class RemoveRedundantExpressionQualifierProcessing : NewJ2kPostProcessing {
         private fun check(qualifiedExpression: KtQualifiedExpression): Boolean {
             val qualifier = (qualifiedExpression.receiverExpression as? KtNameReferenceExpression)

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

@@ -19,8 +19,11 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.asAssignment
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getTargetFunction
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.resolve.calls.components.isVararg
 import org.jetbrains.kotlin.resolve.descriptorUtil.firstOverridden
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
+import org.jetbrains.kotlin.resolve.scopes.getDescriptorsFiltered
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class ConstraintsCollector(
@@ -97,10 +100,6 @@ internal class ConstraintsCollector(
                 expression.condition?.addEqualsNullabilityConstraint(Nullability.NOT_NULL, ConstraintCameFrom.INITIALIZER)
             }
 
-            expression is KtForExpression -> {
-                expression.loopRange?.addEqualsNullabilityConstraint(Nullability.NOT_NULL, ConstraintCameFrom.INITIALIZER)
-            }
-
             expression is KtCallExpression -> {
                 collectConstraintsForCallExpression(
                     expression,
@@ -123,6 +122,40 @@ internal class ConstraintsCollector(
             expression is KtNamedFunction -> {
                 collectSuperDeclarationConstraintsForFunction(expression)
             }
+
+            expression is KtForExpression -> {
+                expression.loopRange?.addEqualsNullabilityConstraint(Nullability.NOT_NULL, ConstraintCameFrom.INITIALIZER)
+
+                val loopParameterTypeVariable =
+                    expression.loopParameter?.typeReference?.typeElement?.let { typeElement ->
+                        analysisContext.typeElementToTypeVariable[typeElement]
+                    }
+                if (loopParameterTypeVariable != null) {
+                    val loopRangeBoundType = boundTypeStorage.boundTypeFor(expression.loopRange!!)
+                    val loopRangeType = expression.loopRange?.getType(expression.analyze()) ?: return
+                    val loopRangeItemType = loopRangeType
+                        .constructor
+                        .declarationDescriptor
+                        ?.safeAs<ClassDescriptor>()
+                        ?.getMemberScope(loopRangeType.arguments)
+                        ?.getDescriptorsFiltered(DescriptorKindFilter.FUNCTIONS) {
+                            it.asString() == "iterator"
+                        }?.filterIsInstance<FunctionDescriptor>()
+                        ?.firstOrNull { it.valueParameters.isEmpty() }
+                        ?.original
+                        ?.returnType
+                        ?: return
+                    val boundType = boundTypeStorage.boundTypeForType(
+                        loopRangeItemType,
+                        loopRangeBoundType,
+                        emptyMap()
+                    ) ?: return
+                    loopParameterTypeVariable.addEqualsNullabilityConstraint(
+                        boundType.typeParameters.firstOrNull()?.boundType ?: return,
+                        ConstraintCameFrom.ASSIGNMENT_TARGET
+                    )
+                }
+            }
         }
         Unit
     }

nj2k/src/org/jetbrains/kotlin/nj2k/ConversionsRunner.kt

@@ -17,9 +17,7 @@
 package org.jetbrains.kotlin.nj2k
 
 import org.jetbrains.kotlin.nj2k.conversions.*
-import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.JKTreeRoot
-import org.jetbrains.kotlin.nj2k.tree.impl.JKTreeRootImpl
 
 object ConversionsRunner {
 
@@ -68,7 +66,6 @@ object ConversionsRunner {
             +LiteralConversion()
             +ForConversion(context)
             +LabeledStatementConversion()
-            +ForInConversion(context)
             +TypeParametersNullabilityConversion()
             +ArrayOperationsConversion(context)
             +BuiltinMembersConversion(context)

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/ForInConversion.kt

@@ -1,29 +0,0 @@
-/*
- * 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.nj2k.conversions
-
-import com.intellij.psi.*
-import org.jetbrains.kotlin.nj2k.*
-import org.jetbrains.kotlin.nj2k.tree.*
-import org.jetbrains.kotlin.nj2k.tree.impl.*
-import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.name.ClassId
-
-
-class ForInConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
-    override fun applyToElement(element: JKTreeElement): JKTreeElement {
-        if (element !is JKForInStatement) return recurse(element)
-
-        val parameterDeclaration = element.declaration as? JKVariable
-        if (parameterDeclaration != null
-            && !context.converter.settings.specifyLocalVariableTypeByDefault
-        ) {
-            parameterDeclaration.type = JKTypeElementImpl(JKNoTypeImpl)
-        }
-
-        return recurse(element)
-    }
-}

nj2k/testData/nullabilityAnalysis/loops.kt

@@ -0,0 +1,16 @@
+fun bar(map: HashMap<String, Int>, list1: List<Int>, list2: List<String>) {
+    for (entry: MutableMap.MutableEntry<String, Int> in map.entries) {
+        val value: Int = entry.value
+        if (entry.key == null) {
+            println(value + 1)
+        }
+    }
+
+    for (i: Int in list1) {
+        i + 1
+    }
+
+    for (i: String in list2) {
+        i == null
+    }
+}

nj2k/testData/nullabilityAnalysis/loops.kt.after

@@ -0,0 +1,16 @@
+fun bar(map: HashMap<String?, Int>, list1: List<Int>, list2: List<String?>) {
+    for (entry: MutableMap.MutableEntry<String?, Int> in map.entries) {
+        val value: Int = entry.value
+        if (entry.key == null) {
+            println(value + 1)
+        }
+    }
+
+    for (i: Int in list1) {
+        i + 1
+    }
+
+    for (i: String? in list2) {
+        i == null
+    }
+}

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

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