Skip to main content

vyre_libs/parsing/python/
source_cache.rs

1//! Content-keyed Python 3.12 pipeline Program cache.
2
3use std::sync::Arc;
4
5use crate::parsing::python::lex::python312_lexer;
6use crate::parsing::python::parse::calls::python312_extract_calls;
7use crate::parsing::python::parse::structure::{
8    python312_extract_imports, python312_extract_structure,
9};
10use crate::parsing::source_cache::{source_len_u32_nonzero, ParsedSourceLru};
11use vyre::ir::Program;
12
13/// Cached Python pipeline stage bundle for one source shape.
14#[derive(Debug, Clone)]
15pub struct Python312PipelinePrograms {
16    /// Python lexer Program.
17    pub lex: Program,
18    /// Definition / class / span structural extractor.
19    pub structure: Program,
20    /// Import extractor.
21    pub imports: Program,
22    /// Call-site extractor.
23    pub calls: Program,
24}
25
26/// Bounded content-keyed Python pipeline cache.
27pub type Python312PipelineCache = ParsedSourceLru<Python312PipelinePrograms>;
28
29/// Build or fetch the Python 3.12 pipeline bundle for `source`.
30pub fn get_or_build_python312_pipeline(
31    cache: &Python312PipelineCache,
32    source: &[u8],
33    extra: &[u8],
34) -> Arc<Python312PipelinePrograms> {
35    cache.get_or_parse(source, extra, |bytes| {
36        let len = source_len_u32_nonzero(bytes);
37        Python312PipelinePrograms {
38            lex: python312_lexer(
39                "haystack",
40                "tok_types",
41                "tok_starts",
42                "tok_lens",
43                "tok_counts",
44                len,
45            ),
46            structure: python312_extract_structure(
47                "tok_types",
48                "tok_starts",
49                "tok_lens",
50                "structure_records",
51                "structure_counts",
52                len,
53            ),
54            imports: python312_extract_imports(
55                "tok_types",
56                "tok_starts",
57                "tok_lens",
58                "import_records",
59                "import_counts",
60                len,
61            ),
62            calls: python312_extract_calls(
63                "tok_types",
64                "tok_starts",
65                "tok_lens",
66                "call_records",
67                "call_counts",
68                "kwarg_records",
69                "kwarg_counts",
70                len,
71            ),
72        }
73    })
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn python_pipeline_cache_reuses_same_source_bundle() {
82        let cache = Python312PipelineCache::with_capacity(4);
83        let a = get_or_build_python312_pipeline(&cache, b"def f():\n    return 1\n", b"py312");
84        let b = get_or_build_python312_pipeline(&cache, b"def f():\n    return 1\n", b"py312");
85        assert!(Arc::ptr_eq(&a, &b));
86        assert_eq!(cache.len(), 1);
87    }
88
89    #[test]
90    fn python_pipeline_cache_separates_options() {
91        let cache = Python312PipelineCache::with_capacity(4);
92        let a = get_or_build_python312_pipeline(&cache, b"x = 1\n", b"mode=a");
93        let b = get_or_build_python312_pipeline(&cache, b"x = 1\n", b"mode=b");
94        assert!(!Arc::ptr_eq(&a, &b));
95        assert_eq!(cache.len(), 2);
96    }
97}