Create from Usage: Make property lateinit where possible

#KT-17651 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt

@@ -31,6 +31,7 @@ import com.intellij.psi.*
 import com.intellij.psi.codeStyle.CodeStyleManager
 import com.intellij.psi.codeStyle.JavaCodeStyleManager
 import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
 import org.jetbrains.kotlin.cfg.pseudocode.Pseudocode
 import org.jetbrains.kotlin.cfg.pseudocode.getContainingPseudocode
@@ -41,6 +42,7 @@ import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl
 import org.jetbrains.kotlin.descriptors.impl.TypeParameterDescriptorImpl
 import org.jetbrains.kotlin.idea.caches.resolve.analyzeFullyAndGetResult
 import org.jetbrains.kotlin.idea.caches.resolve.getJavaClassDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
 import org.jetbrains.kotlin.idea.core.*
 import org.jetbrains.kotlin.idea.imports.importableFqName
@@ -68,6 +70,7 @@ import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier
 import org.jetbrains.kotlin.resolve.scopes.utils.memberScopeAsImportingScope
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeProjectionImpl
+import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.typeUtil.isAnyOrNullableAny
@@ -608,6 +611,16 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
             return typeRefsToShorten
         }
 
+        private fun postprocessDeclaration(declaration: KtNamedDeclaration) {
+            if (callableInfo is PropertyInfo && callableInfo.isLateinitPreferred) {
+                if (declaration.containingClassOrObject == null) return
+                val propertyDescriptor = declaration.resolveToDescriptor() as? PropertyDescriptor ?: return
+                val returnType = propertyDescriptor.returnType ?: return
+                if (TypeUtils.isNullableType(returnType) || KotlinBuiltIns.isPrimitiveType(returnType)) return
+                declaration.addModifier(KtTokens.LATEINIT_KEYWORD)
+            }
+        }
+
         private fun setupDeclarationBody(func: KtDeclarationWithBody) {
             val oldBody = func.bodyExpression ?: return
             val templateKind = when (func) {
@@ -910,6 +923,8 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                                                                                     false) ?: return
 
                         runWriteAction {
+                            postprocessDeclaration(newDeclaration)
+
                             // file templates
                             if (newDeclaration is KtNamedFunction || newDeclaration is KtSecondaryConstructor) {
                                 setupDeclarationBody(newDeclaration as KtFunction)

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableInfo.kt

@@ -220,7 +220,8 @@ class PropertyInfo(name: String,
                    val writable: Boolean,
                    possibleContainers: List<KtElement> = Collections.emptyList(),
                    typeParameterInfos: List<TypeInfo> = Collections.emptyList(),
-                   isAbstract: Boolean = false
+                   isAbstract: Boolean = false,
+                   val isLateinitPreferred: Boolean = false
 ) : CallableInfo(name, receiverTypeInfo, returnTypeInfo, possibleContainers, typeParameterInfos, isAbstract) {
     override val kind: CallableKind get() = CallableKind.PROPERTY
     override val parameterInfos: List<ParameterInfo> get() = Collections.emptyList()
@@ -232,6 +233,7 @@ class PropertyInfo(name: String,
             writable,
             possibleContainers,
             typeParameterInfos,
-            isAbstract
+            isAbstract,
+            isLateinitPreferred
     )
 }

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createCallable/CreateCallableFromCallActionFactory.kt

@@ -20,6 +20,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiClass
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
@@ -187,11 +188,14 @@ sealed class CreateCallableFromCallActionFactory<E : KtExpression>(
         ): CallableInfo? {
             val fullCallExpr = expression.getQualifiedExpressionForSelectorOrThis()
             val varExpected = fullCallExpr.getAssignmentByLHS() != null
-            val returnType = TypeInfo(
-                    fullCallExpr.getExpressionForTypeGuess(),
-                    if (varExpected) Variance.INVARIANT else Variance.OUT_VARIANCE
-            )
-            return PropertyInfo(name, receiverType, returnType, varExpected, possibleContainers)
+            val expressionForTypeGuess = fullCallExpr.getExpressionForTypeGuess()
+            val returnTypes = expressionForTypeGuess.guessTypes(analysisResult.bindingContext, analysisResult.moduleDescriptor)
+            val returnTypeInfo = TypeInfo(expressionForTypeGuess, if (varExpected) Variance.INVARIANT else Variance.OUT_VARIANCE)
+            val canBeLateinit =
+                    varExpected
+                    && returnTypes.any { !it.isMarkedNullable && !KotlinBuiltIns.isPrimitiveType(it) }
+                    && fullCallExpr.parents.firstOrNull { it is KtDeclarationWithBody || it is KtClassInitializer } is KtDeclarationWithBody
+            return PropertyInfo(name, receiverType, returnTypeInfo, varExpected, possibleContainers, isLateinitPreferred = canBeLateinit)
         }
 
         object Default : Property() {

idea/testData/quickfix/createFromUsage/createVariable/property/varOnUserType.kt.after

@@ -1,8 +1,7 @@
 // "Create member property 'A.foo'" "true"
-// ERROR: Property must be initialized or be abstract
 
 class A<T>(val n: T) {
-    var foo: String
+    lateinit var foo: String
 }
 
 fun test() {