Stop expanding UseTrees during ItemTree lowering

Author: jonas-schievink

crates/hir/src/lib.rs

@@ -472,27 +472,13 @@ impl Module {
                     });
                 }
 
-                DefDiagnosticKind::UnresolvedImport { ast, index } => {
-                    let use_item = ast.to_node(db.upcast());
-                    let hygiene = Hygiene::new(db.upcast(), ast.file_id);
-                    let mut cur = 0;
-                    let mut tree = None;
-                    ModPath::expand_use_item(
-                        db.upcast(),
-                        InFile::new(ast.file_id, use_item),
-                        &hygiene,
-                        |_mod_path, use_tree, _is_glob, _alias| {
-                            if cur == *index {
-                                tree = Some(use_tree.clone());
-                            }
-
-                            cur += 1;
-                        },
-                    );
+                DefDiagnosticKind::UnresolvedImport { id, index } => {
+                    let file_id = id.file_id();
+                    let item_tree = id.item_tree(db.upcast());
+                    let import = &item_tree[id.value];
 
-                    if let Some(tree) = tree {
-                        sink.push(UnresolvedImport { file: ast.file_id, node: AstPtr::new(&tree) });
-                    }
+                    let use_tree = import.use_tree_to_ast(db.upcast(), file_id, *index);
+                    sink.push(UnresolvedImport { file: file_id, node: AstPtr::new(&use_tree) });
                 }
 
                 DefDiagnosticKind::UnconfiguredCode { ast, cfg, opts } => {

crates/hir_def/src/item_tree.rs

@@ -523,21 +523,38 @@ impl<N: ItemTreeNode> Index<FileItemTreeId<N>> for ItemTree {
     }
 }
 
-/// A desugared `use` import.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Import {
-    pub path: Interned<ModPath>,
-    pub alias: Option<ImportAlias>,
     pub visibility: RawVisibilityId,
-    pub is_glob: bool,
-    /// AST ID of the `use` item this import was derived from. Note that many `Import`s can map to
-    /// the same `use` item.
     pub ast_id: FileAstId<ast::Use>,
-    /// Index of this `Import` when the containing `Use` is visited via `ModPath::expand_use_item`.
-    ///
-    /// This can be used to get the `UseTree` this `Import` corresponds to and allows emitting
-    /// precise diagnostics.
-    pub index: usize,
+    pub use_tree: UseTree,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct UseTree {
+    pub index: Idx<ast::UseTree>,
+    kind: UseTreeKind,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum UseTreeKind {
+    /// ```ignore
+    /// use path::to::Item;
+    /// use path::to::Item as Renamed;
+    /// use path::to::Trait as _;
+    /// ```
+    Single { path: ModPath, alias: Option<ImportAlias> },
+
+    /// ```ignore
+    /// use *;  // (invalid, but can occur in nested tree)
+    /// use path::*;
+    /// ```
+    Glob { path: Option<ModPath> },
+
+    /// ```ignore
+    /// use prefix::{self, Item, ...};
+    /// ```
+    Prefixed { prefix: Option<ModPath>, list: Vec<UseTree> },
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
@@ -711,6 +728,97 @@ pub struct MacroDef {
     pub ast_id: FileAstId<ast::MacroDef>,
 }
 
+impl Import {
+    /// Maps a `UseTree` contained in this import back to its AST node.
+    pub fn use_tree_to_ast(
+        &self,
+        db: &dyn DefDatabase,
+        file_id: HirFileId,
+        index: Idx<ast::UseTree>,
+    ) -> ast::UseTree {
+        // Re-lower the AST item and get the source map.
+        // Note: The AST unwraps are fine, since if they fail we should have never obtained `index`.
+        let ast = InFile::new(file_id, self.ast_id).to_node(db.upcast());
+        let ast_use_tree = ast.use_tree().expect("missing `use_tree`");
+        let hygiene = Hygiene::new(db.upcast(), file_id);
+        let (_, source_map) =
+            lower::lower_use_tree(db, &hygiene, ast_use_tree).expect("failed to lower use tree");
+        source_map[index].clone()
+    }
+}
+
+impl UseTree {
+    /// Expands the `UseTree` into individually imported `ModPath`s.
+    pub fn expand(
+        &self,
+        mut cb: impl FnMut(Idx<ast::UseTree>, ModPath, /* is_glob */ bool, Option<ImportAlias>),
+    ) {
+        self.expand_impl(None, &mut cb)
+    }
+
+    fn expand_impl(
+        &self,
+        prefix: Option<ModPath>,
+        cb: &mut dyn FnMut(
+            Idx<ast::UseTree>,
+            ModPath,
+            /* is_glob */ bool,
+            Option<ImportAlias>,
+        ),
+    ) {
+        fn concat_mod_paths(prefix: Option<ModPath>, path: &ModPath) -> Option<ModPath> {
+            match (prefix, &path.kind) {
+                (None, _) => Some(path.clone()),
+                (Some(mut prefix), PathKind::Plain) => {
+                    for segment in path.segments() {
+                        prefix.push_segment(segment.clone());
+                    }
+                    Some(prefix)
+                }
+                (Some(prefix), PathKind::Super(0)) => {
+                    // `some::path::self` == `some::path`
+                    if path.segments().is_empty() {
+                        Some(prefix)
+                    } else {
+                        None
+                    }
+                }
+                (Some(_), _) => None,
+            }
+        }
+
+        match &self.kind {
+            UseTreeKind::Single { path, alias } => {
+                if let Some(path) = concat_mod_paths(prefix, path) {
+                    cb(self.index, path, false, alias.clone());
+                }
+            }
+            UseTreeKind::Glob { path: Some(path) } => {
+                if let Some(path) = concat_mod_paths(prefix, path) {
+                    cb(self.index, path, true, None);
+                }
+            }
+            UseTreeKind::Glob { path: None } => {
+                if let Some(prefix) = prefix {
+                    cb(self.index, prefix, true, None);
+                }
+            }
+            UseTreeKind::Prefixed { prefix: additional_prefix, list } => {
+                let prefix = match additional_prefix {
+                    Some(path) => match concat_mod_paths(prefix, path) {
+                        Some(path) => Some(path),
+                        None => return,
+                    },
+                    None => prefix,
+                };
+                for tree in list {
+                    tree.expand_impl(prefix.clone(), cb);
+                }
+            }
+        }
+    }
+}
+
 macro_rules! impl_froms {
     ($e:ident { $($v:ident ($t:ty)),* $(,)? }) => {
         $(

crates/hir_def/src/item_tree/lower.rs

@@ -35,7 +35,6 @@ pub(super) struct Ctx<'a> {
     db: &'a dyn DefDatabase,
     tree: ItemTree,
     hygiene: Hygiene,
-    file: HirFileId,
     source_ast_id_map: Arc<AstIdMap>,
     body_ctx: crate::body::LowerCtx<'a>,
     forced_visibility: Option<RawVisibilityId>,
@@ -47,7 +46,6 @@ impl<'a> Ctx<'a> {
             db,
             tree: ItemTree::default(),
             hygiene,
-            file,
             source_ast_id_map: db.ast_id_map(file),
             body_ctx: crate::body::LowerCtx::new(db, file),
             forced_visibility: None,
@@ -561,30 +559,13 @@ impl<'a> Ctx<'a> {
         Some(id(self.data().impls.alloc(res)))
     }
 
-    fn lower_use(&mut self, use_item: &ast::Use) -> Vec<FileItemTreeId<Import>> {
+    fn lower_use(&mut self, use_item: &ast::Use) -> Option<FileItemTreeId<Import>> {
         let visibility = self.lower_visibility(use_item);
         let ast_id = self.source_ast_id_map.ast_id(use_item);
+        let (use_tree, _) = lower_use_tree(self.db, &self.hygiene, use_item.use_tree()?)?;
 
-        // Every use item can expand to many `Import`s.
-        let mut imports = Vec::new();
-        let tree = self.tree.data_mut();
-        ModPath::expand_use_item(
-            self.db,
-            InFile::new(self.file, use_item.clone()),
-            &self.hygiene,
-            |path, _use_tree, is_glob, alias| {
-                imports.push(id(tree.imports.alloc(Import {
-                    path: Interned::new(path),
-                    alias,
-                    visibility,
-                    is_glob,
-                    ast_id,
-                    index: imports.len(),
-                })));
-            },
-        );
-
-        imports
+        let res = Import { visibility, ast_id, use_tree };
+        Some(id(self.data().imports.alloc(res)))
     }
 
     fn lower_extern_crate(
@@ -884,3 +865,76 @@ fn lower_abi(abi: ast::Abi) -> Interned<str> {
         }
     }
 }
+
+struct UseTreeLowering<'a> {
+    db: &'a dyn DefDatabase,
+    hygiene: &'a Hygiene,
+    mapping: Arena<ast::UseTree>,
+}
+
+impl UseTreeLowering<'_> {
+    fn lower_use_tree(&mut self, tree: ast::UseTree) -> Option<UseTree> {
+        if let Some(use_tree_list) = tree.use_tree_list() {
+            let prefix = match tree.path() {
+                // E.g. use something::{{{inner}}};
+                None => None,
+                // E.g. `use something::{inner}` (prefix is `None`, path is `something`)
+                // or `use something::{path::{inner::{innerer}}}` (prefix is `something::path`, path is `inner`)
+                Some(path) => {
+                    match ModPath::from_src(self.db, path, &self.hygiene) {
+                        Some(it) => Some(it),
+                        None => return None, // FIXME: report errors somewhere
+                    }
+                }
+            };
+
+            let list =
+                use_tree_list.use_trees().filter_map(|tree| self.lower_use_tree(tree)).collect();
+
+            Some(self.use_tree(UseTreeKind::Prefixed { prefix, list }, tree))
+        } else {
+            let is_glob = tree.star_token().is_some();
+            let path = match tree.path() {
+                Some(path) => Some(ModPath::from_src(self.db, path, &self.hygiene)?),
+                None => None,
+            };
+            let alias = tree.rename().map(|a| {
+                a.name().map(|it| it.as_name()).map_or(ImportAlias::Underscore, ImportAlias::Alias)
+            });
+            if alias.is_some() && is_glob {
+                return None;
+            }
+
+            match (path, alias, is_glob) {
+                (path, None, true) => {
+                    if path.is_none() {
+                        cov_mark::hit!(glob_enum_group);
+                    }
+                    Some(self.use_tree(UseTreeKind::Glob { path }, tree))
+                }
+                // Globs can't be renamed
+                (_, Some(_), true) | (None, None, false) => None,
+                // `bla::{ as Name}` is invalid
+                (None, Some(_), false) => None,
+                (Some(path), alias, false) => {
+                    Some(self.use_tree(UseTreeKind::Single { path, alias }, tree))
+                }
+            }
+        }
+    }
+
+    fn use_tree(&mut self, kind: UseTreeKind, ast: ast::UseTree) -> UseTree {
+        let index = self.mapping.alloc(ast);
+        UseTree { index, kind }
+    }
+}
+
+pub(super) fn lower_use_tree(
+    db: &dyn DefDatabase,
+    hygiene: &Hygiene,
+    tree: ast::UseTree,
+) -> Option<(UseTree, Arena<ast::UseTree>)> {
+    let mut lowering = UseTreeLowering { db, hygiene, mapping: Arena::new() };
+    let tree = lowering.lower_use_tree(tree)?;
+    Some((tree, lowering.mapping))
+}

crates/hir_def/src/item_tree/pretty.rs

@@ -163,21 +163,46 @@ impl<'a> Printer<'a> {
         }
     }
 
+    fn print_use_tree(&mut self, use_tree: &UseTree) {
+        match &use_tree.kind {
+            UseTreeKind::Single { path, alias } => {
+                w!(self, "{}", path);
+                if let Some(alias) = alias {
+                    w!(self, " as {}", alias);
+                }
+            }
+            UseTreeKind::Glob { path } => {
+                if let Some(path) = path {
+                    w!(self, "{}::", path);
+                }
+                w!(self, "*");
+            }
+            UseTreeKind::Prefixed { prefix, list } => {
+                if let Some(prefix) = prefix {
+                    w!(self, "{}::", prefix);
+                }
+                w!(self, "{{");
+                for (i, tree) in list.iter().enumerate() {
+                    if i != 0 {
+                        w!(self, ", ");
+                    }
+                    self.print_use_tree(tree);
+                }
+                w!(self, "}}");
+            }
+        }
+    }
+
     fn print_mod_item(&mut self, item: ModItem) {
         self.print_attrs_of(item);
 
         match item {
             ModItem::Import(it) => {
-                let Import { visibility, path, is_glob, alias, ast_id: _, index } = &self.tree[it];
+                let Import { visibility, use_tree, ast_id: _ } = &self.tree[it];
                 self.print_visibility(*visibility);
-                w!(self, "use {}", path);
-                if *is_glob {
-                    w!(self, "::*");
-                }
-                if let Some(alias) = alias {
-                    w!(self, " as {}", alias);
-                }
-                wln!(self, ";  // {}", index);
+                w!(self, "use ");
+                self.print_use_tree(use_tree);
+                wln!(self, ";");
             }
             ModItem::ExternCrate(it) => {
                 let ExternCrate { name, alias, visibility, ast_id: _ } = &self.tree[it];