remove explicit sort_params

flux-rs · Nov 2, 2023 · aee782b · aee782b
1 parent 3ab9b60
commit aee782b
Show file tree

Hide file tree

Showing 8 changed files with 66 additions and 26 deletions.
diff --git a/crates/flux-desugar/src/desugar.rs b/crates/flux-desugar/src/desugar.rs
@@ -10,6 +10,7 @@ use flux_middle::{
 };
 use flux_syntax::surface;
 use hir::{def::DefKind, ItemKind, PrimTy};
+use itertools::Itertools;
 use rustc_data_structures::{
     fx::{FxIndexMap, IndexEntry},
     unord::UnordMap,
@@ -53,9 +54,9 @@ pub fn desugar_qualifier(
 
 pub fn desugar_defn(genv: &GlobalEnv, defn: surface::FuncDef) -> Result<Option<fhir::Defn>> {
     if let Some(body) = defn.body {
-        let sort_params = &defn.sort_vars[..];
+        let sort_params = defn.sort_vars.iter().map(|ident| ident.name).collect_vec();
         let sort_resolver =
-            SortResolver::with_sort_params(genv.sess, genv.map().sort_decls(), sort_params);
+            SortResolver::with_sort_params(genv.sess, genv.map().sort_decls(), &sort_params);
         let mut binders = Binders::from_params(genv, &sort_resolver, &defn.args)?;
         let local_id_gen = IndexGen::new();
         let cx = ExprCtxt::new(genv, FluxOwnerId::Flux(defn.name.name), &local_id_gen);
@@ -76,8 +77,8 @@ pub fn func_def_to_func_decl(
     defn: &surface::FuncDef,
 ) -> Result<fhir::FuncDecl> {
     let params = defn.sort_vars.len();
-    let sort_vars = &defn.sort_vars[..];
-    let sr = SortResolver::with_sort_params(sess, sort_decls, sort_vars);
+    let sort_vars = defn.sort_vars.iter().map(|ident| ident.name).collect_vec();
+    let sr = SortResolver::with_sort_params(sess, sort_decls, &sort_vars);
     let inputs: Vec<fhir::Sort> = defn
         .args
         .iter()
@@ -89,6 +90,57 @@ pub fn func_def_to_func_decl(
     Ok(fhir::FuncDecl { name: defn.name.name, sort, kind })
 }
 
+fn gather_base_sort_vars(
+    generics: &FxHashSet<Symbol>,
+    base_sort: &surface::BaseSort,
+    sort_vars: &mut FxHashSet<Symbol>,
+) {
+    match base_sort {
+        surface::BaseSort::Ident(x) => {
+            if generics.contains(&x.name) {
+                sort_vars.insert(x.name);
+            }
+        }
+        surface::BaseSort::BitVec(_) => {}
+        surface::BaseSort::App(_, base_sorts) => {
+            for base_sort in base_sorts {
+                gather_base_sort_vars(generics, base_sort, sort_vars)
+            }
+        }
+    }
+}
+fn gather_sort_vars(
+    generics: &FxHashSet<Symbol>,
+    sort: &surface::Sort,
+    sort_vars: &mut FxHashSet<Symbol>,
+) {
+    match sort {
+        surface::Sort::Base(base_sort) => gather_base_sort_vars(generics, base_sort, sort_vars),
+        surface::Sort::Func { inputs, output } => {
+            for base_sort in inputs {
+                gather_base_sort_vars(generics, base_sort, sort_vars);
+            }
+            gather_base_sort_vars(generics, output, sort_vars);
+        }
+        surface::Sort::Infer => {}
+    }
+}
+
+fn gather_refined_by_sort_vars(
+    generics: &rustc_middle::ty::Generics,
+    refined_by: &surface::RefinedBy,
+) -> Vec<Symbol> {
+    let generics_syms: FxHashSet<Symbol> = generics.params.iter().map(|param| param.name).collect();
+    let mut sort_idents = FxHashSet::default();
+    for refine_param in &refined_by.index_params {
+        gather_sort_vars(&generics_syms, &refine_param.sort, &mut sort_idents);
+    }
+    generics
+        .params
+        .iter()
+        .filter_map(|param| if sort_idents.contains(&param.name) { Some(param.name) } else { None })
+        .collect()
+}
 pub fn desugar_refined_by(
     sess: &FluxSession,
     sort_decls: &fhir::SortDecls,
@@ -105,7 +157,8 @@ pub fn desugar_refined_by(
         }
     })?;
 
-    let sr = SortResolver::with_sort_params(sess, sort_decls, &refined_by.sort_vars);
+    let sort_vars = gather_refined_by_sort_vars(generics, refined_by);
+    let sr = SortResolver::with_sort_params(sess, sort_decls, &sort_vars);
 
     let early_bound_params: Vec<_> = refined_by
         .early_bound_params
@@ -125,11 +178,7 @@ pub fn desugar_refined_by(
         .enumerate()
         .map(|(i, param)| (param.name, i))
         .collect();
-    let sort_params = refined_by
-        .sort_vars
-        .iter()
-        .map(|ident| generic_idx[&ident.name])
-        .collect();
+    let sort_params = sort_vars.iter().map(|sym| generic_idx[&sym]).collect();
 
     Ok(fhir::RefinedBy::new(
         owner_id.def_id,
@@ -1366,12 +1415,12 @@ impl<'a> SortResolver<'a> {
     pub fn with_sort_params(
         sess: &'a FluxSession,
         sort_decls: &'a fhir::SortDecls,
-        sort_params: &[surface::Ident],
+        sort_params: &[Symbol],
     ) -> Self {
         let sort_params = sort_params
             .iter()
             .enumerate()
-            .map(|(i, v)| (v.name, i))
+            .map(|(i, v)| (*v, i))
             .collect();
         Self { sess, sort_decls, generic_params: Default::default(), sort_params }
     }

diff --git a/crates/flux-syntax/src/grammar.lalrpop b/crates/flux-syntax/src/grammar.lalrpop
@@ -48,7 +48,6 @@ pub TyAlias: surface::TyAlias = {
     <ty:Ty>
     <hi:@R> => {
         let refined_by = surface::RefinedBy {
-            sort_vars: vec![],
             early_bound_params: early_bound_params.unwrap_or_default(),
             index_params: index_params.unwrap_or_default(),
             span: cx.map_span(refined_by_lo, refined_by_hi)
@@ -64,14 +63,7 @@ pub TyAlias: surface::TyAlias = {
 }
 
 pub RefinedBy: surface::RefinedBy = {
-    <lo:@L>  <sort_vars:SortVars> "{" <index_params:RefineParams<"!">> "}" <hi:@R> => surface::RefinedBy {
-        sort_vars,
-        index_params,
-        early_bound_params: vec![],
-        span: cx.map_span(lo, hi)
-    },
     <lo:@L>  <index_params:RefineParams<"!">> <hi:@R> => surface::RefinedBy {
-        sort_vars: vec![],
         index_params,
         early_bound_params: vec![],
         span: cx.map_span(lo, hi)

diff --git a/crates/flux-syntax/src/surface.rs b/crates/flux-syntax/src/surface.rs
@@ -108,7 +108,6 @@ pub struct VariantRet {
 #[derive(Debug, Default)]
 pub struct RefinedBy {
     pub early_bound_params: Vec<RefineParam>,
-    pub sort_vars: Vec<Ident>,
     pub index_params: Vec<RefineParam>,
     pub span: Span,
 }

diff --git a/crates/flux-tests/tests/lib/rmap.rs b/crates/flux-tests/tests/lib/rmap.rs
@@ -9,7 +9,7 @@ use std::hash::Hash;
 
 /// define a type indexed by a map
 #[flux::opaque]
-#[flux::refined_by(<K, V> { vals: Map<K, V> } )]
+#[flux::refined_by(vals: Map<K, V>)]
 pub struct RMap<K, V> {
     inner: std::collections::HashMap<K, V>,
 }

diff --git a/crates/flux-tests/tests/lib/rmapk.rs b/crates/flux-tests/tests/lib/rmapk.rs
@@ -12,7 +12,7 @@ use std::hash::Hash;
 
 /// define a type indexed by a map
 #[flux::opaque]
-#[flux::refined_by(<K,V>{keys: Set<K>, vals: Map<K, V>})]
+#[flux::refined_by(keys: Set<K>, vals: Map<K, V>)]
 pub struct RMap<K, V> {
     inner: std::collections::HashMap<K, V>,
 }

diff --git a/crates/flux-tests/tests/lib/rset.rs b/crates/flux-tests/tests/lib/rset.rs
@@ -3,7 +3,7 @@
 use std::hash::Hash;
 
 #[flux::opaque]
-#[flux::refined_by(<T> {elems: Set<T>})]
+#[flux::refined_by(elems: Set<T>)]
 pub struct RSet<T> {
     pub inner: std::collections::HashSet<T>,
 }

diff --git a/crates/flux-tests/tests/neg/surface/rset00.rs b/crates/flux-tests/tests/neg/surface/rset00.rs
@@ -1,6 +1,6 @@
 use std::hash::Hash;
 #[flux::opaque]
-#[flux::refined_by(<Tiger> { elems: Set<Tiger> } )]
+#[flux::refined_by(elems: Set<Tiger>)]
 pub struct RSet<Tiger> {
     pub inner: std::collections::HashSet<Tiger>,
 }

diff --git a/crates/flux-tests/tests/pos/surface/rset00.rs b/crates/flux-tests/tests/pos/surface/rset00.rs
@@ -1,6 +1,6 @@
 use std::hash::Hash;
 #[flux::opaque]
-#[flux::refined_by(<Tiger> { elems: Set<Tiger> } )]
+#[flux::refined_by( elems: Set<Tiger> )]
 pub struct RSet<Tiger> {
     pub inner: std::collections::HashSet<Tiger>,
 }