rustc_utils/mir/
body.rs

1//! Utilities for [`Body`].
2
3use std::{
4  io::Write,
5  path::Path,
6  process::{Command, Stdio},
7};
8
9use anyhow::{Result, ensure};
10use rustc_data_structures::fx::FxHashMap as HashMap;
11use rustc_hir::{CoroutineDesugaring, CoroutineKind, HirId, def_id::DefId};
12use rustc_middle::{
13  mir::{
14    BasicBlock, Body, Local, Location, Place, SourceInfo, TerminatorKind,
15    VarDebugInfoContents, pretty::MirWriter,
16  },
17  ty::{Region, Ty, TyCtxt},
18};
19use smallvec::SmallVec;
20
21use super::control_dependencies::ControlDependencies;
22use crate::{PlaceExt, TyExt};
23
24/// Extension trait for [`Body`].
25pub trait BodyExt<'tcx> {
26  /// Returns an iterator over the locations of [`TerminatorKind::Return`] instructions in a body.
27  fn all_returns(&self) -> impl Iterator<Item = Location> + '_;
28
29  /// Returns an iterator over all the locations in a body.
30  fn all_locations(&self) -> impl Iterator<Item = Location> + '_;
31
32  /// Returns all the locations in a [`BasicBlock`].
33  fn locations_in_block(&self, block: BasicBlock) -> impl Iterator<Item = Location>;
34
35  /// Returns a mapping from source-level variable names to [`Local`]s.
36  fn debug_info_name_map(&self) -> HashMap<String, Local>;
37
38  /// Converts a Body to a debug representation.
39  fn to_string(&self, tcx: TyCtxt<'tcx>) -> Result<String>;
40
41  /// Returns the [`HirId`] corresponding to a MIR [`Location`].
42  ///
43  /// You **MUST** use the `-Zmaximize-hir-to-mir-mapping` flag for this
44  /// function to work.
45  fn location_to_hir_id(&self, location: Location) -> HirId;
46
47  fn source_info_to_hir_id(&self, info: &SourceInfo) -> HirId;
48
49  /// Returns all the control dependencies within the CFG.
50  ///
51  /// See the [`control_dependencies`][super::control_dependencies] module documentation
52  /// for details.
53  fn control_dependencies(&self) -> ControlDependencies<BasicBlock>;
54
55  /// If this body is an async function, then return the type of the context that holds
56  /// locals across await calls.
57  fn async_context(&self, tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Ty<'tcx>>;
58
59  /// Returns an iterator over all projections of all local variables in the body.
60  fn all_places(
61    &self,
62    tcx: TyCtxt<'tcx>,
63    def_id: DefId,
64  ) -> impl Iterator<Item = Place<'tcx>> + '_;
65
66  /// Returns an iterator over all the regions that appear in argument types to the body.
67  fn regions_in_args(&self) -> impl Iterator<Item = Region<'tcx>> + '_;
68
69  /// Returns an iterator over all the regions that appear in the body's return type.
70  fn regions_in_return(&self) -> impl Iterator<Item = Region<'tcx>> + '_;
71}
72
73impl<'tcx> BodyExt<'tcx> for Body<'tcx> {
74  fn all_returns(&self) -> impl Iterator<Item = Location> + '_ {
75    self
76      .basic_blocks
77      .iter_enumerated()
78      .filter_map(|(block, data)| match data.terminator().kind {
79        TerminatorKind::Return => Some(Location {
80          block,
81          statement_index: data.statements.len(),
82        }),
83        _ => None,
84      })
85  }
86
87  fn all_locations(&self) -> impl Iterator<Item = Location> + '_ {
88    self
89      .basic_blocks
90      .iter_enumerated()
91      .flat_map(|(block, data)| {
92        (0 ..= data.statements.len()).map(move |statement_index| Location {
93          block,
94          statement_index,
95        })
96      })
97  }
98
99  fn locations_in_block(&self, block: BasicBlock) -> impl Iterator<Item = Location> {
100    let num_stmts = self.basic_blocks[block].statements.len();
101    (0 ..= num_stmts).map(move |statement_index| Location {
102      block,
103      statement_index,
104    })
105  }
106
107  fn debug_info_name_map(&self) -> HashMap<String, Local> {
108    self
109      .var_debug_info
110      .iter()
111      .filter_map(|info| match info.value {
112        VarDebugInfoContents::Place(place) => Some((info.name.to_string(), place.local)),
113        _ => None,
114      })
115      .collect()
116  }
117
118  fn to_string(&self, tcx: TyCtxt<'tcx>) -> Result<String> {
119    let mut buffer = Vec::new();
120    MirWriter::new(tcx).write_mir_fn(self, &mut buffer)?;
121    Ok(String::from_utf8(buffer)?)
122  }
123
124  fn location_to_hir_id(&self, location: Location) -> HirId {
125    let source_info = self.source_info(location);
126    self.source_info_to_hir_id(source_info)
127  }
128
129  fn source_info_to_hir_id(&self, info: &SourceInfo) -> HirId {
130    let scope = &self.source_scopes[info.scope];
131    let local_data = scope.local_data.as_ref().unwrap_crate_local();
132    local_data.lint_root
133  }
134
135  fn control_dependencies(&self) -> ControlDependencies<BasicBlock> {
136    ControlDependencies::build_many(
137      &self.basic_blocks,
138      self.all_returns().map(|loc| loc.block),
139    )
140  }
141
142  fn async_context(&self, tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Ty<'tcx>> {
143    if matches!(
144      tcx.coroutine_kind(def_id),
145      Some(CoroutineKind::Desugared(CoroutineDesugaring::Async, _))
146    ) {
147      Some(self.local_decls[Local::from_usize(2)].ty)
148    } else {
149      None
150    }
151  }
152
153  fn regions_in_args(&self) -> impl Iterator<Item = Region<'tcx>> + '_ {
154    self
155      .args_iter()
156      .flat_map(|arg_local| self.local_decls[arg_local].ty.inner_regions())
157  }
158
159  fn regions_in_return(&self) -> impl Iterator<Item = Region<'tcx>> + '_ {
160    self
161      .return_ty()
162      .inner_regions()
163      .collect::<SmallVec<[Region<'tcx>; 8]>>()
164      .into_iter()
165  }
166
167  fn all_places(
168    &self,
169    tcx: TyCtxt<'tcx>,
170    def_id: DefId,
171  ) -> impl Iterator<Item = Place<'tcx>> + '_ {
172    self.local_decls.indices().flat_map(move |local| {
173      Place::from_local(local, tcx).interior_paths(tcx, self, def_id)
174    })
175  }
176}
177
178pub fn run_dot(path: &Path, buf: &[u8]) -> Result<()> {
179  let mut p = Command::new("dot")
180    .args(["-Tpdf", "-o", &path.display().to_string()])
181    .stdin(Stdio::piped())
182    .spawn()?;
183
184  p.stdin.as_mut().unwrap().write_all(buf)?;
185
186  let status = p.wait()?;
187  ensure!(status.success(), "dot for {} failed", path.display());
188
189  Ok(())
190}
191
192#[cfg(test)]
193mod test {
194  use super::BodyExt;
195  use crate::test_utils;
196
197  #[test]
198  fn test_body_ext() {
199    let input = r"
200fn foobar<'a>(x: &'a i32, y: &'a i32) -> &'a i32 {
201  if *x > 0 {
202    return x;
203  }
204
205  y
206}";
207
208    test_utils::compile_body(input, |_, _, body| {
209      let body = &body.body;
210      assert_eq!(body.regions_in_args().count(), 2);
211      assert_eq!(body.regions_in_return().count(), 1);
212    });
213  }
214}
rustc_utils/mir/body.rs

rustc_utils/mir/
body.rs
