"Surround with" uses partial body resolve

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/KotlinSurrounderUtils.java

@@ -20,14 +20,9 @@ import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.idea.JetBundle;
-import org.jetbrains.kotlin.idea.caches.resolve.ResolvePackage;
 import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils;
 import org.jetbrains.kotlin.psi.JetBlockExpression;
-import org.jetbrains.kotlin.psi.JetExpression;
-import org.jetbrains.kotlin.resolve.BindingContext;
-import org.jetbrains.kotlin.types.JetType;
 
 public class KotlinSurrounderUtils {
     public static String SURROUND_WITH = JetBundle.message("surround.with");
@@ -44,12 +39,6 @@ public class KotlinSurrounderUtils {
         block.addRangeAfter(statements[0], statements[statements.length - 1], lBrace);
     }
 
-    @Nullable
-    public static JetType getExpressionType(JetExpression expression) {
-        BindingContext expressionBindingContext = ResolvePackage.analyze(expression);
-        return expressionBindingContext.get(BindingContext.EXPRESSION_TYPE, expression);
-    }
-
     public static void showErrorHint(@NotNull Project project, @NotNull Editor editor, @NotNull String message) {
         CodeInsightUtils.showErrorHint(project, editor, message, SURROUND_WITH, null);
     }

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/MoveDeclarationsOutHelper.java

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.idea.util.ShortenReferences;
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.types.JetType;
 
 import java.util.ArrayList;
@@ -116,9 +117,9 @@ public class MoveDeclarationsOutHelper {
 
     @NotNull
     private static JetType getPropertyType(@NotNull JetProperty property) {
-        BindingContext expressionBindingContext = ResolvePackage.analyze(property);
+        BindingContext bindingContext = ResolvePackage.analyze(property, BodyResolveMode.PARTIAL);
 
-        VariableDescriptor propertyDescriptor = expressionBindingContext.get(BindingContext.VARIABLE, property);
+        VariableDescriptor propertyDescriptor = bindingContext.get(BindingContext.VARIABLE, property);
         assert propertyDescriptor != null : "Couldn't resolve property to property descriptor " + property.getText();
         return propertyDescriptor.getType();
     }

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/expression/KotlinExpressionSurrounder.java

@@ -24,10 +24,12 @@ import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
-import org.jetbrains.kotlin.idea.codeInsight.surroundWith.KotlinSurrounderUtils;
+import org.jetbrains.kotlin.idea.caches.resolve.ResolvePackage;
 import org.jetbrains.kotlin.psi.JetCallExpression;
 import org.jetbrains.kotlin.psi.JetExpression;
 import org.jetbrains.kotlin.psi.JetQualifiedExpression;
+import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.types.JetType;
 
 public abstract class KotlinExpressionSurrounder implements Surrounder {
@@ -42,7 +44,7 @@ public abstract class KotlinExpressionSurrounder implements Surrounder {
         if (expression instanceof JetCallExpression && expression.getParent() instanceof JetQualifiedExpression) {
             return false;
         }
-        JetType type = KotlinSurrounderUtils.getExpressionType(expression);
+        JetType type = ResolvePackage.analyze(expression, BodyResolveMode.PARTIAL).get(BindingContext.EXPRESSION_TYPE, expression);
         if (type == null || type.equals(KotlinBuiltIns.getInstance().getUnitType())) {
             return false;
         }

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/expression/KotlinNotSurrounder.java

@@ -24,10 +24,12 @@ import com.intellij.openapi.util.TextRange;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
-import org.jetbrains.kotlin.idea.codeInsight.surroundWith.KotlinSurrounderUtils;
+import org.jetbrains.kotlin.idea.caches.resolve.ResolvePackage;
 import org.jetbrains.kotlin.psi.JetExpression;
 import org.jetbrains.kotlin.psi.JetParenthesizedExpression;
 import org.jetbrains.kotlin.psi.JetPrefixExpression;
+import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.types.JetType;
 
 import static org.jetbrains.kotlin.psi.PsiPackage.JetPsiFactory;
@@ -40,7 +42,7 @@ public class KotlinNotSurrounder extends KotlinExpressionSurrounder {
 
     @Override
     public boolean isApplicable(@NotNull JetExpression expression) {
-        JetType type = KotlinSurrounderUtils.getExpressionType(expression);
+        JetType type = ResolvePackage.analyze(expression, BodyResolveMode.PARTIAL).get(BindingContext.EXPRESSION_TYPE, expression);
         return KotlinBuiltIns.getInstance().getBooleanType().equals(type);
     }
 

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/expression/KotlinWhenSurrounder.java

@@ -26,11 +26,13 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor;
 import org.jetbrains.kotlin.descriptors.ClassKind;
 import org.jetbrains.kotlin.descriptors.ClassifierDescriptor;
 import org.jetbrains.kotlin.idea.JetBundle;
-import org.jetbrains.kotlin.idea.codeInsight.surroundWith.KotlinSurrounderUtils;
+import org.jetbrains.kotlin.idea.caches.resolve.ResolvePackage;
 import org.jetbrains.kotlin.psi.JetExpression;
 import org.jetbrains.kotlin.psi.JetWhenCondition;
 import org.jetbrains.kotlin.psi.JetWhenEntry;
 import org.jetbrains.kotlin.psi.JetWhenExpression;
+import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.types.JetType;
 
 import static org.jetbrains.kotlin.psi.PsiPackage.JetPsiFactory;
@@ -68,7 +70,7 @@ public class KotlinWhenSurrounder extends KotlinExpressionSurrounder {
     }
 
     private String getCodeTemplate(JetExpression expression) {
-        JetType type = KotlinSurrounderUtils.getExpressionType(expression);
+        JetType type = ResolvePackage.analyze(expression, BodyResolveMode.PARTIAL).get(BindingContext.EXPRESSION_TYPE, expression);
         if (type != null) {
             ClassifierDescriptor descriptor = type.getConstructor().getDeclarationDescriptor();
             if (descriptor instanceof ClassDescriptor && ((ClassDescriptor) descriptor).getKind() == ClassKind.ENUM_CLASS) {