Workaround Closure model via manually replacing type

Author: coord-e

src/analyze/annot.rs

@@ -89,6 +89,14 @@ pub fn array_model_path() -> [Symbol; 3] {
     ]
 }
 
+pub fn closure_model_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("closure_model"),
+    ]
+}
+
 /// A [`annot::Resolver`] implementation for resolving function parameters.
 ///
 /// The parameter names and their sorts needs to be configured via

src/analyze/did_cache.rs

@@ -18,6 +18,7 @@ struct DefIds {
     mut_model: OnceCell<Option<DefId>>,
     box_model: OnceCell<Option<DefId>>,
     array_model: OnceCell<Option<DefId>>,
+    closure_model: OnceCell<Option<DefId>>,
 }
 
 /// Retrieves and caches well-known [`DefId`]s.
@@ -127,4 +128,11 @@ impl<'tcx> DefIdCache<'tcx> {
             .array_model
             .get_or_init(|| self.annotated_def(&crate::analyze::annot::array_model_path()))
     }
+
+    pub fn closure_model(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .closure_model
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::closure_model_path()))
+    }
 }

src/refine/template.rs

@@ -110,7 +110,50 @@ impl<'tcx> TypeBuilder<'tcx> {
         rty::ParamType::new(index).into()
     }
 
-    fn resolve_model_ty(&self, ty: mir_ty::Ty<'tcx>) -> mir_ty::Ty<'tcx> {
+    /// Replaces {closure} types with thrust_models::Closure<{closure}>.
+    ///
+    /// Ideally, we want to have `impl<F> Model for F where F: Fn` instead of this and let
+    /// normalization do the work. However, it doesn't work:
+    /// 1. it simply conflicts with other blanket impls
+    /// 2. attempts to fake the impl via override_queries have failed so far
+    fn replace_closure_model(&self, ty: mir_ty::Ty<'tcx>) -> mir_ty::Ty<'tcx> {
+        let Some(closure_model_id) = self.def_ids.closure_model() else {
+            return ty;
+        };
+
+        struct ReplaceClosureModel<'tcx> {
+            tcx: mir_ty::TyCtxt<'tcx>,
+            closure_model_id: DefId,
+        }
+
+        use mir_ty::fold::TypeFoldable;
+        impl<'tcx> mir_ty::fold::TypeFolder<mir_ty::TyCtxt<'tcx>> for ReplaceClosureModel<'tcx> {
+            fn interner(&self) -> mir_ty::TyCtxt<'tcx> {
+                self.tcx
+            }
+
+            fn fold_ty(&mut self, ty: mir_ty::Ty<'tcx>) -> mir_ty::Ty<'tcx> {
+                use mir_ty::fold::TypeSuperFoldable;
+                if let mir_ty::TyKind::Closure(_, args) = ty.kind() {
+                    let args = self
+                        .tcx
+                        .mk_args(&[args.as_closure().tupled_upvars_ty().into()]);
+                    let adt_def = self.tcx.adt_def(self.closure_model_id);
+                    return mir_ty::Ty::new_adt(self.tcx, adt_def, args);
+                }
+                ty.super_fold_with(self)
+            }
+        }
+
+        ty.fold_with(&mut ReplaceClosureModel {
+            tcx: self.tcx,
+            closure_model_id,
+        })
+    }
+
+    fn resolve_model_ty(&self, orig_ty: mir_ty::Ty<'tcx>) -> mir_ty::Ty<'tcx> {
+        let ty = self.replace_closure_model(orig_ty);
+
         let Some(model_ty_def_id) = self.def_ids.model_ty() else {
             return ty;
         };
@@ -151,6 +194,11 @@ impl<'tcx> TypeBuilder<'tcx> {
             return Some(rty::ArrayType::new(idx_ty, elem_ty).into());
         }
 
+        if Some(adt.did()) == self.def_ids.closure_model() {
+            let tupled_upvars_ty = args.type_at(0);
+            return Some(self.build(tupled_upvars_ty));
+        }
+
         None
     }
 
@@ -210,7 +258,6 @@ impl<'tcx> TypeBuilder<'tcx> {
                     unimplemented!("unsupported ADT: {:?}", ty);
                 }
             }
-            mir_ty::TyKind::Closure(_, args) => self.build(args.as_closure().tupled_upvars_ty()),
             kind => unimplemented!("unrefined_ty: {:?}", kind),
         }
     }
@@ -310,6 +357,11 @@ where
             return Some(rty::ArrayType::new(idx_ty, elem_ty).into());
         }
 
+        if Some(adt.did()) == self.inner.def_ids.closure_model() {
+            let tupled_upvars_ty = args.type_at(0);
+            return Some(self.build(tupled_upvars_ty));
+        }
+
         None
     }
 
@@ -361,7 +413,6 @@ where
                     unimplemented!("unsupported ADT: {:?}", ty);
                 }
             }
-            mir_ty::TyKind::Closure(_, args) => self.build(args.as_closure().tupled_upvars_ty()),
             kind => unimplemented!("ty: {:?}", kind),
         }
     }

std.rs

@@ -21,6 +21,9 @@ mod thrust_models {
         #[thrust::def::array_model]
         pub struct Array<I, T>(PhantomData<I>, PhantomData<T>);
 
+        #[thrust::def::closure_model]
+        pub struct Closure<T>(PhantomData<T>);
+
         pub struct Vec<T>(pub Array<Int, T>, pub usize);
     }
 
@@ -56,6 +59,10 @@ mod thrust_models {
         type Ty = ();
     }
 
+    impl<T> Model for model::Closure<T> {
+        type Ty = model::Closure<T>;
+    }
+
     impl<T0> Model for (T0,) where T0: Model {
         type Ty = (<T0 as Model>::Ty,);
     }