Support polymorphic functions

src/analyze.rs

@@ -17,7 +17,7 @@ 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 +103,17 @@ impl<'tcx> ReplacePlacesVisitor<'tcx> {
     }
 }
 
+#[derive(Debug, Clone)]
+struct DeferredDefTy {
+    cache: Rc<RefCell<HashMap<rty::TypeArgs, rty::RefinedType>>>,
+}
+
+#[derive(Debug, Clone)]
+enum DefTy {
+    Concrete(rty::RefinedType),
+    Deferred(DeferredDefTy),
+}
+
 #[derive(Clone)]
 pub struct Analyzer<'tcx> {
     tcx: TyCtxt<'tcx>,
@@ -112,7 +123,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>,
 
     /// Resulting CHC system.
     system: Rc<RefCell<chc::System>>,
@@ -207,11 +218,65 @@ 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 def_ty(&self, def_id: DefId) -> Option<&rty::RefinedType> {
-        self.defs.get(&def_id)
+    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_with_args(
+        &mut self,
+        def_id: DefId,
+        rty_args: rty::TypeArgs,
+    ) -> Option<rty::RefinedType> {
+        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 deferred_ty_cache = Rc::clone(&deferred_ty.cache); // to cut reference to allow &mut self
+        if let Some(rty) = deferred_ty_cache.borrow().get(&rty_args) {
+            return Some(rty.clone());
+        }
+
+        let type_builder = TypeBuilder::new(self.tcx, def_id).with_param_mapper({
+            let rty_args = rty_args.clone();
+            move |ty: rty::ParamType| rty_args[ty.idx].clone()
+        });
+        let mut analyzer = self.local_def_analyzer(def_id.as_local()?);
+        analyzer.type_builder(type_builder);
+
+        let expected = analyzer.expected_ty();
+        deferred_ty_cache
+            .borrow_mut()
+            .insert(rty_args, expected.clone());
+
+        analyzer.run(&expected);
+        Some(expected)
     }
 
     pub fn register_basic_block_ty(

src/analyze/basic_block.rs

@@ -33,6 +33,7 @@ pub struct Analyzer<'tcx, 'ctx> {
     basic_block: BasicBlock,
     body: Cow<'tcx, Body<'tcx>>,
 
+    type_builder: TypeBuilder<'tcx>,
     env: Env,
     local_decls: IndexVec<Local, mir::LocalDecl<'tcx>>,
     // TODO: remove this
@@ -56,10 +57,6 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         self.ctx.basic_block_ty(self.local_def_id, bb)
     }
 
-    fn type_builder(&self) -> TypeBuilder<'tcx> {
-        TypeBuilder::new(self.tcx, self.local_def_id.to_def_id())
-    }
-
     fn bind_local(&mut self, local: Local, rty: rty::RefinedType<Var>) {
         let rty = if self.is_mut_local(local) {
             // elaboration:
@@ -226,7 +223,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                         let rty_args: IndexVec<_, _> = args
                             .types()
                             .map(|ty| {
-                                self.type_builder()
+                                self.type_builder
                                     .for_template(&mut self.ctx)
                                     .with_scope(&self.env)
                                     .build_refined(ty)
@@ -267,10 +264,13 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             ) => {
                 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()
+                        let rty_args: IndexVec<_, _> =
+                            args.types().map(|ty| self.type_builder.build(ty)).collect();
+                        self.ctx
+                            .def_ty_with_args(def_id, rty_args)
+                            .expect("unknown def")
+                            .ty
+                            .clone()
                     }
                     _ => unimplemented!(),
                 };
@@ -440,7 +440,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
             // TODO: move this to well-known defs?
             Some((def_id, args)) if self.is_box_new(def_id) => {
                 let inner_ty = self
-                    .type_builder()
+                    .type_builder
                     .for_template(&mut self.ctx)
                     .build(args.type_at(0))
                     .vacuous();
@@ -454,7 +454,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 rty::FunctionType::new([param].into_iter().collect(), ret).into()
             }
             Some((def_id, args)) if self.is_mem_swap(def_id) => {
-                let inner_ty = self.type_builder().build(args.type_at(0)).vacuous();
+                let inner_ty = self.type_builder.build(args.type_at(0)).vacuous();
                 let param1 =
                     rty::RefinedType::unrefined(rty::PointerType::mut_to(inner_ty.clone()).into());
                 let param2 =
@@ -472,14 +472,11 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 rty::FunctionType::new([param1, param2].into_iter().collect(), ret).into()
             }
             Some((def_id, args)) => {
-                if !args.is_empty() {
-                    tracing::warn!(?args, ?def_id, "generic args ignored");
-                }
+                let rty_args = args.types().map(|ty| self.type_builder.build(ty)).collect();
                 self.ctx
-                    .def_ty(def_id)
+                    .def_ty_with_args(def_id, rty_args)
                     .expect("unknown def")
                     .ty
-                    .clone()
                     .vacuous()
             }
             _ => self.operand_type(func.clone()).ty,
@@ -541,7 +538,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     }
 
     fn add_prophecy_var(&mut self, statement_index: usize, ty: mir_ty::Ty<'tcx>) {
-        let ty = self.type_builder().build(ty);
+        let ty = self.type_builder.build(ty);
         let temp_var = self.env.push_temp_var(ty.vacuous());
         self.prophecy_vars.insert(statement_index, temp_var);
         tracing::debug!(stmt_idx = %statement_index, temp_var = ?temp_var, "add_prophecy_var");
@@ -562,7 +559,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         referent: mir::Place<'tcx>,
         prophecy_ty: mir_ty::Ty<'tcx>,
     ) -> rty::RefinedType<Var> {
-        let prophecy_ty = self.type_builder().build(prophecy_ty);
+        let prophecy_ty = self.type_builder.build(prophecy_ty);
         let prophecy = self.env.push_temp_var(prophecy_ty.vacuous());
         let place = self.elaborate_place_for_borrow(&referent);
         self.env.borrow_place(place, prophecy).into()
@@ -675,7 +672,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
 
             let decl = self.local_decls[destination].clone();
             let rty = self
-                .type_builder()
+                .type_builder
                 .for_template(&mut self.ctx)
                 .with_scope(&self.env)
                 .build_refined(decl.ty);
@@ -749,7 +746,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     #[tracing::instrument(skip(self))]
     fn ret_template(&mut self) -> rty::RefinedType<Var> {
         let ret_ty = self.body.local_decls[mir::RETURN_PLACE].ty;
-        self.type_builder()
+        self.type_builder
             .for_template(&mut self.ctx)
             .with_scope(&self.env)
             .build_refined(ret_ty)
@@ -955,13 +952,15 @@ 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());
         Self {
             ctx,
             tcx,
             local_def_id,
             drop_points,
             basic_block,
             body,
+            type_builder,
             env,
             local_decls,
             prophecy_vars,
@@ -989,6 +988,11 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         self
     }
 
+    pub fn type_builder(&mut self, type_builder: TypeBuilder<'tcx>) -> &mut Self {
+        self.type_builder = type_builder;
+        self
+    }
+
     pub fn run(&mut self, expected: &BasicBlockType) {
         let span = tracing::info_span!("bb", bb = ?self.basic_block);
         let _guard = span.enter();