From c84ff1d8b1cb57d0173138bdc23dbcdc86e0871a Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Thu, 21 Mar 2019 16:29:08 +0700
Subject: [PATCH] Inspection "MoveVariableDeclarationIntoWhen" should inline
 variable declaration

---
 ...veVariableDeclarationIntoWhenInspection.kt | 44 +++++++++++++++----
 .../withComment.kt.after                      |  2 +-
 .../withNewLine.kt.after                      |  2 +-
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt b/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt
index 401c09d82dd..b90df0a41f4 100644
--- a/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/inspections/MoveVariableDeclarationIntoWhenInspection.kt
@@ -23,33 +23,61 @@ class MoveVariableDeclarationIntoWhenInspection : AbstractKotlinInspection(), Cl
         whenExpressionVisitor(fun(expression: KtWhenExpression) {
             val subjectExpression = expression.subjectExpression ?: return
             val property = expression.findDeclarationNear() ?: return
-            if (!property.isUsedOnlyIn(expression)) return
+            val action = property.action(expression)
+            if (action == Action.NOTHING) return
+
             val identifier = property.nameIdentifier ?: return
             holder.registerProblem(
                 property,
                 TextRange.from(identifier.startOffsetInParent, identifier.textLength),
-                "Variable declaration could be moved inside `when`",
-                MoveVariableDeclarationIntoWhenFix(subjectExpression.createSmartPointer())
+                action.description, action.createFix(subjectExpression.createSmartPointer())
             )
         })
 }
 
-private fun KtProperty.isUsedOnlyIn(element: KtElement): Boolean = countUsages() == countUsages(element)
+private enum class Action {
+    NOTHING,
+    MOVE,
+    INLINE;
+
+    val description: String
+        get() = when (this) {
+            MOVE -> "Variable declaration could be moved into `when`"
+            INLINE -> "Variable declaration could be inlined"
+            NOTHING -> "Nothing to do"
+        }
+
+    fun createFix(subjectExpressionPointer: SmartPsiElementPointer<KtExpression>): VariableDeclarationIntoWhenFix = when (this) {
+        MOVE -> VariableDeclarationIntoWhenFix("Move variable declaration into `when`", subjectExpressionPointer) { it }
+        INLINE -> VariableDeclarationIntoWhenFix("Inline variable", subjectExpressionPointer) { it.initializer }
+        else -> error("Illegal action")
+    }
+}
+
+private fun KtProperty.action(element: KtElement): Action = when (val elementUsages = countUsages(element)) {
+    countUsages() -> if (elementUsages == 1) Action.INLINE else Action.MOVE
+    else -> Action.NOTHING
+}
 
 private fun KtWhenExpression.findDeclarationNear(): KtProperty? {
     val previousProperty = previousStatement() as? KtProperty ?: return null
     return previousProperty.takeIf { !it.isVar && it.hasInitializer() && it.nameIdentifier?.text == subjectExpression?.text }
 }
 
-private class MoveVariableDeclarationIntoWhenFix(val subjectExpressionPointer: SmartPsiElementPointer<KtExpression>) : LocalQuickFix {
-    override fun getName() = "Move the variable into `when`"
+private class VariableDeclarationIntoWhenFix(
+    private val actionName: String,
+    private val subjectExpressionPointer: SmartPsiElementPointer<KtExpression>,
+    private val transform: (KtProperty) -> KtExpression?
+) : LocalQuickFix {
+    override fun getName() = actionName
 
     override fun getFamilyName() = name
 
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
         val property = descriptor.psiElement as? KtProperty ?: return
         val subjectExpression = subjectExpressionPointer.element ?: return
-        subjectExpression.replace(property.copy())
+        val newElement = transform(property)?.copy() ?: return
+        subjectExpression.replace(newElement)
         property.delete()
     }
-}
+}
\ No newline at end of file
diff --git a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment.kt.after b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment.kt.after
index a2340155488..e680b931378 100644
--- a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment.kt.after
+++ b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withComment.kt.after
@@ -1,7 +1,7 @@
 fun foo() {
 
     // comment
-    when (val a = 1) {
+    when (1) {
         1 -> {
         }
         else -> {
diff --git a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withNewLine.kt.after b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withNewLine.kt.after
index 3c250ade30c..9c87694d24d 100644
--- a/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withNewLine.kt.after
+++ b/idea/testData/inspectionsLocal/moveVariableDeclarationIntoWhen/withNewLine.kt.after
@@ -1,7 +1,7 @@
 fun foo() {
 
 
-    when (val a = 1) {
+    when (1) {
         1 -> {
         }
         else -> {