1 files changed, 54 insertions, 64 deletions

diff --git a/crates/ra_mbe/src/mbe_expander.rs b/crates/ra_mbe/src/mbe_expander.rs
index 01641fdee..78df96880 100644
--- a/crates/ra_mbe/src/mbe_expander.rs
+++ b/crates/ra_mbe/src/mbe_expander.rs
@@ -81,21 +81,26 @@ struct Bindings {
 
 #[derive(Debug)]
 enum Binding {
-    Simple(tt::TokenTree),
+    Fragment(Fragment),
     Nested(Vec<Binding>),
     Empty,
 }
 
+#[derive(Debug, Clone)]
+enum Fragment {
+    /// token fragments are just copy-pasted into the output
+    Tokens(tt::TokenTree),
+    /// Ast fragments are inserted with fake delimiters, so as to make things
+    /// like `$i * 2` where `$i = 1 + 1` work as expectd.
+    Ast(tt::TokenTree),
+}
+
 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(),
-            ),
-        );
+        let tt = tt::Subtree { delimiter: tt::Delimiter::None, token_trees: vec![] }.into();
+        self.inner.insert(name.clone(), Binding::Fragment(Fragment::Tokens(tt)));
     }
 
     fn push_empty(&mut self, name: &SmolStr) {
@@ -106,13 +111,13 @@ impl Bindings {
         self.inner.contains_key(name)
     }
 
-    fn get(&self, name: &SmolStr, nesting: &[usize]) -> Result<&tt::TokenTree, ExpandError> {
+    fn get(&self, name: &SmolStr, nesting: &[usize]) -> Result<&Fragment, ExpandError> {
         let mut b = self.inner.get(name).ok_or_else(|| {
             ExpandError::BindingError(format!("could not find binding `{}`", name))
         })?;
         for &idx in nesting.iter() {
             b = match b {
-                Binding::Simple(_) => break,
+                Binding::Fragment(_) => break,
                 Binding::Nested(bs) => bs.get(idx).ok_or_else(|| {
                     ExpandError::BindingError(format!("could not find nested binding `{}`", name))
                 })?,
@@ -125,7 +130,7 @@ impl Bindings {
             };
         }
         match b {
-            Binding::Simple(it) => Ok(it),
+            Binding::Fragment(it) => Ok(it),
             Binding::Nested(_) => Err(ExpandError::BindingError(format!(
                 "expected simple binding, found nested binding `{}`",
                 name
@@ -195,8 +200,8 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
                 crate::Leaf::Var(crate::Var { text, kind }) => {
                     let kind = kind.as_ref().ok_or(ExpandError::UnexpectedToken)?;
                     match match_meta_var(kind.as_str(), input)? {
-                        Some(tt) => {
-                            res.inner.insert(text.clone(), Binding::Simple(tt));
+                        Some(fragment) => {
+                            res.inner.insert(text.clone(), Binding::Fragment(fragment));
                         }
                         None => res.push_optional(text),
                     }
@@ -292,7 +297,7 @@ fn match_lhs(pattern: &crate::Subtree, input: &mut TtCursor) -> Result<Bindings,
     Ok(res)
 }
 
-fn match_meta_var(kind: &str, input: &mut TtCursor) -> Result<Option<tt::TokenTree>, ExpandError> {
+fn match_meta_var(kind: &str, input: &mut TtCursor) -> Result<Option<Fragment>, ExpandError> {
     let fragment = match kind {
         "path" => Path,
         "expr" => Expr,
@@ -303,7 +308,7 @@ fn match_meta_var(kind: &str, input: &mut TtCursor) -> Result<Option<tt::TokenTr
         "meta" => MetaItem,
         "item" => Item,
         _ => {
-            let binding = match kind {
+            let tt = match kind {
                 "ident" => {
                     let ident = input.eat_ident().ok_or(ExpandError::UnexpectedToken)?.clone();
                     tt::Leaf::from(ident).into()
@@ -321,11 +326,12 @@ fn match_meta_var(kind: &str, input: &mut TtCursor) -> Result<Option<tt::TokenTr
                 },
                 _ => return Err(ExpandError::UnexpectedToken),
             };
-            return Ok(Some(binding));
+            return Ok(Some(Fragment::Tokens(tt)));
         }
     };
-    let binding = input.eat_fragment(fragment).ok_or(ExpandError::UnexpectedToken)?;
-    Ok(Some(binding))
+    let tt = input.eat_fragment(fragment).ok_or(ExpandError::UnexpectedToken)?;
+    let fragment = if kind == "expr" { Fragment::Ast(tt) } else { Fragment::Tokens(tt) };
+    Ok(Some(fragment))
 }
 
 #[derive(Debug)]
@@ -339,45 +345,20 @@ fn expand_subtree(
     template: &crate::Subtree,
     ctx: &mut ExpandCtx,
 ) -> Result<tt::Subtree, ExpandError> {
-    let token_trees = template
-        .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 })
-}
-
-/// Reduce single token subtree to single token
-/// In `tt` matcher case, all tt tokens will be braced by a Delimiter::None
-/// which makes all sort of problems.
-fn reduce_single_token(mut subtree: tt::Subtree) -> tt::TokenTree {
-    if subtree.delimiter != tt::Delimiter::None || subtree.token_trees.len() != 1 {
-        return subtree.into();
+    let mut buf: Vec<tt::TokenTree> = Vec::new();
+    for tt in template.token_trees.iter() {
+        let tt = expand_tt(tt, ctx)?;
+        push_fragment(&mut buf, tt);
     }
 
-    match subtree.token_trees.pop().unwrap() {
-        tt::TokenTree::Subtree(subtree) => reduce_single_token(subtree),
-        tt::TokenTree::Leaf(token) => token.into(),
-    }
+    Ok(tt::Subtree { delimiter: template.delimiter, token_trees: buf })
 }
 
-fn expand_tt(
-    template: &crate::TokenTree,
-    ctx: &mut ExpandCtx,
-) -> Result<tt::TokenTree, ExpandError> {
+fn expand_tt(template: &crate::TokenTree, ctx: &mut ExpandCtx) -> Result<Fragment, ExpandError> {
     let res: tt::TokenTree = match template {
         crate::TokenTree::Subtree(subtree) => expand_subtree(subtree, ctx)?.into(),
         crate::TokenTree::Repeat(repeat) => {
-            let mut token_trees: Vec<tt::TokenTree> = Vec::new();
+            let mut buf: Vec<tt::TokenTree> = Vec::new();
             ctx.nesting.push(0);
             // Dirty hack to make macro-expansion terminate.
             // This should be replaced by a propper macro-by-example implementation
@@ -418,23 +399,23 @@ fn expand_tt(
 
                 let idx = ctx.nesting.pop().unwrap();
                 ctx.nesting.push(idx + 1);
-                token_trees.push(reduce_single_token(t));
+                push_subtree(&mut buf, t);
 
                 if let Some(ref sep) = repeat.separator {
                     match sep {
                         crate::Separator::Ident(ident) => {
                             has_seps = 1;
-                            token_trees.push(tt::Leaf::from(ident.clone()).into());
+                            buf.push(tt::Leaf::from(ident.clone()).into());
                         }
                         crate::Separator::Literal(lit) => {
                             has_seps = 1;
-                            token_trees.push(tt::Leaf::from(lit.clone()).into());
+                            buf.push(tt::Leaf::from(lit.clone()).into());
                         }
 
                         crate::Separator::Puncts(puncts) => {
                             has_seps = puncts.len();
                             for punct in puncts {
-                                token_trees.push(tt::Leaf::from(*punct).into());
+                                buf.push(tt::Leaf::from(*punct).into());
                             }
                         }
                     }
@@ -450,16 +431,16 @@ fn expand_tt(
 
             ctx.nesting.pop().unwrap();
             for _ in 0..has_seps {
-                token_trees.pop();
+                buf.pop();
             }
 
             if crate::RepeatKind::OneOrMore == repeat.kind && counter == 0 {
                 return Err(ExpandError::UnexpectedToken);
             }
 
-            // Check if it is a singel token subtree without any delimiter
+            // Check if it is a single token subtree without any delimiter
             // e.g {Delimiter:None> ['>'] /Delimiter:None>}
-            reduce_single_token(tt::Subtree { token_trees, delimiter: tt::Delimiter::None })
+            tt::Subtree { delimiter: tt::Delimiter::None, token_trees: buf }.into()
         }
         crate::TokenTree::Leaf(leaf) => match leaf {
             crate::Leaf::Ident(ident) => {
@@ -500,20 +481,15 @@ fn expand_tt(
                     }
                     .into()
                 } else {
-                    let tkn = ctx.bindings.get(&v.text, &ctx.nesting)?.clone();
+                    let fragment = ctx.bindings.get(&v.text, &ctx.nesting)?.clone();
                     ctx.var_expanded = true;
-
-                    if let tt::TokenTree::Subtree(subtree) = tkn {
-                        reduce_single_token(subtree)
-                    } else {
-                        tkn
-                    }
+                    return Ok(fragment);
                 }
             }
             crate::Leaf::Literal(l) => tt::Leaf::from(tt::Literal { text: l.text.clone() }).into(),
         },
     };
-    Ok(res)
+    Ok(Fragment::Tokens(res))
 }
 
 #[cfg(test)]
@@ -586,3 +562,17 @@ mod tests {
         expand_rule(&rules.rules[0], &invocation_tt)
     }
 }
+
+fn push_fragment(buf: &mut Vec<tt::TokenTree>, fragment: Fragment) {
+    match fragment {
+        Fragment::Tokens(tt::TokenTree::Subtree(tt)) => push_subtree(buf, tt),
+        Fragment::Tokens(tt) | Fragment::Ast(tt) => buf.push(tt),
+    }
+}
+
+fn push_subtree(buf: &mut Vec<tt::TokenTree>, tt: tt::Subtree) {
+    match tt.delimiter {
+        tt::Delimiter::None => buf.extend(tt.token_trees),
+        _ => buf.push(tt.into()),
+    }
+}