INT: generalize MatchToIfLetIntention to arbitrary arm count

src/main/kotlin/org/rust/ide/intentions/MatchToIfLetIntention.kt

@@ -8,9 +8,11 @@ package org.rust.ide.intentions
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
+import org.rust.ide.refactoring.findBinding
 import org.rust.lang.core.psi.*
 import org.rust.lang.core.psi.ext.ancestorStrict
-import org.rust.lang.core.psi.ext.getNextNonCommentSibling
+import org.rust.lang.core.types.ty.TyUnit
+import org.rust.lang.core.types.type
 
 class MatchToIfLetIntention : RsElementBaseIntentionAction<MatchToIfLetIntention.Context>() {
     override fun getText() = "Convert match statement to if let"
@@ -19,38 +21,92 @@ class MatchToIfLetIntention : RsElementBaseIntentionAction<MatchToIfLetIntention
     data class Context(
         val match: RsMatchExpr,
         val matchTarget: RsExpr,
-        val nonVoidArm: RsMatchArm,
-        val pat: RsPat
+        val matchBody: RsMatchBody
     )
 
     override fun findApplicableContext(project: Project, editor: Editor, element: PsiElement): Context? {
         val matchExpr = element.ancestorStrict<RsMatchExpr>() ?: return null
         if (element != matchExpr.match) return null
         val matchTarget = matchExpr.expr ?: return null
         val matchBody = matchExpr.matchBody ?: return null
-        val matchArmList = matchBody.matchArmList
+        if (matchBody.matchArmList.isEmpty()) return null
+        if (matchBody.matchArmList.any { it.matchArmGuard != null || it.outerAttrList.isNotEmpty() }) return null
 
-        val nonVoidArm = matchArmList.singleOrNull { it.expr?.isVoid == false } ?: return null
-        if (nonVoidArm.matchArmGuard != null || nonVoidArm.outerAttrList.isNotEmpty()) return null
-        val pat = nonVoidArm.pat
-
-        return Context(matchExpr, matchTarget, nonVoidArm, pat)
+        return Context(matchExpr, matchTarget, matchBody)
     }
 
     override fun invoke(project: Project, editor: Editor, ctx: Context) {
-        val (matchExpr, matchTarget, arm, pat) = ctx
+        val arms = ctx.matchBody.matchArmList
+        val lastArm = arms.lastOrNull() ?: return
+
+        // else is required for `if` expressions if the resulting type is not ()
+        val hasUnitType = ctx.match.type is TyUnit
+        val lastArmHasBinding = lastArm.pat.findBinding() != null
 
-        var bodyText = arm.expr?.text ?: return
-        if (arm.expr !is RsBlockExpr) {
-            bodyText = "{\n$bodyText\n}"
+        fun isExprEmpty(expr: RsExpr?): Boolean {
+            return expr?.let {
+                when {
+                    it is RsUnitExpr -> true
+                    it is RsBlockExpr && it.block.children.isEmpty() -> true
+                    else -> false
+                }
+            } ?: false
         }
 
-        val exprText = "if let ${pat.text} = ${matchTarget.text} $bodyText"
-        val rustIfLetExprElement = RsPsiFactory(project).createExpression(exprText) as RsIfExpr
-        matchExpr.replace(rustIfLetExprElement)
-    }
+        val lastArmExprEmpty = isExprEmpty(lastArm.expr)
+
+        val text = buildString {
+            var hasElseBlock = false
+
+            for ((index, arm) in arms.withIndex()) {
+                val expr = arm.expr ?: continue
+
+                var armText = "if let ${arm.pat.text} = ${ctx.matchTarget.text}"
+                when (index) {
+                    // First, nothing special
+                    0 -> {}
+                    // Last, handle else block
+                    arms.size - 1 -> {
+                        // We don't need an else block and the last arm is empty, skip it
+                        if (hasUnitType && lastArmExprEmpty) {
+                            break
+                        }
+                        // We can coerce the last arm to an else block, because it has no bindings
+                        else if (!lastArmHasBinding) {
+                            armText = " else"
+                            hasElseBlock = true
+                        } else {
+                            armText = " else $armText"
+                        }
+                    }
+                    else -> armText = " else $armText"
+                }
+                append(armText)
+                append(' ')
 
-    private val RsExpr.isVoid: Boolean
-        get() = (this is RsBlockExpr && block.lbrace.getNextNonCommentSibling() == block.rbrace)
-            || this is RsUnitExpr
+                val innerExprText = when {
+                    !hasUnitType && isExprEmpty(expr) -> "unreachable!()"
+                    expr is RsBlockExpr -> {
+                        val text = expr.block.text
+                        val start = expr.block.lbrace.startOffsetInParent + 1
+                        val end = expr.block.rbrace?.startOffsetInParent ?: text.length
+                        text.substring(start, end)
+                    }
+                    else -> expr.text
+                }
+
+                val exprText = "{\n    $innerExprText\n}"
+                append(exprText)
+            }
+
+            // We need to add an extra else arm in this case
+            if (!hasUnitType && !hasElseBlock) {
+                append(" else {\n    unreachable!()\n}")
+            }
+        }
+
+        val factory = RsPsiFactory(project)
+        val ifLetExpr = factory.createExpression(text)
+        ctx.match.replace(ifLetExpr)
+    }
 }