Skip to main content

zen_expression/intellisense/
mod.rs

1use crate::compiler::Compiler;
2use crate::intellisense::completion::Completions;
3use crate::intellisense::dependency::DependencyResolutionWalker;
4use crate::intellisense::diagnostic::{
5    collect_parser_diagnostics, collect_type_diagnostics, compiler_error_to_diagnostic,
6    lexer_error_to_diagnostic, Diagnostic, DiagnosticSource, Severity,
7};
8use crate::intellisense::inspection::{inspect_at, InspectionResult};
9use crate::intellisense::scope::IntelliSenseScope;
10use crate::intellisense::type_provider::TypesProvider;
11use crate::lexer::Lexer;
12use crate::nl::project::Projector;
13use crate::nl::{NlRequest, NlResult};
14use crate::parser::{Node, NodeMetadata, Parser};
15use crate::variable::VariableType;
16use bumpalo::Bump;
17use nohash_hasher::BuildNoHashHasher;
18use serde::Serialize;
19use std::cell::RefCell;
20use std::collections::HashMap;
21use std::rc::Rc;
22
23pub mod completion;
24pub mod dependency;
25pub mod diagnostic;
26mod discriminant;
27mod entity_flow;
28mod inspection;
29mod scope;
30pub(crate) mod type_provider;
31
32pub use dependency::{DependencyResult, ReadDependency, Reference};
33pub use discriminant::{ArmTest, NumberCover};
34pub use entity_flow::FlowSource;
35
36pub type AstMetadata = HashMap<usize, NodeMetadata, BuildNoHashHasher<usize>>;
37
38#[derive(Debug, Serialize)]
39#[serde(rename_all = "camelCase")]
40pub struct IntelliSenseToken {
41    pub span: (u32, u32),
42    pub kind: VariableType,
43    pub node_kind: &'static str,
44    pub error: Option<String>,
45}
46
47#[derive(Debug, Clone, Serialize)]
48#[serde(rename_all = "camelCase")]
49pub struct ExpressionAnalysis {
50    pub return_type: VariableType,
51    pub reads: Vec<ReadDependency>,
52    pub references: Vec<Reference>,
53    pub diagnostics: Vec<Diagnostic>,
54}
55
56pub type NlLabelResolver = Rc<dyn Fn(&str, &str) -> Option<String>>;
57
58pub struct IntelliSense {
59    arena: Bump,
60    lexer: Lexer,
61    strict: bool,
62    nl_labels: Option<NlLabelResolver>,
63}
64
65impl IntelliSense {
66    pub fn new() -> Self {
67        Self {
68            arena: Bump::new(),
69            lexer: Lexer::new(),
70            strict: false,
71            nl_labels: None,
72        }
73    }
74
75    pub fn with_strict(mut self, strict: bool) -> Self {
76        self.strict = strict;
77        self
78    }
79
80    pub fn set_nl_labels(&mut self, labels: Option<NlLabelResolver>) {
81        self.nl_labels = labels;
82    }
83
84    pub fn completions(
85        &mut self,
86        source: &str,
87        pos: u32,
88        data: &VariableType,
89    ) -> Vec<completion::Completion> {
90        let tokens = match self.type_check(source, data) {
91            Some(t) => t,
92            None => return Completions::build_scope(data),
93        };
94
95        Completions::build(source, pos, data, &tokens)
96    }
97
98    pub fn inspect(
99        &mut self,
100        source: &str,
101        pos: u32,
102        data: &VariableType,
103    ) -> Option<InspectionResult> {
104        let tokens = self.type_check(source, data)?;
105        inspect_at(source, pos, &tokens)
106    }
107
108    pub fn analyze(&mut self, source: &str, data: &VariableType) -> Rc<ExpressionAnalysis> {
109        Rc::new(self.analyze_standard_inner(source, data))
110    }
111
112    fn analyze_standard_inner(&mut self, source: &str, data: &VariableType) -> ExpressionAnalysis {
113        self.arena.reset();
114        let arena = &self.arena;
115        let mut diagnostics = Vec::new();
116
117        let tokens = match self.lexer.tokenize(arena, source) {
118            Ok(tokens) => tokens,
119            Err(err) => {
120                diagnostics.push(lexer_error_to_diagnostic(&err));
121                return ExpressionAnalysis {
122                    return_type: VariableType::Any,
123                    reads: Vec::new(),
124                    references: Vec::new(),
125                    diagnostics,
126                };
127            }
128        };
129
130        let Ok(parser) = Parser::try_new(&tokens, arena) else {
131            return ExpressionAnalysis {
132                return_type: VariableType::Any,
133                reads: Vec::new(),
134                references: Vec::new(),
135                diagnostics,
136            };
137        };
138
139        let parser = parser.standard().with_metadata();
140        let parser_result = parser.parse();
141        let ast = parser_result.root;
142
143        if !parser_result.is_complete || ast.has_error() {
144            if !parser_result.is_complete {
145                diagnostics.push(Diagnostic {
146                    span: (0, 0),
147                    message: "Incomplete expression".to_string(),
148                    severity: Severity::Error,
149                    source: DiagnosticSource::Parser,
150                });
151            }
152            collect_parser_diagnostics(ast, &mut diagnostics);
153            return ExpressionAnalysis {
154                return_type: VariableType::Any,
155                reads: Vec::new(),
156                references: Vec::new(),
157                diagnostics,
158            };
159        }
160
161        let metadata = parser_result.metadata.unwrap_or_default();
162
163        let scope = IntelliSenseScope {
164            pointer_data: data.shallow_clone(),
165            root_data: data.shallow_clone(),
166            current_data: data.shallow_clone(),
167            ..Default::default()
168        };
169
170        let type_data = TypesProvider::generate(ast, scope, self.strict);
171
172        let return_type = type_data
173            .get_type(ast)
174            .map(|t| t.kind.clone())
175            .unwrap_or(VariableType::Any);
176
177        collect_type_diagnostics(ast, &type_data, &metadata, &mut diagnostics);
178
179        let dep_result = DependencyResolutionWalker::walk(ast, &metadata);
180
181        let mut compiler = Compiler::new();
182        if let Err(err) = compiler.compile(ast) {
183            diagnostics.push(compiler_error_to_diagnostic(&err));
184        }
185
186        ExpressionAnalysis {
187            return_type,
188            reads: dep_result.reads,
189            references: dep_result.references,
190            diagnostics,
191        }
192    }
193
194    pub fn nl_tokenize_batch(
195        &mut self,
196        requests: &[NlRequest],
197        root_type: &VariableType,
198    ) -> Vec<NlResult> {
199        requests
200            .iter()
201            .map(|request| self.nl_tokenize(request, root_type))
202            .collect()
203    }
204
205    pub fn nl_tokenize(&mut self, request: &NlRequest, root_type: &VariableType) -> NlResult {
206        let scope = if request.unary {
207            Self::unary_scope(root_type, request.subject_type.as_ref())
208        } else {
209            root_type.shallow_clone()
210        };
211        let expected = (!request.unary)
212            .then_some(request.subject_type.as_ref())
213            .flatten();
214        let mut result = self.nl_tokenize_scoped(
215            &request.id,
216            &request.expression,
217            request.unary,
218            &scope,
219            expected,
220        );
221        if request.unary {
222            let subject = scope.get("$");
223            result.subject_options = self.nl_subject_options(&subject);
224            result.subject_type = Some(subject);
225        } else if let Some(expected) = expected {
226            result.subject_options = self.nl_subject_options(expected);
227            result.subject_type = Some(expected.shallow_clone());
228        }
229        result
230    }
231
232    pub fn nl_subject_options(&self, subject: &VariableType) -> Option<Vec<crate::nl::EnumOption>> {
233        crate::nl::subject_enum_options(subject, self.nl_labels.as_ref())
234    }
235
236    pub fn nl_tokenize_scoped(
237        &mut self,
238        id: &str,
239        source: &str,
240        unary: bool,
241        scope_type: &VariableType,
242        expected: Option<&VariableType>,
243    ) -> NlResult {
244        let mut result = NlResult {
245            id: id.to_string(),
246            tokens: Vec::new(),
247            enums: Vec::new(),
248            diagnostics: Vec::new(),
249            subject_type: None,
250            subject_options: None,
251        };
252
253        self.arena.reset();
254        let arena = &self.arena;
255
256        let tokens = match self.lexer.tokenize(arena, source) {
257            Ok(tokens) => tokens,
258            Err(err) => {
259                result.diagnostics.push(lexer_error_to_diagnostic(&err));
260                return result;
261            }
262        };
263
264        let Ok(parser) = Parser::try_new(&tokens, arena) else {
265            return result;
266        };
267
268        let parser_result = if unary {
269            parser.unary().with_metadata().parse()
270        } else {
271            parser.standard().with_metadata().parse()
272        };
273        let ast = parser_result.root;
274
275        if !parser_result.is_complete || ast.has_error() {
276            if !parser_result.is_complete {
277                result.diagnostics.push(Diagnostic {
278                    span: (0, 0),
279                    message: "Incomplete expression".to_string(),
280                    severity: Severity::Error,
281                    source: DiagnosticSource::Parser,
282                });
283            }
284            collect_parser_diagnostics(ast, &mut result.diagnostics);
285            return result;
286        }
287
288        let metadata = parser_result.metadata.unwrap_or_default();
289
290        let scope = IntelliSenseScope {
291            pointer_data: scope_type.shallow_clone(),
292            root_data: scope_type.shallow_clone(),
293            current_data: scope_type.shallow_clone(),
294            ..Default::default()
295        };
296
297        let type_data = TypesProvider::generate(ast, scope, self.strict);
298        collect_type_diagnostics(ast, &type_data, &metadata, &mut result.diagnostics);
299
300        let (tokens, enums) =
301            Projector::new(source, &type_data, &metadata, unary, self.nl_labels.clone())
302                .run(ast, expected.map(|e| e.shallow_clone()));
303        result.tokens = tokens;
304        result.enums = enums;
305        result
306    }
307
308    fn unary_scope(root_type: &VariableType, subject_type: Option<&VariableType>) -> VariableType {
309        let subject = subject_type
310            .map(|s| s.shallow_clone())
311            .unwrap_or(VariableType::Any);
312
313        let object = VariableType::empty_object();
314        if let VariableType::Object(target) = &object {
315            if let VariableType::Object(source) = root_type {
316                for (key, value) in source.borrow().iter() {
317                    target
318                        .borrow_mut()
319                        .insert(key.clone(), value.shallow_clone());
320                }
321            }
322            target.borrow_mut().insert(Rc::from("$"), subject);
323        }
324        object
325    }
326
327    pub fn with_ast<T>(
328        &mut self,
329        source: &str,
330        unary: bool,
331        f: impl for<'arena> FnOnce(&'arena Node<'arena>, &AstMetadata) -> T,
332    ) -> Option<T> {
333        self.arena.reset();
334        let arena = &self.arena;
335        let tokens = self.lexer.tokenize(arena, source).ok()?;
336        let parser = Parser::try_new(&tokens, arena).ok()?;
337        let parser_result = if unary {
338            parser.unary().with_metadata().parse()
339        } else {
340            parser.standard().with_metadata().parse()
341        };
342        let ast = parser_result.root;
343        if !parser_result.is_complete || ast.has_error() {
344            return None;
345        }
346        let metadata = parser_result.metadata.unwrap_or_default();
347        Some(f(ast, &metadata))
348    }
349
350    pub fn field_reads(
351        &mut self,
352        source: &str,
353        field_path: &[&str],
354    ) -> Option<Vec<ReadDependency>> {
355        self.arena.reset();
356        let arena = &self.arena;
357        let tokens = self.lexer.tokenize(arena, source).ok()?;
358        let parser = Parser::try_new(&tokens, arena)
359            .ok()?
360            .standard()
361            .with_metadata();
362        let parser_result = parser.parse();
363        let ast = parser_result.root;
364        if !parser_result.is_complete || ast.has_error() {
365            return None;
366        }
367        let metadata = parser_result.metadata.unwrap_or_default();
368        DependencyResolutionWalker::field_dependencies(ast, &metadata, field_path)
369    }
370
371    pub fn arm_test(&mut self, source: &str) -> ArmTest {
372        if source.trim().is_empty() {
373            return ArmTest::Default;
374        }
375        self.arena.reset();
376        let arena = &self.arena;
377        let result = (|| {
378            let tokens = self.lexer.tokenize(arena, source).ok()?;
379            let parser = Parser::try_new(&tokens, arena).ok()?;
380            let parser_result = parser.standard().with_metadata().parse();
381            let ast = parser_result.root;
382            if !parser_result.is_complete || ast.has_error() {
383                return None;
384            }
385            Some(ArmTest::from_node(ast))
386        })();
387        result.unwrap_or(ArmTest::Unrecognized)
388    }
389
390    pub fn cell_test(&mut self, source: &str) -> ArmTest {
391        if source.trim().is_empty() {
392            return ArmTest::Default;
393        }
394        self.arena.reset();
395        let arena = &self.arena;
396        let result = (|| {
397            let tokens = self.lexer.tokenize(arena, source).ok()?;
398            let parser = Parser::try_new(&tokens, arena).ok()?;
399            let parser_result = parser.unary().with_metadata().parse();
400            let ast = parser_result.root;
401            if !parser_result.is_complete || ast.has_error() {
402                return None;
403            }
404            let test = ArmTest::from_node(ast);
405            let on_reference = match &test {
406                ArmTest::Enum { path, .. }
407                | ArmTest::Bool { path, .. }
408                | ArmTest::Number { path, .. } => {
409                    matches!(path.as_slice(), [p] if p.as_ref() == "$")
410                }
411                _ => true,
412            };
413            on_reference.then_some(test)
414        })();
415        result.unwrap_or(ArmTest::Unrecognized)
416    }
417
418    pub fn flow_source(&mut self, source: &str) -> Option<FlowSource> {
419        if source.trim().is_empty() {
420            return None;
421        }
422        self.arena.reset();
423        let arena = &self.arena;
424        let tokens = self.lexer.tokenize(arena, source).ok()?;
425        let parser = Parser::try_new(&tokens, arena).ok()?;
426        let parser_result = parser.standard().parse();
427        let ast = parser_result.root;
428        if !parser_result.is_complete || ast.has_error() {
429            return None;
430        }
431        FlowSource::from_node(ast)
432    }
433
434    pub fn reads(&mut self, source: &str) -> Vec<ReadDependency> {
435        self.reads_inner(source, false).reads
436    }
437
438    pub fn reads_unary(&mut self, source: &str) -> Vec<ReadDependency> {
439        self.reads_inner(source, true).reads
440    }
441
442    pub fn dependencies(&mut self, source: &str) -> DependencyResult {
443        self.reads_inner(source, false)
444    }
445
446    fn reads_inner(&mut self, source: &str, unary: bool) -> DependencyResult {
447        self.arena.reset();
448        let arena = &self.arena;
449        let result = (|| {
450            let tokens = self.lexer.tokenize(arena, source).ok()?;
451            let parser = Parser::try_new(&tokens, arena).ok()?;
452            let parser_result = if unary {
453                parser.unary().with_metadata().parse()
454            } else {
455                parser.standard().with_metadata().parse()
456            };
457            let ast = parser_result.root;
458            if !parser_result.is_complete || ast.has_error() {
459                return None;
460            }
461            let metadata = parser_result.metadata.unwrap_or_default();
462            let dep = if unary {
463                DependencyResolutionWalker::walk_with_locals(ast, &metadata, &["$"])
464            } else {
465                DependencyResolutionWalker::walk(ast, &metadata)
466            };
467            Some(dep)
468        })();
469        result.unwrap_or_default()
470    }
471
472    pub fn analyze_unary(&mut self, source: &str, data: &VariableType) -> Rc<ExpressionAnalysis> {
473        Rc::new(self.analyze_unary_inner(source, data))
474    }
475
476    fn analyze_unary_inner(&mut self, source: &str, data: &VariableType) -> ExpressionAnalysis {
477        self.arena.reset();
478        let arena = &self.arena;
479        let mut diagnostics = Vec::new();
480
481        let tokens = match self.lexer.tokenize(arena, source) {
482            Ok(tokens) => tokens,
483            Err(err) => {
484                diagnostics.push(lexer_error_to_diagnostic(&err));
485                return ExpressionAnalysis {
486                    return_type: VariableType::Bool,
487                    reads: Vec::new(),
488                    references: Vec::new(),
489                    diagnostics,
490                };
491            }
492        };
493
494        let Ok(parser) = Parser::try_new(&tokens, arena) else {
495            return ExpressionAnalysis {
496                return_type: VariableType::Bool,
497                reads: Vec::new(),
498                references: Vec::new(),
499                diagnostics,
500            };
501        };
502
503        let parser = parser.unary().with_metadata();
504        let parser_result = parser.parse();
505        let ast = parser_result.root;
506
507        if !parser_result.is_complete || ast.has_error() {
508            if !parser_result.is_complete {
509                diagnostics.push(Diagnostic {
510                    span: (0, 0),
511                    message: "Incomplete expression".to_string(),
512                    severity: Severity::Error,
513                    source: DiagnosticSource::Parser,
514                });
515            }
516            collect_parser_diagnostics(ast, &mut diagnostics);
517            return ExpressionAnalysis {
518                return_type: VariableType::Bool,
519                reads: Vec::new(),
520                references: Vec::new(),
521                diagnostics,
522            };
523        }
524
525        let metadata = parser_result.metadata.unwrap_or_default();
526
527        let scope = IntelliSenseScope {
528            pointer_data: data.shallow_clone(),
529            root_data: data.shallow_clone(),
530            current_data: data.shallow_clone(),
531            ..Default::default()
532        };
533
534        let type_data = TypesProvider::generate(ast, scope, self.strict);
535        collect_type_diagnostics(ast, &type_data, &metadata, &mut diagnostics);
536
537        let dep_result = DependencyResolutionWalker::walk_with_locals(ast, &metadata, &["$"]);
538
539        let mut compiler = Compiler::new();
540        if let Err(err) = compiler.compile(ast) {
541            diagnostics.push(compiler_error_to_diagnostic(&err));
542        }
543
544        ExpressionAnalysis {
545            return_type: VariableType::Bool,
546            reads: dep_result.reads,
547            references: dep_result.references,
548            diagnostics,
549        }
550    }
551
552    pub fn type_check(
553        &mut self,
554        source: &str,
555        data: &VariableType,
556    ) -> Option<Vec<IntelliSenseToken>> {
557        self.arena.reset();
558        let arena = &self.arena;
559
560        let tokens = self.lexer.tokenize(arena, source).ok()?;
561        let parser = Parser::try_new(&tokens, arena).map(|p| p.standard()).ok()?;
562
563        let parser_result = parser.with_metadata().parse();
564        let ast = parser_result.root;
565        let metadata = parser_result.metadata?;
566
567        let type_data = TypesProvider::generate(
568            ast,
569            IntelliSenseScope {
570                pointer_data: data.shallow_clone(),
571                root_data: data.shallow_clone(),
572                current_data: data.shallow_clone(),
573                ..Default::default()
574            },
575            self.strict,
576        );
577
578        let results = RefCell::new(Vec::new());
579        ast.walk(|node| {
580            let addr = node as *const Node as usize;
581            let mut r = results.borrow_mut();
582            let typ = type_data.get_type(node);
583
584            r.push(IntelliSenseToken {
585                span: node
586                    .span()
587                    .or_else(|| metadata.get(&addr).map(|s| s.span))
588                    .unwrap_or_default(),
589                node_kind: node.into(),
590                error: typ.map(|t| t.error.clone()).flatten(),
591                kind: typ
592                    .map(|t| t.kind.clone())
593                    .unwrap_or_else(|| VariableType::Any),
594            });
595        });
596        Some(results.into_inner())
597    }
598
599    pub fn type_check_unary(
600        &mut self,
601        source: &str,
602        data: &VariableType,
603    ) -> Option<Vec<IntelliSenseToken>> {
604        self.arena.reset();
605        let arena = &self.arena;
606
607        let tokens = self.lexer.tokenize(arena, source).ok()?;
608        let parser = Parser::try_new(&tokens, arena).map(|p| p.unary()).ok()?;
609
610        let parser_result = parser.with_metadata().parse();
611        let ast = parser_result.root;
612        let metadata = parser_result.metadata?;
613
614        let type_data = TypesProvider::generate(
615            ast,
616            IntelliSenseScope {
617                pointer_data: data.shallow_clone(),
618                root_data: data.shallow_clone(),
619                current_data: data.shallow_clone(),
620                ..Default::default()
621            },
622            self.strict,
623        );
624
625        let results = RefCell::new(Vec::new());
626        ast.walk(|node| {
627            let addr = node as *const Node as usize;
628            let mut r = results.borrow_mut();
629            let typ = type_data.get_type(node);
630
631            r.push(IntelliSenseToken {
632                span: metadata.get(&addr).map(|s| s.span).unwrap_or_default(),
633                node_kind: node.into(),
634                error: typ.map(|t| t.error.clone()).flatten(),
635                kind: typ
636                    .map(|t| t.kind.clone())
637                    .unwrap_or_else(|| VariableType::Any),
638            });
639        });
640        Some(results.into_inner())
641    }
642}