From b39e9ebdb2fdc13a4e81aa8700279c209ceb4184 Mon Sep 17 00:00:00 2001 From: Valentin Kipyatkov Date: Sat, 23 May 2015 19:24:55 +0300 Subject: [PATCH] Implemented special CommentSaver helper for not loosing comments when doing code transformation --- .../org/jetbrains/kotlin/psi/JetPsiFactory.kt | 2 +- .../jetbrains/kotlin/psi/psiUtil/psiUtils.kt | 2 + .../kotlin/idea/core/CommentSaver.kt | 217 ++++++++++++++++++ 3 files changed, 220 insertions(+), 1 deletion(-) create mode 100644 idea/idea-core/src/org/jetbrains/kotlin/idea/core/CommentSaver.kt diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiFactory.kt b/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiFactory.kt index 6b57f2987d8..3e1718f8671 100644 --- a/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiFactory.kt +++ b/compiler/frontend/src/org/jetbrains/kotlin/psi/JetPsiFactory.kt @@ -38,7 +38,7 @@ import java.io.PrintWriter import java.io.StringWriter public fun JetPsiFactory(project: Project?): JetPsiFactory = JetPsiFactory(project!!) -public fun JetPsiFactory(elementForProject: JetElement): JetPsiFactory = JetPsiFactory(elementForProject.getProject()) +public fun JetPsiFactory(elementForProject: PsiElement): JetPsiFactory = JetPsiFactory(elementForProject.getProject()) public var JetFile.doNotAnalyze: String? by UserDataProperty(Key.create("DO_NOT_ANALYZE")) public var JetFile.analysisContext: PsiElement? by UserDataProperty(Key.create("ANALYSIS_CONTEXT")) diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt b/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt index 65a44521c89..2c6366cd23a 100644 --- a/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt +++ b/compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt @@ -53,6 +53,8 @@ public data class PsiChildRange(public val first: PsiElement?, public val last: companion object { public val EMPTY: PsiChildRange = PsiChildRange(null, null) + + public fun singleElement(element: PsiElement): PsiChildRange = PsiChildRange(element, element) } } diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/CommentSaver.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/CommentSaver.kt new file mode 100644 index 00000000000..52aa250c1af --- /dev/null +++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/CommentSaver.kt @@ -0,0 +1,217 @@ +/* + * Copyright 2010-2015 JetBrains s.r.o. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.jetbrains.kotlin.idea.core + +import com.intellij.openapi.util.Key +import com.intellij.psi.PsiComment +import com.intellij.psi.PsiElement +import com.intellij.psi.PsiRecursiveElementVisitor +import com.intellij.psi.PsiWhiteSpace +import org.jetbrains.kotlin.psi.JetPsiFactory +import org.jetbrains.kotlin.psi.psiUtil.* +import java.util.* + +public class CommentSaver(originalElements: PsiChildRange) { + public constructor(originalElement: PsiElement) : this(PsiChildRange.singleElement(originalElement)) + + private val SAVED_TREE_KEY = Key("SAVED_TREE") + + private class TreeElement( + var parent: TreeElement? = null, + var prev: TreeElement? = null, + var next: TreeElement? = null, + var firstChild: TreeElement? = null, + var lastChild: TreeElement? = null + + ) { + val children: Sequence + get() = sequence({ firstChild }, { it.next }) + + val reverseChildren: Sequence + get() = sequence({ lastChild }, { it.prev }) + + val prevSiblings: Sequence + get() = sequence({ prev }, { it.prev }) + + val nextSiblings: Sequence + get() = sequence({ next }, { it.next }) + + fun withDescendants(leftToRight: Boolean): Sequence { + val children = if (leftToRight) children else reverseChildren + return sequenceOf(this) + children.flatMap { it.withDescendants(leftToRight) } + } + + val prevElements: Sequence + get() = prevSiblings.flatMap { it.withDescendants(leftToRight = false) } + + val nextElements: Sequence + get() = nextSiblings.flatMap { it.withDescendants(leftToRight = true) } + } + + private data class CommentData( + val commentText: String, + val spaceBefore: String, + val spaceAfter: String + ) + + private fun PsiComment.buildData(): CommentData { + val spaceBefore = (prevLeaf(skipEmptyElements = true) as? PsiWhiteSpace)?.getText() ?: "" + val spaceAfter = (nextLeaf(skipEmptyElements = true) as? PsiWhiteSpace)?.getText() ?: "" + return CommentData(getText(), spaceBefore, spaceAfter) + } + + private val commentsToRestore = LinkedHashMap() + + init { + if (originalElements.any { it.anyDescendantOfTypeTemp() }) { + originalElements.save(null) + + for (element in originalElements) { + element.accept(object : PsiRecursiveElementVisitor(){ + override fun visitComment(comment: PsiComment) { + commentsToRestore.put(comment.savedTreeElement!!, comment.buildData()) + } + }) + } + } + } + + private fun PsiChildRange.save(parentTreeElement: TreeElement?) { + var first: TreeElement? = null + var last: TreeElement? = null + for (child in this) { + if (!child.shouldSave) continue + assert(child.savedTreeElement == null) + + val savedChild = TreeElement(parent = parentTreeElement, prev = last) + child.putCopyableUserData(SAVED_TREE_KEY, savedChild) + last?.next = savedChild + last = savedChild + + if (first == null) { + first = savedChild + } + + child.allChildren.save(savedChild) + } + + parentTreeElement?.firstChild = first + parentTreeElement?.lastChild = last + } + + private val PsiElement.shouldSave: Boolean + get() = this !is PsiWhiteSpace + + private val PsiElement.savedTreeElement: TreeElement? + get() = getCopyableUserData(SAVED_TREE_KEY) + + public fun deleteCommentsInside(element: PsiElement) { + element.accept(object : PsiRecursiveElementVisitor() { + override fun visitComment(comment: PsiComment) { + val treeElement = comment.savedTreeElement + if (treeElement != null) { + commentsToRestore.remove(treeElement) + } + } + }) + } + + public fun restoreComments(resultElement: PsiElement) { + restoreComments(PsiChildRange.singleElement(resultElement)) + } + + public fun restoreComments(resultElements: PsiChildRange) { + doRestoreComments(resultElements) + + // clear user data + resultElements.forEach { + it.accept(object : PsiRecursiveElementVisitor() { + override fun visitElement(element: PsiElement) { + element.putCopyableUserData(SAVED_TREE_KEY, null) + super.visitElement(element) + } + }) + } + } + + private fun doRestoreComments(resultElements: PsiChildRange) { + assert(!resultElements.isEmpty) + + if (commentsToRestore.isEmpty()) return + + // remove comments that present inside resultElements from commentsToRestore + resultElements.forEach { deleteCommentsInside(it) } + if (commentsToRestore.isEmpty()) return + + val toNewPsiElementMap = HashMap>() + for (element in resultElements) { + element.accept(object : PsiRecursiveElementVisitor() { + override fun visitElement(element: PsiElement) { + val treeElement = element.savedTreeElement + if (treeElement != null) { + toNewPsiElementMap.getOrPut(treeElement) { ArrayList(1) }.add(element) + } + super.visitElement(element) + } + }) + } + + val psiFactory = JetPsiFactory(resultElements.first!!) + for ((treeElement, commentData) in commentsToRestore) { + val comment = psiFactory.createComment(commentData.commentText) + + //TODO: should we restore multiple? + val putAfter: PsiElement? = treeElement.prevElements.map { toNewPsiElementMap[it]?.first() }.filterNotNull().firstOrNull() + if (putAfter != null) { + val parent = putAfter.getParent() + parent.addAfter(comment, putAfter) + if (commentData.spaceBefore.isNotEmpty()) { + parent.addAfter(psiFactory.createWhiteSpace(commentData.spaceBefore), putAfter) + } + continue + } + + val putBefore: PsiElement? = treeElement.nextElements.map { toNewPsiElementMap[it]?.first() }.filterNotNull().firstOrNull() + if (putBefore != null) { + val parent = putBefore.getParent() + parent.addBefore(comment, putBefore) + if (commentData.spaceAfter.isNotEmpty()) { + parent.addBefore(psiFactory.createWhiteSpace(commentData.spaceAfter), putBefore) + } + continue + } + + // TODO: everything deleted, don't restore anything? + } + } + + //TODO + private inline fun PsiElement.anyDescendantOfTypeTemp(): Boolean { + var result = false + this.accept(object : PsiRecursiveElementVisitor(){ + override fun visitElement(element: PsiElement) { + if (result) return + if (element is T) { + result = true + return + } + super.visitElement(element) + } + }) + return result + } +} \ No newline at end of file