From 791ed01de2efea13d3e19ea2a13e81f3596f476e Mon Sep 17 00:00:00 2001 From: Justus Adam Date: Thu, 8 Feb 2024 14:01:55 -0500 Subject: [PATCH 1/3] Fix async return handling --- crates/flowistry/src/pdg/construct.rs | 94 ++++++++++++++++----------- 1 file changed, 57 insertions(+), 37 deletions(-) diff --git a/crates/flowistry/src/pdg/construct.rs b/crates/flowistry/src/pdg/construct.rs index f59f154c5..58bbb71dc 100644 --- a/crates/flowistry/src/pdg/construct.rs +++ b/crates/flowistry/src/pdg/construct.rs @@ -14,9 +14,9 @@ use rustc_hash::{FxHashMap, FxHashSet}; use rustc_hir::def_id::{DefId, LocalDefId}; use rustc_middle::{ mir::{ - visit::Visitor, AggregateKind, BasicBlock, Body, Location, Operand, Place, PlaceElem, - Rvalue, Statement, StatementKind, Terminator, TerminatorEdges, TerminatorKind, - RETURN_PLACE, + visit::Visitor, AggregateKind, BasicBlock, Body, HasLocalDecls, Location, Operand, + Place, PlaceElem, Rvalue, Statement, StatementKind, Terminator, TerminatorEdges, + TerminatorKind, RETURN_PLACE, }, ty::{GenericArg, GenericArgsRef, List, ParamEnv, TyCtxt, TyKind}, }; @@ -240,7 +240,14 @@ impl<'tcx> GraphConstructor<'tcx> { }; let body = utils::try_monomorphize(tcx, params.root, param_env, &body_with_facts.body); - debug!("{}", body.to_string(tcx).unwrap()); + + if log::log_enabled!(log::Level::Debug) { + use std::io::Write; + let path = tcx.def_path_str(def_id) + ".mir"; + let mut f = std::fs::File::create(path.as_str()).unwrap(); + write!(f, "{}", body.to_string(tcx).unwrap()).unwrap(); + debug!("Dumped debug MIR {path}"); + } let place_info = PlaceInfo::build(tcx, def_id.to_def_id(), body_with_facts); let control_dependencies = body.control_dependencies(); @@ -662,41 +669,54 @@ impl<'tcx> GraphConstructor<'tcx> { let parent_body = &self.body; let translate_to_parent = |child: Place<'tcx>| -> Option> { trace!(" Translating child place: {child:?}"); - let (parent_place, child_projection) = if child.local == RETURN_PLACE { - (destination, &child.projection[..]) - } else { - match call_kind { - // Map arguments to the argument array - CallKind::Direct => ( - args[child.local.as_usize() - 1].place()?, + let (parent_place, child_projection) = match call_kind { + // Async return must be handled special, because it gets wrapped in `Poll::Ready` + CallKind::AsyncPoll if child.local == RETURN_PLACE => { + let in_poll = + destination.project_deeper(&[PlaceElem::Downcast(None, 0_u32.into())], tcx); + let field_idx = 0_u32.into(); + let child_inner_return_type = in_poll + .ty(parent_body.local_decls(), tcx) + .field_ty(tcx, field_idx); + ( + in_poll.project_deeper( + &[PlaceElem::Field(field_idx, child_inner_return_type)], + tcx, + ), &child.projection[..], - ), - // Map arguments to projections of the future, the poll's first argument - CallKind::AsyncPoll => { - if child.local.as_usize() == 1 { - let PlaceElem::Field(idx, _) = child.projection[0] else { - panic!("Unexpected non-projection of async context") - }; - (args[idx.as_usize()].place()?, &child.projection[1 ..]) - } else { - return None; - } + ) + } + _ if child.local == RETURN_PLACE => (destination, &child.projection[..]), + // Map arguments to the argument array + CallKind::Direct => ( + args[child.local.as_usize() - 1].place()?, + &child.projection[..], + ), + // Map arguments to projections of the future, the poll's first argument + CallKind::AsyncPoll => { + if child.local.as_usize() == 1 { + let PlaceElem::Field(idx, _) = child.projection[0] else { + panic!("Unexpected non-projection of async context") + }; + (args[idx.as_usize()].place()?, &child.projection[1 ..]) + } else { + return None; } - // Map closure captures to the first argument. - // Map formal parameters to the second argument. - CallKind::Indirect => { - if child.local.as_usize() == 1 { - (args[0].place()?, &child.projection[..]) - } else { - let tuple_arg = args[1].place()?; - let _projection = child.projection.to_vec(); - let field = FieldIdx::from_usize(child.local.as_usize() - 2); - let field_ty = tuple_arg.ty(parent_body.as_ref(), tcx).field_ty(tcx, field); - ( - tuple_arg.project_deeper(&[PlaceElem::Field(field, field_ty)], tcx), - &child.projection[..], - ) - } + } + // Map closure captures to the first argument. + // Map formal parameters to the second argument. + CallKind::Indirect => { + if child.local.as_usize() == 1 { + (args[0].place()?, &child.projection[..]) + } else { + let tuple_arg = args[1].place()?; + let _projection = child.projection.to_vec(); + let field = FieldIdx::from_usize(child.local.as_usize() - 2); + let field_ty = tuple_arg.ty(parent_body.as_ref(), tcx).field_ty(tcx, field); + ( + tuple_arg.project_deeper(&[PlaceElem::Field(field, field_ty)], tcx), + &child.projection[..], + ) } } }; From 8ce816d30fd1387c36100ad0e833ebaa50e864c2 Mon Sep 17 00:00:00 2001 From: Justus Adam Date: Thu, 8 Feb 2024 14:30:47 -0500 Subject: [PATCH 2/3] Use Language items --- crates/flowistry/src/pdg/construct.rs | 56 +++++++++++++++++++++++---- 1 file changed, 48 insertions(+), 8 deletions(-) diff --git a/crates/flowistry/src/pdg/construct.rs b/crates/flowistry/src/pdg/construct.rs index 58bbb71dc..6d6d3813f 100644 --- a/crates/flowistry/src/pdg/construct.rs +++ b/crates/flowistry/src/pdg/construct.rs @@ -6,7 +6,7 @@ use flowistry_pdg::{CallString, GlobalLocation, RichLocation}; use itertools::Itertools; use log::{debug, trace}; use petgraph::graph::DiGraph; -use rustc_abi::FieldIdx; +use rustc_abi::{FieldIdx, VariantIdx}; use rustc_borrowck::consumers::{ places_conflict, BodyWithBorrowckFacts, PlaceConflictBias, }; @@ -21,6 +21,7 @@ use rustc_middle::{ ty::{GenericArg, GenericArgsRef, List, ParamEnv, TyCtxt, TyKind}, }; use rustc_mir_dataflow::{self as df}; +use rustc_span::sym::poll; use rustc_utils::{ mir::{borrowck_facts, control_dependencies::ControlDependencies}, BodyExt, PlaceExt, @@ -201,6 +202,25 @@ struct CallingContext<'tcx> { call_stack: Vec, } +/// Stores ids that are needed to construct projections around async functions. +struct AsyncInfo { + poll_ready_variant_idx: VariantIdx, + poll_ready_field_idx: FieldIdx, +} + +impl AsyncInfo { + fn make(tcx: TyCtxt) -> Option> { + let lang_items = tcx.lang_items(); + let poll_def = tcx.adt_def(lang_items.poll()?); + let ready_vid = lang_items.poll_ready_variant()?; + assert_eq!(poll_def.variant_with_id(ready_vid).fields.len(), 1); + Some(Rc::new(Self { + poll_ready_variant_idx: poll_def.variant_index_with_id(ready_vid), + poll_ready_field_idx: 0_u32.into(), + })) + } +} + pub struct GraphConstructor<'tcx> { tcx: TyCtxt<'tcx>, params: PdgParams<'tcx>, @@ -212,6 +232,7 @@ pub struct GraphConstructor<'tcx> { body_assignments: utils::BodyAssignments, calling_context: Option>, start_loc: FxHashSet, + async_info: Rc, } macro_rules! trylet { @@ -226,11 +247,20 @@ macro_rules! trylet { impl<'tcx> GraphConstructor<'tcx> { /// Creates a [`GraphConstructor`] at the root of the PDG. pub fn root(params: PdgParams<'tcx>) -> Self { - GraphConstructor::new(params, None) + let tcx = params.tcx; + GraphConstructor::new( + params, + None, + AsyncInfo::make(tcx).expect("async functions are not defined"), + ) } /// Creates [`GraphConstructor`] for a function resolved as `fn_resolution` in a given `calling_context`. - fn new(params: PdgParams<'tcx>, calling_context: Option>) -> Self { + fn new( + params: PdgParams<'tcx>, + calling_context: Option>, + async_info: Rc, + ) -> Self { let tcx = params.tcx; let def_id = params.root.def_id().expect_local(); let body_with_facts = borrowck_facts::get_body_with_borrowck_facts(tcx, def_id); @@ -268,6 +298,7 @@ impl<'tcx> GraphConstructor<'tcx> { def_id, calling_context, body_assignments, + async_info, } } @@ -672,9 +703,12 @@ impl<'tcx> GraphConstructor<'tcx> { let (parent_place, child_projection) = match call_kind { // Async return must be handled special, because it gets wrapped in `Poll::Ready` CallKind::AsyncPoll if child.local == RETURN_PLACE => { - let in_poll = - destination.project_deeper(&[PlaceElem::Downcast(None, 0_u32.into())], tcx); - let field_idx = 0_u32.into(); + let async_info = self.async_info.as_ref(); + let in_poll = destination.project_deeper( + &[PlaceElem::Downcast(None, async_info.poll_ready_variant_idx)], + tcx, + ); + let field_idx = async_info.poll_ready_field_idx; let child_inner_return_type = in_poll .ty(parent_body.local_decls(), tcx) .field_ty(tcx, field_idx); @@ -750,7 +784,8 @@ impl<'tcx> GraphConstructor<'tcx> { param_env, call_stack, }; - let child_constructor = GraphConstructor::new(params, Some(calling_context)); + let child_constructor = + GraphConstructor::new(params, Some(calling_context), self.async_info.clone()); if let Some(callback) = &self.params.call_change_callback { let info = CallInfo { @@ -961,7 +996,12 @@ impl<'tcx> GraphConstructor<'tcx> { call_string, call_stack, }; - return GraphConstructor::new(params, Some(calling_context)).construct_partial(); + return GraphConstructor::new( + params, + Some(calling_context), + self.async_info.clone(), + ) + .construct_partial(); } let mut analysis = DfAnalysis(self) From 5bc82f676577fb64d0e04064ca0e61598e345e43 Mon Sep 17 00:00:00 2001 From: Justus Adam Date: Thu, 8 Feb 2024 14:35:09 -0500 Subject: [PATCH 3/3] Cleaning changes --- crates/flowistry/src/pdg/construct.rs | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/crates/flowistry/src/pdg/construct.rs b/crates/flowistry/src/pdg/construct.rs index 6d6d3813f..1b735407e 100644 --- a/crates/flowistry/src/pdg/construct.rs +++ b/crates/flowistry/src/pdg/construct.rs @@ -21,7 +21,6 @@ use rustc_middle::{ ty::{GenericArg, GenericArgsRef, List, ParamEnv, TyCtxt, TyKind}, }; use rustc_mir_dataflow::{self as df}; -use rustc_span::sym::poll; use rustc_utils::{ mir::{borrowck_facts, control_dependencies::ControlDependencies}, BodyExt, PlaceExt, @@ -270,14 +269,7 @@ impl<'tcx> GraphConstructor<'tcx> { }; let body = utils::try_monomorphize(tcx, params.root, param_env, &body_with_facts.body); - - if log::log_enabled!(log::Level::Debug) { - use std::io::Write; - let path = tcx.def_path_str(def_id) + ".mir"; - let mut f = std::fs::File::create(path.as_str()).unwrap(); - write!(f, "{}", body.to_string(tcx).unwrap()).unwrap(); - debug!("Dumped debug MIR {path}"); - } + debug!("{}", body.to_string(tcx).unwrap()); let place_info = PlaceInfo::build(tcx, def_id.to_def_id(), body_with_facts); let control_dependencies = body.control_dependencies();