Remember files opened as Parametric Rascal LSP before language registration

rascal-lsp/src/main/java/org/rascalmpl/vscode/lsp/parametric/ParametricTextDocumentService.java

@@ -28,7 +28,6 @@
 
 import java.io.IOException;
 import java.io.Reader;
-import java.net.URI;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -43,6 +42,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import org.apache.commons.lang3.tuple.Pair;
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -165,6 +165,7 @@ public class ParametricTextDocumentService implements IBaseTextDocumentService,
     private @MonotonicNonNull BaseWorkspaceService workspaceService;
 
     private final Map<ISourceLocation, TextDocumentState> files;
+    private final Map<ISourceLocation, Pair<TextDocumentItem, Long>> unregisteredFiles = new HashMap<>();
     private final ColumnMaps columns;
 
     /** extension to language */
@@ -280,6 +281,9 @@ public void didOpen(DidOpenTextDocumentParams params) {
         var timestamp = System.currentTimeMillis();
         logger.debug("Did Open file: {}", params.getTextDocument());
         TextDocumentState file = open(params.getTextDocument(), timestamp);
+        if (file == null) {
+            return;
+        }
         handleParsingErrors(file, file.getCurrentDiagnosticsAsync());
         triggerAnalyzer(params.getTextDocument(), NORMAL_DEBOUNCE);
     }
@@ -297,18 +301,21 @@ public void didSave(DidSaveTextDocumentParams params) {
         logger.debug("Did Save file: {}", params.getTextDocument());
         // on save we don't get new file contents, that already came in via didChange
         // but we do trigger the builder on save (if a builder exists)
-        triggerBuilder(params.getTextDocument());
+        if (isOpenAndRegistered(Locations.toLoc(params.getTextDocument()))) {
+            triggerBuilder(params.getTextDocument());
+        }
     }
 
     @Override
     public void didClose(DidCloseTextDocumentParams params) {
         logger.debug("Did Close file: {}", params.getTextDocument());
         var loc = Locations.toLoc(params.getTextDocument());
-        if (files.remove(loc) == null) {
+        if (files.remove(loc) != null) {
+            facts(loc).close(loc);
+        } else if (unregisteredFiles.remove(loc) == null) {
             throw new ResponseErrorException(new ResponseError(ResponseErrorCode.InternalError,
                 "Unknown file: " + loc, params));
         }
-        facts(loc).close(loc);
     }
 
     @Override
@@ -327,6 +334,7 @@ public void didDeleteFiles(DeleteFilesParams params) {
     private void triggerAnalyzer(TextDocumentItem doc, Duration delay) {
         triggerAnalyzer(new VersionedTextDocumentIdentifier(doc.getUri(), doc.getVersion()), delay);
     }
+
     private void triggerAnalyzer(VersionedTextDocumentIdentifier doc, Duration delay) {
         logger.trace("Triggering analyzer for {}", doc.getUri());
         var location = Locations.toLoc(doc);
@@ -343,12 +351,20 @@ private void triggerBuilder(TextDocumentIdentifier doc) {
         fileFacts.calculateBuilder(location, getFile(location).getCurrentTreeAsync(true));
     }
 
-    private TextDocumentState updateContents(VersionedTextDocumentIdentifier doc, String newContents, long timestamp) {
-        TextDocumentState file = getFile(Locations.toLoc(doc));
+    private void updateContents(VersionedTextDocumentIdentifier doc, String newContents, long timestamp) {
         logger.trace("New contents for {}", doc);
+        var loc = Locations.toLoc(doc);
+        if (!isOpenAndRegistered(loc)) {
+            var current = unregisteredFiles.get(loc);
+            var file = current.getLeft();
+            file.setText(newContents);
+            file.setVersion(doc.getVersion());
+            current.setValue(timestamp);
+            return;
+        }
+        TextDocumentState file = getFile(loc);
         columns.clear(file.getLocation());
         handleParsingErrors(file, file.update(doc.getVersion(), newContents, timestamp));
-        return file;
     }
 
     private void handleParsingErrors(TextDocumentState file, CompletableFuture<Versioned<List<Diagnostics.Template>>> diagnosticsAsync) {
@@ -366,6 +382,9 @@ private void handleParsingErrors(TextDocumentState file, CompletableFuture<Versi
     public CompletableFuture<List<? extends CodeLens>> codeLens(CodeLensParams params) {
         logger.trace("codeLens for: {}", params.getTextDocument().getUri());
         ISourceLocation loc = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
         TextDocumentState file = getFile(loc);
         ILanguageContributions contrib = contributions(loc);
 
@@ -385,6 +404,9 @@ public CompletableFuture<Either3<Range, PrepareRenameResult, PrepareRenameDefaul
             PrepareRenameParams params) {
         logger.trace("prepareRename for: {}", params.getTextDocument().getUri());
         ISourceLocation location = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(location)) {
+            return CompletableFuture.completedFuture(null); // intentional `null` - invalid rename
+        }
         ILanguageContributions contribs = contributions(location);
         Position pos = params.getPosition();
         return getFile(location)
@@ -413,6 +435,9 @@ private CompletableFuture<ISourceLocation> computeRenameRange(final ILanguageCon
     public CompletableFuture<WorkspaceEdit> rename(RenameParams params) {
         logger.trace("rename for: {}, new name: {}", params.getTextDocument().getUri(), params.getNewName());
         ISourceLocation loc = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(null); // intentional `null` - invalid rename
+        }
         ILanguageContributions contribs = contributions(loc);
         Position rascalPos = Locations.toRascalPosition(loc, params.getPosition(), columns);
         return getFile(loc)
@@ -539,6 +564,9 @@ private IConstructor fileRenameToDocumentEdit(FileRename rename) {
     public CompletableFuture<List<InlayHint>> inlayHint(InlayHintParams params) {
         logger.trace("inlayHint for: {}", params.getTextDocument().getUri());
         ISourceLocation loc = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
         TextDocumentState file = getFile(loc);
         ILanguageContributions contrib = contributions(loc);
         return recoverExceptions(file.getLastTreeAsync(false)
@@ -592,38 +620,66 @@ private static <T> T last(List<T> l) {
         return l.get(l.size() - 1);
     }
 
+    private boolean isLanguageRegistered(ISourceLocation loc) {
+        return registeredExtensions.containsKey(extension(loc));
+    }
+
+    private String language(ISourceLocation loc) {
+        var ext = extension(loc);
+        var language = registeredExtensions.get(ext);
+        if (language == null) {
+            throw new UnsupportedOperationException(String.format("Rascal Parametric LSP has no support for this file, since no language is registered for extension '%s': %s", ext, loc));
+        }
+        return language;
+    }
+
     private ILanguageContributions contributions(ISourceLocation doc) {
-        String language = registeredExtensions.get(extension(doc));
-        if (language != null) {
-            ILanguageContributions contrib = contributions.get(language);
+        ILanguageContributions contrib = contributions.get(language(doc));
 
-            if (contrib != null) {
-                return contrib;
-            }
+        if (contrib == null) {
+            throw new UnsupportedOperationException("Rascal Parametric LSP has no support for this file: " + doc);
         }
 
-        throw new UnsupportedOperationException("Rascal Parametric LSP has no support for this file: " + doc);
+        return contrib;
     }
 
     private static String extension(ISourceLocation doc) {
         return URIUtil.getExtension(doc);
     }
 
     private ParametricFileFacts facts(ISourceLocation doc) {
-        String language = registeredExtensions.get(extension(doc));
-        if (language != null) {
-            ParametricFileFacts fact = facts.get(language);
-            if (fact != null) {
-                return fact;
-            }
+        ParametricFileFacts fact = facts.get(language(doc));
+
+        if (fact == null) {
+            throw new UnsupportedOperationException("Rascal Parametric LSP has no support for this file: " + doc);
         }
-        throw new UnsupportedOperationException("Rascal Parametric LSP has no support for this file: " + doc);
+
+        return fact;
     }
 
-    private TextDocumentState open(TextDocumentItem doc, long timestamp) {
-        return files.computeIfAbsent(Locations.toLoc(doc),
-            l -> new TextDocumentState(contributions(l)::parsing, l, doc.getVersion(), doc.getText(), timestamp)
-        );
+    private @Nullable TextDocumentState open(TextDocumentItem doc, long timestamp) {
+        var loc = Locations.toLoc(doc);
+        if (isLanguageRegistered(loc)) {
+            return files.computeIfAbsent(loc, l ->
+                new TextDocumentState(contributions(l)::parsing, l, doc.getVersion(), doc.getText(), timestamp));
+        } else {
+            // Language not (yet) supported; postpone state computation
+            // Let's remember this file so later, if someone registers the language for this, we can recover
+            unregisteredFiles.putIfAbsent(loc, Pair.of(doc, timestamp));
+            return null;
+        }
+    }
+
+    private boolean isOpenAndRegistered(ISourceLocation loc) {
+        if (files.containsKey(loc)) {
+            return true;
+        }
+        if (unregisteredFiles.containsKey(loc)) {
+            return false;
+        }
+
+        // This is unexpected; we somehow did not keep track of this file since opening it.
+        throw new ResponseErrorException(new ResponseError(ResponseErrorCode.RequestFailed, "Unknown file: " + loc, loc));
     }
 
     private TextDocumentState getFile(ISourceLocation loc) {
@@ -640,6 +696,9 @@ public void shutdown() {
 
     private CompletableFuture<SemanticTokens> getSemanticTokens(TextDocumentIdentifier doc) {
         var loc = Locations.toLoc(doc);
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(new SemanticTokens());
+        }
         var specialCaseHighlighting = contributions(loc).specialCaseHighlighting();
         return recoverExceptions(getFile(loc).getCurrentTreeAsync(true)
                 .thenApply(Versioned::get)
@@ -674,6 +733,9 @@ public CompletableFuture<List<Either<SymbolInformation, DocumentSymbol>>> docume
         logger.debug("Outline/documentSymbol: {}", params.getTextDocument());
 
         ISourceLocation location = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(location)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
         TextDocumentState file = getFile(location);
         ILanguageContributions contrib = contributions(location);
         return recoverExceptions(file.getCurrentTreeAsync(true)
@@ -689,6 +751,10 @@ public CompletableFuture<List<Either<Command, CodeAction>>> codeAction(CodeActio
         logger.debug("codeAction: {}", params);
 
         var location = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(location)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
+
         final ILanguageContributions contribs = contributions(location);
 
         var range = Locations.toRascalRange(location, params.getRange(), columns);
@@ -727,6 +793,9 @@ private CompletableFuture<IList> computeCodeActions(final ILanguageContributions
 
     private <T> CompletableFuture<List<T>> lookup(SummaryLookup<T> lookup, TextDocumentIdentifier doc, Position cursor) {
         var loc = Locations.toLoc(doc);
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
         return getFile(loc)
             .getCurrentTreeAsync(true)
             .thenApply(tree -> facts(loc).lookupInSummaries(lookup, loc, tree, cursor))
@@ -776,7 +845,11 @@ public CompletableFuture<Hover> hover(HoverParams params) {
     @Override
     public CompletableFuture<List<FoldingRange>> foldingRange(FoldingRangeRequestParams params) {
         logger.debug("Folding range: {}", params.getTextDocument());
-        TextDocumentState file = getFile(Locations.toLoc(params.getTextDocument()));
+        var loc = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
+        TextDocumentState file = getFile(loc);
         return recoverExceptions(file.getCurrentTreeAsync(true).thenApply(Versioned::get).thenApplyAsync(FoldingRanges::getFoldingRanges)
             .whenComplete((r, e) ->
                 logger.trace("Folding regions success, reporting {} regions back", r == null ? 0 : r.size())
@@ -787,6 +860,9 @@ public CompletableFuture<List<FoldingRange>> foldingRange(FoldingRangeRequestPar
     public CompletableFuture<List<SelectionRange>> selectionRange(SelectionRangeParams params) {
         logger.debug("Selection range: {} at {}", params.getTextDocument(), params.getPositions());
         ISourceLocation loc = Locations.toLoc(params.getTextDocument());
+        if (!isOpenAndRegistered(loc)) {
+            return CompletableFuture.completedFuture(Collections.emptyList());
+        }
         ILanguageContributions contrib = contributions(loc);
         TextDocumentState file = getFile(loc);
 
@@ -841,6 +917,21 @@ public synchronized void registerLanguage(LanguageParameter lang) {
         multiplexer.addContributor(buildContributionKey(lang),
             new InterpretedLanguageContributions(lang, this, workspaceService, (IBaseLanguageClient) client, ownExecuter));
 
+        for (var extension: lang.getExtensions()) {
+            // If we opened any files with this extension before, now associate them with contributions
+            for (var f : unregisteredFiles.keySet()) {
+                if (extension(f).equals(extension)) {
+                    var opened = unregisteredFiles.remove(f);
+                    var doc = opened.getLeft();
+                    var file = new TextDocumentState(contributions(f)::parsing, f, doc.getVersion(), doc.getText(), opened.getRight());
+                    files.put(f, file);
+                    // Update open editor
+                    triggerAnalyzer(doc, NORMAL_DEBOUNCE);
+                    handleParsingErrors(file, file.getCurrentDiagnosticsAsync());
+                }
+            }
+        }
+
         fact.reloadContributions();
         if (client != null) {
             fact.setClient(client);