Towards polymorphic sorts

README.md

@@ -16,3 +16,6 @@ For an overview, take a look at the [`flux` website](https://flux-rs.github.io).
 
 Documentation, including installation and usage guides can be found on the
 [website](https://flux-rs.github.io/flux).
+
+
+TRACE: insert_fn_sig DefId(0:14 ~ maps00[5caf]::rmap::{impl#0}::new) fn() -> {a0:RMap<sortof(K), sortof(V)>. RMap[K, V][a0] | true}

crates/flux-desugar/src/lib.rs

@@ -1,6 +1,6 @@
 //! Desugaring from types in [`flux_syntax::surface`] to types in [`flux_middle::fhir`]
 //!
-//! # NOTE
+//! # Generics and Desugaring
 //!
 //! Desugaring requires knowing the sort of each type so we can correctly resolve binders declared with
 //! @ syntax or arg syntax. In particular, to know the sort of a type parameter we need to know its
@@ -53,10 +53,9 @@ pub fn desugar_struct_def(
     let mut cx = DesugarCtxt::new(genv, owner_id, resolver_output, None);
 
     // Desugar and insert generics
-    let (generics, predicates) = cx.as_lift_cx().lift_generics_with_predicates()?;
-    genv.map().insert_generics(def_id, generics);
+    let predicates = cx.as_lift_cx().lift_predicates()?;
 
-    // Desugar of struct_def needs to happen AFTER inserting generics. See crate level comment
+    // Desugar of struct_def needs to happen AFTER inserting generics. See #generics-and-desugaring
     let struct_def = cx.desugar_struct_def(struct_def, &mut Binders::new())?;
     if config::dump_fhir() {
         dbg::dump_item_info(genv.tcx, owner_id, "fhir", &struct_def).unwrap();
@@ -173,15 +172,25 @@ pub fn desugar_fn_sig(
 }
 
 /// HACK(nilehmann) this is a bit of a hack. We use it to properly register generics and predicates
-/// for items that don't have surface syntax (impl blocks, traits, ...). In this cases we just [lift]
-/// them from hir.
+/// for items that don't have surface syntax (impl blocks, traits, ...), or for `impl` blocks with
+/// explicit `generics` annotations. In the former case, we use `desugar`; in the latter cases we
+/// just [lift] them from hir.
 pub fn desugar_generics_and_predicates(
     genv: &mut GlobalEnv,
     owner_id: OwnerId,
+    resolver_output: &ResolverOutput,
+    generics: Option<&surface::Generics>,
 ) -> Result<(), ErrorGuaranteed> {
-    let def_id = owner_id.def_id;
-    let (generics, predicates) =
+    let (lifted_generics, predicates) =
         LiftCtxt::new(genv.tcx, genv.sess, owner_id, None).lift_generics_with_predicates()?;
+
+    let generics = if let Some(generics) = generics {
+        let cx = DesugarCtxt::new(genv, owner_id, resolver_output, None);
+        cx.desugar_generics(generics)?
+    } else {
+        lifted_generics
+    };
+    let def_id = owner_id.def_id;
     genv.map().insert_generics(def_id, generics);
     genv.map_mut().insert_generic_predicates(def_id, predicates);
     Ok(())

crates/flux-driver/src/callbacks.rs

@@ -153,15 +153,19 @@ fn stage1_desugar(genv: &mut GlobalEnv, specs: &Specs) -> Result<(), ErrorGuaran
         .err()
         .or(err);
 
-    // Register RefinedBys
+    // Register RefinedBys (for structs and enums, which also registers their Generics)
     err = specs
         .refined_bys()
         .try_for_each_exhaust(|(owner_id, refined_by)| {
+            let generics = lift::lift_generics(tcx, sess, owner_id)?;
             let refined_by = if let Some(refined_by) = refined_by {
-                desugar::desugar_refined_by(sess, map.sort_decls(), owner_id, refined_by)?
+                let def_id = owner_id.to_def_id();
+                let generics = tcx.generics_of(def_id);
+                desugar::desugar_refined_by(sess, map.sort_decls(), owner_id, generics, refined_by)?
             } else {
                 lift::lift_refined_by(tcx, owner_id)
             };
+            map.insert_generics(owner_id.def_id, generics.with_refined_by(&refined_by));
             map.insert_refined_by(owner_id.def_id, refined_by);
             Ok(())
         })
@@ -273,7 +277,9 @@ fn desugar_item(
             let ty_alias = specs.ty_aliases[&owner_id].as_ref();
             desugar::desugar_type_alias(genv, owner_id, ty_alias, resolver_output)?;
         }
-        hir::ItemKind::OpaqueTy(_) => desugar::desugar_generics_and_predicates(genv, owner_id)?,
+        hir::ItemKind::OpaqueTy(_) => {
+            desugar::desugar_generics_and_predicates(genv, owner_id, resolver_output, None)?;
+        }
         hir::ItemKind::Enum(..) => {
             let enum_def = &specs.enums[&owner_id];
             desugar::desugar_enum_def(genv, owner_id, enum_def, resolver_output)?;
@@ -283,7 +289,7 @@ 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)?;
+            desugar::desugar_generics_and_predicates(genv, owner_id, resolver_output, None)?;
             items.iter().try_for_each_exhaust(|trait_item| {
                 desugar_assoc_item(
                     genv,
@@ -295,7 +301,8 @@ fn desugar_item(
             })?;
         }
         hir::ItemKind::Impl(impl_) => {
-            desugar::desugar_generics_and_predicates(genv, owner_id)?;
+            let generics = specs.impls.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(
                     genv,
@@ -316,11 +323,13 @@ fn desugar_assoc_item(
     specs: &mut Specs,
     owner_id: OwnerId,
     kind: hir::AssocItemKind,
-    resolver_outpt: &ResolverOutput,
+    resolver_output: &ResolverOutput,
 ) -> Result<(), ErrorGuaranteed> {
     match kind {
-        hir::AssocItemKind::Fn { .. } => desugar_fn_sig(genv, specs, owner_id, resolver_outpt),
-        hir::AssocItemKind::Type => desugar::desugar_generics_and_predicates(genv, owner_id),
+        hir::AssocItemKind::Fn { .. } => desugar_fn_sig(genv, specs, owner_id, resolver_output),
+        hir::AssocItemKind::Type => {
+            desugar::desugar_generics_and_predicates(genv, owner_id, resolver_output, None)
+        }
         hir::AssocItemKind::Const => Ok(()),
     }
 }

crates/flux-driver/src/collector.rs

@@ -41,6 +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 enums: FxHashMap<OwnerId, surface::EnumDef>,
     pub qualifs: Vec<surface::Qualifier>,
     pub func_defs: Vec<surface::FuncDef>,
@@ -101,6 +102,7 @@ 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),
                 _ => Ok(()),
             };
         }
@@ -185,6 +187,21 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
         }
     }
 
+    fn parse_impl_spec(
+        &mut self,
+        owner_id: OwnerId,
+        attrs: &[Attribute],
+    ) -> Result<(), ErrorGuaranteed> {
+        let mut attrs = self.parse_flux_attrs(attrs)?;
+        self.report_dups(&attrs)?;
+
+        if let Some(generics) = attrs.generics() {
+            self.specs.impls.insert(owner_id, generics);
+        }
+
+        Ok(())
+    }
+
     fn parse_tyalias_spec(
         &mut self,
         owner_id: OwnerId,
@@ -209,6 +226,8 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
 
         let refined_by = attrs.refined_by();
 
+        let generics = attrs.generics();
+
         let fields = data
             .fields()
             .iter()
@@ -228,9 +247,10 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
                 .insert(extern_def_id, owner_id.def_id);
         }
 
-        self.specs
-            .structs
-            .insert(owner_id, surface::StructDef { refined_by, fields, opaque, invariants });
+        self.specs.structs.insert(
+            owner_id,
+            surface::StructDef { refined_by, generics, fields, opaque, invariants },
+        );
 
         Ok(())
     }
@@ -361,6 +381,9 @@ impl<'tcx, 'a> SpecCollector<'tcx, 'a> {
             ("refined_by", AttrArgs::Delimited(dargs)) => {
                 self.parse(dargs, ParseSess::parse_refined_by, FluxAttrKind::RefinedBy)?
             }
+            ("generics", AttrArgs::Delimited(dargs)) => {
+                self.parse(dargs, ParseSess::parse_generics, FluxAttrKind::Generics)?
+            }
             ("field", AttrArgs::Delimited(dargs)) => {
                 self.parse(dargs, ParseSess::parse_type, FluxAttrKind::Field)?
             }
@@ -489,6 +512,7 @@ impl Specs {
     fn new() -> Specs {
         Specs {
             fn_sigs: Default::default(),
+            impls: Default::default(),
             structs: Default::default(),
             enums: Default::default(),
             qualifs: Vec::default(),
@@ -546,6 +570,7 @@ enum FluxAttrKind {
     Opaque,
     FnSig(surface::FnSig),
     RefinedBy(surface::RefinedBy),
+    Generics(surface::Generics),
     QualNames(surface::QualNames),
     Items(Vec<surface::Item>),
     TypeAlias(surface::TyAlias),
@@ -636,6 +661,10 @@ impl FluxAttrs {
         read_attr!(self, RefinedBy)
     }
 
+    fn generics(&mut self) -> Option<surface::Generics> {
+        read_attr!(self, Generics)
+    }
+
     fn field(&mut self) -> Option<surface::Ty> {
         read_attr!(self, Field)
     }
@@ -674,6 +703,7 @@ impl FluxAttrKind {
             FluxAttrKind::FnSig(_) => attr_name!(FnSig),
             FluxAttrKind::ConstSig(_) => attr_name!(ConstSig),
             FluxAttrKind::RefinedBy(_) => attr_name!(RefinedBy),
+            FluxAttrKind::Generics(_) => attr_name!(Generics),
             FluxAttrKind::Items(_) => attr_name!(Items),
             FluxAttrKind::QualNames(_) => attr_name!(QualNames),
             FluxAttrKind::Field(_) => attr_name!(Field),

crates/flux-fhir-analysis/locales/en-US.ftl

@@ -80,6 +80,9 @@ fhir_analysis_expected_numeric =
 fhir_analysis_no_equality =
     values of sort `{$sort}` cannot be compared for equality
 
+fhir_analysis_invalid_base_instance =
+    values of this type cannot be used as base sorted instances
+
 fhir_analysis_param_not_determined =
     parameter `{$sym}` cannot be determined
     .label = undetermined parameter
@@ -153,4 +156,3 @@ fhir_analysis_assoc_type_not_found =
     associated type not found
     .label = cannot resolve this associated type
     .note = flux cannot resolved associated types if they are defined in a super trait
-

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

@@ -162,6 +162,7 @@ pub(crate) fn conv_generics(
                 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,
             };
@@ -193,12 +194,24 @@ pub(crate) fn conv_generics(
     })
 }
 
+fn sort_args_for_adt(genv: &GlobalEnv, def_id: impl Into<DefId>) -> List<fhir::Sort> {
+    let mut sort_args = vec![];
+    for param in &genv.tcx.generics_of(def_id.into()).params {
+        if let rustc_middle::ty::GenericParamDefKind::Type { .. } = param.kind {
+            sort_args.push(fhir::Sort::Param(param.def_id));
+        }
+    }
+    List::from_vec(sort_args)
+}
+
 pub(crate) fn adt_def_for_struct(
     genv: &GlobalEnv,
     invariants: Vec<rty::Invariant>,
     struct_def: &fhir::StructDef,
 ) -> rty::AdtDef {
-    let sort = rty::Sort::tuple(conv_sorts(genv, genv.index_sorts_of(struct_def.owner_id)));
+    let def_id = struct_def.owner_id;
+    let sort_args = sort_args_for_adt(genv, def_id);
+    let sort = rty::Sort::tuple(conv_sorts(genv, &genv.index_sorts_of(def_id, &sort_args)));
     let adt_def = lowering::lower_adt_def(&genv.tcx.adt_def(struct_def.owner_id));
     rty::AdtDef::new(adt_def, sort, invariants, struct_def.is_opaque())
 }
@@ -208,7 +221,9 @@ pub(crate) fn adt_def_for_enum(
     invariants: Vec<rty::Invariant>,
     enum_def: &fhir::EnumDef,
 ) -> rty::AdtDef {
-    let sort = rty::Sort::tuple(conv_sorts(genv, genv.index_sorts_of(enum_def.owner_id)));
+    let def_id = enum_def.owner_id;
+    let sort_args = sort_args_for_adt(genv, def_id);
+    let sort = rty::Sort::tuple(conv_sorts(genv, &genv.index_sorts_of(def_id, &sort_args)));
     let adt_def = lowering::lower_adt_def(&genv.tcx.adt_def(enum_def.owner_id));
     rty::AdtDef::new(adt_def, sort, invariants, false)
 }
@@ -381,8 +396,9 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
                 let kind = rty::BoundRegionKind::BrNamed(def_id.to_def_id(), name);
                 Ok(rty::BoundVariableKind::Region(kind))
             }
-            fhir::GenericParamKind::Type { default: _ } => bug!("unexpected!"),
-            fhir::GenericParamKind::BaseTy => bug!("unexpected!"),
+            fhir::GenericParamKind::Type { default: _ }
+            | fhir::GenericParamKind::BaseTy
+            | fhir::GenericParamKind::SplTy => bug!("unexpected!"),
         }
     }
 
@@ -621,6 +637,9 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
 
     fn conv_base_ty(&self, env: &mut Env, bty: &fhir::BaseTy) -> QueryResult<rty::Ty> {
         let sort = self.genv.sort_of_bty(bty);
+        // if sort.is_none() {
+        //     println!("TRACE: conv_base_ty {bty:?} ==> None");
+        // }
 
         if let fhir::BaseTyKind::Path(fhir::QPath::Resolved(self_ty, path)) = &bty.kind {
             if let fhir::Res::Def(DefKind::AssocTy, def_id) = path.res {
@@ -640,7 +659,6 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
                 return Ok(rty::Ty::param(param_ty));
             }
         }
-
         let sort = conv_sort(self.genv, &sort.unwrap());
         if sort.is_unit() {
             let idx = rty::Index::from(rty::Expr::unit());
@@ -710,7 +728,7 @@ impl<'a, 'tcx> ConvCtxt<'a, 'tcx> {
                 let expr = self.add_coercions(rty::Expr::abs(body), *fhir_id);
                 (expr, rty::TupleTree::Leaf(false))
             }
-            fhir::RefineArg::Record(_, flds, ..) => {
+            fhir::RefineArg::Record(_, _, flds, ..) => {
                 let mut exprs = vec![];
                 let mut is_binder = vec![];
                 for arg in flds {
@@ -1087,10 +1105,12 @@ impl LookupResult<'_> {
     fn is_record(&self) -> Option<DefId> {
         match &self.kind {
             LookupResultKind::LateBoundList {
-                entry: Entry::Sort { sort: fhir::Sort::Record(def_id), .. },
+                entry: Entry::Sort { sort: fhir::Sort::Record(def_id, _), .. },
                 ..
             } => Some(*def_id),
-            LookupResultKind::EarlyBound { sort: fhir::Sort::Record(def_id), .. } => Some(*def_id),
+            LookupResultKind::EarlyBound { sort: fhir::Sort::Record(def_id, _), .. } => {
+                Some(*def_id)
+            }
             _ => None,
         }
     }
@@ -1142,8 +1162,8 @@ pub fn conv_sort(genv: &GlobalEnv, sort: &fhir::Sort) -> rty::Sort {
         fhir::Sort::Loc => rty::Sort::Loc,
         fhir::Sort::Unit => rty::Sort::unit(),
         fhir::Sort::Func(fsort) => rty::Sort::Func(conv_func_sort(genv, fsort)),
-        fhir::Sort::Record(def_id) => {
-            rty::Sort::tuple(conv_sorts(genv, genv.index_sorts_of(*def_id)))
+        fhir::Sort::Record(def_id, sort_args) => {
+            rty::Sort::tuple(conv_sorts(genv, &genv.index_sorts_of(*def_id, sort_args)))
         }
         fhir::Sort::App(ctor, args) => {
             let ctor = conv_sort_ctor(ctor);

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

@@ -160,7 +160,7 @@ fn generics_of(genv: &GlobalEnv, local_id: LocalDefId) -> QueryResult<rty::Gener
             let is_trait = (def_kind == DefKind::Trait).then_some(local_id);
             let generics = genv
                 .map()
-                .get_generics(local_id)
+                .get_generics(local_id.to_def_id())
                 .unwrap_or_else(|| bug!("no generics for {:?}", def_id));
             let refine_params = genv
                 .map()
@@ -294,7 +294,6 @@ fn check_wf_rust_item(genv: &GlobalEnv, def_id: LocalDefId) -> QueryResult<Rc<fh
         }
         DefKind::Fn | DefKind::AssocFn => {
             let owner_id = OwnerId { def_id };
-
             let fn_sig = genv.map().get_fn_sig(def_id);
             let mut wfckresults = wf::check_fn_sig(genv, fn_sig, owner_id)?;
             annot_check::check_fn_sig(genv.tcx, genv.sess, &mut wfckresults, owner_id, fn_sig)?;

crates/flux-fhir-analysis/src/wf/errors.rs

@@ -196,6 +196,20 @@ impl<'a> InvalidPrimitiveDotAccess<'a> {
     }
 }
 
+#[derive(Diagnostic)]
+#[diag(fhir_analysis_invalid_base_instance, code = "FLUX")]
+pub(super) struct InvalidBaseInstance<'a> {
+    #[primary_span]
+    span: Span,
+    ty: &'a fhir::Ty,
+}
+
+impl<'a> InvalidBaseInstance<'a> {
+    pub(super) fn new(ty: &'a fhir::Ty) -> Self {
+        Self { ty, span: ty.span }
+    }
+}
+
 #[derive(Diagnostic)]
 #[diag(fhir_analysis_no_equality, code = "FLUX")]
 pub(super) struct NoEquality<'a> {