diff --git a/spec/compiler/prepare_grammar/extract_tokens_spec.cpp b/spec/compiler/prepare_grammar/extract_tokens_spec.cpp index 35deb055..3b64dd52 100644 --- a/spec/compiler/prepare_grammar/extract_tokens_spec.cpp +++ b/spec/compiler/prepare_grammar/extract_tokens_spec.cpp @@ -54,6 +54,30 @@ describe("preparing a grammar", []() { character('b') }) }, }))); }); + + it("looks inside sequences, choices and repeats", [&]() { + auto result = extract_tokens(Grammar({ + { "rule1", seq({ + choice({ + repeat(choice({ str("stuff"), sym("a") })), + sym("b"), + }), + sym("c") }) } + })); + + AssertThat(result.first, Equals(Grammar({ + { "rule1", seq({ + choice({ + repeat(choice({ sym("1"), sym("a") })), + sym("b"), + }), + sym("c") }) } + }))); + + AssertThat(result.second, Equals(Grammar("", { + { "1", str("stuff") }, + }))); + }); }); END_TEST \ No newline at end of file diff --git a/src/compiler/prepare_grammar/extract_tokens.cpp b/src/compiler/prepare_grammar/extract_tokens.cpp index 92547f05..6287e32c 100644 --- a/src/compiler/prepare_grammar/extract_tokens.cpp +++ b/src/compiler/prepare_grammar/extract_tokens.cpp @@ -53,6 +53,10 @@ namespace tree_sitter { void visit(const Seq *rule) { value = seq({ apply(rule->left), apply(rule->right) }); } + + void visit(const Repeat *rule) { + value = repeat(apply(rule->content)); + } }; pair extract_tokens(const Grammar &input_grammar) { diff --git a/src/compiler/prepare_grammar/search_for_symbols.cpp b/src/compiler/prepare_grammar/search_for_symbols.cpp index d73b9d10..715559ed 100644 --- a/src/compiler/prepare_grammar/search_for_symbols.cpp +++ b/src/compiler/prepare_grammar/search_for_symbols.cpp @@ -26,6 +26,10 @@ namespace tree_sitter { void visit(const rules::Seq *seq) { value = apply(seq->left) || apply(seq->right); } + + void visit(const rules::Repeat *rule) { + value = apply(rule->content); + } }; bool search_for_symbols(const rules::rule_ptr &rule) {