Add semantic highlighting for let-statements

examples/associated_functions.reid

@@ -28,6 +28,7 @@ fn main() -> u32 {
     let mut list = u64::malloc(15);
     list[4] = 17;
 
+
     print(from_str("value: ") + list[4]);
 
     return i32::sizeof() as u32;

reid-lsp/client/src/extension.ts

@@ -4,7 +4,7 @@
  * ------------------------------------------------------------------------------------------ */
 
 import * as path from 'path';
-import { workspace, ExtensionContext, window } from 'vscode';
+import { workspace, ExtensionContext, window, languages, SemanticTokensBuilder } from 'vscode';
 
 import {
 	Executable,
@@ -53,7 +53,7 @@ export function activate(context: ExtensionContext) {
 		synchronize: {
 			// Notify the server about file changes to '.clientrc files contained in the workspace
 			fileEvents: workspace.createFileSystemWatcher('**/.clientrc')
-		}
+		},
 	};
 
 	// Create the language client and start the client.
@@ -68,8 +68,22 @@ export function activate(context: ExtensionContext) {
 	client.info(`Loaded Reid Language Server from ${server_path}`);
 
 
-	workspace.onDidOpenTextDocument((e) => {
+	workspace.onDidOpenTextDocument((e) => { });
-	});
+
+	client.info("Registering semantic tokens provide");
+	context.subscriptions.push(languages.registerDocumentSemanticTokensProvider({
+		language: 'reid',
+		scheme: 'file'
+	}, {
+		provideDocumentSemanticTokens: () => {
+			client.info("hello!");
+			const builder = new SemanticTokensBuilder();
+			return builder.build();
+		}
+	}, {
+		tokenTypes: [],
+		tokenModifiers: [],
+	}));
 
 	// Start the client. This will also launch the server
 	client.start();

reid-lsp/src/analysis.rs

@@ -1,16 +1,33 @@
 use std::{collections::HashMap, fmt::format, path::PathBuf};
 
 use reid::{
-    ast::{self, FunctionDefinition, lexer::FullToken, token_stream::TokenRange},
+    ast::{
+        self, FunctionDefinition,
+        lexer::{FullToken, Token},
+        token_stream::TokenRange,
+    },
     codegen::intrinsics::get_intrinsic_assoc_functions,
     compile_module,
     error_raporting::{ErrorModules, ReidError},
     mir::{
-        self, Context, FunctionCall, FunctionParam, IfExpression, SourceModuleId, StructType, TypeKind, WhileStatement,
+        self, Context, FunctionCall, FunctionParam, IfExpression, Metadata, SourceModuleId, StructType, TypeKind,
-        typecheck::typerefs::TypeRefs,
+        WhileStatement, typecheck::typerefs::TypeRefs,
     },
     perform_all_passes,
 };
+use tower_lsp::lsp_types::SemanticTokenType;
+
+pub const TOKEN_LEGEND: [SemanticTokenType; 9] = [
+    SemanticTokenType::VARIABLE,
+    SemanticTokenType::FUNCTION,
+    SemanticTokenType::STRUCT,
+    SemanticTokenType::KEYWORD,
+    SemanticTokenType::NUMBER,
+    SemanticTokenType::STRING,
+    SemanticTokenType::OPERATOR,
+    SemanticTokenType::COMMENT,
+    SemanticTokenType::PROPERTY,
+];
 
 #[derive(Debug, Clone)]
 pub struct StaticAnalysis {
@@ -60,8 +77,22 @@ pub struct SymbolId(usize);
 
 #[derive(Debug, Clone)]
 pub struct AnalysisState {
+    /// TokenID -> Analysis map, containing SymbolIDs
     pub map: HashMap<usize, TokenAnalysis>,
-    pub symbol_table: Vec<SymbolKind>,
+    /// SymbolID -> Symbol
+    symbol_table: Vec<Symbol>,
+}
+
+impl AnalysisState {
+    pub fn get_symbol(&self, id: SymbolId) -> &Symbol {
+        self.symbol_table.get(id.0).unwrap()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Symbol {
+    pub kind: SemanticKind,
+    pub definition: usize,
 }
 
 impl AnalysisState {
@@ -93,12 +124,12 @@ impl AnalysisState {
         }
     }
 
-    pub fn set_symbol(&mut self, token_idx: usize, symbol: SymbolId) {
+    pub fn set_symbol(&mut self, idx: usize, symbol: SymbolId) {
-        if let Some(token) = self.map.get_mut(&token_idx) {
+        if let Some(token) = self.map.get_mut(&idx) {
             token.symbol = Some(symbol);
         } else {
             self.map.insert(
-                token_idx,
+                idx,
                 TokenAnalysis {
                     ty: None,
                     autocomplete: Vec::new(),
@@ -108,26 +139,66 @@ impl AnalysisState {
         }
     }
 
-    pub fn new_symbol(&mut self, kind: SymbolKind) -> SymbolId {
+    pub fn new_symbol(&mut self, definition: usize, kind: SemanticKind) -> SymbolId {
         let id = SymbolId(self.symbol_table.len());
-        self.symbol_table.push(kind);
+        self.symbol_table.push(Symbol { kind, definition });
         id
     }
 }
 
 pub struct AnalysisScope<'a> {
     state: &'a mut AnalysisState,
+    tokens: &'a Vec<FullToken>,
     variables: HashMap<String, SymbolId>,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+impl<'a> AnalysisScope<'a> {
-pub enum SymbolKind {
+    pub fn inner(&mut self) -> AnalysisScope {
-    Default,
+        AnalysisScope {
+            state: self.state,
+            tokens: self.tokens,
+            variables: self.variables.clone(),
+        }
+    }
+
+    pub fn token_idx<T: Copy>(&self, meta: &Metadata, pred: T) -> usize
+    where
+        T: FnOnce(&Token) -> bool,
+    {
+        for idx in meta.range.start..=meta.range.end {
+            if let Some(token) = self.tokens.get(idx) {
+                if pred(&token.token) {
+                    return idx;
+                }
+            }
+        }
+        return meta.range.end;
+    }
 }
 
-impl Default for SymbolKind {
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum SemanticKind {
+    Default,
+    Variable,
+}
+
+impl Default for SemanticKind {
     fn default() -> Self {
-        SymbolKind::Default
+        SemanticKind::Default
+    }
+}
+
+impl SemanticKind {
+    pub fn into_token_idx(&self) -> Option<u32> {
+        let token_type = match self {
+            SemanticKind::Variable => SemanticTokenType::VARIABLE,
+            SemanticKind::Default => return None,
+        };
+        TOKEN_LEGEND
+            .iter()
+            .enumerate()
+            .find(|(_, t)| token_type == **t)
+            .map(|(i, _)| i as u32)
     }
 }
 
@@ -174,6 +245,7 @@ pub fn analyze_context(context: &mir::Context, module: &mir::Module, error: Opti
 
     let mut scope = AnalysisScope {
         state: &mut state,
+        tokens: &module.tokens,
         variables: HashMap::new(),
     };
     for import in &module.imports {
@@ -272,6 +344,8 @@ pub fn analyze_block(
     block: &mir::Block,
     scope: &mut AnalysisScope,
 ) {
+    let scope = &mut scope.inner();
+
     for statement in &block.statements {
         match &statement.0 {
             mir::StmtKind::Let(named_variable_ref, _, expression) => {
@@ -282,7 +356,10 @@ pub fn analyze_block(
                         .ok()
                         .map(|(_, ty)| ty),
                 );
-                // return analyze_in_expr(&expression, module_id, token_idx);
+                let idx = scope.token_idx(&named_variable_ref.2, |t| matches!(t, Token::Identifier(_)));
+                let symbol = scope.state.new_symbol(idx, SemanticKind::Variable);
+                scope.state.set_symbol(idx, symbol);
+                scope.variables.insert(named_variable_ref.1.clone(), symbol);
             }
             mir::StmtKind::Set(lhs, rhs) => {
                 analyze_expr(context, source_module, lhs, scope);

reid-lsp/src/main.rs

@@ -3,18 +3,22 @@ use std::path::PathBuf;
 use dashmap::DashMap;
 use reid::ast::lexer::{FullToken, Position};
 use reid::error_raporting::{self, ErrorModules, ReidError};
-use reid::mir::{SourceModuleId, TypeKind};
+use reid::mir::SourceModuleId;
 use reid::parse_module;
 use tower_lsp::lsp_types::{
     self, CompletionItem, CompletionOptions, CompletionParams, CompletionResponse, Diagnostic, DiagnosticSeverity,
-    DidChangeTextDocumentParams, DidOpenTextDocumentParams, Hover, HoverContents, HoverParams, HoverProviderCapability,
+    DidChangeTextDocumentParams, DidOpenTextDocumentParams, DocumentFilter, DocumentSelector, Hover, HoverContents,
-    InitializeParams, InitializeResult, InitializedParams, MarkupContent, MarkupKind, MessageType, OneOf, Range,
+    HoverParams, HoverProviderCapability, InitializeParams, InitializeResult, InitializedParams, MarkupContent,
-    ServerCapabilities, TextDocumentItem, TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
+    MarkupKind, MessageType, OneOf, Range, SemanticToken, SemanticTokenModifier, SemanticTokenType,
-    WorkspaceFoldersServerCapabilities, WorkspaceServerCapabilities,
+    SemanticTokensLegend, SemanticTokensOptions, SemanticTokensParams, SemanticTokensRangeParams,
+    SemanticTokensRangeResult, SemanticTokensResult, SemanticTokensServerCapabilities, ServerCapabilities,
+    StaticRegistrationOptions, TextDocumentItem, TextDocumentRegistrationOptions, TextDocumentSyncCapability,
+    TextDocumentSyncKind, TextDocumentSyncOptions, WorkDoneProgressOptions, WorkspaceFoldersServerCapabilities,
+    WorkspaceServerCapabilities,
 };
 use tower_lsp::{Client, LanguageServer, LspService, Server, jsonrpc};
 
-use crate::analysis::{StaticAnalysis, analyze};
+use crate::analysis::{StaticAnalysis, TOKEN_LEGEND, analyze};
 
 mod analysis;
 
@@ -38,20 +42,43 @@ impl LanguageServer for Backend {
             will_save_wait_until: None,
             save: None,
         };
-        Ok(InitializeResult {
+
-            capabilities: ServerCapabilities {
+        let capabilities = ServerCapabilities {
-                hover_provider: Some(HoverProviderCapability::Simple(true)),
+            hover_provider: Some(HoverProviderCapability::Simple(true)),
-                completion_provider: Some(CompletionOptions { ..Default::default() }),
+            completion_provider: Some(CompletionOptions { ..Default::default() }),
-                text_document_sync: Some(TextDocumentSyncCapability::Options(sync)),
+            text_document_sync: Some(TextDocumentSyncCapability::Options(sync)),
-                workspace: Some(WorkspaceServerCapabilities {
+            workspace: Some(WorkspaceServerCapabilities {
-                    workspace_folders: Some(WorkspaceFoldersServerCapabilities {
+                workspace_folders: Some(WorkspaceFoldersServerCapabilities {
-                        supported: Some(true),
+                    supported: Some(true),
-                        change_notifications: Some(OneOf::Left(true)),
+                    change_notifications: Some(OneOf::Left(true)),
-                    }),
-                    file_operations: None,
                 }),
-                ..Default::default()
+                file_operations: None,
-            },
+            }),
+            semantic_tokens_provider: Some(SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
+                lsp_types::SemanticTokensRegistrationOptions {
+                    text_document_registration_options: TextDocumentRegistrationOptions {
+                        document_selector: Some(vec![DocumentFilter {
+                            language: Some("reid".to_owned()),
+                            scheme: Some("file".to_owned()),
+                            pattern: None,
+                        }]),
+                    },
+                    semantic_tokens_options: SemanticTokensOptions {
+                        work_done_progress_options: Default::default(),
+                        legend: SemanticTokensLegend {
+                            token_types: TOKEN_LEGEND.into(),
+                            token_modifiers: vec![],
+                        },
+                        range: None,
+                        full: Some(lsp_types::SemanticTokensFullOptions::Bool(true)),
+                    },
+                    static_registration_options: Default::default(),
+                },
+            )),
+            ..Default::default()
+        };
+        Ok(InitializeResult {
+            capabilities,
             ..Default::default()
         })
     }
@@ -82,7 +109,7 @@ impl LanguageServer for Backend {
             None
         };
 
-        dbg!(position, token);
+        // dbg!(position, token);
 
         let list = if let Some((idx, _)) = token {
             if let Some(analysis) = self.analysis.get(&file_name).unwrap().token_analysis.map.get(&idx) {
@@ -99,7 +126,7 @@ impl LanguageServer for Backend {
             Vec::new()
         };
 
-        dbg!(&list);
+        // dbg!(&list);
         Ok(Some(CompletionResponse::Array(list)))
     }
 
@@ -172,6 +199,100 @@ impl LanguageServer for Backend {
         })
         .await
     }
+
+    async fn semantic_tokens_full(
+        &self,
+        params: SemanticTokensParams,
+    ) -> jsonrpc::Result<Option<SemanticTokensResult>> {
+        let path = PathBuf::from(params.text_document.uri.path());
+        let file_name = path.file_name().unwrap().to_str().unwrap().to_owned();
+        let analysis = self.analysis.get(&file_name);
+
+        let mut semantic_tokens = Vec::new();
+        if let Some(analysis) = analysis {
+            let mut prev_line = 0;
+            let mut prev_start = 0;
+            for (i, token) in analysis.tokens.iter().enumerate() {
+                let delta_line = (token.position.1.max(1) - 1).max(prev_line) - prev_line;
+                let delta_start = if delta_line == 0 {
+                    (token.position.0.max(1) - 1).max(prev_start) - prev_start
+                } else {
+                    token.position.0.max(1) - 1
+                };
+
+                if let Some(token_analysis) = analysis.token_analysis.map.get(&i) {
+                    if let Some(symbol_id) = token_analysis.symbol {
+                        let symbol = analysis.token_analysis.get_symbol(symbol_id);
+                        if let Some(idx) = symbol.kind.into_token_idx() {
+                            let semantic_token = SemanticToken {
+                                delta_line,
+                                delta_start,
+                                length: token.token.len() as u32,
+                                token_type: idx,
+                                token_modifiers_bitset: 0,
+                            };
+                            semantic_tokens.push(semantic_token);
+                            dbg!(semantic_token, prev_line, prev_start, token);
+                            prev_line = token.position.1.max(1) - 1;
+                            prev_start = token.position.0.max(1) - 1;
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(Some(SemanticTokensResult::Tokens(lsp_types::SemanticTokens {
+            result_id: None,
+            data: semantic_tokens,
+        })))
+    }
+
+    async fn semantic_tokens_range(
+        &self,
+        params: SemanticTokensRangeParams,
+    ) -> jsonrpc::Result<Option<SemanticTokensRangeResult>> {
+        let path = PathBuf::from(params.text_document.uri.path());
+        let file_name = path.file_name().unwrap().to_str().unwrap().to_owned();
+        let analysis = self.analysis.get(&file_name);
+        dbg!("semantic_token_range");
+
+        let mut semantic_tokens = Vec::new();
+        if let Some(analysis) = analysis {
+            let mut prev_line = 0;
+            let mut prev_start = 0;
+            for (i, token) in analysis.tokens.iter().enumerate() {
+                let delta_line = token.position.1 - prev_line;
+                let delta_start = if delta_line == 0 {
+                    token.position.0
+                } else {
+                    token.position.0 - prev_start
+                };
+                prev_line = token.position.1;
+                prev_start = token.position.0;
+
+                if let Some(token_analysis) = analysis.token_analysis.map.get(&i) {
+                    if let Some(symbol_id) = token_analysis.symbol {
+                        let symbol = analysis.token_analysis.get_symbol(symbol_id);
+                        if let Some(idx) = symbol.kind.into_token_idx() {
+                            semantic_tokens.push(SemanticToken {
+                                delta_line,
+                                delta_start,
+                                length: token.token.len() as u32,
+                                token_type: idx,
+                                token_modifiers_bitset: 0,
+                            });
+                        }
+                    }
+                }
+            }
+        }
+
+        dbg!(&semantic_tokens);
+        Ok(Some(SemanticTokensRangeResult::Tokens(lsp_types::SemanticTokens {
+            result_id: None,
+            data: semantic_tokens,
+        })))
+    }
 }
 
 impl Backend {