Parse Array ConstArg (#888)

* parse const for array

* only parse _ ConstArg

* remove the ident from const-arg

crates/flux-desugar/src/desugar.rs

@@ -1072,9 +1072,8 @@ trait DesugarCtxt<'genv, 'tcx: 'genv> {
             }
             surface::TyKind::Array(ty, len) => {
                 let ty = self.desugar_ty(ty)?;
-                let const_arg =
-                    fhir::ConstArg { kind: fhir::ConstArgKind::Lit(len.val), span: len.span };
-                fhir::TyKind::Array(self.genv().alloc(ty), const_arg)
+                let len = Self::desugar_const_arg(len)?;
+                fhir::TyKind::Array(self.genv().alloc(ty), len)
             }
             surface::TyKind::ImplTrait(node_id, bounds) => {
                 self.desugar_impl_trait(*node_id, bounds)?
@@ -1084,6 +1083,14 @@ trait DesugarCtxt<'genv, 'tcx: 'genv> {
         Ok(fhir::Ty { kind, span })
     }
 
+    fn desugar_const_arg(const_arg: &surface::ConstArg) -> Result<fhir::ConstArg> {
+        let kind = match const_arg.kind {
+            surface::ConstArgKind::Lit(val) => fhir::ConstArgKind::Lit(val),
+            surface::ConstArgKind::Infer => fhir::ConstArgKind::Infer,
+        };
+        Ok(fhir::ConstArg { kind, span: const_arg.span })
+    }
+
     fn desugar_bty(&mut self, bty: &surface::BaseTy) -> Result<fhir::BaseTy<'genv>> {
         match &bty.kind {
             surface::BaseTyKind::Path(qself, path) => {

crates/flux-syntax/src/grammar.lalrpop

@@ -368,6 +368,25 @@ pub Ty: surface::Ty = {
     }
 }
 
+ConstArg: surface::ConstArg = {
+   <lo:@L> <lit:Lit> <hi:@R> =>? {
+       let span = cx.map_span(lo, hi);
+       let mut kind = surface::ConstArgKind::Infer;
+       if let surface::LitKind::Integer = lit.kind {
+        if let Ok(val) = lit.symbol.as_str().parse::<usize>() {
+            let kind = surface::ConstArgKind::Lit(val);
+            return Ok(surface::ConstArg { kind, span });
+        }
+       }
+       Err( ParseError::User { error: UserParseError::UnexpectedToken(lo, hi) })
+   },
+   <lo:@L> "_" <hi:@R> => {
+       let span = cx.map_span(lo, hi);
+       let kind = surface::ConstArgKind::Infer;
+       surface::ConstArg { kind, span }
+   }
+}
+
 TyKind: surface::TyKind = {
     "_"                                               => surface::TyKind::Hole,
     <bty:BaseTy>                                      => surface::TyKind::Base(<>),
@@ -381,14 +400,8 @@ TyKind: surface::TyKind = {
     "&" <ty:Ty>                     => surface::TyKind::Ref(surface::Mutability::Not, Box::new(ty)),
     "&" "mut" <ty:Ty>               => surface::TyKind::Ref(surface::Mutability::Mut, Box::new(ty)),
 
-    "[" <ty:Ty> ";" <lo:@L> <lit:Lit> <hi:@R> "]" =>? {
-        let span = cx.map_span(lo, hi);
-        if let surface::LitKind::Integer = lit.kind {
-            if let Ok(val) = lit.symbol.as_str().parse::<usize>() {
-                return Ok(surface::TyKind::Array(Box::new(ty), surface::ArrayLen { val, span }));
-            }
-        }
-        Err(ParseError::User { error: UserParseError::UnexpectedToken(lo, hi) })
+    "[" <ty:Ty> ";" <len:ConstArg> "]" =>? {
+        return Ok(surface::TyKind::Array(Box::new(ty), len));
     },
 
     "impl" <bounds:GenericBounds> => surface::TyKind::ImplTrait(cx.next_node_id(), bounds),

crates/flux-syntax/src/surface.rs

@@ -337,7 +337,7 @@ pub enum TyKind {
     /// Constrained type: an exists without binder
     Constr(Expr, Box<Ty>),
     Tuple(Vec<Ty>),
-    Array(Box<Ty>, ArrayLen),
+    Array(Box<Ty>, ConstArg),
     /// The `NodeId` is used to resolve the type to a corresponding `OpaqueTy`
     ImplTrait(NodeId, GenericBounds),
     Hole,
@@ -385,12 +385,18 @@ pub enum BaseTyKind {
     Slice(Box<Ty>),
 }
 
-#[derive(Debug, Clone, Copy)]
-pub struct ArrayLen {
-    pub val: usize,
+#[derive(PartialEq, Eq, Clone, Debug, Copy)]
+pub struct ConstArg {
+    pub kind: ConstArgKind,
     pub span: Span,
 }
 
+#[derive(PartialEq, Eq, Clone, Debug, Copy)]
+pub enum ConstArgKind {
+    Lit(usize),
+    Infer,
+}
+
 #[derive(Debug)]
 pub struct Indices {
     pub indices: Vec<RefineArg>,

crates/flux-syntax/src/surface/visit.rs

@@ -1,7 +1,7 @@
 use rustc_span::symbol::Ident;
 
 use super::{
-    AliasReft, ArrayLen, Async, BaseSort, BaseTy, BaseTyKind, Ensures, EnumDef, Expr, ExprKind,
+    AliasReft, Async, BaseSort, BaseTy, BaseTyKind, ConstArg, Ensures, EnumDef, Expr, ExprKind,
     ExprPath, ExprPathSegment, FnInput, FnOutput, FnRetTy, FnSig, GenericArg, GenericArgKind,
     GenericParam, Generics, Impl, ImplAssocReft, Indices, Lit, Path, PathSegment, Qualifier,
     RefineArg, RefineParam, Sort, SortPath, SpecFunc, StructDef, Trait, TraitAssocReft, TraitRef,
@@ -135,7 +135,7 @@ pub trait Visitor: Sized {
         walk_ty(self, ty);
     }
 
-    fn visit_array_len(&mut self, _array_len: &ArrayLen) {}
+    fn visit_const_arg(&mut self, _const_arg: &ConstArg) {}
 
     fn visit_bty(&mut self, bty: &BaseTy) {
         walk_bty(self, bty);
@@ -426,7 +426,7 @@ pub fn walk_ty<V: Visitor>(vis: &mut V, ty: &Ty) {
             walk_list!(vis, visit_ty, tys);
         }
         TyKind::Array(ty, len) => {
-            vis.visit_array_len(len);
+            vis.visit_const_arg(len);
             vis.visit_ty(ty);
         }
         TyKind::ImplTrait(_node_id, trait_ref) => {

tests/tests/lib/rvec.rs

@@ -79,6 +79,26 @@ impl<T> RVec<T> {
         self.inner.as_mut_slice()
     }
 
+
+    #[flux::trusted]
+    #[flux::sig(fn(arr:_) -> RVec<T>[N])]
+    pub fn from_array<const N: usize>(arr: [T; N]) -> Self {
+        Self {
+            inner: Vec::from(arr),
+        }
+    }
+
+    #[flux::trusted]
+    #[flux::sig(fn(xs:&[T][@n]) -> RVec<T>[n])]
+    pub fn from_slice(xs: &[T]) -> Self
+    where
+        T: Clone,
+    {
+        Self {
+            inner: Vec::from(xs),
+        }
+    }
+
     #[flux::trusted]
     #[flux::sig(fn(T, n: usize) -> RVec<T>[n])]
     pub fn from_elem_n(elem: T, n: usize) -> Self

tests/tests/pos/surface/const06.rs

@@ -0,0 +1,2 @@
+#[flux::sig(fn (head: [i32; _]))]
+pub fn test2<const N: usize>(_head: [i32; N]) {}