Adding support for type aliases to associated types

prusti-encoder/src/encoders/mir_fn/mod.rs

@@ -6,7 +6,7 @@ pub use function::*;
 pub use method::*;
 pub use signature::*;
 
-use crate::encoders::ty::generics::{GArgs, GParams};
+use crate::encoders::ty::generics::{GArgs, GParams, trait_impls::TraitImplEnc};
 
 use prusti_interface::specs::specifications::SpecQuery;
 use prusti_rustc_interface::{hir, middle::ty, span::def_id::DefId};
@@ -61,4 +61,19 @@ pub fn encode_all_in_crate<'tcx>(tcx: ty::TyCtxt<'tcx>) {
             }
         }
     }
+
+    for def_id in tcx.hir_crate_items(()).definitions() {
+        if let hir::def::DefKind::Impl { of_trait: true } = tcx.def_kind(def_id) {
+            TraitImplEnc::encode(def_id.to_def_id(), false).unwrap();
+        }
+    }
+
+    // for trait_id in tcx.visible_traits() {
+    //         for impl_id in tcx.all_impls(trait_id) {
+    //             {
+    //                 TraitImplEnc::encode(impl_id, false).unwrap();
+    //             }
+
+    //     }
+    // }
 }

prusti-encoder/src/encoders/ty/data.rs

@@ -6,6 +6,8 @@ use std::{
 
 use prusti_rustc_interface::abi;
 
+use prusti_rustc_interface::span::def_id::DefId;
+
 pub trait TyDatas<'vir>: Debug + Clone + Copy {
     type TyData: Debug + Clone + 'vir = ();
 
@@ -46,12 +48,14 @@ pub struct StructData<'vir, D: TyDatas<'vir>> {
     pub data: D::StructData,
     pub inhabited: bool,
     pub fields: Vec<D::FieldData>,
+    pub did: Option<DefId>,
 }
 
 pub struct EnumData<'vir, D: TyDatas<'vir>> {
     pub data: D::EnumData,
     pub inhabited: bool,
     pub variants: Vec<VariantData<'vir, D>>,
+    pub did: Option<DefId>,
 }
 
 pub struct VariantData<'vir, D: TyDatas<'vir>> {
@@ -99,16 +103,22 @@ impl<'vir, D: TyDatas<'vir>> TySpecifics<'vir, D> {
         Self::MutRef(data)
     }
 
-    pub fn mk_structlike(data: D::StructData, inhabited: bool, fields: Vec<D::FieldData>) -> Self {
-        Self::StructLike(StructData::new(data, inhabited, fields))
+    pub fn mk_structlike(
+        data: D::StructData,
+        inhabited: bool,
+        fields: Vec<D::FieldData>,
+        did: Option<DefId>,
+    ) -> Self {
+        Self::StructLike(StructData::new(data, inhabited, fields, did))
     }
 
     pub fn mk_enumlike(
         data: D::EnumData,
         inhabited: bool,
         variants: Vec<VariantData<'vir, D>>,
+        did: Option<DefId>,
     ) -> Self {
-        Self::EnumLike(EnumData::new(data, inhabited, variants))
+        Self::EnumLike(EnumData::new(data, inhabited, variants, did))
     }
 
     pub fn is_param(&self) -> bool {
@@ -253,6 +263,7 @@ impl<'vir, D: TyDatas<'vir>> StructData<'vir, D> {
             data: (&self.data, &other.data),
             inhabited: self.inhabited,
             fields: fields.collect(),
+            did: self.did,
         }
     }
 }
@@ -269,6 +280,7 @@ impl<'vir, D: TyDatas<'vir>> EnumData<'vir, D> {
             data: (&self.data, &other.data),
             inhabited: self.inhabited,
             variants: variants.map(|(v1, v2)| v1.zip(v2)).collect(),
+            did: self.did,
         }
     }
 }
@@ -291,10 +303,10 @@ impl<'vir, D1: TyDatas<'vir>, D2: TyDatas<'vir>> TyDatas<'vir> for (D1, D2) {
 // Deref implementations
 
 macro_rules! impls {
-    ($container:ident$( { $field:ident: $ty:ty })?) => {
+    ($container:ident$( { $field:ident: $ty:ty $(, $did:ident)?})?) => {
 impl<'vir, D: TyDatas<'vir>> $container<'vir, D> {
-    pub fn new(data: D::$container, inhabited: bool $(, $field: $ty)?) -> Self {
-        Self { data, inhabited, $($field,)? }
+    pub fn new(data: D::$container, inhabited: bool $(, $field: $ty $(, $did: Option<DefId>)?)?) -> Self {
+        Self { data, inhabited, $($field, $($did,)?)? }
     }
 }
 
@@ -306,7 +318,7 @@ impl<'vir, D: TyDatas<'vir>> Debug for $container<'vir, D> {
 
 impl<'vir, D: TyDatas<'vir>> Clone for $container<'vir, D> {
     fn clone(&self) -> Self {
-        Self { data: self.data.clone(), inhabited: self.inhabited, $($field: self.$field.clone())? }
+        Self { data: self.data.clone(), inhabited: self.inhabited, $($field: self.$field.clone() $(, $did: self.$did.clone())?)? }
     }
 }
 
@@ -368,8 +380,8 @@ impl<'vir, D: TyDatas<'vir>> $container<'vir, D> {
 
 impls!(TyData { specifics: TySpecifics<'vir, D> });
 impl_zip!(TyData.specifics);
-impls!(StructData { fields: Vec<D::FieldData> });
-impls!(EnumData { variants: Vec<VariantData<'vir, D>> });
+impls!(StructData { fields: Vec<D::FieldData>, did });
+impls!(EnumData { variants: Vec<VariantData<'vir, D>>, did });
 impls!(VariantData { inner: StructData<'vir, D> });
 impl_zip!(VariantData.inner);
 

prusti-encoder/src/encoders/ty/generics/args.rs

@@ -10,6 +10,11 @@ pub struct GArgs<'tcx> {
     pub(super) args: &'tcx [ty::GenericArg<'tcx>],
 }
 
+pub enum GParamVariant<'tcx> {
+    Param(ty::ParamTy),
+    Alias(ty::AliasTy<'tcx>),
+}
+
 impl<'tcx> GArgs<'tcx> {
     pub fn new(context: impl Into<GParams<'tcx>>, args: &'tcx [ty::GenericArg<'tcx>]) -> Self {
         GArgs {
@@ -34,12 +39,11 @@ impl<'tcx> GArgs<'tcx> {
         self.context.normalize(ty)
     }
 
-    pub fn expect_param(self) -> ty::ParamTy {
+    pub fn expect_param(self) -> GParamVariant<'tcx> {
         assert_eq!(self.args.len(), 1);
         match self.args[0].expect_ty().kind() {
-            ty::TyKind::Param(p) => *p,
-            // TODO: this needs to be changed to support type aliases
-            ty::TyKind::Alias(..) => panic!("type aliases are not currently supported"),
+            ty::TyKind::Param(p) => GParamVariant::Param(*p),
+            ty::TyKind::Alias(_k, t) => GParamVariant::Alias(*t),
             other => panic!("expected type parameter, {other:?}"),
         }
     }

prusti-encoder/src/encoders/ty/generics/mod.rs

@@ -3,6 +3,8 @@ mod params;
 mod casters;
 mod args_ty;
 mod args;
+pub mod traits;
+pub mod trait_impls;
 
 pub use args::*;
 pub use args_ty::*;

prusti-encoder/src/encoders/ty/generics/params.rs

@@ -1,14 +1,19 @@
 use prusti_interface::specs::typed::ExternSpecKind;
 use prusti_rustc_interface::{
-    middle::ty,
+    middle::{ty, ty::TyKind},
     span::{def_id::DefId, symbol},
 };
 use task_encoder::{EncodeFullResult, TaskEncoder, TaskEncoderDependencies};
 use vir::{CastType, HasType};
 
 use crate::encoders::{
     TyUsePureEnc,
-    ty::{RustTyDecomposition, data::TySpecifics, generics::GArgsTyEnc, lifted::TyConstructorEnc},
+    ty::{
+        RustTyDecomposition,
+        data::TySpecifics,
+        generics::{GArgsTyEnc, GParamVariant, traits::TraitEnc},
+        lifted::TyConstructorEnc,
+    },
 };
 
 /// The list of defined parameters in a given context. E.g. the type parameters
@@ -176,7 +181,6 @@ impl<'vir> From<DefId> for GParams<'vir> {
 /// `fn foo<T, U>(x: U)` into the Viper `method foo(x: Ref, T: Type, U: Type)`
 /// (handles the type parameters).
 pub struct GenericParamsEnc;
-
 #[derive(Debug, Clone)]
 pub struct GenericParams<'vir> {
     ty_args: &'vir [vir::TypeTyVal<'vir>],
@@ -250,7 +254,26 @@ impl<'vir> GenericParams<'vir> {
     ) -> vir::ExprTyVal<'vir> {
         if let TySpecifics::Param(()) = &ty.ty.specifics {
             let param = ty.args.expect_param();
-            return self.ty_exprs[self.map_idx(param.index).unwrap()];
+            return match param {
+                GParamVariant::Param(p) => self.ty_exprs[self.map_idx(p.index).unwrap()],
+                GParamVariant::Alias(a) => vir::with_vcx(|vcx| {
+                    let tcx = vcx.tcx();
+                    let trait_did = tcx.associated_item(a.def_id).container_id(tcx);
+                    let trait_data = deps.require_dep::<TraitEnc>(trait_did).unwrap();
+                    let tys = &a
+                        .args
+                        .iter()
+                        .map(|arg| match arg.expect_ty().kind() {
+                            TyKind::Param(p) => self.ty_exprs[self.map_idx(p.index).unwrap()],
+                            _ => self.ty_expr(
+                                deps,
+                                RustTyDecomposition::from_ty(arg.expect_ty(), tcx, ty.args.context),
+                            ),
+                        })
+                        .collect::<Vec<_>>();
+                    (trait_data.type_did_fun_mapping.get(&a.def_id).unwrap())(tys)
+                }),
+            };
         }
         let ty_constructor = deps
             .require_ref::<TyConstructorEnc>(ty.ty)

prusti-encoder/src/encoders/ty/generics/trait_impls.rs

@@ -0,0 +1,143 @@
+use std::iter;
+
+use prusti_rustc_interface::{middle::ty::AssocKind, span::def_id::DefId};
+use task_encoder::{EncodeFullResult, TaskEncoder, TaskEncoderDependencies};
+use vir::{CastType, Domain, vir_format_identifier};
+
+use crate::encoders::ty::{
+    RustTyDecomposition,
+    generics::{GArgs, GArgsTyEnc, GParams, GenericParamsEnc, traits::TraitEnc},
+};
+
+pub struct TraitImplEnc;
+
+impl TaskEncoder for TraitImplEnc {
+    task_encoder::encoder_cache!(TraitImplEnc);
+
+    fn task_to_key<'vir>(task: &Self::TaskDescription<'vir>) -> Self::TaskKey<'vir> {
+        *task
+    }
+
+    fn emit_outputs<'vir>(program: &mut task_encoder::Program<'vir>) {
+        for dom in TraitImplEnc::all_outputs_local_no_errors() {
+            program.add_domain(dom);
+        }
+    }
+
+    type TaskDescription<'vir> = DefId;
+    type OutputFullLocal<'vir> = Domain<'vir>;
+
+    fn do_encode_full<'vir>(
+        task_key: &Self::TaskKey<'vir>,
+        deps: &mut TaskEncoderDependencies<'vir, Self>,
+    ) -> EncodeFullResult<'vir, Self> {
+        deps.emit_output_ref(*task_key, ())?;
+
+        vir::with_vcx(|vcx| {
+            let tcx = vcx.tcx();
+
+            let ctx = GParams::from(*task_key);
+
+            let params = deps.require_dep::<GenericParamsEnc>(ctx)?;
+
+            let trait_ref = tcx.impl_trait_ref(task_key).unwrap().instantiate_identity();
+            let trait_did = trait_ref.def_id;
+            let trait_data = deps.require_dep::<TraitEnc>(trait_did)?;
+
+            // for some reason, just using all args (which includes the struct at index 0)
+            // leads to a cycle for simple test cases -> split up and add struct manually
+            let args = deps.require_dep::<GArgsTyEnc>(GArgs::new(ctx, &trait_ref.args[1..]))?;
+
+            let struct_ty = tcx.type_of(task_key).instantiate_identity();
+            let struct_ty_expr =
+                params.ty_expr(deps, RustTyDecomposition::from_ty(struct_ty, tcx, ctx));
+
+            let mut axs = Vec::new();
+
+            let struct_ty = tcx.type_of(task_key).instantiate_identity();
+
+            let impl_fun = trait_data.impl_fun;
+            let trait_ty_decls = params
+                .ty_decls()
+                .iter()
+                .map(|dec| dec.upcast_ty())
+                .collect::<Vec<_>>();
+            let trait_tys = vcx.alloc_slice(
+                &iter::once(struct_ty_expr)
+                    .chain(args.get_ty().to_owned())
+                    .collect::<Vec<_>>(),
+            );
+            axs.push(
+                vcx.mk_domain_axiom(
+                    vir_format_identifier!(vcx, "{}_impl_{}", trait_data.trait_name, struct_ty),
+                    vir::expr! {forall ..[trait_ty_decls] :: {[impl_fun(trait_tys)]} [impl_fun(trait_tys)]}
+                )
+            );
+
+            tcx.associated_items(*task_key)
+                .in_definition_order()
+                .filter(|item| matches!(item.kind, AssocKind::Type { data: _ }))
+                .for_each(|impl_item| {
+                    let assoc_fun = trait_data.type_did_fun_mapping.get(&impl_item.trait_item_def_id.unwrap()).unwrap();
+                    // construct arguments for assoc_item function
+                    // parameters of the trait are substituted 
+                    // by the arguments used in the impl
+                    // parameters of the associated type are kept
+
+                    // parameters of assoc item include already substituted arguments
+                    let assoc_params = deps
+                        .require_dep::<GenericParamsEnc>(GParams::from(impl_item.def_id))
+                        .unwrap();
+
+                    // the type we want to resolve the type alias to
+                    let assoc_type_expr = assoc_params.ty_expr(
+                        deps,
+                        RustTyDecomposition::from_ty(
+                            tcx.type_of(impl_item.def_id).instantiate_identity(),
+                            tcx,
+                            GParams::from(impl_item.def_id),
+                        ),
+                    );
+                    let assoc_decls = assoc_params
+                        .ty_decls()
+                        .iter()
+                        .map(|dec| dec.upcast_ty())
+                        .collect::<Vec<_>>();
+
+                    // Combine substituted trait params with the params of the associated type
+                    let mut trait_ty_decls = trait_ty_decls.clone();
+                    trait_ty_decls.extend_from_slice(&assoc_decls[params.ty_exprs().len()..]);
+
+                    // Combine substituted trait params with the params of the associated type
+                    let trait_tys = vcx.alloc_slice(&iter::once(struct_ty_expr).chain(args.get_ty().to_owned()).chain(assoc_params.ty_exprs()[params.ty_exprs().len()..].to_owned()).collect::<Vec<_>>());
+
+                    axs.push(vcx.mk_domain_axiom(
+                        vir_format_identifier!(
+                            vcx,
+                            "{}_Assoc_{}_{}",
+                            trait_data.trait_name,
+                            tcx.item_name(impl_item.def_id),
+                            struct_ty
+                        ),
+                        vir::expr! {forall ..[trait_ty_decls] :: {[assoc_fun(trait_tys)]} ([assoc_fun(trait_tys)]) == (assoc_type_expr)}
+                    ));
+                });
+
+            Ok((
+                vcx.mk_domain(
+                    vir_format_identifier!(
+                        vcx,
+                        "t_{}_{}",
+                        trait_data.trait_name,
+                        tcx.type_of(*task_key).instantiate_identity().to_string()
+                    ),
+                    &[],
+                    vcx.alloc_slice(&axs),
+                    &[],
+                    None,
+                ),
+                (),
+            ))
+        })
+    }
+}