KT-1379 Override Methods does not respect caret position

#KT-1379 fixed

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetClassBody.java

@@ -68,8 +68,15 @@ public class JetClassBody extends JetElement {
         return (JetClassObject) findChildByType(JetNodeTypes.CLASS_OBJECT);
     }
 
+    @Nullable
     public PsiElement getRBrace() {
         final ASTNode[] children = getNode().getChildren(TokenSet.create(JetTokens.RBRACE));
         return children.length == 1 ? children[0].getPsi() : null;
     }
+
+    @Nullable
+    public PsiElement getLBrace() {
+        final ASTNode[] children = getNode().getChildren(TokenSet.create(JetTokens.LBRACE));
+        return children.length == 1 ? children[0].getPsi() : null;
+    }
 }

idea/src/org/jetbrains/jet/plugin/codeInsight/OverrideImplementMethodsHandler.java

@@ -19,15 +19,16 @@ package org.jetbrains.jet.plugin.codeInsight;
 import com.intellij.codeInsight.hint.HintManager;
 import com.intellij.ide.util.MemberChooser;
 import com.intellij.lang.LanguageCodeInsightActionHandler;
-import com.intellij.openapi.application.Result;
+import com.intellij.openapi.application.ApplicationManager;
-import com.intellij.openapi.command.WriteCommandAction;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.DialogWrapper;
+import com.intellij.openapi.util.Condition;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
@@ -54,34 +55,71 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
     }
 
     public Set<CallableMemberDescriptor> collectMethodsToGenerate(JetClassOrObject classOrObject) {
-        BindingContext bindingContext = WholeProjectAnalyzerFacade.analyzeProjectWithCacheOnAFile((JetFile) classOrObject.getContainingFile());
+        BindingContext bindingContext =
+            WholeProjectAnalyzerFacade.analyzeProjectWithCacheOnAFile((JetFile)classOrObject.getContainingFile());
         final DeclarationDescriptor descriptor = bindingContext.get(BindingContext.DECLARATION_TO_DESCRIPTOR, classOrObject);
         if (descriptor instanceof MutableClassDescriptor) {
-            return collectMethodsToGenerate((MutableClassDescriptor) descriptor);
+            return collectMethodsToGenerate((MutableClassDescriptor)descriptor);
         }
         return Collections.emptySet();
     }
 
     protected abstract Set<CallableMemberDescriptor> collectMethodsToGenerate(MutableClassDescriptor descriptor);
 
-    public static void generateMethods(Project project, Editor editor, JetClassOrObject classOrObject, List<DescriptorClassMember> selectedElements) {
+    public static void generateMethods(Project project,
+                                       Editor editor,
+                                       JetClassOrObject classOrObject,
+                                       List<DescriptorClassMember> selectedElements) {
         final JetClassBody body = classOrObject.getBody();
         if (body == null) {
             return;
         }
 
-       for (DescriptorClassMember selectedElement : selectedElements) {
+        // NOTE + TODO: If you try to cache findInsertBeforeAnchor element, there will be failed assertion
-           final DeclarationDescriptor descriptor = selectedElement.getDescriptor();
+        // "PSI/document inconsistency before reparse: file=" from DocumentCommitThread after inserting two overriding
-           JetFile containingFile = (JetFile) classOrObject.getContainingFile();
+        // with the caret right before existing function start.
-           if (descriptor instanceof SimpleFunctionDescriptor) {
+
-               JetElement target = overrideFunction(project, containingFile, (SimpleFunctionDescriptor) descriptor);
+        PsiElement afterAnchor = findInsertAfterAnchor(editor, body);
-               body.addBefore(target, body.getRBrace());
+
-           }
+        if (afterAnchor == null) {
-           else if (descriptor instanceof PropertyDescriptor) {
+            return;
-               JetElement target = overrideProperty(project, containingFile, (PropertyDescriptor) descriptor);
+        }
-               body.addBefore(target, body.getRBrace());
+
-           }
+        for (DescriptorClassMember selectedElement : selectedElements) {
-       }
+            final DeclarationDescriptor descriptor = selectedElement.getDescriptor();
+            JetFile containingFile = (JetFile)classOrObject.getContainingFile();
+            if (descriptor instanceof SimpleFunctionDescriptor) {
+                JetElement target = overrideFunction(project, containingFile, (SimpleFunctionDescriptor)descriptor);
+                afterAnchor = body.addAfter(target, afterAnchor);
+            }
+            else if (descriptor instanceof PropertyDescriptor) {
+                JetElement target = overrideProperty(project, containingFile, (PropertyDescriptor)descriptor);
+                afterAnchor = body.addBefore(target, afterAnchor);
+            }
+        }
+    }
+
+    @Nullable
+    private static PsiElement findInsertAfterAnchor(Editor editor, final JetClassBody body) {
+        PsiElement afterAnchor = body.getLBrace();
+        if (afterAnchor == null) {
+            return null;
+        }
+
+        int offset = editor.getCaretModel().getOffset();
+        PsiElement offsetCursorElement = PsiTreeUtil.findFirstParent(body.getContainingFile().findElementAt(offset),
+                                                                     new Condition<PsiElement>() {
+                                                                         @Override
+                                                                         public boolean value(PsiElement element) {
+                                                                             return element.getParent() == body;
+                                                                         }
+                                                                     });
+
+        if (offsetCursorElement != null) {
+            afterAnchor = offsetCursorElement;
+        }
+
+        return afterAnchor;
     }
 
     private static JetElement overrideProperty(Project project, JetFile file, PropertyDescriptor descriptor) {
@@ -149,7 +187,7 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
     }
 
     private MemberChooser<DescriptorClassMember> showOverrideImplementChooser(Project project,
-                                                                                     DescriptorClassMember[] members) {
+                                                                              DescriptorClassMember[] members) {
         final MemberChooser<DescriptorClassMember> chooser = new MemberChooser<DescriptorClassMember>(members, true, true, project);
         chooser.setTitle(getChooserTitle());
         chooser.show();
@@ -185,9 +223,12 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
         final List<DescriptorClassMember> selectedElements;
         if (implementAll) {
             selectedElements = members;
-        } else {
+        }
+        else {
             final MemberChooser<DescriptorClassMember> chooser = showOverrideImplementChooser(project,
-                                                                                              members.toArray(new DescriptorClassMember[members.size()]));
+                                                                                              members.toArray(
+                                                                                                  new DescriptorClassMember[members
+                                                                                                      .size()]));
             if (chooser == null) {
                 return;
             }
@@ -196,12 +237,12 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
             if (selectedElements == null || selectedElements.isEmpty()) return;
         }
 
-        new WriteCommandAction(project, file) {
+        ApplicationManager.getApplication().runWriteAction(new Runnable() {
-          protected void run(final Result result) throws Throwable {
+            @Override
-            generateMethods(project, editor, classOrObject, selectedElements);
+            public void run() {
-          }
+                generateMethods(project, editor, classOrObject, selectedElements);
-        }.execute();
+            }
-
+        });
     }
 
     @Override

idea/testData/codeInsight/overrideImplement/property.kt.after

@@ -3,6 +3,6 @@ trait T {
 }
 
 class GC() : T {
-
     override val v : Int = 0
+
 }

idea/testData/codeInsight/overrideImplement/respectCaretPosition.kt

@@ -0,0 +1,15 @@
+trait Test {
+    public open fun test()
+    protected open val testProp : Int
+}
+
+class SomeTest : Test {
+    val hello = 12
+    <caret>
+    /**
+     * test
+     */
+    fun some() {
+
+    }
+}

idea/testData/codeInsight/overrideImplement/respectCaretPosition.kt.after

@@ -0,0 +1,19 @@
+trait Test {
+    public open fun test()
+    protected open val testProp : Int
+}
+
+class SomeTest : Test {
+    val hello = 12
+    override val testProp : Int = 0
+    override fun test() {
+        throw UnsupportedOperationException()
+    }
+
+    /**
+    * test
+    */
+    fun some() {
+
+    }
+}

idea/tests/org/jetbrains/jet/plugin/codeInsight/OverrideImplementTest.java

@@ -69,14 +69,18 @@ public class OverrideImplementTest extends LightCodeInsightFixtureTestCase {
         doFileTest();
     }
 
-    public void testTraitNullableFunction() {
+    public void testRespectCaretPosition() {
-        doFileTest();
+        doMultiFileTest();
     }
 
     public void testGenerateMulti() {
         doMultiFileTest();
     }
 
+    public void testTraitNullableFunction() {
+        doFileTest();
+    }
+
     private void doFileTest() {
         myFixture.configureByFile(getTestName(true) + ".kt");
         doImplement();
@@ -100,7 +104,7 @@ public class OverrideImplementTest extends LightCodeInsightFixtureTestCase {
         final PsiElement elementAtCaret = myFixture.getFile().findElementAt(myFixture.getEditor().getCaretModel().getOffset());
         final JetClassOrObject classOrObject = PsiTreeUtil.getParentOfType(elementAtCaret, JetClassOrObject.class);
         final Set<CallableMemberDescriptor> descriptors = new ImplementMethodsHandler().collectMethodsToGenerate(classOrObject);
-        assertEquals(1, descriptors.size());
+        assertEquals("Invalid number of available descriptors for override", 1, descriptors.size());
         new WriteCommandAction(myFixture.getProject(), myFixture.getFile()) {
             @Override
             protected void run(Result result) throws Throwable {