Translate formula fns after type params are instantiated

Author: coord-e

src/analyze.rs

@@ -190,6 +190,11 @@ impl refine::EnumDefProvider for Rc<RefCell<EnumDefs>> {
 
 pub type Env = refine::Env<Rc<RefCell<EnumDefs>>>;
 
+#[derive(Debug, Clone)]
+struct DeferredFormulaFnDef<'tcx> {
+    cache: Rc<RefCell<HashMap<mir_ty::GenericArgsRef<'tcx>, annot_fn::FormulaFn<'tcx>>>>,
+}
+
 #[derive(Clone)]
 pub struct Analyzer<'tcx> {
     tcx: TyCtxt<'tcx>,
@@ -202,7 +207,7 @@ pub struct Analyzer<'tcx> {
     defs: HashMap<DefId, DefTy<'tcx>>,
 
     /// Collection of functions with `#[thrust::formula_fn]` attribute.
-    formula_fns: HashMap<DefId, annot_fn::FormulaFn<'tcx>>,
+    formula_fns: HashMap<LocalDefId, DeferredFormulaFnDef<'tcx>>,
 
     /// Resulting CHC system.
     system: Rc<RefCell<chc::System>>,
@@ -391,6 +396,29 @@ impl<'tcx> Analyzer<'tcx> {
         })
     }
 
+    pub fn formula_fn_with_args(
+        &self,
+        local_def_id: LocalDefId,
+        generic_args: mir_ty::GenericArgsRef<'tcx>,
+    ) -> Option<annot_fn::FormulaFn<'tcx>> {
+        let deferred_formula_fn = self.formula_fns.get(&local_def_id)?;
+
+        let deferred_formula_fn_cache = Rc::clone(&deferred_formula_fn.cache);
+        if let Some(formula_fn) = deferred_formula_fn_cache.borrow().get(&generic_args) {
+            return Some(formula_fn.clone());
+        }
+
+        let translator = annot_fn::AnnotFnTranslator::new(self.tcx, self.def_ids(), local_def_id)
+            .with_generic_args(generic_args);
+        let formula_fn = translator.to_formula_fn();
+        deferred_formula_fn_cache
+            .borrow_mut()
+            .insert(generic_args, formula_fn.clone());
+
+        tracing::info!(?local_def_id, formula_fn = %formula_fn.display(), ?generic_args, "formula_fn_with_args");
+        Some(formula_fn)
+    }
+
     pub fn def_ty_with_args(
         &mut self,
         def_id: DefId,
@@ -443,9 +471,14 @@ 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_formula_fn(&mut self, local_def_id: LocalDefId) {
+        tracing::info!(?local_def_id, "register_formula_fn");
+        self.formula_fns.insert(
+            local_def_id,
+            DeferredFormulaFnDef {
+                cache: Rc::new(RefCell::new(HashMap::new())),
+            },
+        );
     }
 
     pub fn register_basic_block_ty(
@@ -585,11 +618,13 @@ impl<'tcx> Analyzer<'tcx> {
         None
     }
 
+    // TODO: reduce number of args
     fn extract_require_annot<T>(
         &self,
         local_def_id: LocalDefId,
         resolver: T,
         self_type_name: Option<String>,
+        generic_args: mir_ty::GenericArgsRef<'tcx>,
     ) -> Option<AnnotFormula<T::Output>>
     where
         T: Resolver<Output = rty::FunctionParamIdx>,
@@ -611,10 +646,16 @@ impl<'tcx> Analyzer<'tcx> {
         if let Some(formula_def_id) =
             self.extract_path_with_attr(local_def_id, &analyze::annot::requires_path_path())
         {
+            let Some(formula_def_id) = formula_def_id.as_local() else {
+                panic!(
+                    "require annotation with path is expected to refer to a local def, but found: {:?}",
+                    formula_def_id
+                );
+            };
             if require_annot.is_some() {
                 unimplemented!();
             }
-            let Some(formula_fn) = self.formula_fns.get(&formula_def_id) else {
+            let Some(formula_fn) = self.formula_fn_with_args(formula_def_id, generic_args) else {
                 panic!(
                     "require annotation {:?} is not a formula function",
                     formula_def_id
@@ -626,11 +667,13 @@ impl<'tcx> Analyzer<'tcx> {
         require_annot
     }
 
+    // TODO: reduce number of args
     fn extract_ensure_annot<T>(
         &self,
         local_def_id: LocalDefId,
         resolver: T,
         self_type_name: Option<String>,
+        generic_args: mir_ty::GenericArgsRef<'tcx>,
     ) -> Option<AnnotFormula<T::Output>>
     where
         T: Resolver<Output = rty::RefinedTypeVar<rty::FunctionParamIdx>>,
@@ -653,10 +696,16 @@ impl<'tcx> Analyzer<'tcx> {
         if let Some(formula_def_id) =
             self.extract_path_with_attr(local_def_id, &analyze::annot::ensures_path_path())
         {
+            let Some(formula_def_id) = formula_def_id.as_local() else {
+                panic!(
+                    "require annotation with path is expected to refer to a local def, but found: {:?}",
+                    formula_def_id
+                );
+            };
             if ensure_annot.is_some() {
                 unimplemented!();
             }
-            let Some(formula_fn) = self.formula_fns.get(&formula_def_id) else {
+            let Some(formula_fn) = self.formula_fn_with_args(formula_def_id, generic_args) else {
                 panic!(
                     "ensure annotation {:?} is not a formula function",
                     formula_def_id

src/analyze/annot_fn.rs

@@ -124,6 +124,7 @@ pub struct AnnotFnTranslator<'tcx> {
 
     typeck: &'tcx mir_ty::TypeckResults<'tcx>,
     body: &'tcx rustc_hir::Body<'tcx>,
+    generic_args: mir_ty::GenericArgsRef<'tcx>,
 
     def_ids: DefIdCache<'tcx>,
     env: HashMap<HirId, chc::Term<rty::FunctionParamIdx>>,
@@ -134,20 +135,27 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         let map = tcx.hir();
         let body_id = map.body_owned_by(local_def_id);
         let body = map.body(body_id);
+        let generic_args = tcx.mk_args(&[]);
 
         let typeck = tcx.typeck(local_def_id);
         let mut translator = Self {
             tcx,
             local_def_id,
             typeck,
             body,
+            generic_args,
             def_ids,
             env: HashMap::default(),
         };
         translator.build_env_from_params();
         translator
     }
 
+    pub fn with_generic_args(mut self, generic_args: mir_ty::GenericArgsRef<'tcx>) -> Self {
+        self.generic_args = generic_args;
+        self
+    }
+
     fn build_env_from_params(&mut self) {
         for (idx, param) in self.body.params.iter().enumerate() {
             let param_idx = rty::FunctionParamIdx::from(idx);
@@ -177,12 +185,19 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         }
     }
 
+    fn expr_ty(&self, expr: &'tcx rustc_hir::Expr<'tcx>) -> mir_ty::Ty<'tcx> {
+        let ty = self.typeck.expr_ty(expr);
+        let instantiated = mir_ty::EarlyBinder::bind(ty).instantiate(self.tcx, self.generic_args);
+        let param_env = mir_ty::ParamEnv::reveal_all();
+        self.tcx.normalize_erasing_regions(param_env, instantiated)
+    }
+
     pub fn to_formula_fn(&self) -> FormulaFn<'tcx> {
         let formula = self.to_formula(self.body.value);
         let params = self
             .tcx
             .fn_sig(self.local_def_id.to_def_id())
-            .instantiate_identity()
+            .instantiate(self.tcx, self.generic_args)
             .skip_binder()
             .inputs()
             .to_vec();
@@ -260,7 +275,7 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                     FormulaOrTerm::Term(operand.neg())
                 }
                 rustc_hir::UnOp::Not => {
-                    let operand_ty = self.typeck.expr_ty(operand);
+                    let operand_ty = self.expr_ty(operand);
                     match operand_ty.ty_adt_def() {
                         Some(adt) if Some(adt.did()) == self.def_ids.mut_model() => {
                             let operand = self.to_term(operand);
@@ -273,7 +288,7 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                     }
                 }
                 rustc_hir::UnOp::Deref => {
-                    let operand_ty = self.typeck.expr_ty(operand);
+                    let operand_ty = self.expr_ty(operand);
                     let adt = operand_ty
                         .ty_adt_def()
                         .expect("deref operand must be a model type");

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::LocalDefId;
 
-use crate::analyze::{self, annot_fn::AnnotFnTranslator};
+use crate::analyze;
 use crate::chc;
 use crate::rty::{self, ClauseBuilderExt as _};
 
@@ -95,10 +95,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         }
 
         if analyzer.is_annotated_as_formula_fn() {
-            let formula_fn =
-                AnnotFnTranslator::new(self.tcx, self.ctx.def_ids(), local_def_id).to_formula_fn();
-            self.ctx
-                .register_formula_fn(local_def_id.to_def_id(), formula_fn);
+            self.ctx.register_formula_fn(local_def_id);
             self.skip_analysis.insert(local_def_id);
             return;
         }

src/analyze/local_def.rs

@@ -48,6 +48,8 @@ pub struct Analyzer<'tcx, 'ctx> {
     local_def_id: LocalDefId,
 
     body: Body<'tcx>,
+    /// to substitute HIR types during translation in [`crate::analyze::annot_fn`]
+    generic_args: mir_ty::GenericArgsRef<'tcx>,
     drop_points: HashMap<BasicBlock, analyze::basic_block::DropPoints>,
     type_builder: TypeBuilder<'tcx>,
 }
@@ -310,12 +312,14 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             self.local_def_id,
             &param_resolver,
             self_type_name.clone(),
+            self.generic_args,
         );
 
         let mut ensure_annot = self.ctx.extract_ensure_annot(
             self.local_def_id,
             &result_param_resolver,
             self_type_name.clone(),
+            self.generic_args,
         );
 
         if let Some(trait_item_id) = self.trait_item_id() {
@@ -324,11 +328,13 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 trait_item_id,
                 &param_resolver,
                 self_type_name.clone(),
+                self.generic_args,
             );
             let trait_ensure_annot = self.ctx.extract_ensure_annot(
                 trait_item_id,
                 &result_param_resolver,
                 self_type_name.clone(),
+                self.generic_args,
             );
 
             assert!(require_annot.is_none() || trait_require_annot.is_none());
@@ -851,17 +857,20 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         let body = tcx.optimized_mir(local_def_id.to_def_id()).clone();
         let drop_points = Default::default();
         let type_builder = TypeBuilder::new(tcx, ctx.def_ids(), local_def_id.to_def_id());
+        let generic_args = tcx.mk_args(&[]);
         Self {
             ctx,
             tcx,
             local_def_id,
             body,
+            generic_args,
             drop_points,
             type_builder,
         }
     }
 
     pub fn generic_args(&mut self, generic_args: mir_ty::GenericArgsRef<'tcx>) -> &mut Self {
+        self.generic_args = generic_args;
         self.body =
             mir_ty::EarlyBinder::bind(self.body.clone()).instantiate(self.tcx, generic_args);
         self

tests/ui/fail/annot_mut_term_formula_fn_poly.rs

@@ -0,0 +1,31 @@
+//@error-in-other-file: Unsat
+
+#[allow(unused_variables)]
+#[thrust::formula_fn]
+fn _thrust_requires_swap(x: thrust_models::model::Mut<i64>, y: i64) -> bool {
+    true
+}
+
+#[allow(unused_variables)]
+#[thrust::formula_fn]
+fn _thrust_ensures_swap(result: (), x: thrust_models::model::Mut<i64>, y: thrust_models::model::Mut<i64>) -> bool {
+    *x == *y && *y == *x
+}
+
+#[allow(path_statements)]
+fn swap<T>(x: &mut T, y: &mut T) {
+    #[thrust::requires_path]
+    _thrust_requires_swap;
+
+    #[thrust::ensures_path]
+    _thrust_ensures_swap;
+
+    std::mem::swap(x, y)
+}
+
+fn main() {
+    let mut a = 1;
+    let mut b = 2;
+    swap(&mut a, &mut b);
+    assert!(a == 2 && b == 1);
+}

tests/ui/pass/annot_mut_term_formula_fn_poly.rs

@@ -0,0 +1,31 @@
+//@check-pass
+
+#[allow(unused_variables)]
+#[thrust::formula_fn]
+fn _thrust_requires_swap(x: thrust_models::model::Mut<i64>, y: i64) -> bool {
+    true
+}
+
+#[allow(unused_variables)]
+#[thrust::formula_fn]
+fn _thrust_ensures_swap(result: (), x: thrust_models::model::Mut<i64>, y: thrust_models::model::Mut<i64>) -> bool {
+    !x == *y && !y == *x
+}
+
+#[allow(path_statements)]
+fn swap<T>(x: &mut T, y: &mut T) {
+    #[thrust::requires_path]
+    _thrust_requires_swap;
+
+    #[thrust::ensures_path]
+    _thrust_ensures_swap;
+
+    std::mem::swap(x, y)
+}
+
+fn main() {
+    let mut a = 1;
+    let mut b = 2;
+    swap(&mut a, &mut b);
+    assert!(a == 2 && b == 1);
+}