macros: Add remaining context and improve parsing macro dispatch

This allows us to expand macor invocations in more places, as macro
calls are not limited to statements or expressions. It is quite common
to use macros to abstract writing repetitive boilerplate for type
implementations, for example.

4 files changed, 182 insertions, 69 deletions

diff --git a/gcc/rust/expand/rust-attribute-visitor.cc b/gcc/rust/expand/rust-attribute-visitor.cc
index 255e9ba..50821ed 100644
--- a/gcc/rust/expand/rust-attribute-visitor.cc
+++ b/gcc/rust/expand/rust-attribute-visitor.cc
@@ -1114,8 +1114,6 @@ AttrVisitor::visit (AST::ClosureExprInner &expr)
 void
 AttrVisitor::visit (AST::BlockExpr &expr)
 {
-  expander.push_context (MacroExpander::BLOCK);
-
   // initial strip test based on outer attrs
   expander.expand_cfg_attrs (expr.get_outer_attrs ());
   if (expander.fails_cfg_with_expand (expr.get_outer_attrs ()))
@@ -1135,31 +1133,13 @@ AttrVisitor::visit (AST::BlockExpr &expr)
       return;
     }
 
-  // strip all statements
-  auto &stmts = expr.get_statements ();
-  for (auto it = stmts.begin (); it != stmts.end ();)
-    {
-      auto &stmt = *it;
-
-      stmt->accept_vis (*this);
+  std::function<std::unique_ptr<AST::Stmt> (AST::SingleASTNode)> extractor
+    = [] (AST::SingleASTNode node) { return node.take_stmt (); };
 
-      auto fragment = expander.take_expanded_fragment (*this);
-      if (fragment.should_expand ())
-	{
-	  // Remove the current expanded invocation
-	  it = stmts.erase (it);
-	  for (auto &node : fragment.get_nodes ())
-	    {
-	      it = stmts.insert (it, node.take_stmt ());
-	      it++;
-	    }
-	}
+  expand_macro_children (MacroExpander::BLOCK, expr.get_statements (),
+			 extractor);
 
-      else if (stmt->is_marked_for_strip ())
-	it = stmts.erase (it);
-      else
-	it++;
-    }
+  expander.push_context (MacroExpander::BLOCK);
 
   // strip tail expression if exists - can actually fully remove it
   if (expr.has_tail_expr ())
@@ -2489,8 +2469,11 @@ AttrVisitor::visit (AST::Trait &trait)
   if (trait.has_where_clause ())
     expand_where_clause (trait.get_where_clause ());
 
-  // strip trait items if required
-  expand_pointer_allow_strip (trait.get_trait_items ());
+  std::function<std::unique_ptr<AST::TraitItem> (AST::SingleASTNode)> extractor
+    = [] (AST::SingleASTNode node) { return node.take_trait_item (); };
+
+  expand_macro_children (MacroExpander::TRAIT, trait.get_trait_items (),
+			 extractor);
 }
 void
 AttrVisitor::visit (AST::InherentImpl &impl)
@@ -2523,8 +2506,11 @@ AttrVisitor::visit (AST::InherentImpl &impl)
   if (impl.has_where_clause ())
     expand_where_clause (impl.get_where_clause ());
 
-  // strip inherent impl items if required
-  expand_pointer_allow_strip (impl.get_impl_items ());
+  std::function<std::unique_ptr<AST::InherentImplItem> (AST::SingleASTNode)>
+    extractor = [] (AST::SingleASTNode node) { return node.take_impl_item (); };
+
+  expand_macro_children (MacroExpander::IMPL, impl.get_impl_items (),
+			 extractor);
 }
 void
 AttrVisitor::visit (AST::TraitImpl &impl)
@@ -2659,8 +2645,12 @@ AttrVisitor::visit (AST::ExternBlock &block)
       return;
     }
 
-  // strip external items if required
-  expand_pointer_allow_strip (block.get_extern_items ());
+  std::function<std::unique_ptr<AST::ExternalItem> (AST::SingleASTNode)>
+    extractor
+    = [] (AST::SingleASTNode node) { return node.take_external_item (); };
+
+  expand_macro_children (MacroExpander::EXTERN, block.get_extern_items (),
+			 extractor);
 }
 
 // I don't think it would be possible to strip macros without expansion
diff --git a/gcc/rust/expand/rust-attribute-visitor.h b/gcc/rust/expand/rust-attribute-visitor.h
index f6327d6..6da6583 100644
--- a/gcc/rust/expand/rust-attribute-visitor.h
+++ b/gcc/rust/expand/rust-attribute-visitor.h
@@ -40,8 +40,24 @@ public:
   void expand_trait_function_decl (AST::TraitFunctionDecl &decl);
   void expand_trait_method_decl (AST::TraitMethodDecl &decl);
 
-  template <typename T> void expand_pointer_allow_strip (T &values)
+  /**
+   * Expand a set of values, erasing them if they are marked for strip, and
+   * replacing them with expanded macro nodes if necessary.
+   * This function is slightly different from `expand_pointer_allow_strip` as
+   * it can only be called in certain expansion contexts - where macro
+   * invocations are allowed.
+   *
+   * @param ctx Context to use for macro expansion
+   * @param values Iterable reference over values to replace or erase
+   * @param extractor Function to call when replacing values with the content
+   * 		of an expanded AST node
+   */
+  template <typename T, typename U>
+  void expand_macro_children (MacroExpander::ContextType ctx, T &values,
+			      std::function<U (AST::SingleASTNode)> extractor)
   {
+    expander.push_context (ctx);
+
     for (auto it = values.begin (); it != values.end ();)
       {
 	auto &value = *it;
@@ -49,10 +65,45 @@ public:
 	// mark for stripping if required
 	value->accept_vis (*this);
 
+	auto fragment = expander.take_expanded_fragment (*this);
+	if (fragment.should_expand ())
+	  {
+	    it = values.erase (it);
+	    for (auto &node : fragment.get_nodes ())
+	      {
+		it = values.insert (it, extractor (node));
+		it++;
+	      }
+	  }
+	else if (value->is_marked_for_strip ())
+	  {
+	    it = values.erase (it);
+	  }
+	else
+	  {
+	    ++it;
+	  }
+      }
+
+    expander.pop_context ();
+  }
+
+  template <typename T> void expand_pointer_allow_strip (T &values)
+  {
+    for (auto it = values.begin (); it != values.end ();)
+      {
+	auto &value = *it;
+
+	// mark for stripping if required
+	value->accept_vis (*this);
 	if (value->is_marked_for_strip ())
-	  it = values.erase (it);
+	  {
+	    it = values.erase (it);
+	  }
 	else
-	  ++it;
+	  {
+	    ++it;
+	  }
       }
   }
 
diff --git a/gcc/rust/expand/rust-macro-expand.cc b/gcc/rust/expand/rust-macro-expand.cc
index 3389f20..80f822b 100644
--- a/gcc/rust/expand/rust-macro-expand.cc
+++ b/gcc/rust/expand/rust-macro-expand.cc
@@ -786,7 +786,6 @@ parse_many (Parser<MacroInvocLexer> &parser, TokenId &delimiter,
 	    std::function<AST::SingleASTNode ()> parse_fn)
 {
   std::vector<AST::SingleASTNode> nodes;
-
   while (true)
     {
       if (parser.peek_current_token ()->get_id () == delimiter)
@@ -815,6 +814,52 @@ transcribe_many_items (Parser<MacroInvocLexer> &parser, TokenId &delimiter)
 }
 
 /**
+ * Transcribe 0 or more external items from a macro invocation
+ *
+ * @param parser Parser to extract items from
+ * @param delimiter Id of the token on which parsing should stop
+ */
+static std::vector<AST::SingleASTNode>
+transcribe_many_ext (Parser<MacroInvocLexer> &parser, TokenId &delimiter)
+{
+  return parse_many (parser, delimiter, [&parser] () {
+    auto item = parser.parse_external_item ();
+    return AST::SingleASTNode (std::move (item));
+  });
+}
+
+/**
+ * Transcribe 0 or more trait items from a macro invocation
+ *
+ * @param parser Parser to extract items from
+ * @param delimiter Id of the token on which parsing should stop
+ */
+static std::vector<AST::SingleASTNode>
+transcribe_many_trait_items (Parser<MacroInvocLexer> &parser,
+			     TokenId &delimiter)
+{
+  return parse_many (parser, delimiter, [&parser] () {
+    auto item = parser.parse_trait_item ();
+    return AST::SingleASTNode (std::move (item));
+  });
+}
+
+/**
+ * Transcribe 0 or more impl items from a macro invocation
+ *
+ * @param parser Parser to extract items from
+ * @param delimiter Id of the token on which parsing should stop
+ */
+static std::vector<AST::SingleASTNode>
+transcribe_many_impl_items (Parser<MacroInvocLexer> &parser, TokenId &delimiter)
+{
+  return parse_many (parser, delimiter, [&parser] () {
+    auto item = parser.parse_inherent_impl_item ();
+    return AST::SingleASTNode (std::move (item));
+  });
+}
+
+/**
  * Transcribe 0 or more statements from a macro invocation
  *
  * @param parser Parser to extract statements from
@@ -845,6 +890,57 @@ transcribe_expression (Parser<MacroInvocLexer> &parser)
   return {AST::SingleASTNode (std::move (expr))};
 }
 
+static std::vector<AST::SingleASTNode>
+transcribe_on_delimiter (Parser<MacroInvocLexer> &parser, bool semicolon,
+			 AST::DelimType delimiter, TokenId last_token_id)
+{
+  if (semicolon || delimiter == AST::DelimType::CURLY)
+    return transcribe_many_stmts (parser, last_token_id);
+  else
+    return transcribe_expression (parser);
+} // namespace Rust
+
+static std::vector<AST::SingleASTNode>
+transcribe_context (MacroExpander::ContextType ctx,
+		    Parser<MacroInvocLexer> &parser, bool semicolon,
+		    AST::DelimType delimiter, TokenId last_token_id)
+{
+  // The flow-chart in order to choose a parsing function is as follows:
+  //
+  // [switch special context]
+  //     -- Item --> parser.parse_item();
+  //     -- Trait --> parser.parse_trait_item();
+  //     -- Impl --> parser.parse_impl_item();
+  //     -- Extern --> parser.parse_extern_item();
+  //     -- None --> [has semicolon?]
+  //                 -- Yes --> parser.parse_stmt();
+  //                 -- No --> [switch invocation.delimiter()]
+  //                             -- { } --> parser.parse_stmt();
+  //                             -- _ --> parser.parse_expr(); // once!
+
+  // If there is a semicolon OR we are expanding a MacroInvocationSemi, then
+  // we can parse multiple items. Otherwise, parse *one* expression
+
+  switch (ctx)
+    {
+    case MacroExpander::ContextType::ITEM:
+      return transcribe_many_items (parser, last_token_id);
+      break;
+    case MacroExpander::ContextType::TRAIT:
+      return transcribe_many_trait_items (parser, last_token_id);
+      break;
+    case MacroExpander::ContextType::IMPL:
+      return transcribe_many_impl_items (parser, last_token_id);
+      break;
+    case MacroExpander::ContextType::EXTERN:
+      return transcribe_many_ext (parser, last_token_id);
+      break;
+    default:
+      return transcribe_on_delimiter (parser, semicolon, delimiter,
+				      last_token_id);
+    }
+}
+
 AST::ASTFragment
 MacroExpander::transcribe_rule (
   AST::MacroRule &match_rule, AST::DelimTokenTree &invoc_token_tree,
@@ -864,18 +960,15 @@ MacroExpander::transcribe_rule (
   std::vector<std::unique_ptr<AST::Token>> substituted_tokens
     = substitute_context.substitute_tokens ();
 
+  // handy for debugging
+  // for (auto &tok : substituted_tokens)
+  //   rust_debug ("[tok] %s", tok->as_string ().c_str ());
+
   // parse it to an ASTFragment
   MacroInvocLexer lex (std::move (substituted_tokens));
   Parser<MacroInvocLexer> parser (std::move (lex));
 
-  // handy for debugging
-  // for (auto &tok : substituted_tokens)
-  //   {
-  //     rust_debug ("tok: [%s]", tok->as_string ().c_str ());
-  //   }
-
   auto last_token_id = TokenId::RIGHT_CURLY;
-  std::vector<AST::SingleASTNode> nodes;
 
   // this is used so we can check that we delimit the stream correctly.
   switch (transcribe_tree.get_delim_type ())
@@ -905,33 +998,9 @@ MacroExpander::transcribe_rule (
   //   as a statement (either via ExpressionStatement or
   //   MacroInvocationWithSemi)
 
-  // The flow-chart in order to choose a parsing function is as follows:
-  //
-  // [is in item context?]
-  //     -- Yes --> parser.parse_item();
-  //     -- No --> [has semicolon?]
-  //                 -- Yes --> parser.parse_stmt();
-  //                 -- No --> [switch invocation.delimiter()]
-  //                             -- { } --> parser.parse_stmt();
-  //                             -- _ --> parser.parse_expr();
-
-  // If there is a semicolon OR we are expanding a MacroInvocationSemi, then
-  // we can parse multiple items. Otherwise, parse *one* expression
-
-  if (ctx == ContextType::ITEM)
-    nodes = transcribe_many_items (parser, last_token_id);
-  else if (semicolon)
-    nodes = transcribe_many_stmts (parser, last_token_id);
-  else
-    switch (invoc_token_tree.get_delim_type ())
-      {
-      case AST::CURLY:
-	nodes = transcribe_many_stmts (parser, last_token_id);
-	break;
-      default:
-	nodes = transcribe_expression (parser);
-	break;
-      }
+  auto nodes
+    = transcribe_context (ctx, parser, semicolon,
+			  invoc_token_tree.get_delim_type (), last_token_id);
 
   // emit any errors
   if (parser.has_errors ())
diff --git a/gcc/rust/expand/rust-macro-expand.h b/gcc/rust/expand/rust-macro-expand.h
index 0d0282e..61b69e4 100644
--- a/gcc/rust/expand/rust-macro-expand.h
+++ b/gcc/rust/expand/rust-macro-expand.h
@@ -187,6 +187,9 @@ struct MacroExpander
   {
     ITEM,
     BLOCK,
+    TRAIT,
+    IMPL,
+    EXTERN,
   };
 
   ExpansionCfg cfg;