[Feature] JSON Serialize Method with C++ reflection (#335)

This PR provides functionality similar to #277, but leverages C++
reflection to eliminate the need for manually defining
serialization/deserialization rules for each class. It builds on top of
#334, and should be merged after #334 is approved.

### What's new in this PR:

1. Introduced `AutoJSONSerialize` and `AutoJSONDeserialize` to enable
reflection-based JSON serialization.
2. Refactored `tokenizer_info.cc`: moved the definition of
`TokenizerInfo::Impl` into a separate file to support serialization.
3. Added Python bindings for JSON serialization (via `JSONSerializer`)
and corresponding tests.

Author: DarkSharpness

3rdparty/picojson/picojson.h

@@ -50,7 +50,6 @@
 #include <iostream>
 #include <iterator>
 #include <limits>
-#include <map>
 #include <stdexcept>
 #include <string>
 #include <unordered_map>
@@ -267,6 +266,7 @@ class object_with_ordered_keys : private std::unordered_map<std::string, value>
   using std::unordered_map<std::string, value>::at;
   using std::unordered_map<std::string, value>::count;
   using std::unordered_map<std::string, value>::find;
+  using std::unordered_map<std::string, value>::reserve;
 
   value& operator[](const std::string& key) {
     if (count(key) == 0) {

cpp/compiled_grammar_data_structure.h

@@ -18,7 +18,9 @@
 // matcher_data_structure.h is included to use StackElement
 #include "persistent_stack.h"
 #include "support/dynamic_bitset.h"
+#include "support/reflection.h"
 #include "support/utils.h"
+#include "xgrammar/compiler.h"
 
 namespace xgrammar {
 
@@ -74,6 +76,15 @@ struct AdaptiveTokenMask {
   friend std::size_t MemorySize(const AdaptiveTokenMask& mask);
 };
 
+XGRAMMAR_MEMBER_ARRAY(
+    AdaptiveTokenMask,
+    &AdaptiveTokenMask::store_type,
+    &AdaptiveTokenMask::accepted_indices,
+    &AdaptiveTokenMask::rejected_indices,
+    &AdaptiveTokenMask::accepted_bitset,
+    &AdaptiveTokenMask::uncertain_indices
+);
+
 /*!
  * \brief All information that we need to match tokens in the tokenizer to the specified grammar.
  * It is the result of preprocessing.
@@ -119,8 +130,20 @@ class CompiledGrammar::Impl {
   TokenizerInfo GetTokenizerInfo() const { return tokenizer_info; }
 
   std::size_t MemorySize() const;
+
+  friend struct member_trait<Impl>;
 };
 
+XGRAMMAR_MEMBER_TABLE(
+    CompiledGrammar::Impl,
+    "grammar",
+    &CompiledGrammar::Impl::grammar,
+    "tokenizer_info",
+    &CompiledGrammar::Impl::tokenizer_info,
+    "adaptive_token_mask_cache",
+    &CompiledGrammar::Impl::adaptive_token_mask_cache
+);
+
 }  // namespace xgrammar
 
 #endif  // XGRAMMAR_COMPILED_GRAMMAR_DATA_STRUCTURE_H_

cpp/fsm.cc

@@ -12,6 +12,7 @@
 #include <cstring>
 #include <iostream>
 #include <list>
+#include <memory>
 #include <queue>
 #include <set>
 #include <string>
@@ -21,8 +22,11 @@
 #include <utility>
 #include <vector>
 
+#include "picojson.h"
 #include "support/encoding.h"
+#include "support/json.h"
 #include "support/logging.h"
+#include "support/reflection.h"
 #include "support/union_find_set.h"
 
 namespace xgrammar {
@@ -65,6 +69,7 @@ class FSMImplBase {
 
  protected:
   ContainerType edges_;
+  friend struct member_trait<CompactFSM::Impl>;
 };
 
 template <typename ContainerType>
@@ -157,6 +162,8 @@ void FSMImplBase<ContainerType>::GetReachableStates(
 
 class FSM::Impl : public FSMImplBase<std::vector<std::vector<FSMEdge>>> {
  public:
+  Impl() = default;
+
   Impl(int num_states = 0) { edges_.resize(num_states); }
 
   using FSMImplBase<std::vector<std::vector<FSMEdge>>>::FSMImplBase;
@@ -372,6 +379,8 @@ class CompactFSM::Impl : public FSMImplBase<CSRArray<FSMEdge>> {
   friend std::size_t MemorySize(const Impl& self) { return MemorySize(self.edges_); }
 };
 
+XGRAMMAR_MEMBER_ARRAY(CompactFSM::Impl, &CompactFSM::Impl::edges_);
+
 int CompactFSM::Impl::GetNextState(int from, int16_t character) const {
   for (const auto& edge : edges_[from]) {
     if (edge.min == -1) {
@@ -1410,4 +1419,13 @@ FSMWithStartEnd CompactFSMWithStartEnd::ToFSM() const {
   return FSMWithStartEnd(fsm_.ToFSM(), start_, ends_);
 }
 
+picojson::value CompactFSM::SerializeJSONValue() const { return AutoSerializeJSONValue(**this); }
+
+void DeserializeJSONValue(CompactFSM& fsm, const picojson::value& v) {
+  if (!fsm.pimpl_) {
+    fsm.pimpl_ = std::make_unique<CompactFSM::Impl>();
+  }
+  return AutoDeserializeJSONValue(*fsm, v);
+}
+
 }  // namespace xgrammar

cpp/fsm.h

@@ -16,10 +16,11 @@
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
-#include <variant>
 #include <vector>
 
-#include "../cpp/support/csr_array.h"
+#include "picojson.h"
+#include "support/csr_array.h"
+#include "support/reflection.h"
 
 namespace xgrammar {
 
@@ -45,6 +46,9 @@ struct FSMEdge {
         << "Invalid FSMEdge: min > max. min=" << min << ", max=" << max;
   }
 
+  // for serialization only
+  FSMEdge() = default;
+
   /*!
    * \brief Compare the edges. Used to sort the edges in the FSM.
    */
@@ -84,6 +88,8 @@ struct FSMEdge {
   int GetRefRuleId() const { return IsRuleRef() ? max : -1; }
 };
 
+XGRAMMAR_MEMBER_ARRAY(FSMEdge, &FSMEdge::min, &FSMEdge::max, &FSMEdge::target);
+
 }  // namespace xgrammar
 
 namespace std {
@@ -268,6 +274,9 @@ class FSM {
  */
 class CompactFSM {
  public:
+  // for serialization only
+  CompactFSM() = default;
+
   CompactFSM(const CSRArray<FSMEdge>& edges);
 
   CompactFSM(CSRArray<FSMEdge>&& edges);
@@ -346,9 +355,14 @@ class CompactFSM {
    */
   FSM ToFSM() const;
 
+  picojson::value SerializeJSONValue() const;
+  friend void DeserializeJSONValue(CompactFSM& fsm, const picojson::value& v);
+
   XGRAMMAR_DEFINE_PIMPL_METHODS(CompactFSM);
 };
 
+class CompactFSMWithStartEnd;
+
 /*!
  * \brief The base class for FSMWithStartEnd and CompactFSMWithStartEnd. It defines the
  * common constructor and visitor methods.
@@ -361,6 +375,9 @@ class FSMWithStartEndBase {
   );
 
  public:
+  // for serialization only
+  FSMWithStartEndBase() = default;
+
   /*! \brief Constructs an FSMWithStartEnd with a given FSM, start state, and end states. */
   FSMWithStartEndBase(
       const FSMType& fsm, int start, const std::unordered_set<int>& ends, bool is_dfa = false
@@ -456,9 +473,9 @@ class FSMWithStartEndBase {
   std::unordered_set<int> ends_;
   /*! \brief Whether this FSM is a deterministic finite automaton. */
   bool is_dfa_ = false;
-};
 
-class CompactFSMWithStartEnd;
+  friend struct member_trait<CompactFSMWithStartEnd>;
+};
 
 /*!
  * \brief FSMWithStartEnd represents a FSM with start and end states.
@@ -591,6 +608,9 @@ class CompactFSMWithStartEnd : public FSMWithStartEndBase<CompactFSM> {
  public:
   using FSMWithStartEndBase<CompactFSM>::FSMWithStartEndBase;
 
+  // for serialization only
+  CompactFSMWithStartEnd() = default;
+
   /*!
    * \brief Print the FSM.
    * \return The string representation of the FSM.
@@ -613,6 +633,14 @@ class CompactFSMWithStartEnd : public FSMWithStartEndBase<CompactFSM> {
   FSMWithStartEnd ToFSM() const;
 };
 
+XGRAMMAR_MEMBER_ARRAY(
+    CompactFSMWithStartEnd,
+    &CompactFSMWithStartEnd::fsm_,
+    &CompactFSMWithStartEnd::start_,
+    &CompactFSMWithStartEnd::ends_,
+    &CompactFSMWithStartEnd::is_dfa_
+);
+
 /****************** FSMWithStartEndBase Template Implementation ******************/
 
 template <typename FSMType>

cpp/grammar_data_structure.h

@@ -15,6 +15,7 @@
 
 #include "fsm.h"
 #include "support/logging.h"
+#include "support/reflection.h"
 #include "xgrammar/grammar.h"
 
 namespace xgrammar {
@@ -175,14 +176,38 @@ class Grammar::Impl {
   std::vector<int32_t> allow_empty_rule_ids;
 
   friend class GrammarBuilder;
-  friend class GrammarSerializer;
-  friend class GrammarDeserializer;
   friend class GrammarCompiler;
 
   std::size_t MemorySize() const;
   friend std::size_t MemorySize(const Impl& impl);
+  friend struct member_trait<Impl>;
 };
 
+XGRAMMAR_MEMBER_ARRAY(
+    Grammar::Impl::Rule,
+    &Grammar::Impl::Rule::name,
+    &Grammar::Impl::Rule::body_expr_id,
+    &Grammar::Impl::Rule::lookahead_assertion_id
+);
+
+XGRAMMAR_MEMBER_TABLE(
+    Grammar::Impl,
+    "rules_",
+    &Grammar::Impl::rules_,
+    "rule_expr_data_",
+    &Grammar::Impl::rule_expr_data_,
+    "rule_expr_indptr_",
+    &Grammar::Impl::rule_expr_indptr_,
+    "root_rule_id_",
+    &Grammar::Impl::root_rule_id_,
+    "root_tag_dispatch_fsm",
+    &Grammar::Impl::root_tag_dispatch_fsm,
+    "tag_dispatch_end_node_to_rule_id",
+    &Grammar::Impl::tag_dispatch_end_node_to_rule_id,
+    "allow_empty_rule_ids",
+    &Grammar::Impl::allow_empty_rule_ids
+);
+
 }  // namespace xgrammar
 
 #endif  // XGRAMMAR_GRAMMAR_DATA_STRUCTURE_H_

cpp/grammar_functor.cc

@@ -701,7 +701,7 @@ class GrammarUnionFunctorImpl : public SubGrammarCombiner {
   }
 
   // Avoid hiding the original Apply(const Grammar&)
-  Grammar Apply(const Grammar& grammar) final { XGRAMMAR_LOG(FATAL) << "Should not be called"; }
+  Grammar Apply(const Grammar&) final { XGRAMMAR_LOG(FATAL) << "Should not be called"; }
 };
 
 /*!
@@ -735,7 +735,7 @@ class GrammarConcatFunctorImpl : public SubGrammarCombiner {
   }
 
   // Avoid hiding the original Apply(const Grammar&)
-  Grammar Apply(const Grammar& grammar) final { XGRAMMAR_LOG(FATAL) << "Should not be called"; }
+  Grammar Apply(const Grammar&) final { XGRAMMAR_LOG(FATAL) << "Should not be called"; }
 };
 
 /*!
@@ -955,7 +955,7 @@ class StructuralTagGrammarCreatorImpl : public SubGrammarCombiner {
   }
 
   // Avoid hiding the original Apply(const Grammar&)
-  Grammar Apply(const Grammar& grammar) final { XGRAMMAR_LOG(FATAL) << "Should not be called"; }
+  Grammar Apply(const Grammar&) final { XGRAMMAR_LOG(FATAL) << "Should not be called"; }
 };
 
 /*************************** Forward grammar functors to their impl ***************************/

cpp/grammar_serializer.cc

@@ -140,82 +140,4 @@ std::string GrammarPrinter::ToString() {
   return result;
 }
 
-std::string GrammarSerializer::Serialize() {
-  picojson::object grammar_json_obj;
-
-  picojson::array rules_json;
-  for (const auto& rule : grammar_->rules_) {
-    picojson::object rule_json;
-    rule_json["name"] = picojson::value(rule.name);
-    rule_json["body_expr_id"] = picojson::value(static_cast<int64_t>(rule.body_expr_id));
-    rules_json.push_back(picojson::value(rule_json));
-  }
-  grammar_json_obj["rules"] = picojson::value(rules_json);
-
-  picojson::array rule_expr_data_json;
-  for (const auto& data : grammar_->rule_expr_data_) {
-    rule_expr_data_json.push_back(picojson::value(static_cast<int64_t>(data)));
-  }
-  grammar_json_obj["rule_expr_data"] = picojson::value(rule_expr_data_json);
-  picojson::array rule_expr_indptr_json;
-  for (const auto& index_ptr : grammar_->rule_expr_indptr_) {
-    rule_expr_indptr_json.push_back(picojson::value(static_cast<int64_t>(index_ptr)));
-  }
-  grammar_json_obj["rule_expr_indptr"] = picojson::value(rule_expr_indptr_json);
-
-  auto grammar_json = picojson::value(grammar_json_obj);
-  return grammar_json.serialize(prettify_);
-}
-
-Grammar GrammarDeserializer::Deserialize(std::string json_string) {
-  auto node = std::make_shared<Grammar::Impl>();
-
-  auto checker = [&](bool condition) {
-    XGRAMMAR_CHECK(condition) << "Failed to deserialize XGrammar object: " << json_string;
-  };
-
-  picojson::value serialized_value;
-  std::string err = picojson::parse(serialized_value, json_string);
-
-  checker(err.empty() && serialized_value.is<picojson::object>());
-  auto serialized_obj = serialized_value.get<picojson::object>();
-
-  // rules
-  checker(serialized_obj.count("rules") && serialized_obj["rules"].is<picojson::array>());
-  auto rules_array = serialized_obj["rules"].get<picojson::array>();
-
-  checker(rules_array.size() > 0);
-  for (const auto& rule_value : rules_array) {
-    checker(rule_value.is<picojson::object>());
-    auto rule_obj = rule_value.get<picojson::object>();
-    checker(rule_obj.count("name") && rule_obj["name"].is<std::string>());
-    auto name = rule_obj["name"].get<std::string>();
-    checker(rule_obj.count("body_expr_id") && rule_obj["body_expr_id"].is<int64_t>());
-    auto rule_expr = static_cast<int32_t>(rule_obj["body_expr_id"].get<int64_t>());
-    node->rules_.push_back(Grammar::Impl::Rule({name, rule_expr}));
-  }
-
-  // rule_expr_data
-  checker(
-      serialized_obj.count("rule_expr_data") &&
-      serialized_obj["rule_expr_data"].is<picojson::array>()
-  );
-  auto rule_expr_data_array = serialized_obj["rule_expr_data"].get<picojson::array>();
-  for (const auto& data_json : rule_expr_data_array) {
-    node->rule_expr_data_.push_back(static_cast<int32_t>(data_json.get<int64_t>()));
-  }
-
-  // rule_expr_indptr
-  checker(
-      serialized_obj.count("rule_expr_indptr") &&
-      serialized_obj["rule_expr_indptr"].is<picojson::array>()
-  );
-  auto rule_expr_indptr_array = serialized_obj["rule_expr_indptr"].get<picojson::array>();
-  for (const auto& index_ptr_json : rule_expr_indptr_array) {
-    node->rule_expr_indptr_.push_back(static_cast<int32_t>(index_ptr_json.get<int64_t>()));
-  }
-
-  return Grammar(node);
-}
-
 }  // namespace xgrammar