Normalize lexical grammar rules before constructing lex table

project.gyp

@@ -36,6 +36,7 @@
         'src/compiler/prepare_grammar/flatten_grammar.cc',
         'src/compiler/prepare_grammar/intern_symbols.cc',
         'src/compiler/prepare_grammar/is_token.cc',
+        'src/compiler/prepare_grammar/normalize_rules.cc',
         'src/compiler/prepare_grammar/parse_regex.cc',
         'src/compiler/prepare_grammar/prepare_grammar.cc',
         'src/compiler/prepare_grammar/token_description.cc',

spec/compiler/prepare_grammar/extract_choices_spec.cc

@@ -61,12 +61,37 @@ describe("extract_choices", []() {
     })));
   });
 
-  it("handles repeats", [&]() {
-    auto rule = repeat(choice({ sym("a"), sym("b") }));
+  it("does not move choices outside of repeats", [&]() {
+    auto rule = seq({
+      choice({ sym("a"), sym("b") }),
+      repeat(seq({
+        sym("c"),
+        choice({
+          sym("d"),
+          sym("e"),
+        }),
+        sym("f"),
+      })),
+      sym("g"),
+    });
 
     AssertThat(extract_choices(rule), Equals(rule_vector({
-      repeat(sym("a")),
-      repeat(sym("b")),
+      seq({
+        sym("a"),
+        repeat(choice({
+          seq({ sym("c"), sym("d"), sym("f") }),
+          seq({ sym("c"), sym("e"), sym("f") }),
+        })),
+        sym("g"),
+      }),
+      seq({
+        sym("b"),
+        repeat(choice({
+          seq({ sym("c"), sym("d"), sym("f") }),
+          seq({ sym("c"), sym("e"), sym("f") }),
+        })),
+        sym("g"),
+      }),
     })));
   });
 });

src/compiler/prepare_grammar/extract_choices.cc

@@ -42,10 +42,9 @@ class ExtractChoices : public rules::RuleFn<vector<rule_ptr>> {
   }
 
   vector<rule_ptr> apply_to(const rules::Repeat *rule) {
-    vector<rule_ptr> result;
-    for (auto element : apply(rule->content))
-      result.push_back(make_shared<rules::Repeat>(element));
-    return result;
+    return vector<rule_ptr>({
+      rules::Repeat::build(rules::Choice::build(apply(rule->content))),
+    });
   }
 };
 

src/compiler/prepare_grammar/flatten_grammar.h

@@ -1,3 +1,6 @@
+#ifndef COMPILER_PREPARE_GRAMMAR_FLATTEN_GRAMMAR_H_
+#define COMPILER_PREPARE_GRAMMAR_FLATTEN_GRAMMAR_H_
+
 #include <string>
 #include "tree_sitter/compiler.h"
 #include "compiler/syntax_grammar.h"
@@ -11,3 +14,5 @@ SyntaxGrammar flatten_grammar(const InitialSyntaxGrammar &);
 
 }  // namespace prepare_grammar
 }  // namespace tree_sitter
+
+#endif  // COMPILER_PREPARE_GRAMMAR_FLATTEN_GRAMMAR_H_

src/compiler/prepare_grammar/normalize_rules.cc

@@ -0,0 +1,19 @@
+#include "compiler/prepare_grammar/normalize_rules.h"
+#include "compiler/prepare_grammar/extract_choices.h"
+#include "compiler/rules/choice.h"
+
+namespace tree_sitter {
+namespace prepare_grammar {
+
+LexicalGrammar normalize_rules(const LexicalGrammar &input_grammar) {
+  LexicalGrammar result(input_grammar);
+
+  for (Variable &variable : result.variables) {
+    variable.rule = rules::Choice::build(extract_choices(variable.rule));
+  }
+
+  return result;
+}
+
+}  // namespace prepare_grammar
+}  // namespace tree_sitter

src/compiler/prepare_grammar/normalize_rules.h

@@ -0,0 +1,14 @@
+#ifndef COMPILER_PREPARE_GRAMMAR_NORMALIZE_RULES_H_
+#define COMPILER_PREPARE_GRAMMAR_NORMALIZE_RULES_H_
+
+#include "compiler/lexical_grammar.h"
+
+namespace tree_sitter {
+namespace prepare_grammar {
+
+LexicalGrammar normalize_rules(const LexicalGrammar &);
+
+}  // namespace prepare_grammar
+}  // namespace tree_sitter
+
+#endif  // COMPILER_PREPARE_GRAMMAR_NORMALIZE_RULES_H_

src/compiler/prepare_grammar/prepare_grammar.cc

@@ -4,6 +4,7 @@
 #include "compiler/prepare_grammar/extract_tokens.h"
 #include "compiler/prepare_grammar/intern_symbols.h"
 #include "compiler/prepare_grammar/flatten_grammar.h"
+#include "compiler/prepare_grammar/normalize_rules.h"
 #include "compiler/lexical_grammar.h"
 #include "compiler/prepare_grammar/initial_syntax_grammar.h"
 #include "compiler/syntax_grammar.h"
@@ -39,7 +40,7 @@ tuple<SyntaxGrammar, LexicalGrammar, const GrammarError *> prepare_grammar(
   if (error)
     return make_tuple(SyntaxGrammar(), LexicalGrammar(), error);
 
-  return make_tuple(flatten_grammar(syntax_grammar), lex_grammar, nullptr);
+  return make_tuple(flatten_grammar(syntax_grammar), normalize_rules(lex_grammar), nullptr);
 }
 
 }  // namespace prepare_grammar