Feat/lsp improvements

src/lib.rs

@@ -10,6 +10,7 @@ use tx3_lang::Protocol;
 
 mod cmds;
 mod server;
+mod visitor;
 
 #[derive(Error, Debug)]
 pub enum Error {
@@ -62,6 +63,22 @@ pub fn char_index_to_line_col(rope: &Rope, idx: usize) -> (usize, usize) {
     (line, col)
 }
 
+pub fn position_to_offset(text: &str, position: Position) -> usize {
+    let mut offset = 0;
+    for (line_num, line) in text.lines().enumerate() {
+        if line_num == position.line as usize {
+            offset += position.character.min(line.len() as u32) as usize;
+            break;
+        }
+        offset += line.len() + 1;
+    }
+    offset
+}
+
+pub fn span_contains(span: &tx3_lang::ast::Span, offset: usize) -> bool {
+    offset >= span.start && offset < span.end
+}
+
 pub fn span_to_lsp_range(rope: &Rope, loc: &tx3_lang::ast::Span) -> Range {
     let (start_line, start_col) = char_index_to_line_col(rope, loc.start);
     let (end_line, end_col) = char_index_to_line_col(rope, loc.end);
@@ -117,6 +134,166 @@ pub struct Context {
 }
 
 impl Context {
+    fn is_type_field_reference(
+        ast: &tx3_lang::ast::Program,
+        identifier: &str,
+        offset: usize,
+    ) -> bool {
+        for type_def in &ast.types {
+            if crate::span_contains(&type_def.span, offset) {
+                for case in &type_def.cases {
+                    for field in &case.fields {
+                        if identifier == field.r#type.clone().to_str() {
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
+        false
+    }
+    fn collect_semantic_tokens(
+        &self,
+        ast: &tx3_lang::ast::Program,
+        rope: &Rope,
+    ) -> Vec<SemanticToken> {
+        const TOKEN_TYPE: u32 = 0;
+        const TOKEN_PARAMETER: u32 = 1;
+        const TOKEN_VARIABLE: u32 = 2;
+        const TOKEN_CLASS: u32 = 3;
+        const TOKEN_PARTY: u32 = 4;
+        const TOKEN_POLICY: u32 = 5;
+        const TOKEN_FUNCTION: u32 = 6;
+        // const TOKEN_KEYWORD: u32 = 7;
+        // const TOKEN_PROPERTY: u32 = 8;
+
+        const MOD_DECLARATION: u32 = 1 << 0;
+        const MOD_DEFINITION: u32 = 1 << 1;
+
+        #[derive(Debug, Clone)]
+        struct TokenInfo {
+            range: Range,
+            token_type: u32,
+            token_modifiers: u32,
+        }
+
+        let mut token_infos: Vec<TokenInfo> = Vec::new();
+        let text = rope.to_string();
+
+        let mut processed_spans = std::collections::HashSet::new();
+
+        for offset in 0..text.len() {
+            if let Some(symbol) = crate::visitor::find_symbol_in_program(ast, offset) {
+                match symbol {
+                    crate::visitor::SymbolAtOffset::Identifier(identifier) => {
+                        // Skip if we've already processed this exact span
+                        let span_key = (identifier.span.start, identifier.span.end);
+                        if processed_spans.contains(&span_key) {
+                            continue;
+                        }
+                        processed_spans.insert(span_key);
+
+                        let token_type = if ast
+                            .parties
+                            .iter()
+                            .any(|p| p.name.value == identifier.value)
+                        {
+                            TOKEN_PARTY
+                        } else if ast
+                            .policies
+                            .iter()
+                            .any(|p| p.name.value == identifier.value)
+                        {
+                            TOKEN_POLICY
+                        } else if ast.types.iter().any(|t| t.name.value == identifier.value) {
+                            TOKEN_TYPE
+                        } else if Context::is_type_field_reference(ast, &identifier.value, offset) {
+                            TOKEN_TYPE
+                        } else if ast.assets.iter().any(|a| a.name.value == identifier.value) {
+                            TOKEN_CLASS
+                        } else {
+                            let mut found_type = None;
+
+                            for tx in &ast.txs {
+                                if tx.name.value == identifier.value {
+                                    found_type = Some(TOKEN_FUNCTION);
+                                    break;
+                                }
+
+                                if crate::span_contains(&tx.span, offset) {
+                                    for param in &tx.parameters.parameters {
+                                        if param.name.value == identifier.value {
+                                            found_type = Some(TOKEN_PARAMETER);
+                                            break;
+                                        }
+                                    }
+                                }
+
+                                if found_type.is_some() {
+                                    break;
+                                }
+                            }
+                            found_type.unwrap_or(TOKEN_VARIABLE)
+                        };
+
+                        token_infos.push(TokenInfo {
+                            range: crate::span_to_lsp_range(rope, &identifier.span),
+                            token_type,
+                            token_modifiers: MOD_DECLARATION | MOD_DEFINITION,
+                        });
+                    }
+                    visitor::SymbolAtOffset::TypeIdentifier(x) => {
+                        token_infos.push(TokenInfo {
+                            range: crate::span_to_lsp_range(rope, &x.span),
+                            token_type: TOKEN_TYPE,
+                            token_modifiers: MOD_DECLARATION | MOD_DEFINITION,
+                        });
+                    }
+                }
+            }
+        }
+        token_infos.sort_by(|a, b| match a.range.start.line.cmp(&b.range.start.line) {
+            std::cmp::Ordering::Equal => a.range.start.character.cmp(&b.range.start.character),
+            other => other,
+        });
+
+        token_infos.dedup_by(|a, b| a.range.start == b.range.start && a.range.end == b.range.end);
+
+        let mut semantic_tokens = Vec::new();
+        let mut prev_line = 0;
+        let mut prev_start = 0;
+
+        for token_info in token_infos {
+            let line = token_info.range.start.line;
+            let start = token_info.range.start.character;
+            let length = token_info.range.end.character.saturating_sub(start);
+
+            if length == 0 {
+                continue;
+            }
+
+            let delta_line = line.saturating_sub(prev_line);
+            let delta_start = if delta_line == 0 {
+                start.saturating_sub(prev_start)
+            } else {
+                start
+            };
+
+            semantic_tokens.push(SemanticToken {
+                delta_line,
+                delta_start,
+                length,
+                token_type: token_info.token_type,
+                token_modifiers_bitset: token_info.token_modifiers,
+            });
+
+            prev_line = line;
+            prev_start = start;
+        }
+
+        semantic_tokens
+    }
+
     pub fn new_for_client(client: Client) -> Self {
         Self {
             client,