Handle "_has_next" and "_index" as specified types

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/common/util/ForDirectiveUtil.kt

@@ -20,6 +20,7 @@ import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
+import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
@@ -225,7 +226,8 @@ class ForDirectiveUtil {
                 var isBatchAnnotation = false
                 val forItemDeclarationBlocks =
                     if (forDirectiveDeclaration.element is SqlElStaticFieldAccessExpr) {
-                        val staticFieldAccessExpr = forDirectiveDeclaration.element as SqlElStaticFieldAccessExpr
+                        val staticFieldAccessExpr =
+                            forDirectiveDeclaration.element as SqlElStaticFieldAccessExpr
                         staticFieldAccessExpr.accessElements
                     } else {
                         forDirectiveDeclaration.getDeclarationChildren()
@@ -234,18 +236,25 @@ class ForDirectiveUtil {
                 // Defined by StaticFieldAccess
                 if (forDirectiveDeclaration.element is SqlElStaticFieldAccessExpr) {
                     val file = topForDirectiveItem.containingFile
-                    val staticFieldAccessExpr = forDirectiveDeclaration.element as SqlElStaticFieldAccessExpr
+                    val staticFieldAccessExpr =
+                        forDirectiveDeclaration.element as SqlElStaticFieldAccessExpr
                     val clazz = staticFieldAccessExpr.elClass
                     val staticElement = PsiStaticElement(clazz.elIdExprList, file)
                     val referenceClazz = staticElement.getRefClazz() ?: return null
 
                     // In the case of staticFieldAccess, the property that is called first is retrieved.
-                    fieldAccessTopParentClass = getStaticFieldAccessTopElementClassType(staticFieldAccessExpr, referenceClazz)
+                    fieldAccessTopParentClass =
+                        getStaticFieldAccessTopElementClassType(
+                            staticFieldAccessExpr,
+                            referenceClazz,
+                        )
                 } else {
                     // Defined by Dao parameter
                     val file = topForDirectiveItem.containingFile ?: return null
                     val daoMethod = findDaoMethod(file) ?: return null
-                    val topElementText = forDirectiveDeclaration.getDeclarationChildren().firstOrNull()?.text ?: return null
+                    val topElementText =
+                        forDirectiveDeclaration.getDeclarationChildren().firstOrNull()?.text
+                            ?: return null
                     isBatchAnnotation = PsiDaoMethod(project, daoMethod).daoType.isBatchAnnotation()
 
                     val matchParam = daoMethod.findParameter(cleanString(topElementText))
@@ -450,5 +459,28 @@ class ForDirectiveUtil {
             } else {
                 ""
             }
+
+        fun resolveForDirectiveClassTypeIfSuffixExists(
+            project: Project,
+            searchName: String,
+        ): PsiType? {
+            if (searchName.endsWith("_has_next")) {
+                return PsiType.getTypeByName(
+                    "java.lang.Boolean",
+                    project,
+                    GlobalSearchScope.allScope(project),
+                )
+            }
+
+            if (searchName.endsWith("_index")) {
+                return PsiType.getTypeByName(
+                    "java.lang.Integer",
+                    project,
+                    GlobalSearchScope.allScope(project),
+                )
+            }
+
+            return null
+        }
     }
 }

src/main/kotlin/org/domaframework/doma/intellij/contributor/sql/provider/SqlParameterCompletionProvider.kt

@@ -419,20 +419,28 @@ class SqlParameterCompletionProvider : CompletionProvider<CompletionParameters>(
         positionText: String,
         result: CompletionResultSet,
     ): Boolean {
+        val searchWord = cleanString(positionText)
         val project = top.project
         val forDirectiveBlocks = ForDirectiveUtil.getForDirectiveBlocks(top)
         ForDirectiveUtil.findForItem(top, forDirectives = forDirectiveBlocks) ?: return false
 
         val forItemClassType = ForDirectiveUtil.getForDirectiveItemClassType(project, forDirectiveBlocks) ?: return false
+        val specifiedClassType = ForDirectiveUtil.resolveForDirectiveClassTypeIfSuffixExists(project, top.text)
+        val topClassType =
+            if (specifiedClassType != null) {
+                PsiParentClass(specifiedClassType)
+            } else {
+                forItemClassType
+            }
+
         val result =
             ForDirectiveUtil.getFieldAccessLastPropertyClassType(
                 elements,
                 project,
-                forItemClassType,
+                topClassType,
                 shortName = "",
                 dropLastIndex = 1,
                 complete = { lastType ->
-                    val searchWord = cleanString(positionText)
                     setFieldsAndMethodsCompletionResultSet(
                         lastType.searchField(searchWord)?.toTypedArray() ?: emptyArray(),
                         lastType.searchMethod(searchWord)?.toTypedArray() ?: emptyArray(),

src/main/kotlin/org/domaframework/doma/intellij/document/generator/DocumentDaoParameterGenerator.kt

@@ -46,7 +46,14 @@ class DocumentDaoParameterGenerator(
 
         var isBatchAnnotation = false
         if (ForDirectiveUtil.findForItem(searchElement, forDirectives = forDirectives) != null) {
-            topParentType = ForDirectiveUtil.getForDirectiveItemClassType(project, forDirectives)
+            val forItemClassType = ForDirectiveUtil.getForDirectiveItemClassType(project, forDirectives)
+            val specifiedClassType = ForDirectiveUtil.resolveForDirectiveClassTypeIfSuffixExists(project, searchElement.text)
+            topParentType =
+                if (specifiedClassType != null) {
+                    PsiParentClass(specifiedClassType)
+                } else {
+                    forItemClassType
+                }
         } else {
             val daoMethod = findDaoMethod(originalElement.containingFile) ?: return
             val param = daoMethod.findParameter(originalElement.text) ?: return

src/main/kotlin/org/domaframework/doma/intellij/inspection/sql/visitor/SqlInspectionVisitor.kt

@@ -135,7 +135,13 @@ class SqlInspectionVisitor(
                     errorHighlight(topElement, daoMethod, holder)
                     return
                 }
-                result
+                val specifiedClassType =
+                    ForDirectiveUtil.resolveForDirectiveClassTypeIfSuffixExists(project, topElement.text)
+                if (specifiedClassType != null) {
+                    PsiParentClass(specifiedClassType)
+                } else {
+                    result
+                }
             } else {
                 val paramType = daoMethod.findParameter(cleanString(topElement.text))?.type
                 if (paramType == null) {