macro_rules - Fix bad "codegen" for $()+

1 files changed, 38 insertions, 2 deletions

diff --git a/src/macro_rules/parse.cpp b/src/macro_rules/parse.cpp
index 689a81f1..98163998 100644
--- a/src/macro_rules/parse.cpp
+++ b/src/macro_rules/parse.cpp
@@ -470,20 +470,27 @@ namespace {
                     size_t start = rv.size();
                     macro_pattern_to_simple_inner(sp, rv, ent.subpats);
                     push( SimplePatEnt::make_LoopNext({}) );
+                    size_t rewrite_start = rv.size();
                     if( ent.tok != TOK_NULL )
                     {
-                        size_t end = rv.size() + 3;
+                        // If the loop terminator is also valid after the loop
+                        // - Terminate the loop if the next two tokens aren't `<SEP>` folled by any of `<ENTRY>`
                         if( ::std::find(skip_pats.begin(), skip_pats.end(), ExpTok(MacroPatEnt::PAT_TOKEN, &ent.tok)) != skip_pats.end() ) {
+                            size_t expect_and_jump_pos = rv.size() + entry_pats.size() + 1;
                             for(const auto& p : entry_pats)
                             {
                                 auto v = ::make_vec2<SimplePatIfCheck>(
                                     { MacroPatEnt::PAT_TOKEN, ent.tok},
                                     { p.ty, *p.tok }
                                     );
-                                push_ifv(false, mv$(v), ~0u);
+                                // Jump if equal
+                                push_ifv(true, mv$(v), expect_and_jump_pos);
                             }
+                            // If any of the above succeeded, they'll jump past this jump
+                            push( SimplePatEnt::make_Jump({ ~0u }) );
                         }
                         else {
+                            size_t end = rv.size() + 3;
                             push_if( false, MacroPatEnt::PAT_TOKEN, ent.tok, end );
                         }
                         push( SimplePatEnt::make_ExpectTok(ent.tok) );
@@ -497,6 +504,21 @@ namespace {
                             push_if(true, p.ty, *p.tok, start);
                         }
                     }
+
+                    size_t post_loop = rv.size();
+                    for(size_t i = rewrite_start; i < post_loop; i++)
+                    {
+                        if( auto* pe = rv[i].opt_If() ) {
+                            if(pe->jump_target == ~0u) {
+                                pe->jump_target = post_loop;
+                            }
+                        }
+                        if( auto* pe = rv[i].opt_Jump() ) {
+                            if(pe->jump_target == ~0u) {
+                                pe->jump_target = post_loop;
+                            }
+                        }
+                    }
                     push( SimplePatEnt::make_LoopEnd({}) );
                 }
                 else
@@ -589,6 +611,20 @@ namespace {
                 break;
             }
         }
+
+        for(size_t i = level_start; i < rv.size(); i ++)
+        {
+            TU_MATCH_HDRA( (rv[i]), { )
+            default:
+                // Ignore
+            TU_ARMA(If, e) {
+                ASSERT_BUG(sp, e.jump_target < rv.size(), "If target out of bounds, " << e.jump_target << " >= " << rv.size());
+                }
+            TU_ARMA(Jump, e) {
+                ASSERT_BUG(sp, e.jump_target < rv.size(), "Jump target out of bounds, " << e.jump_target << " >= " << rv.size());
+                }
+            }
+        }
     }
 
     ::std::vector<SimplePatEnt> macro_pattern_to_simple(const Span& sp, const ::std::vector<MacroPatEnt>& pattern)