Add generic params to snapshot constructors

prusti-encoder/src/encoders/type/domain.rs

@@ -14,6 +14,8 @@ use vir::{
     AdtDestructor, Arity, CallableIdn, CastType, CompType, DomainAxiomData, DomainIdnCSnap, FunctionIdn, Type
 };
 
+use crate::encoders::lifted::generic::LiftedGeneric;
+
 use super::{
     most_generic_ty::{extract_type_params, get_vir_base_name_kind, MostGenericTy},
     rust_ty_snapshots::RustTySnapshotsEnc,
@@ -54,9 +56,11 @@ pub struct DomainDataMutRef<'vir> {
 pub struct DomainDataStruct<'vir> {
     /// Construct domain from snapshots of fields or for primitive types
     /// from the single Viper primitive value.
-    pub field_snaps_to_snap: FunctionIdn<'vir, vir::ManySnap, vir::CSnap>,
+    pub field_snaps_to_snap: FunctionIdn<'vir, (vir::ManySnap, vir::ManyTyVal), vir::CSnap>,
     /// Functions to access the fields.
     pub field_access: &'vir [AdtDestructor<'vir, vir::CSnap, vir::Snap>],
+    /// Functions to access the type parameters.
+    pub ty_param_accessors: &'vir [AdtDestructor<'vir, vir::CSnap, vir::TyVal>],
 }
 #[derive(Clone, Copy, Debug)]
 pub struct DomainDataEnum<'vir> {

prusti-encoder/src/encoders/type/kinds/adt.rs

@@ -2,7 +2,7 @@ use crate::encoders::{
     domain::{
         AdtBuilder, DomainDataEnum, DomainDataStruct, DomainDataVariant, DomainEnc, DomainEncOutputRef, DomainEncSpecifics, FieldTy, PureTypeBuilder, PureTypeCommon
     },
-    lifted::ty::{EncodeGenericsAsParamTy, LiftedTyEnc},
+    lifted::ty::{EncodeGenericsAsLifted, EncodeGenericsAsParamTy, LiftedTyEnc},
     predicate::{
         PredicateBuilder, PredicateEncData, PredicateEncDataEnum, PredicateEncDataStruct,
         PredicateEncDataVariant, RefToIndirectPred,
@@ -33,7 +33,7 @@ pub(crate) fn domain<'vir>(
         .iter()
         .flat_map(ty::GenericArg::as_type)
         .map(|ty| {
-            deps.require_local::<LiftedTyEnc<EncodeGenericsAsParamTy>>(ty)
+            deps.require_local::<LiftedTyEnc<EncodeGenericsAsLifted>>(ty)
                 .unwrap()
                 .expect_generic()
         })
@@ -48,32 +48,35 @@ pub(crate) fn domain<'vir>(
                 rust_ty_data: None,
             }], task_key, output_ref, &generics, deps, builder)?;
             */
-            let (field_snaps_to_snap, field_access) = super::structlike::domain(
+            let (field_snaps_to_snap, field_access, ty_param_accessors) = super::structlike::domain(
                 "",
                 &[FieldTy::from_ty(
                     deps,
-                    generics[0].to_ty(builder.vcx.tcx()),
+                    params[0].expect_ty(),
                 )?],
+                &generics,
                 &mut builder,
                 None,
             );
 
             Ok((DomainEncSpecifics::StructLike(DomainDataStruct {
                 field_snaps_to_snap,
                 field_access,
+                ty_param_accessors,
             }), Ok(builder)))
         }
         ty::AdtKind::Struct => {
             let variant = adt.non_enum_variant();
             let fields = FieldTy::mk_field_tys(builder.vcx, deps, variant, params)?;
 
-            let (field_snaps_to_snap, field_access) = super::structlike::domain(
-                "", &fields, &mut builder, None,
+            let (field_snaps_to_snap, field_access, ty_param_accessors) = super::structlike::domain(
+                "", &fields, &generics, &mut builder, None,
             );
 
             Ok((DomainEncSpecifics::StructLike(DomainDataStruct {
                 field_snaps_to_snap,
                 field_access,
+                ty_param_accessors,
             }), Ok(builder)))
         }
         ty::AdtKind::Enum => {
@@ -101,9 +104,10 @@ pub(crate) fn domain<'vir>(
 
                     let fields = FieldTy::mk_field_tys(builder.vcx, deps, variant, params)?;
 
-                    let (field_snaps_to_snap, field_access) = super::structlike::domain(
+                    let (field_snaps_to_snap, field_access, ty_param_accessors) = super::structlike::domain(
                         &format!("{var_idx_num}_"),
                         &fields,
+                        &generics,
                         &mut builder,
                         Some(discr),
                     );
@@ -115,6 +119,7 @@ pub(crate) fn domain<'vir>(
                         fields: DomainDataStruct {
                             field_snaps_to_snap,
                             field_access,
+                            ty_param_accessors,
                         },
                     })
                 })

prusti-encoder/src/encoders/type/kinds/closure.rs

@@ -31,12 +31,13 @@ pub(crate) fn domain<'vir>(
         .map(|ty| FieldTy::from_ty(deps, ty))
         .collect::<Result<Vec<_>, _>>()?;
 
-    let (field_snaps_to_snap, field_access) =
-        super::structlike::domain("", &fields, &mut builder, None);
+    let (field_snaps_to_snap, field_access, ty_param_accessors) =
+        super::structlike::domain("", &fields, &[], &mut builder, None);
 
     Ok((DomainEncSpecifics::StructLike(DomainDataStruct {
         field_snaps_to_snap,
         field_access,
+        ty_param_accessors,
     }), Ok(builder)))
 
     /*

prusti-encoder/src/encoders/type/kinds/str.rs

@@ -18,11 +18,12 @@ pub(crate) fn domain<'vir>(
     let ty_kind = ty.kind();
     assert_eq!(*ty_kind, ty::TyKind::Str);
 
-    let dummy_cons_ident = builder.function("cons", &[][..], builder.self_type());
+    let dummy_cons_ident = builder.function("cons", (&[][..], &[][..]), builder.self_type());
 
     Ok((DomainEncSpecifics::StructLike(DomainDataStruct {
         field_snaps_to_snap: dummy_cons_ident,
         field_access: &[],
+        ty_param_accessors: &[],
     }), Err(builder)))
 }
 

prusti-encoder/src/encoders/type/kinds/structlike.rs

@@ -1,8 +1,5 @@
 use crate::encoders::{
-    domain::{AdtBuilder, FieldTy},
-    predicate::PredicateBuilder,
-    rust_ty_predicates::RustTyPredicatesEncOutputRef,
-    snapshot::SnapshotEncOutput, PredicateEnc,
+    domain::{AdtBuilder, FieldTy}, lifted::generic::LiftedGeneric, predicate::PredicateBuilder, rust_ty_predicates::RustTyPredicatesEncOutputRef, snapshot::SnapshotEncOutput, PredicateEnc
 };
 use prusti_rustc_interface::middle::ty::ParamTy;
 use task_encoder::{EncodeFullError, TaskEncoder, TaskEncoderDependencies};
@@ -11,23 +8,29 @@ use vir::{vir_format, CastType, FunctionIdn, HasType, PredicateIdn};
 pub fn domain<'vir>(
     prefix: &str,
     fields: &[FieldTy<'vir>],
+    generics: &[LiftedGeneric<'vir>],
     builder: &mut AdtBuilder<'vir>,
     discr: Option<vir::ExprCSnap<'vir>>,
 ) -> (
-    FunctionIdn<'vir, vir::ManySnap, vir::CSnap>,
+    FunctionIdn<'vir, (vir::ManySnap, vir::ManyTyVal), vir::CSnap>,
     &'vir [vir::AdtDestructor<'vir, vir::CSnap, vir::Snap>],
+    &'vir [vir::AdtDestructor<'vir, vir::CSnap, vir::TyVal>],
 ) {
     let field_tys = builder.vcx.alloc_slice(&fields.iter().map(|f| f.ty).collect::<Vec<_>>());
-    let (cons, des) = builder.constructor(prefix, field_tys, discr);
-    (cons, builder.vcx.alloc_slice(&des).downcast_ty())
+    let generic_tys = builder.vcx.alloc_slice(
+        generics.iter().map(LiftedGeneric::ty).collect::<Vec<_>>().as_slice(),
+    );
+    let (cons, des) = builder.constructor(prefix, (field_tys, generic_tys), discr);
+    let (snap_fields, gen_fields) = builder.vcx.alloc_slice(&des).split_at(fields.len());
+    (cons, snap_fields.downcast_ty(), gen_fields.downcast_ty())
 }
 
 pub(crate) fn predicate<'vir>(
     prefix: &str,
     fields: &[RustTyPredicatesEncOutputRef<'vir>],
     _task_key: <PredicateEnc as TaskEncoder>::TaskKey<'vir>,
     _snap: &SnapshotEncOutput<'vir>,
-    variant_field_snaps_to_snap: FunctionIdn<'vir, vir::ManySnap, vir::CSnap>,
+    variant_field_snaps_to_snap: FunctionIdn<'vir, (vir::ManySnap, vir::ManyTyVal), vir::CSnap>,
     _deps: &mut TaskEncoderDependencies<'vir, PredicateEnc>,
     generic_decls: &[vir::LocalDeclTyVal<'vir>],
     generic_exprs: &[vir::ExprTyVal<'vir>],
@@ -109,7 +112,7 @@ pub(crate) fn predicate<'vir>(
         })
         .collect::<Vec<_>>();
     let variant_snap_expr = vir::expr! {
-        unfolding ([pred_owned](ref_self, ..[generic_exprs])) in ([variant_field_snaps_to_snap](..[snap_args.as_slice()]))
+        unfolding ([pred_owned](ref_self, ..[generic_exprs])) in ([variant_field_snaps_to_snap]([[snap_args.as_slice()]], ..[generic_exprs]))
     };
     /*
     let pred_owned_expr = vir::expr! {

prusti-encoder/src/encoders/type/kinds/tuple.rs

@@ -2,7 +2,7 @@ use crate::encoders::{
     domain::{
         AdtBuilder, DomainDataStruct, DomainEnc, DomainEncOutputRef, DomainEncSpecifics, FieldTy, PureTypeBuilder, PureTypeCommon
     },
-    lifted::ty::{EncodeGenericsAsParamTy, LiftedTyEnc},
+    lifted::ty::{EncodeGenericsAsLifted, EncodeGenericsAsParamTy, LiftedTyEnc},
     predicate::{PredicateBuilder, PredicateEncData, PredicateEncDataStruct},
     rust_ty_predicates::RustTyPredicatesEnc,
     snapshot::SnapshotEncOutput,
@@ -28,7 +28,7 @@ pub(crate) fn domain<'vir>(
     let generics = params
         .iter()
         .map(|ty| {
-            deps.require_local::<LiftedTyEnc<EncodeGenericsAsParamTy>>(ty)
+            deps.require_local::<LiftedTyEnc<EncodeGenericsAsLifted>>(ty)
                 .unwrap()
                 .expect_generic()
         })
@@ -39,13 +39,14 @@ pub(crate) fn domain<'vir>(
         .map(|ty| FieldTy::from_ty(deps, ty))
         .collect::<Result<Vec<_>, _>>()?;
 
-    let (field_snaps_to_snap, field_access) = super::structlike::domain(
-        "", &fields, &mut builder, None,
+    let (field_snaps_to_snap, field_access, ty_param_accessors) = super::structlike::domain(
+        "", &fields, &generics, &mut builder, None,
     );
 
     Ok((DomainEncSpecifics::StructLike(DomainDataStruct {
         field_snaps_to_snap,
         field_access,
+        ty_param_accessors,
     }), Ok(builder)))
 
     /*