Add special class for getting methods from traits

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 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.
@@ -32,8 +32,11 @@ import org.jetbrains.jet.lang.psi.JetClassOrObject
 import com.intellij.psi.impl.light.LightTypeParameterListBuilder
 import com.intellij.psi.search.SearchScope
 
-public class KotlinLightMethodForDeclaration(
-        manager: PsiManager, override val delegate: PsiMethod, override val origin: JetDeclaration, containingClass: PsiClass
+open public class KotlinLightMethodForDeclaration(
+        manager: PsiManager,
+        override val delegate: PsiMethod,
+        override val origin: JetDeclaration,
+        containingClass: PsiClass
 ): LightMethod(manager, delegate, containingClass), KotlinLightMethod {
 
     private val paramsList: CachedValue<PsiParameterList> by Delegates.blockingLazy {

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2014 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.PsiManager
+import com.intellij.psi.PsiMethod
+import org.jetbrains.jet.lang.psi.JetDeclaration
+import com.intellij.psi.PsiClass
+import org.jetbrains.jet.asJava.KotlinLightMethodFromTrait
+import com.intellij.psi.PsiElement
+
+
+public class KotlinLightMethodFromTrait(manager: PsiManager,
+                                            override val delegate: PsiMethod,
+                                            override val origin: JetDeclaration,
+                                            containingClass: PsiClass) :
+        KotlinLightMethodForDeclaration(manager, delegate, origin, containingClass) {
+
+    override fun copy(): PsiElement {
+        return KotlinLightMethodFromTrait(getManager()!!, delegate, origin.copy() as JetDeclaration, getContainingClass()!!)
+    }
+}

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

@@ -24,9 +24,12 @@ import com.intellij.psi.impl.light.LightField;
 import com.intellij.psi.impl.light.LightMethod;
 import com.intellij.psi.impl.source.ClassInnerStuffCache;
 import com.intellij.psi.impl.source.PsiExtensibleClass;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.scope.PsiScopeProcessor;
 import com.intellij.util.Function;
 import com.intellij.util.containers.ContainerUtil;
+import kotlin.Function1;
+import kotlin.KotlinPackage;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.*;
 
@@ -125,13 +128,18 @@ public abstract class KotlinWrappingLightClass extends AbstractLightClass implem
     @NotNull
     @Override
     public List<PsiMethod> getOwnMethods() {
-        return ContainerUtil.map(getDelegate().getMethods(), new Function<PsiMethod, PsiMethod>() {
+        return KotlinPackage.map(getDelegate().getMethods(), new Function1<PsiMethod, PsiMethod>() {
             @Override
-            public PsiMethod fun(PsiMethod method) {
+            public PsiMethod invoke(PsiMethod method) {
                 JetDeclaration declaration = ClsWrapperStubPsiFactory.getOriginalDeclaration(method);
-                return declaration != null
-                       ? new KotlinLightMethodForDeclaration(myManager, method, declaration, KotlinWrappingLightClass.this)
-                       : new LightMethod(myManager, method, KotlinWrappingLightClass.this);
+
+                if (declaration != null) {
+                    return !isMethodFromTrait(declaration) ?
+                           new KotlinLightMethodForDeclaration(myManager, method, declaration, KotlinWrappingLightClass.this) :
+                           new KotlinLightMethodFromTrait(myManager, method, declaration, KotlinWrappingLightClass.this);
+                }
+
+                return new LightMethod(myManager, method, KotlinWrappingLightClass.this);
             }
         });
     }
@@ -152,4 +160,18 @@ public abstract class KotlinWrappingLightClass extends AbstractLightClass implem
         JetClassOrObject origin = getOrigin();
         return origin == null ? null : origin.getText();
     }
+
+    private boolean isMethodFromTrait(@NotNull JetDeclaration originMethodDeclaration) {
+        if (!(originMethodDeclaration instanceof JetNamedFunction ||
+              originMethodDeclaration instanceof JetPropertyAccessor ||
+              originMethodDeclaration instanceof JetProperty)) {
+            return false;
+        }
+
+        JetClassOrObject parentOfMethodOrigin = PsiTreeUtil.getParentOfType(originMethodDeclaration, JetClassOrObject.class);
+        JetClassOrObject thisClassDeclaration = getOrigin();
+
+        // Method was generated from declaration in some other trait
+        return (parentOfMethodOrigin != null && thisClassDeclaration != parentOfMethodOrigin && JetPsiUtil.isTrait(parentOfMethodOrigin));
+    }
 }