substitute_repetition: Correctly insert sub-fragments

2 files changed, 119 insertions, 71 deletions

diff --git a/gcc/rust/expand/rust-macro-expand.cc b/gcc/rust/expand/rust-macro-expand.cc
index 45d2b74..b54aa01 100644
--- a/gcc/rust/expand/rust-macro-expand.cc
+++ b/gcc/rust/expand/rust-macro-expand.cc
@@ -3118,7 +3118,7 @@ MacroExpander::expand_decl_macro (Location invoc_locus,
 
   // find matching arm
   AST::MacroRule *matched_rule = nullptr;
-  std::map<std::string, MatchedFragment> matched_fragments;
+  std::map<std::string, std::vector<MatchedFragment>> matched_fragments;
   for (auto &rule : rules_def.get_rules ())
     {
       sub_stack.push ();
@@ -3127,9 +3127,9 @@ MacroExpander::expand_decl_macro (Location invoc_locus,
 
       if (did_match_rule)
 	{
-	  for (auto &frag : matched_fragments)
-	    rust_debug ("matched fragment: %s",
-			frag.second.as_string ().c_str ());
+	  for (auto &kv : matched_fragments)
+	    rust_debug ("[fragment]: %s (%ld)", kv.first.c_str (),
+			kv.second.size ());
 
 	  matched_rule = &rule;
 	  break;
@@ -3535,9 +3535,8 @@ MacroExpander::match_matcher (Parser<MacroInvocLexer> &parser,
 
 	    // matched fragment get the offset in the token stream
 	    size_t offs_end = source.get_offs ();
-	    sub_stack.peek ().insert (
-	      {fragment->get_ident (),
-	       MatchedFragment (fragment->get_ident (), offs_begin, offs_end)});
+	    sub_stack.insert_fragment (
+	      MatchedFragment (fragment->get_ident (), offs_begin, offs_end));
 	  }
 	  break;
 
@@ -3632,10 +3631,9 @@ MacroExpander::match_n_matches (
 
 		// matched fragment get the offset in the token stream
 		size_t offs_end = source.get_offs ();
-		sub_stack.peek ().insert (
-		  {fragment->get_ident (),
-		   MatchedFragment (fragment->get_ident (), offs_begin,
-				    offs_end)});
+		sub_stack.insert_fragment (
+		  MatchedFragment (fragment->get_ident (), offs_begin,
+				   offs_end));
 	      }
 	      break;
 
@@ -3729,17 +3727,19 @@ MacroExpander::match_repetition (Parser<MacroInvocLexer> &parser,
 	  auto fragment = static_cast<AST::MacroMatchFragment *> (match.get ());
 	  auto it = stack_map.find (fragment->get_ident ());
 
-	  // If we can't find the fragment, but the result was valid, then it's
-	  // a zero-matched fragment and we can insert it
+	  // If we can't find the fragment, but the result was valid, then
+	  // it's a zero-matched fragment and we can insert it
 	  if (it == stack_map.end ())
 	    {
-	      stack_map.insert (
-		{fragment->get_ident (),
-		 MatchedFragment::zero (fragment->get_ident ())});
+	      sub_stack.insert_fragment (
+		MatchedFragment::zero (fragment->get_ident ()));
 	    }
 	  else
 	    {
-	      it->second.set_match_amount (match_amount);
+	      // We can just set the repetition amount on the first match
+	      // FIXME: Make this more ergonomic and similar to what we fetch
+	      // in `substitute_repetition`
+	      it->second[0].set_match_amount (match_amount);
 	    }
 	}
     }
@@ -3750,8 +3750,8 @@ MacroExpander::match_repetition (Parser<MacroInvocLexer> &parser,
 AST::ASTFragment
 MacroExpander::transcribe_rule (
   AST::MacroRule &match_rule, AST::DelimTokenTree &invoc_token_tree,
-  std::map<std::string, MatchedFragment> &matched_fragments, bool semicolon,
-  ContextType ctx)
+  std::map<std::string, std::vector<MatchedFragment>> &matched_fragments,
+  bool semicolon, ContextType ctx)
 {
   // we can manipulate the token tree to substitute the dollar identifiers so
   // that when we call parse its already substituted for us
@@ -3765,10 +3765,10 @@ MacroExpander::transcribe_rule (
     = substitute_tokens (invoc_stream, macro_rule_tokens, matched_fragments);
 
   // // handy for debugging
-  for (auto &tok : substituted_tokens)
-    {
-      rust_debug ("tok: [%s]", tok->as_string ().c_str ());
-    }
+  // 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));
@@ -3888,7 +3888,7 @@ MacroExpander::transcribe_rule (
 std::vector<std::unique_ptr<AST::Token>>
 MacroExpander::substitute_metavar (
   std::vector<std::unique_ptr<AST::Token>> &input,
-  std::map<std::string, MatchedFragment> &fragments,
+  std::map<std::string, std::vector<MatchedFragment>> &fragments,
   std::unique_ptr<AST::Token> &metavar)
 {
   auto metavar_name = metavar->get_str ();
@@ -3903,7 +3903,10 @@ MacroExpander::substitute_metavar (
   else
     {
       // Replace
-      MatchedFragment &frag = it->second;
+      // We only care about the vector when expanding repetitions. Just access
+      // the first element of the vector.
+      // FIXME: Clean this up so it makes more sense
+      auto &frag = it->second[0];
       for (size_t offs = frag.token_offset_begin; offs < frag.token_offset_end;
 	   offs++)
 	{
@@ -3919,8 +3922,8 @@ std::vector<std::unique_ptr<AST::Token>>
 MacroExpander::substitute_repetition (
   std::vector<std::unique_ptr<AST::Token>> &input,
   std::vector<std::unique_ptr<AST::Token>> &macro,
-  std::map<std::string, MatchedFragment> &fragments, size_t pattern_start,
-  size_t pattern_end)
+  std::map<std::string, std::vector<MatchedFragment>> &fragments,
+  size_t pattern_start, size_t pattern_end)
 {
   rust_assert (pattern_end < macro.size ());
 
@@ -3929,10 +3932,6 @@ MacroExpander::substitute_repetition (
 
   std::vector<std::unique_ptr<AST::Token>> expanded;
 
-  for (size_t i = pattern_start; i < pattern_end; i++)
-    rust_debug ("[repetition pattern]: %s",
-		macro.at (i)->as_string ().c_str ());
-
   // Find the first fragment and get the amount of repetitions that we should
   // perform
   size_t repeat_amount = 0;
@@ -3951,29 +3950,56 @@ MacroExpander::substitute_repetition (
 		  // fragment), we can just error out. No need to paste the
 		  // tokens as if nothing had happened.
 		  rust_error_at (frag_token->get_locus (),
-				 "metavar used in repetition does not exist");
+				 "metavar %s used in repetition does not exist",
+				 frag_token->get_str ().c_str ());
+		  // FIXME:
 		  return expanded;
 		}
 
-	      repeat_amount = it->second.match_amount;
+	      // FIXME: Refactor, ugly
+	      repeat_amount = it->second[0].match_amount;
 	    }
 	}
     }
 
+  rust_debug ("repetition amount to use: %lu", repeat_amount);
   std::vector<std::unique_ptr<AST::Token>> new_macro;
-  for (size_t tok_idx = pattern_start; tok_idx < pattern_end; tok_idx++)
-    {
-      new_macro.emplace_back (macro.at (tok_idx)->clone_token ());
-      rust_debug ("new macro token: %s",
-		  macro.at (tok_idx)->as_string ().c_str ());
-    }
 
-  // FIXME: We have to be careful and not push the repetition token
-  auto new_tokens = substitute_tokens (input, new_macro, fragments);
+  // We want to generate a "new macro" to substitute with. This new macro
+  // should contain only the tokens inside the pattern
+  for (size_t tok_idx = pattern_start; tok_idx < pattern_end; tok_idx++)
+    new_macro.emplace_back (macro.at (tok_idx)->clone_token ());
+
+  // Then, we want to create a subset of the matches so that
+  // `substitute_tokens()` can only see one fragment per metavar. Let's say we
+  // have the following user input: (1 145 'h')
+  // on the following match arm: ($($lit:literal)*)
+  // which causes the following matches: { "lit": [1, 145, 'h'] }
+  //
+  // The pattern (new_macro) is `$lit:literal`
+  // The first time we expand it, we want $lit to have the following token: 1
+  // The second time, 145
+  // The third and final time, 'h'
+  //
+  // In order to do so we must create "sub maps", which only contain parts of
+  // the original matches
+  // sub-maps: [ { "lit": 1 }, { "lit": 145 }, { "lit": 'h' } ]
+  //
+  // and give them to `substitute_tokens` one by one.
 
-  rust_debug ("repetition amount to use: %lu", repeat_amount);
   for (size_t i = 0; i < repeat_amount; i++)
     {
+      std::map<std::string, std::vector<MatchedFragment>> sub_map;
+      for (auto &kv_match : fragments)
+	{
+	  std::vector<MatchedFragment> sub_vec;
+	  sub_vec.emplace_back (kv_match.second[i]);
+
+	  sub_map.insert ({kv_match.first, sub_vec});
+	}
+
+      auto new_tokens = substitute_tokens (input, new_macro, sub_map);
+
       for (auto &new_token : new_tokens)
 	expanded.emplace_back (new_token->clone_token ());
     }
@@ -3988,7 +4014,8 @@ std::pair<std::vector<std::unique_ptr<AST::Token>>, size_t>
 MacroExpander::substitute_token (
   std::vector<std::unique_ptr<AST::Token>> &input,
   std::vector<std::unique_ptr<AST::Token>> &macro,
-  std::map<std::string, MatchedFragment> &fragments, size_t token_idx)
+  std::map<std::string, std::vector<MatchedFragment>> &fragments,
+  size_t token_idx)
 {
   auto &token = macro.at (token_idx);
   switch (token->get_id ())
@@ -4020,7 +4047,7 @@ MacroExpander::substitute_token (
 	return {
 	  substitute_repetition (input, macro, fragments, pattern_start,
 				 pattern_end),
-	  // + 2 for the opening and closing parenthesis which are mandatory
+	  // + 2 for the opening and closing parentheses which are mandatory
 	  // + 1 for the repetitor (+, *, ?)
 	  pattern_end - pattern_start + 3};
       }
@@ -4044,7 +4071,7 @@ std::vector<std::unique_ptr<AST::Token>>
 MacroExpander::substitute_tokens (
   std::vector<std::unique_ptr<AST::Token>> &input,
   std::vector<std::unique_ptr<AST::Token>> &macro,
-  std::map<std::string, MatchedFragment> &fragments)
+  std::map<std::string, std::vector<MatchedFragment>> &fragments)
 {
   std::vector<std::unique_ptr<AST::Token>> replaced_tokens;
 
diff --git a/gcc/rust/expand/rust-macro-expand.h b/gcc/rust/expand/rust-macro-expand.h
index f3cb36d..eeafdb8 100644
--- a/gcc/rust/expand/rust-macro-expand.h
+++ b/gcc/rust/expand/rust-macro-expand.h
@@ -89,17 +89,38 @@ public:
 
   void push () { stack.push_back ({}); }
 
-  std::map<std::string, MatchedFragment> pop ()
+  std::map<std::string, std::vector<MatchedFragment>> pop ()
   {
     auto top = stack.back ();
     stack.pop_back ();
     return top;
   }
 
-  std::map<std::string, MatchedFragment> &peek () { return stack.back (); }
+  std::map<std::string, std::vector<MatchedFragment>> &peek ()
+  {
+    return stack.back ();
+  }
+
+  void insert_fragment (MatchedFragment fragment)
+  {
+    auto &current_map = stack.back ();
+    auto it = current_map.find (fragment.fragment_ident);
+
+    if (it == current_map.end ())
+      {
+	auto new_frags = std::vector<MatchedFragment> ();
+	new_frags.emplace_back (fragment);
+	current_map.insert ({fragment.fragment_ident, new_frags});
+      }
+    else
+      {
+	auto &frags = it->second;
+	frags.emplace_back (fragment);
+      }
+  }
 
 private:
-  std::vector<std::map<std::string, MatchedFragment>> stack;
+  std::vector<std::map<std::string, std::vector<MatchedFragment>>> stack;
 };
 
 // Object used to store shared data (between functions) for macro expansion.
@@ -151,11 +172,10 @@ struct MacroExpander
   bool try_match_rule (AST::MacroRule &match_rule,
 		       AST::DelimTokenTree &invoc_token_tree);
 
-  AST::ASTFragment
-  transcribe_rule (AST::MacroRule &match_rule,
-		   AST::DelimTokenTree &invoc_token_tree,
-		   std::map<std::string, MatchedFragment> &matched_fragments,
-		   bool semicolon, ContextType ctx);
+  AST::ASTFragment transcribe_rule (
+    AST::MacroRule &match_rule, AST::DelimTokenTree &invoc_token_tree,
+    std::map<std::string, std::vector<MatchedFragment>> &matched_fragments,
+    bool semicolon, ContextType ctx);
 
   bool match_fragment (Parser<MacroInvocLexer> &parser,
 		       AST::MacroMatchFragment &fragment);
@@ -204,10 +224,10 @@ struct MacroExpander
    * @return A token containing the associated fragment expanded into tokens if
    * any, or the cloned token if no fragment was associated
    */
-  static std::vector<std::unique_ptr<AST::Token>>
-  substitute_metavar (std::vector<std::unique_ptr<AST::Token>> &input,
-		      std::map<std::string, MatchedFragment> &fragments,
-		      std::unique_ptr<AST::Token> &metavar);
+  static std::vector<std::unique_ptr<AST::Token>> substitute_metavar (
+    std::vector<std::unique_ptr<AST::Token>> &input,
+    std::map<std::string, std::vector<MatchedFragment>> &fragments,
+    std::unique_ptr<AST::Token> &metavar);
 
   /**
    * Substitute a macro repetition by its given fragments
@@ -219,11 +239,11 @@ struct MacroExpander
    *
    * @return A vector containing the repeated pattern
    */
-  static std::vector<std::unique_ptr<AST::Token>>
-  substitute_repetition (std::vector<std::unique_ptr<AST::Token>> &input,
-			 std::vector<std::unique_ptr<AST::Token>> &macro,
-			 std::map<std::string, MatchedFragment> &fragments,
-			 size_t pattern_start, size_t pattern_end);
+  static std::vector<std::unique_ptr<AST::Token>> substitute_repetition (
+    std::vector<std::unique_ptr<AST::Token>> &input,
+    std::vector<std::unique_ptr<AST::Token>> &macro,
+    std::map<std::string, std::vector<MatchedFragment>> &fragments,
+    size_t pattern_start, size_t pattern_end);
 
   /**
    * Substitute a given token by its appropriate representation
@@ -240,15 +260,16 @@ struct MacroExpander
    * ahead of the input to avoid mis-substitutions
    */
   static std::pair<std::vector<std::unique_ptr<AST::Token>>, size_t>
-  substitute_token (std::vector<std::unique_ptr<AST::Token>> &input,
-		    std::vector<std::unique_ptr<AST::Token>> &macro,
-		    std::map<std::string, MatchedFragment> &fragments,
-		    size_t token_idx);
-
-  static std::vector<std::unique_ptr<AST::Token>>
-  substitute_tokens (std::vector<std::unique_ptr<AST::Token>> &input,
-		     std::vector<std::unique_ptr<AST::Token>> &macro,
-		     std::map<std::string, MatchedFragment> &fragments);
+  substitute_token (
+    std::vector<std::unique_ptr<AST::Token>> &input,
+    std::vector<std::unique_ptr<AST::Token>> &macro,
+    std::map<std::string, std::vector<MatchedFragment>> &fragments,
+    size_t token_idx);
+
+  static std::vector<std::unique_ptr<AST::Token>> substitute_tokens (
+    std::vector<std::unique_ptr<AST::Token>> &input,
+    std::vector<std::unique_ptr<AST::Token>> &macro,
+    std::map<std::string, std::vector<MatchedFragment>> &fragments);
 
   void push_context (ContextType t) { context.push_back (t); }