sqruff_lib/utils/reflow/
depth_map.rs

1use ahash::{AHashMap, AHashSet};
2use nohash_hasher::{IntMap, IntSet};
3use sqruff_lib_core::dialects::syntax::SyntaxSet;
4use sqruff_lib_core::parser::segments::{ErasedSegment, PathStep};
5
6/// An element of the stack_positions property of DepthInfo.
7#[derive(Debug, PartialEq, Eq, Clone)]
8pub struct StackPosition {
9    pub idx: usize,
10    pub len: usize,
11    pub type_: Option<StackPositionType>,
12}
13
14#[derive(Debug, PartialEq, Eq, Clone)]
15pub enum StackPositionType {
16    Solo,
17    Start,
18    End,
19}
20
21impl StackPosition {
22    /// Interpret a path step for stack_positions.
23    fn stack_pos_interpreter(path_step: &PathStep) -> Option<StackPositionType> {
24        if path_step.code_idxs.is_empty() {
25            None
26        } else if path_step.code_idxs.len() == 1 {
27            Some(StackPositionType::Solo)
28        } else if path_step.idx == *path_step.code_idxs.first().unwrap() {
29            Some(StackPositionType::Start)
30        } else if path_step.idx == *path_step.code_idxs.last().unwrap() {
31            Some(StackPositionType::End)
32        } else {
33            None
34        }
35    }
36
37    /// Interpret a PathStep to construct a StackPosition
38    fn from_path_step(path_step: &PathStep) -> StackPosition {
39        StackPosition {
40            idx: path_step.idx,
41            len: path_step.len,
42            type_: StackPosition::stack_pos_interpreter(path_step),
43        }
44    }
45}
46
47pub struct DepthMap {
48    depth_info: AHashMap<u32, DepthInfo>,
49}
50
51impl DepthMap {
52    fn new<'a>(raws_with_stack: impl Iterator<Item = &'a (ErasedSegment, Vec<PathStep>)>) -> Self {
53        let depth_info = raws_with_stack
54            .into_iter()
55            .map(|(raw, stack)| (raw.id(), DepthInfo::from_stack(stack)))
56            .collect();
57        Self { depth_info }
58    }
59
60    pub fn get_depth_info(&self, seg: &ErasedSegment) -> DepthInfo {
61        self.depth_info[&seg.id()].clone()
62    }
63
64    pub fn copy_depth_info(
65        &mut self,
66        anchor: &ErasedSegment,
67        new_segment: &ErasedSegment,
68        trim: u32,
69    ) {
70        self.depth_info.insert(
71            new_segment.id(),
72            self.get_depth_info(anchor).trim(trim.try_into().unwrap()),
73        );
74    }
75
76    pub fn from_parent(parent: &ErasedSegment) -> Self {
77        Self::new(parent.raw_segments_with_ancestors().iter())
78    }
79
80    pub fn from_raws_and_root(
81        raw_segments: impl Iterator<Item = ErasedSegment>,
82        root_segment: &ErasedSegment,
83    ) -> DepthMap {
84        let depth_info = raw_segments
85            .into_iter()
86            .map(|raw| {
87                let stack = root_segment.path_to(&raw);
88                (raw.id(), DepthInfo::from_stack(&stack))
89            })
90            .collect();
91
92        DepthMap { depth_info }
93    }
94}
95
96/// An object to hold the depth information for a specific raw segment.
97#[derive(Debug, PartialEq, Eq, Clone)]
98pub struct DepthInfo {
99    pub stack_depth: usize,
100    pub stack_hashes: Vec<u64>,
101    /// This is a convenience cache to speed up operations.
102    pub stack_hash_set: IntSet<u64>,
103    pub stack_class_types: Vec<SyntaxSet>,
104    pub stack_positions: IntMap<u64, StackPosition>,
105}
106
107impl DepthInfo {
108    fn from_stack(stack: &[PathStep]) -> DepthInfo {
109        // Build all structures in a single pass to avoid repeated iteration and
110        // intermediate allocations.
111        let mut stack_hashes = Vec::with_capacity(stack.len());
112        let mut stack_hash_set: IntSet<u64> = IntSet::default();
113        stack_hash_set.reserve(stack.len());
114        let mut stack_class_types = Vec::with_capacity(stack.len());
115        let mut stack_positions: IntMap<u64, StackPosition> = IntMap::default();
116        stack_positions.reserve(stack.len());
117
118        for path in stack {
119            let hash = path.segment.hash_value();
120            stack_hashes.push(hash);
121            stack_hash_set.insert(hash);
122            stack_class_types.push(path.segment.class_types().clone());
123            stack_positions.insert(hash, StackPosition::from_path_step(path));
124        }
125
126        DepthInfo {
127            stack_depth: stack_hashes.len(),
128            stack_hashes,
129            stack_hash_set,
130            stack_class_types,
131            stack_positions,
132        }
133    }
134
135    pub fn trim(self, amount: usize) -> DepthInfo {
136        // Return a DepthInfo object with some amount trimmed.
137        if amount == 0 {
138            // The trivial case.
139            return self;
140        }
141
142        let slice_set: IntSet<_> = IntSet::from_iter(
143            self.stack_hashes[self.stack_hashes.len() - amount..]
144                .iter()
145                .copied(),
146        );
147
148        let new_hash_set: IntSet<_> = self
149            .stack_hash_set
150            .difference(&slice_set)
151            .copied()
152            .collect();
153
154        let stack_positions = self
155            .stack_positions
156            .into_iter()
157            .filter(|(hash, _)| new_hash_set.contains(hash))
158            .collect();
159
160        DepthInfo {
161            stack_depth: self.stack_depth - amount,
162            stack_hashes: self.stack_hashes[..self.stack_hashes.len() - amount].to_vec(),
163            stack_hash_set: new_hash_set,
164            stack_class_types: self.stack_class_types[..self.stack_class_types.len() - amount]
165                .to_vec(),
166            stack_positions,
167        }
168    }
169
170    pub fn common_with(&self, other: &DepthInfo) -> Vec<u64> {
171        // Get the common depth and hashes with the other.
172        // We use AHashSet intersection because it's efficient and hashes should be
173        // unique.
174
175        let common_hashes: AHashSet<_> = self
176            .stack_hash_set
177            .intersection(&other.stack_hash_set)
178            .copied()
179            .collect();
180
181        // We should expect there to be _at least_ one common ancestor, because
182        // they should share the same file segment. If that's not the case we
183        // should error because it's likely a bug or programming error.
184        assert!(
185            !common_hashes.is_empty(),
186            "DepthInfo comparison shares no common ancestor!"
187        );
188
189        let common_depth = common_hashes.len();
190        self.stack_hashes
191            .iter()
192            .take(common_depth)
193            .copied()
194            .collect()
195    }
196}
sqruff_lib/utils/reflow/depth_map.rs

sqruff_lib/utils/reflow/
depth_map.rs
