From cdf1104fbc2ebf0ba517732cb114bf1bc0e80fd5 Mon Sep 17 00:00:00 2001
From: Son Ho <hosonmarc@gmail.com>
Date: Mon, 4 Dec 2023 15:44:12 +0100
Subject: [PATCH] Add generics to AggregateKind::Closure

---
 charon/src/expressions.rs                  |  2 +-
 charon/src/expressions_utils.rs            | 14 +++++++++++---
 charon/src/translate_functions_to_ullbc.rs | 15 +++++++++++++--
 charon/src/translate_traits.rs             |  1 +
 4 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/charon/src/expressions.rs b/charon/src/expressions.rs
index 06e3a8356..0e02d4f1e 100644
--- a/charon/src/expressions.rs
+++ b/charon/src/expressions.rs
@@ -382,5 +382,5 @@ pub enum AggregateKind {
     Array(Ty, ConstGeneric),
     /// Aggregated values for closures group the function id together with its
     /// state.
-    Closure(FunDeclId::Id),
+    Closure(FunDeclId::Id, GenericArgs),
 }
diff --git a/charon/src/expressions_utils.rs b/charon/src/expressions_utils.rs
index 565fe74ed..57f9c9b6b 100644
--- a/charon/src/expressions_utils.rs
+++ b/charon/src/expressions_utils.rs
@@ -276,8 +276,13 @@ impl Rvalue {
                     AggregateKind::Array(_, len) => {
                         format!("[{}; {}]", ops_s.join(", "), len.fmt_with_ctx(ctx))
                     }
-                    AggregateKind::Closure(fn_id) => {
-                        format!("{{{}}} {{{}}}", ctx.format_object(*fn_id), ops_s.join(", "))
+                    AggregateKind::Closure(fn_id, generics) => {
+                        format!(
+                            "{{{}{}}} {{{}}}",
+                            ctx.format_object(*fn_id),
+                            generics.fmt_with_ctx_split_trait_refs(ctx),
+                            ops_s.join(", ")
+                        )
                     }
                 }
             }
@@ -464,7 +469,10 @@ pub trait ExprVisitor: crate::types::TypeVisitor {
                 self.visit_ty(ty);
                 self.visit_const_generic(cg);
             }
-            Closure(fn_id) => self.visit_fun_decl_id(fn_id),
+            Closure(fn_id, generics) => {
+                self.visit_fun_decl_id(fn_id);
+                self.visit_generic_args(generics);
+            }
         }
     }
 
diff --git a/charon/src/translate_functions_to_ullbc.rs b/charon/src/translate_functions_to_ullbc.rs
index b11b85f82..913cd0686 100644
--- a/charon/src/translate_functions_to_ullbc.rs
+++ b/charon/src/translate_functions_to_ullbc.rs
@@ -781,10 +781,21 @@ impl<'tcx, 'ctx, 'ctx1> BodyTransCtx<'tcx, 'ctx, 'ctx1> {
                         let akind = AggregateKind::Adt(type_id, variant_id, generics);
                         Ok(Rvalue::Aggregate(akind, operands_t))
                     }
-                    hax::AggregateKind::Closure(def_id, sig) => {
+                    hax::AggregateKind::Closure(def_id, substs, trait_refs, sig) => {
                         trace!("Closure:\n\n- def_id: {:?}\n\n- sig:\n{:?}", def_id, sig);
+
+                        // Translate the substitution
+                        let generics = self.translate_substs_and_trait_refs(
+                            span,
+                            erase_regions,
+                            None,
+                            substs,
+                            trait_refs,
+                        )?;
+
                         let def_id = self.translate_fun_decl_id(def_id.rust_def_id.unwrap());
-                        let akind = AggregateKind::Closure(def_id);
+                        let akind = AggregateKind::Closure(def_id, generics);
+
                         Ok(Rvalue::Aggregate(akind, operands_t))
                     }
                     hax::AggregateKind::Generator(_def_id, _subst, _movability) => {
diff --git a/charon/src/translate_traits.rs b/charon/src/translate_traits.rs
index 59dd34912..ae87d6c3f 100644
--- a/charon/src/translate_traits.rs
+++ b/charon/src/translate_traits.rs
@@ -499,6 +499,7 @@ impl<'tcx, 'ctx> TransCtx<'tcx, 'ctx> {
                 tcx.param_env(rust_id),
                 rust_trait_ref.def_id,
                 rust_trait_ref.substs,
+                None,
             );
             let parent_trait_refs: Vec<TraitRef> =
                 bt_ctx.translate_trait_impl_sources(span, erase_regions, &parent_trait_refs)?;