Inheritors in smart completion work for the classes in the same file too

idea/src/org/jetbrains/jet/plugin/completion/CompletionSession.kt

@@ -87,8 +87,11 @@ abstract class CompletionSessionBase(protected val configuration: CompletionSess
 
     protected val project: Project = position.getProject()
 
+
+    protected val originalSearchScope: GlobalSearchScope = parameters.getOriginalFile().getResolveScope()
+
     // we need to exclude the original file from scope because our resolve session is built with this file replaced by synthetic one
-    protected val searchScope: GlobalSearchScope = object : DelegatingGlobalSearchScope(parameters.getOriginalFile().getResolveScope()) {
+    protected val searchScope: GlobalSearchScope = object : DelegatingGlobalSearchScope(originalSearchScope) {
         override fun contains(file: VirtualFile) = super.contains(file) && file != parameters.getOriginalFile().getVirtualFile()
     }
 
@@ -226,9 +229,10 @@ class SmartCompletionSession(configuration: CompletionSessionConfiguration, para
 
     override fun doComplete() {
         if (jetReference != null) {
+            val converter = ToFromOriginalFileConverter(parameters.getOriginalFile() as JetFile, position.getContainingFile() as JetFile, parameters.getOffset())
             val completion = SmartCompletion(jetReference.expression, resolutionFacade, moduleDescriptor, 
-                                             bindingContext!!, { isVisibleDescriptor(it) }, searchScope,
-                                             parameters.getOriginalFile() as JetFile, boldImmediateLookupElementFactory)
+                                             bindingContext!!, { isVisibleDescriptor(it) }, originalSearchScope,
+                                             converter, boldImmediateLookupElementFactory)
             val result = completion.execute()
             if (result != null) {
                 collector.addElements(result.additionalItems)

idea/src/org/jetbrains/jet/plugin/completion/ToFromOriginalFileConverter.kt

@@ -0,0 +1,70 @@
+/*
+ * 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.plugin.completion
+
+import org.jetbrains.jet.lang.psi.JetFile
+import org.jetbrains.jet.lang.psi.JetDeclaration
+import com.intellij.psi.util.PsiTreeUtil
+
+class ToFromOriginalFileConverter(
+        val originalFile: JetFile,
+        val syntheticFile: JetFile,
+        val completionOffset: Int
+) {
+    private val syntheticLength: Int
+    private val originalLength: Int
+    private val tailLength: Int
+    private val shift: Int
+
+    //TODO: lazy initialization?
+    ;{
+        val originalText = originalFile.getText()
+        val syntheticText = syntheticFile.getText()
+        assert(originalText.subSequence(0, completionOffset) == syntheticText.subSequence(0, completionOffset)) //TODO: drop it
+
+        syntheticLength = syntheticText.length
+        originalLength = originalText.length
+        tailLength = originalText.indices.first { syntheticText[syntheticLength - it - 1] != originalText[originalLength - it - 1] }
+        shift = syntheticLength - originalLength
+    }
+
+    public fun toOriginalFile(offset: Int): Int? {
+        return when {
+            offset <= completionOffset -> offset
+            offset >= syntheticLength - tailLength -> offset - shift
+            else -> null
+        }
+    }
+
+    public fun toSyntheticFile(offset: Int): Int? {
+        return when {
+            offset <= completionOffset -> offset
+            offset >= originalLength - tailLength -> offset + shift
+            else -> null
+        }
+    }
+
+    public fun toOriginalFile(declaration: JetDeclaration): JetDeclaration? {
+        val offset = toOriginalFile(declaration.getTextRange().getStartOffset()) ?: return null
+        return PsiTreeUtil.findElementOfClassAtOffset(originalFile, offset, javaClass<JetDeclaration>(), true)
+    }
+
+    public fun toSyntheticFile(declaration: JetDeclaration): JetDeclaration? {
+        val offset = toSyntheticFile(declaration.getTextRange().getStartOffset()) ?: return null
+        return PsiTreeUtil.findElementOfClassAtOffset(syntheticFile, offset, javaClass<JetDeclaration>(), true)
+    }
+}

idea/src/org/jetbrains/jet/plugin/completion/smart/SmartCompletion.kt

@@ -45,8 +45,8 @@ class SmartCompletion(val expression: JetSimpleNameExpression,
                       val moduleDescriptor: ModuleDescriptor,
                       val bindingContext: BindingContext,
                       val visibilityFilter: (DeclarationDescriptor) -> Boolean,
-                      val searchScope: GlobalSearchScope,
-                      val originalFile: JetFile,
+                      val inheritorSearchScope: GlobalSearchScope,
+                      val toFromOriginalFileConverter: ToFromOriginalFileConverter,
                       val boldImmediateLookupElementFactory: LookupElementFactory) {
     private val project = expression.getProject()
 
@@ -145,7 +145,7 @@ class SmartCompletion(val expression: JetSimpleNameExpression,
         val additionalItems = ArrayList<LookupElement>()
         val inheritanceSearchers = ArrayList<InheritanceItemsSearcher>()
         if (receiver == null) {
-            TypeInstantiationItems(resolutionFacade, moduleDescriptor, bindingContext, visibilityFilter, searchScope)
+            TypeInstantiationItems(resolutionFacade, moduleDescriptor, bindingContext, visibilityFilter, toFromOriginalFileConverter, inheritorSearchScope)
                     .add(additionalItems, inheritanceSearchers, expectedInfos)
 
             StaticMembers(bindingContext, resolutionFacade).addToCollection(additionalItems, expectedInfos, expression, itemsToSkip)
@@ -174,8 +174,7 @@ class SmartCompletion(val expression: JetSimpleNameExpression,
         // if our expression is initializer of implicitly typed variable - take type of variable from original file (+ the same for function)
         val declaration = implicitlyTypedDeclarationFromInitializer(expression)
         if (declaration != null) {
-            val offset = declaration.getTextRange()!!.getStartOffset()
-            val originalDeclaration = PsiTreeUtil.findElementOfClassAtOffset(originalFile, offset, javaClass<JetDeclaration>(), true)
+            val originalDeclaration = toFromOriginalFileConverter.toOriginalFile(declaration)
             if (originalDeclaration != null) {
                 val originalDescriptor = originalDeclaration.resolveToDescriptor() as? CallableDescriptor
                 val returnType = originalDescriptor?.getReturnType()

idea/src/org/jetbrains/jet/plugin/completion/smart/StaticMembers.kt

@@ -77,11 +77,6 @@ class StaticMembers(val bindingContext: BindingContext, val resolutionFacade: Re
             else if (DescriptorUtils.isEnumEntry(descriptor) && !enumEntriesToSkip.contains(descriptor)) {
                 classifier = { ExpectedInfoClassification.MATCHES } /* we do not need to check type of enum entry because it's taken from proper enum */
             }
-            else if (descriptor is ClassDescriptor && DescriptorUtils.isObject(descriptor)) {
-                classifier = { expectedInfo ->
-                    if (descriptor.getDefaultType().isSubtypeOf(expectedInfo.type)) ExpectedInfoClassification.MATCHES else ExpectedInfoClassification.NOT_MATCHES
-                }
-            }
             else {
                 return
             }

idea/src/org/jetbrains/jet/plugin/completion/smart/TypeInstantiationItems.kt

@@ -59,13 +59,15 @@ import org.jetbrains.jet.lang.descriptors.ModuleDescriptor
 import org.jetbrains.jet.lang.resolve.resolveTopLevelClass
 import org.jetbrains.jet.lang.types.TypeProjectionImpl
 import org.jetbrains.jet.lang.types.Variance
+import org.jetbrains.jet.lang.psi.JetDeclaration
 
 class TypeInstantiationItems(
         val resolutionFacade: ResolutionFacade,
         val moduleDescriptor: ModuleDescriptor,
         val bindingContext: BindingContext,
         val visibilityFilter: (DeclarationDescriptor) -> Boolean,
-        val searchScope: GlobalSearchScope
+        val toFromOriginalFileConverter: ToFromOriginalFileConverter,
+        val inheritorSearchScope: GlobalSearchScope
 ) {
     public fun add(
             items: MutableCollection<LookupElement>,
@@ -109,13 +111,17 @@ class TypeInstantiationItems(
     private fun MutableCollection<InheritanceItemsSearcher>.addInheritorSearcher(
             descriptor: ClassDescriptor, kotlinClassDescriptor: ClassDescriptor, typeArgs: List<TypeProjection>, tail: Tail?
     ) {
-        val declaration = DescriptorToSourceUtils.descriptorToDeclaration(descriptor)
+        val _declaration = DescriptorToSourceUtils.descriptorToDeclaration(descriptor) ?: return
+        val declaration = if (_declaration.getContainingFile() == toFromOriginalFileConverter.syntheticFile)
+            toFromOriginalFileConverter.toOriginalFile(_declaration as JetDeclaration) ?: return
+        else
+            _declaration
+
         val psiClass: PsiClass = when (declaration) {
             is PsiClass -> declaration
             is JetClassOrObject -> LightClassUtil.getPsiClass(declaration) ?: return
             else -> return
         }
-        if (declaration.getContainingFile().getVirtualFile() == null) return //TODO!
         add(InheritanceSearcher(psiClass, kotlinClassDescriptor, typeArgs, tail))
     }
 
@@ -269,11 +275,15 @@ class TypeInstantiationItems(
         private val expectedType = JetTypeImpl(Annotations.EMPTY, typeConstructor, false, typeArgs, classDescriptor.getMemberScope(typeArgs))
 
         override fun search(nameFilter: (String) -> Boolean, consumer: (LookupElement) -> Unit) {
-            val parameters = ClassInheritorsSearch.SearchParameters(psiClass, searchScope, true, true, false, nameFilter)
+            val parameters = ClassInheritorsSearch.SearchParameters(psiClass, inheritorSearchScope, true, true, false, nameFilter)
             for (inheritor in ClassInheritorsSearch.search(parameters)) {
                 val descriptor = if (inheritor is KotlinLightClass) {
                     val origin = inheritor.origin ?: continue
-                    resolutionFacade.resolveToDescriptor(origin)
+                    val declaration = if (origin.getContainingFile() == toFromOriginalFileConverter.originalFile)
+                        toFromOriginalFileConverter.toSyntheticFile(origin) ?: continue
+                    else
+                        origin
+                    resolutionFacade.resolveToDescriptor(declaration)
                 }
                 else {
                     resolutionFacade.get(JavaResolveExtension)(inheritor).first.resolveClass(JavaClassImpl(inheritor))

idea/testData/completion/smart/Inheritors3.kt

@@ -0,0 +1,31 @@
+trait T
+
+open class A : T
+
+class B : A()
+
+fun foo(): T {
+    open class Local1 : T
+
+    class Local2 : Local1()
+
+    return <caret>
+}
+
+fun x() {
+    /*TODO*/
+    /*class Local3 : T*/
+}
+
+open class C : T
+
+class D : C()
+
+
+// EXIST: A
+// EXIST: B
+// EXIST: Local1
+// EXIST: Local2
+// ABSENT: Local3
+// EXIST: C
+// EXIST: D

idea/testData/completion/smart/ObjectsFromType.kt

@@ -10,5 +10,8 @@ fun foo(): T {
     return <caret>
 }
 
-// EXIST: { lookupString:"Null", itemText:"T.Null", tailText:" (p)", typeText:"T" }
+// EXIST: { lookupString:"Null", itemText:"Null", tailText:" (p.T)" }
+// EXIST: foo
+// EXIST: object
 // ABSENT: Other
+// NUMBER: 3

idea/testData/completion/smartMultiFile/InheritorInTheSameFile/Declarations.kt

@@ -0,0 +1,5 @@
+trait T
+
+class C : A()
+
+// ALLOW_AST_ACCESS

idea/testData/completion/smartMultiFile/InheritorInTheSameFile/InheritorInTheSameFile.kt

@@ -0,0 +1,11 @@
+open class A : T
+
+fun foo(): T {
+    return <caret>
+}
+
+class B : T
+
+// EXIST: A
+// EXIST: B
+// EXIST: C

idea/tests/org/jetbrains/jet/completion/JvmSmartCompletionTestGenerated.java

@@ -19,7 +19,6 @@ package org.jetbrains.jet.completion;
 import com.intellij.testFramework.TestDataPath;
 import org.jetbrains.jet.JUnit3RunnerWithInners;
 import org.jetbrains.jet.JetTestUtils;
-import org.jetbrains.jet.test.InnerTestClasses;
 import org.jetbrains.jet.test.TestMetadata;
 import org.junit.runner.RunWith;
 
@@ -402,6 +401,12 @@ public class JvmSmartCompletionTestGenerated extends AbstractJvmSmartCompletionT
         doTest(fileName);
     }
 
+    @TestMetadata("Inheritors3.kt")
+    public void testInheritors3() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/smart/Inheritors3.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("InsideIdentifier.kt")
     public void testInsideIdentifier() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/smart/InsideIdentifier.kt");

idea/tests/org/jetbrains/jet/completion/MultiFileSmartCompletionTestGenerated.java

@@ -66,6 +66,12 @@ public class MultiFileSmartCompletionTestGenerated extends AbstractMultiFileSmar
         doTest(fileName);
     }
 
+    @TestMetadata("InheritorInTheSameFile")
+    public void testInheritorInTheSameFile() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/smartMultiFile/InheritorInTheSameFile/");
+        doTest(fileName);
+    }
+
     @TestMetadata("Inheritors")
     public void testInheritors() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/smartMultiFile/Inheritors/");