Extract Function: In case of top-level and member declarations place extracted function after the original container

idea/src/org/jetbrains/jet/plugin/debugger/evaluate/KotlinEvaluationBuilder.kt

@@ -356,10 +356,10 @@ private fun getFunctionForExtractedFragment(
             val newDebugExpression = addDebugExpressionBeforeContextElement(codeFragment, contextElement)
             if (newDebugExpression == null) return null
 
-            val nextSibling = tmpFile.getDeclarations().firstOrNull()
-            if (nextSibling == null) return null
+            val targetSibling = tmpFile.getDeclarations().firstOrNull()
+            if (targetSibling == null) return null
 
-            val analysisResult = ExtractionData(tmpFile, Collections.singletonList(newDebugExpression), nextSibling).performAnalysis()
+            val analysisResult = ExtractionData(tmpFile, Collections.singletonList(newDebugExpression), targetSibling).performAnalysis()
             if (analysisResult.status != Status.SUCCESS) {
                 throw EvaluateExceptionUtil.createEvaluateException(getErrorMessageForExtractFunctionResult(analysisResult))
             }

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ExtractKotlinFunctionHandler.kt

@@ -56,11 +56,11 @@ public class ExtractKotlinFunctionHandler : RefactoringActionHandler {
             editor: Editor,
             file: JetFile,
             elements: List<PsiElement>,
-            nextSibling: PsiElement
+            targetSibling: PsiElement
     ) {
         val project = file.getProject()
 
-        val analysisResult = ExtractionData(file, elements, nextSibling).performAnalysis()
+        val analysisResult = ExtractionData(file, elements, targetSibling).performAnalysis()
 
         if (ApplicationManager.getApplication()!!.isUnitTestMode() && analysisResult.status != Status.SUCCESS) {
             throw ConflictsInTestsException(analysisResult.messages.map { it.renderMessage() })
@@ -119,8 +119,8 @@ public class ExtractKotlinFunctionHandler : RefactoringActionHandler {
     override fun invoke(project: Project, editor: Editor, file: PsiFile, dataContext: DataContext?) {
         if (file !is JetFile) return
 
-        selectElements(editor, file) { (elements, targetNextSibling) ->
-            doInvoke(editor, file, elements, targetNextSibling)
+        selectElements(editor, file) { (elements, targetSibling) ->
+            doInvoke(editor, file, elements, targetSibling)
         }
     }
 
@@ -142,7 +142,7 @@ private fun showErrorHintByKey(project: Project, editor: Editor, key: String) {
 fun selectElements(
         editor: Editor,
         file: PsiFile,
-        continuation: (elements: List<PsiElement>, targetNextSibling: PsiElement) -> Unit
+        continuation: (elements: List<PsiElement>, targetSibling: PsiElement) -> Unit
 ) {
     fun noExpressionError() {
         showErrorHintByKey(file.getProject(), editor, "cannot.refactor.no.expression")

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ExtractionData.kt

@@ -33,7 +33,6 @@ import org.jetbrains.jet.plugin.codeInsight.DescriptorToDeclarationUtil
 import java.util.Collections
 import org.jetbrains.jet.lang.psi.JetBlockExpression
 import org.jetbrains.jet.renderer.DescriptorRenderer
-import org.jetbrains.jet.lang.psi.psiUtil.getParentByTypeAndBranch
 import org.jetbrains.jet.lang.psi.JetQualifiedExpression
 import org.jetbrains.jet.lang.psi.psiUtil.isInsideOf
 import java.util.ArrayList
@@ -41,6 +40,9 @@ import com.intellij.psi.PsiNamedElement
 import org.jetbrains.jet.lang.psi.JetSuperExpression
 import org.jetbrains.jet.lang.types.JetType
 import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache
+import org.jetbrains.jet.lang.psi.psiUtil.getParentByType
+import org.jetbrains.jet.lang.psi.JetDeclaration
+import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
 
 data class ResolveResult(
         val originalRefExpr: JetSimpleNameExpression,
@@ -58,10 +60,12 @@ data class ResolvedReferenceInfo(
 class ExtractionData(
         val originalFile: JetFile,
         val originalElements: List<PsiElement>,
-        val nextSibling: PsiElement
+        val targetSibling: PsiElement
 ) {
     val project: Project = originalFile.getProject()
 
+    val insertBefore: Boolean = targetSibling.getParentByType(javaClass<JetDeclaration>(), true) is JetDeclarationWithBody
+
     fun getExpressions(): List<JetExpression> = originalElements.filterIsInstance(javaClass<JetExpression>())
 
     fun getCodeFragmentTextRange(): TextRange? {

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractFunctionUtils.kt

@@ -97,7 +97,6 @@ import org.jetbrains.jet.lang.diagnostics.Errors
 import org.jetbrains.jet.lang.psi.JetTypeReference
 import org.jetbrains.jet.lang.psi.JetTypeParameterListOwner
 import org.jetbrains.jet.plugin.refactoring.extractFunction.AnalysisResult.Status
-import org.jetbrains.jet.lang.descriptors.FunctionDescriptor
 import org.jetbrains.jet.lang.psi.codeFragmentUtil.skipVisibilityCheck
 import org.jetbrains.jet.lang.psi.codeFragmentUtil.setSkipVisibilityCheck
 import org.jetbrains.jet.plugin.refactoring.extractFunction.AnalysisResult.ErrorMessage
@@ -232,15 +231,32 @@ private fun List<Instruction>.analyzeControlFlow(
     return Pair(DefaultControlFlow, null)
 }
 
-private fun ExtractionData.createTemporaryCodeBlock(): JetBlockExpression {
-    val position = nextSibling.getTextRange()!!.getStartOffset()
-    val tmpFile = originalFile.createTempCopy { text ->
-        StringBuilder(text).insert(position, "fun() {\n${getCodeFragmentText()}\n}\n").toString()
+fun ExtractionData.createTemporaryFunction(functionText: String): JetNamedFunction {
+    val textRange = targetSibling.getTextRange()!!
+
+    val insertText: String
+    val insertPosition: Int
+    val lookupPosition: Int
+    if (insertBefore) {
+        insertPosition = textRange.getStartOffset()
+        lookupPosition = insertPosition
+        insertText = functionText
     }
-    val tmpFunction = tmpFile.findElementAt(position)?.getParentByType(javaClass<JetNamedFunction>())!!
-    return tmpFunction.getBodyExpression() as JetBlockExpression
+    else {
+        insertPosition = textRange.getEndOffset()
+        lookupPosition = insertPosition + 1
+        insertText = "\n$functionText"
+    }
+
+    val tmpFile = originalFile.createTempCopy { text ->
+        StringBuilder(text).insert(insertPosition, insertText).toString()
+    }
+    return tmpFile.findElementAt(lookupPosition)?.getParentByType(javaClass<JetNamedFunction>())!!
 }
 
+private fun ExtractionData.createTemporaryCodeBlock(): JetBlockExpression =
+        createTemporaryFunction("fun() {\n${getCodeFragmentText()}\n}\n").getBodyExpression() as JetBlockExpression
+
 private fun JetType.collectReferencedTypes(): List<JetType> {
     return DFS.dfsFromNode(
             this,
@@ -534,8 +550,8 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
     checkDeclarationsMovingOutOfScope(controlFlow)?.let { messages.add(it) }
 
     val functionNameValidator = JetNameValidatorImpl(
-            nextSibling.getParent(),
-            nextSibling,
+            targetSibling.getParent(),
+            targetSibling,
             JetNameValidatorImpl.Target.FUNCTIONS_AND_CLASSES
     )
     val functionName = JetNameSuggester.suggestNames(controlFlow.returnType, functionNameValidator, DEFAULT_FUNCTION_NAME).first()
@@ -673,14 +689,11 @@ fun ExtractionDescriptor.generateFunction(
     fun createFunction(): JetNamedFunction {
         return with(extractionData) {
             if (inTempFile) {
-                val position = nextSibling.getTextRange()!!.getStartOffset()
-                val tmpFile = originalFile.createTempCopy { text ->
-                    StringBuilder(text).insert(position, getFunctionText() + "\n").toString()
-                }
+                val function = createTemporaryFunction("${getFunctionText()}\n")
                 if (originalFile.skipVisibilityCheck()) {
-                    tmpFile.setSkipVisibilityCheck(true)
+                    function.getContainingJetFile().setSkipVisibilityCheck(true)
                 }
-                tmpFile.findElementAt(position)?.getParentByType(javaClass<JetNamedFunction>())!!
+                function
             }
             else {
                 JetPsiFactory.createFunction(project, getFunctionText())
@@ -763,11 +776,20 @@ fun ExtractionDescriptor.generateFunction(
 
     fun insertFunction(function: JetNamedFunction): JetNamedFunction {
         return with(extractionData) {
-            val targetContainer = nextSibling.getParent()!!
-            val functionInFile = targetContainer.addBefore(function, nextSibling) as JetNamedFunction
-            targetContainer.addBefore(JetPsiFactory.createWhiteSpace(project, "\n\n"), nextSibling)
+            val targetContainer = targetSibling.getParent()!!
+            val emptyLines = JetPsiFactory.createWhiteSpace(project, "\n\n")
+            if (insertBefore) {
+                val functionInFile = targetContainer.addBefore(function, targetSibling) as JetNamedFunction
+                targetContainer.addBefore(emptyLines, targetSibling)
 
-            functionInFile
+                functionInFile
+            }
+            else {
+                val functionInFile = targetContainer.addAfter(function, targetSibling) as JetNamedFunction
+                targetContainer.addAfter(emptyLines, targetSibling)
+
+                functionInFile
+            }
         }
     }
 

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ui/KotlinExtractFunctionDialog.java

@@ -26,7 +26,6 @@ import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.containers.MultiMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.JetClassBody;
-import org.jetbrains.jet.lang.psi.JetClassOrObject;
 import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.plugin.refactoring.JetNameSuggester;
 import org.jetbrains.jet.plugin.refactoring.JetRefactoringBundle;
@@ -72,7 +71,7 @@ public class KotlinExtractFunctionDialog extends DialogWrapper {
     }
 
     private boolean isVisibilitySectionAvailable() {
-        PsiElement target = originalDescriptor.getDescriptor().getExtractionData().getNextSibling().getParent();
+        PsiElement target = originalDescriptor.getDescriptor().getExtractionData().getTargetSibling().getParent();
         return target instanceof JetClassBody || target instanceof JetFile;
     }