feat: be a bit smarter about hir lowering, store hir ids in the eleme…

…nts they point to
dnbln · Feb 7, 2025 · 56dc1fd · 56dc1fd
1 parent 0c4d019
commit 56dc1fd
Show file tree

Hide file tree

Showing 10 changed files with 477 additions and 528 deletions.
diff --git a/crates/narxia-hir-typechk/src/sema/mod.rs b/crates/narxia-hir-typechk/src/sema/mod.rs
@@ -160,15 +160,15 @@ struct ProgramStructureVisitor<'hir> {
 }
 
 macro_rules! scope_creating_elements {
-    ($($vis_name:ident ($id:ty, $t:ty) => $walk_name:ident $($id_use:expr)? ;)*) => {
+    ($($vis_name:ident ($t:ty) => $walk_name:ident;)*) => {
         $(
-            fn $vis_name(&mut self, id: $id, t: &'hir $t) {
+            fn $vis_name(&mut self, t: &'hir $t) {
                 let last_scope = self.stack.last().copied();
-                let self_elem = self.program_structure.scope_tree.push_scope_element(None, scope_elem(id));
+                let self_elem = self.program_structure.scope_tree.push_scope_element(None, scope_elem(t.hir_id));
                 let scope = self.program_structure.scope_tree.push_scope(last_scope, Some(self_elem));
                 self.stack.push(scope);
 
-                vis::$walk_name(self, $(if $id_use == () {id}else{id},)? t);
+                vis::$walk_name(self, t);
 
                 let s = self.stack.pop();
 
@@ -179,13 +179,13 @@ macro_rules! scope_creating_elements {
 }
 
 macro_rules! scope_adding_elements {
-    ($($vis_name:ident ($id:ty, $t:ty) => $walk_name:ident $($id_use:expr)? ;)*) => {
+    ($($vis_name:ident ($t:ty) => $walk_name:ident;)*) => {
         $(
-            fn $vis_name(&mut self, id: $id, t: &'hir $t) {
+            fn $vis_name(&mut self, t: &'hir $t) {
                 let last_scope = self.stack.last().copied();
-                let _self_elem = self.program_structure.scope_tree.push_scope_element(last_scope, scope_elem(id));
+                let _self_elem = self.program_structure.scope_tree.push_scope_element(last_scope, scope_elem(t.hir_id));
 
-                vis::$walk_name(self, $(if $id_use == () {id}else{id},)? t);
+                vis::$walk_name(self, t);
             }
         )*
     };
@@ -197,13 +197,13 @@ impl<'hir> vis::HirVisitor<'hir> for ProgramStructureVisitor<'hir> {
     }
 
     scope_creating_elements! {
-        visit_mod_def(hir::ModId, hir::ModDef) => walk_mod_def;
-        visit_fn_def(hir::FnId, hir::FnDef) => walk_fn_def;
-        visit_block(hir::BlockId, hir::Block) => walk_block;
+        visit_mod_def(hir::ModDef) => walk_mod_def;
+        visit_fn_def(hir::FnDef) => walk_fn_def;
+        visit_block(hir::Block) => walk_block;
     }
 
     scope_adding_elements! {
-        visit_use_stmt(hir::UseStmtId, hir::UseStmt) => walk_use_stmt;
+        visit_use_stmt(hir::UseStmt) => walk_use_stmt;
     }
 }
 

diff --git a/crates/narxia-hir/src/hir/hir_debug.rs b/crates/narxia-hir/src/hir/hir_debug.rs
@@ -179,7 +179,7 @@ pub fn display_mod_def(
     write!(f, "{:indent$}", "", indent = hdc.depth)?;
     write!(f, "{} {}", "module".keyword(), mod_def.name.text,)?;
     if let Some(body) = &mod_def.body {
-        writeln!(f, "{}", "{".punctuation());
+        writeln!(f, "{}", "{".punctuation())?;
         display_item_list(f, &body.items, hdc.make_child())?;
         writeln!(f)?;
         write!(f, "{:indent$}", "", indent = hdc.depth)?;
@@ -614,7 +614,7 @@ fn display_ty(f: &mut fmt::Formatter, ty: &TyRef, hdc: HirDisplayContext) -> fmt
                         write!(f, "{} ", ",".punctuation())?;
                     }
 
-                    display_ty_generic_arg(f, arg, hdc.make_child())?;
+                    display_ty_generic_arg_id(f, *arg, hdc.make_child())?;
                 }
                 write!(f, "{}", ">".punctuation())?;
             }
@@ -649,6 +649,16 @@ impl fmt::Display for TyRef {
     }
 }
 
+fn display_ty_generic_arg_id(
+    f: &mut fmt::Formatter,
+    id: TyGenericArgId,
+    hdc: HirDisplayContext,
+) -> fmt::Result {
+    write!(f, "{}", id.0)?;
+
+    Ok(())
+}
+
 fn display_ty_generic_arg(
     f: &mut fmt::Formatter,
     arg: &TyGenericArg,

diff --git a/crates/narxia-hir/src/hir/mod.rs b/crates/narxia-hir/src/hir/mod.rs
@@ -14,10 +14,19 @@ mod hir_debug;
 pub use hir_debug::*;
 
 pub trait HirIdNewtype {
+    fn new(hir_id: HirId) -> Self
+    where
+        Self: Sized;
     fn hir_id(&self) -> HirId;
 }
 
 impl HirIdNewtype for HirId {
+    fn new(hir_id: HirId) -> Self
+    where
+        Self: Sized,
+    {
+        hir_id
+    }
     fn hir_id(&self) -> HirId {
         *self
     }
@@ -29,6 +38,13 @@ macro_rules! hir_id_newtype {
         pub struct $name(pub HirId);
 
         impl HirIdNewtype for $name {
+            fn new(hir_id: HirId) -> Self
+            where
+                Self: Sized,
+            {
+                Self(hir_id)
+            }
+
             fn hir_id(&self) -> HirId {
                 self.0
             }
@@ -104,6 +120,8 @@ pub struct ModDef {
     pub mod_kw: ModuleKw,
     pub name: Ident,
     pub body: Option<ModBody>,
+
+    pub hir_id: ModId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -117,6 +135,7 @@ pub struct ModBody {
 pub struct Item {
     pub attrs: AttrList,
     pub kind: ItemKind,
+    pub hir_id: ItemId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -192,6 +211,7 @@ pub struct UseStmt {
     pub use_kw: UseKw,
     pub span: HirSpan,
     pub path: UsePath,
+    pub hir_id: UseStmtId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -211,6 +231,7 @@ hir_id_newtype!(UsePathSegmentId, UsePathSegment);
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct UsePathSegment {
     pub ident: Ident,
+    pub hir_id: UsePathSegmentId,
 }
 
 #[derive(Eq, PartialEq, PartialOrd, Ord, Clone)]
@@ -256,6 +277,8 @@ pub struct FnDef {
     pub params: Option<FnParamList>,
     pub ret_ty: Option<FnRetTy>,
     pub body: BlockId,
+
+    pub hir_id: FnId,
 }
 
 hir_id_newtype!(BlockId, Block);
@@ -330,6 +353,7 @@ hir_id_newtype!(ExprId, Expr);
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct Expr {
     pub kind: ExprKind,
+    pub hir_id: ExprId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -405,6 +429,13 @@ pub struct LambdaExpr {
 pub struct LambdaExprId(pub ExprId);
 
 impl HirIdNewtype for LambdaExprId {
+    fn new(hir_id: HirId) -> Self
+    where
+        Self: Sized,
+    {
+        Self(ExprId::new(hir_id))
+    }
+
     fn hir_id(&self) -> HirId {
         self.0.hir_id()
     }
@@ -733,6 +764,8 @@ pub struct Block {
     pub lbrace: LBrace,
     pub items: ItemList,
     pub rbrace: RBrace,
+
+    pub hir_id: BlockId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -747,6 +780,7 @@ hir_id_newtype!(StmtId, Stmt);
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct Stmt {
     pub kind: StmtKind,
+    pub hir_id: StmtId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
@@ -832,6 +866,7 @@ pub enum PatKind {
 pub struct TyRef {
     pub span: HirSpan,
     pub kind: TyRefKind,
+    pub hir_id: TyRefId,
 }
 
 hir_id_newtype!(TyRefId, TyRef);
@@ -870,13 +905,16 @@ pub enum PrimitiveTyKind {
 
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct TyGenericArgs {
-    pub args: Vec<TyGenericArg>,
+    pub args: Vec<TyGenericArgId>,
 }
 
+hir_id_newtype!(TyGenericArgId, TyGenericArg);
+
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct TyGenericArg {
     pub kind: TyGenericArgKind,
     pub span: HirSpan,
+    pub hir_id: TyGenericArgId,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]

diff --git a/crates/narxia-hir/src/hir_map.rs b/crates/narxia-hir/src/hir_map.rs
@@ -34,6 +34,41 @@ pub enum HirElem {
     StrLiteral(StrLiteral),
     StrLiteralDisplayFragment(StrLiteralDisplayFragment),
     StrLiteralDebugFragment(StrLiteralDebugFragment),
+    #[doc(hidden)]
+    __Allocated(HirSpan),
+}
+
+impl HirElem {
+    fn get_hir_id_in_self(&self) -> HirId {
+        match self {
+            HirElem::Mod(mod_def) => mod_def.hir_id.hir_id(),
+            HirElem::Item(item) => item.hir_id.hir_id(),
+            HirElem::Ident(ident) => todo!(),
+            HirElem::Fn(fn_def) => fn_def.hir_id.hir_id(),
+            HirElem::FnParam(fn_param) => todo!(),
+            HirElem::FnRetTy(fn_ret_ty) => todo!(),
+            HirElem::Expr(expr) => expr.hir_id.hir_id(),
+            HirElem::LoopExpr(loop_expr) => todo!(),
+            HirElem::BreakExpr(break_expr) => todo!(),
+            HirElem::ContinueExpr(continue_expr) => todo!(),
+            HirElem::ReturnExpr(return_expr) => todo!(),
+            HirElem::Pat(pat) => todo!(),
+            HirElem::Stmt(stmt) => stmt.hir_id.hir_id(),
+            HirElem::ForStmt(for_stmt) => todo!(),
+            HirElem::WhileStmt(while_stmt) => todo!(),
+            HirElem::UseStmt(use_stmt) => use_stmt.hir_id.hir_id(),
+            HirElem::UsePathSegment(use_path_segment) => use_path_segment.hir_id.hir_id(),
+            HirElem::Block(block) => block.hir_id.hir_id(),
+            HirElem::TyRef(ty_ref) => ty_ref.hir_id.hir_id(),
+            HirElem::TyGenericArg(ty_generic_arg) => ty_generic_arg.hir_id.hir_id(),
+            HirElem::LetStmt(let_stmt) => todo!(),
+            HirElem::AssignmentStmt(assignment_stmt) => todo!(),
+            HirElem::StrLiteral(str_literal) => todo!(),
+            HirElem::StrLiteralDisplayFragment(str_literal_display_fragment) => todo!(),
+            HirElem::StrLiteralDebugFragment(str_literal_debug_fragment) => todo!(),
+            HirElem::__Allocated(_) => unreachable!(),
+        }
+    }
 }
 
 impl fmt::Display for HirElem {
@@ -64,6 +99,7 @@ impl fmt::Display for HirElem {
             Self::StrLiteral(s) => write!(f, "{}", s),
             Self::StrLiteralDisplayFragment(s) => write!(f, "{}", s),
             Self::StrLiteralDebugFragment(s) => write!(f, "{}", s),
+            Self::__Allocated(_) => write!(f, "<Allocated>"),
         }
     }
 }
@@ -94,6 +130,18 @@ impl HirMap {
         HirId::new(self.buffer.len())
     }
 
+    pub fn allocate_hir_id(&mut self, span: HirSpan) -> HirId {
+        let mut id = self.next_hir_id();
+        #[cfg(hir_id_span)]
+        {
+            id.span = span;
+        }
+        self.buffer.push(HirElem::__Allocated(span));
+        self.parents.push(HirId::ORPHAN_HIRID);
+        self.files.push(self.current_file.unwrap());
+        id
+    }
+
     pub fn push_ref(&mut self, r: HirElem, span: HirSpan) -> HirId {
         let mut id = self.next_hir_id();
         #[cfg(hir_id_span)]
@@ -106,6 +154,12 @@ impl HirMap {
         id
     }
 
+    pub fn push_ref_at_allocation(&mut self, r: HirElem, allocation: HirId) {
+        debug_assert_eq!(self.buffer[allocation.id], HirElem::__Allocated(allocation.span));
+        debug_assert_eq!(r.get_hir_id_in_self(), allocation);
+        self.buffer[allocation.id] = r;
+    }
+
     pub fn get(&self, at: HirId) -> &HirElem {
         &self.buffer[at.id]
     }
@@ -184,6 +238,13 @@ impl HirMap {
         }
     }
 
+    pub fn get_ty_generic_arg(&self, at: TyGenericArgId) -> &TyGenericArg {
+        match self.get(at.0) {
+            HirElem::TyGenericArg(t) => t,
+            x => panic!("Expected TyGenericArg, found {x:?}"),
+        }
+    }
+
     fn update_parent(&mut self, at: HirId, parent: HirId) {
         self.parents[at.id] = parent;
     }
@@ -195,6 +256,27 @@ impl HirMap {
     pub fn get_file(&self, at: HirId) -> SrcFile {
         self.files[at.id]
     }
+
+    pub fn __get_allocated_hirids(&self) -> Vec<HirId> {
+        self.buffer
+            .iter()
+            .enumerate()
+            .filter_map(|(i, x)| {
+                if let HirElem::__Allocated(span) = x {
+                    let mut hir_id = HirId::new(i);
+
+                    #[cfg(hir_id_span)]
+                    {
+                        hir_id.span = *span;
+                    }
+
+                    Some(hir_id)
+                } else {
+                    None
+                }
+            })
+            .collect()
+    }
 }
 
 struct ParentUpdateVisitor<'hir> {