Get PsiMethod wrappers for properties and property accessors

compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/LightClassUtil.java

@@ -31,10 +31,13 @@ import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.stubs.PsiFileStub;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.SmartList;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.codegen.binding.PsiCodegenPredictor;
 import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.resolve.java.JetClsMethod;
+import org.jetbrains.jet.lang.resolve.java.JvmAbi;
 import org.jetbrains.jet.lang.resolve.java.JvmClassName;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
@@ -42,6 +45,7 @@ import org.jetbrains.jet.utils.KotlinVfsUtil;
 
 import java.net.MalformedURLException;
 import java.net.URL;
+import java.util.*;
 
 public class LightClassUtil {
     private static final Logger LOG = Logger.getInstance(LightClassUtil.class);
@@ -116,42 +120,70 @@ public class LightClassUtil {
         return LightClassGenerationSupport.getInstance(classOrObject.getProject()).getPsiClass(classOrObject);
     }
 
+    @Nullable
+    public static PsiMethod getLightClassAccessorMethod(@NotNull JetPropertyAccessor accessor) {
+        return getPsiMethodWrapper(accessor);
+    }
+
+    @NotNull
+    public static PropertyAccessorsPsiMethods getLightClassPropertyMethods(@NotNull JetProperty property) {
+        JetPropertyAccessor getter = property.getGetter();
+        JetPropertyAccessor setter = property.getSetter();
+
+        PsiMethod getterWrapper = getter != null ? getLightClassAccessorMethod(getter) : null;
+        PsiMethod setterWrapper = setter != null ? getLightClassAccessorMethod(setter) : null;
+
+        if (getterWrapper == null || setterWrapper == null) {
+            List<PsiMethod> wrappers = getPsiMethodWrappers(property, true);
+            assert wrappers.size() <= 2 : "Maximum two wrappers are expected to be generated for property";
+
+            for (PsiMethod wrapper : wrappers) {
+                if (wrapper.getName().startsWith(JvmAbi.SETTER_PREFIX)) {
+                    assert setterWrapper == null : String.format(
+                            "Setter accessor isn't expected to be reassigned (old: %s, new: %s)", setterWrapper, wrapper);
+
+                    setterWrapper = wrapper;
+                }
+                else {
+                    assert getterWrapper == null : String.format(
+                            "Getter accessor isn't expected to be reassigned (old: %s, new: %s)", getterWrapper, wrapper);
+
+                    getterWrapper = wrapper;
+                }
+            }
+        }
+
+        return new PropertyAccessorsPsiMethods(getterWrapper, setterWrapper);
+    }
+
     @Nullable
     public static PsiMethod getLightClassMethod(@NotNull JetNamedFunction function) {
-        //noinspection unchecked
+        return getPsiMethodWrapper(function);
-        if (PsiTreeUtil.getParentOfType(function, JetFunction.class, JetProperty.class) != null) {
+    }
-            // Don't genClassOrObject method wrappers for internal declarations. Their classes are not generated during calcStub
+
-            // with ClassBuilderMode.SIGNATURES mode, and this produces "Class not found exception" in getDelegate()
+    @Nullable
-            return null;
+    private static PsiMethod getPsiMethodWrapper(@NotNull JetDeclaration declaration) {
+        List<PsiMethod> wrappers = getPsiMethodWrappers(declaration, false);
+        return !wrappers.isEmpty() ? wrappers.get(0) : null;
+    }
+
+    @NotNull
+    private static List<PsiMethod> getPsiMethodWrappers(@NotNull JetDeclaration declaration, boolean collectAll) {
+        PsiClass psiClass = getWrappingClass(declaration);
+        if (psiClass == null) {
+            return Collections.emptyList();
         }
 
-        Project project = function.getProject();
+        List<PsiMethod> methods = new SmartList<PsiMethod>();
-
-        PsiElement parent = function.getParent();
-        PsiClass psiClass;
-        if (parent == function.getContainingFile()) {
-            // top-level function
-            JvmClassName jvmClassName = PsiCodegenPredictor.getPredefinedJvmClassName((JetFile) parent, true);
-            if (jvmClassName == null) return null;
-
-            String fqName = jvmClassName.getFqName().getFqName();
-            psiClass = JavaElementFinder.getInstance(project).findClass(fqName, GlobalSearchScope.allScope(project));
-        }
-        else {
-            if (!(parent instanceof JetClassBody)) return null;
-            assert parent.getParent() instanceof JetClassOrObject;
-
-            // function in a class
-            JetClassOrObject classOrObject = (JetClassOrObject) parent.getParent();
-            psiClass = getPsiClass(classOrObject);
-        }
-
-        if (psiClass == null) return null;
-
         for (PsiMethod method : psiClass.getMethods()) {
             try {
-                if (method instanceof PsiCompiledElement && ((PsiCompiledElement) method).getMirror() == function) {
+                if (method instanceof JetClsMethod) {
-                    return method;
+                    if (method instanceof PsiCompiledElement && ((PsiCompiledElement) method).getMirror() == declaration) {
+                        methods.add(method);
+                        if (!collectAll) {
+                            return methods;
+                        }
+                    }
                 }
             }
             catch (ProcessCanceledException e) {
@@ -159,24 +191,94 @@ public class LightClassUtil {
             }
             catch (Throwable e) {
                 throw new IllegalStateException(
-                        "Error while wrapping function " + function.getName() +
+                        "Error while wrapping declaration " + declaration.getName() +
                         "Context\n:" +
                         String.format("=== In file ===\n" +
-                                        "%s\n" +
+                                      "%s\n" +
-                                        "===On element===\n" +
+                                      "=== On element ===\n" +
-                                        "%s\n" +
+                                      "%s\n" +
-                                        "===WrappedElement===\n" +
+                                      "=== WrappedElement ===\n" +
-                                        "%s\n",
+                                      "%s\n",
-                                        function.getContainingFile().getText(),
+                                      declaration.getContainingFile().getText(),
-                                        function.getText(),
+                                      declaration.getText(),
-                                        method.toString()),
+                                      method.toString()),
                         e
                 );
             }
         }
 
+        return methods;
+    }
+
+    @Nullable
+    private static PsiClass getWrappingClass(@NotNull JetDeclaration declaration) {
+        if (declaration instanceof JetPropertyAccessor) {
+            PsiElement propertyParent = declaration.getParent();
+            assert propertyParent instanceof JetProperty : "JetProperty is expected to be parent of accessor";
+
+            declaration = (JetProperty) propertyParent;
+        }
+
+        //noinspection unchecked
+        if (PsiTreeUtil.getParentOfType(declaration, JetFunction.class, JetProperty.class) != null) {
+            // Can't get wrappers for internal declarations. Their classes are not generated during calcStub
+            // with ClassBuilderMode.SIGNATURES mode, and this produces "Class not found exception" in getDelegate()
+            return null;
+        }
+
+        PsiElement parent = declaration.getParent();
+
+        if (parent instanceof JetFile) {
+            // top-level declaration
+            JvmClassName jvmName = PsiCodegenPredictor.getPredefinedJvmClassName((JetFile) parent, true);
+            if (jvmName != null) {
+                Project project = declaration.getProject();
+
+                String fqName = jvmName.getFqName().getFqName();
+                return JavaElementFinder.getInstance(project).findClass(fqName, GlobalSearchScope.allScope(project));
+            }
+        }
+        else if (parent instanceof JetClassBody) {
+            assert parent.getParent() instanceof JetClassOrObject;
+            return getPsiClass((JetClassOrObject) parent.getParent());
+        }
+
         return null;
     }
 
+    public static class PropertyAccessorsPsiMethods implements Iterable<PsiMethod> {
+        private final PsiMethod getter;
+        private final PsiMethod setter;
+        private final Collection<PsiMethod> accessors = new ArrayList<PsiMethod>(2);
+
+        PropertyAccessorsPsiMethods(@Nullable PsiMethod getter, @Nullable PsiMethod setter) {
+            this.getter = getter;
+            if (getter != null) {
+                accessors.add(getter);
+            }
+
+            this.setter = setter;
+            if (setter != null) {
+                accessors.add(setter);
+            }
+        }
+
+        @Nullable
+        public PsiMethod getGetter() {
+            return getter;
+        }
+
+        @Nullable
+        public PsiMethod getSetter() {
+            return setter;
+        }
+
+        @NotNull
+        @Override
+        public Iterator<PsiMethod> iterator() {
+            return accessors.iterator();
+        }
+    }
+
     private LightClassUtil() {}
 }

idea/testData/javaFacade/wrapValTopLevelProperty.kt

@@ -0,0 +1 @@
+val <caret>test: Int? = null

idea/testData/javaFacade/wrapVarPropertyInClass.kt

@@ -0,0 +1,3 @@
+class Test {
+  var <caret>test = "Test"
+}

idea/testData/javaFacade/wrapVarPropertyInLocalClass.kt

@@ -0,0 +1,5 @@
+fun test() {
+  class Test {
+    var <caret>someProperty = "No wrapper is expected"
+  }
+}

idea/testData/javaFacade/wrapVarPropertyWithAccessorsInTrait.kt

@@ -0,0 +1,5 @@
+trait Hello {
+    var <caret>some: String
+        get() = "Hi"
+        set(value) {}
+}

idea/testData/javaFacade/wrapVarTopLevelAccessor.kt

@@ -0,0 +1,3 @@
+var some: String
+    get() = "Hi"
+    <caret>set(value) {}

idea/testData/javaFacade/wrapVarTopLevelProperty.kt

@@ -0,0 +1 @@
+var <caret>test: Int? = null

idea/tests/org/jetbrains/jet/plugin/javaFacade/JetJavaFacadeTest.java

@@ -24,8 +24,7 @@ import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.asJava.KotlinLightClass;
 import org.jetbrains.jet.asJava.LightClassUtil;
-import org.jetbrains.jet.lang.psi.JetClass;
+import org.jetbrains.jet.lang.psi.*;
-import org.jetbrains.jet.lang.psi.JetNamedFunction;
 import org.jetbrains.jet.lang.resolve.java.JvmAbi;
 import org.jetbrains.jet.plugin.JetLightProjectDescriptor;
 import org.jetbrains.jet.plugin.PluginTestCaseBase;
@@ -99,6 +98,30 @@ public class JetJavaFacadeTest extends LightCodeInsightFixtureTestCase {
         doTestWrapMethod(true);
     }
 
+    public void testWrapValTopLevelProperty() {
+        doTestWrapProperty(true, false);
+    }
+
+    public void testWrapVarPropertyInClass() {
+        doTestWrapProperty(true, true);
+    }
+
+    public void testWrapVarPropertyWithAccessorsInTrait() {
+        doTestWrapProperty(true, true);
+    }
+
+    public void testWrapVarTopLevelProperty() {
+        doTestWrapProperty(true, true);
+    }
+
+    public void testWrapVarPropertyInLocalClass() {
+        doTestWrapProperty(false, false);
+    }
+
+    public void testWrapVarTopLevelAccessor() {
+        doTestWrapPropertyAccessor(true);
+    }
+
     public void testEa38770() {
         myFixture.configureByFile(getTestName(true) + ".kt");
 
@@ -166,6 +189,37 @@ public class JetJavaFacadeTest extends LightCodeInsightFixtureTestCase {
     }
 
     private void doTestWrapMethod(boolean shouldBeWrapped) {
+        JetNamedFunction jetFunction = getPreparedElement(JetNamedFunction.class);
+
+        // Should not fail!
+        PsiMethod psiMethod = LightClassUtil.getLightClassMethod(jetFunction);
+
+        checkDeclarationMethodWrapped(shouldBeWrapped, jetFunction, psiMethod);
+    }
+
+    private void doTestWrapProperty(boolean shouldWrapGetter, boolean shouldWrapSetter) {
+        JetProperty jetProperty = getPreparedElement(JetProperty.class);
+
+        // Should not fail!
+        LightClassUtil.PropertyAccessorsPsiMethods propertyAccessors = LightClassUtil.getLightClassPropertyMethods(jetProperty);
+
+        JetPropertyAccessor getter = jetProperty.getGetter();
+        JetPropertyAccessor setter = jetProperty.getSetter();
+
+        checkDeclarationMethodWrapped(shouldWrapGetter, getter != null ? getter : jetProperty, propertyAccessors.getGetter());
+        checkDeclarationMethodWrapped(shouldWrapSetter, setter != null ? setter : jetProperty, propertyAccessors.getSetter());
+    }
+
+    private void doTestWrapPropertyAccessor(boolean shouldWrapAccessor) {
+        JetPropertyAccessor jetPropertyAccessor = getPreparedElement(JetPropertyAccessor.class);
+
+        // Should not fail!
+        PsiMethod propertyAccessors = LightClassUtil.getLightClassAccessorMethod(jetPropertyAccessor);
+        checkDeclarationMethodWrapped(shouldWrapAccessor, jetPropertyAccessor, propertyAccessors);
+    }
+
+    @NotNull
+    private <T extends JetElement> T getPreparedElement(Class<T> elementClass) {
         myFixture.configureByFile(getTestName(true) + ".kt");
 
         int offset = myFixture.getEditor().getCaretModel().getOffset();
@@ -173,16 +227,20 @@ public class JetJavaFacadeTest extends LightCodeInsightFixtureTestCase {
 
         assertNotNull("Caret should be set for tested file", elementAt);
 
-        JetNamedFunction jetFunction = PsiTreeUtil.getParentOfType(elementAt, JetNamedFunction.class);
+        T caretElement = PsiTreeUtil.getParentOfType(elementAt, elementClass);
-        assertNotNull("Caret should be placed to function definition", jetFunction);
+        assertNotNull(
+                String.format("Caret should be placed to element of type: %s, but was at element '%s' of type %s",
+                              elementClass, elementAt, elementAt.getClass()),
+                caretElement);
 
-        // Should not fail!
+        return caretElement;
-        PsiMethod psiMethod = LightClassUtil.getLightClassMethod(jetFunction);
+    }
 
+    private static void checkDeclarationMethodWrapped(boolean shouldBeWrapped, JetDeclaration declaration, PsiMethod psiMethod) {
         if (shouldBeWrapped) {
-            assertNotNull(String.format("Failed to wrap jetFunction '%s' to method", jetFunction.getText()), psiMethod);
+            assertNotNull(String.format("Failed to wrap declaration '%s' to method", declaration.getText()), psiMethod);
             assertInstanceOf(psiMethod, PsiCompiledElement.class);
-            assertEquals("Invalid original element for generated method", ((PsiCompiledElement) psiMethod).getMirror(), jetFunction);
+            assertEquals("Invalid original element for generated method", ((PsiCompiledElement) psiMethod).getMirror(), declaration);
         }
         else {
             assertNull("There should be no wrapper for given method", psiMethod);