diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt index e1c6f4fbf68..91c261cf259 100644 --- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt +++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/CommentSaver.kt @@ -28,7 +28,10 @@ import java.util.* import kotlin.properties.Delegates class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: Boolean = false/*TODO?*/) { - constructor(originalElement: PsiElement, saveLineBreaks: Boolean = false/*TODO?*/) : this(PsiChildRange.singleElement(originalElement), saveLineBreaks) + constructor(originalElement: PsiElement, saveLineBreaks: Boolean = false/*TODO?*/) : this( + PsiChildRange.singleElement(originalElement), + saveLineBreaks + ) private val SAVED_TREE_KEY = Key("SAVED_TREE") private val psiFactory = KtPsiFactory(originalElements.first!!) @@ -234,15 +237,25 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: putNewElementIntoMap(newPsiElement, treeElement) } - fun restore(resultElement: PsiElement, isCommentBeneathSingleLine: Boolean, forceAdjustIndent: Boolean = false) { - restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent, isCommentBeneathSingleLine) + fun restore( + resultElement: PsiElement, + isCommentBeneathSingleLine: Boolean, + isCommentInside: Boolean, + forceAdjustIndent: Boolean + ) { + restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent, isCommentBeneathSingleLine, isCommentInside) } fun restore(resultElement: PsiElement, forceAdjustIndent: Boolean = false) { restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent) } - fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false, isCommentBeneathSingleLine: Boolean = false) { + fun restore( + resultElements: PsiChildRange, + forceAdjustIndent: Boolean = false, + isCommentBeneathSingleLine: Boolean = false, + isCommentInside: Boolean = false + ) { assert(!isFinished) assert(!resultElements.isEmpty) @@ -264,7 +277,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: }) } - restoreComments(resultElements, isCommentBeneathSingleLine) + restoreComments(resultElements, isCommentBeneathSingleLine, isCommentInside) restoreLineBreaks() @@ -295,7 +308,11 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: isFinished = true } - private fun restoreComments(resultElements: PsiChildRange, isCommentBeneathSingleLine: Boolean = false) { + private fun restoreComments( + resultElements: PsiChildRange, + isCommentBeneathSingleLine: Boolean = false, + isCommentInside: Boolean = false + ) { var putAbandonedCommentsAfter = resultElements.last!! for (commentTreeElement in commentsToRestore) { @@ -333,6 +350,14 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: } } else { restored = putAbandonedCommentsAfter.parent.addBefore(comment, putAbandonedCommentsAfter) as PsiComment + + if (isCommentInside) { + val element = resultElements.first + val innerExpression = element?.lastChild?.getPrevSiblingIgnoringWhitespace() + innerExpression?.add(psiFactory.createWhiteSpace()) + innerExpression?.add(restored) + restored.delete() + } putAbandonedCommentsAfter = restored } diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt index 28d326d8dcc..094d159e1fa 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt @@ -18,6 +18,8 @@ package org.jetbrains.kotlin.idea.intentions import com.intellij.openapi.editor.Editor import com.intellij.psi.* +import org.jetbrains.kotlin.idea.intentions.loopToCallChain.nextStatement +import org.jetbrains.kotlin.idea.intentions.loopToCallChain.previousStatement import org.jetbrains.kotlin.idea.refactoring.getLineCount import org.jetbrains.kotlin.idea.refactoring.getLineNumber import org.jetbrains.kotlin.idea.util.CommentSaver @@ -53,6 +55,7 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja if (editor == null) throw IllegalArgumentException("This intention requires an editor") val expression = element.getTargetExpression(editor.caretModel.offset)!! var isCommentBeneath = false + var isCommentInside = false val psiFactory = KtPsiFactory(element) val semicolon = element.getNextSiblingIgnoringWhitespaceAndComments()?.takeIf { it.node.elementType == KtTokens.SEMICOLON } @@ -79,12 +82,31 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja } } + // Check for nested if/else + if (element is KtIfExpression && expression.isInSingleLine() && element.`else` != null && + element.parent.nextSibling is PsiWhiteSpace && + element.parent.nextSibling.nextSibling is PsiComment + ) { + isCommentInside = true + } + val nextComment = when { element is KtDoWhileExpression -> null // bound to the closing while element is KtIfExpression && expression === element.then && element.`else` != null -> null // bound to else else -> element.getNextSiblingIgnoringWhitespace().takeIf { it is PsiComment } } - val saver = if (nextComment == null) CommentSaver(element) else CommentSaver(PsiChildRange(element, nextComment)) + + val saver = when { + isCommentInside -> { + CommentSaver(element.parent.nextSibling.nextSibling) + } + else -> if (nextComment == null) CommentSaver(element) else CommentSaver(PsiChildRange(element, nextComment)) + } + + if (isCommentInside) { + element.parent.nextSibling.nextSibling.delete() + } + element.allChildren.filterIsInstance().toList().forEach { it.delete() } @@ -101,7 +123,7 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja } // Check for single line expression with comment beneath. - saver.restore(result, isCommentBeneath, forceAdjustIndent = false) + saver.restore(result, isCommentBeneath, isCommentInside, forceAdjustIndent = false) } private fun KtElement.getTargetExpression(caretLocation: Int): KtExpression? { diff --git a/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt new file mode 100644 index 00000000000..eb12d953266 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt @@ -0,0 +1,5 @@ +fun test(x: Int, b: Boolean) { + val foo = if (x == 1) + if (b) 1 else 2 // comment + else 0 +} \ No newline at end of file diff --git a/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt.after b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt.after new file mode 100644 index 00000000000..6e9454beef2 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt.after @@ -0,0 +1,7 @@ +fun test(x: Int, b: Boolean) { + val foo = if (x == 1) + if (b) 1 else { + 2 // comment + } + else 0 +} \ No newline at end of file diff --git a/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt new file mode 100644 index 00000000000..2016825c2f4 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt @@ -0,0 +1,8 @@ +fun test(x: Int, b: Boolean) { + val bar = when (x) { + 1 -> + if (b) 1 else 2 // comment + else -> + 0 + } +} \ No newline at end of file diff --git a/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt.after b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt.after new file mode 100644 index 00000000000..e8f44c4241c --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt.after @@ -0,0 +1,10 @@ +fun test(x: Int, b: Boolean) { + val bar = when (x) { + 1 -> + if (b) 1 else { + 2 // comment + } + else -> + 0 + } +} \ No newline at end of file diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java index 82b8b65c238..95e2c28c7a9 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java @@ -656,6 +656,16 @@ public class IntentionTestGenerated extends AbstractIntentionTest { runTest("idea/testData/intentions/addBraces/addBracesForElse.kt"); } + @TestMetadata("addBracesForElseWithCommentInsideIf.kt") + public void testAddBracesForElseWithCommentInsideIf() throws Exception { + runTest("idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideIf.kt"); + } + + @TestMetadata("addBracesForElseWithCommentInsideWhen.kt") + public void testAddBracesForElseWithCommentInsideWhen() throws Exception { + runTest("idea/testData/intentions/addBraces/addBracesForElseWithCommentInsideWhen.kt"); + } + @TestMetadata("addBracesForFor.kt") public void testAddBracesForFor() throws Exception { runTest("idea/testData/intentions/addBraces/addBracesForFor.kt");