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 59a326ef294..e1c6f4fbf68 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 @@ -123,9 +123,9 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: private class LineBreakTreeElement() : TreeElement() private class CommentTreeElement( - val commentText: String, - val spaceBefore: String, - val spaceAfter: String + val commentText: String, + val spaceBefore: String, + val spaceAfter: String ) : TreeElement() { companion object { fun create(comment: PsiComment): CommentTreeElement { @@ -234,11 +234,15 @@ 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, forceAdjustIndent: Boolean = false) { restore(PsiChildRange.singleElement(resultElement), forceAdjustIndent) } - fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false) { + fun restore(resultElements: PsiChildRange, forceAdjustIndent: Boolean = false, isCommentBeneathSingleLine: Boolean = false) { assert(!isFinished) assert(!resultElements.isEmpty) @@ -260,7 +264,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: }) } - restoreComments(resultElements) + restoreComments(resultElements, isCommentBeneathSingleLine) restoreLineBreaks() @@ -291,7 +295,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: isFinished = true } - private fun restoreComments(resultElements: PsiChildRange) { + private fun restoreComments(resultElements: PsiChildRange, isCommentBeneathSingleLine: Boolean = false) { var putAbandonedCommentsAfter = resultElements.last!! for (commentTreeElement in commentsToRestore) { @@ -316,25 +320,30 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: val whiteSpace = restored.nextLeaf(skipEmptyElements = true) as? PsiWhiteSpace if (whiteSpace == null) { parent.addAfter(psiFactory.createWhiteSpace("\n"), restored) - } - else if (!whiteSpace.textContains('\n')) { + } else if (!whiteSpace.textContains('\n')) { val newWhiteSpace = psiFactory.createWhiteSpace("\n" + whiteSpace.text) whiteSpace.replace(newWhiteSpace) } } - } - else { + } else { restored = parent.addBefore(comment, anchorElement) as PsiComment if (commentTreeElement.spaceAfter.isNotEmpty()) { parent.addBefore(psiFactory.createWhiteSpace(commentTreeElement.spaceAfter), anchorElement) } } - } - else { + } else { restored = putAbandonedCommentsAfter.parent.addBefore(comment, putAbandonedCommentsAfter) as PsiComment putAbandonedCommentsAfter = restored } + // shift (possible contained) comment in expression underneath braces + if (isCommentBeneathSingleLine && resultElements.count() == 1) { + val element = resultElements.first + element?.add(psiFactory.createWhiteSpace("\n")) + element?.add(restored) + restored.delete() + } + bindNewElement(restored, commentTreeElement) // will be used when restoring line breaks if (restored.tokenType == KtTokens.EOL_COMMENT) { @@ -352,9 +361,9 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: fun findRestored(leaf: TreeElement): PsiElement? { if (leaf is LineBreakTreeElement) return null return leaf.parentsWithSelf - .takeWhile { it != lineBreakParent } - .mapNotNull { toNewPsiElementMap[it]?.first() } //TODO: what about multiple? - .firstOrNull() + .takeWhile { it != lineBreakParent } + .mapNotNull { toNewPsiElementMap[it]?.first() } //TODO: what about multiple? + .firstOrNull() } @@ -365,8 +374,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: psiElement = skipTokensForward(psiElement, tokensToMatch.asReversed()) psiElement?.restoreLineBreakAfter() break - } - else { + } else { if (leaf !is TokenTreeElement) break tokensToMatch.add(leaf.tokenType) } @@ -391,8 +399,7 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: if (whitespace == null) { addAfter.parent.addAfter(psiFactory.createNewLine(), addAfter) - } - else { + } else { whitespace.replace(psiFactory.createWhiteSpace("\n" + whitespace.text)) } @@ -419,8 +426,8 @@ class CommentSaver(originalElements: PsiChildRange, private val saveLineBreaks: private fun PsiElement.anchorToAddCommentOrSpace(before: Boolean): PsiElement { return parentsWithSelf - .dropWhile { it.parent !is PsiFile && (if (before) it.prevSibling else it.nextSibling) == null } - .first() + .dropWhile { it.parent !is PsiFile && (if (before) it.prevSibling else it.nextSibling) == null } + .first() } private fun chooseAnchor(anchorBefore: Anchor?, anchorAfter: Anchor?): Anchor? { diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt index 08741a48933..28d326d8dcc 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt @@ -17,14 +17,15 @@ package org.jetbrains.kotlin.idea.intentions import com.intellij.openapi.editor.Editor -import com.intellij.psi.PsiComment -import com.intellij.psi.PsiWhiteSpace +import com.intellij.psi.* +import org.jetbrains.kotlin.idea.refactoring.getLineCount import org.jetbrains.kotlin.idea.refactoring.getLineNumber import org.jetbrains.kotlin.idea.util.CommentSaver +import org.jetbrains.kotlin.j2k.isInSingleLine import org.jetbrains.kotlin.lexer.KtTokens import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.* -import java.lang.IllegalArgumentException +import org.jetbrains.kotlin.resolve.calls.CallExpressionElement class AddBracesIntention : SelfTargetingIntention(KtElement::class.java, "Add braces") { override fun isApplicableTo(element: KtElement, caretOffset: Int): Boolean { @@ -51,7 +52,7 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja override fun applyTo(element: KtElement, editor: Editor?) { if (editor == null) throw IllegalArgumentException("This intention requires an editor") val expression = element.getTargetExpression(editor.caretModel.offset)!! - + var isCommentBeneath = false val psiFactory = KtPsiFactory(element) val semicolon = element.getNextSiblingIgnoringWhitespaceAndComments()?.takeIf { it.node.elementType == KtTokens.SEMICOLON } @@ -63,6 +64,21 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja semicolon.delete() } + // Check for single line if expression + if (element is KtIfExpression && expression.isInSingleLine() && element.`else` == null) { + // Check if a comment is actually underneath (\n) the expression + val allElements = element.siblings(withItself = false).filterIsInstance() + val sibling = allElements.firstOrNull { it is PsiComment } + if (sibling is PsiComment) { + // Check if \n before first received comment sibling + // if false, the normal procedure of adding braces occurs. + isCommentBeneath = + sibling.prevSibling is PsiWhiteSpace && + sibling.prevSibling.textContains('\n') && + (sibling.prevSibling.prevSibling is PsiComment || sibling.prevSibling.prevSibling is PsiElement) + } + } + 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 @@ -83,7 +99,9 @@ class AddBracesIntention : SelfTargetingIntention(KtElement::class.ja is KtIfExpression -> (result?.parent?.nextSibling as? PsiWhiteSpace)?.delete() } - saver.restore(result) + + // Check for single line expression with comment beneath. + saver.restore(result, isCommentBeneath, forceAdjustIndent = false) } private fun KtElement.getTargetExpression(caretLocation: Int): KtExpression? { diff --git a/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt new file mode 100644 index 00000000000..d91ca239af4 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt @@ -0,0 +1,6 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) foo() + // ccc +} \ No newline at end of file diff --git a/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt.after b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt.after new file mode 100644 index 00000000000..e7751d7be52 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt.after @@ -0,0 +1,8 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) { + foo() + } + // ccc +} diff --git a/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt new file mode 100644 index 00000000000..aa51dd97243 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt @@ -0,0 +1,6 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) foo() // bbb + // ccc +} \ No newline at end of file diff --git a/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt.after b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt.after new file mode 100644 index 00000000000..24c6a5ee701 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt.after @@ -0,0 +1,8 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) { + foo() // bbb + } + // ccc +} diff --git a/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt new file mode 100644 index 00000000000..b8742f760b6 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt @@ -0,0 +1,6 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) /* aaa */ foo() // bbb + // ccc +} \ No newline at end of file diff --git a/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt.after b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt.after new file mode 100644 index 00000000000..6b55320bb47 --- /dev/null +++ b/idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt.after @@ -0,0 +1,8 @@ +fun foo() {} + +fun test(b: Boolean) { + if (b) { + /* aaa */ foo() // bbb + } + // ccc +} diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java index 6d56269ed2c..98c03653002 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java @@ -696,6 +696,21 @@ public class IntentionTestGenerated extends AbstractIntentionTest { runTest("idea/testData/intentions/addBraces/addBracesForIfWithSemicolonAndExpression.kt"); } + @TestMetadata("addBracesForSingleLineIfWithCommentBeneath.kt") + public void testAddBracesForSingleLineIfWithCommentBeneath() throws Exception { + runTest("idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath.kt"); + } + + @TestMetadata("addBracesForSingleLineIfWithCommentBeneath2.kt") + public void testAddBracesForSingleLineIfWithCommentBeneath2() throws Exception { + runTest("idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath2.kt"); + } + + @TestMetadata("addBracesForSingleLineIfWithCommentBeneath3.kt") + public void testAddBracesForSingleLineIfWithCommentBeneath3() throws Exception { + runTest("idea/testData/intentions/addBraces/addBracesForSingleLineIfWithCommentBeneath3.kt"); + } + @TestMetadata("addBracesForWhile.kt") public void testAddBracesForWhile() throws Exception { runTest("idea/testData/intentions/addBraces/addBracesForWhile.kt");