PSI Pattern Patching: Use precomputed context for matching of substitution parameters

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractableAnalysisUtil.kt

@@ -58,7 +58,6 @@ import org.jetbrains.jet.lang.cfg.pseudocodeTraverser.TraversalOrder
 import org.jetbrains.jet.lang.resolve.bindingContextUtil.getTargetFunctionDescriptor
 import org.jetbrains.jet.lang.resolve.OverridingUtil
 import org.jetbrains.jet.lang.resolve.bindingContextUtil.isUsedAsStatement
-import org.jetbrains.jet.lang.psi.psiUtil.isAncestor
 import org.jetbrains.jet.lang.resolve.DescriptorToSourceUtils
 import org.jetbrains.jet.plugin.imports.importableFqNameSafe
 import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValue.Initializer
@@ -67,6 +66,7 @@ import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValue.Expressi
 import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValue.Jump
 import org.jetbrains.jet.lang.cfg.pseudocodeTraverser.traverseFollowingInstructions
 import org.jetbrains.jet.plugin.refactoring.extractFunction.OutputValueBoxer.AsList
+import org.jetbrains.jet.plugin.refactoring.getContextForContainingDeclarationBody
 
 private val DEFAULT_FUNCTION_NAME = "myFun"
 private val DEFAULT_RETURN_TYPE = KotlinBuiltIns.getInstance().getUnitType()
@@ -655,21 +655,8 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
 
     val commonParent = PsiTreeUtil.findCommonParent(originalElements) as JetElement
 
-    val enclosingDeclaration = commonParent.getParentByType(javaClass<JetDeclaration>(), true)
+    val bindingContext = commonParent.getContextForContainingDeclarationBody()
-    val bodyElement = when (enclosingDeclaration) {
+    if (bindingContext == null) return noContainerError
-        is JetDeclarationWithBody -> enclosingDeclaration.getBodyExpression()
-        is JetWithExpressionInitializer -> enclosingDeclaration.getInitializer()
-        is JetMultiDeclaration -> enclosingDeclaration.getInitializer()
-        is JetParameter -> enclosingDeclaration.getDefaultValue()
-        is JetClassInitializer -> enclosingDeclaration.getBody()
-        is JetClass -> {
-            val delegationSpecifierList = enclosingDeclaration.getDelegationSpecifierList()
-            if (delegationSpecifierList.isAncestor(commonParent)) commonParent else return noContainerError
-        }
-        else -> return noContainerError
-    }
-    val resolveSession = originalFile.getLazyResolveSession()
-    val bindingContext = resolveSession.resolveToElement(bodyElement)
 
     val pseudocodeDeclaration = PsiTreeUtil.getParentOfType(
             commonParent, javaClass<JetDeclarationWithBody>(), javaClass<JetClassOrObject>()
@@ -693,7 +680,7 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
             analyzeControlFlow(
                     localInstructions,
                     pseudocode,
-                    resolveSession.getModuleDescriptor(),
+                    originalFile.getLazyResolveSession().getModuleDescriptor(),
                     bindingContext,
                     modifiedVarDescriptorsForControlFlow,
                     options,
@@ -712,7 +699,8 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
         )
     }
 
-    checkDeclarationsMovingOutOfScope(enclosingDeclaration!!, controlFlow, bindingContext)?.let { messages.add(it) }
+    val enclosingDeclaration = commonParent.getParentByType(javaClass<JetDeclaration>(), true)!!
+    checkDeclarationsMovingOutOfScope(enclosingDeclaration, controlFlow, bindingContext)?.let { messages.add(it) }
 
     val functionNameValidator =
             JetNameValidatorImpl(

idea/src/org/jetbrains/jet/plugin/refactoring/jetRefactoringUtil.kt

@@ -68,6 +68,10 @@ import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.openapi.util.Computable
 import com.intellij.openapi.ui.DialogWrapper
+import org.jetbrains.jet.lang.resolve.BindingContext
+import org.jetbrains.jet.lang.psi.psiUtil.getParentByType
+import org.jetbrains.jet.lang.psi.psiUtil.isAncestor
+import org.jetbrains.jet.plugin.caches.resolve.getLazyResolveSession
 
 /**
  * Replace [[JetSimpleNameExpression]] (and its enclosing qualifier) with qualified element given by FqName
@@ -282,4 +286,21 @@ fun PsiElement.getLineCount(): Int {
     return (getText() ?: "").count { it == '\n' } + 1
 }
 
-fun PsiElement.isMultiLine(): Boolean = getLineCount() > 1
+fun PsiElement.isMultiLine(): Boolean = getLineCount() > 1
+
+public fun JetElement.getContextForContainingDeclarationBody(): BindingContext? {
+    val enclosingDeclaration = getParentByType(javaClass<JetDeclaration>(), true)
+    val bodyElement = when (enclosingDeclaration) {
+        is JetDeclarationWithBody -> enclosingDeclaration.getBodyExpression()
+        is JetWithExpressionInitializer -> enclosingDeclaration.getInitializer()
+        is JetMultiDeclaration -> enclosingDeclaration.getInitializer()
+        is JetParameter -> enclosingDeclaration.getDefaultValue()
+        is JetClassInitializer -> enclosingDeclaration.getBody()
+        is JetClass -> {
+            val delegationSpecifierList = enclosingDeclaration.getDelegationSpecifierList()
+            if (delegationSpecifierList.isAncestor(this)) this else null
+        }
+        else -> null
+    }
+    return bodyElement?.let { getContainingJetFile().getLazyResolveSession().resolveToElement(it) }
+}

idea/src/org/jetbrains/jet/plugin/util/psi/patternMatching/JetPsiUnifier.kt

@@ -56,7 +56,6 @@ import org.jetbrains.jet.lang.psi.JetDeclaration
 import org.jetbrains.jet.lang.types.ErrorUtils
 import com.intellij.lang.ASTNode
 import com.intellij.util.containers.MultiMap
-import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache
 import org.jetbrains.jet.lang.psi.JetCallableReferenceExpression
 import org.jetbrains.jet.lang.resolve.scopes.receivers.ThisReceiver
 import org.jetbrains.jet.lang.psi.JetThisExpression
@@ -88,6 +87,7 @@ import org.jetbrains.jet.lang.psi.JetReturnExpression
 import org.jetbrains.jet.lang.psi.JetProperty
 import org.jetbrains.jet.lang.psi.JetDelegatorToSuperClass
 import org.jetbrains.jet.lang.psi.JetDelegationSpecifier
+import org.jetbrains.jet.plugin.refactoring.getContextForContainingDeclarationBody
 
 public trait UnificationResult {
     public enum class Status {
@@ -139,11 +139,19 @@ public class JetPsiUnifier(
             val originalTarget: JetPsiRange,
             val originalPattern: JetPsiRange
     ) {
+        val patternContext: BindingContext = originalPattern.getBindingContext()
+        val targetContext: BindingContext = originalTarget.getBindingContext()
         val substitution = HashMap<UnifierParameter, JetExpression>()
         val declarationPatternsToTargets = MultiMap<DeclarationDescriptor, DeclarationDescriptor>()
         val weakMatches = HashMap<JetElement, JetElement>()
         var checkEquivalence: Boolean = false
 
+        private fun JetPsiRange.getBindingContext(): BindingContext {
+            val element = (this as? JetPsiRange.ListRange)?.startElement as? JetElement
+            if ((element?.getContainingFile() as? JetFile)?.doNotAnalyze != null) return BindingContext.EMPTY
+            return element?.getContextForContainingDeclarationBody() ?: BindingContext.EMPTY
+        }
+
         private fun matchDescriptors(d1: DeclarationDescriptor?, d2: DeclarationDescriptor?): Boolean {
             if (d1 == d2 || d2 in declarationPatternsToTargets[d1] || d1 in declarationPatternsToTargets[d2]) return true
             if (d1 == null || d2 == null) return false
@@ -275,10 +283,7 @@ public class JetPsiUnifier(
             }
         }
 
-        private val JetElement.bindingContext: BindingContext get() {
+        private val JetElement.bindingContext: BindingContext get() = if (this in originalPattern) patternContext else targetContext
-            if ((getContainingFile() as? JetFile)?.doNotAnalyze != null) return BindingContext.EMPTY
-            return AnalyzerFacadeWithCache.getContextForElement(this)
-        }
 
         private fun JetElement.getAdjustedResolvedCall(): ResolvedCall<*>? {
             val rc = getResolvedCall(bindingContext)?.let {