support impl blocks

flux-rs · Oct 21, 2023 · 01d021d · 01d021d
1 parent 90aec9f
commit 01d021d
Show file tree

Hide file tree

Showing 9 changed files with 112 additions and 58 deletions.
diff --git a/README.md b/README.md
@@ -15,25 +15,4 @@ For an overview, take a look at the [`flux` website](https://flux-rs.github.io).
 # Docs
 
 Documentation, including installation and usage guides can be found on the
-[website](https://flux-rs.github.io/flux).
-
------
-fn sort_params(
-    generics: &rustc_middle::ty::Generics,
-    refined_by: &surface::RefinedBy,
-) -> Vec<DefId> {
-    let sort_params: FxHashSet<_> = refined_by
-        .sort_params
-        .iter()
-        .map(|ident| ident.name)
-        .collect();
-    let mut params = vec![];
-    for param in &generics.params {
-        if let rustc_middle::ty::GenericParamDefKind::Type { .. } = param.kind &&
-           sort_params.contains(&param.name)
-        {
-            params.push(param.def_id);
-        }
-    }
-    params
-}
+[website](https://flux-rs.github.io/flux).
diff --git a/crates/flux-desugar/src/desugar.rs b/crates/flux-desugar/src/desugar.rs
@@ -196,9 +196,25 @@ impl<'a, 'tcx> DesugarCtxt<'a, 'tcx> {
 
     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();
-        let generics_map: FxHashMap<_, _> = hir_generics
-            .params
-            .iter()
+
+        let hir_generic_params = hir_generics.params.iter();
+        // let parent_generic_params = self
+        //     .genv
+        //     .tcx
+        //     .opt_local_parent(self.owner.def_id)
+        //     .map(|parent_def_id| {
+        //         self.genv
+        //             .hir()
+        //             .get_generics(parent_def_id)
+        //             .unwrap()
+        //             .params
+        //             .iter()
+        //     })
+        //     .into_iter()
+        //     .flatten();
+
+        let generics_map: FxHashMap<_, _> = hir_generic_params
+            // .chain(parent_generic_params)
             .flat_map(|param| {
                 if let hir::ParamName::Plain(name) = param.name {
                     Some((name, param.def_id))
@@ -1076,27 +1092,13 @@ impl DesugarCtxt<'_, '_> {
                 }
             })
     }
-    // ORIGINAL
-    // fn gather_params_variant_ret(
-    //     &self,
-    //     ret: &surface::VariantRet,
-    //     binders: &mut Binders,
-    // ) -> Result {
-    //     self.gather_params_path(&ret.path, TypePos::Other, binders)?;
-    //     let res = self.resolver_output.path_res_map[&ret.path.node_id];
-    //     let Some(sort) = self.genv.sort_of_res(res) else {
-    //         return Err(self.emit_err(errors::RefinedUnrefinableType::new(ret.path.span)));
-    //     };
-    //     self.gather_params_indices(sort, &ret.indices, TypePos::Other, binders)
-    // }
 
     fn gather_params_variant_ret(
         &self,
         ret: &surface::VariantRet,
         binders: &mut Binders,
     ) -> Result {
         self.gather_params_path(&ret.path, TypePos::Other, binders)?;
-        // let Some(sort) = self.resolver_output.sort_of_surface_path(&ret.path) else {
         let Some(sort) = sort_of_surface_path(self.genv, self.resolver_output, &ret.path) else {
             return Err(self.emit_err(errors::RefinedUnrefinableType::new(ret.path.span)));
         };
@@ -1520,14 +1522,6 @@ impl Binders {
         Binder::Refined(self.fresh(), sort, true)
     }
 
-    // fn binder_from_res(&self, genv: &GlobalEnv, res: fhir::Res) -> Binder {
-    //     if let Some(sort) = genv.sort_of_res(res) {
-    //         self.binder_from_sort(sort)
-    //     } else {
-    //         Binder::Unrefined
-    //     }
-    // }
-
     fn binder_from_bty(
         &self,
         genv: &GlobalEnv,
@@ -1637,14 +1631,15 @@ fn index_sort(
     bty: &surface::BaseTy,
 ) -> Option<fhir::Sort> {
     // CODESYNC(sort-of, 4) sorts should be given consistently
-    match &bty.kind {
+    let res = match &bty.kind {
         surface::BaseTyKind::Path(path) => {
             sort_of_surface_path(genv, resolver_output, path)
             // let res = resolver_output.path_res_map[&path.node_id];
             // genv.sort_of_res(res)
         }
         surface::BaseTyKind::Slice(_) => Some(fhir::Sort::Int),
-    }
+    };
+    res
 }
 
 fn sort_of_surface_path(
@@ -1673,6 +1668,7 @@ fn sort_of_surface_path(
         }
         fhir::Res::Def(DefKind::TyParam, def_id) => {
             let param = genv.get_generic_param(def_id.expect_local());
+
             match &param.kind {
                 fhir::GenericParamKind::BaseTy => Some(fhir::Sort::Param(def_id)),
                 fhir::GenericParamKind::Type { .. } | fhir::GenericParamKind::Lifetime => None,

diff --git a/crates/flux-desugar/src/lib.rs b/crates/flux-desugar/src/lib.rs
@@ -178,10 +178,19 @@ pub fn desugar_fn_sig(
 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(())

diff --git a/crates/flux-driver/src/callbacks.rs b/crates/flux-driver/src/callbacks.rs
@@ -275,7 +275,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)?;
@@ -285,7 +287,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,
@@ -297,7 +299,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,
@@ -318,11 +321,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(()),
     }
 }

diff --git a/crates/flux-driver/src/collector.rs b/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,
@@ -361,6 +378,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 +509,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 +567,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 +658,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 +700,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),

diff --git a/crates/flux-syntax/src/grammar.lalrpop b/crates/flux-syntax/src/grammar.lalrpop
@@ -10,7 +10,7 @@ use lalrpop_util::ParseError;
 
 grammar(cx: &mut ParseCtxt<'_>);
 
-Generics: surface::Generics = {
+pub Generics: surface::Generics = {
     <lo:@L> "<" <params:Comma<GenericParam>> ">" <hi:@R> => {
         surface::Generics {
             params,

diff --git a/crates/flux-syntax/src/lib.rs b/crates/flux-syntax/src/lib.rs
@@ -41,6 +41,14 @@ impl ParseSess {
         parse!(self, grammar::RefinedByParser, tokens, span)
     }
 
+    pub fn parse_generics(
+        &mut self,
+        tokens: &TokenStream,
+        span: Span,
+    ) -> ParseResult<surface::Generics> {
+        parse!(self, grammar::GenericsParser, tokens, span)
+    }
+
     pub fn parse_type_alias(
         &mut self,
         tokens: &TokenStream,

diff --git a/crates/flux-tests/tests/neg/surface/rset01.rs b/crates/flux-tests/tests/neg/surface/rset01.rs
@@ -0,0 +1,15 @@
+#[path = "../../lib/rset.rs"]
+mod rset;
+use rset::RSet;
+
+#[flux::sig(fn (bool[true]))]
+fn assert(_b: bool) {}
+
+pub fn test() {
+    let mut s = RSet::new();
+    let v0 = 666;
+    let v1 = 667;
+    s.insert(v0);
+    assert(s.contains(&v0));
+    assert(s.contains(&v1)); //~ ERROR refinement type
+}
diff --git a/crates/flux-tests/tests/pos/surface/rset01.rs b/crates/flux-tests/tests/pos/surface/rset01.rs
@@ -0,0 +1,15 @@
+#[path = "../../lib/rset.rs"]
+mod rset;
+use rset::RSet;
+
+#[flux::sig(fn (bool[true]))]
+fn assert(_b: bool) {}
+
+pub fn test() {
+    let mut s = RSet::new();
+    let v0 = 666;
+    let v1 = 667;
+    s.insert(v0);
+    assert(s.contains(&v0));
+    assert(!s.contains(&v1));
+}