Skip to main content

ralph/
symbol_index.rs

1/// Lightweight symbol index built by regex-scanning workspace source files.
2/// Supports Rust, Python, JS/TS, Go, Java/Kotlin.
3use ignore::WalkBuilder;
4use regex::Regex;
5use std::collections::HashMap;
6use std::path::{Path, PathBuf};
7
8// ── Types ─────────────────────────────────────────────────────────────────────
9
10#[derive(Debug, Clone)]
11pub enum SymbolKind {
12    Function,
13    AsyncFunction,
14    Struct,
15    Enum,
16    Trait,
17    Impl,
18    Class,
19    Method,
20    Type,
21    Const,
22    Interface,
23    Module,
24}
25
26impl std::fmt::Display for SymbolKind {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        let s = match self {
29            SymbolKind::Function => "fn",
30            SymbolKind::AsyncFunction => "async fn",
31            SymbolKind::Struct => "struct",
32            SymbolKind::Enum => "enum",
33            SymbolKind::Trait => "trait",
34            SymbolKind::Impl => "impl",
35            SymbolKind::Class => "class",
36            SymbolKind::Method => "method",
37            SymbolKind::Type => "type",
38            SymbolKind::Const => "const",
39            SymbolKind::Interface => "interface",
40            SymbolKind::Module => "mod",
41        };
42        write!(f, "{}", s)
43    }
44}
45
46#[derive(Debug, Clone)]
47pub struct Symbol {
48    pub name: String,
49    pub kind: SymbolKind,
50    /// Path relative to workspace root.
51    pub file: PathBuf,
52    /// 1-based line where the symbol definition starts.
53    pub line_start: usize,
54    /// 1-based line where the symbol body ends (estimated).
55    pub line_end: usize,
56}
57
58// ── Index ─────────────────────────────────────────────────────────────────────
59
60pub struct SymbolIndex {
61    /// Lowercase symbol name → all definitions with that name.
62    index: HashMap<String, Vec<Symbol>>,
63    workspace: PathBuf,
64}
65
66impl SymbolIndex {
67    /// Walk all source files under `workspace` and build the index.
68    pub fn build(workspace: &Path) -> Self {
69        let mut index: HashMap<String, Vec<Symbol>> = HashMap::new();
70
71        let walker = WalkBuilder::new(workspace)
72            .hidden(false)
73            .ignore(true)
74            .git_ignore(true)
75            .build();
76
77        for entry in walker.flatten() {
78            let path = entry.path();
79            if !path.is_file() {
80                continue;
81            }
82            let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
83            if !is_source_ext(ext) {
84                continue;
85            }
86            let content = match std::fs::read_to_string(path) {
87                Ok(c) => c,
88                Err(_) => continue,
89            };
90            let rel = path.strip_prefix(workspace).unwrap_or(path).to_path_buf();
91            for sym in extract_symbols(&content, &rel, ext) {
92                index.entry(sym.name.to_lowercase()).or_default().push(sym);
93            }
94        }
95
96        Self {
97            index,
98            workspace: workspace.to_path_buf(),
99        }
100    }
101
102    /// Find all symbols whose name contains `query` (case-insensitive).
103    /// Exact matches are returned first.
104    pub fn find(&self, query: &str, limit: usize) -> Vec<&Symbol> {
105        let q = query.to_lowercase();
106        let mut results: Vec<&Symbol> = Vec::new();
107
108        // Exact match first
109        if let Some(exact) = self.index.get(&q) {
110            results.extend(exact.iter());
111        }
112
113        // Substring matches
114        if results.len() < limit {
115            for (key, syms) in &self.index {
116                if key == &q {
117                    continue; // already added
118                }
119                if key.contains(q.as_str()) {
120                    for sym in syms {
121                        if results.len() >= limit {
122                            break;
123                        }
124                        results.push(sym);
125                    }
126                }
127                if results.len() >= limit {
128                    break;
129                }
130            }
131        }
132
133        results.truncate(limit);
134        results
135    }
136
137    /// Read the source body of the first symbol matching `name` (exact, case-insensitive).
138    pub fn read_body(&self, name: &str) -> Option<String> {
139        let sym = self.index.get(&name.to_lowercase())?.first()?;
140        let full = self.workspace.join(&sym.file);
141        let content = std::fs::read_to_string(&full).ok()?;
142        let lines: Vec<&str> = content.lines().collect();
143        let start = sym.line_start.saturating_sub(1).min(lines.len());
144        let end = sym.line_end.min(lines.len());
145        if start > end {
146            return None;
147        }
148        Some(format!(
149            "// {}:{}\n{}",
150            sym.file.display(),
151            sym.line_start,
152            lines[start..end].join("\n")
153        ))
154    }
155
156    pub fn symbol_count(&self) -> usize {
157        self.index.values().map(|v| v.len()).sum()
158    }
159}
160
161// ── Extraction ────────────────────────────────────────────────────────────────
162
163fn is_source_ext(ext: &str) -> bool {
164    matches!(
165        ext,
166        "rs" | "py"
167            | "js"
168            | "jsx"
169            | "ts"
170            | "tsx"
171            | "go"
172            | "java"
173            | "kt"
174            | "rb"
175            | "c"
176            | "cc"
177            | "cpp"
178            | "cxx"
179            | "h"
180            | "hpp"
181            | "hxx"
182    )
183}
184
185struct PatternDef {
186    re: Regex,
187    kind: SymbolKind,
188}
189
190fn extract_symbols(content: &str, file: &Path, ext: &str) -> Vec<Symbol> {
191    let pats = patterns_for(ext);
192    let lines: Vec<&str> = content.lines().collect();
193    let mut out: Vec<Symbol> = Vec::new();
194    let mut seen_positions = std::collections::HashSet::new();
195
196    for pd in pats {
197        for cap in pd.re.captures_iter(content) {
198            let name = match cap.get(1) {
199                Some(m) => m.as_str().to_string(),
200                None => continue,
201            };
202            if name.len() > 64 {
203                continue;
204            }
205            let byte_pos = cap.get(0).unwrap().start();
206            let line_start = content[..byte_pos].matches('\n').count() + 1;
207            if !seen_positions.insert(line_start) {
208                continue; // duplicate at same line
209            }
210            let line_end = estimate_end(&lines, line_start - 1, ext);
211            out.push(Symbol {
212                name,
213                kind: pd.kind.clone(),
214                file: file.to_path_buf(),
215                line_start,
216                line_end,
217            });
218        }
219    }
220
221    out.sort_by_key(|s| s.line_start);
222    out
223}
224
225/// Build patterns for a given extension group (called at most once per group).
226fn build_patterns(specs: &[(&str, SymbolKind)]) -> Vec<PatternDef> {
227    specs
228        .iter()
229        .filter_map(|(re_str, kind)| {
230            Regex::new(re_str).ok().map(|re| PatternDef {
231                re,
232                kind: kind.clone(),
233            })
234        })
235        .collect()
236}
237
238fn patterns_for(ext: &str) -> &'static [PatternDef] {
239    use std::sync::OnceLock;
240
241    static RS: OnceLock<Vec<PatternDef>> = OnceLock::new();
242    static PY: OnceLock<Vec<PatternDef>> = OnceLock::new();
243    static JS: OnceLock<Vec<PatternDef>> = OnceLock::new();
244    static GO: OnceLock<Vec<PatternDef>> = OnceLock::new();
245    static JV: OnceLock<Vec<PatternDef>> = OnceLock::new();
246    static CC: OnceLock<Vec<PatternDef>> = OnceLock::new();
247
248    match ext {
249        "rs" => RS.get_or_init(|| build_patterns(&[
250            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?async\s+fn\s+(\w+)", SymbolKind::AsyncFunction),
251            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?fn\s+(\w+)",         SymbolKind::Function),
252            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?struct\s+(\w+)",     SymbolKind::Struct),
253            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?enum\s+(\w+)",       SymbolKind::Enum),
254            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?trait\s+(\w+)",      SymbolKind::Trait),
255            (r"(?m)^[ \t]*impl(?:<[^>]*>)?\s+(\w+)",                    SymbolKind::Impl),
256            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?type\s+(\w+)",       SymbolKind::Type),
257            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?(?:const|static)\s+(\w+)", SymbolKind::Const),
258            (r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?mod\s+(\w+)",        SymbolKind::Module),
259        ])),
260        "py" => PY.get_or_init(|| build_patterns(&[
261            (r"(?m)^[ \t]*async\s+def\s+(\w+)", SymbolKind::AsyncFunction),
262            (r"(?m)^[ \t]*def\s+(\w+)",         SymbolKind::Function),
263            (r"(?m)^[ \t]*class\s+(\w+)",        SymbolKind::Class),
264        ])),
265        "js" | "jsx" | "ts" | "tsx" => JS.get_or_init(|| build_patterns(&[
266            (r"(?m)^[ \t]*(?:export\s+)?(?:default\s+)?async\s+function\s+(\w+)", SymbolKind::AsyncFunction),
267            (r"(?m)^[ \t]*(?:export\s+)?(?:default\s+)?function\s+(\w+)",         SymbolKind::Function),
268            (r"(?m)^[ \t]*(?:export\s+)?(?:default\s+)?class\s+(\w+)",            SymbolKind::Class),
269            (r"(?m)^[ \t]*(?:export\s+)?interface\s+(\w+)",                       SymbolKind::Interface),
270            (r"(?m)^[ \t]*(?:export\s+)?type\s+(\w+)\s*=",                        SymbolKind::Type),
271            (r"(?m)^[ \t]*(?:export\s+)?(?:const|let)\s+(\w+)\s*[=:]",           SymbolKind::Const),
272        ])),
273        "go" => GO.get_or_init(|| build_patterns(&[
274            (r"(?m)^func\s+(?:\([^)]*\)\s+)?(\w+)", SymbolKind::Function),
275            (r"(?m)^type\s+(\w+)\s+struct",         SymbolKind::Struct),
276            (r"(?m)^type\s+(\w+)\s+interface",      SymbolKind::Interface),
277            (r"(?m)^type\s+(\w+)\b",                SymbolKind::Type),
278            (r"(?m)^const\s+(\w+)",                 SymbolKind::Const),
279        ])),
280        "java" | "kt" | "rb" => JV.get_or_init(|| build_patterns(&[
281            (r"(?m)^[ \t]*(?:public\s+|private\s+|protected\s+)?(?:static\s+)?(?:abstract\s+)?(?:class|record)\s+(\w+)", SymbolKind::Class),
282            (r"(?m)^[ \t]*(?:public\s+|private\s+|protected\s+)?interface\s+(\w+)", SymbolKind::Interface),
283        ])),
284        "c" | "cc" | "cpp" | "cxx" | "h" | "hpp" | "hxx" => CC.get_or_init(|| build_patterns(&[
285            (r"(?m)^[ \t]*struct\s+(\w+)",                         SymbolKind::Struct),
286            (r"(?m)^[ \t]*class\s+(\w+)",                          SymbolKind::Class),
287            (r"(?m)^[ \t]*enum\s+(?:class\s+)?(\w+)",              SymbolKind::Enum),
288            (r"(?m)^[ \t]*union\s+(\w+)",                          SymbolKind::Struct),
289            (r"(?m)^[ \t]*namespace\s+(\w+)",                      SymbolKind::Module),
290            (r"(?m)^[ \t]*using\s+(\w+)\s*=",                      SymbolKind::Type),
291            (r"(?m)^[ \t]*typedef\s+[\w\s*&:<>]+\s+(\w+)\s*;",    SymbolKind::Type),
292            (r"(?m)^[ \t]*template\s*<[^>]*>\s*(?:class|struct)\s+(\w+)", SymbolKind::Struct),
293            (r"(?m)^[ \t]*(?:(?:static|inline|extern|virtual|explicit|constexpr|override|friend|__attribute__\S*)\s+)*(?:[\w:*&<>\[\]~]+\s+)+(\w+)\s*\(", SymbolKind::Function),
294        ])),
295        _ => {
296            static EMPTY: OnceLock<Vec<PatternDef>> = OnceLock::new();
297            EMPTY.get_or_init(Vec::new)
298        }
299    }
300}
301
302/// Estimate the end line of a symbol starting at `start_idx` (0-based).
303fn estimate_end(lines: &[&str], start_idx: usize, ext: &str) -> usize {
304    let total = lines.len();
305    if start_idx >= total {
306        return total;
307    }
308
309    if ext == "py" {
310        // Python: end when indentation returns to or below the def/class line
311        let base = leading_spaces(lines[start_idx]);
312        for i in (start_idx + 1)..total {
313            let line = lines[i];
314            if line.trim().is_empty() || line.trim().starts_with('#') {
315                continue;
316            }
317            if leading_spaces(line) <= base {
318                return i;
319            }
320        }
321        return total;
322    }
323
324    // Brace languages: scan for balanced braces
325    let mut depth: i32 = 0;
326    let mut found_open = false;
327    let limit = (start_idx + 250).min(total);
328    for i in start_idx..limit {
329        for ch in lines[i].chars() {
330            match ch {
331                '{' => {
332                    depth += 1;
333                    found_open = true;
334                }
335                '}' => {
336                    depth -= 1;
337                    if found_open && depth == 0 {
338                        return i + 1;
339                    }
340                }
341                _ => {}
342            }
343        }
344    }
345    limit
346}
347
348fn leading_spaces(s: &str) -> usize {
349    s.len() - s.trim_start().len()
350}