From 3c36795984048aa3922c6d3fa27bb11f2bc757ec Mon Sep 17 00:00:00 2001
From: Alexey Sedunov <alexey.sedunov@jetbrains.com>
Date: Mon, 13 Jul 2015 15:58:46 +0300
Subject: [PATCH] Refactoring: Extract JetClassOrObject.getOrCreateBody()
 function

---
 .../org/jetbrains/kotlin/psi/JetClassOrObject.kt   |  2 ++
 .../codeInsight/OverrideImplementMethodsHandler.kt |  7 +------
 .../actions/JetAddFunctionToClassifierAction.java  |  7 +------
 .../callableBuilder/CallableBuilder.kt             | 14 ++++----------
 4 files changed, 8 insertions(+), 22 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/JetClassOrObject.kt b/compiler/frontend/src/org/jetbrains/kotlin/psi/JetClassOrObject.kt
index ccd7a473169..a76aac43c8b 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/psi/JetClassOrObject.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/psi/JetClassOrObject.kt
@@ -44,6 +44,8 @@ abstract public class JetClassOrObject : JetTypeParameterListOwnerStub<KotlinCla
 
     public fun getBody(): JetClassBody? = getStubOrPsiChild(JetStubElementTypes.CLASS_BODY)
 
+    public fun getOrCreateBody(): JetClassBody = getBody() ?: add(JetPsiFactory(this).createEmptyClassBody()) as JetClassBody
+
     public fun isTopLevel(): Boolean = getStub()?.isTopLevel() ?: (getParent() is JetFile)
 
     public fun isLocal(): Boolean = getStub()?.isLocal() ?: JetPsiUtil.isLocal(this)
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/codeInsight/OverrideImplementMethodsHandler.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/codeInsight/OverrideImplementMethodsHandler.kt
index f31bd49781f..ab571fcc557 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/codeInsight/OverrideImplementMethodsHandler.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/codeInsight/OverrideImplementMethodsHandler.kt
@@ -137,12 +137,7 @@ public abstract class OverrideImplementMethodsHandler : LanguageCodeInsightActio
 
         public fun generateMethods(editor: Editor, classOrObject: JetClassOrObject, selectedElements: List<DescriptorClassMember>) {
             runWriteAction {
-                var body = classOrObject.getBody()
-                if (body == null) {
-                    val psiFactory = JetPsiFactory(classOrObject)
-                    classOrObject.add(psiFactory.createWhiteSpace())
-                    body = classOrObject.add(psiFactory.createEmptyClassBody()) as JetClassBody
-                }
+                val body = classOrObject.getOrCreateBody()
 
                 var afterAnchor = findInsertAfterAnchor(editor, body)
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/actions/JetAddFunctionToClassifierAction.java b/idea/src/org/jetbrains/kotlin/idea/actions/JetAddFunctionToClassifierAction.java
index b26c66c04a8..c0bd3609157 100644
--- a/idea/src/org/jetbrains/kotlin/idea/actions/JetAddFunctionToClassifierAction.java
+++ b/idea/src/org/jetbrains/kotlin/idea/actions/JetAddFunctionToClassifierAction.java
@@ -86,12 +86,7 @@ public class JetAddFunctionToClassifierAction implements QuestionAction {
                     @Override
                     public void run() {
                         JetPsiFactory psiFactory = JetPsiFactory(classifierDeclaration);
-                        JetClassBody body = classifierDeclaration.getBody();
-                        if (body == null) {
-                            PsiElement whitespaceBefore = classifierDeclaration.add(psiFactory.createWhiteSpace());
-                            body = (JetClassBody) classifierDeclaration.addAfter(psiFactory.createEmptyClassBody(), whitespaceBefore);
-                            classifierDeclaration.addAfter(psiFactory.createNewLine(), body);
-                        }
+                        JetClassBody body = classifierDeclaration.getOrCreateBody();
 
                         String functionBody = "";
                         if (typeDescriptor.getKind() != ClassKind.INTERFACE && functionDescriptor.getModality() != Modality.ABSTRACT) {
diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt
index 524a029fb02..1ddce6ccdd0 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt
@@ -566,7 +566,7 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                     containingElement is PsiClass -> {
                         if (declaration is JetSecondaryConstructor) {
                             val wrappingClass = psiFactory.createClass("class ${containingElement.getName()} {\n}")
-                            addDeclarationToClassOrObject(wrappingClass, declaration, psiFactory)
+                            addDeclarationToClassOrObject(wrappingClass, declaration)
                             (jetFileToEdit.add(wrappingClass) as JetClass).getDeclarations().first() as JetNamedDeclaration
                         }
                         else {
@@ -575,7 +575,7 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                     }
 
                     containingElement is JetClassOrObject -> {
-                        addDeclarationToClassOrObject(containingElement, declaration, psiFactory)
+                        addDeclarationToClassOrObject(containingElement, declaration)
                     }
                     else -> throw AssertionError("Invalid containing element: ${containingElement.getText()}")
                 }
@@ -593,14 +593,8 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
         }
 
         private fun addDeclarationToClassOrObject(classOrObject: JetClassOrObject,
-                                                  declaration: JetNamedDeclaration,
-                                                  psiFactory: JetPsiFactory): JetNamedDeclaration {
-            var classBody = classOrObject.getBody()
-            if (classBody == null) {
-                classBody = classOrObject.add(psiFactory.createEmptyClassBody()) as JetClassBody
-                classOrObject.addBefore(psiFactory.createWhiteSpace(), classBody)
-            }
-
+                                                  declaration: JetNamedDeclaration): JetNamedDeclaration {
+            val classBody = classOrObject.getOrCreateBody()
             return if (declaration is JetNamedFunction) {
                 val anchor = PsiTreeUtil.skipSiblingsBackward(
                         classBody.getRBrace() ?: classBody.getLastChild()!!,