sqruff_lib_core/parser/
segments.rs

1pub mod bracketed;
2pub mod file;
3pub mod fix;
4pub mod from;
5pub mod generator;
6pub mod join;
7pub mod meta;
8pub mod object_reference;
9pub mod select;
10pub mod test_functions;
11
12use std::cell::{Cell, OnceCell};
13use std::fmt::Debug;
14use std::hash::{Hash, Hasher};
15use std::rc::Rc;
16
17use itertools::enumerate;
18use rustc_hash::FxHashMap;
19use smol_str::SmolStr;
20
21use crate::dialects::init::DialectKind;
22use crate::dialects::syntax::{SyntaxKind, SyntaxSet};
23use crate::lint_fix::LintFix;
24use crate::parser::markers::PositionMarker;
25use crate::parser::segments::fix::{FixPatch, SourceFix};
26use crate::parser::segments::object_reference::{ObjectReferenceKind, ObjectReferenceSegment};
27use crate::segments::AnchorEditInfo;
28use crate::templaters::TemplatedFile;
29
30pub struct SegmentBuilder {
31    node_or_token: NodeOrToken,
32}
33
34impl SegmentBuilder {
35    pub fn whitespace(id: u32, raw: &str) -> ErasedSegment {
36        SegmentBuilder::token(id, raw, SyntaxKind::Whitespace).finish()
37    }
38
39    pub fn newline(id: u32, raw: &str) -> ErasedSegment {
40        SegmentBuilder::token(id, raw, SyntaxKind::Newline).finish()
41    }
42
43    pub fn keyword(id: u32, raw: &str) -> ErasedSegment {
44        SegmentBuilder::token(id, raw, SyntaxKind::Keyword).finish()
45    }
46
47    pub fn comma(id: u32) -> ErasedSegment {
48        SegmentBuilder::token(id, ",", SyntaxKind::Comma).finish()
49    }
50
51    pub fn symbol(id: u32, raw: &str) -> ErasedSegment {
52        SegmentBuilder::token(id, raw, SyntaxKind::Symbol).finish()
53    }
54
55    pub fn node(
56        id: u32,
57        syntax_kind: SyntaxKind,
58        dialect: DialectKind,
59        segments: Vec<ErasedSegment>,
60    ) -> Self {
61        SegmentBuilder {
62            node_or_token: NodeOrToken {
63                id,
64                syntax_kind,
65                class_types: class_types(syntax_kind),
66                position_marker: None,
67                code_idx: OnceCell::new(),
68                kind: NodeOrTokenKind::Node(NodeData {
69                    dialect,
70                    segments,
71                    raw: Default::default(),
72                    source_fixes: vec![],
73                    descendant_type_set: Default::default(),
74                    raw_segments_with_ancestors: Default::default(),
75                }),
76                hash: OnceCell::new(),
77            },
78        }
79    }
80
81    pub fn token(id: u32, raw: &str, syntax_kind: SyntaxKind) -> Self {
82        SegmentBuilder {
83            node_or_token: NodeOrToken {
84                id,
85                syntax_kind,
86                code_idx: OnceCell::new(),
87                class_types: class_types(syntax_kind),
88                position_marker: None,
89                kind: NodeOrTokenKind::Token(TokenData { raw: raw.into() }),
90                hash: OnceCell::new(),
91            },
92        }
93    }
94
95    pub fn position_from_segments(mut self) -> Self {
96        let segments = match &self.node_or_token.kind {
97            NodeOrTokenKind::Node(node) => &node.segments[..],
98            NodeOrTokenKind::Token(_) => &[],
99        };
100
101        self.node_or_token.position_marker = pos_marker(segments).into();
102        self
103    }
104
105    pub fn with_position(mut self, position: PositionMarker) -> Self {
106        self.node_or_token.position_marker = Some(position);
107        self
108    }
109
110    pub fn finish(self) -> ErasedSegment {
111        ErasedSegment {
112            value: Rc::new(self.node_or_token),
113        }
114    }
115}
116
117#[derive(Debug, Default)]
118pub struct Tables {
119    counter: Cell<u32>,
120}
121
122impl Tables {
123    pub fn next_id(&self) -> u32 {
124        let id = self.counter.get();
125        self.counter.set(id + 1);
126        id
127    }
128}
129
130#[derive(Debug, Clone)]
131pub struct ErasedSegment {
132    pub(crate) value: Rc<NodeOrToken>,
133}
134
135impl Hash for ErasedSegment {
136    fn hash<H: Hasher>(&self, state: &mut H) {
137        self.hash_value().hash(state);
138    }
139}
140
141impl Eq for ErasedSegment {}
142
143impl ErasedSegment {
144    pub fn raw(&self) -> &SmolStr {
145        match &self.value.kind {
146            NodeOrTokenKind::Node(node) => node.raw.get_or_init(|| {
147                SmolStr::from_iter(self.segments().iter().map(|segment| segment.raw().as_str()))
148            }),
149            NodeOrTokenKind::Token(token) => &token.raw,
150        }
151    }
152
153    pub fn segments(&self) -> &[ErasedSegment] {
154        match &self.value.kind {
155            NodeOrTokenKind::Node(node) => &node.segments,
156            NodeOrTokenKind::Token(_) => &[],
157        }
158    }
159
160    pub fn get_type(&self) -> SyntaxKind {
161        self.value.syntax_kind
162    }
163
164    pub fn is_type(&self, kind: SyntaxKind) -> bool {
165        self.get_type() == kind
166    }
167
168    pub fn is_meta(&self) -> bool {
169        matches!(
170            self.value.syntax_kind,
171            SyntaxKind::Indent | SyntaxKind::Implicit | SyntaxKind::Dedent | SyntaxKind::EndOfFile
172        )
173    }
174
175    pub fn is_code(&self) -> bool {
176        match &self.value.kind {
177            NodeOrTokenKind::Node(node) => node.segments.iter().any(|s| s.is_code()),
178            NodeOrTokenKind::Token(_) => {
179                !self.is_comment() && !self.is_whitespace() && !self.is_meta()
180            }
181        }
182    }
183
184    pub fn get_raw_segments(&self) -> Vec<ErasedSegment> {
185        self.recursive_crawl_all(false)
186            .into_iter()
187            .filter(|it| it.segments().is_empty())
188            .collect()
189    }
190
191    #[cfg(feature = "stringify")]
192    pub fn stringify(&self, code_only: bool) -> String {
193        serde_yaml::to_string(&self.to_serialised(code_only, true)).unwrap()
194    }
195
196    pub fn child(&self, seg_types: &SyntaxSet) -> Option<ErasedSegment> {
197        self.segments()
198            .iter()
199            .find(|seg| seg_types.contains(seg.get_type()))
200            .cloned()
201    }
202
203    pub fn recursive_crawl(
204        &self,
205        types: &SyntaxSet,
206        recurse_into: bool,
207        no_recursive_types: &SyntaxSet,
208        allow_self: bool,
209    ) -> Vec<ErasedSegment> {
210        let mut acc = Vec::new();
211
212        let matches = allow_self && self.class_types().intersects(types);
213        if matches {
214            acc.push(self.clone());
215        }
216
217        if !self.descendant_type_set().intersects(types) {
218            return acc;
219        }
220
221        if recurse_into || !matches {
222            for seg in self.segments() {
223                if no_recursive_types.is_empty() || !no_recursive_types.contains(seg.get_type()) {
224                    let segments =
225                        seg.recursive_crawl(types, recurse_into, no_recursive_types, true);
226                    acc.extend(segments);
227                }
228            }
229        }
230
231        acc
232    }
233}
234
235impl ErasedSegment {
236    #[track_caller]
237    pub fn new(&self, segments: Vec<ErasedSegment>) -> ErasedSegment {
238        match &self.value.kind {
239            NodeOrTokenKind::Node(node) => SegmentBuilder::node(
240                self.value.id,
241                self.value.syntax_kind,
242                node.dialect,
243                segments,
244            )
245            .with_position(self.get_position_marker().unwrap().clone())
246            .finish(),
247            NodeOrTokenKind::Token(_) => self.deep_clone(),
248        }
249    }
250
251    fn change_segments(&self, segments: Vec<ErasedSegment>) -> ErasedSegment {
252        let NodeOrTokenKind::Node(node) = &self.value.kind else {
253            unimplemented!()
254        };
255
256        ErasedSegment {
257            value: Rc::new(NodeOrToken {
258                id: self.value.id,
259                syntax_kind: self.value.syntax_kind,
260                class_types: self.value.class_types.clone(),
261                position_marker: None,
262                code_idx: OnceCell::new(),
263                kind: NodeOrTokenKind::Node(NodeData {
264                    dialect: node.dialect,
265                    segments,
266                    raw: node.raw.clone(),
267                    source_fixes: node.source_fixes.clone(),
268                    descendant_type_set: node.descendant_type_set.clone(),
269                    raw_segments_with_ancestors: node.raw_segments_with_ancestors.clone(),
270                }),
271                hash: OnceCell::new(),
272            }),
273        }
274    }
275
276    pub fn indent_val(&self) -> i8 {
277        self.value.syntax_kind.indent_val()
278    }
279
280    pub fn can_start_end_non_code(&self) -> bool {
281        matches!(
282            self.value.syntax_kind,
283            SyntaxKind::File | SyntaxKind::Unparsable
284        )
285    }
286
287    pub(crate) fn dialect(&self) -> DialectKind {
288        match &self.value.kind {
289            NodeOrTokenKind::Node(node) => node.dialect,
290            NodeOrTokenKind::Token(_) => todo!(),
291        }
292    }
293
294    pub fn get_start_loc(&self) -> (usize, usize) {
295        match self.get_position_marker() {
296            Some(pos_marker) => pos_marker.working_loc(),
297            None => unreachable!("{self:?} has no PositionMarker"),
298        }
299    }
300
301    pub fn get_end_loc(&self) -> (usize, usize) {
302        match self.get_position_marker() {
303            Some(pos_marker) => pos_marker.working_loc_after(self.raw()),
304            None => {
305                unreachable!("{self:?} has no PositionMarker")
306            }
307        }
308    }
309
310    pub fn is_templated(&self) -> bool {
311        if let Some(pos_marker) = self.get_position_marker() {
312            pos_marker.source_slice.start != pos_marker.source_slice.end && !pos_marker.is_literal()
313        } else {
314            panic!("PosMarker must be set");
315        }
316    }
317
318    pub fn iter_segments(&self, expanding: &SyntaxSet, pass_through: bool) -> Vec<ErasedSegment> {
319        let capacity = if expanding.is_empty() {
320            self.segments().len()
321        } else {
322            0
323        };
324        let mut result = Vec::with_capacity(capacity);
325        for segment in self.segments() {
326            if expanding.contains(segment.get_type()) {
327                let expanding = if pass_through {
328                    expanding
329                } else {
330                    &SyntaxSet::EMPTY
331                };
332                result.append(&mut segment.iter_segments(expanding, false));
333            } else {
334                result.push(segment.clone());
335            }
336        }
337        result
338    }
339
340    pub(crate) fn code_indices(&self) -> Rc<Vec<usize>> {
341        self.value
342            .code_idx
343            .get_or_init(|| {
344                Rc::from(
345                    self.segments()
346                        .iter()
347                        .enumerate()
348                        .filter(|(_, seg)| seg.is_code())
349                        .map(|(idx, _)| idx)
350                        .collect::<Vec<_>>(),
351                )
352            })
353            .clone()
354    }
355
356    pub fn children(
357        &self,
358        seg_types: &'static SyntaxSet,
359    ) -> impl Iterator<Item = &ErasedSegment> + '_ {
360        self.segments()
361            .iter()
362            .filter(move |seg| seg_types.contains(seg.get_type()))
363    }
364
365    pub fn iter_patches(&self, templated_file: &TemplatedFile) -> Vec<FixPatch> {
366        let mut acc = Vec::new();
367
368        let templated_raw = &templated_file.templated_str.as_ref().unwrap()
369            [self.get_position_marker().unwrap().templated_slice.clone()];
370        if self.raw() == templated_raw {
371            acc.extend(self.iter_source_fix_patches(templated_file));
372            return acc;
373        }
374
375        if self.get_position_marker().is_none() {
376            return Vec::new();
377        }
378
379        let pos_marker = self.get_position_marker().unwrap();
380        if pos_marker.is_literal() {
381            acc.extend(self.iter_source_fix_patches(templated_file));
382            acc.push(FixPatch::new(
383                pos_marker.templated_slice.clone(),
384                self.raw().clone(),
385                // SyntaxKind::Literal.into(),
386                pos_marker.source_slice.clone(),
387                templated_file.templated_str.as_ref().unwrap()[pos_marker.templated_slice.clone()]
388                    .to_string(),
389                templated_file.source_str[pos_marker.source_slice.clone()].to_string(),
390            ));
391        } else if self.segments().is_empty() {
392            return acc;
393        } else {
394            let mut segments = self.segments();
395
396            while !segments.is_empty()
397                && matches!(
398                    segments.last().unwrap().get_type(),
399                    SyntaxKind::EndOfFile
400                        | SyntaxKind::Indent
401                        | SyntaxKind::Dedent
402                        | SyntaxKind::Implicit
403                )
404            {
405                segments = &segments[..segments.len() - 1];
406            }
407
408            let pos = self.get_position_marker().unwrap();
409            let mut source_idx = pos.source_slice.start;
410            let mut templated_idx = pos.templated_slice.start;
411            let mut insert_buff = String::new();
412
413            for segment in segments {
414                let pos_marker = segment.get_position_marker().unwrap();
415                if !segment.raw().is_empty() && pos_marker.is_point() {
416                    insert_buff.push_str(segment.raw().as_ref());
417                    continue;
418                }
419
420                let start_diff = pos_marker.templated_slice.start - templated_idx;
421
422                if start_diff > 0 || !insert_buff.is_empty() {
423                    let fixed_raw = std::mem::take(&mut insert_buff);
424                    let raw_segments = segment.get_raw_segments();
425                    let first_segment_pos = raw_segments[0].get_position_marker().unwrap();
426
427                    acc.push(FixPatch::new(
428                        templated_idx..first_segment_pos.templated_slice.start,
429                        fixed_raw.into(),
430                        source_idx..first_segment_pos.source_slice.start,
431                        String::new(),
432                        String::new(),
433                    ));
434                }
435
436                acc.extend(segment.iter_patches(templated_file));
437
438                source_idx = pos_marker.source_slice.end;
439                templated_idx = pos_marker.templated_slice.end;
440            }
441
442            let end_diff = pos.templated_slice.end - templated_idx;
443            if end_diff != 0 || !insert_buff.is_empty() {
444                let source_slice = source_idx..pos.source_slice.end;
445                let templated_slice = templated_idx..pos.templated_slice.end;
446
447                let templated_str = templated_file.templated_str.as_ref().unwrap()
448                    [templated_slice.clone()]
449                .to_owned();
450                let source_str = templated_file.source_str[source_slice.clone()].to_owned();
451
452                acc.push(FixPatch::new(
453                    templated_slice,
454                    insert_buff.into(),
455                    source_slice,
456                    templated_str,
457                    source_str,
458                ));
459            }
460        }
461
462        acc
463    }
464
465    pub fn descendant_type_set(&self) -> &SyntaxSet {
466        match &self.value.kind {
467            NodeOrTokenKind::Node(node) => node.descendant_type_set.get_or_init(|| {
468                self.segments()
469                    .iter()
470                    .flat_map(|segment| {
471                        segment
472                            .descendant_type_set()
473                            .clone()
474                            .union(segment.class_types())
475                    })
476                    .collect()
477            }),
478            NodeOrTokenKind::Token(_) => const { &SyntaxSet::EMPTY },
479        }
480    }
481
482    pub fn is_comment(&self) -> bool {
483        matches!(
484            self.value.syntax_kind,
485            SyntaxKind::Comment | SyntaxKind::InlineComment | SyntaxKind::BlockComment
486        )
487    }
488
489    pub fn is_whitespace(&self) -> bool {
490        matches!(
491            self.value.syntax_kind,
492            SyntaxKind::Whitespace | SyntaxKind::Newline
493        )
494    }
495
496    pub fn is_indent(&self) -> bool {
497        matches!(
498            self.value.syntax_kind,
499            SyntaxKind::Indent | SyntaxKind::Implicit | SyntaxKind::Dedent
500        )
501    }
502
503    pub fn get_position_marker(&self) -> Option<&PositionMarker> {
504        self.value.position_marker.as_ref()
505    }
506
507    pub(crate) fn iter_source_fix_patches(&self, templated_file: &TemplatedFile) -> Vec<FixPatch> {
508        let source_fixes = self.get_source_fixes();
509        let mut patches = Vec::with_capacity(source_fixes.len());
510
511        for source_fix in &source_fixes {
512            patches.push(FixPatch::new(
513                source_fix.templated_slice.clone(),
514                source_fix.edit.clone(),
515                // String::from("source"),
516                source_fix.source_slice.clone(),
517                templated_file.templated_str.clone().unwrap()[source_fix.templated_slice.clone()]
518                    .to_string(),
519                templated_file.source_str[source_fix.source_slice.clone()].to_string(),
520            ));
521        }
522
523        patches
524    }
525
526    pub fn id(&self) -> u32 {
527        self.value.id
528    }
529
530    /// Return any source fixes as list.
531    pub fn get_source_fixes(&self) -> Vec<SourceFix> {
532        match &self.value.kind {
533            NodeOrTokenKind::Node(node) => node.source_fixes.clone(),
534            NodeOrTokenKind::Token(_) => Vec::new(),
535        }
536    }
537
538    pub fn edit(
539        &self,
540        id: u32,
541        raw: Option<String>,
542        _source_fixes: Option<Vec<SourceFix>>,
543    ) -> ErasedSegment {
544        match &self.value.kind {
545            NodeOrTokenKind::Node(_node) => {
546                todo!()
547            }
548            NodeOrTokenKind::Token(token) => {
549                let raw = raw.as_deref().unwrap_or(token.raw.as_ref());
550                SegmentBuilder::token(id, raw, self.value.syntax_kind)
551                    .with_position(self.get_position_marker().unwrap().clone())
552                    .finish()
553            }
554        }
555    }
556
557    pub fn class_types(&self) -> &SyntaxSet {
558        &self.value.class_types
559    }
560
561    pub(crate) fn first_non_whitespace_segment_raw_upper(&self) -> Option<String> {
562        for seg in self.get_raw_segments() {
563            if !seg.raw().is_empty() {
564                return Some(seg.raw().to_uppercase());
565            }
566        }
567        None
568    }
569
570    pub fn is(&self, other: &ErasedSegment) -> bool {
571        Rc::ptr_eq(&self.value, &other.value)
572    }
573
574    pub fn addr(&self) -> usize {
575        Rc::as_ptr(&self.value).addr()
576    }
577
578    pub fn direct_descendant_type_set(&self) -> SyntaxSet {
579        self.segments()
580            .iter()
581            .fold(SyntaxSet::EMPTY, |set, it| set.union(it.class_types()))
582    }
583
584    pub fn is_keyword(&self, p0: &str) -> bool {
585        self.is_type(SyntaxKind::Keyword) && self.raw().eq_ignore_ascii_case(p0)
586    }
587
588    pub fn hash_value(&self) -> u64 {
589        *self.value.hash.get_or_init(|| {
590            let mut hasher = ahash::AHasher::default();
591            self.get_type().hash(&mut hasher);
592            self.raw().hash(&mut hasher);
593
594            if let Some(marker) = &self.get_position_marker() {
595                marker.source_position().hash(&mut hasher);
596            } else {
597                None::<usize>.hash(&mut hasher);
598            }
599
600            hasher.finish()
601        })
602    }
603
604    pub fn deep_clone(&self) -> Self {
605        Self {
606            value: Rc::new(self.value.as_ref().clone()),
607        }
608    }
609
610    #[track_caller]
611    pub(crate) fn get_mut(&mut self) -> &mut NodeOrToken {
612        Rc::get_mut(&mut self.value).unwrap()
613    }
614
615    #[track_caller]
616    pub(crate) fn make_mut(&mut self) -> &mut NodeOrToken {
617        Rc::make_mut(&mut self.value)
618    }
619
620    pub fn reference(&self) -> ObjectReferenceSegment {
621        ObjectReferenceSegment(
622            self.clone(),
623            match self.get_type() {
624                SyntaxKind::TableReference => ObjectReferenceKind::Table,
625                SyntaxKind::WildcardIdentifier => ObjectReferenceKind::WildcardIdentifier,
626                _ => ObjectReferenceKind::Object,
627            },
628        )
629    }
630
631    pub fn recursive_crawl_all(&self, reverse: bool) -> Vec<ErasedSegment> {
632        let mut result = Vec::with_capacity(self.segments().len() + 1);
633
634        if reverse {
635            for seg in self.segments().iter().rev() {
636                result.append(&mut seg.recursive_crawl_all(reverse));
637            }
638            result.push(self.clone());
639        } else {
640            result.push(self.clone());
641            for seg in self.segments() {
642                result.append(&mut seg.recursive_crawl_all(reverse));
643            }
644        }
645
646        result
647    }
648
649    pub fn raw_segments_with_ancestors(&self) -> &[(ErasedSegment, Vec<PathStep>)] {
650        match &self.value.kind {
651            NodeOrTokenKind::Node(node) => node.raw_segments_with_ancestors.get_or_init(|| {
652                let mut buffer: Vec<(ErasedSegment, Vec<PathStep>)> =
653                    Vec::with_capacity(self.segments().len());
654                let code_idxs = self.code_indices();
655
656                for (idx, seg) in self.segments().iter().enumerate() {
657                    let new_step = vec![PathStep {
658                        segment: self.clone(),
659                        idx,
660                        len: self.segments().len(),
661                        code_idxs: code_idxs.clone(),
662                    }];
663
664                    // Use seg.get_segments().is_empty() as a workaround to check if the segment is
665                    // a SyntaxKind::Raw type. In the original Python code, this was achieved
666                    // using seg.is_type(SyntaxKind::Raw). Here, we assume that a SyntaxKind::Raw
667                    // segment is characterized by having no sub-segments.
668
669                    if seg.segments().is_empty() {
670                        buffer.push((seg.clone(), new_step));
671                    } else {
672                        let extended =
673                            seg.raw_segments_with_ancestors()
674                                .iter()
675                                .map(|(raw_seg, stack)| {
676                                    let mut new_step = new_step.clone();
677                                    new_step.extend_from_slice(stack);
678                                    (raw_seg.clone(), new_step)
679                                });
680
681                        buffer.extend(extended);
682                    }
683                }
684
685                buffer
686            }),
687            NodeOrTokenKind::Token(_) => &[],
688        }
689    }
690
691    pub fn path_to(&self, other: &ErasedSegment) -> Vec<PathStep> {
692        let midpoint = other;
693
694        for (idx, seg) in enumerate(self.segments()) {
695            let mut steps = vec![PathStep {
696                segment: self.clone(),
697                idx,
698                len: self.segments().len(),
699                code_idxs: self.code_indices(),
700            }];
701
702            if seg.eq(midpoint) {
703                return steps;
704            }
705
706            let res = seg.path_to(midpoint);
707
708            if !res.is_empty() {
709                steps.extend(res);
710                return steps;
711            }
712        }
713
714        Vec::new()
715    }
716
717    pub fn apply_fixes(
718        &self,
719        fixes: &mut FxHashMap<u32, AnchorEditInfo>,
720    ) -> (ErasedSegment, Vec<ErasedSegment>, Vec<ErasedSegment>) {
721        if fixes.is_empty() || self.segments().is_empty() {
722            return (self.clone(), Vec::new(), Vec::new());
723        }
724
725        let mut seg_buffer = Vec::new();
726        let mut has_applied_fixes = false;
727        let mut _requires_validate = false;
728
729        for seg in self.segments() {
730            // Look for uuid match.
731            // This handles potential positioning ambiguity.
732
733            let Some(mut anchor_info) = fixes.remove(&seg.id()) else {
734                seg_buffer.push(seg.clone());
735                continue;
736            };
737
738            if anchor_info.fixes.len() == 2
739                && matches!(anchor_info.fixes[0], LintFix::CreateAfter { .. })
740            {
741                anchor_info.fixes.reverse();
742            }
743
744            let fixes_count = anchor_info.fixes.len();
745            for lint_fix in anchor_info.fixes {
746                has_applied_fixes = true;
747
748                // Deletes are easy.
749                if matches!(lint_fix, LintFix::Delete { .. }) {
750                    // We're just getting rid of this segment.
751                    _requires_validate = true;
752                    // NOTE: We don't add the segment in this case.
753                    continue;
754                }
755
756                // Otherwise it must be a replace or a create.
757                assert!(matches!(
758                    lint_fix,
759                    LintFix::Replace { .. }
760                        | LintFix::CreateBefore { .. }
761                        | LintFix::CreateAfter { .. }
762                ));
763
764                match lint_fix {
765                    LintFix::CreateAfter { edit, .. } => {
766                        if fixes_count == 1 {
767                            // In the case of a creation after that is not part
768                            // of a create_before/create_after pair, also add
769                            // this segment before the edit.
770                            seg_buffer.push(seg.clone());
771                        }
772                        for s in edit {
773                            seg_buffer.push(s);
774                        }
775                        _requires_validate = true;
776                    }
777                    LintFix::CreateBefore { edit, .. } => {
778                        for s in edit {
779                            seg_buffer.push(s);
780                        }
781                        seg_buffer.push(seg.clone());
782                        _requires_validate = true;
783                    }
784                    LintFix::Replace { edit, .. } => {
785                        let mut consumed_pos = false;
786                        let is_single_same_type =
787                            edit.len() == 1 && edit[0].class_types() == seg.class_types();
788
789                        for mut s in edit {
790                            if !consumed_pos && s.raw() == seg.raw() {
791                                consumed_pos = true;
792                                s.make_mut()
793                                    .set_position_marker(seg.get_position_marker().cloned());
794                            }
795                            seg_buffer.push(s);
796                        }
797
798                        if !is_single_same_type {
799                            _requires_validate = true;
800                        }
801                    }
802                    LintFix::Delete { .. } => {
803                        // Already handled above
804                        unreachable!()
805                    }
806                }
807            }
808        }
809
810        if has_applied_fixes {
811            seg_buffer =
812                position_segments(&seg_buffer, self.get_position_marker().as_ref().unwrap());
813        }
814
815        let seg_queue = seg_buffer;
816        let mut seg_buffer = Vec::new();
817        for seg in seg_queue {
818            let (mid, pre, post) = seg.apply_fixes(fixes);
819
820            seg_buffer.extend(pre);
821            seg_buffer.push(mid);
822            seg_buffer.extend(post);
823        }
824
825        let seg_buffer =
826            position_segments(&seg_buffer, self.get_position_marker().as_ref().unwrap());
827        (self.new(seg_buffer), Vec::new(), Vec::new())
828    }
829}
830
831#[cfg(any(test, feature = "serde"))]
832pub mod serde {
833    use serde::ser::SerializeMap;
834    use serde::{Deserialize, Serialize};
835
836    use crate::parser::segments::ErasedSegment;
837
838    #[derive(Serialize, Deserialize)]
839    #[serde(untagged)]
840    pub enum SerialisedSegmentValue {
841        Single(String),
842        Nested(Vec<TupleSerialisedSegment>),
843    }
844
845    #[derive(Deserialize)]
846    pub struct TupleSerialisedSegment(String, SerialisedSegmentValue);
847
848    impl Serialize for TupleSerialisedSegment {
849        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
850        where
851            S: serde::Serializer,
852        {
853            let mut map = serializer.serialize_map(None)?;
854            map.serialize_key(&self.0)?;
855            map.serialize_value(&self.1)?;
856            map.end()
857        }
858    }
859
860    impl TupleSerialisedSegment {
861        pub fn sinlge(key: String, value: String) -> Self {
862            Self(key, SerialisedSegmentValue::Single(value))
863        }
864
865        pub fn nested(key: String, segments: Vec<TupleSerialisedSegment>) -> Self {
866            Self(key, SerialisedSegmentValue::Nested(segments))
867        }
868    }
869
870    impl ErasedSegment {
871        pub fn to_serialised(&self, code_only: bool, show_raw: bool) -> TupleSerialisedSegment {
872            if show_raw && self.segments().is_empty() {
873                TupleSerialisedSegment::sinlge(
874                    self.get_type().as_str().to_string(),
875                    self.raw().to_string(),
876                )
877            } else if code_only {
878                let segments = self
879                    .segments()
880                    .iter()
881                    .filter(|seg| seg.is_code() && !seg.is_meta())
882                    .map(|seg| seg.to_serialised(code_only, show_raw))
883                    .collect::<Vec<_>>();
884
885                TupleSerialisedSegment::nested(self.get_type().as_str().to_string(), segments)
886            } else {
887                let segments = self
888                    .segments()
889                    .iter()
890                    .map(|seg| seg.to_serialised(code_only, show_raw))
891                    .collect::<Vec<_>>();
892
893                TupleSerialisedSegment::nested(self.get_type().as_str().to_string(), segments)
894            }
895        }
896    }
897}
898
899impl PartialEq for ErasedSegment {
900    fn eq(&self, other: &Self) -> bool {
901        if self.id() == other.id() {
902            return true;
903        }
904
905        let pos_self = self.get_position_marker();
906        let pos_other = other.get_position_marker();
907        if let Some((pos_self, pos_other)) = pos_self.zip(pos_other) {
908            self.get_type() == other.get_type()
909                && pos_self.working_loc() == pos_other.working_loc()
910                && self.raw() == other.raw()
911        } else {
912            false
913        }
914    }
915}
916
917pub fn position_segments(
918    segments: &[ErasedSegment],
919    parent_pos: &PositionMarker,
920) -> Vec<ErasedSegment> {
921    if segments.is_empty() {
922        return Vec::new();
923    }
924
925    let (mut line_no, mut line_pos) = { (parent_pos.working_line_no, parent_pos.working_line_pos) };
926
927    let mut segment_buffer: Vec<ErasedSegment> = Vec::new();
928    for (idx, segment) in enumerate(segments) {
929        let old_position = segment.get_position_marker();
930
931        let mut new_position = match old_position {
932            Some(pos_marker) => pos_marker.clone(),
933            None => {
934                let start_point = if idx > 0 {
935                    let prev_seg = segment_buffer[idx - 1].clone();
936                    Some(prev_seg.get_position_marker().unwrap().end_point_marker())
937                } else {
938                    Some(parent_pos.start_point_marker())
939                };
940
941                let mut end_point = None;
942                for fwd_seg in &segments[idx + 1..] {
943                    if fwd_seg.get_position_marker().is_some() {
944                        end_point = Some(
945                            fwd_seg.get_raw_segments()[0]
946                                .get_position_marker()
947                                .unwrap()
948                                .start_point_marker(),
949                        );
950                        break;
951                    }
952                }
953
954                if let Some((start_point, end_point)) = start_point
955                    .as_ref()
956                    .zip(end_point.as_ref())
957                    .filter(|(start_point, end_point)| start_point != end_point)
958                {
959                    PositionMarker::from_points(start_point, end_point)
960                } else if let Some(start_point) = start_point.as_ref() {
961                    start_point.clone()
962                } else if let Some(end_point) = end_point.as_ref() {
963                    end_point.clone()
964                } else {
965                    unimplemented!("Unable to position new segment")
966                }
967            }
968        };
969
970        new_position = new_position.with_working_position(line_no, line_pos);
971        (line_no, line_pos) = PositionMarker::infer_next_position(segment.raw(), line_no, line_pos);
972
973        let mut new_seg = if !segment.segments().is_empty() && old_position != Some(&new_position) {
974            let child_segments = position_segments(segment.segments(), &new_position);
975            segment.change_segments(child_segments)
976        } else {
977            segment.deep_clone()
978        };
979
980        new_seg.get_mut().set_position_marker(new_position.into());
981        segment_buffer.push(new_seg);
982    }
983
984    segment_buffer
985}
986
987#[derive(Debug, Clone)]
988pub struct NodeOrToken {
989    id: u32,
990    syntax_kind: SyntaxKind,
991    class_types: SyntaxSet,
992    position_marker: Option<PositionMarker>,
993    kind: NodeOrTokenKind,
994    code_idx: OnceCell<Rc<Vec<usize>>>,
995    hash: OnceCell<u64>,
996}
997
998#[derive(Debug, Clone)]
999pub enum NodeOrTokenKind {
1000    Node(NodeData),
1001    Token(TokenData),
1002}
1003
1004impl NodeOrToken {
1005    pub fn set_position_marker(&mut self, position_marker: Option<PositionMarker>) {
1006        self.position_marker = position_marker;
1007    }
1008
1009    pub fn set_id(&mut self, id: u32) {
1010        self.id = id;
1011    }
1012}
1013
1014#[derive(Debug, Clone)]
1015pub struct NodeData {
1016    dialect: DialectKind,
1017    segments: Vec<ErasedSegment>,
1018    raw: OnceCell<SmolStr>,
1019    source_fixes: Vec<SourceFix>,
1020    descendant_type_set: OnceCell<SyntaxSet>,
1021    raw_segments_with_ancestors: OnceCell<Vec<(ErasedSegment, Vec<PathStep>)>>,
1022}
1023
1024#[derive(Debug, Clone, PartialEq)]
1025pub struct TokenData {
1026    raw: SmolStr,
1027}
1028
1029#[track_caller]
1030pub fn pos_marker(segments: &[ErasedSegment]) -> PositionMarker {
1031    let markers = segments.iter().filter_map(|seg| seg.get_position_marker());
1032
1033    PositionMarker::from_child_markers(markers)
1034}
1035
1036#[derive(Debug, Clone)]
1037pub struct PathStep {
1038    pub segment: ErasedSegment,
1039    pub idx: usize,
1040    pub len: usize,
1041    pub code_idxs: Rc<Vec<usize>>,
1042}
1043
1044fn class_types(syntax_kind: SyntaxKind) -> SyntaxSet {
1045    match syntax_kind {
1046        SyntaxKind::ColumnReference => SyntaxSet::new(&[SyntaxKind::ObjectReference, syntax_kind]),
1047        SyntaxKind::WildcardIdentifier => {
1048            SyntaxSet::new(&[SyntaxKind::WildcardIdentifier, SyntaxKind::ObjectReference])
1049        }
1050        SyntaxKind::TableReference => SyntaxSet::new(&[SyntaxKind::ObjectReference, syntax_kind]),
1051        _ => SyntaxSet::single(syntax_kind),
1052    }
1053}
1054
1055#[cfg(test)]
1056mod tests {
1057    use super::*;
1058    use crate::lint_fix::LintFix;
1059    use crate::linter::compute_anchor_edit_info;
1060    use crate::parser::segments::test_functions::{raw_seg, raw_segments};
1061
1062    #[test]
1063    /// Test comparison of raw segments.
1064    fn test_parser_base_segments_raw_compare() {
1065        let template: TemplatedFile = "foobar".into();
1066        let rs1 = SegmentBuilder::token(0, "foobar", SyntaxKind::Word)
1067            .with_position(PositionMarker::new(
1068                0..6,
1069                0..6,
1070                template.clone(),
1071                None,
1072                None,
1073            ))
1074            .finish();
1075        let rs2 = SegmentBuilder::token(0, "foobar", SyntaxKind::Word)
1076            .with_position(PositionMarker::new(
1077                0..6,
1078                0..6,
1079                template.clone(),
1080                None,
1081                None,
1082            ))
1083            .finish();
1084
1085        assert_eq!(rs1, rs2)
1086    }
1087
1088    #[test]
1089    // TODO Implement
1090    /// Test raw segments behave as expected.
1091    fn test_parser_base_segments_raw() {
1092        let raw_seg = raw_seg();
1093
1094        assert_eq!(raw_seg.raw(), "foobar");
1095    }
1096
1097    #[test]
1098    /// Test BaseSegment.compute_anchor_edit_info().
1099    fn test_parser_base_segments_compute_anchor_edit_info() {
1100        let raw_segs = raw_segments();
1101        let tables = Tables::default();
1102
1103        // Construct a fix buffer, intentionally with:
1104        // - one duplicate.
1105        // - two different incompatible fixes on the same segment.
1106        let fixes = vec![
1107            LintFix::replace(
1108                raw_segs[0].clone(),
1109                vec![raw_segs[0].edit(tables.next_id(), Some("a".to_string()), None)],
1110                None,
1111            ),
1112            LintFix::replace(
1113                raw_segs[0].clone(),
1114                vec![raw_segs[0].edit(tables.next_id(), Some("a".to_string()), None)],
1115                None,
1116            ),
1117            LintFix::replace(
1118                raw_segs[0].clone(),
1119                vec![raw_segs[0].edit(tables.next_id(), Some("b".to_string()), None)],
1120                None,
1121            ),
1122        ];
1123
1124        let mut anchor_edit_info = Default::default();
1125        compute_anchor_edit_info(&mut anchor_edit_info, fixes);
1126
1127        // Check the target segment is the only key we have.
1128        assert_eq!(
1129            anchor_edit_info.keys().collect::<Vec<_>>(),
1130            vec![&raw_segs[0].id()]
1131        );
1132
1133        let anchor_info = anchor_edit_info.get(&raw_segs[0].id()).unwrap();
1134
1135        // Check that the duplicate as been deduplicated i.e. this isn't 3.
1136        assert_eq!(anchor_info.replace, 2);
1137
1138        // Check the fixes themselves.
1139        //   Note: There's no duplicated first fix.
1140        assert_eq!(
1141            anchor_info.fixes[0],
1142            LintFix::replace(
1143                raw_segs[0].clone(),
1144                vec![raw_segs[0].edit(tables.next_id(), Some("a".to_string()), None)],
1145                None,
1146            )
1147        );
1148        assert_eq!(
1149            anchor_info.fixes[1],
1150            LintFix::replace(
1151                raw_segs[0].clone(),
1152                vec![raw_segs[0].edit(tables.next_id(), Some("b".to_string()), None)],
1153                None,
1154            )
1155        );
1156
1157        // Check the first replace
1158        assert_eq!(
1159            anchor_info.fixes[anchor_info.first_replace.unwrap()],
1160            LintFix::replace(
1161                raw_segs[0].clone(),
1162                vec![raw_segs[0].edit(tables.next_id(), Some("a".to_string()), None)],
1163                None,
1164            )
1165        );
1166    }
1167}
sqruff_lib_core/parser/segments.rs

sqruff_lib_core/parser/
segments.rs
