ryo_analysis/ast/
registry.rs

1//! ASTRegistry - Complete AST storage for symbols
2//!
3//! Stores complete PureItem per SymbolId for O(1) access.
4//! Enables direct AST mutation without file I/O.
5//!
6//! # Design
7//!
8//! ```text
9//! SymbolRegistry           ASTRegistry
10//! ┌──────────────┐        ┌─────────────────────────────────┐
11//! │ Path ↔ Id    │        │ SecondaryMap<SymbolId, PureItem>│
12//! │ Id → Kind    │        │                                 │
13//! │ Id → Span    │        │ Complete AST per symbol         │
14//! └──────────────┘        └─────────────────────────────────┘
15//!        ↑                           ↑
16//!        └───────────────────────────┘
17//!              Mutation operates on both
18//! ```
19//!
20//! # Usage
21//!
22//! ```ignore
23//! // Get AST for mutation
24//! if let Some(item) = ast_registry.get_mut(symbol_id) {
25//!     if let PureItem::Struct(s) = item {
26//!         s.fields.push(new_field);
27//!     }
28//! }
29//!
30//! // Register new symbol
31//! let id = symbol_registry.register(path, SymbolKind::Struct);
32//! ast_registry.set(id, PureItem::Struct(new_struct));
33//! ```
34
35use ryo_source::pure::{PureFile, PureItem};
36use ryo_symbol::WorkspaceFilePath;
37use slotmap::SecondaryMap;
38use std::collections::HashSet;
39use std::sync::Arc;
40
41use crate::symbol::{SymbolId, SymbolPath, SymbolRegistry};
42
43/// Complete AST storage for symbols
44///
45/// Stores the full PureItem AST for each symbol, enabling:
46/// - Direct AST mutation without file I/O
47/// - Complete AST reconstruction at dump time
48/// - Simplified mutation logic (no type-specific storage)
49///
50/// # SSOT Design
51///
52/// - All items stored in unified `items` map via SymbolId
53/// - Module content stored in PureMod.items (Use statements, Impl blocks, etc.)
54/// - Module children tracked in `module_children` for O(1) access
55/// - Binary entries: individual symbols registered in `items` (no separate storage)
56/// - Inline modules tracked separately in `inline_modules` for generation
57#[derive(Debug, Clone, Default)]
58pub struct ASTRegistry {
59    /// SymbolId → Complete PureItem AST (SSOT: Single Source of Truth)
60    ///
61    /// For modules (PureItem::Mod), the PureMod.items field contains all
62    /// items in that module including Use statements and Impl blocks.
63    items: SecondaryMap<SymbolId, PureItem>,
64    /// ModuleSymbolId → Child SymbolIds (direct children only)
65    ///
66    /// Enables O(1) access to module children without traversing `items`.
67    module_children: SecondaryMap<SymbolId, Vec<SymbolId>>,
68    /// Set of SymbolIds for modules that are defined inline (e.g., `mod tests { ... }`)
69    ///
70    /// This distinguishes inline modules from:
71    /// - External modules (defined in separate files via `mod name;`)
72    /// - Modules created via CreateMod mutation
73    ///
74    /// Used by RegistryGenerator to determine if a module should stay inline
75    /// or get its own file.
76    inline_modules: HashSet<SymbolId>,
77}
78
79impl ASTRegistry {
80    /// Create a new empty registry
81    pub fn new() -> Self {
82        Self {
83            items: SecondaryMap::new(),
84            module_children: SecondaryMap::new(),
85            inline_modules: HashSet::new(),
86        }
87    }
88
89    /// Get AST for a symbol (immutable)
90    pub fn get(&self, id: SymbolId) -> Option<&PureItem> {
91        self.items.get(id)
92    }
93
94    /// Get AST for a symbol (mutable)
95    pub fn get_mut(&mut self, id: SymbolId) -> Option<&mut PureItem> {
96        self.items.get_mut(id)
97    }
98
99    /// Set AST for a symbol
100    ///
101    /// Overwrites existing AST if present.
102    pub fn set(&mut self, id: SymbolId, item: PureItem) {
103        self.items.insert(id, item);
104    }
105
106    /// Remove AST for a symbol
107    ///
108    /// Returns the removed AST if it existed.
109    pub fn remove(&mut self, id: SymbolId) -> Option<PureItem> {
110        self.items.remove(id)
111    }
112
113    /// Check if symbol has AST
114    pub fn contains(&self, id: SymbolId) -> bool {
115        self.items.contains_key(id)
116    }
117
118    /// Get number of stored ASTs
119    pub fn len(&self) -> usize {
120        self.items.len()
121    }
122
123    /// Check if registry is empty
124    pub fn is_empty(&self) -> bool {
125        self.items.is_empty()
126    }
127
128    /// Iterate all (SymbolId, &PureItem) pairs
129    pub fn iter(&self) -> impl Iterator<Item = (SymbolId, &PureItem)> {
130        self.items.iter()
131    }
132
133    /// Iterate all (SymbolId, &mut PureItem) pairs
134    pub fn iter_mut(&mut self) -> impl Iterator<Item = (SymbolId, &mut PureItem)> {
135        self.items.iter_mut()
136    }
137
138    /// Clear all stored ASTs
139    pub fn clear(&mut self) {
140        self.items.clear();
141        self.module_children.clear();
142        self.inline_modules.clear();
143    }
144
145    // === Inline Module Tracking ===
146
147    /// Mark a module as inline (defined with `mod name { ... }` syntax)
148    ///
149    /// Inline modules stay in their parent file during generation,
150    /// while external modules get their own files.
151    pub fn mark_inline_module(&mut self, id: SymbolId) {
152        self.inline_modules.insert(id);
153    }
154
155    /// Check if a module is inline
156    pub fn is_inline_module(&self, id: SymbolId) -> bool {
157        self.inline_modules.contains(&id)
158    }
159
160    /// Get all inline module IDs
161    pub fn inline_module_ids(&self) -> impl Iterator<Item = SymbolId> + '_ {
162        self.inline_modules.iter().copied()
163    }
164
165    // === Module Children API (SSOT) ===
166
167    /// Get child symbol IDs for a module
168    pub fn get_module_children(&self, module_id: SymbolId) -> Option<&Vec<SymbolId>> {
169        self.module_children.get(module_id)
170    }
171
172    /// Get child symbol IDs for a module (mutable)
173    pub fn get_module_children_mut(&mut self, module_id: SymbolId) -> Option<&mut Vec<SymbolId>> {
174        self.module_children.get_mut(module_id)
175    }
176
177    /// Set child symbol IDs for a module
178    pub fn set_module_children(&mut self, module_id: SymbolId, children: Vec<SymbolId>) {
179        self.module_children.insert(module_id, children);
180    }
181
182    /// Add a child to a module
183    ///
184    /// # Panics
185    /// Panics if `module_id` is not a valid SymbolId in the underlying SlotMap.
186    /// Callers are responsible for supplying a registered SymbolId.
187    pub fn add_child_to_module(&mut self, module_id: SymbolId, child_id: SymbolId) {
188        self.module_children
189            .entry(module_id)
190            .expect("caller must supply a valid SymbolId registered in the SlotMap")
191            .or_default()
192            .push(child_id);
193    }
194
195    /// Remove a child from a module
196    pub fn remove_child_from_module(&mut self, module_id: SymbolId, child_id: SymbolId) {
197        if let Some(children) = self.module_children.get_mut(module_id) {
198            children.retain(|&id| id != child_id);
199        }
200    }
201
202    /// Check if module has children tracked
203    pub fn has_module_children(&self, module_id: SymbolId) -> bool {
204        self.module_children.contains_key(module_id)
205    }
206
207    /// Iterate all `(ModuleSymbolId, &Vec<SymbolId>)` pairs
208    pub fn iter_module_children(&self) -> impl Iterator<Item = (SymbolId, &Vec<SymbolId>)> {
209        self.module_children.iter()
210    }
211
212    // === Module Items API ===
213    // These access PureMod.items directly from the items map.
214
215    /// Get all items for a module (including use statements, impl blocks)
216    pub fn get_module_items(&self, module_id: SymbolId) -> Option<&Vec<PureItem>> {
217        match self.items.get(module_id) {
218            Some(PureItem::Mod(m)) => Some(&m.items),
219            _ => None,
220        }
221    }
222
223    /// Get all items for a module (mutable)
224    pub fn get_module_items_mut(&mut self, module_id: SymbolId) -> Option<&mut Vec<PureItem>> {
225        match self.items.get_mut(module_id) {
226            Some(PureItem::Mod(m)) => Some(&mut m.items),
227            _ => None,
228        }
229    }
230
231    /// Set all items for a module.
232    ///
233    /// If the module doesn't exist in the AST registry yet, creates a new
234    /// empty PureMod entry and sets its items. If items is empty and the
235    /// module doesn't exist, no entry is created (mod declarations are
236    /// generated from module hierarchy, not stored in ASTRegistry).
237    pub fn set_module_items(&mut self, module_id: SymbolId, items: Vec<PureItem>) {
238        use ryo_source::pure::{PureMod, PureVis};
239        match self.items.get_mut(module_id) {
240            Some(PureItem::Mod(m)) => {
241                m.items = items;
242            }
243            _ => {
244                // Only create a new module entry if there are actual items.
245                // Empty modules (created via CreateMod with no content) don't need
246                // an ASTRegistry entry - the mod declaration is generated from
247                // the module hierarchy in SymbolRegistry.
248                if items.is_empty() {
249                    return;
250                }
251                // Create a new module entry with the items.
252                // The module name and visibility will be set by RegistryGenerator
253                // based on SymbolRegistry, so we use placeholder values here.
254                self.items.insert(
255                    module_id,
256                    PureItem::Mod(PureMod {
257                        attrs: vec![],
258                        vis: PureVis::Public, // Will be overridden by SymbolRegistry visibility
259                        name: String::new(),  // Will be derived from SymbolPath
260                        items,
261                    }),
262                );
263            }
264        }
265    }
266
267    /// Check if module has items stored
268    pub fn has_module_items(&self, module_id: SymbolId) -> bool {
269        match self.items.get(module_id) {
270            Some(PureItem::Mod(m)) => !m.items.is_empty(),
271            _ => false,
272        }
273    }
274
275    /// Iterate all `(ModuleSymbolId, &Vec<PureItem>)` pairs
276    pub fn iter_module_items(&self) -> impl Iterator<Item = (SymbolId, &Vec<PureItem>)> {
277        self.items.iter().filter_map(|(id, item)| {
278            if let PureItem::Mod(m) = item {
279                Some((id, &m.items))
280            } else {
281                None
282            }
283        })
284    }
285
286    /// Iterate all `(ModuleSymbolId, &mut Vec<PureItem>)` pairs
287    pub fn iter_module_items_mut(
288        &mut self,
289    ) -> impl Iterator<Item = (SymbolId, &mut Vec<PureItem>)> {
290        self.items.iter_mut().filter_map(|(id, item)| {
291            if let PureItem::Mod(m) = item {
292                Some((id, &mut m.items))
293            } else {
294                None
295            }
296        })
297    }
298
299    /// Build ASTRegistry from parsed files.
300    ///
301    /// Iterates over all files and stores PureItem for each symbol found in SymbolRegistry.
302    /// This establishes the initial AST state before mutations.
303    ///
304    /// # Arguments
305    ///
306    /// * `files` - Parsed files keyed by workspace path
307    /// * `registry` - Symbol registry with pre-registered symbols
308    /// * `crate_name` - Crate name for path resolution
309    pub fn build_from_files(
310        files: &im::HashMap<WorkspaceFilePath, Arc<PureFile>>,
311        registry: &SymbolRegistry,
312        _crate_name: &str,
313    ) -> Self {
314        use ryo_symbol::SymbolPathResolver;
315
316        let mut ast_registry = Self::new();
317
318        for (file_path, file) in files {
319            // Use the crate_name from each file's WorkspaceFilePath for correct multi-crate workspace support
320            let file_crate_name = file_path.crate_name().as_str();
321            let resolver = SymbolPathResolver::new(file_crate_name);
322            let module_path_str = resolver.module_path_str(file_path);
323            let module_path = match SymbolPath::parse(&module_path_str) {
324                Ok(p) => p,
325                Err(_) => continue,
326            };
327
328            ast_registry.store_items_from_file(&module_path, file.as_ref(), registry);
329        }
330
331        ast_registry
332    }
333
334    /// Store items from a file using existing SymbolRegistry.
335    ///
336    /// This method merges multiple impl blocks for the same type into one.
337    /// For example, if a file has:
338    /// ```ignore
339    /// impl Foo { fn a() {} }
340    /// impl Foo { fn b() {} }
341    /// ```
342    /// They will be merged into a single impl block with both methods.
343    fn store_items_from_file(
344        &mut self,
345        module_path: &SymbolPath,
346        file: &PureFile,
347        registry: &SymbolRegistry,
348    ) {
349        self.store_items_recursive(module_path, &file.items, registry);
350    }
351
352    /// Recursively store items from a module (handles inline modules).
353    fn store_items_recursive(
354        &mut self,
355        module_path: &SymbolPath,
356        items: &[PureItem],
357        registry: &SymbolRegistry,
358    ) {
359        use ryo_source::pure::PureImpl;
360        use std::collections::HashMap;
361
362        // Phase 1: Merge impl blocks with same (self_ty, trait_) key
363        let mut impl_map: HashMap<(String, Option<String>), PureImpl> = HashMap::new();
364        let mut merged_items: Vec<PureItem> = Vec::new();
365
366        for item in items {
367            if let PureItem::Impl(impl_block) = item {
368                let key = (impl_block.self_ty.clone(), impl_block.trait_.clone());
369                if let Some(existing) = impl_map.get_mut(&key) {
370                    // Merge: append items from this impl block to the existing one
371                    existing.items.extend(impl_block.items.clone());
372                } else {
373                    impl_map.insert(key, impl_block.clone());
374                }
375            } else {
376                merged_items.push(item.clone());
377            }
378        }
379
380        // Add merged impl blocks to the items list
381        for impl_block in impl_map.into_values() {
382            merged_items.push(PureItem::Impl(impl_block));
383        }
384
385        // SSOT: Collect child SymbolIds for module_children
386        let mut child_ids: Vec<SymbolId> = Vec::new();
387
388        // Phase 2: Register individual symbols and collect children
389        for item in &merged_items {
390            // Handle impl blocks: store impl block + individual methods/items
391            // Design:
392            // - Impl blocks are stored as symbols (for trait operations like ExtractTrait)
393            // - Methods are also stored individually for direct access
394            // - Impl block path: <impl Type> or <impl Trait for Type>
395            // - Method path: <impl path>::method_name or Type::method_name (for plain impl)
396            if let PureItem::Impl(impl_block) = item {
397                use ryo_source::pure::PureImplItem;
398
399                // Determine impl block path and method base path
400                let (impl_path_str, method_base_path) =
401                    if let Some(ref trait_name) = &impl_block.trait_ {
402                        // Trait impl: <impl Trait for Type>
403                        let impl_path = format!(
404                            "{}::<impl {} for {}>",
405                            module_path, trait_name, impl_block.self_ty
406                        );
407                        (impl_path.clone(), impl_path)
408                    } else {
409                        // Plain impl: <impl Type> for impl block, Type for methods
410                        let impl_path = format!("{}::<impl {}>", module_path, impl_block.self_ty);
411                        // Strip generic parameters: "Router < S >" → "Router"
412                        // SymbolPath::parse rejects segments with '<' / spaces,
413                        // so method paths must use the base type name.
414                        let base_type = impl_block
415                            .self_ty
416                            .split('<')
417                            .next()
418                            .unwrap_or(&impl_block.self_ty)
419                            .trim();
420                        let method_path = format!("{}::{}", module_path, base_type);
421                        (impl_path, method_path)
422                    };
423
424                // Store impl block itself (required for trait operations)
425                if let Ok(impl_path) = SymbolPath::parse(&impl_path_str) {
426                    if let Some(id) = registry.lookup(&impl_path) {
427                        self.set(id, item.clone());
428                        child_ids.push(id);
429                    }
430                }
431
432                // Store each method/item individually
433                for impl_item in &impl_block.items {
434                    let (item_name, pure_item) = match impl_item {
435                        PureImplItem::Fn(m) => (m.name.clone(), PureItem::Fn(m.clone())),
436                        PureImplItem::Const(c) => (c.name.clone(), PureItem::Const(c.clone())),
437                        PureImplItem::Type(t) => (t.name.clone(), PureItem::Type(t.clone())),
438                        PureImplItem::Other(_) => continue,
439                    };
440
441                    let item_path_str = format!("{}::{}", method_base_path, item_name);
442                    if let Ok(item_path) = SymbolPath::parse(&item_path_str) {
443                        if let Some(id) = registry.lookup(&item_path) {
444                            self.set(id, pure_item);
445                            // Note: method children are under impl block, not module
446                        }
447                    }
448                }
449                continue;
450            }
451
452            // Handle inline modules: recursively process their content
453            if let PureItem::Mod(m) = item {
454                if m.items.is_empty() {
455                    // This is a mod declaration (mod foo;), not an inline module
456                    // Don't store it individually - it would overwrite the actual module content
457                    continue;
458                }
459
460                // Inline module: register the module and recursively process its contents
461                let mod_path = match module_path.child(&m.name) {
462                    Ok(p) => p,
463                    Err(_) => continue,
464                };
465
466                if let Some(id) = registry.lookup(&mod_path) {
467                    child_ids.push(id);
468
469                    // Mark this module as inline (defined with `mod name { ... }` syntax)
470                    // This distinguishes it from external modules created via CreateMod
471                    self.mark_inline_module(id);
472
473                    // Recursively process inline module contents first
474                    // This stores individual items and their children
475                    self.store_items_recursive(&mod_path, &m.items, registry);
476
477                    // After recursion, store the inline module with original attrs
478                    // This overwrites the empty-attrs PureMod created by recursive call
479                    self.set(id, item.clone());
480                }
481                continue;
482            }
483
484            let name = match item_name(item) {
485                Some(n) => n,
486                None => continue,
487            };
488
489            // Build symbol path for this item
490            let item_path = match module_path.child(&name) {
491                Ok(p) => p,
492                Err(_) => continue,
493            };
494
495            // Look up in registry and store
496            if let Some(id) = registry.lookup(&item_path) {
497                self.set(id, item.clone());
498                child_ids.push(id);
499            }
500        }
501
502        // SSOT: Store module with all its items and children
503        // Note: All modules get a PureMod entry with items for proper get_module_items support.
504        // The is_inline_module flag (set during inline module processing above) distinguishes
505        // inline modules from external modules for RegistryGenerator output formatting.
506        if let Some(module_id) = registry.lookup(module_path) {
507            // Determine visibility from SymbolRegistry (preserves original pub/priv)
508            let vis = registry
509                .visibility(module_id)
510                .map(|v| match v {
511                    ryo_symbol::Visibility::Public => ryo_source::pure::PureVis::Public,
512                    _ => ryo_source::pure::PureVis::Private,
513                })
514                .unwrap_or_default();
515
516            // Create PureMod with all items (Use statements, Impl blocks, etc.)
517            let pure_mod = ryo_source::pure::PureMod {
518                attrs: vec![],
519                vis,
520                name: module_path
521                    .segment_refs()
522                    .last()
523                    .map(|s| s.name().to_string())
524                    .unwrap_or_default(),
525                items: merged_items,
526            };
527            self.set(module_id, PureItem::Mod(pure_mod));
528            self.set_module_children(module_id, child_ids);
529        }
530    }
531}
532
533/// Extract name from a PureItem (if it has one).
534fn item_name(item: &PureItem) -> Option<String> {
535    match item {
536        PureItem::Struct(s) => Some(s.name.clone()),
537        PureItem::Enum(e) => Some(e.name.clone()),
538        PureItem::Fn(f) => Some(f.name.clone()),
539        PureItem::Mod(m) => Some(m.name.clone()),
540        PureItem::Trait(t) => Some(t.name.clone()),
541        PureItem::Type(t) => Some(t.name.clone()),
542        PureItem::Const(c) => Some(c.name.clone()),
543        PureItem::Static(s) => Some(s.name.clone()),
544        PureItem::Impl(i) => {
545            // Impl blocks use path format: self_ty::impl or self_ty::impl_TraitName
546            let impl_suffix = match &i.trait_ {
547                Some(trait_name) => format!("impl_{}", trait_name.replace("::", "_")),
548                None => "impl".to_string(),
549            };
550            // Return the combined path segment: e.g., "Counter::impl"
551            Some(format!("{}::{}", i.self_ty, impl_suffix))
552        }
553        PureItem::Use(_) | PureItem::Macro(_) | PureItem::Other(_) => None,
554    }
555}
556
557#[cfg(test)]
558mod tests {
559    use super::*;
560    use ryo_source::pure::{PureFields, PureStruct, PureVis};
561    use slotmap::SlotMap;
562
563    fn make_symbol_id() -> SymbolId {
564        let mut sm: SlotMap<SymbolId, ()> = SlotMap::with_key();
565        sm.insert(())
566    }
567
568    #[test]
569    fn test_inline_module_recursive() {
570        use ryo_source::pure::PureFile;
571
572        let source = r#"
573pub struct Config {
574    pub name: String,
575}
576
577#[cfg(test)]
578mod tests {
579    use super::*;
580
581    #[test]
582    fn test_config() {
583        let c = Config { name: "test".to_string() };
584        assert_eq!(c.name, "test");
585    }
586}
587"#;
588
589        let file = PureFile::from_source(source).unwrap();
590
591        // Verify PureFile structure
592        assert_eq!(file.items.len(), 2, "Should have Config and tests module");
593
594        // Find tests module
595        let tests_mod = file.items.iter().find_map(|item| {
596            if let PureItem::Mod(m) = item {
597                if m.name == "tests" {
598                    return Some(m);
599                }
600            }
601            None
602        });
603
604        let tests_mod = tests_mod.expect("Should have tests module");
605
606        // Verify inline module has content
607        assert!(
608            !tests_mod.items.is_empty(),
609            "tests module should have items"
610        );
611
612        // Check for test_config function (may be wrapped in Use or Fn)
613        let has_fn = tests_mod
614            .items
615            .iter()
616            .any(|item| matches!(item, PureItem::Fn(f) if f.name == "test_config"));
617
618        assert!(
619            has_fn,
620            "tests module should have test_config function, got: {:?}",
621            tests_mod
622                .items
623                .iter()
624                .map(std::mem::discriminant)
625                .collect::<Vec<_>>()
626        );
627    }
628
629    /// Test that #[cfg(test)] attributes are preserved in ASTRegistry
630    #[test]
631    fn test_cfg_test_attr_preserved_in_registry() {
632        use crate::SymbolKind;
633        use ryo_source::pure::{PureAttrMeta, PureFile};
634
635        let source = r#"
636pub struct Config {
637    pub name: String,
638}
639
640#[cfg(test)]
641mod tests {
642    use super::*;
643
644    #[test]
645    fn test_config() {
646        let c = Config { name: "test".to_string() };
647        assert_eq!(c.name, "test");
648    }
649}
650"#;
651
652        let file = PureFile::from_source(source).unwrap();
653
654        // Create SymbolRegistry with all symbols pre-registered
655        let mut symbol_registry = SymbolRegistry::new();
656
657        symbol_registry
658            .register(SymbolPath::parse("my_crate").unwrap(), SymbolKind::Mod)
659            .unwrap();
660        symbol_registry
661            .register(
662                SymbolPath::parse("my_crate::Config").unwrap(),
663                SymbolKind::Struct,
664            )
665            .unwrap();
666        symbol_registry
667            .register(
668                SymbolPath::parse("my_crate::tests").unwrap(),
669                SymbolKind::Mod,
670            )
671            .unwrap();
672        symbol_registry
673            .register(
674                SymbolPath::parse("my_crate::tests::test_config").unwrap(),
675                SymbolKind::Function,
676            )
677            .unwrap();
678
679        // Create ASTRegistry and store items
680        let mut ast_registry = ASTRegistry::new();
681        let module_path = SymbolPath::parse("my_crate").unwrap();
682        ast_registry.store_items_from_file(&module_path, &file, &symbol_registry);
683
684        // Verify tests module is registered
685        let tests_path = SymbolPath::parse("my_crate::tests").unwrap();
686        let tests_id = symbol_registry
687            .lookup(&tests_path)
688            .expect("tests should exist");
689
690        // Critical check: verify #[cfg(test)] attribute is preserved
691        let tests_item = ast_registry
692            .get(tests_id)
693            .expect("tests AST should be registered");
694        if let PureItem::Mod(m) = tests_item {
695            // Check attrs contain cfg(test)
696            let has_cfg_test = m.attrs.iter().any(|attr| {
697                attr.path == "cfg" && matches!(&attr.meta, PureAttrMeta::List(s) if s == "test")
698            });
699            assert!(
700                has_cfg_test,
701                "tests module should have #[cfg(test)] attribute, got attrs: {:?}",
702                m.attrs
703            );
704        } else {
705            panic!("tests should be a module, got: {:?}", tests_item);
706        }
707
708        // Verify inline module is marked
709        assert!(
710            ast_registry.is_inline_module(tests_id),
711            "tests should be marked as inline module"
712        );
713
714        // Verify test_config function is registered
715        let test_config_path = SymbolPath::parse("my_crate::tests::test_config").unwrap();
716        let test_config_id = symbol_registry
717            .lookup(&test_config_path)
718            .expect("test_config should exist");
719        assert!(
720            ast_registry.contains(test_config_id),
721            "test_config AST should be registered"
722        );
723
724        // Verify test_config has #[test] attribute
725        if let Some(PureItem::Fn(f)) = ast_registry.get(test_config_id) {
726            let has_test_attr = f.attrs.iter().any(|attr| attr.path == "test");
727            assert!(
728                has_test_attr,
729                "test_config should have #[test] attribute, got attrs: {:?}",
730                f.attrs
731            );
732        } else {
733            panic!(
734                "test_config should be a function, got: {:?}",
735                ast_registry.get(test_config_id)
736            );
737        }
738
739        // Verify module_items contains use statement
740        let module_items = ast_registry
741            .get_module_items(tests_id)
742            .expect("tests module should have items");
743
744        // Should contain: use super::*, test_config fn
745        let has_use = module_items
746            .iter()
747            .any(|item| matches!(item, PureItem::Use(_)));
748        assert!(
749            has_use,
750            "tests module items should contain use statement, got: {:?}",
751            module_items.iter().map(item_name).collect::<Vec<_>>()
752        );
753
754        let has_fn = module_items
755            .iter()
756            .any(|item| matches!(item, PureItem::Fn(f) if f.name == "test_config"));
757        assert!(
758            has_fn,
759            "tests module items should contain test_config function"
760        );
761    }
762
763    /// Test nested inline test modules
764    #[test]
765    fn test_nested_inline_test_modules() {
766        use crate::SymbolKind;
767        use ryo_source::pure::{PureAttrMeta, PureFile};
768
769        let source = r#"
770pub struct Outer;
771
772#[cfg(test)]
773mod tests {
774    use super::*;
775
776    mod nested {
777        use super::*;
778
779        #[test]
780        fn test_nested() {}
781    }
782
783    #[test]
784    fn test_outer() {}
785}
786"#;
787
788        let file = PureFile::from_source(source).unwrap();
789
790        let mut symbol_registry = SymbolRegistry::new();
791
792        symbol_registry
793            .register(SymbolPath::parse("my_crate").unwrap(), SymbolKind::Mod)
794            .unwrap();
795        symbol_registry
796            .register(
797                SymbolPath::parse("my_crate::Outer").unwrap(),
798                SymbolKind::Struct,
799            )
800            .unwrap();
801        symbol_registry
802            .register(
803                SymbolPath::parse("my_crate::tests").unwrap(),
804                SymbolKind::Mod,
805            )
806            .unwrap();
807        symbol_registry
808            .register(
809                SymbolPath::parse("my_crate::tests::nested").unwrap(),
810                SymbolKind::Mod,
811            )
812            .unwrap();
813        symbol_registry
814            .register(
815                SymbolPath::parse("my_crate::tests::test_outer").unwrap(),
816                SymbolKind::Function,
817            )
818            .unwrap();
819        symbol_registry
820            .register(
821                SymbolPath::parse("my_crate::tests::nested::test_nested").unwrap(),
822                SymbolKind::Function,
823            )
824            .unwrap();
825
826        let mut ast_registry = ASTRegistry::new();
827        let module_path = SymbolPath::parse("my_crate").unwrap();
828        ast_registry.store_items_from_file(&module_path, &file, &symbol_registry);
829
830        // Verify tests module has #[cfg(test)] attribute
831        let tests_path = SymbolPath::parse("my_crate::tests").unwrap();
832        let tests_id = symbol_registry.lookup(&tests_path).unwrap();
833
834        if let Some(PureItem::Mod(m)) = ast_registry.get(tests_id) {
835            let has_cfg_test = m.attrs.iter().any(|attr| {
836                attr.path == "cfg" && matches!(&attr.meta, PureAttrMeta::List(s) if s == "test")
837            });
838            assert!(
839                has_cfg_test,
840                "tests module should have #[cfg(test)], got: {:?}",
841                m.attrs
842            );
843        } else {
844            panic!("tests should be a module");
845        }
846
847        // Verify both tests and nested are marked as inline
848        assert!(ast_registry.is_inline_module(tests_id));
849
850        let nested_path = SymbolPath::parse("my_crate::tests::nested").unwrap();
851        let nested_id = symbol_registry.lookup(&nested_path).unwrap();
852        assert!(ast_registry.is_inline_module(nested_id));
853
854        // Verify nested function is registered
855        let nested_fn_path = SymbolPath::parse("my_crate::tests::nested::test_nested").unwrap();
856        let nested_fn_id = symbol_registry.lookup(&nested_fn_path).unwrap();
857        assert!(
858            ast_registry.contains(nested_fn_id),
859            "nested test function should be registered"
860        );
861
862        // Verify module_children hierarchy
863        let tests_children = ast_registry.get_module_children(tests_id).unwrap();
864        assert!(
865            tests_children.contains(&nested_id),
866            "tests children should contain nested module"
867        );
868    }
869
870    /// Test inline test module with impl block
871    #[test]
872    fn test_inline_test_module_with_impl() {
873        use crate::SymbolKind;
874        use ryo_source::pure::{PureAttrMeta, PureFile};
875
876        let source = r#"
877pub struct Config {
878    pub name: String,
879}
880
881#[cfg(test)]
882mod tests {
883    use super::*;
884
885    struct TestHelper {
886        value: i32,
887    }
888
889    impl TestHelper {
890        fn new(value: i32) -> Self {
891            Self { value }
892        }
893    }
894
895    #[test]
896    fn test_with_helper() {
897        let helper = TestHelper::new(42);
898        assert_eq!(helper.value, 42);
899    }
900}
901"#;
902
903        let file = PureFile::from_source(source).unwrap();
904
905        let mut symbol_registry = SymbolRegistry::new();
906
907        symbol_registry
908            .register(SymbolPath::parse("my_crate").unwrap(), SymbolKind::Mod)
909            .unwrap();
910        symbol_registry
911            .register(
912                SymbolPath::parse("my_crate::Config").unwrap(),
913                SymbolKind::Struct,
914            )
915            .unwrap();
916        symbol_registry
917            .register(
918                SymbolPath::parse("my_crate::tests").unwrap(),
919                SymbolKind::Mod,
920            )
921            .unwrap();
922        symbol_registry
923            .register(
924                SymbolPath::parse("my_crate::tests::TestHelper").unwrap(),
925                SymbolKind::Struct,
926            )
927            .unwrap();
928        symbol_registry
929            .register(
930                SymbolPath::parse("my_crate::tests::<impl TestHelper>").unwrap(),
931                SymbolKind::Impl,
932            )
933            .unwrap();
934        symbol_registry
935            .register(
936                SymbolPath::parse("my_crate::tests::TestHelper::new").unwrap(),
937                SymbolKind::Function,
938            )
939            .unwrap();
940        symbol_registry
941            .register(
942                SymbolPath::parse("my_crate::tests::test_with_helper").unwrap(),
943                SymbolKind::Function,
944            )
945            .unwrap();
946
947        let mut ast_registry = ASTRegistry::new();
948        let module_path = SymbolPath::parse("my_crate").unwrap();
949        ast_registry.store_items_from_file(&module_path, &file, &symbol_registry);
950
951        // Verify tests module has #[cfg(test)] attribute
952        let tests_path = SymbolPath::parse("my_crate::tests").unwrap();
953        let tests_id = symbol_registry.lookup(&tests_path).unwrap();
954
955        if let Some(PureItem::Mod(m)) = ast_registry.get(tests_id) {
956            let has_cfg_test = m.attrs.iter().any(|attr| {
957                attr.path == "cfg" && matches!(&attr.meta, PureAttrMeta::List(s) if s == "test")
958            });
959            assert!(
960                has_cfg_test,
961                "tests module should have #[cfg(test)], got: {:?}",
962                m.attrs
963            );
964        } else {
965            panic!("tests should be a module");
966        }
967
968        // Verify TestHelper struct is registered
969        let helper_path = SymbolPath::parse("my_crate::tests::TestHelper").unwrap();
970        let helper_id = symbol_registry.lookup(&helper_path).unwrap();
971        assert!(
972            ast_registry.contains(helper_id),
973            "TestHelper struct should be registered"
974        );
975
976        // Verify impl block is registered
977        let impl_path = SymbolPath::parse("my_crate::tests::<impl TestHelper>").unwrap();
978        let impl_id = symbol_registry.lookup(&impl_path).unwrap();
979        assert!(
980            ast_registry.contains(impl_id),
981            "impl TestHelper should be registered"
982        );
983
984        // Verify method is registered
985        let new_path = SymbolPath::parse("my_crate::tests::TestHelper::new").unwrap();
986        let new_id = symbol_registry.lookup(&new_path).unwrap();
987        assert!(
988            ast_registry.contains(new_id),
989            "TestHelper::new method should be registered"
990        );
991
992        // Verify module_items contains all expected items
993        let module_items = ast_registry
994            .get_module_items(tests_id)
995            .expect("tests module should have items");
996
997        // Should contain: use, TestHelper struct, impl block, test_with_helper fn
998        assert!(
999            module_items
1000                .iter()
1001                .any(|item| matches!(item, PureItem::Use(_))),
1002            "module_items should contain use statement"
1003        );
1004        assert!(
1005            module_items
1006                .iter()
1007                .any(|item| matches!(item, PureItem::Struct(s) if s.name == "TestHelper")),
1008            "module_items should contain TestHelper struct"
1009        );
1010        assert!(
1011            module_items
1012                .iter()
1013                .any(|item| matches!(item, PureItem::Impl(i) if i.self_ty == "TestHelper")),
1014            "module_items should contain impl TestHelper"
1015        );
1016        assert!(
1017            module_items
1018                .iter()
1019                .any(|item| matches!(item, PureItem::Fn(f) if f.name == "test_with_helper")),
1020            "module_items should contain test_with_helper function"
1021        );
1022    }
1023
1024    #[test]
1025    fn test_store_items_recursive_with_registry() {
1026        use crate::SymbolKind;
1027        use ryo_source::pure::PureFile;
1028
1029        let source = r#"
1030pub struct Config {
1031    pub name: String,
1032}
1033
1034mod tests {
1035    fn test_config() {}
1036}
1037"#;
1038
1039        let file = PureFile::from_source(source).unwrap();
1040
1041        // Create SymbolRegistry with all symbols pre-registered
1042        let mut symbol_registry = SymbolRegistry::new();
1043
1044        // Register symbols manually
1045        symbol_registry
1046            .register(SymbolPath::parse("my_crate").unwrap(), SymbolKind::Mod)
1047            .unwrap();
1048        symbol_registry
1049            .register(
1050                SymbolPath::parse("my_crate::Config").unwrap(),
1051                SymbolKind::Struct,
1052            )
1053            .unwrap();
1054        symbol_registry
1055            .register(
1056                SymbolPath::parse("my_crate::tests").unwrap(),
1057                SymbolKind::Mod,
1058            )
1059            .unwrap();
1060        symbol_registry
1061            .register(
1062                SymbolPath::parse("my_crate::tests::test_config").unwrap(),
1063                SymbolKind::Function,
1064            )
1065            .unwrap();
1066
1067        // Create ASTRegistry and store items
1068        let mut ast_registry = ASTRegistry::new();
1069        let module_path = SymbolPath::parse("my_crate").unwrap();
1070        ast_registry.store_items_from_file(&module_path, &file, &symbol_registry);
1071
1072        // Verify Config is registered
1073        let config_path = SymbolPath::parse("my_crate::Config").unwrap();
1074        let config_id = symbol_registry
1075            .lookup(&config_path)
1076            .expect("Config should exist");
1077        assert!(
1078            ast_registry.contains(config_id),
1079            "Config AST should be registered"
1080        );
1081
1082        // Verify tests module is registered
1083        let tests_path = SymbolPath::parse("my_crate::tests").unwrap();
1084        let tests_id = symbol_registry
1085            .lookup(&tests_path)
1086            .expect("tests should exist");
1087        assert!(
1088            ast_registry.contains(tests_id),
1089            "tests AST should be registered"
1090        );
1091
1092        // Verify test_config function is registered (this is the key test!)
1093        let test_config_path = SymbolPath::parse("my_crate::tests::test_config").unwrap();
1094        let test_config_id = symbol_registry
1095            .lookup(&test_config_path)
1096            .expect("test_config should exist");
1097        assert!(
1098            ast_registry.contains(test_config_id),
1099            "test_config AST should be registered via recursive processing"
1100        );
1101
1102        // Verify test_config is a function
1103        if let Some(PureItem::Fn(f)) = ast_registry.get(test_config_id) {
1104            assert_eq!(f.name, "test_config");
1105        } else {
1106            panic!(
1107                "test_config should be a function, got: {:?}",
1108                ast_registry.get(test_config_id)
1109            );
1110        }
1111    }
1112
1113    #[test]
1114    fn test_basic_operations() {
1115        let mut registry = ASTRegistry::new();
1116        let id = make_symbol_id();
1117
1118        // Initially empty
1119        assert!(registry.is_empty());
1120        assert!(!registry.contains(id));
1121        assert!(registry.get(id).is_none());
1122
1123        // Set AST
1124        let item = PureItem::Struct(PureStruct {
1125            attrs: vec![],
1126            vis: PureVis::Public,
1127            name: "TestStruct".to_string(),
1128            generics: Default::default(),
1129            fields: PureFields::Unit,
1130        });
1131        registry.set(id, item);
1132
1133        // Now contains
1134        assert!(!registry.is_empty());
1135        assert!(registry.contains(id));
1136        assert!(registry.get(id).is_some());
1137        assert_eq!(registry.len(), 1);
1138
1139        // Get mutable and modify
1140        if let Some(PureItem::Struct(s)) = registry.get_mut(id) {
1141            s.name = "ModifiedStruct".to_string();
1142        }
1143
1144        // Verify modification
1145        if let Some(PureItem::Struct(s)) = registry.get(id) {
1146            assert_eq!(s.name, "ModifiedStruct");
1147        } else {
1148            panic!("Expected struct");
1149        }
1150
1151        // Remove
1152        let removed = registry.remove(id);
1153        assert!(removed.is_some());
1154        assert!(registry.is_empty());
1155    }
1156
1157    #[test]
1158    fn test_iteration() {
1159        let mut registry = ASTRegistry::new();
1160        let mut sm: SlotMap<SymbolId, ()> = SlotMap::with_key();
1161
1162        let id1 = sm.insert(());
1163        let id2 = sm.insert(());
1164
1165        registry.set(
1166            id1,
1167            PureItem::Struct(PureStruct {
1168                attrs: vec![],
1169                vis: PureVis::Public,
1170                name: "Struct1".to_string(),
1171                generics: Default::default(),
1172                fields: PureFields::Unit,
1173            }),
1174        );
1175
1176        registry.set(
1177            id2,
1178            PureItem::Struct(PureStruct {
1179                attrs: vec![],
1180                vis: PureVis::Private,
1181                name: "Struct2".to_string(),
1182                generics: Default::default(),
1183                fields: PureFields::Unit,
1184            }),
1185        );
1186
1187        let items: Vec<_> = registry.iter().collect();
1188        assert_eq!(items.len(), 2);
1189    }
1190
1191    #[test]
1192    fn test_is_binary_entry() {
1193        // main.rs should be binary entry
1194        let main_rs = WorkspaceFilePath::new_for_test("src/main.rs", "/workspace", "my_crate");
1195        assert!(main_rs.is_binary_entry());
1196
1197        // lib.rs should NOT be binary entry
1198        let lib_rs = WorkspaceFilePath::new_for_test("src/lib.rs", "/workspace", "my_crate");
1199        assert!(!lib_rs.is_binary_entry());
1200
1201        // src/bin/foo.rs should be binary entry
1202        let bin_foo = WorkspaceFilePath::new_for_test("src/bin/foo.rs", "/workspace", "my_crate");
1203        assert!(bin_foo.is_binary_entry());
1204
1205        // Regular module should NOT be binary entry
1206        let module = WorkspaceFilePath::new_for_test("src/utils/mod.rs", "/workspace", "my_crate");
1207        assert!(!module.is_binary_entry());
1208
1209        // Workspace path with main.rs
1210        let workspace_main =
1211            WorkspaceFilePath::new_for_test("crates/my_app/src/main.rs", "/workspace", "my_app");
1212        assert!(workspace_main.is_binary_entry());
1213    }
1214}
ryo_analysis/ast/registry.rs

ryo_analysis/ast/
registry.rs
