Find the top migration file faster (#3108)

Alec Strong · AlecKazakova · commit 97eff26120bd · 2022-04-23T14:17:27.000Z
diff --git a/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/SqlDelightCompiler.kt b/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/SqlDelightCompiler.kt
@@ -24,6 +24,7 @@ import app.cash.sqldelight.core.lang.SqlDelightQueriesFile
 import app.cash.sqldelight.core.lang.queriesName
 import app.cash.sqldelight.core.lang.util.sqFile
 import app.cash.sqldelight.dialect.api.SqlDelightDialect
+import com.alecstrong.sql.psi.core.psi.InvalidElementDetectedException
 import com.alecstrong.sql.psi.core.psi.NamedElement
 import com.alecstrong.sql.psi.core.psi.SqlCreateViewStmt
 import com.intellij.openapi.module.Module
@@ -176,7 +177,7 @@ object SqlDelightCompiler {
   }
 
   private fun List<NamedQuery>.writeQueryInterfaces(file: SqlDelightFile, output: FileAppender) {
-    return filter { tryWithElement(it.select) { it.needsInterface() } }
+    return filter { tryWithElement(it.select) { it.needsInterface() } == true }
       .forEach { namedQuery ->
         val fileSpec = FileSpec.builder(namedQuery.interfaceType.packageName, namedQuery.name)
           .apply {
@@ -213,11 +214,14 @@ object SqlDelightCompiler {
 internal fun <T> tryWithElement(
   element: PsiElement,
   block: () -> T
-): T {
+): T? {
   try {
     return block()
   } catch (e: ProcessCanceledException) {
     throw e
+  } catch (e: InvalidElementDetectedException) {
+    // It's okay if compilation is cut short, we can just quit out.
+    return null
   } catch (e: Throwable) {
     val exception = IllegalStateException("Failed to compile $element :\n${element.text}")
     exception.initCause(e)
diff --git a/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/TreeUtil.kt b/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/TreeUtil.kt
@@ -39,6 +39,7 @@ import com.alecstrong.sql.psi.core.psi.SqlTypeName
 import com.alecstrong.sql.psi.core.psi.SqlTypes
 import com.alecstrong.sql.psi.core.psi.mixins.ColumnDefMixin
 import com.intellij.lang.ASTNode
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.tree.IElementType
@@ -53,7 +54,7 @@ internal inline fun <reified R : PsiElement> PsiElement.parentOfType(): R {
   return PsiTreeUtil.getParentOfType(this, R::class.java)!!
 }
 
-internal fun PsiElement.type(): IntermediateType = when (this) {
+internal fun PsiElement.type(): IntermediateType = if (!isValid) throw ProcessCanceledException() else when (this) {
   is ExposableType -> type()
   is SqlTypeName -> sqFile().typeResolver.definitionType(this)
   is AliasElement -> source().type().copy(name = name)
diff --git a/sqldelight-idea-plugin/src/main/kotlin/app/cash/sqldelight/intellij/inspections/SchemaNeedsMigrationInspection.kt b/sqldelight-idea-plugin/src/main/kotlin/app/cash/sqldelight/intellij/inspections/SchemaNeedsMigrationInspection.kt
@@ -5,7 +5,6 @@ import app.cash.sqldelight.core.SqlDelightProjectService
 import app.cash.sqldelight.core.lang.MigrationFile
 import app.cash.sqldelight.core.lang.SqlDelightQueriesFile
 import app.cash.sqldelight.core.lang.psi.parameterValue
-import app.cash.sqldelight.core.lang.util.findChildrenOfType
 import app.cash.sqldelight.intellij.refactoring.SqlDelightSignatureBuilder
 import app.cash.sqldelight.intellij.refactoring.SqlDelightSuggestedRefactoringExecution
 import app.cash.sqldelight.intellij.refactoring.SqlDelightSuggestedRefactoringExecution.SuggestedMigrationData
@@ -16,6 +15,7 @@ import com.intellij.codeInspection.LocalQuickFixOnPsiElement
 import com.intellij.codeInspection.ProblemHighlightType.GENERIC_ERROR
 import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.PsiFile
@@ -34,10 +34,20 @@ internal class SchemaNeedsMigrationInspection : LocalInspectionTool() {
       val file = createTable.containingFile as? SqlDelightQueriesFile ?: return
       val module = file.module ?: return
 
+      fun PsiDirectory.migrationFiles(): Sequence<MigrationFile> {
+        return children.asSequence().flatMap {
+          when (it) {
+            is PsiDirectory -> it.migrationFiles()
+            is MigrationFile -> sequenceOf(it)
+            else -> emptySequence()
+          }
+        }
+      }
+
       val dbFile = file.findDbFile() ?: return
       val fileIndex = SqlDelightFileIndex.getInstance(module)
-      val topMigrationFile = fileIndex.sourceFolders(file)
-        .flatMap { it.findChildrenOfType<MigrationFile>() }
+      val topMigrationFile = fileIndex.sourceFolders(file).asSequence()
+        .flatMap { it.migrationFiles() }
         .maxByOrNull { it.version }
 
       val tables = (topMigrationFile ?: dbFile).tables(true)
