pkl_lsp_core/
index.rs

1use pklr::error::Error as PklError;
2use pklr::lexer::{self, Token, TokenKind};
3use pklr::parser::{self, Annotation, Entry, Expr, Import, Module, Property, TypeExpr};
4use rkyv::{Archive, Deserialize, Serialize};
5use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
6
7use crate::Document;
8use crate::protocol;
9
10#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
11#[rkyv(compare(PartialEq), derive(Debug))]
12pub enum DiagnosticSource {
13    Lex,
14    Parse,
15    Eval,
16    Io,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
20#[rkyv(compare(PartialEq), derive(Debug))]
21pub struct PklDiagnostic {
22    pub uri: String,
23    pub range: TextRange,
24    pub message: String,
25    pub source: DiagnosticSource,
26}
27
28#[derive(
29    Debug,
30    Clone,
31    Copy,
32    PartialEq,
33    Eq,
34    Archive,
35    Serialize,
36    Deserialize,
37    SerdeSerialize,
38    SerdeDeserialize,
39)]
40#[rkyv(compare(PartialEq), derive(Debug))]
41pub struct TextPosition {
42    pub line: u32,
43    pub character: u32,
44}
45
46#[derive(
47    Debug,
48    Clone,
49    Copy,
50    PartialEq,
51    Eq,
52    Archive,
53    Serialize,
54    Deserialize,
55    SerdeSerialize,
56    SerdeDeserialize,
57)]
58#[rkyv(compare(PartialEq), derive(Debug))]
59pub struct TextRange {
60    pub start: TextPosition,
61    pub end: TextPosition,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
65#[rkyv(compare(PartialEq), derive(Debug))]
66pub enum SymbolKind {
67    Module,
68    Import,
69    Property,
70    Class,
71    TypeAlias,
72    Annotation,
73}
74
75impl SymbolKind {
76    pub fn as_protocol(&self) -> u32 {
77        match self {
78            Self::Module => protocol::symbol_kind::MODULE,
79            Self::Import => protocol::symbol_kind::NAMESPACE,
80            Self::Property => protocol::symbol_kind::PROPERTY,
81            Self::Class => protocol::symbol_kind::CLASS,
82            Self::TypeAlias => protocol::symbol_kind::TYPE_PARAMETER,
83            Self::Annotation => protocol::symbol_kind::EVENT,
84        }
85    }
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
89#[rkyv(compare(PartialEq), derive(Debug))]
90pub struct PklSymbol {
91    pub uri: String,
92    pub name: String,
93    pub kind: SymbolKind,
94    pub range: TextRange,
95    pub selection_range: TextRange,
96    pub detail: Option<String>,
97    pub container_name: Option<String>,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
101#[rkyv(compare(PartialEq), derive(Debug))]
102pub enum ImportKind {
103    Import,
104    GlobImport,
105    Amends,
106    Extends,
107}
108
109#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
110#[rkyv(compare(PartialEq), derive(Debug))]
111pub struct ImportEdge {
112    pub uri: String,
113    pub target: String,
114    pub alias: Option<String>,
115    pub kind: ImportKind,
116    pub range: TextRange,
117}
118
119#[derive(Debug, Clone, PartialEq, Eq, Archive, Serialize, Deserialize)]
120#[rkyv(compare(PartialEq), derive(Debug))]
121pub struct WorkspaceIndex {
122    pub diagnostics: Vec<PklDiagnostic>,
123    pub imports: Vec<ImportEdge>,
124    pub symbols: Vec<PklSymbol>,
125}
126
127impl WorkspaceIndex {
128    pub fn build<'a>(documents: impl IntoIterator<Item = &'a Document>) -> Self {
129        let mut index = Self {
130            diagnostics: Vec::new(),
131            imports: Vec::new(),
132            symbols: Vec::new(),
133        };
134
135        for document in documents {
136            index.index_document(document);
137        }
138
139        index
140    }
141
142    pub fn diagnostics_for_protocol(&self, uri: &str) -> Vec<protocol::Diagnostic> {
143        self.diagnostics
144            .iter()
145            .filter(|diagnostic| diagnostic.uri == uri)
146            .map(|diagnostic| protocol::Diagnostic {
147                range: diagnostic.range,
148                severity: Some(match diagnostic.source {
149                    DiagnosticSource::Lex
150                    | DiagnosticSource::Parse
151                    | DiagnosticSource::Eval
152                    | DiagnosticSource::Io => protocol::diagnostic_severity::ERROR,
153                }),
154                source: Some("pkl-lsp".to_string()),
155                message: diagnostic.message.clone(),
156            })
157            .collect()
158    }
159
160    pub fn document_symbols(&self, uri: &str) -> Vec<protocol::DocumentSymbol> {
161        self.symbols
162            .iter()
163            .filter(|symbol| symbol.uri == uri)
164            .map(|symbol| protocol::DocumentSymbol {
165                name: symbol.name.clone(),
166                detail: symbol.detail.clone(),
167                kind: symbol.kind.as_protocol(),
168                range: symbol.range,
169                selection_range: symbol.selection_range,
170                children: None,
171            })
172            .collect()
173    }
174
175    pub fn symbols_in_uri(&self, uri: &str) -> impl Iterator<Item = &PklSymbol> {
176        self.symbols.iter().filter(move |symbol| symbol.uri == uri)
177    }
178
179    fn index_document(&mut self, document: &Document) {
180        let module_symbol = PklSymbol {
181            uri: document.uri.clone(),
182            name: module_name(&document.uri),
183            kind: SymbolKind::Module,
184            range: entire_document_range(&document.text),
185            selection_range: position_range(0, 0, 0, 0),
186            detail: Some("Pkl module".to_string()),
187            container_name: None,
188        };
189        self.symbols.push(module_symbol);
190
191        let tokens = match lexer::lex_named(&document.text, &document.uri) {
192            Ok(tokens) => tokens,
193            Err(error) => {
194                self.diagnostics
195                    .push(error_to_diagnostic(&document.uri, &document.text, error));
196                return;
197            }
198        };
199
200        let module = match parser::parse_named(&tokens, &document.text, &document.uri) {
201            Ok(module) => module,
202            Err(error) => {
203                self.diagnostics
204                    .push(error_to_diagnostic(&document.uri, &document.text, error));
205                self.collect_fast_imports(document, &tokens);
206                return;
207            }
208        };
209
210        self.collect_module(document, &tokens, &module);
211    }
212
213    fn collect_fast_imports(&mut self, document: &Document, tokens: &[Token]) {
214        for uri in parser::collect_imports(tokens) {
215            if let Some(token) = find_string_token(tokens, &uri) {
216                self.imports.push(ImportEdge {
217                    uri: document.uri.clone(),
218                    target: uri,
219                    alias: None,
220                    kind: ImportKind::Import,
221                    range: token_range(token),
222                });
223            }
224        }
225    }
226
227    fn collect_module(&mut self, document: &Document, tokens: &[Token], module: &Module) {
228        if let Some(amends) = &module.amends {
229            self.push_module_edge(document, tokens, amends, ImportKind::Amends, None);
230        }
231        if let Some(extends) = &module.extends {
232            self.push_module_edge(document, tokens, extends, ImportKind::Extends, None);
233        }
234        for import in &module.imports {
235            self.push_import(document, tokens, import);
236        }
237        let mut cursor = TokenCursor::new(tokens);
238        self.collect_annotations(document, &module.annotations, &mut cursor, None);
239        self.collect_entries(document, &module.body, &mut cursor, None);
240    }
241
242    fn push_import(&mut self, document: &Document, tokens: &[Token], import: &Import) {
243        let kind = if import.is_glob {
244            ImportKind::GlobImport
245        } else {
246            ImportKind::Import
247        };
248        self.push_module_edge(document, tokens, &import.uri, kind, import.alias.clone());
249        let name = import
250            .alias
251            .clone()
252            .unwrap_or_else(|| import_name_from_uri(&import.uri));
253        if let Some(token) = find_ident_or_string_token(tokens, &name, &import.uri) {
254            self.symbols.push(PklSymbol {
255                uri: document.uri.clone(),
256                name,
257                kind: SymbolKind::Import,
258                range: token_range(token),
259                selection_range: token_range(token),
260                detail: Some(format!("import {}", import.uri)),
261                container_name: None,
262            });
263        }
264    }
265
266    fn push_module_edge(
267        &mut self,
268        document: &Document,
269        tokens: &[Token],
270        target: &str,
271        kind: ImportKind,
272        alias: Option<String>,
273    ) {
274        let range = find_string_token(tokens, target)
275            .map_or_else(|| entire_document_range(&document.text), token_range);
276        self.imports.push(ImportEdge {
277            uri: document.uri.clone(),
278            target: target.to_string(),
279            alias,
280            kind,
281            range,
282        });
283    }
284
285    fn collect_annotations(
286        &mut self,
287        document: &Document,
288        annotations: &[Annotation],
289        cursor: &mut TokenCursor<'_>,
290        container_name: Option<&str>,
291    ) {
292        for annotation in annotations {
293            if let Some(token) = cursor.next_ident(&annotation.name) {
294                self.symbols.push(PklSymbol {
295                    uri: document.uri.clone(),
296                    name: format!("@{}", annotation.name),
297                    kind: SymbolKind::Annotation,
298                    range: token_range(token),
299                    selection_range: token_range(token),
300                    detail: Some("annotation".to_string()),
301                    container_name: container_name.map(ToOwned::to_owned),
302                });
303            }
304            self.collect_entries(document, &annotation.body, cursor, Some(&annotation.name));
305        }
306    }
307
308    fn collect_entries(
309        &mut self,
310        document: &Document,
311        entries: &[Entry],
312        cursor: &mut TokenCursor<'_>,
313        container_name: Option<&str>,
314    ) {
315        for entry in entries {
316            match entry {
317                Entry::Property(property) => {
318                    self.collect_property(document, property, cursor, container_name)
319                }
320                Entry::ClassDef(name, _, parent, body) => {
321                    if let Some(token) = cursor.next_ident(name) {
322                        self.symbols.push(PklSymbol {
323                            uri: document.uri.clone(),
324                            name: name.clone(),
325                            kind: SymbolKind::Class,
326                            range: token_range(token),
327                            selection_range: token_range(token),
328                            detail: parent.as_ref().map(|parent| format!("extends {parent}")),
329                            container_name: container_name.map(ToOwned::to_owned),
330                        });
331                    }
332                    self.collect_entries(document, body, cursor, Some(name));
333                }
334                Entry::TypeAlias(name, ty) => {
335                    if let Some(token) = cursor.next_ident(name) {
336                        self.symbols.push(PklSymbol {
337                            uri: document.uri.clone(),
338                            name: name.clone(),
339                            kind: SymbolKind::TypeAlias,
340                            range: token_range(token),
341                            selection_range: token_range(token),
342                            detail: Some(type_expr_label(ty)),
343                            container_name: container_name.map(ToOwned::to_owned),
344                        });
345                    }
346                }
347                Entry::DynProperty(key, value) => {
348                    collect_expr(cursor, key);
349                    collect_expr(cursor, value);
350                }
351                Entry::ForGenerator(generator) => {
352                    self.collect_entries(document, &generator.body, cursor, container_name)
353                }
354                Entry::WhenGenerator(generator) => {
355                    self.collect_entries(document, &generator.body, cursor, container_name);
356                    if let Some(else_body) = &generator.else_body {
357                        self.collect_entries(document, else_body, cursor, container_name);
358                    }
359                }
360                Entry::Spread(expr) | Entry::Elem(expr) => collect_expr(cursor, expr),
361            }
362        }
363    }
364
365    fn collect_property(
366        &mut self,
367        document: &Document,
368        property: &Property,
369        cursor: &mut TokenCursor<'_>,
370        container_name: Option<&str>,
371    ) {
372        self.collect_annotations(
373            document,
374            &property.annotations,
375            cursor,
376            Some(&property.name),
377        );
378        if let Some(token) = cursor.next_ident(&property.name) {
379            let detail = property.type_ann.as_ref().map(type_expr_label);
380            self.symbols.push(PklSymbol {
381                uri: document.uri.clone(),
382                name: property.name.clone(),
383                kind: SymbolKind::Property,
384                range: token_range(token),
385                selection_range: token_range(token),
386                detail,
387                container_name: container_name.map(ToOwned::to_owned),
388            });
389        }
390        if let Some(value) = &property.value {
391            collect_expr(cursor, value);
392        }
393        if let Some(body) = &property.body {
394            self.collect_entries(document, body, cursor, Some(&property.name));
395        }
396    }
397}
398
399fn collect_expr(_cursor: &mut TokenCursor<'_>, _expr: &Expr) {}
400
401fn error_to_diagnostic(uri: &str, text: &str, error: PklError) -> PklDiagnostic {
402    match &error {
403        PklError::Lex { message, .. } => PklDiagnostic {
404            uri: uri.to_string(),
405            range: offset_range(text, error.source_offset().unwrap_or_default(), 1),
406            message: message.clone(),
407            source: DiagnosticSource::Lex,
408        },
409        PklError::Parse { message, .. } => PklDiagnostic {
410            uri: uri.to_string(),
411            range: offset_range(text, error.source_offset().unwrap_or_default(), 1),
412            message: message.clone(),
413            source: DiagnosticSource::Parse,
414        },
415        PklError::Io(path, error) => PklDiagnostic {
416            uri: uri.to_string(),
417            range: position_range(0, 0, 0, 1),
418            message: format!("IO error reading {}: {error}", path.display()),
419            source: DiagnosticSource::Io,
420        },
421        PklError::Eval(message)
422        | PklError::ImportNotFound(message)
423        | PklError::Unsupported(message) => PklDiagnostic {
424            uri: uri.to_string(),
425            range: position_range(0, 0, 0, 1),
426            message: message.clone(),
427            source: DiagnosticSource::Eval,
428        },
429    }
430}
431
432pub fn position_at_offset(text: &str, offset: usize) -> TextPosition {
433    let mut line = 0u32;
434    let mut line_start = 0usize;
435    for (idx, ch) in text.char_indices() {
436        if idx >= offset {
437            break;
438        }
439        if ch == '\n' {
440            line += 1;
441            line_start = idx + ch.len_utf8();
442        }
443    }
444    TextPosition {
445        line,
446        character: text[line_start..offset.min(text.len())].chars().count() as u32,
447    }
448}
449
450pub fn offset_at_position(text: &str, position: TextPosition) -> usize {
451    let mut line = 0u32;
452    let mut character = 0u32;
453    for (idx, ch) in text.char_indices() {
454        if line == position.line && character == position.character {
455            return idx;
456        }
457        if ch == '\n' {
458            line += 1;
459            character = 0;
460            if line > position.line {
461                return idx;
462            }
463        } else {
464            character += 1;
465        }
466    }
467    text.len()
468}
469
470pub fn word_at_position(text: &str, position: TextPosition) -> Option<String> {
471    let offset = offset_at_position(text, position);
472    let bytes = text.as_bytes();
473    let mut start = offset.min(bytes.len());
474    while start > 0 && is_ident_byte(bytes[start - 1]) {
475        start -= 1;
476    }
477    let mut end = offset.min(bytes.len());
478    while end < bytes.len() && is_ident_byte(bytes[end]) {
479        end += 1;
480    }
481    (start < end).then(|| text[start..end].to_string())
482}
483
484fn is_ident_byte(byte: u8) -> bool {
485    byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'-'
486}
487
488fn offset_range(text: &str, offset: usize, len: usize) -> TextRange {
489    TextRange {
490        start: position_at_offset(text, offset),
491        end: position_at_offset(text, offset.saturating_add(len).min(text.len())),
492    }
493}
494
495fn token_range(token: &Token) -> TextRange {
496    let start_line = token.line.saturating_sub(1) as u32;
497    let start_col = token.col.saturating_sub(1) as u32;
498    let len = token_len(token) as u32;
499    position_range(
500        start_line,
501        start_col,
502        start_line,
503        start_col.saturating_add(len.max(1)),
504    )
505}
506
507fn token_len(token: &Token) -> usize {
508    match &token.kind {
509        TokenKind::Ident(value) | TokenKind::StringLit(value) => value.chars().count(),
510        TokenKind::IntLit(value) => value.to_string().len(),
511        TokenKind::FloatLit(value) => value.to_string().len(),
512        TokenKind::BoolLit(value) => value.to_string().len(),
513        TokenKind::Null => 4,
514        _ => 1,
515    }
516}
517
518pub fn position_range(start_line: u32, start_col: u32, end_line: u32, end_col: u32) -> TextRange {
519    TextRange {
520        start: TextPosition {
521            line: start_line,
522            character: start_col,
523        },
524        end: TextPosition {
525            line: end_line,
526            character: end_col,
527        },
528    }
529}
530
531fn entire_document_range(text: &str) -> TextRange {
532    let end = position_at_offset(text, text.len());
533    TextRange {
534        start: TextPosition {
535            line: 0,
536            character: 0,
537        },
538        end,
539    }
540}
541
542fn find_string_token<'a>(tokens: &'a [Token], value: &str) -> Option<&'a Token> {
543    tokens
544        .iter()
545        .find(|token| matches!(&token.kind, TokenKind::StringLit(text) if text == value))
546}
547
548fn find_ident_or_string_token<'a>(
549    tokens: &'a [Token],
550    ident: &str,
551    string: &str,
552) -> Option<&'a Token> {
553    tokens.iter().find(|token| match &token.kind {
554        TokenKind::Ident(text) => text == ident,
555        TokenKind::StringLit(text) => text == string,
556        _ => false,
557    })
558}
559
560fn module_name(uri: &str) -> String {
561    uri.rsplit('/')
562        .next()
563        .filter(|name| !name.is_empty())
564        .unwrap_or(uri)
565        .trim_end_matches(".pkl")
566        .to_string()
567}
568
569fn import_name_from_uri(uri: &str) -> String {
570    uri.rsplit('/')
571        .next()
572        .unwrap_or(uri)
573        .trim_end_matches(".pkl")
574        .replace('-', "_")
575}
576
577fn type_expr_label(ty: &TypeExpr) -> String {
578    match ty {
579        TypeExpr::Named(name) => name.clone(),
580        TypeExpr::Nullable(inner) => format!("{}?", type_expr_label(inner)),
581        TypeExpr::Union(items) => items
582            .iter()
583            .map(type_expr_label)
584            .collect::<Vec<_>>()
585            .join(" | "),
586        TypeExpr::Generic(name, args) => {
587            format!(
588                "{name}<{}>",
589                args.iter()
590                    .map(type_expr_label)
591                    .collect::<Vec<_>>()
592                    .join(", ")
593            )
594        }
595        TypeExpr::Constrained(name, _) => format!("{name}(...)"),
596    }
597}
598
599struct TokenCursor<'a> {
600    tokens: &'a [Token],
601    pos: usize,
602}
603
604impl<'a> TokenCursor<'a> {
605    fn new(tokens: &'a [Token]) -> Self {
606        Self { tokens, pos: 0 }
607    }
608
609    fn next_ident(&mut self, name: &str) -> Option<&'a Token> {
610        let found = self.tokens[self.pos..]
611            .iter()
612            .position(|token| matches!(&token.kind, TokenKind::Ident(text) if text == name))?;
613        self.pos += found + 1;
614        self.tokens.get(self.pos - 1)
615    }
616}
617
618#[cfg(test)]
619mod tests {
620    use super::{ImportKind, SymbolKind, WorkspaceIndex};
621    use crate::Document;
622
623    #[test]
624    fn extracts_symbols_and_imports() {
625        let document = Document {
626            uri: "file:///demo.pkl".to_string(),
627            version: 1,
628            text: r#"
629import "base.pkl" as base
630amends "parent.pkl"
631
632@Deprecated { message = "old" }
633class AppConfig {
634  name: String = "demo"
635}
636
637typealias Name = String
638"#
639            .to_string(),
640        };
641
642        let index = WorkspaceIndex::build([&document]);
643
644        assert!(index.diagnostics.is_empty(), "{:?}", index.diagnostics);
645        assert!(
646            index
647                .symbols
648                .iter()
649                .any(|symbol| symbol.name == "AppConfig" && symbol.kind == SymbolKind::Class)
650        );
651        assert!(
652            index
653                .symbols
654                .iter()
655                .any(|symbol| symbol.name == "Name" && symbol.kind == SymbolKind::TypeAlias)
656        );
657        assert!(
658            index
659                .imports
660                .iter()
661                .any(|edge| edge.target == "parent.pkl" && edge.kind == ImportKind::Amends)
662        );
663        assert!(
664            index
665                .imports
666                .iter()
667                .any(|edge| edge.alias.as_deref() == Some("base"))
668        );
669    }
670
671    #[test]
672    fn reports_parse_diagnostics_with_ranges() {
673        let document = Document {
674            uri: "file:///bad.pkl".to_string(),
675            version: 1,
676            text: "name = ".to_string(),
677        };
678
679        let index = WorkspaceIndex::build([&document]);
680        assert_eq!(index.diagnostics.len(), 1);
681        assert_eq!(index.diagnostics[0].range.start.line, 0);
682    }
683}
pkl_lsp_core/index.rs

pkl_lsp_core/
index.rs
