Merge pull request #20 from domaframework/fix/reference-checks-for-repeated-elements-for-directive

Fix: Reference checks for repeated elements in the for directive do not work correctly

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/action/sql/JumpToDaoFromSQLAction.kt

@@ -37,7 +37,7 @@ class JumpToDaoFromSQLAction : AnAction() {
         currentFile = e.getData(CommonDataKeys.PSI_FILE) ?: return
         findDaoMethod(currentFile!!) ?: return
         e.presentation.isEnabledAndVisible =
-            isSupportFileType(currentFile!!.virtualFile)
+            isSupportFileType(currentFile!!)
     }
 
     override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.BGT
@@ -53,7 +53,7 @@ class JumpToDaoFromSQLAction : AnAction() {
         )
 
         val project = e.project ?: return
-        val daoFile = findDaoFile(project, currentFile?.virtualFile!!) ?: return
+        val daoFile = findDaoFile(project, currentFile!!) ?: return
 
         jumpToDaoMethod(project, currentFile?.virtualFile?.nameWithoutExtension!!, daoFile)
         PluginLoggerUtil.countLoggingByAction(

src/main/kotlin/org/domaframework/doma/intellij/action/sql/JumpToDeclarationFromSqlAction.kt

@@ -15,12 +15,14 @@
  */
 package org.domaframework.doma.intellij.action.sql
 
+import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.actionSystem.ActionUpdateThread
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiLiteralExpression
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiType
 import com.intellij.psi.tree.IFileElementType
@@ -35,7 +37,6 @@ import org.domaframework.doma.intellij.common.psi.PsiStaticElement
 import org.domaframework.doma.intellij.extension.psi.findParameter
 import org.domaframework.doma.intellij.extension.psi.getDomaAnnotationType
 import org.domaframework.doma.intellij.extension.psi.getIterableClazz
-import org.domaframework.doma.intellij.extension.psi.initPsiFileAndElement
 import org.domaframework.doma.intellij.extension.psi.isNotWhiteSpace
 import org.domaframework.doma.intellij.extension.psi.methodParameters
 import org.domaframework.doma.intellij.psi.SqlElClass
@@ -54,15 +55,23 @@ class JumpToDeclarationFromSqlAction : AnAction() {
     override fun update(e: AnActionEvent) {
         e.presentation.isEnabledAndVisible = false
         val editor = e.getData(CommonDataKeys.EDITOR) ?: return
-        val caretOffset = editor.caretModel.offset
+        val caretOffset = editor.caretModel.primaryCaret.selectionStart
 
         currentFile = e.getData(CommonDataKeys.PSI_FILE) ?: return
         currentFile?.let { element = it.findElementAt(caretOffset) ?: return }
         if (element == null) return
 
         val project = element?.project ?: return
         if (isJavaOrKotlinFileType(currentFile ?: return)) {
-            currentFile = currentFile!!.initPsiFileAndElement(project, caretOffset)
+            val injectedLanguageManager =
+                InjectedLanguageManager.getInstance(project)
+            val literal = PsiTreeUtil.getParentOfType(element, PsiLiteralExpression::class.java)
+            currentFile =
+                injectedLanguageManager
+                    .getInjectedPsiFiles(literal!!)
+                    ?.firstOrNull()
+                    ?.first as? PsiFile
+            element = currentFile?.findElementAt(countInjectionOffset(literal, caretOffset))
         }
         findDaoMethod(currentFile ?: return) ?: return
 
@@ -77,6 +86,31 @@ class JumpToDeclarationFromSqlAction : AnAction() {
         if (targetElement.isNotEmpty()) e.presentation.isEnabledAndVisible = true
     }
 
+    private fun countInjectionOffset(
+        literal: PsiLiteralExpression,
+        caretOffset: Int,
+    ): Int {
+        val quoteCount = countLeadingDoubleQuotes(literal.text)
+        val literalOffset = caretOffset - literal.textOffset
+        val indentCount =
+            literal.text.substring(0, literalOffset)
+        val indentRegex = Regex("(?<=\\n)[ \\t]+")
+        val indentSpacesCount = indentRegex.findAll(indentCount).sumOf { it.value.length }
+        return literalOffset - quoteCount - indentSpacesCount
+    }
+
+    private fun countLeadingDoubleQuotes(s: String): Int {
+        var count = 0
+        for (char in s) {
+            if (char == '"') {
+                count++
+            } else {
+                break
+            }
+        }
+        return count
+    }
+
     override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.BGT
 
     override fun actionPerformed(e: AnActionEvent) {

src/main/kotlin/org/domaframework/doma/intellij/common/FileTypeCheck.kt

@@ -48,14 +48,24 @@ fun isJavaOrKotlinFileType(daoFile: PsiFile): Boolean {
 /*
  * Determine whether the open file is an SQL template file extension
  */
-fun isSupportFileType(file: VirtualFile): Boolean {
-    val extension = file.extension
+fun isSupportFileType(file: PsiFile): Boolean {
+    val extension = file.fileType.defaultExtension
     return when (extension) {
         "sql", "script" -> true
         else -> false
     }
 }
 
+fun isInjectionSqlFile(file: PsiFile): Boolean {
+    val extension = file.fileType.defaultExtension
+    val filePath = file.virtualFile.path
+    return when (extension) {
+        "sql" -> true
+        else -> false
+    } &&
+        !(filePath.endsWith(".sql") || filePath.endsWith(".script"))
+}
+
 /**
  * Dao file search for SQL files
  */

src/main/kotlin/org/domaframework/doma/intellij/common/dao/DaoMethodUtil.kt

@@ -15,7 +15,6 @@
  */
 package org.domaframework.doma.intellij.common.dao
 
-import com.intellij.openapi.fileTypes.FileTypeManager
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
@@ -27,6 +26,7 @@ import com.intellij.psi.util.PsiTreeUtil
 import org.domaframework.doma.intellij.common.CommonPathParameter.Companion.RESOURCES_META_INF_PATH
 import org.domaframework.doma.intellij.common.CommonPathParameter.Companion.RESOURCES_PATH
 import org.domaframework.doma.intellij.common.getExtension
+import org.domaframework.doma.intellij.common.isInjectionSqlFile
 import org.domaframework.doma.intellij.common.isSupportFileType
 import org.domaframework.doma.intellij.common.searchDaoFile
 import org.domaframework.doma.intellij.extension.findFile
@@ -40,15 +40,18 @@ import org.domaframework.doma.intellij.extension.getModule
 fun findDaoMethod(originalFile: PsiFile): PsiMethod? {
     val project = originalFile.project
     val module = project.getModule(originalFile.virtualFile) ?: return null
-    val fileType = FileTypeManager.getInstance().getFileTypeByFile(originalFile.virtualFile)
 
-    if (isSupportFileType(originalFile.virtualFile)) {
+    if (isInjectionSqlFile(originalFile)) {
+        originalFile.let {
+            return PsiTreeUtil.getParentOfType(originalFile.context, PsiMethod::class.java)
+        }
+    } else if (isSupportFileType(originalFile)) {
         // TODO: Add Support Kotlin
         val fileTypeName = "JAVA"
-        val daoFile = findDaoFile(project, originalFile.virtualFile) ?: return null
+        val daoFile = findDaoFile(project, originalFile) ?: return null
         val relativePath =
             formatDaoPathFromSqlFilePath(
-                originalFile.virtualFile,
+                originalFile,
                 project.getContentRoot(originalFile.virtualFile)?.path ?: "",
                 fileTypeName,
             )
@@ -71,10 +74,6 @@ fun findDaoMethod(originalFile: PsiFile): PsiMethod? {
                 }
             return daoMethod
         }
-    } else if (fileType.defaultExtension == "sql") {
-        originalFile.let {
-            return PsiTreeUtil.getParentOfType(originalFile.context, PsiMethod::class.java)
-        }
     }
     return null
 }
@@ -84,14 +83,15 @@ fun findDaoMethod(originalFile: PsiFile): PsiMethod? {
  */
 fun findDaoFile(
     project: Project,
-    sqlFile: VirtualFile,
+    sqlFile: PsiFile,
 ): VirtualFile? {
-    project.getModule(sqlFile) ?: return null
-    val contentRoot = project.getContentRoot(sqlFile) ?: return null
+    val virtualFile = sqlFile.virtualFile ?: return null
+    project.getModule(virtualFile) ?: return null
+    val contentRoot = project.getContentRoot(virtualFile) ?: return null
     // TODO: Add Support Kotlin
     val relativeFilePath =
         formatDaoPathFromSqlFilePath(sqlFile, contentRoot.path, "JAVA")
-    return searchDaoFile(contentRoot, sqlFile.path, relativeFilePath)
+    return searchDaoFile(contentRoot, virtualFile.path, relativeFilePath)
 }
 
 private fun findDaoClass(
@@ -103,11 +103,15 @@ private fun findDaoClass(
  * Generate Dao deployment path from SQL file path
  */
 fun formatDaoPathFromSqlFilePath(
-    relativeBaseSqlFile: VirtualFile,
+    relativeBaseSqlFile: PsiFile,
     projectRootPath: String,
     extension: String,
 ): String {
-    var relativeFilePath = relativeBaseSqlFile.path.substring(projectRootPath.length)
+    if (isInjectionSqlFile(relativeBaseSqlFile)) {
+        return ""
+    }
+    val sqlPath = relativeBaseSqlFile.virtualFile.path
+    var relativeFilePath = sqlPath.substring(projectRootPath.length)
     if (!relativeFilePath.startsWith("/")) {
         relativeFilePath = "/$relativeFilePath"
     }

src/main/kotlin/org/domaframework/doma/intellij/gutter/sql/SqlLineMakerProvider.kt

@@ -29,6 +29,7 @@ import org.domaframework.doma.intellij.common.PluginLoggerUtil
 import org.domaframework.doma.intellij.common.dao.findDaoFile
 import org.domaframework.doma.intellij.common.dao.findDaoMethod
 import org.domaframework.doma.intellij.common.dao.jumpToDaoMethod
+import org.domaframework.doma.intellij.common.isInjectionSqlFile
 import org.domaframework.doma.intellij.common.isSupportFileType
 import org.jetbrains.kotlin.idea.core.util.toPsiFile
 import java.awt.event.MouseEvent
@@ -43,18 +44,18 @@ class SqlLineMakerProvider : RelatedItemLineMarkerProvider() {
         result: MutableCollection<in RelatedItemLineMarkerInfo<*>>,
     ) {
         val project = e.project
-        val virtualFile = e.containingFile.virtualFile
-        if (!isSupportFileType(virtualFile)) return
+        val file = e.containingFile ?: return
+        if (!isSupportFileType(file) || isInjectionSqlFile(file)) return
         // Display only on the first line
         if (e.originalElement.parent.originalElement !is PsiFile ||
-            e.textRange.startOffset != e.containingFile.textRange.startOffset
+            e.textRange.startOffset != file.textRange.startOffset
         ) {
             return
         }
 
         val identifier = e.firstChild ?: e
         val daoFile =
-            findDaoFile(project, virtualFile)?.let {
+            findDaoFile(project, file)?.let {
                 findDaoMethod(e.containingFile) ?: return
                 it
             } ?: return
@@ -65,7 +66,7 @@ class SqlLineMakerProvider : RelatedItemLineMarkerProvider() {
                 identifier.textRange,
                 getIcon(daoFile.toPsiFile(project)),
                 getToolTipTitle(daoFile.toPsiFile(project)),
-                getHandler(daoFile, identifier, virtualFile.nameWithoutExtension),
+                getHandler(daoFile, identifier, file.virtualFile.nameWithoutExtension),
                 GutterIconRenderer.Alignment.RIGHT,
             ) {
                 ArrayList<GotoRelatedItem>()