[Serving][Grammar] Porting the json schema converter from python to C++ (#2112)

[Serve][Grammar] Porting the json schema converter from python to C++

This PR ports the json schema converter from python to C++. It defines
the interface:
```
std::string JSONSchemaToEBNF(
    std::string schema, std::optional<int> indent = std::nullopt,
    std::optional<std::pair<std::string, std::string>> separators = std::nullopt,
    bool strict_mode = true);
```

And uses it in BNFGrammar::FromSchema.

This helps cases where python cannot be deployed.

Author: Ubospica

cpp/serve/grammar/grammar.cc

@@ -8,6 +8,7 @@
 #include "grammar_parser.h"
 #include "grammar_serializer.h"
 #include "grammar_simplifier.h"
+#include "json_schema_converter.h"
 
 namespace mlc {
 namespace llm {
@@ -20,7 +21,7 @@ std::ostream& operator<<(std::ostream& os, const BNFGrammar& grammar) {
   return os;
 }
 
-BNFGrammar BNFGrammar::FromEBNFString(const String& ebnf_string, const String& main_rule,
+BNFGrammar BNFGrammar::FromEBNFString(const std::string& ebnf_string, const std::string& main_rule,
                                       bool normalize, bool simplify) {
   auto grammar = EBNFParser::Parse(ebnf_string, main_rule);
   if (normalize) {
@@ -34,41 +35,39 @@ TVM_REGISTER_GLOBAL("mlc.serve.BNFGrammarFromEBNFString")
       return BNFGrammar::FromEBNFString(ebnf_string, main_rule, normalize, simplify);
     });
 
-BNFGrammar BNFGrammar::FromJSON(const String& json_string) {
+BNFGrammar BNFGrammar::FromJSON(const std::string& json_string) {
   return BNFJSONParser::Parse(json_string);
 }
 
 TVM_REGISTER_GLOBAL("mlc.serve.BNFGrammarFromJSON").set_body_typed([](String json_string) {
   return BNFGrammar::FromJSON(json_string);
 });
 
-BNFGrammar BNFGrammar::FromSchema(const String& schema, int indent,
-                                  Optional<Array<String>> separators, bool strict_mode) {
-  static const PackedFunc* json_schema_to_ebnf = Registry::Get("mlc.serve.json_schema_to_ebnf");
-  CHECK(json_schema_to_ebnf != nullptr) << "mlc.serve.json_schema_to_ebnf is not registered.";
-
-  String ebnf_string;
-
-  // Convert the indent parameter to NullOpt for sending it to the PackedFunc.
-  if (indent == -1) {
-    // The conversion from TVMRetValue to String is ambiguous, so we call the conversion function
-    // explicitly
-    ebnf_string =
-        ((*json_schema_to_ebnf)(schema, Optional<ObjectRef>(NullOpt), separators, strict_mode)
-             .
-             operator String());
+BNFGrammar BNFGrammar::FromSchema(const std::string& schema, std::optional<int> indent,
+                                  std::optional<std::pair<std::string, std::string>> separators,
+                                  bool strict_mode) {
+  return FromEBNFString(JSONSchemaToEBNF(schema, indent, separators, strict_mode));
+}
+
+TVM_REGISTER_GLOBAL("mlc.serve.BNFGrammarFromSchema").set_body([](TVMArgs args, TVMRetValue* rv) {
+  std::optional<int> indent;
+  if (args[1].type_code() != kTVMNullptr) {
+    indent = args[1];
   } else {
-    ebnf_string = (*json_schema_to_ebnf)(schema, indent, separators, strict_mode).operator String();
-    ;
+    indent = std::nullopt;
   }
-  return FromEBNFString(ebnf_string);
-}
 
-TVM_REGISTER_GLOBAL("mlc.serve.BNFGrammarFromSchema")
-    .set_body_typed([](const String& schema, int indent, Optional<Array<String>> separators,
-                       bool strict_mode) {
-      return BNFGrammar::FromSchema(schema, indent, separators, strict_mode);
-    });
+  std::optional<std::pair<std::string, std::string>> separators;
+  if (args[2].type_code() != kTVMNullptr) {
+    Array<String> separators_arr = args[2];
+    CHECK(separators_arr.size() == 2);
+    separators = std::make_pair(separators_arr[0], separators_arr[1]);
+  } else {
+    separators = std::nullopt;
+  }
+
+  *rv = BNFGrammar::FromSchema(args[0], indent, separators, args[3]);
+});
 
 const std::string kJSONGrammarString = R"(
 main ::= (

cpp/serve/grammar/grammar.h

@@ -11,6 +11,7 @@
 #include <tvm/runtime/registry.h>
 
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -183,33 +184,38 @@ class BNFGrammar : public ObjectRef {
    * \param simplify Whether to simplify the grammar to make matching more efficient. Default: true.
    * Not implemented yet.
    */
-  static BNFGrammar FromEBNFString(const String& ebnf_string, const String& main_rule = "main",
-                                   bool normalize = true, bool simplify = true);
+  static BNFGrammar FromEBNFString(const std::string& ebnf_string,
+                                   const std::string& main_rule = "main", bool normalize = true,
+                                   bool simplify = true);
 
   /*!
    * \brief Construct a BNF grammar from the dumped JSON string.
    * \param json_string The JSON-formatted string. This string should have the same format as
    * the result of BNFGrammarJSONSerializer::ToString.
    */
-  static BNFGrammar FromJSON(const String& json_string);
+  static BNFGrammar FromJSON(const std::string& json_string);
 
   /*!
    * \brief Construct a BNF grammar from the json schema string. The schema string should be in the
    * format of the schema of a JSON file. We will parse the schema and generate a BNF grammar.
    * \param schema The schema string.
-   * \param indent The number of spaces for indentation. If -1, the output will be in one line.
-   * Default: -1.
+   * \param indent The number of spaces for indentation. If set to std::nullopt, the output will be
+   * in one line. Default: std::nullopt.
    * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"},
-   * {", ", ": "}. If NullOpt, the default separators will be used: {",", ": "} when the indent
-   * is not -1, and {", ", ": "} otherwise. Default: NullOpt.
+   * {", ", ": "}. If std::nullopt, the default separators will be used: {",", ": "} when the
+   * indent is not -1, and {", ", ": "} otherwise. This follows the convention in python
+   * json.dumps(). Default: std::nullopt.
    * \param strict_mode Whether to use strict mode. In strict mode, the generated grammar will not
-   * allow unevaluatedProperties and unevaluatedItems, i.e. these will be set to false by default.
+   * allow properties and items that is not specified in the schema. This is equivalent to
+   * setting unevaluatedProperties and unevaluatedItems to false.
+   *
    * This helps LLM to generate accurate output in the grammar-guided generation with JSON
    * schema. Default: true.
    */
-  static BNFGrammar FromSchema(const String& schema, int indent = -1,
-                               Optional<Array<String>> separators = NullOpt,
-                               bool strict_mode = true);
+  static BNFGrammar FromSchema(
+      const std::string& schema, std::optional<int> indent = std::nullopt,
+      std::optional<std::pair<std::string, std::string>> separators = std::nullopt,
+      bool strict_mode = true);
 
   /*!
    * \brief Get the grammar of standard JSON format. We have built-in support for JSON.

cpp/serve/grammar/grammar_parser.cc

@@ -16,7 +16,7 @@ namespace serve {
 class EBNFParserImpl {
  public:
   /*! \brief The logic of parsing the grammar string. */
-  BNFGrammar DoParse(String ebnf_string, String main_rule);
+  BNFGrammar DoParse(std::string ebnf_string, std::string main_rule);
 
  private:
   using Rule = BNFGrammarNode::Rule;
@@ -192,7 +192,7 @@ int32_t EBNFParserImpl::ParseString() {
   std::vector<int32_t> character_classes;
   while (Peek() && Peek() != '\"') {
     if (Peek() == '\r' || Peek() == '\n') {
-      ThrowParseError("String should not contain newline");
+      ThrowParseError("There should be no newline character in a string literal");
     }
     auto [codepoint, len] = Utf8OrEscapeToCodepoint(cur_);
     if (codepoint == static_cast<TCodepoint>(CharHandlingError::kInvalidUtf8)) {
@@ -391,7 +391,7 @@ void EBNFParserImpl::ResetStringIterator(const char* cur) {
   in_parentheses_ = false;
 }
 
-BNFGrammar EBNFParserImpl::DoParse(String ebnf_string, String main_rule) {
+BNFGrammar EBNFParserImpl::DoParse(std::string ebnf_string, std::string main_rule) {
   ResetStringIterator(ebnf_string.c_str());
   BuildRuleNameToId();
 
@@ -412,12 +412,12 @@ BNFGrammar EBNFParserImpl::DoParse(String ebnf_string, String main_rule) {
   return builder_.Get(main_rule);
 }
 
-BNFGrammar EBNFParser::Parse(String ebnf_string, String main_rule) {
+BNFGrammar EBNFParser::Parse(std::string ebnf_string, std::string main_rule) {
   EBNFParserImpl parser;
   return parser.DoParse(ebnf_string, main_rule);
 }
 
-BNFGrammar BNFJSONParser::Parse(String json_string) {
+BNFGrammar BNFJSONParser::Parse(std::string json_string) {
   auto node = make_object<BNFGrammarNode>();
   auto grammar_json = json::ParseToJsonObject(json_string);
   auto rules_json = json::Lookup<picojson::array>(grammar_json, "rules");

cpp/serve/grammar/grammar_parser.h

@@ -37,7 +37,7 @@ class EBNFParser {
    * \param main_rule The name of the main rule. Default is "main".
    * \return The parsed grammar.
    */
-  static BNFGrammar Parse(String ebnf_string, String main_rule = "main");
+  static BNFGrammar Parse(std::string ebnf_string, std::string main_rule = "main");
 
   /*!
    * \brief The exception thrown when parsing fails.
@@ -58,7 +58,7 @@ class BNFJSONParser {
    * \param json_string The JSON string.
    * \return The parsed BNF grammar.
    */
-  static BNFGrammar Parse(String json_string);
+  static BNFGrammar Parse(std::string json_string);
 };
 
 }  // namespace serve

cpp/serve/grammar/grammar_serializer.cc

@@ -107,7 +107,7 @@ std::string BNFGrammarPrinter::PrintCharacterClassStar(const RuleExpr& rule_expr
   return PrintRuleExpr(rule_expr[0]) + "*";
 }
 
-String BNFGrammarPrinter::ToString() {
+std::string BNFGrammarPrinter::ToString() {
   std::string result;
   auto num_rules = grammar_->NumRules();
   for (auto i = 0; i < num_rules; ++i) {
@@ -120,7 +120,7 @@ TVM_REGISTER_GLOBAL("mlc.serve.BNFGrammarToString").set_body_typed([](const BNFG
   return BNFGrammarPrinter(grammar).ToString();
 });
 
-String BNFGrammarJSONSerializer::ToString() {
+std::string BNFGrammarJSONSerializer::ToString() {
   picojson::object grammar_json;
 
   picojson::array rules_json;

cpp/serve/grammar/grammar_serializer.h

@@ -27,7 +27,7 @@ class BNFGrammarSerializer {
   explicit BNFGrammarSerializer(const BNFGrammar& grammar) : grammar_(grammar) {}
 
   /*! \brief Serialize the grammar to string. */
-  virtual String ToString() = 0;
+  virtual std::string ToString() = 0;
 
  protected:
   const BNFGrammar& grammar_;
@@ -50,7 +50,7 @@ class BNFGrammarPrinter : public BNFGrammarSerializer {
   explicit BNFGrammarPrinter(const BNFGrammar& grammar) : BNFGrammarSerializer(grammar) {}
 
   /*! \brief Print the complete grammar. */
-  String ToString() final;
+  std::string ToString() final;
 
   /*! \brief Print a rule. */
   std::string PrintRule(const Rule& rule);
@@ -102,7 +102,7 @@ class BNFGrammarJSONSerializer : public BNFGrammarSerializer {
    * \brief Dump the raw representation of the AST to a JSON file.
    * \param prettify Whether to format the JSON string. If false, all whitespaces will be removed.
    */
-  String ToString() final;
+  std::string ToString() final;
 
  private:
   bool prettify_;