Insert raw commands into .smt2 file with #![raw_command()] 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_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,37 @@ 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 => lit.symbol.to_string(),
+                _ => panic!("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 +219,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/chc.rs

@@ -1606,6 +1606,14 @@ impl Clause {
     }
 }
 
+/// A command specified using `thrust::raw_command` attribute
+///
+/// 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 +1663,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 +1674,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

@@ -161,6 +161,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 +170,6 @@ pub fn unbox(system: System) -> System {
         clauses,
         pred_vars,
         datatypes,
+        raw_commands,
     }
 }

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);
+}

tests/ui/pass/annot_raw_command_multi.rs

@@ -0,0 +1,31 @@
+//@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
+    )
+)")]
+
+// multiple raw commands can be inserted.
+#![thrust::raw_command("(define-fun is_triple ((x Int) (tripled_x Int)) Bool
+    (=
+        (* x 3)
+        tripled_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);
+}