rustc_utils/source_map/spanner/
mod.rs

1//! Mapping source ranges to/from the HIR and MIR.
2
3use either::Either;
4use log::trace;
5use rustc_hir::{self as hir, BodyId, ExprKind, MatchSource, Node};
6use rustc_middle::{
7  mir::{
8    self, Body, RETURN_PLACE, StatementKind, TerminatorKind, visit::Visitor as MirVisitor,
9  },
10  ty::TyCtxt,
11};
12use rustc_span::{Span, SpanData, source_map::Spanned};
13
14pub use self::hir_span::EnclosingHirSpans;
15use self::{
16  mir_span::{MirSpanCollector, MirSpannedPlace},
17  span_tree::SpanTree,
18};
19use crate::{BodyExt, SpanDataExt, SpanExt, mir::location_or_arg::LocationOrArg};
20
21mod hir_span;
22mod mir_span;
23mod span_tree;
24
25/// Converts MIR locations to source spans using HIR information.
26pub struct Spanner<'tcx> {
27  pub(super) tcx: TyCtxt<'tcx>,
28  pub(super) mir_spans: Vec<MirSpannedPlace<'tcx>>,
29  pub mir_span_tree: SpanTree<MirSpannedPlace<'tcx>>,
30  pub body_span: Span,
31  pub item_span: Span,
32  pub ret_span: Span,
33}
34
35impl<'tcx> Spanner<'tcx> {
36  pub fn new(tcx: TyCtxt<'tcx>, body_id: BodyId, body: &Body<'tcx>) -> Self {
37    let hir_body = tcx.hir_body(body_id);
38    let owner = tcx.hir_body_owner(body_id);
39    let item_span = tcx.hir_span_with_body(owner);
40    let ret_span = tcx.hir_fn_decl_by_hir_id(owner).unwrap().output.span();
41
42    let mut spanner = Spanner {
43      mir_spans: Vec::new(),
44      mir_span_tree: SpanTree::new([]),
45      body_span: hir_body.value.span,
46      item_span,
47      ret_span,
48      tcx,
49    };
50    trace!(
51      "Body span: {:?}, item span: {:?}",
52      spanner.body_span, spanner.item_span
53    );
54
55    let mut mir_collector = MirSpanCollector(&mut spanner, body);
56    mir_collector.visit_body(body);
57
58    spanner.mir_span_tree =
59      SpanTree::new(spanner.mir_spans.drain(..).map(|node| Spanned {
60        span: node.span.span(),
61        node,
62      }));
63
64    spanner
65  }
66
67  pub(super) fn invalid_span(&self, span: Span) -> bool {
68    span.is_dummy()
69      || span.source_equal(self.body_span)
70      || span.source_equal(self.item_span)
71  }
72
73  fn find_matching<T>(
74    predicate: impl Fn(SpanData) -> bool,
75    query: SpanData,
76    spans: &'_ SpanTree<T>,
77  ) -> impl ExactSizeIterator<Item = &'_ T> + '_ {
78    let mut matching = spans
79      .overlapping(query)
80      .filter(|(span, _)| predicate(*span))
81      .collect::<Vec<_>>();
82    matching.sort_by_key(|(span, _)| span.size());
83    matching.into_iter().map(|(_, t)| t)
84  }
85
86  pub fn location_to_spans(
87    &self,
88    location: LocationOrArg,
89    body: &Body,
90    _span_type: EnclosingHirSpans,
91  ) -> Vec<Span> {
92    let (source_info, stmt) = match location {
93      LocationOrArg::Arg(local) => (&body.local_decls[local].source_info, None),
94      LocationOrArg::Location(location) => {
95        (body.source_info(location), Some(body.stmt_at(location)))
96      }
97    };
98
99    let hir_id = body.source_info_to_hir_id(source_info);
100
101    let mir_span = match source_info.span.as_local(self.item_span) {
102      Some(span) if !self.invalid_span(span) => span,
103      _ => {
104        return vec![];
105      }
106    };
107
108    let mut hir_spans = Vec::new();
109
110    // Include the MIR span, skipping spans that get mapped to the end brace of a body
111    if mir_span != body.span.shrink_to_hi() {
112      hir_spans.push(mir_span);
113    }
114
115    // Include the span for the immediately enclosing HIR node
116    if let Some(spans) = self.hir_spans(hir_id, EnclosingHirSpans::OuterOnly) {
117      hir_spans.extend(spans);
118    }
119
120    let enclosing_hir = self.tcx.hir_parent_iter(hir_id).collect::<Vec<_>>();
121    macro_rules! add_first_matching {
122      ($p:pat) => {
123        if let Some((id, _)) = enclosing_hir.iter().find(|(_, node)| matches!(node, $p)) {
124          if let Some(spans) = self.hir_spans(*id, EnclosingHirSpans::OuterOnly) {
125            hir_spans.extend(spans);
126          }
127        }
128      };
129    }
130
131    // Add the spans of the first enclosing statement
132    add_first_matching!(Node::Stmt(..));
133
134    // Include `return` keyword if the location is an expression under a return.
135    add_first_matching!(Node::Expr(hir::Expr {
136      kind: hir::ExprKind::Ret(..),
137      ..
138    }));
139
140    if let Some(Either::Right(mir::Terminator {
141      kind: TerminatorKind::SwitchInt { .. },
142      ..
143    })) = stmt
144    {
145      // If the location is a switch, then include the closest enclosing if or match
146      add_first_matching!(Node::Expr(hir::Expr {
147        kind: ExprKind::If(..) | ExprKind::Match(.., MatchSource::Normal),
148        ..
149      }));
150
151      // Also include enclosing loops
152      add_first_matching!(Node::Expr(hir::Expr {
153        kind: ExprKind::Loop(..),
154        ..
155      }));
156    }
157
158    if let Some(Either::Left(mir::Statement {
159      kind: StatementKind::Assign(box (lhs, _)),
160      ..
161    })) = stmt
162      && lhs.local == RETURN_PLACE
163    {
164      hir_spans.push(self.ret_span);
165    }
166
167    let format_spans = |spans: &[Span]| -> String {
168      spans
169        .iter()
170        .map(|span| span.to_string(self.tcx))
171        .collect::<Vec<_>>()
172        .join(" -- ")
173    };
174
175    trace!(
176      "Location {location:?} ({})\n  has loc span:\n  {}\n  and HIR spans:\n  {}",
177      location.to_string(body),
178      format_spans(&[mir_span]),
179      format_spans(&hir_spans)
180    );
181
182    hir_spans
183  }
184
185  pub fn span_to_places<'this>(
186    &'this self,
187    span: Span,
188  ) -> Vec<&'this MirSpannedPlace<'tcx>> {
189    // Note that MIR does not have granular source maps around projections.
190    // So in the expression `let x = z.0`, the MIR Body only contains the place
191    // z.0 with a span for the string "z.0". If the user selects only "z", there
192    // is no way to determine map that selection back to a subset of the projection.
193    //
194    // At least, we can conservatively include the containing span "z.0" and slice on that.
195
196    let span_data = span.data();
197
198    let mut contained = Self::find_matching(
199      move |mir_span| span_data.contains(mir_span),
200      span_data,
201      &self.mir_span_tree,
202    );
203    let mut vec = if let Some(first) = contained.next() {
204      contained
205        .take_while(|other| other.span.size() == first.span.size())
206        .chain([first])
207        .collect()
208    } else {
209      let mut containing = Self::find_matching(
210        move |mir_span| mir_span.contains(span_data),
211        span_data,
212        &self.mir_span_tree,
213      );
214      if let Some(first) = containing.next() {
215        containing
216          .take_while(|other| other.span.size() == first.span.size())
217          .chain([first])
218          .collect()
219      } else {
220        Vec::new()
221      }
222    };
223
224    vec.dedup();
225    vec
226  }
227}
228
229#[cfg(test)]
230mod test {
231  use rustc_data_structures::fx::FxHashSet as HashSet;
232  use rustc_middle::mir::BasicBlock;
233  use test_log::test;
234
235  use super::*;
236  use crate::{mir::borrowck_facts, source_map::range::ToSpan, test_utils};
237
238  fn harness(
239    src: &str,
240    f: impl for<'tcx> FnOnce(TyCtxt<'tcx>, BodyId, &Body<'tcx>, Vec<Span>) + Send,
241  ) {
242    let (input, _) = test_utils::parse_ranges(src, [("`(", ")`")]).unwrap();
243    test_utils::compile_body(input, move |tcx, body_id, body_with_facts| {
244      let (_, mut ranges) = test_utils::parse_ranges(src, [("`(", ")`")]).unwrap();
245      let spans = ranges
246        .remove("`(")
247        .unwrap()
248        .into_iter()
249        .map(|range| range.to_span(tcx).unwrap())
250        .collect::<Vec<_>>();
251      f(tcx, body_id, &body_with_facts.body, spans);
252    });
253  }
254
255  #[test]
256  fn test_span_to_places() {
257    let src = r#"fn foo(`(z)`: i32){
258      let `(x)` = 1;
259      let y = 1;
260      `(x + y)`;
261      `(x)` + y;
262      `(x + )`y;
263      print!("{} {}", x, `(y)`);
264      let w = (0, 0);
265      `(w)`.0;
266      `(w.0)`;
267      `(w.)`0;
268    }"#;
269    harness(src, |tcx, body_id, body, spans| {
270      let source_map = tcx.sess.source_map();
271      let spanner = Spanner::new(tcx, body_id, body);
272      let expected: &[&[_]] = &[
273        &["z"],
274        &["x"],
275        &["x", "y"],
276        &["x"],
277        &["x"],
278        &["y"],
279        &["w.0"],
280        &["w.0"],
281        &["w.0"],
282      ];
283      for (input_span, desired) in spans.into_iter().zip(expected) {
284        let outputs = spanner.span_to_places(input_span);
285        let snippets = outputs
286          .into_iter()
287          .map(|spanned| source_map.span_to_snippet(spanned.span.span()).unwrap())
288          .collect::<HashSet<_>>();
289
290        println!("input_span={input_span:?}");
291        compare_sets(&desired.iter().collect::<HashSet<_>>(), &snippets);
292      }
293    });
294  }
295
296  fn compare_sets(desired: &HashSet<impl AsRef<str>>, actual: &HashSet<impl AsRef<str>>) {
297    let desired = desired.iter().map(AsRef::as_ref).collect::<HashSet<_>>();
298    let actual = actual.iter().map(AsRef::as_ref).collect::<HashSet<_>>();
299    let missing_desired = &desired - &actual;
300    let missing_actual = &actual - &desired;
301
302    let check = |key: &str, set: HashSet<&str>| {
303      if let Some(el) = set.iter().next() {
304        panic!("Missing {key}: {el}. Actual = {actual:?}. Desired = {desired:?}",);
305      }
306    };
307
308    check("desired", missing_desired);
309    check("actual", missing_actual);
310  }
311
312  #[test]
313  fn test_location_to_spans() {
314    let src = r"fn foo() {
315  let mut x: i32 = 1;
316  let y = x + 2;
317  let w = if true {
318    let z = 0;
319    z
320  } else {
321    3
322  };
323  let z = &mut x;
324  *z = 4;
325  let q = x
326    .leading_ones()
327    .trailing_zeros();
328}";
329
330    // This affects source mapping, and this feature is primarily used by Flowistry, so
331    // we enable MIR simplification for consistency with Flowistry.
332    borrowck_facts::enable_mir_simplification();
333
334    let (input, _ranges) = test_utils::parse_ranges(src, [("`(", ")`")]).unwrap();
335    test_utils::compile_body(input, move |tcx, body_id, body_with_facts| {
336      let body = &body_with_facts.body;
337      let source_map = tcx.sess.source_map();
338
339      let spanner = Spanner::new(tcx, body_id, body);
340
341      // These locations are just selected by inspecting the actual body, so this test might break
342      // if the compiler is updated. Run with RUST_LOG=debug to see the body.
343      let pairs: &[(_, &[&str])] = &[
344        // Variable assignment
345        ((0, 0), &["let mut x: i32 = ", "1", ";"]),
346        // Expression RHS
347        ((0, 3), &["let y = ", "x", ";"]),
348        ((0, 4), &["let y = ", " + ", "x + 2", ";"]),
349        // If expression
350        ((1, 2), &["let w = ", "true", ";"]),
351        ((1, 3), &[
352          "let w = ",
353          "if ",
354          "true",
355          " {\n    ",
356          "\n    ",
357          "\n  } else {\n    ",
358          "\n  }",
359          ";",
360        ]),
361        // Reference
362        ((4, 1), &["let z = ", "&mut ", "&mut x", ";"]),
363        // Reference assignment
364        ((4, 3), &[" = ", ";", "*z = 4"]),
365        // Method chain
366        ((4, 4), &["let q = ", "x", ";"]),
367        ((4, 5), &[
368          "let q = ",
369          "x\n    .leading_ones()",
370          "\n    .leading_ones()",
371          ";",
372        ]),
373        ((5, 0), &[
374          "let q = ",
375          "x\n    .leading_ones()\n    .trailing_zeros()",
376          "\n    .trailing_zeros()",
377          ";",
378        ]),
379      ];
380
381      for ((i, j), outp) in pairs {
382        let loc = mir::Location {
383          block: BasicBlock::from_usize(*i),
384          statement_index: *j,
385        };
386        let spans = spanner.location_to_spans(
387          LocationOrArg::Location(loc),
388          body,
389          EnclosingHirSpans::OuterOnly,
390        );
391        let desired = outp.iter().collect::<HashSet<_>>();
392        let actual = spans
393          .into_iter()
394          .map(|s| source_map.span_to_snippet(s).unwrap())
395          .collect::<HashSet<_>>();
396        compare_sets(&desired, &actual);
397      }
398    });
399  }
400}
rustc_utils/source_map/spanner/mod.rs

rustc_utils/source_map/spanner/
mod.rs
