diff --git a/crates/squawk_wasm/src/lib.rs b/crates/squawk_wasm/src/lib.rs index 88e0105c..d6f0cc42 100644 --- a/crates/squawk_wasm/src/lib.rs +++ b/crates/squawk_wasm/src/lib.rs @@ -102,6 +102,45 @@ impl SemanticTokenEncoder { } } +fn encode_semantic_tokens(db: &Database, file: File) -> Vec { + 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, + formatted_file: Option, } #[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, 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 { + 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 { diff --git a/playground/src/App.tsx b/playground/src/App.tsx index 485bc85f..b5af7295 100644 --- a/playground/src/App.tsx +++ b/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() diff --git a/playground/src/providers.tsx b/playground/src/providers.tsx index 32d85226..0a008195 100644 --- a/playground/src/providers.tsx +++ b/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, diff --git a/playground/src/squawk.tsx b/playground/src/squawk.tsx index 138b42b5..93184796 100644 --- a/playground/src/squawk.tsx +++ b/playground/src/squawk.tsx @@ -35,6 +35,14 @@ 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 @@ -42,10 +50,7 @@ let db: SquawkDatabase | null = null // // 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() }