Use ClassTrie for completing import paths (#225)

build-tools/projects/src/main/java/com/itsaky/androidide/projects/util/ClassTrie.kt

@@ -100,7 +100,6 @@ open class ClassTrie(val root: Node = Node()) {
 
       node = node.children[segment]
     }
-
     return node != null
   }
 
@@ -142,6 +141,24 @@ open class ClassTrie(val root: Node = Node()) {
     return findInPackage(packageName).map { it.qualifiedName }
   }
 
+  /**
+   * Finds node with the given qualified name. Or `null` if none was found.
+   *
+   * @param qualifiedName The fully qualified name of the node to find.
+   */
+  open fun findNode(qualifiedName: String): Node? {
+    val segments = segments(qualifiedName)
+    var node: Node? = root
+    for (segment in segments) {
+      if (node == null) {
+        break
+      }
+      node = node.children[segment]
+    }
+
+    return node
+  }
+
   /** Returns all class nodes available in this trie. */
   fun allClassNodes(): Set<Node> {
     return root.allClassNodes()

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/ClassNamesCompletionProvider.kt

@@ -19,7 +19,7 @@ package com.itsaky.androidide.lsp.java.providers.completion
 
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
-import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.java.providers.CompletionProvider
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
 import com.sun.source.tree.ClassTree
@@ -37,7 +37,7 @@ import java.util.*
 class ClassNamesCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings,
   val root: CompilationUnitTree
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {
@@ -63,7 +63,7 @@ class ClassNamesCompletionProvider(
         continue
       }
 
-      list.add(classItem(imports, file, className, partial, matchLevel))
+      list.add(classItem(imports, file, className, matchLevel))
       uniques.add(className)
     }
 
@@ -81,7 +81,7 @@ class ClassNamesCompletionProvider(
         break
       }
 
-      list.add(classItem(imports, file, className, partial, matchLevel))
+      list.add(classItem(imports, file, className, matchLevel))
       uniques.add(className)
     }
 
@@ -97,7 +97,7 @@ class ClassNamesCompletionProvider(
       }
 
       val name = packageName + "." + t.simpleName
-      list.add(classItem(name, partial, matchLevel))
+      list.add(classItem(name, matchLevel))
 
       if (list.size > CompletionProvider.MAX_COMPLETION_ITEMS) {
         break

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/IJavaCompletionProvider.kt

@@ -177,19 +177,14 @@ abstract class IJavaCompletionProvider(
     return item
   }
 
-  protected open fun classItem(
-    className: String,
-    partialName: String,
-    matchLevel: MatchLevel
-  ): CompletionItem {
-    return classItem(emptySet(), null, className, partialName, matchLevel)
+  protected open fun classItem(className: String, matchLevel: MatchLevel): CompletionItem {
+    return classItem(emptySet(), null, className, matchLevel)
   }
 
   protected open fun classItem(
     imports: Set<String>,
     file: Path?,
     className: String,
-    partialName: String,
     matchLevel: MatchLevel,
   ): CompletionItem {
     val item = CompletionItem()
@@ -201,7 +196,9 @@ abstract class IJavaCompletionProvider(
     val data = CompletionData()
     data.className = className
     item.data = data
-    item.additionalEditHandler = ClassImportEditHandler(imports, file!!)
+
+    // If file is not provided, we are probably completing an import path
+    item.additionalEditHandler = if (file == null) null else ClassImportEditHandler(imports, file)
     return item
   }
 
@@ -210,11 +207,7 @@ abstract class IJavaCompletionProvider(
     return if (dot == -1) className else className.subSequence(dot + 1, className.length)
   }
 
-  protected open fun packageItem(
-    name: String,
-    partialName: String,
-    matchLevel: MatchLevel
-  ): CompletionItem =
+  protected open fun packageItem(name: String, matchLevel: MatchLevel): CompletionItem =
     CompletionItem().apply {
       setLabel(name)
       this.kind = MODULE

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/IdentifierCompletionProvider.kt

@@ -20,14 +20,15 @@ package com.itsaky.androidide.lsp.java.providers.completion
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
 import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.sun.source.util.TreePath
 import java.nio.file.Path
 
 /** @author Akash Yadav */
 class IdentifierCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {
 

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/ImportCompletionProvider.kt

@@ -19,12 +19,13 @@ package com.itsaky.androidide.lsp.java.providers.completion
 
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
-import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
 import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.java.providers.CompletionProvider.MAX_COMPLETION_ITEMS
 import com.itsaky.androidide.lsp.models.CompletionItem
 import com.itsaky.androidide.lsp.models.CompletionResult
+import com.itsaky.androidide.lsp.models.MatchLevel.CASE_SENSITIVE_EQUAL
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
+import com.itsaky.androidide.projects.util.ClassTrie.Node
 import com.sun.source.util.TreePath
 import java.nio.file.Path
 
@@ -48,10 +49,81 @@ class ImportCompletionProvider(
     partial: String,
     endsWithParen: Boolean,
   ): CompletionResult {
-    log.info("...complete import")
+    log.info("...complete import for path:", importPath)
 
     val names: MutableSet<String> = HashSet()
     val list = mutableListOf<CompletionItem>()
+
+    var pkgName = importPath
+    var incomplete: String = ""
+    if (!pkgName.contains(".")) {
+      pkgName = ""
+      incomplete = importPath
+    } else if (pkgName.endsWith(".")) {
+      pkgName = pkgName.substring(0, pkgName.lastIndex)
+      incomplete = ""
+    } else {
+      pkgName = pkgName.substringBeforeLast(delimiter = '.')
+      incomplete = pkgName.substringAfterLast(delimiter = '.')
+    }
+
+    val module = compiler.module
+    if (module == null) {
+      legacyImportPathCompletion(partial, names, list)
+      return CompletionResult(list)
+    }
+
+    val sourceNode =
+      if (pkgName.isEmpty()) module.compileJavaSourceClasses.root
+      else module.compileJavaSourceClasses.findNode(pkgName)
+    if (sourceNode != null) {
+      addDirectChildNodes(sourceNode, incomplete, list, names)
+    }
+
+    val classpathNode =
+      if (pkgName.isEmpty()) module.compileClasspathClasses.root
+      else module.compileClasspathClasses.findNode(pkgName)
+    if (classpathNode != null) {
+      addDirectChildNodes(classpathNode, incomplete, list, names)
+    }
+
+    // TODO Add from bootstrap classes
+
+    return CompletionResult(list)
+  }
+
+  private fun addDirectChildNodes(
+    sourceNode: Node,
+    incomplete: String,
+    list: MutableList<CompletionItem>,
+    names: MutableSet<String>
+  ) {
+    for (child in sourceNode.children.values) {
+      val match =
+        if (incomplete.isEmpty()) {
+          CASE_SENSITIVE_EQUAL
+        } else {
+          matchLevel(child.name, incomplete)
+        }
+
+      if (match == NO_MATCH || names.contains(child.name)) {
+        continue
+      }
+
+      if (child.isClass) {
+        list.add(classItem(child.qualifiedName, match))
+      } else {
+        list.add(packageItem(child.name, match))
+      }
+      names.add(child.name)
+    }
+  }
+
+  private fun legacyImportPathCompletion(
+    partial: String,
+    names: MutableSet<String>,
+    list: MutableList<CompletionItem>
+  ) {
     for (className in compiler.publicTopLevelTypes()) {
       val matchLevel = matchLevel(className, partial)
       if (matchLevel == NO_MATCH) {
@@ -70,16 +142,14 @@ class ImportCompletionProvider(
       names.add(segment)
       val isClass = end == importPath.length
       if (isClass) {
-        list.add(classItem(className, importPath, matchLevel))
+        list.add(classItem(className, matchLevel))
       } else {
-        list.add(packageItem(segment, importPath, matchLevel))
+        list.add(packageItem(segment, matchLevel))
       }
 
       if (list.size > MAX_COMPLETION_ITEMS) {
         break
       }
     }
-
-    return CompletionResult(list)
   }
 }

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/KeywordCompletionProvider.kt

@@ -20,6 +20,7 @@ package com.itsaky.androidide.lsp.java.providers.completion
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
 import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
 import com.sun.source.tree.ClassTree
 import com.sun.source.tree.CompilationUnitTree
@@ -36,7 +37,7 @@ import java.nio.file.Path
 class KeywordCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {
 

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/MemberReferenceCompletionProvider.kt

@@ -20,6 +20,7 @@ package com.itsaky.androidide.lsp.java.providers.completion
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
 import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
 import com.sun.source.tree.MemberReferenceTree
 import com.sun.source.tree.Scope
@@ -42,7 +43,7 @@ import javax.lang.model.type.TypeVariable
 class MemberReferenceCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings,
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {
 

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/MemberSelectCompletionProvider.kt

@@ -20,6 +20,7 @@ package com.itsaky.androidide.lsp.java.providers.completion
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
 import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.java.utils.ScopeHelper
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
 import com.sun.source.tree.MemberSelectTree
@@ -44,7 +45,7 @@ import javax.lang.model.type.TypeVariable
 class MemberSelectCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings,
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {
 

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/StaticImportCompletionProvider.kt

@@ -20,6 +20,7 @@ package com.itsaky.androidide.lsp.java.providers.completion
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
 import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.java.providers.CompletionProvider
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
 import com.sun.source.tree.CompilationUnitTree
@@ -42,7 +43,7 @@ import javax.lang.model.element.TypeElement
 class StaticImportCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings,
   val root: CompilationUnitTree,
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/completion/SwitchConstantCompletionProvider.kt

@@ -20,6 +20,7 @@ package com.itsaky.androidide.lsp.java.providers.completion
 import com.itsaky.androidide.lsp.api.IServerSettings
 import com.itsaky.androidide.lsp.java.compiler.CompileTask
 import com.itsaky.androidide.lsp.java.compiler.CompilerProvider
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService
 import com.itsaky.androidide.lsp.models.MatchLevel.NO_MATCH
 import com.sun.source.tree.SwitchTree
 import com.sun.source.util.TreePath
@@ -38,7 +39,7 @@ import javax.lang.model.type.DeclaredType
 class SwitchConstantCompletionProvider(
   completingFile: Path,
   cursor: Long,
-  compiler: CompilerProvider,
+  compiler: JavaCompilerService,
   settings: IServerSettings,
 ) : IJavaCompletionProvider(completingFile, cursor, compiler, settings) {