vscode: add show tokens to get lexer output (#579)

Author: sbdchd

.vscode/settings.json

@@ -6,13 +6,10 @@
   "[rust]": {
     "editor.defaultFormatter": "rust-lang.rust-analyzer"
   },
-  "files.associations": {
-    "*.sql": "postgres"
-  },
-  "[postgres]": {
+  "[ungrammar]": {
     "editor.tabSize": 2
   },
-  "[ungrammar]": {
+  "[sql]": {
     "editor.tabSize": 2
   }
 }

Cargo.lock

@@ -15,9 +15,10 @@ lsp-server.workspace = true
 lsp-types.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+squawk_lexer.workspace = true
 squawk_linter.workspace = true
 squawk_syntax.workspace = true
 line-index.workspace = true
 
 [lints]
-workspace = true
+workspace = true

crates/squawk_server/Cargo.toml

@@ -58,17 +58,20 @@ fn main_loop(connection: Connection, params: serde_json::Value) -> Result<()> {
                     return Ok(());
                 }
 
-                if req.method == GotoDefinition::METHOD {
-                    handle_goto_definition(&connection, req)?;
-                    continue;
+                match req.method.as_ref() {
+                    GotoDefinition::METHOD => {
+                        handle_goto_definition(&connection, req)?;
+                    }
+                    "squawk/syntaxTree" => {
+                        handle_syntax_tree(&connection, req)?;
+                    }
+                    "squawk/tokens" => {
+                        handle_tokens(&connection, req)?;
+                    }
+                    _ => {
+                        info!("Ignoring unhandled request: {}", req.method);
+                    }
                 }
-
-                if req.method == "squawk/syntaxTree" {
-                    handle_syntax_tree(&connection, req)?;
-                    continue;
-                }
-
-                info!("Ignoring unhandled request: {}", req.method);
             }
             Message::Response(resp) => {
                 info!("Received response: id={:?}", resp.id);
@@ -269,3 +272,46 @@ fn handle_syntax_tree(connection: &Connection, req: lsp_server::Request) -> Resu
     connection.sender.send(Message::Response(resp))?;
     Ok(())
 }
+
+#[derive(serde::Deserialize)]
+struct TokensParams {
+    #[serde(rename = "textDocument")]
+    text_document: lsp_types::TextDocumentIdentifier,
+    // TODO: once we start storing the text doc on the server we won't need to
+    // send the content across the wire
+    text: String,
+}
+
+fn handle_tokens(connection: &Connection, req: lsp_server::Request) -> Result<()> {
+    let params: TokensParams = serde_json::from_value(req.params)?;
+    let uri = params.text_document.uri;
+    let content = params.text;
+
+    info!("Generating tokens for: {}", uri);
+
+    let tokens = squawk_lexer::tokenize(&content);
+
+    let mut output = Vec::new();
+    let mut char_pos = 0;
+    for token in tokens {
+        let token_start = char_pos;
+        let token_end = token_start + token.len as usize;
+        let token_text = &content[token_start..token_end];
+        output.push(format!(
+            "{:?}@{}..{} {:?}",
+            token.kind, token_start, token_end, token_text
+        ));
+        char_pos = token_end;
+    }
+
+    let tokens_output = output.join("\n");
+
+    let resp = Response {
+        id: req.id,
+        result: Some(serde_json::to_value(&tokens_output).unwrap()),
+        error: None,
+    };
+
+    connection.sender.send(Message::Response(resp))?;
+    Ok(())
+}

crates/squawk_server/src/lib.rs

@@ -47,6 +47,11 @@
         "title": "Show Syntax Tree",
         "category": "Squawk"
       },
+      {
+        "command": "squawk.showTokens",
+        "title": "Show Tokens",
+        "category": "Squawk"
+      },
       {
         "command": "squawk.showClientLogs",
         "title": "Show Client Logs",

squawk-vscode/package.json

@@ -31,6 +31,14 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
+  const tokensProvider = new TokensProvider(context)
+  context.subscriptions.push(
+    vscode.workspace.registerTextDocumentContentProvider(
+      "squawk-tokens",
+      tokensProvider,
+    ),
+  )
+
   context.subscriptions.push(
     vscode.commands.registerCommand("squawk.serverVersion", () => {
       try {
@@ -82,9 +90,7 @@ export async function activate(context: vscode.ExtensionContext) {
   await startServer(context)
 }
 
-export async function deactivate() {
-  await client?.stop()
-}
+export async function deactivate() {}
 
 function isSqlDocument(document: vscode.TextDocument): boolean {
   return document.languageId === "sql" || document.languageId === "postgres"
@@ -166,14 +172,17 @@ function getSquawkPath(context: vscode.ExtensionContext): vscode.Uri {
 }
 
 async function startServer(context: vscode.ExtensionContext) {
-  if (client?.state === State.Running) {
-    log.info("Server is already running")
-    return
-  }
-
-  if (client?.state === State.Starting) {
-    log.info("Server is already starting")
-    return
+  const state = client?.state
+  switch (state) {
+    case State.Running:
+    case State.Starting:
+      log.info("Server is already running")
+      break
+    case State.Stopped:
+    case undefined:
+      break
+    default:
+      assertNever(state)
   }
 
   log.info("Starting Squawk Language Server...")
@@ -212,6 +221,7 @@ async function startServer(context: vscode.ExtensionContext) {
       onClientStateChange.fire(event)
     }),
   )
+  context.subscriptions.push(client)
 
   log.info("server starting...")
   try {
@@ -315,6 +325,78 @@ class SyntaxTreeProvider implements vscode.TextDocumentContentProvider {
     }
   }
 }
+
+class TokensProvider implements vscode.TextDocumentContentProvider {
+  _eventEmitter = new vscode.EventEmitter<vscode.Uri>()
+  _activeEditor: vscode.TextEditor | undefined
+  _uri = vscode.Uri.parse("squawk-tokens://tokens/tokens.rast")
+
+  constructor(context: vscode.ExtensionContext) {
+    context.subscriptions.push(
+      vscode.window.onDidChangeActiveTextEditor((editor) => {
+        this._onDidChangeActiveTextEditor(editor)
+      }),
+    )
+    context.subscriptions.push(
+      vscode.workspace.onDidChangeTextDocument((event) => {
+        this._onDidChangeTextDocument(event.document)
+      }),
+    )
+    context.subscriptions.push(
+      vscode.commands.registerCommand("squawk.showTokens", async () => {
+        const doc = await vscode.workspace.openTextDocument(this._uri)
+        await vscode.window.showTextDocument(doc, vscode.ViewColumn.Beside)
+      }),
+    )
+
+    // initial kick off to make sure we have the editor set
+    this._onDidChangeActiveTextEditor(vscode.window.activeTextEditor)
+  }
+
+  onDidChange = this._eventEmitter.event
+
+  _onDidChangeActiveTextEditor(editor: vscode.TextEditor | undefined) {
+    if (editor && isSqlEditor(editor)) {
+      this._activeEditor = editor
+      this._eventEmitter.fire(this._uri)
+    }
+  }
+
+  _onDidChangeTextDocument(document: vscode.TextDocument) {
+    if (
+      isSqlDocument(document) &&
+      this._activeEditor &&
+      document === this._activeEditor.document
+    ) {
+      this._eventEmitter.fire(this._uri)
+    }
+  }
+
+  async provideTextDocumentContent(_uri: vscode.Uri): Promise<string> {
+    try {
+      const document = this._activeEditor?.document
+      if (!document) {
+        return "Error: no active editor found"
+      }
+      if (!client) {
+        return "Error: no client found"
+      }
+      const text = document.getText()
+      const uri = document.uri.toString()
+      log.info(`Requesting tokens for: ${uri}`)
+      const response = await client.sendRequest<string>("squawk/tokens", {
+        textDocument: { uri },
+        text,
+      })
+      log.info("Tokens received")
+      return response
+    } catch (error) {
+      log.error(`Failed to get tokens:`, error)
+      return `Error: Failed to get tokens: ${String(error)}`
+    }
+  }
+}
+
 function assertNever(param: never): never {
   throw new Error(`should never get here, but got ${String(param)}`)
 }