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;
@@ -196,6 +197,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> {
@@ -223,6 +231,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> {
@@ -246,6 +256,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,
@@ -254,6 +265,7 @@ impl<'tcx> Analyzer<'tcx> {
             basic_blocks,
             def_ids: did_cache::DefIdCache::new(tcx),
             enum_defs,
+            type_params,
         }
     }
 
@@ -287,7 +299,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 {
@@ -386,14 +398,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> {
@@ -407,6 +418,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)?;
 
@@ -415,9 +427,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(), owner_fn_id);
         let formula_fn = translator.to_formula_fn();
         deferred_formula_fn_cache
             .borrow_mut()
@@ -431,8 +449,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) => {
@@ -603,8 +622,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) {
@@ -698,6 +728,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>,
@@ -728,7 +759,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
@@ -747,6 +780,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>>,
@@ -778,7 +812,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;
 
@@ -141,15 +146,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,
@@ -159,6 +179,8 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
             def_ids,
             type_builder,
             env: HashMap::default(),
+            type_params,
+            system,
         };
         translator.build_env_from_params();
         translator
@@ -169,12 +191,14 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         self
     }
 
-    pub fn with_def_id_cache(mut self, def_ids: DefIdCache<'tcx>) -> Self {
+    pub fn with_def_id_cache(mut self, def_ids: DefIdCache<'tcx>, owner_fn_id: DefId) -> Self {
         self.def_ids = def_ids;
         self.type_builder = TypeBuilder::new(
             self.tcx,
             self.def_ids.clone(),
-            self.local_def_id.to_def_id(),
+            owner_fn_id,
+            self.type_params.clone(),
+            self.system.clone(),
         );
         self
     }
@@ -208,29 +232,50 @@ 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)
+        self.tcx
+            .try_normalize_erasing_regions(typing_env, instantiated)
+            .unwrap_or(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,
@@ -499,8 +544,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,
@@ -511,11 +560,35 @@ 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 typeck_results = self.tcx.typeck(self.local_def_id);
+                            let pred = if is_unresolved_args {
+                                let pred = refine::forall_pred(self.tcx, pred_def_id);
+                                let sig = args
+                                    .iter()
+                                    .map(|e| {
+                                        let ty = typeck_results.expr_ty(e);
+                                        self.type_builder.build(ty).to_sort()
+                                    })
+                                    .collect();
+                                tracing::debug!(
+                                    "register ForallPred {:?} with signature {:?}",
+                                    pred,
+                                    sig
+                                );
+                                self.system
+                                    .borrow_mut()
+                                    .register_forall_pred(pred.clone(), sig);
+                                pred.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));
                         }
                     }