Refactor: Use validator classes in code completion

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/common/sql/PsiClassTypeUtil.kt

@@ -15,8 +15,10 @@
  */
 package org.domaframework.doma.intellij.common.sql
 
+import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiType
+import com.intellij.psi.search.GlobalSearchScope
 
 class PsiClassTypeUtil {
     companion object {
@@ -29,5 +31,45 @@ class PsiClassTypeUtil {
             }
             return classType
         }
+
+        fun isIterableType(
+            type: PsiClassType,
+            project: Project,
+        ): Boolean {
+            val iterableType =
+                PsiType.getTypeByName(
+                    "java.lang.Iterable",
+                    project,
+                    GlobalSearchScope.allScope(project),
+                )
+            return iterableType.isAssignableFrom(type)
+        }
+
+        fun getParameterType(
+            project: Project,
+            baseType: PsiType?,
+            preReturnListType: PsiType,
+            index: Int,
+        ): PsiClassType? {
+            val returnType = baseType as? PsiClassType ?: return null
+            val preReturnType = preReturnListType as? PsiClassType ?: return null
+
+            if (returnType.name == "E" && isIterableType(preReturnListType, project)) {
+                var count = 1
+                var type: PsiType? = preReturnType.parameters.firstOrNull()
+                while (index > count && type != null && type is PsiClassType) {
+                    type = type.parameters.firstOrNull()
+                    count++
+                }
+                return type as? PsiClassType
+            }
+            return null
+        }
+
+        fun isCollect(type: PsiType): Boolean =
+            type.superTypes
+                .firstOrNull()
+                ?.presentableText
+                ?.startsWith("Collection") == true
     }
 }

src/main/kotlin/org/domaframework/doma/intellij/common/sql/foritem/ForDeclarationDaoBaseItem.kt

@@ -16,11 +16,12 @@
 package org.domaframework.doma.intellij.common.sql.foritem
 
 import com.intellij.psi.PsiClassType
-import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiParameter
 import org.domaframework.doma.intellij.common.psi.PsiParentClass
+import org.domaframework.doma.intellij.common.sql.PsiClassTypeUtil
+import org.domaframework.doma.intellij.common.sql.validator.SqlElFieldAccessorChildElementValidator
 import org.domaframework.doma.intellij.extension.psi.DomaAnnotationType
-import org.domaframework.doma.intellij.extension.psi.getIterableClazz
+import org.domaframework.doma.intellij.psi.SqlElIdExpr
 
 /**
  * In the For directive, define the element based on the Dao parameter.
@@ -29,21 +30,39 @@ import org.domaframework.doma.intellij.extension.psi.getIterableClazz
  *    ...
  */
 class ForDeclarationDaoBaseItem(
-    override val element: PsiElement,
-    val daoParameter: PsiParameter,
+    private val blocks: List<SqlElIdExpr>,
     val nestIndex: Int,
     val domaAnnotationType: DomaAnnotationType,
-) : ForDeclarationItem(element) {
+    val daoParameter: PsiParameter? = null,
+) : ForDeclarationItem(blocks.first()) {
     fun getPsiParentClass(): PsiParentClass? {
+        val daoParamFieldValidator =
+            SqlElFieldAccessorChildElementValidator(
+                blocks,
+                element.containingFile,
+                "",
+                daoParameter,
+            )
+
+        var lastType: PsiParentClass? = null
+        val errorElement =
+            daoParamFieldValidator.validateChildren(
+                complete = { parent ->
+                    lastType = parent
+                },
+            )
+        if (errorElement != null || lastType == null) return null
+
+        val topElm = blocks.first()
+        (lastType.type as? PsiClassType)?.let { if (!PsiClassTypeUtil.isIterableType(it, topElm.project)) return null }
+
+        var nestClassType: PsiClassType? = (lastType.type as? PsiClassType)
         var i = 0
-        var nestClassType: PsiClassType? =
-            (daoParameter.getIterableClazz(domaAnnotationType.isBatchAnnotation()).type as PsiClassType)
         if (domaAnnotationType.isBatchAnnotation()) {
-            nestClassType?.parameters?.firstOrNull() as PsiClassType?
+            nestClassType?.parameters?.firstOrNull() as? PsiClassType?
         }
-
-        while (i <= nestIndex) {
-            nestClassType = nestClassType?.parameters?.firstOrNull() as PsiClassType?
+        while (nestClassType != null && i <= nestIndex && PsiClassTypeUtil.isIterableType(nestClassType, topElm.project)) {
+            nestClassType = nestClassType.parameters.firstOrNull() as? PsiClassType?
             i++
         }
         return nestClassType?.let { PsiParentClass(it) }

src/main/kotlin/org/domaframework/doma/intellij/common/sql/foritem/ForDeclarationItem.kt

@@ -18,6 +18,7 @@ package org.domaframework.doma.intellij.common.sql.foritem
 import com.intellij.psi.PsiElement
 import org.domaframework.doma.intellij.extension.expr.accessElements
 import org.domaframework.doma.intellij.psi.SqlElFieldAccessExpr
+import org.domaframework.doma.intellij.psi.SqlElIdExpr
 
 /**
  *  definition source in the For directive
@@ -26,7 +27,7 @@ import org.domaframework.doma.intellij.psi.SqlElFieldAccessExpr
 open class ForDeclarationItem(
     override val element: PsiElement,
 ) : ForDirectiveItemBase(element) {
-    fun getDeclarationChildren(): List<PsiElement> =
-        (element as? SqlElFieldAccessExpr)?.accessElements?.mapNotNull { it as PsiElement }
-            ?: listOf(element)
+    fun getDeclarationChildren(): List<SqlElIdExpr> =
+        (element as? SqlElFieldAccessExpr)?.accessElements?.mapNotNull { it as SqlElIdExpr }
+            ?: listOf(element as SqlElIdExpr)
 }

src/main/kotlin/org/domaframework/doma/intellij/common/sql/validator/SqlElChildElementValidator.kt

@@ -16,49 +16,117 @@
 package org.domaframework.doma.intellij.common.sql.validator
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
+import com.intellij.psi.util.elementType
 import org.domaframework.doma.intellij.common.psi.PsiParentClass
+import org.domaframework.doma.intellij.common.sql.PsiClassTypeUtil
+import org.domaframework.doma.intellij.common.sql.cleanString
+import org.domaframework.doma.intellij.common.sql.validator.result.ValidationCompleteResult
 import org.domaframework.doma.intellij.common.sql.validator.result.ValidationPropertyResult
 import org.domaframework.doma.intellij.common.sql.validator.result.ValidationResult
 import org.domaframework.doma.intellij.extension.psi.getMethodReturnType
+import org.domaframework.doma.intellij.extension.psi.getParameterType
+import org.domaframework.doma.intellij.psi.SqlElIdExpr
+import org.domaframework.doma.intellij.psi.SqlElParameters
+import org.domaframework.doma.intellij.psi.SqlTypes
 
 abstract class SqlElChildElementValidator(
     open val blocks: List<PsiElement>,
     open val shorName: String,
 ) {
     abstract fun validateChildren(): ValidationResult?
 
-    protected fun validateFieldAccess(topParent: PsiParentClass): ValidationResult? {
+    open fun validateChildren(
+        findFieldMethod: (PsiType) -> PsiParentClass = { type: PsiType -> PsiParentClass(type) },
+        complete: (PsiParentClass) -> Unit,
+    ): ValidationResult? = null
+
+    protected fun validateFieldAccess(
+        topParent: PsiParentClass,
+        findFieldMethod: ((PsiType) -> PsiParentClass)? = { type -> PsiParentClass(type) },
+        complete: ((PsiParentClass) -> Unit) = { parent: PsiParentClass? -> },
+    ): ValidationResult? {
         var parent = topParent
-        val topElementType = topParent.type
-        var listTypeSearchIndex = 0
-        for (eml in blocks.drop(1)) {
-            var isExist = false
-            parent
-                .findField(eml.text)
-                ?.let { match ->
-                    isExist = true
-                    parent = PsiParentClass(match.type)
-                }
-            if (isExist) continue
-            parent
-                .findMethod(eml.text)
-                ?.let { match ->
-                    val returnType =
-                        match.getMethodReturnType(topElementType, listTypeSearchIndex)
-                            ?: return null
-                    isExist = true
-                    parent = PsiParentClass(returnType)
-                }
-            listTypeSearchIndex++
-            if (!isExist) {
+        var competeResult: ValidationCompleteResult? = null
+
+        var getMethodReturnType: PsiType? =
+            if (PsiClassTypeUtil.isCollect(topParent.type)) {
+                topParent.type
+            } else {
+                null
+            }
+        var listParamIndex = 0
+        for (element in blocks.drop(1)) {
+            val searchElm = cleanString(getSearchElementText(element))
+            if (searchElm.isEmpty()) {
+                complete.invoke(parent)
+                return ValidationCompleteResult(
+                    element,
+                    parent,
+                )
+            }
+
+            val field =
+                parent
+                    .findField(element.text)
+                    ?.let { match ->
+                        val type = match.type
+                        val methodReturnType =
+                            getMethodReturnType?.let { match.getParameterType(it, type, listParamIndex) }
+                                ?: type
+                        if (PsiClassTypeUtil.isCollect(type)) {
+                            getMethodReturnType = type
+                            listParamIndex = 0
+                        }
+                        findFieldMethod?.invoke(methodReturnType)
+                    }
+            val method =
+                parent
+                    .findMethod(element.text)
+                    ?.let { match ->
+                        val returnType = match.returnType ?: return null
+                        val methodReturnType =
+                            getMethodReturnType?.let { match.getMethodReturnType(it, listParamIndex) }
+                                ?: returnType
+                        if (PsiClassTypeUtil.isCollect(returnType)) {
+                            getMethodReturnType = returnType
+                            listParamIndex = 0
+                        }
+                        findFieldMethod?.invoke(methodReturnType)
+                    }
+            listParamIndex++
+            if (field == null && method == null) {
                 return ValidationPropertyResult(
-                    eml,
+                    element,
                     parent,
                     shorName,
-                    eml.textRange,
                 )
+            } else {
+                if (field != null && element.nextSibling !is SqlElParameters) {
+                    parent = field
+                    competeResult =
+                        ValidationCompleteResult(
+                            element,
+                            parent,
+                        )
+                } else if (method != null) {
+                    parent = method
+                    competeResult =
+                        ValidationCompleteResult(
+                            element,
+                            parent,
+                        )
+                }
             }
         }
-        return null
+        complete.invoke(parent)
+        return competeResult
     }
+
+    private fun getSearchElementText(elm: PsiElement): String =
+        if (elm is SqlElIdExpr || elm.elementType == SqlTypes.EL_IDENTIFIER) {
+            elm.text
+        } else {
+            ""
+        }
 }

src/main/kotlin/org/domaframework/doma/intellij/common/sql/validator/SqlElFieldAccessorChildElementValidator.kt

@@ -17,7 +17,10 @@ package org.domaframework.doma.intellij.common.sql.validator
 
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiParameter
+import com.intellij.psi.PsiType
 import org.domaframework.doma.intellij.common.dao.findDaoMethod
+import org.domaframework.doma.intellij.common.psi.PsiParentClass
 import org.domaframework.doma.intellij.common.sql.validator.result.ValidationDaoParamResult
 import org.domaframework.doma.intellij.common.sql.validator.result.ValidationResult
 import org.domaframework.doma.intellij.extension.psi.findParameter
@@ -27,25 +30,54 @@ import org.domaframework.doma.intellij.extension.psi.getIterableClazz
 class SqlElFieldAccessorChildElementValidator(
     override val blocks: List<PsiElement>,
     private val file: PsiFile,
-    override val shorName: String,
+    override val shorName: String = "",
+    private val topDaoParameter: PsiParameter? = null,
 ) : SqlElChildElementValidator(blocks, shorName) {
-    override fun validateChildren(): ValidationResult? {
+    override fun validateChildren(
+        findFieldMethod: (PsiType) -> PsiParentClass,
+        complete: (PsiParentClass) -> Unit,
+    ): ValidationResult? {
         val daoMethod = findDaoMethod(file) ?: return null
         val topElement: PsiElement = blocks.firstOrNull() ?: return null
         val validDaoParam =
-            daoMethod.findParameter(topElement.text)
-                ?: return ValidationDaoParamResult(
-                    topElement,
-                    daoMethod.name,
-                    shorName,
-                    topElement.textRange,
-                )
+            topDaoParameter
+                ?: daoMethod.findParameter(topElement.text)
+
+        if (validDaoParam == null) {
+            return ValidationDaoParamResult(
+                topElement,
+                daoMethod.name,
+                shorName,
+            )
+        }
 
         val parentClass =
             validDaoParam.getIterableClazz(daoMethod.getDomaAnnotationType())
 
         return validateFieldAccess(
             parentClass,
+            complete = complete,
         )
     }
+
+    override fun validateChildren(): ValidationResult? {
+        val daoMethod = findDaoMethod(file) ?: return null
+        val topElement: PsiElement = blocks.firstOrNull() ?: return null
+        val validDaoParam =
+            topDaoParameter
+                ?: daoMethod.findParameter(topElement.text)
+
+        if (validDaoParam == null) {
+            return ValidationDaoParamResult(
+                topElement,
+                daoMethod.name,
+                shorName,
+            )
+        }
+
+        val parentClass =
+            validDaoParam.getIterableClazz(daoMethod.getDomaAnnotationType())
+
+        return validateFieldAccess(parentClass)
+    }
 }

src/main/kotlin/org/domaframework/doma/intellij/common/sql/validator/SqlElForItemFieldAccessorChildElementValidator.kt

@@ -16,13 +16,23 @@
 package org.domaframework.doma.intellij.common.sql.validator
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
 import org.domaframework.doma.intellij.common.psi.PsiParentClass
 import org.domaframework.doma.intellij.common.sql.validator.result.ValidationResult
 
 class SqlElForItemFieldAccessorChildElementValidator(
     override val blocks: List<PsiElement>,
     private val declarationType: PsiParentClass,
-    override val shorName: String,
+    override val shorName: String = "",
 ) : SqlElChildElementValidator(blocks, shorName) {
+    override fun validateChildren(
+        findFieldMethod: (PsiType) -> PsiParentClass,
+        complete: (PsiParentClass) -> Unit,
+    ): ValidationResult? =
+        validateFieldAccess(
+            declarationType,
+            complete = complete,
+        )
+
     override fun validateChildren(): ValidationResult? = validateFieldAccess(declarationType)
 }