Skip to main content

ruff_db/
parsed.rs

1use std::fmt::Formatter;
2use std::sync::Arc;
3
4use arc_swap::ArcSwapOption;
5use get_size2::GetSize;
6use ruff_python_ast::{
7    AnyRootNodeRef, HasNodeIndex, ModExpression, ModModule, NodeIndex, NodeIndexError,
8    PythonVersion, StringLiteral,
9};
10use ruff_python_parser::{
11    ParseError, ParseErrorType, ParseOptions, Parsed, parse_cells_unchecked,
12    parse_string_annotation, parse_unchecked,
13};
14
15use crate::files::File;
16use crate::source::source_text;
17use crate::{Db, PythonFile};
18
19/// Returns the parsed AST of `file`, including its token stream.
20///
21/// The query uses Ruff's error-resilient parser. That means that the parser always succeeds to produce an
22/// AST even if the file contains syntax errors. The parse errors
23/// are then accessible through [`Parsed::errors`].
24///
25/// The query is only cached when the [`source_text()`] hasn't changed. This is because
26/// comparing two ASTs is a non-trivial operation and every offset change is directly
27/// reflected in the changed AST offsets.
28/// The other reason is that Ruff's AST doesn't implement `Eq` which Salsa requires
29/// for determining if a query result is unchanged.
30///
31/// The LRU capacity of 200 was picked without any empirical evidence that it's optimal,
32/// instead it's a wild guess that it should be unlikely that incremental changes involve
33/// more than 200 modules. Parsed ASTs within the same revision are never evicted by Salsa.
34#[salsa::tracked(returns(ref), no_eq, heap_size=ruff_memory_usage::heap_size, lru=200)]
35pub fn parsed_module(db: &dyn Db, file: PythonFile<'_>) -> ParsedModule {
36    let source_file = file.file(db);
37    let python_version = file.python_version(db);
38    let _span = tracing::trace_span!("parsed_module", ?source_file, %python_version).entered();
39
40    let parsed = parsed_module_impl(db, source_file, python_version);
41
42    ParsedModule::new(source_file, python_version, parsed)
43}
44
45pub(super) fn disable_lru(db: &mut dyn Db) {
46    parsed_module::set_lru_capacity(db, 0);
47}
48
49fn parsed_module_impl(db: &dyn Db, file: File, target_version: PythonVersion) -> Parsed<ModModule> {
50    let source = source_text(db, file);
51    let ty = file.source_type(db);
52
53    let options = ParseOptions::from(ty).with_target_version(target_version);
54
55    // Notebooks parse each cell as an independent module so a syntax error confined to one cell is
56    // surfaced instead of being masked by a later cell's content. Regular files take the existing
57    // single-parse path.
58    if let Some(notebook) = source.as_notebook() {
59        parse_cells_unchecked(&source, notebook.cell_offsets().content_ranges(), &options)
60    } else {
61        parse_unchecked(&source, options)
62            .try_into_module()
63            .expect("PySourceType always parses into a module")
64    }
65}
66
67pub fn parsed_string_annotation(
68    source: &str,
69    string: &StringLiteral,
70) -> Result<Parsed<ModExpression>, ParseError> {
71    let expr = parse_string_annotation(source, string)?;
72
73    // We need the sub-ast of the string annotation to be indexed
74    indexed::ensure_indexed(&expr, string.node_index().load()).map_err(|err| {
75        let message = match err {
76            NodeIndexError::NoParent => {
77                "Internal error: string annotation's parent had no NodeIndex"
78            }
79            NodeIndexError::TooNested => {
80                "Too many levels of nested string annotations; \
81                remove the redundant nested quotes"
82            }
83            NodeIndexError::OverflowedIndices => {
84                "File too long for string annotations; either break up the file \
85                or don't use string annotations"
86            }
87            NodeIndexError::OverflowedSubIndices => {
88                "File too long for nested string annotations; remove the redundant nested quotes"
89            }
90            NodeIndexError::ExhaustedSubIndices => {
91                "String annotation is too long; consider introducing type aliases to simplify"
92            }
93            NodeIndexError::ExhaustedSubSubIndices => {
94                "Nested string annotation is too long; remove the redundant nested quotes"
95            }
96        };
97
98        ParseError {
99            error: ParseErrorType::StringAnnotationError(message),
100            location: string.range,
101        }
102    })?;
103
104    Ok(expr)
105}
106
107/// A wrapper around a parsed module.
108///
109/// This type manages instances of the module AST. A particular instance of the AST
110/// is represented with the [`ParsedModuleRef`] type.
111#[derive(Clone, get_size2::GetSize)]
112pub struct ParsedModule {
113    file: File,
114    python_version: PythonVersion,
115    #[get_size(size_fn = arc_swap_size)]
116    inner: Arc<ArcSwapOption<indexed::IndexedModule>>,
117}
118
119impl ParsedModule {
120    pub fn new(file: File, python_version: PythonVersion, parsed: Parsed<ModModule>) -> Self {
121        Self {
122            file,
123            python_version,
124            inner: Arc::new(ArcSwapOption::new(Some(indexed::IndexedModule::new(
125                parsed,
126            )))),
127        }
128    }
129    /// Loads a reference to the parsed module.
130    ///
131    /// Note that holding on to the reference will prevent garbage collection
132    /// of the AST. This method will reparse the module if it has been collected.
133    pub fn load(&self, db: &dyn Db) -> ParsedModuleRef {
134        let parsed = match self.inner.load_full() {
135            Some(parsed) => parsed,
136            None => {
137                // Re-parse the file.
138                let parsed = indexed::IndexedModule::new(parsed_module_impl(
139                    db,
140                    self.file,
141                    self.python_version,
142                ));
143                tracing::debug!(
144                    "File `{}` was reparsed after being collected in the current Salsa revision",
145                    self.file.path(db)
146                );
147
148                self.inner.store(Some(parsed.clone()));
149                parsed
150            }
151        };
152
153        ParsedModuleRef {
154            module: self.clone(),
155            indexed: parsed,
156        }
157    }
158
159    /// Clear the parsed module, dropping the AST once all references to it are dropped.
160    pub fn clear(&self) {
161        self.inner.store(None);
162    }
163
164    /// Returns the file to which this module belongs.
165    pub fn file(&self) -> File {
166        self.file
167    }
168
169    /// Returns the Python version used to parse this module.
170    pub fn python_version(&self) -> PythonVersion {
171        self.python_version
172    }
173}
174
175impl std::fmt::Debug for ParsedModule {
176    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
177        f.debug_tuple("ParsedModule").field(&self.inner).finish()
178    }
179}
180
181impl PartialEq for ParsedModule {
182    fn eq(&self, other: &Self) -> bool {
183        Arc::ptr_eq(&self.inner, &other.inner)
184    }
185}
186
187impl Eq for ParsedModule {}
188
189/// Cheap cloneable wrapper around an instance of a module AST.
190#[derive(Clone)]
191pub struct ParsedModuleRef {
192    module: ParsedModule,
193    indexed: Arc<indexed::IndexedModule>,
194}
195
196impl ParsedModuleRef {
197    /// Returns a reference to the [`ParsedModule`] that this instance was loaded from.
198    pub fn module(&self) -> &ParsedModule {
199        &self.module
200    }
201
202    /// Returns a reference to the AST node at the given index.
203    pub fn get_by_index<'ast>(&'ast self, index: NodeIndex) -> AnyRootNodeRef<'ast> {
204        self.indexed.get_by_index(index)
205    }
206}
207
208impl std::ops::Deref for ParsedModuleRef {
209    type Target = Parsed<ModModule>;
210
211    fn deref(&self) -> &Self::Target {
212        &self.indexed.parsed
213    }
214}
215
216/// Returns the heap-size of the currently stored `T` in the `ArcSwap`.
217fn arc_swap_size<T>(arc_swap: &Arc<ArcSwapOption<T>>) -> usize
218where
219    T: GetSize,
220{
221    if let Some(value) = &*arc_swap.load() {
222        T::get_heap_size(value)
223    } else {
224        0
225    }
226}
227
228mod indexed {
229    use std::sync::Arc;
230
231    use ruff_python_ast::visitor::source_order::*;
232    use ruff_python_ast::*;
233    use ruff_python_parser::Parsed;
234
235    /// A wrapper around the AST that allows access to AST nodes by index.
236    #[derive(Debug, get_size2::GetSize)]
237    pub struct IndexedModule {
238        index: IndexedNodes,
239        pub parsed: Parsed<ModModule>,
240    }
241
242    /// Compact storage for the address and [`RootNodeKind`] of every indexed AST node.
243    ///
244    /// This stores the information needed to reconstruct an [`AnyRootNodeRef`] without retaining
245    /// a fat pointer per node. Entries are divided into fixed-size chunks so that unrelated AST
246    /// allocations do not force every node into a wider representation. Each chunk starts on a
247    /// word boundary in `words`.
248    ///
249    /// # Safety invariant
250    ///
251    /// Every entry preserves the exact exposed address and [`RootNodeKind`] obtained from the same
252    /// [`AnyRootNodeRef`]. Relative entries use lossless address arithmetic and wide entries store
253    /// the full address. The parsed AST is placed in its final [`Arc`] before those addresses are
254    /// collected. Installing the completed index does not move or mutate the parsed AST, and no
255    /// API moves, replaces, or mutably exposes it while the index exists. Lookups pair each stored
256    /// address with `NonNull::with_exposed_provenance` and its original kind.
257    ///
258    /// # Memory reporting
259    ///
260    /// The actual number of words used by the index depends on the relative addresses of the AST
261    /// nodes. Allocator placement can vary between processes, which makes exact accounting noisy
262    /// in CI memory comparisons even when the indexed AST is unchanged. Memory reports normalize
263    /// the payload to a fixed 32 bits per entry. This conservatively covers 99% of entries in the
264    /// measured Ruff corpus while preserving the size reduction over storing a full
265    /// [`AnyRootNodeRef`] per node. The actual encoding is often narrower than 32 bits. This only
266    /// affects memory reporting; the index itself continues to use the narrowest lossless
267    /// representation for each chunk.
268    #[derive(Debug, Default)]
269    struct IndexedNodes {
270        chunks: Box<[IndexChunk]>,
271        words: Box<[u64]>,
272    }
273
274    /// Describes the entries for one consecutive group of node indices.
275    #[derive(Debug, get_size2::GetSize)]
276    struct IndexChunk {
277        /// Minimum node address in a relative chunk; unused for a wide chunk.
278        base: usize,
279        /// Index of this chunk's first word in [`IndexedNodes::words`].
280        word_start: u32,
281        /// Number of bits per packed entry.
282        entry_bits: u8,
283        /// Number of entries, which is only less than [`IndexedNodes::CHUNK_LEN`] in the last
284        /// chunk.
285        entry_count: u8,
286        layout: IndexChunkLayout,
287    }
288
289    impl get_size2::GetSize for IndexedNodes {
290        fn get_heap_size_with_tracker<T: get_size2::GetSizeTracker>(
291            &self,
292            tracker: T,
293        ) -> (usize, T) {
294            let (chunks_size, tracker) =
295                get_size2::GetSize::get_heap_size_with_tracker(&self.chunks, tracker);
296            let words = self
297                .chunks
298                .iter()
299                .map(|chunk| {
300                    (usize::from(chunk.entry_count) * Self::REPORTED_ENTRY_BITS)
301                        .div_ceil(u64::BITS as usize)
302                })
303                .sum::<usize>();
304
305            (chunks_size + words * size_of::<u64>(), tracker)
306        }
307    }
308
309    #[derive(Copy, Clone, Debug, get_size2::GetSize)]
310    #[repr(u8)]
311    enum IndexChunkLayout {
312        /// Packs each scaled address offset together with its root-node kind:
313        ///
314        /// ```text
315        /// | address offset | root-node kind |
316        ///   entry_bits - 5       5 bits
317        /// ```
318        ///
319        /// The original address is `base + address_offset * ALIGNMENT`. Entries may cross `u64`
320        /// boundaries and the unused bits at the end of the chunk are padding.
321        Relative,
322        /// Stores full addresses followed by a packed stream of root-node kinds:
323        ///
324        /// ```text
325        /// | address 0 | ... | address n - 1 | kind 0 | ... | kind n - 1 |
326        ///     64 bits           64 bits        5 bits          5 bits
327        /// ```
328        Wide,
329    }
330
331    #[derive(Default)]
332    struct IndexedNodesBuilder<'ast> {
333        chunks: Vec<IndexChunk>,
334        words: Vec<u64>,
335        pending: Vec<AnyRootNodeRef<'ast>>,
336        #[cfg(test)]
337        all_nodes: Vec<AnyRootNodeRef<'ast>>,
338    }
339
340    impl<'ast> IndexedNodesBuilder<'ast> {
341        fn new() -> Self {
342            Self {
343                pending: Vec::with_capacity(IndexedNodes::CHUNK_LEN),
344                ..Self::default()
345            }
346        }
347
348        fn push(&mut self, node: AnyRootNodeRef<'ast>) {
349            #[cfg(test)]
350            self.all_nodes.push(node);
351
352            self.pending.push(node);
353
354            if self.pending.len() == IndexedNodes::CHUNK_LEN {
355                self.flush();
356            }
357        }
358
359        fn finish(mut self) -> IndexedNodes {
360            self.flush();
361
362            IndexedNodes {
363                chunks: self.chunks.into_boxed_slice(),
364                words: self.words.into_boxed_slice(),
365            }
366        }
367
368        fn flush(&mut self) {
369            IndexedNodes::extend_from_nodes(&mut self.chunks, &mut self.words, &self.pending);
370            self.pending.clear();
371        }
372    }
373
374    impl IndexedNodes {
375        const ALIGNMENT: usize = std::mem::align_of::<AtomicNodeIndex>();
376        const CHUNK_LEN: usize = 64;
377        const KIND_BITS: u8 = 5;
378        const KIND_MASK: u64 = (1 << Self::KIND_BITS) - 1;
379        const REPORTED_ENTRY_BITS: usize = 32;
380
381        fn extend_from_nodes(
382            chunks: &mut Vec<IndexChunk>,
383            words: &mut Vec<u64>,
384            nodes: &[AnyRootNodeRef<'_>],
385        ) {
386            for node_chunk in nodes.chunks(Self::CHUNK_LEN) {
387                let (base, max, aligned) =
388                    node_chunk
389                        .iter()
390                        .fold((usize::MAX, 0, true), |(base, max, aligned), node| {
391                            let (_, pointer) = node.into_raw_parts();
392                            let address = pointer.as_ptr().expose_provenance();
393                            (
394                                base.min(address),
395                                max.max(address),
396                                aligned && address.is_multiple_of(Self::ALIGNMENT),
397                            )
398                        });
399                let offset_bits = usize::BITS - ((max - base) / Self::ALIGNMENT).leading_zeros();
400                let relative_bits = u8::try_from(offset_bits)
401                    .expect("an address offset cannot require more than u8::MAX bits")
402                    + Self::KIND_BITS;
403                let word_start = u32::try_from(words.len())
404                    .expect("indexed AST bitstream should fit in u32 words");
405
406                if aligned && relative_bits <= 64 {
407                    let entry_count = u8::try_from(node_chunk.len())
408                        .expect("an index chunk contains at most 64 entries");
409                    chunks.push(IndexChunk {
410                        base,
411                        word_start,
412                        entry_bits: relative_bits,
413                        entry_count,
414                        layout: IndexChunkLayout::Relative,
415                    });
416                    for (entry, node) in node_chunk.iter().enumerate() {
417                        let (kind, pointer) = node.into_raw_parts();
418                        let address = pointer.as_ptr().expose_provenance();
419                        let offset = (address - base) / Self::ALIGNMENT;
420                        let offset = u64::try_from(offset)
421                            .expect("relative address offset was checked to fit in 64 bits");
422                        Self::write_bits(
423                            words,
424                            word_start as usize * 64 + entry * usize::from(relative_bits),
425                            (offset << Self::KIND_BITS) | u64::from(kind as u8),
426                            relative_bits,
427                        );
428                    }
429                } else {
430                    // Wide chunks store one address word per entry followed by packed node kinds.
431                    let entry_count = u8::try_from(node_chunk.len())
432                        .expect("an index chunk contains at most 64 entries");
433                    chunks.push(IndexChunk {
434                        base: 0,
435                        word_start,
436                        entry_bits: Self::KIND_BITS,
437                        entry_count,
438                        layout: IndexChunkLayout::Wide,
439                    });
440                    words.extend(node_chunk.iter().map(|node| {
441                        let (_, pointer) = node.into_raw_parts();
442                        u64::try_from(pointer.as_ptr().expose_provenance())
443                            .expect("AST node addresses should fit in a bitstream word")
444                    }));
445                    for (entry, node) in node_chunk.iter().enumerate() {
446                        let (kind, _) = node.into_raw_parts();
447                        Self::write_bits(
448                            words,
449                            (word_start as usize + node_chunk.len()) * 64
450                                + entry * usize::from(Self::KIND_BITS),
451                            u64::from(kind as u8),
452                            Self::KIND_BITS,
453                        );
454                    }
455                }
456            }
457        }
458
459        fn write_bits(words: &mut Vec<u64>, bit: usize, value: u64, bits: u8) {
460            debug_assert!((1..=64).contains(&bits));
461            let word = bit / 64;
462            let shift = bit % 64;
463            let end = bit + usize::from(bits);
464            words.resize(words.len().max(end.div_ceil(64)), 0);
465            words[word] |= value << shift;
466            if end > (word + 1) * 64 {
467                words[word + 1] |= value >> (64 - shift);
468            }
469        }
470
471        fn read_bits(words: &[u64], bit: usize, bits: u8) -> u64 {
472            debug_assert!((1..=64).contains(&bits));
473            let word = bit / 64;
474            let shift = bit % 64;
475            let low = words[word] >> shift;
476            let value = if shift + usize::from(bits) <= 64 {
477                low
478            } else {
479                low | (words[word + 1] << (64 - shift))
480            };
481            if bits == 64 {
482                value
483            } else {
484                value & ((1 << bits) - 1)
485            }
486        }
487
488        #[cfg(test)]
489        fn len(&self) -> usize {
490            self.chunks
491                .iter()
492                .map(|chunk| usize::from(chunk.entry_count))
493                .sum()
494        }
495
496        fn get(&self, index: usize) -> (usize, RootNodeKind) {
497            let chunk_index = index / Self::CHUNK_LEN;
498            let entry_index = index % Self::CHUNK_LEN;
499            let chunk = &self.chunks[chunk_index];
500
501            // A partial chunk's trailing bits are padding, not indexed nodes.
502            assert!(
503                entry_index < usize::from(chunk.entry_count),
504                "index out of bounds: the len is {} but the index is {index}",
505                chunk_index * Self::CHUNK_LEN + usize::from(chunk.entry_count),
506            );
507
508            let words = &self.words[chunk.word_start as usize..];
509
510            match chunk.layout {
511                IndexChunkLayout::Relative => {
512                    let entry = Self::read_bits(
513                        words,
514                        entry_index * usize::from(chunk.entry_bits),
515                        chunk.entry_bits,
516                    );
517                    let offset = (entry >> Self::KIND_BITS) as usize;
518                    let kind = RootNodeKind::from_u8((entry & Self::KIND_MASK) as u8)
519                        .expect("packed node kind should be valid");
520                    (chunk.base + offset * Self::ALIGNMENT, kind)
521                }
522                IndexChunkLayout::Wide => {
523                    let address = usize::try_from(words[entry_index])
524                        .expect("stored AST node address should fit in usize");
525                    let kind_bit = usize::from(chunk.entry_count) * 64
526                        + entry_index * usize::from(Self::KIND_BITS);
527                    let kind =
528                        RootNodeKind::from_u8(
529                            Self::read_bits(words, kind_bit, Self::KIND_BITS) as u8
530                        )
531                        .expect("packed node kind should be valid");
532                    (address, kind)
533                }
534            }
535        }
536    }
537
538    const _: () = assert!(RootNodeKind::ALL.len() <= 1 << IndexedNodes::KIND_BITS);
539
540    /// Ensure the following sub-AST is indexed, using the parent node's index
541    /// as a basis for unambiguous AST node indices.
542    pub fn ensure_indexed(
543        parsed: &Parsed<ModExpression>,
544        parent_node_index: NodeIndex,
545    ) -> Result<(), NodeIndexError> {
546        let parent_index = parent_node_index.as_u32().ok_or(NodeIndexError::NoParent)?;
547        let (index, max_index) = sub_indices(parent_index)?;
548        let mut visitor = Visitor {
549            overflowed: false,
550            nodes: None,
551            index,
552            max_index,
553        };
554
555        AnyNodeRef::from(parsed.syntax()).visit_source_order(&mut visitor);
556
557        if visitor.overflowed {
558            let level = sub_ast_level(parent_index);
559            if level == 0 {
560                return Err(NodeIndexError::ExhaustedSubIndices);
561            } else {
562                return Err(NodeIndexError::ExhaustedSubSubIndices);
563            }
564        }
565
566        Ok(())
567    }
568
569    impl IndexedModule {
570        /// Create a new [`IndexedModule`] from the given AST.
571        pub fn new(parsed: Parsed<ModModule>) -> Arc<Self> {
572            let mut visitor = Visitor {
573                nodes: Some(IndexedNodesBuilder::new()),
574                index: 0,
575                max_index: MAX_REAL_INDEX,
576                overflowed: false,
577            };
578
579            let mut inner = Arc::new(IndexedModule {
580                parsed,
581                index: IndexedNodes::default(),
582            });
583
584            AnyNodeRef::from(inner.parsed.syntax()).visit_source_order(&mut visitor);
585
586            let index = visitor
587                .nodes
588                .expect("top-level AST visitor should collect indexed nodes")
589                .finish();
590            Arc::get_mut(&mut inner)
591                .expect("newly created indexed module should have a unique Arc")
592                .index = index;
593
594            inner
595        }
596
597        /// Returns the node at the given index.
598        pub fn get_by_index<'ast>(&'ast self, index: NodeIndex) -> AnyRootNodeRef<'ast> {
599            let index = index
600                .as_u32()
601                .expect("attempted to access uninitialized `NodeIndex`");
602
603            let index = index as usize;
604            let (address, kind) = self.index.get(index);
605
606            // SAFETY: By the `IndexedNodes` safety invariant, this is the exact exposed address and
607            // root-node kind recorded from the same node after `parsed` reached its stable address.
608            // `self` keeps the AST alive and immutable for the returned reference's lifetime.
609            unsafe {
610                AnyRootNodeRef::from_raw_parts(
611                    kind,
612                    std::ptr::NonNull::with_exposed_provenance(
613                        std::num::NonZeroUsize::new(address)
614                            .expect("recorded AST node address should be non-null"),
615                    ),
616                )
617            }
618        }
619    }
620
621    /// A visitor that indexes nodes in source order.
622    struct Visitor<'ast> {
623        index: u32,
624        max_index: u32,
625        nodes: Option<IndexedNodesBuilder<'ast>>,
626        overflowed: bool,
627    }
628
629    impl<'ast> Visitor<'ast> {
630        fn visit_node<T>(&mut self, node: &'ast T)
631        where
632            T: HasNodeIndex,
633            AnyRootNodeRef<'ast>: From<&'ast T>,
634        {
635            // Only check on write (the maximum is orders of magnitude less than u32::MAX)
636            if self.index > self.max_index {
637                self.overflowed = true;
638            } else {
639                node.node_index().set(NodeIndex::from(self.index));
640            }
641
642            if let Some(nodes) = &mut self.nodes {
643                nodes.push(AnyRootNodeRef::from(node));
644            }
645            self.index += 1;
646        }
647    }
648
649    impl<'a> SourceOrderVisitor<'a> for Visitor<'a> {
650        #[inline]
651        fn visit_stmt(&mut self, stmt: &'a Stmt) {
652            self.visit_node(stmt);
653            walk_stmt(self, stmt);
654        }
655
656        #[inline]
657        fn visit_annotation(&mut self, expr: &'a Expr) {
658            // `walk_annotation` delegates to `visit_expr`, which indexes the expression once.
659            walk_annotation(self, expr);
660        }
661
662        #[inline]
663        fn visit_expr(&mut self, expr: &'a Expr) {
664            self.visit_node(expr);
665            walk_expr(self, expr);
666        }
667
668        #[inline]
669        fn visit_decorator(&mut self, decorator: &'a Decorator) {
670            self.visit_node(decorator);
671            walk_decorator(self, decorator);
672        }
673
674        #[inline]
675        fn visit_comprehension(&mut self, comprehension: &'a Comprehension) {
676            self.visit_node(comprehension);
677            walk_comprehension(self, comprehension);
678        }
679
680        #[inline]
681        fn visit_except_handler(&mut self, except_handler: &'a ExceptHandler) {
682            self.visit_node(except_handler);
683            walk_except_handler(self, except_handler);
684        }
685
686        #[inline]
687        fn visit_arguments(&mut self, arguments: &'a Arguments) {
688            self.visit_node(arguments);
689            walk_arguments(self, arguments);
690        }
691
692        #[inline]
693        fn visit_parameters(&mut self, parameters: &'a Parameters) {
694            self.visit_node(parameters);
695            walk_parameters(self, parameters);
696        }
697
698        #[inline]
699        fn visit_parameter(&mut self, arg: &'a Parameter) {
700            self.visit_node(arg);
701            walk_parameter(self, arg);
702        }
703
704        fn visit_parameter_with_default(
705            &mut self,
706            parameter_with_default: &'a ParameterWithDefault,
707        ) {
708            self.visit_node(parameter_with_default);
709            walk_parameter_with_default(self, parameter_with_default);
710        }
711
712        #[inline]
713        fn visit_keyword(&mut self, keyword: &'a Keyword) {
714            self.visit_node(keyword);
715            walk_keyword(self, keyword);
716        }
717
718        #[inline]
719        fn visit_alias(&mut self, alias: &'a Alias) {
720            self.visit_node(alias);
721            walk_alias(self, alias);
722        }
723
724        #[inline]
725        fn visit_with_item(&mut self, with_item: &'a WithItem) {
726            self.visit_node(with_item);
727            walk_with_item(self, with_item);
728        }
729
730        #[inline]
731        fn visit_type_params(&mut self, type_params: &'a TypeParams) {
732            self.visit_node(type_params);
733            walk_type_params(self, type_params);
734        }
735
736        #[inline]
737        fn visit_type_param(&mut self, type_param: &'a TypeParam) {
738            self.visit_node(type_param);
739            walk_type_param(self, type_param);
740        }
741
742        #[inline]
743        fn visit_match_case(&mut self, match_case: &'a MatchCase) {
744            self.visit_node(match_case);
745            walk_match_case(self, match_case);
746        }
747
748        #[inline]
749        fn visit_pattern(&mut self, pattern: &'a Pattern) {
750            self.visit_node(pattern);
751            walk_pattern(self, pattern);
752        }
753
754        #[inline]
755        fn visit_pattern_arguments(&mut self, pattern_arguments: &'a PatternArguments) {
756            self.visit_node(pattern_arguments);
757            walk_pattern_arguments(self, pattern_arguments);
758        }
759
760        #[inline]
761        fn visit_pattern_keyword(&mut self, pattern_keyword: &'a PatternKeyword) {
762            self.visit_node(pattern_keyword);
763            walk_pattern_keyword(self, pattern_keyword);
764        }
765
766        #[inline]
767        fn visit_elif_else_clause(&mut self, elif_else_clause: &'a ElifElseClause) {
768            self.visit_node(elif_else_clause);
769            walk_elif_else_clause(self, elif_else_clause);
770        }
771
772        #[inline]
773        fn visit_f_string(&mut self, f_string: &'a FString) {
774            self.visit_node(f_string);
775            walk_f_string(self, f_string);
776        }
777
778        #[inline]
779        fn visit_interpolated_string_element(
780            &mut self,
781            interpolated_string_element: &'a InterpolatedStringElement,
782        ) {
783            self.visit_node(interpolated_string_element);
784            walk_interpolated_string_element(self, interpolated_string_element);
785        }
786
787        #[inline]
788        fn visit_t_string(&mut self, t_string: &'a TString) {
789            self.visit_node(t_string);
790            walk_t_string(self, t_string);
791        }
792
793        #[inline]
794        fn visit_string_literal(&mut self, string_literal: &'a StringLiteral) {
795            self.visit_node(string_literal);
796            walk_string_literal(self, string_literal);
797        }
798
799        #[inline]
800        fn visit_bytes_literal(&mut self, bytes_literal: &'a BytesLiteral) {
801            self.visit_node(bytes_literal);
802            walk_bytes_literal(self, bytes_literal);
803        }
804
805        #[inline]
806        fn visit_identifier(&mut self, identifier: &'a Identifier) {
807            self.visit_node(identifier);
808            walk_identifier(self, identifier);
809        }
810    }
811
812    #[cfg(test)]
813    mod tests {
814        use super::*;
815
816        #[test]
817        #[should_panic(expected = "index out of bounds: the len is 1 but the index is 1")]
818        fn indexed_nodes_relative_tail_bounds() {
819            let index = IndexedNodes {
820                chunks: vec![IndexChunk {
821                    base: 0x1000,
822                    word_start: 0,
823                    entry_bits: IndexedNodes::KIND_BITS,
824                    entry_count: 1,
825                    layout: IndexChunkLayout::Relative,
826                }]
827                .into_boxed_slice(),
828                words: vec![u64::from(RootNodeKind::Stmt as u8)].into_boxed_slice(),
829            };
830
831            assert_eq!(index.get(0), (0x1000, RootNodeKind::Stmt));
832            index.get(1);
833        }
834
835        #[test]
836        #[should_panic(expected = "index out of bounds: the len is 1 but the index is 1")]
837        fn indexed_nodes_wide_tail_bounds() {
838            let index = IndexedNodes {
839                chunks: vec![IndexChunk {
840                    base: 0,
841                    word_start: 0,
842                    entry_bits: IndexedNodes::KIND_BITS,
843                    entry_count: 1,
844                    layout: IndexChunkLayout::Wide,
845                }]
846                .into_boxed_slice(),
847                words: vec![0x1000, u64::from(RootNodeKind::Stmt as u8)].into_boxed_slice(),
848            };
849
850            assert_eq!(index.get(0), (0x1000, RootNodeKind::Stmt));
851            index.get(1);
852        }
853
854        #[test]
855        #[should_panic(expected = "index out of bounds: the len is 65 but the index is 65")]
856        fn indexed_nodes_chunk_boundary() {
857            let parsed = ruff_python_parser::parse_module(&"pass\n".repeat(65)).unwrap();
858            let indexed = IndexedModule::new(parsed);
859
860            assert_eq!(indexed.index.len(), 65);
861            assert_eq!(indexed.index.get(63).1, RootNodeKind::Stmt);
862            assert_eq!(indexed.index.get(64).1, RootNodeKind::Stmt);
863            indexed.index.get(65);
864        }
865
866        #[test]
867        fn indexed_nodes_round_trip() {
868            let parsed = ruff_python_parser::parse_module(
869                r#"
870import os as imported_os
871
872@decorator
873class C[T](Base, metaclass=Meta):
874    def method(self, value: int = 1, *args, keyword=2, **kwargs):
875        try:
876            with context() as items:
877                return [item for item in items if item]
878        except Error as error:
879            match error:
880                case Error(code=code):
881                    if code:
882                        return f"{code!r}"
883                    elif code is None:
884                        return t"{code}"
885                    else:
886                        return "string"
887                case _:
888                    return b"bytes"
889"#,
890            )
891            .expect("test source should parse");
892            let indexed = IndexedModule::new(parsed);
893            let mut visitor = Visitor {
894                nodes: Some(IndexedNodesBuilder::new()),
895                index: 0,
896                max_index: MAX_REAL_INDEX,
897                overflowed: false,
898            };
899            AnyNodeRef::from(indexed.parsed.syntax()).visit_source_order(&mut visitor);
900            let nodes = visitor
901                .nodes
902                .expect("test visitor should collect indexed nodes")
903                .all_nodes;
904
905            assert_eq!(indexed.index.len(), nodes.len());
906            let mut seen_kinds = [false; 1 << IndexedNodes::KIND_BITS];
907
908            for (raw_index, expected_node) in nodes.into_iter().enumerate() {
909                let (kind, pointer) = expected_node.into_raw_parts();
910                let address = pointer.as_ptr().expose_provenance();
911                let index = NodeIndex::from(
912                    u32::try_from(raw_index).expect("node index should fit in u32"),
913                );
914                seen_kinds[usize::from(kind as u8)] = true;
915                assert_eq!(indexed.index.get(raw_index), (address, kind));
916
917                let node = indexed.get_by_index(index);
918                let (actual_kind, actual_pointer) = node.into_raw_parts();
919                assert_eq!(actual_kind, kind);
920                assert_eq!(actual_pointer.as_ptr().expose_provenance(), address);
921                assert_eq!(node.node_index().load(), index);
922            }
923            for kind in RootNodeKind::ALL {
924                let is_indexed = !matches!(
925                    kind,
926                    RootNodeKind::Mod | RootNodeKind::InterpolatedStringFormatSpec
927                );
928                assert_eq!(seen_kinds[usize::from(*kind as u8)], is_indexed);
929            }
930        }
931    }
932}
933
934#[cfg(test)]
935mod tests {
936    use crate::Db;
937    use crate::PythonFile;
938    use crate::files::{system_path_to_file, vendored_path_to_file};
939    use crate::parsed::parsed_module;
940    use crate::system::{
941        DbWithTestSystem, DbWithWritableSystem as _, SystemPath, SystemVirtualPath,
942    };
943    use crate::tests::TestDb;
944    use crate::vendored::{VendoredFileSystemBuilder, VendoredPath};
945    use ruff_python_ast::PythonVersion;
946    use zip::CompressionMethod;
947
948    #[test]
949    fn python_file() -> crate::system::Result<()> {
950        let mut db = TestDb::new();
951        let path = "test.py";
952
953        db.write_file(path, "x = 10")?;
954
955        let file = system_path_to_file(&db, path).unwrap();
956
957        let file = PythonFile::new(&db, file, PythonVersion::latest_ty());
958        let parsed = parsed_module(&db, file).load(&db);
959
960        assert!(parsed.has_valid_syntax());
961
962        Ok(())
963    }
964
965    #[test]
966    fn python_ipynb_file() -> crate::system::Result<()> {
967        let mut db = TestDb::new();
968        let path = SystemPath::new("test.ipynb");
969
970        db.write_file(path, "%timeit a = b")?;
971
972        let file = system_path_to_file(&db, path).unwrap();
973
974        let file = PythonFile::new(&db, file, PythonVersion::latest_ty());
975        let parsed = parsed_module(&db, file).load(&db);
976
977        assert!(parsed.has_valid_syntax());
978
979        Ok(())
980    }
981
982    #[test]
983    fn virtual_python_file() -> crate::system::Result<()> {
984        let mut db = TestDb::new();
985        let path = SystemVirtualPath::new("untitled:Untitled-1");
986
987        db.write_virtual_file(path, "x = 10");
988
989        let virtual_file = db.files().virtual_file(&db, path);
990
991        let file = PythonFile::new(&db, virtual_file.file(), PythonVersion::latest_ty());
992        let parsed = parsed_module(&db, file).load(&db);
993
994        assert!(parsed.has_valid_syntax());
995
996        Ok(())
997    }
998
999    #[test]
1000    fn virtual_ipynb_file() -> crate::system::Result<()> {
1001        let mut db = TestDb::new();
1002        let path = SystemVirtualPath::new("untitled:Untitled-1.ipynb");
1003
1004        db.write_virtual_file(path, "%timeit a = b");
1005
1006        let virtual_file = db.files().virtual_file(&db, path);
1007
1008        let file = PythonFile::new(&db, virtual_file.file(), PythonVersion::latest_ty());
1009        let parsed = parsed_module(&db, file).load(&db);
1010
1011        assert!(parsed.has_valid_syntax());
1012
1013        Ok(())
1014    }
1015
1016    #[test]
1017    fn vendored_file() {
1018        let mut db = TestDb::new();
1019
1020        let mut vendored_builder = VendoredFileSystemBuilder::new(CompressionMethod::Stored);
1021        vendored_builder
1022            .add_file(
1023                "path.pyi",
1024                r#"
1025import sys
1026
1027if sys.platform == "win32":
1028    from ntpath import *
1029    from ntpath import __all__ as __all__
1030else:
1031    from posixpath import *
1032    from posixpath import __all__ as __all__"#,
1033            )
1034            .unwrap();
1035        let vendored = vendored_builder.finish().unwrap();
1036        db.with_vendored(vendored);
1037
1038        let file = vendored_path_to_file(&db, VendoredPath::new("path.pyi")).unwrap();
1039
1040        let file = PythonFile::new(&db, file, PythonVersion::latest_ty());
1041        let parsed = parsed_module(&db, file).load(&db);
1042
1043        assert!(parsed.has_valid_syntax());
1044    }
1045
1046    #[test]
1047    fn same_file_at_different_python_versions() -> crate::system::Result<()> {
1048        let mut db = TestDb::new();
1049        db.write_file("test.py", "type Alias = int")?;
1050        let file = system_path_to_file(&db, "test.py").unwrap();
1051
1052        let py311 = PythonFile::new(&db, file, PythonVersion::PY311);
1053        let py312 = PythonFile::new(&db, file, PythonVersion::PY312);
1054        let parsed_py311 = parsed_module(&db, py311);
1055        let parsed_py312 = parsed_module(&db, py312);
1056
1057        for _ in 0..2 {
1058            assert!(
1059                !parsed_py311
1060                    .load(&db)
1061                    .unsupported_syntax_errors()
1062                    .is_empty()
1063            );
1064            assert!(
1065                parsed_py312
1066                    .load(&db)
1067                    .unsupported_syntax_errors()
1068                    .is_empty()
1069            );
1070
1071            parsed_py311.clear();
1072            parsed_py312.clear();
1073        }
1074
1075        Ok(())
1076    }
1077}