Skip to main content

tessera_codegraph/
snapshot.rs

1use std::collections::HashMap;
2use std::fs::{self, File};
3use std::io::Write;
4use std::path::Path;
5
6use anyhow::Result;
7use memmap2::Mmap;
8use rusqlite::Connection;
9use serde::{Deserialize, Serialize};
10
11use crate::db;
12use crate::types::SymbolRecord;
13
14/// Memory-mapped, derived view of the graph. Snapshot is a derived artifact —
15/// SQLite is always the source of truth. Building the snapshot lets the MCP
16/// server skip SQLite round-trips on the hot path: every `find_definition`,
17/// `find_references`, and `impact` query becomes pure-memory work.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct GraphSnapshot {
20    pub version: u32,
21    pub symbols: Vec<SymbolEntry>,
22    pub edges: Vec<(u32, u32)>,
23    pub name_index: HashMap<String, Vec<u32>>,
24    pub qualified_index: HashMap<String, Vec<u32>>,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct SymbolEntry {
29    pub id: i64,
30    pub name: String,
31    pub qualified_name: String,
32    pub kind: String,
33    pub file_id: i64,
34    pub path: String,
35    pub language: String,
36    pub start_line: u32,
37    pub end_line: u32,
38    pub signature: String,
39    pub exported: bool,
40}
41
42impl SymbolEntry {
43    pub fn to_record(&self) -> SymbolRecord {
44        SymbolRecord {
45            id: self.id,
46            name: self.name.clone(),
47            qualified_name: self.qualified_name.clone(),
48            kind: self.kind.clone(),
49            file_id: self.file_id,
50            path: self.path.clone(),
51            language: self.language.clone(),
52            start_line: self.start_line as usize,
53            end_line: self.end_line as usize,
54            signature: self.signature.clone(),
55            exported: self.exported,
56        }
57    }
58}
59
60pub const SNAPSHOT_VERSION: u32 = 1;
61
62pub fn build(conn: &Connection, path: &Path) -> Result<()> {
63    if let Some(parent) = path.parent() {
64        if !parent.as_os_str().is_empty() {
65            fs::create_dir_all(parent)?;
66        }
67    }
68
69    let mut id_to_index: HashMap<i64, u32> = HashMap::new();
70    let mut symbols: Vec<SymbolEntry> = Vec::new();
71    let mut name_index: HashMap<String, Vec<u32>> = HashMap::new();
72    let mut qualified_index: HashMap<String, Vec<u32>> = HashMap::new();
73
74    {
75        let mut stmt = conn.prepare(
76            "
77            SELECT s.id, s.name, s.qualified_name, s.kind, s.file_id, f.path, f.language,
78                   s.start_line, s.end_line, s.signature, s.exported
79            FROM symbols s
80            JOIN files f ON f.id = s.file_id
81            ORDER BY s.id
82            ",
83        )?;
84        let rows = stmt.query_map([], db::map_symbol)?;
85        for row in rows {
86            let record = row?;
87            let idx = symbols.len() as u32;
88            id_to_index.insert(record.id, idx);
89            name_index.entry(record.name.clone()).or_default().push(idx);
90            qualified_index
91                .entry(record.qualified_name.clone())
92                .or_default()
93                .push(idx);
94            symbols.push(SymbolEntry {
95                id: record.id,
96                name: record.name,
97                qualified_name: record.qualified_name,
98                kind: record.kind,
99                file_id: record.file_id,
100                path: record.path,
101                language: record.language,
102                start_line: record.start_line as u32,
103                end_line: record.end_line as u32,
104                signature: record.signature,
105                exported: record.exported,
106            });
107        }
108    }
109
110    let mut edges: Vec<(u32, u32)> = Vec::new();
111    {
112        let mut stmt = conn.prepare(
113            "
114            SELECT e.from_symbol_id, e.to_symbol_name
115            FROM edges e
116            ",
117        )?;
118        let rows = stmt.query_map([], |row| {
119            let from: i64 = row.get(0)?;
120            let to_name: String = row.get(1)?;
121            Ok((from, to_name))
122        })?;
123        for row in rows {
124            let (from, to_name) = row?;
125            let Some(&from_idx) = id_to_index.get(&from) else {
126                continue;
127            };
128            // Resolve to-name to one or more target indices.
129            let candidates = qualified_index
130                .get(&to_name)
131                .or_else(|| name_index.get(&to_name));
132            if let Some(targets) = candidates {
133                for &to_idx in targets {
134                    edges.push((from_idx, to_idx));
135                }
136            }
137        }
138    }
139
140    let snapshot = GraphSnapshot {
141        version: SNAPSHOT_VERSION,
142        symbols,
143        edges,
144        name_index,
145        qualified_index,
146    };
147    let encoded = bincode::serialize(&snapshot)?;
148    let tmp_path = path.with_extension("bin.tmp");
149    {
150        let mut file = File::create(&tmp_path)?;
151        file.write_all(&encoded)?;
152        file.sync_all()?;
153    }
154    fs::rename(&tmp_path, path)?;
155    Ok(())
156}
157
158pub struct MappedSnapshot {
159    _mmap: Mmap,
160    snapshot: GraphSnapshot,
161}
162
163impl MappedSnapshot {
164    pub fn open(path: &Path) -> Result<Self> {
165        let file = File::open(path)?;
166        let mmap = unsafe { Mmap::map(&file)? };
167        let snapshot: GraphSnapshot = bincode::deserialize(&mmap[..])?;
168        Ok(Self {
169            _mmap: mmap,
170            snapshot,
171        })
172    }
173
174    pub fn graph(&self) -> &GraphSnapshot {
175        &self.snapshot
176    }
177}
178
179pub fn try_open(path: &Path) -> Option<MappedSnapshot> {
180    MappedSnapshot::open(path).ok()
181}