span/
map.rs

1//! A map that maps a span to every position in a file. Usually maps a span to some range of positions.
2//! Allows bidirectional lookup.
3
4use std::{fmt, hash::Hash};
5
6use stdx::{always, itertools::Itertools};
7
8use crate::{
9    EditionedFileId, ErasedFileAstId, ROOT_ERASED_FILE_AST_ID, Span, SpanAnchor, SyntaxContext,
10    TextRange, TextSize,
11};
12
13/// Maps absolute text ranges for the corresponding file to the relevant span data.
14#[derive(Debug, PartialEq, Eq, Clone, Hash)]
15pub struct SpanMap {
16    /// The offset stored here is the *end* of the node.
17    spans: Vec<(TextSize, Span)>,
18    /// Index of the matched macro arm on successful expansion for declarative macros.
19    // FIXME: Does it make sense to have this here?
20    pub matched_arm: Option<u32>,
21}
22
23impl SpanMap {
24    /// Creates a new empty [`SpanMap`].
25    pub fn empty() -> Self {
26        Self { spans: Vec::new(), matched_arm: None }
27    }
28
29    /// Finalizes the [`SpanMap`], shrinking its backing storage and validating that the offsets are
30    /// in order.
31    pub fn finish(&mut self) {
32        always!(
33            self.spans.iter().tuple_windows().all(|(a, b)| a.0 < b.0),
34            "spans are not in order"
35        );
36        self.spans.shrink_to_fit();
37    }
38
39    /// Pushes a new span onto the [`SpanMap`].
40    pub fn push(&mut self, offset: TextSize, span: Span) {
41        if cfg!(debug_assertions)
42            && let Some(&(last_offset, _)) = self.spans.last()
43        {
44            assert!(
45                last_offset < offset,
46                "last_offset({last_offset:?}) must be smaller than offset({offset:?})"
47            );
48        }
49        self.spans.push((offset, span));
50    }
51
52    /// Returns all [`TextRange`]s that correspond to the given span.
53    ///
54    /// Note this does a linear search through the entire backing vector.
55    pub fn ranges_with_span_exact(
56        &self,
57        span: Span,
58    ) -> impl Iterator<Item = (TextRange, SyntaxContext)> + '_ {
59        self.spans.iter().enumerate().filter_map(move |(idx, &(end, s))| {
60            if !s.eq_ignoring_ctx(span) {
61                return None;
62            }
63            let start = idx.checked_sub(1).map_or(TextSize::new(0), |prev| self.spans[prev].0);
64            Some((TextRange::new(start, end), s.ctx))
65        })
66    }
67
68    /// Returns all [`TextRange`]s whose spans contain the given span.
69    ///
70    /// Note this does a linear search through the entire backing vector.
71    pub fn ranges_with_span(
72        &self,
73        span: Span,
74    ) -> impl Iterator<Item = (TextRange, SyntaxContext)> + '_ {
75        self.spans.iter().enumerate().filter_map(move |(idx, &(end, s))| {
76            if s.anchor != span.anchor {
77                return None;
78            }
79            if !s.range.contains_range(span.range) {
80                return None;
81            }
82            let start = idx.checked_sub(1).map_or(TextSize::new(0), |prev| self.spans[prev].0);
83            Some((TextRange::new(start, end), s.ctx))
84        })
85    }
86
87    /// Returns the span at the given position.
88    pub fn span_at(&self, offset: TextSize) -> Span {
89        let entry = self.spans.partition_point(|&(it, _)| it <= offset);
90        self.spans[entry].1
91    }
92
93    /// Returns the spans associated with the given range.
94    /// In other words, this will return all spans that correspond to all offsets within the given range.
95    pub fn spans_for_range(&self, range: TextRange) -> impl Iterator<Item = Span> + '_ {
96        let (start, end) = (range.start(), range.end());
97        let start_entry = self.spans.partition_point(|&(it, _)| it <= start);
98        let end_entry = self.spans[start_entry..].partition_point(|&(it, _)| it <= end); // FIXME: this might be wrong?
99        self.spans[start_entry..][..end_entry].iter().map(|&(_, s)| s)
100    }
101
102    pub fn iter(&self) -> impl Iterator<Item = (TextSize, Span)> + '_ {
103        self.spans.iter().copied()
104    }
105
106    /// Merges this span map with another span map, where `other` is inserted at (and replaces) `other_range`.
107    ///
108    /// The length of the replacement node needs to be `other_size`.
109    pub fn merge(&mut self, other_range: TextRange, other_size: TextSize, other: &SpanMap) {
110        // I find the following diagram helpful to illustrate the bounds and why we use `<` or `<=`:
111        // --------------------------------------------------------------------
112        //   1   3   5   6   7   10    11          <-- offsets we store
113        // 0-1 1-3 3-5 5-6 6-7 7-10 10-11          <-- ranges these offsets refer to
114        //       3   ..      7                     <-- other_range
115        //         3-5 5-6 6-7                     <-- ranges we replace (len = 7-3 = 4)
116        //         ^^^^^^^^^^^ ^^^^^^^^^^
117        //           remove       shift
118        //   2   3   5   9                         <-- offsets we insert
119        // 0-2 2-3 3-5 5-9                         <-- ranges we insert (other_size = 9-0 = 9)
120        // ------------------------------------
121        //   1   3
122        // 0-1 1-3                                 <-- these remain intact
123        //           5   6   8   12
124        //         3-5 5-6 6-8 8-12                <-- we shift these by other_range.start() and insert them
125        //                             15    16
126        //                          12-15 15-16    <-- we shift these by other_size-other_range.len() = 9-4 = 5
127        // ------------------------------------
128        //   1   3   5   6   8   12    15    16    <-- final offsets we store
129        // 0-1 1-3 3-5 5-6 6-8 8-12 12-15 15-16    <-- final ranges
130
131        self.spans.retain_mut(|(offset, _)| {
132            if other_range.start() < *offset && *offset <= other_range.end() {
133                false
134            } else {
135                if *offset > other_range.end() {
136                    *offset += other_size;
137                    *offset -= other_range.len();
138                }
139                true
140            }
141        });
142
143        self.spans
144            .extend(other.spans.iter().map(|&(offset, span)| (offset + other_range.start(), span)));
145
146        self.spans.sort_unstable_by_key(|&(offset, _)| offset);
147
148        // Matched arm info is no longer correct once we have multiple macros.
149        self.matched_arm = None;
150    }
151}
152
153#[cfg(not(no_salsa_async_drops))]
154impl Drop for SpanMap {
155    fn drop(&mut self) {
156        let spans = std::mem::take(&mut self.spans);
157        static SPAN_MAP_DROP_THREAD: std::sync::OnceLock<
158            std::sync::mpsc::Sender<Vec<(TextSize, Span)>>,
159        > = std::sync::OnceLock::new();
160
161        SPAN_MAP_DROP_THREAD
162            .get_or_init(|| {
163                let (sender, receiver) = std::sync::mpsc::channel::<Vec<(TextSize, Span)>>();
164                std::thread::Builder::new()
165                    .name("SpanMapDropper".to_owned())
166                    .spawn(move || {
167                        loop {
168                            // block on a receive
169                            _ = receiver.recv();
170                            // then drain the entire channel
171                            while receiver.try_recv().is_ok() {}
172                            // and sleep for a bit
173                            std::thread::sleep(std::time::Duration::from_millis(100));
174                        }
175                        // why do this over just a `receiver.iter().for_each(drop)`? To reduce contention on the channel lock.
176                        // otherwise this thread will constantly wake up and sleep again.
177                    })
178                    .unwrap();
179                sender
180            })
181            .send(spans)
182            .unwrap();
183    }
184}
185
186#[derive(PartialEq, Eq, Hash, Debug)]
187pub struct RealSpanMap {
188    file_id: EditionedFileId,
189    /// Invariant: Sorted vec over TextSize
190    // FIXME: SortedVec<(TextSize, ErasedFileAstId)>?
191    pairs: Box<[(TextSize, ErasedFileAstId)]>,
192    end: TextSize,
193}
194
195impl fmt::Display for RealSpanMap {
196    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
197        writeln!(f, "RealSpanMap({:?}):", self.file_id)?;
198        for span in self.pairs.iter() {
199            writeln!(f, "{}: {:#?}", u32::from(span.0), span.1)?;
200        }
201        Ok(())
202    }
203}
204
205impl RealSpanMap {
206    /// Creates a real file span map that returns absolute ranges (relative ranges to the root ast id).
207    pub fn absolute(file_id: EditionedFileId) -> Self {
208        RealSpanMap {
209            file_id,
210            pairs: Box::from([(TextSize::new(0), ROOT_ERASED_FILE_AST_ID)]),
211            end: TextSize::new(!0),
212        }
213    }
214
215    pub fn from_file(
216        file_id: EditionedFileId,
217        pairs: Box<[(TextSize, ErasedFileAstId)]>,
218        end: TextSize,
219    ) -> Self {
220        Self { file_id, pairs, end }
221    }
222
223    pub fn span_for_range(&self, range: TextRange) -> Span {
224        assert!(
225            range.end() <= self.end,
226            "range {range:?} goes beyond the end of the file {:?}",
227            self.end
228        );
229        let start = range.start();
230        let idx = self
231            .pairs
232            .binary_search_by(|&(it, _)| it.cmp(&start).then(std::cmp::Ordering::Less))
233            .unwrap_err();
234        let (offset, ast_id) = self.pairs[idx - 1];
235        Span {
236            range: range - offset,
237            anchor: SpanAnchor { file_id: self.file_id, ast_id },
238            ctx: SyntaxContext::root(self.file_id.edition()),
239        }
240    }
241}