Add basic support for augmentsSyntaxTokens and non-standard semantic token config

crates/ide/src/syntax_highlighting.rs

@@ -469,25 +469,8 @@ fn traverse(
             }
 
             // apply config filtering
-            match &mut highlight.tag {
-                HlTag::StringLiteral if !config.strings => continue,
-                // If punctuation is disabled, make the macro bang part of the macro call again.
-                tag @ HlTag::Punctuation(HlPunct::MacroBang) => {
-                    if !config.macro_bang {
-                        *tag = HlTag::Symbol(SymbolKind::Macro);
-                    } else if !config.specialize_punctuation {
-                        *tag = HlTag::Punctuation(HlPunct::Other);
-                    }
-                }
-                HlTag::Punctuation(_) if !config.punctuation => continue,
-                tag @ HlTag::Punctuation(_) if !config.specialize_punctuation => {
-                    *tag = HlTag::Punctuation(HlPunct::Other);
-                }
-                HlTag::Operator(_) if !config.operator && highlight.mods.is_empty() => continue,
-                tag @ HlTag::Operator(_) if !config.specialize_operator => {
-                    *tag = HlTag::Operator(HlOperator::Other);
-                }
-                _ => (),
+            if !filter_by_config(&mut highlight, config) {
+                continue;
             }
 
             if inside_attribute {
@@ -498,3 +481,27 @@ fn traverse(
         }
     }
 }
+
+fn filter_by_config(highlight: &mut Highlight, config: HighlightConfig) -> bool {
+    match &mut highlight.tag {
+        HlTag::StringLiteral if !config.strings => return false,
+        // If punctuation is disabled, make the macro bang part of the macro call again.
+        tag @ HlTag::Punctuation(HlPunct::MacroBang) => {
+            if !config.macro_bang {
+                *tag = HlTag::Symbol(SymbolKind::Macro);
+            } else if !config.specialize_punctuation {
+                *tag = HlTag::Punctuation(HlPunct::Other);
+            }
+        }
+        HlTag::Punctuation(_) if !config.punctuation => return false,
+        tag @ HlTag::Punctuation(_) if !config.specialize_punctuation => {
+            *tag = HlTag::Punctuation(HlPunct::Other);
+        }
+        HlTag::Operator(_) if !config.operator && highlight.mods.is_empty() => return false,
+        tag @ HlTag::Operator(_) if !config.specialize_operator => {
+            *tag = HlTag::Operator(HlOperator::Other);
+        }
+        _ => (),
+    }
+    true
+}

crates/rust-analyzer/src/config.rs

@@ -483,6 +483,8 @@ config_data! {
         /// When enabled, rust-analyzer will highlight rust source in doc comments as well as intra
         /// doc links.
         semanticHighlighting_doc_comment_inject_enable: bool = "true",
+        /// Whether the server is allowed to emit non-standard tokens and modifiers.
+        semanticHighlighting_nonStandardTokens: bool = "true",
         /// Use semantic tokens for operators.
         ///
         /// When disabled, rust-analyzer will emit semantic tokens only for operator tokens when
@@ -1028,6 +1030,11 @@ impl Config {
         .is_some()
     }
 
+    pub fn semantics_tokens_augments_syntax_tokens(&self) -> bool {
+        try_!(self.caps.text_document.as_ref()?.semantic_tokens.as_ref()?.augments_syntax_tokens?)
+            .unwrap_or(false)
+    }
+
     pub fn position_encoding(&self) -> PositionEncoding {
         negotiated_encoding(&self.caps)
     }
@@ -1459,6 +1466,10 @@ impl Config {
         }
     }
 
+    pub fn highlighting_non_standard_tokens(&self) -> bool {
+        self.data.semanticHighlighting_nonStandardTokens
+    }
+
     pub fn highlighting_config(&self) -> HighlightConfig {
         HighlightConfig {
             strings: self.data.semanticHighlighting_strings_enable,

crates/rust-analyzer/src/handlers/request.rs

@@ -1472,7 +1472,13 @@ pub(crate) fn handle_semantic_tokens_full(
         snap.workspaces.is_empty() || !snap.proc_macros_loaded;
 
     let highlights = snap.analysis.highlight(highlight_config, file_id)?;
-    let semantic_tokens = to_proto::semantic_tokens(&text, &line_index, highlights);
+    let semantic_tokens = to_proto::semantic_tokens(
+        &text,
+        &line_index,
+        highlights,
+        snap.config.semantics_tokens_augments_syntax_tokens(),
+        snap.config.highlighting_non_standard_tokens(),
+    );
 
     // Unconditionally cache the tokens
     snap.semantic_tokens_cache.lock().insert(params.text_document.uri, semantic_tokens.clone());
@@ -1496,7 +1502,13 @@ pub(crate) fn handle_semantic_tokens_full_delta(
         snap.workspaces.is_empty() || !snap.proc_macros_loaded;
 
     let highlights = snap.analysis.highlight(highlight_config, file_id)?;
-    let semantic_tokens = to_proto::semantic_tokens(&text, &line_index, highlights);
+    let semantic_tokens = to_proto::semantic_tokens(
+        &text,
+        &line_index,
+        highlights,
+        snap.config.semantics_tokens_augments_syntax_tokens(),
+        snap.config.highlighting_non_standard_tokens(),
+    );
 
     let mut cache = snap.semantic_tokens_cache.lock();
     let cached_tokens = cache.entry(params.text_document.uri).or_default();
@@ -1530,7 +1542,13 @@ pub(crate) fn handle_semantic_tokens_range(
         snap.workspaces.is_empty() || !snap.proc_macros_loaded;
 
     let highlights = snap.analysis.highlight_range(highlight_config, frange)?;
-    let semantic_tokens = to_proto::semantic_tokens(&text, &line_index, highlights);
+    let semantic_tokens = to_proto::semantic_tokens(
+        &text,
+        &line_index,
+        highlights,
+        snap.config.semantics_tokens_augments_syntax_tokens(),
+        snap.config.highlighting_non_standard_tokens(),
+    );
     Ok(Some(semantic_tokens.into()))
 }
 

crates/rust-analyzer/src/semantic_tokens.rs

@@ -13,7 +13,7 @@ macro_rules! define_semantic_token_types {
             $($standard:ident),*$(,)?
         }
         custom {
-            $(($custom:ident, $string:literal)),*$(,)?
+            $(($custom:ident, $string:literal) $(=> $fallback:ident)?),*$(,)?
         }
 
     ) => {
@@ -24,6 +24,15 @@ macro_rules! define_semantic_token_types {
             $(SemanticTokenType::$standard,)*
             $($custom),*
         ];
+
+        pub(crate) fn standard_fallback_type(token: SemanticTokenType) -> Option<SemanticTokenType> {
+            $(
+                if token == $custom {
+                    None $(.or(Some(SemanticTokenType::$fallback)))?
+                } else
+            )*
+            { Some(token )}
+        }
     };
 }
 
@@ -51,42 +60,46 @@ define_semantic_token_types![
 
     custom {
         (ANGLE, "angle"),
-        (ARITHMETIC, "arithmetic"),
-        (ATTRIBUTE, "attribute"),
-        (ATTRIBUTE_BRACKET, "attributeBracket"),
-        (BITWISE, "bitwise"),
+        (ARITHMETIC, "arithmetic") => OPERATOR,
+        (ATTRIBUTE, "attribute") => DECORATOR,
+        (ATTRIBUTE_BRACKET, "attributeBracket") => DECORATOR,
+        (BITWISE, "bitwise") => OPERATOR,
         (BOOLEAN, "boolean"),
         (BRACE, "brace"),
         (BRACKET, "bracket"),
-        (BUILTIN_ATTRIBUTE, "builtinAttribute"),
+        (BUILTIN_ATTRIBUTE, "builtinAttribute") => DECORATOR,
         (BUILTIN_TYPE, "builtinType"),
-        (CHAR, "character"),
+        (CHAR, "character") => STRING,
         (COLON, "colon"),
         (COMMA, "comma"),
-        (COMPARISON, "comparison"),
+        (COMPARISON, "comparison") => OPERATOR,
         (CONST_PARAMETER, "constParameter"),
-        (DERIVE, "derive"),
-        (DERIVE_HELPER, "deriveHelper"),
+        (DERIVE, "derive") => DECORATOR,
+        (DERIVE_HELPER, "deriveHelper") => DECORATOR,
         (DOT, "dot"),
-        (ESCAPE_SEQUENCE, "escapeSequence"),
-        (FORMAT_SPECIFIER, "formatSpecifier"),
-        (GENERIC, "generic"),
+        (ESCAPE_SEQUENCE, "escapeSequence") => STRING,
+        (FORMAT_SPECIFIER, "formatSpecifier") => STRING,
+        (GENERIC, "generic") => TYPE_PARAMETER,
         (LABEL, "label"),
         (LIFETIME, "lifetime"),
-        (LOGICAL, "logical"),
-        (MACRO_BANG, "macroBang"),
+        (LOGICAL, "logical") => OPERATOR,
+        (MACRO_BANG, "macroBang") => MACRO,
         (PARENTHESIS, "parenthesis"),
         (PUNCTUATION, "punctuation"),
-        (SELF_KEYWORD, "selfKeyword"),
-        (SELF_TYPE_KEYWORD, "selfTypeKeyword"),
+        (SELF_KEYWORD, "selfKeyword") => KEYWORD,
+        (SELF_TYPE_KEYWORD, "selfTypeKeyword") => KEYWORD,
         (SEMICOLON, "semicolon"),
         (TYPE_ALIAS, "typeAlias"),
-        (TOOL_MODULE, "toolModule"),
+        (TOOL_MODULE, "toolModule") => DECORATOR,
         (UNION, "union"),
         (UNRESOLVED_REFERENCE, "unresolvedReference"),
     }
 ];
 
+macro_rules! count_tts {
+    () => {0usize};
+    ($_head:tt $($tail:tt)*) => {1usize + count_tts!($($tail)*)};
+}
 macro_rules! define_semantic_token_modifiers {
     (
         standard {
@@ -105,6 +118,8 @@ macro_rules! define_semantic_token_modifiers {
             $(SemanticTokenModifier::$standard,)*
             $($custom),*
         ];
+
+        const LAST_STANDARD_MOD: usize = count_tts!($($standard)*);
     };
 }
 
@@ -137,6 +152,13 @@ define_semantic_token_modifiers![
 #[derive(Default)]
 pub(crate) struct ModifierSet(pub(crate) u32);
 
+impl ModifierSet {
+    pub(crate) fn standard_fallback(&mut self) {
+        // Remove all non standard modifiers
+        self.0 = self.0 & !(!0u32 << LAST_STANDARD_MOD)
+    }
+}
+
 impl ops::BitOrAssign<SemanticTokenModifier> for ModifierSet {
     fn bitor_assign(&mut self, rhs: SemanticTokenModifier) {
         let idx = SUPPORTED_MODIFIERS.iter().position(|it| it == &rhs).unwrap();

crates/rust-analyzer/src/to_proto.rs

@@ -24,7 +24,7 @@ use crate::{
     line_index::{LineEndings, LineIndex, PositionEncoding},
     lsp_ext,
     lsp_utils::invalid_params_error,
-    semantic_tokens,
+    semantic_tokens::{self, standard_fallback_type},
 };
 
 pub(crate) fn position(line_index: &LineIndex, offset: TextSize) -> lsp_types::Position {
@@ -586,6 +586,8 @@ pub(crate) fn semantic_tokens(
     text: &str,
     line_index: &LineIndex,
     highlights: Vec<HlRange>,
+    semantics_tokens_augments_syntax_tokens: bool,
+    non_standard_tokens: bool,
 ) -> lsp_types::SemanticTokens {
     let id = TOKEN_RESULT_COUNTER.fetch_add(1, Ordering::SeqCst).to_string();
     let mut builder = semantic_tokens::SemanticTokensBuilder::new(id);
@@ -595,7 +597,35 @@ pub(crate) fn semantic_tokens(
             continue;
         }
 
-        let (ty, mods) = semantic_token_type_and_modifiers(highlight_range.highlight);
+        if semantics_tokens_augments_syntax_tokens {
+            match highlight_range.highlight.tag {
+                HlTag::BoolLiteral
+                | HlTag::ByteLiteral
+                | HlTag::CharLiteral
+                | HlTag::Comment
+                | HlTag::Keyword
+                | HlTag::NumericLiteral
+                | HlTag::Operator(_)
+                | HlTag::Punctuation(_)
+                | HlTag::StringLiteral
+                | HlTag::None
+                    if highlight_range.highlight.mods.is_empty() =>
+                {
+                    continue
+                }
+                _ => (),
+            }
+        }
+
+        let (mut ty, mut mods) = semantic_token_type_and_modifiers(highlight_range.highlight);
+
+        if !non_standard_tokens {
+            ty = match standard_fallback_type(ty) {
+                Some(ty) => ty,
+                None => continue,
+            };
+            mods.standard_fallback();
+        }
         let token_index = semantic_tokens::type_index(ty);
         let modifier_bitset = mods.0;
 

docs/user/generated_config.adoc

@@ -753,6 +753,11 @@ Inject additional highlighting into doc comments.
 When enabled, rust-analyzer will highlight rust source in doc comments as well as intra
 doc links.
 --
+[[rust-analyzer.semanticHighlighting.nonStandardTokens]]rust-analyzer.semanticHighlighting.nonStandardTokens (default: `true`)::
++
+--
+Whether the server is allowed to emit non-standard tokens and modifiers.
+--
 [[rust-analyzer.semanticHighlighting.operator.enable]]rust-analyzer.semanticHighlighting.operator.enable (default: `true`)::
 +
 --

editors/code/package.json

@@ -1395,6 +1395,11 @@
                     "default": true,
                     "type": "boolean"
                 },
+                "rust-analyzer.semanticHighlighting.nonStandardTokens": {
+                    "markdownDescription": "Whether the server is allowed to emit non-standard tokens and modifiers.",
+                    "default": true,
+                    "type": "boolean"
+                },
                 "rust-analyzer.semanticHighlighting.operator.enable": {
                     "markdownDescription": "Use semantic tokens for operators.\n\nWhen disabled, rust-analyzer will emit semantic tokens only for operator tokens when\nthey are tagged with modifiers.",
                     "default": true,