Support user-defined predicate call in annotation

src/analyze/annot.rs

@@ -37,6 +37,10 @@ pub fn extern_spec_fn_path() -> [Symbol; 2] {
     [Symbol::intern("thrust"), Symbol::intern("extern_spec_fn")]
 }
 
+pub fn raw_command_path() -> [Symbol; 2] {
+    [Symbol::intern("thrust"), Symbol::intern("raw_command")]
+}
+
 /// A [`annot::Resolver`] implementation for resolving function parameters.
 ///
 /// The parameter names and their sorts needs to be configured via

src/analyze/crate_.rs

@@ -2,6 +2,7 @@
 
 use std::collections::HashSet;
 
+use rustc_hir::def_id::CRATE_DEF_ID;
 use rustc_middle::ty::{self as mir_ty, TyCtxt};
 use rustc_span::def_id::{DefId, LocalDefId};
 
@@ -26,6 +27,38 @@ pub struct Analyzer<'tcx, 'ctx> {
 }
 
 impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
+    fn analyze_raw_command_annot(&mut self) {
+        for attrs in self.tcx.get_attrs_by_path(
+            CRATE_DEF_ID.to_def_id(),
+            &analyze::annot::raw_command_path(),
+        ) {
+            use rustc_ast::token::{LitKind, Token, TokenKind};
+            use rustc_ast::tokenstream::TokenTree;
+
+            let ts = analyze::annot::extract_annot_tokens(attrs.clone());
+            let tt = ts.trees().next().expect("string literal");
+
+            let raw_command = match tt {
+                TokenTree::Token(
+                    Token {
+                        kind: TokenKind::Literal(lit),
+                        ..
+                    },
+                    _,
+                ) if lit.kind == LitKind::Str => Some(lit.symbol.to_string()),
+                _ => None,
+            }
+            .expect("invalid raw_command annotation");
+
+            self.ctx
+                .system
+                .borrow_mut()
+                .push_raw_command(chc::RawCommand {
+                    command: raw_command,
+                });
+        }
+    }
+
     fn refine_local_defs(&mut self) {
         for local_def_id in self.tcx.mir_keys(()) {
             if self.tcx.def_kind(*local_def_id).is_fn_like() {
@@ -187,6 +220,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         let span = tracing::debug_span!("crate", krate = %self.tcx.crate_name(rustc_span::def_id::LOCAL_CRATE));
         let _guard = span.enter();
 
+        self.analyze_raw_command_annot();
         self.refine_local_defs();
         self.analyze_local_defs();
         self.assert_callable_entry();

src/annot.rs

@@ -420,7 +420,7 @@ where
         Ok(AnnotPath { segments })
     }
 
-    fn parse_datatype_ctor_args(&mut self) -> Result<Vec<chc::Term<T::Output>>> {
+    fn parse_arg_terms(&mut self) -> Result<Vec<chc::Term<T::Output>>> {
         if self.look_ahead_token(0).is_none() {
             return Ok(Vec::new());
         }
@@ -478,6 +478,30 @@ where
                         FormulaOrTerm::Term(var, sort.clone())
                     }
                     _ => {
+                        // If the single-segment identifier is followed by parethesized arguments,
+                        // parse them as user-defined predicate calls.
+                        let next_tt = self.look_ahead_token_tree(0);
+
+                        if let Some(TokenTree::Delimited(_, _, Delimiter::Parenthesis, args)) =
+                            next_tt
+                        {
+                            let args = args.clone();
+                            self.consume();
+
+                            let pred_symbol = chc::UserDefinedPred::new(ident.name.to_string());
+                            let pred = chc::Pred::UserDefined(pred_symbol);
+
+                            let mut parser = Parser {
+                                resolver: self.boxed_resolver(),
+                                cursor: args.trees(),
+                                formula_existentials: self.formula_existentials.clone(),
+                            };
+                            let args = parser.parse_arg_terms()?;
+
+                            let atom = chc::Atom::new(pred, args);
+                            let formula = chc::Formula::Atom(atom);
+                            return Ok(FormulaOrTerm::Formula(formula));
+                        }
                         let (v, sort) = self.resolve(*ident)?;
                         FormulaOrTerm::Term(chc::Term::var(v), sort)
                     }
@@ -497,7 +521,7 @@ where
                     cursor: s.trees(),
                     formula_existentials: self.formula_existentials.clone(),
                 };
-                let args = parser.parse_datatype_ctor_args()?;
+                let args = parser.parse_arg_terms()?;
                 parser.end_of_input()?;
                 let (term, sort) = path.to_datatype_ctor(args);
                 FormulaOrTerm::Term(term, sort)

src/chc.rs

@@ -902,12 +902,39 @@ impl MatcherPred {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct UserDefinedPred {
+    inner: String,
+}
+
+impl std::fmt::Display for UserDefinedPred {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        self.inner.fmt(f)
+    }
+}
+
+impl<'a, 'b, D> Pretty<'a, D, termcolor::ColorSpec> for &'b UserDefinedPred
+where
+    D: pretty::DocAllocator<'a, termcolor::ColorSpec>,
+{
+    fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, termcolor::ColorSpec> {
+        allocator.text(self.inner.clone())
+    }
+}
+
+impl UserDefinedPred {
+    pub fn new(inner: String) -> Self {
+        Self { inner }
+    }
+}
+
 /// A predicate.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Pred {
     Known(KnownPred),
     Var(PredVarId),
     Matcher(MatcherPred),
+    UserDefined(UserDefinedPred),
 }
 
 impl std::fmt::Display for Pred {
@@ -916,6 +943,7 @@ impl std::fmt::Display for Pred {
             Pred::Known(p) => p.fmt(f),
             Pred::Var(p) => p.fmt(f),
             Pred::Matcher(p) => p.fmt(f),
+            Pred::UserDefined(p) => p.fmt(f),
         }
     }
 }
@@ -930,6 +958,7 @@ where
             Pred::Known(p) => p.pretty(allocator),
             Pred::Var(p) => p.pretty(allocator),
             Pred::Matcher(p) => p.pretty(allocator),
+            Pred::UserDefined(p) => p.pretty(allocator),
         }
     }
 }
@@ -958,6 +987,7 @@ impl Pred {
             Pred::Known(p) => p.name().into(),
             Pred::Var(p) => p.to_string().into(),
             Pred::Matcher(p) => p.name().into(),
+            Pred::UserDefined(p) => p.to_string().into(),
         }
     }
 
@@ -966,6 +996,7 @@ impl Pred {
             Pred::Known(p) => p.is_negative(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 
@@ -974,6 +1005,7 @@ impl Pred {
             Pred::Known(p) => p.is_infix(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 
@@ -982,6 +1014,7 @@ impl Pred {
             Pred::Known(p) => p.is_top(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 
@@ -990,6 +1023,7 @@ impl Pred {
             Pred::Known(p) => p.is_bottom(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 }
@@ -1606,6 +1640,14 @@ impl Clause {
     }
 }
 
+/// A command specified using #![thrust::define_raw()]
+///
+/// Those will be directly inserted into the generated SMT-LIB2 file.
+#[derive(Debug, Clone)]
+pub struct RawCommand {
+    pub command: String,
+}
+
 /// A selector for a datatype constructor.
 ///
 /// A selector is a function that extracts a field from a datatype value.
@@ -1655,6 +1697,7 @@ pub struct PredVarDef {
 /// A CHC system.
 #[derive(Debug, Clone, Default)]
 pub struct System {
+    pub raw_commands: Vec<RawCommand>,
     pub datatypes: Vec<Datatype>,
     pub clauses: IndexVec<ClauseId, Clause>,
     pub pred_vars: IndexVec<PredVarId, PredVarDef>,
@@ -1665,6 +1708,10 @@ impl System {
         self.pred_vars.push(PredVarDef { sig, debug_info })
     }
 
+    pub fn push_raw_command(&mut self, raw_command: RawCommand) {
+        self.raw_commands.push(raw_command)
+    }
+
     pub fn push_clause(&mut self, clause: Clause) -> Option<ClauseId> {
         if clause.is_nop() {
             return None;

src/chc/smtlib2.rs

@@ -370,6 +370,24 @@ impl<'ctx, 'a> Clause<'ctx, 'a> {
     }
 }
 
+/// A wrapper around a [`chc::RawCommand`] that provides a [`std::fmt::Display`] implementation in SMT-LIB2 format.
+#[derive(Debug, Clone)]
+pub struct RawCommand<'a> {
+    inner: &'a chc::RawCommand,
+}
+
+impl<'a> std::fmt::Display for RawCommand<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.inner.command,)
+    }
+}
+
+impl<'a> RawCommand<'a> {
+    pub fn new(inner: &'a chc::RawCommand) -> Self {
+        Self { inner }
+    }
+}
+
 /// A wrapper around a [`chc::DatatypeSelector`] that provides a [`std::fmt::Display`] implementation in SMT-LIB2 format.
 #[derive(Debug, Clone)]
 pub struct DatatypeSelector<'ctx, 'a> {
@@ -555,6 +573,11 @@ impl<'a> std::fmt::Display for System<'a> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         writeln!(f, "(set-logic HORN)\n")?;
 
+        // insert command from #![thrust::raw_command()] here
+        for raw_command in &self.inner.raw_commands {
+            writeln!(f, "{}\n", RawCommand::new(raw_command))?;
+        }
+
         writeln!(f, "{}\n", Datatypes::new(&self.ctx, self.ctx.datatypes()))?;
         for datatype in self.ctx.datatypes() {
             writeln!(f, "{}", DatatypeDiscrFun::new(&self.ctx, datatype))?;

src/chc/unbox.rs

@@ -42,6 +42,7 @@ fn unbox_pred(pred: Pred) -> Pred {
         Pred::Known(pred) => Pred::Known(pred),
         Pred::Var(pred) => Pred::Var(pred),
         Pred::Matcher(pred) => unbox_matcher_pred(pred),
+        Pred::UserDefined(pred) => Pred::UserDefined(pred),
     }
 }
 
@@ -161,6 +162,7 @@ pub fn unbox(system: System) -> System {
         clauses,
         pred_vars,
         datatypes,
+        raw_commands,
     } = system;
     let datatypes = datatypes.into_iter().map(unbox_datatype).collect();
     let clauses = clauses.into_iter().map(unbox_clause).collect();
@@ -169,5 +171,6 @@ pub fn unbox(system: System) -> System {
         clauses,
         pred_vars,
         datatypes,
+        raw_commands,
     }
 }

tests/ui/pass/annot_preds_raw_define.rs

@@ -0,0 +1,21 @@
+//@check-pass
+//@compile-flags: -Adead_code -C debug-assertions=off
+
+#![feature(custom_inner_attributes)]
+#![thrust::raw_command("(define-fun is_double ((x Int) (doubled_x Int)) Bool
+    (=
+        (* x 2)
+        doubled_x
+    )
+)")]
+
+#[thrust::requires(true)]
+#[thrust::ensures(is_double(x, result))]
+fn double(x: i64) -> i64 {
+    x + x
+}
+
+fn main() {
+    let a = 3;
+    assert!(double(a) == 6);
+}

tests/ui/pass/annot_preds_raw_define_multi.rs

@@ -0,0 +1,36 @@
+//@check-pass
+//@compile-flags: -Adead_code -C debug-assertions=off
+
+#![feature(custom_inner_attributes)]
+#![thrust::raw_command("(define-fun is_double ((x Int) (doubled_x Int)) Bool
+    (=
+        (* x 2)
+        doubled_x
+    )
+)")]
+
+#![thrust::raw_command("(define-fun is_triple ((x Int) (tripled_x Int)) Bool
+    (=
+        (* x 3)
+        tripled_x
+    )
+)")]
+
+#[thrust::requires(true)]
+#[thrust::ensures(is_double(x, result))]
+fn double(x: i64) -> i64 {
+    x + x
+}
+
+#[thrust::require(is_double(x, doubled_x))]
+#[thrust::ensures(is_triple(x, result))]
+fn triple(x: i64, doubled_x: i64) -> i64 {
+    x + doubled_x
+}
+
+fn main() {
+    let a = 3;
+    let double_a = double(a);
+    assert!(double_a == 6);
+    assert!(triple(a, double_a) == 9);
+}

tests/ui/pass/annot_raw_command.rs

@@ -0,0 +1,23 @@
+//@check-pass
+//@compile-flags: -Adead_code -C debug-assertions=off
+
+// Insert commands written in SMT-LIB2 format into .smt2 file directly.
+// This feature is intended for debug or experiment purpose.
+#![feature(custom_inner_attributes)]
+#![thrust::raw_command("(define-fun is_double ((x Int) (doubled_x Int)) Bool
+    (=
+        (* x 2)
+        doubled_x
+    )
+)")]
+
+#[thrust::requires(true)]
+#[thrust::ensures(result == 2 * x)] 
+fn double(x: i64) -> i64 {
+    x + x
+}
+
+fn main() {
+    let a = 3;
+    assert!(double(a) == 6);
+}