1 files changed, 123 insertions, 29 deletions

diff --git a/crates/ra_mbe/src/mbe_expander.rs b/crates/ra_mbe/src/mbe_expander.rs
index d5189b537..4b007647c 100644
--- a/crates/ra_mbe/src/mbe_expander.rs
+++ b/crates/ra_mbe/src/mbe_expander.rs
@@ -81,9 +81,29 @@ struct Bindings {
 enum Binding {
     Simple(tt::TokenTree),
     Nested(Vec<Binding>),
+    Empty,
 }
 
 impl Bindings {
+    fn push_optional(&mut self, name: &SmolStr) {
+        // FIXME: Do we have a better way to represent an empty token ?
+        // Insert an empty subtree for empty token
+        self.inner.insert(
+            name.clone(),
+            Binding::Simple(
+                tt::Subtree { delimiter: tt::Delimiter::None, token_trees: vec![] }.into(),
+            ),
+        );
+    }
+
+    fn push_empty(&mut self, name: &SmolStr) {
+        self.inner.insert(name.clone(), Binding::Empty);
+    }
+
+    fn contains(&self, name: &SmolStr) -> bool {
+        self.inner.contains_key(name)
+    }
+
     fn get(&self, name: &SmolStr, nesting: &[usize]) -> Result<&tt::TokenTree, ExpandError> {
         let mut b = self
             .inner
@@ -96,6 +116,12 @@ impl Bindings {
                     "could not find nested binding `{}`",
                     name
                 )))?,
+                Binding::Empty => {
+                    return Err(ExpandError::BindingError(format!(
+                        "could not find empty binding `{}`",
+                        name
+                    )))
+                }
             };
         }
         match b {
@@ -104,16 +130,26 @@ impl Bindings {
                 "expected simple binding, found nested binding `{}`",
                 name
             ))),
+            Binding::Empty => Err(ExpandError::BindingError(format!(
+                "expected simple binding, found empty binding `{}`",
+                name
+            ))),
         }
     }
 
-    fn push_nested(&mut self, nested: Bindings) -> Result<(), ExpandError> {
+    fn push_nested(&mut self, idx: usize, nested: Bindings) -> Result<(), ExpandError> {
         for (key, value) in nested.inner {
             if !self.inner.contains_key(&key) {
                 self.inner.insert(key.clone(), Binding::Nested(Vec::new()));
             }
             match self.inner.get_mut(&key) {
-                Some(Binding::Nested(it)) => it.push(value),
+                Some(Binding::Nested(it)) => {
+                    // insert empty nested bindings before this one
+                    while it.len() < idx {
+                        it.push(Binding::Nested(vec![]));
+                    }
+                    it.push(value);
+                }
                 _ => {
                     return Err(ExpandError::BindingError(format!(
                         "could not find binding `{}`",
@@ -130,6 +166,24 @@ impl Bindings {
     }
 }
 
+fn collect_vars(subtree: &crate::Subtree) -> Vec<SmolStr> {
+    let mut res = vec![];
+
+    for tkn in subtree.token_trees.iter() {
+        match tkn {
+            crate::TokenTree::Leaf(crate::Leaf::Var(crate::Var { text, .. })) => {
+                res.push(text.clone());
+            }
+            crate::TokenTree::Subtree(subtree) => {
+                res.extend(collect_vars(subtree));
+            }
+            _ => {}
+        }
+    }
+
+    res
+}
+
 fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings, ExpandError> {
     let mut res = Bindings::default();
     for pat in pattern.token_trees.iter() {
@@ -178,10 +232,6 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                                 input.eat_meta().ok_or(ExpandError::UnexpectedToken)?.clone();
                             res.inner.insert(text.clone(), Binding::Simple(meta.into()));
                         }
-                        // FIXME:
-                        // Enable followiing code when everything is fixed
-                        // At least we can dogfood itself to not stackoverflow
-                        //
                         "tt" => {
                             let token = input.eat().ok_or(ExpandError::UnexpectedToken)?.clone();
                             res.inner.insert(text.clone(), Binding::Simple(token.into()));
@@ -206,8 +256,13 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                             );
                         }
                         "vis" => {
-                            let vis = input.eat_vis().ok_or(ExpandError::UnexpectedToken)?.clone();
-                            res.inner.insert(text.clone(), Binding::Simple(vis.into()));
+                            // `vis` is optional
+                            if let Some(vis) = input.try_eat_vis() {
+                                let vis = vis.clone();
+                                res.inner.insert(text.clone(), Binding::Simple(vis.into()));
+                            } else {
+                                res.push_optional(&text);
+                            }
                         }
 
                         _ => return Err(ExpandError::UnexpectedToken),
@@ -236,7 +291,6 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                 loop {
                     match match_lhs(subtree, input) {
                         Ok(nested) => {
-                            counter += 1;
                             limit -= 1;
                             if limit == 0 {
                                 log::warn!("match_lhs excced in repeat pattern exceed limit => {:#?}\n{:#?}\n{:#?}\n{:#?}", subtree, input, kind, separator);
@@ -244,7 +298,8 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                             }
 
                             memento = input.save();
-                            res.push_nested(nested)?;
+                            res.push_nested(counter, nested)?;
+                            counter += 1;
                             if counter == 1 {
                                 if let crate::RepeatKind::ZeroOrOne = kind {
                                     break;
@@ -252,20 +307,9 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                             }
 
                             if let Some(separator) = separator {
-                                use crate::Separator::*;
-
                                 if !input
                                     .eat_seperator()
-                                    .map(|sep| match (sep, separator) {
-                                        (Ident(ref a), Ident(ref b)) => a.text == b.text,
-                                        (Literal(ref a), Literal(ref b)) => a.text == b.text,
-                                        (Puncts(ref a), Puncts(ref b)) if a.len() == b.len() => {
-                                            let a_iter = a.iter().map(|a| a.char);
-                                            let b_iter = b.iter().map(|b| b.char);
-                                            a_iter.eq(b_iter)
-                                        }
-                                        _ => false,
-                                    })
+                                    .map(|sep| sep == *separator)
                                     .unwrap_or(false)
                                 {
                                     input.rollback(memento);
@@ -284,6 +328,10 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                     crate::RepeatKind::OneOrMore if counter == 0 => {
                         return Err(ExpandError::UnexpectedToken);
                     }
+                    _ if counter == 0 => {
+                        // Collect all empty variables in subtrees
+                        collect_vars(subtree).iter().for_each(|s| res.push_empty(s));
+                    }
                     _ => {}
                 }
             }
@@ -322,6 +370,14 @@ fn expand_subtree(
         .token_trees
         .iter()
         .map(|it| expand_tt(it, ctx))
+        .filter(|it| {
+            // Filter empty subtree
+            if let Ok(tt::TokenTree::Subtree(subtree)) = it {
+                subtree.delimiter != tt::Delimiter::None || !subtree.token_trees.is_empty()
+            } else {
+                true
+            }
+        })
         .collect::<Result<Vec<_>, ExpandError>>()?;
 
     Ok(tt::Subtree { token_trees, delimiter: template.delimiter })
@@ -356,14 +412,23 @@ fn expand_tt(
             let mut has_seps = 0;
             let mut counter = 0;
 
+            // We store the old var expanded value, and restore it later
+            // It is because before this `$repeat`,
+            // it is possible some variables already expanad in the same subtree
+            //
+            // `some_var_expanded` keep check if the deeper subtree has expanded variables
             let mut some_var_expanded = false;
+            let old_var_expanded = ctx.var_expanded;
             ctx.var_expanded = false;
 
             while let Ok(t) = expand_subtree(&repeat.subtree, ctx) {
-                // if no var expaned in the child, we count it as a fail
+                // if no var expanded in the child, we count it as a fail
                 if !ctx.var_expanded {
                     break;
                 }
+
+                // Reset `ctx.var_expandeded` to see if there is other expanded variable
+                // in the next matching
                 some_var_expanded = true;
                 ctx.var_expanded = false;
 
@@ -407,7 +472,8 @@ fn expand_tt(
                 }
             }
 
-            ctx.var_expanded = some_var_expanded;
+            // Restore the `var_expanded` by combining old one and the new one
+            ctx.var_expanded = some_var_expanded || old_var_expanded;
 
             ctx.nesting.pop().unwrap();
             for _ in 0..has_seps {
@@ -433,6 +499,33 @@ fn expand_tt(
                     // FIXME: Properly handle $crate token
                     tt::Leaf::from(tt::Ident { text: "$crate".into(), id: TokenId::unspecified() })
                         .into()
+                } else if !ctx.bindings.contains(&v.text) {
+                    // Note that it is possible to have a `$var` inside a macro which is not bound.
+                    // For example:
+                    // ```
+                    // macro_rules! foo {
+                    //     ($a:ident, $b:ident, $c:tt) => {
+                    //         macro_rules! bar {
+                    //             ($bi:ident) => {
+                    //                 fn $bi() -> u8 {$c}
+                    //             }
+                    //         }
+                    //     }
+                    // ```
+                    // We just treat it a normal tokens
+                    tt::Subtree {
+                        delimiter: tt::Delimiter::None,
+                        token_trees: vec![
+                            tt::Leaf::from(tt::Punct { char: '$', spacing: tt::Spacing::Alone })
+                                .into(),
+                            tt::Leaf::from(tt::Ident {
+                                text: v.text.clone(),
+                                id: TokenId::unspecified(),
+                            })
+                            .into(),
+                        ],
+                    }
+                    .into()
                 } else {
                     let tkn = ctx.bindings.get(&v.text, &ctx.nesting)?.clone();
                     ctx.var_expanded = true;
@@ -459,11 +552,12 @@ mod tests {
 
     #[test]
     fn test_expand_rule() {
-        assert_err(
-            "($i:ident) => ($j)",
-            "foo!{a}",
-            ExpandError::BindingError(String::from("could not find binding `j`")),
-        );
+        // FIXME: The missing $var check should be in parsing phase
+        // assert_err(
+        //     "($i:ident) => ($j)",
+        //     "foo!{a}",
+        //     ExpandError::BindingError(String::from("could not find binding `j`")),
+        // );
 
         assert_err(
             "($($i:ident);*) => ($i)",