meta_ast/extractor/

mod.rs

//! Parallel file extraction orchestration.
//!
//! Uses rayon `par_iter` to read, parse, and extract symbols/imports/
//! references across files concurrently. Each file is processed
//! independently; errors are accumulated as diagnostics per file.
//!
//! Set `ExtractOptions::skip_imports_and_refs` to `true` when only
//! symbol listing is needed (e.g. inspect mode); skips the import and
//! reference query passes, roughly halving per-file extraction time.

use std::path::{Path, PathBuf};

use rayon::prelude::*;

use crate::error::{Diagnostic, Severity};
use crate::language::LangId;
use crate::model::{IdGenerator, Symbol, SymbolId};
use crate::parser;

pub use crate::model::FileExtraction;

/// Controls what the extraction pass produces.
#[derive(Debug, Clone, Copy, Default)]
pub struct ExtractOptions {
    /// Skip import and reference extraction entirely; only extract symbols
    /// and AST node counts. Halves per-file extraction cost for pure
    /// symbol-inspection workflows.
    pub skip_imports_and_refs: bool,
}

pub struct ExtractionResult {
    pub files: Vec<FileExtraction>,
}

/// ID allocation state shared by disk and in-memory extraction.
#[derive(Debug, Default)]
pub struct ExtractionIdGenerators {
    symbols: IdGenerator<SymbolId>,
    #[cfg(feature = "dataflow")]
    data_nodes: IdGenerator<crate::model::DataNodeId>,
}

impl ExtractionIdGenerators {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_symbol_start(symbol_start: u32) -> Self {
        Self {
            symbols: IdGenerator::with_start(symbol_start),
            #[cfg(feature = "dataflow")]
            data_nodes: IdGenerator::new(),
        }
    }

    #[cfg(feature = "dataflow")]
    pub fn with_starts(symbol_start: u32, data_node_start: u32) -> Self {
        Self {
            symbols: IdGenerator::with_start(symbol_start),
            data_nodes: IdGenerator::with_start(data_node_start),
        }
    }

    pub fn symbols(&self) -> &IdGenerator<SymbolId> {
        &self.symbols
    }

    #[cfg(feature = "dataflow")]
    pub fn data_nodes(&self) -> &IdGenerator<crate::model::DataNodeId> {
        &self.data_nodes
    }
}

/// Source text supplied by an editor buffer.
#[derive(Debug, Clone, Copy)]
pub struct InMemorySource<'a> {
    pub uri: &'a str,
    pub text: &'a str,
    pub version: i32,
    pub language: LangId,
}

/// Extraction result tied to the editor document version that produced it.
#[derive(Debug, Clone)]
pub struct VersionedExtraction {
    pub uri: String,
    pub version: i32,
    pub file: FileExtraction,
}

pub fn extract(files: &[(std::path::PathBuf, LangId)]) -> ExtractionResult {
    extract_with_options(files, &ExtractOptions::default())
}

pub fn extract_with_options(
    files: &[(std::path::PathBuf, LangId)],
    opts: &ExtractOptions,
) -> ExtractionResult {
    let id_generators = ExtractionIdGenerators::new();
    extract_with_id_gen(files, opts, &id_generators)
}

pub fn extract_with_id_gen(
    files: &[(PathBuf, LangId)],
    opts: &ExtractOptions,
    id_generators: &ExtractionIdGenerators,
) -> ExtractionResult {
    let mut file_extractions: Vec<_> = files
        .par_iter()
        .map(|(path, lang)| extract_single_file(path, lang, id_generators, opts))
        .collect();

    file_extractions.sort_by(|a, b| a.path.cmp(&b.path));

    ExtractionResult {
        files: file_extractions,
    }
}

fn extract_single_file(
    path: &Path,
    lang: &LangId,
    id_generators: &ExtractionIdGenerators,
    opts: &ExtractOptions,
) -> FileExtraction {
    let source = match std::fs::read(path) {
        Ok(source) => source,
        Err(error) => {
            return failed_extraction(path, *lang, format!("failed to read file: {error}"));
        }
    };

    extract_source(path, *lang, &source, id_generators, opts)
}

/// Extract an open editor buffer without reading its backing file.
pub fn extract_text_with_id_gen(
    source: InMemorySource<'_>,
    opts: &ExtractOptions,
    id_generators: &ExtractionIdGenerators,
) -> Result<VersionedExtraction, crate::Error> {
    let parsed_uri =
        url::Url::parse(source.uri).map_err(|error| crate::Error::InvalidSourceUri {
            uri: source.uri.to_string(),
            message: error.to_string(),
        })?;
    let path = parsed_uri
        .to_file_path()
        .map_err(|()| crate::Error::InvalidSourceUri {
            uri: source.uri.to_string(),
            message: "URI must use the file scheme and contain an absolute path".to_string(),
        })?;
    let path = crate::input::simplified_path(&path).to_path_buf();
    let file = extract_source(
        &path,
        source.language,
        source.text.as_bytes(),
        id_generators,
        opts,
    );

    Ok(VersionedExtraction {
        uri: source.uri.to_string(),
        version: source.version,
        file,
    })
}

fn extract_source(
    path: &Path,
    lang: LangId,
    source: &[u8],
    id_generators: &ExtractionIdGenerators,
    opts: &ExtractOptions,
) -> FileExtraction {
    let tree = match crate::parser::parse_tree(lang, source) {
        Ok(t) => t,
        Err(e) => {
            return failed_extraction(path, lang, e.to_string());
        }
    };

    let metrics = parser::tree_metrics(&tree, source);
    let mut diags = Vec::new();

    if metrics.error_ratio > 0.5 {
        diags.push(Diagnostic {
            path: path.to_path_buf(),
            severity: Severity::Warning,
            message: format!(
                "file has {:.0}% parse errors, results may be incomplete",
                metrics.error_ratio * 100.0
            ),
            source_range: None,
        });
    }

    let raw_symbols = crate::language::extract_symbols_for(lang, &tree, source);
    let symbols = raw_symbols
        .into_iter()
        .map(|raw| Symbol {
            id: id_generators.symbols.next(),
            name: raw.name.into_owned(),
            kind: raw.kind,
            language: lang,
            file_path: path.to_path_buf(),
            source_range: raw.source_range,
            visibility: raw.visibility,
            signature: raw.signature.map(|s| s.into_owned()),
            docstring: raw.docstring.map(|s| s.into_owned()),
            is_async: raw.is_async,
        })
        .collect();

    let (imports, references) = if opts.skip_imports_and_refs {
        (Vec::new(), Vec::new())
    } else {
        crate::language::extract_imports_and_references_for(lang, &tree, source, path)
    };

    #[cfg(feature = "metacall-deploy")]
    let call_sites = crate::deploy::scanner::scan_file(lang, &tree, source, path);

    #[cfg(feature = "dataflow")]
    let (data_nodes, flow_edges) =
        crate::language::dataflow::extract_dataflow(lang, &tree, source, &id_generators.data_nodes);

    let mut out = FileExtraction::empty(path.to_path_buf(), lang);
    out.symbols = symbols;
    out.imports = imports;
    out.references = references;
    out.diagnostics = diags;
    out.ast_node_count = metrics.node_count;
    #[cfg(feature = "metacall-deploy")]
    {
        out.call_sites = call_sites;
    }
    #[cfg(feature = "dataflow")]
    {
        out.data_nodes = data_nodes;
        out.flow_edges = flow_edges;
    }
    out
}

fn failed_extraction(path: &Path, lang: LangId, message: String) -> FileExtraction {
    FileExtraction::failed(path.to_path_buf(), lang, message)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    fn test_dir() -> PathBuf {
        let dir = std::env::temp_dir().join("meta_ast_test_extractor");
        let _ = std::fs::create_dir_all(&dir);
        dir
    }

    fn write_temp(name: &str, content: &[u8]) -> PathBuf {
        let path = test_dir().join(name);
        std::fs::write(&path, content).unwrap();
        path
    }

    #[test]
    fn extract_single_python_file() {
        let path = write_temp("single.py", b"def hello(): pass\n");
        let result = extract(&[(path.clone(), LangId::Python)]);
        assert_eq!(result.files.len(), 1);
        assert!(!result.files[0].symbols.is_empty());
        assert!(result.files[0].diagnostics.is_empty());
        let names: Vec<&str> = result.files[0]
            .symbols
            .iter()
            .map(|s| s.name.as_str())
            .collect();
        assert!(names.contains(&"hello"));
    }

    #[test]
    fn extract_multiple_files_parallel() {
        let p1 = write_temp("file_a.py", b"def alpha(): pass\n");
        let p2 = write_temp("file_b.py", b"def beta(): pass\ndef gamma(): pass\n");
        let p3 = write_temp("file_c.py", b"class Delta: pass\n");

        let files = vec![
            (p1.clone(), LangId::Python),
            (p2.clone(), LangId::Python),
            (p3.clone(), LangId::Python),
        ];
        let result = extract(&files);
        let all_names: Vec<&str> = result
            .files
            .iter()
            .flat_map(|f| f.symbols.iter().map(|s| s.name.as_str()))
            .collect();
        assert!(all_names.contains(&"alpha"), "missing alpha: {all_names:?}");
        assert!(all_names.contains(&"beta"), "missing beta: {all_names:?}");
        assert!(all_names.contains(&"gamma"), "missing gamma: {all_names:?}");
        assert!(all_names.contains(&"Delta"), "missing Delta: {all_names:?}");
    }

    #[test]
    fn accumulate_diagnostics_on_malformed() {
        let path = test_dir().join("nonexistent_broken.py");
        let _ = std::fs::remove_file(&path);
        let result = extract(&[(path, LangId::Python)]);
        assert!(!result.files[0].diagnostics.is_empty());
    }

    #[test]
    fn partial_extraction_on_errors() {
        let valid = write_temp("valid_partial.py", b"def works(): pass\n");
        let broken = write_temp(
            "broken_partial.py",
            b"def broken(\n   # missing close paren and colon\n",
        );
        let result = extract(&[(valid.clone(), LangId::Python), (broken, LangId::Python)]);
        let names: Vec<&str> = result
            .files
            .iter()
            .flat_map(|f| f.symbols.iter().map(|s| s.name.as_str()))
            .collect();
        assert!(
            names.contains(&"works"),
            "valid file symbols should be present: {names:?}"
        );
    }

    #[test]
    fn output_deterministic() {
        let path = write_temp("deterministic.py", b"def foo(): pass\ndef bar(): pass\n");
        let files = vec![(path.clone(), LangId::Python)];

        let r1 = extract(&files);
        let r2 = extract(&files);

        let names1: Vec<String> = r1
            .files
            .iter()
            .flat_map(|f| f.symbols.iter().map(|s| s.name.clone()))
            .collect();
        let names2: Vec<String> = r2
            .files
            .iter()
            .flat_map(|f| f.symbols.iter().map(|s| s.name.clone()))
            .collect();
        assert_eq!(names1, names2);
    }

    #[test]
    fn in_memory_extraction_uses_unsaved_text_and_preserves_version() {
        let path = test_dir().join("buffer.py");
        let _ = std::fs::remove_file(&path);
        let uri = url::Url::from_file_path(&path).unwrap().to_string();
        let id_generators = ExtractionIdGenerators::with_symbol_start(40);

        let result = extract_text_with_id_gen(
            InMemorySource {
                uri: &uri,
                text: "def unsaved(): pass\n",
                version: 7,
                language: LangId::Python,
            },
            &ExtractOptions::default(),
            &id_generators,
        )
        .unwrap();

        assert_eq!(result.version, 7);
        assert_eq!(result.uri, uri);
        assert_eq!(result.file.path, path);
        assert_eq!(result.file.symbols[0].name, "unsaved");
        assert_eq!(result.file.symbols[0].id, SymbolId::new(40).unwrap());
    }

    #[test]
    fn in_memory_extraction_rejects_non_file_uri() {
        let id_generators = ExtractionIdGenerators::new();
        let error = extract_text_with_id_gen(
            InMemorySource {
                uri: "untitled:buffer.py",
                text: "def value(): pass\n",
                version: 1,
                language: LangId::Python,
            },
            &ExtractOptions::default(),
            &id_generators,
        )
        .unwrap_err();

        assert!(matches!(error, crate::Error::InvalidSourceUri { .. }));
    }

    #[test]
    fn extraction_id_generators_accessors() {
        let id_generators = ExtractionIdGenerators::with_symbol_start(100);
        assert_eq!(id_generators.symbols().next(), SymbolId::new(100).unwrap());
        #[cfg(feature = "dataflow")]
        {
            let dual = ExtractionIdGenerators::with_starts(200, 300);
            assert_eq!(dual.symbols().next(), SymbolId::new(200).unwrap());
            assert_eq!(
                dual.data_nodes().next(),
                crate::model::DataNodeId::new(300).unwrap()
            );
        }
    }

    #[test]
    fn symbols_assigned_ids() {
        let path = write_temp("ids.py", b"def a(): pass\ndef b(): pass\ndef c(): pass\n");
        let result = extract(&[(path, LangId::Python)]);
        let ids: Vec<u32> = result.files[0]
            .symbols
            .iter()
            .map(|s| s.id.to_raw())
            .collect();
        let mut sorted_ids = ids.clone();
        sorted_ids.sort();
        assert_eq!(ids, sorted_ids, "IDs should be sequential");

        for window in sorted_ids.windows(2) {
            assert_eq!(window[1] - window[0], 1, "IDs should be consecutive");
        }

        let unique: std::collections::HashSet<u32> = ids.iter().copied().collect();
        assert_eq!(
            unique.len(),
            result.files[0].symbols.len(),
            "all IDs must be unique"
        );
    }
}