Haven experiments

crates/flowistry_pdg/src/pdg.rs

@@ -331,13 +331,40 @@ pub enum Constant {
     Uint(u64),
     Char(char),
     // Placeholder. Floats in the rust compiler are a bit weird so I'll skip them for now.
-    //Float(f64),
+    Float(FloatWrapper),
     Bool(bool),
     String(Intern<String>),
     //Unknown(Intern<String>),
     Zst(Intern<String>),
 }
 
+/// This is an unsafe wrapper around f64 in that it defines hash and equality
+/// based on bit representation. This is not in line with float semantics.
+///
+/// But I really need this to be hash and eq, hence the hack.
+#[derive(Clone, Copy, Serialize, Deserialize, Debug)]
+#[serde(transparent)]
+pub struct FloatWrapper(pub f64);
+
+impl std::cmp::PartialEq for FloatWrapper {
+    fn eq(&self, other: &Self) -> bool {
+        unsafe {
+            std::mem::transmute::<&f64, &u64>(&self.0)
+                == std::mem::transmute::<&f64, &u64>(&other.0)
+        }
+    }
+}
+
+impl std::cmp::Eq for FloatWrapper {}
+
+impl std::hash::Hash for FloatWrapper {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        unsafe {
+            std::mem::transmute::<&f64, &u64>(&self.0).hash(state);
+        }
+    }
+}
+
 impl Constant {
     pub fn int(i: impl Into<i64>) -> Self {
         Self::Int(i.into())
@@ -368,6 +395,7 @@ impl std::fmt::Display for Constant {
             Self::Uint(u) => Display::fmt(u, f),
             Self::Char(c) => Display::fmt(c, f),
             Self::String(s) => Debug::fmt(s, f),
+            Self::Float(fl) => Display::fmt(&fl.0, f),
             Self::Zst(s) => f.write_str(s),
             //Self::Unknown(u) => write!(f, "Unsupported constant: {u}"),
         }

crates/flowistry_pdg_construction/src/analysis/async_support.rs

@@ -6,10 +6,10 @@ use rustc_abi::{FieldIdx, VariantIdx};
 use rustc_hir::def_id::DefId;
 use rustc_middle::{
     mir::{
-        AggregateKind, BasicBlock, Body, Location, Operand, Place, Rvalue, Statement,
+        visit, AggregateKind, BasicBlock, Body, Location, Operand, Place, Rvalue, Statement,
         StatementKind, Terminator, TerminatorKind,
     },
-    ty::{GenericArgsRef, Instance, TyCtxt},
+    ty::{self, GenericArgsRef, Instance, TyCtxt},
 };
 use rustc_span::{source_map::Spanned, Span};
 
@@ -26,6 +26,17 @@ use crate::utils;
 pub enum AsyncType {
     Fn,
     Trait,
+    Tool,
+}
+
+impl AsyncType {
+    pub fn describe(&self) -> &'static str {
+        match self {
+            AsyncType::Fn => "async function",
+            AsyncType::Trait => "async trait",
+            AsyncType::Tool => "async tool",
+        }
+    }
 }
 
 /// Context for a call to [`Future::poll`](std::future::Future::poll), when
@@ -70,37 +81,22 @@ pub fn try_as_async_trait_function<'tcx>(
     if !has_async_trait_signature(tcx, def_id) {
         return None;
     }
-    let mut matching_statements =
-        body.basic_blocks
-            .iter_enumerated()
-            .flat_map(|(block, bbdat)| {
-                bbdat.statements.iter().enumerate().filter_map(
-                    move |(statement_index, statement)| {
-                        let (def_id, generics) = match_async_trait_assign(statement)?;
-                        Some((
-                            def_id,
-                            generics,
-                            Location {
-                                block,
-                                statement_index,
-                            },
-                        ))
-                    },
-                )
-            })
-            .collect::<Vec<_>>();
-    assert_eq!(matching_statements.len(), 1);
-    matching_statements.pop()
+    let (def_id, generics, loc, _) = find_coroutine_assign(body);
+    Some((def_id, generics, loc))
 }
 
-pub fn match_async_trait_assign<'tcx>(
-    statement: &Statement<'tcx>,
-) -> Option<(DefId, GenericArgsRef<'tcx>)> {
+pub fn match_coroutine_assign<'tcx, 'a>(
+    statement: &'a Statement<'tcx>,
+) -> Option<(
+    DefId,
+    GenericArgsRef<'tcx>,
+    &'a rustc_index::IndexVec<FieldIdx, Operand<'tcx>>,
+)> {
     match &statement.kind {
         StatementKind::Assign(box (
             _,
-            Rvalue::Aggregate(box AggregateKind::Coroutine(def_id, generic_args), _args),
-        )) => Some((*def_id, *generic_args)),
+            Rvalue::Aggregate(box AggregateKind::Coroutine(def_id, generic_args), args),
+        )) => Some((*def_id, *generic_args, args)),
         _ => None,
     }
 }
@@ -121,7 +117,6 @@ fn has_async_trait_signature(tcx: TyCtxt, def_id: DefId) -> bool {
     }
 }
 
-use rustc_middle::ty;
 fn match_pin_box_dyn_ty(lang_items: &rustc_hir::LanguageItems, t: ty::Ty) -> bool {
     let ty::TyKind::Adt(pin_ty, args) = t.kind() else {
         return false;
@@ -166,6 +161,79 @@ fn get_async_generator<'tcx>(body: &Body<'tcx>) -> (DefId, GenericArgsRef<'tcx>,
     (*def_id, generic_args, location)
 }
 
+// matches std::pin::Pin<std::boxed::Box<dyn std::future::Future<_>>
+fn has_async_tool_signature(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
+    use rustc_hir::def::DefKind;
+    if !matches!(tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+        return false;
+    }
+    let sig = tcx.fn_sig(def_id);
+    let lang_items = tcx.lang_items();
+    let ty::TyKind::Adt(adt_def, inner) = sig.skip_binder().output().skip_binder().kind() else {
+        return false;
+    };
+    if !lang_items.pin_type().is_some_and(|p| p == adt_def.did()) {
+        return false;
+    }
+    let [b_ty] = inner.as_slice() else {
+        return false;
+    };
+    let Some(f_ty) = b_ty.as_type().and_then(ty::Ty::boxed_ty) else {
+        return false;
+    };
+    let ty::TyKind::Dynamic(preds, _, ty::DynKind::Dyn) = f_ty.kind() else {
+        return false;
+    };
+    let Some(ty::ExistentialPredicate::Trait(t)) = preds.first().map(|b| b.skip_binder()) else {
+        return false;
+    };
+    lang_items.future_trait().is_some_and(|f| f == t.def_id)
+}
+
+pub fn find_coroutine_assign<'tcx, 'a>(
+    body: &'a Body<'tcx>,
+) -> (
+    DefId,
+    GenericArgsRef<'tcx>,
+    Location,
+    &'a rustc_index::IndexVec<FieldIdx, Operand<'tcx>>,
+) {
+    let mut matching_statements =
+        body.basic_blocks
+            .iter_enumerated()
+            .flat_map(|(block, bbdat)| {
+                bbdat.statements.iter().enumerate().filter_map(
+                    move |(statement_index, statement)| {
+                        let (def_id, generics, args) = match_coroutine_assign(statement)?;
+                        Some((
+                            def_id,
+                            generics,
+                            Location {
+                                block,
+                                statement_index,
+                            },
+                            args,
+                        ))
+                    },
+                )
+            })
+            .collect::<Vec<_>>();
+    assert_eq!(matching_statements.len(), 1);
+    matching_statements.pop().unwrap()
+}
+
+fn try_as_async_tool<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+    body: &Body<'tcx>,
+) -> Option<(DefId, GenericArgsRef<'tcx>, Location)> {
+    if !has_async_tool_signature(tcx, def_id) {
+        return None;
+    }
+    let (def_id, gargs, loc, _) = find_coroutine_assign(body);
+    Some((def_id, gargs, loc))
+}
+
 /// Try to interpret this function as an async function.
 ///
 /// If this is an async function it returns the [`Instance`] of the generator,
@@ -178,11 +246,10 @@ pub fn determine_async<'tcx>(
 ) -> Option<(Instance<'tcx>, Location, AsyncType)> {
     let ((generator_def_id, args, loc), asyncness) = if is_async(tcx, def_id) {
         (get_async_generator(body), AsyncType::Fn)
+    } else if let Some(g) = try_as_async_trait_function(tcx, def_id, body) {
+        (g, AsyncType::Trait)
     } else {
-        (
-            try_as_async_trait_function(tcx, def_id, body)?,
-            AsyncType::Trait,
-        )
+        (try_as_async_tool(tcx, def_id, body)?, AsyncType::Tool)
     };
     let typing_env = body.typing_env(tcx).with_post_analysis_normalized(tcx);
     let generator_fn = utils::try_resolve_function(tcx, generator_def_id, typing_env, args)?;

crates/flowistry_pdg_construction/src/analysis/global/call_tree_visit.rs

@@ -12,8 +12,12 @@
 
 use std::{borrow::Cow, hash::Hash, rc::Rc};
 
+use flowistry::mir::FlowistryInput;
 use flowistry_pdg::{CallString, GlobalLocation, RichLocation};
-use rustc_middle::{mir::Location, ty::Instance};
+use rustc_middle::{
+    mir::{self, Location},
+    ty::Instance,
+};
 
 use crate::{analysis::global::partial_graph::NodeKey, DepNodeKind, MemoPdgConstructor};
 
@@ -80,6 +84,13 @@ impl<'tcx, 'c, K: Clone> VisitDriver<'tcx, 'c, K> {
             },
         )
     }
+
+    pub fn current_body(&self) -> &'tcx mir::Body<'tcx> {
+        self.memo
+            .body_cache()
+            .get(self.current_function().def_id())
+            .body()
+    }
 }
 
 pub trait Visitor<'tcx, K: Hash + Eq + Clone> {

crates/flowistry_pdg_construction/src/analysis/global/mod.rs

@@ -25,7 +25,7 @@ use crate::{
     callback::DefaultCallback,
     constants::PlaceOrConst,
     source_access::{self, BodyCache, CachedBody},
-    utils::{manufacture_substs_for, try_resolve_function, TwoLevelCache},
+    utils::{manufacture_substs_for, try_resolve_function, PlaceConflictContext, TwoLevelCache},
     CallChangeCallback,
 };
 
@@ -84,6 +84,7 @@ pub struct MemoPdgConstructor<'tcx, K> {
     pub(crate) body_cache: Rc<source_access::BodyCache<'tcx>>,
     disable_cache: bool,
     relaxed: bool,
+    pub(crate) place_conflict_context: PlaceConflictContext<'tcx>,
 }
 
 impl<'tcx, K: Default> MemoPdgConstructor<'tcx, K> {
@@ -105,6 +106,7 @@ impl<'tcx, K: Default> MemoPdgConstructor<'tcx, K> {
             body_cache,
             disable_cache: false,
             relaxed: false,
+            place_conflict_context: PlaceConflictContext::new(tcx),
         }
     }
 }
@@ -124,6 +126,7 @@ impl<'tcx, K> MemoPdgConstructor<'tcx, K> {
             body_cache: Rc::new(BodyCache::new(tcx)),
             disable_cache: false,
             relaxed: false,
+            place_conflict_context: PlaceConflictContext::new(tcx),
         }
     }
 
@@ -658,7 +661,7 @@ enum Input<'tcx> {
     },
 }
 
-struct GraphSizeEstimator {
+pub struct GraphSizeEstimator {
     nodes: usize,
     edges: usize,
     functions: usize,
@@ -667,7 +670,7 @@ struct GraphSizeEstimator {
 
 #[allow(dead_code)]
 impl GraphSizeEstimator {
-    fn new() -> Self {
+    pub fn new() -> Self {
         Self {
             nodes: 0,
             edges: 0,
@@ -676,7 +679,7 @@ impl GraphSizeEstimator {
         }
     }
 
-    fn format_size(&self) -> String {
+    pub fn format_size(&self) -> String {
         format!(
             "nodes: {}, edges: {}, functions: {}, call_string_length: {}",
             HumanInt(self.nodes),

crates/flowistry_pdg_construction/src/analysis/local/mod.rs

@@ -285,7 +285,11 @@ impl<'tcx, 'a, K> LocalAnalysis<'tcx, 'a, K> {
                             place.local == alias.local
                         } else {
                             trace!("Checking conflict status of {place:?} and {alias:?}");
-                            utils::places_conflict(self.tcx(), &self.mono_body, *place, alias)
+                            self.memo.place_conflict_context.places_conflict(
+                                &self.mono_body,
+                                *place,
+                                alias,
+                            )
                         }
                     });
 
@@ -366,7 +370,7 @@ impl<'tcx, 'a, K> LocalAnalysis<'tcx, 'a, K> {
     }
 
     fn fmt_fn(&self, def_id: DefId) -> String {
-        self.tcx().def_path_str(def_id)
+        format!("{} ({}:{})", self.tcx().def_path_str(def_id), def_id.krate.as_u32(), def_id.index.as_u32())
     }
 }
 
@@ -757,7 +761,17 @@ impl<'tcx, 'a, K: Hash + Eq + Clone> LocalAnalysis<'tcx, 'a, K> {
                     } else {
                         continue;
                     };
-                    let src = final_state.get_place_node(*place, location.into()).unwrap();
+                    let Some(src) = final_state.get_place_node(*place, location.into()) else {
+                        let span = match location {
+                            RichLocation::Location(l) => self.mono_body.source_info(*l).span,
+                            _ => self.mono_body.span,
+                        };
+                        // CORNER CUTTING: we should investigate why this happens.
+                        self.tcx()
+                            .dcx()
+                            .span_warn(span, format!("could not find reference to {place:?} here"));
+                        continue;
+                    };
                     let eloc = RichLocation::End;
                     let key = NodeKey::for_place(*place, eloc.into());
                     let prior = final_state.get_node(&key);