Skip to main content

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        let mut ts: Vec<Option<SemanticTokenType>> = Vec::with_capacity(rope.len_chars());
65        ts.resize(rope.len_bytes(), None);
66        types.iter().for_each(|Spanned(ty, r)| {
67            r.clone().for_each(|j| {
68                if j < ts.len() {
69                    ts[j] = Some(ty.clone())
70                } else {
71                    tracing::error!(
72                        "Semantic tokens type {} (index={}) falls outside of rope size (chars: {} bytes: {})",
73                        ty.as_str(),
74                        j,
75                        rope.len_chars(),
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_chars() - start,
104                ty: res.get(&t).cloned().unwrap_or(0),
105            });
106        }
107
108        let mut pre_line = 0;
109        let mut pre_start = 0;
110        req.0 = out_tokens
111            .into_iter()
112            .flat_map(|token| {
113                let line = rope.try_byte_to_line(token.start as usize).ok()? as u32;
114                let first = rope.try_line_to_char(line as usize).ok()? as u32;
115                let start = rope.try_byte_to_char(token.start as usize).ok()? as u32 - first;
116                let delta_line = line - pre_line;
117                let delta_start = if delta_line == 0 {
118                    start - pre_start
119                } else {
120                    start
121                };
122                let ret = Some(SemanticToken {
123                    delta_line,
124                    delta_start,
125                    length: token.length as u32,
126                    token_type: token.ty as u32,
127                    token_modifiers_bitset: 0,
128                });
129                pre_line = line;
130                pre_start = start;
131                ret
132            })
133            .collect();
134    }
135}