Translate and register formula_fn

Author: coord-e

src/analyze.rs

@@ -15,6 +15,7 @@ use rustc_index::IndexVec;
 use rustc_middle::mir::{self, BasicBlock, Local};
 use rustc_middle::ty::{self as mir_ty, TyCtxt};
 use rustc_span::def_id::{DefId, LocalDefId};
+use rustc_span::symbol::Symbol;
 
 use crate::analyze;
 use crate::annot::{AnnotFormula, AnnotParser, Resolver};
@@ -24,11 +25,38 @@ use crate::refine::{self, BasicBlockType, TypeBuilder};
 use crate::rty;
 
 mod annot;
+mod annot_fn;
 mod basic_block;
 mod crate_;
 mod did_cache;
 mod local_def;
 
+pub fn mir_borrowck_skip_formula_fn(
+    tcx: rustc_middle::ty::TyCtxt<'_>,
+    local_def_id: rustc_span::def_id::LocalDefId,
+) -> rustc_middle::query::queries::mir_borrowck::ProvidedValue {
+    // TODO: unify impl with local_def::Analyzer
+    let is_annotated_as_formula_fn = tcx
+        .get_attrs_by_path(local_def_id.to_def_id(), &analyze::annot::formula_fn_path())
+        .next()
+        .is_some();
+
+    if is_annotated_as_formula_fn {
+        tracing::debug!(?local_def_id, "skipping borrow check for formula fn");
+        let dummy_result = rustc_middle::mir::BorrowCheckResult {
+            concrete_opaque_types: Default::default(),
+            closure_requirements: None,
+            used_mut_upvars: Default::default(),
+            tainted_by_errors: None,
+        };
+        return tcx.arena.alloc(dummy_result);
+    }
+
+    (rustc_interface::DEFAULT_QUERY_PROVIDERS
+        .queries
+        .mir_borrowck)(tcx, local_def_id)
+}
+
 pub fn local_of_function_param(idx: rty::FunctionParamIdx) -> Local {
     Local::from(idx.index() + 1)
 }
@@ -155,6 +183,9 @@ pub struct Analyzer<'tcx> {
     /// (at least for every defs referenced by local def bodies)
     defs: HashMap<DefId, DefTy<'tcx>>,
 
+    /// Collection of functions with `#[thrust::formula_fn]` attribute.
+    formula_fns: HashMap<DefId, annot_fn::FormulaFn<'tcx>>,
+
     /// Resulting CHC system.
     system: Rc<RefCell<chc::System>>,
 
@@ -181,12 +212,14 @@ impl<'tcx> Analyzer<'tcx> {
 impl<'tcx> Analyzer<'tcx> {
     pub fn new(tcx: TyCtxt<'tcx>) -> Self {
         let defs = Default::default();
+        let formula_fns = Default::default();
         let system = Default::default();
         let basic_blocks = Default::default();
         let enum_defs = Default::default();
         Self {
             tcx,
             defs,
+            formula_fns,
             system,
             basic_blocks,
             def_ids: did_cache::DefIdCache::new(tcx),
@@ -346,6 +379,11 @@ impl<'tcx> Analyzer<'tcx> {
         Some(expected)
     }
 
+    pub fn register_formula_fn(&mut self, def_id: DefId, formula_fn: annot_fn::FormulaFn<'tcx>) {
+        tracing::info!(def_id = ?def_id, formula_fn = %formula_fn.display(), "register_formula_fn");
+        self.formula_fns.insert(def_id, formula_fn);
+    }
+
     pub fn register_basic_block_ty(
         &mut self,
         def_id: LocalDefId,

src/analyze/annot.rs

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

src/analyze/crate_.rs

@@ -6,7 +6,7 @@ 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::analyze::{self, annot_fn::AnnotFnTranslator};
 use crate::chc;
 use crate::rty::{self, ClauseBuilderExt as _};
 
@@ -88,6 +88,14 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             analyzer.analyze_predicate_definition(local_def_id);
         }
 
+        if analyzer.is_annotated_as_formula_fn() {
+            self.trusted.insert(local_def_id.to_def_id());
+            let formula_fn = AnnotFnTranslator::new(self.tcx, local_def_id).to_formula_fn();
+            self.ctx
+                .register_formula_fn(local_def_id.to_def_id(), formula_fn);
+            return;
+        }
+
         use mir_ty::TypeVisitableExt as _;
         if sig.has_param() && !analyzer.is_fully_annotated() {
             self.ctx.register_deferred_def(local_def_id.to_def_id());

src/analyze/local_def.rs

@@ -202,6 +202,16 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             .is_some()
     }
 
+    pub fn is_annotated_as_formula_fn(&self) -> bool {
+        self.tcx
+            .get_attrs_by_path(
+                self.local_def_id.to_def_id(),
+                &analyze::annot::formula_fn_path(),
+            )
+            .next()
+            .is_some()
+    }
+
     // TODO: unify this logic with extraction functions above
     pub fn is_fully_annotated(&self) -> bool {
         let has_require = self

src/lib.rs

@@ -1,9 +1,11 @@
 #![feature(rustc_private)]
 
 extern crate rustc_ast;
+extern crate rustc_borrowck;
 extern crate rustc_data_structures;
 extern crate rustc_hir;
 extern crate rustc_index;
+extern crate rustc_interface;
 extern crate rustc_middle;
 extern crate rustc_mir_dataflow;
 extern crate rustc_span;
@@ -23,4 +25,5 @@ mod rty;
 // utility
 mod pretty;
 
+pub use analyze::mir_borrowck_skip_formula_fn;
 pub use analyze::Analyzer;

src/main.rs

@@ -18,6 +18,10 @@ impl Callbacks for CompilerCalls {
         let attrs = &mut config.opts.unstable_opts.crate_attr;
         attrs.push("feature(register_tool)".to_owned());
         attrs.push("register_tool(thrust)".to_owned());
+
+        config.override_queries = Some(|_sess, providers| {
+            providers.mir_borrowck = thrust::mir_borrowck_skip_formula_fn;
+        });
     }
 
     fn after_crate_root_parsing<'tcx>(