squawk_line_index/
lib.rs

1//! See [`LineIndex`].
2//!
3//! Forked from rust-analyzer's `line-index`. Our changes are tagged with a
4//! `SQUAWK:` comment so they can be reapplied when resyncing with upstream:
5//!
6//! - `\r` and `\r\n` start a new line, not just `\n`. Upstream only breaks on
7//!   `\n`, which leaves a classic Mac file looking like one enormous line.
8//! - [`newlines`] adds helpers upstream doesn't have.
9
10#![deny(missing_debug_implementations, missing_docs, rust_2018_idioms)]
11
12#[allow(missing_debug_implementations, missing_docs)]
13mod newlines;
14#[cfg(test)]
15mod tests;
16
17use nohash_hasher::IntMap;
18
19pub use newlines::{
20    Line, LineEnding, NewlineWithTrailingNewline, UniversalNewlineIterator, UniversalNewlines,
21    find_newline,
22};
23pub use text_size::{TextRange, TextSize};
24
25/// `(line, column)` information in the native, UTF-8 encoding.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
27pub struct LineCol {
28    /// Zero-based.
29    pub line: u32,
30    /// Zero-based UTF-8 offset.
31    pub col: u32,
32}
33
34/// A kind of wide character encoding.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
36#[non_exhaustive]
37pub enum WideEncoding {
38    /// UTF-16.
39    Utf16,
40    /// UTF-32.
41    Utf32,
42}
43
44impl WideEncoding {
45    /// Returns the number of code units it takes to encode `text` in this encoding.
46    pub fn measure(&self, text: &str) -> usize {
47        match self {
48            WideEncoding::Utf16 => text.encode_utf16().count(),
49            WideEncoding::Utf32 => text.chars().count(),
50        }
51    }
52}
53
54/// `(line, column)` information in wide encodings.
55///
56/// See [`WideEncoding`] for the kinds of wide encodings available.
57//
58// Deliberately not a generic type and different from `LineCol`.
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
60pub struct WideLineCol {
61    /// Zero-based.
62    pub line: u32,
63    /// Zero-based.
64    pub col: u32,
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68struct WideChar {
69    /// Start offset of a character inside a line, zero-based.
70    start: TextSize,
71    /// End offset of a character inside a line, zero-based.
72    end: TextSize,
73}
74
75impl WideChar {
76    /// Returns the length in 8-bit UTF-8 code units.
77    fn len(&self) -> TextSize {
78        self.end - self.start
79    }
80
81    /// Returns the length in UTF-16 or UTF-32 code units.
82    fn wide_len(&self, enc: WideEncoding) -> u32 {
83        match enc {
84            WideEncoding::Utf16 => {
85                if self.len() == TextSize::from(4) {
86                    2
87                } else {
88                    1
89                }
90            }
91            WideEncoding::Utf32 => 1,
92        }
93    }
94}
95
96/// Maps flat [`TextSize`] offsets to/from `(line, column)` representation.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct LineIndex {
99    /// Offset the beginning of each line (except the first, which always has offset 0).
100    newlines: Box<[TextSize]>,
101    /// List of non-ASCII characters on each line.
102    line_wide_chars: IntMap<u32, Box<[WideChar]>>,
103    /// The length of the entire text.
104    len: TextSize,
105}
106
107impl LineIndex {
108    /// Returns a `LineIndex` for the `text`.
109    pub fn new(text: &str) -> LineIndex {
110        let (newlines, line_wide_chars) = analyze_source_file(text);
111        LineIndex {
112            newlines: newlines.into_boxed_slice(),
113            line_wide_chars,
114            len: TextSize::of(text),
115        }
116    }
117
118    /// Transforms the `TextSize` into a `LineCol`.
119    ///
120    /// # Panics
121    ///
122    /// If the offset is invalid. See [`Self::try_line_col`].
123    pub fn line_col(&self, offset: TextSize) -> LineCol {
124        self.try_line_col(offset).expect("invalid offset")
125    }
126
127    /// Transforms the `TextSize` into a `LineCol`.
128    ///
129    /// Returns `None` if the `offset` was invalid, e.g. if it extends past the end of the text or
130    /// points to the middle of a multi-byte character.
131    pub fn try_line_col(&self, offset: TextSize) -> Option<LineCol> {
132        if offset > self.len {
133            return None;
134        }
135        let line = self.newlines.partition_point(|&it| it <= offset);
136        let start = self.start_offset(line)?;
137        let col = offset - start;
138        let ret = LineCol {
139            line: line as u32,
140            col: col.into(),
141        };
142        self.line_wide_chars
143            .get(&ret.line)
144            .into_iter()
145            .flat_map(|it| it.iter())
146            .all(|it| col <= it.start || it.end <= col)
147            .then_some(ret)
148    }
149
150    /// Transforms the `LineCol` into a `TextSize`.
151    pub fn offset(&self, line_col: LineCol) -> Option<TextSize> {
152        self.start_offset(line_col.line as usize)
153            .map(|start| start + TextSize::from(line_col.col))
154    }
155
156    fn start_offset(&self, line: usize) -> Option<TextSize> {
157        match line.checked_sub(1) {
158            None => Some(TextSize::from(0)),
159            Some(it) => self.newlines.get(it).copied(),
160        }
161    }
162
163    /// Transforms the `LineCol` with the given `WideEncoding` into a `WideLineCol`.
164    pub fn to_wide(&self, enc: WideEncoding, line_col: LineCol) -> Option<WideLineCol> {
165        let mut col = line_col.col;
166        if let Some(wide_chars) = self.line_wide_chars.get(&line_col.line) {
167            for c in wide_chars {
168                if u32::from(c.end) <= line_col.col {
169                    col = col.checked_sub(u32::from(c.len()) - c.wide_len(enc))?;
170                } else {
171                    // From here on, all utf16 characters come *after* the character we are mapping,
172                    // so we don't need to take them into account
173                    break;
174                }
175            }
176        }
177        Some(WideLineCol {
178            line: line_col.line,
179            col,
180        })
181    }
182
183    /// Transforms the `WideLineCol` with the given `WideEncoding` into a `LineCol`.
184    pub fn to_utf8(&self, enc: WideEncoding, line_col: WideLineCol) -> Option<LineCol> {
185        let mut col = line_col.col;
186        if let Some(wide_chars) = self.line_wide_chars.get(&line_col.line) {
187            for c in wide_chars {
188                if col > u32::from(c.start) {
189                    col = col.checked_add(u32::from(c.len()) - c.wide_len(enc))?;
190                } else {
191                    // From here on, all utf16 characters come *after* the character we are mapping,
192                    // so we don't need to take them into account
193                    break;
194                }
195            }
196        }
197        Some(LineCol {
198            line: line_col.line,
199            col,
200        })
201    }
202
203    /// Returns the given line's range.
204    pub fn line(&self, line: u32) -> Option<TextRange> {
205        let start = self.start_offset(line as usize)?;
206        let next_newline = self
207            .newlines
208            .get(line as usize)
209            .copied()
210            .unwrap_or(self.len);
211        let line_length = next_newline - start;
212        Some(TextRange::new(start, start + line_length))
213    }
214
215    /// Given a range [start, end), returns a sorted iterator of non-empty ranges [start, x1), [x1,
216    /// x2), ..., [xn, end) where all the xi, which are positions of newlines, are inside the range
217    /// [start, end).
218    pub fn lines(&self, range: TextRange) -> impl Iterator<Item = TextRange> + '_ {
219        let lo = self.newlines.partition_point(|&it| it < range.start());
220        let hi = self.newlines.partition_point(|&it| it <= range.end());
221        let all = std::iter::once(range.start())
222            .chain(self.newlines[lo..hi].iter().copied())
223            .chain(std::iter::once(range.end()));
224
225        all.clone()
226            .zip(all.skip(1))
227            .map(|(lo, hi)| TextRange::new(lo, hi))
228            .filter(|it| !it.is_empty())
229    }
230
231    /// Returns the length of the original text.
232    pub fn len(&self) -> TextSize {
233        self.len
234    }
235}
236
237/// This is adapted from the `rustc_span` crate, <https://github.com/rust-lang/rust/blob/de59844c98f7925242a798a72c59dc3610dd0e2c/compiler/rustc_span/src/analyze_source_file.rs>
238fn analyze_source_file(src: &str) -> (Vec<TextSize>, IntMap<u32, Box<[WideChar]>>) {
239    assert!(src.len() < !0u32 as usize);
240    let mut lines = vec![];
241    let mut line_wide_chars = IntMap::<u32, Vec<WideChar>>::default();
242
243    // Calls the right implementation, depending on hardware support available.
244    analyze_source_file_dispatch(src, &mut lines, &mut line_wide_chars);
245
246    (
247        lines,
248        line_wide_chars
249            .into_iter()
250            .map(|(k, v)| (k, v.into_boxed_slice()))
251            .collect(),
252    )
253}
254
255#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
256fn analyze_source_file_dispatch(
257    src: &str,
258    lines: &mut Vec<TextSize>,
259    multi_byte_chars: &mut IntMap<u32, Vec<WideChar>>,
260) {
261    if is_x86_feature_detected!("sse2") {
262        // SAFETY: SSE2 support was checked
263        unsafe {
264            analyze_source_file_sse2(src, lines, multi_byte_chars);
265        }
266    } else {
267        analyze_source_file_generic(src, src.len(), TextSize::from(0), lines, multi_byte_chars);
268    }
269}
270
271#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
272fn analyze_source_file_dispatch(
273    src: &str,
274    lines: &mut Vec<TextSize>,
275    multi_byte_chars: &mut IntMap<u32, Vec<WideChar>>,
276) {
277    if std::arch::is_aarch64_feature_detected!("neon") {
278        // SAFETY: NEON support was checked
279        unsafe {
280            analyze_source_file_neon(src, lines, multi_byte_chars);
281        }
282    } else {
283        analyze_source_file_generic(src, src.len(), TextSize::from(0), lines, multi_byte_chars);
284    }
285}
286
287/// Checks 16 byte chunks of text at a time. If the chunk contains
288/// something other than printable ASCII characters and newlines, the
289/// function falls back to the generic implementation. Otherwise it uses
290/// SSE2 intrinsics to quickly find all newlines.
291#[target_feature(enable = "sse2")]
292#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
293// This can be removed once 1.87 is stable due to some intrinsics switching to safe.
294#[allow(unsafe_op_in_unsafe_fn)]
295unsafe fn analyze_source_file_sse2(
296    src: &str,
297    lines: &mut Vec<TextSize>,
298    multi_byte_chars: &mut IntMap<u32, Vec<WideChar>>,
299) {
300    #[cfg(target_arch = "x86")]
301    use std::arch::x86::*;
302    #[cfg(target_arch = "x86_64")]
303    use std::arch::x86_64::*;
304
305    const CHUNK_SIZE: usize = 16;
306
307    let src_bytes = src.as_bytes();
308
309    let chunk_count = src.len() / CHUNK_SIZE;
310
311    // This variable keeps track of where we should start decoding a
312    // chunk. If a multi-byte character spans across chunk boundaries,
313    // we need to skip that part in the next chunk because we already
314    // handled it.
315    let mut intra_chunk_offset = 0;
316
317    for chunk_index in 0..chunk_count {
318        let ptr = src_bytes.as_ptr() as *const __m128i;
319        // We don't know if the pointer is aligned to 16 bytes, so we
320        // use `loadu`, which supports unaligned loading.
321        let chunk = unsafe { _mm_loadu_si128(ptr.add(chunk_index)) };
322
323        // For character in the chunk, see if its byte value is < 0, which
324        // indicates that it's part of a UTF-8 char.
325        let multibyte_test = _mm_cmplt_epi8(chunk, _mm_set1_epi8(0));
326        // Create a bit mask from the comparison results.
327        let multibyte_mask = _mm_movemask_epi8(multibyte_test);
328
329        // SQUAWK: `\r` needs to look at the following byte to tell `\r` from
330        // `\r\n`, which the chunked mask can't do, so we send any chunk holding
331        // one down the generic path instead.
332        let cr_test = _mm_cmpeq_epi8(chunk, _mm_set1_epi8(b'\r' as i8));
333        let cr_mask = _mm_movemask_epi8(cr_test);
334
335        // If the bit mask is all zero, we only have ASCII chars here:
336        if multibyte_mask == 0 && cr_mask == 0 {
337            assert!(intra_chunk_offset == 0);
338
339            // Check for newlines in the chunk
340            let newlines_test = _mm_cmpeq_epi8(chunk, _mm_set1_epi8(b'\n' as i8));
341            let newlines_mask = _mm_movemask_epi8(newlines_test);
342
343            if newlines_mask != 0 {
344                // All control characters are newlines, record them
345                let mut newlines_mask = 0xFFFF0000 | newlines_mask as u32;
346                let output_offset = TextSize::from((chunk_index * CHUNK_SIZE + 1) as u32);
347
348                loop {
349                    let index = newlines_mask.trailing_zeros();
350
351                    if index >= CHUNK_SIZE as u32 {
352                        // We have arrived at the end of the chunk.
353                        break;
354                    }
355
356                    lines.push(TextSize::from(index) + output_offset);
357
358                    // Clear the bit, so we can find the next one.
359                    newlines_mask &= (!1) << index;
360                }
361            }
362            continue;
363        }
364
365        // The slow path.
366        // There are control chars in here, fallback to generic decoding.
367        let scan_start = chunk_index * CHUNK_SIZE + intra_chunk_offset;
368        intra_chunk_offset = analyze_source_file_generic(
369            &src[scan_start..],
370            CHUNK_SIZE - intra_chunk_offset,
371            TextSize::from(scan_start as u32),
372            lines,
373            multi_byte_chars,
374        );
375    }
376
377    // There might still be a tail left to analyze
378    let tail_start = chunk_count * CHUNK_SIZE + intra_chunk_offset;
379    if tail_start < src.len() {
380        analyze_source_file_generic(
381            &src[tail_start..],
382            src.len() - tail_start,
383            TextSize::from(tail_start as u32),
384            lines,
385            multi_byte_chars,
386        );
387    }
388}
389
390#[target_feature(enable = "neon")]
391#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
392#[inline]
393// See https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon
394//
395// The mask is a 64-bit integer, where each 4-bit corresponds to a u8 in the
396// input vector. The least significant 4 bits correspond to the first byte in
397// the vector.
398// This can be removed once 1.87 is stable due to some intrinsics switching to safe.
399#[allow(unsafe_op_in_unsafe_fn)]
400unsafe fn move_mask(v: std::arch::aarch64::uint8x16_t) -> u64 {
401    use std::arch::aarch64::*;
402
403    let nibble_mask = vshrn_n_u16(vreinterpretq_u16_u8(v), 4);
404    vget_lane_u64(vreinterpret_u64_u8(nibble_mask), 0)
405}
406
407#[target_feature(enable = "neon")]
408#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
409// This can be removed once 1.87 is stable due to some intrinsics switching to safe.
410#[allow(unsafe_op_in_unsafe_fn)]
411unsafe fn analyze_source_file_neon(
412    src: &str,
413    lines: &mut Vec<TextSize>,
414    multi_byte_chars: &mut IntMap<u32, Vec<WideChar>>,
415) {
416    use std::arch::aarch64::*;
417
418    const CHUNK_SIZE: usize = 16;
419
420    let src_bytes = src.as_bytes();
421
422    let chunk_count = src.len() / CHUNK_SIZE;
423
424    let newline = vdupq_n_s8(b'\n' as i8);
425    // SQUAWK: see the matching comment in the sse2 version
426    let carriage_return = vdupq_n_s8(b'\r' as i8);
427
428    // This variable keeps track of where we should start decoding a
429    // chunk. If a multi-byte character spans across chunk boundaries,
430    // we need to skip that part in the next chunk because we already
431    // handled it.
432    let mut intra_chunk_offset = 0;
433
434    for chunk_index in 0..chunk_count {
435        let ptr = src_bytes.as_ptr() as *const i8;
436        let chunk = unsafe { vld1q_s8(ptr.add(chunk_index * CHUNK_SIZE)) };
437
438        // For character in the chunk, see if its byte value is < 0, which
439        // indicates that it's part of a UTF-8 char.
440        let multibyte_test = vcltzq_s8(chunk);
441        // Create a bit mask from the comparison results.
442        let multibyte_mask = unsafe { move_mask(multibyte_test) };
443
444        // SQUAWK: any chunk with a `\r` goes down the generic path
445        let cr_test = vceqq_s8(chunk, carriage_return);
446        let cr_mask = unsafe { move_mask(cr_test) };
447
448        // If the bit mask is all zero, we only have ASCII chars here:
449        if multibyte_mask == 0 && cr_mask == 0 {
450            assert!(intra_chunk_offset == 0);
451
452            // Check for newlines in the chunk
453            let newlines_test = vceqq_s8(chunk, newline);
454            let mut newlines_mask = unsafe { move_mask(newlines_test) };
455
456            // If the bit mask is not all zero, there are newlines in this chunk.
457            if newlines_mask != 0 {
458                let output_offset = TextSize::from((chunk_index * CHUNK_SIZE + 1) as u32);
459
460                while newlines_mask != 0 {
461                    let trailing_zeros = newlines_mask.trailing_zeros();
462                    let index = trailing_zeros / 4;
463
464                    lines.push(TextSize::from(index) + output_offset);
465
466                    // Clear the current 4-bit, so we can find the next one.
467                    newlines_mask &= (!0xF) << trailing_zeros;
468                }
469            }
470            continue;
471        }
472
473        let scan_start = chunk_index * CHUNK_SIZE + intra_chunk_offset;
474        intra_chunk_offset = analyze_source_file_generic(
475            &src[scan_start..],
476            CHUNK_SIZE - intra_chunk_offset,
477            TextSize::from(scan_start as u32),
478            lines,
479            multi_byte_chars,
480        );
481    }
482
483    let tail_start = chunk_count * CHUNK_SIZE + intra_chunk_offset;
484    if tail_start < src.len() {
485        analyze_source_file_generic(
486            &src[tail_start..],
487            src.len() - tail_start,
488            TextSize::from(tail_start as u32),
489            lines,
490            multi_byte_chars,
491        );
492    }
493}
494
495#[cfg(not(any(
496    target_arch = "x86",
497    target_arch = "x86_64",
498    all(target_arch = "aarch64", target_endian = "little")
499)))]
500// The target (or compiler version) does not support SSE2 ...
501fn analyze_source_file_dispatch(
502    src: &str,
503    lines: &mut Vec<TextSize>,
504    multi_byte_chars: &mut IntMap<u32, Vec<WideChar>>,
505) {
506    analyze_source_file_generic(src, src.len(), TextSize::from(0), lines, multi_byte_chars);
507}
508
509// `scan_len` determines the number of bytes in `src` to scan. Note that the
510// function can read past `scan_len` if a multi-byte character start within the
511// range but extends past it. The overflow is returned by the function.
512fn analyze_source_file_generic(
513    src: &str,
514    scan_len: usize,
515    output_offset: TextSize,
516    lines: &mut Vec<TextSize>,
517    multi_byte_chars: &mut IntMap<u32, Vec<WideChar>>,
518) -> usize {
519    assert!(src.len() >= scan_len);
520    let mut i = 0;
521    let src_bytes = src.as_bytes();
522
523    while i < scan_len {
524        let byte = unsafe {
525            // We verified that i < scan_len <= src.len()
526            *src_bytes.get_unchecked(i)
527        };
528
529        // How much to advance in order to get to the next UTF-8 char in the
530        // string.
531        let mut char_len = 1;
532
533        if byte == b'\n' {
534            lines.push(TextSize::from(i as u32 + 1) + output_offset);
535        // SQUAWK: `\r` and `\r\n` start a new line too. A `\r\n` pair is a
536        // single break, so we skip the `\r` and let the `\n` push the start.
537        // The `\n` can live past `scan_len` when the pair straddles a chunk
538        // boundary, in which case the next chunk pushes it.
539        } else if byte == b'\r' && src_bytes.get(i + 1) != Some(&b'\n') {
540            lines.push(TextSize::from(i as u32 + 1) + output_offset);
541        } else if byte >= 127 {
542            // The slow path: Just decode to `char`.
543            let c = src[i..].chars().next().unwrap();
544            char_len = c.len_utf8();
545
546            // The last element of `lines` represents the offset of the start of
547            // current line. To get the offset inside the line, we subtract it.
548            let pos = TextSize::from(i as u32) + output_offset
549                - lines.last().unwrap_or(&TextSize::default());
550
551            if char_len > 1 {
552                assert!((2..=4).contains(&char_len));
553                let mbc = WideChar {
554                    start: pos,
555                    end: pos + TextSize::from(char_len as u32),
556                };
557                multi_byte_chars
558                    .entry(lines.len() as u32)
559                    .or_default()
560                    .push(mbc);
561            }
562        }
563
564        i += char_len;
565    }
566
567    i - scan_len
568}
squawk_line_index/lib.rs

squawk_line_index/
lib.rs
