From 6f269708e819cc01aad24be0b7456875d1a4041f Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Mon, 12 Jul 2021 21:13:43 +0300
Subject: [PATCH] internal: get rid of a call to slow O(N) visibility_of
 function

Instead of inferring module's declared visibility by looking at the
scope of its parent, let's just remeber the declared visibility in the
DefMap.
---
 crates/hir/src/lib.rs                   | 10 +++++++++-
 crates/hir_def/src/nameres.rs           | 24 ++++++++++++++++++------
 crates/hir_def/src/nameres/collector.rs |  2 +-
 crates/ide_db/src/defs.rs               |  6 +-----
 4 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index f601fd6276..9c74d8c5f7 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -430,8 +430,16 @@ impl Module {
             .collect()
     }
 
+    pub fn visibility(self, db: &dyn HirDatabase) -> Visibility {
+        let def_map = self.id.def_map(db.upcast());
+        let module_data = &def_map[self.id.local_id];
+        module_data.visibility
+    }
+
     pub fn visibility_of(self, db: &dyn HirDatabase, def: &ModuleDef) -> Option<Visibility> {
-        self.id.def_map(db.upcast())[self.id.local_id].scope.visibility_of((*def).into())
+        let def_map = self.id.def_map(db.upcast());
+        let module_data = &def_map[self.id.local_id];
+        module_data.scope.visibility_of((*def).into())
     }
 
     pub fn diagnostics(self, db: &dyn HirDatabase, acc: &mut Vec<AnyDiagnostic>) {
diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index c5f8d45d68..cd314c8c9a 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -72,6 +72,7 @@ use crate::{
     nameres::{diagnostics::DefDiagnostic, path_resolution::ResolveMode},
     path::ModPath,
     per_ns::PerNs,
+    visibility::Visibility,
     AstId, BlockId, BlockLoc, LocalModuleId, ModuleDefId, ModuleId,
 };
 
@@ -192,12 +193,14 @@ impl ModuleOrigin {
 
 #[derive(Debug, PartialEq, Eq)]
 pub struct ModuleData {
+    /// Where does this module come from?
+    pub origin: ModuleOrigin,
+    /// Declared visibility of this module.
+    pub visibility: Visibility,
+
     pub parent: Option<LocalModuleId>,
     pub children: FxHashMap<Name, LocalModuleId>,
     pub scope: ItemScope,
-
-    /// Where does this module come from?
-    pub origin: ModuleOrigin,
 }
 
 impl DefMap {
@@ -243,7 +246,15 @@ impl DefMap {
 
     fn empty(krate: CrateId, edition: Edition, root_module_origin: ModuleOrigin) -> DefMap {
         let mut modules: Arena<ModuleData> = Arena::default();
-        let root = modules.alloc(ModuleData::new(root_module_origin));
+
+        let local_id = LocalModuleId::from_raw(la_arena::RawIdx::from(0));
+        // NB: we use `None` as block here, which would be wrong for implicit
+        // modules declared by blocks with items. At the moment, we don't use
+        // this visibility for anything outside IDE, so that's probably OK.
+        let visibility = Visibility::Module(ModuleId { krate, local_id, block: None });
+        let root = modules.alloc(ModuleData::new(root_module_origin, visibility));
+        assert_eq!(local_id, root);
+
         DefMap {
             _c: Count::new(),
             block: None,
@@ -448,12 +459,13 @@ impl DefMap {
 }
 
 impl ModuleData {
-    pub(crate) fn new(origin: ModuleOrigin) -> Self {
+    pub(crate) fn new(origin: ModuleOrigin, visibility: Visibility) -> Self {
         ModuleData {
+            origin,
+            visibility,
             parent: None,
             children: FxHashMap::default(),
             scope: ItemScope::default(),
-            origin,
         }
     }
 
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index dd31365a80..79452bba54 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -1629,7 +1629,7 @@ impl ModCollector<'_, '_> {
                 ModuleOrigin::File { declaration, definition, is_mod_rs }
             }
         };
-        let res = modules.alloc(ModuleData::new(origin));
+        let res = modules.alloc(ModuleData::new(origin, vis));
         modules[res].parent = Some(self.module_id);
         for (name, mac) in modules[self.module_id].scope.collect_legacy_macros() {
             modules[res].scope.define_legacy_macro(name, mac)
diff --git a/crates/ide_db/src/defs.rs b/crates/ide_db/src/defs.rs
index 198b568ad8..b638e35112 100644
--- a/crates/ide_db/src/defs.rs
+++ b/crates/ide_db/src/defs.rs
@@ -45,11 +45,7 @@ impl Definition {
         match self {
             Definition::Field(sf) => Some(sf.visibility(db)),
             Definition::ModuleDef(def) => match def {
-                ModuleDef::Module(it) => {
-                    // FIXME: should work like other cases here.
-                    let parent = it.parent(db)?;
-                    parent.visibility_of(db, def)
-                }
+                ModuleDef::Module(it) => Some(it.visibility(db)),
                 ModuleDef::Function(it) => Some(it.visibility(db)),
                 ModuleDef::Adt(it) => Some(it.visibility(db)),
                 ModuleDef::Const(it) => Some(it.visibility(db)),