1 files changed, 81 insertions, 39 deletions

diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index ff62928df..3233b1957 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -21,9 +21,10 @@ use hir_expand::{
     HirFileId, InFile,
 };
 use la_arena::{Arena, Idx, RawIdx};
+use profile::Count;
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
-use syntax::{ast, match_ast};
+use syntax::{ast, match_ast, SyntaxKind};
 use test_utils::mark;
 
 use crate::{
@@ -65,22 +66,27 @@ impl GenericParamsId {
 }
 
 /// The item tree of a source file.
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug, Default, Eq, PartialEq)]
 pub struct ItemTree {
+    _c: Count<Self>,
+
     top_level: SmallVec<[ModItem; 1]>,
     attrs: FxHashMap<AttrOwner, RawAttrs>,
-    inner_items: FxHashMap<FileAstId<ast::Item>, SmallVec<[ModItem; 1]>>,
 
     data: Option<Box<ItemTreeData>>,
 }
 
 impl ItemTree {
-    pub fn item_tree_query(db: &dyn DefDatabase, file_id: HirFileId) -> Arc<ItemTree> {
+    pub(crate) fn item_tree_query(db: &dyn DefDatabase, file_id: HirFileId) -> Arc<ItemTree> {
         let _p = profile::span("item_tree_query").detail(|| format!("{:?}", file_id));
         let syntax = if let Some(node) = db.parse_or_expand(file_id) {
+            if node.kind() == SyntaxKind::ERROR {
+                // FIXME: not 100% sure why these crop up, but return an empty tree to avoid a panic
+                return Default::default();
+            }
             node
         } else {
-            return Arc::new(Self::empty());
+            return Default::default();
         };
 
         let hygiene = Hygiene::new(db.upcast(), file_id);
@@ -96,15 +102,17 @@ impl ItemTree {
                     ctx.lower_module_items(&items)
                 },
                 ast::MacroStmts(stmts) => {
-                    ctx.lower_inner_items(stmts.syntax())
+                    // The produced statements can include items, which should be added as top-level
+                    // items.
+                    ctx.lower_macro_stmts(stmts)
                 },
-                // Macros can expand to expressions. We return an empty item tree in this case, but
-                // still need to collect inner items.
                 ast::Expr(e) => {
+                    // Macros can expand to expressions. We return an empty item tree in this case, but
+                    // still need to collect inner items.
                     ctx.lower_inner_items(e.syntax())
                 },
                 _ => {
-                    panic!("cannot create item tree from {:?}", syntax);
+                    panic!("cannot create item tree from {:?} {}", syntax, syntax);
                 },
             }
         };
@@ -116,15 +124,6 @@ impl ItemTree {
         Arc::new(item_tree)
     }
 
-    fn empty() -> Self {
-        Self {
-            top_level: Default::default(),
-            attrs: Default::default(),
-            inner_items: Default::default(),
-            data: Default::default(),
-        }
-    }
-
     fn shrink_to_fit(&mut self) {
         if let Some(data) = &mut self.data {
             let ItemTreeData {
@@ -147,6 +146,8 @@ impl ItemTree {
                 macro_defs,
                 vis,
                 generics,
+                type_refs,
+                inner_items,
             } = &mut **data;
 
             imports.shrink_to_fit();
@@ -169,6 +170,10 @@ impl ItemTree {
 
             vis.arena.shrink_to_fit();
             generics.arena.shrink_to_fit();
+            type_refs.arena.shrink_to_fit();
+            type_refs.map.shrink_to_fit();
+
+            inner_items.shrink_to_fit();
         }
     }
 
@@ -191,16 +196,18 @@ impl ItemTree {
         self.raw_attrs(of).clone().filter(db, krate)
     }
 
-    /// Returns the lowered inner items that `ast` corresponds to.
-    ///
-    /// Most AST items are lowered to a single `ModItem`, but some (eg. `use` items) may be lowered
-    /// to multiple items in the `ItemTree`.
-    pub fn inner_items(&self, ast: FileAstId<ast::Item>) -> &[ModItem] {
-        &self.inner_items[&ast]
+    pub fn all_inner_items(&self) -> impl Iterator<Item = ModItem> + '_ {
+        match &self.data {
+            Some(data) => Some(data.inner_items.values().flatten().copied()).into_iter().flatten(),
+            None => None.into_iter().flatten(),
+        }
     }
 
-    pub fn all_inner_items(&self) -> impl Iterator<Item = ModItem> + '_ {
-        self.inner_items.values().flatten().copied()
+    pub fn inner_items_of_block(&self, block: FileAstId<ast::BlockExpr>) -> &[ModItem] {
+        match &self.data {
+            Some(data) => data.inner_items.get(&block).map(|it| &**it).unwrap_or(&[]),
+            None => &[],
+        }
     }
 
     pub fn source<S: ItemTreeNode>(&self, db: &dyn DefDatabase, of: ItemTreeId<S>) -> S::Source {
@@ -233,7 +240,7 @@ impl ItemVisibilities {
     fn alloc(&mut self, vis: RawVisibility) -> RawVisibilityId {
         match &vis {
             RawVisibility::Public => RawVisibilityId::PUB,
-            RawVisibility::Module(path) if path.segments.is_empty() => match &path.kind {
+            RawVisibility::Module(path) if path.segments().is_empty() => match &path.kind {
                 PathKind::Super(0) => RawVisibilityId::PRIV,
                 PathKind::Crate => RawVisibilityId::PUB_CRATE,
                 _ => RawVisibilityId(self.arena.alloc(vis).into_raw().into()),
@@ -244,10 +251,8 @@ impl ItemVisibilities {
 }
 
 static VIS_PUB: RawVisibility = RawVisibility::Public;
-static VIS_PRIV: RawVisibility =
-    RawVisibility::Module(ModPath { kind: PathKind::Super(0), segments: Vec::new() });
-static VIS_PUB_CRATE: RawVisibility =
-    RawVisibility::Module(ModPath { kind: PathKind::Crate, segments: Vec::new() });
+static VIS_PRIV: RawVisibility = RawVisibility::Module(ModPath::from_kind(PathKind::Super(0)));
+static VIS_PUB_CRATE: RawVisibility = RawVisibility::Module(ModPath::from_kind(PathKind::Crate));
 
 #[derive(Default, Debug, Eq, PartialEq)]
 struct GenericParamsStorage {
@@ -275,6 +280,32 @@ static EMPTY_GENERICS: GenericParams = GenericParams {
     where_predicates: Vec::new(),
 };
 
+/// `TypeRef` interner.
+#[derive(Default, Debug, Eq, PartialEq)]
+struct TypeRefStorage {
+    arena: Arena<Arc<TypeRef>>,
+    map: FxHashMap<Arc<TypeRef>, Idx<Arc<TypeRef>>>,
+}
+
+impl TypeRefStorage {
+    // Note: We lie about the `Idx<TypeRef>` to hide the interner details.
+
+    fn intern(&mut self, ty: TypeRef) -> Idx<TypeRef> {
+        if let Some(id) = self.map.get(&ty) {
+            return Idx::from_raw(id.into_raw());
+        }
+
+        let ty = Arc::new(ty);
+        let idx = self.arena.alloc(ty.clone());
+        self.map.insert(ty, idx);
+        Idx::from_raw(idx.into_raw())
+    }
+
+    fn lookup(&self, id: Idx<TypeRef>) -> &TypeRef {
+        &self.arena[Idx::from_raw(id.into_raw())]
+    }
+}
+
 #[derive(Default, Debug, Eq, PartialEq)]
 struct ItemTreeData {
     imports: Arena<Import>,
@@ -297,6 +328,9 @@ struct ItemTreeData {
 
     vis: ItemVisibilities,
     generics: GenericParamsStorage,
+    type_refs: TypeRefStorage,
+
+    inner_items: FxHashMap<FileAstId<ast::BlockExpr>, SmallVec<[ModItem; 1]>>,
 }
 
 #[derive(Debug, Eq, PartialEq, Hash)]
@@ -483,6 +517,14 @@ impl Index<GenericParamsId> for ItemTree {
     }
 }
 
+impl Index<Idx<TypeRef>> for ItemTree {
+    type Output = TypeRef;
+
+    fn index(&self, id: Idx<TypeRef>) -> &Self::Output {
+        self.data().type_refs.lookup(id)
+    }
+}
+
 impl<N: ItemTreeNode> Index<FileItemTreeId<N>> for ItemTree {
     type Output = N;
     fn index(&self, id: FileItemTreeId<N>) -> &N {
@@ -526,9 +568,9 @@ pub struct Function {
     /// Whether the function is located in an `extern` block (*not* whether it is an
     /// `extern "abi" fn`).
     pub is_extern: bool,
-    pub params: Box<[TypeRef]>,
+    pub params: Box<[Idx<TypeRef>]>,
     pub is_varargs: bool,
-    pub ret_type: TypeRef,
+    pub ret_type: Idx<TypeRef>,
     pub ast_id: FileAstId<ast::Fn>,
 }
 
@@ -575,7 +617,7 @@ pub struct Const {
     /// const _: () = ();
     pub name: Option<Name>,
     pub visibility: RawVisibilityId,
-    pub type_ref: TypeRef,
+    pub type_ref: Idx<TypeRef>,
     pub ast_id: FileAstId<ast::Const>,
 }
 
@@ -586,7 +628,7 @@ pub struct Static {
     pub mutable: bool,
     /// Whether the static is in an `extern` block.
     pub is_extern: bool,
-    pub type_ref: TypeRef,
+    pub type_ref: Idx<TypeRef>,
     pub ast_id: FileAstId<ast::Static>,
 }
 
@@ -603,8 +645,8 @@ pub struct Trait {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Impl {
     pub generic_params: GenericParamsId,
-    pub target_trait: Option<TypeRef>,
-    pub target_type: TypeRef,
+    pub target_trait: Option<Idx<TypeRef>>,
+    pub target_type: Idx<TypeRef>,
     pub is_negative: bool,
     pub items: Box<[AssocItem]>,
     pub ast_id: FileAstId<ast::Impl>,
@@ -617,7 +659,7 @@ pub struct TypeAlias {
     /// Bounds on the type alias itself. Only valid in trait declarations, eg. `type Assoc: Copy;`.
     pub bounds: Box<[TypeBound]>,
     pub generic_params: GenericParamsId,
-    pub type_ref: Option<TypeRef>,
+    pub type_ref: Option<Idx<TypeRef>>,
     pub is_extern: bool,
     pub ast_id: FileAstId<ast::TypeAlias>,
 }
@@ -800,6 +842,6 @@ pub enum Fields {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Field {
     pub name: Name,
-    pub type_ref: TypeRef,
+    pub type_ref: Idx<TypeRef>,
     pub visibility: RawVisibilityId,
 }