Track param kind in fhir

crates/flux-desugar/src/desugar.rs

@@ -207,7 +207,7 @@ type Env = env::Env<Param>;
 struct Param {
     name: fhir::Name,
     sort: fhir::Sort,
-    implicit: bool,
+    kind: fhir::ParamKind,
 }
 
 struct FluxItemCtxt<'a, 'tcx> {
@@ -959,7 +959,7 @@ impl Env {
             env.insert(
                 genv.sess,
                 param.name,
-                Param { name: name_gen.fresh(), sort, implicit: false },
+                Param { name: name_gen.fresh(), sort, kind: fhir::ParamKind::Explicit },
             )?;
         }
         Ok(env)
@@ -1015,12 +1015,7 @@ impl Scope<Param> {
         for (ident, param) in self.into_iter() {
             let ident = fhir::Ident::new(param.name, ident);
             let fhir_id = cx.next_fhir_id();
-            params.push(fhir::RefineParam {
-                ident,
-                sort: param.sort,
-                implicit: param.implicit,
-                fhir_id,
-            });
+            params.push(fhir::RefineParam { ident, sort: param.sort, kind: param.kind, fhir_id });
         }
         params
     }

crates/flux-desugar/src/desugar/gather.rs

@@ -81,13 +81,15 @@ enum Param {
     Pound,
     /// A parameter declared with `x: T` syntax.
     Colon,
-    /// A parameter which we know syntactically cannot be used, e.g., when writing
+    /// A parameter that we know *syntactically* cannot be used inside a refinement. We track these
+    /// parameters to report errors at the use site. For example, consider the following function:
     ///
     /// ```ignore
     /// fn(x: {v. i32[v] | v > 0}) -> i32[x]
     /// ```
-    /// We know syntatically that `x` binds to a non-base type. We track the parameter to report
-    /// errors at the use site.
+    ///
+    /// In this definition, we know syntatically that `x` binds to a non-base type so it's an error
+    /// to use `x` as an index in the return type.
     SyntaxError,
 }
 
@@ -410,30 +412,20 @@ impl Env {
     fn into_desugar_env(self) -> env::Env<super::Param> {
         let name_gen = IndexGen::default();
         self.filter_map(|param, used| {
-            match param {
-                Param::Explicit(sort) => {
-                    Some(super::Param { name: name_gen.fresh(), sort, implicit: false })
-                }
-                Param::At | Param::Pound => {
-                    Some(super::Param {
-                        name: name_gen.fresh(),
-                        sort: fhir::Sort::Wildcard,
-                        implicit: true,
-                    })
-                }
+            let (sort, kind) = match param {
+                Param::Explicit(sort) => (sort, fhir::ParamKind::Explicit),
+                Param::At => (fhir::Sort::Wildcard, fhir::ParamKind::At),
+                Param::Pound => (fhir::Sort::Wildcard, fhir::ParamKind::Pound),
                 Param::Colon => {
                     if used {
-                        Some(super::Param {
-                            name: name_gen.fresh(),
-                            sort: fhir::Sort::Wildcard,
-                            implicit: true,
-                        })
+                        (fhir::Sort::Wildcard, fhir::ParamKind::Colon)
                     } else {
-                        None
+                        return None;
                     }
                 }
-                Param::SyntaxError => None,
-            }
+                Param::SyntaxError => return None,
+            };
+            Some(super::Param { name: name_gen.fresh(), sort, kind })
         })
     }
 }

crates/flux-middle/src/fhir.rs

@@ -448,11 +448,32 @@ pub enum Res {
 pub struct RefineParam {
     pub ident: Ident,
     pub sort: Sort,
-    /// Whether the parameter was implicitly scoped with `@n`, `#n` or `x: T` syntax.
-    pub implicit: bool,
+    pub kind: ParamKind,
     pub fhir_id: FhirId,
 }
 
+/// How the declared parameter in the surface syntax. This is used to adjust how errors are reported
+/// and to control the [inference mode].
+///
+/// [inference mode]: InferMode
+#[derive(Debug, Clone, Copy)]
+pub enum ParamKind {
+    /// A parameter declared in an explicit scope
+    Explicit,
+    /// An implicitly scoped parameter declared with `@a` syntax
+    At,
+    /// An implicitly scoped parameter declared with `#a` syntax
+    Pound,
+    /// An implicitly scoped parameter declared with `x: T` syntax
+    Colon,
+}
+
+impl ParamKind {
+    fn is_implicit(&self) -> bool {
+        matches!(self, ParamKind::At | ParamKind::Pound | ParamKind::Colon)
+    }
+}
+
 /// *Infer*ence *mode* for parameter at function calls
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Encodable, Decodable)]
 pub enum InferMode {
@@ -953,7 +974,7 @@ impl RefineParam {
     }
 
     pub fn infer_mode(&self) -> InferMode {
-        if self.sort.is_pred() && !self.implicit {
+        if self.sort.is_pred() && !self.kind.is_implicit() {
             InferMode::KVar
         } else {
             InferMode::EVar

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

@@ -140,7 +140,7 @@ pub fn walk_fn_sig<V: Visitor>(vis: &mut V, sig: &FnSig) {
 }
 
 pub fn walk_refine_param<V: Visitor>(vis: &mut V, param: &RefineParam) {
-    let RefineParam { ident, sort, implicit: _, fhir_id: _ } = param;
+    let RefineParam { ident, sort, kind: _, fhir_id: _ } = param;
     vis.visit_ident(*ident);
     vis.visit_sort(sort);
 }