Insert raw definition into .smt2 file with #![raw_definie()] attribute

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;
@@ -1076,6 +1076,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 +1234,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

@@ -1606,6 +1606,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 +1663,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 +1674,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))?;

src/chc/unbox.rs

@@ -161,6 +161,7 @@ pub fn unbox(system: System) -> System {
         clauses,
         pred_vars,
         datatypes,
+        raw_definitions,
     } = system;
     let datatypes = datatypes.into_iter().map(unbox_datatype).collect();
     let clauses = clauses.into_iter().map(unbox_clause).collect();
@@ -169,5 +170,6 @@ pub fn unbox(system: System) -> System {
         clauses,
         pred_vars,
         datatypes,
+        raw_definitions,
     }
 }

tests/ui/fail/annot_raw_define.rs

@@ -0,0 +1,20 @@
+//@error-in-other-file: UnexpectedToken
+//@compile-flags: -Adead_code -C debug-assertions=off
+
+// Insert definitions written in SMT-LIB2 format into .smt file directly.
+// This feature is intended for debug or experiment purpose.
+#![feature(custom_inner_attributes)]
+#![thrust::raw_define(true)] // argument must be single string literal
+
+#[thrust::requires(true)]
+#[thrust::ensures(result == 2 * x)]
+// #[thrust::ensures(is_double(x, result))]
+fn double(x: i64) -> i64 {
+    x + x
+}
+
+fn main() {
+    let a = 3;
+    assert!(double(a) == 6);
+    // assert!(is_double(a, double(a)));
+}

tests/ui/fail/annot_raw_define_without_params.rs

@@ -0,0 +1,20 @@
+//@error-in-other-file: invalid attribute
+//@compile-flags: -Adead_code -C debug-assertions=off
+
+// Insert definitions written in SMT-LIB2 format into .smt file directly.
+// This feature is intended for debug or experiment purpose.
+#![feature(custom_inner_attributes)]
+#![thrust::raw_define] // argument must be single string literal
+
+#[thrust::requires(true)]
+#[thrust::ensures(result == 2 * x)]
+// #[thrust::ensures(is_double(x, result))]
+fn double(x: i64) -> i64 {
+    x + x
+}
+
+fn main() {
+    let a = 3;
+    assert!(double(a) == 6);
+    // assert!(is_double(a, double(a)));
+}

tests/ui/pass/annot_raw_define.rs

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

tests/ui/pass/annot_raw_define_multi.rs

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