internal: move visibility to a prefix entry point

4 files changed, 62 insertions, 8 deletions

diff --git a/crates/mbe/src/expander/matcher.rs b/crates/mbe/src/expander/matcher.rs
index e589dc759f8..3636979d007 100644
--- a/crates/mbe/src/expander/matcher.rs
+++ b/crates/mbe/src/expander/matcher.rs
@@ -895,7 +895,7 @@ impl<'a> TtIter<'a> {
     }
 
     fn eat_vis(&mut self) -> Option<tt::TokenTree> {
-        self.expect_fragment(ParserEntryPoint::Visibility).value
+        self.expect_fragment2(parser::PrefixEntryPoint::Vis).value
     }
 
     fn eat_char(&mut self, c: char) -> Option<tt::TokenTree> {
diff --git a/crates/mbe/src/tt_iter.rs b/crates/mbe/src/tt_iter.rs
index 2d2dbd8994f..03d0cbee8b1 100644
--- a/crates/mbe/src/tt_iter.rs
+++ b/crates/mbe/src/tt_iter.rs
@@ -146,6 +146,63 @@ impl<'a> TtIter<'a> {
         ExpandResult { value: res, err }
     }
 
+    pub(crate) fn expect_fragment2(
+        &mut self,
+        entry_point: parser::PrefixEntryPoint,
+    ) -> ExpandResult<Option<tt::TokenTree>> {
+        let buffer = TokenBuffer::from_tokens(self.inner.as_slice());
+        let parser_input = to_parser_input(&buffer);
+        let tree_traversal = entry_point.parse(&parser_input);
+
+        let mut cursor = buffer.begin();
+        let mut error = false;
+        for step in tree_traversal.iter() {
+            match step {
+                parser::Step::Token { kind, mut n_input_tokens } => {
+                    if kind == SyntaxKind::LIFETIME_IDENT {
+                        n_input_tokens = 2;
+                    }
+                    for _ in 0..n_input_tokens {
+                        cursor = cursor.bump_subtree();
+                    }
+                }
+                parser::Step::Enter { .. } | parser::Step::Exit => (),
+                parser::Step::Error { .. } => error = true,
+            }
+        }
+
+        let mut err = if !cursor.is_root() || error {
+            Some(err!("expected {:?}", entry_point))
+        } else {
+            None
+        };
+
+        let mut curr = buffer.begin();
+        let mut res = vec![];
+
+        if cursor.is_root() {
+            while curr != cursor {
+                if let Some(token) = curr.token_tree() {
+                    res.push(token);
+                }
+                curr = curr.bump();
+            }
+        }
+        self.inner = self.inner.as_slice()[res.len()..].iter();
+        if res.is_empty() && err.is_none() {
+            err = Some(err!("no tokens consumed"));
+        }
+        let res = match res.len() {
+            1 => Some(res[0].cloned()),
+            0 => None,
+            _ => Some(tt::TokenTree::Subtree(tt::Subtree {
+                delimiter: None,
+                token_trees: res.into_iter().map(|it| it.cloned()).collect(),
+            })),
+        };
+        ExpandResult { value: res, err }
+    }
+
     pub(crate) fn peek_n(&self, n: usize) -> Option<&tt::TokenTree> {
         self.inner.as_slice().get(n)
     }
diff --git a/crates/parser/src/grammar.rs b/crates/parser/src/grammar.rs
index 86dce61da2d..1c58f217a30 100644
--- a/crates/parser/src/grammar.rs
+++ b/crates/parser/src/grammar.rs
@@ -86,10 +86,6 @@ pub(crate) mod entry_points {
         expressions::stmt(p, expressions::StmtWithSemi::Optional, false);
     }
 
-    pub(crate) fn visibility(p: &mut Parser) {
-        let _ = opt_visibility(p, false);
-    }
-
     // Parse a meta item , which excluded [], e.g : #[ MetaItem ]
     pub(crate) fn meta_item(p: &mut Parser) {
         attributes::meta(p);
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index f0b93c4511a..f4ce5a21e80 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -49,12 +49,13 @@ pub use crate::{
 /// `Option<Output>`. The way MBE work, by the time we *try* to parse `$e:expr`
 /// we already commit to expression. In other words, this API by design can't be
 /// used to implement "rollback and try another alternative" logic.
+#[derive(Debug)]
 pub enum PrefixEntryPoint {
     Vis,
 }
 
 impl PrefixEntryPoint {
-    pub fn parse(self, input: &Input) -> Output {
+    pub fn parse(&self, input: &Input) -> Output {
         let entry_point: fn(&'_ mut parser::Parser) = match self {
             PrefixEntryPoint::Vis => grammar::entry::prefix::vis,
         };
@@ -80,7 +81,7 @@ pub enum ParserEntryPoint {
     Pattern,
     Item,
     Block,
-    Visibility,
+    // Visibility,
     MetaItem,
     Items,
     Statements,
@@ -109,7 +110,7 @@ pub fn parse(inp: &Input, entry_point: ParserEntryPoint) -> Output {
         ParserEntryPoint::Pattern => grammar::entry_points::pattern,
         ParserEntryPoint::Item => grammar::entry_points::item,
         ParserEntryPoint::Block => grammar::entry_points::block_expr,
-        ParserEntryPoint::Visibility => grammar::entry_points::visibility,
+        // ParserEntryPoint::Visibility => grammar::entry_points::visibility,
         ParserEntryPoint::MetaItem => grammar::entry_points::meta_item,
         ParserEntryPoint::Statement => grammar::entry_points::stmt,
         ParserEntryPoint::StatementOptionalSemi => grammar::entry_points::stmt_optional_semi,