Auto merge of #15677 - Veykril:param-ast-id, r=Veykril

Allocate ast ids for parameters Since these can have attributes attached to them, we'll need this sooner or later (sooner being me tinkering with the token map right now)
rust-lang · Sep 28, 2023 · b3f4574 · b3f4574
2 parents f93b6ac + 2b9dde1
commit b3f4574
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 20 deletions.
diff --git a/crates/hir-def/src/data.rs b/crates/hir-def/src/data.rs
@@ -15,9 +15,7 @@ use crate::{
     attr::Attrs,
     db::DefDatabase,
     expander::{Expander, Mark},
-    item_tree::{
-        self, AssocItem, FnFlags, ItemTree, ItemTreeId, MacroCall, ModItem, Param, TreeId,
-    },
+    item_tree::{self, AssocItem, FnFlags, ItemTree, ItemTreeId, MacroCall, ModItem, TreeId},
     macro_call_as_call_id, macro_id_to_def_id,
     nameres::{
         attr_resolution::ResolvedAttr,
@@ -69,7 +67,7 @@ impl FunctionData {
         let is_varargs = enabled_params
             .clone()
             .next_back()
-            .map_or(false, |param| matches!(item_tree[param], Param::Varargs));
+            .map_or(false, |param| item_tree[param].type_ref.is_none());
 
         let mut flags = func.flags;
         if is_varargs {
@@ -105,10 +103,7 @@ impl FunctionData {
             name: func.name.clone(),
             params: enabled_params
                 .clone()
-                .filter_map(|id| match &item_tree[id] {
-                    Param::Normal(ty) => Some(ty.clone()),
-                    Param::Varargs => None,
-                })
+                .filter_map(|id| item_tree[id].type_ref.clone())
                 .collect(),
             ret_type: func.ret_type.clone(),
             attrs: item_tree.attrs(db, krate, ModItem::from(loc.id.value).into()),

diff --git a/crates/hir-def/src/item_tree.rs b/crates/hir-def/src/item_tree.rs
@@ -613,10 +613,17 @@ pub struct Function {
     pub(crate) flags: FnFlags,
 }
 
-#[derive(Debug, Clone, Eq, PartialEq)]
-pub enum Param {
-    Normal(Interned<TypeRef>),
-    Varargs,
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Param {
+    /// This is [`None`] for varargs
+    pub type_ref: Option<Interned<TypeRef>>,
+    pub ast_id: ParamAstId,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ParamAstId {
+    Param(FileAstId<ast::Param>),
+    SelfParam(FileAstId<ast::SelfParam>),
 }
 
 bitflags::bitflags! {

diff --git a/crates/hir-def/src/item_tree/lower.rs b/crates/hir-def/src/item_tree/lower.rs
@@ -295,8 +295,12 @@ impl<'a> Ctx<'a> {
                         }
                     }
                 };
-                let ty = Interned::new(self_type);
-                let idx = self.data().params.alloc(Param::Normal(ty));
+                let type_ref = Interned::new(self_type);
+                let ast_id = self.source_ast_id_map.ast_id(&self_param);
+                let idx = self.data().params.alloc(Param {
+                    type_ref: Some(type_ref),
+                    ast_id: ParamAstId::SelfParam(ast_id),
+                });
                 self.add_attrs(
                     idx.into(),
                     RawAttrs::new(self.db.upcast(), &self_param, self.hygiene()),
@@ -305,11 +309,19 @@ impl<'a> Ctx<'a> {
             }
             for param in param_list.params() {
                 let idx = match param.dotdotdot_token() {
-                    Some(_) => self.data().params.alloc(Param::Varargs),
+                    Some(_) => {
+                        let ast_id = self.source_ast_id_map.ast_id(&param);
+                        self.data()
+                            .params
+                            .alloc(Param { type_ref: None, ast_id: ParamAstId::Param(ast_id) })
+                    }
                     None => {
                         let type_ref = TypeRef::from_ast_opt(&self.body_ctx, param.ty());
                         let ty = Interned::new(type_ref);
-                        self.data().params.alloc(Param::Normal(ty))
+                        let ast_id = self.source_ast_id_map.ast_id(&param);
+                        self.data()
+                            .params
+                            .alloc(Param { type_ref: Some(ty), ast_id: ParamAstId::Param(ast_id) })
                     }
                 };
                 self.add_attrs(idx.into(), RawAttrs::new(self.db.upcast(), &param, self.hygiene()));

diff --git a/crates/hir-def/src/item_tree/pretty.rs b/crates/hir-def/src/item_tree/pretty.rs
@@ -261,15 +261,15 @@ impl Printer<'_> {
                     self.indented(|this| {
                         for param in params.clone() {
                             this.print_attrs_of(param, "\n");
-                            match &this.tree[param] {
-                                Param::Normal(ty) => {
+                            match &this.tree[param].type_ref {
+                                Some(ty) => {
                                     if flags.contains(FnFlags::HAS_SELF_PARAM) {
                                         w!(this, "self: ");
                                     }
                                     this.print_type_ref(ty);
                                     wln!(this, ",");
                                 }
-                                Param::Varargs => {
+                                None => {
                                     wln!(this, "...");
                                 }
                             };

diff --git a/crates/hir-expand/src/ast_id_map.rs b/crates/hir-expand/src/ast_id_map.rs
@@ -99,7 +99,7 @@ register_ast_id_node! {
         TraitAlias,
         TypeAlias,
         Use,
-    AssocItem, BlockExpr, Variant, RecordField, TupleField, ConstArg
+    AssocItem, BlockExpr, Variant, RecordField, TupleField, ConstArg, Param, SelfParam
 }
 
 /// Maps items' `SyntaxNode`s to `ErasedFileAstId`s and back.