Support for directive definition items in reference resolution

Author: xterao

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

@@ -34,29 +34,51 @@ abstract class SqlElChildElementValidator(
     open val blocks: List<PsiElement>,
     open val shorName: String,
 ) {
-    abstract fun validateChildren(): ValidationResult?
+    abstract fun validateChildren(dropIndex: Int = 0): ValidationResult?
 
     open fun validateChildren(
+        dropIndex: Int = 0,
         findFieldMethod: (PsiType) -> PsiParentClass = { type: PsiType -> PsiParentClass(type) },
         complete: (PsiParentClass) -> Unit,
     ): ValidationResult? = null
 
     protected fun validateFieldAccess(
         topParent: PsiParentClass,
+        dropLastIndex: Int = 0,
+        findFieldMethod: ((PsiType) -> PsiParentClass)? = { type -> PsiParentClass(type) },
+        complete: ((PsiParentClass) -> Unit) = { parent: PsiParentClass? -> },
+    ): ValidationResult? =
+        getFieldAccessParentClass(
+            topParent,
+            dropLastIndex,
+            findFieldMethod = findFieldMethod,
+            complete = complete,
+        )
+
+    protected fun getFieldAccessParentClass(
+        topParent: PsiParentClass,
+        dropLastIndex: Int = 0,
         findFieldMethod: ((PsiType) -> PsiParentClass)? = { type -> PsiParentClass(type) },
         complete: ((PsiParentClass) -> Unit) = { parent: PsiParentClass? -> },
     ): ValidationResult? {
         var parent = topParent
         var competeResult: ValidationCompleteResult? = null
 
+        if (dropLastIndex > 0 && blocks.drop(1).dropLast(dropLastIndex).isEmpty()) {
+            return ValidationCompleteResult(
+                blocks.last(),
+                parent,
+            )
+        }
+
         var getMethodReturnType: PsiType? =
             if (PsiClassTypeUtil.isCollect(topParent.type)) {
                 topParent.type
             } else {
                 null
             }
         var listParamIndex = 0
-        for (element in blocks.drop(1)) {
+        for (element in blocks.drop(1).dropLast(dropLastIndex)) {
             val searchElm = cleanString(getSearchElementText(element))
             if (searchElm.isEmpty()) {
                 complete.invoke(parent)

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

@@ -21,6 +21,7 @@ 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.ValidationCompleteResult
 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
@@ -34,33 +35,33 @@ class SqlElFieldAccessorChildElementValidator(
     private val topDaoParameter: PsiParameter? = null,
 ) : SqlElChildElementValidator(blocks, shorName) {
     override fun validateChildren(
+        dropIndex: Int,
         findFieldMethod: (PsiType) -> PsiParentClass,
         complete: (PsiParentClass) -> Unit,
     ): 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 parentClassResult = getParentClass()
+        val parentClass = parentClassResult?.parentClass
+        if (parentClassResult is ValidationCompleteResult && parentClass != null) {
+            return validateFieldAccess(
+                parentClass,
+                complete = complete,
             )
         }
 
-        val parentClass =
-            validDaoParam.getIterableClazz(daoMethod.getDomaAnnotationType())
+        return parentClassResult
+    }
 
-        return validateFieldAccess(
-            parentClass,
-            complete = complete,
-        )
+    override fun validateChildren(dropIndex: Int): ValidationResult? {
+        val parentClassResult = getParentClass()
+        val parentClass = parentClassResult?.parentClass
+        if (parentClassResult is ValidationCompleteResult && parentClass != null) {
+            return validateFieldAccess(parentClass)
+        }
+
+        return parentClassResult
     }
 
-    override fun validateChildren(): ValidationResult? {
+    private fun getParentClass(): ValidationResult? {
         val daoMethod = findDaoMethod(file) ?: return null
         val topElement: PsiElement = blocks.firstOrNull() ?: return null
         val validDaoParam =
@@ -75,9 +76,9 @@ class SqlElFieldAccessorChildElementValidator(
             )
         }
 
-        val parentClass =
-            validDaoParam.getIterableClazz(daoMethod.getDomaAnnotationType())
-
-        return validateFieldAccess(parentClass)
+        return ValidationCompleteResult(
+            topElement,
+            validDaoParam.getIterableClazz(daoMethod.getDomaAnnotationType()),
+        )
     }
 }

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

@@ -26,13 +26,15 @@ class SqlElForItemFieldAccessorChildElementValidator(
     override val shorName: String = "",
 ) : SqlElChildElementValidator(blocks, shorName) {
     override fun validateChildren(
+        dropIndex: Int,
         findFieldMethod: (PsiType) -> PsiParentClass,
         complete: (PsiParentClass) -> Unit,
     ): ValidationResult? =
         validateFieldAccess(
             declarationType,
+            dropIndex,
             complete = complete,
         )
 
-    override fun validateChildren(): ValidationResult? = validateFieldAccess(declarationType)
+    override fun validateChildren(dropIndex: Int): ValidationResult? = validateFieldAccess(declarationType, dropIndex)
 }

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

@@ -35,6 +35,7 @@ class SqlElStaticFieldAccessorChildElementValidator(
     override val shorName: String,
 ) : SqlElChildElementValidator(blocks, shorName) {
     override fun validateChildren(
+        dropIndex: Int,
         findFieldMethod: (PsiType) -> PsiParentClass,
         complete: (PsiParentClass) -> Unit,
     ): ValidationResult? {
@@ -52,7 +53,7 @@ class SqlElStaticFieldAccessorChildElementValidator(
         }
     }
 
-    override fun validateChildren(): ValidationResult? {
+    override fun validateChildren(dropIndex: Int): ValidationResult? {
         val getParentResult = getParent()
         when (getParentResult) {
             is ValidationCompleteResult -> {

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

@@ -265,9 +265,29 @@ class SqlParameterCompletionProvider : CompletionProvider<CompletionParameters>(
         }
 
         if (topElementType == null) {
-            topElementType = getTopElementClassType(top, elements, originalFile, topText, result) ?: return
+            val forDeclaration = ForDirectiveInspection("")
+            val forItem = forDeclaration.getForItem(top)
+            if (forItem != null) {
+                val errorElement = forDeclaration.checkForItem(elements)
+                if (errorElement is ValidationCompleteResult) {
+                    topElementType = errorElement.parentClass.type
+                    val parentClass = errorElement.parentClass
+                    val searchWord = cleanString(positionText)
+                    setFieldsAndMethodsCompletionResultSet(
+                        parentClass.searchField(searchWord)?.toTypedArray() ?: emptyArray(),
+                        parentClass.searchMethod(searchWord)?.toTypedArray() ?: emptyArray(),
+                        result,
+                    )
+                    return
+                }
+                topElementType =
+                    getElementTypeByFieldAccess(originalFile, topText, elements, result)
+                        ?: return
+            }
         }
 
+        if (topElementType == null) return
+
         val fieldAccessorChildElementValidator =
             SqlElForItemFieldAccessorChildElementValidator(
                 elements,
@@ -304,25 +324,6 @@ class SqlParameterCompletionProvider : CompletionProvider<CompletionParameters>(
         }
     }
 
-    private fun getTopElementClassType(
-        top: PsiElement,
-        elements: List<PsiElement>,
-        originalFile: PsiFile,
-        topText: String,
-        result: CompletionResultSet,
-    ): PsiType? {
-        val forDeclaration = ForDirectiveInspection("")
-        val forItem = forDeclaration.getForItem(top)
-        if (forItem != null) {
-            val errorElement = forDeclaration.checkForItem(elements)
-            if (errorElement is ValidationCompleteResult) {
-                return errorElement.parentClass.type
-            }
-            return null
-        }
-        return getElementTypeByFieldAccess(originalFile, topText, elements, result)
-    }
-
     private fun setStaticFieldAccess(
         top: PsiElement,
         prevWord: String,

src/main/kotlin/org/domaframework/doma/intellij/inspection/ForDirectiveInspection.kt

@@ -97,6 +97,32 @@ class ForDirectiveInspection(
         return errorElement
     }
 
+    fun getFieldAccessParentClass(blockElements: List<PsiElement>): ValidationResult? {
+        val targetElement: PsiElement = blockElements.firstOrNull() ?: return null
+        val file = targetElement.containingFile ?: return null
+
+        val forItem = getForItem(targetElement)
+        var errorElement: ValidationResult? = ValidationDaoParamResult(targetElement, "", shorName)
+        if (forItem != null) {
+            val declarationItem =
+                getDeclarationItem(forItem, file)
+
+            if (declarationItem != null && declarationItem is ForDeclarationDaoBaseItem) {
+                val forItemElementsParentClass = declarationItem.getPsiParentClass()
+                if (forItemElementsParentClass != null) {
+                    val validator =
+                        SqlElForItemFieldAccessorChildElementValidator(
+                            blockElements,
+                            forItemElementsParentClass,
+                            shorName,
+                        )
+                    errorElement = validator.validateChildren(dropIndex = 1)
+                }
+            }
+        }
+        return errorElement
+    }
+
     /**
      * Count the `for`, `if`, and `end` elements from the beginning
      * to the target element (`targetElement`)

src/main/kotlin/org/domaframework/doma/intellij/reference/SqlElForDirectiveIdExprReference.kt

@@ -61,11 +61,18 @@ class SqlElForDirectiveIdExprReference(
 
     override fun getVariants(): Array<Any> = emptyArray()
 
+    /**
+     * In the for directive, set the reference on the left side to the element on the right side.
+     */
     private fun getDeclarationItem(): PsiElement? {
-        val forDirectiveParent = PsiTreeUtil.getParentOfType(element, SqlElForDirective::class.java)
-        if (forDirectiveParent == null) return null
+        val forDirectiveParent =
+            PsiTreeUtil.getParentOfType(element, SqlElForDirective::class.java) ?: return null
 
-        return PsiTreeUtil.getChildrenOfType(forDirectiveParent, SqlElFieldAccessExpr::class.java)?.last()
-            ?: PsiTreeUtil.getChildrenOfType(forDirectiveParent, SqlElFieldAccessExpr::class.java)?.last()
+        return PsiTreeUtil
+            .getChildrenOfType(forDirectiveParent, SqlElFieldAccessExpr::class.java)
+            ?.last()
+            ?: PsiTreeUtil
+                .getChildrenOfType(forDirectiveParent, SqlElFieldAccessExpr::class.java)
+                ?.last()
     }
 }