Support Rust syntax in thrust::predicate bodies

src/analyze/crate_.rs

@@ -83,13 +83,20 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 self.skip_analysis.insert(*local_def_id);
                 keys.swap_remove(local_def_id);
             }
-            if analyzer.is_annotated_as_predicate() {
-                analyzer.analyze_predicate_definition();
-                self.skip_analysis.insert(*local_def_id);
-                keys.swap_remove(local_def_id);
-            }
-            if analyzer.is_annotated_as_formula_fn() {
-                self.ctx.register_formula_fn(*local_def_id);
+            let is_predicate = analyzer.is_annotated_as_predicate();
+            let is_formula_fn = analyzer.is_annotated_as_formula_fn();
+            if is_predicate || is_formula_fn {
+                // A predicate whose body is a Rust expression is also marked
+                // `formula_fn`; register it first so `analyze_predicate_definition`
+                // can pull the translated formula via `formula_fn_with_args`.
+                if is_formula_fn {
+                    self.ctx.register_formula_fn(*local_def_id);
+                }
+                if is_predicate {
+                    self.ctx
+                        .local_def_analyzer(*local_def_id)
+                        .analyze_predicate_definition();
+                }
                 self.skip_analysis.insert(*local_def_id);
                 keys.swap_remove(local_def_id);
             }

src/analyze/local_def.rs

@@ -60,6 +60,40 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     }
 
     fn define_as_predicate(&self, pred: chc::UserDefinedPred) {
+        let sig = self.ctx.fn_sig(self.local_def_id.to_def_id());
+        let arg_sorts = sig
+            .inputs()
+            .iter()
+            .map(|input_ty| self.type_builder.build(*input_ty).to_sort());
+
+        // A predicate marked `formula_fn` carries a Rust-expression body that has
+        // been translated into a `chc::Formula`; otherwise the body is a raw
+        // SMT-LIB2 string literal.
+        if self.is_annotated_as_formula_fn() {
+            // Name the parameters `v{i}` to match how `chc::TermVarIdx` renders the
+            // formula's variables (see `chc::UserDefinedPredBody::Formula`).
+            let arg_name_and_sorts = arg_sorts
+                .enumerate()
+                .map(|(i, sort)| (format!("v{i}"), sort))
+                .collect::<Vec<_>>();
+
+            let formula_fn = self
+                .ctx
+                .formula_fn_with_args(self.local_def_id, self.tcx.mk_args(&[]))
+                .expect("predicate formula function is not registered");
+            let formula = formula_fn
+                .formula()
+                .clone()
+                .map_var(|idx| chc::TermVarIdx::from(idx.index()));
+
+            self.ctx.system.borrow_mut().push_pred_define_formula(
+                pred,
+                chc::UserDefinedPredSig::from(arg_name_and_sorts),
+                formula,
+            );
+            return;
+        }
+
         // function's body
         use rustc_hir::{Block, Expr, ExprKind};
 
@@ -88,12 +122,6 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                     .to_string()
             });
 
-        let sig = self.ctx.fn_sig(self.local_def_id.to_def_id());
-        let arg_sorts = sig
-            .inputs()
-            .iter()
-            .map(|input_ty| self.type_builder.build(*input_ty).to_sort());
-
         let arg_name_and_sorts = arg_names.into_iter().zip(arg_sorts).collect::<Vec<_>>();
 
         self.ctx.system.borrow_mut().push_pred_define(

src/chc.rs

@@ -1918,9 +1918,32 @@ impl Clause {
     pub fn is_nop(&self) -> bool {
         self.head.is_top() || self.body.is_bottom()
     }
+}
+
+/// Resolves the sort of term-level variables.
+///
+/// Rendering a [`Term`] to SMT-LIB2 needs the sort of each variable, because
+/// box/mut/tuple constructors and selectors are sort-indexed. Anything that can
+/// provide variable sorts (a [`Clause`] via its `vars`, or a bare list of sorts
+/// for a `define-fun` signature) can drive term rendering through this trait,
+/// so callers need not fabricate a [`Clause`].
+pub trait TermSortEnv {
+    fn var_sort(&self, var: TermVarIdx) -> Sort;
 
     fn term_sort(&self, term: &Term<TermVarIdx>) -> Sort {
-        term.sort(|v| self.vars[*v].clone())
+        term.sort(|v| self.var_sort(*v))
+    }
+}
+
+impl TermSortEnv for Clause {
+    fn var_sort(&self, var: TermVarIdx) -> Sort {
+        self.vars[var].clone()
+    }
+}
+
+impl TermSortEnv for IndexVec<TermVarIdx, Sort> {
+    fn var_sort(&self, var: TermVarIdx) -> Sort {
+        self[var].clone()
     }
 }
 
@@ -1980,11 +2003,22 @@ pub struct PredVarDef {
 
 pub type UserDefinedPredSig = Vec<(String, Sort)>;
 
+/// The body of a user-defined predicate.
+///
+/// A predicate can be defined either by a raw SMT-LIB2 string (inserted into the
+/// generated `define-fun` verbatim) or by a [`Formula`] translated from a Rust
+/// expression via the `formula_fn` infrastructure.
+#[derive(Debug, Clone)]
+pub enum UserDefinedPredBody {
+    Raw(String),
+    Formula(Formula<TermVarIdx>),
+}
+
 #[derive(Debug, Clone)]
 pub struct UserDefinedPredDef {
     symbol: UserDefinedPred,
     sig: UserDefinedPredSig,
-    body: String,
+    body: UserDefinedPredBody,
 }
 
 /// A CHC system.
@@ -2012,8 +2046,24 @@ impl System {
         sig: UserDefinedPredSig,
         body: String,
     ) {
-        self.user_defined_pred_defs
-            .push(UserDefinedPredDef { symbol, sig, body })
+        self.user_defined_pred_defs.push(UserDefinedPredDef {
+            symbol,
+            sig,
+            body: UserDefinedPredBody::Raw(body),
+        })
+    }
+
+    pub fn push_pred_define_formula(
+        &mut self,
+        symbol: UserDefinedPred,
+        sig: UserDefinedPredSig,
+        formula: Formula<TermVarIdx>,
+    ) {
+        self.user_defined_pred_defs.push(UserDefinedPredDef {
+            symbol,
+            sig,
+            body: UserDefinedPredBody::Formula(formula),
+        })
     }
 
     pub fn push_clause(&mut self, clause: Clause) -> Option<ClauseId> {

src/chc/format_context.rs

@@ -7,7 +7,7 @@
 
 use std::collections::BTreeSet;
 
-use crate::chc::{self, hoice::HoiceDatatypeRenamer};
+use crate::chc::{self, hoice::HoiceDatatypeRenamer, TermSortEnv as _};
 
 /// A context for formatting a CHC system.
 ///