Fix ast viewer not working with synthetic files

kralicky · Oct 27, 2023 · b0e6957 · b0e6957
1 parent 3a4ba5a
commit b0e6957
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 13 deletions.
diff --git a/editors/vscode/client/src/astviewer.ts b/editors/vscode/client/src/astviewer.ts
@@ -14,8 +14,9 @@ export class ASTViewer implements vscode.TextDocumentContentProvider {
           this.virtualUrisByFile.delete(doc.uri.toString())
           break
         }
+        case "protoast2":
         case "protoast": {
-          this.virtualUrisByFile.delete(fromProtoAstURi(doc.uri).toString())
+          this.virtualUrisByFile.delete(fromProtoAstUri(doc.uri).toString())
           break
         }
       }
@@ -57,7 +58,7 @@ export class ASTViewer implements vscode.TextDocumentContentProvider {
     uri: vscode.Uri,
     token: vscode.CancellationToken,
   ): vscode.ProviderResult<string> {
-    const fileUri = fromProtoAstURi(uri)
+    const fileUri = fromProtoAstUri(uri)
     if (!this.virtualUrisByFile.has(fileUri.toString())) {
       return ""
     }
@@ -71,15 +72,33 @@ export class ASTViewer implements vscode.TextDocumentContentProvider {
 }
 
 export function toProtoAstUri(uri: vscode.Uri): vscode.Uri {
-  return uri.with({
-    scheme: "protoast",
-    path: uri.path + " [AST]",
-  })
+  if (uri.scheme === "file") {
+    return uri.with({
+      scheme: "protoast",
+      path: uri.path + " [AST]",
+    })
+  } else if (uri.scheme === "proto") {
+    return uri.with({
+      scheme: "protoast2",
+      path: uri.path + " [AST]",
+    })
+  } else {
+    return uri
+  }
 }
 
-export function fromProtoAstURi(uri: vscode.Uri): vscode.Uri {
-  return uri.with({
-    scheme: "file",
-    path: uri.path.replace(/ \[AST\]$/, ""),
-  })
+export function fromProtoAstUri(uri: vscode.Uri): vscode.Uri {
+  if (uri.scheme === "protoast") {
+    return uri.with({
+      scheme: "file",
+      path: uri.path.replace(/ \[AST\]$/, ""),
+    })
+  } else if (uri.scheme === "protoast2") {
+    return uri.with({
+      scheme: "proto",
+      path: uri.path.replace(/ \[AST\]$/, ""),
+    })
+  } else {
+    return uri
+  }
 }
diff --git a/editors/vscode/client/src/extension.ts b/editors/vscode/client/src/extension.ts
@@ -1,6 +1,6 @@
 import * as vscode from "vscode"
 import { LanguageClient } from "vscode-languageclient/node"
-import { ASTViewer, fromProtoAstURi } from "./astviewer"
+import { ASTViewer, fromProtoAstUri } from "./astviewer"
 import { buildLanguageClient } from "./client"
 import { initCommands } from "./commands"
 
@@ -13,12 +13,19 @@ export async function activate(context: vscode.ExtensionContext) {
   client.start()
 
   const astViewer = new ASTViewer((uri) => {
-    return client.sendRequest("protols/ast", fromProtoAstURi(uri).toString())
+    return client.sendRequest("protols/ast", fromProtoAstUri(uri).toString())
   })
   context.subscriptions.push(
     astViewer,
     vscode.workspace.registerTextDocumentContentProvider("protoast", astViewer),
   )
+  context.subscriptions.push(
+    astViewer,
+    vscode.workspace.registerTextDocumentContentProvider(
+      "protoast2",
+      astViewer,
+    ),
+  )
 
   context.subscriptions.push(
     vscode.commands.registerCommand("protols.restart", async () => {