meta_ast/output/shard/
file.rs

1//! Shard file record representation and JSONL read/write operations.
2
3use std::io::{BufRead, Write};
4use std::path::{Path, PathBuf};
5
6use petgraph::visit::EdgeRef;
7use serde::{Deserialize, Serialize};
8
9use crate::error::Diagnostic;
10use crate::graph::{CodeGraph, NodeData};
11use crate::language::LangId;
12use crate::model::{
13    FileExtraction, IdGenerator, SourceRange, Symbol, SymbolId, SymbolKind, UnresolvedImport,
14    UnresolvedReference, Visibility,
15};
16use crate::output::shard::edge::{ShardEdge, validate_edge};
17use crate::output::shard::error::ShardError;
18use crate::output::shard::name::{StableNameIndex, node_owner_path, normalized_path};
19
20/// Shard payload version.
21///
22/// Version 5 stores the import specifier without its surrounding quotes and
23/// drops a specifier that is not valid UTF-8, so version 4 records no longer
24/// parse. Regenerate the index after an upgrade.
25pub const SHARD_SCHEMA_VERSION: u32 = 5;
26
27/// A per-file shard record stored in `.meta-ast/shards/<n>.jsonl`.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct ShardFile {
30    pub schema_version: u32,
31    pub path: PathBuf,
32    pub language: LangId,
33    pub symbols: Vec<ShardSymbol>,
34    pub imports: Vec<UnresolvedImport>,
35    pub references: Vec<UnresolvedReference>,
36    pub diagnostics: Vec<Diagnostic>,
37    pub ast_node_count: usize,
38    /// MetaCall load and client-call sites, persisted for cold-start fidelity.
39    #[cfg(feature = "metacall-deploy")]
40    #[serde(default, skip_serializing_if = "Vec::is_empty")]
41    pub call_sites: Vec<crate::deploy::scanner::CallSite>,
42    pub edges: Vec<ShardEdge>,
43}
44
45/// Symbol representation in a shard file without runtime-local identifiers.
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ShardSymbol {
48    pub name: String,
49    pub kind: SymbolKind,
50    pub source_range: SourceRange,
51    /// Range of the symbol name, absent in records written before the field
52    /// existed and when the language pack reported none. Consumers fall back to
53    /// the declaration range themselves.
54    #[serde(default, skip_serializing_if = "Option::is_none")]
55    pub name_range: Option<SourceRange>,
56    pub visibility: Option<Visibility>,
57    pub signature: Option<String>,
58    pub docstring: Option<String>,
59    pub is_async: bool,
60}
61
62/// Extraction result reconstituted from a shard record.
63#[derive(Debug)]
64pub struct LoadedShard {
65    pub file: FileExtraction,
66    pub edges: Vec<ShardEdge>,
67}
68
69impl ShardFile {
70    /// Convert an in-memory `FileExtraction` and graph into a shard record.
71    ///
72    /// An edge is stored by the shard of its source owner only, so a
73    /// cross-file edge exists once in the index instead of once per endpoint.
74    pub fn from_extraction(
75        extraction: &FileExtraction,
76        graph: &CodeGraph,
77    ) -> Result<Self, ShardError> {
78        normalized_path(&extraction.path)?;
79        for diagnostic in &extraction.diagnostics {
80            normalized_path(&diagnostic.path)?;
81        }
82        let names = StableNameIndex::new(graph)?;
83        let symbols = extraction.symbols.iter().map(ShardSymbol::from).collect();
84        #[cfg(feature = "dataflow")]
85        let dropped_payload = dataflow_drop_diagnostic(extraction);
86        #[cfg(not(feature = "dataflow"))]
87        let dropped_payload: Option<Diagnostic> = None;
88        let diagnostics: Vec<Diagnostic> = extraction
89            .diagnostics
90            .iter()
91            .cloned()
92            .chain(dropped_payload)
93            .collect();
94        let mut edges = graph
95            .graph()
96            .edge_references()
97            .filter(|edge| {
98                !matches!(graph.graph()[edge.source()], NodeData::Data(_))
99                    && !matches!(graph.graph()[edge.target()], NodeData::Data(_))
100                    && node_owner_path(graph, edge.source())
101                        .or_else(|| node_owner_path(graph, edge.target()))
102                        == Some(extraction.path.as_path())
103            })
104            .map(|edge| {
105                let source_name = names
106                    .name_of(edge.source())
107                    .ok_or(ShardError::MissingNodeOwner {
108                        node_index: edge.source().index(),
109                    })?
110                    .to_string();
111                let target_name = names
112                    .name_of(edge.target())
113                    .ok_or(ShardError::MissingNodeOwner {
114                        node_index: edge.target().index(),
115                    })?
116                    .to_string();
117                Ok(ShardEdge {
118                    source_name,
119                    target_name,
120                    kind: edge.weight().kind.into(),
121                    confidence: edge.weight().confidence,
122                    flow_kind: edge.weight().flow_kind.map(Into::into),
123                })
124            })
125            .collect::<Result<Vec<_>, ShardError>>()?;
126        edges.sort_by(|left, right| {
127            (
128                &left.source_name,
129                &left.target_name,
130                left.kind,
131                left.flow_kind,
132            )
133                .cmp(&(
134                    &right.source_name,
135                    &right.target_name,
136                    right.kind,
137                    right.flow_kind,
138                ))
139        });
140        for (edge_index, edge) in edges.iter().enumerate() {
141            validate_edge(edge, 1, edge_index)?;
142        }
143
144        Ok(Self {
145            schema_version: SHARD_SCHEMA_VERSION,
146            path: extraction.path.clone(),
147            language: extraction.lang,
148            symbols,
149            imports: extraction.imports.clone(),
150            references: extraction.references.clone(),
151            diagnostics,
152            ast_node_count: extraction.ast_node_count,
153            #[cfg(feature = "metacall-deploy")]
154            call_sites: extraction.call_sites.clone(),
155            edges,
156        })
157    }
158
159    /// Reconstitute a `FileExtraction` and assign fresh IDs using the provided generator.
160    pub fn load(self, id_gen: &IdGenerator<SymbolId>) -> Result<LoadedShard, ShardError> {
161        if self.schema_version != SHARD_SCHEMA_VERSION {
162            return Err(ShardError::SchemaVersion {
163                line: 1,
164                found: self.schema_version,
165                expected: SHARD_SCHEMA_VERSION,
166            });
167        }
168
169        let symbols = self
170            .symbols
171            .into_iter()
172            .map(|symbol| symbol.into_symbol(&self.path, self.language, id_gen.next()))
173            .collect();
174        let mut file = FileExtraction::empty(self.path, self.language);
175        file.symbols = symbols;
176        file.imports = self.imports;
177        file.references = self.references;
178        file.diagnostics = self.diagnostics;
179        file.ast_node_count = self.ast_node_count;
180        #[cfg(feature = "metacall-deploy")]
181        {
182            file.call_sites = self.call_sites;
183        }
184
185        Ok(LoadedShard {
186            file,
187            edges: self.edges,
188        })
189    }
190}
191
192/// Warning for a dataflow payload that the shard schema cannot store.
193///
194/// Schema version 3 keeps symbols, imports, references and edges. A record
195/// with dataflow nodes or flow edges loses them on write, so the record says
196/// so instead of dropping them without a word.
197#[cfg(feature = "dataflow")]
198fn dataflow_drop_diagnostic(extraction: &FileExtraction) -> Option<Diagnostic> {
199    if extraction.data_nodes.is_empty() && extraction.flow_edges.is_empty() {
200        return None;
201    }
202    Some(Diagnostic {
203        path: extraction.path.clone(),
204        severity: crate::error::Severity::Warning,
205        message: format!(
206            "dataflow payload not persisted: shard schema version {SHARD_SCHEMA_VERSION} stores symbols and edges only"
207        ),
208        source_range: None,
209    })
210}
211
212impl From<&Symbol> for ShardSymbol {
213    fn from(symbol: &Symbol) -> Self {
214        Self {
215            name: symbol.name.clone(),
216            kind: symbol.kind,
217            source_range: symbol.source_range.clone(),
218            name_range: symbol.name_range.clone(),
219            visibility: symbol.visibility,
220            signature: symbol.signature.clone(),
221            docstring: symbol.docstring.clone(),
222            is_async: symbol.is_async,
223        }
224    }
225}
226
227impl ShardSymbol {
228    fn into_symbol(self, file_path: &Path, language: LangId, id: SymbolId) -> Symbol {
229        Symbol {
230            id,
231            name: self.name,
232            kind: self.kind,
233            language,
234            file_path: file_path.to_path_buf(),
235            source_range: self.source_range,
236            name_range: self.name_range,
237            visibility: self.visibility,
238            signature: self.signature,
239            docstring: self.docstring,
240            is_async: self.is_async,
241        }
242    }
243}
244
245/// Write shard records to a writer in JSONL format.
246///
247/// Use `std::io::BufWriter` when writing to a file to prevent frequent write system calls.
248pub fn write_shard<W: Write>(mut writer: W, files: &[ShardFile]) -> Result<(), ShardError> {
249    for (line_index, file) in files.iter().enumerate() {
250        validate_file(file, line_index + 1)?;
251        serde_json::to_writer(&mut writer, file).map_err(ShardError::Encode)?;
252        writer.write_all(b"\n")?;
253    }
254    writer.flush()?;
255    Ok(())
256}
257
258pub fn read_shard<R: BufRead>(reader: R) -> Result<Vec<ShardFile>, ShardError> {
259    let mut files = Vec::new();
260    for (line_index, line) in reader.lines().enumerate() {
261        let line_number = line_index + 1;
262        let line = line?;
263        if line.trim().is_empty() {
264            continue;
265        }
266        let file: ShardFile = serde_json::from_str(&line).map_err(|source| ShardError::Decode {
267            line: line_number,
268            source,
269        })?;
270        validate_file(&file, line_number)?;
271        files.push(file);
272    }
273    Ok(files)
274}
275
276pub(crate) fn validate_file(file: &ShardFile, line: usize) -> Result<(), ShardError> {
277    if file.schema_version != SHARD_SCHEMA_VERSION {
278        return Err(ShardError::SchemaVersion {
279            line,
280            found: file.schema_version,
281            expected: SHARD_SCHEMA_VERSION,
282        });
283    }
284    normalized_path(&file.path)?;
285    for diagnostic in &file.diagnostics {
286        normalized_path(&diagnostic.path)?;
287    }
288    for (edge_index, edge) in file.edges.iter().enumerate() {
289        validate_edge(edge, line, edge_index)?;
290    }
291    Ok(())
292}
meta_ast/output/shard/file.rs

meta_ast/output/shard/
file.rs
