Declare the sorts of quantifier binders in the emitted SMT-LIB2

src/analyze.rs

@@ -7,7 +7,7 @@
 //! CHC solver with the [`Analyzer::solve`] and subsequently reports the result.
 
 use std::cell::RefCell;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
 use rustc_hir::lang_items::LangItem;
@@ -378,6 +378,45 @@ impl<'tcx> Analyzer<'tcx> {
         enum_def
     }
 
+    /// Registers the definitions of the enums `ty` is built from, including the ones reachable
+    /// through the fields of the datatypes it mentions.
+    pub fn register_enum_defs_in_ty(&self, type_builder: &TypeBuilder<'tcx>, ty: mir_ty::Ty<'tcx>) {
+        use mir_ty::{TypeSuperVisitable as _, TypeVisitable as _};
+        struct EnumCollector<'a, 'tcx> {
+            tcx: TyCtxt<'tcx>,
+            builder: &'a TypeBuilder<'tcx>,
+            enums: HashSet<DefId>,
+            visited: HashSet<mir_ty::Ty<'tcx>>,
+        }
+        impl<'tcx> mir_ty::TypeVisitor<TyCtxt<'tcx>> for EnumCollector<'_, 'tcx> {
+            fn visit_ty(&mut self, ty: mir_ty::Ty<'tcx>) {
+                let ty = self.builder.resolve_model_ty(ty);
+                if let mir_ty::TyKind::Adt(def, args) = ty.kind() {
+                    if self.visited.insert(ty) {
+                        if def.is_enum() {
+                            self.enums.insert(def.did());
+                        }
+                        for field in def.all_fields() {
+                            field.ty(self.tcx, args).visit_with(self);
+                        }
+                    }
+                }
+                ty.super_visit_with(self);
+            }
+        }
+
+        let mut visitor = EnumCollector {
+            tcx: self.tcx,
+            builder: type_builder,
+            enums: HashSet::new(),
+            visited: HashSet::new(),
+        };
+        ty.visit_with(&mut visitor);
+        for def_id in visitor.enums {
+            self.get_or_register_enum_def(def_id);
+        }
+    }
+
     pub fn register_def(&mut self, def_id: DefId, rty: rty::RefinedType) {
         tracing::info!(def_id = ?def_id, rty = %rty.display(), "register_def");
         self.defs.insert(def_id, DefTy::Concrete(rty));

src/analyze/annot_fn.rs

@@ -617,6 +617,8 @@ impl<'a, 'tcx> AnnotFnTranslator<'a, 'tcx> {
                 );
             };
             let param_ty = self.pat_ty(param.pat);
+            self.analyzer
+                .register_enum_defs_in_ty(&self.type_builder, param_ty);
             let sort = self.type_builder.build(param_ty).to_sort();
             let var_term = chc::Term::FormulaQuantifiedVar(sort.clone(), ident.name.to_string());
             inner_translator.env.insert(hir_id, var_term);

src/analyze/basic_block.rs

@@ -1433,40 +1433,9 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     }
 
     fn register_enum_defs(&mut self) {
-        use mir_ty::{TypeSuperVisitable as _, TypeVisitable as _};
-        struct EnumCollector<'tcx> {
-            tcx: mir_ty::TyCtxt<'tcx>,
-            builder: TypeBuilder<'tcx>,
-            enums: std::collections::HashSet<DefId>,
-            visited: std::collections::HashSet<mir_ty::Ty<'tcx>>,
-        }
-        impl<'tcx> mir_ty::TypeVisitor<mir_ty::TyCtxt<'tcx>> for EnumCollector<'tcx> {
-            fn visit_ty(&mut self, ty: mir_ty::Ty<'tcx>) {
-                let ty = self.builder.resolve_model_ty(ty);
-                if let mir_ty::TyKind::Adt(def, args) = ty.kind() {
-                    if self.visited.insert(ty) {
-                        if def.is_enum() {
-                            self.enums.insert(def.did());
-                        }
-                        for field in def.all_fields() {
-                            field.ty(self.tcx, args).visit_with(self);
-                        }
-                    }
-                }
-                ty.super_visit_with(self);
-            }
-        }
-        let mut visitor = EnumCollector {
-            tcx: self.tcx,
-            builder: self.type_builder.clone(),
-            enums: std::collections::HashSet::new(),
-            visited: std::collections::HashSet::new(),
-        };
         for local_decl in &self.local_decls {
-            local_decl.ty.visit_with(&mut visitor);
-        }
-        for def_id in visitor.enums {
-            self.ctx.get_or_register_enum_def(def_id);
+            self.ctx
+                .register_enum_defs_in_ty(&self.type_builder, local_decl.ty);
         }
     }
 }

src/chc.rs

@@ -1365,6 +1365,11 @@ impl<V> Atom<V> {
         self.args.iter().flat_map(|t| t.fv()).chain(guard_fvs)
     }
 
+    /// Iterates over the variables bound by the quantifiers of the guard.
+    pub fn iter_quantified_vars(&self) -> impl Iterator<Item = &(String, Sort)> {
+        self.guard.iter().flat_map(|g| g.iter_quantified_vars())
+    }
+
     pub fn guarded(self, new_guard: Formula<V>) -> Atom<V> {
         let Atom {
             guard: self_guard,
@@ -1651,6 +1656,26 @@ impl<V> Formula<V> {
         }
     }
 
+    /// Iterates over the variables bound by the quantifiers occurring in this formula.
+    pub fn iter_quantified_vars(&self) -> impl Iterator<Item = &(String, Sort)> {
+        self.iter_quantified_vars_impl()
+    }
+
+    fn iter_quantified_vars_impl(&self) -> Box<dyn Iterator<Item = &(String, Sort)> + '_> {
+        match self {
+            Formula::Atom(atom) => Box::new(atom.iter_quantified_vars()),
+            Formula::Not(fo) => Box::new(fo.iter_quantified_vars()),
+            Formula::And(fs) => Box::new(fs.iter().flat_map(Formula::iter_quantified_vars)),
+            Formula::Or(fs) => Box::new(fs.iter().flat_map(Formula::iter_quantified_vars)),
+            Formula::Implies(lhs, rhs) => {
+                Box::new(lhs.iter_quantified_vars().chain(rhs.iter_quantified_vars()))
+            }
+            Formula::Exists(vars, fo) | Formula::Forall(vars, fo) => {
+                Box::new(vars.iter().chain(fo.iter_quantified_vars()))
+            }
+        }
+    }
+
     pub fn push_conj(&mut self, other: Self) {
         match self {
             Formula::And(fs) => {
@@ -1823,6 +1848,13 @@ impl<V> Body<V> {
     pub fn iter_atoms(&self) -> impl Iterator<Item = &Atom<V>> {
         self.formula.iter_atoms().chain(&self.atoms)
     }
+
+    /// Iterates over the variables bound by the quantifiers occurring in this body.
+    pub fn iter_quantified_vars(&self) -> impl Iterator<Item = &(String, Sort)> {
+        self.formula
+            .iter_quantified_vars()
+            .chain(self.atoms.iter().flat_map(Atom::iter_quantified_vars))
+    }
 }
 
 impl<V> Body<V>

src/chc/format_context.rs

@@ -226,6 +226,11 @@ fn collect_sorts(system: &chc::System) -> BTreeSet<chc::Sort> {
         for a in clause.body.iter_atoms() {
             atom_sorts(clause, a, &mut sorts);
         }
+        let quantified_vars = clause
+            .head
+            .iter_quantified_vars()
+            .chain(clause.body.iter_quantified_vars());
+        sorts.extend(quantified_vars.map(|(_, sort)| sort.clone()));
     }
 
     for sort in sorts.clone() {

tests/ui/fail/annot_exists_enum_binder.rs

@@ -0,0 +1,22 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -Adead_code -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+use thrust_models::exists;
+
+pub enum X {
+    A(i64),
+    B(bool),
+}
+
+impl thrust_models::Model for X {
+    type Ty = X;
+}
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(exists(|x: X| result < 0))]
+fn f() -> i64 {
+    1
+}
+
+fn main() {}

tests/ui/fail/annot_exists_tuple_binder.rs

@@ -0,0 +1,13 @@
+//@error-in-other-file: Unsat
+//@compile-flags: -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+use thrust_models::exists;
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(exists(|p: (i64, bool)| result < 0))]
+fn f() -> i64 {
+    1
+}
+
+fn main() {}

tests/ui/pass/annot_exists_enum_binder.rs

@@ -0,0 +1,22 @@
+//@check-pass
+//@compile-flags: -Adead_code -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+use thrust_models::exists;
+
+pub enum X {
+    A(i64),
+    B(bool),
+}
+
+impl thrust_models::Model for X {
+    type Ty = X;
+}
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(exists(|x: X| result >= 0))]
+fn f() -> i64 {
+    1
+}
+
+fn main() {}

tests/ui/pass/annot_exists_tuple_binder.rs

@@ -0,0 +1,13 @@
+//@check-pass
+//@compile-flags: -C debug-assertions=off
+//@rustc-env: THRUST_SOLVER=tests/thrust-pcsat-wrapper
+
+use thrust_models::exists;
+
+#[thrust_macros::requires(true)]
+#[thrust_macros::ensures(exists(|p: (i64, bool)| result >= 0))]
+fn f() -> i64 {
+    1
+}
+
+fn main() {}