Enable to handle unannotated polymorphic function

coord-e · coord-e · commit 62b3a6a51a36 · 2025-10-25T11:43:02.000+09:00
diff --git a/src/analyze.rs b/src/analyze.rs
@@ -11,13 +11,14 @@ use std::collections::HashMap;
 use std::rc::Rc;
 
 use rustc_hir::lang_items::LangItem;
+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 crate::chc;
 use crate::pretty::PrettyDisplayExt as _;
-use crate::refine::{self, BasicBlockType};
+use crate::refine::{self, BasicBlockType, TypeBuilder};
 use crate::rty;
 
 mod annot;
@@ -103,6 +104,17 @@ impl<'tcx> ReplacePlacesVisitor<'tcx> {
     }
 }
 
+#[derive(Debug, Clone)]
+struct DeferredDefTy<'tcx> {
+    cache: Rc<RefCell<HashMap<Vec<mir_ty::Ty<'tcx>>, rty::RefinedType>>>,
+}
+
+#[derive(Debug, Clone)]
+enum DefTy<'tcx> {
+    Concrete(rty::RefinedType),
+    Deferred(DeferredDefTy<'tcx>),
+}
+
 #[derive(Clone)]
 pub struct Analyzer<'tcx> {
     tcx: TyCtxt<'tcx>,
@@ -112,7 +124,7 @@ pub struct Analyzer<'tcx> {
     /// currently contains only local-def templates,
     /// but will be extended to contain externally known def's refinement types
     /// (at least for every defs referenced by local def bodies)
-    defs: HashMap<DefId, rty::RefinedType>,
+    defs: HashMap<DefId, DefTy<'tcx>>,
 
     /// Resulting CHC system.
     system: Rc<RefCell<chc::System>>,
@@ -207,11 +219,72 @@ impl<'tcx> Analyzer<'tcx> {
 
     pub fn register_def(&mut self, def_id: DefId, rty: rty::RefinedType) {
         tracing::info!(def_id = ?def_id, rty = %rty.display(), "register_def");
-        self.defs.insert(def_id, rty);
+        self.defs.insert(def_id, DefTy::Concrete(rty));
+    }
+
+    pub fn register_deferred_def(&mut self, def_id: DefId) {
+        tracing::info!(def_id = ?def_id, "register_deferred_def");
+        self.defs.insert(
+            def_id,
+            DefTy::Deferred(DeferredDefTy {
+                cache: Rc::new(RefCell::new(HashMap::new())),
+            }),
+        );
+    }
+
+    pub fn concrete_def_ty(&self, def_id: DefId) -> Option<&rty::RefinedType> {
+        self.defs.get(&def_id).and_then(|def_ty| match def_ty {
+            DefTy::Concrete(rty) => Some(rty),
+            DefTy::Deferred(_) => None,
+        })
     }
 
-    pub fn def_ty(&self, def_id: DefId) -> Option<&rty::RefinedType> {
-        self.defs.get(&def_id)
+    pub fn def_ty_with_args(
+        &mut self,
+        def_id: DefId,
+        args: mir_ty::GenericArgsRef<'tcx>,
+    ) -> Option<rty::RefinedType> {
+        let type_builder = TypeBuilder::new(self.tcx, def_id);
+        let rty_args: IndexVec<_, _> = args.types().map(|ty| type_builder.build(ty)).collect();
+
+        let deferred_ty = match self.defs.get(&def_id)? {
+            DefTy::Concrete(rty) => {
+                let mut def_ty = rty.clone();
+                def_ty.instantiate_ty_params(
+                    rty_args
+                        .clone()
+                        .into_iter()
+                        .map(rty::RefinedType::unrefined)
+                        .collect(),
+                );
+                return Some(def_ty);
+            }
+            DefTy::Deferred(deferred) => deferred,
+        };
+
+        let ty_args: Vec<_> = args.types().collect();
+        let deferred_ty_cache = Rc::clone(&deferred_ty.cache); // to cut reference to allow &mut self
+        if let Some(rty) = deferred_ty_cache.borrow().get(&ty_args) {
+            return Some(rty.clone());
+        }
+        let local_def_id = def_id.as_local()?;
+
+        let sig = self
+            .tcx
+            .fn_sig(def_id)
+            .instantiate(self.tcx, args)
+            .skip_binder();
+        let expected = self
+            .crate_analyzer()
+            .fn_def_ty_with_sig(local_def_id.to_def_id(), sig)
+            .unwrap();
+        self.local_def_analyzer(local_def_id)
+            .type_builder(type_builder.with_param_mapper(move |ty| rty_args[ty.idx].clone()))
+            .run(&expected);
+        deferred_ty_cache
+            .borrow_mut()
+            .insert(ty_args, expected.clone());
+        Some(expected)
     }
 
     pub fn register_basic_block_ty(
diff --git a/src/analyze/basic_block.rs b/src/analyze/basic_block.rs
@@ -263,12 +263,12 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 _ty,
             ) => {
                 let func_ty = match operand.const_fn_def() {
-                    Some((def_id, args)) => {
-                        if !args.is_empty() {
-                            tracing::warn!(?args, ?def_id, "generic args ignored");
-                        }
-                        self.ctx.def_ty(def_id).expect("unknown def").ty.clone()
-                    }
+                    Some((def_id, args)) => self
+                        .ctx
+                        .def_ty_with_args(def_id, args)
+                        .expect("unknown def")
+                        .ty
+                        .clone(),
                     _ => unimplemented!(),
                 };
                 PlaceType::with_ty_and_term(func_ty.vacuous(), chc::Term::null())
@@ -468,18 +468,12 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 let ret = rty::RefinedType::new(rty::Type::unit(), ret_formula.into());
                 rty::FunctionType::new([param1, param2].into_iter().collect(), ret).into()
             }
-            Some((def_id, args)) => {
-                if args.consts().next().is_some() {
-                    tracing::warn!(?args, ?def_id, "const generic args ignored");
-                }
-                let ty_args = args
-                    .types()
-                    .map(|ty| rty::RefinedType::unrefined(self.type_builder.build(ty)))
-                    .collect();
-                let mut def_ty = self.ctx.def_ty(def_id).expect("unknown def").clone();
-                def_ty.instantiate_ty_params(ty_args);
-                def_ty.ty.vacuous()
-            }
+            Some((def_id, args)) => self
+                .ctx
+                .def_ty_with_args(def_id, args)
+                .expect("unknown def")
+                .ty
+                .vacuous(),
             _ => self.operand_type(func.clone()).ty,
         };
         let expected_args: IndexVec<_, _> = args
diff --git a/src/analyze/crate_.rs b/src/analyze/crate_.rs
@@ -128,8 +128,19 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     #[tracing::instrument(skip(self), fields(def_id = %self.tcx.def_path_str(def_id)))]
     fn refine_fn_def(&mut self, def_id: DefId) {
         let sig = self.tcx.fn_sig(def_id);
-        let sig = sig.instantiate_identity().skip_binder(); // TODO: is it OK?
+        let sig = sig.instantiate_identity().skip_binder();
+        if let Some(rty) = self.fn_def_ty_with_sig(def_id, sig) {
+            self.ctx.register_def(def_id, rty);
+        } else {
+            self.ctx.register_deferred_def(def_id);
+        }
+    }
 
+    pub fn fn_def_ty_with_sig(
+        &mut self,
+        def_id: DefId,
+        sig: mir_ty::FnSig<'tcx>,
+    ) -> Option<rty::RefinedType> {
         let mut param_resolver = analyze::annot::ParamResolver::default();
         for (input_ident, input_ty) in self.tcx.fn_arg_names(def_id).iter().zip(sig.inputs()) {
             let input_ty = TypeBuilder::new(self.tcx, def_id).build(*input_ty);
@@ -198,8 +209,14 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         if let Some(ret_rty) = ret_annot {
             builder.ret_rty(ret_rty);
         }
-        let rty = rty::RefinedType::unrefined(builder.build().into());
-        self.ctx.register_def(def_id, rty);
+
+        // can't generate template with type parameter...
+        use mir_ty::TypeVisitableExt as _;
+        if builder.would_contain_template() && sig.has_param() {
+            None
+        } else {
+            Some(rty::RefinedType::unrefined(builder.build().into()))
+        }
     }
 
     fn analyze_local_defs(&mut self) {
@@ -211,10 +228,16 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 tracing::info!(?local_def_id, "trusted");
                 continue;
             }
+            let Some(expected) = self.ctx.concrete_def_ty(local_def_id.to_def_id()) else {
+                // when the local_def_id is deferred it would be skipped
+                continue;
+            };
+
             // check polymorphic function def by replacing type params with some opaque type
+            // (and this is no-op if the function is mono)
             let type_builder = TypeBuilder::new(self.tcx, local_def_id.to_def_id())
                 .with_param_mapper(|_| rty::Type::int());
-            let mut expected = self.ctx.def_ty(local_def_id.to_def_id()).unwrap().clone();
+            let mut expected = expected.clone();
             let subst = rty::TypeParamSubst::new(
                 expected
                     .free_ty_params()
@@ -236,7 +259,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             // TODO: replace code here with relate_* in Env + Refine context (created with empty env)
             let entry_ty = self
                 .ctx
-                .def_ty(def_id)
+                .concrete_def_ty(def_id)
                 .unwrap()
                 .ty
                 .as_function()
diff --git a/src/refine/template.rs b/src/refine/template.rs
@@ -83,9 +83,12 @@ where
 pub struct TypeBuilder<'tcx> {
     tcx: mir_ty::TyCtxt<'tcx>,
     /// Maps index in [`mir_ty::ParamTy`] to [`rty::TypeParamIdx`].
-    /// These indices may differ because we skip lifetime parameters.
+    /// These indices may differ because we skip lifetime parameters and they always need to be
+    /// mapped when we translate a [`mir_ty::ParamTy`] to [`rty::ParamType`].
     /// See [`rty::TypeParamIdx`] for more details.
     param_idx_mapping: HashMap<u32, rty::TypeParamIdx>,
+    /// Optionally also want to further map rty::ParamType to other rty::Type before generating
+    /// templates. This is no-op by default.
     param_type_mapper: std::rc::Rc<dyn ParamTypeMapper>,
 }
 
@@ -100,7 +103,7 @@ impl<'tcx> TypeBuilder<'tcx> {
                 mir_ty::GenericParamDefKind::Type { .. } => {
                     param_idx_mapping.insert(i as u32, param_idx_mapping.len().into());
                 }
-                mir_ty::GenericParamDefKind::Const { .. } => unimplemented!(),
+                mir_ty::GenericParamDefKind::Const { .. } => {}
             }
         }
         Self {
@@ -397,6 +400,17 @@ impl<'tcx, 'a, R> FunctionTemplateTypeBuilder<'tcx, 'a, R> {
         self.ret_rty = Some(rty);
         self
     }
+
+    pub fn would_contain_template(&self) -> bool {
+        if self.param_tys.is_empty() {
+            return self.ret_rty.is_none();
+        }
+
+        let last_param_idx = rty::FunctionParamIdx::from(self.param_tys.len() - 1);
+        let param_annotated =
+            self.param_refinement.is_some() || self.param_rtys.contains_key(&last_param_idx);
+        self.ret_rty.is_none() || !param_annotated
+    }
 }
 
 impl<'tcx, 'a, R> FunctionTemplateTypeBuilder<'tcx, 'a, R>
diff --git a/tests/ui/fail/fn_poly.rs b/tests/ui/fail/fn_poly.rs
@@ -0,0 +1,9 @@
+//@error-in-other-file: Unsat
+
+fn left<T, U>(x: (T, U)) -> T {
+    x.0
+}
+
+fn main() {
+    assert!(left((42, 0)) == 0);
+}
diff --git a/tests/ui/pass/fn_poly.rs b/tests/ui/pass/fn_poly.rs
@@ -0,0 +1,9 @@
+//@check-pass
+
+fn left<T, U>(x: (T, U)) -> T {
+    x.0
+}
+
+fn main() {
+    assert!(left((42, 0)) == 42);
+}
