fix: Update metavariable expression implementation

Author: Veykril

crates/mbe/src/benchmark.rs

@@ -199,7 +199,7 @@ fn invocation_fixtures(
                 });
                 parent.token_trees.push(subtree.into());
             }
-            Op::Ignore { .. } | Op::Index { .. } | Op::Count { .. } => {}
+            Op::Ignore { .. } | Op::Index { .. } | Op::Count { .. } | Op::Length { .. } => {}
         };
 
         // Simple linear congruential generator for deterministic result

crates/mbe/src/expander/matcher.rs

@@ -588,7 +588,9 @@ fn match_loop_inner<'t, S: Span>(
                 item.is_error = true;
                 error_items.push(item);
             }
-            OpDelimited::Op(Op::Ignore { .. } | Op::Index { .. } | Op::Count { .. }) => {
+            OpDelimited::Op(
+                Op::Ignore { .. } | Op::Index { .. } | Op::Count { .. } | Op::Length { .. },
+            ) => {
                 stdx::never!("metavariable expression in lhs found");
             }
             OpDelimited::Open => {
@@ -851,7 +853,7 @@ fn collect_vars<S: Span>(collector_fun: &mut impl FnMut(SmolStr), pattern: &Meta
             Op::Subtree { tokens, .. } => collect_vars(collector_fun, tokens),
             Op::Repeat { tokens, .. } => collect_vars(collector_fun, tokens),
             Op::Literal(_) | Op::Ident(_) | Op::Punct(_) => {}
-            Op::Ignore { .. } | Op::Index { .. } | Op::Count { .. } => {
+            Op::Ignore { .. } | Op::Index { .. } | Op::Count { .. } | Op::Length { .. } => {
                 stdx::never!("metavariable expression in lhs found");
             }
         }

crates/mbe/src/expander/transcriber.rs

@@ -232,6 +232,21 @@ fn expand_subtree<S: Span>(
                     .into(),
                 );
             }
+            Op::Length { depth } => {
+                let length = ctx.nesting.get(ctx.nesting.len() - 1 - depth).map_or(0, |_nest| {
+                    // FIXME: to be implemented
+                    0
+                });
+                arena.push(
+                    tt::Leaf::Literal(tt::Literal {
+                        text: length.to_string().into(),
+                        // FIXME
+                        #[allow(deprecated)]
+                        span: S::DUMMY,
+                    })
+                    .into(),
+                );
+            }
             Op::Count { name, depth } => {
                 let mut binding = match ctx.bindings.get(name.as_str()) {
                     Ok(b) => b,
@@ -518,28 +533,18 @@ fn fix_up_and_push_path_tt<S: Span>(buf: &mut Vec<tt::TokenTree<S>>, subtree: tt
 fn count<S>(
     ctx: &ExpandCtx<'_, S>,
     binding: &Binding<S>,
-    our_depth: usize,
-    count_depth: Option<usize>,
+    depth_curr: usize,
+    depth_max: usize,
 ) -> Result<usize, CountError> {
     match binding {
-        Binding::Nested(bs) => match count_depth {
-            None => bs.iter().map(|b| count(ctx, b, our_depth + 1, None)).sum(),
-            Some(0) => Ok(bs.len()),
-            Some(d) => bs.iter().map(|b| count(ctx, b, our_depth + 1, Some(d - 1))).sum(),
-        },
-        Binding::Empty => Ok(0),
-        Binding::Fragment(_) | Binding::Missing(_) => {
-            if our_depth == 0 {
-                // `${count(t)}` is placed inside the innermost repetition. This includes cases
-                // where `t` is not a repeated fragment.
-                Err(CountError::Misplaced)
-            } else if count_depth.is_none() {
-                Ok(1)
+        Binding::Nested(bs) => {
+            if depth_curr == depth_max {
+                Ok(bs.len())
             } else {
-                // We've reached at the innermost repeated fragment, but the user wants us to go
-                // further!
-                Err(CountError::OutOfBounds)
+                bs.iter().map(|b| count(ctx, b, depth_curr + 1, depth_max)).sum()
             }
         }
+        Binding::Empty => Ok(0),
+        Binding::Fragment(_) | Binding::Missing(_) => Ok(1),
     }
 }

crates/mbe/src/parser.rs

@@ -54,7 +54,8 @@ pub(crate) enum Op<S> {
     Var { name: SmolStr, kind: Option<MetaVarKind>, id: S },
     Ignore { name: SmolStr, id: S },
     Index { depth: usize },
-    Count { name: SmolStr, depth: Option<usize> },
+    Length { depth: usize },
+    Count { name: SmolStr, depth: usize },
     Repeat { tokens: MetaTemplate<S>, kind: RepeatKind, separator: Option<Separator<S>> },
     Subtree { tokens: MetaTemplate<S>, delimiter: tt::Delimiter<S> },
     Literal(tt::Literal<S>),
@@ -299,15 +300,16 @@ fn parse_metavar_expr<S: Span>(src: &mut TtIter<'_, S>) -> Result<Op<S>, ()> {
 
     let op = match &*func.text {
         "ignore" => {
+            args.expect_dollar()?;
             let ident = args.expect_ident()?;
             Op::Ignore { name: ident.text.clone(), id: ident.span }
         }
         "index" => Op::Index { depth: parse_depth(&mut args)? },
+        "length" => Op::Length { depth: parse_depth(&mut args)? },
         "count" => {
+            args.expect_dollar()?;
             let ident = args.expect_ident()?;
-            // `${count(t)}` and `${count(t,)}` have different meanings. Not sure if this is a bug
-            // but that's how it's implemented in rustc as of this writing. See rust-lang/rust#111904.
-            let depth = if try_eat_comma(&mut args) { Some(parse_depth(&mut args)?) } else { None };
+            let depth = if try_eat_comma(&mut args) { parse_depth(&mut args)? } else { 0 };
             Op::Count { name: ident.text.clone(), depth }
         }
         _ => return Err(()),

crates/mbe/src/tt_iter.rs

@@ -51,6 +51,13 @@ impl<'a, S: Span> TtIter<'a, S> {
         }
     }
 
+    pub(crate) fn expect_dollar(&mut self) -> Result<(), ()> {
+        match self.expect_leaf()? {
+            tt::Leaf::Punct(tt::Punct { char: '$', .. }) => Ok(()),
+            _ => Err(()),
+        }
+    }
+
     pub(crate) fn expect_ident(&mut self) -> Result<&'a tt::Ident<S>, ()> {
         match self.expect_leaf()? {
             tt::Leaf::Ident(it) if it.text != "_" => Ok(it),