swls_core/feature/
semantic.rs

1use bevy_ecs::{prelude::*, schedule::ScheduleLabel};
2use derive_more::{AsMut, AsRef, Deref, DerefMut};
3
4use crate::{
5    lsp_types::{SemanticToken, SemanticTokenType},
6    prelude::*,
7};
8
9/// [`Resource`] mapping a [`SemanticTokenType`] to their used index.
10#[derive(Resource, AsRef, Deref, AsMut, DerefMut, Debug, Default)]
11pub struct SemanticTokensDict(pub std::collections::HashMap<SemanticTokenType, usize>);
12
13/// [`Component`] indicating that the current document is handling a Highlight request.
14#[derive(Component, AsRef, Deref, AsMut, DerefMut, Debug)]
15pub struct HighlightRequest(pub Vec<SemanticToken>);
16
17#[derive(ScheduleLabel, Clone, Eq, PartialEq, Debug, Hash)]
18pub struct Label;
19
20pub fn setup_world(world: &mut World) {
21    // Each language registers its own basic_semantic_tokens::<L> in its own semantic schedule.
22    let mut semantic_tokens = bevy_ecs::schedule::Schedule::new(Label);
23    semantic_tokens.add_systems(semantic_tokens_system);
24    world.add_schedule(semantic_tokens);
25}
26
27struct TokenHelper {
28    start: usize,
29    length: usize,
30    ty: usize,
31}
32
33pub type TokenTypesComponent = Wrapped<Vec<Spanned<SemanticTokenType>>>;
34
35/// Generic CST-based semantic token extraction.  Register this for each language via:
36/// `world.schedule_scope(semantic::Label, |_, sched| { sched.add_systems(basic_semantic_tokens::<L>); })`
37pub fn basic_semantic_tokens<L: Lang + Component>(
38    mut query: Query<(Entity, &CstTokens, &Source), (With<HighlightRequest>, With<L>)>,
39    mut commands: Commands,
40) {
41    for (e, cst_tokens, text) in &mut query {
42        let types: TokenTypesComponent = Wrapped(
43            cst_tokens
44                .0
45                .iter()
46                .flat_map(|Spanned(kind, span)| {
47                    L::semantic_token_spans(*kind, span.clone(), text)
48                        .into_iter()
49                        .map(|(t, s)| spanned(t, s))
50                })
51                .collect(),
52        );
53        commands.entity(e).insert(types);
54    }
55}
56
57pub fn semantic_tokens_system(
58    mut query: Query<(&RopeC, &TokenTypesComponent, &mut HighlightRequest)>,
59    res: Res<SemanticTokensDict>,
60) {
61    tracing::debug!("semantic_tokens_system called");
62    for (rope, types, mut req) in &mut query {
63        let rope = &rope.0;
64        // `ts` is indexed by *byte* offset; spans (`r`) are byte ranges.
65        let mut ts: Vec<Option<SemanticTokenType>> = Vec::with_capacity(rope.len_bytes());
66        ts.resize(rope.len_bytes(), None);
67        types.iter().for_each(|Spanned(ty, r)| {
68            r.clone().for_each(|j| {
69                if j < ts.len() {
70                    ts[j] = Some(ty.clone())
71                } else {
72                    tracing::error!(
73                        "Semantic tokens type {} (index={}) falls outside of document size ({} bytes)",
74                        ty.as_str(),
75                        j,
76                        rope.len_bytes()
77                    );
78                }
79            });
80        });
81
82        let mut last = None;
83        let mut start = 0;
84        let mut out_tokens = Vec::new();
85        for (i, ty) in ts.into_iter().enumerate() {
86            if last != ty {
87                if let Some(t) = last {
88                    out_tokens.push(TokenHelper {
89                        start,
90                        length: i - start,
91                        ty: res.get(&t).cloned().unwrap_or(0),
92                    });
93                }
94
95                last = ty;
96                start = i;
97            }
98        }
99
100        if let Some(t) = last {
101            out_tokens.push(TokenHelper {
102                start,
103                length: rope.len_bytes() - start, // byte length, like the branch above
104                ty: res.get(&t).cloned().unwrap_or(0),
105            });
106        }
107
108        // Emit LSP semantic tokens. `start`/`length` on each `TokenHelper` are byte
109        // offsets/lengths; LSP wants UTF-16 columns and lengths. We walk the source
110        // once with a monotonically advancing cursor (tokens are sorted, non-
111        // overlapping), so this is O(n) overall — crucially avoiding an O(line²)
112        // blow-up on minified single-line documents that a per-token byte→UTF-16
113        // conversion would incur.
114        let text = rope.as_str();
115        let mut cur_byte = 0usize;
116        let mut cur_line = 0u32;
117        let mut cur_col = 0u32; // UTF-16 column within the current line
118        let mut pre_line = 0u32;
119        let mut pre_col = 0u32;
120        req.0 = out_tokens
121            .into_iter()
122            .map(|token| {
123                let advance = |from: usize, to: usize, line: &mut u32, col: &mut u32| -> u32 {
124                    let mut utf16 = 0u32;
125                    for ch in text.get(from..to).unwrap_or("").chars() {
126                        let w = ch.len_utf16() as u32;
127                        utf16 += w;
128                        if ch == '\n' {
129                            *line += 1;
130                            *col = 0;
131                        } else {
132                            *col += w;
133                        }
134                    }
135                    utf16
136                };
137
138                advance(cur_byte, token.start, &mut cur_line, &mut cur_col);
139                let start_line = cur_line;
140                let start_col = cur_col;
141                let length = advance(token.start, token.start + token.length, &mut cur_line, &mut cur_col);
142                cur_byte = token.start + token.length;
143
144                let delta_line = start_line - pre_line;
145                let delta_start = if delta_line == 0 {
146                    start_col - pre_col
147                } else {
148                    start_col
149                };
150                pre_line = start_line;
151                pre_col = start_col;
152
153                SemanticToken {
154                    delta_line,
155                    delta_start,
156                    length,
157                    token_type: token.ty as u32,
158                    token_modifiers_bitset: 0,
159                }
160            })
161            .collect();
162    }
163}
swls_core/feature/semantic.rs

swls_core/feature/
semantic.rs
