Repository: cqfn/diKTat (kotlin)

Issue:
autofix for `WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES` removes values from enum class
## Describe the bug

Applying autofix for the rule `WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES` via the Gradle plugin (`./gradlew diktatFix`) can result in values being removed from an enum class such as the following:

```
enum class Enum {
  FOO,
  BAR,
  ;

  companion object
  fun f() {}
}
```

## Expected behavior

The function definition and the companion object should be swapped, resulting in the following file:
```
enum class Enum {
  FOO,
  BAR,
  ;

  fun f() {}
  companion object
}
```

The above is passes `./gradlew diktatCheck`.

## Observed behavior

The autofix removes the values from the enum class, resulting in the following file:

```
enum class Enum {
  fun f() {}

  companion object
}
```

## Steps to Reproduce

Run autofix against a project containing an enum class such as the one described above, e.g. via `./gradlew diktatFix` when using the Gradle plugin.

## Environment information
* diktat version: 0.5.0
* build tool (maven/gradle): Gradle, v6.8.3
* how is diktat run (CLI, plugin, etc.):
* kotlin version: 1.4.32 (kotlin-compiler-embeddable-1.4.30)
* JDK: Corretto-11.0.7.10.1
* operating system: GNU/Linux; Ubuntu 20.04

Pull request:
### What's done:
* Changed code style guide - properties are considered loggers only if they are private
* Fixed bug in logic
* Added tests
* Added smoke tests
* Refactoring

This pull request closes #843

Patch:
diff --git a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/ClassLikeStructuresOrderRule.kt b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/ClassLikeStructuresOrderRule.kt
index 12a6d7ba4f..75ba3892cf 100644
--- a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/ClassLikeStructuresOrderRule.kt
+++ b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/ClassLikeStructuresOrderRule.kt
@@ -12,13 +12,13 @@ import com.pinterest.ktlint.core.ast.ElementType.CLASS_BODY
 import com.pinterest.ktlint.core.ast.ElementType.CLASS_INITIALIZER
 import com.pinterest.ktlint.core.ast.ElementType.COMPANION_KEYWORD
 import com.pinterest.ktlint.core.ast.ElementType.CONST_KEYWORD
+import com.pinterest.ktlint.core.ast.ElementType.ENUM_ENTRY
 import com.pinterest.ktlint.core.ast.ElementType.EOL_COMMENT
 import com.pinterest.ktlint.core.ast.ElementType.FUN
 import com.pinterest.ktlint.core.ast.ElementType.IDENTIFIER
 import com.pinterest.ktlint.core.ast.ElementType.KDOC
 import com.pinterest.ktlint.core.ast.ElementType.LATEINIT_KEYWORD
 import com.pinterest.ktlint.core.ast.ElementType.LBRACE
-import com.pinterest.ktlint.core.ast.ElementType.MODIFIER_LIST
 import com.pinterest.ktlint.core.ast.ElementType.OBJECT_DECLARATION
 import com.pinterest.ktlint.core.ast.ElementType.PRIVATE_KEYWORD
 import com.pinterest.ktlint.core.ast.ElementType.PROPERTY
@@ -33,7 +33,9 @@ import com.pinterest.ktlint.core.ast.parent
 import com.pinterest.ktlint.core.ast.prevSibling
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.PsiWhiteSpaceImpl
+import org.jetbrains.kotlin.com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.psi.psiUtil.siblings
@@ -53,69 +55,28 @@ class ClassLikeStructuresOrderRule(configRules: List<RulesConfig>) : DiktatRule(
         }
     }
 
-    @Suppress(
-        "UnsafeCallOnNullableType",
-        "LongMethod",
-        "ComplexMethod",
-        "TOO_LONG_FUNCTION"
-    )
     private fun checkDeclarationsOrderInClass(node: ASTNode) {
-        val allProperties = node.getChildren(TokenSet.create(PROPERTY))
-        val constProperties = allProperties.filter { it.findLeafWithSpecificType(CONST_KEYWORD) != null }.toMutableList()
-        val lateInitProperties = allProperties.filter { it.findLeafWithSpecificType(LATEINIT_KEYWORD) != null }.toMutableList()
-        val loggers = allProperties
-            .filter {
-                (it.findChildByType(MODIFIER_LIST) == null || it.findLeafWithSpecificType(PRIVATE_KEYWORD) != null) &&
-                        it.getIdentifierName()!!.text.contains(loggerPropertyRegex)
-            }
-            .toMutableList()
-        val properties = allProperties.filter { it !in lateInitProperties && it !in loggers && it !in constProperties }.toMutableList()
-        val initBlocks = node.getChildren(TokenSet.create(CLASS_INITIALIZER)).toMutableList()
-        val constructors = node.getChildren(TokenSet.create(SECONDARY_CONSTRUCTOR)).toMutableList()
-        val methods = node.getChildren(TokenSet.create(FUN)).toMutableList()
-        val (usedClasses, unusedClasses) = node
-            .getChildren(TokenSet.create(CLASS))
-            .partition { classNode ->
-                classNode.getIdentifierName()?.let { identifierNode ->
-                    node
-                        .parents()
-                        .last()
-                        .findAllDescendantsWithSpecificType(REFERENCE_EXPRESSION)
-                        .any { ref ->
-                            ref.parent({ it == classNode }) == null && ref.text.contains(identifierNode.text)
-                        }
-                } ?: false
-            }
-            .let { it.first.toMutableList() to it.second.toMutableList() }
-        val (companionObject, objects) = node.getChildren(TokenSet.create(OBJECT_DECLARATION))
-            .partition { it.findChildByType(MODIFIER_LIST)?.findLeafWithSpecificType(COMPANION_KEYWORD) != null }
-        val blocks = Blocks(AllProperties(loggers, constProperties, properties, lateInitProperties), objects.toMutableList(),
-            initBlocks, constructors, methods, usedClasses, companionObject.toMutableList(),
-            unusedClasses)
+        val allProperties = AllProperties.fromClassBody(node)
+        val initBlocks = node.getAllChildrenWithType(CLASS_INITIALIZER)
+        val constructors = node.getAllChildrenWithType(SECONDARY_CONSTRUCTOR)
+        val methods = node.getAllChildrenWithType(FUN)
+        val (usedClasses, unusedClasses) = node.getUsedAndUnusedClasses()
+        val (companionObject, objects) = node.getAllChildrenWithType(OBJECT_DECLARATION)
+            .partition { it.hasModifier(COMPANION_KEYWORD) }
+        val blocks = Blocks(
+            (node.psi as KtClassBody).enumEntries.map { it.node },
+            allProperties, objects, initBlocks, constructors,
+            methods, usedClasses, companionObject, unusedClasses
+        )
             .allBlockFlattened()
             .map { astNode ->
-                Pair(astNode, astNode.siblings(false).takeWhile { it.elementType == WHITE_SPACE || it.isPartOfComment() }.toList())
+                listOf(astNode) +
+                        astNode.siblings(false)
+                            .takeWhile { it.elementType == WHITE_SPACE || it.isPartOfComment() }
+                            .toList()
             }
 
-        val classChildren = node.children().filter { it.elementType in childrenTypes }.toList()
-        if (classChildren != blocks.map { it.first }) {
-            blocks.filterIndexed { index, pair -> classChildren[index] != pair.first }
-                .forEach { listOfChildren ->
-                    val astNode = listOfChildren.first
-                    WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnAndFix(configRules, emitWarn, isFixMode,
-                        "${astNode.elementType}: ${astNode.findChildByType(IDENTIFIER)?.text ?: astNode.text}", astNode.startOffset, astNode) {
-                        node.removeRange(node.findChildByType(LBRACE)!!.treeNext, node.findChildByType(RBRACE)!!)
-                        blocks
-                            .reversed()
-                            .map { bodyChild ->
-                                node.addChild(bodyChild.first, node.children().toList()[1])
-                                bodyChild.second.map { node.addChild(it, node.children().toList()[1]) }
-                            }
-                        node.addChild(PsiWhiteSpaceImpl("\n"), node.lastChildNode)
-                        // fixme maybe wrong space between nodes
-                    }
-                }
-        }
+        node.checkAndReorderBlocks(blocks)
     }
 
     @Suppress("UnsafeCallOnNullableType")
@@ -150,6 +111,49 @@ class ClassLikeStructuresOrderRule(configRules: List<RulesConfig>) : DiktatRule(
         }
     }
 
+    /**
+     * Returns nested classes grouped by whether they are used inside [this] file.
+     * [this] ASTNode should have elementType [CLASS_BODY]
+     */
+    private fun ASTNode.getUsedAndUnusedClasses() = getAllChildrenWithType(CLASS)
+        .partition { classNode ->
+            classNode.getIdentifierName()?.let { identifierNode ->
+                parents()
+                    .last()
+                    .findAllDescendantsWithSpecificType(REFERENCE_EXPRESSION)
+                    .any { ref ->
+                        ref.parent({ it == classNode }) == null && ref.text.contains(identifierNode.text)
+                    }
+            } ?: false
+        }
+
+    /**
+     * Checks whether all class elements in [this] node are correctly ordered and reorders them in fix mode.
+     * [this] ASTNode should have elementType [CLASS_BODY]
+     *
+     * @param blocks list of class elements with leading whitespaces and comments
+     */
+    @Suppress("UnsafeCallOnNullableType")
+    private fun ASTNode.checkAndReorderBlocks(blocks: List<List<ASTNode>>) {
+        val classChildren = this.children().filter { it.elementType in childrenTypes }.toList()
+        if (classChildren != blocks.map { it.first() }) {
+            blocks.filterIndexed { index, pair -> classChildren[index] != pair.first() }
+                .forEach { listOfChildren ->
+                    val astNode = listOfChildren.first()
+                    WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnAndFix(configRules, emitWarn, isFixMode,
+                        "${astNode.elementType}: ${astNode.findChildByType(IDENTIFIER)?.text ?: astNode.text}", astNode.startOffset, astNode) {
+                        removeRange(findChildByType(LBRACE)!!.treeNext, findChildByType(RBRACE)!!)
+                        blocks.reversed()
+                            .forEach { bodyChild ->
+                                bodyChild.forEach { this.addChild(it, this.children().take(2).last()) }
+                            }
+                        // Add newline before the closing `}`. All other newlines will be properly formatted by `NewlinesRule`.
+                        this.addChild(PsiWhiteSpaceImpl("\n"), this.lastChildNode)
+                    }
+                }
+        }
+    }
+
     /**
      * Data class containing different groups of properties in file
      *
@@ -158,12 +162,33 @@ class ClassLikeStructuresOrderRule(configRules: List<RulesConfig>) : DiktatRule(
      * @property properties all other properties
      * @property lateInitProperties `lateinit var`s
      */
-    private data class AllProperties(val loggers: MutableList<ASTNode>,
-                                     val constProperties: MutableList<ASTNode>,
-                                     val properties: MutableList<ASTNode>,
-                                     val lateInitProperties: MutableList<ASTNode>)
+    private data class AllProperties(val loggers: List<ASTNode>,
+                                     val constProperties: List<ASTNode>,
+                                     val properties: List<ASTNode>,
+                                     val lateInitProperties: List<ASTNode>) {
+        companion object {
+            /**
+             * Create [AllProperties] wrapper from node with type [CLASS_BODY]
+             *
+             * @param node an ASTNode with type [CLASS_BODY]
+             * @return an instance of [AllProperties]
+             */
+            @Suppress("UnsafeCallOnNullableType")
+            fun fromClassBody(node: ASTNode): AllProperties {
+                val allProperties = node.getAllChildrenWithType(PROPERTY)
+                val constProperties = allProperties.filterByModifier(CONST_KEYWORD)
+                val lateInitProperties = allProperties.filterByModifier(LATEINIT_KEYWORD)
+                val loggers = allProperties.filterByModifier(PRIVATE_KEYWORD).filter {
+                    it.getIdentifierName()!!.text.contains(loggerPropertyRegex)
+                }
+                val properties = allProperties.filter { it !in lateInitProperties && it !in loggers && it !in constProperties }
+                return AllProperties(loggers, constProperties, properties, lateInitProperties)
+            }
+        }
+    }
 
     /**
+     * @property enumEntries if this class is a enum class, list of its entries. Otherwise an empty list.
      * @property allProperties an instance of [AllProperties]
      * @property objects objects
      * @property initBlocks `init` blocks
@@ -173,23 +198,24 @@ class ClassLikeStructuresOrderRule(configRules: List<RulesConfig>) : DiktatRule(
      * @property companion `companion object`s
      * @property unusedClasses nested classes that are *not* used in the enclosing class
      */
-    private data class Blocks(val allProperties: AllProperties,
-                              val objects: MutableList<ASTNode>,
-                              val initBlocks: MutableList<ASTNode>,
-                              val constructors: MutableList<ASTNode>,
-                              val methods: MutableList<ASTNode>,
-                              val usedClasses: MutableList<ASTNode>,
-                              val companion: MutableList<ASTNode>,
-                              val unusedClasses: MutableList<ASTNode>) {
+    private data class Blocks(val enumEntries: List<ASTNode>,
+                              val allProperties: AllProperties,
+                              val objects: List<ASTNode>,
+                              val initBlocks: List<ASTNode>,
+                              val constructors: List<ASTNode>,
+                              val methods: List<ASTNode>,
+                              val usedClasses: List<ASTNode>,
+                              val companion: List<ASTNode>,
+                              val unusedClasses: List<ASTNode>) {
         init {
-            require(companion.size in 0..1)
+            require(companion.size in 0..1) { "There is more than one companion object in class" }
         }
 
         /**
          * @return all groups of structures in the class
          */
         fun allBlocks() = with(allProperties) {
-            listOf(loggers, constProperties, properties, lateInitProperties, objects,
+            listOf(enumEntries, loggers, constProperties, properties, lateInitProperties, objects,
                 initBlocks, constructors, methods, usedClasses, companion, unusedClasses)
         }
 
@@ -200,6 +226,10 @@ class ClassLikeStructuresOrderRule(configRules: List<RulesConfig>) : DiktatRule(
     }
 
     companion object {
-        private val childrenTypes = listOf(PROPERTY, CLASS, CLASS_INITIALIZER, SECONDARY_CONSTRUCTOR, FUN, OBJECT_DECLARATION)
+        private val childrenTypes = listOf(PROPERTY, CLASS, CLASS_INITIALIZER, SECONDARY_CONSTRUCTOR, FUN, OBJECT_DECLARATION, ENUM_ENTRY)
     }
 }
+
+private fun Iterable<ASTNode>.filterByModifier(modifier: IElementType) = filter {
+    it.findLeafWithSpecificType(modifier) != null
+}
diff --git a/info/guide/guide-chapter-3.md b/info/guide/guide-chapter-3.md
index 8dc2be2ee0..ed471365d3 100644
--- a/info/guide/guide-chapter-3.md
+++ b/info/guide/guide-chapter-3.md
@@ -76,7 +76,7 @@ If the classes are meant to be used externally, and are not referenced inside th
 - Companion object
 
 **Exception:**
-All variants of a `(private) val` logger should be placed at the beginning of the class (`(private) val log`, `LOG`, `logger`, etc.).
+All variants of a `private val` logger should be placed at the beginning of the class (`private val log`, `LOG`, `logger`, etc.).
 
 #### <a name="r3.1.5"></a> 3.1.5 Order of declaration of top-level code structures
 Kotlin allows several top-level declaration types: classes, objects, interfaces, properties and functions.


Test patch:
diff --git a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderFixTest.kt b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderFixTest.kt
index 8b64b99e2f..1e1bef6fa0 100644
--- a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderFixTest.kt
+++ b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderFixTest.kt
@@ -20,4 +20,10 @@ class ClassLikeStructuresOrderFixTest : FixTestBase("test/paragraph3/file_struct
     fun `should fix order and newlines with comment`() {
         fixAndCompare("OrderWithCommentExpected.kt", "OrderWithCommentTest.kt")
     }
+
+    @Test
+    @Tags(Tag(WarningNames.WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES))
+    fun `regression - should not remove enum members`() {
+        fixAndCompare("OrderWithEnumsExpected.kt", "OrderWithEnumsTest.kt")
+    }
 }
diff --git a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderRuleWarnTest.kt b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderRuleWarnTest.kt
index ef3bae9674..05dbfb2fb0 100644
--- a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderRuleWarnTest.kt
+++ b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/ClassLikeStructuresOrderRuleWarnTest.kt
@@ -8,6 +8,7 @@ import org.cqfn.diktat.util.LintTestBase
 
 import com.pinterest.ktlint.core.LintError
 import generated.WarningNames
+import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Tag
 import org.junit.jupiter.api.Test
 
@@ -50,18 +51,16 @@ class ClassLikeStructuresOrderRuleWarnTest : LintTestBase(::ClassLikeStructuresO
     @Test
     @Tag(WarningNames.WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES)
     fun `should warn if loggers are not on top`() {
-        listOf("private ", "").forEach { modifier ->
-            lintMethod(
-                """
+        lintMethod(
+            """
                     |class Example {
                     |    private val FOO = 42
-                    |    ${modifier}val log = LoggerFactory.getLogger(Example.javaClass)
+                    |    private val log = LoggerFactory.getLogger(Example.javaClass)
                     |}
                 """.trimMargin(),
-                LintError(2, 5, ruleId, "${WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnText()} PROPERTY: FOO", true),
-                LintError(3, 5, ruleId, "${WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnText()} PROPERTY: log", true)
-            )
-        }
+            LintError(2, 5, ruleId, "${WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnText()} PROPERTY: FOO", true),
+            LintError(3, 5, ruleId, "${WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnText()} PROPERTY: log", true)
+        )
     }
 
     // ===== comments on properties ======
@@ -198,4 +197,21 @@ class ClassLikeStructuresOrderRuleWarnTest : LintTestBase(::ClassLikeStructuresO
             LintError(5, 29, ruleId, "${WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES.warnText()} PROPERTY: a", true)
         )
     }
+
+    @Test
+    @Tag(WarningNames.WRONG_ORDER_IN_CLASS_LIKE_STRUCTURES)
+    fun `should correctly check class elements order in enum classes`() {
+        lintMethod(
+            """
+                enum class Enum {
+                    FOO,
+                    BAR,
+                    ;
+                    
+                    fun f() {}
+                    companion object
+                }
+            """.trimMargin()
+        )
+    }
 }
diff --git a/diktat-rules/src/test/resources/test/paragraph3/file_structure/OrderWithEnumsExpected.kt b/diktat-rules/src/test/resources/test/paragraph3/file_structure/OrderWithEnumsExpected.kt
new file mode 100644
index 0000000000..a2c89b1489
--- /dev/null
+++ b/diktat-rules/src/test/resources/test/paragraph3/file_structure/OrderWithEnumsExpected.kt
@@ -0,0 +1,8 @@
+enum class Enum {
+    FOO,
+    BAR,
+    ;
+    fun f() {}
+
+    companion object
+}
\ No newline at end of file
diff --git a/diktat-rules/src/test/resources/test/paragraph3/file_structure/OrderWithEnumsTest.kt b/diktat-rules/src/test/resources/test/paragraph3/file_structure/OrderWithEnumsTest.kt
new file mode 100644
index 0000000000..ee03e65ebe
--- /dev/null
+++ b/diktat-rules/src/test/resources/test/paragraph3/file_structure/OrderWithEnumsTest.kt
@@ -0,0 +1,8 @@
+enum class Enum {
+    FOO,
+    BAR,
+    ;
+
+    companion object
+    fun f() {}
+}
\ No newline at end of file
diff --git a/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Expected.kt b/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Expected.kt
index 8e98e9aa7e..09cfc55dd9 100644
--- a/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Expected.kt
+++ b/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Expected.kt
@@ -37,6 +37,21 @@ enum class IssueType {
     PROJECT_STRUCTURE, TESTS, VCS
 }
 
+enum class IssueType2 {
+    PROJECT_STRUCTURE,
+    TESTS,
+    VCS,
+    ;
+
+    /**
+     * @param bar
+     * @return
+     */
+    fun foo(bar: Int) = bar
+
+    companion object
+}
+
 class Foo {
     /**
      * @implNote lorem ipsum
diff --git a/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Test.kt b/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Test.kt
index ee6d20b539..7ac28f6eec 100644
--- a/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Test.kt
+++ b/diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example3Test.kt
@@ -45,6 +45,13 @@ enum class IssueType {
     VCS, PROJECT_STRUCTURE, TESTS
 }
 
+enum class IssueType2 {
+    VCS, PROJECT_STRUCTURE, TESTS;
+
+    companion object
+    fun foo(bar: Int) = bar
+}
+
 class Foo {
     /**
      * @implNote lorem ipsum