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

sqruff_lib_core/parser/
segments.rs
