use of a set to not have duplicate and prevent cyclic recursion

Alaanor · May 14, 2022 · 8c05dfb · 8c05dfb
1 parent 16eb467
commit 8c05dfb
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 17 deletions.
diff --git a/src/main/kotlin/com/github/alaanor/candid/reference/CandidTypeReference.kt b/src/main/kotlin/com/github/alaanor/candid/reference/CandidTypeReference.kt
@@ -4,11 +4,9 @@ import com.github.alaanor.candid.completion.AddImportInsertHandler
 import com.github.alaanor.candid.psi.CandidElementFactory
 import com.github.alaanor.candid.psi.CandidIdentifierDeclaration
 import com.github.alaanor.candid.psi.CandidIdentifierReference
-import com.github.alaanor.candid.psi.stub.impl.CandidIdentifierDeclarationStub
 import com.github.alaanor.candid.psi.stub.index.CandidStubTypeIndex
 import com.github.alaanor.candid.util.CandidImportUtil
 import com.github.alaanor.candid.util.filePath
-import com.github.alaanor.candid.util.filePathDirectory
 import com.github.alaanor.candid.util.getRelativePath
 import com.intellij.codeInsight.daemon.EmptyResolveMessageProvider
 import com.intellij.codeInsight.lookup.LookupElementBuilder
@@ -76,7 +74,7 @@ class CandidTypeReference(identifierReference: CandidIdentifierReference, privat
             ?.let { return it }
 
         // search further
-        val files = CandidImportUtil.getAllImportedFileFor(element.containingFile, true)
+        val files = CandidImportUtil.getAllImportedFileFor(element.containingFile)
         files.forEach { file ->
             PsiTreeUtil
                 .findChildrenOfType(file, CandidIdentifierDeclaration::class.java)

diff --git a/src/main/kotlin/com/github/alaanor/candid/util/CandidImportUtil.kt b/src/main/kotlin/com/github/alaanor/candid/util/CandidImportUtil.kt
@@ -5,26 +5,25 @@ import com.github.alaanor.candid.psi.importPathString
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
 import com.intellij.psi.util.PsiTreeUtil
-import com.intellij.util.containers.addIfNotNull
 
 object CandidImportUtil {
-    fun getAllImportedFileFor(psiFile: PsiFile, recursively: Boolean = true): List<PsiFile> {
-        val targetFiles = mutableListOf<PsiFile>()
+
+    fun getAllImportedFileFor(psiFile: PsiFile): List<PsiFile> {
+        return getAllImportedFileFor(psiFile, mutableSetOf())
+    }
+
+    private fun getAllImportedFileFor(psiFile: PsiFile, fileSet: MutableSet<PsiFile>): List<PsiFile> {
         PsiTreeUtil.findChildrenOfType(psiFile, CandidImportStatement::class.java).forEach { importStatement ->
-            importStatement.importPathString()?.let { importPath ->
-                targetFiles.addIfNotNull(resolveRelatively(psiFile, importPath))
-            }
-        }
+            val importPath = importStatement.importPathString() ?: return@forEach
+            val file = resolveRelatively(psiFile, importPath) ?: return@forEach
 
-        if (recursively) {
-            targetFiles.addAll(
-                targetFiles
-                    .map { getAllImportedFileFor(it, true) }
-                    .flatten()
-            )
+            if (!fileSet.contains(file)) {
+                fileSet.add(file)
+                getAllImportedFileFor(file, fileSet)
+            }
         }
 
-        return targetFiles
+        return fileSet.toList()
     }
 
     fun resolveRelatively(psiFile: PsiFile, importPath: String): PsiFile? {