c2rust-refactor: Map NodeIds to HirIds using a set of intermediate Span maps

ahomescu · ahomescu · commit ac75014a6610 · 2025-08-27T22:19:23.000-07:00
diff --git a/c2rust-refactor/src/ast_manip/mod.rs b/c2rust-refactor/src/ast_manip/mod.rs
@@ -15,6 +15,7 @@ mod load_modules;
 mod output_exprs;
 mod remove_paren;
 mod seq_edit;
+mod span_maps;
 mod visit;
 mod visit_node;
 
@@ -34,6 +35,7 @@ pub use self::visit::Visit;
 pub use self::visit_node::{visit_nodes, visit_nodes_post, VisitNode};
 pub use self::comments::{collect_comments, gather_comments, Comment, CommentMap, CommentStyle};
 pub use self::load_modules::load_modules;
+pub use self::span_maps::AstSpanMaps;
 
 // Modules with more complex APIs are left as `pub`.
 pub mod comments;
diff --git a/c2rust-refactor/src/ast_manip/span_maps.rs b/c2rust-refactor/src/ast_manip/span_maps.rs
@@ -0,0 +1,124 @@
+use std::collections::HashMap;
+use rustc_ast::*;
+use rustc_ast::node_id::NodeMap;
+use rustc_ast::visit::{self, AssocCtxt, Visitor};
+use rustc_span::Span;
+
+#[derive(Default, Clone)]
+pub struct AstSpanMaps {
+    pub node_id_to_span_map: NodeMap<Span>,
+    pub span_to_node_id_map: HashMap<Span, NodeId>,
+}
+
+impl AstSpanMaps {
+    pub fn new(krate: &Crate) -> Self {
+        let mut mapper = AstSpanMapper::default();
+        mapper.visit_crate(krate);
+        mapper.0
+    }
+}
+
+#[derive(Default, Clone)]
+struct AstSpanMapper(AstSpanMaps);
+
+impl AstSpanMapper {
+    fn insert_mapping(&mut self, id: NodeId, span: Span) {
+        if id == DUMMY_NODE_ID || span.is_dummy() {
+            return;
+        }
+
+        let old_span = self.0.node_id_to_span_map.insert(id, span);
+        let old_id = self.0.span_to_node_id_map.insert(span, id);
+
+        assert!(old_span.is_none(), "id {id:?} already has span {old_span:?} != {span:?}");
+        // Some spans can show up in multiple nodes, e.g., `impl $trait for $ty`
+        //assert!(old_id.is_none(), "span {span:?} already has id {old_id:?} != {id:?}");
+    }
+}
+
+impl Visitor<'_> for AstSpanMapper {
+    fn visit_crate(&mut self, krate: &Crate) {
+        self.insert_mapping(krate.id, krate.spans.inner_span);
+        visit::walk_crate(self, krate);
+    }
+
+    fn visit_block(&mut self, block: &Block) {
+        self.insert_mapping(block.id, block.span);
+        visit::walk_block(self, block);
+    }
+
+    fn visit_pat(&mut self, pat: &Pat) {
+        self.insert_mapping(pat.id, pat.span);
+        visit::walk_pat(self, pat);
+    }
+
+    fn visit_pat_field(&mut self, pf: &PatField) {
+        self.insert_mapping(pf.id, pf.span);
+        visit::walk_pat_field(self, pf);
+    }
+
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        self.insert_mapping(stmt.id, stmt.span);
+        visit::walk_stmt(self, stmt);
+    }
+
+    fn visit_local(&mut self, local: &Local) {
+        self.insert_mapping(local.id, local.span);
+        visit::walk_local(self, local);
+    }
+
+    fn visit_arm(&mut self, arm: &Arm) {
+        self.insert_mapping(arm.id, arm.span);
+        visit::walk_arm(self, arm);
+    }
+
+    fn visit_expr_field(&mut self, ef: &ExprField) {
+        self.insert_mapping(ef.id, ef.span);
+        visit::walk_expr_field(self, ef);
+    }
+
+    fn visit_expr(&mut self, expr: &Expr) {
+        self.insert_mapping(expr.id, expr.span);
+        visit::walk_expr(self, expr);
+    }
+
+    fn visit_assoc_constraint(&mut self, ac: &AssocConstraint) {
+        self.insert_mapping(ac.id, ac.span);
+        visit::walk_assoc_constraint(self, ac);
+    }
+
+    fn visit_ty(&mut self, ty: &Ty) {
+        self.insert_mapping(ty.id, ty.span);
+        visit::walk_ty(self, ty);
+    }
+
+    fn visit_param(&mut self, param: &Param) {
+        self.insert_mapping(param.id, param.span);
+        visit::walk_param(self, param);
+    }
+
+    fn visit_variant(&mut self, var: &Variant) {
+        self.insert_mapping(var.id, var.span);
+        visit::walk_variant(self, var);
+    }
+
+    fn visit_field_def(&mut self, fd: &FieldDef) {
+        self.insert_mapping(fd.id, fd.span);
+        visit::walk_field_def(self, fd);
+    }
+
+    fn visit_item(&mut self, i: &Item) {
+        self.insert_mapping(i.id, i.span);
+        visit::walk_item(self, i);
+    }
+
+    fn visit_foreign_item(&mut self, i: &ForeignItem) {
+        self.insert_mapping(i.id, i.span);
+        visit::walk_foreign_item(self, i);
+    }
+
+    fn visit_assoc_item(&mut self, i: &AssocItem, ctxt: AssocCtxt) {
+        self.insert_mapping(i.id, i.span);
+        visit::walk_assoc_item(self, i, ctxt);
+    }
+}
diff --git a/c2rust-refactor/src/command.rs b/c2rust-refactor/src/command.rs
@@ -28,6 +28,7 @@ use crate::ast_manip::number_nodes::{
 };
 use crate::ast_manip::{load_modules, remove_paren, ListNodeIds, MutVisit, Visit};
 use crate::ast_manip::{collect_comments, gather_comments, Comment, CommentMap};
+use crate::ast_manip::AstSpanMaps;
 use crate::collapse::CollapseInfo;
 use crate::driver::{self, Phase};
 use crate::file_io::FileIO;
@@ -425,6 +426,7 @@ impl RefactorState {
                             node_id_to_def_id,
                             def_id_to_node_id,
                             GenerationalTyCtxt(tcx, tcx_gen.clone()),
+                            AstSpanMaps::new(&expanded),
                         );
                         profile_end!("Lower to HIR");
 
@@ -453,6 +455,7 @@ impl RefactorState {
                             node_id_to_def_id,
                             def_id_to_node_id,
                             GenerationalTyCtxt(tcx, tcx_gen.clone()),
+                            AstSpanMaps::new(&expanded),
                         );
                         profile_end!("Compiler Phase 3");
 
diff --git a/c2rust-refactor/src/context.rs b/c2rust-refactor/src/context.rs
@@ -8,7 +8,7 @@ use rustc_hir::def_id::{CrateNum, DefId, LocalDefId};
 use rustc_hir::{self as hir, BodyId, Node, HirId};
 use rustc_session::Session;
 use rustc_session::config::CrateType;
-use rustc_middle::hir::map as hir_map;
+use rustc_middle::hir::{map as hir_map, nested_filter};
 use rustc_middle::ty::subst::InternalSubsts;
 use rustc_middle::ty::{FnSig, ParamEnv, PolyFnSig, Ty, TyCtxt, TyKind};
 use rustc_errors::{DiagnosticBuilder, Level};
@@ -20,7 +20,7 @@ use rustc_ast::ptr::P;
 use rustc_index::vec::IndexVec;
 use rustc_span::Span;
 
-use crate::ast_manip::AstEquiv;
+use crate::ast_manip::{AstEquiv, AstSpanMaps};
 use crate::command::{GenerationalTyCtxt, TyCtxtGeneration};
 use crate::ast_manip::util::{namespace, is_export_attr};
 use crate::{expect, match_or};
@@ -55,6 +55,57 @@ impl<'a, 'tcx> RefactorCtxt<'a, 'tcx> {
     }
 }
 
+type SpanToHirMap = FxHashMap<Span, HirId>;
+
+struct SpanToHirMapper<'hir> {
+    hir_map: hir_map::Map<'hir>,
+    span_to_hir_map: SpanToHirMap,
+}
+
+impl<'hir> hir::intravisit::Visitor<'hir> for SpanToHirMapper<'hir> {
+    type NestedFilter = nested_filter::OnlyBodies;
+
+    fn nested_visit_map(&mut self) -> Self::Map {
+        self.hir_map
+    }
+
+    fn visit_id(&mut self, id: HirId) {
+        let span = match self.hir_map.find(id) {
+            Some(Node::Param(param)) => param.span,
+            Some(Node::Item(item)) => item.span,
+            Some(Node::ForeignItem(foreign_item)) => foreign_item.span,
+            Some(Node::TraitItem(trait_item)) => trait_item.span,
+            Some(Node::ImplItem(impl_item)) => impl_item.span,
+            Some(Node::Variant(variant)) => variant.span,
+            Some(Node::Field(field)) => field.span,
+            Some(Node::AnonConst(constant)) => self.hir_map.body(constant.body).value.span,
+            // TODO: Expr only if not Block
+            Some(Node::Expr(expr)) => expr.span,
+            Some(Node::Stmt(stmt)) => stmt.span,
+            // Intentionally skip PathSegment to avoid collisions
+            // Intentionally skip Ty because we want the actual definitions
+            Some(Node::TypeBinding(tb)) => tb.span,
+            // Intentionally skip TraitRef to avoid collisions
+            Some(Node::Pat(pat)) => pat.span,
+            Some(Node::Arm(arm)) => arm.span,
+            Some(Node::Block(block)) => block.span,
+            Some(Node::Ctor(..)) => self.hir_map.span_with_body(self.hir_map.get_parent_node(id)),
+            Some(Node::Lifetime(lifetime)) => lifetime.span,
+            Some(Node::GenericParam(param)) => param.span,
+            Some(Node::Infer(i)) => i.span,
+            Some(Node::Local(local)) => local.span,
+            Some(Node::Crate(item)) => item.spans.inner_span,
+            _ => return
+        };
+        if span.is_dummy() {
+            return;
+        }
+
+        let old_id = self.span_to_hir_map.insert(span, id);
+        assert!(old_id.is_none(), "span {span:?} already has id {old_id:?} != {id:?}");
+    }
+}
+
 #[derive(Clone)]
 pub struct HirMap<'hir> {
     map: hir_map::Map<'hir>,
@@ -65,6 +116,9 @@ pub struct HirMap<'hir> {
 
     node_id_to_def_id: FxHashMap<NodeId, LocalDefId>,
     def_id_to_node_id: IndexVec<LocalDefId, NodeId>,
+
+    span_to_hir_map: SpanToHirMap,
+    ast_span_maps: AstSpanMaps,
 }
 
 impl<'hir> HirMap<'hir> {
@@ -73,8 +127,16 @@ impl<'hir> HirMap<'hir> {
         map: hir_map::Map<'hir>,
         node_id_to_def_id: FxHashMap<NodeId, LocalDefId>,
         def_id_to_node_id: IndexVec<LocalDefId, NodeId>,
+        ast_span_maps: AstSpanMaps,
     ) -> Self {
-        Self { map, max_node_id, node_id_to_def_id, def_id_to_node_id }
+        let mut mapper = SpanToHirMapper {
+            hir_map: map,
+            span_to_hir_map: Default::default(),
+        };
+        map.visit_all_item_likes_in_crate(&mut mapper);
+        let SpanToHirMapper { span_to_hir_map, .. } = mapper;
+
+        Self { map, max_node_id, node_id_to_def_id, def_id_to_node_id, span_to_hir_map, ast_span_maps }
     }
 }
 
@@ -611,9 +673,15 @@ impl<'hir> HirMap<'hir> {
         if id > self.max_node_id {
             None
         } else {
-            self.node_id_to_def_id.get(&id).map(|ldid| {
-                self.map.local_def_id_to_hir_id(*ldid)
-            })
+            if let Some(ldid) = self.node_id_to_def_id.get(&id) {
+                return Some(self.map.local_def_id_to_hir_id(*ldid));
+            }
+
+            self
+                .ast_span_maps
+                .node_id_to_span_map
+                .get(&id)
+                .and_then(|span| self.span_to_hir_map.get(&span).copied())
         }
     }
 
@@ -643,6 +711,11 @@ impl<'hir> HirMap<'hir> {
     }
 
     pub fn hir_to_node_id(&self, id: HirId) -> NodeId {
+        let span = self.map.span_with_body(id);
+        if let Some(id) = self.ast_span_maps.span_to_node_id_map.get(&span) {
+            return *id;
+        }
+
         self
             .map
             .opt_local_def_id(id)
diff --git a/c2rust-refactor/src/driver.rs b/c2rust-refactor/src/driver.rs
@@ -42,7 +42,7 @@ use rustc_span::{FileName, Span, DUMMY_SP};
 use rustc_span::edition::Edition;
 use rustc_span::def_id::LocalDefId;
 
-use crate::ast_manip::remove_paren;
+use crate::ast_manip::{AstSpanMaps, remove_paren};
 use crate::command::{GenerationalTyCtxt, RefactorState, Registry};
 use crate::file_io::{ArcFileIO, FileIO};
 // TODO: don't forget to call span_fix after parsing
@@ -76,8 +76,9 @@ impl<'a, 'tcx: 'a> RefactorCtxt<'a, 'tcx> {
         node_id_to_def_id: FxHashMap<NodeId, LocalDefId>,
         def_id_to_node_id: IndexVec<LocalDefId, NodeId>,
         tcx: GenerationalTyCtxt<'tcx>,
+        span_maps: AstSpanMaps,
     ) -> RefactorCtxt<'a, 'tcx> {
-        RefactorCtxt::new(sess, None, Some(HirMap::new(max_node_id, map, node_id_to_def_id, def_id_to_node_id)), Some(tcx))
+        RefactorCtxt::new(sess, None, Some(HirMap::new(max_node_id, map, node_id_to_def_id, def_id_to_node_id, span_maps)), Some(tcx))
     }
 }
 
