ls: Implement semantic tokens

thaliaarchi · thaliaarchi · commit bc04936869a1 · 2024-12-03T04:11:34.000-08:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/omniwsa-ls/Cargo.toml b/omniwsa-ls/Cargo.toml
@@ -10,7 +10,9 @@ keywords = ["whitespace", "language-server"]
 categories = ["compilers"]
 
 [dependencies]
+bstr = "1.11"
 lsp-server = "0.7"
 lsp-types = "0.97"
+omniwsa = { path = ".." }
 serde = "1.0"
 serde_json = "1.0"
diff --git a/omniwsa-ls/src/main.rs b/omniwsa-ls/src/main.rs
@@ -1,18 +1,32 @@
 use std::{
     error::Error,
+    fs,
     io::{self, ErrorKind},
     process::exit,
 };
 
+use bstr::ByteSlice;
 use lsp_server::{Connection, Message, Request, Response};
 use lsp_types::{
     request::{Request as _, SemanticTokensFullRequest},
     InitializeParams, SemanticToken, SemanticTokenModifier, SemanticTokenType, SemanticTokens,
     SemanticTokensFullOptions, SemanticTokensLegend, SemanticTokensOptions, SemanticTokensParams,
     SemanticTokensResult, SemanticTokensServerCapabilities, ServerCapabilities,
 };
+use omniwsa::{
+    dialects::{Dialect, Palaiologos},
+    tokens::{
+        comment::BlockCommentError,
+        string::{CharError, QuotedError, StringError},
+        Token,
+    },
+};
 use serde_json::{from_value as from_json, to_value as to_json};
 
+// TODO:
+// - Implement text document API, instead of reading from disk.
+// - Record spans in tokens.
+
 fn main() {
     if let Err(err) = do_main() {
         eprintln!("Error: {err}");
@@ -39,11 +53,13 @@ fn server_capabilities() -> ServerCapabilities {
                 work_done_progress_options: Default::default(),
                 legend: SemanticTokensLegend {
                     token_types: vec![
+                        SemanticTokenType::VARIABLE,
                         SemanticTokenType::FUNCTION,
                         SemanticTokenType::KEYWORD,
                         SemanticTokenType::COMMENT,
                         SemanticTokenType::STRING,
                         SemanticTokenType::NUMBER,
+                        SemanticTokenType::OPERATOR,
                     ],
                     token_modifiers: vec![
                         SemanticTokenModifier::DECLARATION,
@@ -58,40 +74,119 @@ fn server_capabilities() -> ServerCapabilities {
     }
 }
 
+// The order corresponds to the index in SemanticTokensLegend::token_types.
+#[derive(Clone, Copy, Debug)]
+enum TokenType {
+    Variable,
+    Function,
+    Keyword,
+    Comment,
+    String,
+    Number,
+    Operator,
+}
+
+// The order corresponds to the index in SemanticTokensLegend::token_modifiers.
+#[derive(Clone, Copy, Debug)]
+enum TokenModifier {
+    Declaration,
+    Definition,
+}
+
 fn main_loop(
     connection: Connection,
     _initialize_params: InitializeParams,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
     for msg in &connection.receiver {
         match msg {
             Message::Request(req) => {
-                eprintln!(
-                    "Received request {} #{}: {:?}",
-                    req.method, req.id, req.params,
-                );
+                eprintln!("Receive {req:?}");
                 if connection.handle_shutdown(&req)? {
                     return Ok(());
                 }
                 let Request { id, method, params } = req;
                 match &*method {
                     SemanticTokensFullRequest::METHOD => {
-                        let _params: SemanticTokensParams = from_json(params)?;
-                        let tokens = vec![SemanticToken {
-                            delta_line: 1,
-                            delta_start: 2,
-                            length: 3,
-                            token_type: 0,
-                            token_modifiers_bitset: 0,
-                        }];
+                        // https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens
+                        // https://code.visualstudio.com/api/language-extensions/semantic-highlight-guide
+                        // https://github.com/rust-lang/rust-analyzer/blob/master/crates/rust-analyzer/src/lsp/semantic_tokens.rs
+                        // https://github.com/rust-lang/rust-analyzer/blob/c4e040ea8dc4651569514bd8a8d8b293b49390a6/crates/rust-analyzer/src/lsp/capabilities.rs#L123
+
+                        let params: SemanticTokensParams = from_json(params)?;
+                        // TODO: Implement text document API, instead of reading
+                        // from disk.
+                        let path = params.text_document.uri.as_str();
+                        let path = path.strip_prefix("file://").unwrap_or(path);
+                        let src = fs::read(path)?;
+                        let tokens = Palaiologos::new().lex(&src);
+
+                        let mut tokens_out = Vec::with_capacity(tokens.len());
+                        let (mut curr_line, mut curr_col) = (0, 0);
+                        let (mut prev_line, mut prev_col) = (0, 0);
+                        for tok in &tokens {
+                            eprintln!("{tok:?}");
+                            let ungrouped = tok.ungroup();
+                            let ty = match ungrouped {
+                                Token::Mnemonic(_) => Some(TokenType::Keyword),
+                                Token::Integer(_) => Some(TokenType::Number),
+                                Token::String(_) | Token::Char(_) => Some(TokenType::String),
+                                Token::Variable(_) => Some(TokenType::Variable),
+                                Token::Label(_) => Some(TokenType::Function),
+                                Token::LabelColon(_) => Some(TokenType::Operator),
+                                Token::Space(_) | Token::LineTerm(_) | Token::Eof(_) => None,
+                                Token::InstSep(_) | Token::ArgSep(_) => Some(TokenType::Operator),
+                                Token::LineComment(_) | Token::BlockComment(_) => {
+                                    Some(TokenType::Comment)
+                                }
+                                Token::Word(_) => Some(TokenType::Variable),
+                                Token::Quoted(_) | Token::Spliced(_) => panic!("not ungrouped"),
+                                Token::Error(_) => None,
+                                Token::Placeholder => panic!("placeholder"),
+                            };
+                            let modifiers = match ungrouped {
+                                Token::Label(_) => TokenModifier::Declaration as _,
+                                Token::Variable(_) => TokenModifier::Definition as _,
+                                _ => 0,
+                            };
+                            let (len, hlen, vlen) = token_len(tok);
+                            let (mut next_line, mut next_col) = (curr_line, curr_col);
+                            if vlen != 0 {
+                                next_col = 0;
+                                next_line += vlen;
+                            }
+                            next_col += hlen;
+                            if let Some(ty) = ty {
+                                let token_out = SemanticToken {
+                                    delta_line: (curr_line - prev_line) as _,
+                                    delta_start: if curr_line == prev_line {
+                                        (curr_col - prev_col) as _
+                                    } else {
+                                        curr_col as _
+                                    },
+                                    length: len as _,
+                                    token_type: ty as _,
+                                    token_modifiers_bitset: modifiers,
+                                };
+                                eprintln!("=> {token_out:?}");
+                                tokens_out.push(token_out);
+                                prev_line = curr_line;
+                                prev_col = curr_col;
+                            }
+                            curr_line = next_line;
+                            curr_col = next_col;
+                        }
+
                         let result = Some(SemanticTokensResult::Tokens(SemanticTokens {
                             result_id: None,
-                            data: tokens,
+                            data: tokens_out,
                         }));
                         let resp = Response {
                             id,
                             result: Some(to_json(&result)?),
                             error: None,
                         };
+                        eprintln!("Send {resp:?}");
+                        eprintln!();
                         connection.sender.send(Message::Response(resp))?;
                     }
                     _ => {
@@ -103,12 +198,95 @@ fn main_loop(
                 }
             }
             Message::Response(resp) => {
-                eprintln!("Received response: {resp:?}");
+                eprintln!("Receive {resp:?}");
             }
             Message::Notification(notif) => {
-                eprintln!("Received notification: {notif:?}");
+                eprintln!("Receive {notif:?}");
             }
         }
     }
     Ok(())
 }
+
+/// Computes the length of the token. Returns the linear, horizontal, vertical
+/// lengths in chars.
+// TODO: Record spans in tokens instead of this hack.
+fn token_len(tok: &Token<'_>) -> (usize, usize, usize) {
+    let (text, len_before, len_after): (&[u8], usize, usize) = match tok {
+        Token::Mnemonic(tok) => (&tok.mnemonic, 0, 0),
+        Token::Integer(tok) => (&tok.literal, 0, 0),
+        Token::String(tok) => (
+            &tok.literal,
+            tok.quotes.quote().len(),
+            if tok.errors.contains(StringError::Unterminated) {
+                0
+            } else {
+                tok.quotes.quote().len()
+            },
+        ),
+        Token::Char(tok) => (
+            &tok.literal,
+            tok.quotes.quote().len(),
+            if tok.errors.contains(CharError::Unterminated) {
+                0
+            } else {
+                tok.quotes.quote().len()
+            },
+        ),
+        Token::Variable(tok) => (&tok.ident, tok.style.sigil().len(), 0),
+        Token::Label(tok) => (&tok.label, tok.style.sigil().len(), 0),
+        Token::LabelColon(_) => (b":", 0, 0),
+        Token::Space(tok) => (&tok.space, 0, 0),
+        Token::LineTerm(tok) => (tok.style.as_str().as_bytes(), 0, 0),
+        Token::Eof(_) => (b"", 0, 0),
+        Token::InstSep(tok) => (tok.style.as_str().as_bytes(), 0, 0),
+        Token::ArgSep(tok) => (tok.style.as_str().as_bytes(), 0, 0),
+        Token::LineComment(tok) => (tok.text, tok.style.prefix().len(), 0),
+        Token::BlockComment(tok) => (
+            tok.text,
+            tok.style.open().len(),
+            if tok.errors.contains(BlockCommentError::Unterminated) {
+                0
+            } else {
+                tok.style.close().len()
+            },
+        ),
+        Token::Word(tok) => (&tok.word, 0, 0),
+        Token::Quoted(tok) => {
+            let (len, hlen, vlen) = token_len(&tok.inner);
+            let mut quotes = 0;
+            if vlen != 0 {
+                quotes += tok.quotes.quote().len();
+            }
+            if !tok.errors.contains(QuotedError::Unterminated) {
+                quotes += tok.quotes.quote().len();
+            }
+            return (len + quotes, hlen + quotes, vlen);
+        }
+        Token::Spliced(tok) => {
+            let (mut spliced_len, mut spliced_hlen, mut spliced_vlen) = (0, 0, 0);
+            for tok in &tok.tokens {
+                let (len, hlen, vlen) = token_len(tok);
+                if vlen != 0 {
+                    spliced_hlen = 0;
+                    spliced_vlen += vlen;
+                }
+                spliced_len += len;
+                spliced_hlen += hlen;
+            }
+            return (spliced_len, spliced_hlen, spliced_vlen);
+        }
+        Token::Error(tok) => (&tok.text, 0, 0),
+        Token::Placeholder => panic!("placeholder"),
+    };
+    let (len, hlen, vlen) =
+        text.chars()
+            .fold((len_before, len_before, 0), |(len, hlen, vlen), ch| {
+                if ch == '\n' {
+                    (len + 1, 0, vlen + 1)
+                } else {
+                    (len + 1, hlen + 1, vlen)
+                }
+            });
+    (len + len_after, hlen + len_after, vlen)
+}
diff --git a/src/dialects/burghard/dialect.rs b/src/dialects/burghard/dialect.rs
@@ -2,13 +2,14 @@
 
 use crate::{
     dialects::{
-        burghard::{option::OptionNester, parse::Parser},
+        burghard::{lex::Lexer, option::OptionNester, parse::Parser},
         define_mnemonics,
         dialect::DialectState,
         Dialect,
     },
+    lex::Lex,
     syntax::Cst,
-    tokens::integer::IntegerSyntax,
+    tokens::{integer::IntegerSyntax, Token},
 };
 
 // TODO:
@@ -66,6 +67,19 @@ impl Dialect for Burghard {
         OptionNester::new().nest(&mut Parser::new(src, dialect))
     }
 
+    fn lex<'s>(src: &'s [u8], _dialect: &DialectState<Self>) -> Vec<Token<'s>> {
+        let mut lex = Lexer::new(src);
+        let mut toks = Vec::new();
+        loop {
+            let tok = lex.next_token();
+            if let Token::Eof(_) = tok {
+                break;
+            }
+            toks.push(tok);
+        }
+        toks
+    }
+
     fn make_integers() -> IntegerSyntax {
         IntegerSyntax::haskell()
     }
diff --git a/src/dialects/dialect.rs b/src/dialects/dialect.rs
@@ -34,9 +34,7 @@ pub trait Dialect {
     fn parse<'s>(src: &'s [u8], dialect: &DialectState<Self>) -> Cst<'s>;
 
     /// Lexes a Whitespace assembly program in the dialect.
-    fn lex<'s>(_src: &'s [u8], _dialect: &DialectState<Self>) -> Vec<Token<'s>> {
-        unimplemented!();
-    }
+    fn lex<'s>(src: &'s [u8], dialect: &DialectState<Self>) -> Vec<Token<'s>>;
 
     /// Constructs an integer syntax description for this dialect.
     fn make_integers() -> IntegerSyntax;
diff --git a/src/dialects/palaiologos/dialect.rs b/src/dialects/palaiologos/dialect.rs
@@ -1,9 +1,18 @@
 //! Parsing for the Palaiologos Whitespace assembly dialect.
 
 use crate::{
-    dialects::{define_mnemonics, dialect::DialectState, palaiologos::parse::Parser, Dialect},
+    dialects::{
+        define_mnemonics,
+        dialect::DialectState,
+        palaiologos::{lex::Lexer, parse::Parser},
+        Dialect,
+    },
+    lex::Lex,
     syntax::Cst,
-    tokens::integer::{BaseStyle, DigitSep, Integer, IntegerSyntax, SignStyle},
+    tokens::{
+        integer::{BaseStyle, DigitSep, Integer, IntegerSyntax, SignStyle},
+        Token,
+    },
 };
 
 /// Palaiologos Whitespace assembly dialect.
@@ -55,6 +64,19 @@ impl Dialect for Palaiologos {
         Parser::new(src, dialect).parse()
     }
 
+    fn lex<'s>(src: &'s [u8], dialect: &DialectState<Self>) -> Vec<Token<'s>> {
+        let mut lex = Lexer::new(src, dialect);
+        let mut toks = Vec::new();
+        loop {
+            let tok = lex.next_token();
+            if let Token::Eof(_) = tok {
+                break;
+            }
+            toks.push(tok);
+        }
+        toks
+    }
+
     /// Constructs an integer syntax description for this dialect.
     ///
     /// # Syntax
