Skip to main content

sqry_classpath/stub/
index.rs

1//! Merged classpath index for fast FQN lookup.
2//!
3//! Combines all per-JAR stubs into a single sorted index with secondary
4//! indices for package and annotation lookup. Persisted at
5//! `.sqry/classpath/index.sqry` using postcard binary serialization.
6//!
7//! ## Lookup performance
8//!
9//! - **FQN lookup**: O(log n) via binary search on the sorted `classes` vec.
10//! - **Package lookup**: O(1) via `package_index` (returns a slice range).
11//! - **Annotation lookup**: O(1) via `annotation_index` (returns index list).
12//!
13//! ## String table
14//!
15//! The string table is reserved for future deduplication optimizations within
16//! the persisted index. Currently it is populated but not used for indirection
17//! — class stubs contain their own string data.
18
19use std::collections::HashMap;
20use std::fs;
21use std::path::Path;
22
23use serde::{Deserialize, Serialize};
24
25use crate::stub::model::ClassStub;
26use crate::{ClasspathError, ClasspathResult};
27
28// ---------------------------------------------------------------------------
29// Constants
30// ---------------------------------------------------------------------------
31
32/// Current format version for the classpath index.
33///
34/// Increment this when the serialized format changes in a way that is not
35/// backward-compatible. On load, a version mismatch triggers a full rebuild.
36pub const CLASSPATH_INDEX_VERSION: u32 = 1;
37
38// ---------------------------------------------------------------------------
39// ClasspathIndex
40// ---------------------------------------------------------------------------
41
42/// Merged classpath index for fast FQN lookup.
43///
44/// Combines all per-JAR stubs into a single sorted index with secondary
45/// indices for package and annotation lookup.
46///
47/// Persisted at: `.sqry/classpath/index.sqry`
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ClasspathIndex {
50    /// Format version (independent from graph snapshot version).
51    pub version: u32,
52    /// Local string table for deduplication within the index.
53    ///
54    /// Reserved for future optimisation; currently populated with unique
55    /// FQN strings but not used for indirection.
56    pub string_table: Vec<String>,
57    /// Classes sorted by FQN for binary search.
58    pub classes: Vec<ClassStub>,
59    /// Package index: package name to (start, end) range of indices into
60    /// `classes`. The range is half-open: `classes[start..end]`.
61    pub package_index: HashMap<String, (usize, usize)>,
62    /// Annotation index: annotation FQN to list of class indices in
63    /// `classes` that carry that annotation.
64    pub annotation_index: HashMap<String, Vec<usize>>,
65}
66
67impl ClasspathIndex {
68    /// Build an index from collected stubs.
69    ///
70    /// 1. Sorts classes by FQN.
71    /// 2. Builds the package index (groups consecutive classes by package prefix).
72    /// 3. Builds the annotation index (scans all class-level annotations).
73    /// 4. Populates the string table with unique FQNs.
74    #[must_use]
75    pub fn build(mut stubs: Vec<ClassStub>) -> Self {
76        // Sort by FQN for binary search.
77        stubs.sort_by(|a, b| a.fqn.cmp(&b.fqn));
78
79        // Build package index.
80        let package_index = build_package_index(&stubs);
81
82        // Build annotation index.
83        let annotation_index = build_annotation_index(&stubs);
84
85        // Build string table (unique FQNs, sorted).
86        let mut string_table: Vec<String> = stubs.iter().map(|s| s.fqn.clone()).collect();
87        string_table.dedup();
88
89        Self {
90            version: CLASSPATH_INDEX_VERSION,
91            string_table,
92            classes: stubs,
93            package_index,
94            annotation_index,
95        }
96    }
97
98    /// Look up a class by fully qualified name (binary search).
99    ///
100    /// Returns `None` if no class with the given FQN exists in the index.
101    #[must_use]
102    pub fn lookup_fqn(&self, fqn: &str) -> Option<&ClassStub> {
103        let idx = self
104            .classes
105            .binary_search_by_key(&fqn, |s| s.fqn.as_str())
106            .ok()?;
107        Some(&self.classes[idx])
108    }
109
110    /// Look up all classes in a package.
111    ///
112    /// Returns an empty slice if the package is not found.
113    #[must_use]
114    pub fn lookup_package(&self, package: &str) -> &[ClassStub] {
115        match self.package_index.get(package) {
116            Some(&(start, end)) => &self.classes[start..end],
117            None => &[],
118        }
119    }
120
121    /// Look up classes with a specific annotation.
122    ///
123    /// Returns an empty vec if no classes carry the given annotation.
124    #[must_use]
125    pub fn lookup_annotated(&self, annotation_fqn: &str) -> Vec<&ClassStub> {
126        match self.annotation_index.get(annotation_fqn) {
127            Some(indices) => indices
128                .iter()
129                .filter_map(|&i| self.classes.get(i))
130                .collect(),
131            None => vec![],
132        }
133    }
134
135    /// Persist the index to disk at the given path.
136    ///
137    /// Uses atomic write (temp file + rename) to prevent corrupt reads.
138    ///
139    /// # Errors
140    ///
141    /// Returns [`ClasspathError::IndexError`] on serialization or I/O failure.
142    pub fn save(&self, path: &Path) -> ClasspathResult<()> {
143        // Ensure parent directory exists.
144        if let Some(parent) = path.parent() {
145            fs::create_dir_all(parent).map_err(|e| {
146                ClasspathError::IndexError(format!(
147                    "cannot create index directory {}: {e}",
148                    parent.display()
149                ))
150            })?;
151        }
152
153        let bytes = postcard::to_allocvec(self).map_err(|e| {
154            ClasspathError::IndexError(format!("cannot serialize classpath index: {e}"))
155        })?;
156
157        // Atomic write: temp file + rename.
158        let temp_path = path.with_extension("sqry.tmp");
159        fs::write(&temp_path, &bytes).map_err(|e| {
160            ClasspathError::IndexError(format!(
161                "cannot write temp index file {}: {e}",
162                temp_path.display()
163            ))
164        })?;
165
166        fs::rename(&temp_path, path).map_err(|e| {
167            // Clean up temp file on rename failure.
168            let _ = fs::remove_file(&temp_path);
169            ClasspathError::IndexError(format!(
170                "cannot rename temp index to {}: {e}",
171                path.display()
172            ))
173        })?;
174
175        Ok(())
176    }
177
178    /// Load the index from disk.
179    ///
180    /// # Errors
181    ///
182    /// Returns [`ClasspathError::IndexError`] if:
183    /// - The file cannot be read.
184    /// - The file cannot be deserialized.
185    /// - The format version does not match [`CLASSPATH_INDEX_VERSION`].
186    pub fn load(path: &Path) -> ClasspathResult<Self> {
187        let bytes = fs::read(path).map_err(|e| {
188            ClasspathError::IndexError(format!(
189                "cannot read classpath index {}: {e}",
190                path.display()
191            ))
192        })?;
193
194        let index: Self = postcard::from_bytes(&bytes).map_err(|e| {
195            ClasspathError::IndexError(format!(
196                "cannot deserialize classpath index {}: {e}",
197                path.display()
198            ))
199        })?;
200
201        if index.version != CLASSPATH_INDEX_VERSION {
202            return Err(ClasspathError::IndexError(format!(
203                "classpath index version mismatch: expected {CLASSPATH_INDEX_VERSION}, found {}",
204                index.version
205            )));
206        }
207
208        Ok(index)
209    }
210}
211
212// ---------------------------------------------------------------------------
213// Index building helpers
214// ---------------------------------------------------------------------------
215
216/// Build the package index from sorted stubs.
217///
218/// Groups consecutive classes by their package prefix (everything before the
219/// last `.` in the FQN). For each package, records the `(start, end)` range
220/// into the sorted classes vec.
221fn build_package_index(sorted_classes: &[ClassStub]) -> HashMap<String, (usize, usize)> {
222    let mut index: HashMap<String, (usize, usize)> = HashMap::new();
223    if sorted_classes.is_empty() {
224        return index;
225    }
226
227    let mut current_package = package_of(&sorted_classes[0].fqn);
228    let mut range_start = 0;
229
230    for (i, stub) in sorted_classes.iter().enumerate().skip(1) {
231        let pkg = package_of(&stub.fqn);
232        if pkg != current_package {
233            index.insert(current_package, (range_start, i));
234            current_package = pkg;
235            range_start = i;
236        }
237    }
238
239    // Insert the last group.
240    index.insert(current_package, (range_start, sorted_classes.len()));
241
242    index
243}
244
245/// Build the annotation index from sorted stubs.
246///
247/// For each class-level annotation, records the class index.
248fn build_annotation_index(sorted_classes: &[ClassStub]) -> HashMap<String, Vec<usize>> {
249    let mut index: HashMap<String, Vec<usize>> = HashMap::new();
250
251    for (i, stub) in sorted_classes.iter().enumerate() {
252        for ann in &stub.annotations {
253            index.entry(ann.type_fqn.clone()).or_default().push(i);
254        }
255    }
256
257    index
258}
259
260/// Extract the package name from a fully qualified class name.
261///
262/// For `"java.util.HashMap"` returns `"java.util"`.
263/// For `"HashMap"` (default package) returns `""`.
264fn package_of(fqn: &str) -> String {
265    match fqn.rfind('.') {
266        Some(pos) => fqn[..pos].to_owned(),
267        None => String::new(),
268    }
269}
270
271// ---------------------------------------------------------------------------
272// Tests
273// ---------------------------------------------------------------------------
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278    use crate::stub::model::{AccessFlags, AnnotationStub, ClassKind};
279    use tempfile::TempDir;
280
281    /// Create a minimal `ClassStub` for testing.
282    fn make_stub(fqn: &str) -> ClassStub {
283        ClassStub {
284            fqn: fqn.to_owned(),
285            name: fqn.rsplit('.').next().unwrap_or(fqn).to_owned(),
286            kind: ClassKind::Class,
287            access: AccessFlags::new(0x0021),
288            superclass: Some("java.lang.Object".to_owned()),
289            interfaces: vec![],
290            methods: vec![],
291            fields: vec![],
292            annotations: vec![],
293            generic_signature: None,
294            inner_classes: vec![],
295            lambda_targets: vec![],
296            module: None,
297            record_components: vec![],
298            enum_constants: vec![],
299            source_file: None,
300            source_jar: None,
301            kotlin_metadata: None,
302            scala_signature: None,
303        }
304    }
305
306    /// Create a `ClassStub` with annotations for testing.
307    fn make_annotated_stub(fqn: &str, annotation_fqns: &[&str]) -> ClassStub {
308        let mut stub = make_stub(fqn);
309        stub.annotations = annotation_fqns
310            .iter()
311            .map(|a| AnnotationStub {
312                type_fqn: (*a).to_owned(),
313                elements: vec![],
314                is_runtime_visible: true,
315            })
316            .collect();
317        stub
318    }
319
320    #[test]
321    fn test_roundtrip_save_load() {
322        let tmp = TempDir::new().unwrap();
323        let index_path = tmp.path().join("classpath/index.sqry");
324
325        let stubs = vec![
326            make_stub("com.example.Bar"),
327            make_stub("com.example.Foo"),
328            make_stub("java.util.HashMap"),
329        ];
330
331        let index = ClasspathIndex::build(stubs);
332        index.save(&index_path).unwrap();
333
334        let loaded = ClasspathIndex::load(&index_path).unwrap();
335        assert_eq!(loaded.version, CLASSPATH_INDEX_VERSION);
336        assert_eq!(loaded.classes.len(), 3);
337        // Should be sorted.
338        assert_eq!(loaded.classes[0].fqn, "com.example.Bar");
339        assert_eq!(loaded.classes[1].fqn, "com.example.Foo");
340        assert_eq!(loaded.classes[2].fqn, "java.util.HashMap");
341    }
342
343    #[test]
344    fn test_binary_search_by_fqn() {
345        let stubs = vec![
346            make_stub("com.example.Alpha"),
347            make_stub("com.example.Beta"),
348            make_stub("com.example.Gamma"),
349            make_stub("java.util.List"),
350        ];
351
352        let index = ClasspathIndex::build(stubs);
353
354        let found = index.lookup_fqn("com.example.Beta");
355        assert!(found.is_some());
356        assert_eq!(found.unwrap().fqn, "com.example.Beta");
357
358        let found = index.lookup_fqn("java.util.List");
359        assert!(found.is_some());
360        assert_eq!(found.unwrap().fqn, "java.util.List");
361
362        let not_found = index.lookup_fqn("com.example.DoesNotExist");
363        assert!(not_found.is_none());
364    }
365
366    #[test]
367    fn test_package_index_lookup() {
368        let stubs = vec![
369            make_stub("com.example.Alpha"),
370            make_stub("com.example.Beta"),
371            make_stub("java.util.HashMap"),
372            make_stub("java.util.List"),
373            make_stub("java.util.Map"),
374        ];
375
376        let index = ClasspathIndex::build(stubs);
377
378        let com_example = index.lookup_package("com.example");
379        assert_eq!(com_example.len(), 2);
380        assert_eq!(com_example[0].fqn, "com.example.Alpha");
381        assert_eq!(com_example[1].fqn, "com.example.Beta");
382
383        let java_util = index.lookup_package("java.util");
384        assert_eq!(java_util.len(), 3);
385
386        let empty = index.lookup_package("org.nonexistent");
387        assert!(empty.is_empty());
388    }
389
390    #[test]
391    fn test_annotation_index_lookup() {
392        let stubs = vec![
393            make_annotated_stub(
394                "com.example.MyController",
395                &["org.springframework.stereotype.Controller"],
396            ),
397            make_annotated_stub(
398                "com.example.MyService",
399                &["org.springframework.stereotype.Service"],
400            ),
401            make_annotated_stub(
402                "com.example.AnotherController",
403                &[
404                    "org.springframework.stereotype.Controller",
405                    "org.springframework.web.bind.annotation.RestController",
406                ],
407            ),
408        ];
409
410        let index = ClasspathIndex::build(stubs);
411
412        let controllers = index.lookup_annotated("org.springframework.stereotype.Controller");
413        assert_eq!(controllers.len(), 2);
414        // Sorted by FQN.
415        assert_eq!(controllers[0].fqn, "com.example.AnotherController");
416        assert_eq!(controllers[1].fqn, "com.example.MyController");
417
418        let services = index.lookup_annotated("org.springframework.stereotype.Service");
419        assert_eq!(services.len(), 1);
420        assert_eq!(services[0].fqn, "com.example.MyService");
421
422        let none = index.lookup_annotated("javax.persistence.Entity");
423        assert!(none.is_empty());
424    }
425
426    #[test]
427    fn test_empty_index() {
428        let stubs: Vec<ClassStub> = vec![];
429        let index = ClasspathIndex::build(stubs);
430
431        assert_eq!(index.classes.len(), 0);
432        assert!(index.lookup_fqn("anything").is_none());
433        assert!(index.lookup_package("anything").is_empty());
434        assert!(index.lookup_annotated("anything").is_empty());
435    }
436
437    #[test]
438    fn test_version_mismatch_on_load() {
439        let tmp = TempDir::new().unwrap();
440        let index_path = tmp.path().join("index.sqry");
441
442        let mut index = ClasspathIndex::build(vec![make_stub("com.example.Foo")]);
443        index.version = 999; // Wrong version.
444        index.save(&index_path).unwrap();
445
446        let result = ClasspathIndex::load(&index_path);
447        assert!(result.is_err());
448        let err_msg = result.unwrap_err().to_string();
449        assert!(
450            err_msg.contains("version mismatch"),
451            "expected version mismatch error, got: {err_msg}"
452        );
453    }
454
455    #[test]
456    fn test_large_index_sort_and_search() {
457        let stubs: Vec<ClassStub> = (0..1500)
458            .map(|i| make_stub(&format!("com.example.Class{i:04}")))
459            .collect();
460
461        let index = ClasspathIndex::build(stubs);
462        assert_eq!(index.classes.len(), 1500);
463
464        // Verify sorted order.
465        for window in index.classes.windows(2) {
466            assert!(
467                window[0].fqn <= window[1].fqn,
468                "sort violation: {} > {}",
469                window[0].fqn,
470                window[1].fqn
471            );
472        }
473
474        // Binary search should find specific entries.
475        assert!(index.lookup_fqn("com.example.Class0000").is_some());
476        assert!(index.lookup_fqn("com.example.Class0750").is_some());
477        assert!(index.lookup_fqn("com.example.Class1499").is_some());
478        assert!(index.lookup_fqn("com.example.Class1500").is_none());
479    }
480
481    #[test]
482    fn test_load_corrupt_file() {
483        let tmp = TempDir::new().unwrap();
484        let index_path = tmp.path().join("index.sqry");
485        fs::write(&index_path, b"corrupt garbage data").unwrap();
486
487        let result = ClasspathIndex::load(&index_path);
488        assert!(result.is_err());
489        assert!(matches!(result.unwrap_err(), ClasspathError::IndexError(_)),);
490    }
491
492    #[test]
493    fn test_load_nonexistent_file() {
494        let result = ClasspathIndex::load(Path::new("/nonexistent/index.sqry"));
495        assert!(result.is_err());
496    }
497
498    #[test]
499    fn test_default_package() {
500        let stubs = vec![make_stub("DefaultClass"), make_stub("AnotherDefault")];
501
502        let index = ClasspathIndex::build(stubs);
503
504        // Default package is empty string.
505        let default_pkg = index.lookup_package("");
506        assert_eq!(default_pkg.len(), 2);
507    }
508
509    #[test]
510    fn test_package_of() {
511        assert_eq!(package_of("java.util.HashMap"), "java.util");
512        assert_eq!(package_of("HashMap"), "");
513        assert_eq!(
514            package_of("com.example.deep.nested.Class"),
515            "com.example.deep.nested"
516        );
517    }
518}