Share document accumulator between Java/Kotlin plugins

semanticdb-javac/src/main/java/com/sourcegraph/semanticdb_javac/SemanticdbTaskListener.java

@@ -1,6 +1,7 @@
 package com.sourcegraph.semanticdb_javac;
 
 import com.sourcegraph.semanticdb.Semanticdb;
+import com.sourcegraph.semanticdb.SemanticdbDocumentBuilder;
 import com.sourcegraph.semanticdb.SemanticdbPaths;
 import com.sourcegraph.semanticdb.SemanticdbWriter;
 
@@ -12,15 +13,16 @@
 import javax.lang.model.util.Types;
 
 import javax.tools.JavaFileObject;
-import java.io.*;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintWriter;
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.HashSet;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
-import java.util.stream.Collectors;
 
 /**
  * Callback hook that generates SemanticDB when the compiler has completed typechecking a Java
@@ -33,6 +35,8 @@ public final class SemanticdbTaskListener implements TaskListener {
   private final Types types;
   private final Trees trees;
   private final Elements elements;
+  // Javac fires ANALYZE once per top-level type; accumulate across rounds per output path.
+  private final Map<Path, PerSourceState> perSourceState = new HashMap<>();
   private int noRelativePathCounter = 0;
 
   public SemanticdbTaskListener(
@@ -58,8 +62,10 @@ public void started(TaskEvent e) {
       inferBazelSourceroot(e.getSourceFile());
       Result<Path, String> semanticdbPath = semanticdbOutputPath(options, e);
       if (semanticdbPath.isOk()) {
+        Path output = semanticdbPath.getOrThrow();
+        perSourceState.remove(output);
         try {
-          Files.deleteIfExists(semanticdbPath.getOrThrow());
+          Files.deleteIfExists(output);
         } catch (IOException ex) {
           this.reportException(ex, e);
         }
@@ -116,12 +122,20 @@ private void onFinishedAnalyze(TaskEvent e) {
     Result<Path, String> path = semanticdbOutputPath(options, e);
     if (path != null) {
       if (path.isOk()) {
+        Path output = path.getOrThrow();
+        PerSourceState state = perSourceState.computeIfAbsent(output, k -> new PerSourceState());
         Semanticdb.TextDocument textDocument =
-            new SemanticdbVisitor(globals, e.getCompilationUnit(), options, types, trees, elements)
+            new SemanticdbVisitor(
+                    globals,
+                    state.locals,
+                    e.getCompilationUnit(),
+                    options,
+                    types,
+                    trees,
+                    elements,
+                    state.documentBuilder)
                 .buildTextDocument(e.getCompilationUnit());
-        Path output = path.getOrThrow();
-        if (Files.exists(output)) appendSemanticdb(e, output, textDocument);
-        else writeSemanticdb(e, output, textDocument);
+        writeSemanticdb(e, output, textDocument);
       } else {
         reporter.error(path.getErrorOrThrow(), e);
       }
@@ -136,76 +150,9 @@ private void writeSemanticdb(TaskEvent event, Path output, Semanticdb.TextDocume
     }
   }
 
-  private void appendSemanticdb(
-      TaskEvent event, Path output, Semanticdb.TextDocument textDocument) {
-    /*
-     * If there already is a semanticdb file at the given path,
-     * we do the following:
-     * - Read a documents collection
-     * - Try to find the document with the matching relative path (matching the incoming textDocument)
-     * - Then, depending on whether a matching document already exists in the collection:
-     *   - if YES, mutate it in place to only add entries from the incoming document
-     *   - if NO, simply add the incoming text document to the collection
-     * - Write the collection back to disk
-     * */
-    Semanticdb.TextDocument document = null;
-    int documentIndex = -1;
-    Semanticdb.TextDocuments documents = null;
-
-    try (InputStream is = Files.newInputStream(output.toFile().toPath())) {
-      documents = Semanticdb.TextDocuments.parseFrom(is);
-
-      for (int i = 0; i < documents.getDocumentsCount(); i++) {
-        Semanticdb.TextDocument candidate = documents.getDocuments(i);
-        if (document == null && candidate.getUri().equals(textDocument.getUri())) {
-          document = candidate;
-          documentIndex = i;
-        }
-      }
-
-    } catch (IOException e) {
-      this.reportException(e, event);
-      return;
-    }
-
-    if (document != null) {
-      // If there is a previous semanticdb document at this path, we need
-      // to deduplicate symbols and occurrences and mutate the document in place
-      Set<Semanticdb.SymbolInformation> symbols = new HashSet<>(textDocument.getSymbolsList());
-      Set<Semanticdb.SymbolOccurrence> occurrences =
-          new HashSet<>(textDocument.getOccurrencesList());
-      Set<Semanticdb.Synthetic> synthetics = new HashSet<>(textDocument.getSyntheticsList());
-
-      symbols.addAll(document.getSymbolsList());
-      occurrences.addAll(document.getOccurrencesList());
-      synthetics.addAll(document.getSyntheticsList());
-
-      documents
-          .toBuilder()
-          .addDocuments(
-              documentIndex,
-              document
-                  .toBuilder()
-                  .clearOccurrences()
-                  .addAllOccurrences(occurrences)
-                  .clearSymbols()
-                  .addAllSymbols(symbols)
-                  .clearSynthetics()
-                  .addAllSynthetics(synthetics));
-
-    } else {
-      // If no prior document was found, we can just add the incoming one to the collection
-      documents = documents.toBuilder().addDocuments(textDocument).build();
-    }
-
-    byte[] bytes = documents.toByteArray();
-
-    try {
-      Files.createDirectories(output.getParent());
-      Files.write(output, bytes);
-    } catch (IOException e) {
-      this.reportException(e, event);
-    }
+  private static final class PerSourceState {
+    final SemanticdbDocumentBuilder documentBuilder = new SemanticdbDocumentBuilder();
+    final LocalSymbolsCache locals = new LocalSymbolsCache();
   }
 
   public static Path absolutePathFromUri(SemanticdbJavacOptions options, JavaFileObject file) {

semanticdb-javac/src/main/java/com/sourcegraph/semanticdb_javac/SemanticdbVisitor.java

@@ -2,6 +2,7 @@
 
 import com.sourcegraph.semanticdb.Semanticdb;
 
+import com.sourcegraph.semanticdb.SemanticdbDocumentBuilder;
 import com.sourcegraph.semanticdb.SemanticdbPaths;
 import com.sourcegraph.semanticdb.SemanticdbSymbols;
 
@@ -69,29 +70,29 @@ public class SemanticdbVisitor extends TreePathScanner<Void, Void> {
   private final CompilationUnitTree compUnitTree;
   private final Elements elements;
   private final SemanticdbJavacOptions options;
-  private final ArrayList<Semanticdb.SymbolOccurrence> occurrences;
-  private final ArrayList<Semanticdb.SymbolInformation> symbolInfos;
+  private final SemanticdbDocumentBuilder documentBuilder;
   private String source;
   private String uri;
 
   private final LinkedHashMap<Tree, TreePath> nodes;
 
   public SemanticdbVisitor(
       GlobalSymbolsCache globals,
+      LocalSymbolsCache locals,
       CompilationUnitTree compUnitTree,
       SemanticdbJavacOptions options,
       Types types,
       Trees trees,
-      Elements elements) {
-    this.globals = globals; // Reused cache between compilation units.
-    this.locals = new LocalSymbolsCache(); // Fresh cache per compilation unit.
+      Elements elements,
+      SemanticdbDocumentBuilder documentBuilder) {
+    this.globals = globals;
+    this.locals = locals;
     this.options = options;
     this.types = types;
     this.elements = elements;
     this.trees = trees;
     this.compUnitTree = compUnitTree;
-    this.occurrences = new ArrayList<>();
-    this.symbolInfos = new ArrayList<>();
+    this.documentBuilder = documentBuilder;
     this.source = semanticdbText();
     this.uri = semanticdbUri(compUnitTree, options);
     this.nodes = new LinkedHashMap<>();
@@ -102,15 +103,8 @@ public Semanticdb.TextDocument buildTextDocument(CompilationUnitTree tree) {
 
     resolveNodes();
 
-    return Semanticdb.TextDocument.newBuilder()
-        .setSchema(Semanticdb.Schema.SEMANTICDB4)
-        .setLanguage(Semanticdb.Language.JAVA)
-        .setUri(uri)
-        .setText(options.includeText ? this.source : "")
-        .setMd5(semanticdbMd5())
-        .addAllOccurrences(occurrences)
-        .addAllSymbols(symbolInfos)
-        .build();
+    return documentBuilder.build(
+        Semanticdb.Language.JAVA, uri, options.includeText ? this.source : "", semanticdbMd5());
   }
 
   private Optional<Semanticdb.Range> emitSymbolOccurrence(
@@ -135,7 +129,7 @@ private void emitSymbolOccurrence(
     if (sym == null) return;
     Optional<Semanticdb.SymbolOccurrence> occ =
         semanticdbOccurrence(sym, range, role, enclosingRange);
-    occ.ifPresent(occurrences::add);
+    occ.ifPresent(documentBuilder::addOccurrence);
   }
 
   private void emitSymbolInformation(Element sym, Tree tree) {
@@ -201,7 +195,7 @@ private void emitSymbolInformation(Element sym, Tree tree) {
 
     Semanticdb.SymbolInformation info = builder.build();
 
-    symbolInfos.add(info);
+    documentBuilder.addSymbol(info);
   }
 
   void resolveNodes() {

semanticdb-kotlinc/src/main/kotlin/com/sourcegraph/semanticdb_kotlinc/SemanticdbTextDocumentBuilder.kt

@@ -3,6 +3,7 @@ package com.sourcegraph.semanticdb_kotlinc
 import com.sourcegraph.semanticdb.Semanticdb
 
 import com.sourcegraph.semanticdb.Semanticdb.SymbolOccurrence.Role
+import com.sourcegraph.semanticdb.SemanticdbDocumentBuilder
 import com.sourcegraph.semanticdb.SemanticdbPaths
 import java.nio.file.Path
 import java.nio.file.Paths
@@ -33,21 +34,12 @@ class SemanticdbTextDocumentBuilder(
     private val lineMap: LineMap,
     private val cache: SymbolsCache,
 ) {
-    private val occurrences = mutableListOf<Semanticdb.SymbolOccurrence>()
-    private val symbols = mutableListOf<Semanticdb.SymbolInformation>()
+    private val documentBuilder = SemanticdbDocumentBuilder()
     private val fileText = file.getContentsAsStream().reader().readText()
     private val semanticMd5 = semanticdbMD5()
 
-    fun build() = TextDocument {
-        this.text = fileText
-        this.uri = semanticdbURI()
-        this.md5 = semanticMd5
-        this.schema = Semanticdb.Schema.SEMANTICDB4
-        this.language = Semanticdb.Language.KOTLIN
-        occurrences.sortWith(compareBy({ it.range.startLine }, { it.range.startCharacter }))
-        this.addAllOccurrences(occurrences)
-        this.addAllSymbols(symbols)
-    }
+    fun build(): Semanticdb.TextDocument =
+        documentBuilder.build(Semanticdb.Language.KOTLIN, semanticdbURI(), fileText, semanticMd5)
 
     fun emitSemanticdbData(
         firBasedSymbol: FirBasedSymbol<*>?,
@@ -57,14 +49,10 @@ class SemanticdbTextDocumentBuilder(
         context: CheckerContext,
         enclosingSource: KtSourceElement? = null,
     ) {
-        symbolOccurrence(symbol, element, role, enclosingSource).let {
-            if (!occurrences.contains(it)) {
-                occurrences.add(it)
-            }
+        documentBuilder.addOccurrence(symbolOccurrence(symbol, element, role, enclosingSource))
+        if (role == Role.DEFINITION) {
+            documentBuilder.addSymbol(symbolInformation(firBasedSymbol, symbol, element, context))
         }
-        val symbolInformation = symbolInformation(firBasedSymbol, symbol, element, context)
-        if (role == Role.DEFINITION && !symbols.contains(symbolInformation))
-            symbols.add(symbolInformation)
     }
 
     @OptIn(SymbolInternals::class)

semanticdb-shared/src/main/java/com/sourcegraph/semanticdb/SemanticdbDocumentBuilder.java

@@ -0,0 +1,73 @@
+package com.sourcegraph.semanticdb;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Accumulator for {@link Semanticdb.SymbolOccurrence}/{@link Semanticdb.SymbolInformation} that
+ * assembles a final {@link Semanticdb.TextDocument}. First emission wins: occurrences are
+ * deduplicated by {@code (range, symbol, role)}, symbols by {@code symbol}; occurrences are sorted
+ * by start position. Contains no compiler-API dependencies.
+ */
+public final class SemanticdbDocumentBuilder {
+  private static final Comparator<Semanticdb.SymbolOccurrence> OCCURRENCE_ORDER =
+      Comparator.<Semanticdb.SymbolOccurrence>comparingInt(o -> o.getRange().getStartLine())
+          .thenComparingInt(o -> o.getRange().getStartCharacter());
+
+  private final Map<OccurrenceKey, Semanticdb.SymbolOccurrence> occurrences = new LinkedHashMap<>();
+  private final Map<String, Semanticdb.SymbolInformation> symbols = new LinkedHashMap<>();
+
+  public void addOccurrence(Semanticdb.SymbolOccurrence occurrence) {
+    occurrences.putIfAbsent(new OccurrenceKey(occurrence), occurrence);
+  }
+
+  public void addSymbol(Semanticdb.SymbolInformation symbol) {
+    symbols.putIfAbsent(symbol.getSymbol(), symbol);
+  }
+
+  public Semanticdb.TextDocument build(
+      Semanticdb.Language language, String uri, String text, String md5) {
+    List<Semanticdb.SymbolOccurrence> sortedOccurrences = new ArrayList<>(occurrences.values());
+    sortedOccurrences.sort(OCCURRENCE_ORDER);
+    return Semanticdb.TextDocument.newBuilder()
+        .setSchema(Semanticdb.Schema.SEMANTICDB4)
+        .setLanguage(language)
+        .setUri(uri)
+        .setText(text)
+        .setMd5(md5)
+        .addAllOccurrences(sortedOccurrences)
+        .addAllSymbols(symbols.values())
+        .build();
+  }
+
+  private static final class OccurrenceKey {
+    private final Semanticdb.Range range;
+    private final String symbol;
+    private final Semanticdb.SymbolOccurrence.Role role;
+
+    OccurrenceKey(Semanticdb.SymbolOccurrence occurrence) {
+      this.range = occurrence.hasRange() ? occurrence.getRange() : null;
+      this.symbol = occurrence.getSymbol();
+      this.role = occurrence.getRole();
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      if (this == other) return true;
+      if (!(other instanceof OccurrenceKey)) return false;
+      OccurrenceKey that = (OccurrenceKey) other;
+      return role == that.role
+          && Objects.equals(range, that.range)
+          && Objects.equals(symbol, that.symbol);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(range, symbol, role);
+    }
+  }
+}