Allow refining Self

crates/flux-desugar/src/desugar.rs

@@ -29,7 +29,7 @@
 use crate::{
     errors,
     resolver::ResolverOutput,
-    sort_resolver::{SortResolver, SORTS},
+    sort_resolver::{SelfRes, SortResolver, SORTS},
 };
 
 pub fn desugar_qualifier(
@@ -198,6 +198,25 @@
     NumConst(i128),
 }
 
+fn self_res(genv: &GlobalEnv, owner: OwnerId) -> SelfRes {
+    let def_id = owner.def_id;
+    let parent_id = genv.tcx.opt_parent(def_id.to_def_id());
+    if let Some(alias_to) = parent_id {
+        match genv.tcx.def_kind(alias_to) {
+            DefKind::Trait => return SelfRes::Param(alias_to),
+            DefKind::Impl { .. } => {
+                if let Some(sort) = genv.sort_of_self_ty_alias(alias_to) {
+                    return SelfRes::Alias(sort);
+                } else {
+                    return SelfRes::None;
+                }
+            }
+            _ => return SelfRes::None,
+        }
+    }
+    return SelfRes::None;
+}
+
 impl<'a, 'tcx> RustItemCtxt<'a, 'tcx> {
     pub(crate) fn new(
         genv: &'a GlobalEnv<'a, 'tcx>,
@@ -206,8 +225,9 @@
         opaque_tys: Option<&'a mut UnordMap<LocalDefId, fhir::OpaqueTy>>,
     ) -> RustItemCtxt<'a, 'tcx> {
         let generics = genv.tcx.generics_of(owner);
+        let self_res = self_res(genv, owner);
         let sort_resolver =
-            SortResolver::with_generics(genv.sess, genv.map().sort_decls(), generics);
+            SortResolver::with_generics(genv.sess, genv.map().sort_decls(), generics, self_res);
         RustItemCtxt {
             genv,
             owner,
@@ -232,9 +252,38 @@
         )
     }
 
-    pub(crate) fn desugar_generics(&self, generics: &surface::Generics) -> Result<fhir::Generics> {
-        let hir_generics = self.genv.hir().get_generics(self.owner.def_id).unwrap();
+    /// [desugar_generics] starts with the `lifted_generics` and "updates" it with the surface `generics`
+    pub(crate) fn desugar_generics(
+        &self,
+        lifted_generics: fhir::Generics,
+        generics: &surface::Generics,
+    ) -> Result<fhir::Generics> {
+        // Step 1: desugar the surface generics by themselves
+        let generics = self.desugar_surface_generics(generics)?;
+
+        // Step 2: map each (surface) generic to its specified kind
+        let generic_kinds: FxHashMap<_, _> = generics
+            .params
+            .into_iter()
+            .map(|param| (param.def_id, param.kind))
+            .collect();
 
+        // Step 3: traverse lifted_generics, using the surface-kind, if specified, and lifted kind otherwise
+        let params = lifted_generics
+            .params
+            .iter()
+            .map(|lifted_param| {
+                let def_id = lifted_param.def_id;
+                let kind = generic_kinds.get(&def_id).unwrap_or(&lifted_param.kind);
+                fhir::GenericParam { def_id, kind: kind.clone() }
+            })
+            .collect();
+
+        Ok(fhir::Generics { params, self_kind: generics.self_kind.clone() })
+    }
+
+    fn desugar_surface_generics(&self, generics: &surface::Generics) -> Result<fhir::Generics> {
+        let hir_generics = self.genv.hir().get_generics(self.owner.def_id).unwrap();
         let generics_map: FxHashMap<_, _> = hir_generics
             .params
             .iter()
@@ -248,6 +297,7 @@
             .collect();
 
         let mut params = vec![];
+        let mut self_kind = None;
         for param in &generics.params {
             let kind = match &param.kind {
                 surface::GenericParamKind::Type => fhir::GenericParamKind::Type { default: None },
@@ -258,13 +308,16 @@
                 }
             };
 
-            let def_id = *generics_map
-                .get(&param.name)
-                .ok_or_else(|| self.emit_err(errors::UnresolvedGenericParam::new(param.name)))?;
-
-            params.push(fhir::GenericParam { def_id, kind });
+            if param.name.name == kw::SelfUpper {
+                self_kind = Some(kind);
+            } else {
+                let def_id = *generics_map.get(&param.name).ok_or_else(|| {
+                    self.emit_err(errors::UnresolvedGenericParam::new(param.name))
+                })?;
+                params.push(fhir::GenericParam { def_id, kind });
+            }
         }
-        Ok(fhir::Generics { params })
+        Ok(fhir::Generics { params, self_kind })
     }
 
     fn desugar_predicates(

crates/flux-desugar/src/lib.rs

@@ -138,11 +138,13 @@ pub fn desugar_fn_sig(
         let mut cx = RustItemCtxt::new(genv, owner_id, resolver_output, Some(&mut opaque_tys));
 
         // Desugar and insert generics
+        let lifted_generics = cx.as_lift_cx().lift_generics()?;
         let generics = if let Some(generics) = &fn_sig.generics {
-            cx.desugar_generics(generics)?
+            cx.desugar_generics(lifted_generics, generics)?
         } else {
-            cx.as_lift_cx().lift_generics()?
+            lifted_generics
         };
+
         genv.map().insert_generics(def_id, generics);
 
         // Desugar of fn_sig needs to happen AFTER inserting generics. See crate level comment
@@ -158,6 +160,7 @@ pub fn desugar_fn_sig(
         genv.map_mut().insert_opaque_tys(opaque_tys);
     } else {
         let (generics, fn_info) = lift::lift_fn(genv.tcx, genv.sess, owner_id)?;
+
         if config::dump_fhir() {
             dbg::dump_item_info(genv.tcx, def_id, "fhir", &fn_info.fn_sig).unwrap();
         }
@@ -187,7 +190,7 @@ pub fn desugar_generics_and_predicates(
 
     let generics = if let Some(generics) = generics {
         let cx = RustItemCtxt::new(genv, owner_id, resolver_output, None);
-        cx.desugar_generics(generics)?
+        cx.desugar_generics(lifted_generics, generics)?
     } else {
         lifted_generics
     };

crates/flux-desugar/src/sort_resolver.rs

@@ -9,19 +9,32 @@ use rustc_errors::ErrorGuaranteed;
 use rustc_hash::FxHashMap;
 use rustc_middle::ty::Generics;
 use rustc_span::{
+    def_id::DefId,
     sym::{self},
+    symbol::kw::{self},
     Symbol,
 };
 
 use crate::errors;
 
 type Result<T = ()> = std::result::Result<T, ErrorGuaranteed>;
 
+#[derive(Debug)]
+pub enum SelfRes {
+    /// A `Self` parameter in a trait definition.
+    Param(DefId),
+    /// An alias to another sort, e.g., when used inside an impl block
+    Alias(fhir::Sort),
+    /// It is not valid to use `Self`, e.g., when used in a free function
+    None,
+}
+
 pub(crate) struct SortResolver<'a> {
     sess: &'a FluxSession,
     sort_decls: &'a fhir::SortDecls,
-    generic_params: FxHashMap<Symbol, rustc_span::def_id::DefId>,
+    generic_params: FxHashMap<Symbol, DefId>,
     sort_params: FxHashMap<Symbol, usize>,
+    self_res: SelfRes,
 }
 
 impl<'a> SortResolver<'a> {
@@ -35,16 +48,23 @@ impl<'a> SortResolver<'a> {
             .enumerate()
             .map(|(i, v)| (*v, i))
             .collect();
-        Self { sess, sort_decls, generic_params: Default::default(), sort_params }
+        Self {
+            sess,
+            sort_decls,
+            generic_params: Default::default(),
+            sort_params,
+            self_res: SelfRes::None,
+        }
     }
 
     pub(crate) fn with_generics(
         sess: &'a FluxSession,
         sort_decls: &'a fhir::SortDecls,
         generics: &'a Generics,
+        self_res: SelfRes,
     ) -> Self {
         let generic_params = generics.params.iter().map(|p| (p.name, p.def_id)).collect();
-        Self { sess, sort_decls, sort_params: Default::default(), generic_params }
+        Self { sess, sort_decls, sort_params: Default::default(), generic_params, self_res }
     }
 
     pub(crate) fn resolve_sort(&self, sort: &surface::Sort) -> Result<fhir::Sort> {
@@ -115,6 +135,12 @@ impl<'a> SortResolver<'a> {
             Ok(fhir::Sort::Bool)
         } else if ident.name == SORTS.real {
             Ok(fhir::Sort::Real)
+        } else if ident.name == kw::SelfUpper {
+            match &self.self_res {
+                SelfRes::Param(def_id) => Ok(fhir::Sort::SelfParam(*def_id)),
+                SelfRes::Alias(sort) => Ok(sort.clone()),
+                SelfRes::None => Err(self.sess.emit_err(errors::UnresolvedSort::new(*ident))),
+            }
         } else if let Some(def_id) = self.generic_params.get(&ident.name) {
             Ok(fhir::Sort::Param(*def_id))
         } else if let Some(idx) = self.sort_params.get(&ident.name) {

crates/flux-driver/src/callbacks.rs

@@ -292,7 +292,8 @@ fn desugar_item(
             desugar::desugar_struct_def(genv, owner_id, struct_def, resolver_output)?;
         }
         hir::ItemKind::Trait(.., items) => {
-            desugar::desugar_generics_and_predicates(genv, owner_id, resolver_output, None)?;
+            let generics = specs.generics.get(&owner_id);
+            desugar::desugar_generics_and_predicates(genv, owner_id, resolver_output, generics)?;
             items.iter().try_for_each_exhaust(|trait_item| {
                 desugar_assoc_item(
                     genv,
@@ -304,7 +305,7 @@ fn desugar_item(
             })?;
         }
         hir::ItemKind::Impl(impl_) => {
-            let generics = specs.impls.get(&owner_id);
+            let generics = specs.generics.get(&owner_id);
             desugar::desugar_generics_and_predicates(genv, owner_id, resolver_output, generics)?;
             impl_.items.iter().try_for_each_exhaust(|impl_item| {
                 desugar_assoc_item(

crates/flux-driver/src/collector.rs

@@ -41,7 +41,7 @@ pub type Ignores = UnordSet<IgnoreKey>;
 pub(crate) struct Specs {
     pub fn_sigs: UnordMap<OwnerId, FnSpec>,
     pub structs: FxHashMap<OwnerId, surface::StructDef>,
-    pub impls: FxHashMap<OwnerId, surface::Generics>,
+    pub generics: FxHashMap<OwnerId, surface::Generics>,
     pub enums: FxHashMap<OwnerId, surface::EnumDef>,
     pub qualifs: Vec<surface::Qualifier>,
     pub func_defs: Vec<surface::FuncDef>,
@@ -102,7 +102,8 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
                 ItemKind::Mod(..) => collector.parse_mod_spec(owner_id.def_id, attrs),
                 ItemKind::TyAlias(..) => collector.parse_tyalias_spec(owner_id, attrs),
                 ItemKind::Const(..) => collector.parse_const_spec(item, attrs),
-                ItemKind::Impl(_) => collector.parse_impl_spec(owner_id, attrs),
+                ItemKind::Impl(_) => collector.parse_generics_spec(owner_id, attrs),
+                ItemKind::Trait(..) => collector.parse_generics_spec(owner_id, attrs),
                 _ => Ok(()),
             };
         }
@@ -187,7 +188,7 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
         }
     }
 
-    fn parse_impl_spec(
+    fn parse_generics_spec(
         &mut self,
         owner_id: OwnerId,
         attrs: &[Attribute],
@@ -196,7 +197,7 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
         self.report_dups(&attrs)?;
 
         if let Some(generics) = attrs.generics() {
-            self.specs.impls.insert(owner_id, generics);
+            self.specs.generics.insert(owner_id, generics);
         }
 
         Ok(())
@@ -571,7 +572,7 @@ impl Specs {
     fn new() -> Specs {
         Specs {
             fn_sigs: Default::default(),
-            impls: Default::default(),
+            generics: Default::default(),
             structs: Default::default(),
             enums: Default::default(),
             qualifs: Vec::default(),

crates/flux-fhir-analysis/src/conv.rs

@@ -136,6 +136,17 @@ pub(crate) fn conv_opaque_ty(
         .collect())
 }
 
+fn conv_generic_param_kind(kind: &fhir::GenericParamKind) -> rty::GenericParamDefKind {
+    match kind {
+        fhir::GenericParamKind::Type { default } => {
+            rty::GenericParamDefKind::Type { has_default: default.is_some() }
+        }
+        fhir::GenericParamKind::SplTy => rty::GenericParamDefKind::SplTy,
+        fhir::GenericParamKind::BaseTy => rty::GenericParamDefKind::BaseTy,
+        fhir::GenericParamKind::Lifetime => rty::GenericParamDefKind::Lifetime,
+    }
+}
+
 pub(crate) fn conv_generics(
     genv: &GlobalEnv,
     rust_generics: &rustc::ty::Generics,
@@ -144,12 +155,12 @@ pub(crate) fn conv_generics(
     is_trait: Option<LocalDefId>,
 ) -> QueryResult<rty::Generics> {
     let opt_self = is_trait.map(|def_id| {
-        rty::GenericParamDef {
-            index: 0,
-            name: kw::SelfUpper,
-            def_id: def_id.to_def_id(),
-            kind: rty::GenericParamDefKind::Type { has_default: false },
-        }
+        let kind = generics
+            .self_kind
+            .as_ref()
+            .map(conv_generic_param_kind)
+            .unwrap_or(rty::GenericParamDefKind::Type { has_default: false });
+        rty::GenericParamDef { index: 0, name: kw::SelfUpper, def_id: def_id.to_def_id(), kind }
     });
     let params = opt_self
         .into_iter()
@@ -159,14 +170,7 @@ pub(crate) fn conv_generics(
                 .params
                 .iter()
                 .find(|param| param.def_id.to_def_id() == rust_param.def_id)?;
-            let kind = match &param.kind {
-                fhir::GenericParamKind::Type { default } => {
-                    rty::GenericParamDefKind::Type { has_default: default.is_some() }
-                }
-                fhir::GenericParamKind::SplTy => rty::GenericParamDefKind::SplTy,
-                fhir::GenericParamKind::BaseTy => rty::GenericParamDefKind::BaseTy,
-                fhir::GenericParamKind::Lifetime => rty::GenericParamDefKind::Lifetime,
-            };
+            let kind = conv_generic_param_kind(&param.kind);
             let def_id = param.def_id.to_def_id();
             Some(rty::GenericParamDef {
                 kind,
@@ -680,7 +684,7 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
             if let fhir::Res::SelfTyParam { .. } = path.res
                 && sort.is_none()
             {
-                return Ok(rty::Ty::param(rty::ParamTy { index: 0, name: kw::SelfUpper }));
+                return Ok(rty::Ty::param(self_param_ty()));
             }
             if let fhir::Res::Def(DefKind::TyParam, def_id) = path.res
                 && sort.is_none()
@@ -805,6 +809,7 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
             fhir::Res::Def(DefKind::TyParam, def_id) => {
                 rty::BaseTy::Param(def_id_to_param_ty(self.genv.tcx, def_id.expect_local()))
             }
+            fhir::Res::SelfTyParam { .. } => rty::BaseTy::Param(self_param_ty()),
             fhir::Res::SelfTyAlias { alias_to, .. } => {
                 return Ok(self
                     .genv
@@ -825,7 +830,7 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
                     .instantiate(&generics, &refine)
                     .replace_bound_expr(&idx));
             }
-            fhir::Res::Def(..) | fhir::Res::SelfTyParam { .. } => {
+            fhir::Res::Def(..) => {
                 span_bug!(path.span, "unexpected resolution in conv_indexed_path: {:?}", path.res)
             }
         };
@@ -1193,6 +1198,7 @@ fn conv_sort(genv: &GlobalEnv, sort: &fhir::Sort) -> rty::Sort {
         fhir::Sort::Param(def_id) => {
             rty::Sort::Param(def_id_to_param_ty(genv.tcx, def_id.expect_local()))
         }
+        fhir::Sort::SelfParam(_def_id) => rty::Sort::Param(self_param_ty()),
         fhir::Sort::Var(n) => rty::Sort::Var(rty::SortVar::from(*n)),
         fhir::Sort::Error | fhir::Sort::Wildcard | fhir::Sort::Infer(_) => {
             bug!("unexpected sort `{sort:?}`")
@@ -1230,6 +1236,10 @@ fn def_id_to_param_ty(tcx: TyCtxt, def_id: LocalDefId) -> rty::ParamTy {
     }
 }
 
+fn self_param_ty() -> rty::ParamTy {
+    rty::ParamTy { index: 0, name: kw::SelfUpper }
+}
+
 fn def_id_to_param_index(tcx: TyCtxt, def_id: LocalDefId) -> u32 {
     let item_def_id = tcx.hir().ty_param_owner(def_id);
     let generics = tcx.generics_of(item_def_id);