Handling of foreign functions that may execute local code

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 = "beca0fdda699591052445a1da840c072fffa1f21"
+rev = "d7bde3ae16a7137de594c2dd811030cd761a8993"
 
 [replace."rustc_plugin:0.7.4-nightly-2023-08-25"]
 # path = "../rustc_plugin/crates/rustc_plugin"
 git = "https://github.com/JustusAdam/rustc_plugin"
-rev = "beca0fdda699591052445a1da840c072fffa1f21"
+rev = "d7bde3ae16a7137de594c2dd811030cd761a8993"

crates/flowistry_pdg_construction/src/body_cache.rs

@@ -99,7 +99,7 @@ impl<'tcx> BodyCache<'tcx> {
     /// Serve the body from the cache or read it from the disk.
     ///
     /// Returns `None` if the policy forbids loading from this crate.
-    pub fn get(&self, key: DefId) -> Option<&'tcx CachedBody<'tcx>> {
+    pub fn get(&self, key: DefId) -> &'tcx CachedBody<'tcx> {
         let cbody = self.cache.get(key, |_| load_body_and_facts(self.tcx, key));
         // SAFETY: Theoretically this struct may not outlive the body, but
         // to simplify lifetimes flowistry uses 'tcx anywhere. But if we

crates/flowistry_pdg_construction/src/callback.rs

@@ -1,44 +1,53 @@
-//! CAllbacks to influence graph construction and their supporting types.
+//! Callbacks to influence graph construction and their supporting types.
 
 use flowistry_pdg::{rustc_portable::Location, CallString};
-use rustc_middle::ty::Instance;
+
+use rustc_middle::{
+    mir::{self, Operand},
+    ty::{Instance, ParamEnv},
+};
+use rustc_span::Span;
+
+use crate::calling_convention::CallingConvention;
 
 pub trait CallChangeCallback<'tcx> {
-    fn on_inline(&self, info: CallInfo<'tcx>) -> CallChanges;
+    fn on_inline(&self, info: CallInfo<'tcx, '_>) -> CallChanges<'tcx>;
 
     fn on_inline_miss(
         &self,
         _resolution: Instance<'tcx>,
+        _param_env: ParamEnv<'tcx>,
         _loc: Location,
         _under_analysis: Instance<'tcx>,
-        _call_string: Option<CallString>,
         _reason: InlineMissReason,
+        _call_span: Span,
     ) {
     }
 }
 
 pub struct CallChangeCallbackFn<'tcx> {
-    f: Box<dyn Fn(CallInfo<'tcx>) -> CallChanges + 'tcx>,
+    f: Box<dyn Fn(CallInfo<'tcx, '_>) -> CallChanges<'tcx> + 'tcx>,
 }
 
 impl<'tcx> CallChangeCallbackFn<'tcx> {
-    pub fn new(f: impl Fn(CallInfo<'tcx>) -> CallChanges + 'tcx) -> Self {
+    pub fn new(f: impl Fn(CallInfo<'tcx, '_>) -> CallChanges<'tcx> + 'tcx) -> Self {
         Self { f: Box::new(f) }
     }
 }
 
 impl<'tcx> CallChangeCallback<'tcx> for CallChangeCallbackFn<'tcx> {
-    fn on_inline(&self, info: CallInfo<'tcx>) -> CallChanges {
+    fn on_inline(&self, info: CallInfo<'tcx, '_>) -> CallChanges<'tcx> {
         (self.f)(info)
     }
 }
 
 #[derive(Debug)]
 pub enum InlineMissReason {
     Async(String),
+    TraitMethod,
 }
 
-impl Default for CallChanges {
+impl<'tcx> Default for CallChanges<'tcx> {
     fn default() -> Self {
         CallChanges {
             skip: SkipCall::NoSkip,
@@ -47,7 +56,7 @@ impl Default for CallChanges {
 }
 
 /// Information about the function being called.
-pub struct CallInfo<'tcx> {
+pub struct CallInfo<'tcx, 'mir> {
     /// The potentially-monomorphized resolution of the callee.
     pub callee: Instance<'tcx>,
 
@@ -60,29 +69,42 @@ pub struct CallInfo<'tcx> {
 
     /// Would the PDG for this function be served from the cache.
     pub is_cached: bool,
+
+    pub span: Span,
+
+    pub arguments: &'mir [Operand<'tcx>],
+
+    pub caller_body: &'mir mir::Body<'tcx>,
+    pub param_env: ParamEnv<'tcx>,
 }
 
 /// User-provided changes to the default PDG construction behavior for function calls.
 ///
 /// Construct [`CallChanges`] via [`CallChanges::default`].
 #[derive(Debug)]
-pub struct CallChanges {
-    pub(crate) skip: SkipCall,
+pub struct CallChanges<'tcx> {
+    pub(crate) skip: SkipCall<'tcx>,
 }
 
 /// Whether or not to skip recursing into a function call during PDG construction.
 #[derive(Debug)]
-pub enum SkipCall {
+pub enum SkipCall<'tcx> {
     /// Skip the function, and perform a modular approxmation of its effects.
     Skip,
 
     /// Recurse into the function as normal.
     NoSkip,
+
+    /// Replace with a call to this other function and arguments.
+    Replace {
+        instance: Instance<'tcx>,
+        calling_convention: CallingConvention<'tcx>,
+    },
 }
 
-impl CallChanges {
-    /// Inidicate whether or not to skip recursing into the given function.
-    pub fn with_skip(self, skip: SkipCall) -> Self {
+impl<'tcx> CallChanges<'tcx> {
+    /// Indicate whether or not to skip recursing into the given function.
+    pub fn with_skip(self, skip: SkipCall<'tcx>) -> Self {
         CallChanges { skip }
     }
 }

crates/flowistry_pdg_construction/src/calling_convention.rs

@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use flowistry_pdg::rustc_portable::DefId;
 use log::trace;
 use rustc_abi::FieldIdx;
@@ -9,26 +11,27 @@ use rustc_middle::{
 
 use crate::{async_support::AsyncInfo, local_analysis::CallKind, utils};
 
-pub enum CallingConvention<'tcx, 'a> {
-    Direct(&'a [Operand<'tcx>]),
+#[derive(Debug)]
+pub enum CallingConvention<'tcx> {
+    Direct(Box<[Operand<'tcx>]>),
     Indirect {
-        closure_arg: &'a Operand<'tcx>,
-        tupled_arguments: &'a Operand<'tcx>,
+        closure_arg: Operand<'tcx>,
+        tupled_arguments: Operand<'tcx>,
     },
     Async(Place<'tcx>),
 }
 
-impl<'tcx, 'a> CallingConvention<'tcx, 'a> {
+impl<'tcx> CallingConvention<'tcx> {
     pub fn from_call_kind(
         kind: &CallKind<'tcx>,
-        args: &'a [Operand<'tcx>],
-    ) -> CallingConvention<'tcx, 'a> {
+        args: Cow<'_, [Operand<'tcx>]>,
+    ) -> CallingConvention<'tcx> {
         match kind {
             CallKind::AsyncPoll(poll) => CallingConvention::Async(poll.generator_data),
-            CallKind::Direct => CallingConvention::Direct(args),
+            CallKind::Direct => CallingConvention::Direct(args.into()),
             CallKind::Indirect => CallingConvention::Indirect {
-                closure_arg: &args[0],
-                tupled_arguments: &args[1],
+                closure_arg: args[0].clone(),
+                tupled_arguments: args[1].clone(),
             },
         }
     }

crates/flowistry_pdg_construction/src/construct.rs

@@ -1,4 +1,4 @@
-use std::rc::Rc;
+use std::{borrow::Cow, rc::Rc};
 
 use either::Either;
 use flowistry::mir::FlowistryInput;
@@ -99,21 +99,23 @@ impl<'tcx> MemoPdgConstructor<'tcx> {
         )
         .unwrap();
 
-        self.construct_for(resolution).unwrap()
+        self.construct_for(resolution)
+            .expect("Invariant broken, entrypoint cannot have been recursive.")
     }
 
+    /// Construct a  graph for this instance of return it from the cache.
+    ///
+    /// Returns `None` if this is a recursive call trying to construct the graph again.
     pub(crate) fn construct_for<'a>(
         &'a self,
         resolution: Instance<'tcx>,
     ) -> Option<&'a PartialGraph<'tcx>> {
-        self.pdg_cache
-            .try_retrieve(resolution, |_| {
-                let g = LocalAnalysis::new(self, resolution)?.construct_partial();
-                trace!("Computed new for {resolution:?}");
-                g.check_invariants();
-                Some(g)
-            })
-            .as_success()
+        self.pdg_cache.get_maybe_recursive(resolution, |_| {
+            let g = LocalAnalysis::new(self, resolution).construct_partial();
+            trace!("Computed new for {resolution:?}");
+            g.check_invariants();
+            g
+        })
     }
 
     /// Has a PDG been constructed for this instance before?
@@ -130,7 +132,7 @@ impl<'tcx> MemoPdgConstructor<'tcx> {
         if let Some((generator, loc, _ty)) = determine_async(
             self.tcx,
             function.to_def_id(),
-            self.body_cache.get(function.to_def_id()).unwrap().body(),
+            self.body_cache.get(function.to_def_id()).body(),
         ) {
             // TODO remap arguments
 
@@ -152,7 +154,7 @@ impl<'tcx> MemoPdgConstructor<'tcx> {
     /// Try to retrieve or load a body for this id.
     ///
     /// Returns `None` if the loading policy forbids loading from this crate.
-    pub fn body_for_def_id(&self, key: DefId) -> Option<&'tcx CachedBody<'tcx>> {
+    pub fn body_for_def_id(&self, key: DefId) -> &'tcx CachedBody<'tcx> {
         self.body_cache.get(key)
     }
 
@@ -259,8 +261,8 @@ impl<'mir, 'tcx> ResultsVisitor<'mir, 'tcx, LocalAnalysisResults<'tcx, 'mir>>
             if matches!(
                 constructor.determine_call_handling(
                     location,
-                    func,
-                    args,
+                    Cow::Borrowed(func),
+                    Cow::Borrowed(args),
                     terminator.source_info.span
                 ),
                 Some(CallHandling::Ready { .. })
@@ -331,9 +333,12 @@ impl<'tcx> PartialGraph<'tcx> {
             function: constructor.def_id,
         };
 
-        let Some(handling) =
-            constructor.determine_call_handling(location, func, args, terminator.source_info.span)
-        else {
+        let Some(handling) = constructor.determine_call_handling(
+            location,
+            Cow::Borrowed(func),
+            Cow::Borrowed(args),
+            terminator.source_info.span,
+        ) else {
             return false;
         };
 

crates/flowistry_pdg_construction/src/lib.rs

@@ -31,7 +31,7 @@ use rustc_middle::ty::TyCtxt;
 mod approximation;
 mod async_support;
 pub mod body_cache;
-mod calling_convention;
+pub mod calling_convention;
 mod construct;
 pub mod encoder;
 pub mod graph;