Add Model trait to map logical models with Rust types

src/analyze.rs

@@ -29,6 +29,9 @@ mod crate_;
 mod did_cache;
 mod local_def;
 
+// TODO: organize structure and remove cross dependency between refine
+pub use did_cache::DefIdCache;
+
 pub fn local_of_function_param(idx: rty::FunctionParamIdx) -> Local {
     Local::from(idx.index() + 1)
 }
@@ -209,6 +212,11 @@ impl<'tcx> Analyzer<'tcx> {
         }
     }
 
+    pub fn def_ids(&self) -> did_cache::DefIdCache<'tcx> {
+        // DefIdCache is backed by Rc
+        self.def_ids.clone()
+    }
+
     pub fn add_clause(&mut self, clause: chc::Clause) {
         self.system.borrow_mut().push_clause(clause);
     }
@@ -234,7 +242,7 @@ impl<'tcx> Analyzer<'tcx> {
                     .iter()
                     .map(|field| {
                         let field_ty = self.tcx.type_of(field.did).instantiate_identity();
-                        TypeBuilder::new(self.tcx, def_id).build(field_ty)
+                        TypeBuilder::new(self.tcx, self.def_ids(), def_id).build(field_ty)
                     })
                     .collect();
                 rty::EnumVariantDef {
@@ -355,7 +363,7 @@ impl<'tcx> Analyzer<'tcx> {
         def_id: DefId,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
     ) -> Option<rty::RefinedType> {
-        let type_builder = TypeBuilder::new(self.tcx, def_id);
+        let type_builder = TypeBuilder::new(self.tcx, self.def_ids(), def_id);
 
         let deferred_ty = match self.defs.get(&def_id)? {
             DefTy::Concrete(rty) => {

src/analyze/annot.rs

@@ -49,6 +49,54 @@ pub fn ignored_path() -> [Symbol; 2] {
     [Symbol::intern("thrust"), Symbol::intern("ignored")]
 }
 
+pub fn model_ty_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("model_ty"),
+    ]
+}
+
+pub fn int_model_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("int_model"),
+    ]
+}
+
+pub fn mut_model_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("mut_model"),
+    ]
+}
+
+pub fn box_model_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("box_model"),
+    ]
+}
+
+pub fn array_model_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("array_model"),
+    ]
+}
+
+pub fn closure_model_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("closure_model"),
+    ]
+}
+
 /// A [`annot::Resolver`] implementation for resolving function parameters.
 ///
 /// The parameter names and their sorts needs to be configured via

src/analyze/basic_block.rs

@@ -1116,7 +1116,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         let env = ctx.new_env();
         let local_decls = body.local_decls.clone();
         let prophecy_vars = Default::default();
-        let type_builder = TypeBuilder::new(tcx, local_def_id.to_def_id());
+        let type_builder = TypeBuilder::new(tcx, ctx.def_ids(), local_def_id.to_def_id());
         Self {
             ctx,
             tcx,

src/analyze/did_cache.rs

@@ -1,15 +1,24 @@
 //! Retrieves and caches well-known [`DefId`]s.
 
 use std::cell::OnceCell;
+use std::rc::Rc;
 
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::DefId;
+use rustc_span::symbol::Symbol;
 use rustc_target::abi::FieldIdx;
 
 #[derive(Debug, Clone, Default)]
 struct DefIds {
     unique: OnceCell<Option<DefId>>,
     nonnull: OnceCell<Option<DefId>>,
+
+    model_ty: OnceCell<Option<DefId>>,
+    int_model: OnceCell<Option<DefId>>,
+    mut_model: OnceCell<Option<DefId>>,
+    box_model: OnceCell<Option<DefId>>,
+    array_model: OnceCell<Option<DefId>>,
+    closure_model: OnceCell<Option<DefId>>,
 }
 
 /// Retrieves and caches well-known [`DefId`]s.
@@ -19,14 +28,14 @@ struct DefIds {
 #[derive(Clone)]
 pub struct DefIdCache<'tcx> {
     tcx: TyCtxt<'tcx>,
-    def_ids: DefIds,
+    def_ids: Rc<DefIds>,
 }
 
 impl<'tcx> DefIdCache<'tcx> {
     pub fn new(tcx: TyCtxt<'tcx>) -> Self {
         Self {
             tcx,
-            def_ids: DefIds::default(),
+            def_ids: Rc::new(DefIds::default()),
         }
     }
 
@@ -63,4 +72,67 @@ impl<'tcx> DefIdCache<'tcx> {
             Some(nonnull_def.did())
         })
     }
+
+    fn annotated_def(&self, path: &[Symbol]) -> Option<DefId> {
+        let map = self.tcx.hir();
+        for item_id in map.items() {
+            let def_id = item_id.owner_id.to_def_id();
+            if self.tcx.get_attrs_by_path(def_id, path).next().is_some() {
+                return Some(def_id);
+            }
+
+            let item = map.item(item_id);
+            if let rustc_hir::ItemKind::Trait(_, _, _, _, trait_item_refs) = item.kind {
+                for trait_item_ref in trait_item_refs {
+                    let def_id = trait_item_ref.id.owner_id.to_def_id();
+                    if self.tcx.get_attrs_by_path(def_id, path).next().is_some() {
+                        return Some(def_id);
+                    }
+                }
+            }
+        }
+        None
+    }
+
+    pub fn model_ty(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .model_ty
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::model_ty_path()))
+    }
+
+    pub fn int_model(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .int_model
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::int_model_path()))
+    }
+
+    pub fn mut_model(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .mut_model
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::mut_model_path()))
+    }
+
+    pub fn box_model(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .box_model
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::box_model_path()))
+    }
+
+    pub fn array_model(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .array_model
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::array_model_path()))
+    }
+
+    pub fn closure_model(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .closure_model
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::closure_model_path()))
+    }
 }

src/analyze/local_def.rs

@@ -842,7 +842,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         let tcx = ctx.tcx;
         let body = tcx.optimized_mir(local_def_id.to_def_id()).clone();
         let drop_points = Default::default();
-        let type_builder = TypeBuilder::new(tcx, local_def_id.to_def_id());
+        let type_builder = TypeBuilder::new(tcx, ctx.def_ids(), local_def_id.to_def_id());
         Self {
             ctx,
             tcx,