Encode unknown generic type parameters with (declare-forall-sort)

src/analyze.rs

@@ -8,6 +8,7 @@
 
 use std::cell::RefCell;
 use std::collections::HashMap;
+use std::hash::Hash;
 use std::rc::Rc;
 
 use rustc_hir::lang_items::LangItem;
@@ -19,7 +20,7 @@ use rustc_span::Symbol;
 
 use crate::analyze;
 use crate::annot::{AnnotFormula, AnnotParser, Resolver};
-use crate::chc;
+use crate::chc::{self, ForallSortIdx};
 use crate::pretty::PrettyDisplayExt as _;
 use crate::refine::{self, BasicBlockType, TypeBuilder};
 use crate::rty;
@@ -190,6 +191,13 @@ impl refine::EnumDefProvider for Rc<RefCell<EnumDefs>> {
 }
 
 pub type Env = refine::Env<Rc<RefCell<EnumDefs>>>;
+pub type TypeParamMap = HashMap<TypeParam, ForallSortIdx>;
+
+#[derive(Eq, PartialEq, Hash)]
+pub enum TypeParam {
+    GenericType(DefId, u32),
+    AssocType(DefId),
+}
 
 #[derive(Debug, Clone)]
 struct DeferredFormulaFnDef<'tcx> {
@@ -217,6 +225,8 @@ pub struct Analyzer<'tcx> {
     def_ids: did_cache::DefIdCache<'tcx>,
 
     enum_defs: Rc<RefCell<EnumDefs>>,
+
+    type_params: Rc<RefCell<TypeParamMap>>,
 }
 
 impl<'tcx> crate::refine::TemplateRegistry for Analyzer<'tcx> {
@@ -240,6 +250,7 @@ impl<'tcx> Analyzer<'tcx> {
         let system = Default::default();
         let basic_blocks = Default::default();
         let enum_defs = Default::default();
+        let type_params = Default::default();
         Self {
             tcx,
             defs,
@@ -248,6 +259,7 @@ impl<'tcx> Analyzer<'tcx> {
             basic_blocks,
             def_ids: did_cache::DefIdCache::new(tcx),
             enum_defs,
+            type_params,
         }
     }
 
@@ -281,7 +293,7 @@ impl<'tcx> Analyzer<'tcx> {
                     .iter()
                     .map(|field| {
                         let field_ty = self.tcx.type_of(field.did).instantiate_identity();
-                        TypeBuilder::new(self.tcx, self.def_ids(), def_id).build(field_ty)
+                        self.type_builder(self.def_ids(), def_id).build(field_ty)
                     })
                     .collect();
                 rty::EnumVariantDef {
@@ -380,14 +392,13 @@ impl<'tcx> Analyzer<'tcx> {
             ?mode,
             "register_deferred_def"
         );
-        self.defs.insert(
-            target_def_id,
+        self.defs.entry(target_def_id).or_insert_with(|| {
             DefTy::Deferred(DeferredDefTy {
                 local_def_id,
                 cache: Rc::new(RefCell::new(HashMap::new())),
                 mode,
-            }),
-        );
+            })
+        });
     }
 
     pub fn concrete_def_ty(&self, def_id: DefId) -> Option<&rty::RefinedType> {
@@ -401,6 +412,7 @@ impl<'tcx> Analyzer<'tcx> {
         &self,
         local_def_id: LocalDefId,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        owner_fn_id: DefId,
     ) -> Option<annot_fn::FormulaFn<'tcx>> {
         let deferred_formula_fn = self.formula_fns.get(&local_def_id)?;
 
@@ -409,9 +421,15 @@ impl<'tcx> Analyzer<'tcx> {
             return Some(formula_fn.clone());
         }
 
-        let translator = annot_fn::AnnotFnTranslator::new(self.tcx, local_def_id)
-            .with_generic_args(generic_args)
-            .with_def_id_cache(self.def_ids());
+        let translator = annot_fn::AnnotFnTranslator::new(
+            self.tcx,
+            local_def_id,
+            self.type_params.clone(),
+            self.system.clone(),
+            owner_fn_id,
+        )
+        .with_generic_args(generic_args)
+        .with_def_id_cache(self.def_ids());
         let formula_fn = translator.to_formula_fn();
         deferred_formula_fn_cache
             .borrow_mut()
@@ -425,8 +443,9 @@ impl<'tcx> Analyzer<'tcx> {
         &mut self,
         def_id: DefId,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        caller_def_id: DefId,
     ) -> Option<rty::RefinedType> {
-        let type_builder = TypeBuilder::new(self.tcx, self.def_ids(), def_id);
+        let type_builder = self.type_builder(self.def_ids(), caller_def_id);
 
         let deferred_ty = match self.defs.get(&def_id)? {
             DefTy::Concrete(rty) => {
@@ -531,8 +550,19 @@ impl<'tcx> Analyzer<'tcx> {
         &mut self,
         local_def_id: LocalDefId,
         bb: BasicBlock,
+        owner_fn_id: DefId,
     ) -> basic_block::Analyzer<'tcx, '_> {
-        basic_block::Analyzer::new(self, local_def_id, bb)
+        basic_block::Analyzer::new(self, local_def_id, bb, owner_fn_id)
+    }
+
+    pub fn type_builder(&self, def_ids: DefIdCache<'tcx>, owner_fn_id: DefId) -> TypeBuilder<'tcx> {
+        TypeBuilder::new(
+            self.tcx,
+            def_ids,
+            owner_fn_id,
+            self.type_params.clone(),
+            self.system.clone(),
+        )
     }
 
     pub fn solve(&mut self) {
@@ -626,6 +656,7 @@ impl<'tcx> Analyzer<'tcx> {
         resolver: T,
         self_type_name: Option<String>,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        owner_fn_id: DefId,
     ) -> Option<AnnotFormula<T::Output>>
     where
         T: Resolver<Output = rty::FunctionParamIdx>,
@@ -656,7 +687,9 @@ impl<'tcx> Analyzer<'tcx> {
             if require_annot.is_some() {
                 unimplemented!();
             }
-            let Some(formula_fn) = self.formula_fn_with_args(formula_def_id, generic_args) else {
+            let Some(formula_fn) =
+                self.formula_fn_with_args(formula_def_id, generic_args, owner_fn_id)
+            else {
                 panic!(
                     "require annotation {:?} is not a formula function",
                     formula_def_id
@@ -675,6 +708,7 @@ impl<'tcx> Analyzer<'tcx> {
         resolver: T,
         self_type_name: Option<String>,
         generic_args: mir_ty::GenericArgsRef<'tcx>,
+        owner_fn_id: DefId,
     ) -> Option<AnnotFormula<T::Output>>
     where
         T: Resolver<Output = rty::RefinedTypeVar<rty::FunctionParamIdx>>,
@@ -706,7 +740,9 @@ impl<'tcx> Analyzer<'tcx> {
             if ensure_annot.is_some() {
                 unimplemented!();
             }
-            let Some(formula_fn) = self.formula_fn_with_args(formula_def_id, generic_args) else {
+            let Some(formula_fn) =
+                self.formula_fn_with_args(formula_def_id, generic_args, owner_fn_id)
+            else {
                 panic!(
                     "ensure annotation {:?} is not a formula function",
                     formula_def_id

src/analyze/annot_fn.rs

@@ -1,13 +1,18 @@
+use std::cell::RefCell;
 use std::collections::HashMap;
+use std::rc::Rc;
 
 use pretty::{termcolor, Pretty};
-use rustc_hir::{def_id::LocalDefId, HirId};
+use rustc_hir::{
+    def_id::{DefId, LocalDefId},
+    HirId,
+};
 use rustc_index::IndexVec;
-use rustc_middle::ty::{self as mir_ty, TyCtxt};
+use rustc_middle::ty::{self as mir_ty, TyCtxt, TypeFoldable};
 
-use crate::analyze::{self, did_cache::DefIdCache};
+use crate::analyze::{self, did_cache::DefIdCache, TypeParamMap};
 use crate::annot::AnnotFormula;
-use crate::chc;
+use crate::chc::{self, System};
 use crate::refine::{self, TypeBuilder};
 use crate::rty;
 
@@ -137,15 +142,30 @@ pub struct AnnotFnTranslator<'tcx> {
     def_ids: DefIdCache<'tcx>,
     type_builder: TypeBuilder<'tcx>,
     env: HashMap<HirId, chc::Term<rty::FunctionParamIdx>>,
+
+    type_params: Rc<RefCell<TypeParamMap>>,
+    system: Rc<RefCell<System>>,
 }
 
 impl<'tcx> AnnotFnTranslator<'tcx> {
-    pub fn new(tcx: TyCtxt<'tcx>, local_def_id: LocalDefId) -> Self {
+    pub fn new(
+        tcx: TyCtxt<'tcx>,
+        local_def_id: LocalDefId,
+        type_params: Rc<RefCell<TypeParamMap>>,
+        system: Rc<RefCell<System>>,
+        owner_fn_id: DefId,
+    ) -> Self {
         let body = tcx.hir_body_owned_by(local_def_id);
         let generic_args = tcx.mk_args(&[]);
         let typeck = tcx.typeck(local_def_id);
         let def_ids = DefIdCache::new(tcx);
-        let type_builder = TypeBuilder::new(tcx, def_ids.clone(), local_def_id.to_def_id());
+        let type_builder = TypeBuilder::new(
+            tcx,
+            def_ids.clone(),
+            owner_fn_id,
+            type_params.clone(),
+            system.clone(),
+        );
         let mut translator = Self {
             tcx,
             local_def_id,
@@ -155,6 +175,8 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
             def_ids,
             type_builder,
             env: HashMap::default(),
+            type_params,
+            system,
         };
         translator.build_env_from_params();
         translator
@@ -171,6 +193,8 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
             self.tcx,
             self.def_ids.clone(),
             self.local_def_id.to_def_id(),
+            self.type_params.clone(),
+            self.system.clone(),
         );
         self
     }
@@ -204,29 +228,48 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         }
     }
 
+    fn instantiate_generics<T>(
+        &self,
+        ty: T,
+        generic_args: mir_ty::GenericArgsRef<'tcx>,
+    ) -> Option<T>
+    where
+        T: TypeFoldable<TyCtxt<'tcx>>,
+    {
+        if !self.generic_args.is_empty() {
+            Some(mir_ty::EarlyBinder::bind(ty).instantiate(self.tcx, generic_args))
+        } else {
+            None
+        }
+    }
+
     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 instantiated = self
+            .instantiate_generics(ty, self.generic_args)
+            .unwrap_or(ty);
         let typing_env = mir_ty::TypingEnv::fully_monomorphized();
         self.tcx.normalize_erasing_regions(typing_env, instantiated)
     }
 
     fn pat_ty(&self, pat: &'tcx rustc_hir::Pat<'tcx>) -> mir_ty::Ty<'tcx> {
         let ty = self.typeck.pat_ty(pat);
-        let instantiated = mir_ty::EarlyBinder::bind(ty).instantiate(self.tcx, self.generic_args);
+        let instantiated = self
+            .instantiate_generics(ty, self.generic_args)
+            .unwrap_or(ty);
         let typing_env = mir_ty::TypingEnv::fully_monomorphized();
         self.tcx.normalize_erasing_regions(typing_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(self.tcx, self.generic_args)
-            .skip_binder()
-            .inputs()
-            .to_vec();
+        let fn_sig = self.tcx.fn_sig(self.local_def_id.to_def_id());
+        let binder = if self.generic_args.is_empty() {
+            fn_sig.skip_binder()
+        } else {
+            fn_sig.instantiate(self.tcx, self.generic_args)
+        };
+        let params = binder.skip_binder().inputs().to_vec();
         FormulaFn {
             params: IndexVec::from_raw(params),
             formula,
@@ -495,8 +538,12 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                                 outer_generic_args = ?self.generic_args,
                                 "resolving predicate call in formula"
                             );
-                            let generic_args = mir_ty::EarlyBinder::bind(generic_args)
-                                .instantiate(self.tcx, self.generic_args);
+                            let (mut is_unresolved_args, generic_args) =
+                                match self.instantiate_generics(generic_args, self.generic_args) {
+                                    Some(args) => (false, args),
+                                    None => (true, generic_args),
+                                };
+
                             let instance = mir_ty::Instance::try_resolve(
                                 self.tcx,
                                 typing_env,
@@ -507,11 +554,18 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                             let pred_def_id = if let Some(instance) = instance {
                                 instance.def_id()
                             } else {
+                                is_unresolved_args = true;
                                 def_id
                             };
-                            let pred = refine::user_defined_pred(self.tcx, pred_def_id);
+
+                            let pred = if is_unresolved_args {
+                                refine::forall_pred(self.tcx, pred_def_id).into()
+                            } else {
+                                refine::user_defined_pred(self.tcx, pred_def_id).into()
+                            };
+                            tracing::debug!("resolved predicate call in formula: {:?}", pred);
                             let arg_terms = args.iter().map(|e| self.to_term(e)).collect();
-                            let atom = chc::Atom::new(pred.into(), arg_terms);
+                            let atom = chc::Atom::new(pred, arg_terms);
                             return FormulaOrTerm::Formula(chc::Formula::Atom(atom));
                         }
                     }

src/analyze/basic_block.rs

@@ -636,7 +636,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         def_id: DefId,
         args: mir_ty::GenericArgsRef<'tcx>,
     ) -> rty::Type<rty::Closed> {
-        if let Some(def_ty) = self.ctx.def_ty_with_args(def_id, args) {
+        let caller_def_id = self.type_builder.owner_fn_id;
+        if let Some(def_ty) = self.ctx.def_ty_with_args(def_id, args, caller_def_id) {
             return def_ty.ty;
         }
 
@@ -648,7 +649,10 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             );
         }
         tracing::info!(?def_id, ?resolved_def_id, ?resolved_args, "resolved");
-        let Some(def_ty) = self.ctx.def_ty_with_args(resolved_def_id, resolved_args) else {
+        let Some(def_ty) = self
+            .ctx
+            .def_ty_with_args(resolved_def_id, resolved_args, caller_def_id)
+        else {
             panic!(
                 "unknown def (resolved): {:?}, args: {:?}",
                 resolved_def_id, resolved_args
@@ -1166,14 +1170,15 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         ctx: &'ctx mut analyze::Analyzer<'tcx>,
         local_def_id: LocalDefId,
         basic_block: BasicBlock,
+        owner_fn_id: DefId,
     ) -> Self {
         let tcx = ctx.tcx;
         let drop_points = DropPoints::default();
         let body = Cow::Borrowed(tcx.optimized_mir(local_def_id.to_def_id()));
         let env = ctx.new_env();
         let local_decls = body.local_decls.clone();
         let prophecy_vars = Default::default();
-        let type_builder = TypeBuilder::new(tcx, ctx.def_ids(), local_def_id.to_def_id());
+        let type_builder = ctx.type_builder(ctx.def_ids(), owner_fn_id);
         Self {
             ctx,
             tcx,