From c9707f4db90265b65b811c78a3ea5ce61a1ae0bc Mon Sep 17 00:00:00 2001
From: Justus Adam <justbldr@amazon.com>
Date: Fri, 6 Sep 2024 17:02:03 -0700
Subject: [PATCH] Make sure we check for external functions too

---
 crates/paralegal-flow/src/ann/db.rs | 53 ++++++++++++-----------------
 1 file changed, 21 insertions(+), 32 deletions(-)

diff --git a/crates/paralegal-flow/src/ann/db.rs b/crates/paralegal-flow/src/ann/db.rs
index b0a0f14539..1d25ae7ed6 100644
--- a/crates/paralegal-flow/src/ann/db.rs
+++ b/crates/paralegal-flow/src/ann/db.rs
@@ -207,26 +207,16 @@ impl<'tcx> MarkerCtx<'tcx> {
         res: impl Into<MaybeMonomorphized<'tcx>>,
     ) -> impl Iterator<Item = Identifier> + '_ {
         let res = res.into();
-        // TODO this check is wrong. This should either check whether this gets inlined or be removed altogether.
-        if res.def_id().is_local() {
-            let mut direct_markers = self
-                .combined_markers(res.def_id())
-                .map(|m| m.marker)
-                .peekable();
-            let non_direct = direct_markers
-                .peek()
-                .is_none()
-                .then(|| self.get_reachable_markers(res));
-
-            Either::Right(direct_markers.chain(non_direct.into_iter().flatten().copied()))
-        } else {
-            Either::Left(
-                self.all_function_markers(res)
-                    .map(|m| m.0.marker)
-                    .collect::<Vec<_>>(),
-            )
-        }
-        .into_iter()
+        let mut direct_markers = self
+            .combined_markers(res.def_id())
+            .map(|m| m.marker)
+            .peekable();
+        let non_direct = direct_markers
+            .peek()
+            .is_none()
+            .then(|| self.get_reachable_markers(res));
+
+        direct_markers.chain(non_direct.into_iter().flatten().copied())
     }
 
     /// If the transitive marker cache did not contain the answer, this is what
@@ -309,23 +299,22 @@ impl<'tcx> MarkerCtx<'tcx> {
         );
 
         if let Some(model) = self.has_stub(res.def_id()) {
-            let MaybeMonomorphized::Monomorphized(instance) = &mut res else {
+            if let MaybeMonomorphized::Monomorphized(instance) = &mut res {
+                if let Ok(new_instance) = model.resolve_alternate_instance(
+                    self.tcx(),
+                    *instance,
+                    param_env,
+                    terminator.source_info.span,
+                ) {
+                    v.extend(self.get_reachable_and_self_markers(new_instance));
+                }
+            } else {
                 self.span_err(
                     terminator.source_info.span,
                     "Could not apply stub to an partially resolved function",
                 );
-                return v.into_iter();
             };
-            if let Ok(new_instance) = model.resolve_alternate_instance(
-                self.tcx(),
-                *instance,
-                param_env,
-                terminator.source_info.span,
-            ) {
-                *instance = new_instance;
-            } else {
-                return v.into_iter();
-            }
+            return v.into_iter();
         }
 
         v.extend(self.get_reachable_and_self_markers(res));