Prioritize trait method spec keys

Author: coord-e

src/analyze/crate_.rs

@@ -60,19 +60,41 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     }
 
     fn refine_local_defs(&mut self) {
+        // Prioritize trait method specs so that they are always available when refining trait impl
+        // functions; see the partialeq_impl.rs case.
         let mut keys = self.tcx.mir_keys(()).clone();
+        let mut formula_fn_keys = HashSet::new();
+        let mut trait_method_spec_keys = HashSet::new();
         for local_def_id in self.tcx.mir_keys(()) {
             let analyzer = self.ctx.local_def_analyzer(*local_def_id);
             if analyzer.is_annotated_as_extern_spec_fn() {
-                if let Some(target_def_id) = analyzer.extern_spec_fn_target_def_id().as_local() {
+                let target_def_id = analyzer.extern_spec_fn_target_def_id();
+                if let Some(local_target_def_id) = target_def_id.as_local() {
                     self.skip_analysis.insert(local_target_def_id);
-                    keys.swap_remove(&target_def_id);
+                    keys.swap_remove(&local_target_def_id);
                 }
+                if matches!(
+                    self.tcx.def_kind(target_def_id),
+                    rustc_hir::def::DefKind::AssocFn
+                ) {
+                    trait_method_spec_keys.insert(*local_def_id);
+                    keys.swap_remove(local_def_id);
+                }
+            }
+            if analyzer.is_annotated_as_formula_fn() {
+                formula_fn_keys.insert(*local_def_id);
+                keys.swap_remove(local_def_id);
             }
             if !self.tcx.def_kind(*local_def_id).is_fn_like() {
                 keys.swap_remove(local_def_id);
             }
         }
+        for local_def_id in &formula_fn_keys {
+            self.refine_fn_def(*local_def_id);
+        }
+        for local_def_id in &trait_method_spec_keys {
+            self.refine_fn_def(*local_def_id);
+        }
         for local_def_id in &keys {
             self.refine_fn_def(*local_def_id);
         }

tests/ui/fail/partialeq_impl.rs

@@ -0,0 +1,30 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+
+enum X {
+    A(i32),
+}
+
+impl thrust_models::Model for X {
+    type Ty = Self;
+}
+
+impl PartialEq for X {
+    fn eq(&self, other: &X) -> bool {
+        match (self, other) {
+            (Self::A(lhs), Self::A(rhs)) => !lhs.eq(rhs),
+        }
+    }
+}
+
+#[thrust::trusted]
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+fn rand() -> X {
+    unimplemented!()
+}
+
+fn main() {
+    let x = rand();
+    assert!(x == x);
+}

tests/ui/pass/partialeq_impl.rs

@@ -0,0 +1,30 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+
+enum X {
+    A(i32),
+}
+
+impl thrust_models::Model for X {
+    type Ty = Self;
+}
+
+impl PartialEq for X {
+    fn eq(&self, other: &X) -> bool {
+        match (self, other) {
+            (Self::A(lhs), Self::A(rhs)) => lhs.eq(rhs),
+        }
+    }
+}
+
+#[thrust::trusted]
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(true)]
+fn rand() -> X {
+    unimplemented!()
+}
+
+fn main() {
+    let x = rand();
+    assert!(x == x);
+}