Allow refining Self

crates/flux-desugar/src/desugar.rs

@@ -235,8 +235,11 @@
         opaque_tys: Option<&'a mut UnordMap<LocalDefId, fhir::OpaqueTy>>,
     ) -> RustItemCtxt<'a, 'tcx> {
         let generics = genv.tcx.generics_of(owner);
+
+        let parent_id = genv.tcx.opt_parent(owner.def_id.to_def_id());
+
         let sort_resolver =
-            SortResolver::with_generics(genv.sess, genv.map().sort_decls(), generics);
+            SortResolver::with_generics(genv.sess, genv.map().sort_decls(), generics, parent_id);
         RustItemCtxt {
             genv,
             owner,
@@ -275,6 +278,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 },
@@ -285,13 +289,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

@@ -158,6 +158,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();
         }

crates/flux-desugar/src/sort_resolver.rs

@@ -9,7 +9,9 @@ use rustc_errors::ErrorGuaranteed;
 use rustc_hash::FxHashMap;
 use rustc_middle::ty::Generics;
 use rustc_span::{
+    def_id::DefId,
     sym::{self},
+    symbol::kw::SelfUpper,
     Symbol,
 };
 
@@ -20,8 +22,9 @@ type Result<T = ()> = std::result::Result<T, ErrorGuaranteed>;
 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>,
+    parent_id: Option<DefId>,
 }
 
 impl<'a> SortResolver<'a> {
@@ -30,21 +33,26 @@ impl<'a> SortResolver<'a> {
         sort_decls: &'a fhir::SortDecls,
         sort_params: &[Symbol],
     ) -> Self {
-        let sort_params = sort_params
+        let sort_params: std::collections::HashMap<
+            Symbol,
+            usize,
+            std::hash::BuildHasherDefault<rustc_hash::FxHasher>,
+        > = sort_params
             .iter()
             .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, parent_id: None }
     }
 
     pub(crate) fn with_generics(
         sess: &'a FluxSession,
         sort_decls: &'a fhir::SortDecls,
         generics: &'a Generics,
+        parent_id: Option<DefId>,
     ) -> 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, parent_id }
     }
 
     pub(crate) fn resolve_sort(&self, sort: &surface::Sort) -> Result<fhir::Sort> {
@@ -115,6 +123,10 @@ impl<'a> SortResolver<'a> {
             Ok(fhir::Sort::Bool)
         } else if ident.name == SORTS.real {
             Ok(fhir::Sort::Real)
+        } else if let Some(def_id) = self.parent_id
+            && ident.name == SelfUpper
+        {
+            Ok(fhir::Sort::SelfParam(def_id))
         } 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 @@
         .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 @@
     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(|kind| conv_generic_param_kind(&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 @@
                 .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,
@@ -302,7 +306,7 @@
         .cloned()
         .collect();
 
-    let res = rty::PolyFnSig::new(rty::FnSig::new(requires, args, output), vars);
+    let res = rty::PolyFnSig::new(rty::FnSig::new(requires, args, output, fn_sig.lifted), vars);
     Ok(rty::EarlyBinder(res))
 }
 
@@ -680,7 +684,7 @@
             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()
@@ -815,6 +819,7 @@
             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
@@ -835,7 +840,7 @@
                     .instantiate(&generics, &refine)
                     .replace_bound_expr(&idx.expr));
             }
-            fhir::Res::Def(..) | fhir::Res::SelfTyParam { .. } => {
+            fhir::Res::Def(..) => {
                 span_bug!(path.span, "unexpected resolution in conv_indexed_path: {:?}", path.res)
             }
         };
@@ -1203,6 +1208,10 @@
         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(def_id_to_self_ty(genv.tcx, def_id.expect_local()))
+            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:?}`")
@@ -1240,6 +1249,14 @@
     }
 }
 
+fn self_param_ty() -> rty::ParamTy {
+    rty::ParamTy { index: 0, name: kw::SelfUpper }
+}
+
+fn _def_id_to_self_ty(tcx: TyCtxt, def_id: LocalDefId) -> rty::ParamTy {
+    rty::ParamTy { index: def_id_to_param_index(tcx, def_id), 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);

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

@@ -193,7 +193,10 @@ fn type_of(genv: &GlobalEnv, def_id: LocalDefId) -> QueryResult<rty::EarlyBinder
             conv::expand_type_alias(genv, alias, &wfckresults)?
         }
         DefKind::TyParam => {
-            match &genv.get_generic_param(def_id).kind {
+            let generic_param = genv
+                .get_generic_param(def_id)
+                .unwrap_or_else(|| bug!("no generic param for {:?}", def_id));
+            match &generic_param.kind {
                 fhir::GenericParamKind::Type { default: Some(ty) } => {
                     let wfckresults = genv.check_wf(def_id)?;
                     conv::conv_ty(genv, ty, &wfckresults)?

crates/flux-middle/src/fhir.rs

@@ -54,6 +54,7 @@ use crate::{
 #[derive(Debug)]
 pub struct Generics {
     pub params: Vec<GenericParam>,
+    pub self_kind: Option<GenericParamKind>,
 }
 
 #[derive(Debug)]
@@ -556,6 +557,8 @@ pub enum Sort {
     Record(DefId, List<Sort>),
     /// The sort associated to a type variable
     Param(DefId),
+    /// The sort associated to the `Self` type variable for a trait [DefId]
+    SelfParam(DefId),
     /// A sort that needs to be inferred
     Wildcard,
     /// Sort inference variable generated for a [Sort::Wildcard] during sort checking
@@ -836,8 +839,8 @@ pub struct Defn {
 }
 
 impl Generics {
-    pub(crate) fn get_param(&self, def_id: LocalDefId) -> &GenericParam {
-        self.params.iter().find(|p| p.def_id == def_id).unwrap()
+    pub(crate) fn get_param(&self, def_id: LocalDefId) -> Option<&GenericParam> {
+        self.params.iter().find(|p| p.def_id == def_id)
     }
 
     pub fn with_refined_by(self, refined_by: &RefinedBy) -> Self {
@@ -850,7 +853,7 @@ impl Generics {
             };
             params.push(GenericParam { def_id: param.def_id, kind });
         }
-        Generics { params }
+        Generics { params, self_kind: self.self_kind }
     }
 }
 
@@ -974,6 +977,7 @@ impl Sort {
             | Sort::Unit
             | Sort::BitVec(_)
             | Sort::Param(_)
+            | Sort::SelfParam(_)
             | Sort::Wildcard
             | Sort::Record(_, _)
             | Sort::Infer(_)
@@ -1745,6 +1749,9 @@ impl fmt::Debug for Sort {
                 }
             }
             Sort::Param(def_id) => write!(f, "sortof({})", pretty::def_id_to_string(*def_id)),
+            Sort::SelfParam(def_id) => {
+                write!(f, "sortof({}::Self)", pretty::def_id_to_string(*def_id))
+            }
             Sort::Wildcard => write!(f, "_"),
             Sort::Infer(vid) => write!(f, "{vid:?}"),
             Sort::App(ctor, args) => write!(f, "{ctor}<{}>", args.iter().join(", ")),

crates/flux-middle/src/fhir/lift.rs

@@ -202,7 +202,7 @@ impl<'a, 'tcx> LiftCtxt<'a, 'tcx> {
             .iter()
             .map(|param| self.lift_generic_param(param))
             .try_collect_exhaust()?;
-        Ok(fhir::Generics { params })
+        Ok(fhir::Generics { params, self_kind: None })
     }
 
     pub fn lift_generic_predicates(

crates/flux-middle/src/fhir/visit.rs

@@ -248,6 +248,7 @@ pub fn walk_sort<V: Visitor>(vis: &mut V, sort: &Sort) {
         Sort::BitVec(_)
         | Sort::Int
         | Sort::Param(_)
+        | Sort::SelfParam(_)
         | Sort::Var(_)
         | Sort::Bool
         | Sort::Real