Skip to content

Commit

Permalink
Add left recursion check: quit early instead of going into an infinit…
Browse files Browse the repository at this point in the history
…e loop
  • Loading branch information
nuchi committed May 5, 2024
1 parent 8425001 commit 19d51ba
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 0 deletions.
72 changes: 72 additions & 0 deletions llama.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13085,6 +13085,25 @@ static std::vector<llama_grammar_candidate> llama_grammar_reject_candidates(
// grammar - external
//

enum detect_left_recursion_status {
// haven't searched this nonterminal
LLAMA_LEFT_REC_NOT_SEARCHED = 0,

// searching this nonterminal in progress
LLAMA_LEFT_REC_IN_PROGRESS = 1,

// finished searching this nonterminal
LLAMA_LEFT_REC_FINISHED_SEARCH = 2,

// detected a cycle
LLAMA_LEFT_REC_FOUND_CYCLE = 3,
};

static void detect_left_recursion(
const std::vector<std::vector<llama_grammar_element>> & rules,
size_t rule_index,
std::vector<enum detect_left_recursion_status> * rules_visited);

struct llama_grammar * llama_grammar_init(
const llama_grammar_element ** rules,
size_t n_rules,
Expand All @@ -13100,6 +13119,17 @@ struct llama_grammar * llama_grammar_init(
vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
}

// Check for left recursion
std::vector<enum detect_left_recursion_status> rules_visited(n_rules);
for (size_t i = 0; i < n_rules; i++) {
detect_left_recursion(vec_rules, i, &rules_visited);
}

auto iter = std::find(rules_visited.begin(), rules_visited.end(), LLAMA_LEFT_REC_FOUND_CYCLE);
if (iter != rules_visited.end()) {
throw std::runtime_error(format("unsupported grammar, left recursion detected for nonterminal at index %d", (int)(iter - rules_visited.begin())));
}

// loop over alternates of start rule to build initial stacks
std::vector<std::vector<const llama_grammar_element *>> stacks;
pos = vec_rules[start_rule_index].data();
Expand All @@ -13122,9 +13152,51 @@ struct llama_grammar * llama_grammar_init(
}
} while (true);

// Important: vec_rules has to be moved here, not copied, because stacks contains
// pointers to elements of vec_rules. If vec_rules were copied into llama_grammar
// then the pointers would be invalidated when the local vec_rules goes out of scope.
return new llama_grammar{ std::move(vec_rules), std::move(stacks), {} };
}

static void detect_left_recursion(
const std::vector<std::vector<llama_grammar_element>> & rules,
size_t rule_index,
std::vector<enum detect_left_recursion_status> * rules_visited) {

int visit_status = (*rules_visited)[rule_index];
if (visit_status == LLAMA_LEFT_REC_IN_PROGRESS) {
// in progress -- we're in a cycle
(*rules_visited)[rule_index] = LLAMA_LEFT_REC_FOUND_CYCLE;
return;
} else if (visit_status == LLAMA_LEFT_REC_NOT_SEARCHED) {
// haven't visited yet. mark in progress, recurse, then mark complete.
// mark in progress
(*rules_visited)[rule_index] = LLAMA_LEFT_REC_IN_PROGRESS;

// recurse
const std::vector<llama_grammar_element> & rule = rules[rule_index];
size_t i = 0;
do {
if (rule[i].type == LLAMA_GRETYPE_RULE_REF) {
detect_left_recursion(rules, (size_t)rule[i].value, rules_visited);
}
while (!llama_grammar_is_end_of_sequence(&rule[i])) {
i++;
}
i++;
} while (i < rule.size());

// mark complete, but only if the recursive call didn't mark a cycle.
// that doesn't mean there's definitely no cycle *for this rule* -- the recursive call
// might have found a different cycle and stopped early.
if ((*rules_visited)[rule_index] == LLAMA_LEFT_REC_IN_PROGRESS) {
(*rules_visited)[rule_index] = LLAMA_LEFT_REC_FINISHED_SEARCH;
}
}

return;
}

void llama_grammar_free(struct llama_grammar * grammar) {
delete grammar;
}
Expand Down
37 changes: 37 additions & 0 deletions tests/test-grammar-integration.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,18 @@ static llama_grammar* build_grammar(const std::string & grammar_str) {
return grammar;
}

static bool test_build_grammar_fails(const std::string & grammar_str) {
bool grammar_fails = false;
try {
build_grammar(grammar_str);
fprintf(stderr, "❌ Expected build failure, but succeeded: %s\n", grammar_str.c_str());
} catch (const std::exception & err) {
grammar_fails = true;
fprintf(stdout, "✅︎\n");
}
return grammar_fails;
}

static bool match_string(const std::string & input, llama_grammar* grammar) {
auto decoded = decode_utf8(input, {});

Expand Down Expand Up @@ -320,13 +332,38 @@ number ::= [0-9]+)""";
fprintf(stderr, " ✅︎ Passed\n");
}

static void test_failure_left_recursion() {
fprintf(stderr, "⚫ Testing left recursion detection:\n");

// Test simple left recursion detection
const std::string simple_str = R"""(root ::= "a" | root "a")""";
assert(test_build_grammar_fails(simple_str));

// Test more complicated left recursion detection
const std::string medium_str = R"""(
root ::= asdf
asdf ::= "a" | asdf "a"
)""";
assert(test_build_grammar_fails(medium_str));

// Test even more complicated left recursion detection
const std::string hard_str = R"""(
root ::= asdf
asdf ::= "a" | foo "b"
foo ::= "c" | asdf "d" | "e")""";
assert(test_build_grammar_fails(hard_str));

fprintf(stderr, " ✅︎ Passed\n");
}

int main() {
fprintf(stdout, "Running grammar integration tests...\n");
test_simple_grammar();
test_complex_grammar();
test_quantifiers();
test_failure_missing_root();
test_failure_missing_reference();
test_failure_left_recursion();
fprintf(stdout, "All tests passed.\n");
return 0;
}

0 comments on commit 19d51ba

Please sign in to comment.