From 364007affd9d0223ce585618d433399a2b00e7c7 Mon Sep 17 00:00:00 2001 From: tdevelope Date: Wed, 21 Jan 2026 13:40:32 +0200 Subject: [PATCH 1/4] grammar: fix stack overflow using cycle detection and depth limit --- src/llama-grammar.cpp | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index 64ea2fd00a..321d26f3df 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -10,6 +10,7 @@ #include #define MAX_REPETITION_THRESHOLD 2000 +static constexpr uint32_t MAX_GRAMMAR_RECURSION_DEPTH = 2000; // // helpers // @@ -830,7 +831,16 @@ static bool llama_grammar_match_token( static void llama_grammar_advance_stack( const llama_grammar_rules & rules, const llama_grammar_stack & stack, - llama_grammar_stacks & new_stacks) { + llama_grammar_stacks & new_stacks, + std::vector & history) { + if (history.size() >= MAX_GRAMMAR_RECURSION_DEPTH) { + return; + } + + if (std::find(history.begin(), history.end(), stack) != history.end()) { + return; + } + if (stack.empty()) { if (std::find(new_stacks.begin(), new_stacks.end(), stack) == new_stacks.end()) { new_stacks.emplace_back(stack); @@ -855,7 +865,9 @@ static void llama_grammar_advance_stack( // if alternate is nonempty, add to stack new_stack.push_back(subpos); } - llama_grammar_advance_stack(rules, new_stack, new_stacks); + history.push_back(stack); + llama_grammar_advance_stack(rules, new_stack, new_stacks, history); + history.pop_back(); while (!llama_grammar_is_end_of_sequence(subpos)) { // scan to end of alternate def subpos++; @@ -887,6 +899,14 @@ static void llama_grammar_advance_stack( } } +static void llama_grammar_advance_stack( + const llama_grammar_rules & rules, + const llama_grammar_stack & stack, + llama_grammar_stacks & new_stacks) { + std::vector history; + llama_grammar_advance_stack(rules, stack, new_stacks, history); +} + static llama_grammar_candidates llama_grammar_reject_candidates( const llama_grammar_rules & rules, const llama_grammar_stacks & stacks, From 200776bb515ab1860fc536b625bd23f1f625da55 Mon Sep 17 00:00:00 2001 From: tdevelope Date: Wed, 28 Jan 2026 18:18:24 +0200 Subject: [PATCH 2/4] Use #define for MAX_GRAMMAR_RECURSION_DEPTH for consistency --- src/llama-grammar.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index 321d26f3df..2eacd68659 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -10,7 +10,7 @@ #include #define MAX_REPETITION_THRESHOLD 2000 -static constexpr uint32_t MAX_GRAMMAR_RECURSION_DEPTH = 2000; +#define MAX_GRAMMAR_RECURSION_DEPTH 2000 // // helpers // From 23e1519ebec94d5066061303907272ccc2ade9a5 Mon Sep 17 00:00:00 2001 From: tdevelope Date: Wed, 28 Jan 2026 18:36:27 +0200 Subject: [PATCH 3/4] Use unordered_set instead of vector for cycle detection --- src/llama-grammar.cpp | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index 2eacd68659..25b43ab5b8 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -9,6 +9,8 @@ #include #include +#include + #define MAX_REPETITION_THRESHOLD 2000 #define MAX_GRAMMAR_RECURSION_DEPTH 2000 // @@ -826,18 +828,29 @@ static bool llama_grammar_match_token( return false; } +// Hash function for llama_grammar_stack to enable std::unordered_set usage +struct llama_grammar_stack_hash { + std::size_t operator()(const llama_grammar_stack & stack) const { + std::size_t hash = 0; + for (const auto * elem : stack) { + hash ^= std::hash{}(elem) + 0x9e3779b9 + (hash << 6) + (hash >> 2); + } + return hash; + } +}; + // transforms a grammar pushdown stack into N possible stacks, all ending // at a character range (terminal element) static void llama_grammar_advance_stack( const llama_grammar_rules & rules, const llama_grammar_stack & stack, llama_grammar_stacks & new_stacks, - std::vector & history) { + std::unordered_set & history) { if (history.size() >= MAX_GRAMMAR_RECURSION_DEPTH) { return; } - if (std::find(history.begin(), history.end(), stack) != history.end()) { + if (history.count(stack) > 0) { return; } @@ -865,9 +878,9 @@ static void llama_grammar_advance_stack( // if alternate is nonempty, add to stack new_stack.push_back(subpos); } - history.push_back(stack); + history.insert(stack); llama_grammar_advance_stack(rules, new_stack, new_stacks, history); - history.pop_back(); + history.erase(stack); while (!llama_grammar_is_end_of_sequence(subpos)) { // scan to end of alternate def subpos++; @@ -903,7 +916,7 @@ static void llama_grammar_advance_stack( const llama_grammar_rules & rules, const llama_grammar_stack & stack, llama_grammar_stacks & new_stacks) { - std::vector history; + std::unordered_set history; llama_grammar_advance_stack(rules, stack, new_stacks, history); } From 74e6235f5e80b0a76e80bc135a48047f61dbcdcc Mon Sep 17 00:00:00 2001 From: tdevelope Date: Wed, 28 Jan 2026 18:50:00 +0200 Subject: [PATCH 4/4] Add test case for nested repetition pattern --- tests/test-grammar-integration.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test-grammar-integration.cpp b/tests/test-grammar-integration.cpp index 7aa7e58a5c..f8b3a028db 100644 --- a/tests/test-grammar-integration.cpp +++ b/tests/test-grammar-integration.cpp @@ -784,6 +784,23 @@ static void test_quantifiers() { "0xFF 0x12 0xAB 0x00 0x00 0x00", } ); + test_grammar( + "nested repetition", + // Grammar + R"""(root ::= ("a"* )*)""", + // Passing strings + { + "", + "a", + "aa", + "aaa", + }, + // Failing strings + { + "b", + "ab", + } + ); } static void test_failure_missing_root() {