Enable to write thrust_macros::{ensures,requires} inside impl

src/analyze.rs

@@ -610,16 +610,10 @@ impl<'tcx> Analyzer<'tcx> {
                 );
                 continue;
             };
-            let rustc_hir::QPath::Resolved(_, path) = qpath else {
+            let typeck = self.tcx.typeck(local_def_id);
+            let rustc_hir::def::Res::Def(_, def_id) = typeck.qpath_res(&qpath, expr.hir_id) else {
                 self.tcx.dcx().span_err(
                     expr.span,
-                    "annotated path is expected to be a resolved path",
-                );
-                continue;
-            };
-            let rustc_hir::def::Res::Def(_, def_id) = path.res else {
-                self.tcx.dcx().span_err(
-                    path.span,
                     "annotated path is expected to refer to a definition",
                 );
                 continue;

src/analyze/local_def.rs

@@ -400,11 +400,28 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     /// semicolon) in the function body block.
     pub fn extern_spec_fn_target_def_id(&self) -> DefId {
         let node = self.tcx.hir_node_by_def_id(self.local_def_id);
-        let rustc_hir::Node::Item(item) = node else {
-            panic!("extern_spec_fn must be a function item");
-        };
-        let rustc_hir::ItemKind::Fn(_, _, body_id) = item.kind else {
-            panic!("extern_spec_fn must be a function");
+        let body_id = match node {
+            rustc_hir::Node::Item(item) => {
+                let rustc_hir::ItemKind::Fn(_, _, body_id) = item.kind else {
+                    panic!("extern_spec_fn must be a function");
+                };
+                body_id
+            }
+            rustc_hir::Node::ImplItem(impl_item) => {
+                let rustc_hir::ImplItemKind::Fn(_, body_id) = impl_item.kind else {
+                    panic!("extern_spec_fn must be a function");
+                };
+                body_id
+            }
+            rustc_hir::Node::TraitItem(trait_item) => {
+                let rustc_hir::TraitItemKind::Fn(_, rustc_hir::TraitFn::Provided(body_id)) =
+                    trait_item.kind
+                else {
+                    panic!("extern_spec_fn must be a function with a body");
+                };
+                body_id
+            }
+            _ => panic!("extern_spec_fn must be a function item or impl item"),
         };
 
         let body = self.tcx.hir().body(body_id);

tests/ui/fail/annot_struct_impl.rs

@@ -0,0 +1,53 @@
+//@error-in-other-file: Unsat
+
+struct VecWrap<T> {
+    inner: Vec<T>
+}
+
+impl<T> thrust_models::Model for VecWrap<T> where T: thrust_models::Model {
+  type Ty = (
+      thrust_models::model::Array<
+          thrust_models::model::Int,
+          <T as thrust_models::Model>::Ty
+      >,
+      thrust_models::model::Int,
+  );
+}
+
+#[thrust_macros::context]
+impl<T> VecWrap<T> {
+  #[thrust::trusted]
+  #[thrust_macros::ensures(result.1 == 0)]
+  fn new() -> Self {
+      VecWrap { inner: Vec::new() }
+  }
+
+  #[thrust::trusted]
+  #[thrust_macros::ensures((!self).0 == (*self).0.store((*self).1, elem))]
+  #[thrust_macros::ensures((!self).1 == (*self).1)]
+  fn push(&mut self, elem: T) {
+      self.inner.push(elem);
+  }
+
+  #[thrust::trusted]
+  #[thrust_macros::ensures(result == (*self).1)]
+  fn len(&self) -> usize {
+      self.inner.len()
+  }
+
+  #[thrust::trusted]
+  #[thrust_macros::requires(index < (*self).1)]
+  #[thrust_macros::ensures(*result == (*self).0[index])]
+  fn get(&self, index: usize) -> &T {
+      &self.inner[index]
+  }
+}
+
+fn main() {
+    let mut v = VecWrap::new();
+    v.push(10);
+    v.push(20);
+    assert!(v.len() == 2);
+    assert!(*v.get(0) == 10);
+    assert!(*v.get(1) == 20);
+}

tests/ui/pass/annot_struct_impl.rs

@@ -0,0 +1,53 @@
+//@check-pass
+
+struct VecWrap<T> {
+    inner: Vec<T>
+}
+
+impl<T> thrust_models::Model for VecWrap<T> where T: thrust_models::Model {
+  type Ty = (
+      thrust_models::model::Array<
+          thrust_models::model::Int,
+          <T as thrust_models::Model>::Ty
+      >,
+      thrust_models::model::Int,
+  );
+}
+
+#[thrust_macros::context]
+impl<T> VecWrap<T> {
+  #[thrust::trusted]
+  #[thrust_macros::ensures(result.1 == 0)]
+  fn new() -> Self {
+      VecWrap { inner: Vec::new() }
+  }
+
+  #[thrust::trusted]
+  #[thrust_macros::ensures((!self).0 == (*self).0.store((*self).1, elem))]
+  #[thrust_macros::ensures((!self).1 == (*self).1 + 1)]
+  fn push(&mut self, elem: T) {
+      self.inner.push(elem);
+  }
+
+  #[thrust::trusted]
+  #[thrust_macros::ensures(result == (*self).1)]
+  fn len(&self) -> usize {
+      self.inner.len()
+  }
+
+  #[thrust::trusted]
+  #[thrust_macros::requires(index < (*self).1)]
+  #[thrust_macros::ensures(*result == (*self).0[index])]
+  fn get(&self, index: usize) -> &T {
+      &self.inner[index]
+  }
+}
+
+fn main() {
+    let mut v = VecWrap::new();
+    v.push(10);
+    v.push(20);
+    assert!(v.len() == 2);
+    assert!(*v.get(0) == 10);
+    assert!(*v.get(1) == 20);
+}

thrust-macros/Cargo.toml

@@ -10,4 +10,4 @@ proc-macro = true
 [dependencies]
 proc-macro2 = "1"
 quote = "1"
-syn = { version = "2", features = ["full"] }
+syn = { version = "2", features = ["full", "visit", "visit-mut"] }