Support existential quantification in formula_fn

Author: coord-e

src/analyze/annot.rs

@@ -125,6 +125,14 @@ pub fn box_model_new_path() -> [Symbol; 3] {
     ]
 }
 
+pub fn exists_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("exists"),
+    ]
+}
+
 /// A [`annot::Resolver`] implementation for resolving function parameters.
 ///
 /// The parameter names and their sorts needs to be configured via

src/analyze/annot_fn.rs

@@ -119,6 +119,7 @@ impl<T> FormulaOrTerm<T> {
     }
 }
 
+#[derive(Clone)]
 pub struct AnnotFnTranslator<'tcx> {
     tcx: TyCtxt<'tcx>,
     local_def_id: LocalDefId,
@@ -201,6 +202,13 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         self.tcx.normalize_erasing_regions(param_env, instantiated)
     }
 
+    fn pat_ty(&self, pat: &'tcx rustc_hir::Pat<'tcx>) -> mir_ty::Ty<'tcx> {
+        let ty = self.typeck.pat_ty(pat);
+        let instantiated = mir_ty::EarlyBinder::bind(ty).instantiate(self.tcx, self.generic_args);
+        let param_env = mir_ty::ParamEnv::reveal_all();
+        self.tcx.normalize_erasing_regions(param_env, instantiated)
+    }
+
     pub fn to_formula_fn(&self) -> FormulaFn<'tcx> {
         let formula = self.to_formula(self.body.value);
         let params = self
@@ -377,6 +385,39 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                 if let ExprKind::Path(qpath) = &func_expr.kind {
                     let res = self.typeck.qpath_res(qpath, func_expr.hir_id);
                     if let rustc_hir::def::Res::Def(def_kind, def_id) = res {
+                        if Some(def_id) == self.def_ids.exists() {
+                            assert_eq!(args.len(), 1, "exists takes exactly 1 argument");
+                            let ExprKind::Closure(closure) = args[0].kind else {
+                                panic!("exists argument must be a closure");
+                            };
+                            let closure_body = self.tcx.hir().body(closure.body);
+
+                            let mut inner_translator = self.clone();
+                            let mut vars = Vec::new();
+                            for param in closure_body.params {
+                                let rustc_hir::PatKind::Binding(_, hir_id, ident, None) =
+                                    param.pat.kind
+                                else {
+                                    panic!(
+                                        "exists closure parameter must be a simple binding: {:?}",
+                                        param.pat
+                                    );
+                                };
+                                let param_ty = self.pat_ty(param.pat);
+                                let sort = self.type_builder.build(param_ty).to_sort();
+                                let var_term = chc::Term::FormulaExistentialVar(
+                                    sort.clone(),
+                                    ident.name.to_string(),
+                                );
+                                inner_translator.env.insert(hir_id, var_term);
+                                vars.push((ident.name.to_string(), sort));
+                            }
+                            let body_formula = inner_translator.to_formula(closure_body.value);
+                            return FormulaOrTerm::Formula(chc::Formula::exists(
+                                vars,
+                                body_formula,
+                            ));
+                        }
                         if Some(def_id) == self.def_ids.mut_model_new() {
                             assert_eq!(args.len(), 2, "Mut::new takes exactly 2 arguments");
                             let t1 = self.to_term(&args[0]);

src/analyze/did_cache.rs

@@ -22,6 +22,8 @@ struct DefIds {
 
     mut_model_new: OnceCell<Option<DefId>>,
     box_model_new: OnceCell<Option<DefId>>,
+
+    exists: OnceCell<Option<DefId>>,
 }
 
 /// Retrieves and caches well-known [`DefId`]s.
@@ -160,4 +162,11 @@ impl<'tcx> DefIdCache<'tcx> {
             .box_model_new
             .get_or_init(|| self.annotated_def(&crate::analyze::annot::box_model_new_path()))
     }
+
+    pub fn exists(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .exists
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::exists_path()))
+    }
 }

std.rs

@@ -234,6 +234,12 @@ mod thrust_models {
     impl<T, E> Model for Result<T, E> where T: Model, E: Model {
         type Ty = Result<<T as Model>::Ty, <E as Model>::Ty>;
     }
+
+    #[allow(dead_code)]
+    #[thrust::def::exists]
+    pub fn exists<T>(_x: T) -> bool {
+        unimplemented!()
+    }
 }
 
 #[thrust::extern_spec_fn]

tests/ui/fail/annot_exists_formula_fn.rs

@@ -0,0 +1,31 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+#[thrust::trusted]
+#[thrust::callable]
+fn rand() -> i32 { unimplemented!() }
+
+#[thrust::formula_fn]
+fn _thrust_requires_f() -> bool {
+    true
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_f(result: i32) -> bool {
+    thrust_models::exists(|x: i32| result == 2 * x)
+}
+
+#[allow(path_statements)]
+fn f() -> i32 {
+    #[thrust::requires_path]
+    _thrust_requires_f;
+
+    #[thrust::ensures_path]
+    _thrust_ensures_f;
+
+    let x = rand();
+    x + x + x
+}
+
+fn main() {}

tests/ui/pass/annot_exists_formula_fn.rs

@@ -0,0 +1,31 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+#[thrust::trusted]
+#[thrust::callable]
+fn rand() -> i32 { unimplemented!() }
+
+#[thrust::formula_fn]
+fn _thrust_requires_f() -> bool {
+    true
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_f(result: i32) -> bool {
+    thrust_models::exists(|x: i32| result == 2 * x)
+}
+
+#[allow(path_statements)]
+fn f() -> i32 {
+    #[thrust::requires_path]
+    _thrust_requires_f;
+
+    #[thrust::ensures_path]
+    _thrust_ensures_f;
+
+    let x = rand();
+    x + x
+}
+
+fn main() {}