feat: [EXPERIMENTAL] prune method bodies of files before parsing in JavaCompilerImpl

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/compiler/CompileBatch.java

@@ -30,6 +30,7 @@ import androidx.core.util.Pair;
 import com.itsaky.androidide.builder.model.IJavaCompilerSettings;
 import com.itsaky.androidide.javac.services.compiler.ReusableBorrow;
 import com.itsaky.androidide.javac.services.partial.DiagnosticListenerImpl;
+import com.itsaky.androidide.lsp.java.models.CompilationRequest;
 import com.itsaky.androidide.lsp.java.visitors.MethodRangeScanner;
 import com.itsaky.androidide.models.Range;
 import com.itsaky.androidide.projects.api.ModuleProject;
@@ -43,6 +44,7 @@ import openjdk.source.util.TreePath;
 import openjdk.tools.javac.api.ClientCodeWrapper;
 import openjdk.tools.javac.api.JavacTaskImpl;
 import openjdk.tools.javac.code.Kinds;
+import openjdk.tools.javac.util.Context;
 import openjdk.tools.javac.util.JCDiagnostic;
 
 import java.io.File;
@@ -76,18 +78,26 @@ public class CompileBatch implements AutoCloseable {
   boolean closed;
 
   CompileBatch(
-      JavaCompilerService parent,
-      Collection<? extends JavaFileObject> files,
-      CompilationTaskProcessor taskProcessor) {
+    JavaCompilerService parent,
+    Collection<? extends JavaFileObject> files,
+    CompilationRequest compilationRequest) {
     this.parent = parent;
     this.borrow = batchTask(parent, files);
     this.task = borrow.task;
     this.roots = new ArrayList<>();
-
-    Objects.requireNonNull(taskProcessor, "A task processor is required");
+  
+    final var context = task.getContext();
+    final var config = JavaCompilerConfig.instance(context);
+    config.setFiles(files);
+    
+    if (compilationRequest.configureContext != null) {
+      compilationRequest.configureContext.accept(context);
+    }
+  
+    Objects.requireNonNull(compilationRequest, "A task processor is required");
 
     try {
-      taskProcessor.process(borrow.task, this::processCompilationUnit);
+      compilationRequest.compilationTaskProcessor.process(borrow.task, this::processCompilationUnit);
     } catch (Throwable e) {
       throw new RuntimeException(e);
     }

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/compiler/JavaCompilerConfig.kt

@@ -0,0 +1,60 @@
+/*
+ *  This file is part of AndroidIDE.
+ *
+ *  AndroidIDE is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  AndroidIDE is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *   along with AndroidIDE.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.itsaky.androidide.lsp.java.compiler
+
+import com.itsaky.androidide.models.Position
+import jdkx.tools.JavaFileObject
+import openjdk.tools.javac.util.Context
+
+/**
+ * Configuration for the [JavaCompilerImpl].
+ *
+ * @property files The main files that are being compiled/parsed.
+ * @property completionInfo Information about the completion
+ * @author Akash Yadav
+ */
+class JavaCompilerConfig(context: Context) {
+  init {
+    context.put(compilerConfigKey, this)
+  }
+
+  var files: Collection<JavaFileObject>? = null
+  var completionInfo: CompletionInfo? = null
+
+  companion object {
+
+    @JvmField val compilerConfigKey = Context.Key<JavaCompilerConfig>()
+
+    @JvmStatic
+    fun instance(context: Context): JavaCompilerConfig {
+      var instance = context.get(compilerConfigKey)
+      if (instance == null) {
+        instance = JavaCompilerConfig(context)
+      }
+      return instance
+    }
+  }
+}
+
+/**
+ * Information about the completion request initiated by the Java completion provider.
+ *
+ * @property cursor The cursor position for the completion.
+ * @author Akash Yadav
+ */
+data class CompletionInfo(val cursor: Position)

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/compiler/JavaCompilerImpl.kt

@@ -17,17 +17,61 @@
 
 package com.itsaky.androidide.lsp.java.compiler
 
+import android.util.Log
 import com.itsaky.androidide.javac.services.compiler.ReusableContext
 import com.itsaky.androidide.javac.services.compiler.ReusableJavaCompiler
+import com.itsaky.androidide.lsp.java.parser.ts.TSJavaParser
+import com.itsaky.androidide.lsp.java.parser.ts.TSMethodPruner.prune
+import com.itsaky.androidide.projects.FileManager
+import com.itsaky.androidide.utils.ILogger
+import com.itsaky.androidide.utils.StopWatch
 import jdkx.tools.JavaFileObject
+import jdkx.tools.JavaFileObject.Kind.SOURCE
+import kotlin.io.path.name
+import openjdk.tools.javac.api.ClientCodeWrapper
 import openjdk.tools.javac.tree.JCTree.JCCompilationUnit
 import openjdk.tools.javac.util.Context
 
 class JavaCompilerImpl(context: Context?) : ReusableJavaCompiler(context) {
 
+  private val _log = ILogger.newInstance("JavaCompilerImpl")
+
   override fun parse(filename: JavaFileObject?, content: CharSequence?): JCCompilationUnit {
-    // TODO Prune unnecessary methods to speed up compilation
-    return super.parse(filename, content)
+    val file = ClientCodeWrapper.instance(context).unwrap(filename)
+    val compilerConfig = JavaCompilerConfig.instance(context)
+
+    // Preconditions
+    if (
+      content == null ||
+        compilerConfig.files == null ||
+        filename?.kind != SOURCE ||
+        compilerConfig.files?.contains(file) == false
+    ) {
+      return super.parse(filename, content)
+    }
+
+    // If the file is NOT being parsed for a completion request,
+    // we should not prune method bodies of active documents
+    if (compilerConfig.completionInfo == null && FileManager.isActive(filename.toUri())) {
+      return super.parse(filename, content)
+    }
+    
+    val pruned =
+      StopWatch("${if(file is SourceFileObject) "[${file.path.name}] " else ""}Prune method bodies")
+        .let {
+          val contentBuilder = StringBuilder(content)
+          val parseResult = TSJavaParser.parse(file)
+          prune(
+            contentBuilder,
+            parseResult.tree,
+            compilerConfig.completionInfo?.cursor?.index ?: -1
+          )
+          it.log()
+          Log.d("JavaMethodPruner", "Pruned contents for file: $file\n$contentBuilder")
+          return@let contentBuilder
+        }
+
+    return super.parse(filename, pruned)
   }
 
   companion object {

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/compiler/JavaCompilerService.java

@@ -413,7 +413,7 @@ public class JavaCompilerService implements CompilerProvider {
       throw new RuntimeException("empty sources");
     }
 
-    CompileBatch firstAttempt = new CompileBatch(this, sources, request.compilationTaskProcessor);
+    CompileBatch firstAttempt = new CompileBatch(this, sources, request);
     Set<Path> addFiles = firstAttempt.needsAdditionalSources();
 
     if (addFiles.isEmpty()) {
@@ -430,7 +430,7 @@ public class JavaCompilerService implements CompilerProvider {
       moreSources.add(new SourceFileObject(add));
     }
 
-    return new CompileBatch(this, moreSources, request.compilationTaskProcessor);
+    return new CompileBatch(this, moreSources, request);
   }
 
   private boolean containsWord(Path file, String word) {

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/compiler/SourceFileObject.java

@@ -28,8 +28,10 @@ import java.io.Reader;
 import java.io.StringReader;
 import java.io.Writer;
 import java.net.URI;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.Instant;
+import java.util.Objects;
 
 import jdkx.lang.model.element.Modifier;
 import jdkx.lang.model.element.NestingKind;
@@ -53,19 +55,7 @@ public class SourceFileObject implements JavaFileObject {
     this.contents = contents;
     this.modified = modified;
   }
-
-  @Override
-  public int hashCode() {
-    return this.path.hashCode();
-  }
-
-  @Override
-  public boolean equals(Object other) {
-    if (other.getClass() != SourceFileObject.class) return false;
-    SourceFileObject that = (SourceFileObject) other;
-    return DocumentUtils.isSameFile(this.path, that.path);
-  }
-
+  
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this).add("path", this.path.toString()).toString();
@@ -158,4 +148,28 @@ public class SourceFileObject implements JavaFileObject {
   public boolean delete() {
     throw new UnsupportedOperationException();
   }
+  
+  @Override
+  public boolean equals(final Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof SourceFileObject)) {
+      return false;
+    }
+    final SourceFileObject that = (SourceFileObject) o;
+    try {
+      return this.path != null && that.path != null
+        && Files.isSameFile(this.path, that.path)
+        && Objects.equals(contents, that.contents)
+        && Objects.equals(modified, that.modified);
+    } catch (Exception e) {
+      return false;
+    }
+  }
+  
+  @Override
+  public int hashCode() {
+    return Objects.hash(path, contents, modified);
+  }
 }

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/models/CompilationRequest.kt

@@ -19,14 +19,16 @@ package com.itsaky.androidide.lsp.java.models
 
 import com.itsaky.androidide.lsp.java.compiler.CompilationTaskProcessor
 import com.itsaky.androidide.lsp.java.compiler.DefaultCompilationTaskProcessor
+import java.util.function.Consumer
 import jdkx.tools.JavaFileObject
+import openjdk.tools.javac.util.Context
 
 /**
  * Data sent to compiler to request a compilation.
  *
- * @author Akash Yadav
  * @param sources The source files to compile.
  * @param partialRequest Data that will be used to a partial reparse.
+ * @author Akash Yadav
  */
 data class CompilationRequest
 @JvmOverloads
@@ -34,5 +36,6 @@ constructor(
   @JvmField val sources: Collection<JavaFileObject>,
   @JvmField val partialRequest: PartialReparseRequest? = null,
   @JvmField
-  val compilationTaskProcessor: CompilationTaskProcessor = DefaultCompilationTaskProcessor()
+  val compilationTaskProcessor: CompilationTaskProcessor = DefaultCompilationTaskProcessor(),
+  @JvmField var configureContext: Consumer<Context>? = null
 )

lsp/java/src/main/java/com/itsaky/androidide/lsp/java/providers/CompletionProvider.java

@@ -28,13 +28,13 @@ import com.itsaky.androidide.lsp.api.ICompletionProvider;
 import com.itsaky.androidide.lsp.api.IServerSettings;
 import com.itsaky.androidide.lsp.internal.model.CachedCompletion;
 import com.itsaky.androidide.lsp.java.compiler.CompileTask;
+import com.itsaky.androidide.lsp.java.compiler.CompletionInfo;
+import com.itsaky.androidide.lsp.java.compiler.JavaCompilerConfig;
 import com.itsaky.androidide.lsp.java.compiler.JavaCompilerService;
 import com.itsaky.androidide.lsp.java.compiler.SourceFileObject;
 import com.itsaky.androidide.lsp.java.compiler.SynchronizedTask;
 import com.itsaky.androidide.lsp.java.models.CompilationRequest;
 import com.itsaky.androidide.lsp.java.models.PartialReparseRequest;
-import com.itsaky.androidide.lsp.java.parser.ts.TSJavaParser;
-import com.itsaky.androidide.lsp.java.parser.ts.TSMethodPruner;
 import com.itsaky.androidide.lsp.java.providers.completion.IJavaCompletionProvider;
 import com.itsaky.androidide.lsp.java.providers.completion.IdentifierCompletionProvider;
 import com.itsaky.androidide.lsp.java.providers.completion.ImportCompletionProvider;
@@ -50,8 +50,6 @@ import com.itsaky.androidide.lsp.models.CompletionParams;
 import com.itsaky.androidide.lsp.models.CompletionResult;
 import com.itsaky.androidide.utils.DocumentUtils;
 import com.itsaky.androidide.utils.ILogger;
-import com.itsaky.androidide.utils.StopWatch;
-import com.itsaky.androidide.utils.VMUtils;
 
 import java.nio.file.Path;
 import java.time.Duration;
@@ -156,23 +154,15 @@ public class CompletionProvider extends AbstractServiceProvider implements IComp
     final var sourceObject = new SourceFileObject(file);
     final var contentBuilder = new StringBuilder(sourceObject.getCharContent(true));
     
-    if (false && !VMUtils.isJvm()) {
-      // Cannot use tree sitter in tests
-      // TODO(itsaky): Should we use the legacy method pruner for tests?
-      final StopWatch watch = new StopWatch("Prune method bodies");
-      final var parseResult = TSJavaParser.INSTANCE.parse(sourceObject);
-      TSMethodPruner.INSTANCE.prune(contentBuilder, parseResult.getTree(), (int) cursor);
-      watch.log();
-    }
-    
     int endOfLine = endOfLine(contentBuilder, (int) cursor);
     contentBuilder.insert(endOfLine, ';');
 
     final StringBuilder contents;
-    if (compiler.compiler.currentContext != null) {
+    final var context = compiler.compiler.currentContext;
+    if (context != null) {
       abortIfCancelled();
       abortCompletionIfCancelled();
-      contents = new ASTFixer(compiler.compiler.currentContext).fix(contentBuilder);
+      contents = new ASTFixer(context).fix(contentBuilder);
     } else {
       contents = contentBuilder;
     }
@@ -182,7 +172,8 @@ public class CompletionProvider extends AbstractServiceProvider implements IComp
         new PartialReparseRequest(cursor - params.requirePrefix().length(), contentString);
     abortIfCancelled();
     abortCompletionIfCancelled();
-    CompletionResult result = compileAndComplete(file, contentString, cursor, partialRequest);
+  
+    CompletionResult result = compileAndComplete(contentString, params, partialRequest);
     if (result == null) {
       result = CompletionResult.EMPTY;
     }
@@ -231,17 +222,24 @@ public class CompletionProvider extends AbstractServiceProvider implements IComp
   }
 
   private CompletionResult compileAndComplete(
-      Path file, String contents, final long cursor, PartialReparseRequest partialRequest) {
-    final Instant started = Instant.now();
-    final SourceFileObject source = new SourceFileObject(file, contents, Instant.now());
-    final String partial = partialIdentifier(contents, (int) cursor);
-    final boolean endsWithParen = endsWithParen(contents, (int) cursor);
+    String contents, CompletionParams params, PartialReparseRequest partialRequest) {
+    final long cursor = params.getPosition().requireIndex();
+    final var file = params.getFile();
+    final var started = Instant.now();
+    final var source = new SourceFileObject(file, contents, Instant.now());
+    final var partial = partialIdentifier(contents, (int) cursor);
+    final var endsWithParen = endsWithParen(contents, (int) cursor);
 
     abortIfCancelled();
     abortCompletionIfCancelled();
 
     final CompilationRequest request =
         new CompilationRequest(Collections.singletonList(source), partialRequest);
+    request.configureContext = ctx -> {
+      final var config = JavaCompilerConfig.instance(ctx);
+      config.setCompletionInfo(new CompletionInfo(params.getPosition()));
+    };
+    
     SynchronizedTask synchronizedTask = compiler.compile(request);
     return synchronizedTask.get(
         task -> {
@@ -266,6 +264,13 @@ public class CompletionProvider extends AbstractServiceProvider implements IComp
   
           final var result = doComplete(file, contents, cursor, newPartial, endsWithParen, task, path);
           new TopLevelSnippetsProvider().complete(partial, task.root(), result);
+          
+          // IMPORTANT: Unregister the completion info from the compiler configuration
+          if (task.task.getContext() != null) {
+            final var compilerConfig = JavaCompilerConfig.instance(task.task.getContext());
+            compilerConfig.setCompletionInfo(null);
+          }
+          
           return result;
         });
   }

subprojects/jdk-compiler/src/main/java/openjdk/tools/javac/api/ClientCodeWrapper.java

@@ -129,7 +129,8 @@ public class ClientCodeWrapper {
         return new WrappedFileObject(fo);
     }
 
-    FileObject unwrap(FileObject fo) {
+    // AndroidIDE Changed: Make public
+    public FileObject unwrap(FileObject fo) {
         if (fo instanceof WrappedFileObject)
             return ((WrappedFileObject) fo).clientFileObject;
         else
@@ -148,8 +149,9 @@ public class ClientCodeWrapper {
             wrapped.add(wrap(fo));
         return Collections.unmodifiableList(wrapped);
     }
-
-    JavaFileObject unwrap(JavaFileObject fo) {
+    
+    // AndroidIDE Changed: Make public
+    public JavaFileObject unwrap(JavaFileObject fo) {
           if (fo instanceof WrappedJavaFileObject)
             return ((JavaFileObject) ((WrappedJavaFileObject) fo).clientFileObject);
         else
@@ -181,9 +183,10 @@ public class ClientCodeWrapper {
             c.add(unwrap(l));
         return c;
     }
-
+    
+    // AndroidIDE Changed: Make public
     @SuppressWarnings("unchecked")
-    private <T> Diagnostic<T> unwrap(final Diagnostic<T> diagnostic) {
+    public  <T> Diagnostic<T> unwrap(final Diagnostic<T> diagnostic) {
         if (diagnostic instanceof JCDiagnostic) {
             JCDiagnostic d = (JCDiagnostic) diagnostic;
             return (Diagnostic<T>) new DiagnosticSourceUnwrapper(d);

subprojects/projects/src/main/java/com/itsaky/androidide/projects/FileManager.kt

@@ -35,6 +35,8 @@ import java.time.Instant
 import org.apache.commons.io.FileUtils
 import org.greenrobot.eventbus.Subscribe
 import org.greenrobot.eventbus.ThreadMode.BACKGROUND
+import java.net.URI
+import java.nio.file.Paths
 
 /**
  * Manages active documents.
@@ -46,6 +48,10 @@ object FileManager : EventReceiver {
   private val log = ILogger.newInstance(javaClass.simpleName)
   private val activeDocuments: MutableMap<Path, ActiveDocument> = mutableMapOf()
 
+  fun isActive(uri: URI): Boolean {
+    return isActive(Paths.get(uri))
+  }
+  
   fun isActive(file: Path): Boolean {
     return this.activeDocuments.containsKey(file.normalize())
   }