Use visitor to collect base sorts

flux-rs · Jan 4, 2024 · 5ec4c76 · 5ec4c76
1 parent df232bf
commit 5ec4c76
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 43 deletions.
diff --git a/crates/flux-desugar/src/desugar.rs b/crates/flux-desugar/src/desugar.rs
@@ -92,56 +92,27 @@ 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);
+    struct IdentCollector {
+        found: FxHashSet<Symbol>,
+    }
+    impl surface::visit::Visitor for IdentCollector {
+        fn visit_base_sort(&mut self, bsort: &surface::BaseSort) {
+            if let surface::BaseSort::Ident(x) = bsort {
+                self.found.insert(x.name);
+            }
+            surface::visit::walk_base_sort(self, bsort);
+        }
     }
+    let mut vis = IdentCollector { found: FxHashSet::default() };
+    surface::visit::Visitor::visit_refined_by(&mut vis, refined_by);
     generics
         .params
         .iter()
-        .filter_map(|param| if sort_idents.contains(&param.name) { Some(param.name) } else { None })
+        .filter_map(|param| if vis.found.contains(&param.name) { Some(param.name) } else { None })
         .collect()
 }
 

diff --git a/crates/flux-fhir-analysis/src/wf/sortck.rs b/crates/flux-fhir-analysis/src/wf/sortck.rs
@@ -10,7 +10,6 @@ use flux_middle::{
 use itertools::izip;
 use rustc_data_structures::unord::UnordMap;
 use rustc_errors::IntoDiagnostic;
-use rustc_hir::def_id::DefId;
 use rustc_span::Span;
 
 use super::errors;