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_define_path() -> [Symbol; 2] {
+    [Symbol::intern("thrust"), Symbol::intern("raw_define")]
+}
+
 /// A [`annot::Resolver`] implementation for resolving function parameters.
 ///
 /// The parameter names and their sorts needs to be configured via

src/analyze/crate_.rs

@@ -2,12 +2,14 @@
 
 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};
 
 use crate::analyze;
 use crate::chc;
 use crate::rty::{self, ClauseBuilderExt as _};
+use crate::annot;
 
 /// An implementation of local crate analysis.
 ///
@@ -26,6 +28,21 @@ pub struct Analyzer<'tcx, 'ctx> {
 }
 
 impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
+    fn analyze_raw_define_annot(&mut self) {
+        for attrs in self.tcx.get_attrs_by_path(
+            CRATE_DEF_ID.to_def_id(),
+            &analyze::annot::raw_define_path(),
+        ) {
+            let ts = analyze::annot::extract_annot_tokens(attrs.clone());
+            let parser = annot::AnnotParser::new(
+                // TODO: this resolver is not actually used.
+                analyze::annot::ParamResolver::default()
+            );
+            let raw_definition = parser.parse_raw_definition(ts).unwrap();
+            self.ctx.system.borrow_mut().push_raw_definition(raw_definition);
+        }
+    }
+
     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 +204,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_define_annot();
         self.refine_local_defs();
         self.analyze_local_defs();
         self.assert_callable_entry();

src/annot.rs

@@ -8,7 +8,7 @@
 //! The main entry point is [`AnnotParser`], which parses a [`TokenStream`] into a
 //! [`rty::RefinedType`] or a [`chc::Formula`].
 
-use rustc_ast::token::{BinOpToken, Delimiter, LitKind, Token, TokenKind};
+use rustc_ast::token::{BinOpToken, Delimiter, LitKind, Lit, Token, TokenKind};
 use rustc_ast::tokenstream::{RefTokenTreeCursor, Spacing, TokenStream, TokenTree};
 use rustc_index::IndexVec;
 use rustc_span::symbol::Ident;
@@ -477,7 +477,27 @@ where
                             chc::Term::FormulaExistentialVar(sort.clone(), ident.name.to_string());
                         FormulaOrTerm::Term(var, sort.clone())
                     }
-                    _ => {
+                    _ => {                        
+                        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::UserDefinedPredSymbol::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_datatype_ctor_args()?;
+
+                            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)
                     }
@@ -1076,6 +1096,32 @@ where
             .ok_or_else(|| ParseAttrError::unexpected_term("in annotation"))?;
         Ok(AnnotFormula::Formula(formula))
     }
+
+    pub fn parse_annot_raw_definition(&mut self) -> Result<chc::RawDefinition> {
+        let t = self.next_token("raw CHC definition")?;
+
+        match t {
+            Token {
+                kind: TokenKind::Literal(
+                    Lit { kind, symbol, .. }
+                ),
+                ..
+            } => {
+                match kind {
+                    LitKind::Str => {
+                        let definition = symbol.to_string();
+                        Ok(chc::RawDefinition{ definition })
+                    },
+                    _ =>  Err(ParseAttrError::unexpected_token(
+                        "string literal", t.clone()
+                    ))
+                }
+            },
+            _ => Err(ParseAttrError::unexpected_token(
+                "string literal", t.clone()
+            ))
+        }
+    }
 }
 
 /// A [`Resolver`] implementation for resolving specific variable as [`rty::RefinedTypeVar::Value`].
@@ -1208,4 +1254,15 @@ where
         parser.end_of_input()?;
         Ok(formula)
     }
+
+    pub fn parse_raw_definition(&self, ts: TokenStream) -> Result<chc::RawDefinition> {
+        let mut parser = Parser {
+            resolver: &self.resolver,
+            cursor: ts.trees(),
+            formula_existentials: Default::default(),
+        };
+        let raw_definition = parser.parse_annot_raw_definition()?;
+        parser.end_of_input()?;
+        Ok(raw_definition)
+    }
 }

src/chc.rs

@@ -902,12 +902,87 @@ impl MatcherPred {
     }
 }
 
+// TODO: This struct is almost copy of `DatatypeSymbol`. Two traits maight be unified with aliases.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct UserDefinedPredSymbol {
+    inner: String,
+}
+
+impl std::fmt::Display for UserDefinedPredSymbol {
+    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 UserDefinedPredSymbol
+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 UserDefinedPredSymbol {
+    pub fn new(inner: String) -> Self {
+        Self { inner }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct UserDefinedPred {
+    symbol: UserDefinedPredSymbol,
+    args: Vec<Sort>,
+}
+
+impl<'a> std::fmt::Display for UserDefinedPred {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{}", self.symbol.inner)
+    }
+}
+
+impl<'a, 'b, D> Pretty<'a, D, termcolor::ColorSpec> for &'b UserDefinedPred
+where
+    D: pretty::DocAllocator<'a, termcolor::ColorSpec>,
+    D::Doc: Clone,
+{
+    fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, termcolor::ColorSpec> {
+        let args = allocator.intersperse(
+            self.args.iter().map(|a| a.pretty(allocator)),
+            allocator.text(", "),
+        );
+        allocator
+            .text("user_defined_pred")
+            .append(
+                allocator
+                    .text(self.symbol.inner.clone())
+                    .append(args.angles())
+                    .angles(),
+            )
+            .group()
+    }
+}
+
+impl UserDefinedPred {
+    pub fn new(symbol: UserDefinedPredSymbol, args: Vec<Sort>) -> Self {
+        Self {
+            symbol,
+            args,
+        }
+    }
+
+    pub fn name(&self) -> &str {
+        &self.symbol.inner
+    }
+}
+
 /// A predicate.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Pred {
     Known(KnownPred),
     Var(PredVarId),
     Matcher(MatcherPred),
+    UserDefined(UserDefinedPredSymbol),
 }
 
 impl std::fmt::Display for Pred {
@@ -916,6 +991,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 +1006,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 +1035,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.inner.clone().into(),
         }
     }
 
@@ -966,6 +1044,7 @@ impl Pred {
             Pred::Known(p) => p.is_negative(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 
@@ -974,6 +1053,7 @@ impl Pred {
             Pred::Known(p) => p.is_infix(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 
@@ -982,6 +1062,7 @@ impl Pred {
             Pred::Known(p) => p.is_top(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 
@@ -990,6 +1071,7 @@ impl Pred {
             Pred::Known(p) => p.is_bottom(),
             Pred::Var(_) => false,
             Pred::Matcher(_) => false,
+            Pred::UserDefined(_) => false,
         }
     }
 }
@@ -1606,6 +1688,14 @@ impl Clause {
     }
 }
 
+/// A definition specified using #![thrust::define_raw()]
+///
+/// Those will be directly inserted into the generated SMT-LIB2 file.
+#[derive(Debug, Clone)]
+pub struct RawDefinition {
+    pub definition: String,
+}
+
 /// A selector for a datatype constructor.
 ///
 /// A selector is a function that extracts a field from a datatype value.
@@ -1655,6 +1745,7 @@ pub struct PredVarDef {
 /// A CHC system.
 #[derive(Debug, Clone, Default)]
 pub struct System {
+    pub raw_definitions: Vec<RawDefinition>,
     pub datatypes: Vec<Datatype>,
     pub clauses: IndexVec<ClauseId, Clause>,
     pub pred_vars: IndexVec<PredVarId, PredVarDef>,
@@ -1665,6 +1756,10 @@ impl System {
         self.pred_vars.push(PredVarDef { sig, debug_info })
     }
 
+    pub fn push_raw_definition(&mut self, raw_definition: RawDefinition) {
+        self.raw_definitions.push(raw_definition)
+    }
+
     pub fn push_clause(&mut self, clause: Clause) -> Option<ClauseId> {
         if clause.is_nop() {
             return None;

src/chc/format_context.rs

@@ -21,6 +21,7 @@ use crate::chc::{self, hoice::HoiceDatatypeRenamer};
 pub struct FormatContext {
     renamer: HoiceDatatypeRenamer,
     datatypes: Vec<chc::Datatype>,
+    raw_definitions: Vec<chc::RawDefinition>,
 }
 
 // FIXME: this is obviously ineffective and should be replaced
@@ -273,13 +274,18 @@ impl FormatContext {
             .filter(|d| d.params == 0)
             .collect();
         let renamer = HoiceDatatypeRenamer::new(&datatypes);
-        FormatContext { renamer, datatypes }
+        let raw_definitions = system.raw_definitions.clone();
+        FormatContext { renamer, datatypes, raw_definitions }
     }
 
     pub fn datatypes(&self) -> &[chc::Datatype] {
         &self.datatypes
     }
 
+    pub fn raw_definitions(&self) -> &[chc::RawDefinition] {
+        &self.raw_definitions
+    }
+
     pub fn box_ctor(&self, sort: &chc::Sort) -> impl std::fmt::Display {
         let ss = Sort::new(sort).sorts();
         format!("box{ss}")

src/chc/smtlib2.rs

@@ -370,6 +370,30 @@ impl<'ctx, 'a> Clause<'ctx, 'a> {
     }
 }
 
+/// A wrapper around a [`chc::RawDefinition`] that provides a [`std::fmt::Display`] implementation in SMT-LIB2 format.
+#[derive(Debug, Clone)]
+pub struct RawDefinition<'a> {
+    inner: &'a chc::RawDefinition,
+}
+
+impl<'a> std::fmt::Display for RawDefinition<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}",
+            self.inner.definition,
+        )
+    }
+}
+
+impl<'a> RawDefinition<'a> {
+    pub fn new(inner: &'a chc::RawDefinition) -> 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 +579,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 definition from #![thrust::define_chc()] here
+        for raw_def in self.ctx.raw_definitions() {
+            writeln!(f, "{}\n", RawDefinition::new(raw_def))?;
+        }
+
         writeln!(f, "{}\n", Datatypes::new(&self.ctx, self.ctx.datatypes()))?;
         for datatype in self.ctx.datatypes() {
             writeln!(f, "{}", DatatypeDiscrFun::new(&self.ctx, datatype))?;