playground: semantic syntax highlighting for formatted output

crates/squawk_wasm/src/lib.rs

@@ -102,6 +102,45 @@ impl SemanticTokenEncoder {
     }
 }
 
+fn encode_semantic_tokens(db: &Database, file: File) -> Vec<u32> {
+    let line_index = db::line_index(db, file);
+    let content = file.content(db);
+    let tokens = semantic_tokens(db, file, None);
+
+    let mut encoder = SemanticTokenEncoder::with_capacity(tokens.len());
+
+    // Duplicated from squawk-server, fyi
+    for token in &tokens {
+        // Taken from rust-analyzer, this solves the case where we have a
+        // multi line semantic token which isn't supported by the LSP spec.
+        // see: https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/rust-analyzer/src/lsp/to_proto.rs#L781C28-L781C28
+        for mut text_range in line_index.lines(token.range) {
+            if let Some((index, _)) = find_newline(&content[text_range]) {
+                text_range = TextRange::at(text_range.start(), TextSize::try_from(index).unwrap());
+            }
+            let start_lc = line_index.line_col(text_range.start());
+            let end_lc = line_index.line_col(text_range.end());
+            let start_wide = line_index
+                .to_wide(squawk_line_index::WideEncoding::Utf16, start_lc)
+                .unwrap();
+            let end_wide = line_index
+                .to_wide(squawk_line_index::WideEncoding::Utf16, end_lc)
+                .unwrap();
+
+            encoder.push(EncodedSemanticToken {
+                line: start_wide.line,
+                start: start_wide.col,
+                length: end_wide.col - start_wide.col,
+                token_type: token.token_type,
+                // TODO: once we get modifiers going, we'll need to update this
+                modifiers: 0,
+            });
+        }
+    }
+
+    encoder.finish()
+}
+
 #[wasm_bindgen(start)]
 pub fn run() {
     use log::Level;
@@ -122,6 +161,7 @@ pub fn run() {
 pub struct SquawkDatabase {
     db: Database,
     file: Option<File>,
+    formatted_file: Option<File>,
 }
 
 #[wasm_bindgen]
@@ -132,6 +172,7 @@ impl SquawkDatabase {
         SquawkDatabase {
             db: Database::default(),
             file: None,
+            formatted_file: None,
         }
     }
 
@@ -535,43 +576,22 @@ impl SquawkDatabase {
 
     pub fn semantic_tokens(&self) -> Result<Vec<u32>, Error> {
         let file = self.file()?;
-        let line_index = db::line_index(&self.db, file);
-        let content = file.content(&self.db);
-        let tokens = semantic_tokens(&self.db, file, None);
-
-        let mut encoder = SemanticTokenEncoder::with_capacity(tokens.len());
-
-        // Duplicated from squawk-server, fyi
-        for token in &tokens {
-            // Taken from rust-analyzer, this solves the case where we have a
-            // multi line semantic token which isn't supported by the LSP spec.
-            // see: https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/rust-analyzer/src/lsp/to_proto.rs#L781C28-L781C28
-            for mut text_range in line_index.lines(token.range) {
-                if let Some((index, _)) = find_newline(&content[text_range]) {
-                    text_range =
-                        TextRange::at(text_range.start(), TextSize::try_from(index).unwrap());
-                }
-                let start_lc = line_index.line_col(text_range.start());
-                let end_lc = line_index.line_col(text_range.end());
-                let start_wide = line_index
-                    .to_wide(squawk_line_index::WideEncoding::Utf16, start_lc)
-                    .unwrap();
-                let end_wide = line_index
-                    .to_wide(squawk_line_index::WideEncoding::Utf16, end_lc)
-                    .unwrap();
+        Ok(encode_semantic_tokens(&self.db, file))
+    }
 
-                encoder.push(EncodedSemanticToken {
-                    line: start_wide.line,
-                    start: start_wide.col,
-                    length: end_wide.col - start_wide.col,
-                    token_type: token.token_type,
-                    // TODO: once we get modifiers going, we'll need to update this
-                    modifiers: 0,
-                });
+    pub fn formatted_semantic_tokens(&mut self, content: String) -> Vec<u32> {
+        let file = match self.formatted_file {
+            Some(file) => {
+                file.set_content(&mut self.db).to(content.into());
+                file
             }
-        }
-
-        Ok(encoder.finish())
+            None => {
+                let file = File::new(&self.db, content.into());
+                self.formatted_file = Some(file);
+                file
+            }
+        };
+        encode_semantic_tokens(&self.db, file)
     }
 
     pub fn semantic_tokens_legend() -> Result<JsValue, Error> {

playground/src/App.tsx

@@ -32,6 +32,7 @@ import {
   provideSelectionRanges,
   provideCompletionItems,
   semanticTokensProvider,
+  formattedSemanticTokensProvider,
 } from "./providers"
 import { language as pgsqlMonarchLanguage } from "./pgsql"
 import { breakpoints, colors, transitions } from "./tokens.stylex"
@@ -591,6 +592,11 @@ function registerMonacoProvidersOnce() {
     "pgsql-formatted",
     pgsqlMonarchLanguage,
   )
+  const pgsqlFormattedSemanticTokensProvider =
+    monaco.languages.registerDocumentSemanticTokensProvider(
+      "pgsql-formatted",
+      formattedSemanticTokensProvider,
+    )
 
   monaco.languages.register({ id: "rast" })
   const tokenProvider = monaco.languages.setMonarchTokensProvider("rast", {
@@ -720,6 +726,7 @@ function registerMonacoProvidersOnce() {
     languageConfig.dispose()
     pgsqlTokenProvider.dispose()
     pgsqlFormattedTokenProvider.dispose()
+    pgsqlFormattedSemanticTokensProvider.dispose()
     codeActionProvider.dispose()
     hoverProvider.dispose()
     definitionProvider.dispose()

playground/src/providers.tsx

@@ -11,6 +11,7 @@ import {
   selection_ranges,
   semantic_tokens,
   semantic_tokens_legend,
+  formatted_semantic_tokens,
   DocumentSymbol,
 } from "./squawk"
 
@@ -351,6 +352,27 @@ export const semanticTokensProvider: monaco.languages.DocumentSemanticTokensProv
     releaseDocumentSemanticTokens() {},
   }
 
+export const formattedSemanticTokensProvider: monaco.languages.DocumentSemanticTokensProvider =
+  {
+    getLegend() {
+      return semantic_tokens_legend()
+    },
+    provideDocumentSemanticTokens(model) {
+      if (model.getLineCount() > 2000) return null
+      const content = model.getValue()
+      if (!content) return null
+
+      try {
+        const data = formatted_semantic_tokens(content)
+        return { data, resultId: undefined }
+      } catch (e) {
+        console.error("Error in provideDocumentSemanticTokens:", e)
+        return null
+      }
+    },
+    releaseDocumentSemanticTokens() {},
+  }
+
 export async function provideCompletionItems(
   model: monaco.editor.ITextModel,
   position: monaco.Position,

playground/src/squawk.tsx

@@ -35,17 +35,22 @@ export interface HoverResult {
 
 let db: SquawkDatabase | null = null
 
+function openDb(): SquawkDatabase {
+  if (db == null) {
+    db = new SquawkDatabase()
+    db.open_file("")
+  }
+  return db
+}
+
 // We pass in content and version here so that we:
 // 1. update the database
 // 2. so the react compiler doesn't just cache the functions at their initial
 //    results. We need them to be dependent on their input.
 //
 // We can probably do better than this.
 function getDb(content: string, version: number): SquawkDatabase {
-  if (db == null) {
-    db = new SquawkDatabase()
-    db.open_file(content)
-  }
+  const db = openDb()
   db.update_file(content, version)
 
   return db
@@ -135,6 +140,10 @@ export function semantic_tokens(content: string, version: number): Uint32Array {
   return getDb(content, version).semantic_tokens()
 }
 
+export function formatted_semantic_tokens(formatted: string): Uint32Array {
+  return openDb().formatted_semantic_tokens(formatted)
+}
+
 export function semantic_tokens_legend(): SemanticTokensLegend {
   return SquawkDatabase.semantic_tokens_legend()
 }