Repository: ggml-org/llama.cpp
Pull request: grammar : Fix grammar root symbol check

# The problem

Constructing a GBNF grammar allows the programmer to select a `grammar_root`- the symbol to start the grammar from.

The `llama_grammar_init_impl` function incldued a check to see whether the grammar contains a rule for a symbol named literally "root", instead of checking for a symbol with the named passed in as `grammar_root`. This causes valid grammars with non-"root" root symbols to fail, and invalid grammars with a rule named "root", but a different chosen `grammar_root` symbol to pass the check, and immediately fail hard (see failure case in Tests section).

# The fix

Check whether there is a rule for a symbol with the name passed in as `grammar_root`, not literally `"root"`.

I also replaced some `fprintf` error logs with `LLAMA_LOG_ERROR`.

# Alternate solution

We could also remove the option of selecting a non-"root" root symbol, and enforce the root symbol to always be literally "root". In that case we could keep the check as it is, but it would require breaking the libllama API.

# Tests

Before my change, constructing a grammar which had a rule named "root", but specified the root symbol to be something other than root would cause a hard crash like this:

```
⚫ Testing missing start symbol:
/home/asbjorn/Development/llama.cpp/build/bin/libggml-base.so.0(+0x1859b) [0x7ffff7af259b]
/home/asbjorn/Development/llama.cpp/build/bin/libggml-base.so.0(ggml_print_backtrace+0x21c) [0x7ffff7af2a4c]
/home/asbjorn/Development/llama.cpp/build/bin/libggml-base.so.0(+0x2bdb9) [0x7ffff7b05db9]
/nix/store/97f3gw9vpyxvwjv2i673isvg92q65mwn-gcc-13.3.0-lib/lib/libstdc++.so.6(+0xbc20a) [0x7ffff793220a]
/nix/store/97f3gw9vpyxvwjv2i673isvg92q65mwn-gcc-13.3.0-lib/lib/libstdc++.so.6(+0xbc275) [0x7ffff7932275]
/nix/store/97f3gw9vpyxvwjv2i673isvg92q65mwn-gcc-13.3.0-lib/lib/libstdc++.so.6(+0xbc4c7) [0x7ffff79324c7]
/nix/store/97f3gw9vpyxvwjv2i673isvg92q65mwn-gcc-13.3.0-lib/lib/libstdc++.so.6(_ZSt20__throw_out_of_rangePKc+0x40) [0x7ffff79254f8]
/home/asbjorn/Development/llama.cpp/build/bin/libllama.so.0(_Z23llama_grammar_init_implPK11llama_vocabPKcS3_bPS3_mPKim+0xdcb) [0x7ffff7db4f5b]
/home/asbjorn/Development/llama.cpp/build/bin/test-grammar-integration() [0x417ae9]
/nix/store/pacbfvpzqz2mksby36awvbcn051zcji3-glibc-2.40-36/lib/libc.so.6(+0x2a27e) [0x7ffff759b27e]
/nix/store/pacbfvpzqz2mksby36awvbcn051zcji3-glibc-2.40-36/lib/libc.so.6(__libc_start_main+0x89) [0x7ffff759b339]
/home/asbjorn/Development/llama.cpp/build/bin/test-grammar-integration() [0x419cb5]
terminate called after throwing an instance of 'std::out_of_range'
  what():  map::at
```

This is the first of the two test cases I added.

# AI Usage Disclosure

AI was *not* used to identify or correct this bug.

AI was used to review my code changes, at which point it suggested adding tests. The tests were also written and verified by human hand.

Changed files:
filename: src/llama-grammar.cpp
patch: @@ -1160,13 +1160,13 @@ struct llama_grammar * llama_grammar_init_impl(
     // if there is a grammar, parse it
     // rules will be empty (default) if there are parse errors
     if (!parser.parse(grammar_str) || parser.rules.empty()) {
-        fprintf(stderr, "%s: failed to parse grammar\n", __func__);
+        LLAMA_LOG_ERROR("failed to parse grammar\n");
         return nullptr;
     }
 
-    // Ensure that there is a "root" node.
-    if (parser.symbol_ids.find("root") == parser.symbol_ids.end()) {
-        fprintf(stderr, "%s: grammar does not contain a 'root' symbol\n", __func__);
+    // Ensure that the grammar contains the start symbol
+    if (parser.symbol_ids.find(grammar_root) == parser.symbol_ids.end()) {
+        LLAMA_LOG_ERROR("grammar does not contain a '%s' symbol\n", grammar_root);
         return nullptr;
     }
 
@@ -1195,7 +1195,7 @@ struct llama_grammar * llama_grammar_init_impl(
             continue;
         }
         if (llama_grammar_detect_left_recursion(vec_rules, i, &rules_visited, &rules_in_progress, &rules_may_be_empty)) {
-            LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu", i);
+            LLAMA_LOG_ERROR("unsupported grammar, left recursion detected for nonterminal at index %zu\n", i);
             return nullptr;
         }
     }

filename: tests/test-grammar-integration.cpp
patch: @@ -15,8 +15,12 @@
 
 using json = nlohmann::ordered_json;
 
+static llama_grammar * build_grammar_with_root(const std::string & grammar_str, const char * grammar_root) {
+    return llama_grammar_init_impl(nullptr, grammar_str.c_str(), grammar_root, false, nullptr, 0, nullptr, 0);
+}
+
 static llama_grammar * build_grammar(const std::string & grammar_str) {
-    return llama_grammar_init_impl(nullptr, grammar_str.c_str(), "root", false, nullptr, 0, nullptr, 0);
+    return build_grammar_with_root(grammar_str, "root");
 }
 
 static bool test_build_grammar_fails(const std::string & grammar_str) {
@@ -860,6 +864,36 @@ static void test_failure_left_recursion() {
     fprintf(stderr, "  ✅︎ Passed\n");
 }
 
+static void test_failure_missing_root_symbol() {
+    fprintf(stderr, "⚫ Testing missing root symbol:\n");
+
+    const std::string grammar_str = R"""(
+        root ::= "foobar"
+    )""";
+
+    llama_grammar * failure_result = build_grammar_with_root(grammar_str, "nonexistent");
+    assert(failure_result == nullptr);
+
+    fprintf(stderr, "  ✅︎ Passed\n");
+}
+
+static void test_custom_root_symbol_check() {
+    fprintf(stderr, "⚫ Testing custom root symbol check:\n");
+
+    const std::string custom_root_grammar_str = R"""(
+        foobar ::= "foobar"
+    )""";
+
+    llama_grammar * failure_result = build_grammar_with_root(custom_root_grammar_str, "root");
+    assert(failure_result == nullptr);
+
+    llama_grammar * success_result = build_grammar_with_root(custom_root_grammar_str, "foobar");
+    assert(success_result != nullptr);
+    llama_grammar_free_impl(success_result);
+
+    fprintf(stderr, "  ✅︎ Passed\n");
+}
+
 static void test_json_schema() {
     // Note that this is similar to the regular grammar tests,
     //  but we convert each json schema to a grammar before parsing.
@@ -1433,6 +1467,8 @@ int main() {
     test_failure_missing_root();
     test_failure_missing_reference();
     test_failure_left_recursion();
+    test_failure_missing_root_symbol();
+    test_custom_root_symbol_check();
     test_json_schema();
     fprintf(stdout, "All tests passed.\n");
     return 0;