Refactoring: improve isApplicable methods in surrounders

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

@@ -21,8 +21,13 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.idea.KotlinBundle;
+import org.jetbrains.kotlin.idea.caches.resolve.ResolutionUtils;
 import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils;
 import org.jetbrains.kotlin.psi.KtBlockExpression;
+import org.jetbrains.kotlin.psi.KtExpression;
+import org.jetbrains.kotlin.resolve.BindingContext;
+import org.jetbrains.kotlin.resolve.bindingContextUtil.BindingContextUtilsKt;
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 
 public class KotlinSurrounderUtils {
     public static String SURROUND_WITH = KotlinBundle.message("surround.with");
@@ -42,4 +47,14 @@ public class KotlinSurrounderUtils {
     public static void showErrorHint(@NotNull Project project, @NotNull Editor editor, @NotNull String message) {
         CodeInsightUtils.showErrorHint(project, editor, message, SURROUND_WITH, null);
     }
+
+    public static boolean isUsedAsStatement(@NotNull KtExpression expression) {
+        BindingContext context = ResolutionUtils.analyze(expression, BodyResolveMode.PARTIAL);
+        return BindingContextUtilsKt.isUsedAsStatement(expression, context);
+    }
+
+    public static boolean isUsedAsExpression(@NotNull KtExpression expression) {
+        BindingContext context = ResolutionUtils.analyze(expression, BodyResolveMode.PARTIAL);
+        return BindingContextUtilsKt.isUsedAsExpression(expression, context);
+    }
 }

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/expression/KotlinControlFlowExpressionSurrounderBase.kt

@@ -26,8 +26,6 @@ import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.createExpressionByPattern
 
 abstract class KotlinControlFlowExpressionSurrounderBase : KotlinExpressionSurrounder() {
-
-    override fun isApplicable(expression: KtExpression) = true
     override fun isApplicableToStatements() = false
 
     override fun surroundExpression(project: Project, editor: Editor, expression: KtExpression): TextRange? {

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

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 import org.jetbrains.kotlin.types.KotlinType;
 
 import static org.jetbrains.kotlin.builtins.KotlinBuiltIns.isUnit;
+import static org.jetbrains.kotlin.idea.codeInsight.surroundWith.KotlinSurrounderUtils.isUsedAsStatement;
 
 public abstract class KotlinExpressionSurrounder implements Surrounder {
 
@@ -45,26 +46,28 @@ public abstract class KotlinExpressionSurrounder implements Surrounder {
         if (expression instanceof KtCallExpression && expression.getParent() instanceof KtQualifiedExpression) {
             return false;
         }
+
+        return isApplicable(expression);
+    }
+
+    protected boolean isApplicable(@NotNull KtExpression expression) {
         BindingContext context = ResolutionUtils.analyze(expression, BodyResolveMode.PARTIAL);
         KotlinType type = context.getType(expression);
         if (type == null || (isUnit(type) && isApplicableToStatements())) {
             return false;
         }
 
-        boolean usedAsExpression = Boolean.TRUE.equals(context.get(BindingContext.USED_AS_EXPRESSION, expression));
-        if (!isApplicableToStatements() && !usedAsExpression) {
+        if (!isApplicableToStatements() && isUsedAsStatement(expression)) {
             return false;
         }
 
-        return isApplicable(expression);
+        return true;
     }
 
     protected boolean isApplicableToStatements() {
         return true;
     }
 
-    protected abstract boolean isApplicable(@NotNull KtExpression expression);
-
     @Nullable
     @Override
     public TextRange surroundElements(@NotNull Project project, @NotNull Editor editor, @NotNull PsiElement[] elements) {

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

@@ -40,6 +40,7 @@ public class KotlinNotSurrounder extends KotlinExpressionSurrounder {
 
     @Override
     public boolean isApplicable(@NotNull KtExpression expression) {
+        if (!super.isApplicable(expression)) return false;
         KotlinType type = ResolutionUtils.analyze(expression, BodyResolveMode.PARTIAL).getType(expression);
         return type != null && KotlinBuiltIns.isBoolean(type);
     }

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

@@ -33,11 +33,6 @@ public class KotlinParenthesesSurrounder extends KotlinExpressionSurrounder {
         return CodeInsightBundle.message("surround.with.parenthesis.template");
     }
 
-    @Override
-    public boolean isApplicable(@NotNull KtExpression expression) {
-        return true;
-    }
-
     @Nullable
     @Override
     public TextRange surroundExpression( @NotNull Project project, @NotNull Editor editor, @NotNull KtExpression expression) {

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/expression/KotlinRuntimeTypeCastSurrounder.kt

@@ -31,8 +31,8 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.idea.debugger.evaluate.KotlinRuntimeTypeEvaluator
 import org.jetbrains.kotlin.idea.core.ShortenReferences
+import org.jetbrains.kotlin.idea.debugger.evaluate.KotlinRuntimeTypeEvaluator
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
@@ -43,6 +43,8 @@ import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 class KotlinRuntimeTypeCastSurrounder: KotlinExpressionSurrounder() {
 
     override fun isApplicable(expression: KtExpression): Boolean {
+        if (!super.isApplicable(expression)) return false
+
         if (!expression.isPhysical) return false
         val file = expression.containingFile
         if (file !is KtCodeFragment) return false

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

@@ -33,7 +33,7 @@ public class KotlinStringTemplateSurrounder extends KotlinExpressionSurrounder {
 
     @Override
     public boolean isApplicable(@NotNull KtExpression expression) {
-        return !(expression instanceof KtStringTemplateExpression);
+        return !(expression instanceof KtStringTemplateExpression) && super.isApplicable(expression);
     }
 
     @Nullable

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

@@ -37,11 +37,6 @@ public class KotlinWhenSurrounder extends KotlinExpressionSurrounder {
         return KotlinBundle.message("surround.with.when.template");
     }
 
-    @Override
-    public boolean isApplicable(@NotNull KtExpression expression) {
-        return true;
-    }
-
     @Nullable
     @Override
     public TextRange surroundExpression(@NotNull Project project, @NotNull Editor editor, @NotNull KtExpression expression) {

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/expression/KotlinWithIfExpressionSurrounder.kt

@@ -32,7 +32,7 @@ import org.jetbrains.kotlin.utils.sure
 
 class KotlinWithIfExpressionSurrounder(val withElse: Boolean) : KotlinExpressionSurrounder() {
     override fun isApplicable(expression: KtExpression) =
-            expression.analyze(BodyResolveMode.PARTIAL).getType(expression)?.isBoolean() ?: false
+            super.isApplicable(expression) && (expression.analyze(BodyResolveMode.PARTIAL).getType(expression)?.isBoolean() ?: false)
 
     override fun surroundExpression(project: Project, editor: Editor, expression: KtExpression): TextRange? {
         val factory = KtPsiFactory(project)

idea/src/org/jetbrains/kotlin/idea/codeInsight/surroundWith/statement/KotlinStatementsSurrounder.java

@@ -24,11 +24,8 @@ import com.intellij.psi.PsiElement;
 import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.jetbrains.kotlin.idea.caches.resolve.ResolutionUtils;
-import org.jetbrains.kotlin.psi.KtElement;
+import org.jetbrains.kotlin.idea.codeInsight.surroundWith.KotlinSurrounderUtils;
 import org.jetbrains.kotlin.psi.KtExpression;
-import org.jetbrains.kotlin.resolve.BindingContext;
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode;
 
 public abstract class KotlinStatementsSurrounder implements Surrounder {
 
@@ -38,13 +35,8 @@ public abstract class KotlinStatementsSurrounder implements Surrounder {
             return false;
         }
 
-        if (elements.length == 1 || !(elements[0] instanceof KtExpression) && !isApplicableWhenUsedAsExpression()) {
-            KtElement expression = (KtElement) elements[0];
-
-            Boolean usedAsExpression = ResolutionUtils.analyze(expression, BodyResolveMode.PARTIAL)
-                    .get(BindingContext.USED_AS_EXPRESSION, expression);
-
-            if (usedAsExpression != null && usedAsExpression) {
+        if (elements.length == 1 || elements[0] instanceof KtExpression) {
+            if (!isApplicableWhenUsedAsExpression() && KotlinSurrounderUtils.isUsedAsExpression((KtExpression) elements[0])) {
                 return false;
             }
         }
@@ -61,8 +53,7 @@ public abstract class KotlinStatementsSurrounder implements Surrounder {
     public TextRange surroundElements(
             @NotNull Project project,
             @NotNull Editor editor,
-            @NotNull PsiElement[] elements
-    ) throws IncorrectOperationException {
+            @NotNull PsiElement[] elements) throws IncorrectOperationException {
         PsiElement container = elements[0].getParent();
         if (container == null) return null;
         return surroundStatements(project, editor, container, elements);