Wrap cls-methods with kotlin prototypes

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

@@ -17,23 +17,41 @@
 package org.jetbrains.jet.asJava;
 
 import com.intellij.openapi.util.Key;
+import com.intellij.openapi.util.UserDataHolder;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.compiled.ClsClassImpl;
 import com.intellij.psi.impl.compiled.ClsEnumConstantImpl;
 import com.intellij.psi.impl.compiled.ClsFieldImpl;
+import com.intellij.psi.impl.compiled.ClsRepositoryPsiElement;
 import com.intellij.psi.impl.java.stubs.*;
 import com.intellij.psi.stubs.StubBase;
+import com.intellij.psi.stubs.StubElement;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.psi.JetDeclaration;
 
 class ClsWrapperStubPsiFactory extends StubPsiFactory {
     public static final Key<PsiElement> ORIGIN_ELEMENT = Key.create("ORIGIN_ELEMENT");
     private final StubPsiFactory delegate = new ClsStubPsiFactory();
-    
+
+    public static JetDeclaration getOriginalDeclaration(PsiMember member) {
+        if (member instanceof ClsRepositoryPsiElement<?>) {
+            StubElement stubElement = ((ClsRepositoryPsiElement<?>) member).getStub();
+            if (stubElement instanceof UserDataHolder) {
+                PsiElement original = ((UserDataHolder) stubElement).getUserData(ORIGIN_ELEMENT);
+                if (original instanceof JetDeclaration) {
+                    return (JetDeclaration) original;
+                }
+            }
+        }
+
+        return null;
+    }
+
     @Override
     public PsiClass createClass(PsiClassStub stub) {
         final PsiElement origin = ((StubBase) stub).getUserData(ORIGIN_ELEMENT);
         if (origin == null) return delegate.createClass(stub);
-        
+
         return new ClsClassImpl(stub) {
             @NotNull
             @Override
@@ -94,12 +112,7 @@ class ClsWrapperStubPsiFactory extends StubPsiFactory {
 
     @Override
     public PsiMethod createMethod(PsiMethodStub stub) {
-        PsiElement origin = ((StubBase) stub).getUserData(ORIGIN_ELEMENT);
+        return delegate.createMethod(stub);
-        if (origin == null) {
-            return delegate.createMethod(stub);
-        }
-
-        return new JetClsMethodImpl(stub, origin);
     }
 
     @Override

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

@@ -1,69 +0,0 @@
-/*
- * Copyright 2010-2013 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.jetbrains.jet.asJava;
-
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.impl.compiled.ClsMethodImpl;
-import com.intellij.psi.impl.java.stubs.PsiMethodStub;
-import com.intellij.psi.search.SearchScope;
-import com.intellij.util.IncorrectOperationException;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.jet.lang.psi.JetDeclaration;
-import org.jetbrains.jet.lang.resolve.java.jetAsJava.JetClsMethod;
-
-public class JetClsMethodImpl extends ClsMethodImpl implements JetClsMethod {
-    @NotNull
-    private final PsiElement origin;
-
-    public JetClsMethodImpl(@NotNull PsiMethodStub stub, @NotNull PsiElement origin) {
-        super(stub);
-        this.origin = origin;
-    }
-
-    @Override
-    public PsiElement getMirror() {
-        return origin;
-    }
-
-    @NotNull
-    @Override
-    public PsiElement getNavigationElement() {
-        return origin;
-    }
-
-    @Override
-    public JetDeclaration getOrigin() {
-        return (JetDeclaration) origin;
-    }
-
-    @Override
-    public void delete() throws IncorrectOperationException {
-        origin.delete();
-    }
-
-    @Override
-    public boolean isEquivalentTo(PsiElement another) {
-        if (another instanceof JetClsMethod && getOrigin().equals(((JetClsMethod) another).getOrigin())) return true;
-        return super.isEquivalentTo(another);
-    }
-
-    @NotNull
-    @Override
-    public SearchScope getUseScope() {
-        return origin.getUseScope();
-    }
-}

compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/KotlinLightMethodForDeclaration.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.asJava
+
+import com.intellij.psi.impl.light.LightMethod
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiManager
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiNamedElement
+import org.jetbrains.jet.lang.psi.JetDeclaration
+import org.jetbrains.jet.lang.resolve.java.jetAsJava.KotlinLightMethod
+import com.intellij.psi.PsiParameterList
+import org.jetbrains.jet.plugin.JetLanguage
+import kotlin.properties.Delegates
+
+public class KotlinLightMethodForDeclaration(manager: PsiManager, val method: PsiMethod, val jetDeclaration: JetDeclaration, containingClass: PsiClass):
+        LightMethod(manager, method, containingClass), KotlinLightMethod {
+
+    override fun getNavigationElement() : PsiElement = jetDeclaration
+    override fun getOriginalElement() : PsiElement = jetDeclaration
+    override fun getOrigin(): JetDeclaration? = jetDeclaration
+
+    override fun getParent(): PsiElement? = getContainingClass()
+
+    override fun setName(name: String): PsiElement? {
+        (jetDeclaration as PsiNamedElement).setName(name)
+        return this
+    }
+
+    public override fun delete() {
+        if (jetDeclaration.isValid()) {
+            jetDeclaration.delete()
+        }
+    }
+
+    override fun isEquivalentTo(another: PsiElement?): Boolean {
+        if (another is KotlinLightMethod && getOrigin() == another.getOrigin()) {
+            return true
+        }
+
+        return super<LightMethod>.isEquivalentTo(another)
+    }
+
+    override fun copy(): PsiElement? {
+        return KotlinLightMethodForDeclaration(getManager()!!, method, jetDeclaration.copy() as JetDeclaration, getContainingClass()!!)
+    }
+}

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

@@ -30,6 +30,7 @@ import com.intellij.psi.impl.source.PsiExtensibleClass;
 import com.intellij.util.Function;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.psi.JetDeclaration;
 
 import java.util.List;
 
@@ -115,7 +116,10 @@ public abstract class KotlinWrappingLightClass extends AbstractLightClass implem
         return ContainerUtil.map(getDelegate().getMethods(), new Function<PsiMethod, PsiMethod>() {
             @Override
             public PsiMethod fun(PsiMethod method) {
-                return new LightMethod(myManager, method, KotlinWrappingLightClass.this);
+                JetDeclaration declaration = ClsWrapperStubPsiFactory.getOriginalDeclaration(method);
+                return declaration != null
+                       ? new KotlinLightMethodForDeclaration(myManager, method, declaration, KotlinWrappingLightClass.this)
+                       : new LightMethod(myManager, method, KotlinWrappingLightClass.this);
             }
         });
     }

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

@@ -25,6 +25,7 @@ import org.jetbrains.jet.asJava.KotlinLightClass;
 import org.jetbrains.jet.asJava.LightClassUtil;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.java.JvmAbi;
+import org.jetbrains.jet.lang.resolve.java.jetAsJava.KotlinLightMethod;
 import org.jetbrains.jet.plugin.JetLightCodeInsightFixtureTestCase;
 import org.jetbrains.jet.plugin.JetLightProjectDescriptor;
 import org.jetbrains.jet.plugin.PluginTestCaseBase;
@@ -265,8 +266,8 @@ public class JetJavaFacadeTest extends JetLightCodeInsightFixtureTestCase {
     private static void checkDeclarationMethodWrapped(boolean shouldBeWrapped, JetDeclaration declaration, PsiMethod psiMethod) {
         if (shouldBeWrapped) {
             assertNotNull(String.format("Failed to wrap declaration '%s' to method", declaration.getText()), psiMethod);
-            assertInstanceOf(psiMethod, PsiCompiledElement.class);
+            assertInstanceOf(psiMethod, KotlinLightMethod.class);
-            assertEquals("Invalid original element for generated method", ((PsiCompiledElement) psiMethod).getMirror(), declaration);
+            assertEquals("Invalid original element for generated method", ((KotlinLightMethod) psiMethod).getOrigin(), declaration);
         }
         else {
             assertNull("There should be no wrapper for given method", psiMethod);