Add annotation support for trait methods and verify that implementations satisfy them (#25)

* add: tests for annotations in traits

* add: reference trait-side definitions for require/ensure annotations of
functions in impl blocks

Update src/analyze/local_def.rs

Co-authored-by: Hiromi Ogawa <me@coord-e.com>

* change: move extract_*_annot()s to analyze::Analyzer from
analyze::local_def::Analyzer

* change: insert type names as prefix of name for predicates in impl blocks

* add: tests for identifying struct‑bound predicates using `Self::`

* add: Identify struct-bounded predicates using `Self::` prefix

---------

Co-authored-by: Hiromi Ogawa <me@coord-e.com>

Author: coeff-aij

src/analyze.rs

@@ -16,6 +16,8 @@ use rustc_middle::mir::{self, BasicBlock, Local};
 use rustc_middle::ty::{self as mir_ty, TyCtxt};
 use rustc_span::def_id::{DefId, LocalDefId};
 
+use crate::analyze;
+use crate::annot::{AnnotFormula, AnnotParser, Resolver};
 use crate::chc;
 use crate::pretty::PrettyDisplayExt as _;
 use crate::refine::{self, BasicBlockType, TypeBuilder};
@@ -435,4 +437,54 @@ impl<'tcx> Analyzer<'tcx> {
         let body = self.tcx.optimized_mir(local_def_id);
         self.local_fn_sig_with_body(local_def_id, body)
     }
+
+    fn extract_require_annot<T>(
+        &self,
+        def_id: DefId,
+        resolver: T,
+        self_type_name: Option<String>,
+    ) -> Option<AnnotFormula<T::Output>>
+    where
+        T: Resolver,
+    {
+        let mut require_annot = None;
+        let parser = AnnotParser::new(&resolver, self_type_name);
+        for attrs in self
+            .tcx
+            .get_attrs_by_path(def_id, &analyze::annot::requires_path())
+        {
+            if require_annot.is_some() {
+                unimplemented!();
+            }
+            let ts = analyze::annot::extract_annot_tokens(attrs.clone());
+            let require = parser.parse_formula(ts).unwrap();
+            require_annot = Some(require);
+        }
+        require_annot
+    }
+
+    fn extract_ensure_annot<T>(
+        &self,
+        def_id: DefId,
+        resolver: T,
+        self_type_name: Option<String>,
+    ) -> Option<AnnotFormula<T::Output>>
+    where
+        T: Resolver,
+    {
+        let mut ensure_annot = None;
+        let parser = AnnotParser::new(&resolver, self_type_name);
+        for attrs in self
+            .tcx
+            .get_attrs_by_path(def_id, &analyze::annot::ensures_path())
+        {
+            if ensure_annot.is_some() {
+                unimplemented!();
+            }
+            let ts = analyze::annot::extract_annot_tokens(attrs.clone());
+            let ensure = parser.parse_formula(ts).unwrap();
+            ensure_annot = Some(ensure);
+        }
+        ensure_annot
+    }
 }

src/analyze/local_def.rs

@@ -53,66 +53,38 @@ pub struct Analyzer<'tcx, 'ctx> {
 }
 
 impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
-    fn extract_require_annot<T>(&self, resolver: T) -> Option<AnnotFormula<T::Output>>
-    where
-        T: annot::Resolver,
-    {
-        let mut require_annot = None;
-        for attrs in self.tcx.get_attrs_by_path(
-            self.local_def_id.to_def_id(),
-            &analyze::annot::requires_path(),
-        ) {
-            if require_annot.is_some() {
-                unimplemented!();
-            }
-            let ts = analyze::annot::extract_annot_tokens(attrs.clone());
-            let require = AnnotParser::new(&resolver).parse_formula(ts).unwrap();
-            require_annot = Some(require);
-        }
-        require_annot
-    }
-
-    fn extract_ensure_annot<T>(&self, resolver: T) -> Option<AnnotFormula<T::Output>>
-    where
-        T: annot::Resolver,
-    {
-        let mut ensure_annot = None;
-        for attrs in self.tcx.get_attrs_by_path(
-            self.local_def_id.to_def_id(),
-            &analyze::annot::ensures_path(),
-        ) {
-            if ensure_annot.is_some() {
-                unimplemented!();
-            }
-            let ts = analyze::annot::extract_annot_tokens(attrs.clone());
-            let ensure = AnnotParser::new(&resolver).parse_formula(ts).unwrap();
-            ensure_annot = Some(ensure);
-        }
-        ensure_annot
-    }
-
-    fn extract_param_annots<T>(&self, resolver: T) -> Vec<(Ident, rty::RefinedType<T::Output>)>
+    fn extract_param_annots<T>(
+        &self,
+        resolver: T,
+        self_type_name: Option<String>,
+    ) -> Vec<(Ident, rty::RefinedType<T::Output>)>
     where
         T: annot::Resolver,
     {
         let mut param_annots = Vec::new();
+        let parser = AnnotParser::new(&resolver, self_type_name);
         for attrs in self
             .tcx
             .get_attrs_by_path(self.local_def_id.to_def_id(), &analyze::annot::param_path())
         {
             let ts = analyze::annot::extract_annot_tokens(attrs.clone());
             let (ident, ts) = analyze::annot::split_param(&ts);
-            let param = AnnotParser::new(&resolver).parse_rty(ts).unwrap();
+            let param = parser.parse_rty(ts).unwrap();
             param_annots.push((ident, param));
         }
         param_annots
     }
 
-    fn extract_ret_annot<T>(&self, resolver: T) -> Option<rty::RefinedType<T::Output>>
+    fn extract_ret_annot<T>(
+        &self,
+        resolver: T,
+        self_type_name: Option<String>,
+    ) -> Option<rty::RefinedType<T::Output>>
     where
         T: annot::Resolver,
     {
         let mut ret_annot = None;
+        let parser = AnnotParser::new(&resolver, self_type_name);
         for attrs in self
             .tcx
             .get_attrs_by_path(self.local_def_id.to_def_id(), &analyze::annot::ret_path())
@@ -121,14 +93,34 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 unimplemented!();
             }
             let ts = analyze::annot::extract_annot_tokens(attrs.clone());
-            let ret = AnnotParser::new(&resolver).parse_rty(ts).unwrap();
+            let ret = parser.parse_rty(ts).unwrap();
             ret_annot = Some(ret);
         }
         ret_annot
     }
 
+    fn impl_type(&self) -> Option<rustc_middle::ty::Ty<'tcx>> {
+        use rustc_hir::def::DefKind;
+
+        let parent_def_id = self.tcx.parent(self.local_def_id.to_def_id());
+
+        if !matches!(self.tcx.def_kind(parent_def_id), DefKind::Impl { .. }) {
+            return None;
+        }
+
+        let self_ty = self.tcx.type_of(parent_def_id).instantiate_identity();
+
+        Some(self_ty)
+    }
+
     pub fn analyze_predicate_definition(&self, local_def_id: LocalDefId) {
-        let pred_name = self.tcx.item_name(local_def_id.to_def_id()).to_string();
+        // predicate's name
+        let impl_type = self.impl_type();
+        let pred_item_name = self.tcx.item_name(local_def_id.to_def_id()).to_string();
+        let pred_name = match impl_type {
+            Some(t) => t.to_string() + "_" + &pred_item_name,
+            None => pred_item_name,
+        };
 
         // function's body
         use rustc_hir::{Block, Expr, ExprKind};
@@ -252,6 +244,17 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             || (all_params_annotated && has_ret)
     }
 
+    pub fn trait_item_id(&self) -> Option<LocalDefId> {
+        let impl_item_assoc = self
+            .tcx
+            .opt_associated_item(self.local_def_id.to_def_id())?;
+        let trait_item_id = impl_item_assoc
+            .trait_item_def_id
+            .and_then(|id| id.as_local())?;
+
+        Some(trait_item_id)
+    }
+
     pub fn expected_ty(&mut self) -> rty::RefinedType {
         let sig = self
             .ctx
@@ -268,16 +271,47 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             param_resolver.push_param(input_ident.name, input_ty.to_sort());
         }
 
-        let mut require_annot = self.extract_require_annot(&param_resolver);
-        let mut ensure_annot = {
-            let output_ty = self.type_builder.build(sig.output());
-            let resolver = annot::StackedResolver::default()
-                .resolver(analyze::annot::ResultResolver::new(output_ty.to_sort()))
-                .resolver((&param_resolver).map(rty::RefinedTypeVar::Free));
-            self.extract_ensure_annot(resolver)
-        };
-        let param_annots = self.extract_param_annots(&param_resolver);
-        let ret_annot = self.extract_ret_annot(&param_resolver);
+        let output_ty = self.type_builder.build(sig.output());
+        let result_param_resolver = annot::StackedResolver::default()
+            .resolver(analyze::annot::ResultResolver::new(output_ty.to_sort()))
+            .resolver((&param_resolver).map(rty::RefinedTypeVar::Free));
+
+        let self_type_name = self.impl_type().map(|ty| ty.to_string());
+
+        let mut require_annot = self.ctx.extract_require_annot(
+            self.local_def_id.to_def_id(),
+            &param_resolver,
+            self_type_name.clone(),
+        );
+
+        let mut ensure_annot = self.ctx.extract_ensure_annot(
+            self.local_def_id.to_def_id(),
+            &result_param_resolver,
+            self_type_name.clone(),
+        );
+
+        if let Some(trait_item_id) = self.trait_item_id() {
+            tracing::info!("trait item fonud: {:?}", trait_item_id);
+            let trait_require_annot = self.ctx.extract_require_annot(
+                trait_item_id.into(),
+                &param_resolver,
+                self_type_name.clone(),
+            );
+            let trait_ensure_annot = self.ctx.extract_ensure_annot(
+                trait_item_id.into(),
+                &result_param_resolver,
+                self_type_name.clone(),
+            );
+
+            assert!(require_annot.is_none() || trait_require_annot.is_none());
+            require_annot = require_annot.or(trait_require_annot);
+
+            assert!(ensure_annot.is_none() || trait_ensure_annot.is_none());
+            ensure_annot = ensure_annot.or(trait_ensure_annot);
+        }
+
+        let param_annots = self.extract_param_annots(&param_resolver, self_type_name.clone());
+        let ret_annot = self.extract_ret_annot(&param_resolver, self_type_name);
 
         if self.is_annotated_as_callable() {
             if require_annot.is_some() || ensure_annot.is_some() {

src/annot.rs

@@ -250,6 +250,7 @@ impl<T> FormulaOrTerm<T> {
 /// A parser for refinement type annotations and formula annotations.
 struct Parser<'a, T> {
     resolver: T,
+    self_type_name: Option<String>,
     cursor: RefTokenTreeCursor<'a>,
     formula_existentials: HashMap<String, chc::Sort>,
 }
@@ -453,6 +454,7 @@ where
             TokenTree::Delimited(_, _, Delimiter::Parenthesis, s) => {
                 let mut parser = Parser {
                     resolver: self.boxed_resolver(),
+                    self_type_name: self.self_type_name.clone(),
                     cursor: s.trees(),
                     formula_existentials: self.formula_existentials.clone(),
                 };
@@ -493,6 +495,7 @@ where
 
                             let mut parser = Parser {
                                 resolver: self.boxed_resolver(),
+                                self_type_name: self.self_type_name.clone(),
                                 cursor: args.trees(),
                                 formula_existentials: self.formula_existentials.clone(),
                             };
@@ -518,11 +521,40 @@ where
                 };
                 let mut parser = Parser {
                     resolver: self.boxed_resolver(),
+                    self_type_name: self.self_type_name.clone(),
                     cursor: s.trees(),
                     formula_existentials: self.formula_existentials.clone(),
                 };
                 let args = parser.parse_arg_terms()?;
                 parser.end_of_input()?;
+
+                // Identify struct-bound predicates call such as `Self::pred()`
+                match path.segments.first() {
+                    Some(AnnotPathSegment {
+                        ident: Ident { name: symbol, .. },
+                        generic_args,
+                    }) if symbol.as_str() == "Self" && generic_args.is_empty() => {
+                        if path.segments.len() != 2 {
+                            unimplemented!("long path beginning with `Self::`");
+                        }
+
+                        let func_name = path.segments.get(1).unwrap().ident.name.as_str();
+                        let pred_name = if let Some(self_type_name) = &self.self_type_name {
+                            self_type_name.clone() + "_" + func_name
+                        } else {
+                            func_name.to_string()
+                        };
+
+                        let pred_symbol = chc::UserDefinedPred::new(pred_name);
+                        let pred = chc::Pred::UserDefined(pred_symbol);
+
+                        let atom = chc::Atom::new(pred, args);
+                        let formula = chc::Formula::Atom(atom);
+                        return Ok(FormulaOrTerm::Formula(formula));
+                    }
+                    _ => {}
+                }
+
                 let (term, sort) = path.to_datatype_ctor(args);
                 FormulaOrTerm::Term(term, sort)
             }
@@ -908,6 +940,7 @@ where
             TokenTree::Delimited(_, _, Delimiter::Parenthesis, ts) => {
                 let mut parser = Parser {
                     resolver: self.boxed_resolver(),
+                    self_type_name: self.self_type_name.clone(),
                     cursor: ts.trees(),
                     formula_existentials: self.formula_existentials.clone(),
                 };
@@ -1014,6 +1047,7 @@ where
             TokenTree::Delimited(_, _, Delimiter::Parenthesis, ts) => {
                 let mut parser = Parser {
                     resolver: self.boxed_resolver(),
+                    self_type_name: self.self_type_name.clone(),
                     cursor: ts.trees(),
                     formula_existentials: self.formula_existentials.clone(),
                 };
@@ -1050,6 +1084,7 @@ where
 
         let mut parser = Parser {
             resolver: self.boxed_resolver(),
+            self_type_name: self.self_type_name.clone(),
             cursor: ts.trees(),
             formula_existentials: self.formula_existentials.clone(),
         };
@@ -1074,6 +1109,7 @@ where
 
         let mut parser = Parser {
             resolver: RefinementResolver::new(self.boxed_resolver()),
+            self_type_name: self.self_type_name.clone(),
             cursor: parser.cursor,
             formula_existentials: self.formula_existentials.clone(),
         };
@@ -1199,11 +1235,15 @@ impl<'a, T> StackedResolver<'a, T> {
 #[derive(Debug, Clone)]
 pub struct AnnotParser<T> {
     resolver: T,
+    self_type_name: Option<String>,
 }
 
 impl<T> AnnotParser<T> {
-    pub fn new(resolver: T) -> Self {
-        Self { resolver }
+    pub fn new(resolver: T, self_type_name: Option<String>) -> Self {
+        Self {
+            resolver,
+            self_type_name,
+        }
     }
 }
 
@@ -1214,6 +1254,7 @@ where
     pub fn parse_rty(&self, ts: TokenStream) -> Result<rty::RefinedType<T::Output>> {
         let mut parser = Parser {
             resolver: &self.resolver,
+            self_type_name: self.self_type_name.clone(),
             cursor: ts.trees(),
             formula_existentials: Default::default(),
         };
@@ -1225,6 +1266,7 @@ where
     pub fn parse_formula(&self, ts: TokenStream) -> Result<AnnotFormula<T::Output>> {
         let mut parser = Parser {
             resolver: &self.resolver,
+            self_type_name: self.self_type_name.clone(),
             cursor: ts.trees(),
             formula_existentials: Default::default(),
         };