Type resolution for elements defined in the for directive

Author: xterao

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright Doma Tools Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.domaframework.doma.intellij.common.sql
+
+import com.intellij.psi.PsiClassType
+import com.intellij.psi.PsiType
+
+class PsiClassTypeUtil {
+    companion object {
+        fun getPsiTypeByList(
+            classType: PsiClassType,
+            useListParam: Boolean,
+        ): PsiType? {
+            if (classType.className == "List" && useListParam) {
+                return classType.parameters.firstOrNull()
+            }
+            return classType
+        }
+    }
+}

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

@@ -0,0 +1,51 @@
+/*
+ * Copyright Doma Tools Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+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.extension.psi.DomaAnnotationType
+import org.domaframework.doma.intellij.extension.psi.getIterableClazz
+
+/**
+ * In the For directive, define the element based on the Dao parameter.
+ * %for [ForItem] : [ForDeclarationDaoBaseItem]
+ *   %for [ForItem] : [ForDeclarationItem]
+ *    ...
+ */
+class ForDeclarationDaoBaseItem(
+    override val element: PsiElement,
+    val daoParameter: PsiParameter,
+    val nestIndex: Int,
+    val domaAnnotationType: DomaAnnotationType,
+) : ForDeclarationItem(element) {
+    fun getPsiParentClass(): PsiParentClass? {
+        var i = 0
+        var nestClassType: PsiClassType? =
+            (daoParameter.getIterableClazz(domaAnnotationType.isBatchAnnotation()).type as PsiClassType)
+        if (domaAnnotationType.isBatchAnnotation()) {
+            nestClassType?.parameters?.firstOrNull() as PsiClassType?
+        }
+
+        while (i <= nestIndex) {
+            nestClassType = nestClassType?.parameters?.firstOrNull() as PsiClassType?
+            i++
+        }
+        return nestClassType?.let { PsiParentClass(it) }
+    }
+}

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright Doma Tools Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.domaframework.doma.intellij.common.sql.foritem
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import org.domaframework.doma.intellij.psi.SqlElFieldAccessExpr
+import org.domaframework.doma.intellij.psi.SqlElIdExpr
+
+/**
+ *  definition source in the For directive
+ *  %for [ForItem] : [ForDeclarationItem]
+ */
+open class ForDeclarationItem(
+    override val element: PsiElement,
+) : ForDirectiveItemBase(element) {
+    fun getDeclarationChildren(): List<PsiElement> =
+        if (element is SqlElFieldAccessExpr) {
+            PsiTreeUtil.getChildrenOfType(element, SqlElIdExpr::class.java)?.toList() ?: emptyList()
+        } else {
+            listOf(element)
+        }
+}

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright Doma Tools Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.domaframework.doma.intellij.common.sql.foritem
+
+import com.intellij.psi.PsiElement
+
+abstract class ForDirectiveItemBase(
+    open val element: PsiElement,
+)

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright Doma Tools Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.domaframework.doma.intellij.common.sql.foritem
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import org.domaframework.doma.intellij.psi.SqlElForDirective
+
+/**
+ * Elements defined within the For directive
+ * %for [ForItem] : [ForDeclarationItem]
+ */
+class ForItem(
+    override val element: PsiElement,
+) : ForDirectiveItemBase(element) {
+    fun getParentForDirectiveExpr() = PsiTreeUtil.getParentOfType(element, SqlElForDirective::class.java)
+}

src/main/kotlin/org/domaframework/doma/intellij/extension/psi/PsiMethodExtension.kt

@@ -23,7 +23,6 @@ import com.intellij.psi.PsiType
 import com.intellij.psi.PsiTypeParameterList
 import com.intellij.psi.impl.compiled.ClsClassImpl
 import com.intellij.psi.impl.compiled.ClsTypeParametersListImpl
-import com.intellij.psi.impl.source.PsiClassReferenceType
 
 fun PsiMethod.findParameter(searchName: String): PsiParameter? = this.methodParameters.firstOrNull { it.name == searchName }
 
@@ -32,6 +31,7 @@ val PsiMethod.methodParameters: List<PsiParameter>
 
 fun PsiMethod.searchParameter(searchName: String): List<PsiParameter> = this.methodParameters.filter { it.name.startsWith(searchName) }
 
+@OptIn(ExperimentalStdlibApi::class)
 fun PsiMethod.getDomaAnnotationType(): DomaAnnotationType {
     DomaAnnotationType.entries.forEach {
         if (AnnotationUtil.findAnnotation(this, it.fqdn) != null) {
@@ -49,14 +49,16 @@ fun PsiMethod.getMethodReturnType(
     topElementType: PsiType,
     index: Int,
 ): PsiClassType? {
+    // TODO:When handling properties that return a List type from a field or method,
+    //  the defined type cannot be obtained, so it is not supported.
     val returnType = this.returnType as? PsiClassType
     val cls = returnType?.resolve()?.parent as? ClsTypeParametersListImpl
     val listType = ((cls as? PsiTypeParameterList)?.parent as? ClsClassImpl)
 
     if (returnType?.name == "E" && listType?.qualifiedName == "java.util.List") {
         var count = 1
-        var type: PsiType? = (topElementType as? PsiClassReferenceType)?.parameters?.firstOrNull()
-        while (index >= count && type != null && type is PsiClassReferenceType) {
+        var type: PsiType? = (topElementType as? PsiClassType)?.parameters?.firstOrNull()
+        while (index >= count && type != null && type is PsiClassType) {
             type = type.parameters.firstOrNull()
             count++
         }

src/main/kotlin/org/domaframework/doma/intellij/extension/psi/PsiParameterExtension.kt

@@ -15,42 +15,41 @@
  */
 package org.domaframework.doma.intellij.extension.psi
 
+import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiParameter
-import com.intellij.psi.impl.source.PsiClassReferenceType
 import org.domaframework.doma.intellij.common.psi.PsiParentClass
+import org.domaframework.doma.intellij.common.sql.PsiClassTypeUtil
 
 /**
  * For List type, if the annotation type is Batch type,
  * return the content type.
  */
-fun PsiParameter.getIterableClazz(annotationType: DomaAnnotationType): PsiParentClass {
-    val immediate: PsiClassReferenceType? = this.type as? PsiClassReferenceType
-    if (immediate != null) {
-        if (immediate.name == "List" && annotationType.isBatchAnnotation()) {
-            val listType =
-                (this.type as PsiClassReferenceType).parameters.firstOrNull()
-                    ?: return PsiParentClass(this.type)
-            return PsiParentClass(listType)
-        }
+fun PsiParameter.getIterableClazz(annotationType: DomaAnnotationType): PsiParentClass = getIterableClazz(annotationType.isBatchAnnotation())
+
+fun PsiParameter.getIterableClazz(useListParam: Boolean): PsiParentClass {
+    val immediate = this.type as? PsiClassType
+    val classType = immediate?.let { PsiClassTypeUtil.getPsiTypeByList(it, useListParam) }
+    if (classType != null) {
+        return PsiParentClass(classType)
     }
     return PsiParentClass(this.type)
 }
 
 val PsiParameter.isFunctionClazz: Boolean
     get() =
-        (this.typeElement?.type as? PsiClassReferenceType)
+        (this.typeElement?.type as? PsiClassType)
             ?.resolve()
             ?.qualifiedName
             ?.contains("java.util.function") == true
 
 val PsiParameter.isSelectOption: Boolean
     get() =
-        (this.typeElement?.type as? PsiClassReferenceType)
+        (this.typeElement?.type as? PsiClassType)
             ?.resolve()
             ?.qualifiedName == "org.seasar.doma.jdbc.SelectOptions"
 
 val PsiParameter.isCollector: Boolean
     get() =
-        (this.typeElement?.type as? PsiClassReferenceType)
+        (this.typeElement?.type as? PsiClassType)
             ?.resolve()
             ?.qualifiedName == "java.util.stream.Collector"

src/main/kotlin/org/domaframework/doma/intellij/extension/psi/SqlElForDirectiveExtension.kt

@@ -17,9 +17,18 @@ package org.domaframework.doma.intellij.extension.psi
 
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.PsiTreeUtil
+import org.domaframework.doma.intellij.common.sql.foritem.ForDeclarationItem
+import org.domaframework.doma.intellij.psi.SqlElFieldAccessExpr
 import org.domaframework.doma.intellij.psi.SqlElForDirective
 import org.domaframework.doma.intellij.psi.SqlElIdExpr
 
 fun SqlElForDirective.getForItem(): PsiElement? =
     PsiTreeUtil
         .getChildOfType(this, SqlElIdExpr::class.java)
+
+fun SqlElForDirective.getForItemDeclaration(): ForDeclarationItem? {
+    val declarationElm =
+        PsiTreeUtil.getChildrenOfType(this, SqlElFieldAccessExpr::class.java)?.last()
+            ?: PsiTreeUtil.getChildrenOfType(this, SqlElIdExpr::class.java)?.last()
+    return declarationElm?.let { ForDeclarationItem(it) }
+}