Fixes found when experimenting with Lemmy

Cargo.toml

@@ -35,9 +35,9 @@ debug = true
 [replace."rustc_utils:0.7.4-nightly-2023-08-25"]
 # path = "../rustc_plugin/crates/rustc_utils"
 git = "https://github.com/JustusAdam/rustc_plugin"
-rev = "d7bde3ae16a7137de594c2dd811030cd761a8993"
+rev = "e413907b2ae9a03d2c8e9aca3b72dd451a16b1db"
 
 [replace."rustc_plugin:0.7.4-nightly-2023-08-25"]
 # path = "../rustc_plugin/crates/rustc_plugin"
 git = "https://github.com/JustusAdam/rustc_plugin"
-rev = "d7bde3ae16a7137de594c2dd811030cd761a8993"
+rev = "e413907b2ae9a03d2c8e9aca3b72dd451a16b1db"

crates/flowistry_pdg_construction/src/async_support.rs

@@ -3,7 +3,7 @@ use std::rc::Rc;
 use either::Either;
 use itertools::Itertools;
 use rustc_abi::{FieldIdx, VariantIdx};
-use rustc_hir::def_id::{DefId, LocalDefId};
+use rustc_hir::def_id::DefId;
 use rustc_middle::{
     mir::{
         AggregateKind, BasicBlock, Body, Location, Operand, Place, Rvalue, Statement,
@@ -12,6 +12,8 @@ use rustc_middle::{
     ty::{GenericArgsRef, Instance, TyCtxt},
 };
 
+use crate::utils::is_async;
+
 use super::{
     local_analysis::{CallKind, LocalAnalysis},
     utils,
@@ -63,7 +65,7 @@ pub fn try_as_async_trait_function<'tcx>(
     tcx: TyCtxt,
     def_id: DefId,
     body: &Body<'tcx>,
-) -> Option<(LocalDefId, GenericArgsRef<'tcx>, Location)> {
+) -> Option<(DefId, GenericArgsRef<'tcx>, Location)> {
     if !has_async_trait_signature(tcx, def_id) {
         return None;
     }
@@ -75,7 +77,7 @@ pub fn try_as_async_trait_function<'tcx>(
                     move |(statement_index, statement)| {
                         let (def_id, generics) = match_async_trait_assign(statement)?;
                         Some((
-                            def_id.as_local()?,
+                            def_id,
                             generics,
                             Location {
                                 block,
@@ -147,7 +149,7 @@ fn match_pin_box_dyn_ty(lang_items: &rustc_hir::LanguageItems, t: ty::Ty) -> boo
     })
 }
 
-fn get_async_generator<'tcx>(body: &Body<'tcx>) -> (LocalDefId, GenericArgsRef<'tcx>, Location) {
+fn get_async_generator<'tcx>(body: &Body<'tcx>) -> (DefId, GenericArgsRef<'tcx>, Location) {
     let block = BasicBlock::from_usize(0);
     let location = Location {
         block,
@@ -163,7 +165,7 @@ fn get_async_generator<'tcx>(body: &Body<'tcx>) -> (LocalDefId, GenericArgsRef<'
     else {
         panic!("Async fn should assign to a generator")
     };
-    (def_id.expect_local(), generic_args, location)
+    (*def_id, generic_args, location)
 }
 
 /// Try to interpret this function as an async function.
@@ -176,7 +178,7 @@ pub fn determine_async<'tcx>(
     def_id: DefId,
     body: &Body<'tcx>,
 ) -> Option<(Instance<'tcx>, Location, AsyncType)> {
-    let ((generator_def_id, args, loc), asyncness) = if tcx.asyncness(def_id).is_async() {
+    let ((generator_def_id, args, loc), asyncness) = if is_async(tcx, def_id) {
         (get_async_generator(body), AsyncType::Fn)
     } else {
         (
@@ -185,8 +187,7 @@ pub fn determine_async<'tcx>(
         )
     };
     let param_env = tcx.param_env_reveal_all_normalized(def_id);
-    let generator_fn =
-        utils::try_resolve_function(tcx, generator_def_id.to_def_id(), param_env, args)?;
+    let generator_fn = utils::try_resolve_function(tcx, generator_def_id, param_env, args)?;
     Some((generator_fn, loc, asyncness))
 }
 

crates/flowistry_pdg_construction/src/calling_convention.rs

@@ -11,16 +11,59 @@ use rustc_middle::{
 
 use crate::{async_support::AsyncInfo, local_analysis::CallKind, utils};
 
+/// Describes how the formal parameters of a given function call relate to the
+/// actual parameters.
 #[derive(Debug)]
 pub enum CallingConvention<'tcx> {
+    /// 1 to 1 mapping
     Direct(Box<[Operand<'tcx>]>),
+    /// First argument is the closed-over environment, second argument is a
+    /// tuple that contains the actual argument to the call of the closure
+    /// function.
     Indirect {
+        once_shim: bool,
         closure_arg: Operand<'tcx>,
         tupled_arguments: Operand<'tcx>,
     },
+    /// An async generator, only has one argument which is the generator state.
     Async(Place<'tcx>),
 }
 
+/// The result of calculating a translation from a child place (in a called
+/// function) to a parent place (in the caller).
+///
+/// This is partially translated and thus allows us to either complete the
+/// translation to a precise parent place ([`Self::make_translated_place`]),
+/// e.g. one that corresponds to the child 1-1, or to just use the parent place,
+/// for strategic overtaint, e.g. discarding the child projections
+/// ([`Self::base_place`]).
+pub struct PlaceTranslation<'a, 'tcx> {
+    new_base: Place<'tcx>,
+    additional_projection: &'tcx [PlaceElem<'tcx>],
+    scope: &'a PlaceTranslator<'a, 'tcx>,
+}
+
+impl<'a, 'tcx> PlaceTranslation<'a, 'tcx> {
+    /// Complete the translation and return a precise parent place.
+    pub fn make_translated_place(&self) -> Place<'tcx> {
+        let base_place_projected = self
+            .new_base
+            .project_deeper(self.additional_projection, self.scope.tcx);
+        trace!("  ⮑ Translated to: {base_place_projected:?}");
+        utils::retype_place(
+            base_place_projected,
+            self.scope.tcx,
+            self.scope.parent_body,
+            self.scope.parent_body_def_id,
+        )
+    }
+
+    /// Return the base version of the parent place with no child projections applied.
+    pub fn base_place(&self) -> Place<'tcx> {
+        self.new_base
+    }
+}
+
 impl<'tcx> CallingConvention<'tcx> {
     pub fn from_call_kind(
         kind: &CallKind<'tcx>,
@@ -29,71 +72,109 @@ impl<'tcx> CallingConvention<'tcx> {
         match kind {
             CallKind::AsyncPoll(poll) => CallingConvention::Async(poll.generator_data),
             CallKind::Direct => CallingConvention::Direct(args.into()),
-            CallKind::Indirect => CallingConvention::Indirect {
+            CallKind::Indirect { once_shim } => CallingConvention::Indirect {
+                once_shim: *once_shim,
                 closure_arg: args[0].clone(),
                 tupled_arguments: args[1].clone(),
             },
         }
     }
+}
 
-    pub(crate) fn translate_to_parent(
-        &self,
-        child: Place<'tcx>,
-        async_info: &AsyncInfo,
+/// This struct represents all the information necessary to translate places
+/// from a child (the callee) to its parent (caller) at the boundary of a
+/// particular function call.
+pub struct PlaceTranslator<'a, 'tcx> {
+    async_info: &'a AsyncInfo,
+    parent_body_def_id: DefId,
+    parent_body: &'a Body<'tcx>,
+    tcx: TyCtxt<'tcx>,
+    destination: Place<'tcx>,
+    calling_convention: &'a CallingConvention<'tcx>,
+    /// Governs whether the translation produces precise results (1-1
+    /// child-parent translations) or approximate one's (discarding child
+    /// projections).
+    precise: bool,
+}
+
+impl<'a, 'tcx> PlaceTranslator<'a, 'tcx> {
+    /// `destination` must be the place to which the return is assigned in the
+    /// parent (caller).
+    ///
+    /// The `precise` parameter governs whether the translation produces precise
+    /// results (1-1 child-parent translations) or approximate one's (discarding
+    /// child projections).
+    pub(crate) fn new(
+        async_info: &'a AsyncInfo,
+        parent_body_def_id: DefId,
+        parent_body: &'a Body<'tcx>,
         tcx: TyCtxt<'tcx>,
-        parent_body: &Body<'tcx>,
-        parent_def_id: DefId,
         destination: Place<'tcx>,
-    ) -> Option<Place<'tcx>> {
-        trace!("  Translating child place: {child:?}");
-        let (parent_place, child_projection) =
-            self.handle_translate(async_info, tcx, child, destination, parent_body)?;
-
-        let parent_place_projected = parent_place.project_deeper(child_projection, tcx);
-        trace!("  ⮑ Translated to: {parent_place_projected:?}");
-        Some(utils::retype_place(
-            parent_place_projected,
-            tcx,
+        calling_convention: &'a CallingConvention<'tcx>,
+        precise: bool,
+    ) -> Self {
+        Self {
+            async_info,
             parent_body,
-            parent_def_id,
-        ))
+            parent_body_def_id,
+            tcx,
+            destination,
+            calling_convention,
+            precise,
+        }
     }
 
-    pub(crate) fn handle_translate(
-        &self,
-        async_info: &AsyncInfo,
-        tcx: TyCtxt<'tcx>,
+    /// Returns a fully translated parent place. If `self.precise == true` this
+    /// place will be a precise 1-1 translation, otherwise just the base parent
+    /// place.
+    ///
+    /// Returns `None` if the input child cannot be represented in the parent.
+    pub(crate) fn translate_to_parent(&self, child: Place<'tcx>) -> Option<Place<'tcx>> {
+        let translation = self.handle_translate(child)?;
+        Some(if self.precise {
+            translation.make_translated_place()
+        } else {
+            translation.base_place()
+        })
+    }
+
+    /// Returns a calculated translation that needs to be finished.
+    ///
+    /// Returns `None` if the input child cannot be represented in the parent.
+    pub(crate) fn handle_translate<'b>(
+        &'b self,
         child: Place<'tcx>,
-        destination: Place<'tcx>,
-        parent_body: &Body<'tcx>,
-    ) -> Option<(Place<'tcx>, &[PlaceElem<'tcx>])> {
-        let result = match self {
+    ) -> Option<PlaceTranslation<'b, 'tcx>> {
+        let (new_base, additional_projection) = match self.calling_convention {
             // Async return must be handled special, because it gets wrapped in `Poll::Ready`
-            Self::Async { .. } if child.local == RETURN_PLACE => {
-                let in_poll = destination.project_deeper(
-                    &[PlaceElem::Downcast(None, async_info.poll_ready_variant_idx)],
-                    tcx,
+            CallingConvention::Async { .. } if child.local == RETURN_PLACE => {
+                let in_poll = self.destination.project_deeper(
+                    &[PlaceElem::Downcast(
+                        None,
+                        self.async_info.poll_ready_variant_idx,
+                    )],
+                    self.tcx,
                 );
-                let field_idx = async_info.poll_ready_field_idx;
+                let field_idx = self.async_info.poll_ready_field_idx;
                 let child_inner_return_type = in_poll
-                    .ty(parent_body.local_decls(), tcx)
-                    .field_ty(tcx, field_idx);
+                    .ty(self.parent_body.local_decls(), self.tcx)
+                    .field_ty(self.tcx, field_idx);
                 (
                     in_poll.project_deeper(
                         &[PlaceElem::Field(field_idx, child_inner_return_type)],
-                        tcx,
+                        self.tcx,
                     ),
                     &child.projection[..],
                 )
             }
-            _ if child.local == RETURN_PLACE => (destination, &child.projection[..]),
+            _ if child.local == RETURN_PLACE => (self.destination, &child.projection[..]),
             // Map arguments to the argument array
-            Self::Direct(args) => (
+            CallingConvention::Direct(args) => (
                 args[child.local.as_usize() - 1].place()?,
                 &child.projection[..],
             ),
             // Map arguments to projections of the future, the poll's first argument
-            Self::Async(ctx) => {
+            CallingConvention::Async(ctx) => {
                 if child.local.as_usize() == 1 {
                     (*ctx, &child.projection[..])
                 } else {
@@ -102,24 +183,52 @@ impl<'tcx> CallingConvention<'tcx> {
             }
             // Map closure captures to the first argument.
             // Map formal parameters to the second argument.
-            Self::Indirect {
+            CallingConvention::Indirect {
+                once_shim,
                 closure_arg,
                 tupled_arguments,
             } => {
                 if child.local.as_usize() == 1 {
-                    (closure_arg.place()?, &child.projection[..])
+                    // Accounting for shims
+                    let next_idx = if *once_shim {
+                        // If this is a once shim then the signature of the
+                        // function and its call don't match fully. (We are
+                        // calling a closure that takes it's `self` by reference
+                        // with a `self` by value.)
+                        if let Some(fst) = child.projection.first() {
+                            // If there is a first place it must be a deref
+                            assert_eq!(fst, &PlaceElem::Deref);
+                        } else {
+                            // We cannot remap the raw first place as it is a
+                            // reference that does not exist in the caller (as
+                            // the caller passes `self` by value.)
+                            return None;
+                        }
+                        // We skip the first projection element (a deref) to
+                        // account for the difference in signature
+                        1
+                    } else {
+                        0
+                    };
+                    (closure_arg.place()?, &child.projection[next_idx..])
                 } else {
                     let tuple_arg = tupled_arguments.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, tcx).field_ty(tcx, field);
+                    let field_ty = tuple_arg
+                        .ty(self.parent_body, self.tcx)
+                        .field_ty(self.tcx, field);
                     (
-                        tuple_arg.project_deeper(&[PlaceElem::Field(field, field_ty)], tcx),
+                        tuple_arg.project_deeper(&[PlaceElem::Field(field, field_ty)], self.tcx),
                         &child.projection[..],
                     )
                 }
             }
         };
-        Some(result)
+        Some(PlaceTranslation {
+            new_base,
+            additional_projection,
+            scope: self,
+        })
     }
 }