From b039c69952688798b7f11710e38c60c8c4762bfb Mon Sep 17 00:00:00 2001
From: Andrey Breslav <andrey.breslav@jetbrains.com>
Date: Thu, 10 Jan 2013 19:00:38 +0400
Subject: [PATCH] Class object members supported in PsiCodegenPredictor

---
 .../codegen/binding/PsiCodegenPredictor.java  |  5 ++++
 .../jetbrains/jet/asJava/LightClassUtil.java  | 24 +++++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/binding/PsiCodegenPredictor.java b/compiler/backend/src/org/jetbrains/jet/codegen/binding/PsiCodegenPredictor.java
index 8b6ccbbe8bd..dee7cc18a3d 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/binding/PsiCodegenPredictor.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/binding/PsiCodegenPredictor.java
@@ -79,6 +79,11 @@ public final class PsiCodegenPredictor {
         // TODO: Method won't give correct class name for traits implementations
 
         JetDeclaration parentDeclaration = PsiTreeUtil.getParentOfType(declaration, JetDeclaration.class);
+        if (parentDeclaration instanceof JetClassObject) {
+            assert declaration instanceof JetObjectDeclaration : "Only object declarations can be children of JetClassObject: " + declaration;
+            return getPredefinedJvmClassName(parentDeclaration);
+        }
+
         JvmClassName parentClassName = parentDeclaration != null ?
                                        getPredefinedJvmClassName(parentDeclaration) :
                                        getPredefinedJvmClassName((JetFile) declaration.getContainingFile(), false);
diff --git a/compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/LightClassUtil.java b/compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/LightClassUtil.java
index e0a0e95cced..eaaf244070d 100644
--- a/compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/LightClassUtil.java
+++ b/compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/LightClassUtil.java
@@ -24,6 +24,7 @@ import com.intellij.openapi.util.SystemInfo;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiCompiledElement;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiMethod;
 import com.intellij.psi.impl.java.stubs.PsiClassStub;
 import com.intellij.psi.search.GlobalSearchScope;
@@ -129,11 +130,26 @@ public class LightClassUtil {
         }
 
         Project project = function.getProject();
-        JvmClassName jvmClassName = PsiCodegenPredictor.getPredefinedJvmClassNameForFun(function);
-        if (jvmClassName == null) return null;
 
-        String fqName = jvmClassName.getFqName().getFqName();
-        PsiClass psiClass = JavaElementFinder.getInstance(project).findClass(fqName, GlobalSearchScope.allScope(project));
+        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 = createLightClass(classOrObject);
+        }
+
         if (psiClass == null) return null;
 
         for (PsiMethod method : psiClass.getMethods()) {