Fix: SQL Formatter

src/main/kotlin/org/domaframework/doma/intellij/formatter/block/SqlBlock.kt

@@ -67,21 +67,28 @@ open class SqlBlock(
     open val childBlocks = mutableListOf<SqlBlock>()
     open var prevBlocks = emptyList<SqlBlock>()
 
+    companion object {
+        private const val DEFAULT_INDENT_SIZE = 4
+        private const val DEFAULT_TEXT_LENGTH_INCREMENT = 1
+        private val EXCLUDED_FROM_TEXT_LENGTH = setOf(SqlTypes.DOT, SqlTypes.RIGHT_PAREN)
+        private val SPACING_ONE = Spacing.createSpacing(1, 1, 0, true, 0)
+        private val SPACING_ZERO = Spacing.createSpacing(0, 0, 0, true, 0)
+        private val SPACING_ONE_NO_KEEP = Spacing.createSpacing(1, 1, 0, false, 0)
+    }
+
     fun getChildrenTextLen(): Int = childBlocks.sumOf { child -> calculateChildTextLength(child) }
 
     private fun calculateChildTextLength(child: SqlBlock): Int {
         val nonCommentChildren = child.childBlocks.filterNot { it is SqlDefaultCommentBlock }
 
-        if (nonCommentChildren.isNotEmpty()) {
-            return child.getChildrenTextLen() + child.getNodeText().length
+        return when {
+            nonCommentChildren.isNotEmpty() -> child.getChildrenTextLen() + child.getNodeText().length
+            isExcludedFromTextLength(child) -> 0
+            else -> child.getNodeText().length + DEFAULT_TEXT_LENGTH_INCREMENT
         }
-        if (isExcludedFromTextLength(child)) {
-            return 0
-        }
-        return child.getNodeText().length + 1
     }
 
-    private fun isExcludedFromTextLength(block: SqlBlock): Boolean = block.node.elementType in setOf(SqlTypes.DOT, SqlTypes.RIGHT_PAREN)
+    private fun isExcludedFromTextLength(block: SqlBlock): Boolean = block.node.elementType in EXCLUDED_FROM_TEXT_LENGTH
 
     /**
      * Checks if a conditional loop directive is registered before the parent block.
@@ -112,9 +119,9 @@ open class SqlBlock(
      *
      * @example
      * ```sql
-     * WHERE
-     *      /*%if status == "pending" */
-     *      status = 'pending'
+     * WHERE -- grand
+     *      /*%if status == "pending" */ -- parent
+     *      status = 'pending' -- child
      * ```
      */
     protected fun isElementAfterConditionLoopDirective(): Boolean =
@@ -123,15 +130,33 @@ open class SqlBlock(
                 (parent.parentBlock is SqlNewGroupBlock || parent.parentBlock is SqlElConditionLoopCommentBlock)
         } == true
 
+    protected fun isElementAfterConditionLoopEnd(): Boolean {
+        val prevChildren =
+            prevBlocks
+                .firstOrNull()
+                ?.childBlocks
+
+        val firstConditionBlock = (prevChildren?.firstOrNull() as? SqlElConditionLoopCommentBlock)
+        val endBlock = firstConditionBlock?.conditionEnd
+        if (endBlock == null) return false
+        val lastBlock = prevBlocks.lastOrNull()
+
+        return endBlock.node.startOffset > (lastBlock?.node?.startOffset ?: 0)
+    }
+
     protected fun isFirstChildConditionLoopDirective(): Boolean = childBlocks.firstOrNull() is SqlElConditionLoopCommentBlock
 
     fun getChildBlocksDropLast(
         dropIndex: Int = 1,
         skipCommentBlock: Boolean = true,
+        skipConditionLoopCommentBlock: Boolean = true,
     ): List<SqlBlock> {
-        val children = childBlocks.dropLast(dropIndex)
+        var children = childBlocks.dropLast(dropIndex)
         if (skipCommentBlock) {
-            return children.filter { it !is SqlDefaultCommentBlock }
+            children = children.filter { it !is SqlDefaultCommentBlock }
+        }
+        if (skipConditionLoopCommentBlock) {
+            children = children.filter { it !is SqlElConditionLoopCommentBlock }
         }
         return children
     }
@@ -167,15 +192,14 @@ open class SqlBlock(
     open fun isSaveSpace(lastGroup: SqlBlock?): Boolean =
         when (lastGroup) {
             is SqlNewGroupBlock -> shouldSaveSpaceForNewGroup(lastGroup)
-            else -> {
-                shouldSaveSpaceForConditionLoop()
-            }
-        } == true
+            else -> shouldSaveSpaceForConditionLoop()
+        }
 
     private fun shouldSaveSpaceForConditionLoop(): Boolean =
         isConditionLoopDirectiveRegisteredBeforeParent() ||
             isElementAfterConditionLoopDirective() ||
-            isFirstChildConditionLoopDirective()
+            isFirstChildConditionLoopDirective() ||
+            isElementAfterConditionLoopEnd()
 
     private fun shouldSaveSpaceForNewGroup(parent: SqlNewGroupBlock): Boolean {
         val prevWord = prevBlocks.lastOrNull { it !is SqlCommentBlock }
@@ -184,15 +208,22 @@ open class SqlBlock(
             return false
         }
 
-        return isFollowedByConditionLoop() || isPrecededByConditionLoop(parent)
+        return hasConditionLoopAround(parent)
     }
 
     private fun isNonBreakingKeywordCombination(
         parent: SqlNewGroupBlock,
         prevWord: SqlBlock?,
-    ): Boolean =
-        SqlKeywordUtil.isSetLineKeyword(getNodeText(), parent.getNodeText()) ||
-            SqlKeywordUtil.isSetLineKeyword(getNodeText(), prevWord?.getNodeText() ?: "")
+    ): Boolean {
+        val currentText = getNodeText()
+        val parentText = parent.getNodeText()
+        val prevText = prevWord?.getNodeText() ?: ""
+
+        return SqlKeywordUtil.isSetLineKeyword(currentText, parentText) ||
+            SqlKeywordUtil.isSetLineKeyword(currentText, prevText)
+    }
+
+    private fun hasConditionLoopAround(parent: SqlNewGroupBlock): Boolean = isFollowedByConditionLoop() || isPrecededByConditionLoop(parent)
 
     private fun isFollowedByConditionLoop(): Boolean = childBlocks.lastOrNull() is SqlElConditionLoopCommentBlock
 
@@ -202,6 +233,9 @@ open class SqlBlock(
             lastPrevBlock.node.psi.startOffset > parent.node.psi.startOffset
     }
 
+    protected fun getLastBlockHasConditionLoopDirective(): SqlElConditionLoopCommentBlock? =
+        (prevBlocks.lastOrNull()?.childBlocks?.firstOrNull() as? SqlElConditionLoopCommentBlock)
+
     /**
      * Creates the indentation length for the block.
      *
@@ -249,158 +283,67 @@ open class SqlBlock(
     /**
      * Creates a spacing builder specifically for directive block comments.
      */
-    protected fun createBlockDirectiveCommentSpacingBuilder(): SqlCustomSpacingBuilder =
-        SqlCustomSpacingBuilder()
-            .withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_ID_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_PRIMARY_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_STRING,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_NUMBER,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.BOOLEAN,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_NULL,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_FIELD_ACCESS_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.EL_STATIC_FIELD_ACCESS_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.HASH,
-                Spacing.createSpacing(0, 0, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_ID_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_PRIMARY_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_STRING,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_NUMBER,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.BOOLEAN,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_NULL,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_FIELD_ACCESS_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.HASH,
-                SqlTypes.EL_STATIC_FIELD_ACCESS_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_START,
-                SqlTypes.CARET,
-                Spacing.createSpacing(0, 0, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
+    protected open fun createBlockDirectiveCommentSpacingBuilder(): SqlCustomSpacingBuilder {
+        val builder = SqlCustomSpacingBuilder()
+
+        // Types that need spacing after BLOCK_COMMENT_START
+        val typesNeedingSpaceAfterStart =
+            listOf(
                 SqlTypes.EL_ID_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.EL_PRIMARY_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.EL_STRING,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.EL_NUMBER,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.BOOLEAN,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.EL_NULL,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.EL_FIELD_ACCESS_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.CARET,
                 SqlTypes.EL_STATIC_FIELD_ACCESS_EXPR,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.BLOCK_COMMENT_CONTENT,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(0, 0, 0, true, 0),
-            ).withSpacing(
-                SqlTypes.EL_FIELD_ACCESS_EXPR,
-                SqlTypes.OTHER,
-                Spacing.createSpacing(1, 1, 0, false, 0),
-            ).withSpacing(
-                SqlTypes.EL_STATIC_FIELD_ACCESS_EXPR,
-                SqlTypes.OTHER,
-                Spacing.createSpacing(1, 1, 0, false, 0),
-            ).withSpacing(
+            )
+
+        // Types that need spacing before BLOCK_COMMENT_END
+        val typesNeedingSpaceBeforeEnd =
+            listOf(
                 SqlTypes.EL_ID_EXPR,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.EL_PRIMARY_EXPR,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.STRING,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.EL_NUMBER,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.EL_NULL,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.BOOLEAN,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.EL_FIELD_ACCESS_EXPR,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
-            ).withSpacing(
                 SqlTypes.EL_STATIC_FIELD_ACCESS_EXPR,
-                SqlTypes.BLOCK_COMMENT_END,
-                Spacing.createSpacing(1, 1, 0, true, 0),
             )
 
+        // Add spacing rules for BLOCK_COMMENT_START
+        typesNeedingSpaceAfterStart.forEach { type ->
+            builder.withSpacing(SqlTypes.BLOCK_COMMENT_START, type, SPACING_ONE)
+        }
+
+        // Special cases for BLOCK_COMMENT_START
+        builder.withSpacing(SqlTypes.BLOCK_COMMENT_START, SqlTypes.HASH, SPACING_ZERO)
+        builder.withSpacing(SqlTypes.BLOCK_COMMENT_START, SqlTypes.CARET, SPACING_ZERO)
+
+        // Add spacing rules for HASH
+        typesNeedingSpaceAfterStart.forEach { type ->
+            builder.withSpacing(SqlTypes.HASH, type, SPACING_ONE)
+        }
+
+        // Add spacing rules for CARET
+        typesNeedingSpaceAfterStart.forEach { type ->
+            builder.withSpacing(SqlTypes.CARET, type, SPACING_ONE)
+        }
+
+        // Special spacing rules
+        builder.withSpacing(SqlTypes.BLOCK_COMMENT_CONTENT, SqlTypes.BLOCK_COMMENT_END, SPACING_ZERO)
+        builder.withSpacing(SqlTypes.EL_FIELD_ACCESS_EXPR, SqlTypes.OTHER, SPACING_ONE_NO_KEEP)
+        builder.withSpacing(SqlTypes.EL_STATIC_FIELD_ACCESS_EXPR, SqlTypes.OTHER, SPACING_ONE_NO_KEEP)
+
+        // Add spacing rules before BLOCK_COMMENT_END
+        typesNeedingSpaceBeforeEnd.forEach { type ->
+            builder.withSpacing(type, SqlTypes.BLOCK_COMMENT_END, SPACING_ONE)
+        }
+
+        return builder
+    }
+
     /**
      * Returns the child indentation for the block.
      *
@@ -413,10 +356,6 @@ open class SqlBlock(
             Indent.getSpaceIndent(0)
         }
 
-    companion object {
-        private const val DEFAULT_INDENT_SIZE = 4
-    }
-
     /**
      * Determines whether the block is a leaf node.
      *