aboutsummaryrefslogtreecommitdiff
path: root/crates/ra_hir/src/nameres/collector.rs
diff options
context:
space:
mode:
Diffstat (limited to 'crates/ra_hir/src/nameres/collector.rs')
-rw-r--r--crates/ra_hir/src/nameres/collector.rs165
1 files changed, 154 insertions, 11 deletions
diff --git a/crates/ra_hir/src/nameres/collector.rs b/crates/ra_hir/src/nameres/collector.rs
index 6147b3219..4590a5184 100644
--- a/crates/ra_hir/src/nameres/collector.rs
+++ b/crates/ra_hir/src/nameres/collector.rs
@@ -42,12 +42,40 @@ pub(super) fn collect_defs(db: &impl DefDatabase, mut def_map: CrateDefMap) -> C
42 unresolved_imports: Vec::new(), 42 unresolved_imports: Vec::new(),
43 unexpanded_macros: Vec::new(), 43 unexpanded_macros: Vec::new(),
44 global_macro_scope: FxHashMap::default(), 44 global_macro_scope: FxHashMap::default(),
45 marco_stack_count: 0, 45 macro_stack_monitor: MacroStackMonitor::default(),
46 }; 46 };
47 collector.collect(); 47 collector.collect();
48 collector.finish() 48 collector.finish()
49} 49}
50 50
51#[derive(Default)]
52struct MacroStackMonitor {
53 counts: FxHashMap<MacroDefId, u32>,
54
55 /// Mainly use for test
56 validator: Option<Box<dyn Fn(u32) -> bool>>,
57}
58
59impl MacroStackMonitor {
60 fn increase(&mut self, macro_def_id: MacroDefId) {
61 *self.counts.entry(macro_def_id).or_default() += 1;
62 }
63
64 fn decrease(&mut self, macro_def_id: MacroDefId) {
65 *self.counts.entry(macro_def_id).or_default() -= 1;
66 }
67
68 fn is_poison(&self, macro_def_id: MacroDefId) -> bool {
69 let cur = *self.counts.get(&macro_def_id).unwrap_or(&0);
70
71 if let Some(validator) = &self.validator {
72 validator(cur)
73 } else {
74 cur > 100
75 }
76 }
77}
78
51/// Walks the tree of module recursively 79/// Walks the tree of module recursively
52struct DefCollector<DB> { 80struct DefCollector<DB> {
53 db: DB, 81 db: DB,
@@ -59,7 +87,7 @@ struct DefCollector<DB> {
59 87
60 /// Some macro use `$tt:tt which mean we have to handle the macro perfectly 88 /// Some macro use `$tt:tt which mean we have to handle the macro perfectly
61 /// To prevent stackoverflow, we add a deep counter here for prevent that. 89 /// To prevent stackoverflow, we add a deep counter here for prevent that.
62 marco_stack_count: u32, 90 macro_stack_monitor: MacroStackMonitor,
63} 91}
64 92
65impl<'a, DB> DefCollector<&'a DB> 93impl<'a, DB> DefCollector<&'a DB>
@@ -317,30 +345,40 @@ where
317 let def_map = self.db.crate_def_map(krate); 345 let def_map = self.db.crate_def_map(krate);
318 if let Some(macro_id) = def_map.public_macros.get(&path.segments[1].name).cloned() { 346 if let Some(macro_id) = def_map.public_macros.get(&path.segments[1].name).cloned() {
319 let call_id = MacroCallLoc { def: macro_id, ast_id: *ast_id }.id(self.db); 347 let call_id = MacroCallLoc { def: macro_id, ast_id: *ast_id }.id(self.db);
320 resolved.push((*module_id, call_id)); 348 resolved.push((*module_id, call_id, macro_id));
321 } 349 }
322 false 350 false
323 }); 351 });
324 352
325 for (module_id, macro_call_id) in resolved { 353 for (module_id, macro_call_id, macro_def_id) in resolved {
326 self.collect_macro_expansion(module_id, macro_call_id); 354 self.collect_macro_expansion(module_id, macro_call_id, macro_def_id);
327 } 355 }
328 res 356 res
329 } 357 }
330 358
331 fn collect_macro_expansion(&mut self, module_id: CrateModuleId, macro_call_id: MacroCallId) { 359 fn collect_macro_expansion(
332 self.marco_stack_count += 1; 360 &mut self,
361 module_id: CrateModuleId,
362 macro_call_id: MacroCallId,
363 macro_def_id: MacroDefId,
364 ) {
365 if self.def_map.poison_macros.contains(&macro_def_id) {
366 return;
367 }
368
369 self.macro_stack_monitor.increase(macro_def_id);
333 370
334 if self.marco_stack_count < 300 { 371 if !self.macro_stack_monitor.is_poison(macro_def_id) {
335 let file_id: HirFileId = macro_call_id.into(); 372 let file_id: HirFileId = macro_call_id.into();
336 let raw_items = self.db.raw_items(file_id); 373 let raw_items = self.db.raw_items(file_id);
337 ModCollector { def_collector: &mut *self, file_id, module_id, raw_items: &raw_items } 374 ModCollector { def_collector: &mut *self, file_id, module_id, raw_items: &raw_items }
338 .collect(raw_items.items()) 375 .collect(raw_items.items());
339 } else { 376 } else {
340 log::error!("Too deep macro expansion: {}", macro_call_id.debug_dump(self.db)); 377 log::error!("Too deep macro expansion: {}", macro_call_id.debug_dump(self.db));
378 self.def_map.poison_macros.insert(macro_def_id);
341 } 379 }
342 380
343 self.marco_stack_count -= 1; 381 self.macro_stack_monitor.decrease(macro_def_id);
344 } 382 }
345 383
346 fn finish(self) -> CrateDefMap { 384 fn finish(self) -> CrateDefMap {
@@ -484,7 +522,7 @@ where
484 { 522 {
485 let macro_call_id = MacroCallLoc { def: macro_id, ast_id }.id(self.def_collector.db); 523 let macro_call_id = MacroCallLoc { def: macro_id, ast_id }.id(self.def_collector.db);
486 524
487 self.def_collector.collect_macro_expansion(self.module_id, macro_call_id); 525 self.def_collector.collect_macro_expansion(self.module_id, macro_call_id, macro_id);
488 return; 526 return;
489 } 527 }
490 528
@@ -530,3 +568,108 @@ fn resolve_submodule(
530 None => Err(if is_dir_owner { file_mod } else { file_dir_mod }), 568 None => Err(if is_dir_owner { file_mod } else { file_dir_mod }),
531 } 569 }
532} 570}
571
572#[cfg(test)]
573mod tests {
574 use ra_db::SourceDatabase;
575
576 use crate::{Crate, mock::MockDatabase, DefDatabase};
577 use ra_arena::{Arena};
578 use super::*;
579 use rustc_hash::FxHashSet;
580
581 fn do_collect_defs(
582 db: &impl DefDatabase,
583 def_map: CrateDefMap,
584 monitor: MacroStackMonitor,
585 ) -> CrateDefMap {
586 let mut collector = DefCollector {
587 db,
588 def_map,
589 glob_imports: FxHashMap::default(),
590 unresolved_imports: Vec::new(),
591 unexpanded_macros: Vec::new(),
592 global_macro_scope: FxHashMap::default(),
593 macro_stack_monitor: monitor,
594 };
595 collector.collect();
596 collector.finish()
597 }
598
599 fn do_limited_resolve(code: &str, limit: u32, poison_limit: u32) -> CrateDefMap {
600 let (db, _source_root, _) = MockDatabase::with_single_file(&code);
601 let crate_id = db.crate_graph().iter().next().unwrap();
602 let krate = Crate { crate_id };
603
604 let def_map = {
605 let edition = krate.edition(&db);
606 let mut modules: Arena<CrateModuleId, ModuleData> = Arena::default();
607 let root = modules.alloc(ModuleData::default());
608 CrateDefMap {
609 krate,
610 edition,
611 extern_prelude: FxHashMap::default(),
612 prelude: None,
613 root,
614 modules,
615 public_macros: FxHashMap::default(),
616 poison_macros: FxHashSet::default(),
617 diagnostics: Vec::new(),
618 }
619 };
620
621 let mut monitor = MacroStackMonitor::default();
622 monitor.validator = Some(Box::new(move |count| {
623 assert!(count < limit);
624 count >= poison_limit
625 }));
626
627 do_collect_defs(&db, def_map, monitor)
628 }
629
630 #[test]
631 fn test_macro_expand_limit_width() {
632 do_limited_resolve(
633 r#"
634 macro_rules! foo {
635 ($($ty:ty)*) => { foo!($($ty)*, $($ty)*); }
636 }
637foo!(KABOOM);
638 "#,
639 16,
640 1000,
641 );
642 }
643
644 #[test]
645 fn test_macro_expand_poisoned() {
646 let def = do_limited_resolve(
647 r#"
648 macro_rules! foo {
649 ($ty:ty) => { foo!($ty); }
650 }
651foo!(KABOOM);
652 "#,
653 100,
654 16,
655 );
656
657 assert_eq!(def.poison_macros.len(), 1);
658 }
659
660 #[test]
661 fn test_macro_expand_normal() {
662 let def = do_limited_resolve(
663 r#"
664 macro_rules! foo {
665 ($ident:ident) => { struct $ident {} }
666 }
667foo!(Bar);
668 "#,
669 16,
670 16,
671 );
672
673 assert_eq!(def.poison_macros.len(), 0);
674 }
675}