FIR IDE: register quickfix to add lateinit modifier

idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/MainKtQuickFixRegistrar.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixRegistrar
 import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixesList
 import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixesListBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
+import org.jetbrains.kotlin.idea.quickfix.fixes.AddLateInitFactory
 import org.jetbrains.kotlin.idea.quickfix.fixes.ChangeTypeQuickFix
 import org.jetbrains.kotlin.idea.quickfix.fixes.InitializePropertyQuickFixFactory
 import org.jetbrains.kotlin.idea.quickfix.fixes.ReplaceCallFixFactories
@@ -63,6 +64,7 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
             AddModifierFix.addAbstractModifier,
         )
         registerApplicator(InitializePropertyQuickFixFactory.initializePropertyFactory)
+        registerApplicator(AddLateInitFactory.addLateInitFactory)
     }
 
     private val overrides = KtQuickFixesListBuilder.registerPsiQuickFix {

idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddLateInitFactory.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.idea.quickfix.fixes
+
+import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactory
+import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
+import org.jetbrains.kotlin.idea.frontend.api.types.isPrimitive
+import org.jetbrains.kotlin.idea.quickfix.AddModifierFix
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.KtProperty
+
+object AddLateInitFactory {
+
+    @OptIn(ExperimentalStdlibApi::class)
+    val addLateInitFactory = diagnosticFixFactory<KtFirDiagnostic.MustBeInitializedOrBeAbstract> { diagnostic ->
+        val property: KtProperty = diagnostic.psi
+        if (!property.isVar) return@diagnosticFixFactory emptyList()
+
+        val type = property.getReturnKtType()
+
+        if (type.isPrimitive || type.canBeNull) return@diagnosticFixFactory emptyList()
+
+        listOf(AddModifierFix(property, KtTokens.LATEINIT_KEYWORD))
+    }
+}

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtTypeProvider.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.frontend.api.components
 
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtNamedClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 
@@ -34,6 +33,8 @@ interface KtTypeProviderMixIn : KtAnalysisSessionMixIn {
     fun KtNamedClassOrObjectSymbol.buildSelfClassType(): KtType =
         analysisSession.typeProvider.buildSelfClassType(this)
 
+    val KtType.canBeNull: Boolean
+
 }
 
 @Suppress("PropertyName")

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/types/ktTypeUtils.kt

@@ -27,6 +27,12 @@ fun KtType.isClassTypeWithClassId(classId: ClassId): Boolean {
     return this.classId == classId
 }
 
+val KtType.isPrimitive: Boolean
+    get() {
+        if (this !is KtClassType) return false
+        return this.classId in DefaultTypeClassIds.PRIMITIVES
+    }
+
 private object DefaultTypeClassIds {
     val UNIT = ClassId.topLevel(StandardNames.FqNames.unit.toSafe())
     val INT = ClassId.topLevel(StandardNames.FqNames._int.toSafe())
@@ -38,6 +44,7 @@ private object DefaultTypeClassIds {
     val CHAR = ClassId.topLevel(StandardNames.FqNames._char.toSafe())
     val BOOLEAN = ClassId.topLevel(StandardNames.FqNames._boolean.toSafe())
     val STRING = ClassId.topLevel(StandardNames.FqNames.string.toSafe())
+    val PRIMITIVES = setOf(INT, LONG, SHORT, BYTE, FLOAT, DOUBLE, CHAR, BOOLEAN)
 }
 
 val KtType.defaultInitializer: String?

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSession.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.types.canBeNull
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getFirFile
@@ -20,9 +21,11 @@ import org.jetbrains.kotlin.idea.frontend.api.tokens.ReadActionConfinementValidi
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.components.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirSymbolProvider
+import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.recordCompletionContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.threadLocal
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -59,6 +62,8 @@ private constructor(
 
     override val expressionInfoProviderImpl = KtFirExpressionInfoProvider(this, token)
 
+    override val KtType.canBeNull: Boolean get() = (this as KtFirType).coneType.canBeNull
+
     override val typeProviderImpl = KtFirTypeProvider(this, token)
 
     override val subtypingComponentImpl = KtFirSubtypingComponent(this, token)

idea/testData/quickfix/modifiers/addLateinit.kt

@@ -3,4 +3,3 @@
 class A {
     private var a: String<caret>
 }
-/* IGNORE_FIR */

idea/testData/quickfix/modifiers/addLateinit.kt.after

@@ -3,4 +3,3 @@
 class A {
     private lateinit var a: String<caret>
 }
-/* IGNORE_FIR */