grammar : limit the number of alternatives

Author: o7si

src/llama-grammar.cpp

@@ -10,6 +10,7 @@
 #include <set>
 #include <stdexcept>
 
+#define MAX_ALTERNATE_THRESHOLD 2000
 #define MAX_REPETITION_THRESHOLD 2000
 //
 // helpers
@@ -444,6 +445,9 @@ const char * llama_grammar_parser::parse_alternates(
     llama_grammar_rule rule;
     const char * pos = parse_sequence(src, rule_name, rule, is_nested);
     while (*pos == '|') {
+        if (++n_alternates > MAX_ALTERNATE_THRESHOLD) {
+            throw std::runtime_error("number of grammar alternatives exceeds sane defaults, please reduce the number of alternatives");
+        }
         rule.push_back({LLAMA_GRETYPE_ALT, 0});
         pos = parse_space(pos + 1, true);
         pos = parse_sequence(pos, rule_name, rule, is_nested);
@@ -1527,4 +1531,3 @@ void llama_grammar_accept_token(struct llama_grammar & grammar, llama_token toke
         throw std::runtime_error("Unexpected empty grammar stack after accepting piece: " + piece + " (" + std::to_string(token) + ")");
     }
 }
-

src/llama-grammar.h

@@ -88,6 +88,7 @@ struct llama_grammar_parser {
     std::map<std::string, uint32_t> symbol_ids;
 
     llama_grammar_rules rules;
+    size_t n_alternates = 0;
 
     llama_grammar_parser(const struct llama_vocab * vocab = nullptr) : vocab(vocab) {}
 

tests/test-grammar-parser.cpp

@@ -139,8 +139,20 @@ static void verify_failure(const char * grammar_bytes) {
     assert(result.rules.empty() && "should have failed");
 }
 
+static std::string grammar_with_alternates(size_t n_alternates) {
+    std::string grammar = "root ::= \"a\"";
+    for (size_t i = 0; i < n_alternates; ++i) {
+        grammar += " | \"a\"";
+    }
+    return grammar;
+}
+
 int main()
 {
+    verify_failure(
+        grammar_with_alternates(2001).c_str()
+    );
+
     verify_failure(R"""(
         root ::= "a"{,}"
     )""");