Enable to annotate functions with formula_fn

Author: coord-e

src/analyze.rs

@@ -472,20 +472,71 @@ impl<'tcx> Analyzer<'tcx> {
         self.fn_sig_with_body(def_id, body)
     }
 
+    fn extract_path_with_attr(
+        &self,
+        local_def_id: LocalDefId,
+        attr_path: &[Symbol],
+    ) -> Option<DefId> {
+        let map = self.tcx.hir();
+        let Some(body_id) = map.maybe_body_owned_by(local_def_id) else {
+            return None;
+        };
+        let body = map.body(body_id);
+
+        let rustc_hir::ExprKind::Block(block, _) = body.value.kind else {
+            return None;
+        };
+        for stmt in block.stmts {
+            if map
+                .attrs(stmt.hir_id)
+                .iter()
+                .all(|attr| !attr.path_matches(attr_path))
+            {
+                continue;
+            }
+            let rustc_hir::StmtKind::Semi(expr) = stmt.kind else {
+                panic!(
+                    "annotated path is expected to be a semi statement, but found: {:?}",
+                    stmt.kind
+                );
+            };
+            let rustc_hir::ExprKind::Path(qpath) = expr.kind else {
+                panic!(
+                    "annotated path is expected to be a path expression, but found: {:?}",
+                    expr.kind
+                );
+            };
+            let rustc_hir::QPath::Resolved(_, path) = qpath else {
+                panic!(
+                    "qpath is unexpectedly not resolved in annotated path: {:?}",
+                    qpath
+                );
+            };
+            let rustc_hir::def::Res::Def(_, def_id) = path.res else {
+                panic!(
+                    "annotated path is expected to refer to a def, but found: {:?}",
+                    path.res
+                );
+            };
+            return Some(def_id);
+        }
+        None
+    }
+
     fn extract_require_annot<T>(
         &self,
-        def_id: DefId,
+        local_def_id: LocalDefId,
         resolver: T,
         self_type_name: Option<String>,
     ) -> Option<AnnotFormula<T::Output>>
     where
-        T: Resolver,
+        T: Resolver<Output = rty::FunctionParamIdx>,
     {
         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())
+            .get_attrs_by_path(local_def_id.to_def_id(), &analyze::annot::requires_path())
         {
             if require_annot.is_some() {
                 unimplemented!();
@@ -494,23 +545,47 @@ impl<'tcx> Analyzer<'tcx> {
             let require = parser.parse_formula(ts).unwrap();
             require_annot = Some(require);
         }
+
+        if let Some(formula_def_id) =
+            self.extract_path_with_attr(local_def_id, &analyze::annot::requires_path())
+        {
+            if require_annot.is_some() {
+                unimplemented!();
+            }
+            let Some(formula_fn) = self.formula_fns.get(&formula_def_id) else {
+                panic!(
+                    "require annotation {:?} is not a formula function",
+                    formula_def_id
+                );
+            };
+            let Some(annot) = formula_fn.to_require_annot(self.fn_sig(local_def_id.to_def_id()))
+            else {
+                panic!(
+                    "require annotation {:?} has incompatible signature with {:?}",
+                    formula_def_id, local_def_id
+                );
+            };
+            require_annot = Some(annot);
+        }
+
         require_annot
     }
 
     fn extract_ensure_annot<T>(
         &self,
-        def_id: DefId,
+        local_def_id: LocalDefId,
         resolver: T,
         self_type_name: Option<String>,
     ) -> Option<AnnotFormula<T::Output>>
     where
-        T: Resolver,
+        T: Resolver<Output = rty::RefinedTypeVar<rty::FunctionParamIdx>>,
     {
         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())
+            .get_attrs_by_path(local_def_id.to_def_id(), &analyze::annot::ensures_path())
         {
             if ensure_annot.is_some() {
                 unimplemented!();
@@ -519,6 +594,29 @@ impl<'tcx> Analyzer<'tcx> {
             let ensure = parser.parse_formula(ts).unwrap();
             ensure_annot = Some(ensure);
         }
+
+        if let Some(formula_def_id) =
+            self.extract_path_with_attr(local_def_id, &analyze::annot::ensures_path())
+        {
+            if ensure_annot.is_some() {
+                unimplemented!();
+            }
+            let Some(formula_fn) = self.formula_fns.get(&formula_def_id) else {
+                panic!(
+                    "ensure annotation {:?} is not a formula function",
+                    formula_def_id
+                );
+            };
+            let Some(annot) = formula_fn.to_ensure_annot(self.fn_sig(local_def_id.to_def_id()))
+            else {
+                panic!(
+                    "ensure annotation {:?} has incompatible signature with {:?}",
+                    formula_def_id, local_def_id
+                );
+            };
+            ensure_annot = Some(annot);
+        }
+
         ensure_annot
     }
 

src/analyze/local_def.rs

@@ -214,22 +214,30 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
 
     // TODO: unify this logic with extraction functions above
     pub fn is_fully_annotated(&self) -> bool {
-        let has_require = self
+        let has_requires = self
             .tcx
             .get_attrs_by_path(
                 self.local_def_id.to_def_id(),
                 &analyze::annot::requires_path(),
             )
             .next()
-            .is_some();
-        let has_ensure = self
+            .is_some()
+            || self
+                .ctx
+                .extract_path_with_attr(self.local_def_id, &analyze::annot::requires_path())
+                .is_some();
+        let has_ensures = self
             .tcx
             .get_attrs_by_path(
                 self.local_def_id.to_def_id(),
                 &analyze::annot::ensures_path(),
             )
             .next()
-            .is_some();
+            .is_some()
+            || self
+                .ctx
+                .extract_path_with_attr(self.local_def_id, &analyze::annot::ensures_path())
+                .is_some();
         let annotated_params: Vec<_> = self
             .tcx
             .get_attrs_by_path(self.local_def_id.to_def_id(), &analyze::annot::param_path())
@@ -250,7 +258,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             .iter()
             .all(|ident| annotated_params.contains(ident));
         self.is_annotated_as_callable()
-            || (has_require && has_ensure)
+            || (has_requires && has_ensures)
             || (all_params_annotated && has_ret)
     }
 
@@ -289,26 +297,26 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         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(),
+            self.local_def_id,
             &param_resolver,
             self_type_name.clone(),
         );
 
         let mut ensure_annot = self.ctx.extract_ensure_annot(
-            self.local_def_id.to_def_id(),
+            self.local_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(),
+                trait_item_id,
                 &param_resolver,
                 self_type_name.clone(),
             );
             let trait_ensure_annot = self.ctx.extract_ensure_annot(
-                trait_item_id.into(),
+                trait_item_id,
                 &result_param_resolver,
                 self_type_name.clone(),
             );

tests/ui/fail/annot_formula_fn.rs

@@ -0,0 +1,54 @@
+//@error-in-other-file: Unsat
+
+#[thrust::formula_fn]
+fn _thrust_requires_rand_except(x: i64) -> bool {
+    true
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_rand_except(result: i64, x: i64) -> bool {
+    result != x
+}
+
+fn rand_except(x: i64) -> i64 {
+    #[thrust::requires]
+    _thrust_requires_rand_except;
+    #[thrust::ensures]
+    _thrust_ensures_rand_except;
+
+    if x == 0 {
+        1
+    } else {
+        0
+    }
+}
+
+#[thrust::formula_fn]
+fn _thrust_requires_f(x: i64) -> bool {
+    true
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_f(result: i64, x: i64) -> bool {
+    (result == 1) || (result == -1) && result == 0
+}
+
+fn f(x: i64) -> i64 {
+    #[thrust::requires]
+    _thrust_requires_f;
+    #[thrust::ensures]
+    _thrust_ensures_f;
+
+    let y = rand_except(x);
+    if y > x {
+        1
+    } else if y < x {
+        -1
+    } else {
+        0
+    }
+}
+
+fn main() {
+    assert!(rand_except(1) == 0);
+}

tests/ui/pass/annot_formula_fn.rs

@@ -0,0 +1,54 @@
+//@check-pass
+
+#[thrust::formula_fn]
+fn _thrust_requires_rand_except(x: i64) -> bool {
+    true
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_rand_except(result: i64, x: i64) -> bool {
+    result != x
+}
+
+fn rand_except(x: i64) -> i64 {
+    #[thrust::requires]
+    _thrust_requires_rand_except;
+    #[thrust::ensures]
+    _thrust_ensures_rand_except;
+
+    if x == 0 {
+        1
+    } else {
+        0
+    }
+}
+
+#[thrust::formula_fn]
+fn _thrust_requires_f(x: i64) -> bool {
+    true
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_f(result: i64, x: i64) -> bool {
+    (result == 1) || (result == -1)
+}
+
+fn f(x: i64) -> i64 {
+    #[thrust::requires]
+    _thrust_requires_f;
+    #[thrust::ensures]
+    _thrust_ensures_f;
+
+    let y = rand_except(x);
+    if y > x {
+        1
+    } else if y < x {
+        -1
+    } else {
+        0
+    }
+}
+
+fn main() {
+    assert!(rand_except(1) != 1);
+}