Other way of handling post-insertion actions

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/MutableReplacementCode.kt

@@ -16,31 +16,40 @@
 
 package org.jetbrains.kotlin.idea.replacement
 
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.core.copied
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
-import org.jetbrains.kotlin.utils.SmartList
 import org.jetbrains.kotlin.utils.addToStdlib.singletonOrEmptyList
-import java.util.*
+
+private val POST_INSERTION_ACTION: Key<(KtElement) -> Unit> = Key("POST_INSERTION_ACTION")
 
 internal class MutableReplacementCode(
         var mainExpression: KtExpression?,
         val statementsBefore: MutableList<KtExpression>,
         val fqNamesToImport: MutableCollection<FqName>
 ) {
-    private val _postInsertionActions = HashMap<KtExpression, SmartList<(KtExpression) -> KtExpression>>()
-
-    val postInsertionActions: Map<KtExpression, List<(KtExpression) -> KtExpression>>
-        get() = _postInsertionActions
-
-    fun <TStatement : KtExpression> addPostInsertionAction(statementBefore: TStatement, action: (TStatement) -> TStatement) {
-        assert(statementBefore in statementsBefore)
+    fun <TElement : KtElement> addPostInsertionAction(element: TElement, action: (TElement) -> Unit) {
+        assert(element in this)
         @Suppress("UNCHECKED_CAST")
-        _postInsertionActions.getOrPut(statementBefore) { SmartList() }.add(action as (KtExpression) -> KtExpression)
+        element.putCopyableUserData(POST_INSERTION_ACTION, action as (KtElement) -> Unit)
+    }
+
+    fun performPostInsertionActions(elements: Collection<PsiElement>) {
+        for (element in elements) {
+            element.forEachDescendantOfType<KtElement> {
+                val action = it.getCopyableUserData(POST_INSERTION_ACTION)
+                if (action != null) {
+                    it.putCopyableUserData(POST_INSERTION_ACTION, null)
+                    action.invoke(it)
+                }
+            }
+        }
     }
 
     fun replaceExpression(oldExpression: KtExpression, newExpression: KtExpression): KtExpression {

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ReplacementPerformer.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
 import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.utils.addToStdlib.singletonList
 import java.util.*
 
 internal abstract class ReplacementPerformer<TElement : KtElement>(
@@ -45,6 +46,9 @@ internal class AnnotationEntryReplacementPerformer(
 
         val dummyAnnotationEntry = createByPattern("@Dummy($0)", replacement.mainExpression!!) { psiFactory.createAnnotationEntry(it) }
         val replaced = elementToBeReplaced.replace(dummyAnnotationEntry)
+
+        replacement.performPostInsertionActions(replaced.singletonList())
+
         var range = PsiChildRange.singleElement(replaced)
         range = postProcessing(range)
 
@@ -63,7 +67,6 @@ internal class ExpressionReplacementPerformer(
 
     override fun doIt(postProcessing: (PsiChildRange) -> PsiChildRange): KtExpression {
         val insertedStatements = ArrayList<KtExpression>()
-        val toInsertedStatementMap = HashMap<KtExpression, KtExpression>()
         for (statement in replacement.statementsBefore) {
             // copy the statement if it can get invalidated by findOrCreateBlockToInsertStatement()
             val statementToUse = if (statement.isPhysical) statement.copy() else statement
@@ -73,16 +76,12 @@ internal class ExpressionReplacementPerformer(
             val inserted = block.addBefore(statementToUse, anchor) as KtExpression
             block.addBefore(psiFactory.createNewLine(), anchor)
             block.addBefore(psiFactory.createNewLine(), inserted)
-            toInsertedStatementMap.put(statement, inserted)
             insertedStatements.add(inserted)
         }
 
         val replaced = elementToBeReplaced.replace(replacement.mainExpression!!) //TODO: support null here
 
-        for ((statement, actions) in replacement.postInsertionActions) {
-            val inserted = toInsertedStatementMap[statement]!!
-            actions.forEach { it(inserted) }
-        }
+        replacement.performPostInsertionActions(insertedStatements + replaced.singletonList())
 
         var range = if (insertedStatements.isEmpty()) {
             PsiChildRange.singleElement(replaced)

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/introduceValue.kt

@@ -91,7 +91,7 @@ internal fun MutableReplacementCode.introduceValue(
             val declaration = psiFactory.createDeclarationByPattern<KtVariableDeclaration>("val $0 = $1", name, value)
             statementsBefore.add(0, declaration)
             if (explicitType != null) {
-                addPostInsertionAction(declaration) { it.setType(explicitType!!); it }
+                addPostInsertionAction(declaration) { it.setType(explicitType!!) }
             }
 
             replaceUsages(name)
@@ -123,7 +123,7 @@ internal fun MutableReplacementCode.introduceValue(
 
             appendFixedText("}")
         }
-        statementsBefore.clear() //TODO: postInsertActions!
+        statementsBefore.clear()
     }
 }