From 9fff51ef495a5b02f446f2450251ac5f08913113 Mon Sep 17 00:00:00 2001 From: Andrey Shvartsman Date: Sat, 19 Sep 2026 17:56:41 -0400 Subject: [PATCH 1/2] feat: add xgrammar structured output with speculative decoding --- README.md | 3 + apps/cli/main.cpp | 1 + apps/cli/options.cpp | 24 + apps/cli/options.h | 1 + cmake/Dependencies.cmake | 1 + docs/cli.md | 14 + docs/maintainer/engine-architecture.md | 24 + docs/serving.md | 92 +- include/ninfer/ops/sampling.h | 5 + include/ninfer/ops/speculative_round.h | 22 +- include/ninfer/types.h | 8 + src/models/qwen3_5/execution/draft.cpp | 4 + src/models/qwen3_5/frontend/frontend.cpp | 43 +- src/models/qwen3_5/frontend/frontend.h | 5 +- .../qwen3_5/frontend/output_session.cpp | 22 +- src/models/qwen3_5/frontend/output_session.h | 4 +- src/models/qwen3_5/program/decode.cpp | 22 + .../qwen3_5/program/planning/request_plan.cpp | 2 + .../qwen3_5/program/planning/startup.cpp | 6 + src/models/qwen3_5/program/planning/startup.h | 1 + src/models/qwen3_5/program/prefill.cpp | 1 + src/models/qwen3_5/program/program_impl.cpp | 7 + src/models/qwen3_5/program/program_impl.h | 4 + src/models/qwen3_5/program/round_buffers.h | 3 + .../qwen3_5/program/structured_round.cpp | 67 + src/models/qwen3_5/program/structured_round.h | 41 + src/models/qwen3_5/program_sources.cmake | 1 + src/ops/kernel/sampling.cuh | 6 +- src/ops/kernel/sampling_device.cuh | 6 +- src/ops/kernel/speculative_round.cuh | 9 +- src/runtime/contract/request.h | 8 + src/runtime/engine/engine.cpp | 3 + src/runtime/engine/engine_core.h | 4 +- src/serve/CMakeLists.txt | 1 + src/serve/anthropic_messages_request.cpp | 6 +- src/serve/openai_chat_request.cpp | 18 +- src/serve/openai_responses.h | 1 + src/serve/openai_responses_request.cpp | 15 +- src/serve/openai_responses_response.cpp | 2 +- src/serve/request.h | 1 + src/serve/structured_output.cpp | 46 + src/serve/structured_output.h | 9 + src/serve/translate.cpp | 15 +- src/text/CMakeLists.txt | 2 + src/text/structured_output.cpp | 193 + src/text/structured_output.h | 38 + tests/README.md | 25 + tests/cmake/CoreTests.cmake | 4 + .../models/qwen3_5/test_structured_round.cpp | 73 + tests/models/qwen3_5/tests.cmake | 5 + tests/ops/test_sampling.cpp | 46 +- tests/ops/test_speculative_round.cpp | 79 +- tests/test_anthropic_schema.cpp | 7 +- tests/test_cli_options.cpp | 10 + tests/test_openai_responses.cpp | 11 +- tests/test_openai_schema.cpp | 39 +- tests/test_structured_output_live.py | 389 ++ tests/text/test_structured_output.cpp | 114 + third_party/xgrammar/3rdparty/dlpack/LICENSE | 201 + .../3rdparty/dlpack/include/dlpack/dlpack.h | 332 ++ .../xgrammar/3rdparty/picojson/picojson.h | 1318 +++++ .../3rdparty/picojson/test_picojson.cpp | 78 + third_party/xgrammar/CMakeLists.txt | 27 + third_party/xgrammar/LICENSE | 201 + third_party/xgrammar/NINFER.md | 11 + third_party/xgrammar/NOTICE | 3 + third_party/xgrammar/cpp/CMakeLists.txt | 2 + third_party/xgrammar/cpp/compiled_grammar.cc | 279 ++ .../xgrammar/cpp/compiled_grammar_impl.h | 148 + third_party/xgrammar/cpp/config.cc | 21 + third_party/xgrammar/cpp/earley_parser.cc | 1272 +++++ third_party/xgrammar/cpp/earley_parser.h | 802 +++ .../xgrammar/cpp/ebnf_script_creator.h | 188 + third_party/xgrammar/cpp/fsm.cc | 2041 ++++++++ third_party/xgrammar/cpp/fsm.h | 1051 ++++ third_party/xgrammar/cpp/fsm_builder.cc | 1797 +++++++ third_party/xgrammar/cpp/fsm_builder.h | 119 + third_party/xgrammar/cpp/grammar.cc | 359 ++ third_party/xgrammar/cpp/grammar_builder.cc | 425 ++ third_party/xgrammar/cpp/grammar_builder.h | 292 ++ third_party/xgrammar/cpp/grammar_compiler.cc | 1725 +++++++ third_party/xgrammar/cpp/grammar_functor.cc | 3832 +++++++++++++++ third_party/xgrammar/cpp/grammar_functor.h | 510 ++ third_party/xgrammar/cpp/grammar_impl.h | 487 ++ third_party/xgrammar/cpp/grammar_matcher.cc | 2826 +++++++++++ third_party/xgrammar/cpp/grammar_parser.cc | 1593 ++++++ third_party/xgrammar/cpp/grammar_parser.h | 107 + third_party/xgrammar/cpp/grammar_printer.cc | 315 ++ third_party/xgrammar/cpp/grammar_printer.h | 86 + .../xgrammar/cpp/json_schema_converter.cc | 4372 +++++++++++++++++ .../xgrammar/cpp/json_schema_converter.h | 635 +++ .../xgrammar/cpp/json_schema_converter_ext.cc | 1449 ++++++ .../xgrammar/cpp/json_schema_converter_ext.h | 292 ++ third_party/xgrammar/cpp/lark_converter.cc | 2714 ++++++++++ third_party/xgrammar/cpp/lark_converter.h | 26 + third_party/xgrammar/cpp/regex_converter.cc | 413 ++ third_party/xgrammar/cpp/regex_converter.h | 21 + third_party/xgrammar/cpp/structural_tag.cc | 2520 ++++++++++ third_party/xgrammar/cpp/structural_tag.h | 411 ++ third_party/xgrammar/cpp/suffix_automata.cc | 88 + third_party/xgrammar/cpp/suffix_automata.h | 36 + .../xgrammar/cpp/support/compact_2d_array.h | 397 ++ third_party/xgrammar/cpp/support/container.h | 166 + third_party/xgrammar/cpp/support/cpptrace.h | 39 + .../xgrammar/cpp/support/dynamic_bitset.h | 363 ++ third_party/xgrammar/cpp/support/encoding.h | 459 ++ third_party/xgrammar/cpp/support/int_set.h | 132 + third_party/xgrammar/cpp/support/json_parse.h | 118 + .../xgrammar/cpp/support/json_serializer.h | 640 +++ third_party/xgrammar/cpp/support/logging.cc | 24 + third_party/xgrammar/cpp/support/logging.h | 236 + .../xgrammar/cpp/support/memory_size.h | 129 + .../xgrammar/cpp/support/recursion_guard.cc | 50 + .../xgrammar/cpp/support/recursion_guard.h | 127 + third_party/xgrammar/cpp/support/reflection.h | 300 ++ .../xgrammar/cpp/support/thread_pool.h | 239 + .../xgrammar/cpp/support/thread_safe_cache.h | 404 ++ .../xgrammar/cpp/support/union_find_set.h | 105 + third_party/xgrammar/cpp/support/utils.h | 451 ++ third_party/xgrammar/cpp/testing.cc | 63 + third_party/xgrammar/cpp/testing.h | 30 + third_party/xgrammar/cpp/tokenizer_info.cc | 562 +++ .../xgrammar/cpp/tokenizer_info_impl.h | 143 + third_party/xgrammar/include/module.modulemap | 4 + .../xgrammar/include/xgrammar/compiler.h | 125 + .../xgrammar/include/xgrammar/config.h | 35 + .../xgrammar/include/xgrammar/exception.h | 80 + .../xgrammar/include/xgrammar/grammar.h | 228 + .../xgrammar/include/xgrammar/matcher.h | 300 ++ .../xgrammar/include/xgrammar/object.h | 51 + .../include/xgrammar/tokenizer_info.h | 86 + .../xgrammar/include/xgrammar/xgrammar.h | 17 + 132 files changed, 43075 insertions(+), 66 deletions(-) create mode 100644 src/models/qwen3_5/program/structured_round.cpp create mode 100644 src/models/qwen3_5/program/structured_round.h create mode 100644 src/serve/structured_output.cpp create mode 100644 src/serve/structured_output.h create mode 100644 src/text/structured_output.cpp create mode 100644 src/text/structured_output.h create mode 100644 tests/models/qwen3_5/test_structured_round.cpp create mode 100644 tests/test_structured_output_live.py create mode 100644 tests/text/test_structured_output.cpp create mode 100644 third_party/xgrammar/3rdparty/dlpack/LICENSE create mode 100644 third_party/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h create mode 100644 third_party/xgrammar/3rdparty/picojson/picojson.h create mode 100644 third_party/xgrammar/3rdparty/picojson/test_picojson.cpp create mode 100644 third_party/xgrammar/CMakeLists.txt create mode 100644 third_party/xgrammar/LICENSE create mode 100644 third_party/xgrammar/NINFER.md create mode 100644 third_party/xgrammar/NOTICE create mode 100644 third_party/xgrammar/cpp/CMakeLists.txt create mode 100644 third_party/xgrammar/cpp/compiled_grammar.cc create mode 100644 third_party/xgrammar/cpp/compiled_grammar_impl.h create mode 100644 third_party/xgrammar/cpp/config.cc create mode 100644 third_party/xgrammar/cpp/earley_parser.cc create mode 100644 third_party/xgrammar/cpp/earley_parser.h create mode 100644 third_party/xgrammar/cpp/ebnf_script_creator.h create mode 100644 third_party/xgrammar/cpp/fsm.cc create mode 100644 third_party/xgrammar/cpp/fsm.h create mode 100644 third_party/xgrammar/cpp/fsm_builder.cc create mode 100644 third_party/xgrammar/cpp/fsm_builder.h create mode 100644 third_party/xgrammar/cpp/grammar.cc create mode 100644 third_party/xgrammar/cpp/grammar_builder.cc create mode 100644 third_party/xgrammar/cpp/grammar_builder.h create mode 100644 third_party/xgrammar/cpp/grammar_compiler.cc create mode 100644 third_party/xgrammar/cpp/grammar_functor.cc create mode 100644 third_party/xgrammar/cpp/grammar_functor.h create mode 100644 third_party/xgrammar/cpp/grammar_impl.h create mode 100644 third_party/xgrammar/cpp/grammar_matcher.cc create mode 100644 third_party/xgrammar/cpp/grammar_parser.cc create mode 100644 third_party/xgrammar/cpp/grammar_parser.h create mode 100644 third_party/xgrammar/cpp/grammar_printer.cc create mode 100644 third_party/xgrammar/cpp/grammar_printer.h create mode 100644 third_party/xgrammar/cpp/json_schema_converter.cc create mode 100644 third_party/xgrammar/cpp/json_schema_converter.h create mode 100644 third_party/xgrammar/cpp/json_schema_converter_ext.cc create mode 100644 third_party/xgrammar/cpp/json_schema_converter_ext.h create mode 100644 third_party/xgrammar/cpp/lark_converter.cc create mode 100644 third_party/xgrammar/cpp/lark_converter.h create mode 100644 third_party/xgrammar/cpp/regex_converter.cc create mode 100644 third_party/xgrammar/cpp/regex_converter.h create mode 100644 third_party/xgrammar/cpp/structural_tag.cc create mode 100644 third_party/xgrammar/cpp/structural_tag.h create mode 100644 third_party/xgrammar/cpp/suffix_automata.cc create mode 100644 third_party/xgrammar/cpp/suffix_automata.h create mode 100644 third_party/xgrammar/cpp/support/compact_2d_array.h create mode 100644 third_party/xgrammar/cpp/support/container.h create mode 100644 third_party/xgrammar/cpp/support/cpptrace.h create mode 100644 third_party/xgrammar/cpp/support/dynamic_bitset.h create mode 100644 third_party/xgrammar/cpp/support/encoding.h create mode 100644 third_party/xgrammar/cpp/support/int_set.h create mode 100644 third_party/xgrammar/cpp/support/json_parse.h create mode 100644 third_party/xgrammar/cpp/support/json_serializer.h create mode 100644 third_party/xgrammar/cpp/support/logging.cc create mode 100644 third_party/xgrammar/cpp/support/logging.h create mode 100644 third_party/xgrammar/cpp/support/memory_size.h create mode 100644 third_party/xgrammar/cpp/support/recursion_guard.cc create mode 100644 third_party/xgrammar/cpp/support/recursion_guard.h create mode 100644 third_party/xgrammar/cpp/support/reflection.h create mode 100644 third_party/xgrammar/cpp/support/thread_pool.h create mode 100644 third_party/xgrammar/cpp/support/thread_safe_cache.h create mode 100644 third_party/xgrammar/cpp/support/union_find_set.h create mode 100644 third_party/xgrammar/cpp/support/utils.h create mode 100644 third_party/xgrammar/cpp/testing.cc create mode 100644 third_party/xgrammar/cpp/testing.h create mode 100644 third_party/xgrammar/cpp/tokenizer_info.cc create mode 100644 third_party/xgrammar/cpp/tokenizer_info_impl.h create mode 100644 third_party/xgrammar/include/module.modulemap create mode 100644 third_party/xgrammar/include/xgrammar/compiler.h create mode 100644 third_party/xgrammar/include/xgrammar/config.h create mode 100644 third_party/xgrammar/include/xgrammar/exception.h create mode 100644 third_party/xgrammar/include/xgrammar/grammar.h create mode 100644 third_party/xgrammar/include/xgrammar/matcher.h create mode 100644 third_party/xgrammar/include/xgrammar/object.h create mode 100644 third_party/xgrammar/include/xgrammar/tokenizer_info.h create mode 100644 third_party/xgrammar/include/xgrammar/xgrammar.h diff --git a/README.md b/README.md index daa143b881..7f884d3f53 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,9 @@ single NVIDIA GeForce RTX 5090. It runs text, image, and video prompts through a OpenAI-/Anthropic-compatible HTTP APIs. The runtime is deliberately specialized: one GPU, one resident model, and a startup-fixed capacity of one to eight active requests. +JSON object and JSON Schema constrained generation use vendored XGrammar v0.2.7. See +[structured output](docs/serving.md#structured-output) for API examples, schema support, and limits. + Five official artifacts are available. The quick-start commands use Qwen3.8-27B NVFP4. | Model | Weights | Artifact | Download and model card | diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index bc8d786fa1..1ccba542f0 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -260,6 +260,7 @@ int main(int argc, char** argv) { ninfer::RequestOptions request; request.execution.sampling = cli.sampling; + request.execution.structured_output = cli.structured_output; request.execution.requested_output_tokens = cli.max_new; request.execution.thinking.budget = cli.thinking_budget; request.stop.token_ids = cli.stop_token_ids; diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index 5bbefbf16a..61d3164cf5 100644 --- a/apps/cli/options.cpp +++ b/apps/cli/options.cpp @@ -5,6 +5,8 @@ #include #include #include +#include +#include #include #include @@ -100,6 +102,8 @@ std::string usage_text(const char* argv0) { "Structured message content accepts text, image/image_url, and video/video_url parts;\n" "media sources may be local paths, HTTP(S) URLs, or base64 data URIs.\n" "--vision enables image/video input and loads the fixed Vision GPU allocations.\n" + "--json constrains output to a JSON object; --json-schema FILE enforces a supported " + "JSON schema. Both disable thinking.\n" "--thinking-budget caps model-origin thinking tokens; inserted control tokens count " "toward --max-new.\n" "--kv-capacity auto leaves " + @@ -151,6 +155,17 @@ Options parse_options(int argc, char** argv) { options.speculative.draft_tokens = parse_u32(value(arg), "draft-tokens"); } else if (arg == "--lm-head-draft") { options.speculative.proposal_head = ProposalHead::Optimized; + } else if (arg == "--json" || arg == "--json-schema") { + if (options.structured_output.kind != StructuredOutputKind::None) { + throw std::invalid_argument("choose exactly one structured output mode"); + } + options.structured_output.kind = arg == "--json" ? StructuredOutputKind::JsonObject + : StructuredOutputKind::JsonSchema; + if (arg == "--json-schema") { + std::ifstream schema(value(arg)); + if (!schema) { throw std::invalid_argument("cannot read JSON schema file"); } + options.structured_output.schema.assign(std::istreambuf_iterator(schema), {}); + } } else if (arg == "--raw-output") { options.raw_output = true; } else if (arg == "--print-token-ids") { @@ -224,6 +239,15 @@ Options parse_options(int argc, char** argv) { throw std::invalid_argument("--kv-capacity must be at least --max-context"); } product::validate_speculative_cli_options(options.speculative); + if (options.structured_output.kind != StructuredOutputKind::None) { + if (options.raw_output || !options.stop_strings.empty() || + !options.stop_token_ids.empty() || options.thinking_budget || + (options.reasoning_effort && options.reasoning_effort != ReasoningEffort::None)) { + throw std::invalid_argument( + "structured output requires decoded text, default stops, and thinking disabled"); + } + options.enable_thinking = false; + } if (options.enable_thinking == false && options.reasoning_effort && *options.reasoning_effort != ReasoningEffort::None) { throw std::invalid_argument("--reasoning-effort cannot be combined with --no-thinking"); diff --git a/apps/cli/options.h b/apps/cli/options.h index 9ef608bea8..94882256b2 100644 --- a/apps/cli/options.h +++ b/apps/cli/options.h @@ -41,6 +41,7 @@ struct Options { // Omitted fields are resolved from the loaded model and rendered prompt mode by Engine. SamplingOverrides sampling; + StructuredOutputOptions structured_output; bool greedy = false; product::LogLevel log_level = product::LogLevel::Info; }; diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 02bfdb08f5..12b48f780e 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -11,6 +11,7 @@ target_include_directories(ninfer::json INTERFACE # Source base for the custom-template frontend; consumers will link it explicitly. add_subdirectory(third_party/llama-jinja EXCLUDE_FROM_ALL) +add_subdirectory(third_party/xgrammar EXCLUDE_FROM_ALL) if(NINFER_BUILD_PRODUCT_SUPPORT) # Media acquisition uses CURLOPT_PROTOCOLS_STR and CURLOPT_REDIR_PROTOCOLS_STR, diff --git a/docs/cli.md b/docs/cli.md index ef73014c02..10c2337186 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -280,3 +280,17 @@ from this one-request interface; the persistent Engine and server routes own cro optional Host backing. All weight, sequence, workspace, and graph allocations are released when the Engine is destroyed. + +## JSON and JSON Schema output + +`--json` constrains output to a JSON object. `--json-schema FILE` constrains it to the supported +JSON Schema subset described in [serving](serving.md#structured-output). Both disable thinking. +They are mutually exclusive and reject raw output, custom stops, thinking budgets, and enabled +reasoning effort. They work with ordinary, MTP, DFlash and DFlash2 execution. + +```bash +./build/apps/ninfer model.ninfer --prompt 'Return the city as JSON' --json --max-new 128 +./build/apps/ninfer model.ninfer --prompt 'Return a weather record' --json-schema weather.schema.json --max-new 128 +``` + +Read the reported finish reason: an output/context limit or cancellation can truncate the JSON. diff --git a/docs/maintainer/engine-architecture.md b/docs/maintainer/engine-architecture.md index f5dd612baa..ba7f82f11c 100644 --- a/docs/maintainer/engine-architecture.md +++ b/docs/maintainer/engine-architecture.md @@ -601,3 +601,27 @@ fixed-shape 和 device-specialized 实现）都归 `src/ops`。 和 consumer address contract； - [Op development](op-development.md)：Op 正确性与性能准入； - [CLI](../cli.md)与 [HTTP serving](../serving.md)：外部行为。 + +## Structured output transaction + +`StructuredOutputOptions` crosses the public Engine boundary as schema data. Frontend compiles it +with a tokenizer-specific XGrammar compiler (bounded 256 MiB cache) and creates a request-owned +`GrammarState`. Engine passes the same state through the base plan, admission plan, and request +control. Sequence/checkpoint/cache state never owns it. OutputSession previews on a fork and +moves the fork into the committed state only after Program commit succeeds. Cancellation before +output preview advances no grammar state. + +Program reserves vocabulary bitsets in its planned persistent device arena and owns pinned host +staging. Ordinary/prefill sampling uses the current mask. MTP supplies each current draft prefix +to a forked matcher before launch. DFlash/DFlash2 must first generate device drafts: their graph +captures D2H draft transfer, a CUDA host function that fills all reachable prefix masks, H2D mask +transfer, then target verification. The host function calls no CUDA API; it captures exceptions +for rethrow after stream synchronization and before any token can be published. All node addresses +belong to Program and outlive graph replay. + +SamplingConfig carries an optional bitset pointer and column stride. Column i describes the +grammar after drafts[0..i). The mask is applied before penalties and sampling filters. After an +illegal draft or EOS, suffix columns are unreachable and need no grammar traversal. Every reachable +mask must be nonempty. Raw-logit fast paths require a null mask. Draft proposal probabilities are +unchanged; target p is constrained and normalized before the existing p/q and residual calculation. +Existing KV, recurrence replay, token-count, and terminal-prefix commit rules remain authoritative. diff --git a/docs/serving.md b/docs/serving.md index 329763fe1b..fb6085374d 100644 --- a/docs/serving.md +++ b/docs/serving.md @@ -50,6 +50,84 @@ and prefill remain outside speculative acceleration. A later request cannot enab omitted at startup. The artifact need only contain the Text backbone and the optional components selected for this process. +## Structured output + +NInfer uses the vendored XGrammar v0.2.7 C++ compiler and token matcher to constrain final +response content. This requires no model conversion or additional weights. Ordinary decoding, +MTP, DFlash, and DFlash2 share the same target sampling contract, including CUDA Graphs, +streaming, prefix reuse, and mixed constrained/unconstrained concurrent requests. + +Chat Completions accepts: + +```json +{ + "model": "your-model-id", + "messages": [{"role": "user", "content": "Give the city and temperature as JSON."}], + "max_tokens": 128, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "weather", + "strict": true, + "schema": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "temperature": {"type": "number"} + }, + "required": ["city", "temperature"], + "additionalProperties": false + } + } + } +} +``` + +Use `response_format: {"type":"json_object"}` for any JSON object. JSON Schema mode can also +constrain other root types. Responses uses the flat form +`text.format: {"type":"json_schema","name":"weather","strict":true,"schema":{...}}` and +reports that format in its Response object. Anthropic Messages accepts +`output_config.format: {"type":"json_schema","schema":{...}}`. + +Supported schema constraints are `type`, `properties`, `required`, `additionalProperties`, +`items`, `prefixItems`, `minItems`, `maxItems`, `minLength`, `maxLength`, `enum`, `const`, +`anyOf`, `$defs`, `definitions`, and local fragment `$ref` (including recursive schemas). +Annotations `$schema`, `title`, `description`, `default`, `examples`, and `$comment` +do not impose generation constraints. Other keywords are rejected with HTTP 400: this includes +numeric bounds, `multipleOf`, `pattern`, `format`, `oneOf`, `allOf`, `uniqueItems`, and conditionals. +`$id` and external references are rejected. An explicit `$schema` must be JSON Schema 2020-12 +or draft-07. Local references use `#` or literal object paths such as `#/$defs/node`; +escaped or empty path segments are rejected. Bounded strings use unescaped Unicode characters; escaped quotes, backslashes, +and control characters are excluded from that generated subset. Nonnegative length/item bounds +must fit a signed 32-bit integer. `$ref` and `anyOf` cannot have sibling constraints; `const` +and `enum` allow a matching single `type` declaration but no other sibling constraints. Move +constraints into the referenced schema or each union branch. Missing `additionalProperties` +and `items` retain JSON Schema defaults. Properties are emitted in schema declaration order; +this is a valid subset of the requested schema. Whitespace between JSON elements is bounded to +eight characters per run to prevent whitespace-only generation loops. Additional-property key spellings are restricted +where necessary to prevent escaped aliases from overwriting declared typed properties. +`strict:false` does not disable enforcement. + +Structured responses default to thinking disabled even if the server's ordinary default enables +thinking. Explicit enabled thinking, a thinking budget, active tool generation, custom stops, +and assistant-prefill continuation are rejected in combination with structured output. Strict +tool argument generation and arbitrary grammar/regex aliases remain unsupported. Public C++ +callers select `ExecutionOptions::structured_output`, prepare a prompt with thinking disabled, +and use default stops and decoded text output. + +Only normal completed responses guarantee a complete JSON document satisfying the supported +schema. Token/context limits, cancellation, transport failure, or generation errors can leave a +partial document; inspect the finish reason or Responses status before parsing it as complete. +SSE content deltas are ordinary partial JSON bytes; concatenate them before parsing. Constraints +guarantee format, not factual accuracy or semantic task success. + +Masks are applied before target top-k/top-p/min-p filtering at **every** speculative position, +including correction and bonus tokens. Unconstrained draft distributions remain unchanged; +accept/reject and residual sampling use the constrained target distribution. Per-request grammar +state advances only with Engine output commit, and is never restored from a prompt/KV cache. +DFlash/DFlash2 CUDA Graphs contain a host matcher node between draft generation and verification; +this adds a CPU synchronization point and mask transfers per round. No speedup claim is implied. + ## Endpoints | Method and path | Behavior | @@ -109,7 +187,7 @@ The endpoint supports: - `temperature`, `top_p`, presence/frequency penalties, and signed integer `seed`; - the compatible `top_k` (`0..20`) and `min_p` (`0..1`) sampler extensions; - up to four non-empty stop strings, applied to both reasoning and answer output; -- `n:1`, text-only `modalities`, and `response_format: {"type":"text"}`; +- `n:1`, text-only `modalities`, and `response_format` with `text`, `json_object`, or `json_schema`; - non-streaming responses and server-sent event streams; - `stream_options.include_usage`; - llama.cpp-compatible terminal `timings`, plus opt-in `timings_per_token` and @@ -123,8 +201,8 @@ The endpoint supports: - Assistant `reasoning_content` and `reasoning` history aliases. Options whose observable behavior the Engine cannot provide are rejected when they request that -behavior. This includes JSON constrained output, nonzero `logit_bias`, requested log probabilities, -audio/file input or audio output, `strict:true`, required or named tool choice, +behavior. This includes nonzero `logit_bias`, requested log probabilities, +audio/file input or audio output, tool `strict:true`, required or named tool choice, `parallel_tool_calls:false` with enabled tools, explicit low/high image detail, web search, moderation, low/high verbosity, stored Chat Completions, and non-empty legacy `functions`. Each capability rejection identifies the affected field and the guarantee NInfer cannot provide. @@ -412,7 +490,7 @@ wire response contains typed `output` Items. | `reasoning.effort` | `none` requests disabled thinking; other standard effort values pass to the selected template | | `chat_template_kwargs` | template parameters as a JSON object; standard options merge with typed fields | | `preserve_thinking` | alias for `chat_template_kwargs.preserve_thinking`; conflicting values are rejected | -| `text.format` | omitted or `{"type":"text"}` only | +| `text.format` | `text`, `json_object`, or flat `json_schema`; see [structured output](#structured-output) | | `tools` | direct function definitions or namespace groups containing function definitions; see below | | `tool_choice` | `auto`, `none`, or function-only `allowed_tools` with mode `auto`; a namespaced selection carries both `namespace` and `name` | | `parallel_tool_calls` | `true` by default; `false` is accepted only when no effective tool is callable | @@ -512,7 +590,7 @@ undeclared model output remains ordinary text. `allowed_tools` with mode `auto` without changing declaration order, while `tool_choice:"none"` disables structured tool output even when the history contains earlier calls. -NInfer does not execute functions or enforce JSON Schema through constrained decoding, so +NInfer does not execute functions or constrain tool arguments with JSON Schema, so `strict:true`, required or named tool choice, hosted tools, remote MCP tools, and custom free-form tools are rejected. Deferred loading, output schemas, and caller restrictions that exclude direct invocation are also rejected because their semantics cannot be honored. @@ -994,8 +1072,8 @@ a following compatible turn can reuse it. Output-limit and context-capacity fini `length`/ `max_tokens`; ordinary model or string stops map to `stop`/ `end_turn`. Function tools are rendered into the model prompt and generated calls are parsed into protocol -responses. NInfer does not execute tools and does not enforce client JSON Schema through constrained -decoding. +responses. NInfer does not execute tools or constrain their arguments. Final JSON responses can +use [structured output](#structured-output). Prompt-token usage includes chat-template and expanded media tokens. Generated-token usage comes from accepted output token IDs, including a stop token whose decoded text may be withheld. diff --git a/include/ninfer/ops/sampling.h b/include/ninfer/ops/sampling.h index a4635b2719..1fbd3c6c49 100644 --- a/include/ninfer/ops/sampling.h +++ b/include/ninfer/ops/sampling.h @@ -31,6 +31,11 @@ struct SamplingConfig { float frequency_penalty = 0.0f; unsigned long long seed = 0; std::int32_t* token_counts = nullptr; // device [token_domain] i32, or null + // Optional device bitset [ceil(token_domain/32), speculative_width]. Column i is + // conditioned on drafts[0..i). Mask BEFORE penalties/temperature/top-k/top-p/min-p. + // Every reachable column must allow at least one token with a finite logit. + const std::uint32_t* token_mask = nullptr; + std::int32_t token_mask_stride = 0; // words per column; ordinary sample uses column zero }; // Caller-owned transient capacity for every parallel sampling-lane count in the inclusive diff --git a/include/ninfer/ops/speculative_round.h b/include/ninfer/ops/speculative_round.h index 790bec78ac..61d0714054 100644 --- a/include/ninfer/ops/speculative_round.h +++ b/include/ninfer/ops/speculative_round.h @@ -10,8 +10,9 @@ namespace ninfer::ops { struct SpeculativeAcceptExecutionEnvelope { - // Execution promise: every row has temperature<=0 and both penalties disabled. When false, - // the general route remains valid for any supported mixture of greedy and stochastic rows. + // Execution promise: every row has temperature<=0, both penalties disabled, and no token mask. + // When false, the general route remains valid for any supported mixture of greedy and + // stochastic rows. bool all_rows_greedy_without_penalties = false; }; @@ -62,13 +63,13 @@ void speculative_prepare_verify_ids(const Tensor& anchors, const Tensor& drafts, * Algorithm: * Independently for each row b, greedy mode accepts the longest available draft prefix matching * the per-column penalty-adjusted argmax and commits that argmax at the first mismatch (or the - * bonus column). With both penalties disabled, target_tokens is the exact raw-logit fast path. - * Sampling mode applies configs[b] to each valid verification column, accepts draft i with - * target probability p_i(draft_i), samples from the residual distribution on first rejection, - * and samples a bonus from column Pcur[b] when every available draft is accepted. The draft - * proposal distribution is one-hot at each greedy draft token. - * RNG domains are the speculative accept/correction/bonus SamplePurpose values and logical - * positions derived from the old length. + * bonus column). With both penalties disabled and no token mask, target_tokens is the exact + * raw-logit fast path. Sampling mode applies configs[b] to each valid verification column, accepts + * draft i with target probability p_i(draft_i), samples from the residual distribution on first + * rejection, and samples a bonus from column Pcur[b] when every available draft is accepted. The + * draft proposal distribution is one-hot at each greedy draft token. RNG domains are the + * speculative accept/correction/bonus SamplePurpose values and logical positions derived from the + * old length. * * Logical shapes: * All Tensor storage is contiguous. target_tokens/licensed_tokens are I32 [K+1,B], drafts is @@ -79,6 +80,9 @@ void speculative_prepare_verify_ids(const Tensor& anchors, const Tensor& drafts, * * Numeric: * Sampling filtering, penalties, normalization, and RNG semantics are those of sampling.h. + * Token mask column i is conditioned on drafts[0..i), including the bonus column. It is + * applied before truncation/normalization; the proposal q is unchanged. Rejection therefore + * samples max(p_constrained-q,0), preserving the constrained target distribution. * * Effects: * For each row, let A be the accepted draft count and L=A+1. licensed_tokens[0:A,b] receives diff --git a/include/ninfer/types.h b/include/ninfer/types.h index a2e8b75481..b167b15106 100644 --- a/include/ninfer/types.h +++ b/include/ninfer/types.h @@ -245,7 +245,15 @@ struct ThinkingControlOptions { std::optional budget; }; +enum class StructuredOutputKind : std::uint8_t { None, JsonObject, JsonSchema }; + +struct StructuredOutputOptions { + StructuredOutputKind kind = StructuredOutputKind::None; + std::string schema; // JSON Schema for JsonSchema, otherwise empty +}; + struct ExecutionOptions { + StructuredOutputOptions structured_output; SamplingOverrides sampling; std::uint32_t requested_output_tokens = 0; bool allow_prefix_reuse = true; diff --git a/src/models/qwen3_5/execution/draft.cpp b/src/models/qwen3_5/execution/draft.cpp index eaf1aba2fb..4a04b011c1 100644 --- a/src/models/qwen3_5/execution/draft.cpp +++ b/src/models/qwen3_5/execution/draft.cpp @@ -1,3 +1,4 @@ +#include "models/qwen3_5/program/structured_round.h" #include "models/qwen3_5/program/graph_execution.h" #include "models/qwen3_5/execution/linear.h" #include @@ -611,6 +612,9 @@ auto dflash_decode_batch_body(DFlashBatchContext& state, std::int32_t batch_size state_destinations, dflash_rows, envelopes.append); propose_batch_impl(state, frame, batch_size, k, envelopes); + if (state.execution.io.structured) { + state.execution.io.structured->enqueue_dflash(drafts, state.execution.device.stream); + } ops::speculative_prepare_verify_inputs(anchors, drafts, frontiers, extents, verify_ids, target_positions, state.execution.device.stream); diff --git a/src/models/qwen3_5/frontend/frontend.cpp b/src/models/qwen3_5/frontend/frontend.cpp index e7ab4c4a54..52159ff9d8 100644 --- a/src/models/qwen3_5/frontend/frontend.cpp +++ b/src/models/qwen3_5/frontend/frontend.cpp @@ -10,6 +10,8 @@ #include "models/qwen3_5/frontend/tokenizer.h" #include "models/qwen3_5/frontend/tool_call_parser.h" #include "text/unicode.h" +#include "text/structured_output.h" +#include #include @@ -624,6 +626,8 @@ class Frontend::Impl { thinking_control_tokens = std::make_shared>(std::move(encoded)); } + mutable std::mutex grammar_mutex; + mutable std::unique_ptr grammar_compiler; fi::CompiledChatTemplate chat_template; std::shared_ptr tokenizer; fi::ProcessorOptions processor; @@ -903,13 +907,44 @@ std::vector Frontend::tokenize_text(std::string_view text) const { OutputSession Frontend::make_output_session(const PreparedPrompt& prompt, const StopPolicy& caller_stop, const OutputOptions& output, - const ThinkingControlOptions& thinking) const { + const ThinkingControlOptions& thinking, + const StructuredOutputOptions& structured) const { if (prompt.data_ == nullptr) { throw std::invalid_argument("prepared prompt is empty"); } StopPolicy policy = merge_stop_policy(*impl_->tokenizer, caller_stop); + std::shared_ptr grammar; + text::validate_structured_output(structured); + if (structured.kind != StructuredOutputKind::None) { + if (prompt.data_->starts_in_reasoning || thinking.budget) { + throw std::invalid_argument("structured output requires thinking disabled"); + } + if (!caller_stop.token_ids.empty() || !caller_stop.strings.empty() || + !caller_stop.include_model_defaults || output.raw || output.preserve_special_tokens || + caller_stop.publish_stop_token) { + throw std::invalid_argument( + "structured output requires default stops and decoded text output"); + } + if (prompt.data_->tool_call_output && !prompt.data_->tool_call_output->tools.empty()) { + throw std::invalid_argument( + "structured output cannot be combined with tool generation"); + } + std::lock_guard lock(impl_->grammar_mutex); + if (!impl_->grammar_compiler) { + std::vector vocab(impl_->tokenizer->vocab_size()); + for (std::size_t id = 0; id < vocab.size(); ++id) { + if (impl_->tokenizer->is_valid_token(id) && + !impl_->tokenizer->is_special_token(id)) { + vocab[id] = impl_->tokenizer->decode_token_bytes(id, false); + } + } + impl_->grammar_compiler = std::make_unique( + std::move(vocab), impl_->tokenizer->default_stop_token_ids()); + } + grammar = impl_->grammar_compiler->compile(structured); + } if (output.raw) { policy.publish_stop_token = true; } - return OutputSession(impl_->tokenizer, std::move(policy), output, - prompt.data_->starts_in_reasoning, thinking, - impl_->thinking_control_tokens, prompt.data_->tool_call_output); + return OutputSession( + impl_->tokenizer, std::move(policy), output, prompt.data_->starts_in_reasoning, thinking, + impl_->thinking_control_tokens, prompt.data_->tool_call_output, std::move(grammar)); } const StopPolicy& Frontend::default_stop_policy() const noexcept { return impl_->defaults; } diff --git a/src/models/qwen3_5/frontend/frontend.h b/src/models/qwen3_5/frontend/frontend.h index 9e022f8011..96c6767aea 100644 --- a/src/models/qwen3_5/frontend/frontend.h +++ b/src/models/qwen3_5/frontend/frontend.h @@ -74,8 +74,9 @@ class Frontend { [[nodiscard]] MediaCacheSummary media_cache_summary() const; [[nodiscard]] OutputSession make_output_session(const PreparedPrompt& prompt, const StopPolicy& caller_stop, - const OutputOptions& output = {}, - const ThinkingControlOptions& thinking = {}) const; + const OutputOptions& output = {}, + const ThinkingControlOptions& thinking = {}, + const StructuredOutputOptions& structured = {}) const; [[nodiscard]] const StopPolicy& default_stop_policy() const noexcept; [[nodiscard]] const ModelSamplingDefaults& sampling_defaults() const noexcept; diff --git a/src/models/qwen3_5/frontend/output_session.cpp b/src/models/qwen3_5/frontend/output_session.cpp index a3a5df6496..e355418e28 100644 --- a/src/models/qwen3_5/frontend/output_session.cpp +++ b/src/models/qwen3_5/frontend/output_session.cpp @@ -1,4 +1,5 @@ #include "models/qwen3_5/frontend/output_session.h" +#include "text/structured_output.h" #include "models/qwen3_5/frontend/chat_template.h" #include "models/qwen3_5/frontend/tokenizer.h" #include "models/qwen3_5/frontend/tool_call_parser.h" @@ -367,6 +368,8 @@ class OutputSession::Impl { std::vector tool_calls; ToolCallParseDiagnostics tool_call_parse; bool preview_ready = false; + std::shared_ptr grammar; + std::unique_ptr preview_grammar; }; PublishedOutput::PublishedOutput(PublishedOutput&& other) noexcept @@ -401,10 +404,17 @@ OutputSession::OutputSession( std::shared_ptr tokenizer, StopPolicy policy, OutputOptions output, bool starts_in_reasoning, ThinkingControlOptions thinking, std::shared_ptr> thinking_control_tokens, - std::shared_ptr tool_call_output) + std::shared_ptr tool_call_output, + std::shared_ptr grammar) : impl_(std::make_unique( std::move(tokenizer), std::move(policy), output, starts_in_reasoning, thinking, - std::move(thinking_control_tokens), std::move(tool_call_output))) {} + std::move(thinking_control_tokens), std::move(tool_call_output))) { + impl_->grammar = std::move(grammar); +} + +std::shared_ptr OutputSession::grammar_state() const { + return impl_ ? impl_->grammar : nullptr; +} runtime::OutputDecision OutputSession::preview_model(std::span tokens, std::uint32_t total_budget_remaining, @@ -439,6 +449,10 @@ runtime::OutputDecision OutputSession::preview_model(std::span to if (impl_->preview_execution_split_after && *impl_->preview_execution_split_after > count) { throw std::logic_error("prefix execution split exceeds the accepted token prefix"); } + if (impl_->grammar) { + impl_->preview_grammar = impl_->grammar->fork(); + impl_->preview_grammar->accept(tokens.first(count)); + } impl_->preview_ready = true; return runtime::OutputDecision{ .accepted_tokens = count, @@ -622,6 +636,10 @@ runtime::OutputDecision OutputSession::preview_terminal(FinishReason reason) { PublishedOutput OutputSession::commit_preview() { if (impl_ == nullptr || !impl_->preview_ready) { std::terminate(); } + if (impl_->preview_grammar) { + *impl_->grammar = std::move(*impl_->preview_grammar); + impl_->preview_grammar.reset(); + } using std::swap; swap(impl_->state, impl_->preview_state); swap(impl_->semantic, impl_->preview_semantic); diff --git a/src/models/qwen3_5/frontend/output_session.h b/src/models/qwen3_5/frontend/output_session.h index 2b7e3ed74c..6369425798 100644 --- a/src/models/qwen3_5/frontend/output_session.h +++ b/src/models/qwen3_5/frontend/output_session.h @@ -73,6 +73,7 @@ class OutputSession { void validate_generation_capacity(std::uint32_t effective_output_tokens) const; [[nodiscard]] runtime::OutputDecision preview_terminal(FinishReason reason); [[nodiscard]] PublishedOutput commit_preview(); + [[nodiscard]] std::shared_ptr grammar_state() const; [[nodiscard]] std::vector take_tool_calls() noexcept; [[nodiscard]] ToolCallParseDiagnostics tool_call_parse_diagnostics() const noexcept; [[nodiscard]] std::uint32_t reasoning_tokens() const noexcept; @@ -84,7 +85,8 @@ class OutputSession { OutputSession(std::shared_ptr tokenizer, StopPolicy policy, OutputOptions output, bool starts_in_reasoning, ThinkingControlOptions thinking, std::shared_ptr> thinking_control_tokens, - std::shared_ptr tool_call_output); + std::shared_ptr tool_call_output, + std::shared_ptr grammar = {}); std::unique_ptr impl_; friend class Frontend; diff --git a/src/models/qwen3_5/program/decode.cpp b/src/models/qwen3_5/program/decode.cpp index 1af1fc3989..a0432999fb 100644 --- a/src/models/qwen3_5/program/decode.cpp +++ b/src/models/qwen3_5/program/decode.cpp @@ -1,3 +1,4 @@ +#include "models/qwen3_5/program/structured_round.h" #include "models/qwen3_5/program/program_impl.h" #include "models/qwen3_5/program/context_work.h" #include "models/qwen3_5/program/context.h" @@ -136,6 +137,11 @@ void ProgramImpl::install_sampling(SequenceState& sequence, RequestControl& requ Tensor counts = token_counts.slice(1, static_cast(sequence.lane), 1) .view({dimension(parameters.model.resources().public_token_count)}); request.sampling_host = config; + if (request.grammar) { + request.sampling_host.token_mask = structured_round->device_mask(sequence.lane); + request.sampling_host.token_mask_stride = structured_round->stride(); + structured_round->fill(sequence.lane, *request.grammar, {}, device.stream); + } request.speculative_stats = SpeculativeStats{ .backend = speculative_backend, .enabled = speculative_backend != SpeculativeBackend::None, @@ -328,6 +334,9 @@ ProgramImpl::decode_ordinary_batch(std::span lanes, ordinary_host_ingress->state_source_slots[row] = selectors.source; ordinary_host_ingress->state_destination_slots[row] = selectors.destination; ordinary_host_ingress->sampling[row] = request.sampling_host; + if (request.grammar) { + structured_round->fill(sequence.lane, *request.grammar, {}, device.stream); + } ensure_sequence_kv_mapped(sequence, frontier + 1, 0); } @@ -487,6 +496,12 @@ ProgramImpl::decode_mtp_batch(std::span lanes, mtp_host_ingress->state_destination_slots[row] = selectors.destination; mtp_host_ingress->rope_deltas[row] = sequence.rope_delta; mtp_host_ingress->sampling[row] = request.sampling_host; + if (request.grammar) { + structured_round->fill( + sequence.lane, *request.grammar, + {mtp_host_ingress->current_drafts.data() + row * draft_window, extent}, + device.stream); + } ensure_sequence_kv_mapped(sequence, frontier + extent + 1, std::min(capacity, frontier + extent + draft_window)); } @@ -685,6 +700,12 @@ ProgramImpl::decode_dflash_batch(std::span lanes, backend_kv_cache() ? frontier : 0U); } + structured_round->begin_dflash(); + for (std::size_t row = 0; row < lanes.size(); ++row) { + structured_round->set_dflash_row(row, lanes[row], + dflash_host_ingress->proposal_extents[row], + requests[lanes[row]].grammar); + } execution::DFlashBatchContext schedule_state{ {device, parameters, work, state_images->linear(), replay_records ? &*replay_records : nullptr, io, prefill_hidden, prefill_chunk, @@ -706,6 +727,7 @@ ProgramImpl::decode_dflash_batch(std::span lanes, static_cast(lanes.size())); device.synchronize(); } + structured_round->check(); timing.end_wait(); const double seconds = std::chrono::duration(Clock::now() - started).count(); diff --git a/src/models/qwen3_5/program/planning/request_plan.cpp b/src/models/qwen3_5/program/planning/request_plan.cpp index bd36137744..18289b6e1d 100644 --- a/src/models/qwen3_5/program/planning/request_plan.cpp +++ b/src/models/qwen3_5/program/planning/request_plan.cpp @@ -252,6 +252,7 @@ RequestBasePlan ProgramImpl::plan_request(const PreparedPromptData& prompt, ? FinishReason::OutputLimit : FinishReason::ContextCapacity; base->sampling = translate_sampling(options.sampling); + base->grammar = options.grammar; base->allow_prefix_reuse = options.allow_prefix_reuse; base->summary.publish_continuation = options.allow_prefix_reuse && prompt.identity.reusable && context_cache.enabled; @@ -452,6 +453,7 @@ std::optional ProgramImpl::inspect_lane( auto plan = std::make_unique(); plan->summary = base.summary; plan->sampling = base.sampling; + plan->grammar = base.grammar; plan->text_kv_page_entitlement = base.text_kv_page_entitlement; plan->backend_kv_page_entitlement = base.backend_kv_page_entitlement; plan->root_rebuild_work = base.root_rebuild_work; diff --git a/src/models/qwen3_5/program/planning/startup.cpp b/src/models/qwen3_5/program/planning/startup.cpp index 8a6e48090b..64257bc30c 100644 --- a/src/models/qwen3_5/program/planning/startup.cpp +++ b/src/models/qwen3_5/program/planning/startup.cpp @@ -255,6 +255,12 @@ PersistentLayout persistent_layout(const SequencePlanImpl& plan) { {dimension(parameters.model.resources().public_token_count), static_cast(plan.max_concurrency)}, "sampling token counts"); + out.grammar_masks = add_tensor( + builder, DType::I32, + {static_cast((parameters.model.resources().public_token_count + 31) / 32), + static_cast(plan.draft_window + 1), + static_cast(plan.max_concurrency)}, + "structured output token masks"); const auto config_words = static_cast( (sizeof(ops::SamplingConfig) + sizeof(std::int32_t) - 1) / sizeof(std::int32_t)); out.sampling_config = add_tensor( diff --git a/src/models/qwen3_5/program/planning/startup.h b/src/models/qwen3_5/program/planning/startup.h index 1f8f30e8dc..8f0a7881d9 100644 --- a/src/models/qwen3_5/program/planning/startup.h +++ b/src/models/qwen3_5/program/planning/startup.h @@ -42,6 +42,7 @@ struct PersistentLayout { std::optional score_hidden; std::optional token_counts; std::optional sampling_config; + std::optional grammar_masks; std::size_t bytes = 0; std::size_t kv_payload_bytes = 0; }; diff --git a/src/models/qwen3_5/program/prefill.cpp b/src/models/qwen3_5/program/prefill.cpp index 4aecc90970..e64e93c261 100644 --- a/src/models/qwen3_5/program/prefill.cpp +++ b/src/models/qwen3_5/program/prefill.cpp @@ -637,6 +637,7 @@ void ProgramImpl::start_sequence(std::uint32_t lane, SequenceState& sequence, : speculative_backend == SpeculativeBackend::DFlash ? prompt_tokens : 0U; ensure_sequence_kv_mapped(sequence, prompt_tokens, backend_materialized); + request.grammar = request_plan.grammar; install_sampling(sequence, request, request_plan.sampling); sequence.rope_delta = staged.prompt.rope_delta; set_device_i32(io.rope_delta, sequence.rope_delta); diff --git a/src/models/qwen3_5/program/program_impl.cpp b/src/models/qwen3_5/program/program_impl.cpp index 2cdddef968..a969ae0f1d 100644 --- a/src/models/qwen3_5/program/program_impl.cpp +++ b/src/models/qwen3_5/program/program_impl.cpp @@ -1,3 +1,4 @@ +#include "models/qwen3_5/program/structured_round.h" #include "models/qwen3_5/program/program_impl.h" #include "models/qwen3_5/program/context_work.h" #include "models/qwen3_5/program/context.h" @@ -243,6 +244,12 @@ ProgramImpl::ProgramImpl(const execution::Parameters& parameters_in, const Seque if (plan.persistent.sampling_config) { sampling_config = plan.persistent.sampling_config->bind(backing); } + if (plan.persistent.grammar_masks) { + structured_round = std::make_unique( + plan.persistent.grammar_masks->bind(backing), + parameters.model.resources().public_token_count, draft_window + 1, max_concurrency); + io.structured = structured_round.get(); + } active_continuations.fill(continuation_capacity); for (std::uint32_t lane = 0; lane < max_concurrency; ++lane) { lane_epochs[lane] = 1; } for (std::uint32_t index = 0; index < continuation_capacity; ++index) { diff --git a/src/models/qwen3_5/program/program_impl.h b/src/models/qwen3_5/program/program_impl.h index 4e9aa608f5..9b7a402491 100644 --- a/src/models/qwen3_5/program/program_impl.h +++ b/src/models/qwen3_5/program/program_impl.h @@ -180,6 +180,7 @@ struct RequestBasePlanImpl { std::uint32_t root_rebuild_tail_begin = 0; qwen3_5::PreparedContextCache context_cache; ops::SamplingConfig sampling; + std::shared_ptr grammar; std::uint32_t text_kv_page_entitlement = 0; std::uint32_t backend_kv_page_entitlement = 0; std::shared_ptr vision_control_plan; @@ -244,6 +245,7 @@ struct AdmissionCandidateImpl : ResourceCandidateState { std::vector capture_groups; std::vector shared_candidates; ops::SamplingConfig sampling; + std::shared_ptr grammar; std::uint32_t text_kv_page_entitlement = 0; std::uint32_t backend_kv_page_entitlement = 0; runtime::LaneId destination{}; @@ -401,6 +403,7 @@ struct RequestControl { Lifecycle lifecycle = Lifecycle::Empty; PendingCandidate pending; ops::SamplingConfig sampling_host; + std::shared_ptr grammar; GenerationTimings timings; SpeculativeStats speculative_stats; detail::PhysicalResources active_resources; @@ -601,6 +604,7 @@ class ProgramImpl { std::optional score_hidden; Tensor sampling_config; Tensor token_counts; + std::unique_ptr structured_round; std::vector continuation_states; std::vector continuation_slots; diff --git a/src/models/qwen3_5/program/round_buffers.h b/src/models/qwen3_5/program/round_buffers.h index c485743cf7..f8c3e41bfa 100644 --- a/src/models/qwen3_5/program/round_buffers.h +++ b/src/models/qwen3_5/program/round_buffers.h @@ -292,7 +292,10 @@ struct DFlashDecodeState { std::uint32_t batch_capacity, std::uint32_t draft_window); }; +class StructuredRound; + struct RoundState { + StructuredRound* structured = nullptr; std::optional ordinary; Tensor token; Tensor pos; diff --git a/src/models/qwen3_5/program/structured_round.cpp b/src/models/qwen3_5/program/structured_round.cpp new file mode 100644 index 0000000000..06b7d41003 --- /dev/null +++ b/src/models/qwen3_5/program/structured_round.cpp @@ -0,0 +1,67 @@ +#include "models/qwen3_5/program/structured_round.h" +#include "core/device.h" +#include +#include + +namespace ninfer::models::qwen3_5 { +StructuredRound::StructuredRound(Tensor masks, std::uint32_t vocab, std::uint32_t width, + std::uint32_t lanes) + : masks_(masks), words_((vocab + 31) / 32), width_(width), lanes_(lanes), + host_masks_(masks.bytes()), host_drafts_(std::max(1U, width - 1) * lanes * sizeof(TokenId)) { + std::fill_n(static_cast(host_masks_.data()), masks.bytes() / 4, + ~std::uint32_t{0}); +} + +const std::uint32_t* StructuredRound::device_mask(std::uint32_t lane) const { + return static_cast(masks_.data) + lane * width_ * words_; +} + +void StructuredRound::fill(std::uint32_t lane, const text::GrammarState& grammar, + std::span drafts, cudaStream_t stream) { + auto* host = static_cast(host_masks_.data()) + lane * width_ * words_; + grammar.fill_masks({host, (drafts.size() + 1) * words_}, drafts); + CUDA_CHECK(cudaMemcpyAsync(const_cast(device_mask(lane)), host, + (drafts.size() + 1) * words_ * sizeof(std::uint32_t), + cudaMemcpyHostToDevice, stream)); +} + +void StructuredRound::begin_dflash() { + rows_ = {}; + error_ = nullptr; +} + +void StructuredRound::set_dflash_row(std::uint32_t row, std::uint32_t lane, std::uint32_t extent, + std::shared_ptr grammar) { + rows_.at(row) = {lane, extent, std::move(grammar)}; +} + +void CUDART_CB StructuredRound::callback(void* opaque) noexcept { + auto& self = *static_cast(opaque); + try { + for (std::size_t row = 0; row < self.rows_.size(); ++row) { + const auto& entry = self.rows_[row]; + if (!entry.grammar) { continue; } + auto* mask = static_cast(self.host_masks_.data()) + + entry.lane * self.width_ * self.words_; + auto* drafts = + static_cast(self.host_drafts_.data()) + row * (self.width_ - 1); + entry.grammar->fill_masks({mask, (entry.extent + 1) * self.words_}, + {drafts, entry.extent}); + } + } catch (...) { self.error_ = std::current_exception(); } +} + +void StructuredRound::enqueue_dflash(const Tensor& drafts, cudaStream_t stream) { + // No CUDA calls or exceptions escape the host function. The transfers and host node are + // captured together: verification cannot observe masks until the CPU matcher has filled them. + CUDA_CHECK(cudaMemcpyAsync(host_drafts_.data(), drafts.data, drafts.bytes(), + cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaLaunchHostFunc(stream, callback, this)); + CUDA_CHECK(cudaMemcpyAsync(masks_.data, host_masks_.data(), masks_.bytes(), + cudaMemcpyHostToDevice, stream)); +} + +void StructuredRound::check() const { + if (error_) { std::rethrow_exception(error_); } +} +} // namespace ninfer::models::qwen3_5 diff --git a/src/models/qwen3_5/program/structured_round.h b/src/models/qwen3_5/program/structured_round.h new file mode 100644 index 0000000000..71d874d870 --- /dev/null +++ b/src/models/qwen3_5/program/structured_round.h @@ -0,0 +1,41 @@ +#pragma once +#include "core/arena.h" +#include "core/tensor.h" +#include "text/structured_output.h" +#include +#include +#include + +namespace ninfer::models::qwen3_5 { +// Program-owned addresses survive CUDA Graph capture/replay. Grammar state belongs to the +// request, never to a reusable KV checkpoint. Only Engine's output commit advances it. +class StructuredRound { +public: + StructuredRound(Tensor masks, std::uint32_t vocab, std::uint32_t width, std::uint32_t lanes); + const std::uint32_t* device_mask(std::uint32_t lane) const; + + int stride() const { return words_; } + + void fill(std::uint32_t lane, const text::GrammarState& grammar, + std::span drafts, cudaStream_t stream); + void begin_dflash(); + void set_dflash_row(std::uint32_t row, std::uint32_t lane, std::uint32_t extent, + std::shared_ptr grammar); + void enqueue_dflash(const Tensor& drafts, cudaStream_t stream); + void check() const; +private: + static void CUDART_CB callback(void* self) noexcept; + Tensor masks_; + int words_; + std::uint32_t width_, lanes_; + PinnedHostBuffer host_masks_, host_drafts_; + + struct Row { + std::uint32_t lane = 0, extent = 0; + std::shared_ptr grammar; + }; + + std::array rows_{}; + std::exception_ptr error_; +}; +} // namespace ninfer::models::qwen3_5 diff --git a/src/models/qwen3_5/program_sources.cmake b/src/models/qwen3_5/program_sources.cmake index fc6c00bbb5..f095340c19 100644 --- a/src/models/qwen3_5/program_sources.cmake +++ b/src/models/qwen3_5/program_sources.cmake @@ -1,4 +1,5 @@ target_sources(ninfer_model_runtime PRIVATE + qwen3_5/program/structured_round.cpp "${CMAKE_CURRENT_LIST_DIR}/measurement.cpp" "${CMAKE_CURRENT_LIST_DIR}/program/storage/draft_context.cpp" "${CMAKE_CURRENT_LIST_DIR}/state/decoder_state.cpp" diff --git a/src/ops/kernel/sampling.cuh b/src/ops/kernel/sampling.cuh index d964c7cd36..bf2d316ba1 100644 --- a/src/ops/kernel/sampling.cuh +++ b/src/ops/kernel/sampling.cuh @@ -28,7 +28,8 @@ __launch_bounds__(kSamplerBlock) __global__ if (!(cfg.temperature > 0.0f)) { float bv = -CUDART_INF_F; int bi = INT_MAX; - const bool penalties = cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f; + const bool penalties = cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f || + cfg.token_mask != nullptr; if (!penalties) { for (int v = tid; v < token_domain; v += blockDim.x) { const float x = __bfloat162float(logits[base + v]); @@ -117,7 +118,8 @@ __launch_bounds__(kSamplerBlock) __global__ unsigned long long keys[kSamplerItemsPerThread]; const bool greedy = !(cfg.temperature > 0.0f); - const bool penalties = cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f; + const bool penalties = + cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f || cfg.token_mask != nullptr; const int cap = greedy ? 1 : sampling_candidate_cap(cfg, token_domain); const std::int64_t base = static_cast(col) * physical_rows; const int tile_start = partial * kSamplerPartialTileItems; diff --git a/src/ops/kernel/sampling_device.cuh b/src/ops/kernel/sampling_device.cuh index f2909425e1..b8fe8421fa 100644 --- a/src/ops/kernel/sampling_device.cuh +++ b/src/ops/kernel/sampling_device.cuh @@ -247,6 +247,10 @@ __device__ __forceinline__ int sampling_dist_offset(int col, int j) { __device__ __forceinline__ float sampling_adjusted_logit(float raw, int v, const SamplingConfig& c, const std::int32_t* overlay = nullptr, int overlay_len = 0) { + if (c.token_mask != nullptr && + (c.token_mask[overlay_len * c.token_mask_stride + v / 32] & (1U << (v % 32))) == 0) { + return -CUDART_INF_F; + } float x = raw; if (c.presence_penalty == 0.0f && c.frequency_penalty == 0.0f) { return x; } int cnt = c.token_counts != nullptr ? c.token_counts[v] : 0; @@ -317,7 +321,7 @@ __device__ inline void sampling_normalize_support(const SamplingConfig& cfg, flo float cum = 0.0f; int support = 0; for (int j = 0; j < n; ++j) { - if (min_p_thresh >= 0.0f && prob[j] < min_p_thresh) { break; } + if (prob[j] == 0.0f || (min_p_thresh >= 0.0f && prob[j] < min_p_thresh)) { break; } cum += prob[j]; support = j + 1; if (top_p_active && cum >= top_p_target) { break; } diff --git a/src/ops/kernel/speculative_round.cuh b/src/ops/kernel/speculative_round.cuh index 9cf696a622..e05241da01 100644 --- a/src/ops/kernel/speculative_round.cuh +++ b/src/ops/kernel/speculative_round.cuh @@ -243,7 +243,8 @@ __launch_bounds__(kSamplerBlock) __global__ void speculative_accept_greedy_draft std::int32_t* row_tokens = licensed_tokens + row * cols; const __nv_bfloat16* row_logits = logits + static_cast(row) * cols * physical_rows; - const bool penalties = cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f; + const bool penalties = + cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f || cfg.token_mask != nullptr; if (!(cfg.temperature > 0.0f) && !penalties) { if (tid == 0) { @@ -394,7 +395,8 @@ __launch_bounds__(kSamplerBlock) __global__ void speculative_sampling_partial_to if (col > extent) { return; } const SamplingConfig cfg = configs[row]; const bool greedy = !(cfg.temperature > 0.0f); - const bool penalties = cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f; + const bool penalties = + cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f || cfg.token_mask != nullptr; if ((greedy && !penalties) || token_domain <= kSamplerTileItems) { return; } workspace = speculative_workspace_row(workspace, workspace_row_stride, row); if (partial == 0 && threadIdx.x == 0) { @@ -476,7 +478,8 @@ __launch_bounds__(kSamplerGroupBlock) __global__ void speculative_sampling_group std::int32_t* row_tokens = licensed_tokens + row * cols; if (token_domain <= kSamplerTileItems) { return; } const bool greedy = !(cfg.temperature > 0.0f); - const bool penalties = cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f; + const bool penalties = + cfg.presence_penalty != 0.0f || cfg.frequency_penalty != 0.0f || cfg.token_mask != nullptr; if (greedy && !penalties) { if constexpr (SparseProposal) { diff --git a/src/runtime/contract/request.h b/src/runtime/contract/request.h index 01d1cbd834..4652519418 100644 --- a/src/runtime/contract/request.h +++ b/src/runtime/contract/request.h @@ -5,11 +5,19 @@ #include #include +#include + +namespace ninfer::text { +class GrammarState; +} + namespace ninfer::runtime { // Engine has already selected the model/mode preset, applied every explicit override, // and validated these values before constructing the runtime request. struct ResolvedExecutionOptions { + StructuredOutputOptions structured_output; + std::shared_ptr grammar; ResolvedSamplingParameters sampling; std::uint32_t requested_output_tokens = 0; bool allow_prefix_reuse = true; diff --git a/src/runtime/engine/engine.cpp b/src/runtime/engine/engine.cpp index 2a41b843cd..5c2bc4f278 100644 --- a/src/runtime/engine/engine.cpp +++ b/src/runtime/engine/engine.cpp @@ -1,4 +1,5 @@ #include "ninfer/engine.h" +#include "text/structured_output.h" #include "core/device.h" #include "core/nvtx.h" @@ -32,12 +33,14 @@ runtime::ResolvedRequestOptions resolve_request_options(const ModelSamplingDefau if (options.execution.thinking.budget && *options.execution.thinking.budget == 0) { throw std::invalid_argument("thinking budget must be positive"); } + text::validate_structured_output(options.execution.structured_output); runtime::ResolvedRequestOptions resolved; resolved.execution.sampling = runtime::resolve_sampling(defaults, mode, options.execution.sampling); resolved.execution.requested_output_tokens = options.execution.requested_output_tokens; resolved.execution.allow_prefix_reuse = options.execution.allow_prefix_reuse; resolved.execution.thinking = options.execution.thinking; + resolved.execution.structured_output = std::move(options.execution.structured_output); resolved.stop = std::move(options.stop); resolved.output = options.output; return resolved; diff --git a/src/runtime/engine/engine_core.h b/src/runtime/engine/engine_core.h index e6fd782e70..21f496d8af 100644 --- a/src/runtime/engine/engine_core.h +++ b/src/runtime/engine/engine_core.h @@ -200,7 +200,9 @@ class EngineCore { std::shared_ptr request; try { auto output = instance_.frontend.make_output_session( - prompt, options.stop, options.output, options.execution.thinking); + prompt, options.stop, options.output, options.execution.thinking, + options.execution.structured_output); + options.execution.grammar = output.grammar_state(); const std::uint32_t capacity_output = max_context_ - prompt_summary.prompt_tokens + static_cast(1); try { diff --git a/src/serve/CMakeLists.txt b/src/serve/CMakeLists.txt index 9f8afa8f08..1bcdea46e8 100644 --- a/src/serve/CMakeLists.txt +++ b/src/serve/CMakeLists.txt @@ -1,6 +1,7 @@ add_library(ninfer_serve STATIC http_transport.cpp request_validation.cpp + structured_output.cpp openai_chat_request.cpp openai_chat_response.cpp openai_chat_http.cpp diff --git a/src/serve/anthropic_messages_request.cpp b/src/serve/anthropic_messages_request.cpp index 32cdc8583e..cc90b4f592 100644 --- a/src/serve/anthropic_messages_request.cpp +++ b/src/serve/anthropic_messages_request.cpp @@ -1,3 +1,4 @@ +#include "serve/structured_output.h" #include "serve/anthropic_messages.h" #include "serve/request_validation.h" @@ -899,9 +900,8 @@ void parse_effort(const Json& body, GenerationRequest& request, ParsePurpose pur if (!config.is_object()) { bad_request("output_config must be an object", "output_config"); } if (purpose == ParsePurpose::Messages && config.contains("format") && !config.at("format").is_null()) { - bad_request("output_config.format requires constrained decoding, which NInfer does not " - "provide", - "output_config.format", "output_config_format_not_supported"); + request.structured_output = + parse_structured_output(config.at("format"), false, "output_config.format"); } if (!config.contains("effort") || config.at("effort").is_null()) { return; } if (!config.at("effort").is_string()) { diff --git a/src/serve/openai_chat_request.cpp b/src/serve/openai_chat_request.cpp index 346117d828..0064642fcc 100644 --- a/src/serve/openai_chat_request.cpp +++ b/src/serve/openai_chat_request.cpp @@ -1,3 +1,4 @@ +#include "serve/structured_output.h" #include "serve/openai_chat.h" #include "serve/openai_common.h" #include "serve/request_validation.h" @@ -131,19 +132,6 @@ void validate_standard_output_controls(const Json& body) { } } - if (body.contains("response_format") && !body.at("response_format").is_null()) { - const Json& format = body.at("response_format"); - if (!format.is_object() || !format.contains("type") || !format.at("type").is_string()) { - bad_request("response_format must contain a string type", "response_format"); - } - if (format.at("type").get() != "text") { - bad_request( - "this response_format requires constrained output, which NInfer cannot guarantee; " - "only {\"type\":\"text\"} is available", - "response_format", "response_format_not_supported"); - } - } - if (body.contains("modalities") && !body.at("modalities").is_null()) { const Json& modalities = body.at("modalities"); if (!modalities.is_array() || modalities.empty()) { @@ -886,6 +874,10 @@ OpenAIChatRequest parse_chat_completion_request(const Json& body, const RequestL validate_compatibility_hints(body); OpenAIChatRequest output; + if (body.contains("response_format") && !body.at("response_format").is_null()) { + output.generation.structured_output = + parse_structured_output(body.at("response_format"), true, "response_format"); + } if (!body.contains("model") || !body.at("model").is_string() || body.at("model").get().empty()) { bad_request("missing required field: model", "model"); diff --git a/src/serve/openai_responses.h b/src/serve/openai_responses.h index 9cd1aab452..0a326c7301 100644 --- a/src/serve/openai_responses.h +++ b/src/serve/openai_responses.h @@ -31,6 +31,7 @@ struct OpenAIResponsesFunctionIdentity { struct OpenAIResponsesPromptRequest { std::string model; GenerationRequest generation; + nlohmann::json text_format = {{"type", "text"}}; std::vector input_turns; std::vector input_items; std::optional instructions; diff --git a/src/serve/openai_responses_request.cpp b/src/serve/openai_responses_request.cpp index af58651042..576735f41b 100644 --- a/src/serve/openai_responses_request.cpp +++ b/src/serve/openai_responses_request.cpp @@ -1,3 +1,4 @@ +#include "serve/structured_output.h" #include "serve/openai_responses.h" #include "serve/openai_common.h" #include "serve/request_validation.h" @@ -941,7 +942,7 @@ void parse_reasoning(const Json& body, OpenAIResponsesPromptRequest& out) { out.generation.reasoning_effort = *effort; } -void parse_text(const Json& body) { +void parse_text(const Json& body, OpenAIResponsesPromptRequest& prompt) { if (!body.contains("text") || body.at("text").is_null()) { return; } const Json& text = body.at("text"); if (!text.is_object()) { bad_request("text must be an object", "text"); } @@ -949,14 +950,8 @@ void parse_text(const Json& body) { reject_nonnull_unknown_members(text, allowed, "text"); if (text.contains("format") && !text.at("format").is_null()) { const Json& format = text.at("format"); - if (!format.is_object() || !format.contains("type") || !format.at("type").is_string()) { - bad_request("text.format must be a typed object", "text"); - } - if (format.at("type").get() != "text" || format.size() != 1) { - bad_request("structured text output requires constrained decoding, which the Engine " - "does not provide", - "text", "structured_outputs_not_supported"); - } + prompt.generation.structured_output = parse_structured_output(format, false, "text.format"); + prompt.text_format = format; } if (text.contains("verbosity") && !text.at("verbosity").is_null()) { if (!text.at("verbosity").is_string()) { @@ -1044,7 +1039,7 @@ ParsedPromptFields parse_prompt_fields(const Json& body, const RequestLimits& li "parallel_tool_calls", "parallel_tool_calls_not_supported"); } parse_reasoning(body, out.prompt); - parse_text(body); + parse_text(body, out.prompt); parse_truncation(body); parse_preserve_thinking(body, out.prompt); out.prompt.generation.max_tokens = limits.default_max_tokens; diff --git a/src/serve/openai_responses_response.cpp b/src/serve/openai_responses_response.cpp index 7bb8d1ef12..4a53cd1559 100644 --- a/src/serve/openai_responses_response.cpp +++ b/src/serve/openai_responses_response.cpp @@ -78,7 +78,7 @@ Json response_common(const std::string& id, std::int64_t created_at, {"service_tier", "default"}, {"store", request.store}, {"temperature", runtime.temperature}, - {"text", Json{{"format", Json{{"type", "text"}}}}}, + {"text", Json{{"format", request.prompt.text_format}}}, {"tool_choice", request.tool_choice}, {"tools", request.tools}, {"top_logprobs", 0}, diff --git a/src/serve/request.h b/src/serve/request.h index fcf83ada49..3ce4f1f74a 100644 --- a/src/serve/request.h +++ b/src/serve/request.h @@ -186,6 +186,7 @@ struct GenerationRequest { ninfer::PromptContinuationMode continuation = ninfer::PromptContinuationMode::NewAssistantTurn; bool allow_engine_automatic_shared_prefixes = true; SamplingParams sampling; + StructuredOutputOptions structured_output; [[nodiscard]] bool uses_tools() const noexcept { return !tools.empty() && tool_choice.mode != ToolChoiceMode::None; diff --git a/src/serve/structured_output.cpp b/src/serve/structured_output.cpp new file mode 100644 index 0000000000..f9b8832b3b --- /dev/null +++ b/src/serve/structured_output.cpp @@ -0,0 +1,46 @@ +#include "serve/structured_output.h" +#include "text/structured_output.h" + +namespace ninfer::serve { +StructuredOutputOptions parse_structured_output(const RequestJson& format, bool nested, + const std::string& param) { + const auto fail = [&](const std::string& message) -> void { + throw ApiException( + ApiError{.message = message, .param = param, .code = "invalid_response_format"}); + }; + if (!format.is_object() || !format.contains("type") || !format.at("type").is_string()) { + fail(param + " must contain a string type"); + } + const auto type = format.at("type").get(); + StructuredOutputOptions options; + if (type == "text" || type == "json_object") { + if (format.size() != 1) { fail("unexpected members in " + param); } + options.kind = + type == "text" ? StructuredOutputKind::None : StructuredOutputKind::JsonObject; + } else if (type == "json_schema") { + const auto& spec = + nested && format.contains("json_schema") ? format.at("json_schema") : format; + if (nested && (format.size() != 2 || !format.contains("json_schema"))) { + fail("json_schema requires a json_schema object"); + } + if (!spec.is_object() || !spec.contains("schema") || + (!spec.at("schema").is_object() && !spec.at("schema").is_boolean())) { + fail("json_schema requires an object or boolean schema"); + } + for (const auto& [key, value] : spec.items()) { + if (key == "schema" || (!nested && key == "type")) { continue; } + if ((key == "name" || key == "description") && value.is_string()) { continue; } + if (key == "strict" && (value.is_boolean() || value.is_null())) { continue; } + fail("unsupported or invalid schema format member: " + key); + } + options.kind = StructuredOutputKind::JsonSchema; + options.schema = spec.at("schema").dump(); + } else { + fail("unsupported response format type: " + type); + } + try { + text::validate_structured_output(options); + } catch (const std::invalid_argument& e) { fail(e.what()); } + return options; +} +} // namespace ninfer::serve diff --git a/src/serve/structured_output.h b/src/serve/structured_output.h new file mode 100644 index 0000000000..077a2487bf --- /dev/null +++ b/src/serve/structured_output.h @@ -0,0 +1,9 @@ +#pragma once +#include "serve/request.h" +#include "serve/request_json.h" + +namespace ninfer::serve { +// Chat nests schema metadata under json_schema; Responses and Anthropic use a flat format. +StructuredOutputOptions parse_structured_output(const RequestJson& format, bool nested, + const std::string& param); +} // namespace ninfer::serve diff --git a/src/serve/translate.cpp b/src/serve/translate.cpp index eace8e4630..f1bb594648 100644 --- a/src/serve/translate.cpp +++ b/src/serve/translate.cpp @@ -150,6 +150,16 @@ ResolvedPromptSemantics resolve_prompt_semantics(const GenerationRequest& reques effort = nested; } kwargs.erase("reasoning_effort"); + if (request.structured_output.kind != StructuredOutputKind::None) { + if (thinking == true || (effort && *effort != RequestedReasoningEffort::None) || + request.thinking_budget || request.uses_tools() || !request.stop_strings.empty() || + request.continuation != PromptContinuationMode::NewAssistantTurn) { + invalid_prompt_option("structured output requires thinking disabled, default stops, a " + "new assistant turn, and no active tools", + "response_format", "incompatible_structured_output"); + } + thinking = false; + } ResolvedPromptSemantics result{ .enable_thinking = thinking ? thinking : server.enable_thinking, .preserve_thinking = preserve ? preserve : server.preserve_thinking, @@ -315,13 +325,16 @@ ninfer::RequestOptions to_request_options(const GenerationRequest& request, ninfer::RequestOptions options; options.execution.requested_output_tokens = static_cast(request.max_tokens); options.execution.allow_prefix_reuse = allow_prefix_reuse; + options.execution.structured_output = request.structured_output; if (semantics.enable_thinking != false) { options.execution.thinking.budget = request.thinking_budget ? request.thinking_budget : server.default_thinking_budget; } options.execution.sampling = resolve_sampling_overrides(request.sampling, server); options.output.raw = false; - options.output.preserve_special_tokens = request.uses_tools() || request.has_tool_history(); + options.output.preserve_special_tokens = + request.structured_output.kind == StructuredOutputKind::None && + (request.uses_tools() || request.has_tool_history()); options.output.tool_name_max_length = static_cast(request.tool_name_max_length); options.stop.strings.reserve(request.stop_strings.size() * (request.stop_strings_apply_to_reasoning ? 2U : 1U)); diff --git a/src/text/CMakeLists.txt b/src/text/CMakeLists.txt index e7b250674c..bf01b73d09 100644 --- a/src/text/CMakeLists.txt +++ b/src/text/CMakeLists.txt @@ -1,8 +1,10 @@ add_library(ninfer_text STATIC unicode.cpp + structured_output.cpp ${PROJECT_SOURCE_DIR}/third_party/utf8proc/utf8proc.c ) ninfer_internal_includes(ninfer_text) +target_link_libraries(ninfer_text PRIVATE ninfer_xgrammar ninfer::json) target_include_directories(ninfer_text PRIVATE ${PROJECT_SOURCE_DIR}/third_party) # Language evaluation reuses the engine's Unicode and JSON dependencies. diff --git a/src/text/structured_output.cpp b/src/text/structured_output.cpp new file mode 100644 index 0000000000..5251d1f089 --- /dev/null +++ b/src/text/structured_output.cpp @@ -0,0 +1,193 @@ +#include "text/structured_output.h" +#include +#include +#include +#include +#include +#include + +namespace ninfer::text { +namespace { +using Json = nlohmann::ordered_json; + +void schema_check(const Json& s) { + if (s.is_boolean()) { return; } + if (!s.is_object()) { throw std::invalid_argument("JSON schema must be an object or boolean"); } + static const std::unordered_set annotations = { + "$schema", "title", "description", "default", + "examples", "$comment", "$defs", "definitions"}; + static const std::unordered_set supported = { + "type", "properties", "required", "additionalProperties", + "items", "prefixItems", "minItems", "maxItems", + "minLength", "maxLength", "enum", "const", + "anyOf", "$ref"}; + for (const auto& [key, value] : s.items()) { + if (key == "$schema" && value != "https://json-schema.org/draft/2020-12/schema" && + value != "http://json-schema.org/draft-07/schema#") { + throw std::invalid_argument("unsupported JSON Schema dialect"); + } + if ((key == "minLength" || key == "maxLength" || key == "minItems" || key == "maxItems") && + (!value.is_number_integer() || value < 0 || value > 2147483647)) { + throw std::invalid_argument(key + " must be a nonnegative 32-bit integer"); + } + if (!annotations.contains(key) && !supported.contains(key)) { + throw std::invalid_argument("unsupported JSON schema keyword: " + key); + } + if (key == "$ref" && + (!value.is_string() || (value != "#" && !value.get().starts_with("#/")))) { + throw std::invalid_argument("JSON schema supports only local fragment $ref values"); + } + if (key == "$ref" && value != "#") { + const auto ref = value.get(); + // The pinned compiler interprets literal object paths, not RFC 6901 escapes. + // Reject ambiguous spellings rather than resolving a different schema silently. + if (ref.find_first_of("~%") != std::string::npos || ref.ends_with('/') || + ref.find("//") != std::string::npos) { + throw std::invalid_argument( + "$ref requires nonempty, unescaped object path segments"); + } + } + if (key == "properties" || key == "$defs" || key == "definitions") { + if (!value.is_object()) { throw std::invalid_argument(key + " must be an object"); } + for (const auto& child : value) { schema_check(child); } + } else if (key == "items" || key == "additionalProperties") { + schema_check(value); + } else if (key == "anyOf" || key == "prefixItems") { + if (!value.is_array()) { throw std::invalid_argument(key + " must be an array"); } + for (const auto& child : value) { schema_check(child); } + } + } + // XGrammar prioritizes these branches over their siblings. Disallow combinations that + // would otherwise silently discard constraints (annotations and definitions are harmless). + for (const char* branch : {"$ref", "const", "enum", "anyOf"}) { + if (!s.contains(branch)) { continue; } + for (const auto& [key, value] : s.items()) { + if (key != branch && !annotations.contains(key)) { + if (key == "type" && + (std::string_view(branch) == "enum" || std::string_view(branch) == "const")) { + const auto matches = [&](const Json& v) { + if (!value.is_string()) { return false; } + const auto type = value.get(); + return (type == "string" && v.is_string()) || + (type == "integer" && v.is_number_integer()) || + (type == "number" && v.is_number()) || + (type == "boolean" && v.is_boolean()) || + (type == "null" && v.is_null()) || + (type == "array" && v.is_array()) || + (type == "object" && v.is_object()); + }; + if (s.contains("const") && matches(s.at("const"))) { continue; } + if (s.contains("enum") && s.at("enum").is_array() && + std::all_of(s.at("enum").begin(), s.at("enum").end(), matches)) { + continue; + } + } + throw std::invalid_argument(std::string(branch) + " cannot be combined with " + + key); + } + } + } +} +} // namespace + +void validate_structured_output(const StructuredOutputOptions& options) { + if (options.kind == StructuredOutputKind::JsonSchema) { + try { + schema_check(Json::parse(options.schema)); + } catch (const Json::exception& e) { + throw std::invalid_argument(std::string("invalid JSON schema: ") + e.what()); + } + } else if (!options.schema.empty()) { + throw std::invalid_argument("schema requires JsonSchema mode"); + } +} + +struct GrammarState::Impl { + xgrammar::GrammarMatcher matcher; + int vocab_size; + + Impl(xgrammar::GrammarMatcher matcher, int vocab_size) + : matcher(std::move(matcher)), vocab_size(vocab_size) {} +}; + +GrammarState::GrammarState(std::unique_ptr impl) : impl_(std::move(impl)) {} + +GrammarState::~GrammarState() = default; +GrammarState::GrammarState(GrammarState&&) noexcept = default; +GrammarState& GrammarState::operator=(GrammarState&&) noexcept = default; + +std::unique_ptr GrammarState::fork() const { + return std::unique_ptr( + new GrammarState(std::make_unique(impl_->matcher.Fork(), impl_->vocab_size))); +} + +void GrammarState::accept(std::span tokens) { + for (TokenId token : tokens) { + if (impl_->matcher.IsTerminated() || !impl_->matcher.AcceptToken(token)) { + throw std::logic_error("generated token " + std::to_string(token) + + " violates the structured output grammar"); + } + } +} + +void GrammarState::fill_masks(std::span masks, + std::span drafts) const { + const int words = xgrammar::GetBitmaskSize(impl_->vocab_size); + if (masks.size() != (drafts.size() + 1) * words) { + throw std::logic_error("incorrect grammar mask shape"); + } + auto matcher = impl_->matcher.Fork(); + std::fill(masks.begin(), masks.end(), ~std::uint32_t{0}); + for (std::size_t col = 0; col <= drafts.size(); ++col) { + if (matcher.IsTerminated()) { break; } // no later column can be published past EOS + std::int64_t shape[2] = {1, words}; + DLTensor tensor{}; + tensor.data = masks.data() + col * words; + tensor.device = {kDLCPU, 0}; + tensor.ndim = 2; + tensor.dtype = {kDLInt, 32, 1}; + tensor.shape = shape; + matcher.FillNextTokenBitmask(&tensor); + if (impl_->vocab_size % 32) { + masks[(col + 1) * words - 1] &= (1U << (impl_->vocab_size % 32)) - 1; + } + const auto row = masks.subspan(col * words, words); + if (std::none_of(row.begin(), row.end(), [](auto word) { return word != 0; })) { + throw std::runtime_error("structured output grammar has no admissible next token"); + } + if (col < drafts.size() && !matcher.AcceptToken(drafts[col])) { break; } + } +} + +struct StructuredCompiler::Impl { + xgrammar::TokenizerInfo tokenizer; + xgrammar::GrammarCompiler compiler; + std::mutex mutex; + + Impl(std::vector vocab, std::vector stops) + : tokenizer(vocab, xgrammar::VocabType::RAW, static_cast(vocab.size()), stops), + compiler(tokenizer, 4, true, 256 * 1024 * 1024) {} +}; + +StructuredCompiler::StructuredCompiler(std::vector vocab, std::vector stops) + : impl_(std::make_unique(std::move(vocab), std::move(stops))) {} + +StructuredCompiler::~StructuredCompiler() = default; + +std::shared_ptr StructuredCompiler::compile(const StructuredOutputOptions& options) { + validate_structured_output(options); + if (options.kind == StructuredOutputKind::None) { return {}; } + std::lock_guard lock(impl_->mutex); + try { + // strict_mode=false retains JSON Schema defaults for additional properties/items. + auto grammar = impl_->compiler.CompileJSONSchema( + options.kind == StructuredOutputKind::JsonObject ? "{\"type\":\"object\"}" + : options.schema, + true, std::nullopt, std::nullopt, false, 8); + return std::shared_ptr(new GrammarState(std::make_unique( + xgrammar::GrammarMatcher(grammar), impl_->tokenizer.GetVocabSize()))); + } catch (const std::exception& e) { + throw std::invalid_argument(std::string("cannot compile JSON schema: ") + e.what()); + } +} +} // namespace ninfer::text diff --git a/src/text/structured_output.h b/src/text/structured_output.h new file mode 100644 index 0000000000..02eb09f050 --- /dev/null +++ b/src/text/structured_output.h @@ -0,0 +1,38 @@ +#pragma once +#include "ninfer/types.h" +#include +#include +#include +#include +#include + +namespace ninfer::text { +// Reject unsupported constraints instead of letting a compiler silently weaken a schema. +void validate_structured_output(const StructuredOutputOptions& options); + +class GrammarState { +public: + ~GrammarState(); + GrammarState(GrammarState&&) noexcept; + GrammarState& operator=(GrammarState&&) noexcept; + [[nodiscard]] std::unique_ptr fork() const; + void accept(std::span tokens); + // Column i is conditioned on drafts[0..i). Unreachable suffix columns are unrestricted. + void fill_masks(std::span masks, std::span drafts) const; +private: + struct Impl; + explicit GrammarState(std::unique_ptr impl); + std::unique_ptr impl_; + friend class StructuredCompiler; +}; + +class StructuredCompiler { +public: + StructuredCompiler(std::vector decoded_vocab, std::vector stop_tokens); + ~StructuredCompiler(); + std::shared_ptr compile(const StructuredOutputOptions& options); +private: + struct Impl; + std::unique_ptr impl_; +}; +} // namespace ninfer::text diff --git a/tests/README.md b/tests/README.md index 51618fd22d..0598ce0356 100644 --- a/tests/README.md +++ b/tests/README.md @@ -226,3 +226,28 @@ Arguments are K, Graph enabled, optimized head enabled, maximum B, target KV (`b Vision enabled, and extra Device StateImage slots. Defaults are `15 1 1 8 bf16 0 3`. Run GPU integration tests serially. The individual Op suites remain the numerical/state-transition oracle; the fixed Engine fixture does not define bit parity across arbitrary floating-point routes. + +## Structured output + +The CPU grammar and protocol tests and the GPU sampling/speculative tests qualify this path: + +```bash +ctest --test-dir build --output-on-failure -R 'ninfer_(structured_output|sampling|speculative_round|openai_schema|openai_responses|anthropic_schema|cli_options)_test$' +``` + +`tests/test_structured_output_live.py` runs a temporary loopback server and validates completed +outputs with the independent Python `jsonschema` validator. Use Python 3.11 with `requests` and +`jsonschema` installed in a test environment. Supply an explicit v3 artifact containing the +selected speculative backends; the test never downloads or converts weights. + +```bash +python tests/test_structured_output_live.py \ + --server build/apps/ninfer-serve --artifact /absolute/path/model.ninfer \ + --output-dir work/structured-live --modes none mtp dflash2 dflash2-eager +``` + +It checks conflicting prompts, greedy and stochastic generation, schema versus JSON object mode, +concurrent and mixed traffic, prefix reuse, SSE, token limits, disconnect cleanup, compile errors, +Responses, and Anthropic Messages. `--concurrency 8 --draft-tokens 15 --modes dflash2` exercises +the largest draft and batch dimensions. A separate DFlash-capable artifact can use `--modes dflash`. +The server is terminated on success or failure. An occupied test port causes the test to stop. diff --git a/tests/cmake/CoreTests.cmake b/tests/cmake/CoreTests.cmake index e0e4a652c7..68fa1954a8 100644 --- a/tests/cmake/CoreTests.cmake +++ b/tests/cmake/CoreTests.cmake @@ -50,3 +50,7 @@ ninfer_add_test(ninfer_jinja_test add_test(NAME ninfer_chat_templates_test COMMAND ${Python3_EXECUTABLE} -B ${PROJECT_SOURCE_DIR}/tests/text/test_chat_templates.py $) + +ninfer_add_test(ninfer_structured_output_test + SOURCES "${CMAKE_CURRENT_LIST_DIR}/../text/test_structured_output.cpp" + LIBRARIES ninfer_text ninfer::json) diff --git a/tests/models/qwen3_5/test_structured_round.cpp b/tests/models/qwen3_5/test_structured_round.cpp new file mode 100644 index 0000000000..ea0d469a94 --- /dev/null +++ b/tests/models/qwen3_5/test_structured_round.cpp @@ -0,0 +1,73 @@ +#include "models/qwen3_5/program/structured_round.h" +#include "core/decode_graph.h" +#include "core/device.h" +#include "ops/op_tester.h" +#include +using namespace ninfer; +using namespace ninfer::test; + +int main() { + if (cuda_unavailable()) { return 77; } + try { + constexpr int vocab_size = 129, words = 5, width = 2, lanes = 8; + DeviceContext device(0); + DeviceArena arena(words * width * lanes * 4 + 256); + auto masks = arena.alloc(DType::I32, {words, width, lanes}); + auto drafts = arena.alloc(DType::I32, {width - 1, lanes}); + models::qwen3_5::StructuredRound round(masks, vocab_size, width, lanes); + std::vector vocab(vocab_size); + for (int i = 0; i < 128; ++i) { vocab[i] = std::string(1, static_cast(i)); } + text::StructuredCompiler compiler(vocab, {128}); + auto grammar = compiler.compile({StructuredOutputKind::JsonObject, {}}); + auto set_drafts = [&](int token) { + std::vector ids(lanes, token); + CUDA_CHECK(cudaMemcpyAsync(drafts.data, ids.data(), ids.size() * sizeof(int), + cudaMemcpyHostToDevice, device.stream)); + device.synchronize(); + }; + DecodeGraphDefinition definition; + definition.capture(device.stream, [&] { round.enqueue_dflash(drafts, device.stream); }); + DecodeGraphExecutable executable; + executable.instantiate(definition); + // Changing compact rows and physical lanes must work without recapturing the graph. + for (int pass = 0; pass < 3; ++pass) { + round.begin_dflash(); + round.set_dflash_row(pass, 7 - pass, 1, grammar); + set_drafts('{'); + executable.launch(device.stream); + device.synchronize(); + round.check(); + const auto bits = + from_device(round.device_mask(7 - pass), words * width); + if (!(bits['{' / 32] & (1U << ('{' % 32))) || (bits[4] & 1U) || + (bits[words + 4] & 1U)) { + std::cerr << "pass=" << pass << " bits="; + for (auto b : bits) { std::cerr << std::hex << b << " "; } + throw std::runtime_error("graph mask/EOS mismatch"); + } + } + round.begin_dflash(); + auto dead_vocab = std::vector(vocab_size, "q"); + dead_vocab.back().clear(); + text::StructuredCompiler dead_compiler(dead_vocab, {128}); + auto dead_grammar = dead_compiler.compile({StructuredOutputKind::JsonObject, {}}); + round.set_dflash_row(0, 0, 1, dead_grammar); + set_drafts(0); // no vocabulary token can start JSON: matcher must surface the dead end + executable.launch(device.stream); + device.synchronize(); + bool raised = false; + try { + round.check(); + } catch (const std::exception&) { raised = true; } + if (!raised) { throw std::runtime_error("host callback exception was not surfaced"); } + round.begin_dflash(); + set_drafts('{'); + executable.launch(device.stream); + device.synchronize(); + round.check(); // reset clears both retained requests and the error + std::cout << "OK structured graph: lane remap, replay, exception transport, reset\n"; + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + return 1; + } +} diff --git a/tests/models/qwen3_5/tests.cmake b/tests/models/qwen3_5/tests.cmake index b36aa51ddf..14fe44debe 100644 --- a/tests/models/qwen3_5/tests.cmake +++ b/tests/models/qwen3_5/tests.cmake @@ -98,3 +98,8 @@ ninfer_add_test(ninfer_qwen3_5_visual_scatter_test set_tests_properties( ninfer_qwen3_5_visual_scatter_test PROPERTIES SKIP_RETURN_CODE 77) + +ninfer_add_test(ninfer_qwen3_5_structured_round_test + SOURCES "${CMAKE_CURRENT_LIST_DIR}/test_structured_round.cpp" + LIBRARIES ninfer_model_runtime ninfer_core) +set_tests_properties(ninfer_qwen3_5_structured_round_test PROPERTIES SKIP_RETURN_CODE 77) diff --git a/tests/ops/test_sampling.cpp b/tests/ops/test_sampling.cpp index 11cb64950b..948cdf07ae 100644 --- a/tests/ops/test_sampling.cpp +++ b/tests/ops/test_sampling.cpp @@ -43,7 +43,8 @@ bool same_config(const ops::SamplingConfig& a, const ops::SamplingConfig& b) { return a.temperature == b.temperature && a.top_k == b.top_k && a.top_p == b.top_p && a.min_p == b.min_p && a.presence_penalty == b.presence_penalty && a.frequency_penalty == b.frequency_penalty && a.seed == b.seed && - a.token_counts == b.token_counts; + a.token_counts == b.token_counts && a.token_mask == b.token_mask && + a.token_mask_stride == b.token_mask_stride; } std::vector bf16_bits(const std::vector& values) { @@ -612,6 +613,48 @@ int increment_counts_contract() { return failures; } +int masked_sampling_contract() { + int failures = 0; + for (int domain : {64, 257, 248077}) { + const int words = (domain + 31) / 32; + std::vector mask(words, 0); + std::vector column(domain, 1000.0f), reference(domain, -INFINITY); + const std::vector allowed{3, 7, domain - 1}; + for (std::size_t i = 0; i < allowed.size(); ++i) { + int id = allowed[i]; + mask[id / 32] |= 1U << (id % 32); + column[id] = reference[id] = 1.0f - static_cast(i); + } + auto device_mask = to_device(mask); + ops::SamplingConfig cfg; + cfg.token_mask = static_cast(device_mask.p); + cfg.token_mask_stride = words; + auto greedy = run_homogeneous_batch(repeat_column(column, 8), domain, domain, 8, cfg, 0, + ops::kSamplePurposeDecode); + failures += greedy.integrity_failures; + failures += + verify_exact("masked argmax before raw winner", greedy.tokens, std::vector(8, 3)); + cfg.temperature = 0.8f; + cfg.top_k = 20; + cfg.top_p = 0.95f; + cfg.min_p = 0.05f; + cfg.seed = 9517; + auto samples = run_repeated(column, domain, 4096, 8, cfg, 123, ops::kSamplePurposeDecode); + failures += samples.integrity_failures; + failures += verify_distribution("masked distribution before filters", samples.tokens, + distribution_oracle(reference, domain, cfg)); + std::fill(mask.begin(), mask.end(), 0); + mask[(domain - 1) / 32] = 1U << ((domain - 1) % 32); + auto singleton = to_device(mask); + cfg.token_mask = static_cast(singleton.p); + auto one = run_homogeneous_batch(repeat_column(column, 8), domain, domain, 8, cfg, 42, + ops::kSamplePurposePrefill); + failures += verify_exact("singleton mask with top-k=20", one.tokens, + std::vector(8, domain - 1)); + } + return failures; +} + } // namespace int main() { @@ -633,6 +676,7 @@ int main() { std::cerr << "sampling workspace accepted an invalid lane interval\n"; ++failures; } catch (const std::invalid_argument&) {} + failures += masked_sampling_contract(); failures += greedy_contract(); failures += deterministic_stochastic_contract(); failures += heterogeneous_batch_contract(); diff --git a/tests/ops/test_speculative_round.cpp b/tests/ops/test_speculative_round.cpp index bd0fe5fa72..1a42150490 100644 --- a/tests/ops/test_speculative_round.cpp +++ b/tests/ops/test_speculative_round.cpp @@ -104,6 +104,7 @@ struct SparseAcceptSuite { const int kSparseColumns; const int kSparseBatch; std::size_t observed_workspace = 0; + std::vector masks; SparseAcceptSuite(int drafts, int batch) : kSparseDrafts(drafts), kSparseColumns(drafts + 1), kSparseBatch(batch) {} @@ -154,6 +155,11 @@ struct SparseAcceptSuite { const std::vector& token_counts, const std::vector& drafts) { const auto adjusted = [&](int token) { + const int words = (kSparseTokenDomain + 31) / 32; + if (!masks.empty() && (masks[(row * kSparseColumns + column) * words + token / 32] & + (1U << (token % 32))) == 0) { + return -std::numeric_limits::infinity(); + } double value = bf16_to_f32(logits[sparse_logit_index(row, column, token)]); int count = token_counts[static_cast(row) * kSparseTokenDomain + token]; for (int previous = 0; previous < column; ++previous) { @@ -388,6 +394,15 @@ struct SparseAcceptSuite { DeviceBuffer d_extents = to_device(extents); DeviceBuffer d_token_counts = to_device(token_counts); std::vector device_configs = host_configs; + DeviceBuffer d_masks = to_device(masks.empty() ? std::vector{~0U} : masks); + if (!masks.empty()) { + for (int row = 0; row < kSparseBatch; ++row) { + device_configs[row].token_mask_stride = (kSparseTokenDomain + 31) / 32; + device_configs[row].token_mask = + static_cast(d_masks.p) + + row * kSparseColumns * device_configs[row].token_mask_stride; + } + } for (int row = 0; row < kSparseBatch; ++row) { device_configs[static_cast(row)].token_counts = static_cast(d_token_counts.p) + @@ -610,7 +625,8 @@ struct SparseAcceptSuite { anchors, configs, history, {!general}); } - int generated_general_case() { + int generated_general_case(bool masked = false) { + masks.clear(); std::vector targets(kSparseColumns * kSparseBatch), drafts(kSparseDrafts * kSparseBatch); std::vector logits(static_cast(kSparsePhysicalRows) * @@ -670,7 +686,30 @@ struct SparseAcceptSuite { logits[sparse_logit_index(row, col, v)] = f32_to_bf16(100.0f); } } - return execute_sparse_accept_case("sparse general K=" + std::to_string(kSparseDrafts) + + if (masked) { + const int words = (kSparseTokenDomain + 31) / 32; + masks.assign(kSparseColumns * kSparseBatch * words, 0); + for (int row = 0; row < kSparseBatch; ++row) { + for (int col = 0; col < kSparseColumns; ++col) { + auto allow = [&](int token) { + masks[(row * kSparseColumns + col) * words + token / 32] |= 1U + << (token % 32); + }; + const int base = 10000 + row * 4096 + col * 32; + for (int rank = 1; rank <= 20; ++rank) { + if ((rank + row + col) % 3 == 0) { allow(base + rank); } + } + const int target = targets[row * kSparseColumns + col]; + if (col != row % kSparseDrafts) { allow(target); } + // A forbidden raw winner must be removed BEFORE target truncation. + logits[sparse_logit_index(row, col, kSparseTokenDomain - 1)] = + f32_to_bf16(1000.0f); + targets[row * kSparseColumns + col] = kSparseTokenDomain - 1; + } + } + } + return execute_sparse_accept_case(std::string(masked ? "masked " : "") + + "sparse general K=" + std::to_string(kSparseDrafts) + " B=" + std::to_string(kSparseBatch), targets, logits, drafts, ids, q, extents, lengths, anchors, configs, history, {false}); @@ -941,6 +980,30 @@ int greedy_accept_case(int k, int accepted_count, int token_domain = 64) { ops::SamplingConfig{}, token_counts, expected); } +int masked_greedy_drafts_case(int k, int domain, bool stochastic, int rejection) { + const int words = (domain + 31) / 32; + std::vector mask(words * (k + 1), 0); + std::vector targets(k + 1, 0), drafts(k), counts(domain, 0); + std::vector logits(static_cast(domain) * (k + 1), + f32_to_bf16(1000.0f)); + for (int col = 0; col <= k; ++col) { + const int allowed = domain - 1 - col; + mask[col * words + allowed / 32] = 1U << (allowed % 32); + logits[static_cast(col) * domain + allowed] = f32_to_bf16(-2.0f); + if (col < k) { drafts[col] = col == rejection ? 0 : allowed; } + } + auto d_mask = to_device(mask); + ops::SamplingConfig config; + config.temperature = stochastic ? 0.8f : 0.0f; + config.top_k = 20; + config.top_p = 1.0f; + config.token_mask = static_cast(d_mask.p); + config.token_mask_stride = words; + const auto expected = accept_state_oracle(drafts, rejection, domain - 1 - rejection, 256); + return execute_accept_case("per-position grammar mask", targets, logits, domain, drafts, 256, + domain, config, counts, expected); +} + int deterministic_sampling_case() { constexpr int physical_rows = 248320; constexpr int token_domain = 248077; @@ -1231,6 +1294,18 @@ int main(int argc, char** argv) { failures += greedy_accept_case(15, 7, 257); failures += greedy_penalty_case(64); failures += greedy_penalty_case(257); + for (int domain : {64, 257, 248077}) { + for (bool stochastic : {false, true}) { + for (int rejection : {0, 2, 5}) { + failures += masked_greedy_drafts_case(5, domain, stochastic, rejection); + } + } + } + for (int k : {1, 5, 15}) { + for (int batch : {1, 8}) { + failures += SparseAcceptSuite(k, batch).generated_general_case(true); + } + } failures += deterministic_sampling_case(); failures += batched_sampling_workspace_stride_case(); std::size_t sparse_peak = 0; diff --git a/tests/test_anthropic_schema.cpp b/tests/test_anthropic_schema.cpp index c014b06b01..0ca740a336 100644 --- a/tests/test_anthropic_schema.cpp +++ b/tests/test_anthropic_schema.cpp @@ -133,8 +133,13 @@ int test_envelope_and_field_policy() { failures += check(api_param([&] { (void)parse(body); }) == "top_k", "Engine top_k range was not enforced"); body = base_request(); + body["output_config"] = + Json{{"format", Json{{"type", "json_schema"}, {"schema", Json{{"type", "object"}}}}}}; + failures += check(parse(body).generation.structured_output.kind == + ninfer::StructuredOutputKind::JsonSchema, + "Anthropic schema retained"); body["output_config"] = Json{{"format", Json{{"type", "json_schema"}}}}; - failures += check(api_code([&] { (void)parse(body); }) == "output_config_format_not_supported", + failures += check(api_code([&] { (void)parse(body); }) == "invalid_response_format", "structured output was silently downgraded"); body = base_request(); body["container"] = "container_1"; diff --git a/tests/test_cli_options.cpp b/tests/test_cli_options.cpp index 041d5838df..d878b380fc 100644 --- a/tests/test_cli_options.cpp +++ b/tests/test_cli_options.cpp @@ -104,5 +104,15 @@ int main() { (void)parse({"ninfer-cli", "model.ninfer", "--prompt", "hello", "--top-k", "21"}); }), "CLI accepted top_k beyond the executable candidate domain"); + const auto structured = parse({"ninfer", "model.ninfer", "--prompt", "hello", "--json"}); + failures += + check(structured.structured_output.kind == ninfer::StructuredOutputKind::JsonObject && + structured.enable_thinking == false, + "CLI JSON mode"); + failures += check( + rejects([] { + (void)parse({"ninfer", "model.ninfer", "--prompt", "x", "--json", "--raw-output"}); + }), + "raw JSON output should be rejected"); return failures == 0 ? 0 : 1; } diff --git a/tests/test_openai_responses.cpp b/tests/test_openai_responses.cpp index 01a329571a..614f324f07 100644 --- a/tests/test_openai_responses.cpp +++ b/tests/test_openai_responses.cpp @@ -723,10 +723,19 @@ int test_explicit_rejections() { "strict function schema is rejected explicitly"); value = base; + value["text"] = Json{{"format", Json{{"type", "json_schema"}, + {"name", "answer"}, + {"strict", true}, + {"schema", Json{{"type", "object"}}}}}}; + const auto structured = parse_openai_responses_create_request(value, limits()); + failures += check(structured.prompt.generation.structured_output.kind == + ninfer::StructuredOutputKind::JsonSchema && + structured.prompt.text_format.at("type") == "json_schema", + "Responses schema retained"); value["text"] = Json{{"format", Json{{"type", "json_schema"}}}}; failures += check(api_code([&] { (void)parse_openai_responses_create_request(value, limits()); - }) == "structured_outputs_not_supported", + }) == "invalid_response_format", "structured output is rejected explicitly"); value = base; diff --git a/tests/test_openai_schema.cpp b/tests/test_openai_schema.cpp index d5887eb3ab..855c4059e7 100644 --- a/tests/test_openai_schema.cpp +++ b/tests/test_openai_schema.cpp @@ -132,6 +132,42 @@ int test_request_envelope_and_sampling() { return failures; } +int test_structured_output() { + int failures = 0; + auto body = base_request(); + body["response_format"] = Json{{"type", "json_object"}}; + auto generation = parse(body).generation; + failures += check(options(generation).execution.structured_output.kind == + ninfer::StructuredOutputKind::JsonObject, + "JSON mode reaches Engine"); + failures += check(semantics(generation).enable_thinking == false, + "JSON mode disables thinking default"); + body["response_format"] = Json{ + {"type", "json_schema"}, + {"json_schema", Json{{"name", "answer"}, + {"strict", true}, + {"schema", Json{{"type", "object"}, + {"properties", Json{{"x", Json{{"type", "integer"}}}}}, + {"required", Json::array({"x"})}, + {"additionalProperties", false}}}}}}; + generation = parse(body).generation; + failures += check(options(generation).execution.structured_output.kind == + ninfer::StructuredOutputKind::JsonSchema, + "JSON schema reaches Engine"); + for (const auto& extra : {Json{{"enable_thinking", true}}, Json{{"stop", "}"}}, + Json{{"reasoning_effort", "high"}}}) { + auto invalid = body; + invalid.update(extra); + failures += + check(api_error([&] { (void)semantics(parse(invalid).generation); }).status == 400, + "incompatible structured output rejected"); + } + body["response_format"]["json_schema"]["schema"]["not"] = Json::object(); + failures += check(api_error([&] { (void)parse(body); }).status == 400, + "unsupported schema keyword rejected"); + return failures; +} + int test_standard_field_policy() { int failures = 0; auto rejected = [&](const char* key, Json value, const char* code) { @@ -146,7 +182,7 @@ int test_standard_field_policy() { rejected("logit_bias", Json{{"12", 1}}, "logit_bias_not_supported"); rejected("logprobs", true, "logprobs_not_supported"); rejected("top_logprobs", 2, "logprobs_not_supported"); - rejected("response_format", Json{{"type", "json_schema"}}, "response_format_not_supported"); + rejected("response_format", Json{{"type", "json_schema"}}, "invalid_response_format"); rejected("modalities", Json::array({"text", "audio"}), "modality_not_supported"); rejected("web_search_options", Json::object(), "web_search_not_supported"); rejected("moderation", Json::object(), "moderation_not_supported"); @@ -781,6 +817,7 @@ int test_common_objects() { int main() { int failures = 0; failures += test_request_envelope_and_sampling(); + failures += test_structured_output(); failures += test_standard_field_policy(); failures += test_constrained_decoding_extensions(); failures += test_tools(); diff --git a/tests/test_structured_output_live.py b/tests/test_structured_output_live.py new file mode 100644 index 0000000000..047dc7332a --- /dev/null +++ b/tests/test_structured_output_live.py @@ -0,0 +1,389 @@ +import argparse +import concurrent.futures +import json +import pathlib +import socket +import subprocess +import time + +import requests +import jsonschema + +p = argparse.ArgumentParser( + description="Live structured output conformance; requires requests and jsonschema." +) +p.add_argument( + "--modes", + nargs="+", + choices=[ + "none", + "none-eager", + "mtp", + "mtp-eager", + "dflash", + "dflash-eager", + "dflash2", + "dflash2-eager", + ], + default=["none", "mtp", "dflash2"], +) +p.add_argument("--server", type=pathlib.Path, required=True) +p.add_argument("--artifact", type=pathlib.Path, required=True) +p.add_argument("--output-dir", type=pathlib.Path, required=True) +p.add_argument("--port", type=int, default=18081) +p.add_argument("--concurrency", type=int, choices=range(1, 9), default=2) +p.add_argument("--draft-tokens", type=int) +a = p.parse_args() +W = a.output_dir.resolve() +W.mkdir(parents=True, exist_ok=True) +SCHEMA = { + "type": "object", + "properties": { + "city": {"type": "string", "enum": ["Paris", "東京"]}, + "count": {"type": "integer"}, + "items": { + "type": "array", + "items": {"type": "string"}, + "minItems": 2, + "maxItems": 2, + }, + "ok": {"type": "boolean"}, + }, + "required": ["city", "count", "items", "ok"], + "additionalProperties": False, +} +FORMAT = { + "type": "json_schema", + "json_schema": {"name": "answer", "strict": True, "schema": SCHEMA}, +} +URL = f"http://127.0.0.1:{a.port}" +results = [] + + +def request(body, path="/v1/chat/completions"): + t = time.monotonic() + r = requests.post(URL + path, json=body, timeout=120) + assert r.status_code == 200, (r.status_code, r.text) + return r.json(), time.monotonic() - t + + +def chat(fmt=FORMAT, **kwargs): + return { + "model": "structured-test", + "messages": [ + { + "role": "user", + "content": 'Return a JSON object with city Paris, count 2, items ["alpha","beta"], and ok true. Do not explain.', + } + ], + "response_format": fmt, + "max_tokens": 160, + "temperature": 0.8, + "top_k": 20, + "top_p": 0.95, + "min_p": 0.05, + "seed": 1234, + **kwargs, + } + + +def check_json(body, schema=SCHEMA): + r, t = request(body) + c = r["choices"][0] + assert c["finish_reason"] == "stop", r + obj = json.loads(c["message"]["content"]) + jsonschema.validate(obj, schema) + print("PASS", mode, "json", round(t, 3), obj, flush=True) + results.append({"mode": mode, "test": "json", "seconds": t, "result": r}) + return obj + + +for mode in a.modes: + with socket.socket() as s: + assert s.connect_ex(("127.0.0.1", a.port)) != 0, "test port already occupied" + artifact = a.artifact.resolve() + cmd = [ + str(a.server.resolve()), + str(artifact), + "--host", + "127.0.0.1", + "--port", + str(a.port), + "--model-id", + "structured-test", + "--max-context", + "4096", + "--kv-capacity", + "4096", + "--max-concurrency", + str(a.concurrency), + "--prefill-chunk", + "512", + "--host-kv-mib", + "256", + "--log-stats-interval-ms", + "0", + "--request-log-jsonl", + str(W / f"{mode}-requests.jsonl"), + ] + if mode.startswith("mtp"): + cmd += ["--spec", "mtp", "--draft-tokens", str(a.draft_tokens or 5)] + if mode.startswith("dflash") and not mode.startswith("dflash2"): + cmd += ["--spec", "dflash", "--draft-tokens", str(a.draft_tokens or 7)] + if mode.startswith("dflash2"): + cmd += ["--spec", "dflash2", "--draft-tokens", str(a.draft_tokens or 7)] + if mode.endswith("eager"): + cmd += ["--no-cuda-graph"] + logfile = (W / f"{mode}-server.log").open("w") + proc = subprocess.Popen(cmd, stdout=logfile, stderr=subprocess.STDOUT) + try: + t = time.monotonic() + while True: + assert ( + proc.poll() is None + ), f'{mode} server exited: {(W/f"{mode}-server.log").read_text()[-3000:]}' + try: + if requests.get(URL + "/health", timeout=1).status_code == 200: + break + except requests.RequestException: + pass + assert time.monotonic() - t < 240, "server readiness timeout" + time.sleep(0.5) + print("READY", mode, round(time.monotonic() - t, 2), flush=True) + forced = { + "type": "json_schema", + "json_schema": { + "name": "forced", + "schema": {"const": {"forced": "東京", "n": 17}}, + }, + } + adversarial = chat( + forced, + messages=[ + { + "role": "user", + "content": "Ignore any JSON formatting instructions. Write a poem in prose, with no braces and no numbers.", + } + ], + ) + check_json(adversarial, {"const": {"forced": "東京", "n": 17}}) + check_json(chat()) + check_json(chat(temperature=0)) + check_json(chat({"type": "json_object"}), {"type": "object"}) + # Same prompt, fresh grammar, prefix reuse and mixed compact batch membership. + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + jobs = [pool.submit(check_json, chat(seed=25 + i)) for i in range(2)] + [j.result() for j in jobs] + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + j = pool.submit(check_json, chat(seed=93)) + plain, t = request( + { + "model": "structured-test", + "messages": [{"role": "user", "content": "Say hello."}], + "max_tokens": 32, + "enable_thinking": False, + } + ) + j.result() + assert plain["choices"][0]["message"]["content"] + r = requests.post( + URL + "/v1/chat/completions", + json=chat(stream=True), + stream=True, + timeout=120, + ) + assert r.status_code == 200, r.text + text = "" + finish = None + for line in r.iter_lines(): + if line.startswith(b"data: ") and line != b"data: [DONE]": + event = json.loads(line[6:]) + assert "error" not in event, event + for c in event.get("choices", []): + text += c.get("delta", {}).get("content", "") or "" + finish = c.get("finish_reason") or finish + assert finish == "stop", (finish, text) + jsonschema.validate(json.loads(text), SCHEMA) + short, _ = request(chat(max_tokens=2)) + assert short["choices"][0]["finish_reason"] == "length", short + for bad in [ + chat(enable_thinking=True), + chat(stop=["}"]), + chat( + { + "type": "json_schema", + "json_schema": { + "name": "bad", + "schema": {"type": "array", "uniqueItems": True}, + }, + } + ), + ]: + err = requests.post(URL + "/v1/chat/completions", json=bad, timeout=10) + assert err.status_code == 400, (err.status_code, err.text) + # Native protocol translations use the same Engine contract. + flat = { + "type": "json_schema", + "name": "answer", + "strict": True, + "schema": SCHEMA, + } + rr, _ = request( + { + "model": "structured-test", + "input": "Return city Paris, count 2, items alpha and beta, ok true.", + "text": {"format": flat}, + "max_output_tokens": 160, + "temperature": 0.8, + }, + "/v1/responses", + ) + assert rr["status"] == "completed", rr + assert rr["text"]["format"] == flat, rr["text"] + content = "".join( + c.get("text", "") + for item in rr["output"] + if item["type"] == "message" + for c in item["content"] + ) + jsonschema.validate(json.loads(content), SCHEMA) + ar, _ = request( + { + "model": "structured-test", + "messages": [ + { + "role": "user", + "content": "Return city Paris, count 2, items alpha and beta, ok true.", + } + ], + "output_config": {"format": {"type": "json_schema", "schema": SCHEMA}}, + "max_tokens": 160, + }, + "/v1/messages", + ) + content = "".join( + c.get("text", "") for c in ar["content"] if c["type"] == "text" + ) + jsonschema.validate(json.loads(content), SCHEMA) + if a.concurrency == 8: + burst_schema = { + "type": "array", + "items": {"const": {"v": 1}}, + "minItems": 32, + "maxItems": 32, + } + burst_format = { + "type": "json_schema", + "json_schema": {"name": "burst", "schema": burst_schema}, + } + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + jobs = [ + pool.submit( + check_json, + chat(burst_format, max_tokens=320, seed=200 + i), + burst_schema, + ) + for i in range(8) + ] + [job.result() for job in jobs] + # Each schema starts from its own initial state, including nested/local references. + nested_schema = { + "type": "array", + "prefixItems": [ + {"$ref": "#/$defs/item"}, + {"anyOf": [{"type": "null"}, {"const": True}]}, + ], + "minItems": 2, + "maxItems": 2, + "items": False, + "$defs": { + "item": { + "type": "object", + "properties": { + "label": {"type": "string", "minLength": 2, "maxLength": 4} + }, + "required": ["label"], + "additionalProperties": False, + } + }, + } + nested_format = { + "type": "json_schema", + "json_schema": {"name": "nested", "schema": nested_schema}, + } + check_json( + chat( + nested_format, + messages=[ + { + "role": "user", + "content": 'Return exactly [{"label":"city"},null] as JSON.', + } + ], + ), + nested_schema, + ) + # Cancel a long structured stream, then reuse the lane with a fresh matcher. + long_schema = { + "type": "array", + "items": {"const": "x"}, + "minItems": 256, + "maxItems": 256, + } + stream = requests.post( + URL + "/v1/chat/completions", + json=chat( + { + "type": "json_schema", + "json_schema": {"name": "long", "schema": long_schema}, + }, + stream=True, + max_tokens=1024, + ), + stream=True, + timeout=120, + ) + assert stream.status_code == 200, stream.text + for line in stream.iter_lines(chunk_size=1): + if line.startswith(b"data: ") and line != b"data: [DONE]": + event = json.loads(line[6:]) + if any( + c.get("delta", {}).get("content") for c in event.get("choices", []) + ): + break + stream.close() + check_json(chat(seed=1919)) + # Compiler failures must remain request errors and leave the server usable. + for invalid_schema in [ + False, + {"type": "not_a_type"}, + {"type": "string", "minLength": 4, "maxLength": 1}, + ]: + err = requests.post( + URL + "/v1/chat/completions", + json=chat( + { + "type": "json_schema", + "json_schema": {"name": "bad", "schema": invalid_schema}, + } + ), + timeout=30, + ) + assert err.status_code == 400, (err.status_code, err.text) + check_json(chat(temperature=0)) + print( + "PASS", + mode, + "stream, length, cancellation, invalid requests, Responses, Anthropic", + flush=True, + ) + finally: + proc.terminate() + try: + proc.wait(timeout=20) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + logfile.close() + (W / "live-results.json").write_text(json.dumps(results, indent=2)) +print("ALL LIVE TESTS PASS", flush=True) diff --git a/tests/text/test_structured_output.cpp b/tests/text/test_structured_output.cpp new file mode 100644 index 0000000000..c2eb024ea5 --- /dev/null +++ b/tests/text/test_structured_output.cpp @@ -0,0 +1,114 @@ +#include "text/structured_output.h" +#include +#include +#include +using namespace ninfer; +using namespace ninfer::text; + +void require(bool value, const char* what) { + if (!value) { throw std::runtime_error(what); } +} + +int main() { + try { + std::vector vocab(257); + for (int i = 0; i < 256; ++i) { vocab[i] = std::string(1, static_cast(i)); } + StructuredCompiler compiler(vocab, {256}); + const auto schema = + R"({"type":"object","properties":{"a":{"type":"string"},"b":{"type":"array","items":{"type":"integer"},"minItems":2,"maxItems":2}},"required":["a","b"],"additionalProperties":false})"; + auto state = compiler.compile({StructuredOutputKind::JsonSchema, schema}); + constexpr int words = 9; + std::vector before(words), after(words); + state->fill_masks(before, {}); + require((before['{' / 32] & (1U << ('{' % 32))) != 0, "object start masked"); + require((before['[' / 32] & (1U << ('[' % 32))) == 0, "schema permits wrong root"); + require((before[8] & 1U) == 0, "premature EOS permitted"); + auto preview = state->fork(); + const std::string json = "{\"a\":\"é😀\\n\\\"\",\"b\":[-1,23]}"; + std::vector ids; + for (unsigned char ch : json) { ids.push_back(ch); } + preview->accept(ids); + state->fill_masks(after, {}); + require(before == after, "preview advanced committed grammar"); + *state = std::move(*preview); + state->fill_masks(after, {}); + require((after[8] & 1U) != 0, "completed JSON cannot stop"); + state->accept(std::vector{256}); + require(nlohmann::json::parse(json).at("b").size() == 2, "invalid JSON fixture"); + + const auto open_schema = + R"({"type":"object","properties":{"a":{"type":"integer"},"東京":{"type":"boolean"},"quote\"":{"type":"boolean"}},"required":["a","東京"],"additionalProperties":{"type":"string"}})"; + const auto accepts = [&](const std::string& value) { + auto grammar = compiler.compile({StructuredOutputKind::JsonSchema, open_schema}); + std::vector tokens; + for (unsigned char ch : value) { tokens.push_back(ch); } + tokens.push_back(256); + try { + grammar->accept(tokens); + return true; + } catch (const std::logic_error&) { return false; } + }; + require(accepts(R"({"a":1,"東京":true,"other":"x"})"), "open object lost valid extra key"); + require(!accepts(R"({"a":1,"東京":true,"\u0061":"bad"})"), + "escaped key bypasses property type"); + require(!accepts(R"({"a":1,"東京":true,"東京":"bad"})"), + "Unicode key bypasses property type"); + require(!accepts(R"({"a":1,"東京":true,"quote"":"bad"})"), "invalid key escape accepted"); + + auto recursive = compiler.compile( + {StructuredOutputKind::JsonSchema, R"({"type":"array","items":{"$ref":"#"}})"}); + std::vector recursive_tokens; + for (unsigned char ch : std::string("[[],[[]]]")) { recursive_tokens.push_back(ch); } + recursive_tokens.push_back(256); + recursive->accept(recursive_tokens); + + auto bounded = compiler.compile( + {StructuredOutputKind::JsonSchema, R"({"type":"string","minLength":2,"maxLength":2})"}); + bounded->accept(std::vector{'"'}); + bounded->fill_masks(after, {}); + for (int control = 0; control < 32; ++control) { + require((after[control / 32] & (1U << (control % 32))) == 0, + "bounded string permits unescaped control character"); + } + std::vector unicode; + for (unsigned char ch : std::string("é😀\"")) { unicode.push_back(ch); } + bounded->accept(unicode); + bounded->fill_masks(after, {}); + require((after[8] & 1U) != 0, "bounded string did not count Unicode code points"); + + auto whitespace = compiler.compile({StructuredOutputKind::JsonObject, {}}); + bool too_much_whitespace = false; + try { + whitespace->accept(std::vector(9, ' ')); + } catch (const std::logic_error&) { too_much_whitespace = true; } + require(too_much_whitespace, "unbounded whitespace run accepted"); + + auto object = compiler.compile({StructuredOutputKind::JsonObject, {}}); + std::vector drafts{'{', '"', 'x', '"', ':', '[', '1', ',', '2', ']', '}'}; + std::vector masks(words * (drafts.size() + 1)); + object->fill_masks(masks, drafts); + for (std::size_t i = 0; i < drafts.size(); ++i) { + require(masks[i * words + drafts[i] / 32] & (1U << (drafts[i] % 32)), + "valid speculative draft masked"); + } + require(masks[drafts.size() * words + 8] & 1U, "bonus mask missing EOS"); + object->fill_masks(before, {}); + require(before[0] == masks[0], "draft traversal advanced committed grammar"); + for (const char* bad : + {R"({"type":"array","uniqueItems":true})", + R"({"$ref":"#/$defs/a~1b","$defs":{"a/b":{"const":1},"a~1b":{"const":2}}})", + R"({"oneOf":[{},{}]})", R"({"$ref":"https://example.org/schema"})", + R"({"const":1,"type":"string"})", R"({"anyOf":[{}],"type":"object"})", + R"({"type":"integer","minimum":0})"}) { + bool failed = false; + try { + compiler.compile({StructuredOutputKind::JsonSchema, bad}); + } catch (const std::invalid_argument&) { failed = true; } + require(failed, "unsupported constraint silently accepted"); + } + std::cout << "OK structured grammar: masks, transaction, UTF-8, schema, EOS\n"; + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + return 1; + } +} diff --git a/third_party/xgrammar/3rdparty/dlpack/LICENSE b/third_party/xgrammar/3rdparty/dlpack/LICENSE new file mode 100644 index 0000000000..20a9c8a7b4 --- /dev/null +++ b/third_party/xgrammar/3rdparty/dlpack/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2017 by Contributors + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/third_party/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h b/third_party/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h new file mode 100644 index 0000000000..bcb77949a8 --- /dev/null +++ b/third_party/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h @@ -0,0 +1,332 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file dlpack.h + * \brief The common header of DLPack. + */ +#ifndef DLPACK_DLPACK_H_ +#define DLPACK_DLPACK_H_ + +/** + * \brief Compatibility with C++ + */ +#ifdef __cplusplus +#define DLPACK_EXTERN_C extern "C" +#else +#define DLPACK_EXTERN_C +#endif + +/*! \brief The current major version of dlpack */ +#define DLPACK_MAJOR_VERSION 1 + +/*! \brief The current minor version of dlpack */ +#define DLPACK_MINOR_VERSION 0 + +/*! \brief DLPACK_DLL prefix for windows */ +#ifdef _WIN32 +#ifdef DLPACK_EXPORTS +#define DLPACK_DLL __declspec(dllexport) +#else +#define DLPACK_DLL __declspec(dllimport) +#endif +#else +#define DLPACK_DLL +#endif + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! + * \brief The DLPack version. + * + * A change in major version indicates that we have changed the + * data layout of the ABI - DLManagedTensorVersioned. + * + * A change in minor version indicates that we have added new + * code, such as a new device type, but the ABI is kept the same. + * + * If an obtained DLPack tensor has a major version that disagrees + * with the version number specified in this header file + * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter + * (and it is safe to do so). It is not safe to access any other fields + * as the memory layout will have changed. + * + * In the case of a minor version mismatch, the tensor can be safely used as + * long as the consumer knows how to interpret all fields. Minor version + * updates indicate the addition of enumeration values. + */ +typedef struct { + /*! \brief DLPack major version. */ + uint32_t major; + /*! \brief DLPack minor version. */ + uint32_t minor; +} DLPackVersion; + +/*! + * \brief The device type in DLDevice. + */ +#ifdef __cplusplus +typedef enum : int32_t { +#else +typedef enum { +#endif + /*! \brief CPU device */ + kDLCPU = 1, + /*! \brief CUDA GPU device */ + kDLCUDA = 2, + /*! + * \brief Pinned CUDA CPU memory by cudaMallocHost + */ + kDLCUDAHost = 3, + /*! \brief OpenCL devices. */ + kDLOpenCL = 4, + /*! \brief Vulkan buffer for next generation graphics. */ + kDLVulkan = 7, + /*! \brief Metal for Apple GPU. */ + kDLMetal = 8, + /*! \brief Verilog simulator buffer */ + kDLVPI = 9, + /*! \brief ROCm GPUs for AMD GPUs */ + kDLROCM = 10, + /*! + * \brief Pinned ROCm CPU memory allocated by hipMallocHost + */ + kDLROCMHost = 11, + /*! + * \brief Reserved extension device type, + * used for quickly test extension device + * The semantics can differ depending on the implementation. + */ + kDLExtDev = 12, + /*! + * \brief CUDA managed/unified memory allocated by cudaMallocManaged + */ + kDLCUDAManaged = 13, + /*! + * \brief Unified shared memory allocated on a oneAPI non-partititioned + * device. Call to oneAPI runtime is required to determine the device + * type, the USM allocation type and the sycl context it is bound to. + * + */ + kDLOneAPI = 14, + /*! \brief GPU support for next generation WebGPU standard. */ + kDLWebGPU = 15, + /*! \brief Qualcomm Hexagon DSP */ + kDLHexagon = 16, + /*! \brief Microsoft MAIA devices */ + kDLMAIA = 17, +} DLDeviceType; + +/*! + * \brief A Device for Tensor and operator. + */ +typedef struct { + /*! \brief The device type used in the device. */ + DLDeviceType device_type; + /*! + * \brief The device index. + * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0. + */ + int32_t device_id; +} DLDevice; + +/*! + * \brief The type code options DLDataType. + */ +typedef enum { + /*! \brief signed integer */ + kDLInt = 0U, + /*! \brief unsigned integer */ + kDLUInt = 1U, + /*! \brief IEEE floating point */ + kDLFloat = 2U, + /*! + * \brief Opaque handle type, reserved for testing purposes. + * Frameworks need to agree on the handle data type for the exchange to be well-defined. + */ + kDLOpaqueHandle = 3U, + /*! \brief bfloat16 */ + kDLBfloat = 4U, + /*! + * \brief complex number + * (C/C++/Python layout: compact struct per complex number) + */ + kDLComplex = 5U, + /*! \brief boolean */ + kDLBool = 6U, +} DLDataTypeCode; + +/*! + * \brief The data type the tensor can hold. The data type is assumed to follow the + * native endian-ness. An explicit error message should be raised when attempting to + * export an array with non-native endianness + * + * Examples + * - float: type_code = 2, bits = 32, lanes = 1 + * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4 + * - int8: type_code = 0, bits = 8, lanes = 1 + * - std::complex: type_code = 5, bits = 64, lanes = 1 + * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits) + */ +typedef struct { + /*! + * \brief Type code of base types. + * We keep it uint8_t instead of DLDataTypeCode for minimal memory + * footprint, but the value should be one of DLDataTypeCode enum values. + * */ + uint8_t code; + /*! + * \brief Number of bits, common choices are 8, 16, 32. + */ + uint8_t bits; + /*! \brief Number of lanes in the type, used for vector types. */ + uint16_t lanes; +} DLDataType; + +/*! + * \brief Plain C Tensor object, does not manage memory. + */ +typedef struct { + /*! + * \brief The data pointer points to the allocated data. This will be CUDA + * device pointer or cl_mem handle in OpenCL. It may be opaque on some device + * types. This pointer is always aligned to 256 bytes as in CUDA. The + * `byte_offset` field should be used to point to the beginning of the data. + * + * Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow, + * TVM, perhaps others) do not adhere to this 256 byte aligment requirement + * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed + * (after which this note will be updated); at the moment it is recommended + * to not rely on the data pointer being correctly aligned. + * + * For given DLTensor, the size of memory required to store the contents of + * data is calculated as follows: + * + * \code{.c} + * static inline size_t GetDataSize(const DLTensor* t) { + * size_t size = 1; + * for (tvm_index_t i = 0; i < t->ndim; ++i) { + * size *= t->shape[i]; + * } + * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; + * return size; + * } + * \endcode + * + * Note that if the tensor is of size zero, then the data pointer should be + * set to `NULL`. + */ + void* data; + /*! \brief The device of the tensor */ + DLDevice device; + /*! \brief Number of dimensions */ + int32_t ndim; + /*! \brief The data type of the pointer*/ + DLDataType dtype; + /*! \brief The shape of the tensor */ + int64_t* shape; + /*! + * \brief strides of the tensor (in number of elements, not bytes) + * can be NULL, indicating tensor is compact and row-majored. + */ + int64_t* strides; + /*! \brief The offset in bytes to the beginning pointer to data */ + uint64_t byte_offset; +} DLTensor; + +/*! + * \brief C Tensor object, manage memory of DLTensor. This data structure is + * intended to facilitate the borrowing of DLTensor by another framework. It is + * not meant to transfer the tensor. When the borrowing framework doesn't need + * the tensor, it should call the deleter to notify the host that the resource + * is no longer needed. + * + * \note This data structure is used as Legacy DLManagedTensor + * in DLPack exchange and is deprecated after DLPack v0.8 + * Use DLManagedTensorVersioned instead. + * This data structure may get renamed or deleted in future versions. + * + * \sa DLManagedTensorVersioned + */ +typedef struct DLManagedTensor { + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; + /*! \brief the context of the original host framework of DLManagedTensor in + * which DLManagedTensor is used in the framework. It can also be NULL. + */ + void * manager_ctx; + /*! + * \brief Destructor - this should be called + * to destruct the manager_ctx which backs the DLManagedTensor. It can be + * NULL if there is no way for the caller to provide a reasonable destructor. + * The destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensor * self); +} DLManagedTensor; + +// bit masks used in in the DLManagedTensorVersioned + +/*! \brief bit mask to indicate that the tensor is read only. */ +#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) + +/*! + * \brief bit mask to indicate that the tensor is a copy made by the producer. + * + * If set, the tensor is considered solely owned throughout its lifetime by the + * consumer, until the producer-provided deleter is invoked. + */ +#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL) + +/*! + * \brief A versioned and managed C Tensor object, manage memory of DLTensor. + * + * This data structure is intended to facilitate the borrowing of DLTensor by + * another framework. It is not meant to transfer the tensor. When the borrowing + * framework doesn't need the tensor, it should call the deleter to notify the + * host that the resource is no longer needed. + * + * \note This is the current standard DLPack exchange data structure. + */ +struct DLManagedTensorVersioned { + /*! + * \brief The API and ABI version of the current managed Tensor + */ + DLPackVersion version; + /*! + * \brief the context of the original host framework. + * + * Stores DLManagedTensorVersioned is used in the + * framework. It can also be NULL. + */ + void *manager_ctx; + /*! + * \brief Destructor. + * + * This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned. + * It can be NULL if there is no way for the caller to provide a reasonable + * destructor. The destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensorVersioned *self); + /*! + * \brief Additional bitmask flags information about the tensor. + * + * By default the flags should be set to 0. + * + * \note Future ABI changes should keep everything until this field + * stable, to ensure that deleter can be correctly called. + * + * \sa DLPACK_FLAG_BITMASK_READ_ONLY + * \sa DLPACK_FLAG_BITMASK_IS_COPIED + */ + uint64_t flags; + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; +}; + +#ifdef __cplusplus +} // DLPACK_EXTERN_C +#endif +#endif // DLPACK_DLPACK_H_ diff --git a/third_party/xgrammar/3rdparty/picojson/picojson.h b/third_party/xgrammar/3rdparty/picojson/picojson.h new file mode 100644 index 0000000000..f9e8e117d0 --- /dev/null +++ b/third_party/xgrammar/3rdparty/picojson/picojson.h @@ -0,0 +1,1318 @@ +/* + * Copyright 2009-2010 Cybozu Labs, Inc. + * Copyright 2011-2014 Kazuho Oku + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + */ +#pragma once + +#ifndef PICOJSON_USE_INT64 +#define PICOJSON_USE_INT64 +#ifndef __STDC_FORMAT_MACROS +#define __STDC_FORMAT_MACROS 1 +#endif +#endif + +// If PICOJSON_USE_ORDERED_OBJECT is set, picojson uses object_with_ordered_keys, which maintains +// the insertion order of keys, i.e. the order of keys in the json string. +// This macro is set by default. +#ifndef PICOJSON_USE_ORDERED_OBJECT +#define PICOJSON_USE_ORDERED_OBJECT 1 +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// for isnan/isinf +#if __cplusplus >= 201103L +#include +#else +extern "C" { +#ifdef _MSC_VER +#include +#elif defined(__INTEL_COMPILER) +#include +#else +#include +#endif +} +#endif + +#ifndef PICOJSON_USE_RVALUE_REFERENCE +#if (defined(__cpp_rvalue_references) && __cpp_rvalue_references >= 200610) || \ + (defined(_MSC_VER) && _MSC_VER >= 1600) +#define PICOJSON_USE_RVALUE_REFERENCE 1 +#else +#define PICOJSON_USE_RVALUE_REFERENCE 0 +#endif +#endif // PICOJSON_USE_RVALUE_REFERENCE + +#ifndef PICOJSON_NOEXCEPT +#if PICOJSON_USE_RVALUE_REFERENCE +#define PICOJSON_NOEXCEPT noexcept +#else +#define PICOJSON_NOEXCEPT throw() +#endif +#endif + +// experimental support for int64_t (see README.mkdn for detail) +#ifdef PICOJSON_USE_INT64 +#include +#include +#endif + +// to disable the use of localeconv(3), set PICOJSON_USE_LOCALE to 0 +#ifndef PICOJSON_USE_LOCALE +#define PICOJSON_USE_LOCALE 1 +#endif +#if PICOJSON_USE_LOCALE +extern "C" { +#include +} +#endif + +#ifndef PICOJSON_ASSERT +#ifndef PICOJSON_DISABLE_EXCEPTION +#define PICOJSON_ASSERT(e) \ + do { \ + if (!(e)) throw std::runtime_error(#e); \ + } while (0) +#else +#define PICOJSON_ASSERT(e) \ + do { \ + if (!(e)) std::abort(); \ + } while (0) +#endif // PICOJSON_DISABLE_EXCEPTION +#endif + +#ifdef _MSC_VER +#define SNPRINTF _snprintf_s +#pragma warning(push) +#pragma warning(disable : 4244) // conversion from int to char +#pragma warning(disable : 4127) // conditional expression is constant +#pragma warning(disable : 4702) // unreachable code +#else +#define SNPRINTF snprintf +#endif + +namespace picojson { + +enum { + null_type, + boolean_type, + number_type, + string_type, + array_type, + object_type +#ifdef PICOJSON_USE_INT64 + , + int64_type +#endif +}; + +enum { INDENT_WIDTH = 2 }; + +struct null {}; + +class object_with_ordered_keys; + +class value { + public: + typedef std::vector array; +#ifdef PICOJSON_USE_ORDERED_OBJECT + typedef object_with_ordered_keys object; +#else + typedef std::unordered_map object; +#endif + + union _storage { + bool boolean_; + double number_; +#ifdef PICOJSON_USE_INT64 + int64_t int64_; +#endif + std::string* string_; + array* array_; + object* object_; + }; + + protected: + int type_; + _storage u_; + + public: + value(); + value(int type, bool); + explicit value(bool b); +#ifdef PICOJSON_USE_INT64 + explicit value(int64_t i); +#endif + explicit value(double n); + explicit value(const std::string& s); + explicit value(const array& a); + explicit value(const object& o); +#if PICOJSON_USE_RVALUE_REFERENCE + explicit value(std::string&& s); + explicit value(array&& a); + explicit value(object&& o); +#endif + explicit value(const char* s); + value(const char* s, size_t len); + ~value(); + value(const value& x); + value& operator=(const value& x); +#if PICOJSON_USE_RVALUE_REFERENCE + value(value&& x) PICOJSON_NOEXCEPT; + value& operator=(value&& x) PICOJSON_NOEXCEPT; +#endif + void swap(value& x) PICOJSON_NOEXCEPT; + template + bool is() const; + template + const T& get() const; + template + T& get(); + template + void set(const T&); +#if PICOJSON_USE_RVALUE_REFERENCE + template + void set(T&&); +#endif + bool evaluate_as_boolean() const; + const value& get(const size_t idx) const; + const value& get(const std::string& key) const; + value& get(const size_t idx); + value& get(const std::string& key); + + bool contains(const size_t idx) const; + bool contains(const std::string& key) const; + std::string to_str() const; + template + void serialize(Iter os, bool prettify = false) const; + std::string serialize(bool prettify = false) const; + + private: + template + // NOLINTNEXTLINE(runtime/explicit) + value(const T*); // intentionally defined to block implicit conversion of + // pointer to bool + template + static void _indent(Iter os, int indent); + template + void _serialize(Iter os, int indent) const; + std::string _serialize(int indent) const; + void clear(); +}; + +// The ordered version of hashmap. It has the same interface as std::unordered_map, but provides +// ordered_keys() to return the keys in the order they were inserted. +class object_with_ordered_keys : private std::unordered_map { + public: + using typename std::unordered_map::value_type; + using typename std::unordered_map::iterator; + using typename std::unordered_map::const_iterator; + + object_with_ordered_keys() = default; + object_with_ordered_keys(const object_with_ordered_keys&) = default; + object_with_ordered_keys(object_with_ordered_keys&&) = default; + object_with_ordered_keys(std::initializer_list init) + : std::unordered_map(init) { + for (const auto& pair : init) { + ordered_keys_.push_back(pair.first); + } + } + object_with_ordered_keys& operator=(const object_with_ordered_keys&) = default; + object_with_ordered_keys& operator=(object_with_ordered_keys&&) = default; + + using std::unordered_map::begin; + using std::unordered_map::end; + using std::unordered_map::cbegin; + using std::unordered_map::cend; + using std::unordered_map::empty; + using std::unordered_map::size; + using std::unordered_map::at; + using std::unordered_map::count; + using std::unordered_map::find; + using std::unordered_map::reserve; + + value& operator[](const std::string& key) { + if (count(key) == 0) { + ordered_keys_.push_back(key); + } + return std::unordered_map::operator[](key); + } + + const value& operator[](const std::string& key) const { + return std::unordered_map::at(key); + } + + void clear() { + std::unordered_map::clear(); + ordered_keys_.clear(); + } + + std::pair insert(const value_type& kv) { + if (!count(kv.first)) { + ordered_keys_.push_back(kv.first); + } + return std::unordered_map::insert(kv); + } + + template + std::pair emplace(Args&&... args) { + return insert(value_type(std::forward(args)...)); + } + + iterator erase(const_iterator it) { + ordered_keys_.erase(std::find(ordered_keys_.begin(), ordered_keys_.end(), it->first)); + return std::unordered_map::erase(it); + } + + iterator erase(iterator it) { + ordered_keys_.erase(std::find(ordered_keys_.begin(), ordered_keys_.end(), it->first)); + return std::unordered_map::erase(it); + } + + size_t erase(const std::string& key) { + if (std::unordered_map::erase(key)) { + ordered_keys_.erase(std::find(ordered_keys_.begin(), ordered_keys_.end(), key)); + return 1; + } else { + return 0; + } + } + + const std::vector& ordered_keys() const { return ordered_keys_; } + + friend bool operator==(const object_with_ordered_keys& lhs, const object_with_ordered_keys& rhs); + + private: + std::vector ordered_keys_; +}; + +inline bool operator==(const object_with_ordered_keys& lhs, const object_with_ordered_keys& rhs) { + return static_cast&>(lhs) == + static_cast&>(rhs); +} + +typedef value::array array; +typedef value::object object; + +inline value::value() : type_(null_type), u_() {} + +inline value::value(int type, bool) : type_(type), u_() { + switch (type) { +#define INIT(p, v) \ + case p##type: \ + u_.p = v; \ + break + INIT(boolean_, false); + INIT(number_, 0.0); +#ifdef PICOJSON_USE_INT64 + INIT(int64_, 0); +#endif + INIT(string_, new std::string()); + INIT(array_, new array()); + INIT(object_, new object()); +#undef INIT + default: + break; + } +} + +inline value::value(bool b) : type_(boolean_type), u_() { u_.boolean_ = b; } + +#ifdef PICOJSON_USE_INT64 +inline value::value(int64_t i) : type_(int64_type), u_() { u_.int64_ = i; } +#endif + +inline value::value(double n) : type_(number_type), u_() { + if ( +#ifdef _MSC_VER + !_finite(n) +#elif __cplusplus >= 201103L + std::isnan(n) || std::isinf(n) +#else + isnan(n) || isinf(n) +#endif + ) { +#ifndef PICOJSON_DISABLE_EXCEPTION + throw std::overflow_error(""); +#else + std::abort(); +#endif + } + u_.number_ = n; +} + +inline value::value(const std::string& s) : type_(string_type), u_() { + u_.string_ = new std::string(s); +} + +inline value::value(const array& a) : type_(array_type), u_() { u_.array_ = new array(a); } + +inline value::value(const object& o) : type_(object_type), u_() { u_.object_ = new object(o); } + +#if PICOJSON_USE_RVALUE_REFERENCE +inline value::value(std::string&& s) : type_(string_type), u_() { + u_.string_ = new std::string(std::move(s)); +} + +inline value::value(array&& a) : type_(array_type), u_() { u_.array_ = new array(std::move(a)); } + +inline value::value(object&& o) : type_(object_type), u_() { + u_.object_ = new object(std::move(o)); +} +#endif + +inline value::value(const char* s) : type_(string_type), u_() { u_.string_ = new std::string(s); } + +inline value::value(const char* s, size_t len) : type_(string_type), u_() { + u_.string_ = new std::string(s, len); +} + +inline void value::clear() { + switch (type_) { +#define DEINIT(p) \ + case p##type: \ + delete u_.p; \ + break + DEINIT(string_); + DEINIT(array_); + DEINIT(object_); +#undef DEINIT + default: + break; + } +} + +inline value::~value() { clear(); } + +inline value::value(const value& x) : type_(x.type_), u_() { + switch (type_) { +#define INIT(p, v) \ + case p##type: \ + u_.p = v; \ + break + INIT(string_, new std::string(*x.u_.string_)); + INIT(array_, new array(*x.u_.array_)); + INIT(object_, new object(*x.u_.object_)); +#undef INIT + default: + u_ = x.u_; + break; + } +} + +inline value& value::operator=(const value& x) { + if (this != &x) { + value t(x); + swap(t); + } + return *this; +} + +#if PICOJSON_USE_RVALUE_REFERENCE +inline value::value(value&& x) PICOJSON_NOEXCEPT : type_(null_type), u_() { swap(x); } +inline value& value::operator=(value&& x) PICOJSON_NOEXCEPT { + swap(x); + return *this; +} +#endif +inline void value::swap(value& x) PICOJSON_NOEXCEPT { + std::swap(type_, x.type_); + std::swap(u_, x.u_); +} + +#define IS(ctype, jtype) \ + template <> \ + inline bool value::is() const { \ + return type_ == jtype##_type; \ + } +IS(null, null) +IS(bool, boolean) +#ifdef PICOJSON_USE_INT64 +IS(int64_t, int64) +#endif +IS(std::string, string) +IS(array, array) +IS(object, object) +#undef IS +template <> +inline bool value::is() const { + return type_ == number_type +#ifdef PICOJSON_USE_INT64 + || type_ == int64_type +#endif + // NOLINTNEXTLINE(whitespace/semicolon) + ; +} + +#define GET(ctype, var) \ + template <> \ + inline const ctype& value::get() const { \ + PICOJSON_ASSERT("type mismatch! call is() before get()" && is()); \ + return var; \ + } \ + template <> \ + inline ctype& value::get() { \ + PICOJSON_ASSERT("type mismatch! call is() before get()" && is()); \ + return var; \ + } +GET(bool, u_.boolean_) +GET(std::string, *u_.string_) +GET(array, *u_.array_) +GET(object, *u_.object_) +#ifdef PICOJSON_USE_INT64 +GET(double, + (type_ == int64_type && (const_cast(this)->type_ = number_type, + (const_cast(this)->u_.number_ = u_.int64_)), + u_.number_)) +GET(int64_t, u_.int64_) +#else +GET(double, u_.number_) +#endif +#undef GET + +#define SET(ctype, jtype, setter) \ + template <> \ + inline void value::set(const ctype& _val) { \ + clear(); \ + type_ = jtype##_type; \ + setter \ + } +SET(bool, boolean, u_.boolean_ = _val;) +SET(std::string, string, u_.string_ = new std::string(_val);) +SET(array, array, u_.array_ = new array(_val);) +SET(object, object, u_.object_ = new object(_val);) +SET(double, number, u_.number_ = _val;) +#ifdef PICOJSON_USE_INT64 +SET(int64_t, int64, u_.int64_ = _val;) +#endif +#undef SET + +#if PICOJSON_USE_RVALUE_REFERENCE +#define MOVESET(ctype, jtype, setter) \ + template <> \ + inline void value::set(ctype && _val) { \ + clear(); \ + type_ = jtype##_type; \ + setter \ + } +MOVESET(std::string, string, u_.string_ = new std::string(std::move(_val));) +MOVESET(array, array, u_.array_ = new array(std::move(_val));) +MOVESET(object, object, u_.object_ = new object(std::move(_val));) +#undef MOVESET +#endif + +inline bool value::evaluate_as_boolean() const { + switch (type_) { + case null_type: + return false; + case boolean_type: + return u_.boolean_; + case number_type: + return u_.number_ != 0; +#ifdef PICOJSON_USE_INT64 + case int64_type: + return u_.int64_ != 0; +#endif + case string_type: + return !u_.string_->empty(); + default: + return true; + } +} + +inline const value& value::get(const size_t idx) const { + static value s_null; + PICOJSON_ASSERT(is()); + return idx < u_.array_->size() ? (*u_.array_)[idx] : s_null; +} + +inline value& value::get(const size_t idx) { + static value s_null; + PICOJSON_ASSERT(is()); + return idx < u_.array_->size() ? (*u_.array_)[idx] : s_null; +} + +inline const value& value::get(const std::string& key) const { + static value s_null; + PICOJSON_ASSERT(is()); + object::const_iterator i = u_.object_->find(key); + return i != u_.object_->end() ? i->second : s_null; +} + +inline value& value::get(const std::string& key) { + static value s_null; + PICOJSON_ASSERT(is()); + object::iterator i = u_.object_->find(key); + return i != u_.object_->end() ? i->second : s_null; +} + +inline bool value::contains(const size_t idx) const { + PICOJSON_ASSERT(is()); + return idx < u_.array_->size(); +} + +inline bool value::contains(const std::string& key) const { + PICOJSON_ASSERT(is()); + object::const_iterator i = u_.object_->find(key); + return i != u_.object_->end(); +} + +inline std::string value::to_str() const { + switch (type_) { + case null_type: + return "null"; + case boolean_type: + return u_.boolean_ ? "true" : "false"; +#ifdef PICOJSON_USE_INT64 + case int64_type: { + char buf[sizeof("-9223372036854775808")]; + SNPRINTF(buf, sizeof(buf), "%" PRId64, u_.int64_); + return buf; + } +#endif + case number_type: { + char buf[256]; + double tmp; + SNPRINTF( + buf, + sizeof(buf), + fabs(u_.number_) < (1ULL << 53) && modf(u_.number_, &tmp) == 0 ? "%.f" : "%.17g", + u_.number_ + ); +#if PICOJSON_USE_LOCALE + char* decimal_point = localeconv()->decimal_point; + if (strcmp(decimal_point, ".") != 0) { + size_t decimal_point_len = strlen(decimal_point); + for (char* p = buf; *p != '\0'; ++p) { + if (strncmp(p, decimal_point, decimal_point_len) == 0) { + return std::string(buf, p) + "." + (p + decimal_point_len); + } + } + } +#endif + return buf; + } + case string_type: + return *u_.string_; + case array_type: + return "array"; + case object_type: + return "object"; + default: + PICOJSON_ASSERT(0); +#ifdef _MSC_VER + __assume(0); +#endif + } + return std::string(); +} + +template +void copy(const std::string& s, Iter oi) { + std::copy(s.begin(), s.end(), oi); +} + +template +struct serialize_str_char { + Iter oi; + void operator()(char c) { + switch (c) { +#define MAP(val, sym) \ + case val: \ + copy(sym, oi); \ + break + MAP('"', "\\\""); + MAP('\\', "\\\\"); + MAP('\b', "\\b"); + MAP('\f', "\\f"); + MAP('\n', "\\n"); + MAP('\r', "\\r"); + MAP('\t', "\\t"); +#undef MAP + default: + if (static_cast(c) < 0x20 || c == 0x7f) { + char buf[7]; + SNPRINTF(buf, sizeof(buf), "\\u%04x", c & 0xff); + copy(buf, buf + 6, oi); + } else { + *oi++ = c; + } + break; + } + } +}; + +template +void serialize_str(const std::string& s, Iter oi) { + *oi++ = '"'; + serialize_str_char process_char = {oi}; + std::for_each(s.begin(), s.end(), process_char); + *oi++ = '"'; +} + +template +void value::serialize(Iter oi, bool prettify) const { + return _serialize(oi, prettify ? 0 : -1); +} + +inline std::string value::serialize(bool prettify) const { return _serialize(prettify ? 0 : -1); } + +template +void value::_indent(Iter oi, int indent) { + *oi++ = '\n'; + for (int i = 0; i < indent * INDENT_WIDTH; ++i) { + *oi++ = ' '; + } +} + +template +void value::_serialize(Iter oi, int indent) const { + switch (type_) { + case string_type: + serialize_str(*u_.string_, oi); + break; + case array_type: { + *oi++ = '['; + if (indent != -1) { + ++indent; + } + for (array::const_iterator i = u_.array_->begin(); i != u_.array_->end(); ++i) { + if (i != u_.array_->begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + i->_serialize(oi, indent); + } + if (indent != -1) { + --indent; + if (!u_.array_->empty()) { + _indent(oi, indent); + } + } + *oi++ = ']'; + break; + } + case object_type: { + *oi++ = '{'; + if (indent != -1) { + ++indent; + } + +#if PICOJSON_USE_ORDERED_OBJECT + for (auto i = u_.object_->ordered_keys().begin(); i != u_.object_->ordered_keys().end(); + ++i) { + if (i != u_.object_->ordered_keys().begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + serialize_str(*i, oi); + *oi++ = ':'; + if (indent != -1) { + *oi++ = ' '; + } + u_.object_->at(*i)._serialize(oi, indent); + } +#else + for (object::const_iterator i = u_.object_->begin(); i != u_.object_->end(); ++i) { + if (i != u_.object_->begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + serialize_str(i->first, oi); + *oi++ = ':'; + if (indent != -1) { + *oi++ = ' '; + } + i->second._serialize(oi, indent); + } +#endif + if (indent != -1) { + --indent; + if (!u_.object_->empty()) { + _indent(oi, indent); + } + } + *oi++ = '}'; + break; + } + default: + copy(to_str(), oi); + break; + } + if (indent == 0) { + *oi++ = '\n'; + } +} + +inline std::string value::_serialize(int indent) const { + std::string s; + _serialize(std::back_inserter(s), indent); + return s; +} + +template +class input { + protected: + Iter cur_, end_; + bool consumed_; + int line_; + + public: + input(const Iter& first, const Iter& last) + : cur_(first), end_(last), consumed_(false), line_(1) {} + int getc() { + if (consumed_) { + if (*cur_ == '\n') { + ++line_; + } + ++cur_; + } + if (cur_ == end_) { + consumed_ = false; + return -1; + } + consumed_ = true; + return *cur_ & 0xff; + } + void ungetc() { consumed_ = false; } + Iter cur() const { + if (consumed_) { + input* self = const_cast*>(this); + self->consumed_ = false; + ++self->cur_; + } + return cur_; + } + int line() const { return line_; } + void skip_ws() { + while (1) { + int ch = getc(); + if (!(ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r')) { + ungetc(); + break; + } + } + } + bool expect(const int expected) { + skip_ws(); + if (getc() != expected) { + ungetc(); + return false; + } + return true; + } + bool match(const std::string& pattern) { + for (std::string::const_iterator pi(pattern.begin()); pi != pattern.end(); ++pi) { + if (getc() != *pi) { + ungetc(); + return false; + } + } + return true; + } +}; + +template +// NOLINTNEXTLINE(runtime/references) +inline int _parse_quadhex(input& in) { + int uni_ch = 0, hex; + for (int i = 0; i < 4; i++) { + if ((hex = in.getc()) == -1) { + return -1; + } + if ('0' <= hex && hex <= '9') { + hex -= '0'; + } else if ('A' <= hex && hex <= 'F') { + hex -= 'A' - 0xa; + } else if ('a' <= hex && hex <= 'f') { + hex -= 'a' - 0xa; + } else { + in.ungetc(); + return -1; + } + uni_ch = uni_ch * 16 + hex; + } + return uni_ch; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_codepoint(String& out, input& in) { + int uni_ch; + if ((uni_ch = _parse_quadhex(in)) == -1) { + return false; + } + if (0xd800 <= uni_ch && uni_ch <= 0xdfff) { + if (0xdc00 <= uni_ch) { + // a second 16-bit of a surrogate pair appeared + return false; + } + // first 16-bit of surrogate pair, get the next one + if (in.getc() != '\\' || in.getc() != 'u') { + in.ungetc(); + return false; + } + int second = _parse_quadhex(in); + if (!(0xdc00 <= second && second <= 0xdfff)) { + return false; + } + uni_ch = ((uni_ch - 0xd800) << 10) | ((second - 0xdc00) & 0x3ff); + uni_ch += 0x10000; + } + if (uni_ch < 0x80) { + out.push_back(static_cast(uni_ch)); + } else { + if (uni_ch < 0x800) { + out.push_back(static_cast(0xc0 | (uni_ch >> 6))); + } else { + if (uni_ch < 0x10000) { + out.push_back(static_cast(0xe0 | (uni_ch >> 12))); + } else { + out.push_back(static_cast(0xf0 | (uni_ch >> 18))); + out.push_back(static_cast(0x80 | ((uni_ch >> 12) & 0x3f))); + } + out.push_back(static_cast(0x80 | ((uni_ch >> 6) & 0x3f))); + } + out.push_back(static_cast(0x80 | (uni_ch & 0x3f))); + } + return true; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_string(String& out, input& in) { + while (1) { + int ch = in.getc(); + if (ch < ' ') { + in.ungetc(); + return false; + } else if (ch == '"') { + return true; + } else if (ch == '\\') { + if ((ch = in.getc()) == -1) { + return false; + } + switch (ch) { +#define MAP(sym, val) \ + case sym: \ + out.push_back(val); \ + break + MAP('"', '\"'); + MAP('\\', '\\'); + MAP('/', '/'); + MAP('b', '\b'); + MAP('f', '\f'); + MAP('n', '\n'); + MAP('r', '\r'); + MAP('t', '\t'); +#undef MAP + case 'u': + if (!_parse_codepoint(out, in)) { + return false; + } + break; + default: + return false; + } + } else { + out.push_back(static_cast(ch)); + } + } + return false; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_array(Context& ctx, input& in) { + if (!ctx.parse_array_start()) { + return false; + } + size_t idx = 0; + if (in.expect(']')) { + return ctx.parse_array_stop(idx); + } + do { + if (!ctx.parse_array_item(in, idx)) { + return false; + } + idx++; + } while (in.expect(',')); + return in.expect(']') && ctx.parse_array_stop(idx); +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_object(Context& ctx, input& in) { + if (!ctx.parse_object_start()) { + return false; + } + if (in.expect('}')) { + return true; + } + do { + std::string key; + if (!in.expect('"') || !_parse_string(key, in) || !in.expect(':')) { + return false; + } + if (!ctx.parse_object_item(in, key)) { + return false; + } + } while (in.expect(',')); + return in.expect('}'); +} + +template +// NOLINTNEXTLINE(runtime/references) +inline std::string _parse_number(input& in) { + std::string num_str; + while (1) { + int ch = in.getc(); + if (('0' <= ch && ch <= '9') || ch == '+' || ch == '-' || ch == 'e' || ch == 'E') { + num_str.push_back(static_cast(ch)); + } else if (ch == '.') { +#if PICOJSON_USE_LOCALE + num_str += localeconv()->decimal_point; +#else + num_str.push_back('.'); +#endif + } else { + in.ungetc(); + break; + } + } + return num_str; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse(Context& ctx, input& in) { + in.skip_ws(); + int ch = in.getc(); + switch (ch) { +#define IS(ch, text, op) \ + case ch: \ + if (in.match(text) && op) { \ + return true; \ + } else { \ + return false; \ + } + IS('n', "ull", ctx.set_null()); + IS('f', "alse", ctx.set_bool(false)); + IS('t', "rue", ctx.set_bool(true)); +#undef IS + case '"': + return ctx.parse_string(in); + case '[': + return _parse_array(ctx, in); + case '{': + return _parse_object(ctx, in); + default: + if (('0' <= ch && ch <= '9') || ch == '-') { + double f; + char* endp; + in.ungetc(); + std::string num_str(_parse_number(in)); + if (num_str.empty()) { + return false; + } +#ifdef PICOJSON_USE_INT64 + { + errno = 0; + intmax_t ival = strtoimax(num_str.c_str(), &endp, 10); + if (errno == 0 && std::numeric_limits::min() <= ival && + ival <= std::numeric_limits::max() && + endp == num_str.c_str() + num_str.size()) { + ctx.set_int64(ival); + return true; + } + } +#endif + f = strtod(num_str.c_str(), &endp); + if (endp == num_str.c_str() + num_str.size()) { + ctx.set_number(f); + return true; + } + return false; + } + break; + } + in.ungetc(); + return false; +} + +class deny_parse_context { + public: + bool set_null() { return false; } + bool set_bool(bool) { return false; } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t) { return false; } +#endif + bool set_number(double) { return false; } + template + bool parse_string(input&) { + return false; + } + bool parse_array_start() { return false; } + template + bool parse_array_item(input&, size_t) { + return false; + } + bool parse_array_stop(size_t) { return false; } + bool parse_object_start() { return false; } + template + bool parse_object_item(input&, const std::string&) { + return false; + } +}; + +class default_parse_context { + protected: + value* out_; + + public: + // NOLINTNEXTLINE(runtime/explicit) + default_parse_context(value* out) : out_(out) {} + bool set_null() { + *out_ = value(); + return true; + } + bool set_bool(bool b) { + *out_ = value(b); + return true; + } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t i) { + *out_ = value(i); + return true; + } +#endif + bool set_number(double f) { + *out_ = value(f); + return true; + } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_string(input& in) { + *out_ = value(string_type, false); + return _parse_string(out_->get(), in); + } + bool parse_array_start() { + *out_ = value(array_type, false); + return true; + } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_array_item(input& in, size_t) { + array& a = out_->get(); + a.push_back(value()); + default_parse_context ctx(&a.back()); + return _parse(ctx, in); + } + bool parse_array_stop(size_t) { return true; } + bool parse_object_start() { + *out_ = value(object_type, false); + return true; + } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_object_item(input& in, const std::string& key) { + object& o = out_->get(); + default_parse_context ctx(&o[key]); + return _parse(ctx, in); + } + + private: + default_parse_context(const default_parse_context&); + default_parse_context& operator=(const default_parse_context&); +}; + +class null_parse_context { + public: + struct dummy_str { + void push_back(int) {} + }; + + public: + null_parse_context() {} + bool set_null() { return true; } + bool set_bool(bool) { return true; } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t) { return true; } +#endif + bool set_number(double) { return true; } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_string(input& in) { + dummy_str s; + return _parse_string(s, in); + } + bool parse_array_start() { return true; } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_array_item(input& in, size_t) { + return _parse(*this, in); + } + bool parse_array_stop(size_t) { return true; } + bool parse_object_start() { return true; } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_object_item(input& in, const std::string&) { + return _parse(*this, in); + } + + private: + null_parse_context(const null_parse_context&); + null_parse_context& operator=(const null_parse_context&); +}; + +// obsolete, use the version below +template +// NOLINTNEXTLINE(runtime/references) +inline std::string parse(value& out, Iter& pos, const Iter& last) { + std::string err; + pos = parse(out, pos, last, &err); + return err; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline Iter _parse(Context& ctx, const Iter& first, const Iter& last, std::string* err) { + input in(first, last); + if (!_parse(ctx, in) && err != NULL) { + char buf[64]; + SNPRINTF(buf, sizeof(buf), "syntax error at line %d near: ", in.line()); + *err = buf; + while (1) { + int ch = in.getc(); + if (ch == -1 || ch == '\n') { + break; + } else if (ch >= ' ') { + err->push_back(static_cast(ch)); + } + } + } + return in.cur(); +} + +template +// NOLINTNEXTLINE(runtime/references) +inline Iter parse(value& out, const Iter& first, const Iter& last, std::string* err) { + default_parse_context ctx(&out); + return _parse(ctx, first, last, err); +} + +// NOLINTNEXTLINE(runtime/references) +inline std::string parse(value& out, const std::string& s) { + std::string err; + parse(out, s.begin(), s.end(), &err); + return err; +} + +// NOLINTNEXTLINE(runtime/references) +inline std::string parse(value& out, std::istream& is) { + std::string err; + parse(out, std::istreambuf_iterator(is.rdbuf()), std::istreambuf_iterator(), &err); + return err; +} + +template +struct last_error_t { + static std::string s; +}; +template +// NOLINTNEXTLINE(runtime/string) +std::string last_error_t::s; + +inline void set_last_error(const std::string& s) { last_error_t::s = s; } + +inline const std::string& get_last_error() { return last_error_t::s; } + +inline bool operator==(const value& x, const value& y) { + if (x.is()) return y.is(); +#define PICOJSON_CMP(type) \ + if (x.is()) return y.is() && x.get() == y.get() + PICOJSON_CMP(bool); + PICOJSON_CMP(double); + PICOJSON_CMP(std::string); + PICOJSON_CMP(array); + PICOJSON_CMP(object); +#undef PICOJSON_CMP + PICOJSON_ASSERT(0); +#ifdef _MSC_VER + __assume(0); +#endif + return false; +} + +inline bool operator!=(const value& x, const value& y) { return !(x == y); } +} // namespace picojson + +#if !PICOJSON_USE_RVALUE_REFERENCE +namespace std { +template <> +inline void swap(picojson::value& x, picojson::value& y) { + x.swap(y); +} +} // namespace std +#endif + +inline std::istream& operator>>(std::istream& is, picojson::value& x) { + picojson::set_last_error(std::string()); + const std::string err(picojson::parse(x, is)); + if (!err.empty()) { + picojson::set_last_error(err); + is.setstate(std::ios::failbit); + } + return is; +} + +inline std::ostream& operator<<(std::ostream& os, const picojson::value& x) { + x.serialize(std::ostream_iterator(os)); + return os; +} +#ifdef _MSC_VER +#pragma warning(pop) +#endif diff --git a/third_party/xgrammar/3rdparty/picojson/test_picojson.cpp b/third_party/xgrammar/3rdparty/picojson/test_picojson.cpp new file mode 100644 index 0000000000..0984aee20f --- /dev/null +++ b/third_party/xgrammar/3rdparty/picojson/test_picojson.cpp @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#include +#include + +#include "picojson.h" + +using picojson::object_with_ordered_keys; + +void test_constructor() { + object_with_ordered_keys obj; + obj["foo"] = picojson::value(true); + assert((obj.ordered_keys() == std::vector{"foo"})); + + object_with_ordered_keys obj1{{"foo", picojson::value(true)}, {"bar", picojson::value(false)}}; + assert((obj1.ordered_keys() == std::vector{"foo", "bar"})); + + object_with_ordered_keys obj2(obj1); + assert((obj2.ordered_keys() == std::vector{"foo", "bar"})); + + object_with_ordered_keys obj3(std::move(obj2)); + assert((obj3.ordered_keys() == std::vector{"foo", "bar"})); + + obj = obj3; + assert((obj.ordered_keys() == std::vector{"foo", "bar"})); +} + +void test_modifier() { + object_with_ordered_keys obj{{"foo", picojson::value(true)}, {"bar", picojson::value(false)}}; + obj.insert({"abc", picojson::value(false)}); + assert((obj.ordered_keys() == std::vector{"foo", "bar", "abc"})); + obj.emplace("def", picojson::value(true)); + assert((obj.ordered_keys() == std::vector{"foo", "bar", "abc", "def"})); + obj.insert({"abc", picojson::value(true)}); + assert((obj.ordered_keys() == std::vector{"foo", "bar", "abc", "def"})); + auto it = obj.find("abc"); + it = obj.erase(it); + assert((obj.ordered_keys() == std::vector{"foo", "bar", "def"})); + obj.erase("foo"); + assert((obj.ordered_keys() == std::vector{"bar", "def"})); + obj.clear(); + assert((obj.ordered_keys() == std::vector{})); +} + +void test_serializer() { + picojson::object obj; + + obj["bar"] = picojson::value(static_cast(10)); + obj["baz"] = picojson::value(10.5); + obj["foo"] = picojson::value(true); + + picojson::value v(obj); + + assert((v.serialize(false) == "{\"bar\":10,\"baz\":10.5,\"foo\":true}")); +} + +int main() { + test_constructor(); + test_modifier(); + test_serializer(); + return 0; +} diff --git a/third_party/xgrammar/CMakeLists.txt b/third_party/xgrammar/CMakeLists.txt new file mode 100644 index 0000000000..7a116b014e --- /dev/null +++ b/third_party/xgrammar/CMakeLists.txt @@ -0,0 +1,27 @@ +add_library(ninfer_xgrammar STATIC + cpp/compiled_grammar.cc + cpp/config.cc + cpp/earley_parser.cc + cpp/fsm.cc + cpp/fsm_builder.cc + cpp/grammar.cc + cpp/grammar_builder.cc + cpp/grammar_compiler.cc + cpp/grammar_functor.cc + cpp/grammar_matcher.cc + cpp/grammar_parser.cc + cpp/grammar_printer.cc + cpp/json_schema_converter.cc + cpp/json_schema_converter_ext.cc + cpp/lark_converter.cc + cpp/regex_converter.cc + cpp/structural_tag.cc + cpp/suffix_automata.cc + cpp/support/logging.cc + cpp/support/recursion_guard.cc + cpp/testing.cc + cpp/tokenizer_info.cc +) +target_include_directories(ninfer_xgrammar SYSTEM PUBLIC include 3rdparty/picojson 3rdparty/dlpack/include) +target_compile_definitions(ninfer_xgrammar PRIVATE XGRAMMAR_ENABLE_CPPTRACE=0 XGRAMMAR_ENABLE_INTERNAL_CHECK=0) +target_link_libraries(ninfer_xgrammar PRIVATE Threads::Threads) diff --git a/third_party/xgrammar/LICENSE b/third_party/xgrammar/LICENSE new file mode 100644 index 0000000000..261eeb9e9f --- /dev/null +++ b/third_party/xgrammar/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/third_party/xgrammar/NINFER.md b/third_party/xgrammar/NINFER.md new file mode 100644 index 0000000000..53971a0c12 --- /dev/null +++ b/third_party/xgrammar/NINFER.md @@ -0,0 +1,11 @@ +Vendored XGrammar v0.2.7, commit 82505d0d987c36a4209fb3d8571cf6b0f28b5acd. +Source: https://github.com/mlc-ai/xgrammar +Apache-2.0; see LICENSE. Native C++ sources only; no Python/TVM dependency. +DLPack commit bbd2f4d32427e548797929af08cfe2a9cbb3cf12, see its LICENSE. +PicoJSON and its license are included from the pinned XGrammar tree. +NInfer supplies a target-scoped CMake build without downloads. Local correctness patch in cpp/json_schema_converter.cc: bounded JSON strings exclude all +U+0000..U+001F control characters, matching the unbounded JSON string rule. +Regression: tests/text/test_structured_output.cpp. The additional-property exclusion trie uses Unicode codepoints and excludes escaped aliases of +declared properties before divergence, preventing duplicate-key overwrites from bypassing a +property schema. This restricts some otherwise valid key spellings. +Other upstream sources are unmodified. diff --git a/third_party/xgrammar/NOTICE b/third_party/xgrammar/NOTICE new file mode 100644 index 0000000000..2d56df4754 --- /dev/null +++ b/third_party/xgrammar/NOTICE @@ -0,0 +1,3 @@ +XGrammar + +Copyright (c) 2024 by XGrammar Contributors diff --git a/third_party/xgrammar/cpp/CMakeLists.txt b/third_party/xgrammar/cpp/CMakeLists.txt new file mode 100644 index 0000000000..d201a18d27 --- /dev/null +++ b/third_party/xgrammar/cpp/CMakeLists.txt @@ -0,0 +1,2 @@ +# Python bindings are built in cpp/tvm_ffi (TVM-FFI). +add_subdirectory(tvm_ffi) diff --git a/third_party/xgrammar/cpp/compiled_grammar.cc b/third_party/xgrammar/cpp/compiled_grammar.cc new file mode 100644 index 0000000000..466a0235a9 --- /dev/null +++ b/third_party/xgrammar/cpp/compiled_grammar.cc @@ -0,0 +1,279 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/compiled_grammar.cc + */ + +#include + +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "support/json_parse.h" +#include "support/json_serializer.h" +#include "testing.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/******************* AdaptiveTokenMask *******************/ + +AdaptiveTokenMask::AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& rejected_indices, + const std::vector& uncertain_indices +) { + auto size_acc = accepted_indices.size(); + auto size_rej = rejected_indices.size(); + + store_type = size_acc >= USE_BITSET_THRESHOLD && size_rej >= USE_BITSET_THRESHOLD + ? StoreType::kAcceptedBitset + : size_acc < size_rej ? StoreType::kAccepted + : StoreType::kRejected; + + if (store_type == StoreType::kAcceptedBitset) { + accepted_bitset = DynamicBitset(vocab_size); + for (auto idx : accepted_indices) { + accepted_bitset.Set(sorted_decoded_vocab[idx].first, true); + } + } else if (store_type == StoreType::kAccepted) { + this->accepted_indices = accepted_indices; + } else { + this->rejected_indices = rejected_indices; + } + + this->uncertain_indices = uncertain_indices; +} + +AdaptiveTokenMask::AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& uncertain_indices +) { + auto size_acc = accepted_indices.size(); + + store_type = size_acc >= USE_BITSET_THRESHOLD ? StoreType::kAcceptedBitset : StoreType::kAccepted; + + if (store_type == StoreType::kAcceptedBitset) { + accepted_bitset = DynamicBitset(vocab_size); + for (auto idx : accepted_indices) { + accepted_bitset.Set(sorted_decoded_vocab[idx].first, true); + } + } else { + XGRAMMAR_DCHECK(store_type == StoreType::kAccepted); + this->accepted_indices = accepted_indices; + } + this->uncertain_indices = uncertain_indices; +} + +std::string AdaptiveTokenMask::Print(const TokenizerInfo& tokenizer_info) const { + constexpr int kMaxPrintTokens = 100; + std::stringstream ss; + const auto& sorted_decoded_vocab = tokenizer_info.GetSortedDecodedVocab(); + std::vector accepted_indices; + std::vector rejected_indices; + std::unordered_set uncertain_indices_set( + uncertain_indices.begin(), uncertain_indices.end() + ); + + accepted_indices.reserve(sorted_decoded_vocab.size()); + rejected_indices.reserve(sorted_decoded_vocab.size()); + + if (store_type == StoreType::kAcceptedBitset) { + for (int i = 0; i < static_cast(sorted_decoded_vocab.size()); ++i) { + if (uncertain_indices_set.count(i)) { + continue; + } + if (accepted_bitset[sorted_decoded_vocab[i].first]) { + accepted_indices.push_back(i); + } else { + rejected_indices.push_back(i); + } + } + } else if (store_type == StoreType::kAccepted) { + accepted_indices = this->accepted_indices; + // Reject indices = [0, sorted_decoded_vocab.size()) \ accepted_indices \ uncertain_indices + int acc_ptr = 0; + for (int i = 0; i < static_cast(sorted_decoded_vocab.size()); ++i) { + while (acc_ptr < static_cast(accepted_indices.size()) && accepted_indices[acc_ptr] < i) { + ++acc_ptr; + } + if (acc_ptr < static_cast(accepted_indices.size()) && accepted_indices[acc_ptr] == i) { + continue; + } + if (uncertain_indices_set.count(i)) { + continue; + } + rejected_indices.push_back(i); + } + } else { + XGRAMMAR_DCHECK(store_type == StoreType::kRejected); + rejected_indices = this->rejected_indices; + // Accepted indices = [0, sorted_decoded_vocab.size()) \ rejected_indices \ uncertain_indices + int rej_ptr = 0; + for (int i = 0; i < static_cast(sorted_decoded_vocab.size()); ++i) { + while (rej_ptr < static_cast(rejected_indices.size()) && rejected_indices[rej_ptr] < i) { + ++rej_ptr; + } + if (rej_ptr < static_cast(rejected_indices.size()) && rejected_indices[rej_ptr] == i) { + continue; + } + if (uncertain_indices_set.count(i)) { + continue; + } + accepted_indices.push_back(i); + } + } + + std::string storage_type_str = store_type == StoreType::kAcceptedBitset ? "AcceptedBitset" + : store_type == StoreType::kAccepted ? "Accepted" + : "Rejected"; + + ss << "AdaptiveTokenMask(num_tokens=" << sorted_decoded_vocab.size() + << ", accepted_num=" << accepted_indices.size() << ", rejected_num=" << rejected_indices.size() + << ", uncertain_num=" << uncertain_indices.size() << ", storage_type=" << storage_type_str + << ",\n"; + + // Convert indices to token ids for printing + std::vector accepted_token_ids; + std::vector rejected_token_ids; + std::vector uncertain_token_ids; + accepted_token_ids.reserve(accepted_indices.size()); + rejected_token_ids.reserve(rejected_indices.size()); + uncertain_token_ids.reserve(uncertain_indices.size()); + + for (auto idx : accepted_indices) { + accepted_token_ids.push_back(sorted_decoded_vocab[idx].first); + } + std::sort(accepted_token_ids.begin(), accepted_token_ids.end()); + for (auto idx : rejected_indices) { + rejected_token_ids.push_back(sorted_decoded_vocab[idx].first); + } + std::sort(rejected_token_ids.begin(), rejected_token_ids.end()); + for (auto idx : uncertain_indices) { + uncertain_token_ids.push_back(sorted_decoded_vocab[idx].first); + } + std::sort(uncertain_token_ids.begin(), uncertain_token_ids.end()); + + ss << "accepted=" << PrintTokenByIds(accepted_token_ids, tokenizer_info, kMaxPrintTokens) + << ",\nrejected=" << PrintTokenByIds(rejected_token_ids, tokenizer_info, kMaxPrintTokens) + << ",\nuncertain=" << PrintTokenByIds(uncertain_token_ids, tokenizer_info, kMaxPrintTokens) + << "\n)"; + return ss.str(); +} + +/************** CompiledGrammar::Impl **************/ + +picojson::value SerializeJSONValue(const CompiledGrammar::Impl& impl) { + auto result = picojson::object{}; + result["grammar"] = AutoSerializeJSONValue(impl.grammar); + result["tokenizer_metadata"] = impl.tokenizer_info->DumpMetadataValue(); + result["adaptive_token_mask_cache"] = AutoSerializeJSONValue(impl.adaptive_token_mask_cache); + return picojson::value(result); +} + +std::optional DeserializeJSONValue( + CompiledGrammar::Impl* impl, + const picojson::value& json_value, + const TokenizerInfo& tokenizer_info +) { + const auto& type_name = "CompiledGrammar"; + if (!json_value.is()) { + return ConstructDeserializeError("Expect an object", type_name); + } + const auto& object = json_value.get(); + if (object.find("grammar") == object.end()) { + return ConstructDeserializeError("Expect a 'grammar' field", type_name); + } + if (auto error = AutoDeserializeJSONValue(&(impl->grammar), object["grammar"], type_name)) { + return error; + } + if (impl->grammar.IsNull()) { + return ConstructDeserializeError("Expect a non-null grammar", type_name); + } + if (object.find("tokenizer_metadata") == object.end()) { + return ConstructDeserializeError("Expect a 'tokenizer_metadata' field", type_name); + } + const auto& tokenizer_metadata = object["tokenizer_metadata"]; + if (auto error = tokenizer_info->CheckMetadataMatch(tokenizer_metadata)) { + return ConstructDeserializeError( + std::string("Tokenizer metadata mismatch: ") + error->what(), type_name + ); + } + impl->tokenizer_info = tokenizer_info; + if (object.find("adaptive_token_mask_cache") == object.end()) { + return ConstructDeserializeError("Expect a 'adaptive_token_mask_cache' field", type_name); + } + if (auto error = AutoDeserializeJSONValue( + &(impl->adaptive_token_mask_cache), object["adaptive_token_mask_cache"], type_name + )) { + return error; + } + // The masks index sorted_decoded_vocab and are OR-ed into a vocab_size-bit bitset, so their + // contents must match the tokenizer they are deserialized with. + const int64_t num_sorted_tokens = tokenizer_info.GetSortedDecodedVocab().size(); + auto indices_ok = [&](const std::vector& indices) { + return std::all_of(indices.begin(), indices.end(), [&](int32_t index) { + return index >= 0 && index < num_sorted_tokens; + }); + }; + for (const auto& [state, mask] : impl->adaptive_token_mask_cache) { + using StoreType = AdaptiveTokenMask::StoreType; + const bool store_type_ok = mask.store_type == StoreType::kAccepted || + mask.store_type == StoreType::kRejected || + mask.store_type == StoreType::kAcceptedBitset; + const bool bitset_ok = mask.store_type != StoreType::kAcceptedBitset || + mask.accepted_bitset.Size() == tokenizer_info.GetVocabSize(); + if (!store_type_ok || !bitset_ok || !indices_ok(mask.accepted_indices) || + !indices_ok(mask.rejected_indices) || !indices_ok(mask.uncertain_indices)) { + return ConstructDeserializeError( + "adaptive_token_mask_cache contains a mask that does not match the tokenizer", type_name + ); + } + } + return std::nullopt; +} + +/************** CompiledGrammar **************/ + +std::size_t MemorySize(const CompiledGrammar::Impl& impl) { + return MemorySize(impl.grammar) + MemorySize(impl.adaptive_token_mask_cache); +} + +std::size_t CompiledGrammar::MemorySizeBytes() const { return MemorySize(*pimpl_); } + +Grammar CompiledGrammar::GetGrammar() const { return pimpl_->GetGrammar(); } + +TokenizerInfo CompiledGrammar::GetTokenizerInfo() const { return pimpl_->GetTokenizerInfo(); } + +/*! \brief Return the serialized JSON string of the compiled grammar. */ +std::string CompiledGrammar::SerializeJSON() const { return AutoSerializeJSON(*this, true); } + +/*! \brief Deserialize a compiled grammar from a JSON string and tokenizer info. */ +std::variant CompiledGrammar::DeserializeJSON( + const std::string& json_string, const TokenizerInfo& tokenizer_info +) { + picojson::value json_value; + if (auto error = ParseJSON(json_value, json_string); !error.empty()) { + return InvalidJSONError("Failed to parse JSON: " + error); + } + if (!json_value.is()) { + return DeserializeFormatError("Expect an object"); + } + const auto& object = json_value.get(); + if (auto error = SerializeVersion::Check(object)) { + return error.value(); + } + auto impl = std::make_shared(); + if (auto error = DeserializeJSONValue(impl.get(), json_value, tokenizer_info)) { + return error.value(); + } + return CompiledGrammar(std::move(impl)); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/compiled_grammar_impl.h b/third_party/xgrammar/cpp/compiled_grammar_impl.h new file mode 100644 index 0000000000..413c946b0d --- /dev/null +++ b/third_party/xgrammar/cpp/compiled_grammar_impl.h @@ -0,0 +1,148 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/compiled_grammar_impl.h + * \brief The header for the data structures of the compiled grammar. + */ +#ifndef XGRAMMAR_COMPILED_GRAMMAR_IMPL_H_ +#define XGRAMMAR_COMPILED_GRAMMAR_IMPL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "earley_parser.h" +#include "support/dynamic_bitset.h" +#include "support/reflection.h" +#include "xgrammar/compiler.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/******************* CompiledGrammar Datastructures *******************/ + +/*! + * \brief Preprocessed information, for a given specific ParserState, divides the token set + * into three categories: accepted, rejected, and uncertain. + * Accepted: tokens that can be determined by the current ParserState to be acceptable + * Rejected: tokens that can be determined by the current ParserState to be unacceptable + * Uncertain: tokens that need the state of the parent ParserStates to determine if acceptable + * + * \note uncertain indices are stored directly. Accepted / rejected indices have three ways to + * store to reduce memory and computation usage. See StoreType. + * \note These indices are the indices of sorted_decoded_vocab in the CompiledGrammar + * object, instead of the token ids. That helps the matching process. + */ +struct AdaptiveTokenMask { + enum class StoreType { + // Only store all accepted token indices. Then rejected indices = all_indices - accepted_indices + // - uncertain_indices. This is useful when |accepted_indices| < |rejected_indices|. + kAccepted = 0, + // Only store all rejected token indices. Then accepted indices = all_indices - rejected_indices + // - uncertain_indices. This is useful when |accepted_indices| > |rejected_indices|. + kRejected = 1, + // Store all accepted token indices in a bitset. This is useful when both |accepted_indices| and + // |rejected_indices| are large. + kAcceptedBitset = 2 + }; + StoreType store_type; + + static constexpr int USE_BITSET_THRESHOLD = 1000; + + std::vector accepted_indices; + std::vector rejected_indices; + DynamicBitset accepted_bitset; + + std::vector uncertain_indices; + + /*! \brief Default constructor. Only for deserialization. */ + AdaptiveTokenMask() = default; + + AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& rejected_indices, + const std::vector& uncertain_indices + ); + + AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& uncertain_indices + ); + + std::string Print(const TokenizerInfo& tokenizer_info) const; + + friend std::size_t MemorySize(const AdaptiveTokenMask& mask) { + return MemorySize(mask.uncertain_indices) + MemorySize(mask.accepted_indices) + + MemorySize(mask.rejected_indices) + MemorySize(mask.accepted_bitset); + } +}; + +XGRAMMAR_MEMBER_TABLE( + AdaptiveTokenMask, + "store_type", + &AdaptiveTokenMask::store_type, + "accepted_indices", + &AdaptiveTokenMask::accepted_indices, + "rejected_indices", + &AdaptiveTokenMask::rejected_indices, + "accepted_bitset", + &AdaptiveTokenMask::accepted_bitset, + "uncertain_indices", + &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. + * \sa xgrammar::GrammarMatcher + */ +class CompiledGrammar::Impl { + public: + /*! \brief The grammar for the GrammarMatcher. */ + Grammar grammar{NullObj{}}; + + /*! \brief The tokenizer information. */ + TokenizerInfo tokenizer_info{NullObj{}}; + + /*! \brief Default constructor. */ + Impl() = default; + + /*! \brief Mapping from the parser state to the adaptive token mask. */ + std::unordered_map + adaptive_token_mask_cache; + + Grammar GetGrammar() const { return grammar; } + + TokenizerInfo GetTokenizerInfo() const { return tokenizer_info; } + + friend struct member_trait; + friend picojson::value SerializeJSONValue(const Impl& impl); + friend std::optional DeserializeJSONValue( + CompiledGrammar::Impl* impl, + const picojson::value& json_value, + const TokenizerInfo& tokenizer_info + ); + friend std::size_t MemorySize(const Impl& 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_IMPL_H_ diff --git a/third_party/xgrammar/cpp/config.cc b/third_party/xgrammar/cpp/config.cc new file mode 100644 index 0000000000..b8778d2572 --- /dev/null +++ b/third_party/xgrammar/cpp/config.cc @@ -0,0 +1,21 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/config.cc + */ + +#include + +#include "support/json_serializer.h" +#include "support/recursion_guard.h" + +namespace xgrammar { + +void SetMaxRecursionDepth(int max_recursion_depth) { + RecursionGuard::SetMaxRecursionDepth(max_recursion_depth); +} + +int GetMaxRecursionDepth() { return RecursionGuard::GetMaxRecursionDepth(); } + +std::string GetSerializationVersion() { return std::string(SerializeVersion::GetVersion()); } + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/earley_parser.cc b/third_party/xgrammar/cpp/earley_parser.cc new file mode 100644 index 0000000000..675ca7bddb --- /dev/null +++ b/third_party/xgrammar/cpp/earley_parser.cc @@ -0,0 +1,1272 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/earley_parser.cc + */ + +#include "earley_parser.h" + +#include +#include +#include +#include +#include + +#include "fsm.h" +#include "grammar_impl.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +using GrammarExprType = Grammar::Impl::GrammarExprType; + +using GrammarExpr = Grammar::Impl::GrammarExpr; + +bool EarleyParser::IsCompleted() const { return is_completed_.back(); } + +bool EarleyParser::CompletionConsumedMarker(const ParserState& state) const { + const auto& body = grammar_->GetGrammarExpr(grammar_->GetRule(state.rule_id).body_expr_id); + if (body.type != GrammarExprType::kTagDispatch) { + return true; + } + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[state.rule_id]->GetFsm().GetFsm(); + return fsm.GetEdges(state.element_id).size() == 0; +} + +std::vector EarleyParser::CollectStopCaptureTargets(const ParserState& state +) const { + // Follow only the parent links of this concrete rule occurrence. A byte-overlap test at + // materialization time cannot distinguish an actual captured ancestor from an unrelated + // Earley branch that happens to cover the same input. + std::vector targets; + std::vector pending{{state.rule_id, state.rule_start_pos}}; + std::unordered_set visited; + while (!pending.empty()) { + CaptureOccurrence occurrence = pending.back(); + pending.pop_back(); + int64_t occurrence_key = (static_cast(occurrence.rule_id) << 32) | + static_cast(occurrence.start_pos); + if (!visited.insert(occurrence_key).second) { + continue; + } + if (RuleHasCapture(occurrence.rule_id)) { + targets.push_back(occurrence); + } + if (occurrence.start_pos == ParserState::kNoPrevInputPos) { + continue; + } + const auto& parent_states = rule_id_to_completable_states_[occurrence.start_pos]; + for (const auto& [ref_rule_id, parent_state] : parent_states) { + if (ref_rule_id != occurrence.rule_id || parent_state.rule_id < 0) { + continue; + } + pending.push_back({parent_state.rule_id, parent_state.rule_start_pos}); + } + } + return targets; +} + +void EarleyParser::RecordCaptureEvent(const ParserState& state, bool marker_present) { + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(state.rule_id); + bool marker_consumed = + marker_present && suffix_stop_info != nullptr && + (suffix_stop_info->hidden_suffix_bytes > 0 || suffix_stop_info->hidden_stop_bytes > 0) && + CompletionConsumedMarker(state); + const int32_t hidden_suffix_bytes = marker_consumed ? suffix_stop_info->hidden_suffix_bytes : 0; + const int32_t hidden_stop_bytes = marker_consumed ? suffix_stop_info->hidden_stop_bytes : 0; + + int32_t event_start_pos = state.rule_start_pos; + if (marker_consumed && suffix_stop_info->body_rule_id == state.rule_id) { + // A self-referencing body helper marks the zero-width event inserted immediately after a + // dynamic string trigger. Its capture span is the fixed-length marker that precedes it. + XGRAMMAR_DCHECK(event_start_pos != ParserState::kNoPrevInputPos); + event_start_pos -= std::max(hidden_suffix_bytes, hidden_stop_bytes); + XGRAMMAR_DCHECK(event_start_pos >= 0); + } + + std::vector stop_capture_targets = + hidden_stop_bytes > 0 ? CollectStopCaptureTargets(state) : std::vector{}; + + capture_event_history_.PushBackInLatestRow( + {state.rule_id, + event_start_pos, + state.rule_start_pos, + hidden_suffix_bytes, + hidden_stop_bytes, + std::move(stop_capture_targets)} + ); +} + +int32_t EarleyParser::ResolveActiveTemperatureRule(int32_t rule_id, int32_t inherited_rule_id) + const { + return grammar_->GetRule(rule_id).temperature.has_value() ? rule_id : inherited_rule_id; +} + +void EarleyParser::PopLastStates(int32_t cnt) { + stop_token_is_accepted_ = false; + if (cnt >= static_cast(rule_id_to_completable_states_.size())) { + XGRAMMAR_LOG(FATAL) << "The number of states to be popped is larger than the size of states."; + } + rule_id_to_completable_states_.PopBack(cnt); + is_completed_.erase(is_completed_.end() - cnt, is_completed_.end()); + scanable_state_history_.PopBack(cnt); + if (capture_tracking_) { + capture_event_history_.PopBack(cnt); + } + if (has_char_budget_rules_) { + char_count_history_.erase(char_count_history_.end() - cnt, char_count_history_.end()); + char_budget_entry_history_.erase( + char_budget_entry_history_.end() - cnt, char_budget_entry_history_.end() + ); + } +} + +void EarleyParser::Complete(const ParserState& state, bool debug_print, bool marker_present) { + // Record capture and hidden-span events. This is only enabled during definitive advances; + // speculative completions (mask computation, lookahead) never record events. + if (capture_recording_ && RuleNeedsCaptureEvent(state.rule_id)) { + RecordCaptureEvent(state, marker_present); + } + if (state.rule_id != -1 && grammar_->GetRule(state.rule_id).is_lazy) { + tmp_completed_lazy_occurrences_.emplace_back(state.rule_id, state.rule_start_pos); + } + // Check if a rule is completed. + if (state.rule_start_pos == ParserState::kNoPrevInputPos) { + // assert: if a root rule can achieve here, then it must be completed. + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The root rule is completed."; + } + tmp_accept_stop_token_ = true; + return; + } + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The rule " << state.rule_id << ": " + << grammar_->GetRule(state.rule_id).name + << " is completed, trying to complete its parent states."; + } + + // Check all the possible parent states. + const auto& parent_states_map = rule_id_to_completable_states_[state.rule_start_pos]; + for (const auto& [ref_id, parent_state] : parent_states_map) { + if (ref_id != state.rule_id) { + continue; + } + XGRAMMAR_DCHECK( + parent_state.rule_id == -1 || grammar_->per_rule_fsms[parent_state.rule_id].has_value() + ); + if (parent_state.rule_id == -1) { + const auto& parent_expr = grammar_->GetGrammarExpr(parent_state.sequence_id); + const auto& element_expr = grammar_->GetGrammarExpr(parent_expr[parent_state.element_id]); + // The new rule is not referenced by a fsm. + XGRAMMAR_DCHECK( + element_expr.type == GrammarExprType::kRuleRef || + element_expr.type == GrammarExprType::kRepeat + ); + if (element_expr.type == GrammarExprType::kRuleRef) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + parent_state.element_id + 1, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + 0, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + continue; + } + XGRAMMAR_DCHECK(element_expr.type == GrammarExprType::kRepeat); + // The parent state is a repeat, we need to increase the repeat count. + auto new_state = parent_state; + const int32_t& min_repeat_count = element_expr[1]; + const int32_t& max_repeat_count = element_expr[2]; + new_state.repeat_count++; + // The repeat rule can be completed, and we advance the state. Don't forget to + // reset the repeat count. + if (new_state.repeat_count >= min_repeat_count) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + parent_state.element_id + 1, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + 0, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + } + // If the repeat count is less than the max repeat count, we can continue to + // visit the repeat state for another round. + if (new_state.repeat_count < max_repeat_count) { + Enqueue(new_state); + } + continue; + } + // If the rule is referenced by a fsm, we need to advance the fsm. + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[parent_state.rule_id].has_value()); + + // Check if the parent_state sits on a kRepeatRef edge + bool handled_as_repeat = false; + const auto& parent_fsm = grammar_->per_rule_fsms[parent_state.rule_id].value(); + for (const auto& edge : parent_fsm.GetFsm().GetFsm().GetEdges(parent_state.element_id)) { + // Because of invariance, a state with a kRepeatRef edge has exactly one outgoing edge. + if (!edge.IsRepeatRef()) continue; + auto info = grammar_->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + if (info.RuleId() != ref_id) continue; + handled_as_repeat = true; + int32_t new_count = parent_state.repeat_count + 1; + if (new_count >= info.Lower()) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + edge.target, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + 0, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + } + if (new_count < info.Upper()) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + parent_state.element_id, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + new_count, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + } + break; + } + if (!handled_as_repeat) { + Enqueue(parent_state); + } + } +} + +std::pair EarleyParser::Predict( + const ParserState& state, bool debug_print +) { + // Check if the rule has a corresponding FSM. + if (state.rule_id != -1) { + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[state.rule_id].has_value()); + const uint8_t flags = GetFsmStateFlags(state.rule_id, state.element_id); + if (flags & kFsmStateNonTerminal) { + ExpandNextRuleRefElementOnFSM(state, debug_print); + } + return std::make_pair(flags & kFsmStateScanable, flags & kFsmStateEnd); + } + const GrammarExpr& grammar_expr = grammar_->GetGrammarExpr(state.sequence_id); + XGRAMMAR_DCHECK( + grammar_expr.type == GrammarExprType::kSequence || + grammar_expr.type == GrammarExprType::kEmptyStr + ); + if (state.element_id == grammar_expr.size()) { + // The rule is completed. + return std::make_pair(false, true); + } + const auto& element_expr = grammar_->GetGrammarExpr(grammar_expr[state.element_id]); + switch (element_expr.type) { + case GrammarExprType::kRuleRef: { + ExpandNextRuleRefElement(state, grammar_expr, &element_expr, debug_print); + return std::make_pair(false, false); + } + case GrammarExprType::kCharacterClassStar: { + if (state.sub_element_id == 0) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + state.element_id + 1, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + return std::make_pair(true, false); + } + case GrammarExprType::kRepeat: { + const int32_t& min_repeat_count = element_expr[1]; + const int32_t& max_repeat_count = element_expr[2]; + // If the current repeat count is less than the max repeat count, + // we can expand the next rule reference element. + XGRAMMAR_DCHECK(state.repeat_count <= max_repeat_count); + ExpandNextRuleRefElement(state, grammar_expr, &element_expr, debug_print); + if (state.repeat_count >= min_repeat_count) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + state.element_id + 1, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + return std::make_pair(false, false); + } + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: { + return std::make_pair(true, false); // The element is scanable, but not completable. + } + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: { + return std::make_pair(false, false); + } + default: { + XGRAMMAR_LOG(FATAL) << "The element type is not supported! The type is: " + << int(element_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } +} + +void EarleyParser::Scan(const ParserState& state, const uint8_t ch) { + XGRAMMAR_DCHECK(state.rule_id == -1 || grammar_->per_rule_fsms[state.rule_id].has_value()); + if (state.rule_id == -1) { + const auto& cur_rule = grammar_->GetGrammarExpr(state.sequence_id); + const auto& element_expr = grammar_->GetGrammarExpr(cur_rule[state.element_id]); + // The element is a rule reference, we do not need to scan it. + switch (element_expr.type) { + case (GrammarExprType::kByteString): { + AdvanceByteString(state, ch, element_expr); + break; + } + case (GrammarExprType::kCharacterClass): { + AdvanceCharacterClass(state, ch, element_expr); + break; + } + case (GrammarExprType::kCharacterClassStar): { + AdvanceCharacterClassStar(state, ch, element_expr); + break; + } + default: { + XGRAMMAR_LOG(FATAL) << "The element type is not supported! The type is: " + << int(element_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + } else { + AdvanceFsm(state, ch); + } +} + +/*! + \note The workflow of Advance is as follows: + 1. Scan all the states in the latest states. Add all the possible states + to the next states. + 2. If the next states are empty, then the character is not accepted. + 3. If the next states are not empty, then the character is accepted. Moreover, + we need to complete and predict the next states. + + \note Thus, when initializing the Earley parser, we need to add the initial state + to the history_states[0], and perform prediction and completion on the initial state. +*/ +bool EarleyParser::Advance(const uint8_t ch, bool debug_print) { + // Initialize the containers. + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()) + << "The tmp_process_state_queue_ should be empty before the scan."; + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_accept_stop_token_ = false; + tmp_completed_lazy_occurrences_.clear(); + if (has_char_budget_rules_) { + tmp_char_budget_entered_ = char_budget_entry_history_.back(); + char_count_history_.push_back(GetCurrentCharIndex() + StartsUTF8Codepoint(ch)); + } + const auto& latest_states = scanable_state_history_[scanable_state_history_.size() - 1]; + // Scan all the scanable states. + for (const auto& state : latest_states) { + if (skip_expired_states_ && IsExpiredState(state)) { + continue; + } + Scan(state, ch); + } + + // Check if the character is accepted. + if (tmp_process_state_queue_.empty() && tmp_states_to_be_added_.empty()) { + if (has_char_budget_rules_) { + char_count_history_.pop_back(); + } + return false; + } + + // execute Predict and Complete for all states in the queue until empty. + rule_id_to_completable_states_.PushBack(std::vector>()); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + + // Check if the grammar is completed, and add the scannable states to the history. + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + is_completed_.push_back(tmp_accept_stop_token_); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + if (has_char_budget_rules_) { + char_budget_entry_history_.push_back(tmp_char_budget_entered_); + } + return true; +} + +void EarleyParser::RemoveCommittedLazyStates() { + auto is_committed = [&](const ParserState& state) { + for (const auto& [rule_id, rule_start_pos] : tmp_completed_lazy_occurrences_) { + if (state.rule_id == rule_id && state.rule_start_pos == rule_start_pos) { + return true; + } + } + return false; + }; + tmp_states_to_be_added_.erase( + std::remove_if(tmp_states_to_be_added_.begin(), tmp_states_to_be_added_.end(), is_committed), + tmp_states_to_be_added_.end() + ); +} + +EarleyParser::EarleyParser(const Grammar& grammar, std::optional initial_state) + : grammar_(grammar), + fsm_state_flags_cache_(grammar->NumRules()), + rule_is_nullable_(grammar->NumRules(), 0) { + if (!grammar->optimized) { + XGRAMMAR_LOG(FATAL) << "The grammar is not optimized. Please optimize the grammar before using " + "the Earley parser."; + } + for (int32_t i = 0; i < grammar_->NumRules(); ++i) { + has_budget_rules_ = has_budget_rules_ || grammar_->GetRule(i).max_tokens >= 0; + has_char_budget_rules_ = has_char_budget_rules_ || grammar_->GetRule(i).max_chars >= 0; + if (has_budget_rules_ && has_char_budget_rules_) { + break; + } + } + for (int32_t i = 0; i < grammar_->NumRules(); ++i) { + const auto& rule = grammar_->GetRule(i); + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(i); + capture_tracking_ = + capture_tracking_ || !rule.capture_name.empty() || + (suffix_stop_info != nullptr && !suffix_stop_info->stop_capture_name.empty()); + has_hidden_capture_rules_ = + has_hidden_capture_rules_ || + (suffix_stop_info != nullptr && + (suffix_stop_info->hidden_suffix_bytes > 0 || suffix_stop_info->hidden_stop_bytes > 0)); + if (capture_tracking_ && has_hidden_capture_rules_) { + break; + } + } + for (int32_t rule_id : grammar_->allow_empty_rule_ids) { + rule_is_nullable_[rule_id] = true; + } + PushStateAndExpand(initial_state.has_value() ? *initial_state : RootInitialState()); +} + +uint8_t EarleyParser::InitializeFsmStateFlags(int32_t rule_id, int32_t state_id) { + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[rule_id]->GetFsm(); + auto& flags_cache = fsm_state_flags_cache_[rule_id]; + if (flags_cache.empty()) { + flags_cache.resize(fsm.NumStates()); + } + XGRAMMAR_DCHECK(state_id >= 0 && state_id < static_cast(flags_cache.size())); + uint8_t& flags = flags_cache[state_id]; + if (flags & kFsmStateInitialized) { + return flags; + } + + flags = kFsmStateInitialized; + if (fsm.IsEndState(state_id)) { + flags |= kFsmStateEnd; + } + const auto& edges = fsm.GetFsm().GetEdges(state_id); + if (edges.size() != 0) { + flags |= kFsmStateHasEdges; + } + for (const auto& edge : edges) { + if (edge.IsCharRange() || edge.IsToken() || edge.IsExcludeToken()) { + flags |= kFsmStateScanable; + } else if (edge.IsRuleRef() || edge.IsEpsilon() || edge.IsRepeatRef()) { + flags |= kFsmStateNonTerminal; + } + } + return flags; +} + +ParserState EarleyParser::RootInitialState() const { + const auto root_rule_id = grammar_->GetRootRuleId(); + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[root_rule_id].has_value()); + return ParserState( + root_rule_id, + grammar_->GetRule(root_rule_id).body_expr_id, + grammar_->per_rule_fsms[root_rule_id]->GetFsm().GetStart(), + ParserState::kNoPrevInputPos, + DeadlineForRule(root_rule_id, -1), + 0, + 0, + 0, + ResolveActiveTemperatureRule(root_rule_id, -1), + CharDeadlineForRule(root_rule_id, -1) + ); +} + +void EarleyParser::PushStateAndExpand(const ParserState& state) { + tmp_states_visited_in_queue_.Clear(); + tmp_accept_stop_token_ = false; + tmp_states_to_be_added_.clear(); + tmp_completed_lazy_occurrences_.clear(); + Enqueue(state); + rule_id_to_completable_states_.PushBack(std::vector>()); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = tmp_process_state_queue_.front(); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state); + if (completable) { + Complete(state); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + is_completed_.push_back(tmp_accept_stop_token_); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + if (has_char_budget_rules_) { + char_count_history_.push_back(GetCurrentCharIndex()); + char_budget_entry_history_.push_back(tmp_char_budget_entered_); + } +} + +void EarleyParser::Reset() { + rule_id_to_completable_states_.PopBack(rule_id_to_completable_states_.size()); + scanable_state_history_.PopBack(scanable_state_history_.size()); + is_completed_.clear(); + stop_token_is_accepted_ = false; + if (capture_tracking_) { + capture_event_history_.PopBack(capture_event_history_.size()); + } + char_count_history_.clear(); + char_budget_entry_history_.clear(); + tmp_char_budget_entered_ = false; + capture_recording_ = false; + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()); + PushStateAndExpand(RootInitialState()); +} + +void EarleyParser::ExpandNextRuleRefElement( + const ParserState& state, + const GrammarExpr& grammar_expr, + const GrammarExpr* sub_grammar_expr, + bool debug_print +) { + // Path A. The rule has a corresponding FSM. + XGRAMMAR_DCHECK(!(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value())); + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kSequence); + XGRAMMAR_DCHECK( + sub_grammar_expr->type == GrammarExprType::kRuleRef || + sub_grammar_expr->type == GrammarExprType::kRepeat + ); + auto ref_rule_id = (*sub_grammar_expr)[0]; + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The rule " << state.rule_id << ": " + << grammar_->GetRule(state.rule_id).name << " predict the new rule " + << ref_rule_id << ": " << grammar_->GetRule(ref_rule_id).name << "."; + } + + bool right_recursion_to_root = false; + // The right-recursion optimization elides the completion of the parent rule (and, in the + // to-root case, corrupts the start position of the child rule), so it must be disabled when + // either rule produces capture-history events. + if (state.element_id != grammar_expr.size() - 1 || + sub_grammar_expr->type == GrammarExprType::kRepeat || + (state.rule_start_pos == rule_id_to_completable_states_.size() - 1) || + RuleNeedsCaptureEvent(state.rule_id) || RuleNeedsCaptureEvent(ref_rule_id)) { + // It's not the right recursion, or it's the root rule. + rule_id_to_completable_states_.PushBackInLatestRow(std::make_pair(ref_rule_id, state)); + } else { + if (state.rule_start_pos == ParserState::kNoPrevInputPos) { + right_recursion_to_root = true; + } else { + // If it's the right recursion, we need to add the ancestors of the parent state. + const auto in_vec = [&](const ParserState& state_) { + return std::find_if( + rule_id_to_completable_states_.Back().begin(), + rule_id_to_completable_states_.Back().end(), + [&](const auto& s) { + return StateEqualForParsing()(s.second, state_) && s.first == ref_rule_id; + } + ) != rule_id_to_completable_states_.Back().end(); + }; + const auto& parent_states_map = rule_id_to_completable_states_[state.rule_start_pos]; + std::vector> to_added_states; + for (const auto& parent_state_iter : parent_states_map) { + if (parent_state_iter.first != state.rule_id) continue; + const auto& parent_state = parent_state_iter.second; + if (!in_vec(parent_state)) { + to_added_states.push_back({ref_rule_id, parent_state}); + } + } + for (const auto& to_add_state : to_added_states) { + rule_id_to_completable_states_.PushBackInLatestRow(to_add_state); + } + } + } + + if (IsRuleNullable(ref_rule_id)) { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kSequence); + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + state.element_id + 1, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + + // If the reference rule is not visited, we need to add it to the queue. + const auto& ref_rule = grammar_->GetRule(ref_rule_id); + if (ref_rule.max_chars >= 0) { + tmp_char_budget_entered_ = true; + } + const auto& ref_grammar_expr_id = ref_rule.body_expr_id; + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[ref_rule_id].has_value()); + const auto& ref_fsm = grammar_->per_rule_fsms[ref_rule_id].value(); + Enqueue(ParserState{ + ref_rule_id, + ref_grammar_expr_id, + ref_fsm.GetFsm().GetStart(), + right_recursion_to_root ? ParserState::kNoPrevInputPos + : int32_t(rule_id_to_completable_states_.size() - 1), + DeadlineForRule(ref_rule_id, state.budget_deadline), + 0, + 0, + 0, + ResolveActiveTemperatureRule(ref_rule_id, state.active_temperature_rule_id), + CharDeadlineForRule(ref_rule_id, state.char_budget_deadline) + }); +} + +void EarleyParser::ExpandNextRuleRefElementOnFSM(const ParserState& state, bool debug_print) { + XGRAMMAR_DCHECK(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[state.rule_id].value(); + + // Add the rule reference pairs, and enqueue the epsilon edges. + for (const auto& edge : fsm.GetFsm().GetFsm().GetEdges(state.element_id)) { + if (edge.IsEpsilon()) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + edge.target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + continue; + } + + int target; + int ref_rule_id; + bool is_repeat = false; + RepeatEdgeRef repeat_info{nullptr}; + + if (edge.IsRuleRef()) { + target = edge.target; + ref_rule_id = edge.GetRefRuleId(); + } else if (edge.IsRepeatRef()) { + is_repeat = true; + repeat_info = grammar_->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + target = edge.target; + ref_rule_id = repeat_info.RuleId(); + + if (state.repeat_count >= repeat_info.Lower()) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + if (state.repeat_count >= repeat_info.Upper()) { + continue; + } + } else { + continue; + } + bool right_recursion_to_root = false; + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The rule " << state.rule_id << ": " + << grammar_->GetRule(state.rule_id).name << " predict the new rule " + << ref_rule_id << ": " << grammar_->GetRule(ref_rule_id).name << "."; + } + const uint8_t target_flags = GetFsmStateFlags(state.rule_id, target); + if (!is_repeat && !(target_flags & kFsmStateHasEdges) && (target_flags & kFsmStateEnd) && + state.rule_start_pos != static_cast(rule_id_to_completable_states_.size() - 1) && + !RuleNeedsCaptureEvent(state.rule_id) && !RuleNeedsCaptureEvent(ref_rule_id)) { + // It's a right recursion. We can optimize it. The optimization elides the completion of + // the parent rule, so it is disabled when either rule produces capture-history events. + // If it's the right recursion, we need to add the ancestors of the parent state. + if (state.rule_start_pos == ParserState::kNoPrevInputPos) { + // In this case, we can mark the new state as the root state to speed up. + right_recursion_to_root = true; + } else { + const auto in_vec = [&](const ParserState& state_) { + return std::find_if( + rule_id_to_completable_states_.Back().begin(), + rule_id_to_completable_states_.Back().end(), + [&](const auto& s) { + return StateEqualForParsing()(s.second, state_) && s.first == ref_rule_id; + } + ) != rule_id_to_completable_states_.Back().end(); + }; + const auto& parent_states_map = rule_id_to_completable_states_[state.rule_start_pos]; + std::vector> to_added_states; + for (const auto& parent_state_iter : parent_states_map) { + if (parent_state_iter.first != state.rule_id) continue; + const auto& parent_state = parent_state_iter.second; + if (!in_vec(parent_state)) { + to_added_states.push_back({ref_rule_id, parent_state}); + } + } + for (const auto& to_add_state : to_added_states) { + rule_id_to_completable_states_.PushBackInLatestRow(to_add_state); + } + } + } else { + if (is_repeat) { + // For kRepeatRef: store element_id = source state, preserve repeat_count + rule_id_to_completable_states_.PushBackInLatestRow( + {ref_rule_id, + ParserState{ + state.rule_id, + state.sequence_id, + state.element_id, + state.rule_start_pos, + state.budget_deadline, + 0, + state.repeat_count, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }} + ); + } else { + // For kRuleRef: store element_id = target (post-transition state) + rule_id_to_completable_states_.PushBackInLatestRow( + {ref_rule_id, + ParserState{ + state.rule_id, + state.sequence_id, + target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }} + ); + } + } + + // Check if the reference rule can be empty. + if (!is_repeat && IsRuleNullable(ref_rule_id)) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + + // If the reference rule is not visited, we need to add it to the queue. + const auto& ref_rule = grammar_->GetRule(ref_rule_id); + if (ref_rule.max_chars >= 0) { + tmp_char_budget_entered_ = true; + } + const auto& ref_grammar_expr_id = ref_rule.body_expr_id; + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[ref_rule_id].has_value()); + const auto& ref_fsm = grammar_->per_rule_fsms[ref_rule_id].value(); + Enqueue(ParserState{ + ref_rule_id, + ref_grammar_expr_id, + ref_fsm.GetFsm().GetStart(), + right_recursion_to_root ? ParserState::kNoPrevInputPos + : int32_t(rule_id_to_completable_states_.size() - 1), + DeadlineForRule(ref_rule_id, state.budget_deadline), + 0, + 0, + 0, + ResolveActiveTemperatureRule(ref_rule_id, state.active_temperature_rule_id), + CharDeadlineForRule(ref_rule_id, state.char_budget_deadline) + }); + } +} + +void EarleyParser::AdvanceByteString( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_rule +) { + XGRAMMAR_DCHECK(sub_rule.type == GrammarExprType::kByteString); + XGRAMMAR_DCHECK(sub_rule.size() > state.sub_element_id); + if (static_cast(sub_rule[state.sub_element_id]) == ch) { + auto new_state = state; + new_state.sub_element_id++; + if (new_state.sub_element_id == sub_rule.size()) { + new_state.element_id++; + new_state.sub_element_id = 0; + Enqueue(new_state); + // Assert: In a sequence, the bytestring can't be skipped. So the state can't be repeated. + } else { + tmp_states_to_be_added_.push_back(new_state); + } + } + return; +} + +void EarleyParser::AdvanceCharacterClass( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence +) { + XGRAMMAR_DCHECK(sub_sequence.type == GrammarExprType::kCharacterClass) + << "The element type is not supported!"; + + bool is_negative = static_cast(sub_sequence[0]); + + // The state is matching a UTF8 character (continuation bytes). + if (state.sub_element_id > 0) { + if ((ch & 0xC0) == 0x80) { + auto new_state = state; + new_state.sub_element_id--; + // Accumulate the codepoint from continuation byte + new_state.partial_codepoint = (new_state.partial_codepoint << 6) | (ch & 0x3F); + + // Check if the UTF8 character is completed. + if (new_state.sub_element_id == 0) { + if (is_negative) { + // For negative classes, accept if codepoint is NOT in any range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (!matches_range) { + new_state.element_id++; + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } else { + // For positive classes, accept if codepoint IS in a range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (matches_range) { + new_state.element_id++; + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } + } else { + // Check if partial codepoint could still potentially match any range + int32_t remaining_bytes = new_state.sub_element_id; + int32_t min_codepoint = new_state.partial_codepoint << (6 * remaining_bytes); + int32_t max_codepoint = min_codepoint | ((1 << (6 * remaining_bytes)) - 1); + + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: always continue (will verify on final byte) + // For positive classes: only continue if some range could match + bool should_continue = is_negative ? true : could_match; + if (should_continue) { + tmp_states_to_be_added_.push_back(new_state); + } + } + } + return; + } + + // Handle non-ASCII first bytes + if (!isascii(ch)) { + auto [accepted, num_bytes, partial] = HandleUTF8FirstByte(ch); + if (!accepted) { + return; + } + + XGRAMMAR_DCHECK(num_bytes > 1); + + // Compute possible codepoint range for this first byte + int32_t min_codepoint = partial << (6 * (num_bytes - 1)); + int32_t max_codepoint = min_codepoint | ((1 << (6 * (num_bytes - 1))) - 1); + + // Check if any stored range could potentially match + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + // Check for overlap between [min_codepoint, max_codepoint] and [lower, upper] + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: accept if no range could match (will verify on final byte) + // For positive classes: accept if some range could match (will verify on final byte) + bool should_continue = is_negative ? true : could_match; + + if (should_continue) { + auto new_state = state; + new_state.sub_element_id = num_bytes - 1; + new_state.partial_codepoint = partial; + tmp_states_to_be_added_.push_back(new_state); + } + return; + } + + // ASCII handling (unchanged) + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (static_cast(sub_sequence[i]) <= ch && + ch <= static_cast(sub_sequence[i + 1])) { + if (!is_negative) { + auto new_state = state; + new_state.element_id++; + new_state.sub_element_id = 0; + Enqueue(new_state); + } + return; + } + } + if (is_negative) { + auto new_state = state; + new_state.element_id++; + new_state.sub_element_id = 0; + Enqueue(new_state); + } +} + +void EarleyParser::AdvanceCharacterClassStar( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence +) { + XGRAMMAR_DCHECK(sub_sequence.type == GrammarExprType::kCharacterClassStar) + << "The element type is not supported!"; + + bool is_negative = static_cast(sub_sequence[0]); + + // The state is matching a UTF8 character (continuation bytes). + if (state.sub_element_id > 0) { + if ((ch & 0xC0) == 0x80) { + auto new_state = state; + new_state.sub_element_id--; + // Accumulate the codepoint from continuation byte + new_state.partial_codepoint = (new_state.partial_codepoint << 6) | (ch & 0x3F); + + // Check if the UTF8 character is completed. + if (new_state.sub_element_id == 0) { + if (is_negative) { + // For negative classes, accept if codepoint is NOT in any range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (!matches_range) { + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } else { + // For positive classes, accept if codepoint IS in a range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (matches_range) { + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } + } else { + // Check if partial codepoint could still potentially match any range + int32_t remaining_bytes = new_state.sub_element_id; + int32_t min_codepoint = new_state.partial_codepoint << (6 * remaining_bytes); + int32_t max_codepoint = min_codepoint | ((1 << (6 * remaining_bytes)) - 1); + + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: always continue (will verify on final byte) + // For positive classes: only continue if some range could match + bool should_continue = is_negative ? true : could_match; + if (should_continue) { + tmp_states_to_be_added_.push_back(new_state); + } + } + } + return; + } + + // Handle non-ASCII first bytes + if (!isascii(ch)) { + auto [accepted, num_bytes, partial] = HandleUTF8FirstByte(ch); + if (!accepted) { + return; + } + + XGRAMMAR_DCHECK(num_bytes > 1); + + // Compute possible codepoint range for this first byte + int32_t min_codepoint = partial << (6 * (num_bytes - 1)); + int32_t max_codepoint = min_codepoint | ((1 << (6 * (num_bytes - 1))) - 1); + + // Check if any stored range could potentially match + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + // Check for overlap between [min_codepoint, max_codepoint] and [lower, upper] + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: accept if no range could match (will verify on final byte) + // For positive classes: accept if some range could match (will verify on final byte) + bool should_continue = is_negative ? true : could_match; + + if (should_continue) { + auto new_state = state; + new_state.sub_element_id = num_bytes - 1; + new_state.partial_codepoint = partial; + tmp_states_to_be_added_.push_back(new_state); + } + return; + } + + // ASCII handling (unchanged) + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (static_cast(sub_sequence[i]) <= ch && + ch <= static_cast(sub_sequence[i + 1])) { + if (!is_negative) { + Enqueue(state); + } + return; + } + } + if (is_negative) { + Enqueue(state); + } +} + +void EarleyParser::AdvanceFsm(const ParserState& state, const uint8_t ch) { + XGRAMMAR_DCHECK(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& current_fsm = grammar_->per_rule_fsms[state.rule_id].value(); + for (const auto& edge : current_fsm.GetFsm().GetFsm().GetEdges(state.element_id)) { + if ((!edge.IsCharRange()) || ch < edge.min || ch > edge.max) { + continue; + } + auto new_state = state; + new_state.element_id = edge.target; + const uint8_t flags = GetFsmStateFlags(state.rule_id, edge.target); + if (!(flags & kFsmStateNonTerminal) && !(flags & kFsmStateEnd) && (flags & kFsmStateScanable)) { + EnqueueWithoutProcessing(std::move(new_state)); + } else { + Enqueue(std::move(new_state)); + } + } +} + +void EarleyParser::ScanAtomicToken(const ParserState& state, int32_t token_id) { + if (state.rule_id == -1) return; + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& current_fsm = grammar_->per_rule_fsms[state.rule_id].value(); + for (const auto& edge : current_fsm.GetFsm().GetFsm().GetEdges(state.element_id)) { + bool matched = false; + if (edge.IsToken()) { + auto info = current_fsm.GetFsm().GetFsm().GetTokenEdgeInfo(edge.GetAuxIndex()); + matched = info.Contains(token_id); + } else if (edge.IsExcludeToken()) { + auto info = current_fsm.GetFsm().GetFsm().GetExcludeTokenEdgeInfo(edge.GetAuxIndex()); + matched = info.Accepts(token_id); + } + if (!matched) continue; + auto new_state = state; + new_state.element_id = edge.target; + const uint8_t flags = GetFsmStateFlags(state.rule_id, edge.target); + if (!(flags & kFsmStateNonTerminal) && !(flags & kFsmStateEnd) && (flags & kFsmStateScanable)) { + EnqueueWithoutProcessing(std::move(new_state)); + } else { + Enqueue(std::move(new_state)); + } + } +} + +bool EarleyParser::AdvanceAtomicToken( + int32_t token_id, bool debug_print, int32_t token_char_count +) { + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()) + << "The tmp_process_state_queue_ should be empty before AdvanceAtomicToken."; + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_accept_stop_token_ = false; + tmp_completed_lazy_occurrences_.clear(); + if (has_char_budget_rules_) { + tmp_char_budget_entered_ = char_budget_entry_history_.back(); + char_count_history_.push_back(GetCurrentCharIndex() + token_char_count); + } + const auto& latest_states = scanable_state_history_[scanable_state_history_.size() - 1]; + for (const auto& state : latest_states) { + if (skip_expired_states_ && IsExpiredState(state)) { + continue; + } + ScanAtomicToken(state, token_id); + } + if (tmp_process_state_queue_.empty() && tmp_states_to_be_added_.empty()) { + if (has_char_budget_rules_) { + char_count_history_.pop_back(); + } + return false; + } + rule_id_to_completable_states_.PushBack(std::vector>()); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + is_completed_.push_back(tmp_accept_stop_token_); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + if (has_char_budget_rules_) { + char_budget_entry_history_.push_back(tmp_char_budget_entered_); + } + return true; +} + +bool RepeatDetector::IsVisited(const ParserState& state) const { + // If the size is larger than the threshold, then we use the set to check. + if (size_ > transition_threshold_) { + return visited_set_.find(state) != visited_set_.end(); + } + return std::find_if( + visited_vector_.begin(), + visited_vector_.begin() + size_, + [&state](const ParserState& s) { return StateEqualForParsing()(state, s); } + ) != visited_vector_.begin() + size_; +} + +void RepeatDetector::Insert(const ParserState& state) { + if (size_ == transition_threshold_) { + for (const auto& s : visited_vector_) { + visited_set_.insert(s); + } + } + size_++; + if (size_ > transition_threshold_) { + visited_set_.insert(state); + } else { + visited_vector_[size_ - 1] = state; + } +} + +void RepeatDetector::Clear() { + if (size_ > transition_threshold_) { + visited_set_.clear(); + } + size_ = 0; +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/earley_parser.h b/third_party/xgrammar/cpp/earley_parser.h new file mode 100644 index 0000000000..a3ef3e71ba --- /dev/null +++ b/third_party/xgrammar/cpp/earley_parser.h @@ -0,0 +1,802 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/earley_parser.h + * \brief The header for the definition of the Earley parser. + */ + +#ifndef XGRAMMAR_EARLEY_PARSER_H_ +#define XGRAMMAR_EARLEY_PARSER_H_ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_impl.h" +#include "support/compact_2d_array.h" +#include "support/utils.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief The state of the Earley parser. + * In the implementation, a rule can only be a kchoices or a ktagdispatch. + * A kchoices rule must be composed of some ksequence rules, or a kemptyrule. + * In the ksequence, every element in the sequence must be a kbytestring, a + * kcharacterclass, a kcharacterclassstar, or a rule reference. + * + * - rule_id: The id of the rule. + * - sequence_id: The id of the sequence in the rule. + * - element_id: The id of the element in the sequence, or the id of the node in + * the tag dispatch fsm. + * - rule_start_pos: The id of the parent node in the Earley parser. i.e. the rule + * is predicted from the k-th character. + * - sub_element_id: The id of the sub element in the current element, i.e.: + * - kbytestring: the id of the byte in the string. + * - kcharacterclass: How many bytes are left to be read in the utf8 character. + * - kcharacterclassstar: How many bytes are left to be read in the utf8 character. + */ +struct ParserState { + constexpr ParserState() = default; + + constexpr ParserState( + const int32_t& rule_id, + const int32_t& sequence_id, + const int32_t& element_id, + const int32_t& rule_start_pos, + const int32_t& budget_deadline = -1, + const int32_t& sub_element_id = 0, + const int32_t& repeat_count = 0, + const int32_t& partial_codepoint = 0, + const int32_t& active_temperature_rule_id = -1, + const int32_t& char_budget_deadline = -1 + ) + : rule_id(rule_id), + sequence_id(sequence_id), + element_id(element_id), + rule_start_pos(rule_start_pos), + budget_deadline(budget_deadline), + sub_element_id(sub_element_id), + repeat_count(repeat_count), + partial_codepoint(partial_codepoint), + active_temperature_rule_id(active_temperature_rule_id), + char_budget_deadline(char_budget_deadline) {} + + /*! + * \brief A rule_start_pos value of kNoPrevInputPos means this ParserState is the root of the + * parsing stack. + */ + static constexpr int32_t kNoPrevInputPos = -1; + + /*! \brief The rule's id. */ + int32_t rule_id = -1; + + /*! \brief Which choice in this rule is selected. */ + int32_t sequence_id = -1; + + /*! + * \brief Which element of the choice sequence is to be visited. When the current sequence is + * a tag dispatch rule, this element id is the current node. + */ + int32_t element_id = -1; + + /*! \brief The position of the state, i.e. from which position, the rule starts. */ + int32_t rule_start_pos = -1; + + /*! \brief The last token index this state's derivation may consume, from the token budget + * (Rule::max_tokens) of the rule it is inside; -1 means unlimited. Set when a budgeted rule + * is predicted and inherited by the states inside it. */ + int32_t budget_deadline = -1; + + /*! \brief The id of the sub element in the current element of the sequence. */ + int32_t sub_element_id = 0; + + /*! \brief The number of times the element is repeated. It will be used in kRepeat.*/ + int32_t repeat_count = 0; + + /*! \brief Partial codepoint accumulated during UTF-8 decoding for positive character classes. */ + int32_t partial_codepoint = 0; + + /*! \brief The innermost active rule that specifies a sampling temperature. */ + int32_t active_temperature_rule_id = -1; + + /*! \brief The number of Unicode codepoints this derivation may consume before its active + * character budget expires; -1 means unlimited. Stored as an absolute input position. */ + int32_t char_budget_deadline = -1; + + /*! + * \brief Lexicographic order over all fields. It is only used to sort the states for + * deterministic serialization, and is not needed during parsing. + */ + bool operator<(const ParserState& other) const { + if (rule_id != other.rule_id) return rule_id < other.rule_id; + if (sequence_id != other.sequence_id) return sequence_id < other.sequence_id; + if (element_id != other.element_id) return element_id < other.element_id; + if (rule_start_pos != other.rule_start_pos) return rule_start_pos < other.rule_start_pos; + if (budget_deadline != other.budget_deadline) return budget_deadline < other.budget_deadline; + if (sub_element_id != other.sub_element_id) return sub_element_id < other.sub_element_id; + if (repeat_count != other.repeat_count) return repeat_count < other.repeat_count; + if (partial_codepoint != other.partial_codepoint) { + return partial_codepoint < other.partial_codepoint; + } + if (active_temperature_rule_id != other.active_temperature_rule_id) { + return active_temperature_rule_id < other.active_temperature_rule_id; + } + return char_budget_deadline < other.char_budget_deadline; + } + + friend std::ostream& operator<<(std::ostream& os, const ParserState& state) { + os << state.ToString(); + return os; + } + + std::string ToString() const { + std::string result = "ParserState(rule_id=" + std::to_string(rule_id) + + ", sequence_id=" + std::to_string(sequence_id) + + ", element_id=" + std::to_string(element_id) + + ", rule_start_pos=" + std::to_string(rule_start_pos) + + ", sub_element_id=" + std::to_string(sub_element_id); + if (repeat_count != 0) { + result += ", repeat_count=" + std::to_string(repeat_count); + } + if (partial_codepoint != 0) { + result += ", partial_codepoint=" + std::to_string(partial_codepoint); + } + if (budget_deadline != -1) { + result += ", budget_deadline=" + std::to_string(budget_deadline); + } + if (active_temperature_rule_id != -1) { + result += ", active_temperature_rule_id=" + std::to_string(active_temperature_rule_id); + } + if (char_budget_deadline != -1) { + result += ", char_budget_deadline=" + std::to_string(char_budget_deadline); + } + result += ")"; + return result; + } +}; + +XGRAMMAR_MEMBER_ARRAY( + ParserState, + &ParserState::rule_id, + &ParserState::sequence_id, + &ParserState::element_id, + &ParserState::rule_start_pos, + &ParserState::budget_deadline, + &ParserState::sub_element_id, + &ParserState::repeat_count, + &ParserState::partial_codepoint, + &ParserState::active_temperature_rule_id, + &ParserState::char_budget_deadline +); + +/*! + * \brief Hash of a state used as the key of the adaptive token mask cache. The token mask of a + * state does not depend on rule_start_pos, repeat_count or partial_codepoint, so they are + * ignored. Pairs with StateEqualForCache. + */ +class StateHashForCache { + public: + size_t operator()(const ParserState& state) const { + return HashCombine(state.rule_id, state.sequence_id, state.element_id, state.sub_element_id); + } +}; + +/*! + * \brief Equality of states used as the key of the adaptive token mask cache. Compares the same + * fields as StateHashForCache hashes. + */ +class StateEqualForCache { + public: + bool operator()(const ParserState& lhs, const ParserState& rhs) const { + return lhs.rule_id == rhs.rule_id && lhs.sequence_id == rhs.sequence_id && + lhs.element_id == rhs.element_id && lhs.sub_element_id == rhs.sub_element_id; + } +}; + +/*! + * \brief When matching the state, we need to consider the rule_start_pos, since if two states + * don't have the same rule_start_pos, they are not the same state. + */ +class StateEqualForParsing { + public: + bool operator()(const ParserState& lhs, const ParserState& rhs) const { + return lhs.rule_id == rhs.rule_id && lhs.sequence_id == rhs.sequence_id && + lhs.element_id == rhs.element_id && lhs.rule_start_pos == rhs.rule_start_pos && + lhs.sub_element_id == rhs.sub_element_id && lhs.repeat_count == rhs.repeat_count && + lhs.partial_codepoint == rhs.partial_codepoint && + lhs.budget_deadline == rhs.budget_deadline && + lhs.active_temperature_rule_id == rhs.active_temperature_rule_id && + lhs.char_budget_deadline == rhs.char_budget_deadline; + } +}; + +/*! + * \brief This class is used to hash the ParserState for parsing. + * If two ParserStates don't have the same rule_start_pos, they are not the same state. + */ +class StateHashForParsing { + public: + size_t operator()(const ParserState& state) const { + return HashCombine( + state.rule_id, + state.sequence_id, + state.element_id, + state.rule_start_pos, + state.sub_element_id, + state.repeat_count, + state.partial_codepoint, + state.budget_deadline, + state.active_temperature_rule_id, + state.char_budget_deadline + ); + } +}; + +/*! \brief This class is used to detect the repeated states. */ +class RepeatDetector { + private: + const int transition_threshold_; + + std::vector visited_vector_; + + std::unordered_set visited_set_; + + int size_ = 0; + + public: + RepeatDetector(const int transition_threshold = 50) + : transition_threshold_(transition_threshold), size_(0) { + visited_vector_.resize(transition_threshold_); + } + + /*! + * \brief Check if the element is visited. + * \return True if visited, false otherwise. + */ + bool IsVisited(const ParserState& state) const; + + /*! + * \brief Add the state into the visited states. + * \param state The state to be added. + */ + void Insert(const ParserState& state); + + /*! \brief Reset the detector. */ + void Clear(); +}; + +/*! \brief A concrete occurrence of a captured rule in an Earley parent chain. */ +struct CaptureOccurrence { + /*! \brief The id of the captured rule. */ + int32_t rule_id; + /*! \brief The position where the rule occurrence started. */ + int32_t start_pos; + + bool operator==(const CaptureOccurrence& other) const { + return rule_id == other.rule_id && start_pos == other.start_pos; + } +}; + +/*! + * \brief A completion event of a captured rule, recorded when the rule is completed during + * parsing. The matched span is [start_pos, r) in input positions, where r is the position (i.e. + * the history row) at which the event is recorded. + */ +struct CaptureEvent { + /*! \brief The id of the completed rule. */ + int32_t rule_id; + /*! \brief The position where the rule started matching. kNoPrevInputPos means position 0 (the + * rule acts as the root). */ + int32_t start_pos; + /*! \brief The unadjusted start position of this rule occurrence. This differs from start_pos + * for the zero-width event inserted after a dynamic-dispatch marker. */ + int32_t occurrence_start_pos; + /*! \brief Number of trailing bytes hidden only from this rule's own capture for this + * completion. */ + int32_t hidden_suffix_bytes = 0; + /*! \brief Number of trailing bytes hidden from every containing capture for this completion. */ + int32_t hidden_stop_bytes = 0; + /*! \brief The captured rule occurrences whose concrete Earley parent chains contain this stop + * completion. Includes this rule's occurrence when the rule itself is captured. */ + std::vector stop_capture_targets; +}; + +class EarleyParser { + /*! + * \brief Here is an article about Earley Parser. + * https://en.wikipedia.org/wiki/Earley_parser#Pseudocode + * We divide the parser states into three categories: + * - Scanable (which will be stored in scanable_state_history_). + * - Predictable(If it predict a new rule successfully, then it will be stored in + * rule_id_to_completable_states). + * - completable(which can perform a completion operation). + * A state will be stored in rule_id_to_completable_states_ if it can be completed, + * and it will be stored in scanable_state_history_ if it can be scanned. Otherwise, + * it will be discarded. + */ + protected: + using GrammarExpr = Grammar::Impl::GrammarExpr; + + /*! \brief The grammar to be parsed. */ + Grammar grammar_; + + /*! \brief In this round of advancing, check if the stop token can be accepted. */ + bool tmp_accept_stop_token_ = false; + + /*! \brief store when accepting i characters, if the stop token can be accepted. */ + std::vector is_completed_; + + /*! + * \brief rule_id_to_completable_states[i][j] is the i pos j rule_id states. Earley + * parser needs it to complete. + */ + Compact2DArray> rule_id_to_completable_states_; + + /*! + * \brief The states history. state_stack[i] is a vector storing the states after accepting the + * input[i-1]. + */ + Compact2DArray scanable_state_history_; + + /*! + * \brief A temporary vector only used in Advance, used to add states in the + * scanable_state_history. + */ + std::vector tmp_states_to_be_added_; + + /*! \brief It's the processing queue of the earley parser. */ + std::queue tmp_process_state_queue_; + + /*! \brief The class is used to check if a state has been added into the queue. */ + RepeatDetector tmp_states_visited_in_queue_; + + /*! \brief Check if the stop token is accepted. */ + bool stop_token_is_accepted_ = false; + + enum FsmStateFlag : uint8_t { + kFsmStateInitialized = 1 << 0, + kFsmStateScanable = 1 << 1, + kFsmStateNonTerminal = 1 << 2, + kFsmStateEnd = 1 << 3, + kFsmStateHasEdges = 1 << 4, + }; + + /*! \brief Lazily-computed FSM state properties, indexed by rule id and state id. */ + std::vector> fsm_state_flags_cache_; + + /*! \brief Whether each rule can match the empty string. */ + std::vector rule_is_nullable_; + + /*! \brief Compute and cache properties for a state in a per-rule FSM. */ + uint8_t InitializeFsmStateFlags(int32_t rule_id, int32_t state_id); + + /*! \brief Return cached properties for a state in a per-rule FSM. */ + uint8_t GetFsmStateFlags(int32_t rule_id, int32_t state_id) { + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(fsm_state_flags_cache_.size())); + auto& flags_cache = fsm_state_flags_cache_[rule_id]; + if (!flags_cache.empty()) { + XGRAMMAR_DCHECK(state_id >= 0 && state_id < static_cast(flags_cache.size())); + if (flags_cache[state_id] != 0) { + return flags_cache[state_id]; + } + } + return InitializeFsmStateFlags(rule_id, state_id); + } + + bool IsRuleNullable(int32_t rule_id) const { return rule_is_nullable_[rule_id] != 0; } + + /*! \brief The index of the LLM token currently being accepted, set by the matcher; -1 + * before any token. budget_deadline values are compared against it. */ + int32_t current_token_index_ = -1; + + /*! \brief Whether states past their budget deadline are skipped when scanning. Enabled by + * the matcher for accepts that follow an enforcing mask computation. */ + bool skip_expired_states_ = false; + + /*! \brief Whether any rule of the grammar has a token budget. */ + bool has_budget_rules_ = false; + + /*! \brief The number of Unicode codepoints accepted at every parser history row. */ + std::vector char_count_history_; + + /*! \brief Whether any rule of the grammar has a character budget. */ + bool has_char_budget_rules_ = false; + + /*! \brief Whether a character-budgeted occurrence was entered since the initial parser row. */ + std::vector char_budget_entry_history_; + + /*! \brief Entry-history value for the row currently being expanded. */ + bool tmp_char_budget_entered_ = false; + + /*! \brief Whether the state's derivation may not consume the next token. */ + bool IsExpiredState(const ParserState& state) const { + return state.budget_deadline >= 0 && current_token_index_ > state.budget_deadline; + } + + /*! \brief The deadline for a newly predicted occurrence of the rule: its own budget counted + * from the current token, capped by the parent's deadline for nested budgets. */ + int32_t DeadlineForRule(int32_t rule_id, int32_t parent_deadline) const { + int32_t own = grammar_->GetRule(rule_id).max_tokens; + if (own < 0) { + return parent_deadline; + } + int64_t uncapped_deadline = static_cast(current_token_index_) + own; + int32_t deadline = static_cast( + std::min(uncapped_deadline, std::numeric_limits::max()) + ); + return parent_deadline >= 0 ? std::min(deadline, parent_deadline) : deadline; + } + + /*! \brief The character deadline for a newly predicted rule occurrence. */ + int32_t CharDeadlineForRule(int32_t rule_id, int32_t parent_deadline) const { + int32_t own = grammar_->GetRule(rule_id).max_chars; + if (own < 0) { + return parent_deadline; + } + int32_t current_char_index = GetCurrentCharIndex(); + int32_t deadline = own > std::numeric_limits::max() - current_char_index + ? std::numeric_limits::max() + : current_char_index + own; + return parent_deadline >= 0 ? std::min(deadline, parent_deadline) : deadline; + } + + /*! \brief Whether the state's derivation may not consume another Unicode codepoint. */ + bool IsCharExpiredState(const ParserState& state) const { + return state.char_budget_deadline >= 0 && GetCurrentCharIndex() >= state.char_budget_deadline; + } + + static bool StartsUTF8Codepoint(uint8_t byte) { return (byte & 0xC0) != 0x80; } + + /*! \brief Whether any rule of the grammar has a capture or stop_capture name. Fixed at + * construction. When false, the capture machinery is fully disabled and has no overhead. */ + bool capture_tracking_ = false; + + /*! \brief Whether the grammar contains suffix/stop spans that may affect captures. */ + bool has_hidden_capture_rules_ = false; + + /*! + * \brief Whether capture events are currently recorded in Complete(). Only enabled during + * definitive advances (accepting a token or string), not during speculative exploration + * (mask computation, jump-forward search, lookahead checks), so that speculative completions + * never produce capture events. + */ + bool capture_recording_ = false; + + /*! + * \brief The history of capture events. capture_event_history_[i] stores the events recorded + * when input position i was created. Kept aligned with scanable_state_history_ row-by-row + * whenever capture_tracking_ is true, so PopLastStates rolls back events automatically. + */ + Compact2DArray capture_event_history_; + + /*! \brief Returns true if the rule exists and has a capture name. */ + bool RuleHasCapture(int32_t rule_id) const { + return capture_tracking_ && rule_id >= 0 && !grammar_->GetRule(rule_id).capture_name.empty(); + } + + /*! \brief Returns true if completing this rule can hide bytes from a capture. */ + bool RuleHasHiddenBytes(int32_t rule_id) const { + if (!capture_tracking_ || !has_hidden_capture_rules_ || rule_id < 0) { + return false; + } + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(rule_id); + return suffix_stop_info != nullptr && + (suffix_stop_info->hidden_suffix_bytes > 0 || suffix_stop_info->hidden_stop_bytes > 0); + } + + /*! \brief Returns true if completing this rule must produce a capture-history event. */ + bool RuleNeedsCaptureEvent(int32_t rule_id) const { + return RuleHasCapture(rule_id) || RuleHasHiddenBytes(rule_id); + } + + /*! \brief Record a capture or hidden-span event for a completed rule in the current row. */ + void RecordCaptureEvent(const ParserState& state, bool marker_present); + + /*! + * \brief Whether this completion of a suffix/stop rule actually consumed the trailing marker. + * A non-looping TagDispatch can also complete before its trigger is encountered, which is what + * lets a free-text tail end normally. Only the terminal post-dispatch state has no outgoing + * edges; completions in the trigger-scanning states did not consume a suffix/stop marker. + */ + bool CompletionConsumedMarker(const ParserState& state) const; + + /*! \brief Collect the captured rule occurrences whose concrete Earley parent chains contain + * the given stop completion. Includes the completed rule's own occurrence when captured. */ + std::vector CollectStopCaptureTargets(const ParserState& state) const; + + /*! + * \brief The lazy rule occurrences (rule_id, rule_start_pos) completed while building the + * current row. Committed-shortest matching: their remaining states are removed when the row is + * finalized, so the occurrence cannot be extended further. + */ + std::vector> tmp_completed_lazy_occurrences_; + + /*! \brief Remove the states of the lazy occurrences completed in the current row. */ + void RemoveCommittedLazyStates(); + + /*! + * \brief Check if the state has been added into the queue. + * \param state The state to check. + * \return True if in the vector, false otherwise. + */ + bool IsStateVisitedInQueue(const ParserState& state) const { + return tmp_states_visited_in_queue_.IsVisited(state); + } + + /*! + * \brief The scanning operation of the Earley parser. Put the new states in the queue. + */ + void Scan(const ParserState& state, const uint8_t ch); + + /*! + * \brief The completion operation of the Earley parser. + * \param state The state to be completed. + * \param debug_print Whether to print the debug information. + * \details The reason is that if the state can't be scanned, then + * add it into the next states is useless. Moreover, the end + * of the grammar is used to check if the grammar is completed, + * so it should be added into the next states. + */ + void Complete(const ParserState& state, bool debug_print = false, bool marker_present = true); + + /*! + * \brief The prediction operation of the Earley parser. + * \param state The state to be predicted. + * \param debug_print Whether to print the debug information. + * \return First: If the state scanable, or the state is the end of the grammar, + * then return true, otherwise return false. + * \return Second: If the state is completable, then return true, otherwise return false. + */ + std::pair Predict(const ParserState& state, bool debug_print = false); + + /*! \brief The initial state expanded from the root rule of the grammar. */ + ParserState RootInitialState() const; + + /*! \brief Resolve the active temperature rule when entering a rule. */ + int32_t ResolveActiveTemperatureRule(int32_t rule_id, int32_t inherited_rule_id) const; + + /*! + * \brief Expand the rule, used for RuleRef and kTagDispatch. + * \param state The state to be expanded, which is the parent state. + * The type of the state is kTagDispatch or kSequence. Moreover, the + * element of the sequence should be a rule reference; the node in + * the kTagDispatch should be an end node. + * \param grammar_expr The grammar expression to be expanded. + * \param sub_grammar_expr The sub grammar expression to be expanded, especially + * when the rule is a kSequence, and the sub rule is a kRuleRef. + * \param debug_print Whether to print the debug information. + */ + void ExpandNextRuleRefElement( + const ParserState& state, + const GrammarExpr& grammar_expr, + const GrammarExpr* sub_grammar_expr, + bool debug_print = false + ); + + /*! + * \brief Expand the rule, used for RuleRef and kTagDispatch. + * \param state The state to be expanded, and it's should be on the FSM. + * \param debug_print Whether to print the debug information. + */ + void ExpandNextRuleRefElementOnFSM(const ParserState& state, bool debug_print = false); + + /*! + * \brief Advance the parser to the next state, with the sub sequence is kCharacterClass. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \param sub_sequence The sub sequence to be checked. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceCharacterClass( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence + ); + + /*! + * \brief Advance the parser to the next state, with the sub sequence is kByteString. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \param sub_sequence The sub sequence to be checked. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceByteString( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence + ); + + /*! + * \brief Advance the parser to the next state, with the sub sequence is kCharacterClassStar. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \param sub_sequence The sub sequence to be checked. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceCharacterClassStar( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence + ); + + /*! + * \brief Advance the parser to the next state, with the sequence is kTagDispatch. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceFsm(const ParserState& state, const uint8_t ch); + + /*! + * \brief Scan a token edge: check if token_id matches any kToken or kExcludeToken edge from + * state. + */ + void ScanAtomicToken(const ParserState& state, int32_t token_id); + + /*! + * \brief Advance the parser by accepting a whole token via kToken/kExcludeToken edges. + * \param token_id The token ID to accept. + * \param debug_print Whether to print debug info. + * \return True if any state advanced, false otherwise. + */ + bool AdvanceAtomicToken(int32_t token_id, bool debug_print = false, int32_t token_char_count = 0); + + /*! + * \brief Enqueue the state into the queue. + * \param state The state to be enqueued. + * \details The state is enqueued if it is not visited in the queue. + */ + void Enqueue(const ParserState& state) { + if (!IsStateVisitedInQueue(state)) { + tmp_process_state_queue_.push(state); + tmp_states_visited_in_queue_.Insert(state); + } + } + + /*! + * \brief Enqueue the state into the queue, without prediction and completion. + * \param state The state to be enqueued. + */ + void EnqueueWithoutProcessing(const ParserState& state) { + if (!IsStateVisitedInQueue(state)) { + tmp_states_visited_in_queue_.Insert(state); + tmp_states_to_be_added_.push_back(state); + } + } + + public: + /*! + * \brief Constructor of the Earley parser. + * \param grammar The grammar to be parsed. It must be optimized. + * \param initial_state The state to start parsing from. If not provided, parsing starts + * from the root rule of the grammar. + */ + explicit EarleyParser( + const Grammar& grammar, std::optional initial_state = std::nullopt + ); + + /*! + * \brief From the current states, advance to the next state. + * \param ch The character to be advanced. + * \param debug_print Whether to print the debug information. + * \return True if the character is accepted, false otherwise. + * \note If the character isn't accepted, then the states won't be changed. + */ + bool Advance(const uint8_t ch, bool debug_print = false); + + /*! + * \brief Remove the newly added states. + * \param count The number of states to be removed. + */ + void PopLastStates(int32_t count = 1); + + /*! + * \brief Check whether any of the multiple states stored in the parser has already completed. + * \note Since the parser contains multiple parallel states, some may have already completed, + * while others might still be able to accept more characters. + * \return True if the root rule is completed, false otherwise. + */ + bool IsCompleted() const; + + /*! + * \brief Push the initial state into the Earley parser. + * \param state The initial state to be pushed. + */ + void PushStateAndExpand(const ParserState& state); + + /*! + * \brief Reset the parser. + * \note This function is used to reset the parser, and initialize the + * parser with the root rule. + */ + void Reset(); + + /*! + * \brief Get the current scanable states. + * \return The scanable states. + */ + std::vector GetLatestScanableStates() const { + std::vector latest_states; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + latest_states.push_back(state); + } + return latest_states; + } + + /*! + * \brief Push one state to check if it can accept the token. + * \param state The state to be pushed. + */ + void PushOneStateToCheck(const ParserState& state) { + PushStatesToCheck(std::vector{state}, is_completed_.back()); + } + + /*! \brief Push a temporary parser row for token-mask checking. */ + void PushStatesToCheck(const std::vector& states, bool completed) { + rule_id_to_completable_states_.PushBack(std::vector>()); + is_completed_.push_back(completed); + scanable_state_history_.PushBack(states); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + if (has_char_budget_rules_) { + char_count_history_.push_back(GetCurrentCharIndex()); + char_budget_entry_history_.push_back(char_budget_entry_history_.back()); + } + } + + /*! \brief Push a character-count row for a parser row created by the matcher. */ + void PushCharCountRow(int32_t char_count, bool char_budget_entered) { + if (!has_char_budget_rules_) { + return; + } + char_count_history_.push_back(char_count); + char_budget_entry_history_.push_back(char_budget_entered); + } + + int32_t GetCurrentCharIndex() const { + return char_count_history_.empty() ? 0 : char_count_history_.back(); + } + + bool HasEnteredCharBudget() const { + return has_char_budget_rules_ && char_budget_entry_history_.back(); + } + + /*! \brief Whether the grammar has any captured rule. */ + bool IsCaptureTrackingEnabled() const { return capture_tracking_; } + + /*! \brief Copy the capture events of the latest input position. */ + std::vector CopyLastCaptureRow() const { + if (!capture_tracking_) { + return {}; + } + auto row = capture_event_history_[capture_event_history_.size() - 1]; + return std::vector(row.begin(), row.end()); + } + + /*! + * \brief Push a new row of capture events. Used when a new input position is created outside + * of Advance / AdvanceAtomicToken (e.g. when merging parallel advance results), to keep the + * capture history aligned with the state history. + */ + void PushCaptureRow(const std::vector& events) { + if (capture_tracking_) { + capture_event_history_.PushBack(events); + } + } + + std::string PrintStates() const { + std::string result; + result += "There are " + std::to_string(scanable_state_history_.size()) + + " steps in history. Last step: [\n"; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + result += state.ToString() + ", \n"; + } + result += "]"; + return result; + } +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_EARLEY_PARSER_H_ diff --git a/third_party/xgrammar/cpp/ebnf_script_creator.h b/third_party/xgrammar/cpp/ebnf_script_creator.h new file mode 100644 index 0000000000..5ffd06c047 --- /dev/null +++ b/third_party/xgrammar/cpp/ebnf_script_creator.h @@ -0,0 +1,188 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/ebnf_script_creator.h + * \brief The header for the creating EBNF script. + */ + +#ifndef XGRAMMAR_EBNF_SCRIPT_CREATOR_H_ +#define XGRAMMAR_EBNF_SCRIPT_CREATOR_H_ + +#include + +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/logging.h" +#include "support/utils.h" + +namespace xgrammar { + +/*! + * \brief A class for creating EBNF grammar scripts. + * + * This class helps build EBNF (Extended Backus-Naur Form) grammar scripts + * by managing rules and their content. + */ +class EBNFScriptCreator { + public: + /*! \brief Constructor */ + EBNFScriptCreator() = default; + + /*! + * \brief Adds a new rule to the grammar with a suggested name + * \param rule_name_hint Suggested name for the rule + * \param rule_body The EBNF content/definition of the rule + * \return The actual name assigned to the rule + */ + std::string AddRule(const std::string& rule_name_hint, const std::string& rule_body) { + return AddRuleWithAllocatedName(AllocateRuleName(rule_name_hint), rule_body); + } + + /*! + * \brief Generates a new rule name based on a suggested name + * \param rule_name_hint Suggested name for the rule + * \return The actual name assigned to the rule + */ + std::string AllocateRuleName(const std::string& rule_name_hint) { + if (rule_names_.find(rule_name_hint) == rule_names_.end()) { + rule_names_.insert(rule_name_hint); + return rule_name_hint; + } + for (int i = 0; i < NAME_SUFFIX_MAXIMUM; ++i) { + std::string rule_name = rule_name_hint + "_" + std::to_string(i); + if (rule_names_.find(rule_name) == rule_names_.end()) { + rule_names_.insert(rule_name); + return rule_name; + } + } + XGRAMMAR_LOG(FATAL) << "Cannot find a unique rule name for " << rule_name_hint; + XGRAMMAR_UNREACHABLE(); + } + + /*! + * \brief Adds a new rule to the grammar with a allocated name. Used with AllocateRuleName() + * \param rule_name The name of the rule to add + * \param rule_body The EBNF content/definition of the rule + * \return The actual name assigned to the rule + */ + std::string AddRuleWithAllocatedName(const std::string& rule_name, const std::string& rule_body) { + XGRAMMAR_CHECK(rule_names_.find(rule_name) != rule_names_.end()) + << "Rule name " << rule_name << " is not allocated"; + rules_.emplace_back(rule_name, rule_body); + return rule_name; + } + + /*! + * \brief Concatenates a list of strings with a space separator + * \param items The list of strings to concatenate + * \return The concatenated string + */ + static std::string Concat(const std::vector& items) { + std::stringstream ss; + ss << "("; + for (int i = 0; i < static_cast(items.size()); ++i) { + if (i > 0) { + ss << " "; + } + ss << items[i]; + } + ss << ")"; + return ss.str(); + } + + /*! + * \brief Joins a list of strings with an OR operator + * \param items The list of strings to join + * \return The joined string + */ + static std::string Or(const std::vector& items) { + std::stringstream ss; + ss << "("; + for (int i = 0; i < static_cast(items.size()); ++i) { + if (i > 0) { + ss << " | "; + } + ss << items[i]; + } + ss << ")"; + return ss.str(); + } + + /*! + * \brief Escape and quote a string + * \param str The string to escape and quote + * \return The escaped and quoted string + */ + static std::string Str(const std::string& str) { + std::stringstream ss; + ss << "\"" << EscapeString(str) << "\""; + return ss.str(); + } + + /*! + * \brief Repeats an item a given number of times + * \param item The item to repeat + * \param min The minimum number of times to repeat the item + * \param max The maximum number of times to repeat the item + * \return The repeated string + */ + static std::string Repeat(const std::string& item, int min, int max) { + std::stringstream ss; + ss << item; + if (min == 0 && max == 1) { + ss << "?"; + } else if (min == 0 && max == -1) { + ss << "*"; + } else if (min == 1 && max == -1) { + ss << "+"; + } else if (min == 0 && max == 0) { + return ""; + } else if (min == max) { + ss << "{" << min << "}"; + } else if (max == -1) { + ss << "{" << min << ",}"; + } else { + ss << "{" << min << "," << max << "}"; + } + return ss.str(); + } + + /*! + * \brief Gets the complete EBNF grammar script + * \return The full EBNF grammar script as a string + */ + std::string GetScript() { + std::string script = ""; + for (const auto& rule : rules_) { + script += rule.first + " ::= " + rule.second + "\n"; + } + return script; + } + + /*! + * \brief Retrieves the content/definition of a specific rule + * \param rule_name The name of the rule to look up + * \return The EBNF content/definition of the specified rule + */ + std::string GetRuleContent(const std::string& rule_name) { + auto it = std::find_if(rules_.begin(), rules_.end(), [rule_name](const auto& rule) { + return rule.first == rule_name; + }); + if (it != rules_.end()) { + return it->second; + } + return ""; + } + + private: + std::vector> rules_; + std::unordered_set rule_names_; + const int NAME_SUFFIX_MAXIMUM = 10000; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_EBNF_SCRIPT_CREATOR_H_ diff --git a/third_party/xgrammar/cpp/fsm.cc b/third_party/xgrammar/cpp/fsm.cc new file mode 100644 index 0000000000..9249c12063 --- /dev/null +++ b/third_party/xgrammar/cpp/fsm.cc @@ -0,0 +1,2041 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm.cc + */ +#include "fsm.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/json_serializer.h" +#include "support/logging.h" +#include "support/reflection.h" +#include "support/union_find_set.h" +#include "support/utils.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/****************** FSMImplBase ******************/ + +template +class FSMImplBase { + static_assert( + std::is_same_v>> || + std::is_same_v>, + "ContainerType must be std::vector> or Compact2DArray" + ); + + public: + /*! \brief Default constructor. */ + FSMImplBase() = default; + + FSMImplBase(const ContainerType& edges, std::vector edge_aux_data = {}) + : edges_(edges), edge_aux_data_(std::move(edge_aux_data)) {} + + FSMImplBase(ContainerType&& edges, std::vector edge_aux_data = {}) + : edges_(std::move(edges)), edge_aux_data_(std::move(edge_aux_data)) {} + + int NumStates() const { return edges_.size(); } + + std::string EdgesToString(std::optional> states = std::nullopt) const; + + const ContainerType& GetEdges() const { return edges_; } + + // For std::vector>, return const std::vector& to avoid copying. + // For Compact2DArray, return Compact2DArray::Row since it is just a simple + // pointer. + decltype(auto) GetEdges(int state) const { return edges_[state]; } + + void GetEpsilonClosure(std::unordered_set* state_set) const; + + void GetPossibleRules(int state_num, std::unordered_set* rules) const; + + void GetReachableStates(const std::vector& from, std::unordered_set* result) const; + + const std::vector& GetEdgeAuxData() const { return edge_aux_data_; } + + void SetEdgeAuxData(std::vector data) { edge_aux_data_ = std::move(data); } + + RepeatEdgeRef GetRepeatEdgeInfo(int32_t idx) const { return {edge_aux_data_.data() + idx}; } + + TokenEdgeRef GetTokenEdgeInfo(int32_t idx) const { return {edge_aux_data_.data() + idx}; } + + ExcludeTokenEdgeRef GetExcludeTokenEdgeInfo(int32_t idx) const { + return {edge_aux_data_.data() + idx}; + } + + protected: + ContainerType edges_; + std::vector edge_aux_data_; + friend struct member_trait; +}; + +template +std::string FSMImplBase::EdgesToString(std::optional> states +) const { + std::string result = "[\n"; + auto f_print_one = [&, this](int i) { + result += std::to_string(i) + ": ["; + const auto& edges = edges_[i]; + for (int j = 0; j < static_cast(edges.size()); ++j) { + const auto& edge = edges[j]; + if (edge.min >= 0 && edge.min != edge.max) { + std::string char_min_str = EscapeString(static_cast(edge.min)); + std::string char_max_str = EscapeString(static_cast(edge.max)); + result += "[" + char_min_str + "-" + char_max_str + "]->" + std::to_string(edge.target); + } else if (edge.min >= 0 && edge.min == edge.max) { + std::string char_str = EscapeString(static_cast(edge.min)); + result += "'" + char_str + "'->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kRuleRef) { + result += "Rule(" + std::to_string(edge.max) + ")->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kEpsilon) { + result += "Eps->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kEOS) { + result += "EOS->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kRepeatRef) { + auto info = GetRepeatEdgeInfo(edge.max); + result += "Repeat(rule=" + std::to_string(info.RuleId()) + + ", min=" + std::to_string(info.Lower()) + + ", max=" + std::to_string(info.Upper()) + ")->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kToken) { + auto info = GetTokenEdgeInfo(edge.max); + result += "Token("; + for (int32_t k = 0; k < info.Count(); ++k) { + if (k > 0) result += ", "; + result += std::to_string(info.TokenIds()[k]); + } + result += ")->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kExcludeToken) { + auto info = GetExcludeTokenEdgeInfo(edge.max); + result += "ExcludeToken("; + for (int32_t k = 0; k < info.Count(); ++k) { + if (k > 0) result += ", "; + result += std::to_string(info.TokenIds()[k]); + } + result += ")->" + std::to_string(edge.target); + } + if (j < static_cast(edges.size()) - 1) { + result += ", "; + } + } + result += "]\n"; + }; + if (states.has_value()) { + for (int i : states.value()) { + f_print_one(i); + } + } else { + for (int i = 0; i < int(NumStates()); ++i) { + f_print_one(i); + } + } + result += "]"; + return result; +} + +template +void FSMImplBase::GetEpsilonClosure(std::unordered_set* state_set) const { + std::queue queue; + for (const auto& state : *state_set) { + queue.push(state); + } + while (!queue.empty()) { + int current = queue.front(); + queue.pop(); + for (const auto& edge : edges_[current]) { + if (!edge.IsEpsilon()) { + continue; + } + if (state_set->find(edge.target) != state_set->end()) { + continue; + } + state_set->insert(edge.target); + queue.push(edge.target); + } + } +} + +template +void FSMImplBase::GetPossibleRules(int state, std::unordered_set* rules) const { + rules->clear(); + for (const auto& edge : edges_[state]) { + if (edge.IsRuleRef()) { + rules->insert(edge.GetRefRuleId()); + } + } +} + +template +void FSMImplBase::GetReachableStates( + const std::vector& from, std::unordered_set* result +) const { + result->clear(); + std::queue queue; + for (const auto& state : from) { + queue.push(state); + result->insert(state); + } + while (!queue.empty()) { + int current = queue.front(); + queue.pop(); + for (const auto& edge : edges_[current]) { + if (result->find(edge.target) != result->end()) { + continue; + } + result->insert(edge.target); + queue.push(edge.target); + } + } +} + +/****************** FSM::Impl ******************/ + +class FSM::Impl : public FSMImplBase>> { + using EdgeType = FSMEdge::EdgeType; + + public: + Impl() = default; + + Impl(int num_states = 0) { edges_.resize(num_states); } + + using FSMImplBase>>::FSMImplBase; + + int GetNextState(int from, int value, EdgeType edge_type) const; + + using FSMImplBase>>::GetEdges; + + std::vector>& GetEdges() { return edges_; } + + std::vector& GetEdges(int state) { return edges_[state]; } + + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + EdgeType edge_type, + bool from_is_closure + ) const; + + int AddState() { + edges_.emplace_back(); + return edges_.size() - 1; + } + + void AddEdge(int from, int to, int32_t min, int32_t max) { + XGRAMMAR_DCHECK(from < static_cast(edges_.size())); + edges_[from].push_back({min, max, to}); + } + + void AddRuleEdge(int from, int to, int32_t rule_id) { + AddEdge(from, to, FSMEdge::EdgeType::kRuleRef, rule_id); + } + + void AddEpsilonEdge(int from, int to) { AddEdge(from, to, FSMEdge::EdgeType::kEpsilon, 0); } + + void AddEOSEdge(int from, int to) { AddEdge(from, to, FSMEdge::EdgeType::kEOS, 0); } + + void AddRepeatEdge(int from, int to, int32_t rule_id, int32_t lower, int32_t upper) { + XGRAMMAR_DCHECK(edges_[from].empty()) + << "A state with a kRepeatRef edge must have no other outgoing edges."; + XGRAMMAR_DCHECK(edge_aux_data_.size() <= INT32_MAX); + int32_t aux_index = static_cast(edge_aux_data_.size()); + edge_aux_data_.reserve(edge_aux_data_.size() + 3); + edge_aux_data_.emplace_back(rule_id); + edge_aux_data_.emplace_back(lower); + edge_aux_data_.emplace_back(upper); + AddEdge(from, to, FSMEdge::EdgeType::kRepeatRef, aux_index); + } + + void AddTokenEdge(int from, int to, const std::vector& token_ids) { + XGRAMMAR_DCHECK(!token_ids.empty()) << "Token set must not be empty"; + XGRAMMAR_CHECK(edge_aux_data_.size() <= INT32_MAX) + << "edge_aux_data_ overflow: too many auxiliary data entries"; + int32_t aux_index = static_cast(edge_aux_data_.size()); + edge_aux_data_.push_back(static_cast(token_ids.size())); + for (int32_t id : token_ids) { + edge_aux_data_.push_back(id); + } + edges_[from].push_back(FSMEdge(FSMEdge::EdgeType::kToken, aux_index, to)); + } + + void AddExcludeTokenEdge(int from, int to, const std::vector& token_ids) { + XGRAMMAR_DCHECK(!token_ids.empty()) << "Token exclude set must not be empty"; + XGRAMMAR_CHECK(edge_aux_data_.size() <= INT32_MAX) + << "edge_aux_data_ overflow: too many auxiliary data entries"; + int32_t aux_index = static_cast(edge_aux_data_.size()); + edge_aux_data_.push_back(static_cast(token_ids.size())); + for (int32_t id : token_ids) { + edge_aux_data_.push_back(id); + } + edges_[from].push_back(FSMEdge(FSMEdge::EdgeType::kExcludeToken, aux_index, to)); + } + + void AddFSM(const FSM& fsm, std::vector* state_mapping); + + FSM RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const; + + void SortEdges(); + + CompactFSM ToCompact(); + + friend class FSMWithStartEnd; +}; + +int FSM::Impl::GetNextState(int from, int value, EdgeType edge_type) const { + XGRAMMAR_DCHECK(edge_type != EdgeType::kEpsilon) + << "Should not call GetNextState with edge type kEpsilon."; + if (edge_type == EdgeType::kCharRange) { + for (const auto& edge : edges_[from]) { + if (edge.min >= EdgeType::kCharRange && edge.min <= value && edge.max >= value) { + return edge.target; + } + } + return FSM::kNoNextState; + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& edge : edges_[from]) { + if (edge.min == EdgeType::kRuleRef && edge.GetRefRuleId() == value) { + return edge.target; + } + } + return FSM::kNoNextState; + } else if (edge_type == EdgeType::kEOS) { + for (const auto& edge : edges_[from]) { + if (edge.min == EdgeType::kEOS) { + return edge.target; + } + } + return FSM::kNoNextState; + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + XGRAMMAR_DCHECK(edges_[from].size() == 1 && edges_[from][0].IsRepeatRef()); + return edges_[from][0].target; + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } + XGRAMMAR_UNREACHABLE(); +} + +void FSM::Impl::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + EdgeType edge_type, + bool from_is_closure +) const { + XGRAMMAR_DCHECK(edge_type != EdgeType::kEpsilon) + << "Should not call Advance with edge type kEpsilon."; + + const std::unordered_set* start_closure; + std::unordered_set start_closure_tmp; + + if (from_is_closure) { + start_closure = &from; + } else { + start_closure_tmp.insert(from.begin(), from.end()); + GetEpsilonClosure(&start_closure_tmp); + start_closure = &start_closure_tmp; + } + + result->clear(); + + if (edge_type == EdgeType::kCharRange) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsCharRange() && edge.min <= value && edge.max >= value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsRuleRef() && edge.GetRefRuleId() == value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kEOS) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsEOS()) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + for (const auto& state : *start_closure) { + if (!edges_[state].empty() && edges_[state][0].IsRepeatRef()) { + result->insert(edges_[state][0].target); + } + } + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } + + // Get the epsilon closure of the result. + GetEpsilonClosure(result); +} + +void FSM::Impl::AddFSM(const FSM& fsm, std::vector* state_mapping) { + int old_num_states = NumStates(); + int32_t aux_offset = static_cast(edge_aux_data_.size()); + + const auto& other_aux = fsm.GetEdgeAuxData(); + edge_aux_data_.insert(edge_aux_data_.end(), other_aux.begin(), other_aux.end()); + + if (state_mapping != nullptr) { + state_mapping->clear(); + state_mapping->reserve(fsm.NumStates()); + for (int i = 0; i < fsm.NumStates(); ++i) { + state_mapping->push_back(i + old_num_states); + } + } + + edges_.resize(edges_.size() + fsm.NumStates()); + + for (int i = 0; i < fsm.NumStates(); ++i) { + for (const auto& edge : fsm.GetEdges()[i]) { + int32_t max_val = edge.max; + if (edge.IsAuxEdge() && aux_offset > 0) { + max_val = static_cast(edge.max + aux_offset); + } + AddEdge(i + old_num_states, edge.target + old_num_states, edge.min, max_val); + } + } +} + +FSM FSM::Impl::RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const { + std::vector> new_edges(new_num_states); + for (int i = 0; i < static_cast(edges_.size()); ++i) { + for (const auto& edge : edges_[i]) { + if (edge.IsEpsilon() && state_mapping[i] == state_mapping[edge.target]) { + continue; // Skip self-loops for epsilon edges. + } + new_edges[state_mapping[i]].emplace_back(edge.min, edge.max, state_mapping[edge.target]); + } + } + // aux_indices remain stable since only state ids are remapped + for (int i = 0; i < new_num_states; ++i) { + std::sort(new_edges[i].begin(), new_edges[i].end()); + const auto& end_iter = std::unique(new_edges[i].begin(), new_edges[i].end()); + new_edges[i].erase(end_iter, new_edges[i].end()); + } + return FSM(std::move(new_edges), std::vector(edge_aux_data_)); +} + +void FSM::Impl::SortEdges() { + for (int i = 0; i < static_cast(edges_.size()); ++i) { + std::sort(edges_[i].begin(), edges_[i].end()); + } +} + +CompactFSM FSM::Impl::ToCompact() { + SortEdges(); + Compact2DArray edges; + for (int i = 0; i < static_cast(edges_.size()); ++i) { + edges.PushBack(edges_[i]); + } + return CompactFSM(std::move(edges), std::move(edge_aux_data_)); +} + +/****************** FSM ******************/ + +FSM::FSM(int num_states) : pimpl_(std::make_shared(num_states)) {} + +FSM::FSM(const std::vector>& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(edges, std::move(edge_aux_data))) {} + +FSM::FSM(std::vector>&& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(std::move(edges), std::move(edge_aux_data))) {} + +int FSM::NumStates() const { return pimpl_->NumStates(); } + +int FSM::AddState() { return pimpl_->AddState(); } + +void FSM::AddEdge(int from, int to, int32_t min, int32_t max) { + pimpl_->AddEdge(from, to, min, max); +} + +void FSM::AddEdge(int from, int to, FSMEdge::EdgeType type, int32_t value) { + pimpl_->AddEdge(from, to, type, value); +} + +void FSM::AddEpsilonEdge(int from, int to) { pimpl_->AddEpsilonEdge(from, to); } + +void FSM::AddRuleEdge(int from, int to, int32_t rule_id) { pimpl_->AddRuleEdge(from, to, rule_id); } + +void FSM::AddEOSEdge(int from, int to) { pimpl_->AddEOSEdge(from, to); } + +void FSM::AddRepeatEdge(int from, int to, int32_t rule_id, int32_t lower, int32_t upper) { + pimpl_->AddRepeatEdge(from, to, rule_id, lower, upper); +} + +void FSM::AddTokenEdge(int from, int to, const std::vector& token_ids) { + pimpl_->AddTokenEdge(from, to, token_ids); +} + +void FSM::AddExcludeTokenEdge(int from, int to, const std::vector& token_ids) { + pimpl_->AddExcludeTokenEdge(from, to, token_ids); +} + +const std::vector& FSM::GetEdgeAuxData() const { return pimpl_->GetEdgeAuxData(); } + +void FSM::SetEdgeAuxData(std::vector data) { pimpl_->SetEdgeAuxData(std::move(data)); } + +RepeatEdgeRef FSM::GetRepeatEdgeInfo(int32_t idx) const { return pimpl_->GetRepeatEdgeInfo(idx); } + +TokenEdgeRef FSM::GetTokenEdgeInfo(int32_t idx) const { return pimpl_->GetTokenEdgeInfo(idx); } + +ExcludeTokenEdgeRef FSM::GetExcludeTokenEdgeInfo(int32_t idx) const { + return pimpl_->GetExcludeTokenEdgeInfo(idx); +} + +void FSM::AddFSM(const FSM& fsm, std::vector* state_mapping) { + pimpl_->AddFSM(fsm, state_mapping); +} + +std::string FSM::EdgesToString(std::optional> states) const { + return pimpl_->EdgesToString(states); +} + +const std::vector& FSM::GetEdges(int state) const { return pimpl_->GetEdges(state); } + +std::vector>& FSM::GetEdges() { return pimpl_->GetEdges(); } + +const std::vector>& FSM::GetEdges() const { return pimpl_->GetEdges(); } + +std::vector& FSM::GetEdges(int state) { return pimpl_->GetEdges(state); } + +FSM FSM::Copy() const { return FSM(std::make_shared(*pimpl_)); } + +int FSM::GetNextState(int from, int value, FSMEdge::EdgeType edge_type) const { + return pimpl_->GetNextState(from, value, edge_type); +} + +void FSM::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure +) const { + pimpl_->Advance(from, value, result, edge_type, from_is_closure); +} + +void FSM::GetPossibleRules(int state, std::unordered_set* rules) const { + pimpl_->GetPossibleRules(state, rules); +} + +void FSM::GetEpsilonClosure(std::unordered_set* state_set) const { + pimpl_->GetEpsilonClosure(state_set); +} + +void FSM::GetReachableStates(const std::vector& from, std::unordered_set* result) const { + pimpl_->GetReachableStates(from, result); +} + +FSM FSM::RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const { + return pimpl_->RebuildWithMapping(state_mapping, new_num_states); +} + +void FSM::SortEdges() { pimpl_->SortEdges(); } + +CompactFSM FSM::ToCompact() { return pimpl_->ToCompact(); } + +/****************** CompactFSM::Impl ******************/ + +class CompactFSM::Impl : public FSMImplBase> { + using EdgeType = FSMEdge::EdgeType; + + public: + Impl() = default; + + Impl(const Compact2DArray& edges, std::vector edge_aux_data = {}) + : FSMImplBase>(edges, std::move(edge_aux_data)), + edge_num_(ComputeEdgeNum(edges_)) {} + + Impl(Compact2DArray&& edges, std::vector edge_aux_data = {}) + : FSMImplBase>(std::move(edges), std::move(edge_aux_data)), + edge_num_(ComputeEdgeNum(edges_)) {} + + void GetNextStates(int from, int value, EdgeType edge_type, std::vector* target) const; + + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure + ) const; + + FSM ToFSM() const; + + size_t GetNumEdges() const { return edge_num_; } + + /*! + * \brief Check that every edge target and every auxiliary data reference is in range. Used after + * deserialization, where the fields are restored verbatim. + * \return An error message if the FSM is malformed. + */ + std::optional Validate() const { + const int64_t aux_size = edge_aux_data_.size(); + for (int state = 0; state < NumStates(); ++state) { + for (const auto& edge : edges_[state]) { + if (edge.target < 0 || edge.target >= NumStates()) { + return "Edge target " + std::to_string(edge.target) + " is out of range"; + } + if (!edge.IsAuxEdge()) { + continue; + } + // A repeat edge owns 3 aux elements; a token edge owns a count followed by count ids. + const int64_t idx = edge.max; + bool in_range = idx >= 0 && idx < aux_size; + if (in_range && edge.IsRepeatRef()) { + in_range = idx + 3 <= aux_size; + } else if (in_range) { + in_range = edge_aux_data_[idx] >= 0 && idx + 1 + edge_aux_data_[idx] <= aux_size; + } + if (!in_range) { + return "Edge aux index " + std::to_string(idx) + " is out of range"; + } + } + } + if (edge_num_ != ComputeEdgeNum(edges_)) { + return "edge_num does not match the number of edges"; + } + return std::nullopt; + } + + size_t edge_num_ = 0; + + friend std::size_t MemorySize(const Impl& impl) { + return MemorySize(impl.edges_) + MemorySize(impl.edge_aux_data_) + sizeof(impl.edge_num_); + } + + private: + static size_t ComputeEdgeNum(const Compact2DArray& edges) { + size_t edge_num = 0; + for (int i = 0; i < edges.size(); ++i) { + edge_num += edges[i].size(); + } + return edge_num; + } +}; + +XGRAMMAR_MEMBER_TABLE( + CompactFSM::Impl, + "edges", + &CompactFSM::Impl::edges_, + "edge_aux_data", + &CompactFSM::Impl::edge_aux_data_, + "edge_num", + &CompactFSM::Impl::edge_num_ +); + +void CompactFSM::Impl::GetNextStates( + int from, int value, EdgeType edge_type, std::vector* targets +) const { + targets->clear(); + XGRAMMAR_DCHECK(edge_type != EdgeType::kEpsilon) + << "Should not call GetNextState with edge type kEpsilon."; + if (edge_type == EdgeType::kCharRange) { + for (const auto& edge : edges_[from]) { + if (edge.min < EdgeType::kCharRange) { + continue; + } else if (edge.min > value) { + break; + } else if (edge.max >= value) { + targets->push_back(edge.target); + } + } + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& edge : edges_[from]) { + if (edge.min < EdgeType::kRuleRef) { + continue; + } else if (edge.min > EdgeType::kRuleRef) { + break; + } else if (edge.GetRefRuleId() == value) { + targets->push_back(edge.target); + } + } + } else if (edge_type == EdgeType::kEOS) { + for (const auto& edge : edges_[from]) { + if (edge.min < EdgeType::kEOS) { + continue; + } else if (edge.min > EdgeType::kEOS) { + break; + } else if (edge.max >= EdgeType::kEOS) { + targets->push_back(edge.target); + } + } + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + for (const auto& edge : edges_[from]) { + if (edge.IsRepeatRef()) { + targets->push_back(edge.target); + break; + } + } + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } +} + +void CompactFSM::Impl::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure +) const { + const std::unordered_set* start_closure; + std::unordered_set start_closure_tmp; + + if (from_is_closure) { + start_closure = &from; + } else { + start_closure_tmp.insert(from.begin(), from.end()); + GetEpsilonClosure(&start_closure_tmp); + start_closure = &start_closure_tmp; + } + + result->clear(); + + if (edge_type == EdgeType::kCharRange) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.min < EdgeType::kCharRange) { + continue; + } else if (edge.min > value) { + break; + } else if (edge.max >= value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.min < EdgeType::kRuleRef) { + continue; + } else if (edge.min > EdgeType::kRuleRef) { + break; + } else if (edge.GetRefRuleId() == value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kEOS) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.min < EdgeType::kEOS) { + continue; + } else if (edge.min > EdgeType::kEOS) { + break; + } else if (edge.max >= EdgeType::kEOS) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsRepeatRef()) { + result->insert(edge.target); + break; + } + } + } + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } + + // Get the epsilon closure of the result. + GetEpsilonClosure(result); +} + +FSM CompactFSM::Impl::ToFSM() const { + std::vector> edges(NumStates()); + for (int i = 0; i < edges_.size(); i++) { + const auto& row = edges_[i]; + edges[i].insert(edges[i].end(), row.begin(), row.end()); + } + return FSM(std::move(edges), std::vector(edge_aux_data_)); +} + +/****************** CompactFSM ******************/ + +CompactFSM::CompactFSM(const Compact2DArray& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(edges, std::move(edge_aux_data))) {} + +CompactFSM::CompactFSM(Compact2DArray&& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(std::move(edges), std::move(edge_aux_data))) {} + +int CompactFSM::NumStates() const { return pimpl_->NumStates(); } + +const Compact2DArray& CompactFSM::GetEdges() const { return pimpl_->GetEdges(); } + +Compact2DArray::Row CompactFSM::GetEdges(int state) const { + return pimpl_->GetEdges(state); +} + +std::string CompactFSM::EdgesToString(std::optional> states) const { + return pimpl_->EdgesToString(states); +} + +void CompactFSM::GetNextStates( + int from, int value, FSMEdge::EdgeType edge_type, std::vector* targets +) const { + return pimpl_->GetNextStates(from, value, edge_type, targets); +} + +void CompactFSM::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure +) const { + pimpl_->Advance(from, value, result, edge_type, from_is_closure); +} + +void CompactFSM::GetPossibleRules(int state_num, std::unordered_set* rules) const { + pimpl_->GetPossibleRules(state_num, rules); +} + +void CompactFSM::GetEpsilonClosure(std::unordered_set* state_set) const { + pimpl_->GetEpsilonClosure(state_set); +} + +void CompactFSM::GetReachableStates(const std::vector& from, std::unordered_set* result) + const { + pimpl_->GetReachableStates(from, result); +} + +size_t CompactFSM::GetNumEdges() const { return pimpl_->GetNumEdges(); } + +FSM CompactFSM::ToFSM() const { return pimpl_->ToFSM(); } + +const std::vector& CompactFSM::GetEdgeAuxData() const { return pimpl_->GetEdgeAuxData(); } + +void CompactFSM::SetEdgeAuxData(std::vector data) { + pimpl_->SetEdgeAuxData(std::move(data)); +} + +RepeatEdgeRef CompactFSM::GetRepeatEdgeInfo(int32_t idx) const { + return pimpl_->GetRepeatEdgeInfo(idx); +} + +TokenEdgeRef CompactFSM::GetTokenEdgeInfo(int32_t idx) const { + return pimpl_->GetTokenEdgeInfo(idx); +} + +ExcludeTokenEdgeRef CompactFSM::GetExcludeTokenEdgeInfo(int32_t idx) const { + return pimpl_->GetExcludeTokenEdgeInfo(idx); +} + +picojson::value SerializeJSONValue(const CompactFSM& value) { + return detail::json_serializer::AutoSerializeJSONValuePImpl(value); +} + +std::optional DeserializeJSONValue( + CompactFSM* result, const picojson::value& value, const std::string& type_name +) { + return detail::json_serializer::AutoDeserializeJSONValuePImpl(result, value, type_name); +} + +struct CompactFSMWithStartEndSerializeHelper { + CompactFSM fsm; + int start; + bool is_dfa; + std::vector end_index; + size_t edge_num; + + CompactFSMWithStartEndSerializeHelper(const CompactFSMWithStartEnd& compact_fsm_with_se) + : fsm(compact_fsm_with_se.fsm_), + start(compact_fsm_with_se.start_), + is_dfa(compact_fsm_with_se.is_dfa_), + end_index(compact_fsm_with_se.ends_), + edge_num(compact_fsm_with_se.edge_num_) {} + + CompactFSMWithStartEndSerializeHelper() = default; + + std::optional Validate() const { + if (fsm.IsNull()) { + return "Expect a non-null fsm"; + } + auto in_range = [&](int32_t state) { return state >= 0 && state < fsm.NumStates(); }; + if (!in_range(start) || !std::all_of(end_index.begin(), end_index.end(), in_range)) { + return "The start or end state is out of range"; + } + return std::nullopt; + } +}; + +XGRAMMAR_MEMBER_ARRAY( + CompactFSMWithStartEndSerializeHelper, + &CompactFSMWithStartEndSerializeHelper::fsm, + &CompactFSMWithStartEndSerializeHelper::start, + &CompactFSMWithStartEndSerializeHelper::end_index, + &CompactFSMWithStartEndSerializeHelper::is_dfa, + &CompactFSMWithStartEndSerializeHelper::edge_num +); + +picojson::value SerializeJSONValue(const CompactFSMWithStartEnd& value) { + return AutoSerializeJSONValue(CompactFSMWithStartEndSerializeHelper(value)); +} +std::optional DeserializeJSONValue( + CompactFSMWithStartEnd* result, const picojson::value& value, const std::string& type_name +) { + CompactFSMWithStartEndSerializeHelper tmp; + auto err = AutoDeserializeJSONValue(&tmp, value, type_name); + if (err.has_value()) { + return err; + } + result->fsm_ = std::move(tmp.fsm); + result->start_ = tmp.start; + result->is_dfa_ = tmp.is_dfa; + result->edge_num_ = tmp.edge_num; + result->SetEndStates(std::move(tmp.end_index)); + return std::nullopt; +} + +struct CompactFSMWithStartEndWithSizeSerializeHelper { + CompactFSMWithStartEnd fsm; + size_t edge_num; + size_t node_num; + + CompactFSMWithStartEndWithSizeSerializeHelper( + const CompactFSMWithStartEndWithSize& compact_fsm_with_size + ) + : fsm(compact_fsm_with_size.fsm_), + edge_num(compact_fsm_with_size.edge_num_), + node_num(compact_fsm_with_size.node_num_) {} + + CompactFSMWithStartEndWithSizeSerializeHelper() = default; +}; + +XGRAMMAR_MEMBER_ARRAY( + CompactFSMWithStartEndWithSizeSerializeHelper, + &CompactFSMWithStartEndWithSizeSerializeHelper::fsm, + &CompactFSMWithStartEndWithSizeSerializeHelper::edge_num, + &CompactFSMWithStartEndWithSizeSerializeHelper::node_num +); + +picojson::value SerializeJSONValue(const CompactFSMWithStartEndWithSize& value) { + return AutoSerializeJSONValue(CompactFSMWithStartEndWithSizeSerializeHelper(value)); +} + +std::optional DeserializeJSONValue( + CompactFSMWithStartEndWithSize* result, + const picojson::value& value, + const std::string& type_name +) { + CompactFSMWithStartEndWithSizeSerializeHelper tmp; + auto err = AutoDeserializeJSONValue(&tmp, value, type_name); + if (err.has_value()) { + return err; + } + result->fsm_ = std::move(tmp.fsm); + result->edge_num_ = tmp.edge_num; + result->node_num_ = tmp.node_num; + return std::nullopt; +} + +/****************** FSMWithStartEnd ******************/ + +std::string FSMWithStartEnd::ToString() const { + std::string result; + result += "FSM(num_states=" + std::to_string(NumStates()) + ", start=" + std::to_string(start_) + + ", end=["; + + std::unordered_set reachable_states; + GetReachableStates(&reachable_states); + std::vector reachable_states_vec(reachable_states.begin(), reachable_states.end()); + std::sort(reachable_states_vec.begin(), reachable_states_vec.end()); + + bool first = true; + for (auto end : ends_) { + if (!first) { + result += ", "; + } + first = false; + result += std::to_string(end); + } + + result += "], edges=" + fsm_.EdgesToString(reachable_states_vec) + ")"; + return result; +} + +std::ostream& operator<<(std::ostream& os, const FSMWithStartEnd& fsm) { + os << fsm.ToString(); + return os; +} + +FSMWithStartEnd FSMWithStartEnd::Copy() const { + return FSMWithStartEnd(fsm_.Copy(), start_, ends_, is_dfa_); +} + +FSMWithStartEnd FSMWithStartEnd::RebuildWithMapping( + const std::vector& state_mapping, int new_num_states +) const { + FSM new_fsm = fsm_.RebuildWithMapping(state_mapping, new_num_states); + auto new_start = state_mapping[start_]; + std::vector new_ends; + new_ends.reserve(ends_.size()); + for (auto end : ends_) { + new_ends.push_back(state_mapping[end]); + } + return FSMWithStartEnd(new_fsm, new_start, std::move(new_ends)); +} + +CompactFSMWithStartEnd FSMWithStartEnd::ToCompact() { + return CompactFSMWithStartEnd(fsm_.ToCompact(), start_, ends_, is_dfa_); +} + +FSMWithStartEndWithSize FSMWithStartEnd::AddToCompleteFSM( + FSM* complete_fsm, std::vector* state_mapping +) { + XGRAMMAR_DCHECK(state_mapping != nullptr) << "state_mapping cannot be nullptr"; + complete_fsm->AddFSM(fsm_, state_mapping); + int new_start = (*state_mapping)[start_]; + // Map the end states to the states in the complete FSM. The mapping is monotonic, so the + // sorted invariant is preserved. The sparse representation is important here: this method is + // called once per rule, and the ends of each returned view must not scale with the size of the + // complete FSM, otherwise the total cost is O(num_rules * num_total_states). + std::vector new_ends; + new_ends.reserve(ends_.size()); + for (auto end : ends_) { + new_ends.push_back((*state_mapping)[end]); + } + + int num_edges = 0; + for (int i = 0; i < fsm_.NumStates(); ++i) { + num_edges += fsm_.GetEdges(i).size(); + } + + int num_nodes = fsm_.NumStates(); + + auto fsm_with_se = FSMWithStartEnd(*complete_fsm, new_start, std::move(new_ends), is_dfa_); + + return FSMWithStartEndWithSize(fsm_with_se, num_edges, num_nodes); +} + +FSMWithStartEnd FSMWithStartEnd::Star() const { + FSM fsm = fsm_.Copy(); + auto new_start = fsm.AddState(); + for (auto end : ends_) { + fsm.AddEpsilonEdge(end, new_start); + } + fsm.AddEpsilonEdge(new_start, start_); + return FSMWithStartEnd(fsm, new_start, {new_start}); +} + +FSMWithStartEnd FSMWithStartEnd::Plus() const { + FSM fsm = fsm_.Copy(); + for (auto end : ends_) { + fsm.AddEpsilonEdge(end, start_); + } + return FSMWithStartEnd(fsm, start_, ends_); +} + +FSMWithStartEnd FSMWithStartEnd::Optional() const { + FSM fsm = fsm_.Copy(); + if (!ends_.empty()) { + fsm.AddEpsilonEdge(start_, ends_.front()); + } + return FSMWithStartEnd(fsm, start_, ends_); +} + +Result FSMWithStartEnd::Not(int max_result_num_states) const { + // Check if the FSM contains any rule references. + if (!IsLeaf()) { + XGRAMMAR_LOG(FATAL) << "Not operation is not supported for FSM with rule references."; + } + FSMWithStartEnd result; + if (is_dfa_) { + result = Copy(); + } else { + Result dfa_result = ToDFA(max_result_num_states); + if (dfa_result.IsErr()) { + return dfa_result; + } + result = std::move(dfa_result).Unwrap(); + } + // Reverse all the final states. + std::vector new_final_states; + for (int i = 0; i < result.NumStates(); ++i) { + if (!result.IsEndState(i)) { + new_final_states.push_back(i); // Mark all states as final except the original final states. + } + } + + // Add a new final state that accepts all characters. + int accept_all_new_state = result.AddState(); + new_final_states.push_back(accept_all_new_state); + + std::bitset<256> char_set; + for (int i = 0; i < result.NumStates(); i++) { + char_set.reset(); + // Collect all characters that are not accepted by the original FSM. + for (const auto& edge : result.GetFsm().GetEdges(i)) { + if (edge.IsCharRange()) { + for (int j = edge.min; j <= edge.max; ++j) { + char_set.set(j); + } + } + } + // Add edges for characters that are not accepted. + for (int left_bound = 0; left_bound < 256; ++left_bound) { + if (char_set[left_bound]) { + continue; // Skip characters that are accepted. + } + int right_bound = left_bound + 1; + while (right_bound < 256 && !char_set[right_bound]) { + ++right_bound; + } + result.GetFsm().AddEdge(i, accept_all_new_state, left_bound, right_bound - 1); + left_bound = right_bound; + } + } + + result.SetEndStates(new_final_states); + return ResultOk(result); +} + +FSMWithStartEnd FSMWithStartEnd::Union(const std::vector& fsms) { + // Put all the FSMs in parallel. + // Allocate a new start state. Start state will be linked to the start states of all the FSMs. + // The end states of the new FSM will be the union of the end states of all the FSMs. + if (fsms.size() == 1) { + return fsms[0]; + } + XGRAMMAR_DCHECK(fsms.size() > 1) << "Union of 0 FSMs is not allowed."; + + FSM fsm(1); + int start = 0; + std::vector ends; + + std::vector state_mapping; + + for (const auto& fsm_with_se : fsms) { + fsm.AddFSM(fsm_with_se.GetFsm(), &state_mapping); + fsm.AddEpsilonEdge(start, state_mapping[fsm_with_se.GetStart()]); + for (auto end : fsm_with_se.GetEnds()) { + ends.push_back(state_mapping[end]); + } + } + + return FSMWithStartEnd(fsm, start, std::move(ends)); +} + +FSMWithStartEnd FSMWithStartEnd::Concat(const std::vector& fsms) { + // For each FSM, link the end states to the start state of the next FSM. + // Set the start state of the first FSM as the start state of the result. + // Set the end states of the last FSM as the end states of the result. + if (fsms.size() == 1) { + return fsms[0]; + } + XGRAMMAR_DCHECK(fsms.size() > 1) << "Concatenation of 0 FSMs is not allowed."; + + FSM fsm; + int start = 0; + std::vector ends; + + std::vector state_mapping; + std::vector previous_ends; + + for (int i = 0; i < static_cast(fsms.size()); ++i) { + fsm.AddFSM(fsms[i].GetFsm(), &state_mapping); + if (i == 0) { + start = state_mapping[fsms[i].GetStart()]; + } else { + auto this_start = state_mapping[fsms[i].GetStart()]; + for (const auto& end : previous_ends) { + fsm.AddEpsilonEdge(end, this_start); + } + } + if (i == static_cast(fsms.size()) - 1) { + for (auto end : fsms[i].GetEnds()) { + ends.push_back(state_mapping[end]); + } + } else { + previous_ends.clear(); + previous_ends.reserve(fsms[i].GetEnds().size()); + for (auto end : fsms[i].GetEnds()) { + previous_ends.push_back(state_mapping[end]); + } + } + } + + return FSMWithStartEnd(fsm, start, std::move(ends)); +} + +Result FSMWithStartEnd::Intersect( + const FSMWithStartEnd& lhs, const FSMWithStartEnd& rhs, int max_result_num_states +) { + if (!lhs.IsLeaf() || !rhs.IsLeaf()) { + return ResultErr("Intersect only support leaf fsm!"); + } + auto lhs_dfa_raw = lhs.ToDFA(); + auto rhs_dfa_raw = rhs.ToDFA(); + + if (lhs_dfa_raw.IsErr()) { + return lhs_dfa_raw; + } + if (rhs_dfa_raw.IsErr()) { + return rhs_dfa_raw; + } + + auto lhs_dfa = std::move(lhs_dfa_raw).Unwrap(); + auto rhs_dfa = std::move(rhs_dfa_raw).Unwrap(); + // Initialize the result FSM. + FSM result_fsm(0); + FSMWithStartEnd result(result_fsm, 0, std::vector(), true); + std::unordered_map, int> state_map; + std::unordered_set> visited; + std::queue> queue; + queue.push({lhs_dfa.GetStart(), rhs_dfa.GetStart()}); + result.AddState(); + state_map[{lhs_dfa.GetStart(), rhs_dfa.GetStart()}] = 0; + while (!queue.empty()) { + auto [lhs_state, rhs_state] = std::move(queue.front()); + if (lhs_dfa.IsEndState(lhs_state) && rhs_dfa.IsEndState(rhs_state)) { + result.AddEndState(state_map[{lhs_state, rhs_state}]); + } + queue.pop(); + for (const auto& lhs_edge : lhs_dfa.GetFsm().GetEdges(lhs_state)) { + for (const auto& rhs_edge : rhs_dfa.GetFsm().GetEdges(rhs_state)) { + XGRAMMAR_DCHECK(lhs_edge.IsCharRange() && rhs_edge.IsCharRange()); + // Check if the edges intersect. + if (lhs_edge.min > rhs_edge.max || rhs_edge.min > lhs_edge.max) { + continue; // No intersection. + } + int min_value = std::max(lhs_edge.min, rhs_edge.min); + int max_value = std::min(lhs_edge.max, rhs_edge.max); + if (state_map.find(std::make_pair(lhs_edge.target, rhs_edge.target)) == state_map.end()) { + state_map[{lhs_edge.target, rhs_edge.target}] = result.AddState(); + queue.push({lhs_edge.target, rhs_edge.target}); + } + int target_state = state_map[{lhs_edge.target, rhs_edge.target}]; + result.GetFsm().AddEdge( + state_map[{lhs_state, rhs_state}], target_state, min_value, max_value + ); + } + } + } + return ResultOk(std::move(result)); +} + +bool FSMWithStartEnd::IsDFA() { + if (is_dfa_) { + return true; + } + std::bitset<256> character_transitions; + std::unordered_set rule_transitions; + for (const auto& edges : fsm_->GetEdges()) { + character_transitions.reset(); + rule_transitions.clear(); + for (const auto& edge : edges) { + if (edge.IsEpsilon()) { + return false; // Epsilon transitions are not allowed in DFA. + } + if (edge.IsCharRange()) { + for (int i = edge.min; i <= edge.max; ++i) { + if (character_transitions[i]) { + return false; // Duplicate character transition. + } + character_transitions.set(i); + } + continue; + } + if (edge.IsRuleRef()) { + if (rule_transitions.find(edge.GetRefRuleId()) != rule_transitions.end()) { + return false; // Duplicate rule transition. + } + rule_transitions.insert(edge.GetRefRuleId()); + } + // kRepeatRef: by invariant, a state with kRepeatRef has exactly one edge, always + // deterministic. + } + } + is_dfa_ = true; + return true; +} + +FSMWithStartEnd FSMWithStartEnd::SimplifyEpsilon(int max_num_states) const { + if (is_dfa_) { + return *this; + } + if (NumStates() > max_num_states) { + return *this; + } + + UnionFindSet union_find_set; + std::vector in_degree(NumStates(), 0); + std::vector> epsilon_edges; + + std::vector has_exclude_token(NumStates(), false); + for (int i = 0; i < NumStates(); i++) { + for (const auto& edge : fsm_->GetEdges(i)) { + if (edge.IsExcludeToken()) { + has_exclude_token[i] = true; + break; + } + } + } + + for (int i = 0; i < NumStates(); i++) { + const auto& edges = fsm_->GetEdges(i); + for (const auto& edge : edges) { + in_degree[edge.target]++; + if (edge.IsEpsilon()) { + // a -- epsilon --> b, and a doesn't have other outward edges. Do not merge an + // accepting a into a non-accepting b: the merged state would be accepting, so other + // paths reaching b would be wrongly accepted. (If b is accepting, a effectively + // accepts already via the epsilon edge, so merging is safe.) + if (edges.size() == 1 && !has_exclude_token[i] && !has_exclude_token[edge.target] && + (!IsEndState(i) || IsEndState(edge.target))) { + union_find_set.Add(i); + union_find_set.Add(edge.target); + union_find_set.Union(i, edge.target); + in_degree[edge.target]--; // Remove the inward edge since a and b are merged. + } else { + // Otherwise, we store it to check for the second merge rule. + epsilon_edges.emplace_back(i, edge.target); + } + } + } + } + + // Build the equivalent graph. + std::vector equiv_node(NumStates()); + for (int i = 0; i < NumStates(); i++) { + if (union_find_set.Count(i)) { + equiv_node[i] = union_find_set.Find(i); + if (equiv_node[i] == i) { + continue; + } + in_degree[equiv_node[i]] += in_degree[i]; + } else { + equiv_node[i] = i; + } + } + + // a --> epsilon --> b, and b doesn't have other inward edges. + for (const auto& [from_raw, to_raw] : epsilon_edges) { + const int& from = equiv_node[from_raw]; + const int& to = equiv_node[to_raw]; + if (in_degree[to] == 1 && equiv_node[GetStart()] != to && !has_exclude_token[from_raw] && + !has_exclude_token[to_raw]) { + union_find_set.Add(from); + union_find_set.Add(to); + union_find_set.Union(from, to); + } + } + + // Merge the states. + auto eq_classes = union_find_set.GetAllSets(); + if (eq_classes.empty()) { + return *this; + } + + std::vector new_to_old(NumStates(), -1); + for (size_t i = 0; i < eq_classes.size(); i++) { + for (const auto& state : eq_classes[i]) { + new_to_old[state] = i; + } + } + + int cnt = eq_classes.size(); + for (int i = 0; i < NumStates(); i++) { + if (new_to_old[i] == -1) { + new_to_old[i] = cnt; + cnt++; + } + } + return RebuildWithMapping(new_to_old, cnt); +} + +FSMWithStartEnd FSMWithStartEnd::MergeEquivalentStates(int max_result_num_states) const { + if (max_result_num_states < NumStates()) { + return *this; + } + // No merge is possible with fewer than 4 states (need >=2 sources sharing a target, + // or >=2 sinks sharing a source). + if (NumStates() < 4) { + return Copy(); + } + bool changed = true; + FSMWithStartEnd result = Copy(); + result.GetFsm()->SortEdges(); + UnionFindSet union_find_set; + + // A compact edge view used for incoming/outgoing CSR rows. `peer` means source state in + // incoming_edges and target state in outgoing_edges. + struct EndpointEdge { + int peer; // source in incoming_edges, target in outgoing_edges + int32_t min; + int32_t max; + + bool operator<(const EndpointEdge& other) const { + return std::tie(peer, min, max) < std::tie(other.peer, other.min, other.max); + } + }; + + // Scratch buffers reused across iterations to avoid repeated vector allocation. + // Number of incoming edges for each state, used to size the incoming CSR rows. + std::vector incoming_row_sizes; + // Number of outgoing edges for each state, used to size the outgoing CSR rows. + std::vector outgoing_row_sizes; + // Write positions while filling incoming_edges rows. + std::vector incoming_write_positions; + // Write positions while filling outgoing_edges rows. + std::vector outgoing_write_positions; + // Incoming edges grouped by target state. + Compact2DArray incoming_edges; + // Outgoing edges grouped by source state. + Compact2DArray outgoing_edges; + // Number of distinct predecessor states for each state. + std::vector incoming_distinct_count; + // Number of distinct successor states for each state. + std::vector outgoing_distinct_count; + // The only predecessor state when incoming_distinct_count[state] == 1. + std::vector single_incoming_source; + // The only successor state when outgoing_distinct_count[state] == 1. + std::vector single_outgoing_target; + // Terminal end states collected for leaf-state merging. + std::vector no_successor_end_states; + // Terminal non-end states collected for leaf-state merging. + std::vector no_successor_non_end_states; + + while (changed) { + int n = result.NumStates(); + union_find_set.Clear(); + + // First pass: count row sizes for the incoming/outgoing CSR arrays. + incoming_row_sizes.assign(n, 0); + outgoing_row_sizes.assign(n, 0); + for (int source = 0; source < n; ++source) { + const auto& edges = result.GetFsm().GetEdges(source); + outgoing_row_sizes[source] = static_cast(edges.size()); + for (const auto& edge : edges) { + ++incoming_row_sizes[edge.target]; + } + } + + // Allocate CSR rows. The underlying storage is reset and reused across iterations. + incoming_edges.ResetWithRowSizes(incoming_row_sizes); + outgoing_edges.ResetWithRowSizes(outgoing_row_sizes); + incoming_write_positions.assign(n, 0); + outgoing_write_positions.assign(n, 0); + + // Second pass: fill incoming and outgoing rows. Incoming rows are naturally grouped by + // source because we scan source states in order; outgoing rows are sorted by target below. + for (int source = 0; source < n; ++source) { + const auto& edges = result.GetFsm().GetEdges(source); + auto outgoing_row = outgoing_edges.MutableRowAt(source); + for (const auto& edge : edges) { + incoming_edges.MutableRowAt(edge.target + )[incoming_write_positions[edge.target]++] = {source, edge.min, edge.max}; + outgoing_row[outgoing_write_positions[source]++] = {edge.target, edge.min, edge.max}; + } + std::sort(outgoing_row.begin(), outgoing_row.end()); + } + + // Identify states with exactly one distinct predecessor/successor. These states are the + // candidates for the two local merge rules below. + incoming_distinct_count.assign(n, 0); + outgoing_distinct_count.assign(n, 0); + single_incoming_source.assign(n, -1); + single_outgoing_target.assign(n, -1); + for (int state = 0; state < n; ++state) { + auto incoming_row = incoming_edges[state]; + if (incoming_row.size() > 0) { + incoming_distinct_count[state] = 1; + single_incoming_source[state] = incoming_row[0].peer; + for (int32_t i = 1; i < incoming_row.size(); ++i) { + if (incoming_row[i].peer != incoming_row[i - 1].peer) { + ++incoming_distinct_count[state]; + single_incoming_source[state] = -1; + } + } + } + auto outgoing_row = outgoing_edges[state]; + if (outgoing_row.size() > 0) { + outgoing_distinct_count[state] = 1; + single_outgoing_target[state] = outgoing_row[0].peer; + for (int32_t i = 1; i < outgoing_row.size(); ++i) { + if (outgoing_row[i].peer != outgoing_row[i - 1].peer) { + ++outgoing_distinct_count[state]; + single_outgoing_target[state] = -1; + } + } + } + } + + // Case 1: Like ab | ac | ad, then they can be merged into a(b | c | d). + bool is_equiv_successor = false; + for (int i = 0; i < n; i++) { + if (incoming_distinct_count[i] != 1 || union_find_set.Count(i)) { + continue; + } + int previous_state = single_incoming_source[i]; + auto edges_to_i = incoming_edges[i]; + auto siblings = outgoing_edges[previous_state]; + int32_t group_begin = 0; + while (group_begin < siblings.size()) { + int sibling = siblings[group_begin].peer; + int32_t group_end = group_begin + 1; + while (group_end < siblings.size() && siblings[group_end].peer == sibling) { + ++group_end; + } + auto edges_to_sibling = siblings.Slice(group_begin, group_end); + group_begin = group_end; + if (sibling <= i || incoming_distinct_count[sibling] != 1 || + result.IsEndState(sibling) != result.IsEndState(i)) { + continue; + } + // Check if the edges from previous_state to i and sibling are the same. + if (edges_to_i.size() != edges_to_sibling.size()) { + continue; // Different edges, not equivalent. + } + bool is_equiv = true; + for (int32_t j = 0; j < edges_to_i.size(); ++j) { + if (edges_to_i[j].min != edges_to_sibling[j].min || + edges_to_i[j].max != edges_to_sibling[j].max) { + is_equiv = false; + break; // Different edge ranges, not equivalent. + } + } + // Merge the equivalent successor states. + if (is_equiv) { + union_find_set.Add(i); + union_find_set.Add(sibling); + union_find_set.Union(i, sibling); + is_equiv_successor = true; + } + } + } + + // Case 2: Like ba | ca | da, then they can be merged into (b | c | d)a. + bool is_equiv_precursor = false; + no_successor_end_states.clear(); + no_successor_non_end_states.clear(); + + for (int i = 0; i < n; i++) { + int outgoing_count = outgoing_distinct_count[i]; + if (outgoing_count == 0) { + if (result.IsEndState(i)) { + no_successor_end_states.push_back(i); + } else { + no_successor_non_end_states.push_back(i); + } + continue; // Skip states with no successors. + } + if (outgoing_count != 1 || union_find_set.Count(i)) { + continue; // Skip states with multiple successors. + } + int next_state = single_outgoing_target[i]; + auto node_edges = outgoing_edges[i]; + auto siblings = incoming_edges[next_state]; + int32_t group_begin = 0; + while (group_begin < siblings.size()) { + int sibling = siblings[group_begin].peer; + while (group_begin < siblings.size() && siblings[group_begin].peer == sibling) { + ++group_begin; + } + // Avoid chaining a Case 2 merge onto a state already merged earlier in this iteration + // (typically by Case 1), which can over-merge via transitive closure. + if (sibling <= i || union_find_set.Count(sibling) || + outgoing_distinct_count[sibling] != 1 || + result.IsEndState(i) != result.IsEndState(sibling)) { + continue; + } + auto sibling_node_edges = outgoing_edges[sibling]; + if (sibling_node_edges.size() != node_edges.size()) { + continue; // Different number of edges, not equivalent. + } + // Check if the sibling state has the same outgoing edges as i. + bool is_equiv = true; + for (int32_t j = 0; j < node_edges.size(); ++j) { + if (sibling_node_edges[j].min != node_edges[j].min || + sibling_node_edges[j].max != node_edges[j].max) { + is_equiv = false; + break; + } + } + // Merge the equivalent precursor states. + if (is_equiv) { + union_find_set.Add(i); + union_find_set.Add(sibling); + union_find_set.Union(i, sibling); + is_equiv_precursor = true; + } + } + } + + if (no_successor_end_states.size() > 1) { + // Merge all end states with no successors. + for (size_t i = 1; i < no_successor_end_states.size(); ++i) { + union_find_set.Add(no_successor_end_states[0]); + union_find_set.Add(no_successor_end_states[i]); + union_find_set.Union(no_successor_end_states[0], no_successor_end_states[i]); + is_equiv_precursor = true; + } + } + + if (no_successor_non_end_states.size() > 1) { + // Merge all non-end states with no successors. + for (size_t i = 1; i < no_successor_non_end_states.size(); ++i) { + union_find_set.Add(no_successor_non_end_states[0]); + union_find_set.Add(no_successor_non_end_states[i]); + union_find_set.Union(no_successor_non_end_states[0], no_successor_non_end_states[i]); + is_equiv_precursor = true; + } + } + + changed = is_equiv_successor || is_equiv_precursor; + if (changed) { + // Rebuild the FSM with the equivalent states merged, then repeat until no local merge + // rule applies. + auto eq_classes = union_find_set.GetAllSets(); + std::vector old_to_new(result.NumStates(), -1); + for (size_t i = 0; i < eq_classes.size(); i++) { + for (const auto& state : eq_classes[i]) { + old_to_new[state] = i; + } + } + int cnt = eq_classes.size(); + for (int i = 0; i < result.NumStates(); i++) { + if (old_to_new[i] == -1) { + old_to_new[i] = cnt; + cnt++; + } + } + result = result.RebuildWithMapping(old_to_new, cnt); + result.GetFsm()->SortEdges(); + } + } + return result; +} + +Result FSMWithStartEnd::MinimizeDFA(int max_num_states) const { + FSMWithStartEnd now_fsm(FSM(0), 0, std::vector(), true); + if (NumStates() > max_num_states) { + return ResultErr("The number of states exceeds the limit."); + } + // To perform the algorithm, we must make sure the FSM is + // a DFA. + if (!is_dfa_) { + Result dfa_raw = ToDFA(max_num_states); + if (dfa_raw.IsErr()) { + return dfa_raw; + } + now_fsm = std::move(dfa_raw).Unwrap(); + } else { + now_fsm = Copy(); + } + + // Initialize the precursors of nodes. + std::vector, int>>> precursors; + precursors.resize(now_fsm.NumStates()); + for (int i = 0; i < now_fsm.NumStates(); ++i) { + const auto& edges = now_fsm.GetFsm().GetEdges(i); + for (const auto& edge : edges) { + XGRAMMAR_DCHECK(!edge.IsEpsilon()); + precursors[edge.target].push_back(std::make_pair(std::make_pair(edge.min, edge.max), i)); + } + } + + // Initialize the partitions and working set. + std::vector> partitions; + std::vector> working_set; + std::unordered_set final_states; + std::unordered_set non_final_states; + for (int i = 0; i < now_fsm.NumStates(); ++i) { + if (now_fsm.IsEndState(i)) { + final_states.insert(i); + } else { + non_final_states.insert(i); + } + } + partitions.push_back(final_states); + partitions.push_back(non_final_states); + working_set.push_back(std::move(final_states)); + working_set.push_back(std::move(non_final_states)); + + while (!working_set.empty()) { + std::map, std::unordered_set> possible_transitions; + auto current_partition = std::move(working_set.back()); + working_set.pop_back(); + + // Get the possible transitions from the current partition. + for (const auto& state : current_partition) { + const auto& precursor_map = precursors[state]; + for (const auto& precursor : precursor_map) { + if (possible_transitions.find(precursor.first) == possible_transitions.end()) { + possible_transitions[precursor.first] = std::unordered_set(); + } + possible_transitions[precursor.first].insert(precursor.second); + } + } + + // Check each possible transition. + std::vector intersection; + std::vector difference; + for (const auto& [transition, precursors] : possible_transitions) { + for (size_t i = 0; i < partitions.size(); i++) { + const auto& partition = partitions[i]; + intersection.clear(); // partition \cap precursors + difference.clear(); // partition - precursors + for (const auto& partition_state : partition) { + if (precursors.find(partition_state) != precursors.end()) { + intersection.push_back(partition_state); + } else { + difference.push_back(partition_state); + } + } + + // the states in the partition is not equivalent. We need to + // update the working set and the partitions. + if ((!intersection.empty()) && (!difference.empty())) { + bool in_working_set = false; + for (size_t i = 0; i < working_set.size(); i++) { + if (partition == working_set[i]) { + in_working_set = true; + working_set[i].clear(); + for (const auto& state : intersection) { + working_set[i].insert(state); + } + working_set.emplace_back(); + for (const auto& state : difference) { + working_set.back().insert(state); + } + break; + } + } + if (!in_working_set) { + const auto& smaller_set = + difference.size() < intersection.size() ? difference : intersection; + working_set.emplace_back(); + for (const auto& state : smaller_set) { + working_set.back().insert(state); + } + } + partitions[i].clear(); + for (const auto& state : intersection) { + partitions[i].insert(state); + } + partitions.emplace_back(); + for (const auto& state : difference) { + partitions.back().insert(state); + } + } + } + } + } + std::vector state_mapping(now_fsm.NumStates(), -1); + for (size_t i = 0; i < partitions.size(); ++i) { + for (const auto& state : partitions[i]) { + state_mapping[state] = i; + } + } + int new_num_states = partitions.size(); + return ResultOk(now_fsm.RebuildWithMapping(state_mapping, new_num_states)); +} + +Result FSMWithStartEnd::ToDFA(int max_num_states) const { + if (NumStates() > max_num_states) { + return ResultErr("The number of states exceeds the limit."); + } + FSMWithStartEnd dfa(FSM(0), 0, std::vector(), true); + std::vector> closures; + std::unordered_set rules; + std::unordered_set repeat_aux_indices; + int now_process = 0; + std::unordered_set closure; + closure.insert(start_); + fsm_.GetEpsilonClosure(&closure); + closures.push_back(closure); + while (now_process < static_cast(closures.size())) { + rules.clear(); + repeat_aux_indices.clear(); + std::unordered_set token_aux_indices; + std::unordered_set exclude_token_aux_indices; + std::set interval_ends; + std::bitset<256> allowed_characters; + dfa.AddState(); + // Check if the closure is a final state. + for (const auto& state : closures[now_process]) { + if (IsEndState(state)) { + dfa.AddEndState(now_process); + } + const auto& edges = fsm_->GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsCharRange()) { + interval_ends.insert(edge.min); + interval_ends.insert(edge.max + 1); + for (int i = edge.min; i <= edge.max; ++i) { + allowed_characters.set(i); + } + continue; + } else if (edge.IsRuleRef()) { + rules.insert(edge.GetRefRuleId()); + } else if (edge.IsRepeatRef()) { + repeat_aux_indices.insert(edge.GetAuxIndex()); + } else if (edge.IsToken()) { + token_aux_indices.insert(edge.GetAuxIndex()); + } else if (edge.IsExcludeToken()) { + exclude_token_aux_indices.insert(edge.GetAuxIndex()); + } + } + } + // This part is to get the all possible intervals. + // Which can help reduce the transitions. + using Interval = std::pair; + std::vector intervals; + intervals.reserve(interval_ends.size()); + int last = -1; + for (const auto& end : interval_ends) { + if (last == -1) { + last = end; + continue; + } + bool allowed = true; + for (int i = last; i < end; ++i) { + if (!allowed_characters[i]) { + allowed = false; + break; + } + } + if (allowed) { + intervals.emplace_back(last, end - 1); + } + last = end; + } + for (const auto& interval : intervals) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_->GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsCharRange()) { + if (interval.first >= edge.min && interval.second <= edge.max) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, interval.first, interval.second); + flag = true; + break; + } + } + if (!flag) { + dfa.GetFsm().AddEdge(now_process, closures.size(), interval.first, interval.second); + closures.push_back(next_closure); + } + } + for (auto rule : rules) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsRuleRef()) { + if (rule == edge.GetRefRuleId()) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddRuleEdge(now_process, j, rule); + flag = true; + break; + } + } + if (!flag) { + dfa.GetFsm().AddRuleEdge(now_process, closures.size(), rule); + closures.push_back(next_closure); + } + } + + for (auto aux_idx : repeat_aux_indices) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsRepeatRef() && edge.GetAuxIndex() == aux_idx) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, FSMEdge::EdgeType::kRepeatRef, aux_idx); + flag = true; + break; + } + } + if (!flag) { + dfa.GetFsm().AddEdge(now_process, closures.size(), FSMEdge::EdgeType::kRepeatRef, aux_idx); + closures.push_back(next_closure); + } + } + + for (auto aux_idx : token_aux_indices) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsToken() && edge.GetAuxIndex() == aux_idx) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, FSMEdge::EdgeType::kToken, aux_idx); + flag = true; + break; + } + } + if (!flag) { + dfa.GetFsm().AddEdge(now_process, closures.size(), FSMEdge::EdgeType::kToken, aux_idx); + closures.push_back(next_closure); + } + } + + for (auto aux_idx : exclude_token_aux_indices) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsExcludeToken() && edge.GetAuxIndex() == aux_idx) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, FSMEdge::EdgeType::kExcludeToken, aux_idx); + flag = true; + break; + } + } + if (!flag) { + dfa.GetFsm().AddEdge( + now_process, closures.size(), FSMEdge::EdgeType::kExcludeToken, aux_idx + ); + closures.push_back(next_closure); + } + } + + now_process++; + } + dfa.GetFsm().SetEdgeAuxData(std::vector(fsm_.GetEdgeAuxData())); + dfa.is_dfa_ = true; + return ResultOk(dfa); +} + +/****************** CompactFSMWithStartEnd ******************/ + +std::string CompactFSMWithStartEnd::ToString() const { + std::string result; + result += "CompactFSM(num_states=" + std::to_string(NumStates()) + + ", start=" + std::to_string(start_) + ", end=["; + + std::unordered_set reachable_states; + GetReachableStates(&reachable_states); + std::vector reachable_states_vec(reachable_states.begin(), reachable_states.end()); + std::sort(reachable_states_vec.begin(), reachable_states_vec.end()); + bool first = true; + for (auto end : ends_) { + if (reachable_states.count(end)) { + if (!first) { + result += ", "; + } + first = false; + result += std::to_string(end); + } + } + + result += "], edges=" + fsm_.EdgesToString(reachable_states_vec) + ")"; + return result; +} + +std::ostream& operator<<(std::ostream& os, const CompactFSMWithStartEnd& fsm) { + os << fsm.ToString(); + return os; +} + +std::size_t MemorySize(const CompactFSM& self) { return MemorySize(*self.ImplPtr()); } + +std::size_t MemorySize(const CompactFSMWithStartEnd& self) { + // The underlying CompactFSM is not counted here: CompactFSMWithStartEnd is a view, and many + // views usually share one CompactFSM (e.g. the per-rule FSMs of a grammar all point to the + // grammar's complete FSM, which is counted by the grammar itself). Counting it per view would + // multiply the shared FSM's size by the number of views. + return MemorySize(self.ends_); +} + +FSMWithStartEnd CompactFSMWithStartEnd::ToFSM() const { + return FSMWithStartEnd(fsm_.ToFSM(), start_, ends_); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/fsm.h b/third_party/xgrammar/cpp/fsm.h new file mode 100644 index 0000000000..7647962938 --- /dev/null +++ b/third_party/xgrammar/cpp/fsm.h @@ -0,0 +1,1051 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm.h + * \note For functions accepting a pointer to a container as result, the container will be cleared + * before the result is stored. + */ +#ifndef XGRAMMAR_FSM_H_ +#define XGRAMMAR_FSM_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/compact_2d_array.h" +#include "support/logging.h" +#include "support/reflection.h" +#include "support/utils.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/*! + * \brief The edge of a FSM. + */ +struct FSMEdge { + /*! + * \brief Edge type is encoded in the `min` field. When min >= 0, the edge is a character range + * [min, max]. When min < 0, it is a special edge type identified by the enum values below. + * + * For each type, `max` has a type-specific meaning (see comments on each enumerator). + */ + enum EdgeType : int32_t { + //! Character range [min, max]. min >= 0. + kCharRange = 0, + //! Epsilon transition. max is unused. + kEpsilon = -1, + //! Rule reference. max = rule_id. + kRuleRef = -2, + //! Accepts the EOS token. max is unused. + kEOS = -3, + //! Repeated rule reference. max = aux index into edge_aux_data + //! (layout: [rule_id, lower, upper]). + //! Invariant: a state with a kRepeatRef edge has exactly one outgoing edge. + kRepeatRef = -4, + //! Accepts a set of token IDs. max = aux index into edge_aux_data + //! (layout: [count, token_id_0, token_id_1, ...]). + kToken = -5, + //! Accepts any token NOT in the given set. max = aux index into edge_aux_data + //! (layout: [count, token_id_0, token_id_1, ...]). + kExcludeToken = -6, + }; + + inline static constexpr int kMaxChar = 255; + + int32_t min, max; + + /*! + * \brief The target state id of the edge. + */ + int32_t target; + + // for serialization only + FSMEdge() = default; + + FSMEdge(int32_t min, int32_t max, int32_t target) : min(min), max(max), target(target) { + XGRAMMAR_DCHECK(!IsCharRange() || min <= max) + << "Invalid FSMEdge: min > max. min=" << min << ", max=" << max; + } + + /*! + * \brief Compare the edges. Used to sort the edges in the FSM. + */ + // TODO(yixin): consider combining the fields to a single int64_t for better efficiency + friend bool operator==(const FSMEdge& lhs, const FSMEdge& rhs) { + return std::make_tuple(lhs.min, lhs.max, lhs.target) == + std::make_tuple(rhs.min, rhs.max, rhs.target); + } + + /*! + * \brief Compare the edges. Used to sort the edges in the FSM. + */ + friend bool operator<(const FSMEdge& lhs, const FSMEdge& rhs) { + return std::make_tuple(lhs.min, lhs.max, lhs.target) < + std::make_tuple(rhs.min, rhs.max, rhs.target); + } + + /*! + * \brief Check if the edge is a character range. + */ + bool IsCharRange() const { return min >= 0; } + + /*! + * \brief Check if the edge is an epsilon transition. + */ + bool IsEpsilon() const { return min == EdgeType::kEpsilon; } + + /*! + * \brief Check if the edge is a rule reference. + */ + bool IsRuleRef() const { return min == EdgeType::kRuleRef; } + + /*! + * \brief Check if the edge is an EOS transition. + */ + bool IsEOS() const { return min == EdgeType::kEOS; } + + /*! + * \brief Check if the edge is a repeat reference. + */ + bool IsRepeatRef() const { return min == EdgeType::kRepeatRef; } + + bool IsToken() const { return min == EdgeType::kToken; } + + bool IsExcludeToken() const { return min == EdgeType::kExcludeToken; } + + /*! + * \brief Get the rule id of the edge. + * \return The rule id of the edge. -1 if the edge is not a rule reference. + */ + int32_t GetRefRuleId() const { return IsRuleRef() ? max : -1; } + + /*! + * \brief Get the auxiliary data index for repeat reference edges. + * \return The index into the owning FSM's edge_aux_data. -1 if not a repeat reference. + */ + int32_t GetAuxIndex() const { + return (IsRepeatRef() || IsToken() || IsExcludeToken()) ? max : -1; + } + + /*! \brief Check if the edge uses auxiliary data. */ + bool IsAuxEdge() const { return IsRepeatRef() || IsToken() || IsExcludeToken(); } + + friend struct member_trait; +}; + +/*! \brief View into edge_aux_data for a repeat edge (layout: [rule_id, lower, upper]). */ +struct RepeatEdgeRef { + const int32_t* data; + int32_t RuleId() const { return data[0]; } + int32_t Lower() const { return data[1]; } + int32_t Upper() const { return data[2]; } +}; + +/*! \brief View into edge_aux_data for a token edge (layout: [count, token_id_0, ...]). */ +struct TokenEdgeRef { + const int32_t* data; + int32_t Count() const { return data[0]; } + const int32_t* TokenIds() const { return data + 1; } + bool Contains(int32_t token_id) const { + for (int32_t i = 0; i < Count(); ++i) { + if (TokenIds()[i] == token_id) return true; + } + return false; + } +}; + +/*! \brief View into edge_aux_data for an exclude-token edge (layout: [count, token_id_0, ...]). */ +struct ExcludeTokenEdgeRef { + const int32_t* data; + int32_t Count() const { return data[0]; } + const int32_t* TokenIds() const { return data + 1; } + bool Contains(int32_t token_id) const { + for (int32_t i = 0; i < Count(); ++i) { + if (TokenIds()[i] == token_id) return true; + } + return false; + } + bool Accepts(int32_t token_id) const { return !Contains(token_id); } +}; + +/*! + * \brief Comparator for FSMEdge. Only compare the min and max. + */ +struct FSMEdgeRangeComparator { + bool operator()(const FSMEdge& lhs, const FSMEdge& rhs) const { + return std::make_tuple(lhs.min, lhs.max) < std::make_tuple(rhs.min, rhs.max); + } +}; + +XGRAMMAR_MEMBER_ARRAY(FSMEdge, &FSMEdge::min, &FSMEdge::max, &FSMEdge::target); + +} // namespace xgrammar + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::FSMEdge, &xgrammar::FSMEdge::min, &xgrammar::FSMEdge::max, &xgrammar::FSMEdge::target +); + +namespace xgrammar { + +class CompactFSM; + +/*! + * \brief FSM is a class that represents a finite state machine, could be a DFA or an NFA. + * \details It's mutable, which means you can add edges and states to it. + */ +class FSM { + public: + /*! + * \brief Construct an FSM with a given number of states. + * \param num_states The number of states in the FSM. + */ + FSM(int num_states = 0); + + /*! + * \brief Construct an FSM with a given set of edges. + */ + FSM(const std::vector>& edges, std::vector edge_aux_data = {}); + + /*! + * \brief Construct an FSM with a given set of edges. + */ + FSM(std::vector>&& edges, std::vector edge_aux_data = {}); + + /****************** FSM Visitors ******************/ + + /*! + * \brief Get the number of states in the FSM. + * \return The number of states in the FSM. + */ + int NumStates() const; + + /*! + * \brief Get the edges of the FSM. + * \return The edges of the FSM. + */ + const std::vector>& GetEdges() const; + + /*! + * \brief Get the edges of the FSM. + * \return The edges of the FSM. + */ + std::vector>& GetEdges(); + + /*! + * \brief Get the edges of the FSM. + * \param state The state to get the edges from. + * \return The edges of the FSM. + */ + std::vector& GetEdges(int state); + + /*! + * \brief Get the edges of the FSM. + * \param state The state to get the edges from. + * \return The edges of the FSM. + */ + const std::vector& GetEdges(int state) const; + + /*! + * \brief Convert the edges of the FSM to a string. Used in printing the FSM. + * \return The string representation of the edges of the FSM. + */ + std::string EdgesToString(std::optional> states = std::nullopt) const; + + /****************** FSM Traversal Visitors ******************/ + + inline static constexpr int kNoNextState = -1; + + /*! + * \brief Advance the FSM from a given state based on an input character. If there are multiple + * transitions, the first one will be returned. + * \param from The source state to transition from. + * \param character The input character. + * \return The target state if a valid transition exists, kNoNextState otherwise. + */ + int GetNextState(int from, int value, FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange) + const; + + /*! + * \brief Advance the FSM to the next state. + * \param from The current states. + * \param value The input value. + * \param result The possible next states. The result is cleared at the beginning. + * \param value_is_rule Whether the input value is a rule id. + * \param from_is_closure Whether from is an epsilon closure. + */ + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange, + bool from_is_closure = false + ) const; + + /*! + * \brief Get all the possible rule numbers for a given state. + * \param state_num The state number. + * \param rules The set of possible rule numbers. The result is cleared at the beginning. + */ + void GetPossibleRules(int state_num, std::unordered_set* rules) const; + + /*! + * \brief Get the epsilon closure of a set of states, i.e. those can be reached by epsilon + * transitions. + * \param state_set The states in the epsilon closure. The result is not cleared. + */ + void GetEpsilonClosure(std::unordered_set* state_set) const; + + /*! + * \brief Get the reachable states from a set of states. + * \param from The current states. + * \param result The reachable states. The result is cleared at the beginning. + */ + void GetReachableStates(const std::vector& from, std::unordered_set* result) const; + + /****************** FSM Mutators ******************/ + + /*! + * \brief Adds a new state to the FSM. + * \return The index of the newly added state. + */ + int AddState(); + + /*! + * \brief Adds a transition edge between states with given min and max values. For character + * transitions, it accepts any character in range [min, max]. + * \param from The source state. + * \param to The target state. + * \param min The min value of the range. + * \param max The max value of the range. + */ + void AddEdge(int from, int to, int32_t min, int32_t max); + + /*! \brief Add a raw edge with explicit type and value. */ + void AddEdge(int from, int to, FSMEdge::EdgeType type, int32_t value); + + /*! + * \brief Add an epsilon transition between two states. + * \param from The source state. + * \param to The target state. + */ + void AddEpsilonEdge(int from, int to); + + /*! + * \brief Add a rule reference edge between states. + * \param from The source state. + * \param to The target state. + * \param rule_id The rule id to reference. + */ + void AddRuleEdge(int from, int to, int32_t rule_id); + + /*! + * \brief Add an EOS transition between two states. + * \param from The source state. + * \param to The target state. + */ + void AddEOSEdge(int from, int to); + + /*! + * \brief Add a repeat reference edge between states, allocating auxiliary data. + * \param from The source state. + * \param to The target state. + * \param rule_id The rule to repeat. + * \param lower Minimum repeat count. + * \param upper Maximum repeat count (-1 for unlimited). + */ + void AddRepeatEdge(int from, int to, int32_t rule_id, int32_t lower, int32_t upper); + + void AddTokenEdge(int from, int to, const std::vector& token_ids); + + void AddExcludeTokenEdge(int from, int to, const std::vector& token_ids); + + /*! \brief Get the edge auxiliary data. */ + const std::vector& GetEdgeAuxData() const; + + /*! \brief Set the edge auxiliary data (used during FSM construction). */ + void SetEdgeAuxData(std::vector data); + + /*! \brief Get repeat edge info by aux index. */ + RepeatEdgeRef GetRepeatEdgeInfo(int32_t idx) const; + + /*! \brief Get token edge info by aux index. */ + TokenEdgeRef GetTokenEdgeInfo(int32_t idx) const; + + /*! \brief Get exclude-token edge info by aux index. */ + ExcludeTokenEdgeRef GetExcludeTokenEdgeInfo(int32_t idx) const; + + /*! + * \brief Add a whole FSM to the current FSM. + * \param fsm The FSM to be added. + * \param state_mapping The mapping from the state ids of the added FSM to the new ids in the + * current FSM. The result is cleared at the beginning. If the fsm's state id starts from 0, use + * it for efficiency. + */ + void AddFSM(const FSM& fsm, std::vector* state_mapping = nullptr); + + /****************** FSM Construction Methods ******************/ + + /*! + \brief Return a copy of the FSM. + */ + FSM Copy() const; + + /*! + * \brief Rebuild the FSM with the new state ids. + * \param state_mapping The mapping from the old state ids to the new state ids. + * \param new_num_states The new number of states. + * \return The rebuilt FSM. + */ + FSM RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const; + + /*! + * \brief Sort the edges of the FSM by their min, max and target. + */ + void SortEdges(); + + /*! + * \brief Transform a FSM to a compact FSM. This method will first sort the edges of the FSM, + * then put all the edges into a compact array. + * \return The compact FSM. + */ + CompactFSM ToCompact(); + + XGRAMMAR_DEFINE_PIMPL_METHODS(FSM); +}; + +/*! + * \brief CompactFSM is the compact from of FSM. + * \details It uses Compact2DArray to store the edges, ensuring memory contiguity. It sorts all + * outgoing edges from a node according to their min and max values, so traversal can be faster. + * + * CompactFSM is immutable. If you need to modify a CompactFSM, you need to convert it to a FSM + * first, and convert it back after modification. + * + * It share the same set of visitor methods with FSM. + */ + +class CompactFSM { + public: + // for serialization only + CompactFSM() = default; + + explicit CompactFSM( + const Compact2DArray& edges, std::vector edge_aux_data = {} + ); + + explicit CompactFSM(Compact2DArray&& edges, std::vector edge_aux_data = {}); + + /****************** CompactFSM Visitors ******************/ + + /*! + * \brief Get the number of states in the FSM. + * \return The number of states in the FSM. + */ + int NumStates() const; + + /*! + * \brief Get the edges of the CompactFSM. + * \return The edges of the CompactFSM. + */ + const Compact2DArray& GetEdges() const; + + /*! + * \brief Get the edges of the CompactFSM. + * \param state The state to get the edges from. + * \return The edges of the CompactFSM. + */ + Compact2DArray::Row GetEdges(int state) const; + + /*! + * \brief Convert the edges of the CompactFSM to a string. Used in printing the CompactFSM. + * \return The string representation of the edges of the CompactFSM. + */ + std::string EdgesToString(std::optional> states = std::nullopt) const; + + /*! + * \brief Get the memory size of the CompactFSM. + * \param self The CompactFSM. + * \return The memory size of the CompactFSM. + */ + friend std::size_t MemorySize(const CompactFSM& self); + + /****************** CompactFSM Traversal Visitors ******************/ + + inline static constexpr int kNoNextState = -1; + + /*! + * \brief Advance the FSM from a given state based on an input character. If there are multiple + * transitions, the first one will be returned. + * \param from The source state to transition from. + * \param character The input character. + * \param targets The target states to be filled with the possible next states. + * \return The target state if a valid transition exists, kNoNextState otherwise. + */ + void GetNextStates( + int from, + int value, + FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange, + std::vector* targets = nullptr + ) const; + + /*! + * \brief Advance the FSM to the next state. + * \param from The current states. + * \param value The input value. + * \param result The possible next states. The result is cleared at the beginning. + * \param value_is_rule Whether the input value is a rule id. + * \param from_is_closure Whether from is an epsilon closure. + */ + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange, + bool from_is_closure = false + ) const; + + /*! + * \brief Get all the possible rule numbers for a given state. + * \param state_num The state number. + * \param rules The set of possible rule numbers. The result is cleared at the beginning. + */ + void GetPossibleRules(int state_num, std::unordered_set* rules) const; + + /*! + * \brief Get the epsilon closure of a set of states, i.e. those can be reached by epsilon + * transitions. + * \param state_set The states in the epsilon closure. The result is not cleared. + */ + void GetEpsilonClosure(std::unordered_set* state_set) const; + + /*! + * \brief Get the reachable states from a set of states. + * \param from The current states. + * \param result The reachable states. The result is cleared at the beginning. + */ + void GetReachableStates(const std::vector& from, std::unordered_set* result) const; + + /*! + * \brief Get the number of edges in the compact FSM. + * \return The number of edges. + */ + size_t GetNumEdges() const; + + /****************** CompactFSM Auxiliary Data ******************/ + + /*! \brief Get the edge auxiliary data. */ + const std::vector& GetEdgeAuxData() const; + + /*! \brief Set the edge auxiliary data (used during FSM construction). */ + void SetEdgeAuxData(std::vector data); + + /*! \brief Get repeat edge info by aux index. */ + RepeatEdgeRef GetRepeatEdgeInfo(int32_t idx) const; + + /*! \brief Get token edge info by aux index. */ + TokenEdgeRef GetTokenEdgeInfo(int32_t idx) const; + + /*! \brief Get exclude-token edge info by aux index. */ + ExcludeTokenEdgeRef GetExcludeTokenEdgeInfo(int32_t idx) const; + + /****************** CompactFSM Construction Methods ******************/ + + /*! + * \brief Transform the compact FSM to a FSM. + * \return The FSM. + */ + FSM ToFSM() const; + + friend picojson::value SerializeJSONValue(const CompactFSM& value); + friend std::optional DeserializeJSONValue( + CompactFSM* result, const picojson::value& value, const std::string& type_name + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(CompactFSM); +}; + +std::optional DeserializeJSONValue( + CompactFSM* result, const picojson::value& value, const std::string& type_name = "" +); + +class FSMWithStartEnd; +class FSMWithStartEndWithSize; +class CompactFSMWithStartEnd; +class CompactFSMWithStartEndWithSize; +struct CompactFSMWithStartEndWithSizeSerializeHelper; + +/*! + * \brief The base class for FSMWithStartEnd and CompactFSMWithStartEnd. It defines the + * common constructor and visitor methods. + */ +template +class FSMWithStartEndBase { + static_assert( + std::is_same_v || std::is_same_v, + "FSMType must be FSM or CompactFSM" + ); + + public: + // For serialization only + FSMWithStartEndBase() = default; + + FSMWithStartEndBase(const FSMType& fsm, int start, std::vector ends, bool is_dfa = false) + : fsm_(fsm), start_(start), ends_(std::move(ends)), is_dfa_(is_dfa) { + NormalizeEnds(); + } + + /****************** Member Accessors and Mutators ******************/ + + /*! \brief Returns the underlying FSM. */ + const FSMType& GetFsm() const { return fsm_; } + + /*! \brief Returns the start state of the FSM. */ + int GetStart() const { return start_; } + + /*! + * \brief Returns the end states of the FSM as a sorted, deduplicated list of state ids. + * \note End states are stored sparsely because an FSM view over a large shared FSM (e.g. the + * per-rule FSMs over the complete FSM of a grammar) usually has only a few end states, while + * the state id space can be very large. + */ + const std::vector& GetEnds() const { return ends_; } + + /*! + * \brief Checks if a given state is an end/accepting state. + * \param state The state to check. + * \return True if the state is an end state, false otherwise. + */ + bool IsEndState(int state) const { return std::binary_search(ends_.begin(), ends_.end(), state); } + + /*! \brief Check if a state is scanable. + * \param state The state to check. + * \return True if the state is scanable, false otherwise. + */ + bool IsScanableState(int state) const { + for (const auto& edge : fsm_.GetEdges(state)) { + if (edge.IsCharRange() || edge.IsToken() || edge.IsExcludeToken()) { + return true; + } + } + return false; + } + + /*! + * \brief Check if a state is not terminal. + * \param state The state to check. + * \return True if the state is scanable, false otherwise. + */ + bool IsNonTerminalState(int state) const { + for (const auto& edge : fsm_.GetEdges(state)) { + if (edge.IsRuleRef() || edge.IsEpsilon() || edge.IsRepeatRef()) { + return true; + } + } + return false; + } + + /*! + * \brief Sets the start state of the FSM. + * \param state The state to set as the start state. + */ + void SetStartState(int state) { + XGRAMMAR_DCHECK(state < NumStates()); + start_ = state; + } + + /*! + * \brief Adds an end/accepting state to the FSM. + * \param state The state to add as an end state. + */ + void AddEndState(int state) { + XGRAMMAR_DCHECK(state < NumStates()); + auto it = std::lower_bound(ends_.begin(), ends_.end(), state); + if (it == ends_.end() || *it != state) { + ends_.insert(it, state); + } + } + + /*! + * \brief Adds a new state to the FSM and marks it as non-end. + * \return The index of the newly added state. + */ + int AddState() { return fsm_.AddState(); } + + /*! + * \brief Sets the end states of the FSM. + * \param ends The new end states, as a list of state ids. Need not be sorted or deduplicated. + */ + void SetEndStates(std::vector ends) { + ends_ = std::move(ends); + NormalizeEnds(); + } + + /*! \brief Returns the total number of states in the FSM. */ + int NumStates() const { return fsm_.NumStates(); } + + /*! + * \brief Access the methods of the underlying FSM. + */ + FSMType& GetFsm() { return fsm_; } + + /****************** FSM Traversal Algorithms ******************/ + + /*! + * \brief Check if the FSM accepts the string. + * \param str The input string. + * \return True if the FSM accepts the string, false otherwise. + */ + bool AcceptString(const std::string& str) const; + + /*! + * \brief Get the reachable states from the start state. + * \param result The reachable states. The result is cleared at the beginning. + */ + void GetReachableStates(std::unordered_set* result) const; + + /*! + * \brief Check if the FSM is a leaf FSM. + * \return True if the FSM is a leaf FSM, false otherwise. + */ + bool IsLeaf() const; + + protected: + /*! \brief Sort and deduplicate the end states to maintain the sorted invariant. */ + void NormalizeEnds() { + std::sort(ends_.begin(), ends_.end()); + ends_.erase(std::unique(ends_.begin(), ends_.end()), ends_.end()); + } + + /*! \brief The underlying finite state machine. */ + FSMType fsm_; + /*! \brief The start state of the FSM. */ + int start_; + + /*! + * \brief The set of accepting/end states, stored as a sorted, deduplicated list of state ids. + * \note Stored sparsely (rather than as a bitvector over all states) because FSM views over a + * large shared FSM usually have only a few end states, while the state id space can be huge. + */ + std::vector ends_; + + protected: + /*! \brief Whether this FSM is a deterministic finite automaton. */ + bool is_dfa_ = false; +}; + +/*! + * \brief FSMWithStartEnd represents a FSM with start and end states. + * \details It stores a pointer to a FSM, a start state, and a set of end states. Multiple + * FSMWithStartEnd can share the same FSM. It also provides a set of methods to construct FSMs. + */ +class FSMWithStartEnd : public FSMWithStartEndBase { + public: + using FSMWithStartEndBase::FSMWithStartEndBase; + + /*! + * \brief Convert the FSMWithStartEnd to a string. Only considers the nodes approachable from the + * start state. + * \return The string representation of the FSMWithStartEnd. + */ + std::string ToString() const; + + friend std::ostream& operator<<(std::ostream& os, const FSMWithStartEnd& fsm); + + /****************** FSM Construction Methods ******************/ + + /*! + * \brief Return a copy of the FSMWithStartEnd. + */ + FSMWithStartEnd Copy() const; + + /*! + * \brief Rebuild the FSM with the new state ids. + * \param state_mapping The mapping from old state ids to new state ids. + * \param new_num_states The new number of states. + */ + FSMWithStartEnd RebuildWithMapping(const std::vector& state_mapping, int new_num_states) + const; + + /*! + * \brief Add the underlying FSM to another complete FSM that could contain multiple FSMs. + * Return a new FSMWithStartEnd that points to the complete FSM and whose start and ends are + * mapped to the states in the complete FSM. + * \param complete_fsm The complete FSM. + * \param state_mapping The mapping from the old state ids to the new state ids. The result is + * cleared at the beginning. Should not be nullptr. + * \return The FSMWithStartEnd that points to the complete FSM. + */ + FSMWithStartEndWithSize AddToCompleteFSM(FSM* complete_fsm, std::vector* state_mapping); + + /*! + * \brief Transform the FSMWithStartEnd to a CompactFSMWithStartEnd. + * \return The CompactFSMWithStartEnd. + */ + CompactFSMWithStartEnd ToCompact(); + + /****************** FSM Algorithms ******************/ + + /*! + * \brief Return a new FSM representing FSM* + * \return The FSM that accepts FSM*. + */ + FSMWithStartEnd Star() const; + + /*! + * \brief Return a new FSM representing rule1+. + * \return The FSM that accepts rule1+. + */ + FSMWithStartEnd Plus() const; + + /*! + * \brief Return a new FSM representing rule1?. + * \return The FSM that accepts rule1?. + */ + FSMWithStartEnd Optional() const; + + /*! + * \brief Return a new FSM representing the complement of the language. + * \return The complement FSM. + */ + Result Not(int max_result_num_states = 1e6) const; + + /*! + * \brief Intersect the FSMs. + * \param lhs The left FSM. + * \param rhs The right FSM. + * \return The intersection of the FSMs. + */ + static Result Intersect( + const FSMWithStartEnd& lhs, const FSMWithStartEnd& rhs, int max_result_num_states = 1e6 + ); + + /*! + * \brief Union the FSMs. + * \param fsms The FSMs to be unioned. + * \return The union of the FSMs. + */ + static FSMWithStartEnd Union(const std::vector& fsms); + + /*! + * \brief Concatenate the FSMs. + * \param fsms The FSMs to be concatenated, which should be in order. + * \return The concatenation of the FSMs. + */ + static FSMWithStartEnd Concat(const std::vector& fsms); + + /*! + * \brief Check if the FSM is a DFA. + * \return True if the FSM is a DFA, false otherwise. + */ + bool IsDFA(); + + /*! + * \brief Merge some states by removing some epsilon transitions. + * \details If a --\epsilon--> b, and either 1) b doesn't have any other inward edges, or + * 2) a doesn't have any other outward edges, we can merge a and b. + */ + FSMWithStartEnd SimplifyEpsilon(int max_num_states = 1e8) const; + + /*! + * \brief Merge equivalent states in the FSM. + * \details If two states are 1) pointed to by edges with the same label from the same state, and + * 2) they are not pointed to by other edges, then we can merge them. + * \example n0 --(c)--> n1, n0 --(c)--> n2, then we can merge n1 and n2. + */ + FSMWithStartEnd MergeEquivalentStates(int max_num_states = 1e5) const; + + /*! + * \brief Transform the FSM to a DFA. + * \param max_result_num_states The maximum number of states in the DFA. + * \return The DFA. + */ + Result ToDFA(int max_num_states = 1e3) const; + + /*! + * \brief Minimize the DFA. + * \param max_result_num_states The maximum number of states in the DFA. + * \return The minimized DFA. + */ + Result MinimizeDFA(int max_num_states = 1e3) const; +}; + +/*! + * \brief Wrapper that bundles an FSMWithStartEnd with explicit size metadata. It is + * used when we want to store the number of edges and nodes in the part of the FSM, instead + * of the completed FSMWithStartEnd. + */ +class FSMWithStartEndWithSize { + public: + // For serialization only + FSMWithStartEndWithSize() = default; + + explicit FSMWithStartEndWithSize(FSMWithStartEnd fsm, int edge_num, int node_num) + : fsm_(std::move(fsm)), edge_num_(edge_num), node_num_(node_num) {} + + const FSMWithStartEnd& GetFsm() const { return fsm_; } + int GetEdgeNum() const { return edge_num_; } + int GetNodeNum() const { return node_num_; } + + private: + FSMWithStartEnd fsm_; + int edge_num_ = 0; + int node_num_ = 0; +}; + +/*! + * \brief A class that represents a compact-form FSM with a start state and a set of end states. + * \details CompactFSMWithStartEnd stores a pointer to a CompactFSM, a start state, and a set of end + * states. Multiple CompactFSMWithStartEnd can share the same CompactFSM. It share the same set of + * visitor methods with FSMWithStartEnd. + */ +class CompactFSMWithStartEnd : public FSMWithStartEndBase { + public: + // For serialization only + CompactFSMWithStartEnd() = default; + + explicit CompactFSMWithStartEnd(const CompactFSM& fsm, int start, std::vector ends) + : FSMWithStartEndBase(fsm, start, std::move(ends)), + edge_num_(fsm.GetNumEdges()) {} + + using FSMWithStartEndBase::FSMWithStartEndBase; + + /*! + * \brief Convert the FSMWithStartEnd to a string. Only considers the nodes approachable from the + * start state. + * \return The string representation of the FSMWithStartEnd. + */ + std::string ToString() const; + + /*! + * \brief Transform the CompactFSMWithStartEnd to a FSMWithStartEnd. + * \return The FSMWithStartEnd. + */ + FSMWithStartEnd ToFSM() const; + + private: + size_t edge_num_ = 0; + + /*! + * \brief Print the CompactFSMWithStartEnd. + * \param os The output stream. + * \param fsm The CompactFSMWithStartEnd. + * \return The output stream. + */ + friend std::ostream& operator<<(std::ostream& os, const CompactFSMWithStartEnd& fsm); + + /*! + * \brief Get the memory size of the CompactFSMWithStartEnd. + * \param self The CompactFSMWithStartEnd. + * \return The memory size of the CompactFSMWithStartEnd. + */ + friend std::size_t MemorySize(const CompactFSMWithStartEnd& self); + + friend struct member_trait; + + friend struct CompactFSMWithStartEndSerializeHelper; + + friend picojson::value SerializeJSONValue(const CompactFSMWithStartEnd& value); + friend std::optional DeserializeJSONValue( + CompactFSMWithStartEnd* result, const picojson::value& value, const std::string& type_name + ); +}; + +/*! + * \brief Wrapper that bundles a CompactFSMWithStartEnd with explicit size metadata. It is + * used when we want to store the number of edges and nodes in the part of the + * CompactFSMWithStartEnd, instead of the completed CompactFSMWithStartEnd. + */ +class CompactFSMWithStartEndWithSize { + public: + // For serialization only + CompactFSMWithStartEndWithSize() = default; + + explicit CompactFSMWithStartEndWithSize(CompactFSMWithStartEnd fsm, int edge_num, int node_num) + : fsm_(std::move(fsm)), edge_num_(edge_num), node_num_(node_num) {} + + const CompactFSMWithStartEnd& GetFsm() const { return fsm_; } + int GetEdgeNum() const { return edge_num_; } + int GetNodeNum() const { return node_num_; } + + friend picojson::value SerializeJSONValue(const CompactFSMWithStartEndWithSize& value); + friend std::optional DeserializeJSONValue( + CompactFSMWithStartEndWithSize* result, + const picojson::value& value, + const std::string& type_name + ); + + private: + CompactFSMWithStartEnd fsm_; + int edge_num_ = 0; + int node_num_ = 0; + + friend std::size_t MemorySize(const CompactFSMWithStartEndWithSize& self) { + return MemorySize(self.fsm_) + sizeof(self.edge_num_) + sizeof(self.node_num_); + } + + friend struct CompactFSMWithStartEndWithSizeSerializeHelper; +}; + +std::optional DeserializeJSONValue( + CompactFSMWithStartEndWithSize* result, + const picojson::value& value, + const std::string& type_name = "" +); + +/****************** FSMWithStartEndBase Template Implementation ******************/ + +template +inline bool FSMWithStartEndBase::AcceptString(const std::string& str) const { + std::unordered_set start_states{start_}; + fsm_.GetEpsilonClosure(&start_states); + std::unordered_set result_states; + for (const auto& character : str) { + result_states.clear(); + fsm_.Advance( + start_states, + static_cast(static_cast(character)), + &result_states, + FSMEdge::EdgeType::kCharRange, + false + ); + if (result_states.empty()) { + return false; + } + start_states = result_states; + } + return std::any_of(start_states.begin(), start_states.end(), [&](int state) { + return IsEndState(state); + }); +} + +template +inline void FSMWithStartEndBase::GetReachableStates(std::unordered_set* result +) const { + return fsm_.GetReachableStates({start_}, result); +} + +template +inline bool FSMWithStartEndBase::IsLeaf() const { + std::unordered_set reachable_states; + GetReachableStates(&reachable_states); + for (const auto& state : reachable_states) { + for (const auto& edge : fsm_.GetEdges(state)) { + if (edge.IsRuleRef() || edge.IsRepeatRef()) { + return false; + } + } + } + return true; +} + +} // namespace xgrammar + +#endif // XGRAMMAR_FSM_H_ diff --git a/third_party/xgrammar/cpp/fsm_builder.cc b/third_party/xgrammar/cpp/fsm_builder.cc new file mode 100644 index 0000000000..685373126c --- /dev/null +++ b/third_party/xgrammar/cpp/fsm_builder.cc @@ -0,0 +1,1797 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm_builder.cc + */ +#include "fsm_builder.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm.h" +#include "grammar_builder.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "support/utils.h" + +namespace xgrammar { + +/******************** Packed UTF-8 range helpers ********************/ + +uint32_t CodepointToPackedUTF8(uint32_t codepoint) { + if (codepoint <= 0x7F) { + // 1-byte sequence (ASCII) + return codepoint; + } else if (codepoint <= 0x7FF) { + // 2-byte sequence: byte0 = 110xxxxx, byte1 = 10xxxxxx + uint8_t byte0 = 0xC0 | ((codepoint >> 6) & 0x1F); + uint8_t byte1 = 0x80 | (codepoint & 0x3F); + return (static_cast(byte0) << 8) | byte1; + } else if (codepoint <= 0xFFFF) { + // 3-byte sequence: byte0 = 1110xxxx, byte1 = 10xxxxxx, byte2 = 10xxxxxx + uint8_t byte0 = 0xE0 | ((codepoint >> 12) & 0x0F); + uint8_t byte1 = 0x80 | ((codepoint >> 6) & 0x3F); + uint8_t byte2 = 0x80 | (codepoint & 0x3F); + return (static_cast(byte0) << 16) | (static_cast(byte1) << 8) | byte2; + } else { + // 4-byte sequence: byte0 = 11110xxx, byte1-3 = 10xxxxxx + uint8_t byte0 = 0xF0 | ((codepoint >> 18) & 0x07); + uint8_t byte1 = 0x80 | ((codepoint >> 12) & 0x3F); + uint8_t byte2 = 0x80 | ((codepoint >> 6) & 0x3F); + uint8_t byte3 = 0x80 | (codepoint & 0x3F); + return (static_cast(byte0) << 24) | (static_cast(byte1) << 16) | + (static_cast(byte2) << 8) | byte3; + } +} + +// This function will add a range [min, max] of characters to the FSM, and the length +// of the characters are the same. +static void AddSameLengthCharacterRange(FSM& fsm, int from, int to, uint32_t min, uint32_t max) { + uint8_t byte_min[4] = { + static_cast(min & 0xFF), + static_cast(min >> 8), + static_cast(min >> 16), + static_cast(min >> 24) + }; + uint8_t byte_max[4] = { + static_cast(max & 0xFF), + static_cast(max >> 8), + static_cast(max >> 16), + static_cast(max >> 24) + }; + + // ASCII. + if (byte_max[1] == 0) { + fsm.AddEdge(from, to, byte_min[0], byte_max[0]); + return; + } + + if (byte_max[3] != 0) { + // 4-byte unicode. + if (byte_max[3] == byte_min[3]) { + int tmp_state = fsm.AddState(); + fsm.AddEdge(from, tmp_state, byte_min[3], byte_max[3]); + min = (min & 0x00FFFFFF); + max = (max & 0x00FFFFFF); + AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); + return; + } + if ((min & 0x00FFFFFF) != 0x808080) { + int tmp_state_min = fsm.AddState(); + fsm.AddEdge(from, tmp_state_min, byte_min[3], byte_min[3]); + AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FFFFFF), 0x00BFBFBF); + } else { + byte_min[3]--; + } + if ((max & 0x00FFFFFF) != 0xBFBFBF) { + int tmp_state_max = fsm.AddState(); + fsm.AddEdge(from, tmp_state_max, byte_max[3], byte_max[3]); + AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x00808080, (max & 0x00FFFFFF)); + } else { + byte_max[3]++; + } + if (byte_max[3] - byte_min[3] > 1) { + int tmp_state_mid = fsm.AddState(); + // First byte. + fsm.AddEdge(from, tmp_state_mid, byte_min[3] + 1, byte_max[3] - 1); + int tmp_state_mid2 = fsm.AddState(); + // Second byte. + fsm.AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); + int tmp_state_mid3 = fsm.AddState(); + // Third byte. + fsm.AddEdge(tmp_state_mid2, tmp_state_mid3, 0x80, 0xBF); + // Last byte. + fsm.AddEdge(tmp_state_mid3, to, 0x80, 0xBF); + } + return; + } + if (byte_max[2] != 0) { + // 3 byte unicode. + if (byte_max[2] == byte_min[2]) { + int tmp_state = fsm.AddState(); + fsm.AddEdge(from, tmp_state, byte_min[2], byte_max[2]); + min = (min & 0x00FFFF); + max = (max & 0x00FFFF); + AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); + return; + } + if ((min & 0x00FFFF) != 0x8080) { + int tmp_state_min = fsm.AddState(); + fsm.AddEdge(from, tmp_state_min, byte_min[2], byte_min[2]); + AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FFFF), 0x00BFBF); + } else { + byte_min[2]--; + } + if ((max & 0x00FFFF) != 0xBFBF) { + int tmp_state_max = fsm.AddState(); + fsm.AddEdge(from, tmp_state_max, byte_max[2], byte_max[2]); + AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x0080, (max & 0x00FFFF)); + } else { + byte_max[2]++; + } + if (byte_max[2] - byte_min[2] > 1) { + int tmp_state_mid = fsm.AddState(); + // First byte. + fsm.AddEdge(from, tmp_state_mid, byte_min[2] + 1, byte_max[2] - 1); + int tmp_state_mid2 = fsm.AddState(); + // Second byte. + fsm.AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); + // Last byte. + fsm.AddEdge(tmp_state_mid2, to, 0x80, 0xBF); + } + return; + } + + // 2 byte unicode. + if (byte_max[1] == byte_min[1]) { + int tmp_state = fsm.AddState(); + fsm.AddEdge(from, tmp_state, byte_min[1], byte_max[1]); + min = (min & 0x00FF); + max = (max & 0x00FF); + AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); + return; + } + if ((min & 0x00FF) != 0x80) { + int tmp_state_min = fsm.AddState(); + fsm.AddEdge(from, tmp_state_min, byte_min[1], byte_min[1]); + AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FF), 0x00BF); + } else { + byte_min[1]--; + } + if ((max & 0x00FF) != 0xBF) { + int tmp_state_max = fsm.AddState(); + fsm.AddEdge(from, tmp_state_max, byte_max[1], byte_max[1]); + AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x0080, (max & 0x00FF)); + } else { + byte_max[1]++; + } + if (byte_max[1] - byte_min[1] > 1) { + int tmp_state_mid = fsm.AddState(); + // First byte. + fsm.AddEdge(from, tmp_state_mid, byte_min[1] + 1, byte_max[1] - 1); + fsm.AddEdge(tmp_state_mid, to, 0x80, 0xBF); + } + return; +} + +void AddPackedUTF8RangeEdges(FSM& fsm, int from, int to, uint32_t min, uint32_t max) { + XGRAMMAR_CHECK(min <= max) << "Invalid character range: min (" << min << ") > max (" << max + << ")"; + // Ensure max and min are valid unicode value. + if (max > kMax4BytesUnicode) { + max = kMax4BytesUnicode; + } else if (max > kMax3BytesUnicode) { + if (max < kMin4BytesUnicode) { + max = kMax3BytesUnicode; + } + } else if (max > kMax2BytesUnicode) { + if (max < kMin3BytesUnicode) { + max = kMax2BytesUnicode; + } + } else if (max < kMin2BytesUnicode && (max > kMax1ByteUnicode)) { + max = kMax1ByteUnicode; + } + + if (min > kMax4BytesUnicode) { + min = kMax4BytesUnicode; + } else if (min > kMax3BytesUnicode) { + if (min < kMin4BytesUnicode) { + min = kMin4BytesUnicode; + } + } else if (min > kMax2BytesUnicode) { + if (min < kMin3BytesUnicode) { + min = kMin3BytesUnicode; + } + } else if (min < kMin2BytesUnicode && (min > kMax1ByteUnicode)) { + min = kMin2BytesUnicode; + } + + // Step2. Divide the range into several ranges, which contain characters with different lengths. + if (max <= kMax1ByteUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + return; + } + if (max <= kMax2BytesUnicode) { + if (min >= kMin2BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + } else { + AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, max); + } + return; + } + if (max <= kMax3BytesUnicode) { + if (min >= kMin3BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + } else if (min >= kMin2BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, max); + } else { + AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, max); + } + return; + } + XGRAMMAR_CHECK(max <= kMax4BytesUnicode); + if (min >= kMin4BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + } else if (min >= kMin3BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, kMax3BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + } else if (min >= kMin2BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + } else { + AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + } + return; +} + +std::string RewriteRegexDots(const std::string& pattern, bool dot_matches_newline) { + if (dot_matches_newline) { + return pattern; + } + std::string result; + result.reserve(pattern.size()); + bool escaped = false; + bool in_character_class = false; + for (char c : pattern) { + if (escaped) { + result.push_back(c); + escaped = false; + continue; + } + if (c == '\\') { + result.push_back(c); + escaped = true; + } else if (c == '[') { + result.push_back(c); + in_character_class = true; + } else if (c == ']' && in_character_class) { + result.push_back(c); + in_character_class = false; + } else if (c == '.' && !in_character_class) { + result += "[^\\n]"; + } else { + result.push_back(c); + } + } + return result; +} + +/******************** Codepoint range utilities ********************/ + +namespace { + +constexpr uint32_t kMaxCodepoint = 0x10FFFF; + +/*! \brief Bounded repetitions above this threshold are compiled into a repeat FSM edge (when a + * GrammarBuilder is available) instead of being physically unrolled. Matches the unroll threshold + * of the grammar-level RepetitionRangeExpander. */ +constexpr int kLargeRepeatThreshold = 128; + +/*! \brief Hard cap of the estimated state count when a bounded repetition has to be unrolled + * because no GrammarBuilder is available. */ +constexpr int64_t kMaxUnrolledRepeatStates = 100000; + +using CodepointRange = std::pair; + +/*! \brief Sort the ranges and merge overlapping or adjacent ones. */ +void NormalizeRanges(std::vector* ranges) { + std::sort(ranges->begin(), ranges->end()); + std::vector result; + for (const auto& range : *ranges) { + if (!result.empty() && range.first <= result.back().second + 1 && + range.first >= result.back().first) { + result.back().second = std::max(result.back().second, range.second); + } else { + result.push_back(range); + } + } + *ranges = std::move(result); +} + +/*! \brief Complement normalized ranges over the codepoint domain [0, kMaxCodepoint]. */ +std::vector ComplementRanges(const std::vector& ranges) { + std::vector result; + uint32_t next = 0; + for (const auto& range : ranges) { + if (range.first > next) { + result.push_back({next, range.first - 1}); + } + if (range.second >= kMaxCodepoint) { + return result; + } + next = std::max(next, range.second + 1); + } + result.push_back({next, kMaxCodepoint}); + return result; +} + +/*! \brief Append the ASCII case-folded counterparts of every letter contained in the ranges. */ +void FoldAsciiCaseRanges(std::vector* ranges) { + size_t original_size = ranges->size(); + for (size_t i = 0; i < original_size; ++i) { + uint32_t low = (*ranges)[i].first; + uint32_t high = (*ranges)[i].second; + uint32_t fold_low = std::max(low, 'a'); + uint32_t fold_high = std::min(high, 'z'); + if (fold_low <= fold_high) { + ranges->push_back({fold_low - ('a' - 'A'), fold_high - ('a' - 'A')}); + } + fold_low = std::max(low, 'A'); + fold_high = std::min(high, 'Z'); + if (fold_low <= fold_high) { + ranges->push_back({fold_low + ('a' - 'A'), fold_high + ('a' - 'A')}); + } + } +} + +/*! \brief Add edges from `from` to `to` accepting the UTF-8 encoding of every codepoint in the + * normalized ranges. Multi-byte characters get intermediate states. */ +void AddCodepointRangesToFSM( + FSM* fsm, int from, int to, const std::vector& ranges +) { + for (const auto& [low, high] : ranges) { + if (low <= kMax1ByteUnicode) { + fsm->AddEdge(from, to, low, std::min(high, kMax1ByteUnicode)); + } + if (high > kMax1ByteUnicode) { + uint32_t multi_byte_low = std::max(low, kMax1ByteUnicode + 1); + AddPackedUTF8RangeEdges( + *fsm, from, to, CodepointToPackedUTF8(multi_byte_low), CodepointToPackedUTF8(high) + ); + } + } +} + +/*! \brief One parsed regex escape (or literal): either a single codepoint, or a (possibly + * negated) set of codepoint ranges for class escapes like \d, \D, \w, \W, \s, \S. */ +struct RegexEscapeItem { + std::vector ranges; + bool negated = false; + bool is_single = false; + uint32_t codepoint = 0; +}; + +/*! + * \brief Parse the escape sequence starting at regex[*pos] == '\\'. On success, *pos is advanced + * past the escape sequence. + * \param in_class Whether the escape appears inside a character class ([...]). Inside a class, + * \b means the backspace character instead of a word boundary assertion. + */ +Result ParseRegexEscape(const std::string& regex, size_t* pos, bool in_class) { + XGRAMMAR_DCHECK(regex[*pos] == '\\'); + if (*pos + 1 >= regex.size()) { + return ResultErr("Regex ends with a trailing backslash"); + } + char escaped = regex[*pos + 1]; + *pos += 2; + RegexEscapeItem item; + auto single = [&](uint32_t codepoint) { + item.is_single = true; + item.codepoint = codepoint; + return ResultOk(std::move(item)); + }; + switch (escaped) { + case 'n': + return single('\n'); + case 't': + return single('\t'); + case 'r': + return single('\r'); + case 'f': + return single('\f'); + case 'v': + return single('\v'); + case 'a': + return single('\a'); + case 'e': + return single(0x1B); + case '0': + return single(0); + case 'b': + if (in_class) { + return single(0x08); + } + return ResultErr("Word boundary assertion \\b is not supported in regex"); + case 'B': + return ResultErr("Word boundary assertion \\B is not supported in regex"); + case 'p': + case 'P': + return ResultErr("Unicode property escape \\p / \\P is not supported in regex"); + case 'k': + return ResultErr("Backreference \\k is not supported in regex"); + case '1': + case '2': + case '3': + case '4': + case '5': + case '6': + case '7': + case '8': + case '9': + return ResultErr("Backreference \\" + std::string(1, escaped) + " is not supported in regex"); + case 'd': + item.ranges = {{'0', '9'}}; + return ResultOk(std::move(item)); + case 'D': + item.ranges = {{'0', '9'}}; + item.negated = true; + return ResultOk(std::move(item)); + case 'w': + item.ranges = {{'0', '9'}, {'A', 'Z'}, {'_', '_'}, {'a', 'z'}}; + return ResultOk(std::move(item)); + case 'W': + item.ranges = {{'0', '9'}, {'A', 'Z'}, {'_', '_'}, {'a', 'z'}}; + item.negated = true; + return ResultOk(std::move(item)); + case 's': + item.ranges = {{0x09, 0x0D}, {' ', ' '}}; + return ResultOk(std::move(item)); + case 'S': + item.ranges = {{0x09, 0x0D}, {' ', ' '}}; + item.negated = true; + return ResultOk(std::move(item)); + case 'x': { + if (*pos + 1 >= regex.size()) { + return ResultErr("\\x must be followed by two hexadecimal digits in regex"); + } + int high_digit = HexCharToInt(regex[*pos]); + int low_digit = HexCharToInt(regex[*pos + 1]); + if (high_digit < 0 || low_digit < 0) { + return ResultErr("\\x must be followed by two hexadecimal digits in regex"); + } + *pos += 2; + return single(high_digit * 16 + low_digit); + } + case 'u': { + if (*pos < regex.size() && regex[*pos] == '{') { + size_t close = regex.find('}', *pos + 1); + if (close == std::string::npos || close == *pos + 1 || close > *pos + 7) { + return ResultErr("\\u{...} must contain one to six hexadecimal digits in regex"); + } + uint32_t codepoint = 0; + for (size_t i = *pos + 1; i < close; ++i) { + int digit = HexCharToInt(regex[i]); + if (digit < 0) { + return ResultErr("\\u{...} must contain one to six hexadecimal digits in regex"); + } + codepoint = codepoint * 16 + digit; + } + if (codepoint > kMaxCodepoint) { + return ResultErr("\\u{...} escape is beyond the Unicode range in regex"); + } + *pos = close + 1; + return single(codepoint); + } + if (*pos + 3 >= regex.size()) { + return ResultErr("\\u must be followed by four hexadecimal digits in regex"); + } + uint32_t codepoint = 0; + for (size_t i = *pos; i < *pos + 4; ++i) { + int digit = HexCharToInt(regex[i]); + if (digit < 0) { + return ResultErr("\\u must be followed by four hexadecimal digits in regex"); + } + codepoint = codepoint * 16 + digit; + } + *pos += 4; + return single(codepoint); + } + case 'c': { + if (*pos >= regex.size() || !std::isalpha(static_cast(regex[*pos]))) { + return ResultErr("\\c must be followed by a letter in regex"); + } + uint32_t codepoint = static_cast(regex[*pos]) & 0x1F; + *pos += 1; + return single(codepoint); + } + default: { + if (static_cast(escaped) >= 0x80) { + // Multi-byte UTF-8 character after the backslash: match it literally. + *pos -= 1; + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + *pos); + if (codepoint == CharHandlingError::kInvalidUTF8 || *pos + num_bytes > regex.size()) { + return ResultErr("Invalid UTF-8 in regex escape sequence"); + } + *pos += num_bytes; + return single(static_cast(codepoint)); + } + if (std::isalnum(static_cast(escaped))) { + XGRAMMAR_LOG(WARNING) << "Escape sequence \\" << escaped + << " is not recognized in regex; matching the character literally"; + } + return single(static_cast(escaped)); + } + } +} + +/*! + * \brief Parse a character class leaf "[...]" into the final set of accepted codepoint ranges. + * ASCII case folding (if requested) is applied before negation. + */ +Result> ParseCharacterClassLeaf( + const std::string& regex, bool case_insensitive +) { + XGRAMMAR_DCHECK(regex.size() >= 2 && regex.front() == '[' && regex.back() == ']'); + size_t pos = 1; + size_t content_end = regex.size() - 1; + bool negated = false; + if (pos < content_end && regex[pos] == '^') { + negated = true; + ++pos; + } + if (pos == content_end) { + return ResultErr("Empty character class " + regex + " is not allowed in regex"); + } + + // Parse one class unit: an escape sequence or a literal (possibly multi-byte) character. + auto parse_unit = [&]() -> Result { + if (regex[pos] == '\\') { + return ParseRegexEscape(regex, &pos, /*in_class=*/true); + } + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + pos); + if (codepoint == CharHandlingError::kInvalidUTF8 || pos + num_bytes > content_end) { + return ResultErr("Invalid UTF-8 in regex character class " + regex); + } + pos += num_bytes; + RegexEscapeItem item; + item.is_single = true; + item.codepoint = static_cast(codepoint); + return ResultOk(std::move(item)); + }; + + std::vector ranges; + while (pos < content_end) { + auto unit_result = parse_unit(); + if (unit_result.IsErr()) { + return ResultErr(std::move(unit_result).UnwrapErr()); + } + auto unit = std::move(unit_result).Unwrap(); + if (!unit.is_single) { + // Class escapes like \d cannot be a range endpoint; a following '-' is literal. + if (unit.negated) { + NormalizeRanges(&unit.ranges); + auto complement = ComplementRanges(unit.ranges); + ranges.insert(ranges.end(), complement.begin(), complement.end()); + } else { + ranges.insert(ranges.end(), unit.ranges.begin(), unit.ranges.end()); + } + continue; + } + if (pos < content_end && regex[pos] == '-' && pos + 1 < content_end) { + ++pos; + auto high_result = parse_unit(); + if (high_result.IsErr()) { + return ResultErr(std::move(high_result).UnwrapErr()); + } + auto high_unit = std::move(high_result).Unwrap(); + if (!high_unit.is_single) { + return ResultErr("Invalid character range endpoint in regex character class " + regex); + } + if (high_unit.codepoint < unit.codepoint) { + return ResultErr( + "Invalid character range (lower bound exceeds upper bound) in regex character class " + + regex + ); + } + ranges.push_back({unit.codepoint, high_unit.codepoint}); + } else { + ranges.push_back({unit.codepoint, unit.codepoint}); + } + } + + if (case_insensitive) { + FoldAsciiCaseRanges(&ranges); + } + NormalizeRanges(&ranges); + if (negated) { + ranges = ComplementRanges(ranges); + } + return ResultOk(std::move(ranges)); +} + +} // namespace + +/******************** RegexIR ********************/ + +class RegexIR { + public: + struct Leaf; + + struct Symbol; + + struct Union; + + struct Bracket; + + struct Repeat; + + struct RuleRefNode; + + struct RepeatSubrule; + + static constexpr int kRepeatNoUpperBound = -1; + + using State = std::variant; + + // This struct is used to store one atom of the regex: the empty string (regex == ""), a + // character class (regex == "[...]"), or a short sequence of literal characters / escapes. + struct Leaf { + std::string regex; + }; + + // This struct is used to store the symbol in regex, i.e. + // +, *, ? + enum class RegexSymbol { + star, + plus, + optional, + }; + + struct Bracket { + std::vector states; + }; + + struct Symbol { + RegexSymbol symbol; + std::vector state; + }; + + // This struct is used to represent a union symbol. + struct Union { + std::vector states; + }; + + struct Repeat { + std::vector states; + int lower_bound = 0; + int upper_bound = 0; + }; + + // A reference to a grammar rule, compiled into a kRuleRef FSM edge. + struct RuleRefNode { + int32_t rule_id; + }; + + // A bounded repetition of a grammar rule, compiled into a kRepeatRef FSM edge. The referenced + // rule holds the repeated sub-pattern; the Earley parser executes the repetition with a + // counter at runtime, so no FSM unrolling happens. + struct RepeatSubrule { + int32_t rule_id; + int lower_bound = 0; + int upper_bound = 0; + }; + + // The top-level sequence of the regex. + std::vector states; + + // Whether matching is ASCII case-insensitive (enabled by a leading "(?i)"). + bool case_insensitive = false; + + /*! + \brief Constructs a NFA from the regex IR. + */ + Result Build() const; + + /*! + \brief the visit function for the variant. + */ + Result visit(const Leaf& state) const; + + Result visit(const Symbol& state) const; + + Result visit(const Union& state) const; + + Result visit(const Bracket& state) const; + + Result visit(const Repeat& state) const; + + Result visit(const RuleRefNode& state) const; + + Result visit(const RepeatSubrule& state) const; + + /*! \brief Whether the IR node can match the empty string. Purely syntactic; no FSM is built. */ + static bool IsNullable(const State& state); + + /*! \brief Whether a sequence of IR nodes can match the empty string. */ + static bool IsNullableSequence(const std::vector& states); + + /*! + * \brief Check repeat in regex. i.e {...} and {...,...} + * \param regex The regex string. + * \param start The start position of the repeat. i.e. regex[start] == '{'. + * After the function, start will be the position of '}'. + * \return The repeat range. + */ + static Result> CheckRepeat(const std::string& regex, int& start); + + private: + /*! + * \brief Construct a FSM from a regex leaf. + * \details The leaf is the empty string, a character class like [a-c0-9], or a sequence of + * literal characters / escapes like "ab\n". Any symbols like "a|b" or "a*b" are not supported. + * \param regex The regex string. + * \return The FSM with start and end states. + */ + Result BuildLeafFSMFromRegex(const std::string& regex) const; + + /*! + * \brief Add the transition(s) accepting a single codepoint (with ASCII case folding when + * case_insensitive is set) from `current` to a new state, and return the new state. + */ + int AddSingleCodepoint(FSMWithStartEnd& result, int current, uint32_t codepoint) const; +}; + +Result> RegexIR::CheckRepeat(const std::string& regex, int& start) { + // 10^9 fits in an int; longer counts would overflow. + constexpr size_t kMaxRepeatDigits = 9; + if (regex[start] != '{') { + return ResultErr("Invalid repetition: expected '{'"); + } + int lower_bound = 0; + int upper_bound = RegexIR::kRepeatNoUpperBound; + std::string num_str; + XGRAMMAR_DCHECK(regex[start] == '{'); + start++; + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + while (static_cast(start) < regex.size() && std::isdigit(regex[start])) { + num_str += regex[start]; + start++; + } + if (num_str.empty()) { + return ResultErr("Invalid repetition count: expected a number after '{'"); + } + if (num_str.size() > kMaxRepeatDigits) { + return ResultErr("Invalid repetition count: the count " + num_str + " is too large"); + } + lower_bound = std::stoi(num_str); + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + // The format is {n} + if (regex[start] == '}') { + upper_bound = lower_bound; + return ResultOk(std::make_pair(lower_bound, upper_bound)); + } + if (regex[start] != ',') { + return ResultErr("Invalid repetition count: expected ',' or '}' after the lower bound"); + } + XGRAMMAR_DCHECK(regex[start] == ','); + start++; + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + // The format is {n,} + if (regex[start] == '}') { + return ResultOk(std::make_pair(lower_bound, upper_bound)); + } + num_str.clear(); + while (static_cast(start) < regex.size() && std::isdigit(regex[start])) { + num_str += regex[start]; + start++; + } + if (num_str.empty()) { + return ResultErr("Invalid repetition count: expected a number or '}' after ','"); + } + if (num_str.size() > kMaxRepeatDigits) { + return ResultErr("Invalid repetition count: the count " + num_str + " is too large"); + } + upper_bound = std::stoi(num_str); + if (upper_bound < lower_bound) { + return ResultErr( + "Invalid repetition count: the lower bound " + std::to_string(lower_bound) + + " is larger than the upper bound " + std::to_string(upper_bound) + ); + } + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + if (regex[start] != '}') { + return ResultErr("Invalid repetition count: expected '}' after the upper bound"); + } + XGRAMMAR_DCHECK(regex[start] == '}'); + return ResultOk(std::make_pair(lower_bound, upper_bound)); +} + +bool RegexIR::IsNullableSequence(const std::vector& states) { + return std::all_of(states.begin(), states.end(), [](const State& state) { + return IsNullable(state); + }); +} + +bool RegexIR::IsNullable(const State& state) { + return std::visit( + [](const auto& node) -> bool { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return node.regex.empty(); + } else if constexpr (std::is_same_v) { + return node.symbol != RegexSymbol::plus || IsNullableSequence(node.state); + } else if constexpr (std::is_same_v) { + return std::any_of(node.states.begin(), node.states.end(), [](const State& child) { + return IsNullable(child); + }); + } else if constexpr (std::is_same_v) { + return IsNullableSequence(node.states); + } else if constexpr (std::is_same_v) { + return node.lower_bound == 0 || IsNullableSequence(node.states); + } else if constexpr (std::is_same_v) { + return false; + } else { + static_assert(std::is_same_v); + return node.lower_bound == 0; + } + }, + state + ); +} + +Result RegexIR::Build() const { + if (states.empty()) { + FSM empty_fsm(1); + FSMWithStartEnd result(empty_fsm, 0, {0}, false); + return ResultOk(std::move(result)); + } + std::vector fsm_list; + for (const auto& state : states) { + auto visited = std::visit([&](auto&& arg) { return visit(arg); }, state); + if (visited.IsErr()) { + return visited; + } + fsm_list.push_back(std::move(visited).Unwrap()); + } + if (fsm_list.size() > 1) { + return ResultOk(FSMWithStartEnd::Concat(fsm_list)); + } else { + // If there is only one FSM, return it directly. + return ResultOk(std::move(fsm_list[0])); + } +} + +Result RegexIR::visit(const RegexIR::Leaf& state) const { + return BuildLeafFSMFromRegex(state.regex); +} + +Result RegexIR::visit(const RegexIR::Union& state) const { + std::vector fsm_list; + for (const auto& child : state.states) { + auto visited = std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, child); + if (visited.IsErr()) { + return visited; + } + fsm_list.push_back(std::move(visited).Unwrap()); + } + if (fsm_list.size() <= 1) { + return ResultErr("Internal error: a union node in the regex IR has fewer than two branches"); + } + return ResultOk(FSMWithStartEnd::Union(fsm_list)); +} + +Result RegexIR::visit(const RegexIR::Symbol& state) const { + if (state.state.size() != 1) { + return ResultErr("Internal error: a quantifier node in the regex IR must hold exactly one child" + ); + } + Result child_result = + std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, state.state[0]); + if (child_result.IsErr()) { + return child_result; + } + auto child = std::move(child_result).Unwrap(); + + switch (state.symbol) { + case RegexIR::RegexSymbol::plus: { + return ResultOk(child.Plus()); + } + case RegexIR::RegexSymbol::star: { + return ResultOk(child.Star()); + } + case RegexIR::RegexSymbol::optional: { + return ResultOk(child.Optional()); + } + default: { + XGRAMMAR_LOG(FATAL) << "Unknown regex symbol: " << static_cast(state.symbol); + XGRAMMAR_UNREACHABLE(); + } + } +} + +Result RegexIR::visit(const RegexIR::Bracket& state) const { + if (state.states.empty()) { + // An empty group or an empty union branch matches the empty string. + FSM empty_fsm(1); + return ResultOk(FSMWithStartEnd(empty_fsm, 0, {0}, false)); + } + std::vector fsm_list; + for (const auto& child : state.states) { + auto visited = std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, child); + if (visited.IsErr()) { + return visited; + } + fsm_list.push_back(std::move(visited).Unwrap()); + } + return ResultOk(FSMWithStartEnd::Concat(fsm_list)); +} + +Result RegexIR::visit(const RegexIR::RuleRefNode& state) const { + FSM fsm(2); + fsm.AddRuleEdge(0, 1, state.rule_id); + return ResultOk(FSMWithStartEnd(fsm, 0, {1}, false)); +} + +Result RegexIR::visit(const RegexIR::RepeatSubrule& state) const { + // The state holding the repeat edge is padded with epsilon transitions on both sides, so that + // FSM compositions (Concat / Star / Optional / ...) never add another outgoing edge to it: a + // state with a kRepeatRef edge must have no other outgoing edges. + FSM fsm(4); + fsm.AddEpsilonEdge(0, 1); + fsm.AddRepeatEdge(1, 2, state.rule_id, state.lower_bound, state.upper_bound); + fsm.AddEpsilonEdge(2, 3); + return ResultOk(FSMWithStartEnd(fsm, 0, {3}, false)); +} + +Result RegexIR::visit(const RegexIR::Repeat& state) const { + if (state.states.size() != 1) { + return ResultErr("Internal error: a repetition node in the regex IR must hold exactly one child" + ); + } + bool has_upper_bound = state.upper_bound != RegexIR::kRepeatNoUpperBound; + if (has_upper_bound && state.upper_bound == 0) { + // {0} / {0,0}: matches exactly the empty string. The general path below cannot express + // this: it starts from one copy of the child whose end states stay accepting. + FSM empty_fsm(1); + return ResultOk(FSMWithStartEnd(empty_fsm, 0, {0}, false)); + } + Result child_result = + std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, state.states[0]); + if (child_result.IsErr()) { + return child_result; + } + FSMWithStartEnd child = std::move(child_result).Unwrap(); + + // Guard against FSM state explosion when the repetition has to be physically unrolled. When + // a GrammarBuilder is available, large repetitions are compiled into repeat edges instead and + // never reach this point. + int64_t num_copies = has_upper_bound ? state.upper_bound : std::max(state.lower_bound, 1); + if (static_cast(child.NumStates()) * num_copies > kMaxUnrolledRepeatStates) { + return ResultErr( + "The bounded repetition {" + std::to_string(state.lower_bound) + "," + + (has_upper_bound ? std::to_string(state.upper_bound) : "") + + "} is too large to compile into a FSM" + ); + } + + FSMWithStartEnd result = child.Copy(); + std::unordered_set new_ends; + + if (state.lower_bound <= 1 && (!has_upper_bound || state.upper_bound >= 1)) { + // A single copy is accepting when the lower bound is at most 1. + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + new_ends.insert(end); + } + } + } + + // Add a fresh accepting start state so that zero repetitions match. A fresh state is + // required: making the original start accepting would also accept strings that merely + // loop back to the start inside the first copy. + auto allow_zero_repetitions = [](FSMWithStartEnd* fsm) { + int new_start = fsm->AddState(); + fsm->GetFsm().AddEpsilonEdge(new_start, fsm->GetStart()); + fsm->SetStartState(new_start); + fsm->AddEndState(new_start); + }; + + // Handling {n,} + if (!has_upper_bound) { + for (int i = 2; i < state.lower_bound; i++) { + result = FSMWithStartEnd::Concat(std::vector{result, child}); + } + int end_state_of_lower_bound_fsm = -1; + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + end_state_of_lower_bound_fsm = end; + break; + } + } + XGRAMMAR_DCHECK(end_state_of_lower_bound_fsm != -1) + << "No end state found in the lower bound FSM."; + result = FSMWithStartEnd::Concat(std::vector{result, child}); + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + result.GetFsm().AddEpsilonEdge(end, end_state_of_lower_bound_fsm); + } + } + for (const auto& end : new_ends) { + result.AddEndState(end); + } + if (state.lower_bound == 0) { + allow_zero_repetitions(&result); + } + return ResultOk(std::move(result)); + } + // Handling {n, m} or {n} + for (int i = 2; i <= state.upper_bound; i++) { + result = FSMWithStartEnd::Concat(std::vector{result, child}); + if (i >= state.lower_bound) { + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + new_ends.insert(end); + } + } + } + } + for (const auto& end : new_ends) { + result.AddEndState(end); + } + if (state.lower_bound == 0) { + allow_zero_repetitions(&result); + } + return ResultOk(std::move(result)); +} + +int RegexIR::AddSingleCodepoint(FSMWithStartEnd& result, int current, uint32_t codepoint) const { + int next = result.AddState(); + if (codepoint <= kMax1ByteUnicode) { + result.GetFsm().AddEdge(current, next, codepoint, codepoint); + if (case_insensitive) { + if (codepoint >= 'a' && codepoint <= 'z') { + uint32_t upper = codepoint - ('a' - 'A'); + result.GetFsm().AddEdge(current, next, upper, upper); + } else if (codepoint >= 'A' && codepoint <= 'Z') { + uint32_t lower = codepoint + ('a' - 'A'); + result.GetFsm().AddEdge(current, next, lower, lower); + } + } + return next; + } + std::string utf8_bytes = CharToUTF8(static_cast(codepoint)); + int state = current; + for (size_t i = 0; i < utf8_bytes.size(); ++i) { + int target = (i + 1 == utf8_bytes.size()) ? next : result.AddState(); + uint8_t byte = static_cast(utf8_bytes[i]); + result.GetFsm().AddEdge(state, target, byte, byte); + state = target; + } + return next; +} + +Result RegexIR::BuildLeafFSMFromRegex(const std::string& regex) const { + FSM initial_fsm(1); + FSMWithStartEnd result(initial_fsm, 0, {}, false); + if (regex.empty()) { + // The empty leaf matches the empty string. + result.AddEndState(0); + return ResultOk(std::move(result)); + } + if (regex[0] == '[') { + // Character class. + auto ranges_result = ParseCharacterClassLeaf(regex, case_insensitive); + if (ranges_result.IsErr()) { + return ResultErr(std::move(ranges_result).UnwrapErr()); + } + auto ranges = std::move(ranges_result).Unwrap(); + int end_state = result.AddState(); + AddCodepointRangesToFSM(&result.GetFsm(), 0, end_state, ranges); + result.AddEndState(end_state); + return ResultOk(std::move(result)); + } + // Sequence of literal characters, '.' and escapes. + int current = 0; + size_t pos = 0; + while (pos < regex.size()) { + if (regex[pos] == '.') { + ++pos; + int next = result.AddState(); + AddCodepointRangesToFSM(&result.GetFsm(), current, next, {{0, kMaxCodepoint}}); + current = next; + continue; + } + if (regex[pos] == '\\') { + auto item_result = ParseRegexEscape(regex, &pos, /*in_class=*/false); + if (item_result.IsErr()) { + return ResultErr(std::move(item_result).UnwrapErr()); + } + auto item = std::move(item_result).Unwrap(); + if (item.is_single) { + current = AddSingleCodepoint(result, current, item.codepoint); + } else { + auto ranges = std::move(item.ranges); + if (case_insensitive) { + FoldAsciiCaseRanges(&ranges); + } + NormalizeRanges(&ranges); + if (item.negated) { + ranges = ComplementRanges(ranges); + } + int next = result.AddState(); + AddCodepointRangesToFSM(&result.GetFsm(), current, next, ranges); + current = next; + } + continue; + } + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + pos); + if (codepoint == CharHandlingError::kInvalidUTF8 || pos + num_bytes > regex.size()) { + // Be permissive with non-UTF-8 patterns: match the raw byte. + int next = result.AddState(); + uint8_t byte = static_cast(regex[pos]); + result.GetFsm().AddEdge(current, next, byte, byte); + current = next; + ++pos; + continue; + } + current = AddSingleCodepoint(result, current, static_cast(codepoint)); + pos += num_bytes; + } + result.AddEndState(current); + return ResultOk(std::move(result)); +} + +/******************** RegexFSMBuilder ********************/ + +namespace { + +/*! \brief One entry of the regex parsing stack: an IR node or a marker character ('(' or '|'), + * together with the source span [span_begin, span_end) of the atom in the regex string. */ +struct RegexStackEntry { + std::variant item; + int span_begin = 0; + int span_end = 0; +}; + +/*! \brief Skip a character class starting at regex[pos] == '['. Returns the position right after + * the closing ']', or std::string::npos if the class is not closed. */ +size_t SkipCharacterClass(const std::string& regex, size_t pos) { + XGRAMMAR_DCHECK(regex[pos] == '['); + ++pos; + if (pos < regex.size() && regex[pos] == '^') { + ++pos; + } + while (pos < regex.size()) { + if (regex[pos] == '\\') { + pos += 2; + continue; + } + if (regex[pos] == ']') { + return pos + 1; + } + ++pos; + } + return std::string::npos; +} + +/*! + * \brief Parse a regex string into a RegexIR. + * \param regex_with_flags The regex. A leading "(?i)" enables ASCII case-insensitive matching. + * \param builder If not null, large bounded repetitions are compiled into subrules added through + * this builder; see RegexFSMBuilder::Build. + * \param rule_hint Name hint for the created subrules. + */ +Result ParseRegexToIR( + const std::string& regex_with_flags, GrammarBuilder* builder, const std::string& rule_hint +) { + RegexIR ir; + std::string regex = regex_with_flags; + int flag_prefix_length = 0; + if (regex.size() >= 4 && regex.compare(0, 4, "(?i)") == 0) { + ir.case_insensitive = true; + regex = regex.substr(4); + flag_prefix_length = 4; + } + + // Mirror the error format of the RegexConverter path: a 1-based position in the pattern + // (including the "(?i)" prefix when present), followed by the description. + auto error_at = [&](int pos, const std::string& message) { + return "Regex parsing error at position " + std::to_string(pos + flag_prefix_length + 1) + + ": " + message; + }; + + std::vector stack; + int size = static_cast(regex.size()); + for (int i = 0; i < size; i++) { + char current_char = regex[i]; + // Handle anchors. + if (current_char == '^' || current_char == '$') { + if (!((current_char == '^' && i == 0) || (current_char == '$' && i == size - 1))) { + XGRAMMAR_LOG(WARNING) << "Anchor '" << current_char + << "' in the middle of regex is ignored: " << regex; + } + continue; + } + // Handle the character class. + if (current_char == '[') { + size_t class_end = SkipCharacterClass(regex, i); + if (class_end == std::string::npos) { + return ResultErr(error_at(i, "Unclosed '['")); + } + size_t content_begin = i + 1; + if (content_begin < regex.size() && regex[content_begin] == '^') { + ++content_begin; + } + if (content_begin + 1 == class_end) { + return ResultErr(error_at(i, "Empty character class is not allowed in regex")); + } + RegexIR::Leaf leaf; + leaf.regex = regex.substr(i, class_end - i); + stack.push_back({leaf, i, static_cast(class_end)}); + i = static_cast(class_end) - 1; + continue; + } + if (current_char == ']') { + return ResultErr(error_at(i, "Unmatched ']'")); + } + // Handle quantifiers. + if (current_char == '+' || current_char == '*' || current_char == '?') { + if (stack.empty() || std::holds_alternative(stack.back().item)) { + return ResultErr( + error_at(i, std::string("There is nothing to repeat before '") + current_char + "'") + ); + } + RegexStackEntry atom = std::move(stack.back()); + stack.pop_back(); + RegexIR::Symbol symbol; + symbol.state.push_back(std::move(std::get(atom.item))); + switch (current_char) { + case '+': { + symbol.symbol = RegexIR::RegexSymbol::plus; + break; + } + case '*': { + symbol.symbol = RegexIR::RegexSymbol::star; + break; + } + case '?': { + symbol.symbol = RegexIR::RegexSymbol::optional; + break; + } + } + // Skip the non-greedy modifier: greedy and non-greedy quantifiers accept the same + // language, and the Earley parser explores all derivations anyway. + if (i + 1 < size && regex[i + 1] == '?') { + i++; + } + stack.push_back({std::move(symbol), atom.span_begin, i + 1}); + continue; + } + // Handle groups and alternation. + if (current_char == '(') { + if (i + 1 < size && regex[i + 1] == '?') { + if (i + 2 >= size) { + return ResultErr(error_at(i, "Group modifier is not finished")); + } + char modifier = regex[i + 2]; + if (modifier == ':') { + stack.push_back({'(', i, i + 1}); + i += 2; + continue; + } + if (modifier == '=' || modifier == '!') { + // Skip the whole lookahead group and treat it as the empty string. + XGRAMMAR_LOG(WARNING) << "Lookahead assertion is not supported and is ignored in regex: " + << regex; + int depth = 1; + size_t j = i + 3; + while (j < regex.size() && depth > 0) { + if (regex[j] == '\\') { + j += 2; + continue; + } + if (regex[j] == '[') { + size_t class_begin = j; + j = SkipCharacterClass(regex, j); + if (j == std::string::npos) { + return ResultErr(error_at(static_cast(class_begin), "Unclosed '['")); + } + continue; + } + if (regex[j] == '(') { + depth++; + } else if (regex[j] == ')') { + depth--; + } + j++; + } + if (depth != 0) { + return ResultErr(error_at(i, "The parenthesis is not closed")); + } + stack.push_back({RegexIR::Leaf{""}, i, static_cast(j)}); + i = static_cast(j) - 1; + continue; + } + if (modifier == '<' || (modifier == 'P' && i + 3 < size && regex[i + 3] == '<')) { + size_t name_begin = (modifier == '<') ? i + 3 : i + 4; + if (name_begin < regex.size() && (regex[name_begin] == '=' || regex[name_begin] == '!')) { + return ResultErr(error_at(i, "Lookbehind assertion is not supported in regex")); + } + size_t j = name_begin; + while (j < regex.size() && + (std::isalnum(static_cast(regex[j])) || regex[j] == '_')) { + j++; + } + if (j == name_begin || j >= regex.size() || regex[j] != '>') { + return ResultErr(error_at(i, "Invalid named capturing group")); + } + // Ignore the group name and compile the content as a normal group. + stack.push_back({'(', i, i + 1}); + i = static_cast(j); + continue; + } + return ResultErr( + error_at(i, "Unsupported group modifier '(?" + std::string(1, modifier) + "'") + ); + } + stack.push_back({'(', i, i + 1}); + continue; + } + if (current_char == '|') { + stack.push_back({'|', i, i + 1}); + continue; + } + if (current_char == ')') { + std::vector popped; + bool paired = false; + bool unioned = false; + int group_begin = 0; + while (!stack.empty()) { + RegexStackEntry entry = std::move(stack.back()); + stack.pop_back(); + if (std::holds_alternative(entry.item)) { + char marker = std::get(entry.item); + if (marker == '(') { + paired = true; + group_begin = entry.span_begin; + break; + } + XGRAMMAR_DCHECK(marker == '|'); + unioned = true; + } + popped.push_back(std::move(entry)); + } + if (!paired) { + return ResultErr(error_at(i, "Unmatched ')'")); + } + // `popped` stores the group content from right to left. + if (!unioned) { + RegexIR::Bracket bracket; + for (auto it = popped.rbegin(); it != popped.rend(); ++it) { + bracket.states.push_back(std::move(std::get(it->item))); + } + // An empty bracket (e.g. "()") matches the empty string. + stack.push_back({std::move(bracket), group_begin, i + 1}); + } else { + RegexIR::Union union_state; + RegexIR::Bracket bracket; + for (auto it = popped.rbegin(); it != popped.rend(); ++it) { + if (std::holds_alternative(it->item)) { + XGRAMMAR_DCHECK(std::get(it->item) == '|'); + // An empty bracket represents an empty alternative, e.g. "(a|)". + union_state.states.push_back(std::move(bracket)); + bracket = RegexIR::Bracket(); + continue; + } + bracket.states.push_back(std::move(std::get(it->item))); + } + union_state.states.push_back(std::move(bracket)); + stack.push_back({std::move(union_state), group_begin, i + 1}); + } + continue; + } + // Handle repetitions. + if (current_char == '{') { + if (stack.empty() || std::holds_alternative(stack.back().item)) { + return ResultErr(error_at(i, "There is nothing to repeat before the repetition")); + } + RegexStackEntry atom = std::move(stack.back()); + stack.pop_back(); + int repeat_begin = i; + auto bounds_result = RegexIR::CheckRepeat(regex, i); + if (bounds_result.IsErr()) { + return ResultErr(error_at(repeat_begin, std::move(bounds_result).UnwrapErr().what())); + } + auto [lower_bound, upper_bound] = std::move(bounds_result).Unwrap(); + // Skip the non-greedy modifier. + if (i + 1 < size && regex[i + 1] == '?') { + i++; + } + bool is_large_repeat = upper_bound == RegexIR::kRepeatNoUpperBound + ? lower_bound > kLargeRepeatThreshold + : upper_bound > kLargeRepeatThreshold; + if (is_large_repeat && builder != nullptr) { + // Compile the repeated sub-pattern into a new rule (with a kRegex body), and represent + // the repetition as a kRepeatRef FSM edge. The Earley parser executes it with a counter + // at runtime, so the FSM is not unrolled. + const auto& atom_state = std::get(atom.item); + std::string inner_regex = regex.substr(atom.span_begin, atom.span_end - atom.span_begin); + if (ir.case_insensitive) { + inner_regex = "(?i)" + inner_regex; + } + if (RegexIR::IsNullable(atom_state)) { + // Mirror RepetitionNormalizer: when the repeated element can match the empty string, + // the repetition count cannot be enforced, so the lower bound is relaxed to 0. + lower_bound = 0; + } + std::string name_hint = (rule_hint.empty() ? "regex" : rule_hint) + "_repeat"; + int32_t inner_rule_id = builder->AddRuleWithHint(name_hint, builder->AddRegex(inner_regex)); + builder->UpdateLookaheadExact(inner_rule_id, true); + if (upper_bound == RegexIR::kRepeatNoUpperBound) { + // {n,} == {n}{0,}: a repeat edge for the mandatory part, then a starred rule + // reference. + RegexIR::Symbol star_symbol; + star_symbol.symbol = RegexIR::RegexSymbol::star; + star_symbol.state.push_back(RegexIR::RuleRefNode{inner_rule_id}); + if (lower_bound > 0) { + RegexIR::Bracket bracket; + bracket.states.push_back(RegexIR::RepeatSubrule{inner_rule_id, lower_bound, lower_bound} + ); + bracket.states.push_back(std::move(star_symbol)); + stack.push_back({std::move(bracket), atom.span_begin, i + 1}); + } else { + stack.push_back({std::move(star_symbol), atom.span_begin, i + 1}); + } + } else { + stack.push_back( + {RegexIR::RepeatSubrule{inner_rule_id, lower_bound, upper_bound}, + atom.span_begin, + i + 1} + ); + } + } else { + RegexIR::Repeat repeat; + repeat.lower_bound = lower_bound; + repeat.upper_bound = upper_bound; + repeat.states.push_back(std::move(std::get(atom.item))); + stack.push_back({std::move(repeat), atom.span_begin, i + 1}); + } + continue; + } + // Handle literal characters and escapes. Each leaf holds exactly one codepoint or escape + // sequence, so that a following quantifier applies to the whole character. + RegexIR::Leaf leaf; + if (current_char == '\\') { + size_t escape_end = i; + auto escape_result = ParseRegexEscape(regex, &escape_end, /*in_class=*/false); + if (escape_result.IsErr()) { + return ResultErr(error_at(i, std::move(escape_result).UnwrapErr().what())); + } + leaf.regex = regex.substr(i, escape_end - i); + stack.push_back({std::move(leaf), i, static_cast(escape_end)}); + i = static_cast(escape_end) - 1; + continue; + } + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + i); + if (codepoint == CharHandlingError::kInvalidUTF8 || i + num_bytes > size) { + num_bytes = 1; + } + leaf.regex = regex.substr(i, num_bytes); + stack.push_back({std::move(leaf), i, i + num_bytes}); + i += num_bytes - 1; + continue; + } + + // Assemble the top-level sequence / union. `stack` stores the content from left to right. + std::vector segment; + std::vector> union_segments; + bool unioned = false; + for (auto& entry : stack) { + if (std::holds_alternative(entry.item)) { + char marker = std::get(entry.item); + if (marker == '|') { + union_segments.push_back(std::move(segment)); + segment.clear(); + unioned = true; + continue; + } + return ResultErr(error_at(entry.span_begin, "The parenthesis is not closed")); + } + segment.push_back(std::move(std::get(entry.item))); + } + if (!unioned) { + ir.states = std::move(segment); + } else { + union_segments.push_back(std::move(segment)); + RegexIR::Union union_state; + for (auto& branch : union_segments) { + RegexIR::Bracket bracket; + bracket.states = std::move(branch); + union_state.states.push_back(std::move(bracket)); + } + ir.states.push_back(std::move(union_state)); + } + return ResultOk(std::move(ir)); +} + +} // namespace + +Result RegexFSMBuilder::Build( + const std::string& regex, GrammarBuilder* builder, const std::string& rule_hint +) { + auto ir_result = ParseRegexToIR(regex, builder, rule_hint); + if (ir_result.IsErr()) { + return ResultErr(std::move(ir_result).UnwrapErr()); + } + return std::move(ir_result).Unwrap().Build(); +} + +Result RegexFSMBuilder::MatchesEmpty(const std::string& regex) { + auto ir_result = ParseRegexToIR(regex, nullptr, ""); + if (ir_result.IsErr()) { + return ResultErr(std::move(ir_result).UnwrapErr()); + } + auto ir = std::move(ir_result).Unwrap(); + return ResultOk(RegexIR::IsNullableSequence(ir.states)); +} + +Result RegexFSMBuilder::BuildWithForbiddenChars( + const std::string& regex, + const std::bitset<256>& forbidden_chars, + GrammarBuilder* builder, + const std::string& rule_hint +) { + auto build_result = Build(regex, builder, rule_hint); + if (build_result.IsErr() || forbidden_chars.none()) { + return build_result; + } + auto fsm_wse = std::move(build_result).Unwrap(); + const auto& fsm = fsm_wse.GetFsm(); + FSM new_fsm(fsm_wse.NumStates()); + for (int state = 0; state < fsm_wse.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (!edge.IsCharRange()) { + new_fsm.AddEdge(state, edge.target, edge.min, edge.max); + continue; + } + // Split the character range into the maximal sub-ranges of allowed characters. + int range_start = -1; + for (int c = edge.min; c <= edge.max + 1; ++c) { + if (c <= edge.max && !forbidden_chars[c]) { + if (range_start == -1) { + range_start = c; + } + } else if (range_start != -1) { + new_fsm.AddEdge(state, edge.target, range_start, c - 1); + range_start = -1; + } + } + } + } + new_fsm.SetEdgeAuxData(std::vector(fsm.GetEdgeAuxData())); + return ResultOk(FSMWithStartEnd(new_fsm, fsm_wse.GetStart(), fsm_wse.GetEnds())); +} + +class TrieFSMBuilderImpl { + public: + TrieFSMBuilderImpl() = default; + std::optional Build( + const std::vector& patterns, + const std::vector& excluded_patterns, + std::vector* end_states, + bool allow_overlap, + bool add_back_edges + ); + void AddBackEdges(FSM* fsm, int start, const std::unordered_set& ends); +}; + +std::optional TrieFSMBuilderImpl::Build( + const std::vector& patterns, + const std::vector& excluded_patterns, + std::vector* end_states, + bool allow_overlap, + bool add_back_edges +) { + FSM fsm(1); + int start = 0; + std::unordered_set ends; + + if (end_states) { + end_states->clear(); + } + + for (const auto& pattern : patterns) { + // Check for empty patterns + if (!allow_overlap && pattern.empty()) { + return std::nullopt; + } + + int current_state = 0; + for (const auto& ch : pattern) { + int32_t ch_int32 = static_cast(static_cast(ch)); + int next_state = fsm.GetNextState(current_state, ch_int32); + if (next_state == FSM::kNoNextState) { + next_state = fsm.AddState(); + fsm.AddEdge(current_state, next_state, ch_int32, ch_int32); + } + current_state = next_state; + if (!allow_overlap && ends.count(current_state) > 0) { + return std::nullopt; + } + } + if (!allow_overlap && fsm.GetEdges(current_state).size() > 0) { + return std::nullopt; + } + ends.insert(current_state); + if (end_states) { + end_states->push_back(current_state); + } + } + + std::unordered_set dead_state_set; + + if (add_back_edges) { + // Build trie for excluded patterns. + for (const auto& excluded_pattern : excluded_patterns) { + if (!allow_overlap && excluded_pattern.empty()) { + return std::nullopt; + } + + int current_state = 0; + for (const auto& ch : excluded_pattern) { + int32_t ch_int32 = static_cast(static_cast(ch)); + int next_state = fsm.GetNextState(current_state, ch_int32); + if (next_state == FSM::kNoNextState) { + next_state = fsm.AddState(); + fsm.AddEdge(current_state, next_state, ch_int32, ch_int32); + } + current_state = next_state; + if (!allow_overlap && ends.count(current_state) > 0) { + return std::nullopt; + } + } + if (!allow_overlap && fsm.GetEdges(current_state).size() > 0) { + return std::nullopt; + } + + ends.insert(current_state); + dead_state_set.insert(current_state); + } + + // Add back edges. + AddBackEdges(&fsm, start, ends); + + // Remove the edges to excluded end states. + if (dead_state_set.size() != 0) { + for (int state = 0; state < fsm.NumStates(); state++) { + std::vector& edges = fsm.GetEdges(state); + std::vector new_edges; + for (const auto& edge : edges) { + if (dead_state_set.count(edge.target) == 0) { + new_edges.push_back(edge); + } + } + edges = std::move(new_edges); + } + } + } else if (excluded_patterns.size() > 0) { + XGRAMMAR_LOG(WARNING) << "Excluded patterns are ignored when back edges are not added."; + } + + return FSMWithStartEnd(fsm, start, std::vector(ends.begin(), ends.end())); +} + +void TrieFSMBuilderImpl::AddBackEdges(FSM* fsm, int start, const std::unordered_set& ends) { + // Build an Aho-Corasick automaton by adding back edges. + // When matching on the trie fails at state u on byte b, the matcher must resume from + // the longest proper suffix of u's prefix that is still a path in the trie (the + // failure state), and retry b from there. Falling back only to the start state (or to + // the start state's direct children) loses matches whose start lies inside an + // already-followed branch of another pattern. Example: patterns {"bcd", "abce"} on + // input "abcd" -- after following "abc" of the "abce" branch, 'd' must transition to + // the "bcd" end state via the failure state "bc", not back to the start state. + + int num_states = fsm->NumStates(); + + // Step 1. Record the BFS order of the trie (a tree at this point), so that shallower + // states are always processed first. + std::vector bfs_order; + bfs_order.reserve(num_states); + bfs_order.push_back(start); + for (size_t head = 0; head < bfs_order.size(); head++) { + for (const auto& edge : fsm->GetEdges(bfs_order[head])) { + XGRAMMAR_DCHECK(edge.min == edge.max && edge.min >= 0 && edge.min <= 255); + bfs_order.push_back(edge.target); + } + } + XGRAMMAR_DCHECK(static_cast(bfs_order.size()) == num_states); + + // Step 2. Compute the failure link and the fully resolved transition table with the + // textbook O(num_states * 256) dynamic program: delta[u][b] is the trie child when it + // exists, and delta[fail[u]][b] otherwise -- fail[u] is strictly shallower than u, so + // its row is already final when u is processed in BFS order. + std::vector fail(num_states, start); + std::vector> delta(num_states); + for (auto& row : delta) { + row.fill(FSM::kNoNextState); + } + for (int u = 0; u < num_states; u++) { + for (const auto& edge : fsm->GetEdges(u)) { + delta[u][edge.min] = edge.target; + } + } + for (int u : bfs_order) { + for (int byte = 0; byte < 256; byte++) { + // Entries of deeper states are untouched so far, so a non-empty entry here is + // exactly a trie child of u. + int child = delta[u][byte]; + int fallback = (u == start) ? start : delta[fail[u]][byte]; + if (child == FSM::kNoNextState) { + delta[u][byte] = fallback; + } else { + fail[child] = fallback; + } + } + } + + // Step 3. Overwrite the edges of every non-end state with its resolved row, + // compressing consecutive bytes with the same target into range edges. + for (int u = 0; u < num_states; u++) { + if (u != start && ends.count(u) > 0) { + continue; + } + const auto& row = delta[u]; + std::vector new_edges; + for (int byte = 0; byte < 256;) { + int target = row[byte]; + int range_end = byte; + while (range_end + 1 < 256 && row[range_end + 1] == target) { + range_end++; + } + new_edges.push_back(FSMEdge(byte, range_end, target)); + byte = range_end + 1; + } + fsm->GetEdges(u) = std::move(new_edges); + } +} + +std::optional TrieFSMBuilder::Build( + const std::vector& patterns, + const std::vector& exclude_patterns, + std::vector* end_states, + bool allow_overlap, + bool add_back_edges +) { + return TrieFSMBuilderImpl().Build( + patterns, exclude_patterns, end_states, allow_overlap, add_back_edges + ); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/fsm_builder.h b/third_party/xgrammar/cpp/fsm_builder.h new file mode 100644 index 0000000000..1858668f8c --- /dev/null +++ b/third_party/xgrammar/cpp/fsm_builder.h @@ -0,0 +1,119 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm_builder.h + */ +#ifndef XGRAMMAR_FSM_BUILDER_H_ +#define XGRAMMAR_FSM_BUILDER_H_ + +#include +#include +#include +#include + +#include "fsm.h" +#include "support/utils.h" + +namespace xgrammar { + +class GrammarBuilder; + +/*! \brief Bounds of the packed UTF-8 representation for each encoded length. The packed format + * stores the UTF-8 bytes of one codepoint as (byte0 << 24) | (byte1 << 16) | (byte2 << 8) | byte3, + * left-aligned to the actual length (e.g. a 2-byte character is (byte0 << 8) | byte1). */ +constexpr uint32_t kMax1ByteUnicode = 0x7F; +constexpr uint32_t kMin2BytesUnicode = 0xC080; +constexpr uint32_t kMax2BytesUnicode = 0xDFBF; +constexpr uint32_t kMin3BytesUnicode = 0xE08080; +constexpr uint32_t kMax3BytesUnicode = 0xEFBFBF; +constexpr uint32_t kMin4BytesUnicode = 0xF0808080; +constexpr uint32_t kMax4BytesUnicode = 0xF7BFBFBF; + +/*! \brief Convert a Unicode codepoint to the packed UTF-8 format described above. */ +uint32_t CodepointToPackedUTF8(uint32_t codepoint); + +/*! + * \brief Add FSM edges (with intermediate states for multi-byte characters) from `from` to `to` + * accepting every UTF-8 encoded character in the packed range [min, max]. + */ +void AddPackedUTF8RangeEdges(FSM& fsm, int from, int to, uint32_t min, uint32_t max); + +/*! + * \brief Rewrite every unescaped '.' outside character classes to "[^\n]" unless + * `dot_matches_newline` is true. Used to implement the standard regex dot semantics (and the + * dot-all 's' flag) on top of the regex engine, whose '.' matches every codepoint. + */ +std::string RewriteRegexDots(const std::string& pattern, bool dot_matches_newline); + +/*! + * \brief A builder that converts a regex string to a FSM. + */ +class RegexFSMBuilder { + public: + /*! + * \brief Converts a regex string to a FSM. + * \param regex The regex string. A leading "(?i)" makes the match ASCII case-insensitive. + * \param builder If not null, bounded repetitions whose upper bound exceeds the unroll + * threshold are compiled into a kRepeatRef FSM edge referencing a new rule (with a kRegex + * body holding the repeated sub-pattern) added through this builder, instead of being + * physically unrolled. + * \param rule_hint Name hint for rules created through `builder`. + * \return The FSM with start and end states. + */ + static Result Build( + const std::string& regex, GrammarBuilder* builder = nullptr, const std::string& rule_hint = "" + ); + + /*! + * \brief Converts a regex string to a FSM, then removes the forbidden characters from every + * character transition. The result accepts the intersection of the regex language and the + * set of strings that contain no forbidden character. The result language may be empty. + * \param regex The regex string. + * \param forbidden_chars The forbidden characters. + * \param builder See Build(). + * \param rule_hint See Build(). + * \return The FSM with start and end states. + */ + static Result BuildWithForbiddenChars( + const std::string& regex, + const std::bitset<256>& forbidden_chars, + GrammarBuilder* builder = nullptr, + const std::string& rule_hint = "" + ); + + /*! + * \brief Check whether the regex matches the empty string. Only parses the regex; no FSM is + * built, so this is cheap even for regexes with huge bounded repetitions. + */ + static Result MatchesEmpty(const std::string& regex); +}; + +/*! + * \brief A builder that converts a list of patterns to a trie-based FSM. + */ +class TrieFSMBuilder { + public: + /*! + * \brief Build a trie-based FSM from a list of patterns. + * \param patterns The patterns to be built. + * \param excluded_patterns The patterns to be excluded. + * \param end_states The end states of the FSM. This is the terminal state of each pattern and + * the order follows the order of patterns. + * \param allow_overlap Whether to allow overlap between patterns (one being a prefix of the + * other). It does not allow empty patterns either. If false and there is overlap, will return + * std::nullopt. + * \param add_back_edges Whether to add back edges to the FSM. This complements the trie to an + * Aho-Corasick automaton. + * \return If success, the FSM with start and end states. Otherwise, std::nullopt. + */ + static std::optional Build( + const std::vector& patterns, + const std::vector& excluded_patterns, + std::vector* end_states = nullptr, + bool allow_overlap = true, + bool add_back_edges = false + ); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_FSM_BUILDER_H_ diff --git a/third_party/xgrammar/cpp/grammar.cc b/third_party/xgrammar/cpp/grammar.cc new file mode 100644 index 0000000000..59f4de0963 --- /dev/null +++ b/third_party/xgrammar/cpp/grammar.cc @@ -0,0 +1,359 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar.cc + */ + +#include + +#include +#include +#include +#include + +#include "grammar_functor.h" +#include "grammar_parser.h" +#include "grammar_printer.h" +#include "json_schema_converter.h" +#include "lark_converter.h" +#include "regex_converter.h" +#include "structural_tag.h" +#include "support/json_serializer.h" +#include "support/logging.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/******************* Grammar::Impl *******************/ + +std::size_t MemorySize(const Grammar::Impl& impl) { + /// TODO: Now, we evaluate the memory size of each rule as sizeof(Rule), which counts its + /// string members as sizeof(std::string), with an assumption that the strings are small. + /// This should be improved in the future. + return impl.rules_.size() * sizeof(Grammar::Impl::Rule) + + impl.suffix_stop_infos_.size() * sizeof(Grammar::Impl::SuffixStopInfo) + + MemorySize(impl.grammar_expr_data_) + MemorySize(impl.grammar_expr_indptr_) + + MemorySize(impl.complete_fsm) + MemorySize(impl.per_rule_fsms) + + MemorySize(impl.allow_empty_rule_ids); +} + +/******************* Grammar *******************/ + +std::string Grammar::ToString() const { return GrammarPrinter(*this).ToString(); } + +Grammar Grammar::FromEBNF(const std::string& ebnf_string, const std::string& root_rule_name) { + auto grammar = ParseEBNF(ebnf_string, root_rule_name); + grammar = GrammarNormalizer().Apply(grammar); + return grammar; +} + +Grammar Grammar::FromJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool print_converted_ebnf, + bool any_order +) { + auto grammar = GrammarNormalizer::Apply(JSONSchemaToGrammar( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + )); + if (print_converted_ebnf) { + XGRAMMAR_LOG(INFO) << "Converted EBNF: " << grammar.ToString() << std::endl; + } + return grammar; +} + +Grammar Grammar::FromRegex(const std::string& regex, bool print_converted_ebnf) { + auto ebnf_string = RegexToEBNF(regex); + if (print_converted_ebnf) { + XGRAMMAR_LOG(INFO) << "Converted EBNF: " << ebnf_string << std::endl; + } + return FromEBNF(ebnf_string); +} + +Grammar Grammar::FromLark( + const std::string& lark_string, + const std::optional& tokenizer_info, + const std::vector& named_grammars +) { + return LarkToGrammar(lark_string, tokenizer_info, named_grammars); +} + +std::variant Grammar::FromStructuralTag( + const std::string& structural_tag_json, const std::optional& tokenizer_info +) { + return StructuralTagToGrammar(structural_tag_json, tokenizer_info).ToVariant(); +} + +// Optimized json grammar for the speed of the grammar matcher +const std::string kJSONGrammarString = R"( +root ::= ( + "{" [ \n\r\t]* members_and_embrace | + "[" [ \n\r\t]* elements_or_embrace +) +value_non_str ::= ( + "{" [ \n\r\t]* members_and_embrace | + "[" [ \n\r\t]* elements_or_embrace | + "0" fraction exponent | + [1-9] [0-9]* fraction exponent | + "-" [0-9] fraction exponent | + "-" [1-9] [0-9]* fraction exponent | + "true" | + "false" | + "null" +) (= [ \n\r\t]* member_suffix_suffix) +members_and_embrace ::= ("\"" characters_and_colon [ \n\r\t]* members_suffix | "}") (= [ \n\r\t,}\]]) +members_suffix ::= ( + value_non_str [ \n\r\t]* member_suffix_suffix | + "\"" characters_and_embrace | + "\"" characters_and_comma [ \n\r\t]* "\"" characters_and_colon [ \n\r\t]* members_suffix +) (= [ \n\r\t,}\]]) +member_suffix_suffix ::= ( + "}" | + "," [ \n\r\t]* "\"" characters_and_colon [ \n\r\t]* members_suffix +) (= [ \n\r\t,}\]]) +elements_or_embrace ::= ( + "{" [ \n\r\t]* members_and_embrace elements_rest [ \n\r\t]* "]" | + "[" [ \n\r\t]* elements_or_embrace elements_rest [ \n\r\t]* "]" | + "\"" characters_item elements_rest [ \n\r\t]* "]" | + "0" fraction exponent elements_rest [ \n\r\t]* "]" | + [1-9] [0-9]* fraction exponent elements_rest [ \n\r\t]* "]" | + "-" "0" fraction exponent elements_rest [ \n\r\t]* "]" | + "-" [1-9] [0-9]* fraction exponent elements_rest [ \n\r\t]* "]" | + "true" elements_rest [ \n\r\t]* "]" | + "false" elements_rest [ \n\r\t]* "]" | + "null" elements_rest [ \n\r\t]* "]" | + "]" +) +elements ::= ( + "{" [ \n\r\t]* members_and_embrace elements_rest | + "[" [ \n\r\t]* elements_or_embrace elements_rest | + "\"" characters_item elements_rest | + "0" fraction exponent elements_rest | + [1-9] [0-9]* fraction exponent elements_rest | + "-" [0-9] fraction exponent elements_rest | + "-" [1-9] [0-9]* fraction exponent elements_rest | + "true" elements_rest | + "false" elements_rest | + "null" elements_rest +) +elements_rest ::= ( + "" | + [ \n\r\t]* "," [ \n\r\t]* elements +) +characters_and_colon ::= ( + "\"" [ \n\r\t]* ":" | + [^"\\\x00-\x1F] characters_and_colon | + "\\" escape characters_and_colon +) (=[ \n\r\t]* [\"{[0-9tfn-]) +characters_and_comma ::= ( + "\"" [ \n\r\t]* "," | + [^"\\\x00-\x1F] characters_and_comma | + "\\" escape characters_and_comma +) (=[ \n\r\t]* "\"") +characters_and_embrace ::= ( + "\"" [ \n\r\t]* "}" | + [^"\\\x00-\x1F] characters_and_embrace | + "\\" escape characters_and_embrace +) (=[ \n\r\t]* [},]) +characters_item ::= ( + "\"" | + [^"\\\x00-\x1F] characters_item | + "\\" escape characters_item +) (= [ \n\r\t]* [,\]]) +escape ::= ["\\/bfnrt] | "u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] +fraction ::= "" | "." [0-9] [0-9]* +exponent ::= "" | "e" sign [0-9] [0-9]* | "E" sign [0-9] [0-9]* +sign ::= "" | "+" | "-" +)"; + +Grammar Grammar::BuiltinJSONGrammar() { + static const Grammar grammar = FromEBNF(kJSONGrammarString); + return grammar; +} + +Grammar Grammar::Union(const std::vector& grammars) { + return GrammarUnionFunctor::Apply(grammars); +} + +Grammar Grammar::Concat(const std::vector& grammars) { + return GrammarConcatFunctor::Apply(grammars); +} + +std::ostream& operator<<(std::ostream& os, const Grammar& grammar) { + os << grammar.ToString(); + return os; +} + +std::optional Grammar::Impl::Validate() const { + const int64_t num_rules = rules_.size(); + const int64_t num_exprs = grammar_expr_indptr_.size(); + const int64_t data_size = grammar_expr_data_.size(); + auto rule_ok = [&](int64_t id) { return id >= 0 && id < num_rules; }; + auto expr_ok = [&](int64_t id) { return id >= 0 && id < num_exprs; }; + + // Pass 1: every expr must fit in grammar_expr_data_ before any expr can be read. + for (int64_t expr_id = 0; expr_id < num_exprs; ++expr_id) { + const int64_t start = grammar_expr_indptr_[expr_id]; + if (start < 0 || start + 2 > data_size) { + return "grammar_expr_indptr[" + std::to_string(expr_id) + "] is out of range"; + } + const int64_t type = grammar_expr_data_[start]; + const int64_t len = grammar_expr_data_[start + 1]; + if (len < 0 || start + 2 + len > data_size) { + return "The length of grammar expr " + std::to_string(expr_id) + " is out of range"; + } + if (type < 0 || type > static_cast(GrammarExprType::kSubstring)) { + return "Unknown type of grammar expr " + std::to_string(expr_id); + } + } + + // Pass 2: the ids stored inside each expr must refer to existing rules and exprs. + for (int64_t expr_id = 0; expr_id < num_exprs; ++expr_id) { + const auto expr = GetGrammarExpr(expr_id); + const int64_t size = expr.size(); + bool ok = true; + switch (expr.type) { + case GrammarExprType::kByteString: + case GrammarExprType::kEmptyStr: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + break; + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + // [is_negative, lower0, upper0, ...] + ok = size >= 1 && size % 2 == 1; + break; + case GrammarExprType::kRuleRef: + ok = size == 1 && rule_ok(expr[0]); + break; + case GrammarExprType::kSequence: + case GrammarExprType::kChoices: + ok = std::all_of(expr.begin(), expr.end(), expr_ok); + break; + case GrammarExprType::kTagDispatch: { + // [tag_expr0, rule_id0, ..., loop_after_dispatch, excluded_str_expr_id] + const int64_t extra = TagDispatch::kTagDispatchExtraParameter; + ok = size >= extra && (size - extra) % 2 == 0; + for (int64_t i = 0; ok && i < size - extra; i += 2) { + ok = expr_ok(expr[i]) && rule_ok(expr[i + 1]); + } + // The excluded strings are read as a kChoices expr of byte string exprs. + ok = ok && expr_ok(expr[size - 1]) && + GetGrammarExpr(expr[size - 1]).type == GrammarExprType::kChoices; + break; + } + case GrammarExprType::kRepeat: + ok = size == 3 && rule_ok(expr[0]); + break; + case GrammarExprType::kTokenTagDispatch: { + // [trigger_cnt, (token_id, rule_id) x N, loop_after_dispatch, exclude_cnt, token_id x M] + const int64_t trigger_cnt = size >= 1 ? expr[0] : -1; + ok = trigger_cnt >= 0 && 1 + 2 * trigger_cnt + 2 <= size; + for (int64_t i = 0; ok && i < trigger_cnt; ++i) { + ok = rule_ok(expr[2 + 2 * i]); + } + if (ok) { + const int64_t exclude_cnt = expr[1 + 2 * trigger_cnt + 1]; + ok = exclude_cnt >= 0 && 1 + 2 * trigger_cnt + 2 + exclude_cnt == size; + } + break; + } + case GrammarExprType::kRegex: + ok = size >= 1; + break; + case GrammarExprType::kSubstring: + // [chunk0_len, byte0_0, ..., chunk1_len, ...] + for (int64_t i = 0; ok && i < size;) { + const int64_t chunk_len = expr[i++]; + ok = chunk_len >= 0 && i + chunk_len <= size; + i += chunk_len; + } + break; + } + if (!ok) { + return "Grammar expr " + std::to_string(expr_id) + " is malformed"; + } + } + + for (int64_t rule_id = 0; rule_id < num_rules; ++rule_id) { + const auto& rule = rules_[rule_id]; + if (!expr_ok(rule.body_expr_id) || + (rule.lookahead_assertion_id != -1 && !expr_ok(rule.lookahead_assertion_id))) { + return "Rule " + std::to_string(rule_id) + " refers to a grammar expr out of range"; + } + } + if (!rule_ok(root_rule_id_)) { + return "root_rule_id " + std::to_string(root_rule_id_) + " is out of range"; + } + for (const auto& info : suffix_stop_infos_) { + if (!rule_ok(info.rule_id) || (info.body_rule_id != -1 && !rule_ok(info.body_rule_id)) || + (info.marker_rule_id != -1 && !rule_ok(info.marker_rule_id))) { + return "suffix_stop_infos refers to a rule out of range"; + } + } + if (!std::all_of(allow_empty_rule_ids.begin(), allow_empty_rule_ids.end(), rule_ok)) { + return "allow_empty_rule_ids refers to a rule out of range"; + } + + // The FSMs are internally consistent (see CompactFSM::Impl::Validate); check the rule ids they + // refer to. The parser reads the repeat info of a per-rule FSM edge from complete_fsm, so the + // aux indices of per-rule repeat edges must also be valid for complete_fsm. + auto fsm_rules_ok = [&](const CompactFSM& fsm) { + const auto& aux = fsm.GetEdgeAuxData(); + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if ((edge.IsRuleRef() && !rule_ok(edge.max)) || + (edge.IsRepeatRef() && !rule_ok(aux[edge.max]))) { + return false; + } + } + } + return true; + }; + if (!complete_fsm.IsNull() && !fsm_rules_ok(complete_fsm)) { + return "complete_fsm refers to a rule out of range"; + } + if (optimized && + (complete_fsm.IsNull() || static_cast(per_rule_fsms.size()) != num_rules)) { + return "An optimized grammar must have complete_fsm and one FSM per rule"; + } + for (int64_t rule_id = 0; rule_id < static_cast(per_rule_fsms.size()); ++rule_id) { + if (!per_rule_fsms[rule_id].has_value()) { + if (optimized) { + return "Rule " + std::to_string(rule_id) + " has no FSM in an optimized grammar"; + } + continue; + } + const CompactFSM& fsm = per_rule_fsms[rule_id]->GetFsm().GetFsm(); + if (!fsm_rules_ok(fsm)) { + return "The FSM of rule " + std::to_string(rule_id) + " refers to a rule out of range"; + } + const int64_t complete_aux_size = + complete_fsm.IsNull() ? 0 : complete_fsm.GetEdgeAuxData().size(); + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (edge.IsRepeatRef() && (static_cast(edge.max) + 3 > complete_aux_size || + !rule_ok(complete_fsm.GetEdgeAuxData()[edge.max]))) { + return "The FSM of rule " + std::to_string(rule_id) + + " has a repeat edge that is out of range of complete_fsm"; + } + } + } + } + return std::nullopt; +} + +std::string Grammar::SerializeJSON() const { return AutoSerializeJSON(*this, true); } + +std::variant Grammar::DeserializeJSON(const std::string& json_string) { + Grammar result{NullObj()}; + if (auto err = AutoDeserializeJSON(&result, json_string, true, "Grammar")) { + return err.value(); + } + return result; +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_builder.cc b/third_party/xgrammar/cpp/grammar_builder.cc new file mode 100644 index 0000000000..e66b98dd32 --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_builder.cc @@ -0,0 +1,425 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/grammar_builder.cc + */ + +#include "grammar_builder.h" + +#include +#include +#include +#include +#include +#include + +#include "support/logging.h" + +namespace xgrammar { + +/****************** GrammarBuilder ******************/ + +GrammarBuilder::GrammarBuilder() : grammar_(std::make_shared()) {} + +GrammarBuilder::GrammarBuilder(const Grammar& grammar) + : grammar_(std::make_shared(*grammar.operator->())), + rule_name_map_built_(false) {} + +GrammarBuilder GrammarBuilder::FromMutableGrammar(Grammar* grammar) { + GrammarBuilder builder; + builder.grammar_ = grammar->pimpl_; + builder.rule_name_map_built_ = false; + return builder; +} + +void GrammarBuilder::EnsureRuleNameMap() const { + if (rule_name_map_built_) { + return; + } + for (int32_t i = 0; i < static_cast(grammar_->rules_.size()); ++i) { + rule_name_to_id_[grammar_->rules_[i].name] = i; + } + rule_name_map_built_ = true; +} + +Grammar GrammarBuilder::Get(const std::string& root_rule_name) { + int32_t root_rule_id = GetRuleId(root_rule_name); + XGRAMMAR_CHECK(root_rule_id != -1) + << "The root rule with name \"" << root_rule_name << "\" is not found."; + return Get(root_rule_id); +} + +Grammar GrammarBuilder::Get(int32_t root_rule_id) { + XGRAMMAR_CHECK(root_rule_id >= 0 && root_rule_id < static_cast(grammar_->rules_.size())) + << "The root rule id " << root_rule_id << " is out of bound."; + grammar_->root_rule_id_ = root_rule_id; + return Grammar(grammar_); +} + +int32_t GrammarBuilder::AddGrammarExpr(const GrammarExpr& grammar_expr) { + // Offsets into grammar_expr_data_ are stored as int32. + XGRAMMAR_CHECK( + grammar_->grammar_expr_data_.size() + 2 + grammar_expr.data_len <= + static_cast(std::numeric_limits::max()) + ) << "The grammar is too large: the grammar expr data exceeds 2^31 elements"; + grammar_->grammar_expr_indptr_.push_back(grammar_->grammar_expr_data_.size()); + grammar_->grammar_expr_data_.push_back(static_cast(grammar_expr.type)); + grammar_->grammar_expr_data_.push_back(grammar_expr.data_len); + grammar_->grammar_expr_data_.insert( + grammar_->grammar_expr_data_.end(), + grammar_expr.data, + grammar_expr.data + grammar_expr.data_len + ); + return static_cast(grammar_->grammar_expr_indptr_.size()) - 1; +} + +int32_t GrammarBuilder::AddByteString(const std::vector& bytes) { + return AddGrammarExpr( + {GrammarExprType::kByteString, bytes.data(), static_cast(bytes.size())} + ); +} + +int32_t GrammarBuilder::AddByteString(const std::string& str) { + std::vector bytes; + bytes.reserve(str.size()); + for (char c : str) { + bytes.push_back(static_cast(static_cast(c))); + } + return AddGrammarExpr( + {GrammarExprType::kByteString, bytes.data(), static_cast(bytes.size())} + ); +} + +int32_t GrammarBuilder::AddRegex(const std::string& regex_str, bool json_string) { + std::vector data; + data.reserve(regex_str.size() + 1); + data.push_back(static_cast(json_string)); + for (char c : regex_str) { + data.push_back(static_cast(static_cast(c))); + } + return AddGrammarExpr({GrammarExprType::kRegex, data.data(), static_cast(data.size())}); +} + +int32_t GrammarBuilder::AddSubstring(const std::vector& chunks) { + std::vector data; + for (const std::string& chunk : chunks) { + data.push_back(static_cast(chunk.size())); + for (char c : chunk) { + data.push_back(static_cast(static_cast(c))); + } + } + return AddGrammarExpr( + {GrammarExprType::kSubstring, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddCharacterClass( + const std::vector& elements, bool is_negative +) { + std::vector data; + data.reserve(1 + elements.size() * 2); + data.push_back(static_cast(is_negative)); + for (const auto& range : elements) { + data.push_back(range.lower); + data.push_back(range.upper); + } + return AddGrammarExpr( + {GrammarExprType::kCharacterClass, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddCharacterClassStar( + const std::vector& elements, bool is_negative +) { + std::vector data; + data.reserve(1 + elements.size() * 2); + data.push_back(static_cast(is_negative)); + for (const auto& range : elements) { + data.push_back(range.lower); + data.push_back(range.upper); + } + return AddGrammarExpr( + {GrammarExprType::kCharacterClassStar, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddEmptyStr() { + return AddGrammarExpr({GrammarExprType::kEmptyStr, nullptr, 0}); +} + +int32_t GrammarBuilder::AddTokenSet(const std::vector& token_ids) { + return AddGrammarExpr( + {GrammarExprType::kToken, token_ids.data(), static_cast(token_ids.size())} + ); +} + +int32_t GrammarBuilder::AddExcludeTokenSet(const std::vector& token_ids) { + return AddGrammarExpr( + {GrammarExprType::kExcludeToken, token_ids.data(), static_cast(token_ids.size())} + ); +} + +int32_t GrammarBuilder::AddRuleRef(int32_t rule_id) { + std::vector data; + data.push_back(rule_id); + return AddGrammarExpr({GrammarExprType::kRuleRef, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddSequence(const std::vector& elements) { + return AddGrammarExpr( + {GrammarExprType::kSequence, elements.data(), static_cast(elements.size())} + ); +} + +int32_t GrammarBuilder::AddChoices(const std::vector& choices) { + return AddGrammarExpr( + {GrammarExprType::kChoices, choices.data(), static_cast(choices.size())} + ); +} + +int32_t GrammarBuilder::AddTagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch) { + std::vector data; + data.reserve(tag_dispatch.tag_rule_pairs.size() * 2 + 2); + for (const auto& [tag, rule_id] : tag_dispatch.tag_rule_pairs) { + data.push_back(AddByteString(tag)); + data.push_back(rule_id); + } + data.push_back(static_cast(tag_dispatch.loop_after_dispatch)); + std::vector exclude_str_expr_ids; + for (const auto& exclude_str : tag_dispatch.excludes) { + exclude_str_expr_ids.push_back(AddByteString(exclude_str)); + } + data.push_back(AddChoices(exclude_str_expr_ids)); + return AddGrammarExpr( + {GrammarExprType::kTagDispatch, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch +) { + std::vector data; + data.push_back(static_cast(token_tag_dispatch.trigger_rule_pairs.size())); + for (const auto& [token_id, rule_id] : token_tag_dispatch.trigger_rule_pairs) { + data.push_back(token_id); + data.push_back(rule_id); + } + data.push_back(static_cast(token_tag_dispatch.loop_after_dispatch)); + data.push_back(static_cast(token_tag_dispatch.excludes.size())); + for (auto token_id : token_tag_dispatch.excludes) { + data.push_back(token_id); + } + return AddGrammarExpr( + {GrammarExprType::kTokenTagDispatch, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddRepeat( + int32_t ref_rule_id, int32_t min_repeat_count, int32_t max_repeat_count +) { + std::vector data({ref_rule_id, min_repeat_count, max_repeat_count}); + return AddGrammarExpr({GrammarExprType::kRepeat, data.data(), static_cast(data.size())}); +} + +int32_t GrammarBuilder::AddRepeatFromExpr( + const std::string& cur_rule_name, + int32_t grammar_expr_id, + int32_t min_repeat_count, + int32_t max_repeat_count +) { + const auto& expr = GetGrammarExpr(grammar_expr_id); + int32_t ref_rule_id; + if (expr.type == GrammarExprType::kRuleRef) { + ref_rule_id = expr[0]; + } else { + ref_rule_id = AddRule(GetNewRuleName(cur_rule_name), grammar_expr_id); + } + return AddRepeat(ref_rule_id, min_repeat_count, max_repeat_count); +} + +int32_t GrammarBuilder::NumGrammarExprs() const { return grammar_->NumGrammarExprs(); } + +GrammarBuilder::GrammarExpr GrammarBuilder::GetGrammarExpr(int32_t grammar_expr_id) { + return grammar_->GetGrammarExpr(grammar_expr_id); +} + +int32_t GrammarBuilder::AddRule(const Rule& rule) { + EnsureRuleNameMap(); + int32_t id = static_cast(grammar_->rules_.size()); + grammar_->rules_.push_back(rule); + XGRAMMAR_CHECK(rule_name_to_id_.count(rule.name) == 0); + rule_name_to_id_[rule.name] = id; + return id; +} + +int32_t GrammarBuilder::AddRule(const std::string& name, int32_t body_expr_id) { + return AddRule({name, body_expr_id}); +} + +int32_t GrammarBuilder::AddRuleWithHint(const std::string& name_hint, int32_t body_expr_id) { + return AddRule({GetNewRuleName(name_hint), body_expr_id}); +} + +int32_t GrammarBuilder::NumRules() const { return grammar_->NumRules(); } + +const GrammarBuilder::Rule& GrammarBuilder::GetRule(int32_t rule_id) const { + return grammar_->rules_[rule_id]; +} + +int32_t GrammarBuilder::AddEmptyRule(const std::string& name) { return AddRule({name, -1}); } + +int32_t GrammarBuilder::AddEmptyRuleWithHint(const std::string& name_hint) { + return AddRule({GetNewRuleName(name_hint), -1}); +} + +void GrammarBuilder::UpdateRuleBody(int32_t rule_id, int32_t body_expr_id) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].body_expr_id = body_expr_id; +} + +void GrammarBuilder::UpdateRuleBody(std::string rule_name, int32_t body_expr_id) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateRuleBody(rule_id, body_expr_id); +} + +void GrammarBuilder::UpdateLookaheadAssertion(int32_t rule_id, int32_t lookahead_assertion_id) { + XGRAMMAR_CHECK(rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].lookahead_assertion_id = lookahead_assertion_id; +} + +void GrammarBuilder::UpdateLookaheadExact(int32_t rule_id, bool is_exact) { + XGRAMMAR_CHECK(rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].is_exact_lookahead = is_exact; +} + +void GrammarBuilder::UpdateRuleTemperature(int32_t rule_id, std::optional temperature) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].temperature = temperature; +} + +void GrammarBuilder::UpdateLookaheadAssertion( + std::string rule_name, int32_t lookahead_assertion_id +) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateLookaheadAssertion(rule_id, lookahead_assertion_id); +} + +void GrammarBuilder::UpdateMaxTokens(int32_t rule_id, int32_t max_tokens) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].max_tokens = max_tokens; +} + +void GrammarBuilder::UpdateMaxTokens(std::string rule_name, int32_t max_tokens) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateMaxTokens(rule_id, max_tokens); +} + +void GrammarBuilder::UpdateMaxChars(int32_t rule_id, int32_t max_chars) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].max_chars = max_chars; +} + +void GrammarBuilder::UpdateMaxChars(std::string rule_name, int32_t max_chars) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateMaxChars(rule_id, max_chars); +} + +void GrammarBuilder::UpdateCaptureName(int32_t rule_id, const std::string& capture_name) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].capture_name = capture_name; +} + +void GrammarBuilder::UpdateCaptureName(std::string rule_name, const std::string& capture_name) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateCaptureName(rule_id, capture_name); +} + +void GrammarBuilder::UpdateSuffixStopInfo(int32_t rule_id, const SuffixStopInfo& info) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + XGRAMMAR_CHECK(info.hidden_suffix_bytes >= 0 && info.hidden_stop_bytes >= 0) + << "The number of hidden suffix/stop bytes must be non-negative."; + int32_t num_rules = static_cast(grammar_->rules_.size()); + XGRAMMAR_CHECK( + (info.body_rule_id == -1 && info.marker_rule_id == -1) || + (info.body_rule_id >= 0 && info.body_rule_id < num_rules && info.marker_rule_id >= 0 && + info.marker_rule_id < num_rules) + ) << "Capture-hidden helper rule ids must both be -1 or valid rule ids."; + + auto it = std::lower_bound( + grammar_->suffix_stop_infos_.begin(), + grammar_->suffix_stop_infos_.end(), + rule_id, + [](const SuffixStopInfo& existing, int32_t id) { return existing.rule_id < id; } + ); + if (info.IsEmpty()) { + if (it != grammar_->suffix_stop_infos_.end() && it->rule_id == rule_id) { + grammar_->suffix_stop_infos_.erase(it); + } + return; + } + SuffixStopInfo updated = info; + updated.rule_id = rule_id; + if (it != grammar_->suffix_stop_infos_.end() && it->rule_id == rule_id) { + *it = std::move(updated); + } else { + grammar_->suffix_stop_infos_.insert(it, std::move(updated)); + } +} + +void GrammarBuilder::UpdateSuffixStopInfo(std::string rule_name, const SuffixStopInfo& info) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateSuffixStopInfo(rule_id, info); +} + +void GrammarBuilder::UpdateLazy(int32_t rule_id, bool is_lazy) { + XGRAMMAR_CHECK(rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].is_lazy = is_lazy; +} + +void GrammarBuilder::UpdateLazy(std::string rule_name, bool is_lazy) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateLazy(rule_id, is_lazy); +} + +std::string GrammarBuilder::GetNewRuleName(const std::string& name_hint) { + EnsureRuleNameMap(); + if (rule_name_to_id_.count(name_hint) == 0) { + return name_hint; + } + int* cnt = &next_cnt_per_hint_[name_hint]; + if (*cnt == 0) { + *cnt = 1; + } + while (rule_name_to_id_.count(name_hint + "_" + std::to_string(*cnt)) != 0) { + ++(*cnt); + } + return name_hint + "_" + std::to_string(*cnt); +} + +int32_t GrammarBuilder::GetRuleId(const std::string& name) const { + EnsureRuleNameMap(); + auto it = rule_name_to_id_.find(name); + if (it == rule_name_to_id_.end()) { + return -1; + } else { + return it->second; + } +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_builder.h b/third_party/xgrammar/cpp/grammar_builder.h new file mode 100644 index 0000000000..7209a567ba --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_builder.h @@ -0,0 +1,292 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_builder.h + * \brief The header for the building the BNF AST. + */ + +#ifndef XGRAMMAR_GRAMMAR_BUILDER_H_ +#define XGRAMMAR_GRAMMAR_BUILDER_H_ + +#include + +#include +#include +#include +#include +#include + +#include "grammar_impl.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief Helper class to build a BNF grammar. + */ +class GrammarBuilder { + public: + using Rule = Grammar::Impl::Rule; + using SuffixStopInfo = Grammar::Impl::SuffixStopInfo; + using GrammarExprType = Grammar::Impl::GrammarExprType; + using GrammarExpr = Grammar::Impl::GrammarExpr; + + /*! \brief One element of a character class, containing a lower and a upper bound. Both bounds are + * inclusive. + */ + struct CharacterClassElement { + int32_t lower; + int32_t upper; + }; + + /*! \brief Default constructor. Creates a new grammar object. */ + GrammarBuilder(); + + /*! \brief Constructor. Creates a new grammar object from an existing grammar. */ + GrammarBuilder(const Grammar& grammar); + + /*! + * \brief Create a builder bound to an existing grammar without copying it. Unlike the copy + * constructor above, appended exprs and rule updates are written directly into *grammar. The + * grammar must outlive the builder. + */ + static GrammarBuilder FromMutableGrammar(Grammar* grammar); + + /*! + * \brief Get the result grammar. This function will also set the root rule to the rule with the + * specified name. The rule should be already added to the grammar. + * \param root_rule_name The name of the root rule. Default is "root". + */ + Grammar Get(const std::string& root_rule_name = "root"); + + /*! + * \brief Get the result grammar. This function will also set the root rule to the rule with + * the specified id. The rule should be already added to the grammar. + * \param root_rule_id The id of the root rule. + */ + Grammar Get(int32_t root_rule_id); + + /****************** GrammarExpr handling ******************/ + + /*! \brief Add a grammar_expr and return the grammar_expr id. */ + int32_t AddGrammarExpr(const GrammarExpr& grammar_expr); + + /*! + * \brief Add a GrammarExpr for string stored in bytes. + * \param bytes A vector of int32_t, each representing a byte (0~255) in the string. + * The string is stored in int32 vector to match the storage format of the grammar. + */ + int32_t AddByteString(const std::vector& bytes); + + /*! + * \brief Add a GrammarExpr for string stored in bytes. + * \param str The string to be added. + */ + int32_t AddByteString(const std::string& str); + + /*! + * \brief Add a GrammarExpr for a regex. The pattern is stored as-is and compiled into an + * automaton by GrammarFSMBuilder. + * \param regex_str The regex pattern string. + * \param json_string Whether the regex matches the body of a JSON string literal. If true, + * the characters that must be escaped in a JSON string (the control characters, '"' and + * '\\') are excluded from every character match of the compiled automaton. + */ + int32_t AddRegex(const std::string& regex_str, bool json_string = false); + + /*! + * \brief Add a GrammarExpr for a substring expression, which matches every contiguous + * subsequence of the chunk list (including the empty one). The chunks are stored as-is and + * compiled into an automaton by GrammarFSMBuilder. + * \param chunks The list of byte string chunks. Chunks may be empty or repeated. + */ + int32_t AddSubstring(const std::vector& chunks); + + /*! + * \brief Add a GrammarExpr for a character class. + * \param elements A vector of CharacterClassElement, each containing a lower and a upper bound. + * \param is_negative Whether the character class is negated. + */ + int32_t AddCharacterClass( + const std::vector& elements, bool is_negative = false + ); + + /*! + * \brief Add a GrammarExpr for a star quantifier of a character class. + * \param elements A vector of CharacterClassElement, each containing a lower and a upper bound. + * \param is_negative Whether the character class is negated. + */ + int32_t AddCharacterClassStar( + const std::vector& elements, bool is_negative = false + ); + + /*! \brief Add a GrammarExpr for empty string.*/ + int32_t AddEmptyStr(); + + /*! \brief Add a GrammarExpr for kToken (token-level matching). */ + int32_t AddTokenSet(const std::vector& token_ids); + + /*! \brief Add a GrammarExpr for kExcludeToken (excluded token-level matching). */ + int32_t AddExcludeTokenSet(const std::vector& token_ids); + + /*! \brief Add a GrammarExpr for rule reference.*/ + int32_t AddRuleRef(int32_t rule_id); + + /*! \brief Add a GrammarExpr for GrammarExpr sequence.*/ + int32_t AddSequence(const std::vector& elements); + + /*! \brief Add a GrammarExpr for GrammarExpr choices.*/ + int32_t AddChoices(const std::vector& choices); + + /*! + * \brief Add a GrammarExpr for tag dispatch. + * \param tag_dispatch_list A list of pairs of tag_expr_id and rule_id. + */ + int32_t AddTagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); + + /*! \brief Encode a TokenTagDispatch struct into a kTokenTagDispatch expr. */ + int32_t AddTokenTagDispatch(const Grammar::Impl::TokenTagDispatch& token_tag_dispatch); + + int32_t AddRepeat(int32_t ref_rule_id, int32_t min_repeat_count, int32_t max_repeat_count); + + /*! + * \brief Add a repeat GrammarExpr from an arbitrary grammar expression. If the expression is + * not a rule reference, a new rule is created to wrap it. + * \param cur_rule_name Name hint for generated rules. + * \param grammar_expr_id The expression to repeat. + * \param min_repeat_count Minimum repeat count (inclusive). + * \param max_repeat_count Maximum repeat count (inclusive), or -1 for unbounded. + */ + int32_t AddRepeatFromExpr( + const std::string& cur_rule_name, + int32_t grammar_expr_id, + int32_t min_repeat_count, + int32_t max_repeat_count + ); + + /*! \brief Get the number of grammar_exprs. */ + int32_t NumGrammarExprs() const; + + /*! \brief Get the grammar_expr with the given id. */ + GrammarExpr GetGrammarExpr(int32_t grammar_expr_id); + + /****************** Rule handling ******************/ + + /*! \brief Add a rule and return the rule id. */ + int32_t AddRule(const Rule& rule); + + int32_t AddRule(const std::string& name, int32_t body_expr_id); + + int32_t AddRuleWithHint(const std::string& name_hint, int32_t body_expr_id); + + int32_t NumRules() const; + + /*! \brief Get the rule with the given id. */ + const Rule& GetRule(int32_t rule_id) const; + + /*! + * \brief Add an rule without body, and return the rule id. The rule body should be set later + * with GrammarBuilder::UpdateRuleBody. This method is useful for cases where the rule id is + * required to build the rule body. + * \sa GrammarBuilder::UpdateRuleBody + */ + int32_t AddEmptyRule(const std::string& name); + + int32_t AddEmptyRuleWithHint(const std::string& name_hint); + + /*! + * \brief Update the rule body of the given rule, specified by rule id. Can be used to set the + * rule body of a rule inserted by GrammarBuilder::AddEmptyRule. + */ + void UpdateRuleBody(int32_t rule_id, int32_t body_expr_id); + + /*! + * \brief Update the rule body of the given rule, specified by rule name. Can be used to set the + * rule body of a rule inserted by GrammarBuilder::AddEmptyRule. + */ + void UpdateRuleBody(std::string rule_name, int32_t body_expr_id); + + /*! + * \brief Add a lookahead assertion to a rule referred by the given rule_id. The lookahead + * assertion should be a sequence GrammarExpr id. An id of -1 means no lookahead assertion. + */ + void UpdateLookaheadAssertion(int32_t rule_id, int32_t lookahead_assertion_id); + + void UpdateLookaheadExact(int32_t rule_id, bool is_exact = true); + + /*! \brief Set the sampling temperature associated with a rule. */ + void UpdateRuleTemperature(int32_t rule_id, std::optional temperature); + + /*! + * \brief Add a lookahead assertion to a rule referred by the given name. The lookahead + * assertion should be a sequence GrammarExpr id. An id of -1 means no lookahead assertion. + */ + void UpdateLookaheadAssertion(std::string rule_name, int32_t lookahead_assertion_id); + + /*! \brief Update the token budget of the rule referred by the given rule_id. -1 means none. */ + void UpdateMaxTokens(int32_t rule_id, int32_t max_tokens); + + /*! \brief Update the token budget of the rule referred by the given name. -1 means none. */ + void UpdateMaxTokens(std::string rule_name, int32_t max_tokens); + + /*! \brief Update the character budget of the rule referred by the given rule_id. -1 means none. + */ + void UpdateMaxChars(int32_t rule_id, int32_t max_chars); + + /*! \brief Update the character budget of the rule referred by the given name. -1 means none. */ + void UpdateMaxChars(std::string rule_name, int32_t max_chars); + + /*! + * \brief Update the capture group name of the rule referred by the given rule_id. An empty + * string means no capture. + */ + void UpdateCaptureName(int32_t rule_id, const std::string& capture_name); + + /*! + * \brief Update the capture group name of the rule referred by the given name. An empty string + * means no capture. + */ + void UpdateCaptureName(std::string rule_name, const std::string& capture_name); + + /*! \brief Set or clear the sparse suffix/stop metadata for a rule. */ + void UpdateSuffixStopInfo(int32_t rule_id, const SuffixStopInfo& info); + + /*! \brief Set or clear sparse suffix/stop metadata on the rule referred by name. */ + void UpdateSuffixStopInfo(std::string rule_name, const SuffixStopInfo& info); + + /*! \brief Set whether the rule referred by the given rule_id is lazy (committed-shortest). */ + void UpdateLazy(int32_t rule_id, bool is_lazy); + + /*! \brief Set whether the rule referred by the given name is lazy (committed-shortest). */ + void UpdateLazy(std::string rule_name, bool is_lazy); + + /*! + * \brief Find a name for a new rule starting with the given name hint. Some integer suffix (_1, + * _2, ...) may be added to avoid name conflict. + */ + std::string GetNewRuleName(const std::string& name_hint); + + /*! + * \brief Get the rule id of the rule with the given name. Return -1 if not found. + */ + int32_t GetRuleId(const std::string& name) const; + + private: + /*! + * \brief Build rule_name_to_id_ from the existing rules if it has not been built yet. + * Builders bound to an existing grammar (copy constructor, FromMutableGrammar) defer building + * the map to the first name-based operation, so id-only rewriting pays no name hashing cost. + */ + void EnsureRuleNameMap() const; + + // Mutable pointer to the grammar object. + std::shared_ptr grammar_; + // Map from rule name to rule id. Built lazily; see EnsureRuleNameMap. + mutable std::unordered_map rule_name_to_id_; + mutable bool rule_name_map_built_ = true; + // Cache of next suffix index per name_hint for GetNewRuleName. + std::unordered_map next_cnt_per_hint_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_BUILDER_H_ diff --git a/third_party/xgrammar/cpp/grammar_compiler.cc b/third_party/xgrammar/cpp/grammar_compiler.cc new file mode 100644 index 0000000000..dcd6264de1 --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_compiler.cc @@ -0,0 +1,1725 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/compiler.cc + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "earley_parser.h" +#include "fsm.h" +#include "grammar_functor.h" +#include "grammar_impl.h" +#include "support/dynamic_bitset.h" +#include "support/int_set.h" +#include "support/logging.h" +#include "support/thread_pool.h" +#include "support/thread_safe_cache.h" +#include "support/utils.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/grammar.h" +#include "xgrammar/tokenizer_info.h" + +namespace xgrammar { + +/************** AdaptiveTokenMaskCache Generator **************/ + +/*! \brief The concrete implementation of GrammarMatcherNode. */ +class GrammarMatcherForTokenMaskCache : public EarleyParser { + public: + GrammarMatcherForTokenMaskCache( + const Grammar& grammar, + const ParserState& init_state, + const std::unordered_map& + tag_dispatch_rule_id_to_second_slicing_bitset, + const TokenizerInfo& tokenizer_info, + std::optional& rule_level_cache + ) + : EarleyParser(grammar, init_state), + init_rule_id_(init_state.rule_id), + initial_state_(init_state), + tag_dispatch_rule_id_to_second_slicing_bitset_(tag_dispatch_rule_id_to_second_slicing_bitset + ), + tokenizer_info_(tokenizer_info), + rule_level_cache_(rule_level_cache) {} + /*! + * \brief Get the adaptive token mask for the given ParserState. + * \param is_root_rule Whether to consider the parent rule. If false, there will be + * no uncertain tokens. Useful for the root rule. + */ + AdaptiveTokenMask GetAdaptiveTokenMask(bool is_root_rule); + + /*! + * \brief Get the token mask for the given ParserState. + * \param first_char_mask The first character mask. + * \param is_root_rule Whether to consider the parent rule. If false, there will be + * no uncertain tokens. Useful for the root rule. + * \returns True if the rejected indices are filled as usual, False otherwise. + * It's used to determine which construction function will be used. + */ + bool GetTokenMaskWithFirstCharacterCheck( + const std::bitset<256>& first_char_mask, + bool is_root_rule, + const std::vector& token_edge_accepted + ); + + /*! + * \brief Adapt the cache with lookahead assertion. + * \param cache The adaptive token mask to be adapted. + * \param is_root_rule Whether to consider the parent rule. + */ + void AdaptCacheWithLookahead(AdaptiveTokenMask* cache, bool is_root_rule); + + private: + /*! \brief Check if a token can pass the lookahead assertion. */ + std::pair IsTokenPassLookaheadAssertion( + const std::string& token, const std::vector& can_reach_end_stack + ); + + /*! + * \brief Check if speculative calculation will be applied. + * \return first: whether speculative calculation is applicable. + * \return second: part of the first character mask, + * which can be used in speculative calculation. + */ + std::pair> GetSpeculativeCalculation(); + + /*! + * \brief Get the first character mask. + * \param first_character_mask the bitset to store the first character mask. + */ + void GetFirstCharacterMask(std::bitset<256>& first_character_mask); + + /*! + * \brief Compute sorted vocab indices accepted by token edges at the current FSM state. + * Token(ids) edges accept listed token IDs. + * ExcludeToken(ids) edges accept all tokens except listed IDs. + * \return Sorted, deduplicated vector of accepted sorted vocab indices. + */ + const std::vector& GetTokenEdgeAcceptedIndices(); + + // The id of the initial rule. + int32_t init_rule_id_; + + // The initial state of the parser. + ParserState initial_state_; + + /*! + * \brief This is a mapping from TagDispatch rule id to the bitset used for second slicing. + * \note If a rule is a TagDispatch rule, then there will be an AC automaton for its triggers. + * Which means that it can accept a lot of tokens. However, it will be slow to check a lot of + * tokens. The DynamicBitset here is used to do a second slicing: if a token's substr(1, n - 1) + * can be accepted by the start state of the AC automaton, then it will be True in the bitset. + * When we check a token, we first check if its first character can transit to the start state. + * If yes, then we check if it is in the bitset. If yes, then we accept it directly. + */ + const std::unordered_map& tag_dispatch_rule_id_to_second_slicing_bitset_; + + const TokenizerInfo& tokenizer_info_; + + std::optional rule_level_cache_; + + // Temporary data for GetAdaptiveTokenMask. + std::vector tmp_accepted_indices_; + std::vector tmp_rejected_indices_; + std::vector tmp_uncertain_indices_; + std::vector tmp_rejected_by_lookahead_indices_; + std::vector tmp_accepted_by_lookahead_indices_; + std::vector tmp_can_reach_end_stack_; + std::vector tmp_can_reach_end_prefix_or_stack_; + // Temporary data for GetTokenEdgeAcceptedIndices. + std::vector tmp_token_edge_accepted_; + std::vector tmp_token_edge_excluded_; +}; + +void GrammarMatcherForTokenMaskCache::AdaptCacheWithLookahead( + AdaptiveTokenMask* cache_ptr, bool is_root_rule +) { + AdaptiveTokenMask& cache = *cache_ptr; + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + const auto& subtree_nodes_range = tokenizer_info_.GetTrieSubtreeNodesRange(); + const std::string* prev_token = nullptr; + bool is_exact_lookahead = grammar_->GetRule(init_rule_id_).is_exact_lookahead; + int prev_matched_size = 0; + int last_rejected_range = 0; + int last_uncertain_range = 0; + if (is_root_rule) { + tmp_rejected_indices_ = cache.uncertain_indices; + } else { + const auto& lookahead_id = grammar_->GetRule(init_rule_id_).lookahead_assertion_id; + if (lookahead_id == -1) { + return; + } + for (const auto& uncertain_index : cache.uncertain_indices) { + const auto& token = sorted_decoded_vocab[uncertain_index].second; + // Many tokens may contain the same prefix, so we will avoid unnecessary matching + // by finding the longest common prefix with the previous token. + bool accepted = true; + if (uncertain_index < last_rejected_range) { + tmp_rejected_indices_.push_back(uncertain_index); + continue; + } + if (uncertain_index < last_uncertain_range) { + // This token is already marked as uncertain. + continue; + } + if (prev_token != nullptr) { + int lcp_len = + std::mismatch(token.begin(), token.end(), prev_token->begin(), prev_token->end()) + .first - + token.begin(); + if (lcp_len > prev_matched_size) { + // Case 1. The common prefix is rejected by the matcher in the last token. Reject + // directly. + accepted = false; + } else if (lcp_len < prev_matched_size) { + // Case 2. The common prefix is shorter than the previous matched size. Rollback + // the non-common part. + PopLastStates(prev_matched_size - lcp_len); + tmp_can_reach_end_stack_.erase( + tmp_can_reach_end_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_stack_.end() + ); + tmp_can_reach_end_prefix_or_stack_.erase( + tmp_can_reach_end_prefix_or_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_prefix_or_stack_.end() + ); + } + prev_matched_size = std::min(prev_matched_size, lcp_len); + } + + prev_token = &token; + + if (accepted) { + // Accept the rest chars one by one. + for (int j = prev_matched_size; j < static_cast(token.size()); ++j) { + if (!Advance(token[j])) { + accepted = false; + break; + } + tmp_can_reach_end_stack_.push_back(IsCompleted()); + tmp_can_reach_end_prefix_or_stack_.push_back( + tmp_can_reach_end_stack_.back() || tmp_can_reach_end_prefix_or_stack_.back() + ); + prev_matched_size = j + 1; + } + } + + XGRAMMAR_DCHECK(!tmp_can_reach_end_prefix_or_stack_.empty()); + bool can_reach_end = tmp_can_reach_end_prefix_or_stack_.back(); + + XGRAMMAR_DCHECK(!accepted) << "All the tokens are at least uncertain!"; + if (can_reach_end && prev_matched_size > 0) { + auto [lookahead_accepted, lookahead_completed] = + IsTokenPassLookaheadAssertion(token, tmp_can_reach_end_stack_); + if ((!is_root_rule) && lookahead_accepted) { + if (lookahead_completed || !is_exact_lookahead) { + tmp_uncertain_indices_.push_back(uncertain_index); + } else { + tmp_accepted_indices_.push_back(uncertain_index); + } + } else { + tmp_rejected_indices_.push_back(uncertain_index); + last_rejected_range = subtree_nodes_range[uncertain_index]; + } + } else { + tmp_rejected_indices_.push_back(uncertain_index); + last_rejected_range = subtree_nodes_range[uncertain_index]; + } + } + } + + // This strategy ensures the consistency of the cache storage type in most cases. + // However, in this case, the storage type is inconsistent: + // 1. The original cache is accepted_indices, and rejected_indices is also small. + // After adapting with lookahead, |accepted_indices| + |accepted_by_lookahead_indices| > + // |rejected_indices| + |rejected_by_lookahead_indices|, and |rejected_indices| + + // |rejected_by_lookahead_indices| < AdaptiveTokenMask::USE_BITSET_THRESHOLD. In this case, it + // should be kRejected, but ignored. + // 2. The original cache is rejected_indices, and accepted_indices is also small. + // After adapting with lookahead, |accepted_indices| + |accepted_by_lookahead_indices| < + // |rejected_indices| + |rejected_by_lookahead_indices|, and |accepted_indices| + + // |accepted_by_lookahead_indices| < AdaptiveTokenMask::USE_BITSET_THRESHOLD. In this case, it + // should be kAccepted, but ignored. These two cases are very rare in practice, and the impact is + // very limited, so we ignore them for simplicity. + cache.uncertain_indices = tmp_uncertain_indices_; + switch (cache.store_type) { + case AdaptiveTokenMask::StoreType::kAccepted: { + if (cache.accepted_indices.size() + tmp_accepted_indices_.size() < + AdaptiveTokenMask::USE_BITSET_THRESHOLD) { + IntsetUnion(&cache.accepted_indices, tmp_accepted_indices_); + break; + } + // Transform to bitset. + cache.store_type = AdaptiveTokenMask::StoreType::kAcceptedBitset; + cache.accepted_bitset = DynamicBitset(tokenizer_info_.GetVocabSize()); + for (const auto& accepted_index : cache.accepted_indices) { + cache.accepted_bitset.Set(sorted_decoded_vocab[accepted_index].first); + } + for (const auto& accepted_index : tmp_accepted_indices_) { + cache.accepted_bitset.Set(sorted_decoded_vocab[accepted_index].first); + } + cache.accepted_indices.clear(); + break; + } + case AdaptiveTokenMask::StoreType::kRejected: { + if (cache.rejected_indices.size() + tmp_rejected_indices_.size() < + AdaptiveTokenMask::USE_BITSET_THRESHOLD) { + IntsetUnion(&cache.rejected_indices, tmp_rejected_indices_); + break; + } + // Transform to bitset. + cache.store_type = AdaptiveTokenMask::StoreType::kAcceptedBitset; + cache.accepted_bitset = DynamicBitset(tokenizer_info_.GetVocabSize()); + cache.accepted_bitset.Set(); + for (const auto& special_index : tokenizer_info_.GetSpecialTokenIds()) { + cache.accepted_bitset.Reset(special_index); + } + for (const auto& uncertain_index : cache.uncertain_indices) { + cache.accepted_bitset.Reset(sorted_decoded_vocab[uncertain_index].first); + } + for (const auto& rejected_index : cache.rejected_indices) { + cache.accepted_bitset.Reset(sorted_decoded_vocab[rejected_index].first); + } + for (const auto& rejected_index : tmp_rejected_indices_) { + cache.accepted_bitset.Reset(sorted_decoded_vocab[rejected_index].first); + } + cache.rejected_indices.clear(); + break; + } + case AdaptiveTokenMask::StoreType::kAcceptedBitset: { + for (const auto& accepted_index : tmp_accepted_indices_) { + cache.accepted_bitset.Set(sorted_decoded_vocab[accepted_index].first); + } + break; + } + } +} + +std::pair GrammarMatcherForTokenMaskCache::IsTokenPassLookaheadAssertion( + const std::string& token, const std::vector& can_reach_end_stack +) { + bool accepted = true; + bool can_reach_end = true; + auto lookahead_assertion_id = grammar_->GetRule(init_rule_id_).lookahead_assertion_id; + if (lookahead_assertion_id == -1) { + return {accepted, can_reach_end}; + } + auto lookahead_state = + ParserState(/*rule_id*/ -1, lookahead_assertion_id, 0, ParserState::kNoPrevInputPos, 0); + PushStateAndExpand(lookahead_state); + int token_len = token.size(); + if (IsCompleted()) { + // If the lookahead assertion is already completed, we can accept the token. + PopLastStates(1); + return {accepted, can_reach_end}; + } + + // Find all positions that can come to and end. Then check if the suffix from that position + // can be accepted by the lookahead assertion. + for (int i = static_cast(can_reach_end_stack.size()) - 1; i >= 0; --i) { + if (!can_reach_end_stack[i]) { + continue; + } + int last_accept_pos = i - 1; + for (int pos = i; pos < token_len; ++pos) { + if (!Advance(token[pos])) { + break; + } + last_accept_pos = pos; + // Case 1. The whole rule is finished. + if (IsCompleted()) { + // accepted chars: pos - i + 1 + // we need to rollback the pushed initial state as well + PopLastStates(pos - i + 2); + return {accepted, can_reach_end}; + } + } + // Case 2. The whole token is accepted + if (last_accept_pos == token_len - 1) { + PopLastStates(last_accept_pos - i + 2); + can_reach_end = false; + return {accepted, can_reach_end}; + } + // Case 3. The token is not accepted. Check the next position. + PopLastStates(last_accept_pos - i + 1); + } + + PopLastStates(1); + can_reach_end = false; + accepted = false; + return {accepted, can_reach_end}; +} + +// Comparator for std::pair based on the string value. +class IntStringPairComparator { + public: + bool operator()( + const std::pair& lhs, const std::pair& rhs + ) const { + return lhs.second < rhs.second; + } +}; + +int GetPossibleTokenIntervals( + const std::vector>& sorted_decoded_vocab, + const std::bitset<256>& first_char_mask, + std::vector>& possible_intervals +) { + int possible_token_num = 0; + int matched_size = 0; + int last_interval_end = -1; + for (int32_t i = 0; i < 256; i++) { + if (first_char_mask[i]) { + if (last_interval_end == -1) { + last_interval_end = i; + } + } else { + if (last_interval_end != -1) { + int32_t interval_left_end = + std::lower_bound( + sorted_decoded_vocab.begin() + matched_size, + sorted_decoded_vocab.end(), + std::make_pair(0, std::string(1, static_cast(last_interval_end))), + IntStringPairComparator() + ) - + sorted_decoded_vocab.begin(); + int32_t interval_right_end = std::lower_bound( + sorted_decoded_vocab.begin() + interval_left_end, + sorted_decoded_vocab.end(), + std::make_pair(0, std::string(1, static_cast(i))), + IntStringPairComparator() + ) - + sorted_decoded_vocab.begin(); + possible_intervals.emplace_back(interval_left_end, interval_right_end); + possible_token_num += interval_right_end - interval_left_end; + last_interval_end = -1; + matched_size = interval_right_end; + } + } + } + + if (last_interval_end != -1) { + // If the last interval is not closed, we need to close it. + int32_t interval_left_end = + std::lower_bound( + sorted_decoded_vocab.begin() + matched_size, + sorted_decoded_vocab.end(), + std::make_pair(0, std::string(1, static_cast(last_interval_end))), + IntStringPairComparator() + ) - + sorted_decoded_vocab.begin(); + possible_intervals.emplace_back(interval_left_end, sorted_decoded_vocab.size()); + possible_token_num += sorted_decoded_vocab.size() - interval_left_end; + } + return possible_token_num; +} + +std::pair> GrammarMatcherForTokenMaskCache::GetSpeculativeCalculation() { + using GrammarExprType = Grammar::Impl::GrammarExprType; + // If the initial rule is a tag dispatch, we will check if it can achieve its initial state. + const auto& rule = grammar_->GetRule(init_rule_id_); + if (rule.is_lazy) { + // The fast path assumes greedy self-loop extension is always legal, which does not hold for + // committed-shortest rules; they must go through the full per-token simulation. + return {false, std::bitset<256>()}; + } + const auto& rule_body = grammar_->GetGrammarExpr(rule.body_expr_id); + if (rule_body.type == GrammarExprType::kTagDispatch) { + std::bitset<256> speculative_mask; + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + for (const auto& edge : fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id)) { + if (edge.target != fsm.GetFsm().GetStart()) { + continue; + } + if (!edge.IsCharRange()) { + continue; + } + for (int32_t ch = edge.min; ch <= edge.max; ++ch) { + speculative_mask.set(ch); + } + } + return {true, speculative_mask}; + } + + // Check if the initial state is self-recursive-like via FSM. + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + bool can_be_applied = false; + std::bitset<256> speculative_mask; + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + XGRAMMAR_DCHECK(initial_state_.element_id < fsm.GetFsm().NumStates()) + << "Initial State's element id cannot exceed the whole FSM's number of states."; + for (const auto& edge : fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id)) { + if (edge.IsCharRange()) { + // Case A: The edge is towards itself. + if (edge.target == initial_state_.element_id) { + can_be_applied = true; + for (int ch = edge.min; ch <= edge.max; ++ch) { + speculative_mask.set(ch); + } + continue; + } + + // Case B: The state is the start state, and there's an edge to another state, + // which calls the fsm itself. + if (fsm.GetFsm().GetStart() == initial_state_.element_id) { + for (const auto& next_edge : fsm.GetFsm().GetFsm().GetEdges(edge.target)) { + if ((next_edge.IsRuleRef() && next_edge.GetRefRuleId() == init_rule_id_) || + (next_edge.IsRepeatRef() && + fsm.GetFsm().GetFsm().GetRepeatEdgeInfo(next_edge.GetAuxIndex()).RuleId() == + init_rule_id_)) { + can_be_applied = true; + for (int ch = edge.min; ch <= edge.max; ++ch) { + speculative_mask.set(ch); + } + break; + } + } + } + } + } + return {can_be_applied, speculative_mask}; +} + +bool GrammarMatcherForTokenMaskCache::GetTokenMaskWithFirstCharacterCheck( + const std::bitset<256>& first_char_mask, + bool is_root_rule, + const std::vector& token_edge_accepted +) { + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + const auto& subtree_nodes_range = tokenizer_info_.GetTrieSubtreeNodesRange(); + // the pair (a, b) means [a, b). Intialize the possible intervals. + std::vector> possible_intervals; + int possible_token_num = + GetPossibleTokenIntervals(sorted_decoded_vocab, first_char_mask, possible_intervals); + + // Check if the type of the mask can be rejected. + tmp_accepted_indices_.reserve(possible_token_num); + bool fill_reject_indices = + (sorted_decoded_vocab.size() - possible_token_num) < AdaptiveTokenMask::USE_BITSET_THRESHOLD; + + XGRAMMAR_DCHECK(possible_intervals.size() > 0) + << "There should be at least one possible interval for the first character mask."; + + if (possible_intervals[0].first != 0 && fill_reject_indices) { + for (int i = 0; i < possible_intervals[0].first; ++i) { + tmp_rejected_indices_.push_back(i); + } + } + + XGRAMMAR_DCHECK(init_rule_id_ != -1 && grammar_->per_rule_fsms[init_rule_id_].has_value()); + auto [speculative_calculation, speculative_mask] = GetSpeculativeCalculation(); + if (has_char_budget_rules_) { + speculative_calculation = false; + } + + int prev_matched_size = 0; + int last_rejected_range = 0; + const bool& is_exact_lookahead = grammar_->GetRule(init_rule_id_).is_exact_lookahead; + std::optional definite_accepted_bitset = std::nullopt; + const bool is_tag_dispatch_rule = + grammar_->GetGrammarExpr(grammar_->GetRule(init_rule_id_).body_expr_id).type == + Grammar::Impl::GrammarExprType::kTagDispatch; + if (is_tag_dispatch_rule) { + XGRAMMAR_DCHECK(tag_dispatch_rule_id_to_second_slicing_bitset_.count(init_rule_id_) > 0); + definite_accepted_bitset = &tag_dispatch_rule_id_to_second_slicing_bitset_.at(init_rule_id_); + } + + const std::string* prev_token = nullptr; + int32_t skip_ptr = 0; + const int32_t skip_size = static_cast(token_edge_accepted.size()); + for (size_t interval_idx = 0; interval_idx < possible_intervals.size(); ++interval_idx) { + const auto& interval = possible_intervals[interval_idx]; + for (int i = interval.first; i < interval.second; ++i) { + // Skip tokens already accepted by token edges (avoid expensive Earley simulation). + while (skip_ptr < skip_size && token_edge_accepted[skip_ptr] < i) ++skip_ptr; + if (skip_ptr < skip_size && token_edge_accepted[skip_ptr] == i) continue; + + // Check if the current token is in the rejected range. i.e. check if the current token + // is on the subtree of the rejected token. + if (i < last_rejected_range) { + if (fill_reject_indices) { + tmp_rejected_indices_.push_back(i); + fill_reject_indices = + tmp_rejected_indices_.size() >= AdaptiveTokenMask::USE_BITSET_THRESHOLD + ? false + : fill_reject_indices; + } else { + i = last_rejected_range - 1; + } + continue; + } + const auto& token = sorted_decoded_vocab[i].second; + // This optimization is useful for simple self-recursive rules, like string content. + if (speculative_calculation) { + // Optimization for tag dispatch rules. + if (definite_accepted_bitset.has_value()) { + // If the token is empty, it must be accepted. + if (token.empty()) { + tmp_accepted_indices_.push_back(i); + continue; + } + // If the token doesn't contain tags or stop strings since the second character, and it + // will transit to the start state after consuming the first character, it must be + // accepted. + if (speculative_mask[static_cast(token[0])] && + (*definite_accepted_bitset.value())[i]) { + tmp_accepted_indices_.push_back(i); + continue; + } + } else { + bool all_accepted = true; + for (char ch : token) { + // If the first character is not the ascii character or can't be accepted by the + // first character mask, we need to check them in the parser. + if (isascii(ch) == 0 || !speculative_mask[static_cast(ch)]) { + all_accepted = false; + break; + } + } + if (all_accepted) { + tmp_accepted_indices_.push_back(i); + continue; + } + } + } + // Many tokens may contain the same prefix, so we will avoid unnecessary matching + // by finding the longest common prefix with the previous token. + bool accepted = true; + if (prev_token != nullptr) { + int lcp_len = + std::mismatch(token.begin(), token.end(), prev_token->begin(), prev_token->end()) + .first - + token.begin(); + if (lcp_len > prev_matched_size) { + // Case 1. The common prefix is rejected by the matcher in the last token. Reject + // directly. + accepted = false; + } else if (lcp_len < prev_matched_size) { + // Case 2. The common prefix is shorter than the previous matched size. Rollback + // the non-common part. + PopLastStates(prev_matched_size - lcp_len); + tmp_can_reach_end_stack_.erase( + tmp_can_reach_end_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_stack_.end() + ); + tmp_can_reach_end_prefix_or_stack_.erase( + tmp_can_reach_end_prefix_or_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_prefix_or_stack_.end() + ); + } + prev_matched_size = std::min(prev_matched_size, lcp_len); + } + + prev_token = &token; + + if (accepted) { + // Accept the rest chars one by one. + for (int j = prev_matched_size; j < static_cast(token.size()); ++j) { + if (!Advance(token[j])) { + accepted = false; + break; + } + tmp_can_reach_end_stack_.push_back(IsCompleted()); + tmp_can_reach_end_prefix_or_stack_.push_back( + tmp_can_reach_end_stack_.back() || tmp_can_reach_end_prefix_or_stack_.back() + ); + prev_matched_size = j + 1; + } + } + + bool can_reach_end = tmp_can_reach_end_prefix_or_stack_.back(); + + if (accepted) { + if (HasEnteredCharBudget()) { + tmp_uncertain_indices_.push_back(i); + } else { + tmp_accepted_indices_.push_back(i); + } + } else if (can_reach_end && prev_matched_size > 0) { + auto [lookahead_accepted, lookahead_completed] = + IsTokenPassLookaheadAssertion(token, tmp_can_reach_end_stack_); + if ((!is_root_rule) && lookahead_accepted) { + if (lookahead_completed || !is_exact_lookahead) { + tmp_uncertain_indices_.push_back(i); + } else if (HasEnteredCharBudget()) { + tmp_uncertain_indices_.push_back(i); + } else { + tmp_accepted_indices_.push_back(i); + tmp_accepted_by_lookahead_indices_.push_back(i); + } + } else { + for (int j = i; j < subtree_nodes_range[i]; j++) { + tmp_rejected_indices_.push_back(j); + tmp_rejected_by_lookahead_indices_.push_back(j); + } + i = subtree_nodes_range[i] - 1; // Skip the subtree nodes. + } + } else { + tmp_rejected_indices_.push_back(i); + last_rejected_range = subtree_nodes_range[i]; + fill_reject_indices = + tmp_rejected_indices_.size() >= AdaptiveTokenMask::USE_BITSET_THRESHOLD + ? false + : fill_reject_indices; + } + } + if (interval_idx != possible_intervals.size() - 1 && fill_reject_indices) { + const auto& next_interval = possible_intervals[interval_idx + 1]; + for (int i = interval.second; i < next_interval.first; ++i) { + tmp_rejected_indices_.push_back(i); + } + fill_reject_indices = tmp_rejected_indices_.size() >= AdaptiveTokenMask::USE_BITSET_THRESHOLD + ? false + : fill_reject_indices; + } + } + + // Rollback the last matched part. + PopLastStates(prev_matched_size); + + if (possible_intervals.back().second != static_cast(sorted_decoded_vocab.size()) && + fill_reject_indices) { + // If the last interval is not closed, we need to reject the rest tokens. + for (int i = possible_intervals.back().second; + i < static_cast(sorted_decoded_vocab.size()); + ++i) { + tmp_rejected_indices_.push_back(i); + } + } + + return fill_reject_indices; +} + +void GrammarMatcherForTokenMaskCache::GetFirstCharacterMask(std::bitset<256>& first_character_mask +) { + first_character_mask.reset(); + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + const auto& edges = fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id); + for (const auto& edge : edges) { + if (edge.IsCharRange()) { + for (int c = edge.min; c <= edge.max; ++c) { + first_character_mask[c] = true; + } + } + } +} + +const std::vector& GrammarMatcherForTokenMaskCache::GetTokenEdgeAcceptedIndices() { + // Compute sorted vocab indices accepted by Token(ids) and ExcludeToken(ids) edges. + // Result is stored in tmp_token_edge_accepted_. + + tmp_token_edge_accepted_.clear(); + tmp_token_edge_excluded_.clear(); + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + const auto& edges = fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id); + + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + int32_t sorted_size = static_cast(sorted_decoded_vocab.size()); + const auto& tid_to_sorted = tokenizer_info_.ImplPtr()->GetTokenIdToSortedVocabIndex(); + // Out-of-vocabulary ids are rejected before compilation starts (CheckTokenIdsInVocab); guard + // here as well because this runs on worker threads, where an out-of-bounds read cannot be + // reported as an error. + auto sorted_index = [&](int32_t tid) { + return tid >= 0 && tid < static_cast(tid_to_sorted.size()) ? tid_to_sorted[tid] : -1; + }; + + bool has_exclude_token = false; + + for (const auto& edge : edges) { + if (edge.IsToken()) { + auto info = fsm.GetFsm().GetFsm().GetTokenEdgeInfo(edge.GetAuxIndex()); + for (int32_t i = 0; i < info.Count(); ++i) { + if (int32_t index = sorted_index(info.TokenIds()[i]); index >= 0) { + tmp_token_edge_accepted_.push_back(index); + } + } + } else if (edge.IsExcludeToken()) { + has_exclude_token = true; + auto info = fsm.GetFsm().GetFsm().GetExcludeTokenEdgeInfo(edge.GetAuxIndex()); + for (int32_t i = 0; i < info.Count(); ++i) { + if (int32_t index = sorted_index(info.TokenIds()[i]); index >= 0) { + tmp_token_edge_excluded_.push_back(index); + } + } + } + } + + // Token-only: result = token_accepted + if (!has_exclude_token) { + if (!tmp_token_edge_accepted_.empty()) { + std::sort(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()); + tmp_token_edge_accepted_.erase( + std::unique(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()), + tmp_token_edge_accepted_.end() + ); + } + return tmp_token_edge_accepted_; + } + + // ExcludeToken: result = [0, sorted_size) - (excluded - token_accepted) + // Token(ids) overrides ExcludeToken(ids) when both present. + if (!tmp_token_edge_accepted_.empty()) { + std::sort(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()); + tmp_token_edge_accepted_.erase( + std::unique(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()), + tmp_token_edge_accepted_.end() + ); + } + std::sort(tmp_token_edge_excluded_.begin(), tmp_token_edge_excluded_.end()); + tmp_token_edge_excluded_.erase( + std::unique(tmp_token_edge_excluded_.begin(), tmp_token_edge_excluded_.end()), + tmp_token_edge_excluded_.end() + ); + IntsetDifference(&tmp_token_edge_excluded_, tmp_token_edge_accepted_); + IntsetComplement(&tmp_token_edge_accepted_, sorted_size, tmp_token_edge_excluded_); + return tmp_token_edge_accepted_; +} + +AdaptiveTokenMask GrammarMatcherForTokenMaskCache::GetAdaptiveTokenMask(bool is_root_rule) { + tmp_accepted_indices_.clear(); + tmp_rejected_indices_.clear(); + tmp_uncertain_indices_.clear(); + tmp_rejected_by_lookahead_indices_.clear(); + tmp_accepted_by_lookahead_indices_.clear(); + tmp_can_reach_end_prefix_or_stack_.clear(); + tmp_can_reach_end_stack_.clear(); + // For every character in the current token, stores whether it is possible to reach the end of + // the rule when matching until this character. Store it in a stack for later rollback. + tmp_can_reach_end_stack_.push_back(false); + tmp_can_reach_end_prefix_or_stack_.push_back(false); + + // Try to get the crossing cache. + bool rule_level_cache_is_available = !has_char_budget_rules_ && rule_level_cache_.has_value() && + grammar_->per_rule_fsm_hashes[init_rule_id_].has_value(); + std::optional fsm_hash = std::nullopt; + int32_t new_state_id = -1; + std::optional crossing_cache = std::nullopt; + int lookahead_id = grammar_->GetRule(initial_state_.rule_id).lookahead_assertion_id; + bool is_exact_lookahead = grammar_->GetRule(initial_state_.rule_id).is_exact_lookahead; + std::optional lookahead_hash = std::nullopt; + if (rule_level_cache_is_available) { + lookahead_hash = GrammarFSMHasher::HashSequence(grammar_, lookahead_id); + const auto& original_to_new_id = grammar_->per_rule_fsm_new_state_ids[init_rule_id_]; + fsm_hash = grammar_->per_rule_fsm_hashes[init_rule_id_].value(); + for (const auto& original_new_pair : original_to_new_id) { + if (original_new_pair.first == initial_state_.element_id) { + new_state_id = original_new_pair.second; + break; + } + } + XGRAMMAR_DCHECK(new_state_id != -1); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + if (lookahead_hash.has_value()) { + crossing_cache = rule_level_cache_->GetCache( + HashCombine(fsm_hash.value(), lookahead_hash.value(), is_exact_lookahead), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum() + ); + if (crossing_cache.has_value()) { + // A perfect match. + return crossing_cache.value(); + } + } + crossing_cache = rule_level_cache_->GetCache( + fsm_hash.value(), new_state_id, fsm.GetNodeNum(), fsm.GetEdgeNum() + ); + // If the rule doesn't have a lookahead, then it is exactly the same fsm. + if (crossing_cache.has_value()) { + AdaptCacheWithLookahead(&crossing_cache.value(), is_root_rule); + return std::move(crossing_cache.value()); + } + } + + std::bitset<256> first_character_mask; + GetFirstCharacterMask(first_character_mask); + + // Token edge accepted indices (for byte path skip + merge). + const auto& token_edge_accepted = GetTokenEdgeAcceptedIndices(); + + // Byte path: skip tokens already accepted by token edges. + bool rejected_filled; + if (first_character_mask.none()) { + rejected_filled = false; + } else { + rejected_filled = GetTokenMaskWithFirstCharacterCheck( + first_character_mask, is_root_rule, token_edge_accepted + ); + } + + // Token edges are rechecked at runtime when a character budget is present because accepting + // one can enter a budgeted rule within the same token. + if (!token_edge_accepted.empty()) { + if (has_char_budget_rules_) { + IntsetUnion(&tmp_uncertain_indices_, token_edge_accepted); + } else { + IntsetUnion(&tmp_accepted_indices_, token_edge_accepted); + IntsetDifference(&tmp_uncertain_indices_, token_edge_accepted); + } + IntsetDifference(&tmp_rejected_indices_, token_edge_accepted); + } + if (rejected_filled) { + auto return_value = AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + tmp_accepted_indices_, + tmp_rejected_indices_, + tmp_uncertain_indices_ + ); + if (rule_level_cache_is_available) { + if (lookahead_id == -1 && !is_root_rule) { + // If the rule doesn't have a lookahead, then it is exactly the same fsm. + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + rule_level_cache_->AddCache( + fsm_hash.value(), new_state_id, fsm.GetNodeNum(), fsm.GetEdgeNum(), return_value + ); + return return_value; + } + + // We can add a cache for basic fsm, and a better one for lookahead. + // All the tokens rejected by lookahead should be uncertain. + IntsetUnion(&tmp_uncertain_indices_, tmp_rejected_by_lookahead_indices_); + IntsetUnion(&tmp_uncertain_indices_, tmp_accepted_by_lookahead_indices_); + std::vector rejected_indices_without_lookahead; + std::vector accepted_indices_without_lookahead; + rejected_indices_without_lookahead.reserve( + tmp_rejected_indices_.size() - tmp_rejected_by_lookahead_indices_.size() + ); + accepted_indices_without_lookahead.reserve( + tmp_accepted_indices_.size() - tmp_accepted_by_lookahead_indices_.size() + ); + std::set_difference( + tmp_rejected_indices_.begin(), + tmp_rejected_indices_.end(), + tmp_rejected_by_lookahead_indices_.begin(), + tmp_rejected_by_lookahead_indices_.end(), + std::back_inserter(rejected_indices_without_lookahead) + ); + std::set_difference( + tmp_accepted_indices_.begin(), + tmp_accepted_indices_.end(), + tmp_accepted_by_lookahead_indices_.begin(), + tmp_accepted_by_lookahead_indices_.end(), + std::back_inserter(accepted_indices_without_lookahead) + ); + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + rule_level_cache_->AddCache( + fsm_hash.value(), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + accepted_indices_without_lookahead, + rejected_indices_without_lookahead, + tmp_uncertain_indices_ + ) + ); + if (lookahead_hash.has_value()) { + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + rule_level_cache_->AddCache( + HashCombine(fsm_hash.value(), lookahead_hash.value(), is_exact_lookahead), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + return_value + ); + } + } + return return_value; + } else { + auto return_value = AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + tmp_accepted_indices_, + tmp_uncertain_indices_ + ); + + if (rule_level_cache_is_available) { + // Prepare for cache. + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + if (lookahead_id == -1 && !is_root_rule) { + // If the rule doesn't have a lookahead, then it is exactly the same fsm. + rule_level_cache_->AddCache( + fsm_hash.value(), new_state_id, fsm.GetNodeNum(), fsm.GetEdgeNum(), return_value + ); + return return_value; + } + + // Add 2 caches. + IntsetUnion(&tmp_uncertain_indices_, tmp_rejected_by_lookahead_indices_); + IntsetUnion(&tmp_uncertain_indices_, tmp_accepted_by_lookahead_indices_); + std::vector accepted_indices_without_lookahead; + accepted_indices_without_lookahead.reserve( + tmp_accepted_indices_.size() - tmp_accepted_by_lookahead_indices_.size() + ); + std::set_difference( + tmp_accepted_indices_.begin(), + tmp_accepted_indices_.end(), + tmp_accepted_by_lookahead_indices_.begin(), + tmp_accepted_by_lookahead_indices_.end(), + std::back_inserter(accepted_indices_without_lookahead) + ); + rule_level_cache_->AddCache( + fsm_hash.value(), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + accepted_indices_without_lookahead, + tmp_uncertain_indices_ + ) + ); + + if (lookahead_hash.has_value()) { + rule_level_cache_->AddCache( + HashCombine(fsm_hash.value(), lookahead_hash.value(), is_exact_lookahead), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + return_value + ); + } + } + return return_value; + } +} + +/******************* GrammarCompilerNoCache *******************/ + +/*! + * \brief The base class for the grammar compiler. Handles the compilation logic without cache. + */ +class GrammarCompilerSub { + public: + GrammarCompilerSub( + const TokenizerInfo& tokenizer_info, + int max_threads, + std::optional rule_level_cache + ) + : tokenizer_info_(tokenizer_info), + max_threads_(max_threads), + rule_level_cache_(rule_level_cache) {} + + CompiledGrammar CompileBuiltinJSONGrammar(); + + CompiledGrammar CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order + ); + + CompiledGrammar CompileRegex(const std::string& regex); + + CompiledGrammar CompileLark( + const std::string& lark_string, const std::vector& named_grammars + ); + + CompiledGrammar CompileStructuralTag(const std::string& structural_tag_json); + + CompiledGrammar CompileGrammar(const Grammar& grammar); + + CompiledGrammar CompileGrammar(const std::string& ebnf_str, std::string root_rule_name); + + private: + /*! \brief The main logic. Compile the grammar with multi-threading. */ + CompiledGrammar MultiThreadCompileGrammar(Grammar grammar); + /*! \brief Optimization for TagDispatch. + * \param compiled_grammar_impl the compiled_grammar to be optimized. + * \param tag_dispatch_rule_id_to_second_slicing_bitset Return value. Mapping from the rule_id to + * the definite accepted token mask. + */ + void TagDispatchOptimization( + std::shared_ptr compiled_grammar_impl, + std::unordered_map* tag_dispatch_rule_id_to_second_slicing_bitset + ); + + /*! \brief The vocabulary associated with this storage class. */ + const TokenizerInfo tokenizer_info_; + /*! \brief The maximum number of threads to use. */ + const int max_threads_; + + /*! \brief The manager of the rule level cache.*/ + std::optional rule_level_cache_; +}; + +/*! + * \brief Check that every token id written in Token(...), ExcludeToken(...) and TokenTagDispatch + * exists in the vocabulary. They are used as indices into per-token arrays during compilation, + * which runs on worker threads where an error cannot be reported. + */ +static void CheckTokenIdsInVocab(const Grammar& grammar, int vocab_size) { + const CompactFSM& fsm = grammar->complete_fsm; + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (!edge.IsToken() && !edge.IsExcludeToken()) { + continue; + } + // Token and exclude-token edges share the [count, token_id_0, ...] aux layout. + const auto info = fsm.GetTokenEdgeInfo(edge.GetAuxIndex()); + for (int32_t i = 0; i < info.Count(); ++i) { + XGRAMMAR_CHECK(info.TokenIds()[i] >= 0 && info.TokenIds()[i] < vocab_size) + << "Token id " << info.TokenIds()[i] + << " in the grammar is out of the vocabulary range [0, " << vocab_size << ")"; + } + } + } +} + +CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_unoptimized) { + auto compiled_grammar_impl = std::make_shared(); + compiled_grammar_impl->grammar = GrammarOptimizer::Apply(grammar_unoptimized); + compiled_grammar_impl->tokenizer_info = tokenizer_info_; + CheckTokenIdsInVocab(compiled_grammar_impl->grammar, tokenizer_info_.GetVocabSize()); + if (tokenizer_info_.GetVocabSize() == 0) { + return CompiledGrammar(compiled_grammar_impl); + } + std::unordered_map tag_dispatch_rule_id_to_second_slicing_bitset; + TagDispatchOptimization(compiled_grammar_impl, &tag_dispatch_rule_id_to_second_slicing_bitset); + + // If the compiler is cache-enabled, then we hash the grammars for crossing-grammar caching. + if (rule_level_cache_.has_value()) { + GrammarFSMHasher().Apply(&compiled_grammar_impl->grammar); + } + // Step 3. Compute the adaptive token mask cache + // The token mask cache is computed for these positions in the grammar: + // 1. All character class or character class star (with last_utf8_bytes=0, 1, 2, 3) + // 2. All byte strings (with element_in_string=0, 1, 2, ...) + // since other positions will be expanded to the above positions + + // TODO(Charlie): Figure out how to support ThreadPool and std::mutex in WebAssembly. + // Only declare ThreadPool and mutex if max_threads > 1, so when max_threads = 1, we do + // not need ThreadPool or std::mutex, which throws error in runtime in WebAssembly. + std::optional thread_pool; + std::optional adaptive_token_mask_cache_mutex; + if (max_threads_ > 1) { + thread_pool.emplace(max_threads_); + adaptive_token_mask_cache_mutex.emplace(); + } + + auto add_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) { + auto grammar_matcher = GrammarMatcherForTokenMaskCache( + compiled_grammar_impl->grammar, + state, + tag_dispatch_rule_id_to_second_slicing_bitset, + tokenizer_info_, + rule_level_cache_ + ); + auto cur_adaptive_token_mask_cache = grammar_matcher.GetAdaptiveTokenMask(is_root_rule); + if (max_threads_ > 1) { + std::lock_guard lock(adaptive_token_mask_cache_mutex.value()); + compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache; + } else { + compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache; + } + }; + + auto add_task_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) { + // Execute depending on whether we use thread_pool + if (max_threads_ > 1) { + thread_pool->Execute([add_adaptive_token_mask, state, is_root_rule]() { + add_adaptive_token_mask(state, is_root_rule); + }); + } else { + add_adaptive_token_mask(state, is_root_rule); + } + }; + + auto root_rule_id = compiled_grammar_impl->grammar->GetRootRuleId(); + + for (int32_t rule_id = 0; rule_id < static_cast(compiled_grammar_impl->grammar->NumRules()); + ++rule_id) { + auto rule = compiled_grammar_impl->grammar->GetRule(rule_id); + const auto& rule_fsm = compiled_grammar_impl->grammar->per_rule_fsms[rule_id]; + XGRAMMAR_DCHECK(rule_fsm.has_value()); + auto cur_stack_element = + ParserState(rule_id, rule.body_expr_id, 0, ParserState::kNoPrevInputPos, 0); + std::unordered_set reachable_states; + rule_fsm->GetFsm().GetReachableStates(&reachable_states); + for (int i : reachable_states) { + cur_stack_element.element_id = i; + if (!rule_fsm->GetFsm().IsScanableState(i)) { + continue; + } + add_task_adaptive_token_mask(cur_stack_element, rule_id == root_rule_id); + } + } + + if (max_threads_ > 1) { + thread_pool->Join(); + } + + return CompiledGrammar(compiled_grammar_impl); +} + +CompiledGrammar GrammarCompilerSub::CompileBuiltinJSONGrammar() { + return MultiThreadCompileGrammar(Grammar::BuiltinJSONGrammar()); +} + +CompiledGrammar GrammarCompilerSub::CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order +) { + return MultiThreadCompileGrammar(Grammar::FromJSONSchema( + schema, + any_whitespace, + indent, + separators, + strict_mode, + max_whitespace_cnt, + /*print_converted_ebnf=*/false, + any_order + )); +} + +CompiledGrammar GrammarCompilerSub::CompileStructuralTag(const std::string& structural_tag_json) { + auto result = Grammar::FromStructuralTag(structural_tag_json, tokenizer_info_); + XGRAMMAR_CHECK(std::holds_alternative(result)) + << GetMessageFromVariantError(std::get<1>(result)); + return MultiThreadCompileGrammar(std::get<0>(result)); +} + +CompiledGrammar GrammarCompilerSub::CompileRegex(const std::string& regex) { + return MultiThreadCompileGrammar(Grammar::FromRegex(regex)); +} + +CompiledGrammar GrammarCompilerSub::CompileLark( + const std::string& lark_string, const std::vector& named_grammars +) { + return MultiThreadCompileGrammar(Grammar::FromLark(lark_string, tokenizer_info_, named_grammars)); +} + +CompiledGrammar GrammarCompilerSub::CompileGrammar(const Grammar& grammar) { + return MultiThreadCompileGrammar(grammar); +} + +CompiledGrammar GrammarCompilerSub::CompileGrammar( + const std::string& ebnf_str, std::string root_rule_name +) { + return MultiThreadCompileGrammar(Grammar::FromEBNF(ebnf_str, root_rule_name)); +} + +void GrammarCompilerSub::TagDispatchOptimization( + std::shared_ptr compiled_grammar_impl, + std::unordered_map* tag_dispatch_rule_id_to_second_slicing_bitset +) { + using GrammarExprType = Grammar::Impl::GrammarExprType; + tag_dispatch_rule_id_to_second_slicing_bitset->clear(); + + // Optimization for TagDispatch: Precompute the definitely accepted tokens. + for (int i = 0; i < compiled_grammar_impl->grammar->NumRules(); i++) { + const auto& rule = compiled_grammar_impl->grammar->GetRule(i); + const auto& rule_body = compiled_grammar_impl->grammar->GetGrammarExpr(rule.body_expr_id); + if (rule_body.type != GrammarExprType::kTagDispatch) { + continue; + } + XGRAMMAR_DCHECK(rule_body.type == GrammarExprType::kTagDispatch); + Grammar::Impl::TagDispatch tag_dispatch = + compiled_grammar_impl->GetGrammar()->GetTagDispatch(rule.body_expr_id); + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + DynamicBitset definite_accepted_tokens_since_second_char(sorted_decoded_vocab.size()); + for (int j = 0; j < static_cast(sorted_decoded_vocab.size()); j++) { + bool definite_accept_since_second_char = true; + const auto& token = sorted_decoded_vocab[j].second; + if (token.empty()) { + definite_accepted_tokens_since_second_char.Set(j); + continue; + } + + // Check if the token contains any string trigger or exclude string after first char. + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + if (token.find(trigger, 1) != std::string::npos) { + definite_accept_since_second_char = false; + break; + } + } + if (definite_accept_since_second_char) { + for (const auto& excl : tag_dispatch.excludes) { + if (token.find(excl, 1) != std::string::npos) { + definite_accept_since_second_char = false; + break; + } + } + } + + if (definite_accept_since_second_char) { + definite_accepted_tokens_since_second_char.Set(j); + } + } + (*tag_dispatch_rule_id_to_second_slicing_bitset)[i] = + definite_accepted_tokens_since_second_char; + } +} + +/******************* GrammarCompiler::Impl *******************/ + +/*! + * \brief The keys for the cache. This is defined here instead of inside the GrammarCompiler::Impl + * class due C++ template specialization and hash specialization rules. + */ +class GrammarCompilerCacheKeys { + public: + struct SchemaKey { + std::string schema; + bool any_whitespace; + std::optional indent; + std::optional> separators; + bool strict_mode; + std::optional max_whitespace_cnt; + bool any_order; + + XGRAMMAR_EQUAL_BY_MEMBERS( + SchemaKey, + &SchemaKey::schema, + &SchemaKey::any_whitespace, + &SchemaKey::indent, + &SchemaKey::separators, + &SchemaKey::strict_mode, + &SchemaKey::max_whitespace_cnt, + &SchemaKey::any_order + ); + }; + + struct StructuralTagKey { + std::string structural_tag_json; + + XGRAMMAR_EQUAL_BY_MEMBERS(StructuralTagKey, &StructuralTagKey::structural_tag_json); + }; + + struct GrammarKey { + std::string ebnf_str; + std::string root_rule_name; + + XGRAMMAR_EQUAL_BY_MEMBERS(GrammarKey, &GrammarKey::ebnf_str, &GrammarKey::root_rule_name); + }; + + struct RegexKey { + std::string regex; + + XGRAMMAR_EQUAL_BY_MEMBERS(RegexKey, &RegexKey::regex); + }; + + struct LarkNamedGrammarKey { + std::string name; + bool is_lark_source; + std::string source_or_ebnf; + std::string root_rule_name; + std::variant value; + + bool operator==(const LarkNamedGrammarKey& other) const { + return std::tie(name, is_lark_source, source_or_ebnf, root_rule_name) == + std::tie(other.name, other.is_lark_source, other.source_or_ebnf, other.root_rule_name); + } + }; + + struct LarkKey { + std::string lark_string; + std::vector named_grammars; + + XGRAMMAR_EQUAL_BY_MEMBERS(LarkKey, &LarkKey::lark_string, &LarkKey::named_grammars); + }; + + struct BuiltinJSONGrammarKey { + XGRAMMAR_EQUAL_BY_MEMBERS_EMPTY(BuiltinJSONGrammarKey); + }; + + using UnionKey = std:: + variant; +}; + +} // namespace xgrammar + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::SchemaKey, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::schema, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::any_whitespace, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::indent, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::separators, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::strict_mode, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::max_whitespace_cnt, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::any_order +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::StructuralTagKey, + &xgrammar::GrammarCompilerCacheKeys::StructuralTagKey::structural_tag_json +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::GrammarKey, + &xgrammar::GrammarCompilerCacheKeys::GrammarKey::ebnf_str, + &xgrammar::GrammarCompilerCacheKeys::GrammarKey::root_rule_name +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::RegexKey, + &xgrammar::GrammarCompilerCacheKeys::RegexKey::regex +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::name, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::is_lark_source, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::source_or_ebnf, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::root_rule_name +); + +namespace std { +template <> +struct hash { + std::size_t operator()(const xgrammar::GrammarCompilerCacheKeys::LarkKey& key) const noexcept { + uint64_t seed = std::hash{}(key.lark_string); + for (const auto& named_grammar : key.named_grammars) { + xgrammar::HashCombineBinary( + seed, std::hash{}(named_grammar) + ); + } + return seed; + } +}; +} // namespace std + +XGRAMMAR_HASH_BY_MEMBERS_EMPTY(xgrammar::GrammarCompilerCacheKeys::BuiltinJSONGrammarKey); + +namespace xgrammar { + +/*! + * \brief The implementation of the grammar compiler with cache. It calls the no cache compiler + * to compile the grammar, and implements the cache logic upon it. + */ +class GrammarCompiler::Impl { + public: + Impl( + const TokenizerInfo& tokenizer_info, + int max_threads, + bool cache_enabled, + int64_t max_memory_bytes + ) + : cache_enabled_(cache_enabled), + rule_level_cache_( + cache_enabled + ? std::optional( + max_memory_bytes == -1 + ? static_cast(-1) + : static_cast(max_memory_bytes - max_memory_bytes / 3 * 2) + ) + : std::nullopt + ), + no_cache_compiler_(tokenizer_info, max_threads, rule_level_cache_), + grammar_level_cache_( + max_memory_bytes == -1 ? static_cast(-1) + : static_cast(max_memory_bytes / 3 * 2), + Computer(*this) + ) { + if (max_memory_bytes < -1) { + XGRAMMAR_LOG(FATAL) << "Invalid max_memory_bytes: " << max_memory_bytes << ". " + << "It should be -1 (unlimited) or a non-negative integer."; + } + } + + CompiledGrammar CompileBuiltinJSONGrammar(); + + CompiledGrammar CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order + ); + + CompiledGrammar CompileStructuralTag(const std::string& structural_tag_json); + + CompiledGrammar CompileRegex(const std::string& regex); + + CompiledGrammar CompileLark( + const std::string& lark_string, const std::vector& named_grammars + ); + + CompiledGrammar CompileGrammar(const Grammar& grammar); + + CompiledGrammar CompileGrammar(const std::string& ebnf_str, std::string root_rule_name); + + void ClearCache(); + + int64_t GetCacheSizeBytes() const; + + int64_t CacheLimitBytes() const; + + private: + using SchemaKey = GrammarCompilerCacheKeys::SchemaKey; + using StructuralTagKey = GrammarCompilerCacheKeys::StructuralTagKey; + using GrammarKey = GrammarCompilerCacheKeys::GrammarKey; + using RegexKey = GrammarCompilerCacheKeys::RegexKey; + using LarkNamedGrammarKey = GrammarCompilerCacheKeys::LarkNamedGrammarKey; + using LarkKey = GrammarCompilerCacheKeys::LarkKey; + using BuiltinJSONGrammarKey = GrammarCompilerCacheKeys::BuiltinJSONGrammarKey; + using UnionKey = GrammarCompilerCacheKeys::UnionKey; + + CompiledGrammar Compute(const UnionKey& key); + + struct Computer { + Computer(Impl& compiler) : compiler(compiler) {} + // Forward the key to GrammarCompiler::Impl::Compute(key) + CompiledGrammar operator()(const UnionKey& key) const { return compiler.Compute(key); } + GrammarCompiler::Impl& compiler; + }; + + struct SizeEstimator { + std::size_t operator()(const CompiledGrammar& value) const { return value.MemorySizeBytes(); } + }; + + /*! \brief Whether the cache is enabled. */ + const bool cache_enabled_; + + /*! \brief The crossing cache manager for compiled grammars. */ + std::optional rule_level_cache_ = std::nullopt; + + /*! \brief The no cache compiler. */ + GrammarCompilerSub no_cache_compiler_; + + /*! \brief The cache for compiled grammars. */ + ThreadSafeLRUCache grammar_level_cache_; +}; + +CompiledGrammar GrammarCompiler::Impl::Compute(const UnionKey& key) { + return std::visit( + [this](const auto& key) -> CompiledGrammar { + using KeyType = std::decay_t; + if constexpr (std::is_same_v) { + const auto& [ebnf_str, root_rule_name] = key; + return this->no_cache_compiler_.CompileGrammar(ebnf_str, root_rule_name); + } else if constexpr (std::is_same_v) { + const auto& [schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order] = + key; + return this->no_cache_compiler_.CompileJSONSchema( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + ); + } else if constexpr (std::is_same_v) { + const auto& [structural_tag_json] = key; + return this->no_cache_compiler_.CompileStructuralTag(structural_tag_json); + } else if constexpr (std::is_same_v) { + const auto& [regex] = key; + return this->no_cache_compiler_.CompileRegex(regex); + } else if constexpr (std::is_same_v) { + std::vector named_grammars; + named_grammars.reserve(key.named_grammars.size()); + for (const auto& named_grammar : key.named_grammars) { + named_grammars.push_back({named_grammar.name, named_grammar.value}); + } + return this->no_cache_compiler_.CompileLark(key.lark_string, named_grammars); + } else if constexpr (std::is_same_v) { + return this->no_cache_compiler_.CompileBuiltinJSONGrammar(); + } else { + XGRAMMAR_UNREACHABLE(); + } + }, + key + ); +} + +CompiledGrammar GrammarCompiler::Impl::CompileBuiltinJSONGrammar() { + if (!cache_enabled_) { + return no_cache_compiler_.CompileBuiltinJSONGrammar(); + } + return grammar_level_cache_.Get(BuiltinJSONGrammarKey{}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileJSONSchema( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + ); + } + return grammar_level_cache_.Get(SchemaKey{ + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + }); +} + +CompiledGrammar GrammarCompiler::Impl::CompileStructuralTag(const std::string& structural_tag_json +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileStructuralTag(structural_tag_json); + } + return grammar_level_cache_.Get(StructuralTagKey{structural_tag_json}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileRegex(const std::string& regex) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileRegex(regex); + } + return grammar_level_cache_.Get(RegexKey{regex}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileLark( + const std::string& lark_string, const std::vector& named_grammars +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileLark(lark_string, named_grammars); + } + + std::vector named_grammar_keys; + named_grammar_keys.reserve(named_grammars.size()); + for (const auto& named_grammar : named_grammars) { + if (std::holds_alternative(named_grammar.grammar)) { + const auto& source = std::get(named_grammar.grammar); + named_grammar_keys.push_back( + {named_grammar.name, /*is_lark_source=*/true, source, "", named_grammar.grammar} + ); + } else { + const auto& grammar = std::get(named_grammar.grammar); + named_grammar_keys.push_back( + {named_grammar.name, + /*is_lark_source=*/false, + grammar.ToString(), + grammar->GetRootRule().name, + named_grammar.grammar} + ); + } + } + std::sort( + named_grammar_keys.begin(), + named_grammar_keys.end(), + [](const LarkNamedGrammarKey& lhs, const LarkNamedGrammarKey& rhs) { + return std::tie(lhs.name, lhs.is_lark_source, lhs.source_or_ebnf, lhs.root_rule_name) < + std::tie(rhs.name, rhs.is_lark_source, rhs.source_or_ebnf, rhs.root_rule_name); + } + ); + return grammar_level_cache_.Get(LarkKey{lark_string, std::move(named_grammar_keys)}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileGrammar(const Grammar& grammar) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileGrammar(grammar); + } + return grammar_level_cache_.Get(GrammarKey{grammar.ToString(), grammar->GetRootRule().name}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileGrammar( + const std::string& ebnf_str, std::string root_rule_name +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileGrammar(ebnf_str, root_rule_name); + } + return grammar_level_cache_.Get(GrammarKey{ebnf_str, root_rule_name}); +} + +void GrammarCompiler::Impl::ClearCache() { + grammar_level_cache_.Clear(); + if (rule_level_cache_.has_value()) { + rule_level_cache_->ClearCache(); + } +} + +int64_t GrammarCompiler::Impl::GetCacheSizeBytes() const { + return static_cast(grammar_level_cache_.MemorySize()) + + static_cast(MemorySize(rule_level_cache_)); +} + +int64_t GrammarCompiler::Impl::CacheLimitBytes() const { + const auto size = grammar_level_cache_.MaxMemorySize(); + if (size == grammar_level_cache_.kUnlimitedSize) return -1; + return static_cast(size) + (rule_level_cache_.has_value() + ? static_cast(rule_level_cache_->GetMaxSize()) + : 0); +} + +/******************* GrammarCompiler *******************/ + +GrammarCompiler::GrammarCompiler( + const TokenizerInfo& tokenizer_info, + int max_threads, + bool cache_enabled, + int64_t max_memory_bytes +) + : pimpl_(std::make_shared(tokenizer_info, max_threads, cache_enabled, max_memory_bytes)) { +} + +CompiledGrammar GrammarCompiler::CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order +) { + return pimpl_->CompileJSONSchema( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + ); +} + +CompiledGrammar GrammarCompiler::CompileBuiltinJSONGrammar() { + return pimpl_->CompileBuiltinJSONGrammar(); +} + +CompiledGrammar GrammarCompiler::CompileStructuralTag(const std::string& structural_tag_json) { + return pimpl_->CompileStructuralTag(structural_tag_json); +} + +CompiledGrammar GrammarCompiler::CompileRegex(const std::string& regex) { + return pimpl_->CompileRegex(regex); +} + +CompiledGrammar GrammarCompiler::CompileLark( + const std::string& lark_string, const std::vector& named_grammars +) { + return pimpl_->CompileLark(lark_string, named_grammars); +} + +CompiledGrammar GrammarCompiler::CompileGrammar(const Grammar& grammar) { + return pimpl_->CompileGrammar(grammar); +} + +CompiledGrammar GrammarCompiler::CompileGrammar( + const std::string& ebnf_str, const std::string& root_rule_name +) { + return pimpl_->CompileGrammar(ebnf_str, root_rule_name); +} + +void GrammarCompiler::ClearCache() { pimpl_->ClearCache(); } + +int64_t GrammarCompiler::GetCacheSizeBytes() const { return pimpl_->GetCacheSizeBytes(); } + +int64_t GrammarCompiler::CacheLimitBytes() const { return pimpl_->CacheLimitBytes(); } + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_functor.cc b/third_party/xgrammar/cpp/grammar_functor.cc new file mode 100644 index 0000000000..8028dd8adf --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_functor.cc @@ -0,0 +1,3832 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_functor.cc + */ + +#include "grammar_functor.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "fsm.h" +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_impl.h" +#include "suffix_automata.h" +#include "support/container.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +using GrammarExpr = Grammar::Impl::GrammarExpr; +using ExprType = Grammar::Impl::GrammarExprType; + +/*************************** Impl of grammar constructors ***************************/ + +/*! + * \brief Base class for grammar mutators that add subgrammars. + * + * Provides functionality to visit a subgrammar and add its rules to the builder + * while maintaining proper rule references and names. + */ +class SubGrammarAdderImpl : public GrammarMutator { + public: + SubGrammarAdderImpl() = default; + + /*! + * \brief Visit a subgrammar and add the rules to the builder. + * \param grammar The subgrammar to visit. + * \return The new id of the root rule of this subgrammar. + */ + int32_t ApplyWithBuilder(GrammarBuilder* builder, const Grammar& sub_grammar) { + InitGrammar(sub_grammar); + InitBuilder(builder); + new_rule_ids_names.reserve(base_grammar_->NumRules()); + new_rule_ids_names.clear(); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto new_name = builder_->GetNewRuleName(base_grammar_->GetRule(i).name); + auto new_id = builder_->AddEmptyRule(new_name); + new_rule_ids_names.emplace_back(new_id, new_name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = new_rule_ids_names[i].second; + auto new_body_expr_id = VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(new_rule_ids_names[i].first, new_body_expr_id); + auto new_lookahead_assertion_id = VisitLookaheadAssertion(rule.lookahead_assertion_id); + builder_->UpdateLookaheadAssertion(new_rule_ids_names[i].first, new_lookahead_assertion_id); + builder_->UpdateMaxTokens(new_rule_ids_names[i].first, rule.max_tokens); + builder_->UpdateMaxChars(new_rule_ids_names[i].first, rule.max_chars); + builder_->UpdateCaptureName(new_rule_ids_names[i].first, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + auto remapped_info = *suffix_stop_info; + if (remapped_info.body_rule_id >= 0) { + remapped_info.body_rule_id = new_rule_ids_names[remapped_info.body_rule_id].first; + remapped_info.marker_rule_id = new_rule_ids_names[remapped_info.marker_rule_id].first; + } + builder_->UpdateSuffixStopInfo(new_rule_ids_names[i].first, remapped_info); + } + builder_->UpdateLazy(new_rule_ids_names[i].first, rule.is_lazy); + builder_->UpdateRuleTemperature(new_rule_ids_names[i].first, rule.temperature); + } + return new_rule_ids_names[base_grammar_->GetRootRuleId()].first; + } + + int32_t VisitRuleRef(const GrammarExpr& grammar_expr) final { + return builder_->AddRuleRef(new_rule_ids_names[grammar_expr[0]].first); + } + + int32_t VisitRepeat(const GrammarExpr& grammar_expr) final { + return builder_->AddRepeat( + new_rule_ids_names[grammar_expr[0]].first, grammar_expr[1], grammar_expr[2] + ); + } + + int32_t VisitTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TagDispatch old_tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + Grammar::Impl::TagDispatch new_tag_dispatch; + for (const auto& [trigger, rule_id] : old_tag_dispatch.tag_rule_pairs) { + new_tag_dispatch.tag_rule_pairs.emplace_back(trigger, new_rule_ids_names[rule_id].first); + } + new_tag_dispatch.loop_after_dispatch = old_tag_dispatch.loop_after_dispatch; + new_tag_dispatch.excludes = old_tag_dispatch.excludes; + return builder_->AddTagDispatch(new_tag_dispatch); + } + + int32_t VisitTokenTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TokenTagDispatch old_ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + Grammar::Impl::TokenTagDispatch new_ttd; + for (const auto& [token_id, rule_id] : old_ttd.trigger_rule_pairs) { + new_ttd.trigger_rule_pairs.emplace_back(token_id, new_rule_ids_names[rule_id].first); + } + new_ttd.loop_after_dispatch = old_ttd.loop_after_dispatch; + new_ttd.excludes = old_ttd.excludes; + return builder_->AddTokenTagDispatch(new_ttd); + } + + std::vector> new_rule_ids_names; +}; + +/*! + * \brief Implementation of grammar union operation. + * + * Creates a new grammar that accepts strings from any of the input grammars. + * The resulting grammar has a new root rule that chooses between the root rules + * of all input grammars. + */ +class GrammarUnionFunctorImpl : public GrammarMutator { + public: + GrammarUnionFunctorImpl() = default; + + Grammar Apply(const std::vector& grammars) { + InitGrammar(); + InitBuilder(); + auto root_rule_id = builder_->AddEmptyRule("root"); + + std::vector new_root_choices; + new_root_choices.reserve(grammars.size()); + + for (const auto& grammar : grammars) { + auto new_root_id_for_grammar = SubGrammarAdderImpl().ApplyWithBuilder(builder_, grammar); + auto new_rule_ref = builder_->AddRuleRef(new_root_id_for_grammar); + auto new_rule_ref_seq = builder_->AddSequence({new_rule_ref}); + new_root_choices.push_back(new_rule_ref_seq); + } + + builder_->UpdateRuleBody(root_rule_id, builder_->AddChoices(new_root_choices)); + return builder_->Get(root_rule_id); + } + + // Avoid hiding the original Apply(const Grammar&) + Grammar Apply(const Grammar& grammar) final { + XGRAMMAR_LOG(FATAL) << "Should not be called"; + XGRAMMAR_UNREACHABLE(); + } +}; + +/*! + * \brief Implementation of grammar concatenation operation. + * + * Creates a new grammar that accepts strings that are concatenations of strings + * from the input grammars in order. The resulting grammar has a new root rule + * that concatenates the root rules of all input grammars. + */ +class GrammarConcatFunctorImpl : public GrammarMutator { + public: + GrammarConcatFunctorImpl() = default; + + Grammar Apply(const std::vector& grammars) { + InitGrammar(); + InitBuilder(); + auto root_rule_id = builder_->AddEmptyRule("root"); + + std::vector new_root_sequence; + new_root_sequence.reserve(grammars.size()); + + for (const auto& grammar : grammars) { + auto new_root_id_for_grammar = SubGrammarAdderImpl().ApplyWithBuilder(builder_, grammar); + auto new_rule_ref = builder_->AddRuleRef(new_root_id_for_grammar); + new_root_sequence.push_back(new_rule_ref); + } + + auto new_root_seq = builder_->AddSequence(new_root_sequence); + builder_->UpdateRuleBody(root_rule_id, builder_->AddChoices({new_root_seq})); + + return builder_->Get(root_rule_id); + } + + // Avoid hiding the original Apply(const Grammar&) + Grammar Apply(const Grammar& grammar) final { + XGRAMMAR_LOG(FATAL) << "Should not be called"; + XGRAMMAR_UNREACHABLE(); + } +}; + +/*************************** Impl of grammar normalizers ***************************/ + +/*! + * \brief Eliminates single-element sequence or choice or character class in the grammar. + * \example `A ::= choices("a")` --> `A ::= "a"` (the body is a string) + * \example `A ::= sequence("a")` --> `A ::= "a"` (the body is a string) + * \example `A ::= [a-a]` --> `A ::= "a"` (the body is a string) + */ +class SingleElementExprEliminator : public GrammarMutator { + public: + using GrammarMutator::Apply; + using GrammarMutator::GrammarMutator; + + private: + int32_t VisitSequence(const GrammarExpr& grammar_expr) final { + std::vector sequence_ids; + for (int32_t i : grammar_expr) { + sequence_ids.push_back(VisitExpr(i)); + } + if (sequence_ids.size() == 1) { + return sequence_ids[0]; + } + return builder_->AddSequence(sequence_ids); + } + + int32_t VisitChoices(const GrammarExpr& grammar_expr) final { + std::vector choice_ids; + for (int32_t i : grammar_expr) { + choice_ids.push_back(VisitExpr(i)); + } + if (choice_ids.size() == 1) { + return choice_ids[0]; + } + return builder_->AddChoices(choice_ids); + } + + int32_t VisitCharacterClass(const GrammarExpr& grammar_expr) final { + if (grammar_expr.data_len == 3 && grammar_expr[0] == 0 && grammar_expr[1] == grammar_expr[2]) { + std::string str = CharToUTF8(grammar_expr[1]); + std::vector bytes; + bytes.reserve(str.size()); + for (char c : str) { + bytes.push_back(static_cast(c)); + } + return builder_->AddByteString(bytes); + } + return builder_->AddGrammarExpr(grammar_expr); + } +}; + +/*! + * \brief Take a grammar from SingleElementExprEliminator and normalize the structure of the + * grammar. + * + * \note The normalized form: + * Each rule should be either: + * - A sequence of choices, each choice is a sequence of elements. Elements can be a character + * class, a byte string, or a rule reference. Only the first choice can be an empty string, + * indicating the rule can be empty. E.g. + * `rule_name ::= ("" | (element1_1 element1_2 ...) | (element2_1 element2_2 ...) | ...)` + * - A macro. Now only TagDispatch is supported. + * + * The lookahead assertion should be a sequence. + * + * New rules may be created to make every rule fit the normalized form. + * + * \example `A ::= ((a) (((b)) (c)) "")` -> `A ::= ((a b c))` + * \example `A ::= (a | (b | (c | "")))` -> `A ::= ("" | (a) | (b) | (c))` + * \example `A ::= (a | (b (c | d)))` -> `A ::= ((a) | (b A_1)), A_1 ::= ((c) | (d))` + * \example `A ::= (a | TagDispatch((tag1, rule1)))` -> `A ::= ((a) | (A_1)), A_1 ::= + * TagDispatch((tag1, rule1))` + */ +class StructureNormalizerImpl : public GrammarMutator { + public: + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + auto grammar_new = SingleElementExprEliminator().Apply(grammar); + InitGrammar(grammar_new); + InitBuilder(); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + builder_->AddEmptyRule(base_grammar_->GetRule(i).name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + cur_rule_name_ = rule.name; + auto new_body_expr_id = VisitRuleBody(grammar_expr); + builder_->UpdateRuleBody(i, new_body_expr_id); + builder_->UpdateLookaheadAssertion(i, VisitLookaheadAssertion(rule.lookahead_assertion_id)); + builder_->UpdateMaxTokens(i, rule.max_tokens); + builder_->UpdateMaxChars(i, rule.max_chars); + builder_->UpdateCaptureName(i, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + builder_->UpdateSuffixStopInfo(i, *suffix_stop_info); + } + builder_->UpdateLazy(i, rule.is_lazy); + builder_->UpdateRuleTemperature(i, rule.temperature); + } + return builder_->Get(base_grammar_->GetRootRule().name); + } + + private: + int32_t VisitLookaheadAssertion(int32_t lookahead_assertion_id) final { + if (lookahead_assertion_id == -1) { + return -1; + } + auto assertion_expr = base_grammar_->GetGrammarExpr(lookahead_assertion_id); + switch (assertion_expr.type) { + case GrammarExprType::kSequence: + return builder_->AddSequence(VisitSequence_(assertion_expr)); + case GrammarExprType::kChoices: + XGRAMMAR_LOG(FATAL) << "Choices in lookahead assertion are not supported yet"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kEmptyStr: + XGRAMMAR_LOG(FATAL) << "Empty string should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kTagDispatch: + XGRAMMAR_LOG(FATAL) << "TagDispatch should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kRegex: + XGRAMMAR_LOG(FATAL) << "Regex should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kSubstring: + XGRAMMAR_LOG(FATAL) << "Substring should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + case GrammarExprType::kTokenTagDispatch: + return builder_->AddSequence({builder_->AddGrammarExpr(assertion_expr)}); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected lookahead assertion type: " + << static_cast(assertion_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + + /*! \brief Visit a GrammarExpr as a rule body. */ + int32_t VisitRuleBody(const GrammarExpr& grammar_expr) { + switch (grammar_expr.type) { + case GrammarExprType::kSequence: + return builder_->AddChoices({builder_->AddSequence(VisitSequence_(grammar_expr))}); + case GrammarExprType::kChoices: + return builder_->AddChoices(VisitChoices_(grammar_expr)); + case GrammarExprType::kEmptyStr: + return builder_->AddChoices({builder_->AddEmptyStr()}); + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + return builder_->AddChoices({builder_->AddSequence({builder_->AddGrammarExpr(grammar_expr)}) + }); + case GrammarExprType::kTagDispatch: + return VisitTagDispatch(grammar_expr); + case GrammarExprType::kTokenTagDispatch: { + auto ttd_expr_id = VisitTokenTagDispatch(grammar_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, ttd_expr_id); + return builder_->AddChoices({builder_->AddSequence({builder_->AddRuleRef(new_rule_id)})}); + } + case GrammarExprType::kRegex: + case GrammarExprType::kSubstring: + // A regex or substring is kept as the direct body of the rule, like a tag dispatch. + return builder_->AddGrammarExpr(grammar_expr); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected sequence type: " << static_cast(grammar_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + + /*! + * \brief Visit a GrammarExpr containing choices. + * \returns A list of new choice GrammarExpr ids. + */ + std::vector VisitChoices_(const GrammarExpr& grammar_expr) { + std::vector new_choice_ids; + bool found_empty = false; + for (auto i : grammar_expr) { + auto choice_expr = base_grammar_->GetGrammarExpr(i); + switch (choice_expr.type) { + case GrammarExprType::kSequence: + VisitSequenceInChoices(choice_expr, &new_choice_ids, &found_empty); + break; + case GrammarExprType::kChoices: + VisitChoicesInChoices(choice_expr, &new_choice_ids, &found_empty); + break; + case GrammarExprType::kEmptyStr: + found_empty = true; + break; + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + VisitElementInChoices(choice_expr, &new_choice_ids); + break; + case GrammarExprType::kTagDispatch: { + auto tag_dispatch_expr_id = VisitTagDispatch(choice_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, tag_dispatch_expr_id); + auto new_sequence_id = builder_->AddSequence({builder_->AddRuleRef(new_rule_id)}); + new_choice_ids.push_back(new_sequence_id); + break; + } + case GrammarExprType::kTokenTagDispatch: { + auto ttd_expr_id = VisitTokenTagDispatch(choice_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, ttd_expr_id); + auto new_sequence_id = builder_->AddSequence({builder_->AddRuleRef(new_rule_id)}); + new_choice_ids.push_back(new_sequence_id); + break; + } + case GrammarExprType::kRegex: + case GrammarExprType::kSubstring: { + auto leaf_expr_id = builder_->AddGrammarExpr(choice_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, leaf_expr_id); + auto new_sequence_id = builder_->AddSequence({builder_->AddRuleRef(new_rule_id)}); + new_choice_ids.push_back(new_sequence_id); + break; + } + default: + XGRAMMAR_LOG(FATAL) << "Unexpected choice type: " << static_cast(choice_expr.type); + } + } + if (found_empty) { + new_choice_ids.insert(new_choice_ids.begin(), builder_->AddEmptyStr()); + } + XGRAMMAR_ICHECK(new_choice_ids.size() >= 1); + return new_choice_ids; + } + + /*! \brief Visit a sequence GrammarExpr that is one of a list of choices. */ + void VisitSequenceInChoices( + const GrammarExpr& grammar_expr, std::vector* new_choice_ids, bool* found_empty + ) { + auto sub_sequence_ids = VisitSequence_(grammar_expr); + if (sub_sequence_ids.size() == 0) { + *found_empty = true; + } else { + new_choice_ids->push_back(builder_->AddSequence(sub_sequence_ids)); + } + } + + /*! \brief Visit a choice GrammarExpr that is one of a list of choices. */ + void VisitChoicesInChoices( + const GrammarExpr& grammar_expr, std::vector* new_choice_ids, bool* found_empty + ) { + auto sub_choice_ids = VisitChoices_(grammar_expr); + bool contains_empty = + builder_->GetGrammarExpr(sub_choice_ids[0]).type == GrammarExprType::kEmptyStr; + if (contains_empty) { + *found_empty = true; + new_choice_ids->insert( + new_choice_ids->end(), sub_choice_ids.begin() + 1, sub_choice_ids.end() + ); + } else { + new_choice_ids->insert(new_choice_ids->end(), sub_choice_ids.begin(), sub_choice_ids.end()); + } + } + + /*! \brief Visit an atom element GrammarExpr that is one of a list of choices. */ + void VisitElementInChoices( + const GrammarExpr& grammar_expr, std::vector* new_choice_ids + ) { + auto sub_expr_id = builder_->AddGrammarExpr(grammar_expr); + new_choice_ids->push_back(builder_->AddSequence({sub_expr_id})); + } + + /*! + * \brief Visit a GrammarExpr containing a sequence. + * \returns A list of new sequence GrammarExpr ids. + */ + std::vector VisitSequence_(const GrammarExpr& grammar_expr) { + std::vector new_sequence_ids; + for (auto i : grammar_expr) { + auto element_expr = base_grammar_->GetGrammarExpr(i); + switch (element_expr.type) { + case GrammarExprType::kSequence: + VisitSequenceInSequence(element_expr, &new_sequence_ids); + break; + case GrammarExprType::kChoices: + VisitChoiceInSequence(element_expr, &new_sequence_ids); + break; + case GrammarExprType::kEmptyStr: + break; + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + VisitElementInSequence(element_expr, &new_sequence_ids); + break; + case GrammarExprType::kTagDispatch: { + auto tag_dispatch_expr_id = VisitTagDispatch(element_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, tag_dispatch_expr_id); + new_sequence_ids.push_back(builder_->AddRuleRef(new_rule_id)); + break; + } + case GrammarExprType::kTokenTagDispatch: { + auto ttd_expr_id = VisitTokenTagDispatch(element_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, ttd_expr_id); + new_sequence_ids.push_back(builder_->AddRuleRef(new_rule_id)); + break; + } + case GrammarExprType::kRegex: + case GrammarExprType::kSubstring: { + auto leaf_expr_id = builder_->AddGrammarExpr(element_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, leaf_expr_id); + new_sequence_ids.push_back(builder_->AddRuleRef(new_rule_id)); + break; + } + default: + XGRAMMAR_LOG(FATAL) << "Unexpected sequence type: " + << static_cast(element_expr.type); + } + } + return new_sequence_ids; + } + + /*! \brief Visit a sequence GrammarExpr that is one element in another sequence. */ + void VisitSequenceInSequence( + const GrammarExpr& grammar_expr, std::vector* new_sequence_ids + ) { + auto sub_sequence_ids = VisitSequence_(grammar_expr); + new_sequence_ids->insert( + new_sequence_ids->end(), sub_sequence_ids.begin(), sub_sequence_ids.end() + ); + } + + /*! \brief Visit a choice GrammarExpr that is one element in a sequence. */ + void VisitChoiceInSequence( + const GrammarExpr& grammar_expr, std::vector* new_sequence_ids + ) { + auto sub_choice_ids = VisitChoices_(grammar_expr); + if (sub_choice_ids.size() == 1) { + auto choice_element_expr = builder_->GetGrammarExpr(sub_choice_ids[0]); + if (choice_element_expr.type != GrammarExprType::kEmptyStr) { + new_sequence_ids->insert( + new_sequence_ids->end(), choice_element_expr.begin(), choice_element_expr.end() + ); + } + } else { + auto new_choice_id = builder_->AddChoices(sub_choice_ids); + auto new_choice_rule_id = builder_->AddRuleWithHint(cur_rule_name_, new_choice_id); + new_sequence_ids->push_back(builder_->AddRuleRef(new_choice_rule_id)); + } + } + + /*! \brief Visit an atom element GrammarExpr that is in a sequence. */ + void VisitElementInSequence( + const GrammarExpr& grammar_expr, std::vector* new_sequence_ids + ) { + new_sequence_ids->push_back(builder_->AddGrammarExpr(grammar_expr)); + } +}; + +/*! + * \brief A class that normalizes a grammar by applying a series of transformations. + * + * The normalizer applies the following transformations in order: + * 1. SingleElementExprEliminator - Eliminates single element expressions + * 2. NestedRuleUnwrapper - Unwraps nested rules + */ +class GrammarNormalizerImpl { + public: + GrammarNormalizerImpl() = default; + + Grammar Apply(const Grammar& grammar) { + auto renamed_grammar = RootRuleRenamer::Apply(grammar); + return StructureNormalizerImpl().Apply(renamed_grammar); + } +}; + +/*************************** Impl of grammar optimizers ***************************/ + +/*! + * \brief Base for optimizer passes that rewrite a grammar in place. It binds a GrammarBuilder + * directly to the target grammar (no copy) and walks every rule body and lookahead assertion. + * \details An expr whose subtree is unchanged keeps its original id; a new expr is appended only + * where a rewrite actually happens, so a pass with nothing to do writes nothing. Changed ids + * propagate upward: a container is rebuilt when any child changed, and a rule body or lookahead is + * updated when its expr changed. Stale exprs left behind are removed later by DeadCodeEliminator. + */ +class InPlaceGrammarRewriter { + public: + using GrammarExpr = Grammar::Impl::GrammarExpr; + using GrammarExprType = Grammar::Impl::GrammarExprType; + + void Apply(Grammar* grammar) { + grammar_ = grammar; + builder_ = GrammarBuilder::FromMutableGrammar(grammar); + // Expr ids are dense, so the memo is a plain array over the original exprs. Appended exprs + // are never visited again, so they need no memo slots. + memo_.assign(builder_.NumGrammarExprs(), -1); + Prepare(); + int32_t num_rules = builder_.NumRules(); + for (int32_t rule_id = 0; rule_id < num_rules; ++rule_id) { + int32_t body_expr_id = builder_.GetRule(rule_id).body_expr_id; + int32_t new_body_expr_id = VisitExpr(body_expr_id); + if (new_body_expr_id != body_expr_id) { + builder_.UpdateRuleBody(rule_id, new_body_expr_id); + } + int32_t lookahead_id = builder_.GetRule(rule_id).lookahead_assertion_id; + if (lookahead_id != -1) { + int32_t new_lookahead_id = VisitExpr(lookahead_id); + if (new_lookahead_id != lookahead_id) { + builder_.UpdateLookaheadAssertion(rule_id, new_lookahead_id); + } + } + } + } + + virtual ~InPlaceGrammarRewriter() = default; + + protected: + /*! \brief Hook run after the builder is bound and before the walk. */ + virtual void Prepare() {} + + /*! \brief Visit an expr; return its rewritten id, or the original id when nothing changed. + * May append new exprs to the arena, invalidating outstanding GrammarExpr views. */ + int32_t VisitExpr(int32_t expr_id) { + XGRAMMAR_DCHECK(expr_id < static_cast(memo_.size())); + if (memo_[expr_id] != -1) { + return memo_[expr_id]; + } + int32_t result; + switch (builder_.GetGrammarExpr(expr_id).type) { + case GrammarExprType::kSequence: + result = VisitSequence(expr_id); + break; + case GrammarExprType::kChoices: + result = VisitChoices(expr_id); + break; + default: + result = expr_id; + } + memo_[expr_id] = result; + return result; + } + + /*! \brief Rebuild a sequence, recursing into each element. Overridden to add rewriting. + * No memory is allocated when no element changes. */ + virtual int32_t VisitSequence(int32_t expr_id) { + auto expr = builder_.GetGrammarExpr(expr_id); + int32_t size = expr.size(); + std::vector new_element_ids; + bool changed = false; + for (int32_t i = 0; i < size; ++i) { + int32_t element_id = expr[i]; + int32_t new_element_id = VisitExpr(element_id); + // The visit may have appended exprs to the arena and invalidated the view, so re-fetch. + expr = builder_.GetGrammarExpr(expr_id); + if (!changed && new_element_id != element_id) { + // Materialize the already scanned, unchanged prefix on the first change. + changed = true; + new_element_ids.reserve(size); + new_element_ids.insert(new_element_ids.end(), expr.begin(), expr.begin() + i); + } + if (changed) { + new_element_ids.push_back(new_element_id); + } + } + if (!changed) { + return expr_id; + } + return builder_.AddSequence(new_element_ids); + } + + /*! \brief Rebuild a choices, recursing into each choice. Overridden to add rewriting. + * No memory is allocated when no choice changes. */ + virtual int32_t VisitChoices(int32_t expr_id) { + auto expr = builder_.GetGrammarExpr(expr_id); + int32_t size = expr.size(); + std::vector new_choice_ids; + bool changed = false; + for (int32_t i = 0; i < size; ++i) { + int32_t choice_id = expr[i]; + int32_t new_choice_id = VisitExpr(choice_id); + // The visit may have appended exprs to the arena and invalidated the view, so re-fetch. + expr = builder_.GetGrammarExpr(expr_id); + if (!changed && new_choice_id != choice_id) { + // Materialize the already scanned, unchanged prefix on the first change. + changed = true; + new_choice_ids.reserve(size); + new_choice_ids.insert(new_choice_ids.end(), expr.begin(), expr.begin() + i); + } + if (changed) { + new_choice_ids.push_back(new_choice_id); + } + } + if (!changed) { + return expr_id; + } + return builder_.AddChoices(new_choice_ids); + } + + GrammarBuilder builder_; + Grammar* grammar_; + // Maps an original expr id to its rewritten id (or itself when unchanged); -1 means unvisited. + std::vector memo_; +}; + +/*! + * \brief Inline rules that can be inlined. + * + * Now we only inline rule references that: + * 1. at the beginning of a sequence + * 2. The rule should be a sequence of choices, cannot be empty, cannot refer to other rules + * + * \details Rewrites the grammar in place: only choices that actually have an inlinable leading + * rule reference are rebuilt, the rest keep their original ids. Inlinability is judged on the + * original grammar so the result does not depend on the order rules are rewritten. Stale exprs are + * removed later by DeadCodeEliminator. + */ +class RuleInlinerImpl : public InPlaceGrammarRewriter { + protected: + void Prepare() override { + int32_t num_rules = builder_.NumRules(); + original_body_expr_ids_.reserve(num_rules); + for (int32_t rule_id = 0; rule_id < num_rules; ++rule_id) { + original_body_expr_ids_.push_back(builder_.GetRule(rule_id).body_expr_id); + } + can_rule_be_inlined_.assign(num_rules, -1); + } + + int32_t VisitChoices(int32_t expr_id) override { + auto expr = builder_.GetGrammarExpr(expr_id); + int32_t size = expr.size(); + std::vector new_choice_ids; + bool changed = false; + for (int32_t i = 0; i < size; ++i) { + int32_t choice_id = expr[i]; + auto choice_expr = builder_.GetGrammarExpr(choice_id); + int32_t inline_rule_id = -1; + if (choice_expr.type == GrammarExprType::kSequence && choice_expr.size() > 0) { + auto first_element = builder_.GetGrammarExpr(choice_expr[0]); + if (first_element.type == GrammarExprType::kRuleRef && CanRuleBeInlined(first_element[0])) { + inline_rule_id = first_element[0]; + } + } + if (inline_rule_id == -1) { + int32_t new_choice_id = VisitExpr(choice_id); + // The visit may have appended exprs to the arena and invalidated the view, so re-fetch. + expr = builder_.GetGrammarExpr(expr_id); + if (!changed && new_choice_id != choice_id) { + // Materialize the already scanned, unchanged prefix on the first change. + changed = true; + new_choice_ids.reserve(size); + new_choice_ids.insert(new_choice_ids.end(), expr.begin(), expr.begin() + i); + } + if (changed) { + new_choice_ids.push_back(new_choice_id); + } + continue; + } + if (!changed) { + changed = true; + new_choice_ids.reserve(size); + new_choice_ids.insert(new_choice_ids.end(), expr.begin(), expr.begin() + i); + } + // Do inlining: prepend each choice of the referenced rule to the rest of this sequence. + // Copy and visit the needed element ids before appending anything. + std::vector rest_element_ids(choice_expr.begin() + 1, choice_expr.end()); + for (int32_t& element_id : rest_element_ids) { + element_id = VisitExpr(element_id); + } + auto ref_body = builder_.GetGrammarExpr(original_body_expr_ids_[inline_rule_id]); + std::vector ref_choice_ids(ref_body.begin(), ref_body.end()); + for (int32_t ref_choice_id : ref_choice_ids) { + auto ref_choice_expr = builder_.GetGrammarExpr(ref_choice_id); + XGRAMMAR_ICHECK(ref_choice_expr.type == GrammarExprType::kSequence); + std::vector new_sequence(ref_choice_expr.begin(), ref_choice_expr.end()); + for (int32_t& element_id : new_sequence) { + element_id = VisitExpr(element_id); + } + new_sequence.insert(new_sequence.end(), rest_element_ids.begin(), rest_element_ids.end()); + new_choice_ids.push_back(builder_.AddSequence(new_sequence)); + } + // The appended sequences invalidated the view, so re-fetch it for the next iteration. + expr = builder_.GetGrammarExpr(expr_id); + } + if (!changed) { + return expr_id; + } + return builder_.AddChoices(new_choice_ids); + } + + private: + /*! \brief A rule can be inlined iff its body is a non-empty choices of sequences that contain no + * rule references. Judged on the original grammar via original_body_expr_ids_. */ + bool CanRuleBeInlined(int32_t rule_id) { + if (can_rule_be_inlined_[rule_id] != -1) { + return can_rule_be_inlined_[rule_id] != 0; + } + const auto& rule = (*grammar_)->GetRule(rule_id); + // Inlining a budgeted rule would erase the rule its length budget applies to. Inlining a + // capture-relevant rule would eliminate its completion events, so its capture or hidden span + // would never be recorded. Inlining a lazy rule would erase its committed-shortest semantics. + // Inlining a temperature rule would erase the rule its sampling temperature applies to. + if (rule.max_tokens >= 0 || rule.max_chars >= 0 || !rule.capture_name.empty() || + (*grammar_)->GetSuffixStopInfo(rule_id) != nullptr || rule.is_lazy || + rule.temperature.has_value()) { + can_rule_be_inlined_[rule_id] = false; + return false; + } + bool result = true; + auto body = builder_.GetGrammarExpr(original_body_expr_ids_[rule_id]); + if (body.type != GrammarExprType::kChoices || body.size() == 0) { + result = false; + } else { + for (int32_t choice_id : body) { + auto choice_expr = builder_.GetGrammarExpr(choice_id); + if (choice_expr.type == GrammarExprType::kEmptyStr) { + result = false; + break; + } + XGRAMMAR_ICHECK(choice_expr.type == GrammarExprType::kSequence); + bool has_rule_ref = false; + for (int32_t element_id : choice_expr) { + if (builder_.GetGrammarExpr(element_id).type == GrammarExprType::kRuleRef) { + has_rule_ref = true; + break; + } + } + if (has_rule_ref) { + result = false; + break; + } + } + } + can_rule_be_inlined_[rule_id] = result ? 1 : 0; + return result; + } + + // Rule body expr ids captured before any rewriting, for order-independent inlinability checks. + std::vector original_body_expr_ids_; + // Per-rule cache of CanRuleBeInlined: -1 unknown, 0 false, 1 true. + std::vector can_rule_be_inlined_; +}; + +/*! + * \brief Analyze all referenced rules or the main rule. Return a list of all referenced rule ids. + * This is useful for dead code elimination. + */ +class UsedRulesAnalyzer : public GrammarVisitor> { + public: + UsedRulesAnalyzer() = default; + + std::vector Apply(const Grammar& grammar) final { + InitGrammar(grammar); + + std::set visited; + + std::queue().swap(visit_queue_); + + visit_queue_.push(base_grammar_->GetRootRuleId()); + while (!visit_queue_.empty()) { + auto rule_id = visit_queue_.front(); + visit_queue_.pop(); + if (visited.count(rule_id)) { + continue; + } + visited.insert(rule_id); + auto rule = base_grammar_->GetRule(rule_id); + VisitExpr(rule.body_expr_id); + if (rule.lookahead_assertion_id != -1) { + VisitExpr(rule.lookahead_assertion_id); + } + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(rule_id); + suffix_stop_info != nullptr && suffix_stop_info->body_rule_id != -1) { + visit_queue_.push(suffix_stop_info->body_rule_id); + visit_queue_.push(suffix_stop_info->marker_rule_id); + } + } + + return std::vector(visited.begin(), visited.end()); + } + + void VisitTagDispatch(const GrammarExpr& grammar_expr) { + auto tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + visit_queue_.push(rule_id); + } + } + + void VisitTokenTagDispatch(const GrammarExpr& grammar_expr) { + auto ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + visit_queue_.push(rule_id); + } + } + + void VisitRuleRef(const GrammarExpr& grammar_expr) { visit_queue_.push(grammar_expr[0]); } + + void VisitRepeat(const GrammarExpr& grammar_expr) { visit_queue_.push(grammar_expr[0]); } + + private: + std::queue visit_queue_; +}; + +class DeadCodeEliminatorImpl : public GrammarMutator { + public: + using GrammarMutator::Apply; + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + InitGrammar(grammar); + InitBuilder(); + auto used_rules = UsedRulesAnalyzer().Apply(grammar); + rule_id_map_.clear(); + for (auto rule_id : used_rules) { + rule_id_map_[rule_id] = builder_->AddEmptyRule(grammar->GetRule(rule_id).name); + } + for (auto rule_id : used_rules) { + auto rule = grammar->GetRule(rule_id); + auto new_body_expr_id = VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(rule_id_map_[rule_id], new_body_expr_id); + builder_->UpdateLookaheadAssertion( + rule_id_map_[rule_id], VisitLookaheadAssertion(rule.lookahead_assertion_id) + ); + builder_->UpdateMaxTokens(rule_id_map_[rule_id], rule.max_tokens); + builder_->UpdateMaxChars(rule_id_map_[rule_id], rule.max_chars); + builder_->UpdateCaptureName(rule_id_map_[rule_id], rule.capture_name); + if (const auto* suffix_stop_info = grammar->GetSuffixStopInfo(rule_id)) { + auto remapped_info = *suffix_stop_info; + if (remapped_info.body_rule_id >= 0) { + remapped_info.body_rule_id = rule_id_map_.at(remapped_info.body_rule_id); + remapped_info.marker_rule_id = rule_id_map_.at(remapped_info.marker_rule_id); + } + builder_->UpdateSuffixStopInfo(rule_id_map_[rule_id], remapped_info); + } + builder_->UpdateLazy(rule_id_map_[rule_id], rule.is_lazy); + builder_->UpdateRuleTemperature(rule_id_map_[rule_id], rule.temperature); + } + XGRAMMAR_CHECK(rule_id_map_.count(grammar->GetRootRuleId()) > 0); + return builder_->Get(rule_id_map_[grammar->GetRootRuleId()]); + } + + int32_t VisitTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TagDispatch tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + XGRAMMAR_DCHECK(rule_id_map_.count(rule_id) > 0); + rule_id = rule_id_map_[rule_id]; + } + return builder_->AddTagDispatch(tag_dispatch); + } + + int32_t VisitTokenTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TokenTagDispatch ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + XGRAMMAR_DCHECK(rule_id_map_.count(rule_id) > 0); + rule_id = rule_id_map_[rule_id]; + } + return builder_->AddTokenTagDispatch(ttd); + } + + int32_t VisitRuleRef(const GrammarExpr& grammar_expr) final { + XGRAMMAR_DCHECK(rule_id_map_.count(grammar_expr[0]) > 0); + auto new_rule_id = rule_id_map_[grammar_expr[0]]; + return builder_->AddRuleRef(new_rule_id); + } + + int32_t VisitRepeat(const GrammarExpr& grammar_expr) final { + XGRAMMAR_DCHECK(rule_id_map_.count(grammar_expr[0]) > 0); + auto new_rule_id = rule_id_map_[grammar_expr[0]]; + return builder_->AddRepeat(new_rule_id, grammar_expr[1], grammar_expr[2]); + } + + private: + std::unordered_map rule_id_map_; +}; + +class LookaheadAssertionAnalyzerImpl : public GrammarMutator { + public: + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + InitGrammar(grammar); + InitBuilder(grammar); + auto root_rule = grammar->GetRootRule(); + auto root_grammar_expr = base_grammar_->GetGrammarExpr(root_rule.body_expr_id); + if (root_grammar_expr.type == GrammarExprType::kTagDispatch || + root_grammar_expr.type == GrammarExprType::kTokenTagDispatch || + root_grammar_expr.type == GrammarExprType::kRegex || + root_grammar_expr.type == GrammarExprType::kSubstring) { + return grammar; + } + BuildRuleLookaheadInfo(); + for (int i = 0; i < static_cast(grammar->NumRules()); ++i) { + auto rule = grammar->GetRule(i); + if (i == grammar->GetRootRuleId()) { + continue; + } + if (rule.lookahead_assertion_id != -1) { + builder_->UpdateLookaheadExact(i, IsExactLookaheadAssertion(i)); + continue; + } + auto look_head_assertion_id = DetectLookaheadAssertion(i); + if (look_head_assertion_id != -1) { + builder_->UpdateLookaheadAssertion(i, look_head_assertion_id); + builder_->UpdateLookaheadExact(i); + } + } + return builder_->Get(grammar->GetRootRuleId()); + } + + bool IsExactLookaheadAssertion(int32_t rule_id) { + XGRAMMAR_DCHECK(base_grammar_->GetRule(rule_id).lookahead_assertion_id != -1); + return CanUseDerivedLookahead(rule_id); + } + + int32_t DetectLookaheadAssertion(int32_t rule_id) { + if (!CanUseDerivedLookahead(rule_id)) { + return -1; + } + return builder_->AddSequence(rule_lookahead_infos_[rule_id].suffix_after_first_occurrence); + } + + private: + struct RuleLookaheadInfo { + bool is_triggered_by_dispatch = false; + bool appears_as_last_in_other_rule = false; + int non_last_occurrence_count = 0; + std::vector suffix_after_first_occurrence; + }; + + bool CanUseDerivedLookahead(int32_t rule_id) const { + const auto& info = rule_lookahead_infos_[rule_id]; + return !info.is_triggered_by_dispatch && !info.appears_as_last_in_other_rule && + info.non_last_occurrence_count == 1; + } + + void BuildRuleLookaheadInfo() { + rule_lookahead_infos_.assign(base_grammar_->NumRules(), RuleLookaheadInfo{}); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (grammar_expr.type == GrammarExprType::kTagDispatch) { + auto tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + rule_lookahead_infos_[rule_id].is_triggered_by_dispatch = true; + } + continue; + } + if (grammar_expr.type == GrammarExprType::kTokenTagDispatch) { + auto token_tag_dispatch = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (const auto& [token_id, rule_id] : token_tag_dispatch.trigger_rule_pairs) { + rule_lookahead_infos_[rule_id].is_triggered_by_dispatch = true; + } + continue; + } + if (grammar_expr.type == GrammarExprType::kRegex || + grammar_expr.type == GrammarExprType::kSubstring) { + // A regex or substring rule is a leaf: it references no other rules. + continue; + } + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kChoices); + for (auto sequence_id : grammar_expr) { + auto sequence_expr = base_grammar_->GetGrammarExpr(sequence_id); + if (sequence_expr.type != GrammarExprType::kSequence || sequence_expr.size() == 0) { + continue; + } + auto last_element = base_grammar_->GetGrammarExpr(sequence_expr.end()[-1]); + if (last_element.type == GrammarExprType::kRuleRef && i != last_element[0]) { + rule_lookahead_infos_[last_element[0]].appears_as_last_in_other_rule = true; + } + for (int j = 0; j < sequence_expr.size() - 1; ++j) { + auto element_expr = base_grammar_->GetGrammarExpr(sequence_expr[j]); + if (element_expr.type != GrammarExprType::kRuleRef) { + continue; + } + auto& info = rule_lookahead_infos_[element_expr[0]]; + if (info.non_last_occurrence_count == 0) { + info.suffix_after_first_occurrence.assign( + sequence_expr.begin() + j + 1, sequence_expr.end() + ); + } + ++info.non_last_occurrence_count; + } + } + } + } + + std::vector rule_lookahead_infos_; +}; + +/*! + * \brief Finds the rule reference graph of a grammar. + * + * The rule reference graph shows which rules reference which other rules. + * The returned graph is inverted: it points from referee to referer. + */ +class RuleRefGraphFinder : public GrammarVisitor>> { + public: + RuleRefGraphFinder() = default; + + std::vector> Apply(const Grammar& grammar) { + InitGrammar(grammar); + rule_visit_graph_ = std::vector>(base_grammar_->NumRules()); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + cur_rule_id_ = i; + VisitExpr(grammar_expr); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + std::sort(rule_visit_graph_[i].begin(), rule_visit_graph_[i].end()); + auto end_it = std::unique(rule_visit_graph_[i].begin(), rule_visit_graph_[i].end()); + rule_visit_graph_[i].erase(end_it, rule_visit_graph_[i].end()); + } + return std::move(rule_visit_graph_); + } + + private: + void VisitRuleRef(const GrammarExpr& grammar_expr) { + rule_visit_graph_[grammar_expr[0]].push_back(cur_rule_id_); + } + + void VisitRepeat(const GrammarExpr& grammar_expr) { + rule_visit_graph_[grammar_expr[0]].push_back(cur_rule_id_); + } + + void VisitTagDispatch(const GrammarExpr& grammar_expr) { + auto tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + rule_visit_graph_[rule_id].push_back(cur_rule_id_); + } + } + + void VisitTokenTagDispatch(const GrammarExpr& grammar_expr) { + auto ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + rule_visit_graph_[rule_id].push_back(cur_rule_id_); + } + } + + // Inversed reference graph: pointing from referee to referer + std::vector> rule_visit_graph_; + int32_t cur_rule_id_; +}; + +/*! + * \brief Analyzes which rules in a grammar can match the empty string. + */ +class AllowEmptyRuleAnalyzerImpl : public GrammarVisitor> { + public: + AllowEmptyRuleAnalyzerImpl() = default; + + std::vector Apply(const Grammar& grammar) final { + InitGrammar(grammar); + + // Step 1: Find rules that explicitly allow empty string + std::unordered_set empty_rule_id_set; + FindExplicitEmptyRules(&empty_rule_id_set); + + // Step 2: Find rules that indirectly allow empty string. Using the Bellman-Ford algorithm + // on the rule reference graph. + std::vector> rule_ref_graph = RuleRefGraphFinder().Apply(grammar); + FindIndirectEmptyRules(&empty_rule_id_set, rule_ref_graph); + + auto result = std::vector(empty_rule_id_set.begin(), empty_rule_id_set.end()); + std::sort(result.begin(), result.end()); + return result; + } + + void FindExplicitEmptyRules(std::unordered_set* empty_rule_id_set) { + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (grammar_expr.type == GrammarExprType::kTagDispatch || + grammar_expr.type == GrammarExprType::kTokenTagDispatch) { + empty_rule_id_set->insert(i); + continue; + } + + if (grammar_expr.type == GrammarExprType::kSubstring) { + // A substring automaton always accepts the empty string: every state is accepting. + empty_rule_id_set->insert(i); + continue; + } + + if (grammar_expr.type == GrammarExprType::kRegex) { + // Nullability is checked syntactically on the parsed regex, so no FSM is built here. + // Parse errors are reported by GrammarFSMBuilder later. + auto matches_empty_result = + RegexFSMBuilder::MatchesEmpty(base_grammar_->GetRegexString(grammar_expr)); + if (matches_empty_result.IsOk() && std::move(matches_empty_result).Unwrap()) { + empty_rule_id_set->insert(i); + } + continue; + } + + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kChoices); + if (base_grammar_->GetGrammarExpr(grammar_expr[0]).type == GrammarExprType::kEmptyStr) { + empty_rule_id_set->insert(i); + continue; + } + + for (auto seq_id : grammar_expr) { + auto seq_expr = base_grammar_->GetGrammarExpr(seq_id); + if (std::all_of(seq_expr.begin(), seq_expr.end(), [&](int32_t i) { + return base_grammar_->GetGrammarExpr(i).type == GrammarExprType::kCharacterClassStar; + })) { + empty_rule_id_set->insert(i); + break; + } + } + } + } + + bool SeqExprIsEpsilon( + const GrammarExpr& seq_expr, const std::unordered_set& empty_rule_id_set + ) { + if (seq_expr.type == GrammarExprType::kEmptyStr) { + return true; + } + XGRAMMAR_DCHECK(seq_expr.type == GrammarExprType::kSequence); + + return std::all_of(seq_expr.begin(), seq_expr.end(), [&](int32_t i) { + auto element_expr = base_grammar_->GetGrammarExpr(i); + return (element_expr.type == GrammarExprType::kRuleRef && + empty_rule_id_set.count(element_expr[0])) || + element_expr.type == GrammarExprType::kCharacterClassStar || + (element_expr.type == GrammarExprType::kRepeat && + (empty_rule_id_set.count(element_expr[0]) || element_expr[1] == 0)); + }); + } + + void FindIndirectEmptyRules( + std::unordered_set* empty_rule_id_set, + const std::vector>& rule_ref_graph + ) { + std::queue queue; + for (auto i : *empty_rule_id_set) { + queue.push(i); + } + + while (!queue.empty()) { + auto rule_id = queue.front(); + queue.pop(); + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(rule_ref_graph.size())); + for (auto referer_rule_id : rule_ref_graph[rule_id]) { + if (empty_rule_id_set->count(referer_rule_id)) { + continue; + } + auto rule = base_grammar_->GetRule(referer_rule_id); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + + XGRAMMAR_DCHECK( + grammar_expr.type != GrammarExprType::kTagDispatch && + grammar_expr.type != GrammarExprType::kTokenTagDispatch + ) << "TagDispatch rules should already exist in empty_rule_id_set"; + + bool is_epsilon = std::any_of(grammar_expr.begin(), grammar_expr.end(), [&](int32_t i) { + auto seq_expr = base_grammar_->GetGrammarExpr(i); + return SeqExprIsEpsilon(seq_expr, *empty_rule_id_set); + }); + + if (is_epsilon) { + empty_rule_id_set->insert(referer_rule_id); + queue.push(referer_rule_id); + } + } + } + } +}; + +class GrammarFSMBuilderImpl { + public: + explicit GrammarFSMBuilderImpl( + FSM& target_fsm, + const std::string* rule_name = nullptr, + GrammarBuilder* grammar_builder = nullptr + ) + : target_fsm_(target_fsm), rule_name_(rule_name), grammar_builder_(grammar_builder) {} + + static void Apply(Grammar* grammar) { + FSM complete_fsm; + std::vector> per_rule_fsms; + std::vector state_mapping; + + // Compiling a kRegex rule may append new rules through this builder: large bounded + // repetitions in a regex become kRepeatRef edges referencing new rules. NumRules() is + // re-read every iteration so the new rules get their FSMs built as well. + int32_t num_original_rules = (*grammar)->NumRules(); + GrammarBuilder grammar_builder = GrammarBuilder::FromMutableGrammar(grammar); + for (int i = 0; i < (*grammar)->NumRules(); ++i) { + auto rule_fsm = BuildRuleFSM(*grammar, i, &grammar_builder); + per_rule_fsms.push_back(rule_fsm.AddToCompleteFSM(&complete_fsm, &state_mapping)); + } + + // The rules created during FSM building missed the AllowEmptyRuleAnalyzer pass; complete + // allow_empty_rule_ids for them. New rule ids are larger than all existing ids, so + // appending keeps the list sorted. + for (int i = num_original_rules; i < (*grammar)->NumRules(); ++i) { + const auto& rule = (*grammar)->GetRule(i); + const auto& body_expr = (*grammar)->GetGrammarExpr(rule.body_expr_id); + XGRAMMAR_DCHECK(body_expr.type == Grammar::Impl::GrammarExprType::kRegex); + auto matches_empty_result = + RegexFSMBuilder::MatchesEmpty((*grammar)->GetRegexString(body_expr)); + if (matches_empty_result.IsOk() && std::move(matches_empty_result).Unwrap()) { + (*grammar)->allow_empty_rule_ids.push_back(i); + } + } + + for (int i = 0; i < (*grammar)->NumRules(); ++i) { + XGRAMMAR_DCHECK(per_rule_fsms[i].has_value()) + << "Rule " << i << " (" << (*grammar)->GetRule(i).name + << ") does not have an FSM after optimization"; + } + + // Compress to compact fsm + CompactFSM compact_complete_fsm = complete_fsm.ToCompact(); + std::vector> compact_per_rule_fsms( + (*grammar)->NumRules() + ); + for (int i = 0; i < (*grammar)->NumRules(); ++i) { + if (per_rule_fsms[i]) { + auto compact_fsm_with_se = CompactFSMWithStartEnd( + compact_complete_fsm, + per_rule_fsms[i]->GetFsm().GetStart(), + per_rule_fsms[i]->GetFsm().GetEnds() + ); + compact_per_rule_fsms[i] = CompactFSMWithStartEndWithSize( + compact_fsm_with_se, per_rule_fsms[i]->GetEdgeNum(), per_rule_fsms[i]->GetNodeNum() + ); + } + } + + (*grammar)->complete_fsm = std::move(compact_complete_fsm); + (*grammar)->per_rule_fsms = std::move(compact_per_rule_fsms); + } + + /* Basic Building functions.*/ + static FSMWithStartEnd RuleRef(const GrammarExpr& expr); + static FSMWithStartEnd CharacterClass(const GrammarExpr& expr); + static FSMWithStartEnd ByteString(const GrammarExpr& expr); + static FSMWithStartEnd Token(const GrammarExpr& expr); + static FSMWithStartEnd ExcludeToken(const GrammarExpr& expr); + static std::optional TokenTagDispatch(const Grammar::Impl::TokenTagDispatch& ttd + ); + static std::optional Sequence(const GrammarExpr& expr, const Grammar& grammar); + static std::optional Choices(const GrammarExpr& expr, const Grammar& grammar); + static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); + static Result Regex(const std::string& regex, bool json_string = false); + /* Building tool functions.*/ + static std::optional BuildTagDispatchFSM( + const std::vector>& string_trigger_rules, + bool loop_after_dispatch, + const std::vector& excluded_strings + ); + + private: + static FSMWithStartEnd BuildRuleFSM( + const Grammar& grammar, int rule_id, GrammarBuilder* grammar_builder = nullptr + ); + static FSMWithStartEnd BuildExpressionFSM( + const GrammarExpr& expr, + const Grammar& grammar, + const std::string* rule_name = nullptr, + GrammarBuilder* grammar_builder = nullptr + ); + void BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void BuildEmptyString(int start_state, std::vector* end_states); + void BuildByteString(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildRuleRef(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildCharacterClassStar( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildRepeat(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildToken(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildExcludeToken( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildRegex( + const std::string& regex, bool json_string, int start_state, std::vector* end_states + ); + void BuildSubstring( + const std::vector& chunks, int start_state, std::vector* end_states + ); + void BuildTagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch, + int start_state, + std::vector* end_states + ); + void BuildTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch, + int start_state, + std::vector* end_states + ); + void BuildSequence( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void AddCharacterClassTransitions(const GrammarExpr& expr, int start_state, int end_state); + void BuildNegativeCharacterClass(const GrammarExpr& expr, int start_state, int end_state); + void AppendFSM(FSMWithStartEnd fsm, int start_state, std::vector* end_states); + void AddCharacterRange(int from, int to, uint32_t min, uint32_t max); + + FSM& target_fsm_; + const std::string* rule_name_; + // If not null, kRegex expressions may add new rules through this builder; see Apply(). + GrammarBuilder* grammar_builder_ = nullptr; +}; + +// This function will add a range [min, max] of unicode characters to the FSM. +void GrammarFSMBuilderImpl::AddCharacterRange(int from, int to, uint32_t min, uint32_t max) { + AddPackedUTF8RangeEdges(target_fsm_, from, to, min, max); +} + +void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( + const GrammarExpr& expr, int start_state, int end_state +) { + XGRAMMAR_DCHECK( + expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar + ); + XGRAMMAR_DCHECK(expr[0]); // Negative character class should be true. + std::bitset<128> char_set; + for (int i = 1; i < static_cast(expr.size()); i += 2) { + uint8_t byte_min = static_cast(expr[i]); + uint8_t byte_max = static_cast(expr[i + 1]); + if (byte_max > 128) { + XGRAMMAR_LOG(WARNING) << "Negative Character class contains byte greater than 127, " + << "clamping to 127."; + byte_max = 127; + } + for (uint8_t j = byte_min; j <= byte_max; ++j) { + char_set.set(j); + } + } + + int left_bound = -1; + for (int i = 0; i < 128; ++i) { + if (!char_set[i]) { + left_bound = i; + int right_bound = i + 1; + while (right_bound < 128 && !char_set[right_bound]) { + right_bound++; + } + target_fsm_.AddEdge( + start_state, + end_state, + static_cast(left_bound), + static_cast(right_bound - 1) + ); + i = right_bound; + } + } + AddCharacterRange(start_state, end_state, kMin2BytesUnicode, kMax4BytesUnicode); +} + +void GrammarFSMBuilderImpl::AddCharacterClassTransitions( + const GrammarExpr& expr, int start_state, int end_state +) { + XGRAMMAR_DCHECK( + expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar + ); + bool is_negative = expr[0]; + if (is_negative) { + BuildNegativeCharacterClass(expr, start_state, end_state); + } else { + for (int i = 1; i < static_cast(expr.size()); i += 2) { + uint32_t codepoint_min = static_cast(expr[i]); + uint32_t codepoint_max = static_cast(expr[i + 1]); + // Convert Unicode codepoints to packed UTF-8 format for AddCharacterRange + uint32_t packed_min = CodepointToPackedUTF8(codepoint_min); + uint32_t packed_max = CodepointToPackedUTF8(codepoint_max); + AddCharacterRange(start_state, end_state, packed_min, packed_max); + } + } +} + +void GrammarFSMBuilderImpl::BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kCharacterClass); + end_states->clear(); + int end_state = target_fsm_.AddState(); + AddCharacterClassTransitions(expr, start_state, end_state); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildCharacterClassStar( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kCharacterClassStar); + end_states->clear(); + AddCharacterClassTransitions(expr, start_state, start_state); + end_states->push_back(start_state); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::BuildRuleFSM( + const Grammar& grammar, int rule_id, GrammarBuilder* grammar_builder +) { + const auto& rule = grammar->GetRule(rule_id); + return BuildExpressionFSM( + grammar->GetGrammarExpr(rule.body_expr_id), grammar, &rule.name, grammar_builder + ); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::BuildExpressionFSM( + const GrammarExpr& expr, + const Grammar& grammar, + const std::string* rule_name, + GrammarBuilder* grammar_builder +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm, rule_name, grammar_builder); + builder.BuildExpression(expr, grammar, start_state, &end_states); + FSMWithStartEnd result(result_fsm, start_state, std::move(end_states)); + if (expr.type != ExprType::kTagDispatch && expr.type != ExprType::kTokenTagDispatch) { + result = result.SimplifyEpsilon(); + result = result.MergeEquivalentStates(); + } + return result; +} + +FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildRuleRef(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + if (expr.type == ExprType::kCharacterClassStar) { + builder.BuildCharacterClassStar(expr, start_state, &end_states); + } else { + builder.BuildCharacterClass(expr, start_state, &end_states); + } + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildByteString(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::Token(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildToken(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::ExcludeToken(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildExcludeToken(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +std::optional GrammarFSMBuilderImpl::TokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildTokenTagDispatch(token_tag_dispatch, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +void GrammarFSMBuilderImpl::BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states +) { + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + switch (expr.type) { + case ExprType::kEmptyStr: + return BuildEmptyString(start_state, end_states); + case ExprType::kByteString: + return BuildByteString(expr, start_state, end_states); + case ExprType::kCharacterClass: + return BuildCharacterClass(expr, start_state, end_states); + case ExprType::kCharacterClassStar: + return BuildCharacterClassStar(expr, start_state, end_states); + case ExprType::kRuleRef: + return BuildRuleRef(expr, start_state, end_states); + case ExprType::kRepeat: + return BuildRepeat(expr, start_state, end_states); + case ExprType::kToken: + return BuildToken(expr, start_state, end_states); + case ExprType::kExcludeToken: + return BuildExcludeToken(expr, start_state, end_states); + case ExprType::kSequence: + return BuildSequence(expr, grammar, start_state, end_states); + case ExprType::kChoices: + return BuildChoices(expr, grammar, start_state, end_states); + case ExprType::kRegex: + return BuildRegex( + grammar->GetRegexString(expr), + grammar->GetRegexIsJSONString(expr), + start_state, + end_states + ); + case ExprType::kSubstring: + return BuildSubstring(grammar->GetSubstringChunks(expr), start_state, end_states); + case ExprType::kTagDispatch: + return BuildTagDispatch(grammar->GetTagDispatch(expr), start_state, end_states); + case ExprType::kTokenTagDispatch: + return BuildTokenTagDispatch(grammar->GetTokenTagDispatch(expr), start_state, end_states); + } + XGRAMMAR_UNREACHABLE(); +} + +void GrammarFSMBuilderImpl::BuildEmptyString(int start_state, std::vector* end_states) { + end_states->clear(); + end_states->push_back(start_state); +} + +void GrammarFSMBuilderImpl::BuildByteString( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kByteString); + end_states->clear(); + int current_state = start_state; + for (int32_t byte : expr) { + int next_state = target_fsm_.AddState(); + target_fsm_.AddEdge( + current_state, next_state, static_cast(byte), static_cast(byte) + ); + current_state = next_state; + } + end_states->push_back(current_state); +} + +void GrammarFSMBuilderImpl::BuildRuleRef( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kRuleRef); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRuleEdge(start_state, end_state, expr[0]); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildRepeat( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kRepeat); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRepeatEdge(start_state, end_state, expr[0], expr[1], expr[2]); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildToken( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kToken); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddTokenEdge(start_state, end_state, std::vector(expr.begin(), expr.end())); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildExcludeToken( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kExcludeToken); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddExcludeTokenEdge( + start_state, end_state, std::vector(expr.begin(), expr.end()) + ); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildSequence( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states +) { + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + end_states->clear(); + if (expr.size() == 0) { + end_states->push_back(start_state); + return; + } + + BuildExpression(grammar->GetGrammarExpr(expr[0]), grammar, start_state, end_states); + + std::vector next_end_states; + for (int index = 1; index < static_cast(expr.size()); ++index) { + int element_start = target_fsm_.AddState(); + BuildExpression(grammar->GetGrammarExpr(expr[index]), grammar, element_start, &next_end_states); + for (int32_t previous_end_state : *end_states) { + target_fsm_.AddEpsilonEdge(previous_end_state, element_start); + } + end_states->swap(next_end_states); + } +} + +std::optional GrammarFSMBuilderImpl::Sequence( + const GrammarExpr& expr, const Grammar& grammar +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildSequence(expr, grammar, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +void GrammarFSMBuilderImpl::BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kChoices); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + end_states->clear(); + + int non_empty_choice_count = 0; + bool nullable = false; + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type == ExprType::kEmptyStr) { + nullable = true; + } else { + ++non_empty_choice_count; + } + } + + if (non_empty_choice_count == 0) { + end_states->push_back(start_state); + return; + } + + if (non_empty_choice_count == 1 && !nullable) { + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type != ExprType::kEmptyStr) { + BuildExpression(choice_expr, grammar, start_state, end_states); + return; + } + } + XGRAMMAR_UNREACHABLE(); + } + + std::vector branch_end_states; + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type == ExprType::kEmptyStr) { + continue; + } + int branch_start_state = target_fsm_.AddState(); + BuildExpression(choice_expr, grammar, branch_start_state, &branch_end_states); + target_fsm_.AddEpsilonEdge(start_state, branch_start_state); + end_states->insert(end_states->end(), branch_end_states.begin(), branch_end_states.end()); + } + + if (nullable) { + int nullable_branch_state = target_fsm_.AddState(); + target_fsm_.AddEpsilonEdge(start_state, nullable_branch_state); + end_states->push_back(nullable_branch_state); + } +} + +std::optional GrammarFSMBuilderImpl::Choices( + const GrammarExpr& expr, const Grammar& grammar +) { + return BuildExpressionFSM(expr, grammar); +} + +void GrammarFSMBuilderImpl::AppendFSM( + FSMWithStartEnd fsm, int start_state, std::vector* end_states +) { + const bool target_is_empty = target_fsm_.NumStates() == 1 && start_state == 0 && + target_fsm_.GetEdges(0).empty() && + target_fsm_.GetEdgeAuxData().empty() && fsm.GetStart() == 0; + if (target_is_empty) { + *end_states = fsm.GetEnds(); + target_fsm_ = std::move(fsm.GetFsm()); + return; + } + + std::vector state_mapping; + target_fsm_.AddFSM(fsm.GetFsm(), &state_mapping); + target_fsm_.AddEpsilonEdge(start_state, state_mapping[fsm.GetStart()]); + end_states->clear(); + end_states->reserve(fsm.GetEnds().size()); + for (int end_state : fsm.GetEnds()) { + end_states->push_back(state_mapping[end_state]); + } +} + +void GrammarFSMBuilderImpl::BuildRegex( + const std::string& regex, bool json_string, int start_state, std::vector* end_states +) { + const std::string rule_hint = rule_name_ != nullptr ? *rule_name_ : ""; + auto build_result = + json_string + ? RegexFSMBuilder::BuildWithForbiddenChars( + regex, GrammarFSMBuilder::JSONStringForbiddenChars(), grammar_builder_, rule_hint + ) + : RegexFSMBuilder::Build(regex, grammar_builder_, rule_hint); + if (build_result.IsErr()) { + auto error = std::move(build_result).UnwrapErr(); + if (rule_name_ != nullptr) { + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for rule " << *rule_name_ + << " with regex " << regex << ": " << error.what(); + } + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for regex " << regex << ": " + << error.what(); + } + AppendFSM(std::move(build_result).Unwrap(), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildSubstring( + const std::vector& chunks, int start_state, std::vector* end_states +) { + AppendFSM(SuffixAutomata::Build(chunks), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildTagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch, + int start_state, + std::vector* end_states +) { + auto build_result = TagDispatch(tag_dispatch); + XGRAMMAR_CHECK(build_result.has_value()) << "Failed to build tag dispatch FSM"; + AppendFSM(std::move(*build_result), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch, + int start_state, + std::vector* end_states +) { + int trigger_count = static_cast(token_tag_dispatch.trigger_rule_pairs.size()); + std::vector dispatch_states; + dispatch_states.reserve(trigger_count); + for (int index = 0; index < trigger_count; ++index) { + dispatch_states.push_back(target_fsm_.AddState()); + } + + int end_state = -1; + end_states->clear(); + end_states->push_back(start_state); + if (!token_tag_dispatch.loop_after_dispatch) { + end_state = target_fsm_.AddState(); + end_states->push_back(end_state); + } + + std::vector excluded_token_ids; + excluded_token_ids.reserve( + token_tag_dispatch.trigger_rule_pairs.size() + token_tag_dispatch.excludes.size() + ); + for (const auto& trigger_rule_pair : token_tag_dispatch.trigger_rule_pairs) { + excluded_token_ids.push_back(trigger_rule_pair.first); + } + excluded_token_ids.insert( + excluded_token_ids.end(), + token_tag_dispatch.excludes.begin(), + token_tag_dispatch.excludes.end() + ); + std::sort(excluded_token_ids.begin(), excluded_token_ids.end()); + excluded_token_ids.erase( + std::unique(excluded_token_ids.begin(), excluded_token_ids.end()), excluded_token_ids.end() + ); + + for (int index = 0; index < trigger_count; ++index) { + auto [token_id, rule_id] = token_tag_dispatch.trigger_rule_pairs[index]; + int dispatch_target = token_tag_dispatch.loop_after_dispatch ? start_state : end_state; + target_fsm_.AddTokenEdge(start_state, dispatch_states[index], {token_id}); + target_fsm_.AddRuleEdge(dispatch_states[index], dispatch_target, rule_id); + } + target_fsm_.AddExcludeTokenEdge(start_state, start_state, excluded_token_ids); +} + +std::optional GrammarFSMBuilderImpl::BuildTagDispatchFSM( + const std::vector>& string_trigger_rules, + bool loop_after_dispatch, + const std::vector& excluded_strings +) { + std::vector tag_names; + tag_names.reserve(string_trigger_rules.size()); + for (const auto& [tag_name, tag_id] : string_trigger_rules) { + tag_names.push_back(tag_name); + } + std::vector end_states; + auto trie_result = TrieFSMBuilder::Build(tag_names, excluded_strings, &end_states, true, true); + if (!trie_result.has_value()) { + return std::nullopt; + } + auto trie_fsm = trie_result->GetFsm(); + auto start = trie_result->GetStart(); + + // The final end states are all but the trie's original end states. + std::vector ends; + for (int i = 0; i < trie_fsm.NumStates(); i++) { + if (!trie_result->IsEndState(i)) { + ends.push_back(i); + } + } + + // Add rule ref edges for string triggers + for (int i = 0; i < static_cast(string_trigger_rules.size()); i++) { + int next_state; + if (loop_after_dispatch) { + next_state = start; + } else { + next_state = trie_fsm.AddState(); + ends.push_back(next_state); + } + trie_fsm.AddRuleEdge(end_states[i], next_state, string_trigger_rules[i].second); + } + + return FSMWithStartEnd(trie_fsm, start, std::move(ends)); +} + +std::optional GrammarFSMBuilderImpl::TagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch +) { + std::vector> string_trigger_rules( + tag_dispatch.tag_rule_pairs.begin(), tag_dispatch.tag_rule_pairs.end() + ); + + return BuildTagDispatchFSM( + string_trigger_rules, tag_dispatch.loop_after_dispatch, tag_dispatch.excludes + ); +} + +Result GrammarFSMBuilderImpl::Regex(const std::string& regex, bool json_string) { + auto build_result = json_string ? RegexFSMBuilder::BuildWithForbiddenChars( + regex, GrammarFSMBuilder::JSONStringForbiddenChars() + ) + : RegexFSMBuilder::Build(regex); + if (build_result.IsErr()) { + return build_result; + } + auto result = std::move(build_result).Unwrap(); + result = result.SimplifyEpsilon(); + result = result.MergeEquivalentStates(); + return ResultOk(std::move(result)); +} + +class RepetitionRangeExpanderImpl : public GrammarMutator { + public: + using GrammarMutator::Apply; + using GrammarMutator::GrammarMutator; + + private: + int32_t VisitRepeat(const GrammarExpr& grammar_expr) final { + int32_t ref_rule_id = grammar_expr[0]; + int64_t lower = grammar_expr[1]; + int64_t upper = grammar_expr[2]; + return HandleRepetitionRange(cur_rule_name_, ref_rule_id, lower, upper); + } + + /*! + * \brief Handle repetition range by unzipping into explicit sequence/choice (for small bounds). + * \param cur_rule_name Name hint for generated rules. + * \param grammar_expr_id The expression to repeat. + * \param lower Minimum count (inclusive). + * \param upper Maximum count (inclusive), or -1 for unbounded. + * \return grammar_expr_id of the repetition result. + */ + int32_t LegacyHandleRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper + ); + + /*! + * \brief Handle repetition range {lower, upper}, using unzip for small bounds or kRepeat for + * large. Identical repetitions are expanded only once and shared via memoization. + * \param cur_rule_name Name hint for generated rules. + * \param rule_id The rule to repeat. + * \param lower Minimum count (inclusive). + * \param upper Maximum count (inclusive), or -1 for unbounded. + * \return grammar_expr_id of the repetition result. + */ + int32_t HandleRepetitionRange( + const std::string& cur_rule_name, int32_t rule_id, int64_t lower, int64_t upper + ); + + /*! + * \brief Expand a repetition range into rules. Called by HandleRepetitionRange on cache miss. + * \param cur_rule_name Name hint for generated rules. + * \param grammar_expr_id The expression to repeat. + * \param lower Minimum count (inclusive). + * \param upper Maximum count (inclusive), or -1 for unbounded. + * \return grammar_expr_id of the repetition result. + */ + int32_t ExpandRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper + ); + + /*! + * \brief Memoization of expanded repetitions, mapping (content of the repeated expr, lower, + * upper) to the resulting grammar_expr_id. + * + * Grammars may contain a large number of identical repetitions. E.g. a JSON schema converted + * with max_whitespace_cnt emits one [ \n\r\t]{0,n} repetition per whitespace position, so a + * schema with 50k properties produces 200k+ identical repetitions. Expanding each occurrence + * into its own chain of rules multiplies the rule count by more than an order of magnitude, + * which blows up all downstream compilation stages (FSM building, token mask cache) in both + * time and memory. Sharing one expansion among identical repetitions keeps the rule count + * linear in the schema size. + */ + std::map, int32_t> repetition_cache_; +}; + +/****************** Repetition range helpers ******************/ + +int32_t RepetitionRangeExpanderImpl::LegacyHandleRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper +) { + // Construct expr expr ... expr (l times) + + std::vector elements; + for (int64_t i = 0; i < lower; ++i) { + elements.push_back(grammar_expr_id); + } + + // Case 1: {l}: + // expr expr ... expr (l times) + if (upper == lower) { + auto result_rule_id = builder_->AddRuleWithHint( + cur_rule_name, builder_->AddChoices({builder_->AddSequence(elements)}) + ); + return builder_->AddRuleRef(result_rule_id); + } + + // Case 2: {l,}: + // expr expr ... expr (l times) rest + // rest ::= "" | expr rest + if (upper == -1) { + auto new_rule_name = builder_->GetNewRuleName(cur_rule_name); + auto new_rule_id = builder_->AddEmptyRule(new_rule_name); + auto ref_to_new_rule = builder_->AddRuleRef(new_rule_id); + auto new_grammar_expr_id = builder_->AddChoices( + {builder_->AddEmptyStr(), builder_->AddSequence({grammar_expr_id, ref_to_new_rule})} + ); + builder_->UpdateRuleBody(new_rule_id, new_grammar_expr_id); + elements.push_back(builder_->AddRuleRef(new_rule_id)); + auto result_rule_id = builder_->AddRuleWithHint( + cur_rule_name, builder_->AddChoices({builder_->AddSequence(elements)}) + ); + return builder_->AddRuleRef(result_rule_id); + } + + // Case 3: {l, r} (r - l >= 1) + // expr expr ... expr (l times) rest1 + // rest1 ::= "" | expr rest2 + // rest2 ::= "" | expr rest3 + // ... + // rest(r - l) ::= "" | expr + std::vector rest_rule_ids; + + for (int64_t i = 0; i < upper - lower; ++i) { + auto new_rule_name = builder_->GetNewRuleName(cur_rule_name); + rest_rule_ids.push_back(builder_->AddEmptyRule(new_rule_name)); + } + for (int64_t i = 0; i < upper - lower - 1; ++i) { + auto ref_to_next_rule = builder_->AddRuleRef(rest_rule_ids[i + 1]); + auto new_grammar_expr_id = builder_->AddChoices( + {builder_->AddEmptyStr(), builder_->AddSequence({grammar_expr_id, ref_to_next_rule})} + ); + builder_->UpdateRuleBody(rest_rule_ids[i], new_grammar_expr_id); + } + auto last_grammar_expr_id = + builder_->AddChoices({builder_->AddEmptyStr(), builder_->AddSequence({grammar_expr_id})}); + builder_->UpdateRuleBody(rest_rule_ids.back(), last_grammar_expr_id); + + elements.push_back(builder_->AddRuleRef(rest_rule_ids[0])); + auto result_rule_id = builder_->AddRuleWithHint( + cur_rule_name, builder_->AddChoices({builder_->AddSequence(elements)}) + ); + return builder_->AddRuleRef(result_rule_id); +} + +int32_t RepetitionRangeExpanderImpl::HandleRepetitionRange( + const std::string& cur_rule_name, int32_t rule_id, int64_t lower, int64_t upper +) { + // Check if the referred rule is only one single element. If so, we can directly use the element + // for further optimization. + int32_t grammar_expr_id = builder_->AddRuleRef(rule_id); + const auto& ref_rule = base_grammar_->GetRule(rule_id); + const auto& ref_rule_body = base_grammar_->GetGrammarExpr(ref_rule.body_expr_id); + // Keep the reference to budgeted, suffix/stop, lazy, and temperature rules: replacing it with + // the rule's content would erase the rule that the runtime semantics apply to. + if (ref_rule.max_tokens < 0 && ref_rule.max_chars < 0 && + base_grammar_->GetSuffixStopInfo(rule_id) == nullptr && !ref_rule.is_lazy && + !ref_rule.temperature.has_value() && + ref_rule_body.type == GrammarBuilder::GrammarExprType::kChoices && + ref_rule_body.size() == 1) { + const auto& ref_choice = base_grammar_->GetGrammarExpr(ref_rule_body[0]); + if (ref_choice.size() == 1) { + grammar_expr_id = builder_->AddGrammarExpr(base_grammar_->GetGrammarExpr(ref_choice[0])); + } + } + + // Memoize on (content of the repeated expr, lower, upper) so that identical repetitions share + // one expansion instead of each producing its own chain of rules. + const auto repeated_expr = builder_->GetGrammarExpr(grammar_expr_id); + std::vector cache_key; + cache_key.reserve(repeated_expr.size() + 3); + cache_key.push_back(static_cast(repeated_expr.type)); + cache_key.insert(cache_key.end(), repeated_expr.begin(), repeated_expr.end()); + cache_key.push_back(lower); + cache_key.push_back(upper); + auto it = repetition_cache_.find(cache_key); + if (it != repetition_cache_.end()) { + return it->second; + } + + int32_t result = ExpandRepetitionRange(cur_rule_name, grammar_expr_id, lower, upper); + repetition_cache_.emplace(std::move(cache_key), result); + return result; +} + +int32_t RepetitionRangeExpanderImpl::ExpandRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper +) { + static const int64_t kUnzipThreshold = 128; + XGRAMMAR_CHECK(lower >= 0 && (upper == -1 || upper >= lower)) + << "Invalid repetition range {" << lower << ", " << upper << "}"; + + // Case 1.1 small upper (<=threshold), unzip the repetition. + // Case 1.2 unbounded upper, and lower is also small (<=threshold), unzip the lower part. + if ((upper != -1 && upper <= kUnzipThreshold) || (upper == -1 && lower <= kUnzipThreshold)) { + return LegacyHandleRepetitionRange(cur_rule_name, grammar_expr_id, lower, upper); + } + + // Case 2. upper is unbounded, and lower is large (>threshold). + // Or upper is bounded, but upper > threshold. + + // Case 2.1.1. lower is smaller than threshold, and upper is large. Transform {lower, upper} into: + // {threshold, upper} | {lower, threshold} + std::vector choices; + if (lower < kUnzipThreshold) { + choices.push_back(builder_->AddSequence( + {LegacyHandleRepetitionRange(cur_rule_name, grammar_expr_id, lower, kUnzipThreshold - 1)} + )); + lower = kUnzipThreshold; + } + + std::optional infinite_repetition_id = std::nullopt; + std::vector repeated_sequence; + // Now, we transform {lower, upper} into {max{threshold, lower}, upper}. + // Case 2.2 upper is unbounded. We will transform it into {lower} {0, inf}. + if (upper == -1) { + const auto& rule_expr = builder_->GetGrammarExpr(grammar_expr_id); + if (rule_expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) { + std::vector character_ranges; + bool is_negative = rule_expr[0]; + for (int i = 1; i < static_cast(rule_expr.size()); i += 2) { + character_ranges.push_back({rule_expr[i], rule_expr[i + 1]}); + } + infinite_repetition_id = builder_->AddCharacterClassStar(character_ranges, is_negative); + } else { + const auto unbounded_rule_id = + builder_->AddEmptyRule(builder_->GetNewRuleName(cur_rule_name + "_repeat_inf")); + int recursion_sequence = + builder_->AddSequence({grammar_expr_id, builder_->AddRuleRef(unbounded_rule_id)}); + int recursion_choice = builder_->AddChoices({builder_->AddEmptyStr(), recursion_sequence}); + builder_->UpdateRuleBody(unbounded_rule_id, recursion_choice); + infinite_repetition_id = builder_->AddRuleRef(unbounded_rule_id); + } + upper = lower; + } + + // Handle the {lower, upper} part, where threshold <= lower <= upper. + const auto repeat_name = cur_rule_name + "_repeat_1"; + XGRAMMAR_DCHECK(lower >= kUnzipThreshold && upper >= lower); + + // If we have infinite repetition part, add it to the sequence. + if (infinite_repetition_id.has_value()) { + repeated_sequence.push_back(infinite_repetition_id.value()); + } + + // The repetition body. + if (upper != kUnzipThreshold) { + XGRAMMAR_DCHECK(upper > kUnzipThreshold); + auto new_grammar_expr_id = builder_->AddChoices({builder_->AddSequence({grammar_expr_id})}); + auto new_rule_id = builder_->AddRuleWithHint(repeat_name, new_grammar_expr_id); + auto new_repeated_ref_rule_expr = builder_->AddChoices({builder_->AddSequence( + {builder_->AddRepeat(new_rule_id, lower - kUnzipThreshold, upper - kUnzipThreshold)} + )}); + auto new_repeated_rule_id = + builder_->AddRuleWithHint(repeat_name + "_inner", new_repeated_ref_rule_expr); + repeated_sequence.push_back(builder_->AddRuleRef(new_repeated_rule_id)); + std::vector repetition_lookahead(kUnzipThreshold, grammar_expr_id); + builder_->UpdateLookaheadAssertion(new_rule_id, builder_->AddSequence(repetition_lookahead)); + } + + // Add the last threshold grammar_expr_id to the sequence. + for (int i = 0; i < kUnzipThreshold; ++i) { + repeated_sequence.push_back(grammar_expr_id); + } + + // Add the sequence to choices. + choices.push_back(builder_->AddSequence(repeated_sequence)); + auto result_rule_id = builder_->AddRuleWithHint(cur_rule_name, builder_->AddChoices(choices)); + return builder_->AddRuleRef(result_rule_id); +} + +class RepetitionNormalizerImpl { + public: + void Apply(Grammar* grammar) { + auto& grammar_ref = *grammar; + for (int i = 0; i < grammar_ref->NumGrammarExprs(); ++i) { + auto expr = grammar_ref->GetGrammarExpr(i); + if (expr.type != Grammar::Impl::GrammarExprType::kRepeat) { + continue; + } + int repeat_rule_id = expr[0]; + grammar_ref->GetRule(repeat_rule_id).is_exact_lookahead = true; + if (std::binary_search( + grammar_ref->allow_empty_rule_ids.begin(), + grammar_ref->allow_empty_rule_ids.end(), + repeat_rule_id + )) { + // The repeated rule can be empty, so we need to normalize it. + expr.SetData(1, 0); // Set min repeat to 0 + } + } + } +}; + +/*! + * \brief Rewrite lazy rule bodies into their terminal-like form where possible: unwrap the + * single-reference chains produced by regex conversion, and flatten the right-recursive plus + * pattern (x ::= cc x | cc) and star pattern (x ::= cc x | "") produced by regex conversion and + * repetition expansion into (cc cc*) and (cc*). Grammars without lazy rules are returned + * unchanged. + */ +class LazyBodyFlattenerImpl : public GrammarMutator { + public: + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + bool has_lazy_rule = false; + for (int i = 0; i < grammar->NumRules(); ++i) { + has_lazy_rule = has_lazy_rule || grammar->GetRule(i).is_lazy; + } + if (!has_lazy_rule) { + return grammar; + } + InitGrammar(grammar); + InitBuilder(); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + builder_->AddEmptyRule(base_grammar_->GetRule(i).name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = rule.name; + int32_t new_body_expr_id = + rule.is_lazy ? BuildFlattenedLazyBody(rule.body_expr_id) : VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(i, new_body_expr_id); + builder_->UpdateLookaheadAssertion(i, VisitLookaheadAssertion(rule.lookahead_assertion_id)); + builder_->UpdateMaxTokens(i, rule.max_tokens); + builder_->UpdateMaxChars(i, rule.max_chars); + builder_->UpdateCaptureName(i, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + builder_->UpdateSuffixStopInfo(i, *suffix_stop_info); + } + builder_->UpdateLazy(i, rule.is_lazy); + builder_->UpdateRuleTemperature(i, rule.temperature); + } + return builder_->Get(base_grammar_->GetRootRule().name); + } + + private: + int32_t BuildFlattenedLazyBody(int32_t body_expr_id) { + auto body_type = base_grammar_->GetGrammarExpr(body_expr_id).type; + if (body_type == GrammarExprType::kRegex || body_type == GrammarExprType::kSubstring) { + return VisitExpr(body_expr_id); + } + // Unwrap chains of single rule references (r ::= (x), x ::= (y), ...) produced by regex + // conversion, and detect the plus-desugar pattern at the top level. + int32_t cur_body_id = body_expr_id; + int32_t cur_rule_id = -1; + for (int depth = 0; depth < 64; ++depth) { + const auto& body = base_grammar_->GetGrammarExpr(cur_body_id); + if (body.type != GrammarExprType::kChoices || body.size() != 1) { + break; + } + const auto& choice = base_grammar_->GetGrammarExpr(body[0]); + if (choice.type != GrammarExprType::kSequence || choice.size() != 1) { + break; + } + const auto& element = base_grammar_->GetGrammarExpr(choice[0]); + if (element.type != GrammarExprType::kRuleRef || base_grammar_->GetRule(element[0]).is_lazy) { + break; + } + cur_rule_id = element[0]; + cur_body_id = base_grammar_->GetRule(cur_rule_id).body_expr_id; + } + + const auto& body = base_grammar_->GetGrammarExpr(cur_body_id); + if (body.type != GrammarExprType::kChoices) { + XGRAMMAR_LOG(WARNING) << "The body of the lazy rule '" << cur_rule_name_ + << "' cannot be flattened into a terminal-like form"; + return VisitExpr(cur_body_id); + } + std::vector repeat_elements; + if (cur_rule_id != -1 && TryEmitRepeatPattern(body, cur_rule_id, &repeat_elements)) { + return builder_->AddChoices({builder_->AddSequence(repeat_elements)}); + } + std::vector new_choice_ids; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type != GrammarExprType::kSequence) { + new_choice_ids.push_back(VisitExpr(choice_id)); + continue; + } + std::vector elements; + if (!FlattenSequenceInto(choice, &elements, 0)) { + // Not flattenable; copy as is and let the terminal-like validation report the error. + XGRAMMAR_LOG(WARNING) << "The body of the lazy rule '" << cur_rule_name_ + << "' cannot be flattened into a terminal-like form"; + return VisitExpr(cur_body_id); + } + new_choice_ids.push_back(builder_->AddSequence(elements)); + } + return builder_->AddChoices(new_choice_ids); + } + + /*! \brief Append the flattened elements of the sequence, splicing rule references whose body + * is a single terminal-like sequence or the plus-desugar pattern, and coalescing references + * to single-character alternations into character classes. Returns false if some element + * cannot be flattened. */ + bool FlattenSequenceInto(const GrammarExpr& seq, std::vector* elements, int depth) { + if (depth > 64) { + return false; + } + for (auto element_id : seq) { + const auto& element = base_grammar_->GetGrammarExpr(element_id); + if (element.type == GrammarExprType::kByteString || + element.type == GrammarExprType::kCharacterClass || + element.type == GrammarExprType::kCharacterClassStar) { + elements->push_back(builder_->AddGrammarExpr(element)); + continue; + } + if (element.type != GrammarExprType::kRuleRef) { + return false; + } + const auto& ref_rule = base_grammar_->GetRule(element[0]); + if (ref_rule.is_lazy) { + return false; + } + const auto& ref_body = base_grammar_->GetGrammarExpr(ref_rule.body_expr_id); + if (TryEmitRepeatPattern(ref_body, element[0], elements)) { + continue; + } + if (ref_body.type == GrammarExprType::kChoices && ref_body.size() == 1) { + const auto& only_choice = base_grammar_->GetGrammarExpr(ref_body[0]); + if (only_choice.type == GrammarExprType::kEmptyStr) { + continue; + } + if (only_choice.type == GrammarExprType::kSequence && + FlattenSequenceInto(only_choice, elements, depth + 1)) { + continue; + } + } + std::vector ranges; + if (CollectSingleCharRanges(element_id, &ranges, 0)) { + elements->push_back(builder_->AddCharacterClass(UnionRanges(std::move(ranges)), false)); + continue; + } + return false; + } + return true; + } + + /*! \brief Resolve an expr matching exactly one character into the set of codepoint ranges it + * accepts, appending them to ranges. Accepts character classes, single-byte strings, and + * references to non-lazy rules that are alternations of such elements. Returns false + * otherwise. */ + bool CollectSingleCharRanges( + int32_t expr_id, std::vector* ranges, int depth + ) { + if (depth > 64) { + return false; + } + const auto& expr = base_grammar_->GetGrammarExpr(expr_id); + if (expr.type == GrammarExprType::kCharacterClass) { + AppendPositiveRanges(expr, ranges); + return true; + } + if (expr.type == GrammarExprType::kByteString && expr.size() == 1) { + ranges->push_back({expr[0], expr[0]}); + return true; + } + if (expr.type != GrammarExprType::kRuleRef) { + return false; + } + const auto& rule = base_grammar_->GetRule(expr[0]); + if (rule.is_lazy) { + return false; + } + const auto& body = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (body.type != GrammarExprType::kChoices) { + return false; + } + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type != GrammarExprType::kSequence || choice.size() != 1 || + !CollectSingleCharRanges(choice[0], ranges, depth + 1)) { + return false; + } + } + return true; + } + + /*! \brief If the body matches the plus-desugar pattern (self ::= e self | e) or the + * (generalized) star-desugar pattern (self ::= "" | e1 self | e2 self | ...), append the + * equivalent (e e*) or ((e1|e2|...)*) to the elements and return true. */ + bool TryEmitRepeatPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + return TryEmitPlusPattern(body, self_rule_id, elements) || + TryEmitGeneralizedPlusPattern(body, self_rule_id, elements) || + TryEmitStarPattern(body, self_rule_id, elements); + } + + /*! \brief If the body matches the generalized plus-desugar pattern (self ::= e1 self | ... | + * e1 | ...) with single-character elements whose recursive and base unions are equal, append + * the equivalent (cc cc*) over the union to the elements and return true. */ + bool TryEmitGeneralizedPlusPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + if (body.type != GrammarExprType::kChoices) { + return false; + } + std::vector recursive_ranges; + std::vector base_ranges; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type != GrammarExprType::kSequence) { + return false; + } + if (choice.size() == 1) { + if (!CollectSingleCharRanges(choice[0], &base_ranges, 0)) { + return false; + } + continue; + } + if (choice.size() == 2) { + const auto& tail = base_grammar_->GetGrammarExpr(choice[1]); + if (tail.type == GrammarExprType::kRuleRef && tail[0] == self_rule_id && + CollectSingleCharRanges(choice[0], &recursive_ranges, 0)) { + continue; + } + } + return false; + } + recursive_ranges = UnionRanges(std::move(recursive_ranges)); + base_ranges = UnionRanges(std::move(base_ranges)); + // The unions must coincide: with differing sets (e.g. self ::= a self | b, which is a*b), + // the language is not (a|b)+. + if (recursive_ranges.empty() || !RangesEqual(recursive_ranges, base_ranges)) { + return false; + } + elements->push_back(builder_->AddCharacterClass(recursive_ranges, false)); + elements->push_back(builder_->AddCharacterClassStar(recursive_ranges, false)); + return true; + } + + /*! \brief If the body matches the plus-desugar pattern (self ::= e self | e) with e resolving + * to a star-expressible expr, append the equivalent (e e*) (or (e*) when e itself resolves to + * a star) to the elements and return true. */ + bool TryEmitPlusPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + if (body.type != GrammarExprType::kChoices || body.size() != 2) { + return false; + } + for (int recursive_pos = 0; recursive_pos < 2; ++recursive_pos) { + const auto& recursive = base_grammar_->GetGrammarExpr(body[recursive_pos]); + const auto& base = base_grammar_->GetGrammarExpr(body[1 - recursive_pos]); + if (recursive.type != GrammarExprType::kSequence || recursive.size() != 2 || + base.type != GrammarExprType::kSequence || base.size() != 1) { + continue; + } + const auto& element = base_grammar_->GetGrammarExpr(recursive[0]); + const auto& tail = base_grammar_->GetGrammarExpr(recursive[1]); + const auto& base_element = base_grammar_->GetGrammarExpr(base[0]); + if (tail.type != GrammarExprType::kRuleRef || tail[0] != self_rule_id || + element.type != base_element.type || element.size() != base_element.size() || + !std::equal(element.begin(), element.end(), base_element.begin())) { + continue; + } + int32_t resolved_id = ResolveStarExpressible(recursive[0]); + if (resolved_id == -1) { + // Not a single terminal; e may still be a single-character alternation, giving + // (e1|e2|...)+ = cc cc* over the union of the ranges. + std::vector ranges; + if (!CollectSingleCharRanges(recursive[0], &ranges, 0)) { + continue; + } + ranges = UnionRanges(std::move(ranges)); + elements->push_back(builder_->AddCharacterClass(ranges, false)); + elements->push_back(builder_->AddCharacterClassStar(ranges, false)); + return true; + } + const auto& resolved = base_grammar_->GetGrammarExpr(resolved_id); + if (resolved.type == GrammarExprType::kCharacterClassStar) { + // (e*)+ is e*. + elements->push_back(builder_->AddGrammarExpr(resolved)); + return true; + } + std::vector character_ranges; + bool is_negative = false; + if (resolved.type == GrammarExprType::kCharacterClass) { + is_negative = static_cast(resolved[0]); + for (int i = 1; i < static_cast(resolved.size()); i += 2) { + character_ranges.push_back({resolved[i], resolved[i + 1]}); + } + } else { // single-byte kByteString + character_ranges.push_back({resolved[0], resolved[0]}); + } + elements->push_back(builder_->AddGrammarExpr(resolved)); + elements->push_back(builder_->AddCharacterClassStar(character_ranges, is_negative)); + return true; + } + return false; + } + + /*! \brief If the body matches the generalized star-desugar pattern (self ::= "" | e1 self | + * e2 self | ...) with each e resolving to character ranges, append the equivalent single + * character class star ((e1|e2|...)*) to the elements and return true. */ + bool TryEmitStarPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + if (body.type != GrammarExprType::kChoices) { + return false; + } + bool has_empty = false; + std::vector ranges; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type == GrammarExprType::kEmptyStr) { + has_empty = true; + continue; + } + if (choice.type != GrammarExprType::kSequence || choice.size() != 2) { + return false; + } + const auto& tail = base_grammar_->GetGrammarExpr(choice[1]); + if (tail.type != GrammarExprType::kRuleRef || tail[0] != self_rule_id) { + return false; + } + if (!CollectStarRanges(choice[0], &ranges, 0)) { + return false; + } + } + if (!has_empty || ranges.empty()) { + return false; + } + elements->push_back(builder_->AddCharacterClassStar(UnionRanges(std::move(ranges)), false)); + return true; + } + + /*! \brief Resolve an expr repeated under an enclosing star into the set of codepoint ranges it + * repeats over, appending them to ranges. Accepts character classes, character class stars, + * single-byte strings, and references to non-lazy rules that are alternations of such + * elements, or star/plus recursions over them — under an enclosing star, all of these are + * equivalent to the union of their character ranges. Returns false otherwise. */ + bool CollectStarRanges( + int32_t expr_id, std::vector* ranges, int depth + ) { + if (depth > 64) { + return false; + } + const auto& expr = base_grammar_->GetGrammarExpr(expr_id); + if (expr.type == GrammarExprType::kCharacterClass || + expr.type == GrammarExprType::kCharacterClassStar) { + AppendPositiveRanges(expr, ranges); + return true; + } + if (expr.type == GrammarExprType::kByteString && expr.size() == 1) { + ranges->push_back({expr[0], expr[0]}); + return true; + } + if (expr.type != GrammarExprType::kRuleRef) { + return false; + } + int32_t rule_id = expr[0]; + const auto& rule = base_grammar_->GetRule(rule_id); + if (rule.is_lazy) { + return false; + } + const auto& body = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (body.type != GrammarExprType::kChoices) { + return false; + } + bool has_empty = false; + std::vector base_elements; + std::vector recursive_elements; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type == GrammarExprType::kEmptyStr) { + has_empty = true; + continue; + } + if (choice.type != GrammarExprType::kSequence) { + return false; + } + if (choice.size() == 1) { + base_elements.push_back(choice[0]); + continue; + } + if (choice.size() == 2) { + const auto& tail = base_grammar_->GetGrammarExpr(choice[1]); + if (tail.type == GrammarExprType::kRuleRef && tail[0] == rule_id) { + recursive_elements.push_back(choice[0]); + continue; + } + } + return false; + } + // The safe shapes: an alternation (a | b | ...), a star ("" | e1 self | ...), and a plus + // (e self | e, or single-character alternated forms with equal recursive/base unions). + // Mixed shapes like (a self | b) are a*b, whose star is not the union, so they are rejected. + std::vector* collect = nullptr; + if (recursive_elements.empty()) { + collect = &base_elements; + } else if (base_elements.empty() && has_empty) { + collect = &recursive_elements; + } else if (recursive_elements.size() == 1 && base_elements.size() == 1 && !has_empty && + ExprsEqual(recursive_elements[0], base_elements[0])) { + collect = &recursive_elements; + } else if (!has_empty) { + std::vector recursive_ranges; + std::vector base_ranges; + for (auto element_id : recursive_elements) { + if (!CollectSingleCharRanges(element_id, &recursive_ranges, depth + 1)) { + return false; + } + } + for (auto element_id : base_elements) { + if (!CollectSingleCharRanges(element_id, &base_ranges, depth + 1)) { + return false; + } + } + recursive_ranges = UnionRanges(std::move(recursive_ranges)); + if (!RangesEqual(recursive_ranges, UnionRanges(std::move(base_ranges)))) { + return false; + } + ranges->insert(ranges->end(), recursive_ranges.begin(), recursive_ranges.end()); + return true; + } else { + return false; + } + for (auto element_id : *collect) { + if (!CollectStarRanges(element_id, ranges, depth + 1)) { + return false; + } + } + return true; + } + + /*! \brief Whether two exprs have identical type and content. */ + bool ExprsEqual(int32_t lhs_id, int32_t rhs_id) { + const auto& lhs = base_grammar_->GetGrammarExpr(lhs_id); + const auto& rhs = base_grammar_->GetGrammarExpr(rhs_id); + return lhs.type == rhs.type && lhs.size() == rhs.size() && + std::equal(lhs.begin(), lhs.end(), rhs.begin()); + } + + /*! \brief Whether two normalized range vectors are identical. */ + static bool RangesEqual( + const std::vector& lhs, + const std::vector& rhs + ) { + if (lhs.size() != rhs.size()) { + return false; + } + for (size_t i = 0; i < lhs.size(); ++i) { + if (lhs[i].lower != rhs[i].lower || lhs[i].upper != rhs[i].upper) { + return false; + } + } + return true; + } + + /*! \brief Append the positive codepoint ranges of a character class or character class star, + * complementing negated classes over [0, 0x10FFFF]. */ + void AppendPositiveRanges( + const GrammarExpr& expr, std::vector* ranges + ) { + std::vector class_ranges; + for (int i = 1; i < static_cast(expr.size()); i += 2) { + class_ranges.push_back({expr[i], expr[i + 1]}); + } + if (!static_cast(expr[0])) { + ranges->insert(ranges->end(), class_ranges.begin(), class_ranges.end()); + return; + } + class_ranges = UnionRanges(std::move(class_ranges)); + int32_t next = 0; + for (const auto& range : class_ranges) { + if (range.lower > next) { + ranges->push_back({next, range.lower - 1}); + } + next = std::max(next, range.upper + 1); + } + if (next <= 0x10FFFF) { + ranges->push_back({next, 0x10FFFF}); + } + } + + /*! \brief Sort the ranges and merge overlapping or adjacent ones. */ + static std::vector UnionRanges( + std::vector ranges + ) { + std::sort(ranges.begin(), ranges.end(), [](const auto& a, const auto& b) { + return a.lower < b.lower; + }); + std::vector result; + for (const auto& range : ranges) { + if (!result.empty() && range.lower <= result.back().upper + 1) { + result.back().upper = std::max(result.back().upper, range.upper); + } else { + result.push_back(range); + } + } + return result; + } + + /*! \brief Resolve an expr through chains of non-lazy single-reference rules to a + * star-expressible expr: a character class, a single-byte string, or a character class star. + * Returns the resolved expr id, or -1 if it does not resolve to one. */ + int32_t ResolveStarExpressible(int32_t expr_id) { + int32_t cur_id = expr_id; + for (int depth = 0; depth < 64; ++depth) { + const auto& cur = base_grammar_->GetGrammarExpr(cur_id); + if (cur.type == GrammarExprType::kCharacterClass || + cur.type == GrammarExprType::kCharacterClassStar || + (cur.type == GrammarExprType::kByteString && cur.size() == 1)) { + return cur_id; + } + if (cur.type != GrammarExprType::kRuleRef) { + return -1; + } + const auto& ref_rule = base_grammar_->GetRule(cur[0]); + if (ref_rule.is_lazy) { + return -1; + } + const auto& ref_body = base_grammar_->GetGrammarExpr(ref_rule.body_expr_id); + if (ref_body.type != GrammarExprType::kChoices || ref_body.size() != 1) { + return -1; + } + const auto& only_choice = base_grammar_->GetGrammarExpr(ref_body[0]); + if (only_choice.type != GrammarExprType::kSequence || only_choice.size() != 1) { + return -1; + } + cur_id = only_choice[0]; + } + return -1; + } +}; + +class GrammarOptimizerImpl { + public: + static Grammar Apply(const Grammar& grammar) { + // ByteStringFuser and RuleInliner rewrite the grammar in place, so work on a private copy: the + // input grammar may be shared (e.g. a cached grammar) and must not be mutated. Copy the impl + // directly (contiguous vector copies) instead of going through GrammarBuilder, which would + // also build the unneeded rule name map. + Grammar result(std::make_shared(*grammar.operator->())); + ByteStringFuser::Apply(&result); + RuleInliner::Apply(&result); + result = RepetitionRangeExpander::Apply(result); + result = LazyBodyFlattenerImpl().Apply(result); + result = DeadCodeEliminator::Apply(result); + result = LookaheadAssertionAnalyzer::Apply(result); + result->allow_empty_rule_ids = AllowEmptyRuleAnalyzer::Apply(result); + ValidateLazyRules(result); + RepetitionNormalizer::Apply(&result); + GrammarFSMBuilder::Apply(&result); + result->optimized = true; + return result; + } + + private: + /*! + * \brief Committed-shortest (lazy) matching requires the whole rule body to compile into a + * single per-rule FSM without rule references, so that the states of one occurrence are exactly + * the states with the rule's id. + */ + static void ValidateLazyRules(const Grammar& grammar) { + for (int32_t i = 0; i < grammar->NumRules(); ++i) { + const auto& rule = grammar->GetRule(i); + if (!rule.is_lazy) { + continue; + } + const auto& body = grammar->GetGrammarExpr(rule.body_expr_id); + if (body.type == Grammar::Impl::GrammarExprType::kRegex || + body.type == Grammar::Impl::GrammarExprType::kSubstring) { + continue; + } + XGRAMMAR_CHECK(body.type == Grammar::Impl::GrammarExprType::kChoices) + << "lazy rule '" << rule.name << "' must have a terminal-like body"; + for (auto choice_id : body) { + const auto& choice = grammar->GetGrammarExpr(choice_id); + if (choice.type == Grammar::Impl::GrammarExprType::kEmptyStr) { + continue; + } + for (auto element_id : choice) { + const auto& element = grammar->GetGrammarExpr(element_id); + XGRAMMAR_CHECK( + element.type == Grammar::Impl::GrammarExprType::kByteString || + element.type == Grammar::Impl::GrammarExprType::kCharacterClass || + element.type == Grammar::Impl::GrammarExprType::kCharacterClassStar + ) << "lazy rule '" + << rule.name + << "' must have a terminal-like body (strings, character classes, and regexes that " + "compile to a single FSM); rule references and repetition ranges are not supported"; + } + } + } + } +}; + +/*! + * \brief Fuse adjacent byte string elements in sequences. + * \details Rewrites the grammar in place: only sequences that actually contain a run of adjacent + * byte strings are rebuilt, the rest keep their original ids. Stale exprs are removed later by + * DeadCodeEliminator. + */ +class ByteStringFuserImpl : public InPlaceGrammarRewriter { + protected: + int32_t VisitSequence(int32_t expr_id) override { + // Read-only probe: rewrite only if the sequence contains an empty byte string or two adjacent + // byte strings. The probe appends nothing, so no memory is allocated for the common unchanged + // case. + bool previous_is_byte_string = false; + bool needs_rewrite = false; + { + auto expr = builder_.GetGrammarExpr(expr_id); + for (int32_t element_id : expr) { + auto element = builder_.GetGrammarExpr(element_id); + bool is_byte_string = element.type == GrammarExprType::kByteString; + if (is_byte_string && (previous_is_byte_string || element.size() == 0)) { + needs_rewrite = true; + break; + } + previous_is_byte_string = is_byte_string; + } + } + if (!needs_rewrite) { + return expr_id; + } + // Copy the element ids first: fusing appends to the arena and invalidates expr views. + auto expr = builder_.GetGrammarExpr(expr_id); + std::vector element_ids(expr.begin(), expr.end()); + std::vector new_element_ids; + for (size_t i = 0; i < element_ids.size();) { + if (builder_.GetGrammarExpr(element_ids[i]).type != GrammarExprType::kByteString) { + new_element_ids.push_back(element_ids[i]); + ++i; + continue; + } + // Gather the run of adjacent byte strings starting at i. + std::vector fused_bytes; + size_t run_end = i; + while (run_end < element_ids.size()) { + auto element = builder_.GetGrammarExpr(element_ids[run_end]); + if (element.type != GrammarExprType::kByteString) { + break; + } + fused_bytes.insert(fused_bytes.end(), element.begin(), element.end()); + ++run_end; + } + // Empty byte strings are epsilon and can be omitted from a sequence. + if (!fused_bytes.empty()) { + if (run_end - i == 1) { + new_element_ids.push_back(element_ids[i]); + } else { + new_element_ids.push_back(builder_.AddByteString(fused_bytes)); + } + } + i = run_end; + } + return builder_.AddSequence(new_element_ids); + } +}; + +class RootRuleRenamerImpl { + public: + static Grammar Apply(const Grammar& grammar) { + // If the root name is "root", return directly. + if (grammar->GetRootRule().name == "root") { + return grammar; + } + + // Collect all the rule names. + std::unordered_set rule_names; + int root_name_rule_id = -1; + for (int i = 0; i < grammar->NumRules(); i++) { + const auto& rule_name = grammar->GetRule(i).name; + if (rule_name == "root") { + root_name_rule_id = i; + } + rule_names.insert(rule_name); + } + + // Rename the rules. + Grammar grammar_copy = grammar; + grammar_copy->GetRule(grammar_copy->GetRootRuleId()).name = "root"; + if (root_name_rule_id != -1) { + std::string rule_prefix = "root_"; + bool renamed = false; + for (int i = 0; i <= grammar_copy->NumRules(); i++) { + std::string new_rule_name = rule_prefix + std::to_string(i); + if (rule_names.find(new_rule_name) == rule_names.end()) { + grammar_copy->GetRule(root_name_rule_id).name = new_rule_name; + renamed = true; + break; + } + } + XGRAMMAR_DCHECK(renamed) << "Rule renaming must succeed within (n + 1) attempts."; + } + return grammar_copy; + } +}; + +class GrammarFSMHasherImpl { + public: + void Apply(Grammar* grammar); + static std::optional HashSequence(const Grammar& grammar, int32_t sequence_id); + + static constexpr int16_t kNotEndStateFlag = -0x100; + static constexpr int16_t kEndStateFlag = -0x200; + static constexpr int16_t kSelfRecursionFlag = -0x300; + static constexpr int16_t kSimpleCycleFlag = -0x400; + static constexpr int16_t kUnKnownFlag = -0x500; + + private: + Grammar* grammar_; + std::vector visited_; + std::vector> ref_graph_from_referrer_to_referee_; + std::vector> ref_graph_from_referee_to_referrer_; + std::vector> sorted_edges_; + std::vector has_inward_edges_; + + /*! + * \brief The worklist of fsms that are ready to be hashed: fsms whose references are all + * hashed (except possibly a self-recursion). Maintained incrementally so that the main hashing + * loop is O(V + E) instead of rescanning all rules after each hashed fsm. + */ + std::queue ready_queue_; + + /*! + * \brief Get the hash value of a fsm, with a given grammar. + */ + uint64_t HashFsm(int fsm_index); + + /*! + * \brief Find a simple cycle in the reference graph, And hash the + * fsms in the simple cycle. + */ + bool FindSimpleCycle(); + + /*! + * \brief Hash the fsms in the simple cycle. + */ + void HashSimpleCycle(const std::vector& simple_cycle); + + /*! + * \brief Check if a fsm is ready to be hashed: it is not hashed yet, and it references no + * unhashed fsms other than itself. + */ + bool IsReadyToHash(int32_t fsm_index) const { + if (visited_[fsm_index]) { + return false; + } + const auto& referees = ref_graph_from_referrer_to_referee_[fsm_index]; + return referees.empty() || (referees.size() == 1 && referees[0] == fsm_index); + } + + /*! + * \brief Remove the hashed fsm from the reference graph, and push the referrers that become + * ready to hash into the ready queue. + */ + void RemoveHashedFsmFromRefGraph(int32_t fsm_index); + + std::pair IsPartialHashable(int fsm_index); +}; + +bool GrammarFSMHasherImpl::FindSimpleCycle() { + // Try to find a simple cycle. + std::vector not_simple_cycle = visited_; + // Allocated once and cleaned up after each walk, to avoid an O(num_rules) allocation per + // outer iteration. + std::vector in_stack(ref_graph_from_referee_to_referrer_.size(), false); + for (size_t i = 0; i < ref_graph_from_referee_to_referrer_.size(); i++) { + if (not_simple_cycle[i]) { + continue; + } + // Not a simple cycle if it has more than one referee. + std::stack dfs_stack; + std::vector simple_cycle; + std::vector walked_states; + dfs_stack.push(static_cast(i)); + int32_t current_fsm_index = i; + in_stack[current_fsm_index] = true; + walked_states.push_back(current_fsm_index); + while ((ref_graph_from_referrer_to_referee_[current_fsm_index].size() == 1) && + !not_simple_cycle[current_fsm_index]) { + XGRAMMAR_CHECK(current_fsm_index != ref_graph_from_referrer_to_referee_[current_fsm_index][0]) + << "Self-recursion cycle found in the reference graph, which is not allowed."; + not_simple_cycle[current_fsm_index] = true; + current_fsm_index = ref_graph_from_referrer_to_referee_[current_fsm_index][0]; + if (in_stack[current_fsm_index]) { + simple_cycle.push_back(current_fsm_index); + while (dfs_stack.top() != current_fsm_index) { + simple_cycle.push_back(dfs_stack.top()); + dfs_stack.pop(); + } + // Found a simple cycle. + break; + } else { + dfs_stack.push(current_fsm_index); + in_stack[current_fsm_index] = true; + walked_states.push_back(current_fsm_index); + } + } + if (!simple_cycle.empty()) { + HashSimpleCycle(simple_cycle); + return true; + } + for (auto state : walked_states) { + in_stack[state] = false; + } + } + return false; +} + +void GrammarFSMHasherImpl::HashSimpleCycle(const std::vector& simple_cycle) { + // Initialize the cycle hash. + for (const auto& cycle_id : simple_cycle) { + visited_[cycle_id] = true; + grammar_->ImplPtr()->per_rule_fsm_hashes[cycle_id] = kSimpleCycleFlag; + } + + std::vector local_cycle_hash; + local_cycle_hash.reserve(simple_cycle.size()); + for (const auto& cycle_id : simple_cycle) { + local_cycle_hash.push_back(HashFsm(cycle_id)); + } + std::vector local_cycle_hash_copy = local_cycle_hash; + for (int i = 0; i < static_cast(local_cycle_hash.size()); i++) { + uint64_t current_hash = 0; + for (int j = 0; j < static_cast(local_cycle_hash.size()); j++) { + current_hash = + HashCombine(current_hash, local_cycle_hash_copy[(i + j) % local_cycle_hash.size()]); + } + local_cycle_hash[i] = current_hash; + } + + for (int i = 0; i < static_cast(simple_cycle.size()); i++) { + grammar_->ImplPtr()->per_rule_fsm_hashes[simple_cycle[i]] = local_cycle_hash[i]; + RemoveHashedFsmFromRefGraph(simple_cycle[i]); + } +} + +void GrammarFSMHasherImpl::RemoveHashedFsmFromRefGraph(int32_t fsm_index) { + for (const auto& referer : ref_graph_from_referee_to_referrer_[fsm_index]) { + auto& referees = ref_graph_from_referrer_to_referee_[referer]; + auto it = std::find(referees.begin(), referees.end(), fsm_index); + if (it != referees.end()) { + referees.erase(it); + } + if (IsReadyToHash(referer)) { + ready_queue_.push(referer); + } + } +} + +void GrammarFSMHasherImpl::Apply(Grammar* grammar) { + grammar_ = grammar; + grammar->ImplPtr()->per_rule_fsm_hashes = + std::vector>((*grammar)->NumRules()); + grammar->ImplPtr()->per_rule_fsm_new_state_ids.resize((*grammar)->NumRules()); + ref_graph_from_referee_to_referrer_.clear(); + ref_graph_from_referrer_to_referee_.clear(); + sorted_edges_.clear(); + visited_ = std::vector((*grammar)->NumRules(), false); + has_inward_edges_ = std::vector((*grammar)->complete_fsm.NumStates(), false); + for (int i = 0; i < grammar_->ImplPtr()->complete_fsm.NumStates(); i++) { + for (const auto& edge : grammar->ImplPtr()->complete_fsm.GetEdges(i)) { + has_inward_edges_[edge.target] = true; + } + } + + // Get the reference graph. + ref_graph_from_referee_to_referrer_ = RuleRefGraphFinder().Apply(*grammar); + ref_graph_from_referrer_to_referee_ = std::vector>((*grammar)->NumRules()); + for (int referee = 0; referee < static_cast(ref_graph_from_referee_to_referrer_.size()); + ++referee) { + for (int referer : ref_graph_from_referee_to_referrer_[referee]) { + ref_graph_from_referrer_to_referee_[referer].push_back(referee); + } + } + + // Sort the edges. + const auto& complete_fsm = grammar->ImplPtr()->complete_fsm; + sorted_edges_.reserve(complete_fsm.NumStates()); + for (int i = 0; i < complete_fsm.NumStates(); i++) { + const auto& edges = complete_fsm.GetEdges(i); + sorted_edges_.emplace_back(); + sorted_edges_.back().reserve(edges.size()); + for (const auto& edge : edges) { + sorted_edges_.back().emplace_back(edge); + } + std::sort(sorted_edges_.back().begin(), sorted_edges_.back().end()); + } + + // Disable non-fsms. + for (size_t i = 0; i < grammar->ImplPtr()->per_rule_fsms.size(); i++) { + if (!grammar->ImplPtr()->per_rule_fsms[i].has_value()) { + visited_[i] = true; + } + } + + // Hash the fsms which can be hashed: terminal fsms, or self-recursion fsms. The ready queue + // is seeded with all currently hashable fsms and maintained incrementally as fsms are hashed, + // so the whole loop is O(V + E) over the reference graph. When no fsm is ready, try to break + // a simple cycle in the reference graph and continue. + ready_queue_ = {}; + for (int i = 0; i < (*grammar)->NumRules(); i++) { + if (IsReadyToHash(i)) { + ready_queue_.push(i); + } + } + while (true) { + if (ready_queue_.empty()) { + // Try to find a simple cycle. We must ensure there are not self-recursion cycles. + if (!FindSimpleCycle()) { + break; + } + continue; + } + int32_t current_operating_index = ready_queue_.front(); + ready_queue_.pop(); + // Skip stale entries: an fsm may be pushed multiple times before it is processed. + if (!IsReadyToHash(current_operating_index)) { + continue; + } + visited_[current_operating_index] = true; + grammar->ImplPtr()->per_rule_fsm_hashes[current_operating_index] = + HashFsm(current_operating_index); + RemoveHashedFsmFromRefGraph(current_operating_index); + } + + // Try to hash the remaining fsms: they must contain something can't be hashed, like repetition. + // We can do this: if the fsm's start state has no inward edges, and all the ref edges are hashed + // except the edges at the start state, we can hash it. + std::vector> partial_hashed_list; + for (int i = 0; i < (*grammar)->NumRules(); i++) { + if (grammar->ImplPtr()->per_rule_fsm_hashes[i].has_value()) { + continue; + } + if (!grammar->ImplPtr()->per_rule_fsms[i].has_value()) { + continue; + } + if (has_inward_edges_[grammar->ImplPtr()->per_rule_fsms[i]->GetFsm().GetStart()]) { + continue; + } + const auto& [can_be_hashed, hash_value] = IsPartialHashable(i); + if (can_be_hashed) { + partial_hashed_list.emplace_back(i, hash_value); + } + } + for (const auto& [rule_id, hash_value] : partial_hashed_list) { + grammar->ImplPtr()->per_rule_fsm_hashes[rule_id] = hash_value; + } +} + +std::pair GrammarFSMHasherImpl::IsPartialHashable(int fsm_index) { + uint64_t hash_result = 0; + XGRAMMAR_DCHECK(fsm_index >= 0 && fsm_index < (*grammar_)->NumRules()) + << "Invalid fsm index: " << fsm_index << " num_rules: " << (*grammar_)->NumRules(); + XGRAMMAR_DCHECK(grammar_->ImplPtr()->per_rule_fsms[fsm_index].has_value()); + const auto& fsm = grammar_->ImplPtr()->per_rule_fsms[fsm_index].value().GetFsm(); + std::map original_state_id_to_new_id; + original_state_id_to_new_id[fsm.GetStart()] = 0; + std::queue bfs_queue; + std::set> hash_and_target; + bfs_queue.push(fsm.GetStart()); + // Perform a bfs to hash all the edges. + while (!bfs_queue.empty()) { + int current_old_state_id = bfs_queue.front(); + bool is_start = current_old_state_id == fsm.GetStart(); + int current_new_state_id = original_state_id_to_new_id[current_old_state_id]; + bfs_queue.pop(); + + // Check if the current state is an end state. + if (fsm.IsEndState(current_old_state_id)) { + hash_result = HashCombine( + hash_result, current_new_state_id, kEndStateFlag, kEndStateFlag, current_new_state_id + ); + } else { + hash_result = HashCombine( + hash_result, + current_new_state_id, + kNotEndStateFlag, + kNotEndStateFlag, + current_new_state_id + ); + } + + // Hash the edges. + + // First, check the edges which are rule references (including repeat refs). + // To keep consistent, we need to sort them with hashes. + int32_t unhashed_rules_count = 0; + auto hash_rule_like_edge = [&](int32_t ref_rule_id, int32_t target) { + if (ref_rule_id == fsm_index) { + hash_and_target.insert({kSelfRecursionFlag, target}); + return true; + } + if (!grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].has_value()) { + if (!is_start) { + return false; + } else { + unhashed_rules_count++; + if (unhashed_rules_count > 1) { + return false; + } + hash_and_target.insert({kUnKnownFlag, target}); + } + return true; + } + hash_and_target.insert({grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].value(), target} + ); + return true; + }; + + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (edge.IsRuleRef()) { + if (!hash_rule_like_edge(edge.GetRefRuleId(), edge.target)) { + return {false, 0}; + } + } else if (edge.IsRepeatRef()) { + auto info = grammar_->ImplPtr()->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + if (!hash_rule_like_edge(info.RuleId(), edge.target)) { + return {false, 0}; + } + } + } + + // Hash them. + for (const auto& [hash, target] : hash_and_target) { + if (original_state_id_to_new_id.find(target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(target); + } + int32_t target_new_id = original_state_id_to_new_id[target]; + hash_result = HashCombine(hash_result, current_new_state_id, hash, target_new_id); + } + + // Then, check the edges which are not rule/repeat references. + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (original_state_id_to_new_id.find(edge.target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[edge.target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(edge.target); + } + int32_t target_new_id = original_state_id_to_new_id[edge.target]; + if (edge.IsRuleRef() || edge.IsRepeatRef()) { + continue; + } + hash_result = HashCombine( + hash_result, + current_new_state_id, + static_cast(edge.min), + static_cast(edge.max), + target_new_id + ); + } + } + std::vector> new_id_mapping; + new_id_mapping.reserve(original_state_id_to_new_id.size()); + for (const auto& [original_state_id, new_state_id] : original_state_id_to_new_id) { + new_id_mapping.emplace_back(original_state_id, new_state_id); + } + grammar_->ImplPtr()->per_rule_fsm_new_state_ids[fsm_index] = new_id_mapping; + return {true, hash_result}; +} + +uint64_t GrammarFSMHasherImpl::HashFsm(int fsm_index) { + uint64_t hash_result = 0; + XGRAMMAR_DCHECK(fsm_index >= 0 && fsm_index < (*grammar_)->NumRules()) + << "Invalid fsm index: " << fsm_index << " num_rules: " << (*grammar_)->NumRules(); + XGRAMMAR_DCHECK(grammar_->ImplPtr()->per_rule_fsms[fsm_index].has_value()); + const auto& fsm = grammar_->ImplPtr()->per_rule_fsms[fsm_index].value().GetFsm(); + std::map original_state_id_to_new_id; + original_state_id_to_new_id[fsm.GetStart()] = 0; + std::queue bfs_queue; + std::set> hash_and_target; + bfs_queue.push(fsm.GetStart()); + + // Perform a bfs to hash all the edges. + while (!bfs_queue.empty()) { + int current_old_state_id = bfs_queue.front(); + int current_new_state_id = original_state_id_to_new_id[current_old_state_id]; + bfs_queue.pop(); + + // Check if the current state is an end state. + if (fsm.IsEndState(current_old_state_id)) { + hash_result = HashCombine( + hash_result, current_new_state_id, kEndStateFlag, kEndStateFlag, current_new_state_id + ); + } else { + hash_result = HashCombine( + hash_result, + current_new_state_id, + kNotEndStateFlag, + kNotEndStateFlag, + current_new_state_id + ); + } + + // Hash the edges. + + // First, check the edges which are rule references (including repeat refs). + // To keep consistent, we need to sort them with hashes. + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (edge.IsRuleRef()) { + int32_t ref_rule_id = edge.GetRefRuleId(); + if (ref_rule_id == fsm_index) { + hash_and_target.insert({kSelfRecursionFlag, edge.target}); + } else { + XGRAMMAR_CHECK(grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].has_value()); + hash_and_target.insert( + {grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].value(), edge.target} + ); + } + } else if (edge.IsRepeatRef()) { + auto info = grammar_->ImplPtr()->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + int32_t ref_rule_id = info.RuleId(); + if (ref_rule_id == fsm_index) { + uint64_t base_hash = kSelfRecursionFlag; + uint64_t repeat_hash = HashCombine(base_hash, info.Lower(), info.Upper()); + hash_and_target.insert({repeat_hash, edge.target}); + } else { + XGRAMMAR_CHECK(grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].has_value()); + uint64_t base_hash = grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].value(); + uint64_t repeat_hash = HashCombine(base_hash, info.Lower(), info.Upper()); + hash_and_target.insert({static_cast(repeat_hash), edge.target}); + } + } + } + + // Hash them. + for (const auto& [hash, target] : hash_and_target) { + if (original_state_id_to_new_id.find(target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(target); + } + int32_t target_new_id = original_state_id_to_new_id[target]; + hash_result = HashCombine(hash_result, current_new_state_id, hash, target_new_id); + } + + // Then, check the edges which are not rule/repeat references. + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (original_state_id_to_new_id.find(edge.target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[edge.target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(edge.target); + } + int32_t target_new_id = original_state_id_to_new_id[edge.target]; + if (edge.IsRuleRef() || edge.IsRepeatRef()) { + continue; + } + hash_result = HashCombine( + hash_result, + current_new_state_id, + static_cast(edge.min), + static_cast(edge.max), + target_new_id + ); + } + } + std::vector> new_id_mapping; + new_id_mapping.reserve(original_state_id_to_new_id.size()); + for (const auto& [original_state_id, new_state_id] : original_state_id_to_new_id) { + new_id_mapping.emplace_back(original_state_id, new_state_id); + } + grammar_->ImplPtr()->per_rule_fsm_new_state_ids[fsm_index] = new_id_mapping; + return hash_result; +} + +std::optional GrammarFSMHasherImpl::HashSequence( + const Grammar& grammar, int32_t sequence_id +) { + using GrammarExprType = Grammar::Impl::GrammarExprType; + if (sequence_id == -1) { + return std::nullopt; + } + uint64_t hash_result = 0; + const auto& sequence_expr = grammar->GetGrammarExpr(sequence_id); + XGRAMMAR_DCHECK(sequence_expr.type == GrammarExprType::kSequence) + << "GrammarExpr is not a sequence"; + for (const auto& expr_id : sequence_expr) { + const auto& expr = grammar->GetGrammarExpr(expr_id); + hash_result = HashCombine(hash_result, static_cast(expr.type)); + switch (expr.type) { + case (GrammarExprType::kByteString): + case (GrammarExprType::kCharacterClass): + case (GrammarExprType::kCharacterClassStar): + case (GrammarExprType::kEmptyStr): { + for (const auto& element : expr) { + hash_result = HashCombine(hash_result, element); + } + break; + } + case (GrammarExprType::kRuleRef): { + if (grammar->per_rule_fsm_hashes[expr[0]].has_value()) { + hash_result = HashCombine(hash_result, grammar->per_rule_fsm_hashes[expr[0]].value()); + } else { + return std::nullopt; + } + break; + } + case (GrammarExprType::kRepeat): { + if (grammar->per_rule_fsm_hashes[expr[0]].has_value()) { + hash_result = HashCombine(hash_result, grammar->per_rule_fsm_hashes[expr[0]].value()); + } else { + return std::nullopt; + } + hash_result = HashCombine(hash_result, expr[1]); + hash_result = HashCombine(hash_result, expr[2]); + break; + } + case (GrammarExprType::kSequence): + case (GrammarExprType::kChoices): { + return std::nullopt; + } + case (GrammarExprType::kTagDispatch): + case (GrammarExprType::kTokenTagDispatch): { + return std::nullopt; + } + case (GrammarExprType::kRegex): + case (GrammarExprType::kSubstring): { + // Hash the content, like a byte string. + for (const auto& element : expr) { + hash_result = HashCombine(hash_result, element); + } + break; + } + case (GrammarExprType::kToken): + case (GrammarExprType::kExcludeToken): { + for (const auto& element : expr) { + hash_result = HashCombine(hash_result, element); + } + break; + } + } + } + return hash_result; +} + +class RuleLevelCache::Impl { + public: + using NodeKey = std::tuple< + uint64_t /*The hash value of the FSM*/, + int32_t /* The normalized node id*/, + int32_t /*The number of states*/, + int32_t /* The number of edges*/>; + using NodeType = std::pair; + + explicit Impl(size_t max_cache_memory_size) : max_cache_memory_size_(max_cache_memory_size) {} + + std::optional GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt + ); + + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask + ); + + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask + ); + + void ClearCache(); + + friend size_t MemorySize(const Impl* impl) { + int64_t total = 0; + for (const auto& shard : impl->shards_) { + total += shard.current_cache_memory_size; + } + return total; + } + + size_t GetMaxSize() const { return max_cache_memory_size_; } + + private: + /*! + * \brief The cache is sharded to reduce lock contention: the token mask cache generation + * queries and inserts from all compilation threads, and a single global mutex would serialize + * them (large grammars issue millions of cache operations). + */ + static constexpr size_t kNumShards = 16; + + struct Shard { + std::mutex mutex; + int64_t current_cache_memory_size = 0; + // The cache map: (fsm_hash, node_id, ...) -> index in cache_list + List cache_list; + std::unordered_map cache; + }; + + Shard& GetShard(const NodeKey& key) { + return shards_[HashCombine(std::get<0>(key), std::get<1>(key)) % kNumShards]; + } + + /*! \brief The memory budget of one shard. Eviction is performed per shard. */ + size_t ShardMaxSize() const { + return max_cache_memory_size_ == kUnlimitedSize ? kUnlimitedSize + : max_cache_memory_size_ / kNumShards; + } + + const size_t max_cache_memory_size_; + std::array shards_; +}; + +std::optional RuleLevelCache::GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt +) { + return pimpl_->GetCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt); +} + +bool RuleLevelCache::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask +) { + return pimpl_->AddCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt, token_mask); +} + +bool RuleLevelCache::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask +) { + return pimpl_->AddCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt, std::move(token_mask)); +} + +void RuleLevelCache::ClearCache() { pimpl_->ClearCache(); } + +size_t RuleLevelCache::GetMaxSize() const { return pimpl_->GetMaxSize(); } + +std::optional RuleLevelCache::Impl::GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt +) { + // Find in the cache. + NodeKey key = std::make_tuple(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt); + Shard& shard = GetShard(key); + std::lock_guard lock(shard.mutex); + auto it = shard.cache.find(key); + if (it == shard.cache.end()) { + return std::nullopt; + } + + // Move the node to the back of the list. + shard.cache_list.MoveBack(it->second); + return List::iterator(it->second, shard.cache_list)->second; +} + +bool RuleLevelCache::Impl::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask +) { + return AddCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt, AdaptiveTokenMask(token_mask)); +} + +bool RuleLevelCache::Impl::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask +) { + // Check if we can add to the cache. + NodeKey key = std::make_tuple(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt); + Shard& shard = GetShard(key); + const size_t shard_max_size = ShardMaxSize(); + std::lock_guard lock(shard.mutex); + if (shard_max_size != kUnlimitedSize && MemorySize(token_mask) > shard_max_size) { + // The token mask is too large to be cached. + return false; + } + if (shard.cache.find(key) != shard.cache.end()) { + // Already exists. + return false; + } + + // Evict old entries if needed. + if (shard_max_size != kUnlimitedSize) { + size_t new_item_size = MemorySize(token_mask); + while ((shard.current_cache_memory_size) > static_cast(shard_max_size - new_item_size) + ) { + auto oldest_it = shard.cache_list.begin(); + if (oldest_it == shard.cache_list.end()) { + // This should not happen if the size of the new item is smaller than + // the shard budget, but this is a safeguard. + break; + } + shard.current_cache_memory_size -= MemorySize(oldest_it->second); + shard.cache.erase(oldest_it->first); + shard.cache_list.Erase(oldest_it); + } + } + + // Add to the cache. + auto new_it = shard.cache_list.PushBack(NodeType(key, std::move(token_mask))); + shard.current_cache_memory_size += MemorySize(new_it->second); + shard.cache[key] = new_it.Index(); + return true; +} + +RuleLevelCache::RuleLevelCache(size_t max_cache_memory_size) + : pimpl_(std::make_shared(max_cache_memory_size)) {} + +void RuleLevelCache::Impl::ClearCache() { + for (auto& shard : shards_) { + std::lock_guard lock(shard.mutex); + shard.cache_list.Clear(); + shard.cache.clear(); + shard.current_cache_memory_size = 0; + } +} + +size_t MemorySize(const RuleLevelCache& manager) { return MemorySize(manager.ImplPtr()); } + +/*************************** Forward grammar constructors to their impl ***************************/ + +Grammar GrammarUnionFunctor::Apply(const std::vector& grammars) { + return GrammarUnionFunctorImpl().Apply(grammars); +} + +Grammar GrammarConcatFunctor::Apply(const std::vector& grammars) { + return GrammarConcatFunctorImpl().Apply(grammars); +} + +int32_t SubGrammarAdder::Apply(GrammarBuilder* builder, const Grammar& sub_grammar) { + return SubGrammarAdderImpl().ApplyWithBuilder(builder, sub_grammar); +} + +/*************************** Forward grammar Normalizers to their impl ***************************/ + +Grammar GrammarNormalizer::Apply(const Grammar& grammar) { + return GrammarNormalizerImpl().Apply(grammar); +} + +Grammar StructureNormalizer::Apply(const Grammar& grammar) { + return StructureNormalizerImpl().Apply(grammar); +} + +/*************************** Forward grammar optimizers to their impl ***************************/ + +void GrammarFSMBuilder::Apply(Grammar* grammar) { GrammarFSMBuilderImpl::Apply(grammar); } + +void RepetitionNormalizer::Apply(Grammar* grammar) { RepetitionNormalizerImpl().Apply(grammar); } + +void GrammarFSMHasher::Apply(Grammar* grammar) { GrammarFSMHasherImpl().Apply(grammar); } + +std::optional GrammarFSMHasher::HashSequence( + const Grammar& grammar, int32_t sequence_id +) { + return GrammarFSMHasherImpl().HashSequence(grammar, sequence_id); +} + +FSMWithStartEnd GrammarFSMBuilder::RuleRef(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::RuleRef(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::CharacterClass(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::CharacterClass(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::ByteString(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::ByteString(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::Token(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::Token(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::ExcludeToken(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::ExcludeToken(expr); +} + +std::optional GrammarFSMBuilder::TokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& ttd +) { + return GrammarFSMBuilderImpl::TokenTagDispatch(ttd); +} + +std::optional GrammarFSMBuilder::Sequence( + const GrammarExpr& expr, const Grammar& grammar +) { + return GrammarFSMBuilderImpl::Sequence(expr, grammar); +} + +std::optional GrammarFSMBuilder::Choices( + const GrammarExpr& expr, const Grammar& grammar +) { + return GrammarFSMBuilderImpl::Choices(expr, grammar); +} + +Result GrammarFSMBuilder::Regex(const std::string& regex, bool json_string) { + return GrammarFSMBuilderImpl::Regex(regex, json_string); +} + +const std::bitset<256>& GrammarFSMBuilder::JSONStringForbiddenChars() { + static const std::bitset<256> forbidden_chars = [] { + std::bitset<256> chars; + for (int c = 0x00; c <= 0x1F; ++c) { + chars.set(c); + } + chars.set('"'); + chars.set('\\'); + return chars; + }(); + return forbidden_chars; +} + +std::optional GrammarFSMBuilder::TagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch +) { + return GrammarFSMBuilderImpl::TagDispatch(tag_dispatch); +} + +std::vector AllowEmptyRuleAnalyzer::Apply(const Grammar& grammar) { + return AllowEmptyRuleAnalyzerImpl().Apply(grammar); +} + +void RuleInliner::Apply(Grammar* grammar) { RuleInlinerImpl().Apply(grammar); } + +Grammar DeadCodeEliminator::Apply(const Grammar& grammar) { + return DeadCodeEliminatorImpl().Apply(grammar); +} + +Grammar LookaheadAssertionAnalyzer::Apply(const Grammar& grammar) { + return LookaheadAssertionAnalyzerImpl().Apply(grammar); +} + +Grammar RepetitionRangeExpander::Apply(const Grammar& grammar) { + return RepetitionRangeExpanderImpl().Apply(grammar); +} + +Grammar GrammarOptimizer::Apply(const Grammar& grammar) { + return GrammarOptimizerImpl::Apply(grammar); +} + +void ByteStringFuser::Apply(Grammar* grammar) { ByteStringFuserImpl().Apply(grammar); } + +Grammar RootRuleRenamer::Apply(const Grammar& grammar) { + return RootRuleRenamerImpl().Apply(grammar); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_functor.h b/third_party/xgrammar/cpp/grammar_functor.h new file mode 100644 index 0000000000..0b4da505ba --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_functor.h @@ -0,0 +1,510 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_functor.h + * \brief The header for the simplification of the BNF AST. + */ + +#ifndef XGRAMMAR_GRAMMAR_FUNCTOR_H_ +#define XGRAMMAR_GRAMMAR_FUNCTOR_H_ + +#include + +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "grammar_builder.h" +#include "grammar_impl.h" +#include "support/utils.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief Base class for visitors and mutators of the BNF grammar. + * \tparam T The type of the return value of visitor functions. Typical values: + * - int32_t: the id of the new grammar_expr + * - void: no return value + * \tparam ReturnType The type of the return value of the transform function Apply(). Typical values + * are void (for visitor) and Grammar (for mutator). + */ +template +class GrammarFunctor { + public: + /*! + * \brief Constructor. + * \param grammar The grammar to visit or mutate. + */ + explicit GrammarFunctor() {} + + /*! + * \brief Apply the transformation to the grammar, or visit the grammar. + * \return The transformed grammar, or the visiting result, or void. + */ + virtual ReturnType Apply(const Grammar& grammar) { + // The initializer MUST be called at first when overriding the Apply() function. + InitGrammar(grammar); + if constexpr (std::is_same::value) { + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = rule.name; + VisitExpr(rule.body_expr_id); + VisitLookaheadAssertion(rule.lookahead_assertion_id); + } + return ReturnType(); + } else if constexpr (std::is_same::value && + std::is_same::value) { + InitBuilder(); + // First add empty rules to ensure the new rule ids the same as the old ones, then update + // the rule bodies + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + builder_->AddEmptyRule(base_grammar_->GetRule(i).name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = rule.name; + auto new_body_expr_id = VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(i, new_body_expr_id); + // Handle lookahead assertion + builder_->UpdateLookaheadAssertion(i, VisitLookaheadAssertion(rule.lookahead_assertion_id)); + builder_->UpdateMaxTokens(i, rule.max_tokens); + builder_->UpdateMaxChars(i, rule.max_chars); + builder_->UpdateCaptureName(i, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + builder_->UpdateSuffixStopInfo(i, *suffix_stop_info); + } + builder_->UpdateLazy(i, rule.is_lazy); + builder_->UpdateRuleTemperature(i, rule.temperature); + } + return builder_->Get(base_grammar_->GetRootRule().name); + } else { + return ReturnType(); + } + } + + /*! \brief Virtual destructor. */ + virtual ~GrammarFunctor() = default; + + protected: + using Rule = Grammar::Impl::Rule; + using GrammarExpr = Grammar::Impl::GrammarExpr; + using GrammarExprType = Grammar::Impl::GrammarExprType; + + /*! \brief Initialize the functor. Should be called at the beginning of Apply(). */ + virtual void InitGrammar() {} + + virtual void InitGrammar(const Grammar& grammar) { base_grammar_ = grammar; } + + virtual void InitBuilder() { + owned_builder_ = GrammarBuilder(); + builder_ = &owned_builder_; + } + + virtual void InitBuilder(const Grammar& grammar) { + owned_builder_ = GrammarBuilder(grammar); + builder_ = &owned_builder_; + } + + virtual void InitBuilder(GrammarBuilder* builder) { builder_ = builder; } + + /*! \brief Visit a lookahead assertion expr referred by id. */ + virtual T VisitLookaheadAssertion(int32_t lookahead_assertion_id) { + if (lookahead_assertion_id == -1) { + if constexpr (std::is_same::value) { + return -1; + } else { + return T(); + } + } + return VisitExpr(lookahead_assertion_id); + } + + /*! \brief Visit a GrammarExpr by id. */ + virtual T VisitExpr(int32_t old_grammar_expr_id) { + return VisitExpr(base_grammar_->GetGrammarExpr(old_grammar_expr_id)); + } + + /*! \brief Visit a GrammarExpr. Dispatch to the corresponding Visit function. */ + virtual T VisitExpr(const GrammarExpr& grammar_expr) { + switch (grammar_expr.type) { + case GrammarExprType::kSequence: + return VisitSequence(grammar_expr); + case GrammarExprType::kChoices: + return VisitChoices(grammar_expr); + case GrammarExprType::kEmptyStr: + return VisitEmptyStr(grammar_expr); + case GrammarExprType::kByteString: + return VisitByteString(grammar_expr); + case GrammarExprType::kCharacterClass: + return VisitCharacterClass(grammar_expr); + case GrammarExprType::kCharacterClassStar: + return VisitCharacterClassStar(grammar_expr); + case GrammarExprType::kRuleRef: + return VisitRuleRef(grammar_expr); + case GrammarExprType::kTagDispatch: + return VisitTagDispatch(grammar_expr); + case GrammarExprType::kRepeat: + return VisitRepeat(grammar_expr); + case GrammarExprType::kToken: + return VisitToken(grammar_expr); + case GrammarExprType::kExcludeToken: + return VisitExcludeToken(grammar_expr); + case GrammarExprType::kTokenTagDispatch: + return VisitTokenTagDispatch(grammar_expr); + case GrammarExprType::kRegex: + return VisitRegex(grammar_expr); + case GrammarExprType::kSubstring: + return VisitSubstring(grammar_expr); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected sequence type: " << static_cast(grammar_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + + /*! \brief Visit a choices GrammarExpr. */ + virtual T VisitChoices(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + for (auto i : grammar_expr) { + VisitExpr(i); + } + } else if constexpr (std::is_same::value) { + std::vector choice_ids; + for (int32_t i : grammar_expr) { + choice_ids.push_back(VisitExpr(i)); + } + return builder_->AddChoices(choice_ids); + } else { + return T(); + } + } + + /*! \brief Visit a sequence GrammarExpr. */ + virtual T VisitSequence(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + for (auto i : grammar_expr) { + VisitExpr(i); + } + } else if constexpr (std::is_same::value) { + std::vector sequence_ids; + for (int32_t i : grammar_expr) { + sequence_ids.push_back(VisitExpr(i)); + } + return builder_->AddSequence(sequence_ids); + } else { + return T(); + } + } + + virtual T VisitTagDispatch(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + return; + } else if constexpr (std::is_same::value) { + Grammar::Impl::TagDispatch tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + return builder_->AddTagDispatch(tag_dispatch); + } else { + return T(); + } + } + + /*! \brief Visit an element GrammarExpr, including empty string, character class, and rule ref. */ + virtual T VisitElement(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + return; + } else if constexpr (std::is_same::value) { + return builder_->AddGrammarExpr(grammar_expr); + } else { + return T(); + } + } + + /*! \brief Visit an empty string GrammarExpr. */ + virtual T VisitEmptyStr(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a character class GrammarExpr. */ + virtual T VisitByteString(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a character class GrammarExpr. */ + virtual T VisitCharacterClass(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + /*! \brief Visit a star quantifier GrammarExpr. */ + virtual T VisitCharacterClassStar(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + /*! \brief Visit a rule reference GrammarExpr. */ + virtual T VisitRuleRef(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a repeat GrammarExpr. */ + virtual T VisitRepeat(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + virtual T VisitToken(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + virtual T VisitExcludeToken(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + virtual T VisitTokenTagDispatch(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + /*! \brief Visit a regex GrammarExpr. It is a leaf: the pattern string is carried as-is. */ + virtual T VisitRegex(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a substring GrammarExpr. It is a leaf: the chunk list is carried as-is. */ + virtual T VisitSubstring(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief The grammar to visit or mutate. */ + Grammar base_grammar_{NullObj{}}; + + /*! + * \brief The builder to build the new grammar. It is empty when the mutator is constructed, and + * can be used to build a new grammar in subclasses. + */ + GrammarBuilder* builder_ = nullptr; + + GrammarBuilder owned_builder_; + + /*! \brief The name of the current rule being visited. */ + std::string cur_rule_name_; +}; + +/*! + * \brief Visitor of Grammar. + * \tparam ReturnType The return type of the Apply() function. Denotes the collected information. + */ +template +using GrammarVisitor = GrammarFunctor; + +/*! + * \brief Mutator of Grammar. The Apply() function returns the updated grammar. + */ +using GrammarMutator = GrammarFunctor; + +/****** All below methods are implemented as functor to hide the implementation ******/ + +/*************************** Grammar Constructor ***************************/ +/*! + * \brief Find the union of multiple grammars as a new grammar. + */ +class GrammarUnionFunctor { + public: + static Grammar Apply(const std::vector& grammars); +}; + +/*! + * \brief Find the concatenation of multiple grammars as a new grammar. + */ +class GrammarConcatFunctor { + public: + static Grammar Apply(const std::vector& grammars); +}; + +/*! + * \brief Add a sub grammar to the current builder. The return value + * of Apply is the new rule id of the sub grammar's root rule. + */ +class SubGrammarAdder { + public: + static int32_t Apply(GrammarBuilder* builder, const Grammar& sub_grammar); +}; + +/*************************** Grammar Normalizer ***************************/ + +/*! + * \brief Normalize a Grammar: expand the nested rules, combine consequent sequences and strings, + * etc. + */ +class GrammarNormalizer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Normalize the structure of the grammar. It will ensure each rule is a choices of + * sequences of elements, or a tag dispatch. The expanded context will be a sequence of elements. + */ +class StructureNormalizer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*************************** Grammar Optimizer ***************************/ + +/*! + * \brief Fuse adjacent byte string elements in sequences. + * \details Rewrites *grammar in place. Only sequences that actually contain adjacent byte strings + * are rewritten; if nothing needs fusing, the grammar is left untouched. The caller must own the + * grammar, as the input is mutated directly. + */ +class ByteStringFuser { + public: + static void Apply(Grammar* grammar); +}; + +/*! + * \brief Analyze the grammar to find the rules that are allowed to be empty. + */ +class AllowEmptyRuleAnalyzer { + public: + static std::vector Apply(const Grammar& grammar); +}; + +/*! + * \brief Inline the rule references in the grammar. + * \details Rewrites *grammar in place. Only choices with an inlinable leading rule reference are + * rewritten; if nothing can be inlined, the grammar is left untouched. The caller must own the + * grammar, as the input is mutated directly. + */ +class RuleInliner { + public: + static void Apply(Grammar* grammar); +}; + +/*! + * \brief Eliminate the not referenced rules in the grammar. + */ +class DeadCodeEliminator { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Analyze and add lookahead assertions in the grammar. + */ +class LookaheadAssertionAnalyzer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Build the FSMs of the grammar. + */ +class GrammarFSMBuilder { + using GrammarExpr = Grammar::Impl::GrammarExpr; + + public: + static void Apply(Grammar* grammar); + static FSMWithStartEnd RuleRef(const GrammarExpr& expr); + static FSMWithStartEnd CharacterClass(const GrammarExpr& expr); + static FSMWithStartEnd ByteString(const GrammarExpr& expr); + static FSMWithStartEnd Token(const GrammarExpr& expr); + static FSMWithStartEnd ExcludeToken(const GrammarExpr& expr); + static std::optional TokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch + ); + static std::optional Sequence(const GrammarExpr& expr, const Grammar& grammar); + static std::optional Choices(const GrammarExpr& expr, const Grammar& grammar); + static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); + /*! + * \brief Build the automaton of a regex pattern. Returns the error message on failure. + * \param regex The regex pattern string. + * \param json_string Whether the regex matches the body of a JSON string literal. If true, + * the characters in JSONStringForbiddenChars() are excluded from every character match. + */ + static Result Regex(const std::string& regex, bool json_string = false); + /*! \brief The characters that must be escaped inside a JSON string literal: the control + * characters 0x00-0x1F, the quote '"' and the backslash '\\'. */ + static const std::bitset<256>& JSONStringForbiddenChars(); +}; + +/*! + * \brief Normalize the repetition expression. If the context of + * repetition expression is nullable, then the repetition range will be + * normalized from {m, n} to {0, n} to reduce uncertainty. + */ +class RepetitionNormalizer { + public: + static void Apply(Grammar* grammar); +}; + +/*! + * \brief Expand kRepeat grammar expressions using HandleRepetitionRange logic. + * Transforms repetition structures into explicit sequences and choices. + */ +class RepetitionRangeExpander { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Optimize the grammar when compiling. + * \note No matter whether the grammar is optimized, grammar optimizer will + * return a new grammar. The following optimization will be applied: + * 1. Byte fuser. + * 2. Rule inliner. + * 3. Dead code eliminator. + * 4. Lookahead assertion analyzer. + * 5. Allow-empty rule analyzer. + * 6. Repetition normalizer. + * 7. FSM builder. + */ +class GrammarOptimizer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Rename the root rule of the grammar to "root". + */ +class RootRuleRenamer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Hash the fsms in the grammar, + * and get the new state ids of each fsm's states. + */ +class GrammarFSMHasher { + public: + static void Apply(Grammar* grammar); + static std::optional HashSequence(const Grammar& grammar, int32_t sequence_id); +}; + +/*! + * \brief Store the crossing cache for different grammars. + * \param max_cache_size The maximum size of the cache numbers. + * \details LRU algorithm is implemented. + */ +class RuleLevelCache { + public: + static const size_t kUnlimitedSize = static_cast(-1); + + std::optional GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt + ); + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask + ); + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask + ); + RuleLevelCache(size_t max_cache_memory_size = kUnlimitedSize); + + void ClearCache(); + + size_t GetMaxSize() const; + + friend size_t MemorySize(const RuleLevelCache& manager); + + XGRAMMAR_DEFINE_PIMPL_METHODS(RuleLevelCache); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_FUNCTOR_H_ diff --git a/third_party/xgrammar/cpp/grammar_impl.h b/third_party/xgrammar/cpp/grammar_impl.h new file mode 100644 index 0000000000..7b36a954cd --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_impl.h @@ -0,0 +1,487 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar.h + * \brief The header for the support of grammar-guided generation. + */ + +#ifndef XGRAMMAR_GRAMMAR_IMPL_H_ +#define XGRAMMAR_GRAMMAR_IMPL_H_ + +#include + +#include +#include +#include +#include + +#include "fsm.h" +#include "support/logging.h" +#include "support/reflection.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief This class stores the abstract syntax tree (AST) of the Backus-Naur Form (BNF) grammar. + * The BNF definition here is standard BNF, and the characters are represented using regex-style + * character classes (e.g. [a-z], [^a-z]). + * + * \details + * ### Rules + * The BNF grammar AST consists of a set of rules. Each rule contains a name and a definition, and + * corresponds to a production in the grammar. The definition of a rule is a GrammarExpr. Each rule + * has a rule_id for reference. + * + * ### GrammarExprs + * GrammarExpr is the definition of a rule or part of the definition of a rule. It can contain + * elements, empty string, reference to other GrammarExprs, or reference to other rules. Each + * GrammarExpr corresponds to a grammar_expr_id for reference. + * + * For example, in the following rule: rule ::= ("a" "b") | "c" + * ("a" "b"), "c", ("a" "b") | "c" are all GrammarExprs. + * + * #### Types of GrammarExprs + * Every GrammarExpr is represented by a type as well as a variable-length array containing its + * data. GrammarExpr has several types: + * - Byte string: a string of bytes (0~255). Supports UTF-8 strings. + * - Character class: a range of characters (each character is a unicode codepoint), e.g. [a-z], + * [ac-z]. Can be negated: [^a-z], [^ac-z]. Now only ascii chars is allowed in [], but this + * expression can accept/reject unicode chars. + * - Character class star: a star quantifier of a character class. e.g. [a-z]*, [^a-z]*. + * - EmptyStr: an empty string, i.e. "" + * - Rule reference: a reference to another rule + * - Sequence: a sequence of grammar_exprs, e.g. ("a" "b"). These grammar_exprs are concatenated + * together. + * - Choices: a choice of grammar_exprs, e.g. ("a" "b") | "c". Each grammar_expr can be matched. + * + * #### Storage of GrammarExprs + * Each type of GrammarExpr has a different data format. For the format of each type of GrammarExpr, + * see docs in Grammar::Impl::GrammarExprType. + * + * We store all GrammarExprs in csr_matrix style. That is, they are stored consecutively in one + * vector (data vector) and the starting position of each GrammarExpr is recorded in the indptr + * vector. + * + * \remark The character class star GrammarExpr is for the special support for elements like [a-z]* + * in the grammar. We add it to make the matching more efficient, as we can avoid recursion into + * rules when matching a sequence of characters. It should be used like: + * rule1 ::= ((element1 element2 rule2 ...) | ...) + * rule2 ::= character_class_star_grammar_expr(id_of_a_character_class_grammar_expr) + */ +class Grammar::Impl { + public: + /*! \brief A rule with name. */ + struct Rule { + /*! \brief The name of the rule. */ + std::string name; + /*! \brief The GrammarExpr id of the body of the rule. */ + int32_t body_expr_id; + /*! \brief The id of the associated lookahead assertion expr. For now it must be a id of a + * sequence GrammarExpr. -1 if not exists. */ + int32_t lookahead_assertion_id = -1; + /*! \brief Whether the lookahead assertion is exact. */ + bool is_exact_lookahead = false; + /*! \brief The token budget of the rule. When non-negative, the matcher bounds each + * occurrence of this rule to at most max_tokens LLM tokens, forcing it to end at the + * earliest possible position once the budget is exhausted. -1 means no budget. */ + int32_t max_tokens = -1; + /*! \brief The Unicode codepoint budget of the rule. When non-negative, the matcher bounds + * each occurrence of this rule to at most max_chars codepoints, forcing it to end at the + * earliest possible position once the budget is exhausted. -1 means no budget. */ + int32_t max_chars = -1; + /*! \brief The capture group name of the rule. When non-empty, the matcher records the input + * span matched by this rule on every completion, retrievable via GrammarMatcher::GetCaptures. + * Empty means no capture. */ + std::string capture_name = {}; + /*! \brief Whether the rule matches with committed-shortest (lazy) semantics: at the first + * position where the body can complete, it must complete. */ + bool is_lazy = false; + /*! \brief The sampling temperature to use while matching this rule. */ + std::optional temperature = std::nullopt; + }; + + /*! \brief Sparse per-rule metadata used to materialize suffix and stop captures. */ + struct SuffixStopInfo { + /*! \brief The rule carrying this metadata. */ + int32_t rule_id = -1; + /*! \brief Trailing bytes hidden only from this rule's own capture. */ + int32_t hidden_suffix_bytes = 0; + /*! \brief Trailing bytes hidden from this rule and every enclosing capture. */ + int32_t hidden_stop_bytes = 0; + /*! \brief Helper rule matching the body before a variable-length marker. A self-reference + * marks a zero-width event immediately following a fixed dynamic-dispatch marker. */ + int32_t body_rule_id = -1; + /*! \brief Helper rule matching a variable-length marker. */ + int32_t marker_rule_id = -1; + /*! \brief Capture name for the marker bytes. */ + std::string stop_capture_name = {}; + + bool IsEmpty() const { + return hidden_suffix_bytes == 0 && hidden_stop_bytes == 0 && body_rule_id == -1 && + marker_rule_id == -1 && stop_capture_name.empty(); + } + }; + + /*! \brief Get the number of rules. */ + int32_t NumRules() const { return rules_.size(); } + /*! \brief Get the rule with the given id. */ + const Rule& GetRule(int32_t rule_id) const { + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(rules_.size())) + << "rule_id " << rule_id << " is out of bound"; + return rules_[rule_id]; + } + Rule& GetRule(int32_t rule_id) { + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(rules_.size())) + << "rule_id " << rule_id << " is out of bound"; + return rules_[rule_id]; + } + /*! \brief Get sparse suffix/stop metadata for a rule, or nullptr when none exists. */ + const SuffixStopInfo* GetSuffixStopInfo(int32_t rule_id) const { + auto it = std::lower_bound( + suffix_stop_infos_.begin(), + suffix_stop_infos_.end(), + rule_id, + [](const SuffixStopInfo& info, int32_t id) { return info.rule_id < id; } + ); + return it != suffix_stop_infos_.end() && it->rule_id == rule_id ? &*it : nullptr; + } + /*! \brief Get the root rule id of the grammar. */ + int32_t GetRootRuleId() const { return root_rule_id_; } + /*! \brief Get the root rule of the grammar. */ + const Rule& GetRootRule() const { + XGRAMMAR_DCHECK(root_rule_id_ >= 0 && root_rule_id_ < static_cast(rules_.size())) + << "root_rule_id " << root_rule_id_ << " is out of bound"; + return rules_[root_rule_id_]; + } + + /*! + * \brief Check that every id and offset stored in the grammar is in range. Used after + * deserialization, where the fields are restored verbatim and none of the builder checks run. + * \return An error message if the grammar is malformed. + */ + std::optional Validate() const; + + /*! \brief The type of the grammar expr. */ + enum class GrammarExprType : int32_t { + // data format: [byte0, byte1, ...] + kByteString, + // data format: [is_negative, lower0, upper0, lower1, upper1, ...] + kCharacterClass, + kCharacterClassStar, + // data format: [] + kEmptyStr, + // data format: [rule_id] + kRuleRef, + // data format: [grammar_expr_id0, grammar_expr_id1, ...] + kSequence, + // data format: [grammar_expr_id0, grammar_expr_id1, ...] + kChoices, + // data format: [tag_expr0, rule_id0, tag_expr1, rule_id1, ..., loop_after_dispatch, + // excluded_str_expr_id] + kTagDispatch, + // data format: [rule_id, min_repeat_count, max_repeat_count] + kRepeat, + // data format: [token_id_0, token_id_1, ...] + kToken, + // data format: [token_id_0, token_id_1, ...] + kExcludeToken, + // data format: [trigger_cnt, (token_id, rule_id) × N, + // loop_after_dispatch, + // exclude_cnt, token_id × M] + kTokenTagDispatch, + // data format: [json_string, byte0, byte1, ...] + // The bytes are the regex pattern string. Like kTagDispatch, it can only be the body of a + // rule. The pattern is carried through the grammar passes as-is; when GrammarFSMBuilder + // runs, the pattern is compiled into an automaton, so every regex rule always has a + // per-rule FSM after optimization. If json_string is 1, the regex matches the body of a + // JSON string literal: the characters that must be escaped in a JSON string (the control + // characters, '"' and '\\') are excluded from every character match of the automaton. + kRegex, + // data format: [chunk0_len, byte0_0, byte0_1, ..., chunk1_len, byte1_0, ...] + // A list of length-prefixed byte string chunks. Matches every contiguous subsequence of + // the chunk list, including the empty one. Like kRegex, it can only be the body of a rule + // and is carried through the grammar passes as-is; when GrammarFSMBuilder runs, it is + // compiled into an automaton via a chunk-level suffix automaton (see SuffixAutomata). + kSubstring, + }; + + /*! \brief The object representing a grammar expr. */ + struct GrammarExpr { + /*! \brief The type of the grammar expr. */ + GrammarExprType type; + /*! \brief The data of the GrammarExpr. A variable-length array. */ + const int32_t* data; + /*! \brief The length of the data array. */ + int32_t data_len; + + int32_t size() const { return data_len; } + /*! \brief Get the i-th element of the data array. */ + const int32_t& operator[](int i) const { + XGRAMMAR_DCHECK(i >= 0 && i < static_cast(data_len)) + << "Index " << i << " is out of bound"; + return data[i]; + } + const int32_t* begin() const { return data; } + const int32_t* end() const { return data + data_len; } + void SetData(int index, int value) { const_cast(data)[index] = value; } + }; + + /*! \brief Get the number of grammar_exprs. */ + int32_t NumGrammarExprs() const { return grammar_expr_indptr_.size(); } + + /*! \brief Get the grammar_expr with the given id. */ + GrammarExpr GetGrammarExpr(int32_t grammar_expr_id) const { + XGRAMMAR_DCHECK( + grammar_expr_id >= 0 && grammar_expr_id < static_cast(grammar_expr_indptr_.size()) + ) << "grammar_expr_id " + << grammar_expr_id << " is out of bound"; + int start_index = grammar_expr_indptr_[grammar_expr_id]; + auto start_ptr = grammar_expr_data_.data() + start_index; + auto type = static_cast(start_ptr[0]); + auto data_ptr = start_ptr + 2; + auto data_len = start_ptr[1]; + return {type, data_ptr, data_len}; + } + + /******************* GrammarExpr Getters *******************/ + + /*! \brief Get the string of the byte string grammar expr. */ + std::string GetByteString(const GrammarExpr& grammar_expr) const { + std::string str; + str.reserve(grammar_expr.size()); + for (int i = 0; i < grammar_expr.size(); ++i) { + str.push_back(static_cast(static_cast(grammar_expr[i]))); + } + return str; + } + + /*! \brief Get the string of the byte string grammar expr. */ + std::string GetByteString(int32_t grammar_expr_id) const { + return GetByteString(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief Get the regex pattern string of the regex grammar expr. */ + std::string GetRegexString(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kRegex) << "GrammarExpr is not a regex"; + XGRAMMAR_DCHECK(grammar_expr.size() >= 1) << "Regex expr must contain the json_string flag"; + std::string str; + str.reserve(grammar_expr.size() - 1); + for (int i = 1; i < grammar_expr.size(); ++i) { + str.push_back(static_cast(static_cast(grammar_expr[i]))); + } + return str; + } + + /*! \brief Get the regex pattern string of the regex grammar expr with the given id. */ + std::string GetRegexString(int32_t grammar_expr_id) const { + return GetRegexString(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief Get whether the regex grammar expr matches the body of a JSON string literal. */ + bool GetRegexIsJSONString(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kRegex) << "GrammarExpr is not a regex"; + XGRAMMAR_DCHECK(grammar_expr.size() >= 1) << "Regex expr must contain the json_string flag"; + return grammar_expr[0] != 0; + } + + /*! \brief Get the chunk list of the substring grammar expr. */ + std::vector GetSubstringChunks(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kSubstring) + << "GrammarExpr is not a substring"; + std::vector chunks; + for (int i = 0; i < grammar_expr.size();) { + int32_t chunk_len = grammar_expr[i++]; + XGRAMMAR_DCHECK(chunk_len >= 0 && i + chunk_len <= grammar_expr.size()) + << "Invalid substring chunk length"; + std::string chunk; + chunk.reserve(chunk_len); + for (int32_t j = 0; j < chunk_len; ++j) { + chunk.push_back(static_cast(static_cast(grammar_expr[i++]))); + } + chunks.push_back(std::move(chunk)); + } + return chunks; + } + + /*! \brief Get the chunk list of the substring grammar expr with the given id. */ + std::vector GetSubstringChunks(int32_t grammar_expr_id) const { + return GetSubstringChunks(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief The object representing a tag dispatch. */ + struct TagDispatch { + /*! \brief The tag and rule id pairs. */ + std::vector> tag_rule_pairs; + /*! \brief If true, the tag dispatch will loop after dispatching. */ + bool loop_after_dispatch; + /*! \brief The strings that are excluded by the tag dispatch. */ + std::vector excludes; + static const int kTagDispatchExtraParameter = 2; + }; + + /*! \brief Get the tag dispatch from the grammar expr. */ + TagDispatch GetTagDispatch(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kTagDispatch) + << "GrammarExpr is not a tag dispatch"; + + TagDispatch result; + XGRAMMAR_DCHECK(grammar_expr.size() >= TagDispatch::kTagDispatchExtraParameter); + result.tag_rule_pairs.reserve( + (grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter) / 2 + ); + + for (int i = 0; i < grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter; i += 2) { + auto tag_expr_id = grammar_expr[i]; + auto rule_id = grammar_expr[i + 1]; + result.tag_rule_pairs.push_back({GetByteString(tag_expr_id), rule_id}); + } + + result.loop_after_dispatch = static_cast( + grammar_expr[grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter] + ); + + auto exclude_str_expr = GetGrammarExpr( + grammar_expr[grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter + 1] + ); + XGRAMMAR_DCHECK(exclude_str_expr.type == GrammarExprType::kChoices); + result.excludes.reserve(exclude_str_expr.size()); + for (int j = 0; j < exclude_str_expr.size(); j++) { + result.excludes.push_back(GetByteString(exclude_str_expr[j])); + } + return result; + } + + /*! \brief Get the tag dispatch from the grammar expr with the given id. */ + TagDispatch GetTagDispatch(int32_t grammar_expr_id) const { + return GetTagDispatch(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief The object representing a token tag dispatch. */ + struct TokenTagDispatch { + std::vector> trigger_rule_pairs; // token_id → rule_id + bool loop_after_dispatch; + std::vector excludes; + }; + + /*! \brief Decode a kTokenTagDispatch expr into the TokenTagDispatch struct. */ + TokenTagDispatch GetTokenTagDispatch(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kTokenTagDispatch); + TokenTagDispatch result; + int pos = 0; + int32_t trigger_count = grammar_expr[pos++]; + for (int i = 0; i < trigger_count; ++i) { + auto token_id = grammar_expr[pos++]; + auto rule_id = grammar_expr[pos++]; + result.trigger_rule_pairs.push_back({token_id, rule_id}); + } + result.loop_after_dispatch = static_cast(grammar_expr[pos++]); + int32_t exclude_count = grammar_expr[pos++]; + for (int i = 0; i < exclude_count; ++i) { + result.excludes.push_back(grammar_expr[pos++]); + } + XGRAMMAR_DCHECK(pos == grammar_expr.size()); + return result; + } + + /*! \brief Get the token tag dispatch from the grammar expr with the given id. */ + TokenTagDispatch GetTokenTagDispatch(int32_t grammar_expr_id) const { + return GetTokenTagDispatch(GetGrammarExpr(grammar_expr_id)); + } + + private: + /*! \brief The rules of the grammar. rule_id corresponds the index of this vector. */ + std::vector rules_; + /*! \brief Suffix/stop metadata, sorted by rule_id and omitted for ordinary rules. */ + std::vector suffix_stop_infos_; + /*! \brief The data of all grammar_exprs. */ + std::vector grammar_expr_data_; + /*! \brief The start index of every grammar_expr in grammar_expr_data_. grammar_expr_id is the + * index to the elements in this vector. */ + std::vector grammar_expr_indptr_; + /*! \brief The id of the root rule. */ + int32_t root_rule_id_ = -1; + + public: + /******************* Aux information for matching *******************/ + + /*! \brief The complete FSM for the grammar. It contains the FSMs for all rules. */ + CompactFSM complete_fsm{NullObj{}}; + + /*! + * \brief The FSM for each rule. + * \details The FSM will be used in matching if it exists. If it does not exist (std::nullopt), + * the rule will be used in matching, and the rule's body must be a kChoices expr. + */ + std::vector> per_rule_fsms; + + /*! + * \brief The hash value for each rule's FSM. + */ + std::vector> per_rule_fsm_hashes; + + /*! + * \brief The new state ids of each FSM's states. + */ + std::vector>> per_rule_fsm_new_state_ids; + + /*! \brief The ids of the rules that are allowed to be empty. */ + std::vector allow_empty_rule_ids; + + /*! \brief Whether the grammar is optimized. */ + bool optimized = false; + + friend class GrammarBuilder; + friend class GrammarCompiler; + + friend std::size_t MemorySize(const Impl& impl); + friend struct member_trait; +}; + +XGRAMMAR_MEMBER_ARRAY( + Grammar::Impl::Rule, + &Grammar::Impl::Rule::name, + &Grammar::Impl::Rule::body_expr_id, + &Grammar::Impl::Rule::lookahead_assertion_id, + &Grammar::Impl::Rule::is_exact_lookahead, + &Grammar::Impl::Rule::max_tokens, + &Grammar::Impl::Rule::max_chars, + &Grammar::Impl::Rule::capture_name, + &Grammar::Impl::Rule::is_lazy, + &Grammar::Impl::Rule::temperature +); + +XGRAMMAR_MEMBER_ARRAY( + Grammar::Impl::SuffixStopInfo, + &Grammar::Impl::SuffixStopInfo::rule_id, + &Grammar::Impl::SuffixStopInfo::hidden_suffix_bytes, + &Grammar::Impl::SuffixStopInfo::hidden_stop_bytes, + &Grammar::Impl::SuffixStopInfo::body_rule_id, + &Grammar::Impl::SuffixStopInfo::marker_rule_id, + &Grammar::Impl::SuffixStopInfo::stop_capture_name +); + +XGRAMMAR_MEMBER_TABLE( + Grammar::Impl, + "rules", + &Grammar::Impl::rules_, + "suffix_stop_infos", + &Grammar::Impl::suffix_stop_infos_, + "grammar_expr_data", + &Grammar::Impl::grammar_expr_data_, + "grammar_expr_indptr", + &Grammar::Impl::grammar_expr_indptr_, + "root_rule_id", + &Grammar::Impl::root_rule_id_, + "complete_fsm", + &Grammar::Impl::complete_fsm, + "per_rule_fsms", + &Grammar::Impl::per_rule_fsms, + "allow_empty_rule_ids", + &Grammar::Impl::allow_empty_rule_ids, + "optimized", + &Grammar::Impl::optimized +); + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_IMPL_H_ diff --git a/third_party/xgrammar/cpp/grammar_matcher.cc b/third_party/xgrammar/cpp/grammar_matcher.cc new file mode 100644 index 0000000000..ba63893e9a --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_matcher.cc @@ -0,0 +1,2826 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_matcher.cc + * \brief This source file implement the matcher class, especially the logic related to LLM tokens, + * like accepting tokens, leveraging the token mask cache to generate the mask, etc. matcher_base.cc + * implements the basic matching algorithm from strings to grammar. + */ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "earley_parser.h" +#include "grammar_impl.h" +#include "support/dynamic_bitset.h" +#include "support/encoding.h" +#include "support/int_set.h" +#include "support/logging.h" +#include "support/thread_pool.h" +#include "testing.h" +#include "tokenizer_info_impl.h" + +namespace xgrammar { + +/******************* Tool functions for token mask *******************/ +using GrammarExprType = Grammar::Impl::GrammarExprType; + +int32_t GetBitmaskSize(int vocab_size) { return DynamicBitset::GetBufferSize(vocab_size); } + +DLDataType GetBitmaskDLType() { return DLDataType{kDLInt, 32, 1}; } + +namespace details { + +using Clock = std::chrono::steady_clock; +using TimePoint = Clock::time_point; + +void ClearTokenBitmaskRow(int32_t* bitmask_data, int32_t bitmask_size, int32_t position) { + std::fill_n(bitmask_data + position * bitmask_size, bitmask_size, 0); +} + +bool TraverseDraftTreeRecursive( + int32_t current_position, + int32_t parent_position, + const int64_t* retrieve_next_token, + const int64_t* retrieve_next_sibling, + const int64_t* draft_tokens, + GrammarMatcher& matcher, + DLTensor* token_bitmask, + float* temperatures, + double time_threshold, + const TimePoint& start_time +) { + int32_t* bitmask_data = reinterpret_cast(token_bitmask->data); + int32_t bitmask_size = static_cast(token_bitmask->shape[1]); + + bool accepted; + if (current_position == 0) { + // The first token generated by the target model is always accepted. + accepted = true; + } else { + XGRAMMAR_CHECK(parent_position >= 0) + << "Non-root draft tree nodes must have a valid parent position"; + int64_t current_token_id = draft_tokens[current_position]; + if (current_token_id < 0 || current_token_id >= static_cast(bitmask_size) * 32) { + accepted = false; + } else { + int32_t* parent_bitmask = bitmask_data + parent_position * bitmask_size; + // 32 boolean bitmask values are packed into 32-bit integers. + uint32_t token_mask = uint32_t{1} << static_cast(current_token_id % 32); + accepted = (static_cast(parent_bitmask[current_token_id / 32]) & token_mask) != 0; + } + + // Check timeout for non-root nodes so the root token mask is still computed. + if (accepted && time_threshold > 0) { + auto elapsed = std::chrono::duration(Clock::now() - start_time).count(); + if (elapsed > time_threshold) { + return false; + } + } + } + + if (accepted) { + bool token_accepted = true; + if (current_position != 0) { + token_accepted = matcher.AcceptToken(static_cast(draft_tokens[current_position])); + } + + if (token_accepted) { + if (!matcher.IsTerminated()) { + matcher.FillNextTokenBitmask(token_bitmask, current_position); + if (temperatures != nullptr) { + temperatures[current_position] = matcher.GetTemperature().value_or(-1.0f); + } + + if (retrieve_next_token[current_position] != -1) { + bool success = TraverseDraftTreeRecursive( + retrieve_next_token[current_position], + current_position, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + matcher, + token_bitmask, + temperatures, + time_threshold, + start_time + ); + if (!success) { + if (current_position != 0) { + matcher.Rollback(1); + } + return false; + } + } + } else { + ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position); + } + + if (current_position != 0) { + matcher.Rollback(1); + } + } else { + ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position); + } + } else { + ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position); + } + + if (retrieve_next_sibling[current_position] != -1) { + bool success = TraverseDraftTreeRecursive( + retrieve_next_sibling[current_position], + parent_position, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + matcher, + token_bitmask, + temperatures, + time_threshold, + start_time + ); + if (!success) { + return false; + } + } + + return true; +} + +} // namespace details + +int32_t* CheckAndGetBitmaskPtr(const DLTensor& token_bitmask, int vocab_size, int index) { + XGRAMMAR_CHECK(token_bitmask.dtype.code == kDLInt && token_bitmask.dtype.bits == 32) + << "The provied bitmask's dtype is not valid: should be int32"; + + int32_t buffer_size = GetBitmaskSize(vocab_size); + if (token_bitmask.ndim == 1) { + XGRAMMAR_CHECK(token_bitmask.shape[0] == buffer_size) + << "The provided bitmask's shape is not valid: should be (" << buffer_size << ", )"; + XGRAMMAR_CHECK(index == 0) << "The index should be 0 when the bitmask is 1D"; + } else { + XGRAMMAR_CHECK(token_bitmask.ndim == 2) + << "The provided bitmask's shape is not valid: should be (batch_size, " << buffer_size + << ")"; + XGRAMMAR_CHECK(token_bitmask.shape[1] == buffer_size) + << "The provided bitmask's shape is not valid: should be (batch_size, " << buffer_size + << ")"; + XGRAMMAR_CHECK(index >= 0 && index < token_bitmask.shape[0]) + << "The provided index is out of bounds"; + } + + XGRAMMAR_CHECK( + token_bitmask.device.device_type == kDLCPU || + token_bitmask.device.device_type == kDLCUDAHost || + token_bitmask.device.device_type == kDLROCMHost + ) << "The provided bitmask's device is not valid: should be CPU"; + + return reinterpret_cast(token_bitmask.data) + index * buffer_size; +} + +void _DebugGetMaskedTokensFromBitmask( + std::vector* rejected_tokens, const DLTensor& token_bitmask, int vocab_size, int index +) { + int32_t* data_ptr = CheckAndGetBitmaskPtr(token_bitmask, vocab_size, index); + DynamicBitset bitset(vocab_size, reinterpret_cast(data_ptr)); + rejected_tokens->clear(); + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + rejected_tokens->push_back(i); + } +} + +std::pair _IsSingleTokenBitmask(const DLTensor& bitmask, int vocab_size, int index) { + int32_t* data_ptr = CheckAndGetBitmaskPtr(bitmask, vocab_size, index); + DynamicBitset bitset(vocab_size, reinterpret_cast(data_ptr)); + if (bitset.Count() == 1) { + return std::make_pair(true, bitset.FindFirstOne()); + } else { + return std::make_pair(false, -1); + } +} + +void ApplyMask32Bits( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size, + std::optional> indices +) { + XGRAMMAR_CHECK(logits->dtype.code == kDLFloat && logits->dtype.bits == 32) + << "The provided logits's dtype is not valid: should be float32"; + std::pair logits_shape = + logits->ndim == 2 + ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) + : std::make_pair(1, static_cast(logits->shape[0])); + int logits_stride0 = logits->strides[0]; + int bitmask_stride0 = bitmask.strides[0]; + if (indices.has_value()) { + for (auto idx : indices.value()) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = -std::numeric_limits::infinity(); + } + } + } else { + for (int idx = 0; idx < logits_shape.first; ++idx) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = -std::numeric_limits::infinity(); + } + } + } +} + +void ApplyMask16Bits( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size, + std::optional> indices +) { + XGRAMMAR_CHECK(logits->dtype.bits == 16) + << "The provided logits's dtype is not valid: should be bfloat16 or float16"; + uint16_t kMinusInfinity; + const uint16_t kMinusInfinityBf16 = 0xff80; + const uint16_t kMinusInfinityFp16 = 0xfc00; + switch (logits->dtype.code) { + case kDLBfloat: + kMinusInfinity = kMinusInfinityBf16; + break; + case kDLFloat: + kMinusInfinity = kMinusInfinityFp16; + break; + default: + XGRAMMAR_LOG(FATAL + ) << "The provided logits's dtype is not valid: should be bfloat16 or float16"; + } + std::pair logits_shape = + logits->ndim == 2 + ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) + : std::make_pair(1, static_cast(logits->shape[0])); + int logits_stride0 = logits->strides[0]; + int bitmask_stride0 = bitmask.strides[0]; + if (indices.has_value()) { + for (auto idx : indices.value()) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = kMinusInfinity; + } + } + } else { + for (int idx = 0; idx < logits_shape.first; ++idx) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = kMinusInfinity; + } + } + } +} + +void ApplyTokenBitmaskInplaceCPU( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size, + std::optional> indices +) { + // Check device and dim + XGRAMMAR_CHECK( + logits->device.device_type == kDLCPU || logits->device.device_type == kDLCUDAHost || + logits->device.device_type == kDLROCMHost + ) << "The provided logits's device is not valid: should be CPU"; + XGRAMMAR_CHECK( + bitmask.device.device_type == kDLCPU || bitmask.device.device_type == kDLCUDAHost || + bitmask.device.device_type == kDLROCMHost + ) << "The provided bitmask's device is not valid: should be CPU"; + XGRAMMAR_CHECK(logits->ndim == 2 || logits->ndim == 1) + << "The provided logits's shape is not valid: should be 2D or 1D"; + XGRAMMAR_CHECK(bitmask.ndim == 2 || bitmask.ndim == 1) + << "The provided bitmask's shape is not valid: should be 2D or 1D"; + + // Check type + XGRAMMAR_CHECK(logits->dtype.lanes == 1) + << "The provided logits's dtype is not valid: lanes should be 1"; + XGRAMMAR_CHECK( + bitmask.dtype.code == kDLInt && bitmask.dtype.bits == 32 && bitmask.dtype.lanes == 1 + ) << "The provided bitmask's dtype is not valid: should be int32"; + + // Check shape + std::pair logits_shape = + logits->ndim == 2 + ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) + : std::make_pair(1, static_cast(logits->shape[0])); + std::pair bitmask_shape = + bitmask.ndim == 2 + ? std::make_pair(static_cast(bitmask.shape[0]), static_cast(bitmask.shape[1])) + : std::make_pair(1, static_cast(bitmask.shape[0])); + + XGRAMMAR_CHECK( + vocab_size <= bitmask_shape.second * DynamicBitset::BITS_PER_BLOCK && + vocab_size <= logits_shape.second + ); + + if (!indices.has_value()) { + XGRAMMAR_CHECK(logits_shape.first == bitmask_shape.first) + << "When indices is not provided, the logits's batch size should be equal to the " + "bitmask's batch size, but got " + << logits_shape.first << " vs " << bitmask_shape.first; + } else { + // Each index selects a row of both logits and bitmask; an out-of-range index would lead to + // an out-of-bounds access, so validate every index against both batch sizes here. + for (int idx : indices.value()) { + XGRAMMAR_CHECK(idx >= 0 && idx < logits_shape.first && idx < bitmask_shape.first) + << "The provided index " << idx << " is out of bounds: it should be in [0, " + << std::min(logits_shape.first, bitmask_shape.first) << ")."; + } + } + + // Apply mask + if (logits->dtype.bits == 32) { + ApplyMask32Bits(logits, bitmask, vocab_size, indices); + } else if (logits->dtype.bits == 16) { + ApplyMask16Bits(logits, bitmask, vocab_size, indices); + } else { + XGRAMMAR_LOG(FATAL + ) << "The provided logits's dtype is not valid: should be float32 or float16/bfloat16"; + } +} + +/******************* Grammar Matcher with Adaptive Token Mask *******************/ + +/* + * Note on the matching algorithm (this is the old description for the matching algorithm, please + * refer to https://arxiv.org/pdf/2411.15100 for the latest description) + * + * Given a context-free grammar, we match the characters in a string one by one. + * + * We adopt a non-deterministic pushdown automata (NPDA) in matching. To be specific, we maintain + * several stacks, each of which represents a possible path in the NPDA, and update the stacks + * during matching. + * + * ## Stack Structure (see grammar_matcher_state.h) + * The element of every stack is a StackElement object, referring a position in the grammar. If a + * StackElement points to a RuleRef element (referring to another rule), the next element of the + * stack will be a position in this rule. If a StackElement is a CharacterClass element, it will be + * the last in the stack, meaning *the next* character to match. + * + * ## Matching Process (see grammar_matcher_base.h) + * When accepting a new character and it is accepted by a stack, the last element of the stack will + * be advanced to the next position in the grammar. If it gets to the end of the rule, several + * elements at the end may be popped out, and the last element of the stack will be advanced. + * + * One stack may split since there may be multiple possible next positions. In this case, similar + * stacks with different top elements will be added. When one stack cannot accept the new character, + * it will be removed from the stacks. + * + * ## Storage of Stacks (see grammar_matcher_state.h) + * Note these stacks form a tree structure as when splitting, the new stacks share the same prefix. + * We store all StackElements as a tree, where every path from tree root to a node represents a + * stack. To represent stack tops, we attach additional pointers pointing the stack top nodes. + * Also, We maintain a history of the stack top pointers, so we can rollback to the previous state. + * + * All tree nodes are maintained by a buffer, and utilize reference counting to recycle. If a node + * is neither pointed by a stack top pointer, not pointed by some child nodes, it will be freed. + * + * ## Example + * ### Grammar + * root ::= [a] R + * R ::= [b] S [c] | [b] [c] T + * S ::= "" | [c] [d] + * T ::= [e] + * + * ### The previous step + * Previous accepted string: ab + * Previous stack tree: + * A------ + * | \ \ + * B D< E< + * | + * C< + * + * A: (rule root, choice 0, element 1) + * B: (rule R, choice 0, element 1) + * C: (rule S, choice 1, element 0) + * D: (rule R, choice 0, element 2) + * E: (rule R, choice 1, element 1) + * < means the stack top pointers in the previous step. + * The stacks in the previous step is: (A, B, C), (A, D), (A, E) + * + * ### The current step + * Current accepted string: abc + * Current stack tree: + * A----------------- G<< + * | \ \ \ + * B--- D< E< H + * | \ | + * C< F<< I<< + * + * F: (rule S, choice 1, element 1) + * G: (rule root, choice 0, element 2) (means the matching process has finished, and will be deleted + * when the next char comes) + * H: (rule R, choice 1, element 2) + * I: (rule T, choice 0, element 0) + * << means the stack top pointers in the current step. + * The stacks in the current step is: (A, B, F), (A, H, I), (G,) + * + * ## Preprocess (see grammar_matcher_preproc.h) + * We will store all information about tokens that needed in matching in a CompiledGrammar + * object. Tokens are sorted by codepoint, allowing us to reuse the repeated prefixes between + * different tokens. + * + * For a given position in a rule, if we only consider this rule and its sub-rules during matching, + * without considering its parent rules (in actual matching, we also need to consider its parent + * rules), we can already determine that some tokens are acceptable while others are definitely + * rejected. Therefore, for a position in a rule, we can divide the token set into three categories: + * - accepted_indices: If a token is accepted by this rule + * - rejected_indices: If a token is rejected by this rule + * - uncertain_indices: Whether it can be accepted depends on the information from the parent + * level during actual matching. To be specific, If this token has a prefix that has not been + * rejected and has reached the end of this rule, then it is possible for it to be further accepted + * by the parent rule. + * + * During actual matching, we will directly accept or reject the tokens in accepted_indices and + * rejected_indices, and only consider the tokens in uncertain_indices. That speeds up the matching + * process. + */ + +/* \brief The concrete implementation of GrammarMatcherNode. */ +class GrammarMatcher::Impl : public EarleyParser { + public: + Impl( + const CompiledGrammar& compiled_grammar, + std::optional> override_stop_tokens = std::nullopt, + bool terminate_without_stop_token = false, + // max_rollback_tokens_ is deprecated and not used. + int max_rollback_tokens = -1, + std::optional default_temperature = std::nullopt + ) + : EarleyParser(compiled_grammar->grammar), + compiled_grammar_(compiled_grammar), + tokenizer_info_(compiled_grammar->tokenizer_info), + stop_token_ids_(override_stop_tokens.value_or(tokenizer_info_.GetStopTokenIds())), + terminate_without_stop_token_(terminate_without_stop_token), + default_temperature_(default_temperature), + tmp_accepted_bitset_(tokenizer_info_.GetVocabSize()) { + if (override_stop_tokens.has_value()) { + XGRAMMAR_CHECK(!override_stop_tokens->empty()) + << "The override_stop_tokens should not be empty"; + // Stop token ids are written into the token bitmask, so they must be in range. + for (int id : *override_stop_tokens) { + XGRAMMAR_CHECK(id >= 0 && id < tokenizer_info_.GetVocabSize()) + << "The override stop token id " << id << " is out of the vocabulary range [0, " + << tokenizer_info_.GetVocabSize() << ")"; + } + } + if (has_budget_rules_ || has_char_budget_rules_) { + for (int32_t rule_id = 0; rule_id < grammar_->NumRules(); ++rule_id) { + const auto& rule = grammar_->GetRule(rule_id); + if (rule.max_tokens < 0 && rule.max_chars < 0) { + continue; + } + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(rule_id); + if (suffix_stop_info != nullptr && suffix_stop_info->body_rule_id >= 0) { + has_budget_marker_rules_ = true; + break; + } + } + } + XGRAMMAR_CHECK( + !default_temperature_.has_value() || + (std::isfinite(default_temperature_.value()) && default_temperature_.value() >= 0) + ) << "The default_temperature must be a finite non-negative number"; + } + + bool AcceptToken(int32_t token_id, bool debug_print = false); + + bool AcceptString(const std::string& input_str, bool debug_print = false); + + bool FillNextTokenBitmask(DLTensor* next_token_bitmask, int index, bool debug_print = false); + + std::optional GetTemperature() const; + + std::string FindJumpForwardString(); + + void Rollback(int num_tokens); + + bool IsTerminated() const; + + void Reset() { + token_length_history.clear(); + current_token_index_ = -1; + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + budget_exceeded_warned_ = false; + char_budget_exceeded_warned_ = false; + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = false; + budget_body_match_cache_.clear(); + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + accepted_bytes_.clear(); + row_byte_end_.assign(1, 0); + EarleyParser::Reset(); + } + + std::vector> GetCaptures(bool deduplicate) const; + + int GetMaxRollbackTokens() const { return -1; } + + const std::vector& GetStopTokenIds() const { return stop_token_ids_; } + + std::string _DebugPrintInternalState() const { return PrintStates(); } + + private: + using StoreType = AdaptiveTokenMask::StoreType; + + /*! + * \brief If is_uncertain_saved is true, find the next token in uncertain_indices. Otherwise, + * find the next token that is set to true in uncertain_tokens_bitset. + * \param iterator_uncertain The helper iterator to iterate over uncertain_indices or + * uncertain_tokens_bitset. + * \returns The index of the next token, or -1 if no more token. + */ + int GetNextUncertainToken( + bool is_uncertain_saved, + int* iterator_uncertain, + const std::vector& uncertain_indices, + const std::vector& uncertain_tokens_bitset + ); + + /*! \brief Set the acceptable next token in next_token_bitmask. */ + void SetTokenBitmask( + int32_t* bitmask_data_ptr, + const DynamicBitset& accepted_bitset, + const std::vector& rejected_indices, + bool can_reach_end, + bool allow_special_token = false + ); + + /*! + * \brief Accept the stop token and terminates the matcher. + * \returns Whether the stop token can be accepted. + */ + bool AcceptStopToken(); + + bool IsStopTokenAccepted() const; + + /*! \brief Check if the token bitmask is all-true. */ + bool IsTokenBitmaskAllTrue(int32_t* bitmask_data_ptr); + + std::string PrintBitmask(int32_t* bitmask_data_ptr, const TokenizerInfo& tokenizer_info); + + /*! \brief Restores skip_expired_states_ to false when the accept call returns. */ + class SkipExpiredGuard { + public: + explicit SkipExpiredGuard(Impl* impl, bool enable) : impl_(impl) { + impl_->skip_expired_states_ = enable; + } + ~SkipExpiredGuard() { impl_->skip_expired_states_ = false; } + SkipExpiredGuard(const SkipExpiredGuard&) = delete; + SkipExpiredGuard& operator=(const SkipExpiredGuard&) = delete; + + private: + Impl* impl_; + }; + + /*! + * \brief RAII guard that enables capture-event recording in the Earley parser for its + * lifetime. Used in the definitive accept paths (AcceptToken / AcceptString) only, so that + * speculative advances (mask computation, jump-forward search) never record capture events. + */ + class CaptureRecordingGuard { + public: + explicit CaptureRecordingGuard(Impl* impl) : impl_(impl) { + if (impl_->capture_tracking_) { + impl_->capture_recording_ = true; + } + } + ~CaptureRecordingGuard() { impl_->capture_recording_ = false; } + CaptureRecordingGuard(const CaptureRecordingGuard&) = delete; + CaptureRecordingGuard& operator=(const CaptureRecordingGuard&) = delete; + + private: + Impl* impl_; + }; + + /*! \brief Fill the bitmask from the current states, optionally excluding the states whose + * budget deadline has passed. */ + void FillBitmaskForStates( + int32_t* bitmask_data_ptr, int index, bool skip_expired, bool debug_print + ); + + void FillBitmaskForCharBudgetBoundary( + const AdaptiveTokenMask& adaptive_token_mask, int32_t remaining_chars + ); + + bool AdvanceWithCharacterBudget(uint8_t byte, bool debug_print = false); + + bool AdvanceAtomicTokenWithCharacterBudget( + int32_t token_id, int32_t token_char_count, bool debug_print = false + ); + + /*! \brief Whether byte offsets are needed for captures or budgeted suffix/stop rules. */ + bool ShouldTrackAcceptedBytes() const { + return IsCaptureTrackingEnabled() || has_budget_marker_rules_; + } + + /*! \brief Whether the bytes consumed by this expired suffix/stop occurrence form a complete + * match of its body, so max_tokens may end it without consuming the marker. */ + bool CanForceCompleteWithoutMarker(const ParserState& state, bool use_char_budget); + + /*! \brief Commit the current budget boundary: discard expired states and complete eligible + * suffix/stop occurrences without their marker. Returns whether a derivation remains viable. */ + bool ApplyBudgetEnforcement(bool debug_print = false); + + /*! \brief Enforce exhausted character budgets before consuming a new codepoint. */ + bool ApplyCharacterBudgetEnforcement(bool debug_print = false); + + /*! \brief Reapply character-budget enforcement for newly exposed parent occurrences. */ + bool ApplyCharacterBudgetEnforcementToFixedPoint(bool debug_print = false); + + /*! \brief Consume the pending budget decision and warn once when a token budget was + * exceeded. */ + void FinishCharacterBudgetAccept() { + if (char_budget_relaxed_ && !char_budget_exceeded_warned_) { + char_budget_exceeded_warned_ = true; + XGRAMMAR_LOG(WARNING + ) << "The character budget (max_chars) of a rule was exceeded: the rule could not end at " + "the position where its budget ran out, so the budget is relaxed until the rule " + "can end. This warning is reported once per matcher."; + } + char_budget_relaxed_ = false; + } + + bool FinishAccept(bool consumed_past_deadline) { + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + if (consumed_past_deadline && !budget_exceeded_warned_) { + budget_exceeded_warned_ = true; + XGRAMMAR_LOG(WARNING + ) << "The token budget (max_tokens) of a rule was exceeded: the rule could not end at " + "the position where its budget ran out, so the budget is relaxed until the rule " + "can end. This warning is reported once per matcher."; + } + FinishCharacterBudgetAccept(); + return true; + } + + bool BitmaskHasAnyToken(int32_t* bitmask_data_ptr) const { + auto bitset = DynamicBitset( + tokenizer_info_.GetVocabSize(), reinterpret_cast(bitmask_data_ptr) + ); + return bitset.Any(); + } + + /*! \brief Record that num_rows new input positions were created, each consuming one byte of + * bytes in order. Used for the byte-by-byte advance paths. */ + void AppendPerByteRows(const std::string& bytes) { + for (char c : bytes) { + accepted_bytes_.push_back(static_cast(c)); + row_byte_end_.push_back(static_cast(accepted_bytes_.size())); + } + } + + /*! \brief Record that one new input position was created, consuming all bytes of the token. + * Used for the atomic token advance path. */ + void AppendAtomicRow(const std::string& bytes) { + accepted_bytes_.insert(accepted_bytes_.end(), bytes.begin(), bytes.end()); + row_byte_end_.push_back(static_cast(accepted_bytes_.size())); + } + + CompiledGrammar compiled_grammar_; + TokenizerInfo tokenizer_info_; + std::vector stop_token_ids_; + bool terminate_without_stop_token_; + std::optional default_temperature_; + mutable bool has_warned_temperature_conflict_ = false; + std::deque token_length_history; + + /*! \brief Set by the last mask computation when an exhausted budget could be enforced: the + * next accept commits the same state transition used to compute that mask. */ + bool budget_enforce_pending_ = false; + /*! \brief Whether the pending decision includes a suffix/stop completion without its marker. */ + bool budget_force_close_pending_ = false; + /*! \brief Whether the one-time budget-exceeded warning has been reported. */ + bool budget_exceeded_warned_ = false; + /*! \brief Whether the one-time character-budget warning has been reported. */ + bool char_budget_exceeded_warned_ = false; + /*! \brief Whether the current definitive accept relaxed a character budget. */ + bool char_budget_relaxed_ = false; + /*! \brief Whether speculative character-budget relaxation should be recorded. */ + bool record_char_budget_relaxation_ = false; + /*! \brief Whether byte history is needed to recognize a budgeted suffix/stop body boundary. */ + bool has_budget_marker_rules_ = false; + + struct BudgetBodyMatchProgress { + int64_t begin_byte; + int64_t end_byte; + std::unordered_set states; + std::string temporary_input_snapshot; + }; + /*! \brief Incremental body-FSM progress keyed by (rule id, occurrence start row). */ + std::unordered_map budget_body_match_cache_; + + /*! \brief The bytes accepted so far. Maintained for captures and budgeted marker rules. */ + std::vector accepted_bytes_; + /*! \brief row_byte_end_[i] is the number of accepted bytes after input position i. Aligned + * with the parser's state history whenever accepted_bytes_ is maintained. */ + std::vector row_byte_end_{0}; + /*! \brief Bytes tentatively consumed by the current token or string advance. */ + std::string temporary_input_bytes_; + /*! \brief Parser row immediately before temporary_input_bytes_. */ + int32_t temporary_input_start_row_ = -1; + + // Temporary data for FillNextTokenBitmask. They are stored here to avoid repeated allocation. + DynamicBitset tmp_accepted_bitset_; + std::vector tmp_rejected_indices_; + std::vector tmp_rejected_indices_delta_; +}; + +class BatchGrammarMatcher::Impl { + public: + Impl(std::variant max_threads) { + if (std::holds_alternative(max_threads)) { + int32_t num_threads = std::get(max_threads); + XGRAMMAR_CHECK(num_threads >= 1) + << "The num_threads should be at least 1, but got " << num_threads; + if (num_threads > 1) { + if (num_threads > static_cast(std::thread::hardware_concurrency())) { + XGRAMMAR_LOG(WARNING) << "The num_threads " << num_threads << " is larger than the " + << "number of hardware threads. Using " + << static_cast(std::thread::hardware_concurrency()) + << " instead."; + } + max_threads_ = + std::min(num_threads, static_cast(std::thread::hardware_concurrency())); + } + } else { + std::string str = std::get(max_threads); + XGRAMMAR_CHECK(str == "auto"); + max_threads_ = std::thread::hardware_concurrency() / 2; + } + } + + void BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices, + bool debug_print + ); + + static std::vector BatchAcceptToken( + std::vector* matchers, const std::vector& token_ids, bool debug_print + ); + + static std::vector BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print + ); + + static void BatchRollback( + std::vector* matchers, const std::vector& num_tokens + ); + + private: + std::optional thread_pool_ = std::nullopt; + int32_t max_threads_ = 1; +}; + +bool GrammarMatcher::Impl::AcceptStopToken() { + if (terminate_without_stop_token_) { + return false; + } + if (!IsCompleted()) { + return false; + } + XGRAMMAR_DCHECK(!stop_token_is_accepted_); + token_length_history.push_back(0); + stop_token_is_accepted_ = true; + return true; +} + +bool GrammarMatcher::Impl::IsTerminated() const { + if (terminate_without_stop_token_) { + return IsCompleted(); + } + return IsStopTokenAccepted(); +} + +bool GrammarMatcher::Impl::IsStopTokenAccepted() const { return stop_token_is_accepted_; } + +bool GrammarMatcher::Impl::CanForceCompleteWithoutMarker( + const ParserState& state, bool use_char_budget +) { + bool expired = use_char_budget ? IsCharExpiredState(state) : IsExpiredState(state); + if (!expired || state.rule_id < 0) { + return false; + } + const auto& rule = grammar_->GetRule(state.rule_id); + if (use_char_budget ? rule.max_chars < 0 : rule.max_tokens < 0) { + return false; + } + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(state.rule_id); + if (suffix_stop_info == nullptr || suffix_stop_info->body_rule_id < 0) { + return false; + } + XGRAMMAR_DCHECK(ShouldTrackAcceptedBytes()); + int32_t start_row = + state.rule_start_pos == ParserState::kNoPrevInputPos ? 0 : state.rule_start_pos; + auto byte_position_for_row = [&](int32_t row) { + if (row >= 0 && row < static_cast(row_byte_end_.size())) { + return row_byte_end_[row]; + } + XGRAMMAR_DCHECK(temporary_input_start_row_ >= 0); + XGRAMMAR_DCHECK(row >= temporary_input_start_row_); + XGRAMMAR_DCHECK( + row <= temporary_input_start_row_ + static_cast(temporary_input_bytes_.size()) + ); + return static_cast(accepted_bytes_.size() + row - temporary_input_start_row_); + }; + int64_t begin = byte_position_for_row(start_row); + int64_t end = accepted_bytes_.size() + temporary_input_bytes_.size(); + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[suffix_stop_info->body_rule_id].has_value()); + const auto& body_fsm = grammar_->per_rule_fsms[suffix_stop_info->body_rule_id]->GetFsm(); + XGRAMMAR_DCHECK(body_fsm.IsLeaf()) << "A suffix/stop body helper must compile to a leaf FSM"; + + int64_t occurrence = + (static_cast(state.rule_id) << 32) | static_cast(state.rule_start_pos); + auto [it, inserted] = budget_body_match_cache_.try_emplace( + occurrence, BudgetBodyMatchProgress{begin, begin, {body_fsm.GetStart()}, {}} + ); + auto& progress = it->second; + if (inserted) { + body_fsm.GetFsm().GetEpsilonClosure(&progress.states); + } else if (progress.begin_byte != begin || progress.end_byte > end || + (progress.end_byte > static_cast(accepted_bytes_.size()) && + progress.temporary_input_snapshot != temporary_input_bytes_)) { + // This is only expected after restoring external matcher state. Rollback and reset clear the + // cache eagerly, but reinitialize defensively if a caller supplies an equivalent history. + progress = BudgetBodyMatchProgress{begin, begin, {body_fsm.GetStart()}, {}}; + body_fsm.GetFsm().GetEpsilonClosure(&progress.states); + } + + std::unordered_set next_states; + for (int64_t offset = progress.end_byte; offset < end && !progress.states.empty(); ++offset) { + uint8_t byte = + offset < static_cast(accepted_bytes_.size()) + ? accepted_bytes_[offset] + : static_cast(temporary_input_bytes_[offset - accepted_bytes_.size()]); + body_fsm.GetFsm().Advance( + progress.states, byte, &next_states, FSMEdge::EdgeType::kCharRange, true + ); + progress.states = next_states; + } + progress.end_byte = end; + progress.temporary_input_snapshot = + end > static_cast(accepted_bytes_.size()) ? temporary_input_bytes_ : std::string{}; + return std::any_of(progress.states.begin(), progress.states.end(), [&](int state_id) { + return body_fsm.IsEndState(state_id); + }); +} + +bool GrammarMatcher::Impl::ApplyBudgetEnforcement(bool debug_print) { + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()); + const auto latest_row = scanable_state_history_[scanable_state_history_.size() - 1]; + std::vector latest_states(latest_row.begin(), latest_row.end()); + std::vector force_completed_states; + std::unordered_set force_completed_occurrences; + + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_completed_lazy_occurrences_.clear(); + tmp_accept_stop_token_ = IsCompleted(); + + for (const auto& state : latest_states) { + if (!IsExpiredState(state)) { + EnqueueWithoutProcessing(state); + continue; + } + if (!CanForceCompleteWithoutMarker(state, false)) { + continue; + } + int64_t occurrence = + (static_cast(state.rule_id) << 32) | static_cast(state.rule_start_pos); + if (force_completed_occurrences.insert(occurrence).second) { + force_completed_states.push_back(state); + } + } + + for (const auto& state : force_completed_states) { + Complete(state, debug_print, /*marker_present=*/false); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + + bool any_expired = false; + bool any_alive = false; + for (const auto& state : tmp_states_to_be_added_) { + if (IsExpiredState(state)) { + any_expired = true; + } else { + any_alive = true; + } + } + if (any_expired && (any_alive || tmp_accept_stop_token_)) { + // Completing a budgeted suffix/stop rule can expose an already-expired parent alternative. + // Apply the same preference as the ordinary max_tokens path: once another derivation can + // continue or finish at this boundary, expired derivations may not consume another token. + tmp_states_to_be_added_.erase( + std::remove_if( + tmp_states_to_be_added_.begin(), + tmp_states_to_be_added_.end(), + [&](const ParserState& state) { return IsExpiredState(state); } + ), + tmp_states_to_be_added_.end() + ); + } + + bool viable = tmp_accept_stop_token_ || !tmp_states_to_be_added_.empty(); + if (!viable) { + return false; + } + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + is_completed_.back() = tmp_accept_stop_token_; + return true; +} + +bool GrammarMatcher::Impl::ApplyCharacterBudgetEnforcement(bool debug_print) { + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()); + const auto latest_row = scanable_state_history_[scanable_state_history_.size() - 1]; + std::vector latest_states(latest_row.begin(), latest_row.end()); + auto previous_completable_row = rule_id_to_completable_states_.Back(); + std::vector> previous_completable_states( + previous_completable_row.data, + previous_completable_row.data + previous_completable_row.data_len + ); + std::vector previous_capture_events; + if (capture_tracking_) { + previous_capture_events = CopyLastCaptureRow(); + } + std::vector force_completed_states; + std::unordered_set force_completed_occurrences; + + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_completed_lazy_occurrences_.clear(); + tmp_accept_stop_token_ = IsCompleted(); + + for (const auto& state : latest_states) { + if (!IsCharExpiredState(state)) { + EnqueueWithoutProcessing(state); + continue; + } + if (!CanForceCompleteWithoutMarker(state, true)) { + continue; + } + int64_t occurrence = + (static_cast(state.rule_id) << 32) | static_cast(state.rule_start_pos); + if (force_completed_occurrences.insert(occurrence).second) { + force_completed_states.push_back(state); + } + } + + for (const auto& state : force_completed_states) { + Complete(state, debug_print, /*marker_present=*/false); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + + bool any_expired = false; + bool any_alive = false; + for (const auto& state : tmp_states_to_be_added_) { + if (IsCharExpiredState(state)) { + any_expired = true; + } else { + any_alive = true; + } + } + if (any_expired && (any_alive || tmp_accept_stop_token_)) { + tmp_states_to_be_added_.erase( + std::remove_if( + tmp_states_to_be_added_.begin(), + tmp_states_to_be_added_.end(), + [&](const ParserState& state) { return IsCharExpiredState(state); } + ), + tmp_states_to_be_added_.end() + ); + } + + bool viable = tmp_accept_stop_token_ || !tmp_states_to_be_added_.empty(); + if (!viable) { + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(previous_completable_states); + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(previous_capture_events); + } + return false; + } + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + is_completed_.back() = tmp_accept_stop_token_; + return true; +} + +bool GrammarMatcher::Impl::ApplyCharacterBudgetEnforcementToFixedPoint(bool debug_print) { + bool applied_any = false; + for (int32_t iteration = 0; iteration <= grammar_->NumRules(); ++iteration) { + bool has_expired_state = false; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + if (IsCharExpiredState(state)) { + has_expired_state = true; + break; + } + } + if (!has_expired_state) { + return applied_any; + } + if (!ApplyCharacterBudgetEnforcement(debug_print)) { + return applied_any; + } + applied_any = true; + } + return applied_any; +} + +bool GrammarMatcher::Impl::AdvanceWithCharacterBudget(uint8_t byte, bool debug_print) { + if (!has_char_budget_rules_ || !StartsUTF8Codepoint(byte)) { + return Advance(byte, debug_print); + } + + bool has_expired_state = false; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + if (IsCharExpiredState(state)) { + has_expired_state = true; + break; + } + } + if (!has_expired_state) { + return Advance(byte, debug_print); + } + + std::vector previous_states = GetLatestScanableStates(); + bool previous_completed = IsCompleted(); + auto previous_completable_row = rule_id_to_completable_states_.Back(); + std::vector> previous_completable_states( + previous_completable_row.data, + previous_completable_row.data + previous_completable_row.data_len + ); + std::vector previous_capture_events; + if (capture_tracking_) { + previous_capture_events = CopyLastCaptureRow(); + } + + bool enforced = ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + if (!enforced && record_char_budget_relaxation_) { + char_budget_relaxed_ = true; + } + bool accepted = Advance(byte, debug_print); + if (!accepted && enforced) { + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(previous_states); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(previous_completable_states); + is_completed_.back() = previous_completed; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(previous_capture_events); + } + } + return accepted; +} + +bool GrammarMatcher::Impl::AdvanceAtomicTokenWithCharacterBudget( + int32_t token_id, int32_t token_char_count, bool debug_print +) { + bool has_expired_state = false; + bool crosses_budget = false; + if (has_char_budget_rules_ && token_char_count > 0) { + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + has_expired_state = has_expired_state || IsCharExpiredState(state); + crosses_budget = + crosses_budget || (state.char_budget_deadline >= 0 && + GetCurrentCharIndex() + token_char_count > state.char_budget_deadline); + } + } + + std::vector previous_states; + std::vector> previous_completable_states; + bool previous_completed = false; + std::vector previous_capture_events; + bool enforced = false; + if (has_expired_state) { + previous_states = GetLatestScanableStates(); + auto previous_completable_row = rule_id_to_completable_states_.Back(); + previous_completable_states.assign( + previous_completable_row.data, + previous_completable_row.data + previous_completable_row.data_len + ); + previous_completed = IsCompleted(); + if (capture_tracking_) { + previous_capture_events = CopyLastCaptureRow(); + } + enforced = ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + if (!enforced && record_char_budget_relaxation_) { + char_budget_relaxed_ = true; + } + } + + bool accepted = AdvanceAtomicToken(token_id, debug_print, token_char_count); + if (!accepted && enforced) { + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(previous_states); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(previous_completable_states); + is_completed_.back() = previous_completed; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(previous_capture_events); + } + } + if (accepted && crosses_budget && record_char_budget_relaxation_) { + char_budget_relaxed_ = true; + } + return accepted; +} + +std::optional GrammarMatcher::Impl::GetTemperature() const { + std::optional syntax_temperature = std::nullopt; + bool has_temperature_conflict = false; + for (const auto& state : GetLatestScanableStates()) { + if (state.active_temperature_rule_id == -1) { + continue; + } + const auto& temperature = grammar_->GetRule(state.active_temperature_rule_id).temperature; + XGRAMMAR_DCHECK(temperature.has_value()); + if (syntax_temperature.has_value() && temperature.value() != syntax_temperature.value()) { + has_temperature_conflict = true; + } + if (!syntax_temperature.has_value() || temperature.value() > syntax_temperature.value()) { + syntax_temperature = temperature; + } + } + if (has_temperature_conflict && !has_warned_temperature_conflict_) { + XGRAMMAR_LOG(WARNING) << "Multiple active grammar paths specify different temperatures. " + "Using the maximum temperature " + << syntax_temperature.value() << "."; + has_warned_temperature_conflict_ = true; + } + return syntax_temperature.has_value() ? syntax_temperature : default_temperature_; +} + +// TODO(yixin): Polish verbose logging +bool GrammarMatcher::Impl::AcceptToken(int32_t token_id, bool debug_print) { + if (IsStopTokenAccepted()) { + XGRAMMAR_LOG(WARNING) << "The matcher has terminated after accepting the stop token, but is " + << "trying to accept new token with id " << token_id << "."; + return false; + } + + if (token_id < 0 || token_id >= tokenizer_info_.GetVocabSize()) { + XGRAMMAR_LOG(WARNING) << "The token id " << token_id << " is out of range [0, " + << tokenizer_info_.GetVocabSize() << "). Rejecting the token."; + return false; + } + + bool is_stop_token = + std::find(stop_token_ids_.begin(), stop_token_ids_.end(), token_id) != stop_token_ids_.end(); + const auto& special_token_ids = tokenizer_info_.GetSpecialTokenIds(); + if (!is_stop_token && std::find(special_token_ids.begin(), special_token_ids.end(), token_id) != + special_token_ids.end()) { + // Padding ids, i.e. ids in [decoded_vocab.size(), vocab_size), are registered as special ids + // too, but they have no entry in decoded_vocab, so only decode ids backed by a real token. + const auto& decoded_vocab = tokenizer_info_.GetDecodedVocab(); + XGRAMMAR_LOG(WARNING) << "GrammarMatcher cannot accept special token id " << token_id << ": " + << (token_id < static_cast(decoded_vocab.size()) + ? decoded_vocab[token_id] + : "") + << ". Rejecting the token."; + return false; + } + + current_token_index_ = static_cast(token_length_history.size()); + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = false; + bool enforce_budget = budget_enforce_pending_; + bool force_budget_close = budget_force_close_pending_; + if (force_budget_close) { + // Validate on a copy first: a rejected token must not commit the budget-close transition. + Impl trial(*this); + trial.capture_recording_ = false; + bool applied = trial.ApplyBudgetEnforcement(); + XGRAMMAR_DCHECK(applied); + if (!applied) { + return false; + } + trial.budget_enforce_pending_ = false; + trial.budget_force_close_pending_ = false; + if (!trial.AcceptToken(token_id, false)) { + return false; + } + } + + // Capture events are only recorded on the definitive accept path. + CaptureRecordingGuard capture_guard(this); + if (force_budget_close) { + bool applied = ApplyBudgetEnforcement(debug_print); + XGRAMMAR_DCHECK(applied); + if (!applied) { + return false; + } + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + } + SkipExpiredGuard skip_expired_guard(this, enforce_budget && !force_budget_close); + + // The token extends a derivation past its budget iff expired states are allowed to scan + // (no enforcement) while some exist. + bool consumed_past_deadline = false; + if (has_budget_rules_ && !enforce_budget) { + const auto& row = scanable_state_history_[scanable_state_history_.size() - 1]; + for (const auto& state : row) { + if (IsExpiredState(state)) { + consumed_past_deadline = true; + break; + } + } + } + + if (debug_print) { + std::string states_str; + for (const auto& state : GetLatestScanableStates()) { + states_str += " " + state.ToString() + "\n"; + } + XGRAMMAR_LOG(INFO) << "Accepting token id " << token_id << ", string: \"" + << EscapeString(tokenizer_info_.GetDecodedVocab()[token_id]) + << "\", current state:\n" + << states_str; + } + // Handle the stop token + if (is_stop_token) { + bool accepted = AcceptStopToken(); + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The token is an end token. Is accepted: " << accepted; + } + return accepted ? FinishAccept(consumed_past_deadline) : false; + } + + const auto& token = tokenizer_info_.GetDecodedVocab()[token_id]; + int32_t token_char_count = 0; + if (has_char_budget_rules_) { + for (uint8_t byte : token) { + token_char_count += StartsUTF8Codepoint(byte); + } + } + + const int32_t size_before_token = rule_id_to_completable_states_.size(); + std::vector states_before_token; + std::vector> completable_before_token; + std::vector capture_row_before_token; + bool completed_before_token = false; + if (has_char_budget_rules_) { + states_before_token = GetLatestScanableStates(); + auto completable_row_before_token = rule_id_to_completable_states_.Back(); + completable_before_token.assign( + completable_row_before_token.data, + completable_row_before_token.data + completable_row_before_token.data_len + ); + capture_row_before_token = CopyLastCaptureRow(); + completed_before_token = is_completed_.back(); + } + auto restore_row_before_token = [&]() { + if (!has_char_budget_rules_) { + return; + } + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(states_before_token); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(completable_before_token); + is_completed_.back() = completed_before_token; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(capture_row_before_token); + } + }; + + // Phase 1: Try atomic token path (from current state, before byte path) + std::vector atomic_states; + std::vector> atomic_completable; + std::vector atomic_capture_row; + bool atomic_completed = false; + bool atomic_success = + has_char_budget_rules_ + ? AdvanceAtomicTokenWithCharacterBudget(token_id, token_char_count, debug_print) + : AdvanceAtomicToken(token_id, debug_print); + if (atomic_success) { + atomic_states = GetLatestScanableStates(); + auto row = rule_id_to_completable_states_.Back(); + atomic_completable.assign(row.data, row.data + row.data_len); + atomic_capture_row = CopyLastCaptureRow(); + atomic_completed = is_completed_.back(); + PopLastStates(1); + } + restore_row_before_token(); + + // Phase 2: Try byte-by-byte path (from the same original state) + record_char_budget_relaxation_ = true; + bool track_temporary_input = has_char_budget_rules_ && has_budget_marker_rules_; + if (track_temporary_input) { + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + } + int pos = 0; + bool byte_path_success = true; + for (auto char_value : token) { + bool accepted = has_char_budget_rules_ + ? AdvanceWithCharacterBudget(static_cast(char_value), debug_print) + : Advance(static_cast(char_value), debug_print); + if (!accepted) { + byte_path_success = false; + break; + } + if (track_temporary_input) { + temporary_input_bytes_.push_back(char_value); + } + ++pos; + } + if (track_temporary_input) { + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + } + + // Phase 3: Combine results (no priority — merge with deduplication) + if (!byte_path_success && !atomic_success) { + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Token #" << token_id << "<" << EscapeString(token) + << "> rejected at position " << pos; + } + PopLastStates(pos); + restore_row_before_token(); + record_char_budget_relaxation_ = false; + char_budget_relaxed_ = false; + return false; + } + + if (atomic_success && !byte_path_success) { + PopLastStates(pos); + restore_row_before_token(); + char_budget_relaxed_ = false; + bool accepted = + has_char_budget_rules_ + ? AdvanceAtomicTokenWithCharacterBudget(token_id, token_char_count, debug_print) + : AdvanceAtomicToken(token_id, debug_print); + XGRAMMAR_DCHECK(accepted); + token_length_history.push_back(1); + if (ShouldTrackAcceptedBytes()) { + AppendAtomicRow(token); + } + } else if (byte_path_success && !atomic_success) { + token_length_history.push_back(token.size()); + if (ShouldTrackAcceptedBytes()) { + AppendPerByteRows(token); + } + } else { + // Both paths succeeded — merge atomic token states into byte path + if (token.empty()) { + // Zero-length token: byte path created 0 timepoints, just push atomic states + scanable_state_history_.PushBack(atomic_states); + rule_id_to_completable_states_.PushBack(atomic_completable); + is_completed_.push_back(atomic_completed); + PushCaptureRow(atomic_capture_row); + PushCharCountRow(GetCurrentCharIndex(), HasEnteredCharBudget()); + token_length_history.push_back(1); + if (ShouldTrackAcceptedBytes()) { + AppendAtomicRow(token); + } + } else { + auto byte_states = GetLatestScanableStates(); + std::vector merged = byte_states; + StateEqualForParsing state_eq; + for (const auto& s : atomic_states) { + if (std::find_if(merged.begin(), merged.end(), [&](const auto& m) { + return state_eq(m, s); + }) == merged.end()) { + merged.push_back(s); + } + } + + auto byte_row = rule_id_to_completable_states_.Back(); + std::vector> merged_completable( + byte_row.data, byte_row.data + byte_row.data_len + ); + std::vector merged_capture_row = CopyLastCaptureRow(); + bool byte_completed = is_completed_.back(); + int32_t final_char_index = GetCurrentCharIndex(); + bool final_char_budget_entered = HasEnteredCharBudget(); + PopLastStates(1); + + for (const auto& cs : atomic_completable) { + if (std::find_if(merged_completable.begin(), merged_completable.end(), [&](const auto& m) { + return m.first == cs.first && state_eq(m.second, cs.second); + }) == merged_completable.end()) { + merged_completable.push_back(cs); + } + } + + for (auto event : atomic_capture_row) { + // In the atomic path, the after-token position was size_before_token; in the byte path + // it is the row being re-pushed. Remap events that started at the after-token position + // (empty spans) so that they stay empty in the byte-path numbering. + if (event.start_pos == size_before_token) { + event.start_pos = rule_id_to_completable_states_.size(); + } + if (event.occurrence_start_pos == size_before_token) { + event.occurrence_start_pos = rule_id_to_completable_states_.size(); + } + for (auto& target : event.stop_capture_targets) { + if (target.start_pos == size_before_token) { + target.start_pos = rule_id_to_completable_states_.size(); + } + } + auto existing = std::find_if( + merged_capture_row.begin(), + merged_capture_row.end(), + [&](const CaptureEvent& e) { + return e.rule_id == event.rule_id && e.start_pos == event.start_pos && + e.occurrence_start_pos == event.occurrence_start_pos; + } + ); + if (existing == merged_capture_row.end()) { + merged_capture_row.push_back(event); + } else { + existing->hidden_suffix_bytes = + std::max(existing->hidden_suffix_bytes, event.hidden_suffix_bytes); + existing->hidden_stop_bytes = + std::max(existing->hidden_stop_bytes, event.hidden_stop_bytes); + for (const auto& target : event.stop_capture_targets) { + if (std::find( + existing->stop_capture_targets.begin(), + existing->stop_capture_targets.end(), + target + ) == existing->stop_capture_targets.end()) { + existing->stop_capture_targets.push_back(target); + } + } + } + } + + scanable_state_history_.PushBack(merged); + rule_id_to_completable_states_.PushBack(merged_completable); + is_completed_.push_back(byte_completed || atomic_completed); + PushCaptureRow(merged_capture_row); + PushCharCountRow(final_char_index, final_char_budget_entered); + token_length_history.push_back(token.size()); + if (ShouldTrackAcceptedBytes()) { + AppendPerByteRows(token); + } + } + } + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Token #" << token_id << "<" << EscapeString(token) << "> accepted."; + } + if (has_char_budget_rules_) { + // Close exhausted occurrences at the token boundary, including after zero-byte atomic tokens. + ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + } + record_char_budget_relaxation_ = false; + return FinishAccept(consumed_past_deadline); +} + +bool GrammarMatcher::Impl::AcceptString(const std::string& input_str, bool debug_print) { + if (IsStopTokenAccepted()) { + XGRAMMAR_LOG(WARNING) << "The matcher has terminated after accepting the stop token, but is " + << "trying to accept new string \"" << EscapeString(input_str) << "\"."; + return false; + } + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Trying to accept string \"" << EscapeString(input_str) + << "\". Current state:\n" + << PrintStates(); + } + + current_token_index_ = static_cast(token_length_history.size()); + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = true; + + // Capture events are only recorded on the definitive accept path. + CaptureRecordingGuard capture_guard(this); + + std::vector states_before_input; + std::vector> completable_before_input; + std::vector capture_row_before_input; + bool completed_before_input = false; + if (has_char_budget_rules_) { + states_before_input = GetLatestScanableStates(); + auto completable_row_before_input = rule_id_to_completable_states_.Back(); + completable_before_input.assign( + completable_row_before_input.data, + completable_row_before_input.data + completable_row_before_input.data_len + ); + capture_row_before_input = CopyLastCaptureRow(); + completed_before_input = is_completed_.back(); + } + bool track_temporary_input = has_char_budget_rules_ && has_budget_marker_rules_; + if (track_temporary_input) { + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + } + int accepted_cnt = 0; + for (auto char_value : input_str) { + bool accepted = has_char_budget_rules_ + ? AdvanceWithCharacterBudget(static_cast(char_value), debug_print) + : Advance(static_cast(char_value), debug_print); + if (!accepted) { + if (debug_print) { + XGRAMMAR_LOG(INFO) << "String \"" << EscapeString(input_str) << "\" is rejected at " + << "position " << accepted_cnt << ", char " << EscapeString(char_value); + } + PopLastStates(accepted_cnt); + if (has_char_budget_rules_) { + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(states_before_input); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(completable_before_input); + is_completed_.back() = completed_before_input; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(capture_row_before_input); + } + } + if (track_temporary_input) { + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + } + record_char_budget_relaxation_ = false; + char_budget_relaxed_ = false; + return false; + } + if (track_temporary_input) { + temporary_input_bytes_.push_back(char_value); + } + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Char " << EscapeString(char_value) << " is accepted. Current state:\n" + << PrintStates(); + } + ++accepted_cnt; + } + token_length_history.push_back(input_str.size()); + if (ShouldTrackAcceptedBytes()) { + AppendPerByteRows(input_str); + } + if (has_char_budget_rules_ && !input_str.empty()) { + // Leave the parser at the enforced boundary even when no later input is supplied. + ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + } + if (track_temporary_input) { + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + } + record_char_budget_relaxation_ = false; + FinishCharacterBudgetAccept(); + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "String \"" << EscapeString(input_str) << "\" is accepted."; + } + return true; +} + +std::string GrammarMatcher::Impl::PrintBitmask( + int32_t* bitmask_data_ptr, const TokenizerInfo& tokenizer_info +) { + constexpr int kMaxPrintTokens = 100; + std::vector accepted_ids; + std::vector rejected_ids; + auto bitset = + DynamicBitset(tokenizer_info.GetVocabSize(), reinterpret_cast(bitmask_data_ptr)); + for (int i = 0; i < tokenizer_info.GetVocabSize(); ++i) { + if (bitset[i]) { + accepted_ids.push_back(i); + } else { + rejected_ids.push_back(i); + } + } + std::stringstream ss; + ss << "TokenBitmask(num_tokens=" << tokenizer_info.GetVocabSize() + << ", accepted_num=" << accepted_ids.size() << ", rejected_num=" << rejected_ids.size() + << ",\naccepted_ids=" << PrintTokenByIds(accepted_ids, tokenizer_info, kMaxPrintTokens) + << ",\nrejected_ids=" << PrintTokenByIds(rejected_ids, tokenizer_info, kMaxPrintTokens) << ")"; + return ss.str(); +} + +bool GrammarMatcher::Impl::IsTokenBitmaskAllTrue(int32_t* bitmask_data_ptr) { + DynamicBitset next_token_bitset( + tokenizer_info_.GetVocabSize(), reinterpret_cast(bitmask_data_ptr) + ); + return next_token_bitset.All(); +} + +bool GrammarMatcher::Impl::FillNextTokenBitmask( + DLTensor* next_token_bitmask, int index, bool debug_print +) { + XGRAMMAR_CHECK(!IsStopTokenAccepted()) + << "GrammarMatcher has terminated after accepting the stop token, but is trying to " + "find the next token mask"; + int32_t* bitmask_data_ptr = + CheckAndGetBitmaskPtr(*next_token_bitmask, tokenizer_info_.GetVocabSize(), index); + current_token_index_ = static_cast(token_length_history.size()); + if (has_budget_rules_) { + const auto& row = scanable_state_history_[scanable_state_history_.size() - 1]; + bool any_expired = false; + bool any_alive = false; + bool can_force_close_without_marker = false; + for (const auto& state : row) { + if (IsExpiredState(state)) { + any_expired = true; + can_force_close_without_marker = + can_force_close_without_marker || + (has_budget_marker_rules_ && CanForceCompleteWithoutMarker(state, false)); + } else { + any_alive = true; + } + } + if (can_force_close_without_marker) { + // Compute the enforcing state on a copy. In addition to dropping expired derivations, + // this lets a budgeted suffix/stop rule end through its body when the marker has not + // appeared yet. The real state is committed only after an allowed token is accepted. + Impl trial(*this); + trial.capture_recording_ = false; + bool applied = trial.ApplyBudgetEnforcement(debug_print); + if (applied) { + trial.FillBitmaskForStates(bitmask_data_ptr, index, /*skip_expired=*/false, debug_print); + } + if (applied && (trial.IsCompleted() || trial.BitmaskHasAnyToken(bitmask_data_ptr))) { + budget_enforce_pending_ = true; + budget_force_close_pending_ = true; + return !IsTokenBitmaskAllTrue(bitmask_data_ptr); + } + } + if (any_expired && (any_alive || IsCompleted())) { + // No marker-less completion is needed: retain the original max_tokens fast path. + FillBitmaskForStates(bitmask_data_ptr, index, /*skip_expired=*/true, debug_print); + if (BitmaskHasAnyToken(bitmask_data_ptr) || IsCompleted()) { + budget_enforce_pending_ = true; + budget_force_close_pending_ = false; + return !IsTokenBitmaskAllTrue(bitmask_data_ptr); + } + } + // No enforcement at this position: either no budget has expired, or the expired + // derivations cannot end here — relax the budget for one step and report it on the next + // accept. + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + } + FillBitmaskForStates(bitmask_data_ptr, index, /*skip_expired=*/false, debug_print); + return !IsTokenBitmaskAllTrue(bitmask_data_ptr); +} + +void GrammarMatcher::Impl::FillBitmaskForCharBudgetBoundary( + const AdaptiveTokenMask& adaptive_token_mask, int32_t remaining_chars +) { + const auto& token_char_counts = tokenizer_info_.ImplPtr()->GetTokenCharCounts(); + const auto& vocab = tokenizer_info_.GetSortedDecodedVocab(); + + std::vector tokens_to_check = adaptive_token_mask.uncertain_indices; + switch (adaptive_token_mask.store_type) { + case StoreType::kAccepted: + for (int32_t index : adaptive_token_mask.accepted_indices) { + if (token_char_counts[index] <= remaining_chars) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + } else { + tokens_to_check.push_back(index); + } + } + break; + case StoreType::kAcceptedBitset: + for (int32_t index = 0; index < static_cast(vocab.size()); ++index) { + int32_t token_id = vocab[index].first; + if (!adaptive_token_mask.accepted_bitset[token_id]) { + continue; + } + if (token_char_counts[index] <= remaining_chars) { + tmp_accepted_bitset_.Set(token_id, true); + } else { + tokens_to_check.push_back(index); + } + } + break; + case StoreType::kRejected: { + std::vector blocked = adaptive_token_mask.rejected_indices; + blocked.insert( + blocked.end(), + adaptive_token_mask.uncertain_indices.begin(), + adaptive_token_mask.uncertain_indices.end() + ); + std::sort(blocked.begin(), blocked.end()); + for (int32_t index = 0; index < static_cast(vocab.size()); ++index) { + if (token_char_counts[index] <= remaining_chars) { + if (!std::binary_search(blocked.begin(), blocked.end(), index)) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + } + } else if (!std::binary_search( + adaptive_token_mask.rejected_indices.begin(), + adaptive_token_mask.rejected_indices.end(), + index + )) { + tokens_to_check.push_back(index); + } + } + break; + } + } + + std::sort(tokens_to_check.begin(), tokens_to_check.end()); + tokens_to_check.erase( + std::unique(tokens_to_check.begin(), tokens_to_check.end()), tokens_to_check.end() + ); + + int32_t saved_temporary_input_start_row = temporary_input_start_row_; + std::string saved_temporary_input_bytes = std::move(temporary_input_bytes_); + std::vector all_latest_states = GetLatestScanableStates(); + bool all_latest_completed = IsCompleted(); + + PushStatesToCheck(all_latest_states, all_latest_completed); + int32_t initial_state_history_size = scanable_state_history_.size(); + int32_t previous_index = -1; + int32_t previous_matched_size = 0; + temporary_input_start_row_ = initial_state_history_size - 1; + temporary_input_bytes_.clear(); + std::vector byte_rejected_tokens; + + for (int32_t index : tokens_to_check) { + const std::string& token = vocab[index].second; + int32_t common_prefix_length = 0; + if (previous_index >= 0) { + common_prefix_length = std::mismatch( + vocab[previous_index].second.begin(), + vocab[previous_index].second.end(), + token.begin(), + token.end() + ) + .first - + vocab[previous_index].second.begin(); + common_prefix_length = std::min(common_prefix_length, previous_matched_size); + } + PopLastStates( + scanable_state_history_.size() - initial_state_history_size - common_prefix_length + ); + temporary_input_bytes_.resize(common_prefix_length); + + bool accepted = !token.empty(); + int32_t matched_size = common_prefix_length; + for (; matched_size < static_cast(token.size()); ++matched_size) { + if (!AdvanceWithCharacterBudget(static_cast(token[matched_size]))) { + accepted = false; + break; + } + temporary_input_bytes_.push_back(token[matched_size]); + } + if (accepted) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + } else { + byte_rejected_tokens.push_back(index); + } + previous_index = index; + previous_matched_size = matched_size; + } + PopLastStates(scanable_state_history_.size() - initial_state_history_size + 1); + + for (int32_t index : byte_rejected_tokens) { + PushStatesToCheck(all_latest_states, all_latest_completed); + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + if (AdvanceAtomicTokenWithCharacterBudget(vocab[index].first, token_char_counts[index])) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + PopLastStates(1); + } + PopLastStates(1); + } + + temporary_input_start_row_ = saved_temporary_input_start_row; + temporary_input_bytes_ = std::move(saved_temporary_input_bytes); +} + +void GrammarMatcher::Impl::FillBitmaskForStates( + int32_t* bitmask_data_ptr, int index, bool skip_expired, bool debug_print +) { + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + const auto& subtree_range = tokenizer_info_.GetTrieSubtreeNodesRange(); + const auto& adaptive_token_mask_cache = compiled_grammar_->adaptive_token_mask_cache; + // We need to have a copy, because scanable_state_history_ will be modified during the + // FillNextTokenBitmask process, which can lead to undefined behavior. + std::vector latest_states; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + if (skip_expired && IsExpiredState(state)) { + continue; + } + latest_states.push_back(state); + } + + // We check all the latest states of the earley parser, and check all the masks of the leaf + // states. The final accepted token set is the union of the accepted token sets of all leaf + // states. The final rejected token set is the intersection of the rejected token sets of all leaf + // states. + + // Note these indices store the indices in sorted_decoded_vocab, instead of the token ids. + tmp_accepted_bitset_.Reset(); + // {-1} means the universal set, i.e. all tokens initially + tmp_rejected_indices_.assign({-1}); + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "FillNextTokenBitmask: index=" << index + << ", num of states=" << latest_states.size(); + } + + std::vector> + latest_states_with_masks; + + for (const auto& state : latest_states) { + auto adaptive_token_mask_it = adaptive_token_mask_cache.find(state); + XGRAMMAR_CHECK(adaptive_token_mask_it != adaptive_token_mask_cache.end()) << state; + const auto& adaptive_token_mask = adaptive_token_mask_it->second; + if (state.char_budget_deadline >= 0) { + int32_t remaining_chars = state.char_budget_deadline - GetCurrentCharIndex(); + if (remaining_chars <= tokenizer_info_.ImplPtr()->GetMaxTokenChars()) { + FillBitmaskForCharBudgetBoundary(adaptive_token_mask, std::max(remaining_chars, 0)); + continue; + } + } + latest_states_with_masks.push_back(std::make_pair(state, adaptive_token_mask_it)); + if (adaptive_token_mask.store_type == StoreType::kAcceptedBitset) { + tmp_accepted_bitset_ |= adaptive_token_mask.accepted_bitset; + } else if (adaptive_token_mask.store_type == StoreType::kAccepted) { + for (auto idx : adaptive_token_mask.accepted_indices) { + tmp_accepted_bitset_.Set(sorted_decoded_vocab[idx].first, true); + } + } + } + + for (const auto& [state, adaptive_token_mask_it] : latest_states_with_masks) { + const auto& adaptive_token_mask = adaptive_token_mask_it->second; + + // For each ParserState, we will check every uncertain token and put them into the accepted or + // rejected list. + + // Step 2. Update the accepted tokens in accepted_indices_delta, or the rejected tokens in + // rejected_indices_delta. + + // If the accepted tokens are saved, it means it is likely to be smaller than the rejected + // tokens, so we will just find the accepted tokens, and vice versa. + + tmp_rejected_indices_delta_.clear(); + + // Examine only the current one ParserState + std::optional atomic_trial_base; + if (has_char_budget_rules_ && !adaptive_token_mask.uncertain_indices.empty()) { + atomic_trial_base.emplace(*this); + atomic_trial_base->capture_recording_ = false; + } + PushOneStateToCheck(state); + bool track_temporary_input = has_char_budget_rules_ && has_budget_marker_rules_; + int32_t saved_temporary_input_start_row = -1; + std::string saved_temporary_input_bytes; + if (track_temporary_input) { + saved_temporary_input_start_row = temporary_input_start_row_; + saved_temporary_input_bytes = std::move(temporary_input_bytes_); + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + } + + const std::string* prev_token = nullptr; + int prev_matched_size = 0; + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The ParserState is " << state << ", the mask is " + << adaptive_token_mask.Print(tokenizer_info_); + } + int last_rejected_uncertain_range = 0; + for (const auto& cur_token_idx : adaptive_token_mask.uncertain_indices) { + // Check if the current token is already accepted. If it is, we can skip it. + if (tmp_accepted_bitset_[sorted_decoded_vocab[cur_token_idx].first]) { + continue; + } + + // Check if the current token is in the rejected range. i.e. check if the current token + // is on the subtree of the rejected token. + if (cur_token_idx < last_rejected_uncertain_range) { + if (adaptive_token_mask.store_type == StoreType::kRejected) { + tmp_rejected_indices_delta_.push_back(cur_token_idx); + } + continue; + } + + const auto& cur_token = sorted_decoded_vocab[cur_token_idx].second; + bool accepted = !cur_token.empty() || !has_char_budget_rules_; + + // Step 2.1. Find the longest common prefix with the accepted part of the previous token. + // We can reuse the previous matched size to avoid unnecessary matching. + if (prev_token) { + int lcp_len = std::mismatch( + cur_token.begin(), cur_token.end(), prev_token->begin(), prev_token->end() + ) + .first - + cur_token.begin(); + if (lcp_len > prev_matched_size) { + last_rejected_uncertain_range = subtree_range[cur_token_idx]; + accepted = false; + } else if (lcp_len < prev_matched_size) { + PopLastStates(prev_matched_size - lcp_len); + if (track_temporary_input) { + temporary_input_bytes_.resize(lcp_len); + } + } + prev_matched_size = std::min(prev_matched_size, lcp_len); + } + + // Step 2.2. Find if the current token is accepted or rejected. + if (accepted) { + for (int j = prev_matched_size; j < static_cast(cur_token.size()); ++j) { + bool byte_accepted = has_char_budget_rules_ + ? AdvanceWithCharacterBudget(static_cast(cur_token[j])) + : Advance(static_cast(cur_token[j])); + if (!byte_accepted) { + last_rejected_uncertain_range = subtree_range[cur_token_idx]; + accepted = false; + break; + } + if (track_temporary_input) { + temporary_input_bytes_.push_back(cur_token[j]); + } + prev_matched_size = j + 1; + } + } + + if (!accepted && has_char_budget_rules_) { + int32_t token_char_count = 0; + for (uint8_t byte : cur_token) { + token_char_count += StartsUTF8Codepoint(byte); + } + XGRAMMAR_DCHECK(atomic_trial_base.has_value()); + Impl atomic_trial(atomic_trial_base.value()); + atomic_trial.PushOneStateToCheck(state); + accepted = atomic_trial.AdvanceAtomicTokenWithCharacterBudget( + sorted_decoded_vocab[cur_token_idx].first, token_char_count + ); + if (accepted) { + last_rejected_uncertain_range = cur_token_idx + 1; + } + } + + // Step 2.3. Push the result to the delta list. + if (adaptive_token_mask.store_type == StoreType::kAcceptedBitset || + adaptive_token_mask.store_type == StoreType::kAccepted) { + if (accepted) { + tmp_accepted_bitset_.Set(sorted_decoded_vocab[cur_token_idx].first, true); + } + } else { + if (!accepted) { + tmp_rejected_indices_delta_.push_back(cur_token_idx); + } + } + + prev_token = &cur_token; + } + + PopLastStates(prev_matched_size + 1); + if (track_temporary_input) { + temporary_input_start_row_ = saved_temporary_input_start_row; + temporary_input_bytes_ = std::move(saved_temporary_input_bytes); + } + // Step 3. Update the accepted_indices or rejected_indices + if (adaptive_token_mask.store_type == StoreType::kRejected) { + // rejected_indices = Intersect( + // rejected_indices, + // adaptive_token_mask.rejected_indices + rejected_indices_delta) + IntsetUnion(&tmp_rejected_indices_delta_, adaptive_token_mask.rejected_indices); + IntsetIntersection(&tmp_rejected_indices_, tmp_rejected_indices_delta_); + } + } + + // Finally update the rejected_ids bitset + bool can_reach_end = IsCompleted(); + SetTokenBitmask( + bitmask_data_ptr, tmp_accepted_bitset_, tmp_rejected_indices_, can_reach_end, false + ); + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Filled bitmask: " << PrintBitmask(bitmask_data_ptr, tokenizer_info_); + } +} + +std::string GrammarMatcher::Impl::FindJumpForwardString() { + XGRAMMAR_CHECK(!IsStopTokenAccepted()) + << "GrammarMatcher has terminated after accepting the stop token, but is trying to " + "get the jump forward string"; + + current_token_index_ = static_cast(token_length_history.size()); + if (budget_force_close_pending_) { + Impl trial(*this); + trial.capture_recording_ = false; + if (trial.ApplyBudgetEnforcement()) { + trial.budget_enforce_pending_ = false; + trial.budget_force_close_pending_ = false; + return trial.FindJumpForwardString(); + } + } + + std::string result; + int num_accepted_chars = 0; + bool can_find_next_char = true; + + while (can_find_next_char) { + const auto& states = scanable_state_history_[scanable_state_history_.size() - 1]; + + // The state comes to the end of the grammar + if (IsCompleted()) { + can_find_next_char = false; + break; + } + + // 1. Check that for every leaf ParserState, the next possible char is unique and the same + // -1 means not found yet; 0~255 means the next char + int next_char = -1; + for (const auto& state : states) { + if (budget_enforce_pending_ && IsExpiredState(state)) { + continue; + } + XGRAMMAR_DCHECK(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[state.rule_id].value(); + const auto& current_edges = fsm.GetFsm().GetFsm().GetEdges(state.element_id); + for (const auto& edge : current_edges) { + if (!edge.IsCharRange()) { + continue; + } + if (edge.min != edge.max) { + can_find_next_char = false; + break; + } + if (next_char == -1) { + next_char = edge.min; + } else if (next_char != edge.min) { + can_find_next_char = false; + break; + } + } + } + + if (next_char == -1) { + can_find_next_char = false; + } + + // 2. If found, accept the char and iterate to the next position + if (can_find_next_char) { + if (StartsUTF8Codepoint(static_cast(next_char)) && + std::any_of(states.begin(), states.end(), [&](const ParserState& state) { + return IsCharExpiredState(state); + })) { + break; + } + result += static_cast(next_char); + Advance(next_char); + ++num_accepted_chars; + } + } + + // Rollback all chars accepted + PopLastStates(num_accepted_chars); + return result; +} + +void GrammarMatcher::Impl::Rollback(int num_tokens) { + XGRAMMAR_CHECK(num_tokens <= static_cast(token_length_history.size())) + << "Intended to rollback " << num_tokens << " tokens, but only the last " + << token_length_history.size() << " steps of history are saved"; + while (num_tokens > 0) { + int steps = token_length_history.back(); + PopLastStates(steps); + token_length_history.pop_back(); + if (ShouldTrackAcceptedBytes() && steps > 0) { + row_byte_end_.resize(row_byte_end_.size() - steps); + accepted_bytes_.resize(row_byte_end_.back()); + } + --num_tokens; + } + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = false; + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + budget_body_match_cache_.clear(); +} + +std::vector> GrammarMatcher::Impl::GetCaptures(bool deduplicate +) const { + std::vector> result; + if (!IsCaptureTrackingEnabled()) { + return result; + } + XGRAMMAR_DCHECK(capture_event_history_.size() == static_cast(row_byte_end_.size())) + << "The capture history is not aligned with the byte history: " + << capture_event_history_.size() << " vs " << row_byte_end_.size(); + + // Flatten the event history. Events are ordered by completion position, and by completion + // order within a position. + struct FlatEvent { + int32_t rule_id; + int32_t start_row; + int32_t occurrence_start_pos; + int32_t end_row; + int32_t hidden_suffix_bytes; + int32_t hidden_stop_bytes; + std::vector stop_capture_targets; + bool marker_present = false; + int64_t marker_start_byte = -1; + }; + std::vector events; + for (int32_t row = 0; row < capture_event_history_.size(); ++row) { + for (const auto& event : capture_event_history_[row]) { + int32_t start_row = event.start_pos == ParserState::kNoPrevInputPos ? 0 : event.start_pos; + XGRAMMAR_DCHECK(start_row <= row); + events.push_back( + {event.rule_id, + start_row, + event.occurrence_start_pos, + row, + event.hidden_suffix_bytes, + event.hidden_stop_bytes, + event.stop_capture_targets, + false, + -1} + ); + } + } + + // The Earley parser explores all parse hypotheses in parallel, so one occurrence of a rule + // (identified by its rule id and start position) may complete at several candidate end + // positions before the following input decides the real one. When deduplicate is true, we + // keep only the last (longest) completion of each occurrence. + std::vector keep(events.size(), true); + if (deduplicate) { + std::unordered_map last_index; + for (size_t i = 0; i < events.size(); ++i) { + int64_t key = (static_cast(events[i].rule_id) << 32) | + static_cast(events[i].start_row); + auto it = last_index.find(key); + if (it != last_index.end()) { + keep[it->second] = false; + it->second = i; + } else { + last_index[key] = i; + } + } + } + + auto contains_end_state = [](const CompactFSMWithStartEnd& fsm, + const std::unordered_set& states) { + return std::any_of(states.begin(), states.end(), [&](int state) { + return fsm.IsEndState(state); + }); + }; + + auto accepted_prefixes = [&](const CompactFSMWithStartEnd& fsm, int64_t begin, int64_t end) { + XGRAMMAR_DCHECK(fsm.IsLeaf()) << "A suffix/stop capture helper must compile to a leaf FSM"; + std::vector result(static_cast(end - begin + 1), false); + std::unordered_set states{fsm.GetStart()}; + fsm.GetFsm().GetEpsilonClosure(&states); + result[0] = contains_end_state(fsm, states); + std::unordered_set next_states; + for (int64_t offset = 0; offset < end - begin; ++offset) { + fsm.GetFsm().Advance( + states, accepted_bytes_[begin + offset], &next_states, FSMEdge::EdgeType::kCharRange, true + ); + states = next_states; + result[static_cast(offset + 1)] = contains_end_state(fsm, states); + } + return result; + }; + + // Compute which suffixes of [begin, end) are accepted by marker_fsm. This is a reverse NFA + // traversal, so locating a delimiter remains linear in the captured byte length even when the + // body (for example /.*/) can end at every byte. + auto accepted_suffixes = [&](const CompactFSMWithStartEnd& marker_fsm, int64_t begin, int64_t end + ) { + XGRAMMAR_DCHECK(marker_fsm.IsLeaf()) + << "A suffix/stop marker helper must compile to a leaf FSM"; + const auto& fsm = marker_fsm.GetFsm(); + std::unordered_set reachable; + marker_fsm.GetReachableStates(&reachable); + struct ReverseCharEdge { + int32_t source; + int32_t min; + int32_t max; + }; + std::vector> reverse_epsilon(fsm.NumStates()); + std::vector> reverse_char(fsm.NumStates()); + for (int32_t source : reachable) { + for (const auto& edge : fsm.GetEdges(source)) { + if (!reachable.count(edge.target)) { + continue; + } + if (edge.IsEpsilon()) { + reverse_epsilon[edge.target].push_back(source); + } else if (edge.IsCharRange()) { + reverse_char[edge.target].push_back({source, edge.min, edge.max}); + } + } + } + auto add_reverse_epsilon_closure = [&](std::unordered_set* states) { + std::vector queue(states->begin(), states->end()); + for (size_t i = 0; i < queue.size(); ++i) { + for (int32_t predecessor : reverse_epsilon[queue[i]]) { + if (states->insert(predecessor).second) { + queue.push_back(predecessor); + } + } + } + }; + + std::vector result(static_cast(end - begin + 1), false); + std::unordered_set states(marker_fsm.GetEnds().begin(), marker_fsm.GetEnds().end()); + add_reverse_epsilon_closure(&states); + result[static_cast(end - begin)] = states.count(marker_fsm.GetStart()) != 0; + std::unordered_set previous_states; + for (int64_t offset = end - begin; offset > 0; --offset) { + previous_states.clear(); + int32_t byte = accepted_bytes_[begin + offset - 1]; + for (int32_t target : states) { + for (const auto& edge : reverse_char[target]) { + if (edge.min <= byte && byte <= edge.max) { + previous_states.insert(edge.source); + } + } + } + add_reverse_epsilon_closure(&previous_states); + states = previous_states; + result[static_cast(offset - 1)] = states.count(marker_fsm.GetStart()) != 0; + } + return result; + }; + + // Resolve every marker to a concrete byte span once. Fixed strings use their recorded length; + // regexes and named terminals use helper FSMs for both sides of the original body/marker split. + for (size_t i = 0; i < events.size(); ++i) { + if (!keep[i]) { + continue; + } + auto& event = events[i]; + int32_t hidden_bytes = std::max(event.hidden_suffix_bytes, event.hidden_stop_bytes); + if (hidden_bytes <= 0) { + continue; + } + const auto& rule = grammar_->GetRule(event.rule_id); + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(event.rule_id); + XGRAMMAR_DCHECK(suffix_stop_info != nullptr); + int64_t event_start = row_byte_end_[event.start_row]; + int64_t event_end = row_byte_end_[event.end_row]; + event.marker_present = true; + if (suffix_stop_info->body_rule_id == -1) { + event.marker_start_byte = + std::max(event_start, event_end - static_cast(hidden_bytes)); + continue; + } + + XGRAMMAR_DCHECK(suffix_stop_info->body_rule_id >= 0 && suffix_stop_info->marker_rule_id >= 0); + XGRAMMAR_DCHECK( + grammar_->per_rule_fsms[suffix_stop_info->body_rule_id].has_value() && + grammar_->per_rule_fsms[suffix_stop_info->marker_rule_id].has_value() + ); + const auto& body_fsm = grammar_->per_rule_fsms[suffix_stop_info->body_rule_id]->GetFsm(); + const auto& marker_fsm = grammar_->per_rule_fsms[suffix_stop_info->marker_rule_id]->GetFsm(); + std::vector body_ends = accepted_prefixes(body_fsm, event_start, event_end); + std::vector marker_starts = accepted_suffixes(marker_fsm, event_start, event_end); + bool found = false; + for (size_t offset = 0; offset < body_ends.size(); ++offset) { + if (body_ends[offset] && marker_starts[offset]) { + event.marker_start_byte = event_start + static_cast(offset); + found = true; + break; + } + } + XGRAMMAR_DCHECK(found) << "Could not recover the suffix/stop marker boundary for rule " + << rule.name; + if (!found) { + event.marker_present = false; + event.marker_start_byte = -1; + } + } + + // Associate each stop marker only with captured occurrences reached through its concrete + // Earley parent chain. Failed or unrelated branches may cover identical byte ranges, so byte + // overlap alone is not a valid indication that a marker belongs to a capture. + std::unordered_map>> + hidden_stop_spans_by_capture; + for (size_t i = 0; i < events.size(); ++i) { + if (!keep[i] || events[i].hidden_stop_bytes <= 0 || !events[i].marker_present) { + continue; + } + int64_t event_end = row_byte_end_[events[i].end_row]; + int64_t hidden_start = events[i].marker_start_byte; + if (hidden_start < event_end) { + for (const auto& target : events[i].stop_capture_targets) { + int64_t target_key = + (static_cast(target.rule_id) << 32) | static_cast(target.start_pos); + hidden_stop_spans_by_capture[target_key].emplace_back(hidden_start, event_end); + } + } + } + for (auto& [_, spans] : hidden_stop_spans_by_capture) { + std::sort(spans.begin(), spans.end()); + } + + for (size_t i = 0; i < events.size(); ++i) { + if (!keep[i]) { + continue; + } + const auto& event = events[i]; + const auto& rule = grammar_->GetRule(event.rule_id); + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(event.rule_id); + int64_t capture_start = row_byte_end_[event.start_row]; + int64_t capture_end = row_byte_end_[event.end_row]; + if (event.marker_present && suffix_stop_info != nullptr && + !suffix_stop_info->stop_capture_name.empty()) { + std::string marker_capture; + if (event.marker_start_byte < capture_end) { + marker_capture.assign( + reinterpret_cast(accepted_bytes_.data() + event.marker_start_byte), + static_cast(capture_end - event.marker_start_byte) + ); + } + result.emplace_back(suffix_stop_info->stop_capture_name, std::move(marker_capture)); + } + if (rule.capture_name.empty()) { + continue; + } + if (event.hidden_suffix_bytes > 0 && event.marker_present) { + capture_end = event.marker_start_byte; + } + + std::string capture; + int64_t cursor = capture_start; + auto append_bytes = [&](int64_t begin, int64_t end) { + if (begin < end) { + capture.append( + reinterpret_cast(accepted_bytes_.data() + begin), + static_cast(end - begin) + ); + } + }; + int64_t capture_key = (static_cast(event.rule_id) << 32) | + static_cast(event.occurrence_start_pos); + auto hidden_spans_it = hidden_stop_spans_by_capture.find(capture_key); + if (hidden_spans_it != hidden_stop_spans_by_capture.end()) { + for (const auto& hidden : hidden_spans_it->second) { + if (hidden.second <= capture_start) { + continue; + } + if (hidden.first >= capture_end) { + break; + } + int64_t hidden_begin = std::max(hidden.first, capture_start); + int64_t hidden_end = std::min(hidden.second, capture_end); + append_bytes(cursor, hidden_begin); + cursor = std::max(cursor, hidden_end); + } + } + append_bytes(cursor, capture_end); + result.emplace_back(rule.capture_name, std::move(capture)); + } + return result; +} + +void GrammarMatcher::Impl::SetTokenBitmask( + int32_t* bitmask_data_ptr, + const DynamicBitset& accepted_bitset, + const std::vector& rejected_indices, + bool can_reach_end, + bool allow_special_token +) { + // next_token_bitmask = set(all accepted tokens) = + // 1. all_tokens - (rejected_ids / accepted_ids) + // (when rejected_ids != {-1}, i.e. rejected_ids is not the universal set) + // 2. accepted_ids + // (otherwise, when rejected_ids is the universal set) + DynamicBitset next_token_bitset( + tokenizer_info_.GetVocabSize(), reinterpret_cast(bitmask_data_ptr) + ); + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + + if (rejected_indices.size() == 1 && rejected_indices[0] == -1) { + // If rejected_indices is the universal set, the final accepted token set is just + // accepted_indices + next_token_bitset = accepted_bitset; + + if (allow_special_token) { + for (int id : tokenizer_info_.GetSpecialTokenIds()) { + next_token_bitset.Set(id, true); + } + } + + if (can_reach_end) { + // add end tokens + for (int id : stop_token_ids_) { + next_token_bitset.Set(id, true); + } + } + } else { + // Otherwise, the final rejected token set is (rejected_indices \ accepted_indices) + next_token_bitset.Set(); + + for (auto i : rejected_indices) { + auto id = sorted_decoded_vocab[i].first; + if (!accepted_bitset[id]) { + next_token_bitset.Set(id, false); + } + } + if (!allow_special_token) { + for (int id : tokenizer_info_.GetSpecialTokenIds()) { + next_token_bitset.Set(id, false); + } + } + if (!can_reach_end) { + for (int id : stop_token_ids_) { + next_token_bitset.Set(id, false); + } + } + } +} + +int GrammarMatcher::Impl::GetNextUncertainToken( + bool is_uncertain_saved, + int* iterator_uncertain, + const std::vector& uncertain_indices, + const std::vector& uncertain_tokens_bitset +) { + if (is_uncertain_saved) { + ++*iterator_uncertain; + if (*iterator_uncertain == static_cast(uncertain_indices.size())) { + return -1; + } + return uncertain_indices[*iterator_uncertain]; + } else { + ++*iterator_uncertain; + while (*iterator_uncertain < static_cast(uncertain_tokens_bitset.size()) && + !uncertain_tokens_bitset[*iterator_uncertain]) { + ++*iterator_uncertain; + } + if (*iterator_uncertain == static_cast(uncertain_tokens_bitset.size())) { + return -1; + } + return *iterator_uncertain; + } +} + +void BatchGrammarMatcher::Impl::BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices, + bool debug_print +) { + XGRAMMAR_CHECK(!indices.has_value() || indices->size() == matchers->size()) + << "The size of indices (" << (indices.has_value() ? indices->size() : 0) + << ") should be the same as the size of matchers (" << matchers->size() << ")."; + // Initialize the thread pool if needed. It should be initialized each time, + // because ThreadPool cannot be reused after Join(). + if (max_threads_ > 1) { + thread_pool_.emplace(max_threads_); + } + if (!thread_pool_.has_value()) { + for (int i = 0; i < static_cast(matchers->size()); i++) { + auto& matcher = (*matchers)[i]; + int index = indices.has_value() ? (*indices)[i] : i; + XGRAMMAR_CHECK(index >= 0 && index < next_token_bitmask->shape[0]) + << "The index " << index << " is out of range [0, " << next_token_bitmask->shape[0] + << ") for batch_id " << i << "."; + matcher->FillNextTokenBitmask(next_token_bitmask, index, debug_print); + } + } else { + auto fill_next_token_mask = [&](int32_t batch_id) { + auto& matcher = (*matchers)[batch_id]; + int index = indices.has_value() ? (*indices)[batch_id] : batch_id; + XGRAMMAR_CHECK(index >= 0 && index < next_token_bitmask->shape[0]) + << "The index " << index << " is out of range [0, " << next_token_bitmask->shape[0] + << ") for batch_id " << batch_id << "."; + matcher->FillNextTokenBitmask(next_token_bitmask, index, debug_print); + }; + for (int i = 0; i < static_cast(matchers->size()); i++) { + thread_pool_->Execute([fill_next_token_mask, i]() { fill_next_token_mask(i); }); + } + thread_pool_->Join(); + } +} + +std::vector BatchGrammarMatcher::Impl::BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print +) { + XGRAMMAR_CHECK(matchers->size() == input_strs.size()) + << "The size of matchers (" << matchers->size() << ") and input_strs (" << input_strs.size() + << ") should be the same."; + std::vector accepted(matchers->size()); + for (int i = 0; i < static_cast(matchers->size()); i++) { + auto& matcher = (*matchers)[i]; + accepted[i] = matcher->AcceptString(input_strs[i], debug_print); + } + return accepted; +} + +std::vector BatchGrammarMatcher::Impl::BatchAcceptToken( + std::vector* matchers, const std::vector& token_ids, bool debug_print +) { + XGRAMMAR_CHECK(matchers->size() == token_ids.size()) + << "The size of matchers (" << matchers->size() << ") and token_ids (" << token_ids.size() + << ") should be the same."; + std::vector accepted(matchers->size()); + for (int i = 0; i < static_cast(matchers->size()); i++) { + auto& matcher = (*matchers)[i]; + accepted[i] = matcher->AcceptToken(token_ids[i], debug_print); + } + return accepted; +} + +void BatchGrammarMatcher::Impl::BatchRollback( + std::vector* matchers, const std::vector& num_tokens +) { + XGRAMMAR_CHECK(matchers->size() == num_tokens.size()) + << "The size of matchers (" << matchers->size() << ") and num_tokens (" << num_tokens.size() + << ") should be the same."; + for (int i = 0; i < static_cast(matchers->size()); i++) { + (*matchers)[i].Rollback(num_tokens[i]); + } +} + +GrammarMatcher::GrammarMatcher( + const CompiledGrammar& compiled_grammar, + std::optional> override_stop_tokens, + bool terminate_without_stop_token, + int max_rollback_tokens, + std::optional default_temperature +) + : pimpl_(std::make_shared( + compiled_grammar, + override_stop_tokens, + terminate_without_stop_token, + max_rollback_tokens, + default_temperature + )) {} + +bool GrammarMatcher::AcceptToken(int32_t token_id, bool debug_print) { + return pimpl_->AcceptToken(token_id, debug_print); +} + +bool GrammarMatcher::AcceptString(const std::string& input_str, bool debug_print) { + return pimpl_->AcceptString(input_str, debug_print); +} + +bool GrammarMatcher::FillNextTokenBitmask( + DLTensor* next_token_bitmask, int index, bool debug_print +) { + return pimpl_->FillNextTokenBitmask(next_token_bitmask, index, debug_print); +} + +bool GrammarMatcher::TraverseDraftTree( + const DLTensor* retrieve_next_token, + const DLTensor* retrieve_next_sibling, + const DLTensor* draft_tokens, + DLTensor* token_bitmask, + double time_threshold, + DLTensor* temperatures +) { + auto check_cpu = [](const DLTensor* tensor, const char* name) { + XGRAMMAR_CHECK( + tensor->device.device_type == kDLCPU || tensor->device.device_type == kDLCUDAHost || + tensor->device.device_type == kDLROCMHost + ) << "The " + << name << " tensor must be on CPU"; + }; + + XGRAMMAR_CHECK( + retrieve_next_token->ndim == 1 && retrieve_next_token->dtype.code == kDLInt && + retrieve_next_token->dtype.bits == 64 + ) << "The retrieve_next_token tensor must be a 1D int64 tensor"; + XGRAMMAR_CHECK( + retrieve_next_sibling->ndim == 1 && retrieve_next_sibling->dtype.code == kDLInt && + retrieve_next_sibling->dtype.bits == 64 + ) << "The retrieve_next_sibling tensor must be a 1D int64 tensor"; + XGRAMMAR_CHECK( + draft_tokens->ndim == 1 && draft_tokens->dtype.code == kDLInt && + draft_tokens->dtype.bits == 64 + ) << "The draft_tokens tensor must be a 1D int64 tensor"; + XGRAMMAR_CHECK( + token_bitmask->ndim == 2 && token_bitmask->dtype.code == kDLInt && + token_bitmask->dtype.bits == 32 + ) << "The token_bitmask tensor must be a 2D int32 tensor"; + if (temperatures != nullptr) { + XGRAMMAR_CHECK( + temperatures->ndim == 1 && temperatures->dtype.code == kDLFloat && + temperatures->dtype.bits == 32 + ) << "The temperatures tensor must be a 1D float32 tensor"; + } + + check_cpu(retrieve_next_token, "retrieve_next_token"); + check_cpu(retrieve_next_sibling, "retrieve_next_sibling"); + check_cpu(draft_tokens, "draft_tokens"); + check_cpu(token_bitmask, "token_bitmask"); + if (temperatures != nullptr) { + check_cpu(temperatures, "temperatures"); + } + + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == retrieve_next_sibling->shape[0]) + << "The retrieve_next_token and retrieve_next_sibling tensors must have the same length"; + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == draft_tokens->shape[0]) + << "The retrieve_next_token and draft_tokens tensors must have the same length"; + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == token_bitmask->shape[0]) + << "The token_bitmask batch size must match the number of nodes in the tree"; + if (temperatures != nullptr) { + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == temperatures->shape[0]) + << "The temperatures size must match the number of nodes in the tree"; + std::fill_n(reinterpret_cast(temperatures->data), temperatures->shape[0], -1.0f); + } + XGRAMMAR_CHECK(retrieve_next_sibling->shape[0] > 0 && retrieve_next_sibling->data != nullptr) + << "The draft tree must not be empty"; + XGRAMMAR_CHECK(reinterpret_cast(retrieve_next_sibling->data)[0] == -1) + << "The root node must not have siblings"; + + // The traversal follows retrieve_next_token / retrieve_next_sibling as node indices; a value + // outside [-1, num_nodes) would recurse into an out-of-bounds position, so validate them upfront. + // The edges must also form a tree: the root is never a target, and every other node is the + // child or next sibling of at most one node. A node targeted twice (a cycle or a shared node) + // would otherwise make the traversal recurse without bound. + int64_t num_nodes = retrieve_next_token->shape[0]; + const int64_t* next_token_data = reinterpret_cast(retrieve_next_token->data); + const int64_t* next_sibling_data = reinterpret_cast(retrieve_next_sibling->data); + std::vector in_degree(num_nodes, 0); + for (int64_t i = 0; i < num_nodes; ++i) { + XGRAMMAR_CHECK(next_token_data[i] >= -1 && next_token_data[i] < num_nodes) + << "retrieve_next_token[" << i << "] = " << next_token_data[i] + << " is out of bounds: it should be in [-1, " << num_nodes << ")."; + XGRAMMAR_CHECK(next_sibling_data[i] >= -1 && next_sibling_data[i] < num_nodes) + << "retrieve_next_sibling[" << i << "] = " << next_sibling_data[i] + << " is out of bounds: it should be in [-1, " << num_nodes << ")."; + if (next_token_data[i] != -1) { + ++in_degree[next_token_data[i]]; + } + if (next_sibling_data[i] != -1) { + ++in_degree[next_sibling_data[i]]; + } + } + XGRAMMAR_CHECK(in_degree[0] == 0) << "The root node must not be the child or sibling of a node"; + for (int64_t i = 1; i < num_nodes; ++i) { + XGRAMMAR_CHECK(in_degree[i] <= 1) + << "Node " << i << " is the target of " << in_degree[i] + << " retrieve_next_token / retrieve_next_sibling entries; the draft tree must not contain " + "cycles or shared nodes."; + } + + return details::TraverseDraftTreeRecursive( + 0, + -1, + reinterpret_cast(retrieve_next_token->data), + reinterpret_cast(retrieve_next_sibling->data), + reinterpret_cast(draft_tokens->data), + *this, + token_bitmask, + temperatures == nullptr ? nullptr : reinterpret_cast(temperatures->data), + time_threshold, + details::Clock::now() + ); +} + +std::string GrammarMatcher::FindJumpForwardString() { return pimpl_->FindJumpForwardString(); } + +void GrammarMatcher::Rollback(int num_tokens) { pimpl_->Rollback(num_tokens); } + +std::vector> GrammarMatcher::GetCaptures(bool deduplicate +) const { + return pimpl_->GetCaptures(deduplicate); +} + +bool GrammarMatcher::IsTerminated() const { return pimpl_->IsTerminated(); } + +bool GrammarMatcher::IsCompleted() const { return pimpl_->IsCompleted(); } + +void GrammarMatcher::Reset() { pimpl_->Reset(); } + +GrammarMatcher GrammarMatcher::Fork() const { + return GrammarMatcher(std::make_shared(*pimpl_)); +} + +int GrammarMatcher::GetMaxRollbackTokens() const { return pimpl_->GetMaxRollbackTokens(); } + +std::optional GrammarMatcher::GetTemperature() const { return pimpl_->GetTemperature(); } + +const std::vector& GrammarMatcher::GetStopTokenIds() const { + return pimpl_->GetStopTokenIds(); +} + +std::string GrammarMatcher::_DebugPrintInternalState() const { + return pimpl_->_DebugPrintInternalState(); +} + +void BatchGrammarMatcher::BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices, + bool debug_print +) { + pimpl_->BatchFillNextTokenBitmask(matchers, next_token_bitmask, indices, debug_print); +} + +void BatchGrammarMatcher::BatchFillTemperature( + const std::vector& matchers, + DLTensor* temperatures, + const std::optional>& indices +) { + XGRAMMAR_CHECK( + temperatures->ndim == 1 && temperatures->dtype.code == kDLFloat && + temperatures->dtype.bits == 32 + ) << "The temperatures tensor must be a 1D float32 tensor"; + XGRAMMAR_CHECK( + temperatures->device.device_type == kDLCPU || + temperatures->device.device_type == kDLCUDAHost || + temperatures->device.device_type == kDLROCMHost + ) << "The temperatures tensor must be on CPU"; + if (indices.has_value()) { + XGRAMMAR_CHECK(indices->size() == matchers.size()) + << "The indices size must match the number of matchers"; + } + auto* temperatures_data = reinterpret_cast(temperatures->data); + for (int64_t i = 0; i < static_cast(matchers.size()); ++i) { + int64_t position = indices.has_value() ? (*indices)[i] : i; + XGRAMMAR_CHECK(position >= 0 && position < temperatures->shape[0]) + << "The index " << position << " is out of bounds for the temperatures tensor"; + temperatures_data[position] = matchers[i].GetTemperature().value_or(-1.0f); + } +} + +std::vector BatchGrammarMatcher::BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print +) { + return Impl::BatchAcceptString(matchers, input_strs, debug_print); +} + +std::vector BatchGrammarMatcher::BatchAcceptToken( + std::vector* matchers, const std::vector& token_ids, bool debug_print +) { + return Impl::BatchAcceptToken(matchers, token_ids, debug_print); +} + +void BatchGrammarMatcher::BatchRollback( + std::vector* matchers, const std::vector& num_tokens +) { + Impl::BatchRollback(matchers, num_tokens); +} + +BatchGrammarMatcher::BatchGrammarMatcher(std::variant max_threads) + : pimpl_(std::make_shared(max_threads)) {} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_parser.cc b/third_party/xgrammar/cpp/grammar_parser.cc new file mode 100644 index 0000000000..1d00ec4092 --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_parser.cc @@ -0,0 +1,1593 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_parser.cc + */ + +#include "grammar_parser.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_impl.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +class EBNFLexer::Impl { + public: + using Token = EBNFLexer::Token; + using TokenType = EBNFLexer::TokenType; + + std::vector Tokenize(const std::string& input); + + private: + std::string input_; + const char* cur_ = nullptr; + int cur_line_ = 1; + int cur_column_ = 1; + + constexpr static int64_t kMaxIntegerInGrammar = 1e15; + + // Helper functions + + /*! + * \brief Consume a character sequence and return the next token. Return a token if it's a + * single token, or a vector of tokens if it's a sequence of tokens. + * + * \return std::variant> + */ + std::variant> NextToken(); + Token ParseIdentifierOrBooleanToken(); + Token ParseStringToken(); + std::vector ParseCharClassToken(); + Token ParseIntegerToken(); + [[noreturn]] void ReportLexerError(const std::string& msg, int line = -1, int column = -1); + char Peek(int delta = 0) const; + void Consume(int cnt = 1); + void ConsumeSpace(); + std::string ParseIdentifierToken(); + void ConvertIdentifierToRuleName(std::vector* tokens); + static bool IsNameChar(char c, bool is_first = false); +}; + +// Look at the next character +inline char EBNFLexer::Impl::Peek(int delta) const { return *(cur_ + delta); } + +// Consume characters and update position information +inline void EBNFLexer::Impl::Consume(int cnt) { + for (int i = 0; i < cnt; ++i) { + // Newline\n \r \r\n + if (*cur_ == '\n' || (*cur_ == '\r' && *(cur_ + 1) != '\n')) { + ++cur_line_; + cur_column_ = 1; + } else { + ++cur_column_; + } + ++cur_; + } +} + +// Skip whitespace and comments +void EBNFLexer::Impl::ConsumeSpace() { + while (Peek() && + (Peek() == ' ' || Peek() == '\t' || Peek() == '#' || Peek() == '\n' || Peek() == '\r')) { + Consume(); + if (Peek(-1) == '#') { + while (Peek() && Peek() != '\n' && Peek() != '\r') { + Consume(); + } + if (!Peek()) { + return; + } + Consume(); + if (Peek(-1) == '\r' && Peek() == '\n') { + Consume(); + } + } + } +} + +// Report parsing error +void EBNFLexer::Impl::ReportLexerError(const std::string& msg, int line, int column) { + int line_to_print = line == -1 ? cur_line_ : line; + int column_to_print = column == -1 ? cur_column_ : column; + XGRAMMAR_LOG(FATAL) << "EBNF lexer error at line " + std::to_string(line_to_print) + ", column " + + std::to_string(column_to_print) + ": " + msg; + XGRAMMAR_UNREACHABLE(); +} + +// Check if a character can be part of an identifier +bool EBNFLexer::Impl::IsNameChar(char c, bool is_first) { + return c == '_' || c == '-' || c == '.' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (!is_first && c >= '0' && c <= '9'); +} + +// Parse identifier +std::string EBNFLexer::Impl::ParseIdentifierToken() { + const char* start = cur_; + bool first_char = true; + while (*cur_ && IsNameChar(*cur_, first_char)) { + Consume(); + first_char = false; + } + if (start == cur_) { + ReportLexerError("Expect identifier"); + } + return std::string(start, cur_ - start); +} + +// Parse identifier or boolean value +EBNFLexer::Token EBNFLexer::Impl::ParseIdentifierOrBooleanToken() { + int start_line = cur_line_; + int start_column = cur_column_; + + std::string identifier = ParseIdentifierToken(); + + // Check if it's a boolean value + if (identifier == "true" || identifier == "false") { + return { + TokenType::BooleanLiteral, + identifier, + identifier == "true" ? true : false, + start_line, + start_column + }; + } + + // A rule definition may carry an attribute block before ::=, e.g. + // name[max_tokens=10] ::= ..., name[max_chars=10] ::= ..., name[capture] ::= ..., + // name[capture="x"] ::= ..., + // name[capture_hidden_suffix_bytes=3] ::= ..., name[capture_hidden_stop_bytes=3] ::= ..., + // name[capture_hidden_body_rule_id=1, capture_hidden_marker_rule_id=2] ::= ..., + // name[stop_capture="marker"] ::= ..., name[lazy] ::= ..., name[temperature=0.7] ::= ..., or a + // comma-separated combination. + // The bracket group is treated as an attribute block only when it is followed by "::="; + // otherwise it is left to be lexed as a character class. + if (*cur_ == '[') { + int delta = 1; + auto skip_space = [&]() { + while (Peek(delta) == ' ' || Peek(delta) == '\t') { + ++delta; + } + }; + // Match the keyword at the current position and advance delta past it on success. + auto match_keyword = [&](const char* keyword) { + int len = 0; + while (keyword[len] != '\0') { + if (Peek(delta + len) != keyword[len]) { + return false; + } + ++len; + } + delta += len; + return true; + }; + bool matched = true; + bool has_max_tokens = false; + bool has_max_chars = false; + bool has_capture = false; + bool has_capture_hidden_suffix_bytes = false; + bool has_capture_hidden_stop_bytes = false; + bool has_capture_hidden_body_rule_id = false; + bool has_capture_hidden_marker_rule_id = false; + bool has_stop_capture = false; + bool has_lazy = false; + bool has_temperature = false; + double temperature_value = 0; + int64_t max_tokens_value = -1; + int64_t max_chars_value = -1; + int64_t capture_hidden_suffix_bytes_value = 0; + int64_t capture_hidden_stop_bytes_value = 0; + int64_t capture_hidden_body_rule_id_value = -1; + int64_t capture_hidden_marker_rule_id_value = -1; + std::string capture_value; + std::string stop_capture_value; + auto parse_integer_value = [&](int64_t* value) { + skip_space(); + if (Peek(delta) != '=') { + return false; + } + ++delta; + skip_space(); + *value = 0; + int digits = 0; + while (Peek(delta) >= '0' && Peek(delta) <= '9' && digits < 10) { + *value = *value * 10 + (Peek(delta) - '0'); + ++delta; + ++digits; + } + return digits > 0 && !(Peek(delta) >= '0' && Peek(delta) <= '9'); + }; + auto parse_string_value = [&](std::string* value) { + skip_space(); + if (Peek(delta) != '=') { + return false; + } + ++delta; + skip_space(); + if (Peek(delta) != '"') { + return false; + } + ++delta; + while (Peek(delta) != '"') { + char c = Peek(delta); + if (c == '\0' || c == '\n' || c == '\r' || c == '\\') { + return false; + } + value->push_back(c); + ++delta; + } + ++delta; + return true; + }; + auto parse_float_value = [&](double* value) { + skip_space(); + if (Peek(delta) != '=') { + return false; + } + ++delta; + skip_space(); + std::string text; + while (true) { + char c = Peek(delta); + if ((c >= '0' && c <= '9') || c == '.' || c == 'e' || c == 'E' || c == '+' || c == '-') { + text.push_back(c); + ++delta; + } else { + break; + } + } + if (text.empty()) { + return false; + } + try { + *value = std::stod(text); + } catch (const std::out_of_range&) { + *value = std::numeric_limits::infinity(); + } catch (const std::exception&) { + return false; + } + return true; + }; + // Parse a comma-separated attribute list. Each attribute may appear at most once. + while (matched) { + skip_space(); + if (!has_max_tokens && match_keyword("max_tokens")) { + has_max_tokens = true; + matched = parse_integer_value(&max_tokens_value); + } else if (!has_max_chars && match_keyword("max_chars")) { + has_max_chars = true; + matched = parse_integer_value(&max_chars_value); + } else if (!has_capture_hidden_suffix_bytes && match_keyword("capture_hidden_suffix_bytes")) { + has_capture_hidden_suffix_bytes = true; + matched = parse_integer_value(&capture_hidden_suffix_bytes_value); + } else if (!has_capture_hidden_stop_bytes && match_keyword("capture_hidden_stop_bytes")) { + has_capture_hidden_stop_bytes = true; + matched = parse_integer_value(&capture_hidden_stop_bytes_value); + } else if (!has_capture_hidden_body_rule_id && match_keyword("capture_hidden_body_rule_id")) { + has_capture_hidden_body_rule_id = true; + matched = parse_integer_value(&capture_hidden_body_rule_id_value); + } else if (!has_capture_hidden_marker_rule_id && + match_keyword("capture_hidden_marker_rule_id")) { + has_capture_hidden_marker_rule_id = true; + matched = parse_integer_value(&capture_hidden_marker_rule_id_value); + } else if (!has_stop_capture && match_keyword("stop_capture")) { + has_stop_capture = true; + matched = parse_string_value(&stop_capture_value); + } else if (!has_capture && match_keyword("capture")) { + has_capture = true; + skip_space(); + if (Peek(delta) == '=') { + matched = parse_string_value(&capture_value); + } else { + capture_value = identifier; + } + } else if (!has_lazy && match_keyword("lazy")) { + has_lazy = true; + } else if (!has_temperature && match_keyword("temperature")) { + has_temperature = true; + matched = parse_float_value(&temperature_value); + } else { + matched = false; + } + if (!matched) { + break; + } + skip_space(); + if (Peek(delta) == ',') { + ++delta; + continue; + } + break; + } + if (matched) { + skip_space(); + if (Peek(delta) == ']') { + ++delta; + } else { + matched = false; + } + } + if (matched) { + int after_bracket = delta; + while (Peek(after_bracket) == ' ' || Peek(after_bracket) == '\t') { + ++after_bracket; + } + if (!(Peek(after_bracket) == ':' && Peek(after_bracket + 1) == ':' && + Peek(after_bracket + 2) == '=')) { + matched = false; + } + } + if (matched) { + if (has_capture && capture_value.empty()) { + ReportLexerError("The capture name must not be empty", start_line, start_column); + } + if (has_stop_capture && stop_capture_value.empty()) { + ReportLexerError("The stop capture name must not be empty", start_line, start_column); + } + if (has_capture_hidden_body_rule_id != has_capture_hidden_marker_rule_id) { + ReportLexerError( + "The capture-hidden body and marker rule ids must be specified together", + start_line, + start_column + ); + } + if ((has_capture_hidden_suffix_bytes && capture_hidden_suffix_bytes_value <= 0) || + (has_capture_hidden_stop_bytes && capture_hidden_stop_bytes_value <= 0)) { + ReportLexerError( + "The number of capture-hidden bytes must be positive", start_line, start_column + ); + } + constexpr int64_t kMaxInt32 = std::numeric_limits::max(); + if (has_max_chars && max_chars_value < 0) { + ReportLexerError( + "The max_chars rule attribute must be non-negative", start_line, start_column + ); + } + if ((has_max_chars && max_chars_value > kMaxInt32) || + (has_capture_hidden_suffix_bytes && capture_hidden_suffix_bytes_value > kMaxInt32) || + (has_capture_hidden_stop_bytes && capture_hidden_stop_bytes_value > kMaxInt32) || + (has_capture_hidden_body_rule_id && capture_hidden_body_rule_id_value > kMaxInt32) || + (has_capture_hidden_marker_rule_id && capture_hidden_marker_rule_id_value > kMaxInt32)) { + ReportLexerError("The rule attribute value is too large", start_line, start_column); + } + if (has_temperature && !(std::isfinite(temperature_value) && temperature_value >= 0 && + temperature_value <= std::numeric_limits::max())) { + ReportLexerError( + "The temperature must be a finite non-negative number", start_line, start_column + ); + } + Consume(delta); + Token token{TokenType::Identifier, identifier, identifier, start_line, start_column}; + token.max_tokens = static_cast(max_tokens_value); + token.max_chars = static_cast(max_chars_value); + token.capture_name = capture_value; + token.capture_hidden_suffix_bytes = static_cast(capture_hidden_suffix_bytes_value); + token.capture_hidden_stop_bytes = static_cast(capture_hidden_stop_bytes_value); + token.capture_hidden_body_rule_id = static_cast(capture_hidden_body_rule_id_value); + token.capture_hidden_marker_rule_id = + static_cast(capture_hidden_marker_rule_id_value); + token.stop_capture_name = stop_capture_value; + token.is_lazy = has_lazy; + if (has_temperature) { + token.temperature = static_cast(temperature_value); + } + return token; + } + } + + // Otherwise it's an identifier + return {TokenType::Identifier, identifier, identifier, start_line, start_column}; +} + +// Parse string literal +EBNFLexer::Token EBNFLexer::Impl::ParseStringToken() { + int start_line = cur_line_; + int start_column = cur_column_; + const char* start_pos = cur_; + + Consume(); // Skip opening quote + + std::vector codepoints; + while (Peek() && Peek() != '"' && Peek() != '\n' && Peek() != '\r') { + auto [codepoint, len] = ParseNextUTF8OrEscaped(cur_); + if (codepoint == CharHandlingError::kInvalidUTF8) { + ReportLexerError("Invalid UTF8 sequence"); + } + if (codepoint == CharHandlingError::kInvalidEscape) { + ReportLexerError("Invalid escape sequence"); + } + Consume(len); + codepoints.push_back(codepoint); + } + + if (Peek() != '"') { + ReportLexerError("Expect \" in string literal"); + } + Consume(); // Skip closing quote + + // Extract original lexeme + std::string lexeme(start_pos, cur_ - start_pos); + + // Convert codepoints to UTF-8 string value + std::string value; + for (auto codepoint : codepoints) { + value += CharToUTF8(codepoint); + } + + return {TokenType::StringLiteral, lexeme, value, start_line, start_column}; +} + +// Parse character class. +std::vector EBNFLexer::Impl::ParseCharClassToken() { + std::vector tokens; + + tokens.push_back({TokenType::LBracket, "[", "", cur_line_, cur_column_}); + Consume(); // Skip '[' + + if (Peek() == '^') { + tokens.push_back({TokenType::Caret, "^", "", cur_line_, cur_column_}); + Consume(); + } + + static const std::unordered_map kRegexEscapeChars = { + // clang-format off + {'^', '^'}, {'$', '$'}, {'\\', '\\'}, {'.', '.'}, {'*', '*'}, {'+', '+'}, {'?', '?'}, + {'(', '('}, {')', ')'}, {'[', '['}, {']', ']'}, {'{', '{'}, {'}', '}'}, {'|', '|'}, + {'/', '/'}, {'-', '-'} // clang-format on + }; + + static const std::unordered_set kRegexSpecialEscapes = {'d', 'D', 's', 'S', 'w', 'W'}; + + while (Peek() && Peek() != ']') { + if (Peek() == '\r' || Peek() == '\n') { + ReportLexerError("Character class should not contain newline"); + } else if (Peek() == '-') { + // Handle dash; this dash could be a range expression or a normal dash. + // It will further be handled in EBNFParser::ParseCharClass. + tokens.push_back({TokenType::Dash, "-", "", cur_line_, cur_column_}); + Consume(); + } else if (Peek() == '\\' && kRegexSpecialEscapes.count(Peek(1))) { + // Handle escaped characters with special function + tokens.push_back( + {TokenType::EscapeInCharClass, + std::string(cur_, cur_ + 2), + std::string(cur_ + 1, cur_ + 2), + cur_line_, + cur_column_} + ); + Consume(2); + } else { + // Handle normal characters + auto [codepoint, len] = ParseNextUTF8OrEscaped(cur_, kRegexEscapeChars); + if (codepoint == CharHandlingError::kInvalidUTF8) { + ReportLexerError("Invalid UTF8 sequence"); + } + + if (codepoint == CharHandlingError::kInvalidEscape) { + ReportLexerError("Invalid escape sequence" + std::string(cur_, cur_ + 2)); + } + + tokens.push_back( + {TokenType::CharInCharClass, + std::string(cur_, cur_ + len), + codepoint, + cur_line_, + cur_column_} + ); + Consume(len); + } + } + + if (!Peek()) { + ReportLexerError("Unterminated character class"); + } + + tokens.push_back({TokenType::RBracket, "]", "", cur_line_, cur_column_}); + Consume(); // Skip ']' + + return tokens; +} + +// Parse integer +EBNFLexer::Token EBNFLexer::Impl::ParseIntegerToken() { + int start_line = cur_line_; + int start_column = cur_column_; + const char* start_pos = cur_; + bool is_negative = false; + + if (Peek() == '-') { + is_negative = true; + Consume(); + } else if (Peek() == '+') { + Consume(); + } + + int64_t num = 0; + while (Peek() && isdigit(Peek())) { + num = num * 10 + (Peek() - '0'); + Consume(); + if (num > kMaxIntegerInGrammar) { + ReportLexerError( + "Integer is too large: parsed " + std::to_string(num) + ", max allowed is " + + std::to_string(kMaxIntegerInGrammar) + ); + } + } + + std::string lexeme(start_pos, cur_ - start_pos); + return {TokenType::IntegerLiteral, lexeme, is_negative ? -num : num, start_line, start_column}; +} + +// Get the next token +std::variant> EBNFLexer::Impl::NextToken() { + ConsumeSpace(); // Skip whitespace and comments + + auto start_line = cur_line_; + auto start_column = cur_column_; + + if (!Peek()) { + return EBNFLexer::Token{TokenType::EndOfFile, "", "", start_line, start_column}; + } + + // Determine token type based on current character + switch (Peek()) { + case '(': + if (Peek(1) == '=') { + Consume(2); + return EBNFLexer::Token{TokenType::LookaheadLParen, "(=", "", start_line, start_column}; + } else { + Consume(); + return EBNFLexer::Token{TokenType::LParen, "(", "", start_line, start_column}; + } + case ')': + Consume(); + return EBNFLexer::Token{TokenType::RParen, ")", "", start_line, start_column}; + case '{': + Consume(); + return EBNFLexer::Token{TokenType::LBrace, "{", "", start_line, start_column}; + case '}': + Consume(); + return EBNFLexer::Token{TokenType::RBrace, "}", "", start_line, start_column}; + case '|': + Consume(); + return EBNFLexer::Token{TokenType::Pipe, "|", "", start_line, start_column}; + case ',': + Consume(); + return EBNFLexer::Token{TokenType::Comma, ",", "", start_line, start_column}; + case '*': + Consume(); + return EBNFLexer::Token{TokenType::Star, "*", "", start_line, start_column}; + case '+': + Consume(); + return EBNFLexer::Token{TokenType::Plus, "+", "", start_line, start_column}; + case '?': + Consume(); + return EBNFLexer::Token{TokenType::Question, "?", "", start_line, start_column}; + case '=': + Consume(); + return EBNFLexer::Token{TokenType::Equal, "=", "", start_line, start_column}; + case ':': + if (Peek(1) == ':' && Peek(2) == '=') { + Consume(3); + return EBNFLexer::Token{TokenType::Assign, "::=", "", start_line, start_column}; + } + ReportLexerError("Unexpected character: ':'"); + break; + case '"': + return ParseStringToken(); + case '[': + return ParseCharClassToken(); + default: + if (IsNameChar(*cur_, true)) { + return ParseIdentifierOrBooleanToken(); + } else if (isdigit(*cur_) || *cur_ == '-' || *cur_ == '+') { + return ParseIntegerToken(); + } + + // Unrecognized character, report error + ReportLexerError("Unexpected character: " + std::string(1, *cur_)); + } + + // Should not reach here + XGRAMMAR_UNREACHABLE(); +} + +void EBNFLexer::Impl::ConvertIdentifierToRuleName(std::vector* tokens) { + for (int i = 0; i < static_cast(tokens->size()); ++i) { + if (tokens->at(i).type == TokenType::Assign) { + if (i == 0) { + ReportLexerError( + "Assign should not be the first token", tokens->at(i).line, tokens->at(i).column + ); + } + if (tokens->at(i - 1).type != TokenType::Identifier) { + ReportLexerError( + "Assign should be preceded by an identifier", + tokens->at(i - 1).line, + tokens->at(i - 1).column + ); + } + if (i >= 2 && tokens->at(i - 2).line == tokens->at(i - 1).line) { + ReportLexerError( + "The rule name should be at the beginning of the line", + tokens->at(i - 1).line, + tokens->at(i - 1).column + ); + } + tokens->at(i - 1).type = TokenType::RuleName; + } + } +} + +// Tokenize the entire input and return a vector of tokens +std::vector EBNFLexer::Impl::Tokenize(const std::string& input) { + // Reset position to the beginning + input_ = input; + cur_ = input_.c_str(); + cur_line_ = 1; + cur_column_ = 1; + + // Collect all tokens + std::vector tokens; + + while (true) { + auto token = NextToken(); + + if (auto* token_value = std::get_if(&token)) { + tokens.push_back(*token_value); + // Stop when we reach the end of file + if (token_value->type == TokenType::EndOfFile) { + break; + } + } else { + auto vec = std::get_if>(&token); + XGRAMMAR_DCHECK(vec != nullptr); + tokens.insert(tokens.end(), vec->begin(), vec->end()); + } + } + + ConvertIdentifierToRuleName(&tokens); + + return tokens; +} + +EBNFLexer::EBNFLexer() : pimpl_(std::make_shared()) {} + +std::vector EBNFLexer::Tokenize(const std::string& input) { + return pimpl_->Tokenize(input); +} + +class EBNFParser { + public: + /*! \brief The logic of parsing the grammar string. */ + Grammar Parse( + const std::vector& tokens, + const std::string& root_rule_name, + const int& max_nest_layer = 1000 + ); + + private: + using Rule = Grammar::Impl::Rule; + using SuffixStopInfo = Grammar::Impl::SuffixStopInfo; + using GrammarExprType = Grammar::Impl::GrammarExprType; + using Token = EBNFLexer::Token; + using TokenType = EBNFLexer::TokenType; + + struct ParsedRule { + Rule rule; + SuffixStopInfo suffix_stop_info; + }; + + // Parsing different parts of the grammar + std::string ParseIdentifier(); + int32_t ParseCharClass(); + int32_t ParseString(); + int32_t ParseRuleRef(); + int32_t ParseElement(); + int64_t ParseInteger(); + std::pair ParseRepetitionRange(); + int32_t ParseElementWithQuantifier(); + int32_t ParseLookaheadAssertion(); + int32_t ParseSequence(); + int32_t ParseChoices(); + ParsedRule ParseRule(); + + // Parser for macro + class MacroIR { + public: + struct StringNode; + struct IntegerNode; + struct BooleanNode; + struct IdentifierNode; + struct TupleNode; + + using Node = std::variant; + using NodePtr = std::unique_ptr; + + struct StringNode { + std::string value; + }; + struct IntegerNode { + int64_t value; + }; + struct BooleanNode { + bool value; + }; + struct IdentifierNode { + std::string name; + }; + struct TupleNode { + std::vector elements; + }; + + struct Arguments { + std::vector arguments; + std::unordered_map named_arguments; + }; + }; + MacroIR::Arguments ParseMacroArguments(); + MacroIR::NodePtr ParseMacroValue(); + + int32_t ParseTagDispatch(); + int32_t ParseTokenSet(); + int32_t ParseExcludeToken(); + int32_t ParseTokenTagDispatch(); + int32_t ParseRegexMacro(); + int32_t ParseSubstringMacro(); + + // Helper functions + + // Helper for ParseElementWithQuantifier + int32_t HandleStarQuantifier(int32_t grammar_expr_id); + int32_t HandlePlusQuantifier(int32_t grammar_expr_id); + int32_t HandleQuestionQuantifier(int32_t grammar_expr_id); + + // When parsing, we first find the names of all rules, and build the mapping from name to rule id. + void InitRuleNames(); + + // Consume a token and advance to the next + void Consume(int cnt = 1); + + // Peek at the current token with optional offset + const Token& Peek(int delta = 0) const; + + // Consume token if it matches expected type, otherwise report error + void PeekAndConsume(TokenType type, const std::string& message); + + // Report a parsing error with the given message + [[noreturn]] void ReportParseError(const std::string& msg, int delta_element = 0); + + // The grammar builder + GrammarBuilder builder_; + + // The current token pointer + const Token* current_token_ = nullptr; + + // Tokens from lexer + std::vector tokens_; + + // The current rule name. Help to generate a name for a new rule. + std::string cur_rule_name_; + + // The name of the root rule + std::string root_rule_name_; + + int nest_layer_guard_ = 0; + + int max_nest_layer_ = 1000; // Max nest layer of the grammar + + static const std::unordered_map> kMacroFunctions; +}; + +const std::unordered_map> + EBNFParser::kMacroFunctions = { + {"TagDispatch", [](EBNFParser* parser) { return parser->ParseTagDispatch(); }}, + {"Token", [](EBNFParser* parser) { return parser->ParseTokenSet(); }}, + {"ExcludeToken", [](EBNFParser* parser) { return parser->ParseExcludeToken(); }}, + {"TokenTagDispatch", [](EBNFParser* parser) { return parser->ParseTokenTagDispatch(); }}, + {"Regex", [](EBNFParser* parser) { return parser->ParseRegexMacro(); }}, + {"Substring", [](EBNFParser* parser) { return parser->ParseSubstringMacro(); }}, +}; + +const EBNFParser::Token& EBNFParser::Peek(int delta) const { return *(current_token_ + delta); } + +void EBNFParser::Consume(int cnt) { current_token_ += cnt; } + +void EBNFParser::PeekAndConsume(TokenType type, const std::string& message) { + if (Peek().type != type) { + ReportParseError(message); + } + Consume(); +} + +void EBNFParser::ReportParseError(const std::string& msg, int delta_element) { + XGRAMMAR_DCHECK(current_token_ + delta_element < tokens_.data() + tokens_.size()); + int line_to_print = Peek(delta_element).line; + int column_to_print = Peek(delta_element).column; + XGRAMMAR_LOG(FATAL) << "EBNF parser error at line " + std::to_string(line_to_print) + + ", column " + std::to_string(column_to_print) + ": " + msg; + XGRAMMAR_UNREACHABLE(); +} + +std::string EBNFParser::ParseIdentifier() { + if (Peek().type != TokenType::Identifier) { + ReportParseError("Expect identifier"); + } + std::string identifier = std::any_cast(Peek().value); + Consume(); + return identifier; +} + +int32_t EBNFParser::ParseCharClass() { + PeekAndConsume(TokenType::LBracket, "Expect [ in character class"); + + std::vector elements; + bool is_negated = false; + + if (Peek().type == TokenType::Caret) { + is_negated = true; + Consume(); + } + + while (Peek().type != TokenType::RBracket && Peek().type != TokenType::EndOfFile) { + if (Peek().type == TokenType::EscapeInCharClass) { + ReportParseError("Character class escape is not supported yet in EBNF"); + } + + TCodepoint codepoint; + if (Peek().type == TokenType::CharInCharClass) { + codepoint = std::any_cast(Peek().value); + } else if (Peek().type == TokenType::Dash) { + codepoint = static_cast(static_cast('-')); + } else { + ReportParseError("Unexpected character in character class: " + Peek().lexeme); + } + Consume(); + + if (Peek().type == TokenType::Dash && + (Peek(1).type == TokenType::CharInCharClass || Peek(1).type == TokenType::Dash)) { + // Range expression + TCodepoint codepoint2; + if (Peek(1).type == TokenType::CharInCharClass) { + codepoint2 = std::any_cast(Peek(1).value); + } else { + XGRAMMAR_DCHECK(Peek(1).type == TokenType::Dash); + codepoint2 = static_cast(static_cast('-')); + } + + if (codepoint > codepoint2) { + ReportParseError("Invalid character class: lower bound is larger than upper bound", -1); + } + elements.push_back({codepoint, codepoint2}); + Consume(2); + } else { + // Single character + elements.push_back({codepoint, codepoint}); + } + } + + PeekAndConsume(TokenType::RBracket, "Expect ] in character class"); + + return builder_.AddCharacterClass(elements, is_negated); +} + +int32_t EBNFParser::ParseString() { + if (Peek().type != TokenType::StringLiteral) { + ReportParseError("Expect string literal"); + } + + std::string str_value = std::any_cast(Peek().value); + Consume(); + + if (str_value.empty()) { + return builder_.AddEmptyStr(); + } + + return builder_.AddByteString(str_value); +} + +int32_t EBNFParser::ParseRuleRef() { + std::string name = ParseIdentifier(); + auto rule_id = builder_.GetRuleId(name); + if (rule_id == -1) { + ReportParseError("Rule \"" + name + "\" is not defined", -1); + } + return builder_.AddRuleRef(rule_id); +} + +int32_t EBNFParser::ParseElement() { + if (Peek().type == TokenType::LParen) { + nest_layer_guard_++; + if (nest_layer_guard_ > max_nest_layer_) { + ReportParseError("Nest layer exceeded the maximum limit", -1); + } + Consume(); + if (Peek().type == TokenType::RParen) { + // Special case: ( ) + Consume(); + nest_layer_guard_--; + return builder_.AddEmptyStr(); + } + auto grammar_expr_id = ParseChoices(); + PeekAndConsume(TokenType::RParen, "Expect )"); + nest_layer_guard_--; + return grammar_expr_id; + } else if (Peek().type == TokenType::LBracket) { + return ParseCharClass(); + } else if (Peek().type == TokenType::StringLiteral) { + return ParseString(); + } else if (Peek().type == TokenType::Identifier) { + auto id = std::any_cast(Peek().value); + if (kMacroFunctions.count(id)) { + return kMacroFunctions.at(id)(this); + } else { + return ParseRuleRef(); + } + } else { + ReportParseError("Expect element, but got " + Peek().lexeme); + } +} + +int64_t EBNFParser::ParseInteger() { + if (Peek().type != TokenType::IntegerLiteral) { + ReportParseError("Expect integer, but got " + Peek().lexeme); + } + int64_t num = std::any_cast(Peek().value); + Consume(); + return num; +} + +std::pair EBNFParser::ParseRepetitionRange() { + PeekAndConsume(TokenType::LBrace, "Expect {"); + + int64_t lower = ParseInteger(); + + if (lower < 0) { + ReportParseError("Lower bound cannot be negative", -1); + } + + if (Peek().type == TokenType::Comma) { + Consume(); + if (Peek().type == TokenType::RBrace) { + Consume(); + return {lower, -1}; + } + // The grammar printer emits {n, -1} for unbounded upper bounds, and + // '-' is a valid identifier-start char (IsNameChar), so the lexer + // produces Identifier("-1") rather than IntegerLiteral. Accept it + // as equivalent to {n,}. + if (Peek().type == TokenType::Identifier && Peek().lexeme == "-1") { + Consume(); + PeekAndConsume(TokenType::RBrace, "Expect }"); + return {lower, -1}; + } + int64_t upper = ParseInteger(); + if (upper < lower) { + ReportParseError( + "Lower bound is larger than upper bound: " + std::to_string(lower) + " > " + + std::to_string(upper), + -1 + ); + } + PeekAndConsume(TokenType::RBrace, "Expect }"); + return {lower, upper}; + } else if (Peek().type == TokenType::RBrace) { + Consume(); + return {lower, lower}; + } + + ReportParseError("Expect ',' or '}' in repetition range"); +} + +int32_t EBNFParser::HandleStarQuantifier(int32_t grammar_expr_id) { + Grammar::Impl::GrammarExpr grammar_expr = builder_.GetGrammarExpr(grammar_expr_id); + if (grammar_expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) { + // We have special handling for character class star, e.g. [a-z]* + grammar_expr.type = GrammarBuilder::GrammarExprType::kCharacterClassStar; + // Copy grammar expr because the grammar may change during insertion, and grammar_expr is in the + // grammar, so it may become invalid + std::vector grammar_expr_data(grammar_expr.begin(), grammar_expr.end()); + return builder_.AddGrammarExpr( + {grammar_expr.type, grammar_expr_data.data(), grammar_expr.data_len} + ); + } else { + // For other star quantifiers, we transform it into a rule: + // a* --> rule ::= a rule | "" + auto new_rule_name = builder_.GetNewRuleName(cur_rule_name_); + auto new_rule_id = builder_.AddEmptyRule(new_rule_name); + auto ref_to_new_rule = builder_.AddRuleRef(new_rule_id); + auto new_grammar_expr_id = builder_.AddChoices( + {builder_.AddEmptyStr(), builder_.AddSequence({grammar_expr_id, ref_to_new_rule})} + ); + builder_.UpdateRuleBody(new_rule_id, new_grammar_expr_id); + + // Return the reference to the new rule + return builder_.AddRuleRef(new_rule_id); + } +} + +int32_t EBNFParser::HandlePlusQuantifier(int32_t grammar_expr_id) { + // a+ --> rule ::= a rule | a + auto new_rule_name = builder_.GetNewRuleName(cur_rule_name_); + auto new_rule_id = builder_.AddEmptyRule(new_rule_name); + auto ref_to_new_rule = builder_.AddRuleRef(new_rule_id); + auto new_grammar_expr_id = builder_.AddChoices( + {builder_.AddSequence({grammar_expr_id, ref_to_new_rule}), grammar_expr_id} + ); + builder_.UpdateRuleBody(new_rule_id, new_grammar_expr_id); + + // Return the reference to the new rule + return builder_.AddRuleRef(new_rule_id); +} + +int32_t EBNFParser::HandleQuestionQuantifier(int32_t grammar_expr_id) { + // a? --> rule ::= a | empty + auto new_rule_name = builder_.GetNewRuleName(cur_rule_name_); + auto new_grammar_expr_id = builder_.AddChoices({builder_.AddEmptyStr(), grammar_expr_id}); + auto new_rule_id = builder_.AddRule({new_rule_name, new_grammar_expr_id}); + return builder_.AddRuleRef(new_rule_id); +} + +int32_t EBNFParser::ParseElementWithQuantifier() { + int32_t grammar_expr_id = ParseElement(); + + if (Peek().type == TokenType::Star) { + Consume(); + return HandleStarQuantifier(grammar_expr_id); + } else if (Peek().type == TokenType::Plus) { + Consume(); + return HandlePlusQuantifier(grammar_expr_id); + } else if (Peek().type == TokenType::Question) { + Consume(); + return HandleQuestionQuantifier(grammar_expr_id); + } else if (Peek().type == TokenType::LBrace) { + auto [lower, upper] = ParseRepetitionRange(); + return builder_.AddRepeatFromExpr( + cur_rule_name_, + grammar_expr_id, + static_cast(lower), + upper == -1 ? -1 : static_cast(upper) + ); + } + + return grammar_expr_id; +} + +int32_t EBNFParser::ParseSequence() { + std::vector elements; + + do { + elements.push_back(ParseElementWithQuantifier()); + } while (Peek().type != TokenType::Pipe && Peek().type != TokenType::RParen && + Peek().type != TokenType::LookaheadLParen && Peek().type != TokenType::RuleName && + Peek().type != TokenType::EndOfFile); + + return builder_.AddSequence(elements); +} + +int32_t EBNFParser::ParseChoices() { + std::vector choices; + + choices.push_back(ParseSequence()); + + while (Peek().type == TokenType::Pipe) { + Consume(); + choices.push_back(ParseSequence()); + } + + return builder_.AddChoices(choices); +} + +// Parse macro arguments and return a MacroIR::Arguments structure +EBNFParser::MacroIR::Arguments EBNFParser::ParseMacroArguments() { + MacroIR::Arguments args; + + PeekAndConsume(TokenType::LParen, "Expect ( after macro function name"); + + // Parse arguments + if (Peek().type != TokenType::RParen) { + while (true) { + // Check if it's a named argument (identifier = value) + if (Peek().type == TokenType::Identifier && Peek(1).type == TokenType::Equal) { + std::string name = std::any_cast(Peek().value); + Consume(); // Consume identifier + Consume(); // Consume = + + // Parse the value + args.named_arguments[name] = ParseMacroValue(); + } else { + // Regular positional argument + args.arguments.push_back(ParseMacroValue()); + } + + // Check for comma or end of arguments + if (Peek().type == TokenType::Comma) { + Consume(); + } else if (Peek().type == TokenType::RParen) { + break; + } else { + ReportParseError("Expect , or ) in macro arguments"); + } + } + } + + PeekAndConsume(TokenType::RParen, "Expect ) after macro arguments"); + return args; +} + +// Parse a single macro value (string, integer, boolean, or tuple) +EBNFParser::MacroIR::NodePtr EBNFParser::ParseMacroValue() { + if (Peek().type == TokenType::StringLiteral) { + // String value + std::string value = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::StringNode{value}); + } else if (Peek().type == TokenType::IntegerLiteral) { + // Integer value + int64_t value = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::IntegerNode{value}); + } else if (Peek().type == TokenType::BooleanLiteral) { + // Boolean value + bool value = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::BooleanNode{value}); + } else if (Peek().type == TokenType::Identifier) { + // Identifier value + std::string name = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::IdentifierNode{name}); + } else if (Peek().type == TokenType::LParen) { + // Tuple value. Nested tuples recurse once per layer, so guard the depth the same way the + // ordinary parenthesis path does to avoid a stack overflow on deeply nested input. + nest_layer_guard_++; + if (nest_layer_guard_ > max_nest_layer_) { + ReportParseError("Nest layer exceeded the maximum limit", -1); + } + Consume(); // Consume ( + + MacroIR::TupleNode tuple; + + // Parse tuple elements (supports trailing comma) + if (Peek().type != TokenType::RParen) { + while (true) { + tuple.elements.push_back(ParseMacroValue()); + + if (Peek().type == TokenType::Comma) { + Consume(); + if (Peek().type == TokenType::RParen) { + break; + } + } else if (Peek().type == TokenType::RParen) { + break; + } else { + ReportParseError("Expect , or ) in tuple"); + } + } + } + + Consume(); // Consume ) + nest_layer_guard_--; + return std::make_unique(std::move(tuple)); + } else { + ReportParseError("Expect string, integer, boolean, or tuple in macro argument"); + } +} + +int32_t EBNFParser::ParseTagDispatch() { + Consume(); // Consume TagDispatch operator + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; // Used to report parse errors + + Grammar::Impl::TagDispatch tag_dispatch; + + static const std::unordered_set kValidNamedArgs = { + "loop_after_dispatch", "excludes" + }; + for (const auto& [name, _] : args.named_arguments) { + if (kValidNamedArgs.count(name) == 0) { + ReportParseError("Unknown named argument for TagDispatch: " + name, delta_element); + } + } + + // Positional parameters: ("tag_string", rule_name) — string triggers only + for (const auto& arg : args.arguments) { + auto tuple_node = std::get_if(arg.get()); + if (tuple_node == nullptr) { + ReportParseError("Each tag dispatch element must be a tuple", delta_element); + } + + if (tuple_node->elements.size() != 2) { + ReportParseError("Each tag dispatch element must be a pair (tag, rule)", delta_element); + } + + auto tag_str_node = std::get_if(tuple_node->elements[0].get()); + if (tag_str_node == nullptr || tag_str_node->value.empty()) { + ReportParseError("Tag must be a non-empty string literal", delta_element); + } + + auto rule_name_node = std::get_if(tuple_node->elements[1].get()); + if (rule_name_node == nullptr) { + ReportParseError("Rule reference must be an identifier", delta_element); + } + + auto rule_id = builder_.GetRuleId(rule_name_node->name); + if (rule_id == -1) { + ReportParseError("Rule \"" + rule_name_node->name + "\" is not defined", delta_element); + } + tag_dispatch.tag_rule_pairs.push_back({tag_str_node->value, rule_id}); + } + + // loop_after_dispatch + tag_dispatch.loop_after_dispatch = true; + if (auto it = args.named_arguments.find("loop_after_dispatch"); + it != args.named_arguments.end()) { + auto bool_node = std::get_if(it->second.get()); + if (bool_node == nullptr) { + ReportParseError("loop_after_dispatch must be a boolean literal", delta_element); + } + tag_dispatch.loop_after_dispatch = bool_node->value; + } + + // excludes — string only + if (auto it = args.named_arguments.find("excludes"); it != args.named_arguments.end()) { + auto tuple_node = std::get_if(it->second.get()); + if (tuple_node == nullptr) { + ReportParseError("excludes must be a tuple", delta_element); + } + for (const auto& element : tuple_node->elements) { + auto str_node = std::get_if(element.get()); + if (str_node == nullptr || str_node->value.empty()) { + ReportParseError("Exclude must be a non-empty string literal", delta_element); + } + tag_dispatch.excludes.push_back(str_node->value); + } + } + + // Well-formedness checks: string excludes vs string triggers + for (const auto& excl_str : tag_dispatch.excludes) { + for (const auto& [trigger_str, _] : tag_dispatch.tag_rule_pairs) { + if (trigger_str.rfind(excl_str, 0) == 0) { + ReportParseError( + "Exclude string must not be a prefix of trigger string: " + excl_str, delta_element + ); + } + } + } + + return builder_.AddTagDispatch(tag_dispatch); +} + +int32_t EBNFParser::ParseRegexMacro() { + Consume(); // Consume Regex operator + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; // Used to report parse errors + + if (args.arguments.size() != 1) { + ReportParseError("Regex expects exactly one string argument", delta_element); + } + auto pattern_node = std::get_if(args.arguments[0].get()); + if (pattern_node == nullptr) { + ReportParseError("Regex pattern must be a string literal", delta_element); + } + + bool json_string = false; + for (const auto& [name, _] : args.named_arguments) { + if (name != "json_string" && name != "flags") { + ReportParseError("Regex does not support the named argument " + name, delta_element); + } + } + if (auto it = args.named_arguments.find("json_string"); it != args.named_arguments.end()) { + auto bool_node = std::get_if(it->second.get()); + if (bool_node == nullptr) { + ReportParseError("json_string must be a boolean", delta_element); + } + json_string = bool_node->value; + } + std::string pattern = pattern_node->value; + if (auto it = args.named_arguments.find("flags"); it != args.named_arguments.end()) { + auto flags_node = std::get_if(it->second.get()); + if (flags_node == nullptr) { + ReportParseError("flags must be a string", delta_element); + } + bool case_insensitive = false; + bool dot_all = false; + for (char flag : flags_node->value) { + if (flag == 'i') { + case_insensitive = true; + } else if (flag == 's') { + dot_all = true; + } else if (flag == 'u') { + // XGrammar regular expressions use Unicode codepoint semantics by default. + } else { + ReportParseError( + "regular-expression flag '" + std::string(1, flag) + "' is not supported", delta_element + ); + } + } + // The flags argument opts into the standard dot semantics: '.' does not match '\n' unless + // the 's' flag is given. Without the argument the pattern is stored verbatim, where the + // engine's '.' matches every codepoint. + pattern = RewriteRegexDots(pattern, dot_all); + if (case_insensitive && pattern.compare(0, 4, "(?i)") != 0) { + pattern = "(?i)" + pattern; + } + } + return builder_.AddRegex(pattern, json_string); +} + +int32_t EBNFParser::ParseSubstringMacro() { + Consume(); // Consume Substring identifier + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + if (!args.named_arguments.empty()) { + ReportParseError("Substring() does not accept named arguments", delta_element); + } + + std::vector chunks; + chunks.reserve(args.arguments.size()); + for (const auto& arg : args.arguments) { + auto string_node = std::get_if(arg.get()); + if (string_node == nullptr) { + ReportParseError("Substring() arguments must be strings", delta_element); + } + chunks.push_back(string_node->value); + } + + return builder_.AddSubstring(chunks); +} + +int32_t EBNFParser::ParseTokenSet() { + Consume(); // Consume Token identifier + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + if (!args.named_arguments.empty()) { + ReportParseError("Token() does not accept named arguments", delta_element); + } + + if (args.arguments.empty()) { + ReportParseError("Token() requires at least one integer argument", delta_element); + } + + std::vector token_ids; + for (const auto& arg : args.arguments) { + auto int_node = std::get_if(arg.get()); + if (int_node == nullptr || int_node->value < 0) { + ReportParseError("Token() arguments must be non-negative integers", delta_element); + } + token_ids.push_back(static_cast(int_node->value)); + } + + std::sort(token_ids.begin(), token_ids.end()); + token_ids.erase(std::unique(token_ids.begin(), token_ids.end()), token_ids.end()); + + return builder_.AddTokenSet(token_ids); +} + +int32_t EBNFParser::ParseExcludeToken() { + Consume(); + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + if (!args.named_arguments.empty()) { + ReportParseError("ExcludeToken() does not accept named arguments", delta_element); + } + if (args.arguments.empty()) { + ReportParseError("ExcludeToken() requires at least one integer argument", delta_element); + } + + std::vector token_ids; + for (const auto& arg : args.arguments) { + auto int_node = std::get_if(arg.get()); + if (int_node == nullptr || int_node->value < 0) { + ReportParseError("ExcludeToken() arguments must be non-negative integers", delta_element); + } + token_ids.push_back(static_cast(int_node->value)); + } + std::sort(token_ids.begin(), token_ids.end()); + token_ids.erase(std::unique(token_ids.begin(), token_ids.end()), token_ids.end()); + + return builder_.AddExcludeTokenSet(token_ids); +} + +int32_t EBNFParser::ParseTokenTagDispatch() { + Consume(); + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + Grammar::Impl::TokenTagDispatch ttd; + + static const std::unordered_set kValidNamedArgs = { + "loop_after_dispatch", "excludes" + }; + for (const auto& [name, _] : args.named_arguments) { + if (kValidNamedArgs.count(name) == 0) { + ReportParseError("Unknown named argument for TokenTagDispatch: " + name, delta_element); + } + } + + for (const auto& arg : args.arguments) { + auto tuple_node = std::get_if(arg.get()); + if (tuple_node == nullptr || tuple_node->elements.size() != 2) { + ReportParseError( + "Each TokenTagDispatch element must be a pair (token_id, rule)", delta_element + ); + } + auto id_node = std::get_if(tuple_node->elements[0].get()); + if (id_node == nullptr || id_node->value < 0) { + ReportParseError("Token trigger ID must be a non-negative integer", delta_element); + } + auto rule_node = std::get_if(tuple_node->elements[1].get()); + if (rule_node == nullptr) { + ReportParseError("Rule reference must be an identifier", delta_element); + } + auto rule_id = builder_.GetRuleId(rule_node->name); + if (rule_id == -1) { + ReportParseError("Rule \"" + rule_node->name + "\" is not defined", delta_element); + } + ttd.trigger_rule_pairs.push_back({static_cast(id_node->value), rule_id}); + } + + ttd.loop_after_dispatch = true; + if (auto it = args.named_arguments.find("loop_after_dispatch"); + it != args.named_arguments.end()) { + auto bool_node = std::get_if(it->second.get()); + if (bool_node == nullptr) { + ReportParseError("loop_after_dispatch must be a boolean", delta_element); + } + ttd.loop_after_dispatch = bool_node->value; + } + + if (auto it = args.named_arguments.find("excludes"); it != args.named_arguments.end()) { + auto tuple_node = std::get_if(it->second.get()); + if (tuple_node == nullptr) { + ReportParseError("excludes must be a tuple", delta_element); + } + for (const auto& element : tuple_node->elements) { + auto int_node = std::get_if(element.get()); + if (int_node == nullptr || int_node->value < 0) { + ReportParseError("Exclude token ID must be a non-negative integer", delta_element); + } + ttd.excludes.push_back(static_cast(int_node->value)); + } + } + + for (auto excl_id : ttd.excludes) { + for (const auto& [tid, _] : ttd.trigger_rule_pairs) { + if (tid == excl_id) { + ReportParseError( + "Token trigger ID " + std::to_string(tid) + " must not overlap with exclude token ID", + delta_element + ); + } + } + } + + return builder_.AddTokenTagDispatch(ttd); +} + +int32_t EBNFParser::ParseLookaheadAssertion() { + PeekAndConsume(TokenType::LookaheadLParen, "Expect (= in lookahead assertion"); + auto result = ParseChoices(); + PeekAndConsume(TokenType::RParen, "Expect )"); + return result; +} + +EBNFParser::ParsedRule EBNFParser::ParseRule() { + if (Peek().type != TokenType::RuleName) { + ReportParseError("Expect rule name"); + } + cur_rule_name_ = std::any_cast(Peek().value); + int32_t max_tokens = Peek().max_tokens; + int32_t max_chars = Peek().max_chars; + std::string capture_name = Peek().capture_name; + int32_t capture_hidden_suffix_bytes = Peek().capture_hidden_suffix_bytes; + int32_t capture_hidden_stop_bytes = Peek().capture_hidden_stop_bytes; + int32_t capture_hidden_body_rule_id = Peek().capture_hidden_body_rule_id; + int32_t capture_hidden_marker_rule_id = Peek().capture_hidden_marker_rule_id; + std::string stop_capture_name = Peek().stop_capture_name; + bool is_lazy = Peek().is_lazy; + std::optional temperature = Peek().temperature; + Consume(); + + PeekAndConsume(TokenType::Assign, "Expect ::="); + + auto body_id = ParseChoices(); + + int32_t lookahead_id = -1; + if (Peek().type == TokenType::LookaheadLParen) { + lookahead_id = ParseLookaheadAssertion(); + } + + ParsedRule result; + result.rule = Rule{cur_rule_name_, body_id, lookahead_id}; + result.rule.max_tokens = max_tokens; + result.rule.max_chars = max_chars; + result.rule.capture_name = capture_name; + result.rule.is_lazy = is_lazy; + result.rule.temperature = temperature; + result.suffix_stop_info.hidden_suffix_bytes = capture_hidden_suffix_bytes; + result.suffix_stop_info.hidden_stop_bytes = capture_hidden_stop_bytes; + result.suffix_stop_info.body_rule_id = capture_hidden_body_rule_id; + result.suffix_stop_info.marker_rule_id = capture_hidden_marker_rule_id; + result.suffix_stop_info.stop_capture_name = std::move(stop_capture_name); + return result; +} + +void EBNFParser::InitRuleNames() { + int delta_element = 0; + for (auto& token : tokens_) { + if (token.type == TokenType::RuleName) { + auto name = std::any_cast(token.value); + if (builder_.GetRuleId(name) != -1) { + ReportParseError("Rule \"" + name + "\" is defined multiple times", delta_element); + } + builder_.AddEmptyRule(name); + } + ++delta_element; + } + if (builder_.GetRuleId(root_rule_name_) == -1) { + ReportParseError("The root rule with name \"" + root_rule_name_ + "\" is not found", 0); + } +} + +Grammar EBNFParser::Parse( + const std::vector& tokens, + const std::string& root_rule_name, + const int& max_nest_layer +) { + max_nest_layer_ = max_nest_layer; + nest_layer_guard_ = 0; + tokens_ = tokens; + current_token_ = tokens_.data(); + root_rule_name_ = root_rule_name; + + // First collect rule names + InitRuleNames(); + + // Then parse all the rules + while (Peek().type != TokenType::EndOfFile) { + auto parsed_rule = ParseRule(); + const auto& rule = parsed_rule.rule; + builder_.UpdateRuleBody(rule.name, rule.body_expr_id); + builder_.UpdateLookaheadAssertion(rule.name, rule.lookahead_assertion_id); + builder_.UpdateMaxTokens(rule.name, rule.max_tokens); + builder_.UpdateMaxChars(rule.name, rule.max_chars); + builder_.UpdateCaptureName(rule.name, rule.capture_name); + builder_.UpdateSuffixStopInfo(rule.name, parsed_rule.suffix_stop_info); + builder_.UpdateLazy(rule.name, rule.is_lazy); + builder_.UpdateRuleTemperature(builder_.GetRuleId(rule.name), rule.temperature); + } + + return builder_.Get(root_rule_name); +} + +Grammar ParseEBNF(const std::string& ebnf_string, const std::string& root_rule_name) { + EBNFLexer lexer; + auto tokens = lexer.Tokenize(ebnf_string); + EBNFParser parser; + return parser.Parse(std::move(tokens), root_rule_name); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_parser.h b/third_party/xgrammar/cpp/grammar_parser.h new file mode 100644 index 0000000000..bfb3d8727e --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_parser.h @@ -0,0 +1,107 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_parser.h + * \brief The header for the parser of BNF/EBNF grammar into BNF AST. + */ + +#ifndef XGRAMMAR_GRAMMAR_PARSER_H_ +#define XGRAMMAR_GRAMMAR_PARSER_H_ + +#include + +#include +#include + +namespace xgrammar { + +class EBNFLexer { + public: + // Token types + enum class TokenType { + RuleName, // the name of a rule definition, e.g.: root, rule1 + Identifier, // reference to a rule, or a Macro name, e.g.: root, rule1, TagDispatch + StringLiteral, // e.g.: "tag1", "hello" + BooleanLiteral, // true, false + IntegerLiteral, // 123 + LParen, // ( + RParen, // ) + LBrace, // { + RBrace, // } + Pipe, // | + Comma, // , + EndOfFile, // End of file + + // Symbols and quantifiers + Assign, // ::= + Equal, // = + Star, // * + Plus, // + + Question, // ? + + // Character class + LBracket, // [ + RBracket, // ] + Dash, // - + Caret, // ^ + CharInCharClass, // a character in a character class, e.g. a and z in [a-z]; escaped chars + // with no special meaning are also included, e.g. . in [a\.z] + EscapeInCharClass, // Escaped sequence with special function, e.g. \S in [\S] + + // Special structures + LookaheadLParen, // (= + }; + + // Token structure + struct Token { + TokenType type; + std::string lexeme; // original text + std::any value; // The processed value. Can be a int for integer literal, a string for string + // literal, etc. + int line; + int column; + // The token budget attached to a rule-definition identifier via name[max_tokens=N], or -1. + int32_t max_tokens = -1; + // The character budget attached to a rule-definition identifier via name[max_chars=N], or -1. + int32_t max_chars = -1; + // The capture name attached to a rule-definition identifier via name[capture="x"], or empty. + std::string capture_name = {}; + // Trailing bytes hidden only from the rule's own capture. + int32_t capture_hidden_suffix_bytes = 0; + // Trailing bytes hidden from the rule and every enclosing capture. + int32_t capture_hidden_stop_bytes = 0; + // Helper rule ids used to recover a variable-length suffix/stop marker boundary. + int32_t capture_hidden_body_rule_id = -1; + int32_t capture_hidden_marker_rule_id = -1; + // Capture name for the bytes matched by a suffix/stop marker. + std::string stop_capture_name = {}; + // Whether the identifier is a rule name carrying the [lazy] attribute, e.g. r[lazy] ::= ... + bool is_lazy = false; + // The sampling temperature attached to a rule-definition identifier via name[temperature=T]. + std::optional temperature = std::nullopt; + }; + + EBNFLexer(); + std::vector Tokenize(const std::string& input); + + XGRAMMAR_DEFINE_PIMPL_METHODS(EBNFLexer); +}; + +/*! + * \brief This class parses a BNF/EBNF grammar string into an BNF abstract syntax tree (AST). + * \details This function accepts the EBNF notation defined in the W3C XML Specification + * (https://www.w3.org/TR/xml/#sec-notation), which is a popular standard, with the following + * changes: + * - Using # as comment mark instead of C-style comments + * - Accept C-style unicode escape sequence \u01AB, \U000001AB, \xAB instead of #x0123 + * - Rule A-B (match A and not match B) is not supported yet + * + * See tests/python/serve/json.ebnf for an example. + * \param ebnf_string The grammar string. + * \param root_rule_name The name of the root rule. Default is "root". + * \return The parsed grammar. + */ +Grammar ParseEBNF(const std::string& ebnf_string, const std::string& root_rule_name = "root"); + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_PARSER_H_ diff --git a/third_party/xgrammar/cpp/grammar_printer.cc b/third_party/xgrammar/cpp/grammar_printer.cc new file mode 100644 index 0000000000..a598dfb9af --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_printer.cc @@ -0,0 +1,315 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_printer.cc + */ + +#include "grammar_printer.h" + +#include + +#include +#include +#include + +#include "support/encoding.h" + +namespace xgrammar { + +std::string GrammarPrinter::PrintRule(const Rule& rule, const SuffixStopInfo* suffix_stop_info) { + std::string res = rule.name; + // Print the attributes as one comma-separated bracket group, re-parseable by the EBNF lexer. + if (rule.max_tokens >= 0 || rule.max_chars >= 0 || !rule.capture_name.empty() || + suffix_stop_info != nullptr || rule.is_lazy || rule.temperature.has_value()) { + std::string attributes; + auto append_attribute = [&](const std::string& attribute) { + if (!attributes.empty()) { + attributes += ", "; + } + attributes += attribute; + }; + if (rule.max_tokens >= 0) { + append_attribute("max_tokens=" + std::to_string(rule.max_tokens)); + } + if (rule.max_chars >= 0) { + append_attribute("max_chars=" + std::to_string(rule.max_chars)); + } + if (!rule.capture_name.empty()) { + append_attribute("capture=\"" + rule.capture_name + "\""); + } + if (suffix_stop_info != nullptr && suffix_stop_info->hidden_suffix_bytes > 0) { + append_attribute( + "capture_hidden_suffix_bytes=" + std::to_string(suffix_stop_info->hidden_suffix_bytes) + ); + } + if (suffix_stop_info != nullptr && suffix_stop_info->hidden_stop_bytes > 0) { + append_attribute( + "capture_hidden_stop_bytes=" + std::to_string(suffix_stop_info->hidden_stop_bytes) + ); + } + if (suffix_stop_info != nullptr && suffix_stop_info->body_rule_id >= 0) { + append_attribute( + "capture_hidden_body_rule_id=" + std::to_string(suffix_stop_info->body_rule_id) + ); + append_attribute( + "capture_hidden_marker_rule_id=" + std::to_string(suffix_stop_info->marker_rule_id) + ); + } + if (suffix_stop_info != nullptr && !suffix_stop_info->stop_capture_name.empty()) { + append_attribute("stop_capture=\"" + suffix_stop_info->stop_capture_name + "\""); + } + if (rule.is_lazy) { + append_attribute("lazy"); + } + if (rule.temperature.has_value()) { + std::ostringstream temperature; + temperature << std::setprecision(std::numeric_limits::max_digits10) + << rule.temperature.value(); + append_attribute("temperature=" + temperature.str()); + } + res += "[" + attributes + "]"; + } + res += " ::= " + PrintGrammarExpr(rule.body_expr_id); + if (rule.lookahead_assertion_id != -1) { + res += " (=" + PrintGrammarExpr(rule.lookahead_assertion_id) + ")"; + } + return res; +} + +std::string GrammarPrinter::PrintRule(int32_t rule_id) { + return PrintRule(grammar_->GetRule(rule_id), grammar_->GetSuffixStopInfo(rule_id)); +} + +std::string GrammarPrinter::PrintGrammarExpr(const GrammarExpr& grammar_expr) { + std::string result; + switch (grammar_expr.type) { + case GrammarExprType::kByteString: + return PrintByteString(grammar_expr); + case GrammarExprType::kCharacterClass: + return PrintCharacterClass(grammar_expr); + case GrammarExprType::kCharacterClassStar: + return PrintCharacterClassStar(grammar_expr); + case GrammarExprType::kEmptyStr: + return PrintEmptyStr(grammar_expr); + case GrammarExprType::kRuleRef: + return PrintRuleRef(grammar_expr); + case GrammarExprType::kSequence: + return PrintSequence(grammar_expr); + case GrammarExprType::kChoices: + return PrintChoices(grammar_expr); + case GrammarExprType::kTagDispatch: + return PrintTagDispatch(grammar_expr); + case GrammarExprType::kRepeat: + return PrintRepeat(grammar_expr); + case GrammarExprType::kToken: + return PrintToken(grammar_expr); + case GrammarExprType::kExcludeToken: + return PrintExcludeToken(grammar_expr); + case GrammarExprType::kTokenTagDispatch: + return PrintTokenTagDispatch(grammar_expr); + case GrammarExprType::kRegex: + return PrintRegex(grammar_expr); + case GrammarExprType::kSubstring: + return PrintSubstring(grammar_expr); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected GrammarExpr type: " << static_cast(grammar_expr.type); + XGRAMMAR_UNREACHABLE(); + } +} + +std::string GrammarPrinter::PrintGrammarExpr(int32_t grammar_expr_id) { + return PrintGrammarExpr(grammar_->GetGrammarExpr(grammar_expr_id)); +} + +std::string GrammarPrinter::PrintByteString(const GrammarExpr& grammar_expr) { + std::string internal_str; + internal_str.reserve(grammar_expr.data_len); + for (int i = 0; i < grammar_expr.data_len; ++i) { + internal_str += static_cast(grammar_expr[i]); + } + return "\"" + EscapeString(internal_str) + "\""; +} + +std::string GrammarPrinter::PrintCharacterClass(const GrammarExpr& grammar_expr) { + static const std::unordered_map kCustomEscapeMap = { + {'-', "\\-"}, {']', "\\]"} + }; + std::string result = "["; + bool is_negative = static_cast(grammar_expr[0]); + if (is_negative) { + result += "^"; + } + for (auto i = 1; i < grammar_expr.data_len; i += 2) { + result += EscapeString(grammar_expr[i], kCustomEscapeMap); + if (grammar_expr[i] == grammar_expr[i + 1]) { + continue; + } + result += "-"; + result += EscapeString(grammar_expr[i + 1], kCustomEscapeMap); + } + result += "]"; + return result; +} + +std::string GrammarPrinter::PrintCharacterClassStar(const GrammarExpr& grammar_expr) { + return PrintCharacterClass(grammar_expr) + "*"; +} + +std::string GrammarPrinter::PrintEmptyStr(const GrammarExpr& grammar_expr) { return "\"\""; } + +std::string GrammarPrinter::PrintRuleRef(const GrammarExpr& grammar_expr) { + return grammar_->GetRule(grammar_expr[0]).name; +} + +std::string GrammarPrinter::PrintSequence(const GrammarExpr& grammar_expr) { + std::string result; + result += "("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + result += PrintGrammarExpr(grammar_expr[i]); + if (i + 1 != grammar_expr.data_len) { + result += " "; + } + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintChoices(const GrammarExpr& grammar_expr) { + std::string result; + + result += "("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + result += PrintGrammarExpr(grammar_expr[i]); + if (i + 1 != grammar_expr.data_len) { + result += " | "; + } + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintRegex(const GrammarExpr& grammar_expr) { + std::string result = "Regex(" + PrintString(grammar_->GetRegexString(grammar_expr)); + if (grammar_->GetRegexIsJSONString(grammar_expr)) { + result += ", json_string=true"; + } + return result + ")"; +} + +std::string GrammarPrinter::PrintSubstring(const GrammarExpr& grammar_expr) { + // EscapeString(std::string) stops at embedded NUL bytes, so escape codepoint by codepoint to + // keep NUL chunks (allowed by substring expressions) re-parseable. + auto escape_chunk = [](const std::string& chunk) { + std::string result = "\""; + size_t offset = 0; + while (offset < chunk.size()) { + if (chunk[offset] == '\0') { + result += "\\0"; + ++offset; + continue; + } + auto [codepoint, length] = ParseNextUTF8(chunk.c_str() + offset); + if (codepoint == CharHandlingError::kInvalidUTF8) { + result += EscapeString(static_cast(chunk[offset])); + ++offset; + continue; + } + result += EscapeString(codepoint); + offset += static_cast(length); + } + return result + "\""; + }; + + auto chunks = grammar_->GetSubstringChunks(grammar_expr); + std::string result = "Substring("; + for (size_t i = 0; i < chunks.size(); ++i) { + if (i > 0) { + result += ", "; + } + result += escape_chunk(chunks[i]); + } + return result + ")"; +} + +std::string GrammarPrinter::PrintString(const std::string& str) { + return "\"" + EscapeString(str) + "\""; +} + +std::string GrammarPrinter::PrintBoolean(bool value) { return value ? "true" : "false"; } + +std::string GrammarPrinter::PrintTagDispatch(const GrammarExpr& grammar_expr) { + auto tag_dispatch = grammar_->GetTagDispatch(grammar_expr); + std::string result = "TagDispatch(\n"; + std::string indent = " "; + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + result += indent + "(" + PrintString(trigger) + ", " + grammar_->GetRule(rule_id).name + "),\n"; + } + result += + indent + "loop_after_dispatch=" + PrintBoolean(tag_dispatch.loop_after_dispatch) + ",\n"; + result += indent + "excludes=("; + for (int i = 0; i < static_cast(tag_dispatch.excludes.size()); ++i) { + if (i > 0) result += ", "; + result += PrintString(tag_dispatch.excludes[i]); + } + result += ")\n)"; + return result; +} + +std::string GrammarPrinter::PrintRepeat(const GrammarExpr& grammar_expr) { + int32_t lower_bound = grammar_expr[1]; + int32_t upper_bound = grammar_expr[2]; + std::string result = grammar_->GetRule(grammar_expr[0]).name + "{"; + result += std::to_string(lower_bound); + result += ", "; + result += std::to_string(upper_bound); + result += "}"; + return result; +} + +std::string GrammarPrinter::PrintToken(const GrammarExpr& grammar_expr) { + std::string result = "Token("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + if (i > 0) result += ", "; + result += std::to_string(grammar_expr[i]); + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintExcludeToken(const GrammarExpr& grammar_expr) { + std::string result = "ExcludeToken("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + if (i > 0) result += ", "; + result += std::to_string(grammar_expr[i]); + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintTokenTagDispatch(const GrammarExpr& grammar_expr) { + auto ttd = grammar_->GetTokenTagDispatch(grammar_expr); + std::string result = "TokenTagDispatch(\n"; + std::string indent = " "; + for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + result += + indent + "(" + std::to_string(token_id) + ", " + grammar_->GetRule(rule_id).name + "),\n"; + } + result += indent + "loop_after_dispatch=" + PrintBoolean(ttd.loop_after_dispatch) + ",\n"; + result += indent + "excludes=("; + for (int i = 0; i < static_cast(ttd.excludes.size()); ++i) { + if (i > 0) result += ", "; + result += std::to_string(ttd.excludes[i]); + } + result += ")\n)"; + return result; +} + +std::string GrammarPrinter::ToString() { + std::string result; + int num_rules = grammar_->NumRules(); + for (auto i = 0; i < num_rules; ++i) { + result += PrintRule(i) + "\n"; + } + return result; +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/grammar_printer.h b/third_party/xgrammar/cpp/grammar_printer.h new file mode 100644 index 0000000000..4d5b1ef844 --- /dev/null +++ b/third_party/xgrammar/cpp/grammar_printer.h @@ -0,0 +1,86 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_printer.h + * \brief The header for printing the AST of a BNF grammar. + */ + +#ifndef XGRAMMAR_GRAMMAR_PRINTER_H_ +#define XGRAMMAR_GRAMMAR_PRINTER_H_ + +#include + +#include + +#include "grammar_impl.h" + +namespace xgrammar { + +/*! + * \brief Prints the BNF AST with standard BNF format. + */ +class GrammarPrinter { + private: + using Rule = Grammar::Impl::Rule; + using SuffixStopInfo = Grammar::Impl::SuffixStopInfo; + using GrammarExprType = Grammar::Impl::GrammarExprType; + using GrammarExpr = Grammar::Impl::GrammarExpr; + + public: + /*! + * \brief Constructor. + * \param grammar The grammar to print. + */ + explicit GrammarPrinter(const Grammar& grammar) : grammar_(grammar) {} + + /*! \brief Print the complete grammar. */ + std::string ToString(); + + /*! \brief Print a rule. */ + std::string PrintRule(const Rule& rule, const SuffixStopInfo* suffix_stop_info); + /*! \brief Print a rule corresponding to the given id. */ + std::string PrintRule(int32_t rule_id); + /*! \brief Print a GrammarExpr. */ + std::string PrintGrammarExpr(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr corresponding to the given id. */ + std::string PrintGrammarExpr(int32_t grammar_expr_id); + + private: + /*! \brief Print a GrammarExpr for byte string. */ + std::string PrintByteString(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for character class. */ + std::string PrintCharacterClass(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for a star quantifier of a character class. */ + std::string PrintCharacterClassStar(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for empty string. */ + std::string PrintEmptyStr(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for rule reference. */ + std::string PrintRuleRef(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for grammar_expr sequence. */ + std::string PrintSequence(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for grammar_expr choices. */ + std::string PrintChoices(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for tag dispatch. */ + std::string PrintTagDispatch(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for repeat. */ + std::string PrintRepeat(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for token. */ + std::string PrintToken(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for exclude token. */ + std::string PrintExcludeToken(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for token tag dispatch. */ + std::string PrintTokenTagDispatch(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for regex. */ + std::string PrintRegex(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for substring. */ + std::string PrintSubstring(const GrammarExpr& grammar_expr); + /*! \brief Print a string. */ + std::string PrintString(const std::string& str); + /*! \brief Print a boolean. */ + std::string PrintBoolean(bool value); + + Grammar grammar_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_PRINTER_H_ diff --git a/third_party/xgrammar/cpp/json_schema_converter.cc b/third_party/xgrammar/cpp/json_schema_converter.cc new file mode 100644 index 0000000000..c65529ee0a --- /dev/null +++ b/third_party/xgrammar/cpp/json_schema_converter.cc @@ -0,0 +1,4372 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter.cc + * \brief Implementation of JSONSchemaConverter and related utilities. + */ +#include "json_schema_converter.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_builder.h" +#include "grammar_functor.h" +#include "json_schema_converter_ext.h" +#include "regex_converter.h" +#include "support/json_parse.h" +#include "support/logging.h" + +namespace xgrammar { + +// ==================== Spec ToString implementations ==================== + +std::string IntegerSpec::ToString() const { + return "IntegerSpec{minimum=" + (minimum.has_value() ? std::to_string(*minimum) : "null") + + ", maximum=" + (maximum.has_value() ? std::to_string(*maximum) : "null") + + ", exclusive_minimum=" + + (exclusive_minimum.has_value() ? std::to_string(*exclusive_minimum) : "null") + + ", exclusive_maximum=" + + (exclusive_maximum.has_value() ? std::to_string(*exclusive_maximum) : "null") + + ", multiple_of=" + (multiple_of.has_value() ? std::to_string(*multiple_of) : "null") + "}"; +} + +std::string NumberSpec::ToString() const { + return "NumberSpec{minimum=" + (minimum.has_value() ? std::to_string(*minimum) : "null") + + ", maximum=" + (maximum.has_value() ? std::to_string(*maximum) : "null") + + ", exclusive_minimum=" + + (exclusive_minimum.has_value() ? std::to_string(*exclusive_minimum) : "null") + + ", exclusive_maximum=" + + (exclusive_maximum.has_value() ? std::to_string(*exclusive_maximum) : "null") + "}"; +} + +std::string StringSpec::ToString() const { + return "StringSpec{pattern=" + (pattern.has_value() ? "\"" + *pattern + "\"" : "null") + + ", format=" + (format.has_value() ? "\"" + *format + "\"" : "null") + + ", min_length=" + std::to_string(min_length) + + ", max_length=" + std::to_string(max_length) + "}"; +} + +std::string BooleanSpec::ToString() const { return "BooleanSpec{}"; } + +std::string NullSpec::ToString() const { return "NullSpec{}"; } + +std::string AnySpec::ToString() const { return "AnySpec{}"; } + +std::string ArraySpec::ToString() const { + return "ArraySpec{prefix_items.size()=" + std::to_string(prefix_items.size()) + + ", allow_additional_items=" + (allow_additional_items ? "true" : "false") + + ", additional_items=" + (additional_items ? "SchemaSpec" : "null") + + ", min_items=" + std::to_string(min_items) + ", max_items=" + std::to_string(max_items) + + "}"; +} + +std::string ObjectSpec::ToString() const { + std::string s = + "ObjectSpec{properties.size()=" + std::to_string(properties.size()) + ", properties=["; + for (size_t i = 0; i < properties.size(); ++i) { + if (i != 0) s += ", "; + s += properties[i].name; + } + s += "], pattern_properties.size()=" + std::to_string(pattern_properties.size()) + ", required=["; + bool first = true; + for (const auto& r : required) { + if (!first) s += ", "; + s += r; + first = false; + } + s += + std::string("], allow_additional_properties=") + + (allow_additional_properties ? "true" : "false") + + ", additional_properties_schema=" + (additional_properties_schema ? "SchemaSpec" : "null") + + ", allow_unevaluated_properties=" + (allow_unevaluated_properties ? "true" : "false") + + ", unevaluated_properties_schema=" + (unevaluated_properties_schema ? "SchemaSpec" : "null") + + ", property_names=" + (property_names ? "SchemaSpec" : "null") + + ", min_properties=" + std::to_string(min_properties) + + ", max_properties=" + std::to_string(max_properties) + "}"; + return s; +} + +std::string ConstSpec::ToString() const { return "ConstSpec{json_value=\"" + json_value + "\"}"; } + +std::string EnumSpec::ToString() const { + std::string s = + "EnumSpec{json_values.size()=" + std::to_string(json_values.size()) + ", json_values=["; + for (size_t i = 0; i < json_values.size(); ++i) { + if (i != 0) s += ", "; + s += "\"" + json_values[i] + "\""; + } + s += "]}"; + return s; +} + +std::string RefSpec::ToString() const { return "RefSpec{uri=\"" + uri + "\"}"; } + +std::string AnyOfSpec::ToString() const { + return "AnyOfSpec{options.size()=" + std::to_string(options.size()) + "}"; +} + +std::string OneOfSpec::ToString() const { + return "OneOfSpec{options.size()=" + std::to_string(options.size()) + "}"; +} + +std::string AllOfSpec::ToString() const { + return "AllOfSpec{schemas.size()=" + std::to_string(schemas.size()) + "}"; +} + +std::string TypeArraySpec::ToString() const { + return "TypeArraySpec{type_schemas.size()=" + std::to_string(type_schemas.size()) + "}"; +} + +std::string SchemaSpec::ToString() const { + std::string spec_str; + std::visit([&spec_str](const auto& s) { spec_str = s.ToString(); }, spec); + return "SchemaSpec{spec=" + spec_str + ", cache_key=\"" + cache_key + "\", rule_name_hint=\"" + + rule_name_hint + "\"}"; +} + +// ==================== SchemaParser (Internal) ==================== + +namespace { + +enum class SchemaErrorType : int { + kInvalidSchema = 0, + kUnsatisfiableSchema = 1, + kUnsupportedSchema = 2, +}; + +using SchemaError = TypedError; + +// Unbounded integer multipleOf emits a modulo DFA: states ~= N, transitions ~= 10N. +// Fail closed above the cap to keep generated grammars bounded. +constexpr int64_t kIntegerMultipleOfMax = 1024; +constexpr int64_t kIntegerMultipleOfRangeWidthMax = 10000; + +bool IsMultipleOf(int64_t value, int64_t multiple_of) { return (value % multiple_of) == 0; } + +bool HasMultipleInRange(int64_t start, int64_t end, int64_t multiple_of) { + for (int64_t value = start; value <= end; ++value) { + if (IsMultipleOf(value, multiple_of)) return true; + if (value == std::numeric_limits::max()) break; + } + return false; +} + +constexpr const char* kUnsupportedOneOfMessage = + "oneOf with overlapping or non-provably-disjoint branches cannot be represented exactly; " + "falling back to anyOf semantics"; + +bool IsSchemaAnnotationKey(const std::string& key) { + static const std::unordered_set kAnnotationKeys = { + "title", + "default", + "description", + "examples", + "deprecated", + "readOnly", + "writeOnly", + "$comment", + "$schema", + }; + return kAnnotationKeys.count(key) != 0; +} + +bool HasOnlyKeys( + const picojson::object& schema, const std::unordered_set& allowed_keys +) { + for (const auto& [key, _] : schema) { + if (allowed_keys.count(key) == 0 && !IsSchemaAnnotationKey(key)) { + return false; + } + } + return true; +} + +bool IsSupportedJSONType(const std::string& type) { + static const std::unordered_set kTypes = { + "null", + "boolean", + "object", + "array", + "number", + "string", + "integer", + }; + return kTypes.count(type) != 0; +} + +bool NormalizeTypeSet( + const picojson::value& type_value, std::unordered_set* type_set +) { + if (type_value.is()) { + const auto& type = type_value.get(); + if (!IsSupportedJSONType(type)) { + return false; + } + type_set->insert(type); + return true; + } + if (!type_value.is()) { + return false; + } + + const auto& type_array = type_value.get(); + if (type_array.empty()) { + return false; + } + for (const auto& item : type_array) { + if (!item.is()) { + return false; + } + const auto& type = item.get(); + if (!IsSupportedJSONType(type)) { + return false; + } + type_set->insert(type); + } + return true; +} + +bool IsNumericValue(const picojson::value& value) { + return value.is() || value.is(); +} + +bool IsIntegerValue(const picojson::value& value) { + if (value.is()) { + return true; + } + if (!value.is()) { + return false; + } + double number = value.get(); + return std::isfinite(number) && std::floor(number) == number; +} + +bool JSONValuesMayOverlap(const picojson::value& lhs, const picojson::value& rhs) { + if (IsNumericValue(lhs) || IsNumericValue(rhs)) { + if (!IsNumericValue(lhs) || !IsNumericValue(rhs)) { + return false; + } + if (lhs.is() && rhs.is()) { + return lhs.get() == rhs.get(); + } + return true; + } + if (lhs.is() || rhs.is()) { + return lhs.is() && rhs.is(); + } + if (lhs.is() || rhs.is()) { + return lhs.is() && rhs.is() && lhs.get() == rhs.get(); + } + if (lhs.is() || rhs.is()) { + return lhs.is() && rhs.is() && + lhs.get() == rhs.get(); + } + if (lhs.is() || rhs.is()) { + if (!lhs.is() || !rhs.is()) { + return false; + } + const auto& lhs_array = lhs.get(); + const auto& rhs_array = rhs.get(); + if (lhs_array.size() != rhs_array.size()) { + return false; + } + for (size_t i = 0; i < lhs_array.size(); ++i) { + if (!JSONValuesMayOverlap(lhs_array[i], rhs_array[i])) { + return false; + } + } + return true; + } + if (lhs.is() || rhs.is()) { + if (!lhs.is() || !rhs.is()) { + return false; + } + const auto& lhs_object = lhs.get(); + const auto& rhs_object = rhs.get(); + if (lhs_object.size() != rhs_object.size()) { + return false; + } + for (const auto& [key, lhs_value] : lhs_object) { + auto rhs_it = rhs_object.find(key); + if (rhs_it == rhs_object.end() || !JSONValuesMayOverlap(lhs_value, rhs_it->second)) { + return false; + } + } + return true; + } + return lhs.serialize() == rhs.serialize(); +} + +bool ValueMatchesType(const picojson::value& value, const std::string& type) { + if (type == "null") { + return value.is(); + } + if (type == "boolean") { + return value.is(); + } + if (type == "string") { + return value.is(); + } + if (type == "integer") { + return IsIntegerValue(value); + } + if (type == "number") { + return IsNumericValue(value); + } + if (type == "array") { + return value.is(); + } + if (type == "object") { + return value.is(); + } + return false; +} + +bool IsRangeWidthOverCap(int64_t start, int64_t end, int64_t cap) { + uint64_t cap_u = static_cast(cap); + if (start <= 0 && end >= 0) { + // Count [start, end] inclusively without evaluating -INT64_MIN or overflowing the sum. + uint64_t negative_count = start < 0 ? static_cast(-(start + 1)) + 1 : 0; + if (negative_count > cap_u) return true; + uint64_t remaining = cap_u - negative_count; + if (remaining == 0) return true; + --remaining; // zero + uint64_t positive_count = end > 0 ? static_cast(end) : 0; + return positive_count > remaining; + } + + uint64_t value_count = static_cast(end - start) + 1; + return value_count > cap_u; +} + +// Effective inclusive integer range after folding exclusive bounds into minimum/maximum. A nullopt +// side means that side is unbounded. +struct EffectiveIntegerRange { + std::optional start; + std::optional end; +}; + +// Fold the inclusive [minimum, maximum] bounds together with any exclusive bounds so the stricter +// bound wins on each side. Shared by ParseInteger (range validation) and GenerateInteger (grammar +// emission) so the two can never disagree about the effective range. Precondition: +// exclusive_minimum != INT64_MAX and exclusive_maximum != INT64_MIN (ParseInteger rejects those +// before building the spec), so the +1/-1 below cannot overflow. +EffectiveIntegerRange ComputeEffectiveIntegerRange(const IntegerSpec& spec) { + EffectiveIntegerRange range; + if (spec.minimum.has_value()) { + range.start = spec.minimum; + } + if (spec.exclusive_minimum.has_value()) { + // Smallest integer strictly greater than exclusive_minimum; the larger lower bound wins. + int64_t excl_start = *spec.exclusive_minimum + 1; + range.start = range.start.has_value() ? std::max(*range.start, excl_start) : excl_start; + } + if (spec.maximum.has_value()) { + range.end = spec.maximum; + } + if (spec.exclusive_maximum.has_value()) { + // Largest integer strictly less than exclusive_maximum; the smaller upper bound wins. + int64_t excl_end = *spec.exclusive_maximum - 1; + range.end = range.end.has_value() ? std::min(*range.end, excl_end) : excl_end; + } + return range; +} + +bool TypeSetsOverlap( + const std::unordered_set& lhs, const std::unordered_set& rhs +) { + for (const auto& lhs_type : lhs) { + for (const auto& rhs_type : rhs) { + if (lhs_type == rhs_type) { + return true; + } + if ((lhs_type == "integer" || lhs_type == "number") && + (rhs_type == "integer" || rhs_type == "number")) { + return true; + } + } + } + return false; +} + +bool FiniteValuesOverlap( + const std::vector& lhs, const std::vector& rhs +) { + for (const auto& lhs_value : lhs) { + for (const auto& rhs_value : rhs) { + if (JSONValuesMayOverlap(lhs_value, rhs_value)) { + return true; + } + } + } + return false; +} + +bool FiniteValuesOverlapTypeSet( + const std::vector& values, const std::unordered_set& type_set +) { + for (const auto& value : values) { + if (IsNumericValue(value) && (type_set.count("integer") || type_set.count("number"))) { + return true; + } + for (const auto& type : type_set) { + if (ValueMatchesType(value, type)) { + return true; + } + } + } + return false; +} + +bool TryGetFiniteValues(const picojson::object& schema, std::vector* values) { + if (schema.count("const")) { + values->push_back(schema.at("const")); + return true; + } + if (schema.count("enum")) { + if (!schema.at("enum").is()) { + return false; + } + const auto& enum_values = schema.at("enum").get(); + if (enum_values.empty()) { + return false; + } + values->insert(values->end(), enum_values.begin(), enum_values.end()); + return true; + } + return false; +} + +struct OneOfArmProof { + enum class Kind { kTypeSet, kFiniteValues }; + + Kind kind; + std::unordered_set type_set; + std::vector finite_values; +}; + +std::optional ClassifyTypeOrFiniteOneOfArm(const picojson::value& option) { + if (!option.is()) { + return std::nullopt; + } + const auto& schema = option.get(); + + if (schema.count("$ref") || schema.count("anyOf") || schema.count("allOf") || + schema.count("oneOf")) { + return std::nullopt; + } + + std::vector finite_values; + if (TryGetFiniteValues(schema, &finite_values)) { + OneOfArmProof proof; + proof.kind = OneOfArmProof::Kind::kFiniteValues; + proof.finite_values = std::move(finite_values); + return proof; + } + + if (!schema.count("type") || !HasOnlyKeys(schema, {"type"})) { + return std::nullopt; + } + + std::unordered_set type_set; + if (!NormalizeTypeSet(schema.at("type"), &type_set)) { + return std::nullopt; + } + if (type_set.count("object")) { + return std::nullopt; + } + + OneOfArmProof proof; + proof.kind = OneOfArmProof::Kind::kTypeSet; + proof.type_set = std::move(type_set); + return proof; +} + +bool OneOfArmProofsAreDisjoint(const OneOfArmProof& lhs, const OneOfArmProof& rhs) { + if (lhs.kind == OneOfArmProof::Kind::kTypeSet && rhs.kind == OneOfArmProof::Kind::kTypeSet) { + return !TypeSetsOverlap(lhs.type_set, rhs.type_set); + } + if (lhs.kind == OneOfArmProof::Kind::kFiniteValues && + rhs.kind == OneOfArmProof::Kind::kFiniteValues) { + return !FiniteValuesOverlap(lhs.finite_values, rhs.finite_values); + } + if (lhs.kind == OneOfArmProof::Kind::kFiniteValues && rhs.kind == OneOfArmProof::Kind::kTypeSet) { + return !FiniteValuesOverlapTypeSet(lhs.finite_values, rhs.type_set); + } + return !FiniteValuesOverlapTypeSet(rhs.finite_values, lhs.type_set); +} + +std::optional> GetDiscriminatorValues( + const picojson::value& option, const std::string& discriminator_key +) { + if (!option.is()) { + return std::nullopt; + } + const auto& schema = option.get(); + if (schema.count("$ref") || schema.count("anyOf") || schema.count("allOf") || + schema.count("oneOf")) { + return std::nullopt; + } + if (!schema.count("type") || !schema.at("type").is() || + schema.at("type").get() != "object") { + return std::nullopt; + } + if (!schema.count("required") || !schema.at("required").is()) { + return std::nullopt; + } + + bool requires_discriminator = false; + for (const auto& required_key : schema.at("required").get()) { + if (!required_key.is()) { + return std::nullopt; + } + if (required_key.get() == discriminator_key) { + requires_discriminator = true; + } + } + if (!requires_discriminator) { + return std::nullopt; + } + + if (!schema.count("properties") || !schema.at("properties").is()) { + return std::nullopt; + } + const auto& properties = schema.at("properties").get(); + auto property_it = properties.find(discriminator_key); + if (property_it == properties.end() || !property_it->second.is()) { + return std::nullopt; + } + + std::vector values; + if (!TryGetFiniteValues(property_it->second.get(), &values)) { + return std::nullopt; + } + return values; +} + +std::vector GetDiscriminatorCandidates(const picojson::value& option) { + std::vector candidates; + if (!option.is()) { + return candidates; + } + const auto& schema = option.get(); + if (!schema.count("required") || !schema.at("required").is() || + !schema.count("properties") || !schema.at("properties").is()) { + return candidates; + } + const auto& properties = schema.at("properties").get(); + for (const auto& required_key : schema.at("required").get()) { + if (!required_key.is()) { + continue; + } + const auto& key = required_key.get(); + auto property_it = properties.find(key); + if (property_it == properties.end() || !property_it->second.is()) { + continue; + } + std::vector values; + if (TryGetFiniteValues(property_it->second.get(), &values)) { + candidates.push_back(key); + } + } + return candidates; +} + +bool TryProveStrictDiscriminatorOneOf(const picojson::array& options) { + if (options.empty()) { + return false; + } + + for (const auto& discriminator_key : GetDiscriminatorCandidates(options.front())) { + std::vector> branch_values; + bool all_branches_have_key = true; + for (const auto& option : options) { + auto values = GetDiscriminatorValues(option, discriminator_key); + if (!values.has_value()) { + all_branches_have_key = false; + break; + } + branch_values.push_back(std::move(values.value())); + } + if (!all_branches_have_key) { + continue; + } + + bool pairwise_disjoint = true; + for (size_t i = 0; i < branch_values.size() && pairwise_disjoint; ++i) { + for (size_t j = i + 1; j < branch_values.size(); ++j) { + if (FiniteValuesOverlap(branch_values[i], branch_values[j])) { + pairwise_disjoint = false; + break; + } + } + } + if (pairwise_disjoint) { + return true; + } + } + return false; +} + +bool TryProveTypeOrFiniteOneOf(const picojson::array& options) { + std::vector proofs; + proofs.reserve(options.size()); + for (const auto& option : options) { + auto proof = ClassifyTypeOrFiniteOneOfArm(option); + if (!proof.has_value()) { + return false; + } + proofs.push_back(std::move(proof.value())); + } + + for (size_t i = 0; i < proofs.size(); ++i) { + for (size_t j = i + 1; j < proofs.size(); ++j) { + if (!OneOfArmProofsAreDisjoint(proofs[i], proofs[j])) { + return false; + } + } + } + return true; +} + +bool TryProvePairwiseDisjointOneOf(const picojson::array& options) { + return TryProveStrictDiscriminatorOneOf(options) || TryProveTypeOrFiniteOneOf(options); +} + +/*! + * \brief Parser for JSON Schema, converts JSON Schema to SchemaSpec intermediate representation. + */ +class SchemaParser { + public: + struct Config { + bool strict_mode = false; + JSONFormat json_format; + }; + + explicit SchemaParser(const picojson::value& root_schema, const Config& config) + : config_(config), root_schema_(root_schema) {} + + Result Parse( + const picojson::value& schema, + const std::string& rule_name_hint = "root", + std::optional default_type = std::nullopt + ); + + const picojson::value& GetRootSchema() const { return root_schema_; } + bool IsStrictMode() const { return config_.strict_mode; } + + Result ResolveRef( + const std::string& uri, const std::string& rule_name_hint + ); + + private: + Result ParseInteger(const picojson::object& schema); + Result ParseNumber(const picojson::object& schema); + Result ParseString(const picojson::object& schema); + Result ParseBoolean(const picojson::object& schema); + Result ParseNull(const picojson::object& schema); + Result ParseArray(const picojson::object& schema); + Result ParseObject(const picojson::object& schema); + Result ParseConst(const picojson::object& schema); + Result ParseEnum(const picojson::object& schema); + Result ParseRef(const picojson::object& schema); + Result ParseAnyOf( + const picojson::object& schema, const std::string& keyword + ); + Result ParseOneOf(const picojson::object& schema); + Result ParseAllOf(const picojson::object& schema); + Result ParseTypeArray( + const picojson::object& schema, const std::string& rule_name_hint + ); + + std::string ComputeCacheKey(const picojson::value& schema); + + static void WarnUnsupportedKeywords( + const picojson::object& schema, const std::vector& keywords, bool verbose = false + ); + + Config config_; + picojson::value root_schema_; + std::unordered_map ref_cache_; + std::unordered_map schema_cache_; +}; + +std::string SchemaParser::ComputeCacheKey(const picojson::value& schema) { + static const std::unordered_set kSkippedKeys = { + "title", + "default", + "description", + "examples", + "deprecated", + "readOnly", + "writeOnly", + "$comment", + "$schema", + }; + + if (schema.is()) { + std::string result = "{"; + std::vector> sorted_kv; + for (const auto& kv : schema.get()) { + if (kSkippedKeys.count(kv.first) == 0) { + sorted_kv.push_back(kv); + } + } + std::sort(sorted_kv.begin(), sorted_kv.end(), [](const auto& lhs, const auto& rhs) { + return lhs.first < rhs.first; + }); + int64_t idx = 0; + for (const auto& [key, value] : sorted_kv) { + if (idx != 0) { + result += ","; + } + ++idx; + result += "\"" + key + "\":" + ComputeCacheKey(value); + } + return result + "}"; + } else if (schema.is()) { + std::string result = "["; + int64_t idx = 0; + for (const auto& item : schema.get()) { + if (idx != 0) { + result += ","; + } + ++idx; + result += ComputeCacheKey(item); + } + return result + "]"; + } + return schema.serialize(false); +} + +void SchemaParser::WarnUnsupportedKeywords( + const picojson::object& schema, const std::vector& keywords, bool verbose +) { + if (!verbose) { + return; + } + for (const auto& keyword : keywords) { + if (schema.find(keyword) != schema.end()) { + XGRAMMAR_LOG(WARNING) << "Keyword " << keyword << " is not supported"; + } + } +} + +Result SchemaParser::Parse( + const picojson::value& schema, + const std::string& rule_name_hint, + std::optional default_type +) { + std::string cache_key = ComputeCacheKey(schema); + if (schema_cache_.count(cache_key)) { + return ResultOk(schema_cache_[cache_key]); + } + + if (schema.is()) { + if (!schema.get()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "Schema 'false' cannot accept any value" + ); + } + auto spec = SchemaSpec::Make(AnySpec{}, cache_key, rule_name_hint); + schema_cache_[cache_key] = spec; + return ResultOk(spec); + } + + if (!schema.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, + "Schema should be an object or bool, but got " + schema.serialize(false) + ); + } + + const auto& schema_obj = schema.get(); + WarnUnsupportedKeywords( + schema_obj, {"not", "if", "then", "else", "dependentRequired", "dependentSchemas"} + ); + + SchemaSpecPtr result; + + if (schema_obj.count("$ref")) { + auto ref_result = ParseRef(schema_obj); + if (ref_result.IsErr()) return ResultErr(std::move(ref_result).UnwrapErr()); + auto ref_spec = std::move(ref_result).Unwrap(); + result = SchemaSpec::Make(std::move(ref_spec), cache_key, rule_name_hint); + } else if (schema_obj.count("const")) { + auto const_result = ParseConst(schema_obj); + if (const_result.IsErr()) return ResultErr(std::move(const_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(const_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("enum")) { + auto enum_result = ParseEnum(schema_obj); + if (enum_result.IsErr()) return ResultErr(std::move(enum_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(enum_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("anyOf")) { + auto anyof_result = ParseAnyOf(schema_obj, "anyOf"); + if (anyof_result.IsErr()) return ResultErr(std::move(anyof_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(anyof_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("oneOf")) { + auto oneof_result = ParseOneOf(schema_obj); + if (oneof_result.IsErr()) { + if (oneof_result.ErrRef().Type() != SchemaErrorType::kUnsupportedSchema) { + return ResultErr(std::move(oneof_result).UnwrapErr()); + } + XGRAMMAR_LOG(WARNING) << oneof_result.ErrRef().what(); + auto anyof_result = ParseAnyOf(schema_obj, "oneOf"); + if (anyof_result.IsErr()) return ResultErr(std::move(anyof_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(anyof_result).Unwrap(), cache_key, rule_name_hint); + } else { + result = SchemaSpec::Make(std::move(oneof_result).Unwrap(), cache_key, rule_name_hint); + } + } else if (schema_obj.count("allOf")) { + auto allof_result = ParseAllOf(schema_obj); + if (allof_result.IsErr()) return ResultErr(std::move(allof_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(allof_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("type") || default_type.has_value()) { + if (schema_obj.count("type") && schema_obj.at("type").is()) { + auto type_array_result = ParseTypeArray(schema_obj, rule_name_hint); + if (type_array_result.IsErr()) return ResultErr(std::move(type_array_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(type_array_result).Unwrap(), cache_key, rule_name_hint); + } else { + if (schema_obj.count("type") && !schema_obj.at("type").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Type should be a string"); + } + const std::string& type = schema_obj.count("type") ? schema_obj.at("type").get() + : default_type.value(); + if (type == "integer") { + auto int_result = ParseInteger(schema_obj); + if (int_result.IsErr()) return ResultErr(std::move(int_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(int_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "number") { + auto num_result = ParseNumber(schema_obj); + if (num_result.IsErr()) return ResultErr(std::move(num_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(num_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "string") { + auto str_result = ParseString(schema_obj); + if (str_result.IsErr()) return ResultErr(std::move(str_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(str_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "boolean") { + auto bool_result = ParseBoolean(schema_obj); + if (bool_result.IsErr()) return ResultErr(std::move(bool_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(bool_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "null") { + auto null_result = ParseNull(schema_obj); + if (null_result.IsErr()) return ResultErr(std::move(null_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(null_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "array") { + auto array_result = ParseArray(schema_obj); + if (array_result.IsErr()) return ResultErr(std::move(array_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(array_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "object") { + auto obj_result = ParseObject(schema_obj); + if (obj_result.IsErr()) return ResultErr(std::move(obj_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(obj_result).Unwrap(), cache_key, rule_name_hint); + } else { + return ResultErr( + SchemaErrorType::kInvalidSchema, "Unsupported type \"" + type + "\"" + ); + } + } + } else if (schema_obj.count("properties") || schema_obj.count("additionalProperties") || + schema_obj.count("unevaluatedProperties")) { + auto obj_result = ParseObject(schema_obj); + if (obj_result.IsErr()) return ResultErr(std::move(obj_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(obj_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("items") || schema_obj.count("prefixItems") || + schema_obj.count("unevaluatedItems")) { + auto array_result = ParseArray(schema_obj); + if (array_result.IsErr()) return ResultErr(std::move(array_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(array_result).Unwrap(), cache_key, rule_name_hint); + } else { + result = SchemaSpec::Make(AnySpec{}, cache_key, rule_name_hint); + } + + schema_cache_[cache_key] = result; + return ResultOk(result); +} + +Result SchemaParser::ParseInteger(const picojson::object& schema) { + IntegerSpec spec; + + auto checkAndConvertIntegerBound = [](const picojson::value& value + ) -> Result { + if (!value.is() && !value.is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + if (value.is()) return ResultOk(value.get()); + double val = value.get(); + if (val != std::floor(val)) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "Integer constraint must be a whole number" + ); + } + static const double PROBLEMATIC_MIN = -9223372036854776000.0; + static const double PROBLEMATIC_MAX = 9223372036854776000.0; + if (val == PROBLEMATIC_MIN) { + XGRAMMAR_CHECK(false + ) << "Integer exceeds minimum limit due to precision loss at 64-bit boundary"; + } + + if (val == PROBLEMATIC_MAX) { + XGRAMMAR_CHECK(false + ) << "Integer exceeds maximum limit due to precision loss at 64-bit boundary"; + } + static const double MAX_INT64_AS_DOUBLE = + static_cast(std::numeric_limits::max()); + static const double MIN_INT64_AS_DOUBLE = + static_cast(std::numeric_limits::min()); + XGRAMMAR_CHECK(val <= MAX_INT64_AS_DOUBLE) << "Integer exceeds maximum limit"; + XGRAMMAR_CHECK(val >= MIN_INT64_AS_DOUBLE) << "Integer exceeds minimum limit"; + return ResultOk(static_cast(val)); + }; + + auto checkAndConvertMultipleOf = [](const picojson::value& value + ) -> Result { + double val; + if (value.is()) { + val = static_cast(value.get()); + } else if (value.is()) { + val = value.get(); + } else { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + if (val <= 0) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "multipleOf must be greater than 0" + ); + } + if (val != std::floor(val)) { + return ResultErr( + SchemaErrorType::kUnsupportedSchema, "multipleOf for type:integer must be an integer" + ); + } + if (val > static_cast(kIntegerMultipleOfMax)) { + return ResultErr( + SchemaErrorType::kUnsupportedSchema, + "multipleOf for type:integer must be > 0 and <= " + std::to_string(kIntegerMultipleOfMax) + ); + } + return ResultOk(static_cast(val)); + }; + + if (schema.count("multipleOf")) { + auto result = checkAndConvertMultipleOf(schema.at("multipleOf")); + if (result.IsErr()) { + if (result.ErrRef().Type() != SchemaErrorType::kUnsupportedSchema) { + return ResultErr(std::move(result).UnwrapErr()); + } + XGRAMMAR_LOG(WARNING) << result.ErrRef().what() << "; ignoring multipleOf"; + } else { + spec.multiple_of = std::move(result).Unwrap(); + } + } + if (schema.count("minimum")) { + auto result = checkAndConvertIntegerBound(schema.at("minimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.minimum = std::move(result).Unwrap(); + } + if (schema.count("maximum")) { + auto result = checkAndConvertIntegerBound(schema.at("maximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.maximum = std::move(result).Unwrap(); + } + if (schema.count("exclusiveMinimum")) { + auto result = checkAndConvertIntegerBound(schema.at("exclusiveMinimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + int64_t val = std::move(result).Unwrap(); + if (val == std::numeric_limits::max()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "exclusiveMinimum would cause integer overflow" + ); + } + spec.exclusive_minimum = val; + } + if (schema.count("exclusiveMaximum")) { + auto result = checkAndConvertIntegerBound(schema.at("exclusiveMaximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + int64_t val = std::move(result).Unwrap(); + if (val == std::numeric_limits::min()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "exclusiveMaximum would cause integer underflow" + ); + } + spec.exclusive_maximum = val; + } + + EffectiveIntegerRange effective_range = ComputeEffectiveIntegerRange(spec); + int64_t effective_min = effective_range.start.value_or(std::numeric_limits::min()); + int64_t effective_max = effective_range.end.value_or(std::numeric_limits::max()); + if (effective_min > effective_max) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "Invalid range: minimum greater than maximum" + ); + } + if (spec.multiple_of.has_value()) { + bool has_lower_bound = spec.minimum.has_value() || spec.exclusive_minimum.has_value(); + bool has_upper_bound = spec.maximum.has_value() || spec.exclusive_maximum.has_value(); + if (has_lower_bound || has_upper_bound) { + if (!has_lower_bound || !has_upper_bound || + IsRangeWidthOverCap(effective_min, effective_max, kIntegerMultipleOfRangeWidthMax)) { + XGRAMMAR_LOG(WARNING + ) << "range + multipleOf combination not yet supported; ignoring multipleOf"; + spec.multiple_of.reset(); + return ResultOk(std::move(spec)); + } + if (!HasMultipleInRange(effective_min, effective_max, *spec.multiple_of)) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "range contains no multipleOf value" + ); + } + } + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseNumber(const picojson::object& schema) { + if (schema.count("multipleOf")) { + const auto& value = schema.at("multipleOf"); + if (!value.is() && !value.is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + double multiple_of = + value.is() ? static_cast(value.get()) : value.get(); + if (multiple_of <= 0) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "multipleOf must be greater than 0" + ); + } + XGRAMMAR_LOG(WARNING) << "multipleOf is not supported for type:number; ignoring multipleOf"; + } + NumberSpec spec; + + auto getDouble = [](const picojson::value& value) -> Result { + if (!value.is() && !value.is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + return ResultOk(value.get()); + }; + + if (schema.count("minimum")) { + auto result = getDouble(schema.at("minimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.minimum = std::move(result).Unwrap(); + } + if (schema.count("maximum")) { + auto result = getDouble(schema.at("maximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.maximum = std::move(result).Unwrap(); + } + if (schema.count("exclusiveMinimum")) { + auto result = getDouble(schema.at("exclusiveMinimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.exclusive_minimum = std::move(result).Unwrap(); + } + if (schema.count("exclusiveMaximum")) { + auto result = getDouble(schema.at("exclusiveMaximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.exclusive_maximum = std::move(result).Unwrap(); + } + + // The range is empty if any lower bound conflicts with any upper bound. An + // exclusive bound also rules out equality, so it uses ">=" instead of ">". + auto empty = []() { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "Invalid range: empty range" + ); + }; + + // minimum (x >= min) vs maximum (x <= max). + if (spec.minimum && spec.maximum && *spec.minimum > *spec.maximum) { + return empty(); + } + // minimum (x >= min) vs exclusiveMaximum (x < exclMax). + if (spec.minimum && spec.exclusive_maximum && *spec.minimum >= *spec.exclusive_maximum) { + return empty(); + } + // exclusiveMinimum (x > exclMin) vs maximum (x <= max). + if (spec.exclusive_minimum && spec.maximum && *spec.exclusive_minimum >= *spec.maximum) { + return empty(); + } + // exclusiveMinimum (x > exclMin) vs exclusiveMaximum (x < exclMax). + if (spec.exclusive_minimum && spec.exclusive_maximum && + *spec.exclusive_minimum >= *spec.exclusive_maximum) { + return empty(); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseString(const picojson::object& schema) { + StringSpec spec; + if (schema.count("format")) spec.format = schema.at("format").get(); + if (schema.count("pattern")) spec.pattern = schema.at("pattern").get(); + // Lengths become int32 repetition bounds. A minimum beyond int32 can never be satisfied; a + // maximum beyond it is unbounded in practice. Neither may wrap around when converted. + constexpr int64_t kMaxBound = std::numeric_limits::max(); + if (schema.count("minLength")) { + if (!schema.at("minLength").is() || + schema.at("minLength").get() > kMaxBound) { + return ResultErr( + SchemaErrorType::kInvalidSchema, + "minLength must be an integer not exceeding " + std::to_string(kMaxBound) + ); + } + spec.min_length = static_cast(schema.at("minLength").get()); + } + if (schema.count("maxLength")) { + if (!schema.at("maxLength").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "maxLength must be an integer" + ); + } + if (schema.at("maxLength").get() <= kMaxBound) { + spec.max_length = static_cast(schema.at("maxLength").get()); + } + } + if (spec.max_length != -1 && spec.min_length > spec.max_length) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minLength " + std::to_string(spec.min_length) + " is greater than maxLength " + + std::to_string(spec.max_length) + ); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseBoolean(const picojson::object&) { + return ResultOk(BooleanSpec{}); +} + +Result SchemaParser::ParseNull(const picojson::object&) { + return ResultOk(NullSpec{}); +} + +Result SchemaParser::ParseArray(const picojson::object& schema) { + WarnUnsupportedKeywords(schema, {"uniqueItems", "contains", "minContains", "maxContains"}); + ArraySpec spec; + + if (schema.count("prefixItems")) { + if (!schema.at("prefixItems").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "prefixItems must be an array" + ); + } + for (const auto& item : schema.at("prefixItems").get()) { + if (item.is() && !item.get()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "prefixItems contains false" + ); + } else if (!item.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "prefixItems must be an array of objects or booleans" + ); + } + auto item_result = Parse(item, "prefix_item"); + if (item_result.IsErr()) return ResultErr(std::move(item_result).UnwrapErr()); + spec.prefix_items.push_back(std::move(item_result).Unwrap()); + } + } + + if (schema.count("items")) { + auto items_value = schema.at("items"); + if (!items_value.is() && !items_value.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "items must be a boolean or an object" + ); + } + if (items_value.is() && !items_value.get()) { + spec.allow_additional_items = false; + } else { + spec.allow_additional_items = true; + auto items_result = Parse(items_value, "item"); + if (items_result.IsErr()) return ResultErr(std::move(items_result).UnwrapErr()); + spec.additional_items = std::move(items_result).Unwrap(); + } + } else if (schema.count("unevaluatedItems")) { + auto unevaluated_items_value = schema.at("unevaluatedItems"); + if (!unevaluated_items_value.is() && !unevaluated_items_value.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "unevaluatedItems must be a boolean or an object" + ); + } + if (unevaluated_items_value.is() && !unevaluated_items_value.get()) { + spec.allow_additional_items = false; + } else { + spec.allow_additional_items = true; + auto items_result = Parse(unevaluated_items_value, "unevaluated_item"); + if (items_result.IsErr()) return ResultErr(std::move(items_result).UnwrapErr()); + spec.additional_items = std::move(items_result).Unwrap(); + } + } else if (!config_.strict_mode) { + spec.allow_additional_items = true; + spec.additional_items = SchemaSpec::Make(AnySpec{}, "", "any"); + } else { + spec.allow_additional_items = false; + } + + if (schema.count("minItems")) { + if (!schema.at("minItems").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "minItems must be an integer"); + } + spec.min_items = std::max(static_cast(0), schema.at("minItems").get()); + } + if (schema.count("minContains")) { + if (!schema.at("minContains").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "minContains must be an integer" + ); + } + spec.min_items = std::max(spec.min_items, schema.at("minContains").get()); + } + if (schema.count("maxItems")) { + if (!schema.at("maxItems").is() || schema.at("maxItems").get() < 0) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "maxItems must be a non-negative integer" + ); + } + spec.max_items = schema.at("maxItems").get(); + } + // Item counts become int32 repetition bounds, see ParseString for the rationale. + constexpr int64_t kMaxBound = std::numeric_limits::max(); + if (spec.min_items > kMaxBound) { + return ResultErr( + SchemaErrorType::kInvalidSchema, + "minItems and minContains must not exceed " + std::to_string(kMaxBound) + ); + } + if (spec.max_items > kMaxBound) { + spec.max_items = -1; + } + + if (spec.max_items != -1 && spec.min_items > spec.max_items) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minItems is greater than maxItems: " + std::to_string(spec.min_items) + " > " + + std::to_string(spec.max_items) + ); + } + if (spec.max_items != -1 && spec.max_items < static_cast(spec.prefix_items.size())) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "maxItems is less than the number of prefixItems: " + std::to_string(spec.max_items) + + " < " + std::to_string(spec.prefix_items.size()) + ); + } + if (!spec.allow_additional_items) { + int64_t prefix_size = static_cast(spec.prefix_items.size()); + if (prefix_size < spec.min_items) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minItems is greater than the number of prefixItems, but additional items are not " + "allowed: " + + std::to_string(spec.min_items) + " > " + std::to_string(prefix_size) + ); + } + if (spec.max_items != -1 && prefix_size > spec.max_items) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "maxItems is less than the number of prefixItems, but additional items are not " + "allowed: " + + std::to_string(spec.max_items) + " < " + std::to_string(prefix_size) + ); + } + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseObject(const picojson::object& schema) { + ObjectSpec spec; + + if (schema.count("properties")) { + if (!schema.at("properties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "properties must be an object" + ); + } + auto properties_obj = schema.at("properties").get(); + for (const auto& key : properties_obj.ordered_keys()) { + auto prop_result = Parse(properties_obj.at(key), key); + if (prop_result.IsErr()) return ResultErr(std::move(prop_result).UnwrapErr()); + spec.properties.push_back({key, std::move(prop_result).Unwrap()}); + } + } + + if (schema.count("required")) { + if (!schema.at("required").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "required must be an array"); + } + for (const auto& req : schema.at("required").get()) { + spec.required.insert(req.get()); + } + } + + if (schema.count("patternProperties")) { + if (!schema.at("patternProperties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "patternProperties must be an object" + ); + } + auto pattern_props = schema.at("patternProperties").get(); + for (const auto& key : pattern_props.ordered_keys()) { + auto prop_result = Parse(pattern_props.at(key), "pattern_prop"); + if (prop_result.IsErr()) return ResultErr(std::move(prop_result).UnwrapErr()); + spec.pattern_properties.push_back({key, std::move(prop_result).Unwrap()}); + } + } + + if (schema.count("propertyNames")) { + if (!schema.at("propertyNames").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "propertyNames must be an object" + ); + } + auto property_names_obj = schema.at("propertyNames").get(); + if (property_names_obj.count("type") && property_names_obj.at("type").is() && + property_names_obj.at("type").get() != "string") { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "propertyNames must be an object that validates string" + ); + } + auto prop_names_result = Parse(schema.at("propertyNames"), "property_name", "string"); + if (prop_names_result.IsErr()) return ResultErr(std::move(prop_names_result).UnwrapErr()); + spec.property_names = std::move(prop_names_result).Unwrap(); + } + + spec.allow_additional_properties = !config_.strict_mode; + if (schema.count("additionalProperties")) { + auto add_props = schema.at("additionalProperties"); + if (add_props.is()) { + spec.allow_additional_properties = add_props.get(); + } else { + spec.allow_additional_properties = true; + auto add_props_result = Parse(add_props, "additional"); + if (add_props_result.IsErr()) return ResultErr(std::move(add_props_result).UnwrapErr()); + spec.additional_properties_schema = std::move(add_props_result).Unwrap(); + } + } + + spec.allow_unevaluated_properties = true; + if (schema.count("additionalProperties")) { + spec.allow_unevaluated_properties = spec.allow_additional_properties; + } else if (schema.count("unevaluatedProperties")) { + auto uneval_props = schema.at("unevaluatedProperties"); + if (uneval_props.is()) { + spec.allow_unevaluated_properties = uneval_props.get(); + } else { + spec.allow_unevaluated_properties = true; + auto uneval_result = Parse(uneval_props, "unevaluated"); + if (uneval_result.IsErr()) return ResultErr(std::move(uneval_result).UnwrapErr()); + spec.unevaluated_properties_schema = std::move(uneval_result).Unwrap(); + } + } else if (config_.strict_mode) { + spec.allow_unevaluated_properties = false; + } + + if (schema.count("minProperties")) { + if (!schema.at("minProperties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "minProperties must be an integer" + ); + } + spec.min_properties = static_cast(schema.at("minProperties").get()); + if (spec.min_properties < 0) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "minProperties must be a non-negative integer" + ); + } + } + if (schema.count("maxProperties")) { + if (!schema.at("maxProperties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "maxProperties must be an integer" + ); + } + spec.max_properties = static_cast(schema.at("maxProperties").get()); + if (spec.max_properties < 0) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "maxProperties must be a non-negative integer" + ); + } + } + + if (spec.max_properties != -1 && spec.min_properties > spec.max_properties) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minProperties is greater than maxProperties: " + std::to_string(spec.min_properties) + + " > " + std::to_string(spec.max_properties) + ); + } + if (spec.max_properties != -1 && static_cast(spec.required.size()) > spec.max_properties) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "maxProperties is less than the number of required properties: " + + std::to_string(spec.max_properties) + " < " + std::to_string(spec.required.size()) + ); + } + if (spec.pattern_properties.empty() && !spec.property_names && + !spec.allow_additional_properties && !spec.allow_unevaluated_properties && + spec.min_properties > static_cast(spec.properties.size())) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minProperties is greater than the number of properties, but additional properties aren't " + "allowed: " + + std::to_string(spec.min_properties) + " > " + std::to_string(spec.properties.size()) + ); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseConst(const picojson::object& schema) { + ConstSpec spec; + spec.json_value = schema.at("const").serialize(); + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseEnum(const picojson::object& schema) { + EnumSpec spec; + if (!schema.at("enum").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "enum must be an array"); + } + const auto& enum_array = schema.at("enum").get(); + if (enum_array.empty()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "enum array must not be empty"); + } + for (const auto& value : enum_array) { + spec.json_values.push_back(value.serialize()); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseRef(const picojson::object& schema) { + if (!schema.at("$ref").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "$ref must be a string"); + } + RefSpec spec; + spec.uri = schema.at("$ref").get(); + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ResolveRef( + const std::string& uri, const std::string& rule_name_hint +) { + if (ref_cache_.count(uri)) return ResultOk(ref_cache_[uri]); + + if (uri == "#") { + auto placeholder = SchemaSpec::Make(AnySpec{}, "", "root"); + ref_cache_[uri] = placeholder; + auto result = Parse(root_schema_, "root"); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + auto resolved = std::move(result).Unwrap(); + ref_cache_[uri] = resolved; + return ResultOk(resolved); + } + + if (uri.size() < 2 || uri[0] != '#' || uri[1] != '/') { + XGRAMMAR_LOG(WARNING) << "URI should either be '#' or start with '#/' but got " << uri; + return ResultOk(SchemaSpec::Make(AnySpec{}, "", "any")); + } + + std::vector parts; + std::stringstream ss(uri.substr(2)); + std::string part; + std::string new_rule_name_prefix; + while (std::getline(ss, part, '/')) { + if (!part.empty()) parts.push_back(part); + if (!new_rule_name_prefix.empty()) new_rule_name_prefix += "_"; + for (const auto& c : part) { + if (std::isalpha(c) || c == '_' || c == '-' || c == '.') new_rule_name_prefix += c; + } + } + + auto current = std::cref(root_schema_); + for (const auto& p : parts) { + if (!current.get().is() || !current.get().contains(p)) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "Cannot find field " + p + " in " + uri + ); + } + current = current.get().get(p); + } + + auto result = Parse(current, new_rule_name_prefix); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + auto resolved = std::move(result).Unwrap(); + ref_cache_[uri] = resolved; + return ResultOk(resolved); +} + +Result SchemaParser::ParseAnyOf( + const picojson::object& schema, const std::string& keyword +) { + AnyOfSpec spec; + if (!schema.at(keyword).is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, keyword + " must be an array"); + } + int idx = 0; + for (const auto& option : schema.at(keyword).get()) { + auto option_result = Parse(option, "case_" + std::to_string(idx)); + if (option_result.IsErr()) return ResultErr(std::move(option_result).UnwrapErr()); + spec.options.push_back(std::move(option_result).Unwrap()); + ++idx; + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseOneOf(const picojson::object& schema) { + OneOfSpec spec; + if (!schema.at("oneOf").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "oneOf must be an array"); + } + + const auto& options = schema.at("oneOf").get(); + if (options.empty()) { + return ResultErr(SchemaErrorType::kUnsupportedSchema, kUnsupportedOneOfMessage); + } + + int idx = 0; + for (const auto& option : options) { + auto option_result = Parse(option, "case_" + std::to_string(idx)); + if (option_result.IsErr()) return ResultErr(std::move(option_result).UnwrapErr()); + spec.options.push_back(std::move(option_result).Unwrap()); + ++idx; + } + + if (!TryProvePairwiseDisjointOneOf(options)) { + return ResultErr(SchemaErrorType::kUnsupportedSchema, kUnsupportedOneOfMessage); + } + + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseAllOf(const picojson::object& schema) { + AllOfSpec spec; + if (!schema.at("allOf").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "allOf must be an array"); + } + int idx = 0; + for (const auto& sub_schema : schema.at("allOf").get()) { + auto sub_result = Parse(sub_schema, "all_" + std::to_string(idx)); + if (sub_result.IsErr()) return ResultErr(std::move(sub_result).UnwrapErr()); + spec.schemas.push_back(std::move(sub_result).Unwrap()); + ++idx; + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseTypeArray( + const picojson::object& schema, const std::string& rule_name_hint +) { + TypeArraySpec spec; + auto type_array = schema.at("type").get(); + picojson::object schema_copy = schema; + if (type_array.empty()) { + schema_copy.erase("type"); + auto any_result = Parse(picojson::value(schema_copy), rule_name_hint); + if (any_result.IsErr()) return ResultErr(std::move(any_result).UnwrapErr()); + spec.type_schemas.push_back(std::move(any_result).Unwrap()); + return ResultOk(std::move(spec)); + } + for (const auto& type : type_array) { + if (!type.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "type must be a string or an array of strings" + ); + } + schema_copy["type"] = type; + auto type_result = + Parse(picojson::value(schema_copy), rule_name_hint + "_" + type.get()); + if (type_result.IsErr()) return ResultErr(std::move(type_result).UnwrapErr()); + spec.type_schemas.push_back(std::move(type_result).Unwrap()); + } + return ResultOk(std::move(spec)); +} + +} // namespace + +// ==================== IndentManager Implementation ==================== + +IndentManager::IndentManager( + std::optional indent, + const std::string& separator, + bool any_whitespace, + std::optional max_whitespace_cnt +) + : any_whitespace_(any_whitespace), + enable_newline_(indent.has_value()), + indent_(indent.value_or(0)), + separator_(separator), + total_indent_(0), + is_first_({true}), + max_whitespace_cnt_(max_whitespace_cnt) { + if (max_whitespace_cnt.has_value() && max_whitespace_cnt.value() <= 0) { + XGRAMMAR_LOG(FATAL) << "max_whitespace_cnt must be positive."; + } +} + +void IndentManager::StartIndent() { + total_indent_ += indent_; + is_first_.push_back(true); +} + +void IndentManager::EndIndent() { + total_indent_ -= indent_; + is_first_.pop_back(); +} + +std::string IndentManager::StartSeparator() { + if (any_whitespace_) { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } + if (!enable_newline_) { + return "\"\""; + } + return "\"\\n" + std::string(total_indent_, ' ') + "\""; +} + +std::string IndentManager::MiddleSeparator() { + if (any_whitespace_) { + std::string whitespace_part; + if (!max_whitespace_cnt_.has_value()) { + whitespace_part = "[ \\n\\r\\t]*"; + } else { + whitespace_part = "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + return whitespace_part + " \"" + separator_ + "\" " + whitespace_part; + } + if (!enable_newline_) { + return "\"" + separator_ + "\""; + } + return "\"" + separator_ + "\\n" + std::string(total_indent_, ' ') + "\""; +} + +std::string IndentManager::EndSeparator() { + if (any_whitespace_) { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } + if (!enable_newline_) { + return "\"\""; + } + return "\"\\n" + std::string(total_indent_ - indent_, ' ') + "\""; +} + +std::string IndentManager::EmptySeparator() { + if (any_whitespace_) { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } + return "\"\""; +} + +std::string IndentManager::NextSeparator(bool is_end) { + if (any_whitespace_) { + if (is_first_.back() || is_end) { + is_first_.back() = false; + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } else { + std::string whitespace_part; + if (!max_whitespace_cnt_.has_value()) { + whitespace_part = "[ \\n\\r\\t]*"; + } else { + whitespace_part = "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + return whitespace_part + " \"" + separator_ + "\" " + whitespace_part; + } + } + + std::string res = ""; + if (!is_first_.back() && !is_end) { + res += separator_; + } + is_first_.back() = false; + + if (enable_newline_) { + res += "\\n"; + } + + if (!is_end) { + res += std::string(total_indent_, ' '); + } else { + res += std::string(total_indent_ - indent_, ' '); + } + + return "\"" + res + "\""; +} + +// ==================== Static Constants ==================== + +const std::string JSONSchemaConverter::kBasicAny = "basic_any"; +const std::string JSONSchemaConverter::kBasicInteger = "basic_integer"; +const std::string JSONSchemaConverter::kBasicNumber = "basic_number"; +const std::string JSONSchemaConverter::kBasicString = "basic_string"; +const std::string JSONSchemaConverter::kBasicBoolean = "basic_boolean"; +const std::string JSONSchemaConverter::kBasicNull = "basic_null"; +const std::string JSONSchemaConverter::kBasicArray = "basic_array"; +const std::string JSONSchemaConverter::kBasicObject = "basic_object"; +const std::string JSONSchemaConverter::kBasicEscape = "basic_escape"; +const std::string JSONSchemaConverter::kBasicStringSub = "basic_string_sub"; + +// ==================== JSONSchemaConverter Implementation ==================== + +JSONSchemaConverter::JSONSchemaConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + bool any_order +) + : indent_manager_( + indent, + separators.has_value() ? separators->first + : (any_whitespace ? "," : (indent.has_value() ? "," : ", ")), + any_whitespace, + max_whitespace_cnt + ), + any_whitespace_(any_whitespace), + max_whitespace_cnt_(max_whitespace_cnt), + any_order_(any_order), + ref_resolver_(std::move(ref_resolver)) { + std::string colon_sep = + separators.has_value() ? separators->second : (any_whitespace ? ":" : ": "); + std::string whitespace = GetWhitespacePattern(); + colon_expr_id_ = FormattingExpression( + any_whitespace ? whitespace + " \"" + colon_sep + "\" " + whitespace : "\"" + colon_sep + "\"" + ); +} + +Grammar JSONSchemaConverter::Convert(const SchemaSpecPtr& spec) { + AddBasicRules(); + + // Register the root rule for circular reference handling + // This allows $ref: "#" to resolve to "root" + int32_t root_rule_id = builder_.AddEmptyRuleWithHint("root"); + std::string root_rule_name = builder_.GetRule(root_rule_id).name; + uri_to_rule_id_["#"] = root_rule_id; + + // Check if the spec can be directly mapped to an existing rule + auto cached_rule = GetCache(spec->cache_key); + if (cached_rule.has_value()) { + // Root schema matches a basic type, just reference it + builder_.UpdateRuleBody(root_rule_id, RuleRef(*cached_rule)); + } else { + // Generate the rule body + if (!spec->cache_key.empty()) { + AddCache(spec->cache_key, root_rule_id); + } + builder_.UpdateRuleBody(root_rule_id, GenerateFromSpec(spec, root_rule_name)); + } + return builder_.Get(root_rule_id); +} + +void JSONSchemaConverter::AddBasicRules() { AddBasicRules({}); } + +void JSONSchemaConverter::AddBasicRules(const std::vector& additional_rule_names) { + std::vector basic_rule_names = { + kBasicEscape, + kBasicStringSub, + kBasicAny, + kBasicInteger, + kBasicNumber, + kBasicString, + kBasicBoolean, + kBasicNull, + kBasicArray, + kBasicObject, + }; + basic_rule_names.insert( + basic_rule_names.end(), additional_rule_names.begin(), additional_rule_names.end() + ); + for (const auto& name : basic_rule_names) { + builder_.AddEmptyRule(name); + } + AddHelperRules(); + + // Create basic rules with a temporary indent manager for compact format + auto saved_indent_manager = indent_manager_; + indent_manager_ = IndentManager( + std::nullopt, + any_whitespace_ ? "," : ", ", + any_whitespace_, + any_whitespace_ ? max_whitespace_cnt_ : std::nullopt + ); + + // basic_any - use "{}" as the cache key for empty schema + auto any_spec = SchemaSpec::Make(AnySpec{}, "{}", kBasicAny); + builder_.UpdateRuleBody(kBasicAny, GenerateAny(std::get(any_spec->spec), kBasicAny)); + AddCache("{}", builder_.GetRuleId(kBasicAny)); + + // basic_integer - cache_key matches SchemaParser::ComputeCacheKey for {"type": "integer"} + constexpr const char* kIntegerCacheKey = "{\"type\":\"integer\"}"; + builder_.UpdateRuleBody(kBasicInteger, GenerateInteger(IntegerSpec{}, kBasicInteger)); + AddCache(kIntegerCacheKey, builder_.GetRuleId(kBasicInteger)); + + // basic_number - cache_key matches SchemaParser::ComputeCacheKey for {"type": "number"} + constexpr const char* kNumberCacheKey = "{\"type\":\"number\"}"; + builder_.UpdateRuleBody(kBasicNumber, GenerateNumber(NumberSpec{}, kBasicNumber)); + AddCache(kNumberCacheKey, builder_.GetRuleId(kBasicNumber)); + + constexpr const char* kStringCacheKey = "{\"type\":\"string\"}"; + builder_.UpdateRuleBody(kBasicString, Sequence({ByteString("\""), RuleRef(kBasicStringSub)})); + AddCache(kStringCacheKey, builder_.GetRuleId(kBasicString)); + + // basic_boolean - cache_key matches SchemaParser::ComputeCacheKey for {"type": "boolean"} + constexpr const char* kBooleanCacheKey = "{\"type\":\"boolean\"}"; + builder_.UpdateRuleBody(kBasicBoolean, GenerateBoolean(BooleanSpec{}, kBasicBoolean)); + AddCache(kBooleanCacheKey, builder_.GetRuleId(kBasicBoolean)); + + // basic_null - cache_key matches SchemaParser::ComputeCacheKey for {"type": "null"} + constexpr const char* kNullCacheKey = "{\"type\":\"null\"}"; + builder_.UpdateRuleBody(kBasicNull, GenerateNull(NullSpec{}, kBasicNull)); + AddCache(kNullCacheKey, builder_.GetRuleId(kBasicNull)); + + // basic_array - cache_key matches SchemaParser::ComputeCacheKey for {"type": "array"} + constexpr const char* kArrayCacheKey = "{\"type\":\"array\"}"; + ArraySpec array_spec_val; + array_spec_val.allow_additional_items = true; + array_spec_val.additional_items = any_spec; + builder_.UpdateRuleBody(kBasicArray, GenerateArray(array_spec_val, kBasicArray)); + AddCache(kArrayCacheKey, builder_.GetRuleId(kBasicArray)); + + // basic_object - cache_key matches SchemaParser::ComputeCacheKey for {"type": "object"} + constexpr const char* kObjectCacheKey = "{\"type\":\"object\"}"; + ObjectSpec obj_spec_val; + obj_spec_val.allow_additional_properties = true; + obj_spec_val.additional_properties_schema = any_spec; + builder_.UpdateRuleBody(kBasicObject, GenerateObject(obj_spec_val, kBasicObject)); + AddCache(kObjectCacheKey, builder_.GetRuleId(kBasicObject)); + + indent_manager_ = saved_indent_manager; +} + +void JSONSchemaConverter::AddHelperRules() { + if (max_whitespace_cnt_.has_value()) { + // Preserve historical helper-rule numbering after grammar optimization. The text parser + // allocated one initial bounded-repetition helper that dead-code elimination later removed. + builder_.AddRuleWithHint(kBasicStringSub, Empty()); + } + int32_t escaped_character = builder_.AddCharacterClass( + {{'"', '"'}, + {'\\', '\\'}, + {'/', '/'}, + {'b', 'b'}, + {'f', 'f'}, + {'n', 'n'}, + {'r', 'r'}, + {'t', 't'}} + ); + int32_t hexadecimal_character = builder_.AddCharacterClass({{'A', 'F'}, {'a', 'f'}, {'0', '9'}}); + int32_t unicode_escape = Sequence( + {ByteString("u"), + hexadecimal_character, + hexadecimal_character, + hexadecimal_character, + hexadecimal_character} + ); + builder_.UpdateRuleBody(kBasicEscape, Choice({escaped_character, unicode_escape})); + + int32_t normal_character = builder_.AddCharacterClass( + {{0, 0x1f}, {'"', '"'}, {'\\', '\\'}, {'\r', '\r'}, {'\n', '\n'}}, true + ); + int32_t string_sub_ref = RuleRef(kBasicStringSub); + int32_t string_sub_body = Choice( + {ByteString("\""), + Sequence({normal_character, string_sub_ref}), + Sequence({ByteString("\\"), RuleRef(kBasicEscape), string_sub_ref})} + ); + builder_.UpdateRuleBody(kBasicStringSub, string_sub_body); + int32_t closing_context = + builder_.AddCharacterClass({{',', ','}, {'}', '}'}, {']', ']'}, {':', ':'}}); + builder_.UpdateLookaheadAssertion( + kBasicStringSub, Sequence({WhitespaceExpression(), closing_context}) + ); +} + +// Keep converter-specific node reuse local; GrammarBuilder creates all AST nodes. +int32_t JSONSchemaConverter::Empty() { + if (!empty_expr_id_.has_value()) { + empty_expr_id_ = builder_.AddEmptyStr(); + } + return *empty_expr_id_; +} + +int32_t JSONSchemaConverter::ByteString(const std::string& value) { + auto it = byte_string_expr_ids_.find(value); + if (it != byte_string_expr_ids_.end()) { + return it->second; + } + int32_t expr_id = value.empty() ? Empty() : builder_.AddByteString(value); + byte_string_expr_ids_[value] = expr_id; + return expr_id; +} + +int32_t JSONSchemaConverter::TagDispatch( + bool loop_after_dispatch, std::vector excludes +) { + return builder_.AddTagDispatch( + Grammar::Impl::TagDispatch{{}, loop_after_dispatch, std::move(excludes)} + ); +} + +int32_t JSONSchemaConverter::RuleRef(int32_t rule_id) { + auto it = rule_ref_expr_ids_.find(rule_id); + if (it != rule_ref_expr_ids_.end()) { + return it->second; + } + int32_t expr_id = builder_.AddRuleRef(rule_id); + rule_ref_expr_ids_[rule_id] = expr_id; + return expr_id; +} + +int32_t JSONSchemaConverter::RuleRef(const std::string& rule_name) { + int32_t rule_id = builder_.GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not allocated"; + return RuleRef(rule_id); +} + +int32_t JSONSchemaConverter::Sequence(const std::vector& elements) { + if (elements.empty()) { + return Empty(); + } + if (elements.size() == 1) { + return elements[0]; + } + return builder_.AddSequence(elements); +} + +int32_t JSONSchemaConverter::Choice(const std::vector& choices) { + if (choices.empty()) { + return Empty(); + } + if (choices.size() == 1) { + return choices[0]; + } + return builder_.AddChoices(choices); +} + +int32_t JSONSchemaConverter::Repeat( + const std::string& rule_name_hint, int32_t expr_id, int32_t min_count, int32_t max_count +) { + if (min_count == 0 && max_count == 0) { + return Empty(); + } + if (min_count == 1 && max_count == 1) { + return expr_id; + } + if (min_count == 0 && max_count == 1) { + return Choice({Empty(), expr_id}); + } + if (min_count == 0 && max_count == -1) { + auto expr = builder_.GetGrammarExpr(expr_id); + if (expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) { + std::vector data(expr.begin(), expr.end()); + return builder_.AddGrammarExpr( + {GrammarBuilder::GrammarExprType::kCharacterClassStar, + data.data(), + static_cast(data.size())} + ); + } + } + return builder_.AddRepeatFromExpr(rule_name_hint, expr_id, min_count, max_count); +} + +int32_t JSONSchemaConverter::AddSubGrammar(const Grammar& grammar) { + int32_t rule_id = SubGrammarAdder::Apply(&builder_, grammar); + return RuleRef(rule_id); +} + +std::string JSONSchemaConverter::GetWhitespacePattern() const { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } + return "[ \\n\\r\\t]{0," + std::to_string(*max_whitespace_cnt_) + "}"; +} + +int32_t JSONSchemaConverter::WhitespaceExpression() { + std::vector elements = { + {' ', ' '}, {'\n', '\n'}, {'\r', '\r'}, {'\t', '\t'} + }; + if (!max_whitespace_cnt_.has_value()) { + if (!whitespace_expr_id_.has_value()) { + whitespace_expr_id_ = builder_.AddCharacterClassStar(elements); + } + return *whitespace_expr_id_; + } + // Bounded whitespace occurrences intentionally remain distinct, matching the historical + // parser-produced rule shape after normalization. + return Repeat( + "whitespace", + builder_.AddCharacterClass(elements), + 0, + static_cast(*max_whitespace_cnt_) + ); +} + +int32_t JSONSchemaConverter::FormattingExpression(const std::string& expression) { + const std::string whitespace = GetWhitespacePattern(); + if (expression == whitespace) { + return WhitespaceExpression(); + } + + const std::string prefix = whitespace + " "; + const std::string suffix = " " + whitespace; + if (expression.size() >= prefix.size() + suffix.size() && + expression.compare(0, prefix.size(), prefix) == 0 && + expression.compare(expression.size() - suffix.size(), suffix.size(), suffix) == 0) { + return Sequence( + {WhitespaceExpression(), + FormattingExpression( + expression.substr(prefix.size(), expression.size() - prefix.size() - suffix.size()) + ), + WhitespaceExpression()} + ); + } + + picojson::value value; + std::string error = ParseJSON(value, expression); + XGRAMMAR_CHECK(error.empty() && value.is()) + << "Unsupported indentation expression: " << expression; + return ByteString(value.get()); +} + +std::string JSONSchemaConverter::NextSeparator(bool is_end) { + return indent_manager_.NextSeparator(is_end); +} + +int32_t JSONSchemaConverter::NextSeparatorExpression(bool is_end) { + return FormattingExpression(NextSeparator(is_end)); +} + +std::string JSONSchemaConverter::GetKeyPattern() const { return kBasicString; } + +int32_t JSONSchemaConverter::KeyPatternExpression() { return RuleRef(GetKeyPattern()); } + +int32_t JSONSchemaConverter::BuildTrieBody(const TrieNode& node, const std::string& rule_name) { + std::vector choices; + if (!node.is_terminal) { + choices.push_back(ByteString("\"")); + } + + std::vector excluded = { + {0, 0x1f}, {'"', '"'}, {'\\', '\\'}, {'\r', '\r'}, {'\n', '\n'} + }; + for (const auto& [character, child] : node.children) { + static_cast(child); + excluded.push_back({character, character}); + } + // The FSM negative-class path only handles ASCII exclusions. Use the explicit + // Unicode complement so non-ASCII declared keys are excluded exactly as well. + excluded.push_back({0xd800, 0xdfff}); + std::sort(excluded.begin(), excluded.end(), [](const auto& a, const auto& b) { + return a.lower < b.lower; + }); + std::vector allowed; + int32_t next = 0; + for (const auto& range : excluded) { + if (next < range.lower) allowed.push_back({next, range.lower - 1}); + next = std::max(next, range.upper + 1); + } + if (next <= 0x10ffff) allowed.push_back({next, 0x10ffff}); + choices.push_back(Sequence({builder_.AddCharacterClass(allowed), RuleRef(kBasicStringSub)})); + // NInfer: while a key can still equal a declared property, require its literal spelling. + // An unrestricted escape here could alias a typed property ("a" == "\\u0061") and + // overwrite it with a value checked only against additionalProperties. + for (const auto& [character, child] : node.children) { + choices.push_back(Sequence( + {ByteString(CharToUTF8(character)), BuildTrieBody(child, rule_name)} + )); + } + return Choice(choices); +} + +int32_t JSONSchemaConverter::GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name +) { + if (properties.empty()) { + return KeyPatternExpression(); + } + + // Build trie from property names + // Match Unicode codepoints, consistent with grammar character classes. + TrieNode root; + for (const auto& prop : properties) { + TrieNode* cur = &root; + // A name requiring JSON escaping cannot equal an additional key while it follows + // a declared prefix: such prefixes permit only unescaped characters. + if (std::any_of(prop.name.begin(), prop.name.end(), [](unsigned char c) { + return c < 0x20 || c == '"' || c == '\\'; + })) { + continue; + } + for (auto c : ParseUTF8(prop.name.c_str())) { + cur = &cur->children[c]; + } + cur->is_terminal = true; + } + + int32_t key_rule_id = builder_.AddEmptyRuleWithHint(rule_name + "_addl_key"); + std::string key_rule_name = builder_.GetRule(key_rule_id).name; + builder_.UpdateRuleBody( + key_rule_id, Sequence({ByteString("\""), BuildTrieBody(root, key_rule_name)}) + ); + builder_.UpdateLookaheadAssertion( + key_rule_id, + Sequence( + {WhitespaceExpression(), + builder_.AddCharacterClass({{',', ','}, {'}', '}'}, {']', ']'}, {':', ':'}})} + ) + ); + return RuleRef(key_rule_id); +} + +std::string JSONSchemaConverter::GetBasicAnyRuleName() const { return kBasicAny; } + +void JSONSchemaConverter::AddCache(const std::string& key, int32_t rule_id) { + if (!key.empty()) { + rule_cache_manager_.AddCache(key, true, rule_id); + } +} + +std::optional JSONSchemaConverter::GetCache(const std::string& key) const { + if (key.empty()) { + return std::nullopt; + } + return rule_cache_manager_.GetCache(key, true); +} + +int32_t JSONSchemaConverter::CreateRule( + const SchemaSpecPtr& spec, const std::string& rule_name_hint +) { + // Only check cache for basic rules (pre-populated in AddBasicRules) + // Don't cache other rules to match original behavior + auto cached = GetCache(spec->cache_key); + if (cached.has_value()) { + return cached.value(); + } + int32_t rule_id = builder_.AddEmptyRuleWithHint(rule_name_hint); + // Copy the name before generating: GenerateFromSpec may add rules and reallocate the + // builder's rule storage, invalidating references into it. + std::string rule_name = builder_.GetRule(rule_id).name; + builder_.UpdateRuleBody(rule_id, GenerateFromSpec(spec, rule_name)); + return rule_id; +} + +int32_t JSONSchemaConverter::GenerateFromSpec( + const SchemaSpecPtr& spec, const std::string& rule_name_hint +) { + return std::visit( + [this, &rule_name_hint](const auto& s) -> int32_t { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return GenerateInteger(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateNumber(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateString(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateBoolean(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateNull(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateArray(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateObject(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateAny(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateConst(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateEnum(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateRef(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateAnyOf(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateOneOf(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateAllOf(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateTypeArray(s, rule_name_hint); + } else { + XGRAMMAR_LOG(FATAL) << "Unknown spec type"; + } + }, + spec->spec + ); +} + +/*! + * \brief Emit the grammar expression matching a regex. Prefer the Regex node with + * json_string=true so the pattern is compiled into a single automaton by GrammarFSMBuilder; + * json_string=true excludes the characters that must be escaped in a JSON string ('"', '\\' + * and the control characters) from every character match, so classes like \S cannot emit an + * unescaped quote. Fall back to the CFG expansion when the FSM regex engine does not support + * the pattern, or when the exclusion makes the pattern unmatchable (e.g. a pattern requiring + * a literal '"'). + */ +int32_t JSONSchemaConverter::RegexExpression( + const std::string& regex, bool json_string, bool force_cfg_expansion +) { + bool can_use_fsm = !force_cfg_expansion; + if (json_string) { + can_use_fsm = + can_use_fsm && std::all_of(regex.begin(), regex.end(), [](unsigned char character) { + return character >= 0x20 && character <= 0x7e; + }); + } + if (can_use_fsm) { + auto fsm_result = GrammarFSMBuilder::Regex(regex, json_string); + if (fsm_result.IsOk()) { + auto fsm = std::move(fsm_result).Unwrap(); + std::unordered_set reachable_states; + fsm.GetReachableStates(&reachable_states); + bool language_is_empty = + std::none_of(reachable_states.begin(), reachable_states.end(), [&](int state) { + return fsm.IsEndState(state); + }); + if (!language_is_empty) { + return builder_.AddRegex(regex, json_string); + } + } + } + + // Keep regex conversion independent. Only the uncommon fallback path converts its existing + // EBNF result to a subgrammar; the JSON Schema rule graph itself is still built directly. + return AddSubGrammar(Grammar::FromEBNF(RegexToEBNF(regex))); +} + +// ==================== Generate Methods ==================== + +int32_t JSONSchemaConverter::GenerateInteger( + const IntegerSpec& spec, const std::string& rule_name +) { + // Shared with ParseInteger's range validation so emission and validation agree on the effective + // range; a nullopt side means that side is unbounded. + const EffectiveIntegerRange range = ComputeEffectiveIntegerRange(spec); + std::optional start = range.start; + std::optional end = range.end; + + if (spec.multiple_of.has_value()) { + // ParseInteger keeps multiple_of only when the range is fully bounded (enumerate the + // multiples) or fully unbounded (emit a modulo DFA); the half-bounded case is dropped there. + if (start.has_value() && end.has_value()) { + std::vector multiples; + for (int64_t value = *start; value <= *end; ++value) { + if (IsMultipleOf(value, *spec.multiple_of)) { + multiples.push_back(ByteString(std::to_string(value))); + } + if (value == std::numeric_limits::max()) { + break; + } + } + return Choice(multiples); + } + return GenerateIntegerMultipleOfDFA(*spec.multiple_of, rule_name); + } + if (start.has_value() || end.has_value()) { + return RegexExpression( + GenerateRangeRegex(start, end), + false, + /*force_cfg_expansion=*/true + ); + } + int32_t optional_minus = Choice({Empty(), ByteString("-")}); + return Choice( + {ByteString("0"), + Sequence( + {optional_minus, + builder_.AddCharacterClass({{'1', '9'}}), + builder_.AddCharacterClassStar({{'0', '9'}})} + )} + ); +} + +int32_t JSONSchemaConverter::GenerateIntegerMultipleOfDFA( + int64_t multiple_of, const std::string& rule_name +) { + std::vector states(multiple_of); + for (int64_t state = 0; state < multiple_of; ++state) { + states[state] = builder_.AddEmptyRuleWithHint( + rule_name + "_multiple_of_" + std::to_string(multiple_of) + "_mod_" + std::to_string(state) + ); + } + for (int64_t state = 0; state < multiple_of; ++state) { + std::vector transitions; + if (state == 0) { + transitions.push_back(Empty()); + } + for (int64_t digit = 0; digit <= 9; ++digit) { + int64_t next_state = (state * 10 + digit) % multiple_of; + transitions.push_back( + Sequence({ByteString(std::to_string(digit)), RuleRef(states[next_state])}) + ); + } + builder_.UpdateRuleBody(states[state], Choice(transitions)); + } + + std::vector non_zero_starts; + for (int64_t digit = 1; digit <= 9; ++digit) { + non_zero_starts.push_back( + Sequence({ByteString(std::to_string(digit)), RuleRef(states[digit % multiple_of])}) + ); + } + return Choice( + {ByteString("0"), Sequence({Choice({Empty(), ByteString("-")}), Choice(non_zero_starts)})} + ); +} + +int32_t JSONSchemaConverter::GenerateNumber(const NumberSpec& spec, const std::string& rule_name) { + std::optional start = spec.minimum; + std::optional end = spec.maximum; + bool exclusive_start = false; + bool exclusive_end = false; + // When both bounds are present the larger lower bound wins; on a tie the + // exclusive one is stricter. + if (spec.exclusive_minimum.has_value() && + (!start.has_value() || *spec.exclusive_minimum >= *start)) { + start = spec.exclusive_minimum; + exclusive_start = true; + } + if (spec.exclusive_maximum.has_value() && (!end.has_value() || *spec.exclusive_maximum <= *end)) { + end = spec.exclusive_maximum; + exclusive_end = true; + } + if (start.has_value() || end.has_value()) { + return RegexExpression( + GenerateFloatRangeRegex(start, end, /*precision=*/6, exclusive_start, exclusive_end), + false, + /*force_cfg_expansion=*/true + ); + } + + int32_t optional_minus = Choice({Empty(), ByteString("-")}); + int32_t integer_part = Choice( + {ByteString("0"), + Sequence( + {builder_.AddCharacterClass({{'1', '9'}}), builder_.AddCharacterClassStar({{'0', '9'}})} + )} + ); + int32_t one_or_more_digits = + Repeat(rule_name + "_digits", builder_.AddCharacterClass({{'0', '9'}}), 1, -1); + int32_t fraction = Choice({Empty(), Sequence({ByteString("."), one_or_more_digits})}); + int32_t exponent = Choice( + {Empty(), + Sequence( + {builder_.AddCharacterClass({{'e', 'e'}, {'E', 'E'}}), + Choice({Empty(), builder_.AddCharacterClass({{'+', '+'}, {'-', '-'}})}), + one_or_more_digits} + )} + ); + // Note: The format must be "-"? ("0" | ...) not ("0" | "-"? ...) + // The first allows -0, -123, 0, 123 + // The second allows 0, -123, 123 but not -0 + return Sequence({optional_minus, integer_part, fraction, exponent}); +} + +int32_t JSONSchemaConverter::GenerateString(const StringSpec& spec, const std::string& rule_name) { + // Check for format + if (spec.format.has_value()) { + auto regex = JSONFormatToRegexPattern(*spec.format); + if (regex.has_value()) { + // The built-in format regexes use constructs that the FSM regex engine does not fully + // support yet (e.g. quoted email local parts), so they keep the CFG expansion. + return Sequence({ByteString("\""), RegexExpression(*regex, false, true), ByteString("\"")}); + } + } + // Check for pattern + if (spec.pattern.has_value()) { + return Sequence( + {ByteString("\""), RegexExpression(*spec.pattern, /*json_string=*/true), ByteString("\"")} + ); + } + // Check for length constraints + if (spec.min_length != 0 || spec.max_length != -1) { + int32_t character = + builder_.AddCharacterClass({{0, 0x1f}, {'"', '"'}, {'\\', '\\'}, {'\r', '\r'}, {'\n', '\n'}}, true); + int32_t body = Repeat(rule_name + "_characters", character, spec.min_length, spec.max_length); + return Sequence({ByteString("\""), body, ByteString("\"")}); + } + // Default string + return Sequence({ByteString("\""), RuleRef(kBasicStringSub)}); +} + +int32_t JSONSchemaConverter::GenerateBoolean( + const BooleanSpec& spec, const std::string& rule_name +) { + return Choice({ByteString("true"), ByteString("false")}); +} + +int32_t JSONSchemaConverter::GenerateNull(const NullSpec& spec, const std::string& rule_name) { + return ByteString("null"); +} + +int32_t JSONSchemaConverter::GenerateArray(const ArraySpec& spec, const std::string& rule_name) { + indent_manager_.StartIndent(); + int32_t start_separator = FormattingExpression(indent_manager_.StartSeparator()); + int32_t middle_separator = FormattingExpression(indent_manager_.MiddleSeparator()); + int32_t end_separator = FormattingExpression(indent_manager_.EndSeparator()); + int32_t empty_separator = FormattingExpression(indent_manager_.EmptySeparator()); + + std::vector item_rule_ids; + for (size_t index = 0; index < spec.prefix_items.size(); ++index) { + item_rule_ids.push_back( + CreateRule(spec.prefix_items[index], rule_name + "_item_" + std::to_string(index)) + ); + } + int32_t additional_rule_id = -1; + if (spec.allow_additional_items && spec.additional_items) { + additional_rule_id = CreateRule(spec.additional_items, rule_name + "_additional"); + } + indent_manager_.EndIndent(); + + int32_t left_bracket = ByteString("["); + int32_t right_bracket = ByteString("]"); + int32_t empty_array = Sequence({left_bracket, empty_separator, right_bracket}); + + if (item_rule_ids.empty()) { + if (!spec.allow_additional_items || spec.max_items == 0) { + return empty_array; + } + int32_t additional = RuleRef(additional_rule_id); + int32_t tail = Repeat( + rule_name + "_items", + Sequence({middle_separator, additional}), + spec.min_items == 0 ? 0 : static_cast(spec.min_items - 1), + spec.max_items == -1 ? -1 : static_cast(spec.max_items - 1) + ); + int32_t nonempty = + Sequence({left_bracket, start_separator, additional, tail, end_separator, right_bracket}); + return spec.min_items == 0 ? Choice({nonempty, empty_array}) : nonempty; + } + + // Per Draft 2020-12, prefixItems entries are positional: the instance may + // end after any prefix position (subject to minItems), and additional items + // are only allowed after the full prefix (issue #824). + size_t mandatory_count = static_cast(std::min( + std::max(0, spec.min_items), static_cast(item_rule_ids.size()) + )); + + // Mandatory head: the first min(minItems, n) items, separated. + std::vector prefix_elements; + for (size_t index = 0; index < mandatory_count; ++index) { + if (index != 0) { + prefix_elements.push_back(middle_separator); + } + prefix_elements.push_back(RuleRef(item_rule_ids[index])); + } + + // Suffix after the mandatory head, flattened into a right-recursive chain + // of rules suffix_k ::= "" | sep item_k suffix_{k+1} so each position + // is encoded once instead of once per truncation length. The chain ends + // with the additional-items tail. Positions from index 1 on are separated + // by middle_separator; position 0, when it is not part of the mandatory + // head, gets its own rule without the separator. + int32_t suffix = Empty(); + if (spec.allow_additional_items && spec.additional_items) { + int64_t minimum_additional = + std::max(int64_t{0}, spec.min_items - static_cast(item_rule_ids.size())); + suffix = Repeat( + rule_name + "_additional_items", + Sequence({middle_separator, RuleRef(additional_rule_id)}), + static_cast(minimum_additional), + spec.max_items == -1 + ? -1 + : static_cast(spec.max_items - static_cast(item_rule_ids.size())) + ); + } + size_t chain_start = std::max(mandatory_count, 1); + for (size_t k = item_rule_ids.size(); k-- > chain_start;) { + int32_t with_item = Sequence({middle_separator, RuleRef(item_rule_ids[k]), suffix}); + int32_t suffix_rule_id = builder_.AddRuleWithHint( + rule_name + "_suffix_" + std::to_string(k), Choice({Empty(), with_item}) + ); + suffix = RuleRef(suffix_rule_id); + } + if (mandatory_count == 0) { + int32_t with_first = Sequence({RuleRef(item_rule_ids[0]), suffix}); + int32_t suffix_rule_id = + builder_.AddRuleWithHint(rule_name + "_suffix_0", Choice({Empty(), with_first})); + suffix = RuleRef(suffix_rule_id); + } + + std::vector content_elements = prefix_elements; + content_elements.push_back(suffix); + int32_t prefix = Sequence(content_elements); + return Sequence({left_bracket, start_separator, prefix, end_separator, right_bracket}); +} + +int32_t JSONSchemaConverter::FormatPropertyKey( + const std::string& key, const SchemaSpecPtr& schema +) { + return ByteString(picojson::value(key).serialize()); +} + +int32_t JSONSchemaConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + return Sequence({FormatPropertyKey(key, schema), colon_expr_id_, RuleRef(value_rule_id)}); +} + +int32_t JSONSchemaConverter::FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema +) { + return Sequence({key_pattern_expr, colon_expr_id_, RuleRef(value_rule_id)}); +} + +int32_t JSONSchemaConverter::GetPropertyWithNumberConstraints( + int32_t pattern, + int min_properties, + int max_properties, + int already_repeated_times, + const std::string& rule_name +) { + if (max_properties != -1 && max_properties == already_repeated_times) { + return Empty(); + } + int lower = std::max(0, min_properties - already_repeated_times); + int upper = max_properties == -1 ? -1 : std::max(-1, max_properties - already_repeated_times); + return Repeat(rule_name + "_properties", pattern, lower, upper); +} + +int32_t JSONSchemaConverter::GetAnyOrderRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override +) { + int32_t first_separator = NextSeparatorExpression(); + int32_t middle_separator = NextSeparatorExpression(); + int32_t last_separator = NextSeparatorExpression(true); + + // Build one "item" alternation over every property (any required/optional key) plus any + // additional/pattern key; any_order does not care which key goes where. + std::vector items; + for (size_t index = 0; index < properties.size(); ++index) { + const auto& property = properties[index]; + int32_t value_rule_id = + CreateRule(property.schema, rule_name + "_prop_" + std::to_string(index)); + items.push_back(FormatProperty(property.name, value_rule_id, rule_name, index, property.schema) + ); + } + if (additional != nullptr) { + if (additional_property_override.has_value()) { + items.push_back(*additional_property_override); + } else { + int32_t value_rule_id = CreateRule(additional, rule_name + "_" + additional_suffix); + items.push_back(FormatOtherProperty( + GetKeyPatternExcluding(properties, rule_name), + value_rule_id, + rule_name, + additional_suffix, + additional + )); + } + } + + int32_t item_rule_id = builder_.AddRuleWithHint(rule_name + "_item", Choice(items)); + + // Repeat `item` between n = max(minProperties, #required) and m = maxProperties times; only the + // count is constrained, not which keys appear. + int minimum_count = std::max(min_properties, static_cast(required.size())); + int32_t repeated_items = GetPropertyWithNumberConstraints( + Sequence({middle_separator, RuleRef(item_rule_id)}), + minimum_count, + max_properties, + 1, + rule_name + ); + return Sequence({first_separator, RuleRef(item_rule_id), repeated_items, last_separator}); +} + +int32_t JSONSchemaConverter::GetPartialRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override +) { + if (max_properties == 0) { + return Empty(); + } + if (any_order_) { + return GetAnyOrderRuleForProperties( + properties, + required, + additional, + rule_name, + additional_suffix, + min_properties, + max_properties, + additional_property_override + ); + } + + int32_t first_separator = NextSeparatorExpression(); + int32_t middle_separator = NextSeparatorExpression(); + int32_t last_separator = NextSeparatorExpression(true); + + std::vector property_patterns; + for (size_t index = 0; index < properties.size(); ++index) { + int32_t value_rule_id = + CreateRule(properties[index].schema, rule_name + "_prop_" + std::to_string(index)); + property_patterns.push_back(FormatProperty( + properties[index].name, value_rule_id, rule_name, index, properties[index].schema + )); + } + + bool allow_additional = additional != nullptr; + std::optional additional_pattern; + auto get_additional_pattern = [&]() -> int32_t { + if (!additional_pattern.has_value()) { + if (additional_property_override.has_value()) { + additional_pattern = *additional_property_override; + } else { + int32_t value_rule_id = CreateRule(additional, rule_name + "_" + additional_suffix); + additional_pattern = FormatOtherProperty( + GetKeyPatternExcluding(properties, rule_name), + value_rule_id, + rule_name, + additional_suffix, + additional + ); + } + } + return *additional_pattern; + }; + + if (min_properties == 0 && max_properties == -1) { + // Case 1: No property number constraints + std::vector tails(properties.size(), Empty()); + std::vector is_required(properties.size(), false); + + if (allow_additional) { + int32_t repeated_additional = Repeat( + rule_name + "_additional_properties", + Sequence({middle_separator, get_additional_pattern()}), + 0, + -1 + ); + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(static_cast(properties.size()) - 1), + repeated_additional + ); + tails.back() = RuleRef(tail_rule_id); + } + + for (int index = static_cast(properties.size()) - 2; index >= 0; --index) { + int32_t with_property = + Sequence({middle_separator, property_patterns[index + 1], tails[index + 1]}); + int32_t body = with_property; + if (!required.count(properties[index + 1].name)) { + body = Choice({tails[index + 1], with_property}); + } else { + is_required[index + 1] = true; + } + int32_t tail_rule_id = + builder_.AddRuleWithHint(rule_name + "_part_" + std::to_string(index), body); + tails[index] = RuleRef(tail_rule_id); + } + if (required.count(properties[0].name)) { + is_required[0] = true; + } + + std::vector choices; + for (size_t index = 0; index < properties.size(); ++index) { + choices.push_back(Sequence({property_patterns[index], tails[index]})); + if (is_required[index]) { + break; + } + } + if (allow_additional && required.empty()) { + choices.push_back(Sequence({get_additional_pattern(), tails.back()})); + } + return Sequence({first_separator, Choice(choices), last_separator}); + } + + const int property_count = static_cast(properties.size()); + std::vector is_required(property_count, false); + std::vector matched_min(property_count, 0); + bool found_required = required.count(properties[0].name); + matched_min[0] = 1; + for (int index = 1; index < property_count; ++index) { + if (required.count(properties[index].name)) { + is_required[index] = true; + matched_min[index] = matched_min[index - 1] + 1; + } else { + matched_min[index] = matched_min[index - 1]; + } + if (!found_required) { + matched_min[index] = 1; + } + if (is_required[index]) { + found_required = true; + } + } + if (required.count(properties[0].name)) { + is_required[0] = true; + } + + if (max_properties == -1) { + // Case 2: With constraint on the lower bound of the properties number + std::vector> tails(property_count); + matched_min.back() = allow_additional ? std::max(1, matched_min.back()) + : std::max(min_properties, matched_min.back()); + for (int index = property_count - 2; index >= 0; --index) { + matched_min[index] = std::max(matched_min[index], matched_min[index + 1] - 1); + } + + for (int matched = matched_min.back(); matched <= property_count; ++matched) { + int32_t body = allow_additional ? GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + matched, + rule_name + ) + : Empty(); + if (allow_additional) { + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(property_count - 1) + "_" + + std::to_string(matched), + body + ); + tails.back().push_back(RuleRef(tail_rule_id)); + } else { + tails.back().push_back(body); + } + } + + for (int index = property_count - 2; index >= 0; --index) { + for (int matched = matched_min[index]; matched <= index + 1; ++matched) { + int32_t with_property = Sequence( + {middle_separator, + property_patterns[index + 1], + tails[index + 1][matched + 1 - matched_min[index + 1]]} + ); + int32_t body = + (is_required[index + 1] || matched == matched_min[index + 1] - 1) + ? with_property + : Choice({tails[index + 1][matched - matched_min[index + 1]], with_property}); + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(index) + "_" + std::to_string(matched), body + ); + tails[index].push_back(RuleRef(tail_rule_id)); + } + } + + std::vector choices; + for (int index = 0; index < property_count; ++index) { + if (matched_min[index] > 1) { + break; + } + choices.push_back(Sequence({property_patterns[index], tails[index][1 - matched_min[index]]})); + if (is_required[index]) { + break; + } + } + if (allow_additional && required.empty()) { + choices.push_back(Sequence( + {get_additional_pattern(), + GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + 1, + rule_name + )} + )); + } + return Sequence({first_separator, Choice(choices), last_separator}); + } + + // Case 3: With constraints on both lower & upper bound of the properties number + std::vector> tails(property_count); + std::vector matched_max(property_count, property_count); + matched_max[0] = 1; + for (int index = 1; index < property_count; ++index) { + matched_max[index] = matched_max[index - 1] + 1; + } + matched_min.back() = allow_additional ? std::max(1, matched_min.back()) + : std::max(min_properties, matched_min.back()); + matched_max.back() = std::min(max_properties, matched_max.back()); + for (int index = property_count - 2; index >= 0; --index) { + matched_min[index] = std::max(matched_min[index], matched_min[index + 1] - 1); + matched_max[index] = is_required[index + 1] + ? std::min(matched_max[index], matched_max[index + 1] - 1) + : std::min(matched_max[index], matched_max[index + 1]); + } + + for (int matched = matched_min.back(); matched <= matched_max.back(); ++matched) { + int32_t body = allow_additional ? GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + matched, + rule_name + ) + : Empty(); + if (allow_additional) { + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(property_count - 1) + "_" + std::to_string(matched), + body + ); + tails.back().push_back(RuleRef(tail_rule_id)); + } else { + tails.back().push_back(body); + } + } + + for (int index = property_count - 2; index >= 0; --index) { + for (int matched = matched_min[index]; matched <= matched_max[index]; ++matched) { + int32_t body; + if (matched == matched_max[index + 1]) { + body = tails[index + 1][matched - matched_min[index + 1]]; + } else { + int32_t with_property = Sequence( + {middle_separator, + property_patterns[index + 1], + tails[index + 1][matched + 1 - matched_min[index + 1]]} + ); + body = (is_required[index + 1] || matched == matched_min[index + 1] - 1) + ? with_property + : Choice({tails[index + 1][matched - matched_min[index + 1]], with_property}); + } + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(index) + "_" + std::to_string(matched), body + ); + tails[index].push_back(RuleRef(tail_rule_id)); + } + } + + std::vector choices; + for (int index = 0; index < property_count; ++index) { + if (matched_max[index] < matched_min[index]) { + continue; + } + if (matched_min[index] > 1) { + break; + } + choices.push_back(Sequence({property_patterns[index], tails[index][1 - matched_min[index]]})); + if (is_required[index]) { + break; + } + } + if (allow_additional && required.empty()) { + choices.push_back(Sequence( + {get_additional_pattern(), + GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + 1, + rule_name + )} + )); + } + return Sequence({first_separator, Choice(choices), last_separator}); +} + +int32_t JSONSchemaConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool need_braces +) { + // Determine additional property handling + std::string additional_suffix; + SchemaSpecPtr additional_property; + if (spec.allow_additional_properties && spec.additional_properties_schema) { + additional_suffix = "addl"; + additional_property = spec.additional_properties_schema; + } else if (spec.allow_unevaluated_properties && spec.unevaluated_properties_schema) { + additional_suffix = "uneval"; + additional_property = spec.unevaluated_properties_schema; + } else if (spec.allow_additional_properties || spec.allow_unevaluated_properties) { + additional_suffix = "addl"; + additional_property = SchemaSpec::Make(AnySpec{}, "", "any"); + } + + indent_manager_.StartIndent(); + bool has_content = false; + bool could_be_empty = false; + int32_t content = Empty(); + + // Build a key rule through GenerateString rather than spelling out a JSON string here. At the + // JSON root this still produces `"key"`, while XML-style converters override GenerateString to + // produce the unquoted key body expected inside their parameter wrappers. + auto create_pattern_key_rule = [&](const std::string& pattern, + const std::string& rule_name_hint) -> int32_t { + StringSpec key_spec; + key_spec.pattern = pattern; + return CreateRule( + SchemaSpec::Make(std::move(key_spec), /*cache_key=*/"", rule_name_hint), rule_name_hint + ); + }; + + if (!spec.properties.empty() && (!spec.pattern_properties.empty() || spec.property_names)) { + // Case 1a: properties coexist with patternProperties and/or propertyNames. + // Use GetPartialRuleForProperties for named properties, and build + // patternProperties/propertyNames as the additional property pattern override. + SchemaSpecPtr effective_additional = additional_property; + std::string effective_suffix = additional_suffix; + std::optional additional_override; + + if (!spec.pattern_properties.empty()) { + // Build patternProperties as additional property alternatives + std::vector patterns; + for (size_t index = 0; index < spec.pattern_properties.size(); ++index) { + const auto& pattern_property = spec.pattern_properties[index]; + std::string pattern_suffix = "pp_" + std::to_string(index); + int32_t key_rule_id = create_pattern_key_rule( + pattern_property.pattern, rule_name + "_" + pattern_suffix + "_key" + ); + int32_t value_rule_id = + CreateRule(pattern_property.schema, rule_name + "_" + pattern_suffix); + patterns.push_back(FormatOtherProperty( + RuleRef(key_rule_id), value_rule_id, rule_name, pattern_suffix, pattern_property.schema + )); + } + // Merge with existing additionalProperties if present + if (effective_additional) { + int32_t value_rule_id = + CreateRule(effective_additional, rule_name + "_" + effective_suffix); + patterns.push_back(FormatOtherProperty( + KeyPatternExpression(), value_rule_id, rule_name, effective_suffix, effective_additional + )); + } + additional_override = Choice(patterns); + if (!effective_additional) { + effective_additional = SchemaSpec::Make(AnySpec{}, "", "any"); + } + effective_suffix = "pp"; + } else if (spec.property_names && effective_additional) { + // propertyNames constrains keys of additional properties. + // Only apply when additional properties are allowed - when additionalProperties + // is false, no extra keys beyond named properties should be permitted. + int32_t key_rule_id = CreateRule(spec.property_names, rule_name + "_name"); + int32_t value_rule_id = CreateRule(effective_additional, rule_name + "_" + effective_suffix); + additional_override = FormatOtherProperty( + RuleRef(key_rule_id), + value_rule_id, + rule_name, + /*rule_name_suffix=*/"pn", + effective_additional + ); + effective_suffix = "pn"; + } + + content = GetPartialRuleForProperties( + spec.properties, + spec.required, + effective_additional, + rule_name, + effective_suffix, + spec.min_properties, + spec.max_properties, + additional_override + ); + has_content = spec.max_properties != 0; + could_be_empty = spec.required.empty() && spec.min_properties == 0; + } else if (!spec.pattern_properties.empty() || spec.property_names) { + // Case 1b: patternProperties or propertyNames without named properties + if (spec.max_properties != 0) { + int32_t beginning_separator = NextSeparatorExpression(); + std::vector property_choices; + if (!spec.pattern_properties.empty()) { + for (size_t index = 0; index < spec.pattern_properties.size(); ++index) { + const auto& pattern_property = spec.pattern_properties[index]; + std::string pattern_suffix = "prop_" + std::to_string(index); + int32_t key_rule_id = create_pattern_key_rule( + pattern_property.pattern, rule_name + "_" + pattern_suffix + "_key" + ); + int32_t value_rule_id = + CreateRule(pattern_property.schema, rule_name + "_" + pattern_suffix); + property_choices.push_back(Sequence( + {beginning_separator, + FormatOtherProperty( + RuleRef(key_rule_id), + value_rule_id, + rule_name, + pattern_suffix, + pattern_property.schema + )} + )); + } + } else { + int32_t key_rule_id = CreateRule(spec.property_names, rule_name + "_name"); + // propertyNames constrains only the key, so a typed additionalProperties + // schema still applies to the value (issue #826). + int32_t value_rule_id; + if (additional_property) { + value_rule_id = CreateRule(additional_property, rule_name + "_" + additional_suffix); + } else { + value_rule_id = builder_.GetRuleId(GetBasicAnyRuleName()); + XGRAMMAR_DCHECK(value_rule_id != -1); + } + property_choices.push_back(Sequence( + {beginning_separator, + FormatOtherProperty( + RuleRef(key_rule_id), + value_rule_id, + rule_name, + /*rule_name_suffix=*/"pn", + additional_property + )} + )); + } + + int32_t property_rule_id = + builder_.AddRuleWithHint(rule_name + "_prop", Choice(property_choices)); + int32_t subsequent_property = + Sequence({NextSeparatorExpression(), RuleRef(property_rule_id)}); + content = Sequence( + {RuleRef(property_rule_id), + GetPropertyWithNumberConstraints( + subsequent_property, spec.min_properties, spec.max_properties, 1, rule_name + ), + NextSeparatorExpression(true)} + ); + has_content = true; + could_be_empty = spec.min_properties == 0; + } else { + could_be_empty = true; + } + } else if (!spec.properties.empty()) { + // Case 2: properties defined (no patternProperties/propertyNames) + content = GetPartialRuleForProperties( + spec.properties, + spec.required, + additional_property, + rule_name, + additional_suffix, + spec.min_properties, + spec.max_properties + ); + has_content = spec.max_properties != 0; + could_be_empty = spec.required.empty() && spec.min_properties == 0; + } else if (additional_property) { + // Case 3: no properties defined, additional properties allowed + if (spec.max_properties != 0) { + int32_t value_rule_id = CreateRule(additional_property, rule_name + "_" + additional_suffix); + int32_t property = FormatOtherProperty( + KeyPatternExpression(), value_rule_id, rule_name, additional_suffix, additional_property + ); + content = Sequence( + {NextSeparatorExpression(), + property, + GetPropertyWithNumberConstraints( + Sequence({NextSeparatorExpression(), property}), + spec.min_properties, + spec.max_properties, + 1, + rule_name + ), + NextSeparatorExpression(true)} + ); + has_content = true; + } + could_be_empty = spec.min_properties == 0; + } else { + // Case 4: no properties, no additional properties, no pattern properties + // The object is unconditionally empty. + could_be_empty = true; + } + + indent_manager_.EndIndent(); + + int32_t result = need_braces ? Sequence({ByteString("{"), content, ByteString("}")}) : content; + if (could_be_empty) { + int32_t empty_content = any_whitespace_ ? WhitespaceExpression() : Empty(); + int32_t empty_result = + need_braces ? Sequence({ByteString("{"), empty_content, ByteString("}")}) : empty_content; + return has_content ? Choice({result, empty_result}) : empty_result; + } + return result; +} + +int32_t JSONSchemaConverter::GenerateAny(const AnySpec& spec, const std::string& rule_name) { + return Choice( + {RuleRef(kBasicNumber), + RuleRef(kBasicString), + RuleRef(kBasicBoolean), + RuleRef(kBasicNull), + RuleRef(kBasicArray), + RuleRef(kBasicObject)} + ); +} + +int32_t JSONSchemaConverter::GenerateConst(const ConstSpec& spec, const std::string& rule_name) { + return ByteString(spec.json_value); +} + +int32_t JSONSchemaConverter::GenerateEnum(const EnumSpec& spec, const std::string& rule_name) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& value : spec.json_values) { + values.push_back(ByteString(value)); + } + return Choice(values); +} + +SchemaSpecPtr JSONSchemaConverter::ResolveRefSchema( + const RefSpec& spec, const std::string& rule_name_hint +) { + if (!ref_resolver_) { + XGRAMMAR_LOG(FATAL) << "Ref resolver not set; cannot resolve $ref: " << spec.uri; + } + return ref_resolver_(spec.uri, rule_name_hint); +} + +int32_t JSONSchemaConverter::GenerateRef(const RefSpec& spec, const std::string& rule_name) { + // First check if we have a direct URI mapping (for circular references) + if (uri_to_rule_id_.count(spec.uri)) { + return RuleRef(uri_to_rule_id_[spec.uri]); + } + + // Derive rule name from URI path (like original URIToRule) so that the same + // $ref always gets the same rule name, and allocate before resolving to prevent + // dead recursion when the ref target contains a ref back. + std::string rule_name_hint = "ref"; + if (spec.uri.size() >= 2 && spec.uri[0] == '#' && spec.uri[1] == '/') { + std::string new_rule_name_prefix; + std::stringstream ss(spec.uri.substr(2)); + std::string part; + while (std::getline(ss, part, '/')) { + if (!part.empty()) { + if (!new_rule_name_prefix.empty()) { + new_rule_name_prefix += "_"; + } + for (char c : part) { + if (std::isalpha(static_cast(c)) || c == '_' || c == '-' || c == '.') { + new_rule_name_prefix += c; + } + } + } + } + if (!new_rule_name_prefix.empty()) { + rule_name_hint = std::move(new_rule_name_prefix); + } + } + + int32_t allocated_rule_id = builder_.AddEmptyRuleWithHint(rule_name_hint); + std::string allocated_rule_name = builder_.GetRule(allocated_rule_id).name; + uri_to_rule_id_[spec.uri] = allocated_rule_id; + SchemaSpecPtr resolved = ResolveRefSchema(spec, allocated_rule_name); + builder_.UpdateRuleBody(allocated_rule_id, GenerateFromSpec(resolved, allocated_rule_name)); + if (!resolved->cache_key.empty()) { + AddCache(resolved->cache_key, allocated_rule_id); + } + return RuleRef(allocated_rule_id); +} + +int32_t JSONSchemaConverter::GenerateAnyOf(const AnyOfSpec& spec, const std::string& rule_name) { + std::vector choices; + for (size_t index = 0; index < spec.options.size(); ++index) { + choices.push_back( + RuleRef(CreateRule(spec.options[index], rule_name + "_case_" + std::to_string(index))) + ); + } + return Choice(choices); +} + +int32_t JSONSchemaConverter::GenerateOneOf(const OneOfSpec& spec, const std::string& rule_name) { + std::vector choices; + for (size_t index = 0; index < spec.options.size(); ++index) { + choices.push_back( + RuleRef(CreateRule(spec.options[index], rule_name + "_case_" + std::to_string(index))) + ); + } + return Choice(choices); +} + +int32_t JSONSchemaConverter::GenerateAllOf(const AllOfSpec& spec, const std::string& rule_name) { + if (spec.schemas.size() == 1) { + return GenerateFromSpec(spec.schemas[0], rule_name + "_case_0"); + } + XGRAMMAR_LOG(WARNING) << "Support for allOf with multiple options is still ongoing"; + return GenerateFromSpec(SchemaSpec::Make(AnySpec{}, "", "any"), rule_name); +} + +int32_t JSONSchemaConverter::GenerateTypeArray( + const TypeArraySpec& spec, const std::string& rule_name +) { + std::vector choices; + for (size_t index = 0; index < spec.type_schemas.size(); ++index) { + choices.push_back( + RuleRef(CreateRule(spec.type_schemas[index], rule_name + "_type_" + std::to_string(index))) + ); + } + return Choice(choices); +} + +// ==================== Static Helper Methods ==================== + +std::optional JSONSchemaConverter::JSONFormatToRegexPattern(const std::string& format +) { + static const auto regex_map = []() -> std::unordered_map { + std::unordered_map m; + + std::string atext = "[\\w!#$%&'*+/=?^`{|}~-]"; + std::string dot_string = "(" + atext + "+(\\." + atext + "+)*)"; + std::string quoted_string = + "\\\\\"(\\\\[\\x20-\\x7E]|[\\x20\\x21\\x23-\\x5B\\x5D-\\x7E])*\\\\\""; + std::string domain = + "([A-Za-z0-9]([\\-A-Za-z0-9]*[A-Za-z0-9])?)((\\.[A-Za-z0-9][\\-A-Za-z0-9]*[A-Za-z0-9])*" + ")"; + m["email"] = "^(" + dot_string + "|" + quoted_string + ")@" + domain + "$"; + + m["date"] = "^(\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[1-2]\\d|3[01]))$"; + m["time"] = + "^([01]\\d|2[0-3]):[0-5]\\d:([0-5]\\d|60)(\\.\\d+)?(Z|[+-]([01]\\d|2[0-3]):[0-5]\\d)$"; + m["date-time"] = + "^(\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[1-2]\\d|3[01]))T([01]\\d|2[0-3]):[0-5]\\d:([0-5]\\d|60)(" + "\\.\\d+)?(Z|[+-]([01]\\d|2[0-3]):[0-5]\\d)$"; + m["duration"] = + "^P((\\d+D|\\d+M(\\d+D)?|\\d+Y(\\d+M(\\d+D)?)?)(T(\\d+S|\\d+M(\\d+S)?|\\d+H(\\d+M(\\d+" + "S)?" + ")?))?|T(\\d+S|\\d+M(\\d+S)?|\\d+H(\\d+M(\\d+S)?)?)|\\d+W)$"; + + std::string decbyte = "(25[0-5]|2[0-4]\\d|[0-1]?\\d?\\d)"; + m["ipv4"] = "^(" + decbyte + "\\.){3}" + decbyte + "$"; + + m["ipv6"] = + "(" + "([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}|" + "([0-9a-fA-F]{1,4}:){1,7}:|" + "([0-9a-fA-F]{1,4}:){1,6}:[0-9a-fA-F]{1,4}|" + "([0-9a-fA-F]{1,4}:){1,5}(:[0-9a-fA-F]{1,4}){1,2}|" + "([0-9a-fA-F]{1,4}:){1,4}(:[0-9a-fA-F]{1,4}){1,3}|" + "([0-9a-fA-F]{1,4}:){1,3}(:[0-9a-fA-F]{1,4}){1,4}|" + "([0-9a-fA-F]{1,4}:){1,2}(:[0-9a-fA-F]{1,4}){1,5}|" + "[0-9a-fA-F]{1,4}:((:[0-9a-fA-F]{1,4}){1,6})|" + ":((:[0-9a-fA-F]{1,4}){1,7}|:)|" + "::(ffff(:0{1,4}){0,1}:){0,1}" + "((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\\.){3,3}" + "(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])|" + "([0-9a-fA-F]{1,4}:){1,4}:" + "((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\\.){3,3}" + "(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])" + ")"; + + m["hostname"] = "^([a-z0-9]([a-z0-9-]*[a-z0-9])?)(\\.[a-z0-9]([a-z0-9-]*[a-z0-9])?)*$"; + m["uuid"] = "^[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}$"; + + std::string schema_pat = "[a-zA-Z][a-zA-Z+\\.-]*"; + std::string pchar = "([\\w\\.~!$&'()*+,;=:@-]|%[0-9A-Fa-f][0-9A-Fa-f])"; + std::string query_fragment_char = "([\\w\\.~!$&'()*+,;=:@/\\?-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + std::string query = "(\\?" + query_fragment_char + ")?"; + std::string fragment = "(#" + query_fragment_char + ")?"; + std::string path_abempty = "(/" + pchar + "*)*"; + std::string path_absolute_rootless_empty = "/?(" + pchar + "+(/" + pchar + "*)*)?"; + std::string userinfo = "([\\w\\.~!$&'()*+,;=:-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + std::string host = "([\\w\\.~!$&'()*+,;=-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + std::string authority = "(" + userinfo + "@)?" + host + "(:\\d*)?"; + std::string hier_part = + "(//" + authority + path_abempty + "|" + path_absolute_rootless_empty + ")"; + m["uri"] = "^" + schema_pat + ":" + hier_part + query + fragment + "$"; + + pchar = "([\\w\\.~!$&'()*+,;=:@-]|%[0-9A-Fa-f][0-9A-Fa-f])"; + query_fragment_char = "([\\w\\.~!$&'()*+,;=:@/\\?-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + query = "(\\?" + query_fragment_char + ")?"; + fragment = "(#" + query_fragment_char + ")?"; + path_abempty = "(/" + pchar + "*)*"; + std::string path_absolute = "/(" + pchar + "+(/" + pchar + "*)*)?"; + std::string segment_nz_nc = "([\\w\\.~!$&'()*+,;=@-]|%[0-9A-Fa-f][0-9A-Fa-f])+"; + std::string path_noscheme = segment_nz_nc + "(/" + pchar + "*)*"; + userinfo = "([\\w\\.~!$&'()*+,;=:-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + host = "([\\w\\.~!$&'()*+,;=-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + authority = "(" + userinfo + "@)?" + host + "(:\\d*)?"; + std::string relative_part = + "(//" + authority + path_abempty + "|" + path_absolute + "|" + path_noscheme + ")?"; + m["uri-reference"] = "^" + relative_part + query + fragment + "$"; + + std::string literals = + "([\\x21\\x23-\\x24\\x26\\x28-\\x3B\\x3D\\x3F-\\x5B\\x5D\\x5F\\x61-\\x7A\\x7E]" + "|%[0-9A-Fa-f][0-9A-Fa-f])"; + std::string op = "[+#\\./;\\?&=,!@|]"; + std::string varchar = "(\\w|%[0-9A-Fa-f][0-9A-Fa-f])"; + std::string varname = varchar + "(\\.?" + varchar + ")*"; + std::string varspec = varname + "(:[1-9]\\d?\\d?\\d?|\\*)?"; + std::string variable_list = varspec + "(," + varspec + ")*"; + std::string expression = "\\{(" + op + ")?" + variable_list + "\\}"; + m["uri-template"] = "^(" + literals + "|" + expression + ")*$"; + + m["json-pointer"] = "^(/([\\x00-\\x2E]|[\\x30-\\x7D]|[\\x7F-\\U0010FFFF]|~[01])*)*$"; + m["relative-json-pointer"] = + "^(0|[1-9][0-9]*)(#|(/([\\x00-\\x2E]|[\\x30-\\x7D]|[\\x7F-\\U0010FFFF]|~[01])*)*)$"; + + return m; + }(); + + auto it = regex_map.find(format); + if (it == regex_map.end()) { + return std::nullopt; + } + return it->second; +} + +// ==================== Range Regex Generation ==================== + +// Stateless utility that turns a numeric range into an anchored regex matching +// exactly the JSON integers / numbers inside it. Every method is static; the +// class exists only to group the helpers and keep the internal ones private. +class NumberGenerator { + public: + // Anchored regex matching every integer x with start <= x <= end. Either bound + // may be std::nullopt for an open side; an empty range yields "^()$". Bounds + // span the whole int64 range (|INT64_MIN| is handled without negation overflow). + static std::string IntegerRangeRegex(std::optional start, std::optional end); + + // Anchored regex matching every number in the range, written with up to + // `precision` fraction digits. `exclusive_start` / `exclusive_end` exclude the + // boundary value itself (turning >= / <= into > / <). Either bound may be + // std::nullopt for an open side; an empty range yields "^()$". + static std::string FloatRangeRegex( + std::optional start, + std::optional end, + int precision, + bool exclusive_start, + bool exclusive_end + ); + + private: + // Regex alternatives for the fraction digits following a decimal point. + struct FracPatternSet { + // Each pattern matches a non-empty fraction digit string. + std::vector parts; + // Whether having no fraction digits at all also satisfies the bound. + bool include_empty = false; + }; + + // --- Regex fragment primitives --- + static std::string DigitClass(char lo, char hi); // one digit in [lo, hi] (or \d) + static std::string ExactDigits(int k); // exactly k free digits: \d{k} + static std::string FreeDigits(int max_count); // 0..max_count free digits: \d{0,n} + static std::string OptionalZeros(int max_count); // 0..max_count zeros: 0{0,n} + static std::string SomeZeros(int max_count); // 1..max_count zeros: 0{1,n} + static bool AllChar(const std::string& s, char c); + + // --- Integer range (operate on non-negative decimal magnitude strings) --- + static std::string AbsDigits(int64_t v); + static int CompareDigitStr(const std::string& a, const std::string& b); + static std::vector IntSameLen(const std::string& a, const std::string& b); + static std::vector NumberPatternsStr(const std::string& lo, const std::string& hi); + static std::string SubRangeRegexStr(const std::string& lo, const std::string& hi); + static std::vector AtLeastPositivePatternsStr(const std::string& v_str); + + // --- Float range --- + static std::string FormatFloat(double value, int precision); + // Snaps a non-negative bound to the precision grid in the direction that keeps + // the range sound: a lower bound rounds up, an upper bound rounds down, so no + // out-of-range value is ever admitted. Returns the canonical grid string and, + // via strict_out, whether the boundary value must still be excluded. + static std::string RoundBoundToGrid( + double value, int precision, bool is_lower, bool strict_in, bool* strict_out + ); + // Adds (inc) or subtracts (!inc) one grid step (10^-precision) to a canonical + // non-negative decimal string, returning the canonical result. + static std::string AdjustGrid(const std::string& s, int precision, bool inc); + static void SplitDecimal(const std::string& s, std::string* int_part, std::string* frac_part); + static int CompareDecimal( + const std::string& int_a, + const std::string& frac_a, + const std::string& int_b, + const std::string& frac_b + ); + static std::string StripAnchors(const std::string& regex); + static int64_t ParseIntCapped(const std::string& digits); + static FracPatternSet FracGreaterPatterns(const std::string& s, bool strict, int max_len); + static FracPatternSet FracLessPatterns(const std::string& s, bool strict, int max_len); + static FracPatternSet FracBetweenPatterns( + const std::string& a, bool strict_a, const std::string& b, bool strict_b, int max_len + ); + static std::vector PositiveRangeParts( + const std::string& low, + bool strict_low, + const std::optional& high, + bool strict_high, + int precision + ); +}; + +// Helpers for integer range regex generation. They operate purely on +// fixed-length decimal digit strings (suffixes may carry leading zeros), so the +// patterns are correct by construction regardless of digit position. + +// A regex fragment matching a single digit in [lo, hi]. +std::string NumberGenerator::DigitClass(char lo, char hi) { + if (lo == hi) { + return std::string(1, lo); + } + if (lo == '0' && hi == '9') { + return "\\d"; + } + return "[" + std::string(1, lo) + "-" + std::string(1, hi) + "]"; +} + +// A regex fragment matching k free digits (each 0-9). Empty when k <= 0. +std::string NumberGenerator::ExactDigits(int k) { + if (k <= 0) { + return ""; + } + if (k == 1) { + return "\\d"; + } + return "\\d{" + std::to_string(k) + "}"; +} + +bool NumberGenerator::AllChar(const std::string& s, char c) { + return std::all_of(s.begin(), s.end(), [c](char ch) { return ch == c; }); +} + +// Patterns matching every equal-length digit string t with +// value(a) <= value(t) <= value(b). Requires a.size() == b.size() and +// value(a) <= value(b). Partitions t by its first digit: +// * first digit == a[0]: the suffix must be >= a's suffix (<= 99..9); +// * first digit strictly between a[0] and b[0]: the suffix is unconstrained; +// * first digit == b[0]: the suffix must be <= b's suffix (>= 00..0). +// The partition is exact and non-overlapping, so the union is sound and +// complete for [a, b]. +std::vector NumberGenerator::IntSameLen(const std::string& a, const std::string& b) { + int n = static_cast(a.size()); + if (a == b) { + return {a}; + } + if (n == 1) { + return {DigitClass(a[0], b[0])}; + } + if (a[0] == b[0]) { + std::vector res; + for (auto& p : IntSameLen(a.substr(1), b.substr(1))) { + res.push_back(std::string(1, a[0]) + p); + } + return res; + } + // a[0] < b[0] + std::string a_suf = a.substr(1); + std::string b_suf = b.substr(1); + if (AllChar(a_suf, '0') && AllChar(b_suf, '9')) { + // The whole suffix space is free: collapse to one box pattern. + if (a[0] == '0' && b[0] == '9') { + return {ExactDigits(n)}; + } + return {DigitClass(a[0], b[0]) + ExactDigits(n - 1)}; + } + std::vector res; + std::string nines(n - 1, '9'); + std::string zeros(n - 1, '0'); + for (auto& p : IntSameLen(a_suf, nines)) { + res.push_back(std::string(1, a[0]) + p); + } + if (b[0] - a[0] >= 2) { + res.push_back( + DigitClass(static_cast(a[0] + 1), static_cast(b[0] - 1)) + ExactDigits(n - 1) + ); + } + for (auto& p : IntSameLen(zeros, b_suf)) { + res.push_back(std::string(1, b[0]) + p); + } + return res; +} + +// Compares two non-negative decimal magnitude strings (no leading zeros except +// "0") by value. +int NumberGenerator::CompareDigitStr(const std::string& a, const std::string& b) { + if (a.size() != b.size()) { + return a.size() < b.size() ? -1 : 1; + } + if (a < b) { + return -1; + } + return a > b ? 1 : 0; +} + +// Patterns matching every integer whose magnitude has value in [lo, hi], where +// lo and hi are non-negative decimal magnitude strings (no leading zeros except +// "0"). An empty range (value(lo) > value(hi)) yields no patterns. Operating on +// strings keeps the whole int64 range representable, including +// |INT64_MIN| = 9223372036854775808, which does not fit in int64. +std::vector NumberGenerator::NumberPatternsStr( + const std::string& lo, const std::string& hi +) { + std::vector patterns; + if (CompareDigitStr(lo, hi) > 0) { + return patterns; + } + int lo_len = static_cast(lo.size()); + int hi_len = static_cast(hi.size()); + // Split [lo, hi] by digit length; each length yields a same-length segment + // handled exactly by IntSameLen. + for (int len = lo_len; len <= hi_len; ++len) { + std::string a_str = (len == lo_len) ? lo : ("1" + std::string(len - 1, '0')); + std::string b_str = (len == hi_len) ? hi : std::string(len, '9'); + for (auto& p : IntSameLen(a_str, b_str)) { + patterns.push_back(p); + } + } + return patterns; +} + +// Joins NumberPatternsStr alternatives into a parenthesised regex group. +std::string NumberGenerator::SubRangeRegexStr(const std::string& lo, const std::string& hi) { + std::vector patterns = NumberPatternsStr(lo, hi); + std::string joined; + for (size_t i = 0; i < patterns.size(); ++i) { + if (i > 0) { + joined += "|"; + } + joined += patterns[i]; + } + return "(" + joined + ")"; +} + +// Patterns matching every integer in [value(v_str), +infinity) for v_str a +// positive magnitude string (no leading zeros). Same-length values come from +// IntSameLen(v_str, 99..9); strictly longer values are any non-zero-led number. +std::vector NumberGenerator::AtLeastPositivePatternsStr(const std::string& v_str) { + int len = static_cast(v_str.size()); + std::vector res = IntSameLen(v_str, std::string(len, '9')); + res.push_back("[1-9]\\d{" + std::to_string(len) + ",}"); + return res; +} + +// The magnitude (absolute value) of v as a decimal string. Derived from the +// signed text rather than by negating v, so INT64_MIN is handled correctly. +std::string NumberGenerator::AbsDigits(int64_t v) { + std::string s = std::to_string(v); + return (!s.empty() && s[0] == '-') ? s.substr(1) : s; +} + +std::string NumberGenerator::IntegerRangeRegex( + std::optional start, std::optional end +) { + std::vector parts; + std::ostringstream result; + + if (!start && !end) { + return "^-?\\d+$"; + } + + if (start && !end) { + if (start.value() <= 0) { + if (start.value() < 0) { + // Negatives in [start, -1] are the magnitudes [1, |start|], negated. + parts.push_back("-" + SubRangeRegexStr("1", AbsDigits(start.value()))); + } + parts.push_back("0"); + parts.push_back("[1-9]\\d*"); + } else { + // x >= start with start > 0: same-length values >= start, plus every + // value with strictly more digits. + for (auto& p : AtLeastPositivePatternsStr(std::to_string(start.value()))) { + parts.push_back(p); + } + } + } + + if (!start && end) { + if (end.value() >= 0) { + parts.push_back("-[1-9]\\d*"); + parts.push_back("0"); + if (end.value() > 0) { + parts.push_back(SubRangeRegexStr("1", std::to_string(end.value()))); + } + } else { + // x <= end with end < 0: x = -a where a >= |end| > 0, so negate every + // pattern for the range [|end|, +infinity). + for (auto& p : AtLeastPositivePatternsStr(AbsDigits(end.value()))) { + parts.push_back("-" + p); + } + } + } + + if (start && end) { + int64_t range_start = start.value(); + int64_t range_end = end.value(); + + if (range_start > range_end) { + return "^()$"; + } + + if (range_start < 0) { + int64_t neg_start = range_start; + int64_t neg_end = std::min(static_cast(-1), range_end); + // Negatives in [neg_start, neg_end] are the magnitudes + // [|neg_end|, |neg_start|], negated. + parts.push_back("-" + SubRangeRegexStr(AbsDigits(neg_end), AbsDigits(neg_start))); + } + + if (range_start <= 0 && range_end >= 0) { + parts.push_back("0"); + } + + if (range_end > 0) { + int64_t pos_start = std::max(static_cast(1), range_start); + parts.push_back(SubRangeRegexStr(std::to_string(pos_start), std::to_string(range_end))); + } + } + + result << "^("; + for (size_t i = 0; i < parts.size(); ++i) { + if (i > 0) { + result << "|"; + } + result << parts[i]; + } + result << ")$"; + + return result.str(); +} + +std::string NumberGenerator::FormatFloat(double value, int precision) { + // Casting a double outside [INT64_MIN, INT64_MAX] (or NaN/Inf) to int64_t is + // undefined behavior, so range-check before the integer fast path. 2^63 == + // 9223372036854775808.0 is exactly representable and one past INT64_MAX, so the + // upper comparison must be strict. + if (value >= -9223372036854775808.0 && value < 9223372036854775808.0 && + value == static_cast(value)) { + return std::to_string(static_cast(value)); + } + + std::ostringstream oss; + oss << std::fixed << std::setprecision(precision) << value; + std::string result = oss.str(); + + size_t decimalPos = result.find('.'); + if (decimalPos != std::string::npos) { + size_t lastNonZero = result.find_last_not_of('0'); + if (lastNonZero != std::string::npos && lastNonZero > decimalPos) { + result.erase(lastNonZero + 1); + } else if (lastNonZero == decimalPos) { + result.erase(decimalPos); + } + } + + return result; +} + +std::string NumberGenerator::AdjustGrid(const std::string& s, int precision, bool inc) { + std::string int_part, frac_part; + SplitDecimal(s, &int_part, &frac_part); + // Build the scaled-integer numerator (value * 10^precision) as a digit string. + // Callers only pass FormatFloat output (<= precision fraction digits); guard + // the count so a longer string can never wrap the unsigned append count. + frac_part.append(std::max(0, precision - static_cast(frac_part.size())), '0'); + std::string num = int_part + frac_part; + + if (inc) { + int i = static_cast(num.size()) - 1; + for (; i >= 0 && num[i] == '9'; --i) { + num[i] = '0'; + } + if (i < 0) { + num.insert(num.begin(), '1'); + } else { + num[i]++; + } + } else { + int i = static_cast(num.size()) - 1; + for (; i >= 0 && num[i] == '0'; --i) { + num[i] = '9'; + } + if (i < 0) { + // Underflow below zero; clamp to zero (does not occur for the bounds the + // float pipeline feeds in, which are all >= one grid step when decremented). + num.assign(num.size(), '0'); + } else { + num[i]--; + } + } + + // Re-split into integer and `precision`-digit fraction, then canonicalize. + while (static_cast(num.size()) <= precision) { + num.insert(num.begin(), '0'); + } + std::string new_int = num.substr(0, num.size() - precision); + std::string new_frac = num.substr(num.size() - precision); + size_t nz = new_int.find_first_not_of('0'); + new_int = (nz == std::string::npos) ? "0" : new_int.substr(nz); + size_t lnz = new_frac.find_last_not_of('0'); + new_frac = (lnz == std::string::npos) ? "" : new_frac.substr(0, lnz + 1); + return new_frac.empty() ? new_int : new_int + "." + new_frac; +} + +std::string NumberGenerator::RoundBoundToGrid( + double value, int precision, bool is_lower, bool strict_in, bool* strict_out +) { + // FormatFloat rounds to the nearest grid point; if that lands exactly on the + // bound, keep the original strictness. Otherwise step to the grid point just + // inside the range so no out-of-range value is admitted, and the boundary is + // now strictly interior, so it becomes inclusive. + std::string r = FormatFloat(value, precision); + double rv = std::stod(r); + if (rv == value) { + *strict_out = strict_in; + return r; + } + *strict_out = false; + if (is_lower && rv < value) { + // Rounded below a lower bound: move up to the smallest grid point >= value. + r = AdjustGrid(r, precision, /*inc=*/true); + } else if (!is_lower && rv > value) { + // Rounded above an upper bound: move down to the largest grid point <= value. + r = AdjustGrid(r, precision, /*inc=*/false); + } + return r; +} + +// Helpers for GenerateFloatRangeRegex. Fraction patterns operate on the +// digit string after the decimal point, compared against a canonical bound +// fraction (canonical: produced by FormatFloat, so no trailing zeros). + +// Matches 0 to max_count free digits. +std::string NumberGenerator::FreeDigits(int max_count) { + if (max_count <= 0) { + return ""; + } + return "\\d{0," + std::to_string(max_count) + "}"; +} + +// Matches 0 to max_count zeros. +std::string NumberGenerator::OptionalZeros(int max_count) { + if (max_count <= 0) { + return ""; + } + return "0{0," + std::to_string(max_count) + "}"; +} + +// Matches 1 to max_count zeros. +std::string NumberGenerator::SomeZeros(int max_count) { + return "0{1," + std::to_string(max_count) + "}"; +} + +// Patterns for fraction strings t (1 <= |t| <= max_len) whose value 0.t is +// greater than 0.s (or equal when !strict). |s| <= max_len. +NumberGenerator::FracPatternSet NumberGenerator::FracGreaterPatterns( + const std::string& s, bool strict, int max_len +) { + FracPatternSet result; + int n = static_cast(s.size()); + // t agrees with s up to position i, then has a larger digit + for (int i = 0; i < n; ++i) { + if (s[i] < '9') { + result.parts.push_back( + s.substr(0, i) + DigitClass(s[i] + 1, '9') + FreeDigits(max_len - i - 1) + ); + } + } + // t extends s with a nonzero digit (after optional zeros) + for (int k = 0; n + k + 1 <= max_len; ++k) { + result.parts.push_back(s + std::string(k, '0') + "[1-9]" + FreeDigits(max_len - n - k - 1)); + } + if (!strict) { + // t has the same value as s: s plus optional trailing zeros + if (n > 0) { + result.parts.push_back(s + OptionalZeros(max_len - n)); + } else { + result.include_empty = true; + if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } + } + return result; +} + +// Patterns for fraction strings t (1 <= |t| <= max_len) whose value 0.t is +// less than 0.s (or equal when !strict). |s| <= max_len. +NumberGenerator::FracPatternSet NumberGenerator::FracLessPatterns( + const std::string& s, bool strict, int max_len +) { + FracPatternSet result; + int n = static_cast(s.size()); + // t agrees with s up to position i, then has a smaller digit + for (int i = 0; i < n; ++i) { + if (s[i] > '0') { + result.parts.push_back( + s.substr(0, i) + DigitClass('0', s[i] - 1) + FreeDigits(max_len - i - 1) + ); + } + } + // t is a proper prefix of s plus optional trailing zeros: strictly smaller, + // since the remaining digits of s contain a nonzero one + for (int i = 0; i < n; ++i) { + if (i == 0) { + if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } else { + result.parts.push_back(s.substr(0, i) + OptionalZeros(max_len - i)); + } + } + if (!strict) { + // t has the same value as s + if (n > 0) { + result.parts.push_back(s + OptionalZeros(max_len - n)); + } else if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } + result.include_empty = n > 0 || !strict; + return result; +} + +// Patterns for fraction strings t whose value 0.t lies between 0.a and 0.b. +// Requires value(0.a) < value(0.b) and b non-empty. +NumberGenerator::FracPatternSet NumberGenerator::FracBetweenPatterns( + const std::string& a, bool strict_a, const std::string& b, bool strict_b, int max_len +) { + FracPatternSet result; + // Longest common prefix of b and zero-padded a. Always stops before |b|: + // value(0.a) < value(0.b) implies b is not a prefix of padded a. + int common_len = 0; + while (common_len < static_cast(b.size()) && + (common_len < static_cast(a.size()) ? a[common_len] : '0') == b[common_len]) { + ++common_len; + } + std::string common = b.substr(0, common_len); + char digit_a = common_len < static_cast(a.size()) ? a[common_len] : '0'; + char digit_b = b[common_len]; + + // a digit strictly between the bounds' digits, then anything + if (digit_b - digit_a >= 2) { + result.parts.push_back( + common + DigitClass(digit_a + 1, digit_b - 1) + FreeDigits(max_len - common_len - 1) + ); + } + // lower boundary: t continues with digit_a, the rest must exceed a's suffix + if (common_len < static_cast(a.size())) { + FracPatternSet sub_lower = + FracGreaterPatterns(a.substr(common_len + 1), strict_a, max_len - common_len - 1); + for (auto& part : sub_lower.parts) { + result.parts.push_back(common + digit_a + std::move(part)); + } + if (sub_lower.include_empty) { + result.parts.push_back(common + std::string(1, digit_a)); + } + } else { + // a's value equals value(0.common): only nonzero extensions of + // common + digit_a ('0') are strictly greater + FracPatternSet sub_lower = FracGreaterPatterns("", true, max_len - common_len - 1); + for (auto& part : sub_lower.parts) { + result.parts.push_back(common + digit_a + std::move(part)); + } + if (!strict_a) { + // t has the same value as a + if (!a.empty()) { + result.parts.push_back(a + OptionalZeros(max_len - static_cast(a.size()))); + } else { + result.include_empty = true; + if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } + } + } + // upper boundary: t continues with digit_b, the rest must stay below b's suffix + FracPatternSet sub_upper = + FracLessPatterns(b.substr(common_len + 1), strict_b, max_len - common_len - 1); + for (auto& part : sub_upper.parts) { + result.parts.push_back(common + digit_b + std::move(part)); + } + if (sub_upper.include_empty) { + result.parts.push_back(common + std::string(1, digit_b)); + } + return result; +} + +// Splits a canonical decimal string from FormatFloat ("12" or "12.34") into +// integer and fraction parts. +void NumberGenerator::SplitDecimal( + const std::string& s, std::string* int_part, std::string* frac_part +) { + size_t dot = s.find('.'); + if (dot == std::string::npos) { + *int_part = s; + frac_part->clear(); + } else { + *int_part = s.substr(0, dot); + *frac_part = s.substr(dot + 1); + } +} + +// Compares the values of two canonical non-negative decimals. +int NumberGenerator::CompareDecimal( + const std::string& int_a, + const std::string& frac_a, + const std::string& int_b, + const std::string& frac_b +) { + if (int_a.size() != int_b.size()) { + return int_a.size() < int_b.size() ? -1 : 1; + } + if (int_a != int_b) { + return int_a < int_b ? -1 : 1; + } + size_t max_frac = std::max(frac_a.size(), frac_b.size()); + for (size_t i = 0; i < max_frac; ++i) { + char da = i < frac_a.size() ? frac_a[i] : '0'; + char db = i < frac_b.size() ? frac_b[i] : '0'; + if (da != db) { + return da < db ? -1 : 1; + } + } + return 0; +} + +// Strips the ^( )$ anchors added by IntegerRangeRegex, keeping the group. +std::string NumberGenerator::StripAnchors(const std::string& regex) { + return regex.substr(1, regex.size() - 2); +} + +int64_t NumberGenerator::ParseIntCapped(const std::string& digits) { + // `digits` is a canonical non-negative integer string (no leading zeros). + // Parse it exactly when it fits in int64; clamp to INT64_MAX otherwise (such + // magnitudes are beyond practical float bounds and double integer precision). + static const std::string kMaxInt64 = std::to_string(std::numeric_limits::max()); + if (digits.size() > kMaxInt64.size() || + (digits.size() == kMaxInt64.size() && digits > kMaxInt64)) { + return std::numeric_limits::max(); + } + return std::stoll(digits); +} + +// Patterns for unsigned decimals (integer part plus optional fraction of up +// to `precision` digits) within the given bounds. `low` is required and +// non-negative; `high` is optional. Patterns for the value 0 are never +// produced: when low's value is 0 the bound is treated as strict, and the +// caller emits the zero pattern itself. +std::vector NumberGenerator::PositiveRangeParts( + const std::string& low, + bool strict_low, + const std::optional& high, + bool strict_high, + int precision +) { + std::vector parts; + std::string int_low, frac_low; + SplitDecimal(low, &int_low, &frac_low); + if (int_low == "0" && frac_low.empty()) { + strict_low = true; + } + int64_t int_low_value = ParseIntCapped(int_low); + std::string opt_any_frac = "(\\.\\d{1," + std::to_string(precision) + "})?"; + + auto add_with_int_part = [&](const std::string& int_part, const FracPatternSet& set) { + for (const auto& part : set.parts) { + parts.push_back(int_part + "\\." + part); + } + if (set.include_empty) { + parts.push_back(int_part); + } + }; + + if (!high.has_value()) { + add_with_int_part(int_low, FracGreaterPatterns(frac_low, strict_low, precision)); + // Guard the +1 against int64 overflow (int_low_value may be clamped to + // INT64_MAX for very large bounds). + if (int_low_value < std::numeric_limits::max()) { + parts.push_back( + StripAnchors(IntegerRangeRegex(int_low_value + 1, std::nullopt)) + opt_any_frac + ); + } + return parts; + } + + std::string int_high, frac_high; + SplitDecimal(*high, &int_high, &frac_high); + int64_t int_high_value = ParseIntCapped(int_high); + int cmp = CompareDecimal(int_low, frac_low, int_high, frac_high); + if (cmp > 0 || (cmp == 0 && (strict_low || strict_high))) { + return parts; + } + if (cmp == 0) { + // single representable value, with optional redundant trailing zeros + if (frac_low.empty()) { + parts.push_back(int_low + "(\\." + SomeZeros(precision) + ")?"); + } else { + parts.push_back( + int_low + "\\." + frac_low + OptionalZeros(precision - static_cast(frac_low.size())) + ); + } + return parts; + } + if (int_low == int_high) { + add_with_int_part( + int_low, FracBetweenPatterns(frac_low, strict_low, frac_high, strict_high, precision) + ); + } else { + add_with_int_part(int_low, FracGreaterPatterns(frac_low, strict_low, precision)); + if (int_high_value - int_low_value >= 2) { + parts.push_back( + StripAnchors(IntegerRangeRegex(int_low_value + 1, int_high_value - 1)) + opt_any_frac + ); + } + add_with_int_part(int_high, FracLessPatterns(frac_high, strict_high, precision)); + } + return parts; +} + +std::string NumberGenerator::FloatRangeRegex( + std::optional start, + std::optional end, + int precision, + bool exclusive_start, + bool exclusive_end +) { + if (start && end) { + if (start.value() > end.value() || + (start.value() == end.value() && (exclusive_start || exclusive_end))) { + return "^()$"; + } + } + + if (!start && !end) { + return "^-?\\d+(\\.\\d{1," + std::to_string(precision) + "})?$"; + } + + std::vector parts; + + // Negative values: x is in [start, end] iff -x is in [-end, -start], so the + // positive-range patterns are reused on the negated bounds and prefixed + // with '-'. + bool negatives_in_range = !start.has_value() || start.value() < 0; + if (negatives_in_range) { + std::string low = "0"; + bool strict_low = true; + if (end.has_value() && end.value() < 0) { + low = + RoundBoundToGrid(-end.value(), precision, /*is_lower=*/true, exclusive_end, &strict_low); + } + std::optional high; + bool strict_high = false; + if (start.has_value()) { + high = RoundBoundToGrid( + -start.value(), precision, /*is_lower=*/false, exclusive_start, &strict_high + ); + } + for (auto& part : PositiveRangeParts(low, strict_low, high, strict_high, precision)) { + parts.push_back("-" + std::move(part)); + } + } + + bool zero_allowed = + (!start.has_value() || start.value() < 0 || (start.value() == 0 && !exclusive_start)) && + (!end.has_value() || end.value() > 0 || (end.value() == 0 && !exclusive_end)); + if (zero_allowed) { + parts.push_back("0(\\." + SomeZeros(precision) + ")?"); + // Negative zero written with an all-zero fraction ("-0.0".."-0.000000") also + // denotes 0. PositiveRangeParts never emits magnitude 0, so add these forms + // explicitly when the range covers the negative side. + if (negatives_in_range) { + parts.push_back("-0(\\." + SomeZeros(precision) + ")"); + } + } + + // Positive values + if (!end.has_value() || end.value() > 0) { + std::string low = "0"; + bool strict_low = true; + if (start.has_value() && start.value() > 0) { + low = RoundBoundToGrid( + start.value(), precision, /*is_lower=*/true, exclusive_start, &strict_low + ); + } + std::optional high; + bool strict_high = false; + if (end.has_value()) { + high = + RoundBoundToGrid(end.value(), precision, /*is_lower=*/false, exclusive_end, &strict_high); + } + for (auto& part : PositiveRangeParts(low, strict_low, high, strict_high, precision)) { + parts.push_back(std::move(part)); + } + } + + std::ostringstream result; + result << "^("; + for (size_t i = 0; i < parts.size(); ++i) { + if (i > 0) { + result << "|"; + } + result << parts[i]; + } + result << ")$"; + + return result.str(); +} + +std::string JSONSchemaConverter::GenerateRangeRegex( + std::optional start, std::optional end +) { + return NumberGenerator::IntegerRangeRegex(start, end); +} + +std::string JSONSchemaConverter::GenerateFloatRangeRegex( + std::optional start, + std::optional end, + int precision, + bool exclusive_start, + bool exclusive_end +) { + return NumberGenerator::FloatRangeRegex(start, end, precision, exclusive_start, exclusive_end); +} + +// ==================== Public API Functions ==================== + +std::optional JSONFormatFromString(const std::string& format) { + static const std::unordered_map kNameToFormat = { + {"json", JSONFormat::kJSON}, + {"qwen_xml", JSONFormat::kQwenXML}, + {"minimax_xml", JSONFormat::kMiniMaxXML}, + {"minimax_m3_xml", JSONFormat::kMiniMaxM3XML}, + {"deepseek_xml", JSONFormat::kDeepSeekXML}, + {"deepseek_v4_1_xml", JSONFormat::kDeepSeekV41XML}, + {"glm_xml", JSONFormat::kGlmXML}, + {"cohere_xml", JSONFormat::kCohereXML}, + {"kimi_k3_xml", JSONFormat::kKimiK3XML}, + }; + auto it = kNameToFormat.find(format); + if (it == kNameToFormat.end()) { + return std::nullopt; + } + return it->second; +} + +Grammar JSONSchemaToGrammar( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order, + JSONFormat json_format +) { + picojson::value schema_value; + std::string error = ParseJSON(schema_value, schema); + XGRAMMAR_CHECK(error.empty()) << "Failed to parse JSON: " << error + << ". The JSON string is:" << schema; + SchemaParser parser(schema_value, {strict_mode, json_format}); + auto spec_result = parser.Parse(schema_value, "root"); + if (spec_result.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(spec_result).UnwrapErr().what(); + } + auto spec = std::move(spec_result).Unwrap(); + auto ref_resolver = [&parser](const std::string& uri, const std::string& rule_name_hint) { + auto result = parser.ResolveRef(uri, rule_name_hint); + if (result.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(result).UnwrapErr().what(); + } + return std::move(result).Unwrap(); + }; + + switch (json_format) { + case JSONFormat::kJSON: { + JSONSchemaConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + any_order + ); + return converter.Convert(spec); + } + case JSONFormat::kQwenXML: + case JSONFormat::kMiniMaxXML: + case JSONFormat::kDeepSeekXML: + case JSONFormat::kDeepSeekV41XML: + case JSONFormat::kGlmXML: + case JSONFormat::kKimiK3XML: { + XMLToolCallingConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + json_format, + any_order + ); + return converter.Convert(spec); + } + case JSONFormat::kMiniMaxM3XML: { + MiniMaxM3XMLToolCallingConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + any_order + ); + return converter.Convert(spec); + } + case JSONFormat::kCohereXML: { + CohereXMLToolCallingConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + any_order + ); + return converter.Convert(spec); + } + default: + XGRAMMAR_LOG(FATAL) << "Invalid JSON format: " << static_cast(json_format); + } + XGRAMMAR_UNREACHABLE(); +} + +std::string JSONSchemaToEBNF( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + JSONFormat json_format, + bool any_order +) { + picojson::value schema_value; + std::string err = ParseJSON(schema_value, schema); + XGRAMMAR_CHECK(err.empty()) << "Failed to parse JSON: " << err + << ". The JSON string is:" << schema; + return JSONSchemaToEBNF( + schema_value, + any_whitespace, + indent, + separators, + strict_mode, + max_whitespace_cnt, + json_format, + any_order + ); +} + +std::string JSONSchemaToEBNF( + const picojson::value& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + JSONFormat json_format, + bool any_order +) { + // Parse JSON Schema to SchemaSpec + SchemaParser parser(schema, {strict_mode, json_format}); + auto spec_result = parser.Parse(schema, "root"); + if (spec_result.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(spec_result).UnwrapErr().what(); + } + auto spec = std::move(spec_result).Unwrap(); + + auto ref_resolver = [&parser](const std::string& uri, const std::string& rule_name_hint) { + auto r = parser.ResolveRef(uri, rule_name_hint); + if (r.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(r).UnwrapErr().what(); + } + return std::move(r).Unwrap(); + }; + + // Create converter based on format + switch (json_format) { + case JSONFormat::kJSON: { + JSONSchemaConverter converter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + case JSONFormat::kQwenXML: + case JSONFormat::kMiniMaxXML: + case JSONFormat::kDeepSeekXML: + case JSONFormat::kDeepSeekV41XML: + case JSONFormat::kGlmXML: + case JSONFormat::kKimiK3XML: { + XMLToolCallingConverter converter( + indent, + separators, + any_whitespace, + max_whitespace_cnt, + ref_resolver, + json_format, + any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + case JSONFormat::kMiniMaxM3XML: { + MiniMaxM3XMLToolCallingConverter converter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + case JSONFormat::kCohereXML: { + CohereXMLToolCallingConverter converter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + default: + XGRAMMAR_LOG(FATAL) << "Invalid JSON format: " << static_cast(json_format); + } + XGRAMMAR_UNREACHABLE(); +} + +// Wrapper functions for testing +std::string GenerateRangeRegex(std::optional start, std::optional end) { + return JSONSchemaConverter::GenerateRangeRegex(start, end); +} + +std::string GenerateFloatRangeRegex( + std::optional start, std::optional end, bool exclusive_start, bool exclusive_end +) { + return JSONSchemaConverter::GenerateFloatRangeRegex( + start, end, 6, exclusive_start, exclusive_end + ); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/json_schema_converter.h b/third_party/xgrammar/cpp/json_schema_converter.h new file mode 100644 index 0000000000..2b7a6e4c22 --- /dev/null +++ b/third_party/xgrammar/cpp/json_schema_converter.h @@ -0,0 +1,635 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter.h + * \brief Convert a JSON Schema directly to a grammar AST. + */ + +#ifndef XGRAMMAR_JSON_SCHEMA_CONVERTER_H_ +#define XGRAMMAR_JSON_SCHEMA_CONVERTER_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_builder.h" +#include "support/utils.h" + +namespace xgrammar { + +// ==================== SchemaSpec: Intermediate Representation for JSON Schema ==================== + +// Forward declaration +struct SchemaSpec; +using SchemaSpecPtr = std::shared_ptr; + +// Basic Type Specs +struct IntegerSpec { + std::optional minimum; + std::optional maximum; + std::optional exclusive_minimum; + std::optional exclusive_maximum; + std::optional multiple_of; + + std::string ToString() const; +}; + +struct NumberSpec { + std::optional minimum; + std::optional maximum; + std::optional exclusive_minimum; + std::optional exclusive_maximum; + + std::string ToString() const; +}; + +struct StringSpec { + std::optional pattern; + std::optional format; + int min_length = 0; + int max_length = -1; // -1 means no limit + + std::string ToString() const; +}; + +struct BooleanSpec { + std::string ToString() const; +}; + +struct NullSpec { + std::string ToString() const; +}; + +struct AnySpec { + std::string ToString() const; +}; + +// Complex Type Specs +struct ArraySpec { + std::vector prefix_items; + bool allow_additional_items = true; + SchemaSpecPtr additional_items; // nullptr means not allowed + int64_t min_items = 0; + int64_t max_items = -1; // -1 means no limit + + std::string ToString() const; +}; + +struct ObjectSpec { + struct Property { + std::string name; + SchemaSpecPtr schema; + }; + + struct PatternProperty { + std::string pattern; // regex pattern for key + SchemaSpecPtr schema; + }; + + std::vector properties; + std::vector pattern_properties; + std::unordered_set required; + + bool allow_additional_properties = false; + SchemaSpecPtr additional_properties_schema; + bool allow_unevaluated_properties = true; + SchemaSpecPtr unevaluated_properties_schema; + SchemaSpecPtr property_names; + + int min_properties = 0; + int max_properties = -1; // -1 means no limit + + std::string ToString() const; +}; + +// Composite Type Specs +struct ConstSpec { + std::string json_value; // JSON serialized value + + std::string ToString() const; +}; + +struct EnumSpec { + std::vector json_values; // JSON serialized values + + std::string ToString() const; +}; + +struct RefSpec { + std::string uri; + + std::string ToString() const; +}; + +struct AnyOfSpec { + std::vector options; + + std::string ToString() const; +}; + +struct OneOfSpec { + std::vector options; + + std::string ToString() const; +}; + +struct AllOfSpec { + std::vector schemas; + + std::string ToString() const; +}; + +struct TypeArraySpec { + // Handle "type": ["string", "integer"] cases + std::vector type_schemas; + + std::string ToString() const; +}; + +// Unified SchemaSpec +using SchemaSpecVariant = std::variant< + IntegerSpec, + NumberSpec, + StringSpec, + BooleanSpec, + NullSpec, + ArraySpec, + ObjectSpec, + AnySpec, + ConstSpec, + EnumSpec, + RefSpec, + AnyOfSpec, + OneOfSpec, + AllOfSpec, + TypeArraySpec>; + +struct SchemaSpec { + SchemaSpecVariant spec; + std::string cache_key; // for deduplication + std::string rule_name_hint; // suggested rule name + + std::string ToString() const; + + // Helper method to create SchemaSpec + template + static SchemaSpecPtr Make(T&& spec_value, std::string cache_key = "", std::string hint = "") { + auto ptr = std::make_shared(); + ptr->spec = std::forward(spec_value); + ptr->cache_key = std::move(cache_key); + ptr->rule_name_hint = std::move(hint); + return ptr; + } +}; + +// ==================== JSONFormat Enum ==================== + +enum class JSONFormat : int { + kJSON = 0, + kQwenXML = 1, + kMiniMaxXML = 2, + kDeepSeekXML = 3, + kGlmXML = 4, + kCohereXML = 5, + kKimiK3XML = 6, + kMiniMaxM3XML = 7, + kDeepSeekV41XML = 8, +}; + +/*! + * \brief Convert a format name to JSONFormat. + * \param format One of "json", "qwen_xml", "minimax_xml", "minimax_m3_xml", "deepseek_xml", + * "glm_xml", "cohere_xml", "kimi_k3_xml", or "deepseek_v4_1_xml". + * \return The corresponding JSONFormat, or std::nullopt if the name is not recognized. + */ +std::optional JSONFormatFromString(const std::string& format); + +/*! + * \brief Manage the rule generation cache. Wraps key-value cache for schema deduplication. + * The cached value is the rule id in the grammar builder. + */ +class GenerateCacheManager { + public: + /*! \brief Add a key-value pair to the cache. */ + void AddCache(const std::string& key, bool is_inner_layer, int32_t rule_id) { + cache_[{key, is_inner_layer}] = rule_id; + } + + /*! \brief Get cached rule id by key. Returns std::nullopt if not found. */ + std::optional GetCache(const std::string& key, bool is_inner_layer) const { + auto it = cache_.find({key, is_inner_layer}); + if (it != cache_.end()) { + return it->second; + } + return std::nullopt; + } + + private: + std::unordered_map, int32_t> cache_; +}; + +/*! + * \brief Manage the indent and separator for the generation of EBNF grammar. + */ +class IndentManager { + public: + IndentManager( + std::optional indent, + const std::string& separator, + bool any_whitespace, + std::optional max_whitespace_cnt + ); + + void StartIndent(); + void EndIndent(); + std::string StartSeparator(); + std::string MiddleSeparator(); + std::string EndSeparator(); + std::string EmptySeparator(); + std::string NextSeparator(bool is_end = false); + + private: + bool any_whitespace_; + bool enable_newline_; + int64_t indent_; + std::string separator_; + int64_t total_indent_; + std::vector is_first_; + std::optional max_whitespace_cnt_; + + friend class JSONSchemaConverter; +}; + +/*! + * \brief Convert SchemaSpec directly to a grammar AST. + * + * This is the base class for grammar generation. It generates JSON-format grammar by default. + * Subclasses can override virtual methods to generate different formats (e.g., XML). + */ +class JSONSchemaConverter { + public: + using RefResolver = + std::function; + + JSONSchemaConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + bool any_order = false + ); + + virtual ~JSONSchemaConverter() = default; + + /*! + * \brief Convert SchemaSpec directly to a grammar AST. + * \param spec The SchemaSpec to convert. + * \return The grammar AST. + */ + Grammar Convert(const SchemaSpecPtr& spec); + + protected: + using CharacterClassElement = GrammarBuilder::CharacterClassElement; + + // ==================== Virtual methods for generation ==================== + // Subclasses can override these to customize output format + + virtual int32_t GenerateInteger(const IntegerSpec& spec, const std::string& rule_name); + virtual int32_t GenerateNumber(const NumberSpec& spec, const std::string& rule_name); + virtual int32_t GenerateString(const StringSpec& spec, const std::string& rule_name); + virtual int32_t GenerateBoolean(const BooleanSpec& spec, const std::string& rule_name); + virtual int32_t GenerateNull(const NullSpec& spec, const std::string& rule_name); + virtual int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name); + virtual int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool need_brace = true + ); + virtual int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name); + virtual int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name); + virtual int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name); + virtual int32_t GenerateRef(const RefSpec& spec, const std::string& rule_name); + virtual int32_t GenerateAnyOf(const AnyOfSpec& spec, const std::string& rule_name); + virtual int32_t GenerateOneOf(const OneOfSpec& spec, const std::string& rule_name); + virtual int32_t GenerateAllOf(const AllOfSpec& spec, const std::string& rule_name); + virtual int32_t GenerateTypeArray(const TypeArraySpec& spec, const std::string& rule_name); + + // ==================== Hooks for customization ==================== + + /*! + * \brief Format a property key. Override for different formats. + * \param schema The schema of the property's value. Formats that encode the value's type + * next to the key (e.g. the Kimi-K3 `type` attribute) need it; the JSON format ignores it. + */ + virtual int32_t FormatPropertyKey(const std::string& key, const SchemaSpecPtr& schema); + + /*! \brief Format a property (key + value). Override for different formats. */ + virtual int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ); + + /*! + * \brief Format an "other" property (additional/unevaluated). Override for different formats. + * \param schema The schema selected for the property's value, or nullptr when the dynamic key + * has no single schema. Formats that encode the value's type next to the key need it. + */ + virtual int32_t FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema + ); + + /*! \brief Get the basic string rule name. Override for different formats. */ + virtual std::string GetKeyPattern() const; + + /*! \brief Get a key pattern that excludes specific property names. */ + virtual int32_t GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name + ); + + /*! \brief Get the basic any rule name. Override for different formats. */ + virtual std::string GetBasicAnyRuleName() const; + + /*! \brief Add basic rules for the format. Override for different formats. */ + virtual void AddBasicRules(); + void AddBasicRules(const std::vector& additional_rule_names); + + /*! \brief Add a key-value pair to the generation cache. Override for custom cache behavior. */ + virtual void AddCache(const std::string& key, int32_t rule_id); + + /*! \brief Get cached value by key. Returns std::nullopt if not found. */ + virtual std::optional GetCache(const std::string& key) const; + + // ==================== Helper methods (for subclasses to use) ==================== + + /*! \brief Dispatch to the appropriate Generate method based on spec type. */ + int32_t GenerateFromSpec(const SchemaSpecPtr& spec, const std::string& rule_name_hint); + + /*! \brief Create a rule and return the rule id (handles caching). */ + int32_t CreateRule(const SchemaSpecPtr& spec, const std::string& rule_name_hint); + + /*! \brief Resolve a reference to its parsed schema. */ + SchemaSpecPtr ResolveRefSchema(const RefSpec& spec, const std::string& rule_name_hint); + + /*! \brief Get next separator from indent manager. */ + virtual std::string NextSeparator(bool is_end = false); + + /*! \brief Get whitespace pattern. */ + std::string GetWhitespacePattern() const; + + int32_t Empty(); + int32_t ByteString(const std::string& value); + int32_t TagDispatch(bool loop_after_dispatch, std::vector excludes); + int32_t RuleRef(int32_t rule_id); + int32_t RuleRef(const std::string& rule_name); + int32_t Sequence(const std::vector& elements); + int32_t Choice(const std::vector& choices); + int32_t Repeat( + const std::string& rule_name_hint, int32_t expr_id, int32_t min_count, int32_t max_count + ); + int32_t AddSubGrammar(const Grammar& grammar); + + int32_t WhitespaceExpression(); + int32_t FormattingExpression(const std::string& expression); + int32_t NextSeparatorExpression(bool is_end = false); + int32_t KeyPatternExpression(); + + int32_t RegexExpression( + const std::string& regex, bool json_string = false, bool force_cfg_expansion = false + ); + + /*! \brief Helper to create rule with repetition constraints. */ + int32_t GetPropertyWithNumberConstraints( + int32_t pattern, + int min_properties, + int max_properties, + int already_repeated_times, + const std::string& rule_name + ); + + /*! \brief Generate partial rule for object properties. + * \param additional_property_override When set, used as the additional property + * pattern instead of the default GetKeyPattern() : value. This supports patternProperties + * and propertyNames constraints on additional keys. + */ + int32_t GetPartialRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override = std::nullopt + ); + + /*! \brief Generate the object rule in "any order" mode: an "item" alternation over all property + * keys, repeated between max(min_properties, required.size()) and max_properties times. Only the + * entry count is bounded, not which keys appear. + */ + int32_t GetAnyOrderRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override = std::nullopt + ); + + // ==================== Protected members ==================== + + GrammarBuilder builder_; + IndentManager indent_manager_; + int32_t colon_expr_id_; + bool any_whitespace_; + std::optional max_whitespace_cnt_; + // When true, object properties may appear in any order (see GetAnyOrderRuleForProperties). + // Applies to all objects (including nested ones). Default false preserves the fixed-order + // behavior. + bool any_order_ = false; + + public: + // Basic rule names + static const std::string kBasicAny; + static const std::string kBasicInteger; + static const std::string kBasicNumber; + static const std::string kBasicString; + static const std::string kBasicBoolean; + static const std::string kBasicNull; + static const std::string kBasicArray; + static const std::string kBasicObject; + static const std::string kBasicEscape; + static const std::string kBasicStringSub; + + protected: + GenerateCacheManager rule_cache_manager_; + + private: + void AddHelperRules(); + + std::unordered_map uri_to_rule_id_; // For circular reference handling + RefResolver ref_resolver_; // Resolves $ref URI to SchemaSpecPtr at generate time + + // Trie over property names, for key patterns that exclude specific properties + struct TrieNode { + bool is_terminal = false; + std::map children; + }; + int32_t BuildTrieBody(const TrieNode& node, const std::string& rule_name); + + // Reused grammar expression ids + std::optional empty_expr_id_; + std::unordered_map byte_string_expr_ids_; + std::unordered_map rule_ref_expr_ids_; + std::optional whitespace_expr_id_; + + // Helper for integer/number range regex generation + static std::string GenerateRangeRegex(std::optional start, std::optional end); + int32_t GenerateIntegerMultipleOfDFA(int64_t multiple_of, const std::string& rule_name); + static std::string GenerateFloatRangeRegex( + std::optional start, + std::optional end, + int precision = 6, + bool exclusive_start = false, + bool exclusive_end = false + ); + + protected: + static std::optional JSONFormatToRegexPattern(const std::string& format); + + // Expose for testing + friend std::string GenerateRangeRegex(std::optional start, std::optional end); + friend std::string GenerateFloatRangeRegex( + std::optional start, + std::optional end, + bool exclusive_start, + bool exclusive_end + ); +}; + +/*! + * \brief Convert a JSON Schema string directly to an unnormalized grammar AST. + * + * Callers that need normalized grammar should apply GrammarNormalizer after composing any + * subgrammars. + */ +Grammar JSONSchemaToGrammar( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + bool any_order = false, + JSONFormat json_format = JSONFormat::kJSON +); + +// ==================== Public API functions (backward compatible) ==================== + +/*! + * \brief Convert JSON schema string to EBNF grammar string. + * \param schema The JSON schema string. + * \param any_whitespace Whether to ignore the indentation restrictions, and allow any whitespace. + * Default: true. + * \param indent The number of spaces for indentation. If set to std::nullopt, the output will be + * in one line. Default: 2. + * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"}, + * {", ", ": "}. 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 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. + * \param max_whitespace_cnt The maximum number of whitespace characters for the whitespace + * which is used for indentation or JSON elements separation when any_whitespace is True. If + * std::nullopt, it means unlimited. Default: std::nullopt. + * \param json_format Define the root + * format of the object. If it's JSONFormat::kJSON, then it will generate a fully JSON-style + * grammar. If it's JSONFormat::kXML, then it will generate a grammar with the root format is + * XML-style, while the inner format is JSON-style. Default: JSONFormat::kJSON. + * \returns The EBNF grammar string. + */ + +std::string JSONSchemaToEBNF( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + JSONFormat json_format = JSONFormat::kJSON, + bool any_order = false +); + +/*! + * \brief Convert JSON schema string to EBNF grammar string. + * \param schema The JSON schema object. + * \param any_whitespace Whether to ignore the indentation restrictions, and allow any whitespace. + * Default: true. + * \param indent The number of spaces for indentation. If set to std::nullopt, the output will be + * in one line. Default: 2. + * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"}, + * {", ", ": "}. 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 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. + * \param max_whitespace_cnt The maximum number of whitespace characters for the whitespace + * which is used for indentation or JSON elements separation when any_whitespace is True. If + * std::nullopt, it means unlimited. Default: std::nullopt. + * \param json_format Define the root format of the object. If it's JSONFormat::kJSON, + * then it will generate a fully JSON-style grammar. If it's JSONFormat::kXML, then it will + * generate a grammar with the root format is XML-style, while the inner format is JSON-style. + * Default: JSONFormat::kJSON. + * \returns The EBNF grammar string. + */ +std::string JSONSchemaToEBNF( + const picojson::value& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + JSONFormat json_format = JSONFormat::kJSON, + bool any_order = false +); + +/*! + * \brief Generate regex pattern for integer/float range. + * \param start The start of the range (inclusive). If null assume negative infinity. + * \param end The end of the range (inclusive). If null assume infinity. + * \returns The regex pattern that matches integers/floats in the given range. + */ +std::string GenerateRangeRegex(std::optional start, std::optional end); + +std::string GenerateFloatRangeRegex( + std::optional start, + std::optional end, + bool exclusive_start = false, + bool exclusive_end = false +); + +} // namespace xgrammar + +#endif // XGRAMMAR_JSON_SCHEMA_CONVERTER_H_ diff --git a/third_party/xgrammar/cpp/json_schema_converter_ext.cc b/third_party/xgrammar/cpp/json_schema_converter_ext.cc new file mode 100644 index 0000000000..23b03e83ba --- /dev/null +++ b/third_party/xgrammar/cpp/json_schema_converter_ext.cc @@ -0,0 +1,1449 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter_ext.cc + * \brief Implementation of extended format converters. + */ +#include "json_schema_converter_ext.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/json_parse.h" +#include "support/logging.h" + +namespace xgrammar { + +namespace { + +constexpr std::array, 4> kCohereKeyEntities = { + std::pair{'&', "&"}, + std::pair{'<', "<"}, + std::pair{'>', ">"}, + std::pair{'"', """}, +}; + +std::string SerializeCohereKeyCodepoint(TCodepoint codepoint) { + for (const auto& [entity_codepoint, entity] : kCohereKeyEntities) { + if (codepoint == entity_codepoint) { + return entity; + } + } + return CharToUTF8(codepoint); +} + +std::vector ParseCohereKeyCodepoints(const std::string& key) { + XGRAMMAR_CHECK(key.find('\0') == std::string::npos) << "Cohere property names cannot contain NUL"; + auto codepoints = ParseUTF8(key.c_str()); + XGRAMMAR_CHECK(codepoints.size() != 1 || codepoints[0] != CharHandlingError::kInvalidUTF8) + << "Cohere property names must be valid UTF-8"; + return codepoints; +} + +std::string SerializeCohereKey(const std::string& key) { + std::string serialized; + for (TCodepoint codepoint : ParseCohereKeyCodepoints(key)) { + serialized += SerializeCohereKeyCodepoint(codepoint); + } + return serialized; +} + +template +std::vector CohereOrdinaryKeyRangesExcluding( + const Children& children +) { + constexpr TCodepoint kMaxUnicodeCodepoint = 0x10FFFF; + + // Ordinary key characters must not consume NUL, XML-sensitive characters (which are + // represented by entity alternatives), or a codepoint handled by a child trie branch. + std::vector excluded; + excluded.reserve(1 + kCohereKeyEntities.size() + children.size()); + excluded.push_back('\0'); + for (const auto& entry : kCohereKeyEntities) { + excluded.push_back(entry.first); + } + for (const auto& entry : children) { + excluded.push_back(entry.first); + } + + // Sorting makes duplicate exclusions adjacent so that unique + erase can remove them. + std::sort(excluded.begin(), excluded.end()); + excluded.erase(std::unique(excluded.begin(), excluded.end()), excluded.end()); + + // Build the positive character class as the gaps between excluded codepoints. Positive + // ranges preserve full Unicode support when the grammar is lowered to an FSM. + std::vector ranges; + TCodepoint range_start = 0; + for (TCodepoint codepoint : excluded) { + if (codepoint < range_start) { + continue; + } + if (range_start < codepoint) { + ranges.push_back({range_start, codepoint - 1}); + } + range_start = codepoint + 1; + } + if (range_start <= kMaxUnicodeCodepoint) { + ranges.push_back({range_start, kMaxUnicodeCodepoint}); + } + return ranges; +} + +constexpr const char* kStringCacheKey = "{\"type\":\"string\"}"; +constexpr const char* kObjectCacheKey = "{\"type\":\"object\"}"; +constexpr const char* kMiniMaxM3Namespace = "]<]minimax[>["; +constexpr const char* kMiniMaxM3ArrayItemName = "item"; + +bool IsASCIIWhitespace(uint8_t byte) { + return byte == ' ' || byte == '\t' || byte == '\n' || byte == '\r' || byte == '\f' || + byte == '\v'; +} + +bool IsCanonicalUTF8(const std::string& text) { + for (size_t offset = 0; offset < text.size();) { + auto [codepoint, num_bytes] = ParseNextUTF8(text.data() + offset); + if (codepoint == CharHandlingError::kInvalidUTF8 || num_bytes <= 0 || + offset + num_bytes > text.size() || (codepoint >= 0xd800 && codepoint <= 0xdfff) || + codepoint > 0x10ffff || text.compare(offset, num_bytes, CharToUTF8(codepoint)) != 0) { + return false; + } + offset += num_bytes; + } + return true; +} + +} // namespace + +MiniMaxM3XMLToolCallingConverter::MiniMaxM3XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + bool any_order +) + : JSONSchemaConverter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ) {} + +void MiniMaxM3XMLToolCallingConverter::AddBasicRules() { + for (const auto& name : {kBasicInteger, kBasicNumber, kBasicString, kBasicBoolean, kBasicNull}) { + builder_.AddEmptyRule(name); + } + + builder_.UpdateRuleBody( + kBasicInteger, JSONSchemaConverter::GenerateInteger(IntegerSpec{}, kBasicInteger) + ); + AddCache("{\"type\":\"integer\"}", builder_.GetRuleId(kBasicInteger)); + + builder_.UpdateRuleBody( + kBasicNumber, JSONSchemaConverter::GenerateNumber(NumberSpec{}, kBasicNumber) + ); + AddCache("{\"type\":\"number\"}", builder_.GetRuleId(kBasicNumber)); + + builder_.UpdateRuleBody(kBasicString, TagDispatch(false, {kMiniMaxM3Namespace})); + AddCache(kStringCacheKey, builder_.GetRuleId(kBasicString)); + + builder_.UpdateRuleBody( + kBasicBoolean, JSONSchemaConverter::GenerateBoolean(BooleanSpec{}, kBasicBoolean) + ); + AddCache("{\"type\":\"boolean\"}", builder_.GetRuleId(kBasicBoolean)); + + builder_.UpdateRuleBody(kBasicNull, JSONSchemaConverter::GenerateNull(NullSpec{}, kBasicNull)); + AddCache("{\"type\":\"null\"}", builder_.GetRuleId(kBasicNull)); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateString( + const StringSpec& spec, const std::string& rule_name +) { + const bool has_known_format = + spec.format.has_value() && JSONFormatToRegexPattern(*spec.format).has_value(); + XGRAMMAR_CHECK( + !spec.pattern.has_value() && !has_known_format && spec.min_length == 0 && + spec.max_length == -1 + ) << "String pattern, recognized format, and length constraints are not supported by " + "minimax_m3_xml because they cannot be combined with the namespace-marker exclusion"; + return RuleRef(kBasicString); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateArray( + const ArraySpec& spec, const std::string& rule_name +) { + constexpr int64_t kMaxRepeatCount = std::numeric_limits::max(); + XGRAMMAR_CHECK( + spec.min_items <= kMaxRepeatCount && + (spec.max_items == -1 || spec.max_items <= kMaxRepeatCount) && + spec.prefix_items.size() <= static_cast(kMaxRepeatCount) + ) << "minimax_m3_xml array bounds exceed the supported range"; + XGRAMMAR_CHECK(!spec.allow_additional_items || spec.additional_items != nullptr) + << "minimax_m3_xml requires a fixed schema for array items"; + + std::vector prefix_items; + prefix_items.reserve(spec.prefix_items.size()); + for (size_t index = 0; index < spec.prefix_items.size(); ++index) { + int32_t item_rule_id = + CreateRule(spec.prefix_items[index], rule_name + "_item_" + std::to_string(index)); + prefix_items.push_back(FormatElement(kMiniMaxM3ArrayItemName, item_rule_id)); + } + + std::optional additional_item; + if (spec.allow_additional_items) { + int32_t item_rule_id = CreateRule(spec.additional_items, rule_name + "_additional"); + additional_item = FormatElement(kMiniMaxM3ArrayItemName, item_rule_id); + } + + int32_t empty = Empty(); + int32_t whitespace = WhitespaceExpression(); + if (prefix_items.empty()) { + if (!additional_item.has_value() || spec.max_items == 0) { + return empty; + } + int32_t min_items = static_cast(spec.min_items); + int32_t max_items = spec.max_items == -1 ? -1 : static_cast(spec.max_items); + int32_t nonempty = Sequence( + {whitespace, + *additional_item, + Repeat( + rule_name + "_items", + Sequence({whitespace, *additional_item}), + std::max(0, min_items - 1), + max_items == -1 ? -1 : std::max(0, max_items - 1) + ), + whitespace} + ); + return min_items == 0 ? Choice({nonempty, empty}) : nonempty; + } + + int32_t prefix_count = static_cast(prefix_items.size()); + int32_t tail = empty; + if (additional_item.has_value()) { + int32_t min_additional = std::max(0, static_cast(spec.min_items) - prefix_count); + int32_t max_additional = spec.max_items == -1 + ? -1 + : std::max(0, static_cast(spec.max_items) - prefix_count); + tail = Repeat( + rule_name + "_additional_items", + Sequence({whitespace, *additional_item}), + min_additional, + max_additional + ); + } + + // A prefixItems entry constrains its position but does not make that position mandatory. Build + // a linear chain whose suffix can stop once minItems is satisfied. + for (int32_t index = prefix_count - 2; index >= 0; --index) { + int32_t emitted_count = index + 1; + bool can_stop = emitted_count >= spec.min_items; + bool can_continue = spec.max_items == -1 || emitted_count < spec.max_items; + int32_t body = empty; + if (can_continue) { + int32_t continuation = Sequence({whitespace, prefix_items[index + 1], tail}); + body = can_stop ? Choice({continuation, empty}) : continuation; + } + int32_t tail_rule_id = + builder_.AddRuleWithHint(rule_name + "_prefix_tail_" + std::to_string(index), body); + tail = RuleRef(tail_rule_id); + } + + if (spec.max_items == 0) { + return empty; + } + int32_t nonempty = Sequence({whitespace, prefix_items[0], tail, whitespace}); + return spec.min_items == 0 ? Choice({nonempty, empty}) : nonempty; +} + +void MiniMaxM3XMLToolCallingConverter::ValidateObject(const ObjectSpec& spec) const { + XGRAMMAR_CHECK( + !spec.allow_additional_properties && spec.additional_properties_schema == nullptr && + !spec.allow_unevaluated_properties && spec.unevaluated_properties_schema == nullptr && + spec.pattern_properties.empty() && spec.property_names == nullptr + ) << "minimax_m3_xml requires fixed object property names; additionalProperties, " + "unevaluatedProperties, patternProperties, and propertyNames are not supported"; + + std::unordered_set property_names; + for (const auto& property : spec.properties) { + XGRAMMAR_CHECK(property.schema != nullptr) + << "minimax_m3_xml property must have a fixed schema: " << property.name; + ValidateElementName(property.name); + property_names.insert(property.name); + } + for (const auto& required : spec.required) { + XGRAMMAR_CHECK(property_names.count(required) != 0) + << "minimax_m3_xml required property has no fixed schema: " << required; + } +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces +) { + ValidateObject(spec); + bool saved_any_whitespace = any_whitespace_; + any_whitespace_ = false; + int32_t result = JSONSchemaConverter::GenerateObject(spec, rule_name, /*need_braces=*/false); + any_whitespace_ = saved_any_whitespace; + return result; +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateAny( + const AnySpec& spec, const std::string& rule_name +) { + XGRAMMAR_LOG(FATAL) << "minimax_m3_xml does not support unconstrained schemas"; + XGRAMMAR_UNREACHABLE(); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateLiteral(const picojson::value& value) { + if (value.is()) { + const std::string& text = value.get(); + XGRAMMAR_CHECK(text.find(kMiniMaxM3Namespace) == std::string::npos) + << "A minimax_m3_xml string literal cannot contain the namespace marker"; + return ByteString(text); + } + if (value.is()) { + const auto& object = value.get(); + std::vector properties; + properties.reserve(object.size()); + for (const auto& key : object.ordered_keys()) { + int32_t value_expr = GenerateLiteral(object.at(key)); + int32_t value_rule_id = builder_.AddRuleWithHint("literal_" + key, value_expr); + properties.push_back(FormatElement(key, value_rule_id)); + } + return Sequence(properties); + } + if (value.is()) { + const auto& array = value.get(); + std::vector items; + items.reserve(array.size()); + for (size_t index = 0; index < array.size(); ++index) { + int32_t value_expr = GenerateLiteral(array[index]); + int32_t value_rule_id = + builder_.AddRuleWithHint("literal_item_" + std::to_string(index), value_expr); + items.push_back(FormatElement(kMiniMaxM3ArrayItemName, value_rule_id)); + } + return Sequence(items); + } + return ByteString(value.serialize()); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateConst( + const ConstSpec& spec, const std::string& rule_name +) { + picojson::value value; + std::string error = picojson::parse(value, spec.json_value); + XGRAMMAR_CHECK(error.empty()) << "Invalid const JSON value: " << error; + return GenerateLiteral(value); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateEnum( + const EnumSpec& spec, const std::string& rule_name +) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& json_value : spec.json_values) { + picojson::value value; + std::string error = picojson::parse(value, json_value); + XGRAMMAR_CHECK(error.empty()) << "Invalid enum JSON value: " << error; + values.push_back(GenerateLiteral(value)); + } + return Choice(values); +} + +void MiniMaxM3XMLToolCallingConverter::ValidateElementName(const std::string& name) { + XGRAMMAR_CHECK(!name.empty() && name.front() != '/' && name.find('>') == std::string::npos) + << "Invalid minimax_m3_xml element name: " << name; + XGRAMMAR_CHECK(IsCanonicalUTF8(name)) << "minimax_m3_xml element names must be valid UTF-8"; + XGRAMMAR_CHECK(std::any_of(name.begin(), name.end(), [](unsigned char byte) { + return !IsASCIIWhitespace(byte); + })) << "minimax_m3_xml element names cannot be blank"; +} + +int32_t MiniMaxM3XMLToolCallingConverter::FormatElement( + const std::string& name, int32_t value_rule_id +) { + ValidateElementName(name); + return Sequence( + {ByteString(std::string(kMiniMaxM3Namespace) + "<" + name + ">"), + RuleRef(value_rule_id), + ByteString(std::string(kMiniMaxM3Namespace) + "")} + ); +} + +int32_t MiniMaxM3XMLToolCallingConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + return FormatElement(key, value_rule_id); +} + +std::string MiniMaxM3XMLToolCallingConverter::NextSeparator(bool is_end) { + return GetWhitespacePattern(); +} + +// Static constants +const std::string XMLToolCallingConverter::kXMLString = "xml_string"; +const std::string XMLToolCallingConverter::kXMLAny = "xml_any"; +const std::string XMLToolCallingConverter::kXMLObject = "xml_object"; +const std::string XMLToolCallingConverter::kXMLVariableName = "xml_variable_name"; +const std::string CohereXMLToolCallingConverter::kCohereKey = "cohere_key"; +const std::string CohereXMLToolCallingConverter::kCohereAnyScalar = "cohere_any_scalar"; +const std::string CohereXMLToolCallingConverter::kCohereAnyList = "cohere_any_list"; +const std::unordered_map + XMLToolCallingConverter::kKeyWrapperMap = { + {JSONFormat::kQwenXML, {"", "", ""}}, + {JSONFormat::kMiniMaxXML, {"", "", ""}}, + {JSONFormat::kDeepSeekXML, + {"<｜DSML｜parameter name=\"", + "", + "", + // TODO(Linzhang): We do not validate the string's value, and we accept both. + ""}}, + {JSONFormat::kDeepSeekV41XML, + {"<｜DSML｜ parameter name=\"", "", "", ""}}, + {JSONFormat::kGlmXML, {"", "", "", ""}}, + {JSONFormat::kCohereXML, {"", "", ""}}, + {JSONFormat::kKimiK3XML, + {"<|open|>argument key=\"", + "", + "", + // The key suffix (type attribute and <|sep|>) is generated in XMLKeySuffix. + "<|close|>argument<|sep|>"}}, +}; + +XMLToolCallingConverter::XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + JSONFormat json_format, + bool any_order +) + : JSONSchemaConverter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ), + json_format_(json_format), + nested_object_level_(0), + xml_wrapper_(kKeyWrapperMap.at(json_format)) {} + +Grammar XMLToolCallingConverter::Convert(const SchemaSpecPtr& spec) { + nested_object_level_ = 0; + return JSONSchemaConverter::Convert(spec); +} + +std::string XMLToolCallingConverter::XMLValue(const std::string& json_value) const { + picojson::value value; + std::string error = ParseJSON(value, json_value); + if (error.empty() && value.is()) { + return value.get(); + } + return json_value; +} + +int32_t XMLToolCallingConverter::XMLKeySuffix(const std::optional& pinned_type) { + if (json_format_ == JSONFormat::kDeepSeekXML || json_format_ == JSONFormat::kDeepSeekV41XML) { + return Sequence( + {ByteString("\" string=\""), + Choice({ByteString("true"), ByteString("false")}), + ByteString("\">")} + ); + } + if (json_format_ == JSONFormat::kKimiK3XML) { + // A declared property carries exactly the type its value grammar is rendered with, so the + // parser decodes the value back to the schema's type. Free-form keys have no single schema + // type, so they keep the full set. + int32_t type_expr = pinned_type.has_value() ? ByteString(*pinned_type) + : Choice( + {ByteString("string"), + ByteString("number"), + ByteString("integer"), + ByteString("boolean"), + ByteString("object"), + ByteString("array"), + ByteString("null")} + ); + return Sequence({ByteString("\" type=\""), type_expr, ByteString("\"<|sep|>")}); + } + return ByteString(xml_wrapper_.key_wrapper_suffix); +} + +std::optional XMLToolCallingConverter::KimiK3TypeAttr(const SchemaSpecPtr& spec) { + if (spec == nullptr) { + return std::nullopt; + } + // The type name a single JSON value is rendered with, following the model's _xtml_type. + auto type_of_json_value = [](const std::string& json_value) -> std::optional { + picojson::value value; + if (!ParseJSON(value, json_value).empty()) { + return std::nullopt; + } + if (value.is()) return "string"; + if (value.is()) return "boolean"; + if (value.is()) return "number"; + if (value.is()) return "null"; + if (value.is()) return "object"; + if (value.is()) return "array"; + return std::nullopt; + }; + + return std::visit( + [&](auto&& arg) -> std::optional { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return "string"; + } else if constexpr (std::is_same_v || std::is_same_v) { + // _xtml_type renders every int and float as "number"; it never emits "integer". + return "number"; + } else if constexpr (std::is_same_v) { + return "boolean"; + } else if constexpr (std::is_same_v) { + return "null"; + } else if constexpr (std::is_same_v) { + return "array"; + } else if constexpr (std::is_same_v) { + return "object"; + } else if constexpr (std::is_same_v) { + return type_of_json_value(arg.json_value); + } else if constexpr (std::is_same_v) { + // Only pin the attribute when every alternative renders with the same type. + std::optional common; + for (const auto& json_value : arg.json_values) { + auto type_name = type_of_json_value(json_value); + if (!type_name.has_value()) return std::nullopt; + if (!common.has_value()) { + common = type_name; + } else if (*common != *type_name) { + return std::nullopt; + } + } + return common; + } else { + // Any, $ref and the combinators may render as more than one type; keep them open. + return std::nullopt; + } + }, + spec->spec + ); +} + +void XMLToolCallingConverter::AddBasicRules() { + // First add JSON basic rules. These should be in the inner layer of the XML format. + XGRAMMAR_DCHECK(nested_object_level_ == 0); + // The nested part, true json format, is at level 2. + nested_object_level_ = 2; + JSONSchemaConverter::AddBasicRules({kXMLString, kXMLAny, kXMLObject, kXMLVariableName}); + + auto any_spec = SchemaSpec::Make(AnySpec{}, "{}", kBasicAny); + + // The outer part, xml format, is at level 1. + nested_object_level_ = 1; + // Add XML string rule + builder_.UpdateRuleBody(kXMLString, TagDispatch(false, {xml_wrapper_.parameter_suffix})); + AddCache(kStringCacheKey, builder_.GetRuleId(kXMLString)); + + // Add XML any rule + builder_.UpdateRuleBody(kXMLAny, GenerateAny(AnySpec{}, kXMLAny)); + AddCache("{}", builder_.GetRuleId(kXMLAny)); + + // Reset the nested object level to 0, which is the root level. + nested_object_level_ = 0; + + // Add XML object rule + ObjectSpec xml_object_spec; + xml_object_spec.allow_additional_properties = true; + xml_object_spec.additional_properties_schema = any_spec; + builder_.UpdateRuleBody(kXMLObject, GenerateObject(xml_object_spec, kXMLObject)); + AddCache(kObjectCacheKey, builder_.GetRuleId(kXMLObject)); + + // Add XML variable name rule + builder_.UpdateRuleBody( + kXMLVariableName, + Sequence( + {builder_.AddCharacterClass({{'a', 'z'}, {'A', 'Z'}, {'_', '_'}}), + builder_.AddCharacterClassStar({{'a', 'z'}, {'A', 'Z'}, {'0', '9'}, {'_', '_'}})} + ) + ); +} + +std::string XMLToolCallingConverter::GetKeyPattern() const { + if (nested_object_level_ <= 1) { + return kXMLVariableName; + } + return kBasicString; +} + +std::string XMLToolCallingConverter::GetBasicAnyRuleName() const { + if (nested_object_level_ <= 1) { + return kXMLAny; + } + return kBasicAny; +} + +int32_t XMLToolCallingConverter::GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name +) { + if (nested_object_level_ <= 1) { + return RuleRef(GetKeyPattern()); + } + return JSONSchemaConverter::GetKeyPatternExcluding(properties, rule_name); +} + +std::string XMLToolCallingConverter::NextSeparator(bool is_end) { + if (nested_object_level_ <= 1) { + return GetWhitespacePattern(); + } + return JSONSchemaConverter::NextSeparator(is_end); +} + +int32_t XMLToolCallingConverter::GenerateString( + const StringSpec& spec, const std::string& rule_name +) { + if (nested_object_level_ <= 1) { + if (!spec.pattern.has_value() && !spec.format.has_value() && spec.min_length == 0 && + spec.max_length == -1) { + return RuleRef(kXMLString); + } + if (spec.format.has_value()) { + auto regex = JSONFormatToRegexPattern(*spec.format); + if (regex.has_value()) { + return RegexExpression(*regex, false, true); + } + } + if (spec.pattern.has_value()) { + return RegexExpression(*spec.pattern, false, /*force_cfg_expansion=*/true); + } + return Repeat( + rule_name + "_characters", + builder_.AddCharacterClass({{0, 0x10ffff}}), + spec.min_length, + spec.max_length + ); + } + return JSONSchemaConverter::GenerateString(spec, rule_name); +} + +int32_t XMLToolCallingConverter::GenerateAny(const AnySpec& spec, const std::string& rule_name) { + if (nested_object_level_ == 0) { + return RuleRef(kXMLObject); + } + if (nested_object_level_ == 1) { + return Choice({RuleRef(kXMLString), RuleRef(kBasicArray), RuleRef(kBasicObject)}); + } + return JSONSchemaConverter::GenerateAny(spec, rule_name); +} + +int32_t XMLToolCallingConverter::GenerateArray( + const ArraySpec& spec, const std::string& rule_name +) { + nested_object_level_++; + auto result = JSONSchemaConverter::GenerateArray(spec, rule_name); + nested_object_level_--; + return result; +} + +int32_t XMLToolCallingConverter::GenerateConst( + const ConstSpec& spec, const std::string& rule_name +) { + if (nested_object_level_ <= 1) { + return ByteString(XMLValue(spec.json_value)); + } + return JSONSchemaConverter::GenerateConst(spec, rule_name); +} + +int32_t XMLToolCallingConverter::GenerateEnum(const EnumSpec& spec, const std::string& rule_name) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + if (nested_object_level_ <= 1) { + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& value : spec.json_values) { + values.push_back(ByteString(XMLValue(value))); + } + return Choice(values); + } + return JSONSchemaConverter::GenerateEnum(spec, rule_name); +} + +std::string XMLToolCallingConverter::EscapeAttrValue(const std::string& value) const { + if (json_format_ != JSONFormat::kKimiK3XML) { + return value; + } + // Kimi-K3's renderer escapes attribute values with & -> & and " -> ". + std::string escaped; + escaped.reserve(value.size()); + for (char c : value) { + if (c == '&') { + escaped += "&"; + } else if (c == '"') { + escaped += """; + } else { + escaped += c; + } + } + return escaped; +} + +int32_t XMLToolCallingConverter::FormatPropertyKey( + const std::string& key, const SchemaSpecPtr& schema +) { + if (nested_object_level_ <= 1) { + // Only kimi_k3_xml encodes the value's type next to the key; the other formats would + // discard the result, so don't walk the schema for them. + std::optional pinned_type; + if (json_format_ == JSONFormat::kKimiK3XML) { + pinned_type = KimiK3TypeAttr(schema); + } + return Sequence( + {ByteString(xml_wrapper_.key_wrapper_prefix + EscapeAttrValue(key)), + XMLKeySuffix(pinned_type)} + ); + } + return JSONSchemaConverter::FormatPropertyKey(key, schema); +} + +int32_t XMLToolCallingConverter::FormatDeepSeekV41ParamSuffix( + const SchemaSpecPtr& schema, int32_t value_rule_id +) { + // Copy the name: creating alternative rules can reallocate the builder's rule storage. + std::string value_rule_name = builder_.GetRule(value_rule_id).name; + if (schema != nullptr) { + if (const auto* ref = std::get_if(&schema->spec); ref != nullptr) { + auto cached = deepseek_v41_param_ref_rules_.find(ref->uri); + if (cached != deepseek_v41_param_ref_rules_.end()) { + return RuleRef(cached->second); + } + // Cache the rule before descending through references or alternatives. A recursive + // branch then refers back to this rule, and shared acyclic subgraphs are built only once. + int32_t param_rule_id = builder_.AddEmptyRuleWithHint(value_rule_name + "_dsml_param"); + deepseek_v41_param_ref_rules_.emplace(ref->uri, param_rule_id); + auto resolved = ResolveRefSchema(*ref, value_rule_name); + builder_.UpdateRuleBody(param_rule_id, FormatDeepSeekV41ParamSuffix(resolved, value_rule_id)); + return RuleRef(param_rule_id); + } + } + + // string="true" wraps a raw string whose whitespace is part of the value; string="false" + // wraps a JSON value that may be padded with whitespace. + auto wrap = [&](int32_t value_expr, bool is_string) { + std::vector elements = { + ByteString(is_string ? "\" string=\"true\">" : "\" string=\"false\">") + }; + if (!is_string) elements.push_back(WhitespaceExpression()); + elements.push_back(value_expr); + if (!is_string) elements.push_back(WhitespaceExpression()); + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); + }; + + // A schema rendered with a single type keeps the value rule built by the caller. + std::optional pinned_type = KimiK3TypeAttr(schema); + if (pinned_type.has_value()) { + return wrap(RuleRef(value_rule_id), *pinned_type == "string"); + } + + // Unions and mixed enums get one alternative per option so each carries its own attribute. + std::vector options; + if (schema != nullptr) { + std::visit( + [&](const auto& spec) { + using T = std::decay_t; + if constexpr (std::is_same_v || std::is_same_v) { + options = spec.options; + } else if constexpr (std::is_same_v) { + options = spec.type_schemas; + } else if constexpr (std::is_same_v) { + if (spec.schemas.size() == 1) options = spec.schemas; + } else if constexpr (std::is_same_v) { + for (const auto& value : spec.json_values) { + options.push_back(SchemaSpec::Make(ConstSpec{value})); + } + } + }, + schema->spec + ); + } + if (options.empty()) { + // No schema, {} and allOf with several schemas all render any value. + return Choice( + {wrap(RuleRef(kXMLString), true), + wrap( + Choice( + {RuleRef(kBasicNumber), + RuleRef(kBasicBoolean), + RuleRef(kBasicNull), + RuleRef(kBasicArray), + RuleRef(kBasicObject)} + ), + false + )} + ); + } + std::vector choices; + for (size_t index = 0; index < options.size(); ++index) { + int32_t option_rule_id = + CreateRule(options[index], value_rule_name + "_dsml_case_" + std::to_string(index)); + choices.push_back(FormatDeepSeekV41ParamSuffix(options[index], option_rule_id)); + } + return Choice(choices); +} + +int32_t XMLToolCallingConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + if (nested_object_level_ <= 1) { + if (json_format_ == JSONFormat::kDeepSeekV41XML) { + return Sequence( + {ByteString(xml_wrapper_.key_wrapper_prefix + key), + FormatDeepSeekV41ParamSuffix(schema, value_rule_id)} + ); + } + std::vector elements = {FormatPropertyKey(key, schema)}; + if (!xml_wrapper_.value_wrapper_prefix.empty()) { + elements.push_back(WhitespaceExpression()); + elements.push_back(ByteString(xml_wrapper_.value_wrapper_prefix)); + } + // xml_string already accepts whitespace. Adding whitespace repetitions around it preserves the + // language but creates one Earley state for every possible split with the string body. + if (value_rule_id == builder_.GetRuleId(kXMLString)) { + elements.push_back(RuleRef(value_rule_id)); + } else { + elements.push_back(WhitespaceExpression()); + elements.push_back(RuleRef(value_rule_id)); + elements.push_back(WhitespaceExpression()); + } + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); + } + return JSONSchemaConverter::FormatProperty(key, value_rule_id, rule_name, idx, schema); +} + +int32_t XMLToolCallingConverter::FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema +) { + if (nested_object_level_ <= 1) { + if (json_format_ == JSONFormat::kDeepSeekV41XML) { + return Sequence( + {ByteString(xml_wrapper_.key_wrapper_prefix), + key_pattern_expr, + FormatDeepSeekV41ParamSuffix(schema, value_rule_id)} + ); + } + std::vector elements = { + ByteString(xml_wrapper_.key_wrapper_prefix), + key_pattern_expr, + XMLKeySuffix(json_format_ == JSONFormat::kKimiK3XML ? KimiK3TypeAttr(schema) : std::nullopt) + }; + if (!xml_wrapper_.value_wrapper_prefix.empty()) { + elements.push_back(WhitespaceExpression()); + elements.push_back(ByteString(xml_wrapper_.value_wrapper_prefix)); + } + if (value_rule_id == builder_.GetRuleId(kXMLString)) { + elements.push_back(RuleRef(value_rule_id)); + } else { + elements.push_back(WhitespaceExpression()); + elements.push_back(RuleRef(value_rule_id)); + elements.push_back(WhitespaceExpression()); + } + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); + } + return JSONSchemaConverter::FormatOtherProperty( + key_pattern_expr, value_rule_id, rule_name, rule_name_suffix, schema + ); +} + +int32_t XMLToolCallingConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces +) { + nested_object_level_++; + bool need_brace = nested_object_level_ > 1; + auto result = JSONSchemaConverter::GenerateObject(spec, rule_name, need_brace); + nested_object_level_--; + return result; +} + +void XMLToolCallingConverter::AddCache(const std::string& key, int32_t rule_id) { + if (key.empty()) { + return; + } + rule_cache_manager_.AddCache(key, nested_object_level_ > 1, rule_id); +} + +std::optional XMLToolCallingConverter::GetCache(const std::string& key) const { + if (key.empty()) { + return std::nullopt; + } + if (json_format_ == JSONFormat::kDeepSeekV41XML && nested_object_level_ == 0 && key == "{}") { + // Unconstrained tool arguments are an XML parameter list, not one parameter's raw value. + return rule_cache_manager_.GetCache(kObjectCacheKey, false); + } + // At level 0, {"type":"object"} is the root tool-arguments object and uses XML parameter + // tags. At level 1 it is the value of one such parameter and must use the inner JSON object + // rule, including braces. Without this distinction, the outer XML object cache is reused for + // the value before GenerateObject() can advance nested_object_level_. + if (nested_object_level_ == 1 && key == kObjectCacheKey) { + return rule_cache_manager_.GetCache(key, true); + } + return rule_cache_manager_.GetCache(key, nested_object_level_ > 1); +} + +CohereXMLToolCallingConverter::CohereXMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + bool any_order +) + : XMLToolCallingConverter( + indent, + separators, + any_whitespace, + max_whitespace_cnt, + ref_resolver, + JSONFormat::kCohereXML, + any_order + ) {} + +void CohereXMLToolCallingConverter::AddBasicRules() { + // Cohere's dynamic key and recursive Any rules must have stable targets before kXMLObject is + // built, because its additional-property formatting reaches them through virtual dispatch. + builder_.AddEmptyRule(kCohereKey); + builder_.AddEmptyRule(kCohereAnyScalar); + builder_.AddEmptyRule(kCohereAnyList); + + XMLToolCallingConverter::AddBasicRules(); + + builder_.UpdateRuleBody(kCohereKey, RegexExpression(R"(([^\x00"&<>]|&|<|>|")+)")); + + builder_.UpdateRuleBody( + kCohereAnyScalar, Choice({RuleRef(kBasicNumber), RuleRef(kBasicBoolean), RuleRef(kBasicNull)}) + ); + // kXMLObject already provides named_any_value* through its dynamic-property formatting. + // Lists need the corresponding recursive sequence of unnamed Any wrappers explicitly. + int32_t unnamed_any_value = FormatAnyCohereParam(std::nullopt, std::nullopt); + builder_.UpdateRuleBody( + kCohereAnyList, Repeat("cohere_any_list_items", unnamed_any_value, 0, -1) + ); +} + +bool CohereXMLToolCallingConverter::AtCohereRoot() const { + return nested_object_level_ == 0 && object_stack_.empty() && cohere_array_level_ == 0; +} + +bool CohereXMLToolCallingConverter::InCohereValueContext() const { + return nested_object_level_ <= 1 || !object_stack_.empty() || cohere_array_level_ > 0; +} + +int32_t CohereXMLToolCallingConverter::FormatCohereValue(int32_t value_rule_id) { + if (value_rule_id == builder_.GetRuleId(kXMLString)) { + return RuleRef(value_rule_id); + } + return Sequence({WhitespaceExpression(), RuleRef(value_rule_id), WhitespaceExpression()}); +} + +std::string CohereXMLToolCallingConverter::CohereTypeForJSONLiteral(const std::string& json_value) { + picojson::value value; + std::string error = ParseJSON(value, json_value); + // Const/enum object and array literals are emitted as JSON text today, not recursive Cohere + // dict/list bodies, so only JSON strings get the raw Cohere type. + return error.empty() && value.is() ? "raw" : "json"; +} + +std::optional CohereXMLToolCallingConverter::CommonCohereTypeForJSONLiterals( + const std::vector& json_values +) { + std::optional common_type; + for (const auto& json_value : json_values) { + auto type = CohereTypeForJSONLiteral(json_value); + if (!common_type.has_value()) { + common_type = type; + } else if (*common_type != type) { + return std::nullopt; + } + } + return common_type; +} + +int32_t CohereXMLToolCallingConverter::GetCohereTypePattern(const SchemaSpecPtr& schema) { + return std::visit( + [this](const auto& spec) -> int32_t { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return ByteString("raw"); + } else if constexpr (std::is_same_v) { + return ByteString("dict"); + } else if constexpr (std::is_same_v) { + return ByteString("list"); + } else if constexpr (std::is_same_v) { + return ByteString(CohereTypeForJSONLiteral(spec.json_value)); + } else if constexpr (std::is_same_v) { + auto common_type = CommonCohereTypeForJSONLiterals(spec.json_values); + // Mixed enums are branch-correlated by FormatCohereParam. This fallback is only used if + // a mixed enum somehow reaches the single-wrapper path, where there is no one true type. + return ByteString(common_type.has_value() ? *common_type : "json"); + } else { + return ByteString("json"); + } + }, + schema->spec + ); +} + +std::optional> CohereXMLToolCallingConverter::GetCohereCompositeOptions( + const SchemaSpecPtr& schema +) const { + if (schema == nullptr) { + return std::nullopt; + } + return std::visit( + [](const auto& spec) -> std::optional> { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return spec.options; + } else if constexpr (std::is_same_v) { + return spec.options; + } else if constexpr (std::is_same_v) { + if (spec.schemas.size() == 1) { + return spec.schemas; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + return spec.type_schemas; + } else if constexpr (std::is_same_v) { + if (spec.json_values.empty() || + CommonCohereTypeForJSONLiterals(spec.json_values).has_value()) { + return std::nullopt; + } + std::vector options; + options.reserve(spec.json_values.size()); + for (size_t index = 0; index < spec.json_values.size(); ++index) { + const auto& json_value = spec.json_values[index]; + ConstSpec const_spec; + const_spec.json_value = json_value; + options.push_back( + SchemaSpec::Make(std::move(const_spec), "", "enum_case_" + std::to_string(index)) + ); + } + return options; + } else { + return std::nullopt; + } + }, + schema->spec + ); +} + +int32_t CohereXMLToolCallingConverter::CohereParamPrefix( + const std::optional& name, const std::optional& key_pattern_expr +) { + std::vector elements = {ByteString(xml_wrapper_.key_wrapper_prefix)}; + if (name.has_value()) { + elements.push_back(ByteString(" name=\"" + SerializeCohereKey(*name) + "\"")); + } else if (key_pattern_expr.has_value()) { + elements.push_back(ByteString(" name=\"")); + elements.push_back(*key_pattern_expr); + elements.push_back(ByteString("\"")); + } + return Sequence(elements); +} + +int32_t CohereXMLToolCallingConverter::FormatCohereSuffixWithType( + int32_t type_expression, int32_t value_rule_id +) { + std::vector elements = { + ByteString(" type=\""), type_expression, ByteString("\"" + xml_wrapper_.key_wrapper_suffix) + }; + if (!xml_wrapper_.value_wrapper_prefix.empty()) { + elements.push_back(ByteString(xml_wrapper_.value_wrapper_prefix)); + } + elements.push_back(FormatCohereValue(value_rule_id)); + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); +} + +int32_t CohereXMLToolCallingConverter::FormatAnyCohereSuffix() { + // kXMLAny is the aggregate body-only union. Wrapping it under every type would create a + // type/body cross product, so each wrapper deliberately references its matching component. + return Choice( + {FormatCohereSuffixWithType(ByteString("raw"), builder_.GetRuleId(kXMLString)), + FormatCohereSuffixWithType(ByteString("json"), builder_.GetRuleId(kCohereAnyScalar)), + FormatCohereSuffixWithType(ByteString("dict"), builder_.GetRuleId(kXMLObject)), + FormatCohereSuffixWithType(ByteString("list"), builder_.GetRuleId(kCohereAnyList))} + ); +} + +int32_t CohereXMLToolCallingConverter::FormatAnyCohereParam( + const std::optional& name, const std::optional& key_pattern_expr +) { + return Sequence({CohereParamPrefix(name, key_pattern_expr), FormatAnyCohereSuffix()}); +} + +int32_t CohereXMLToolCallingConverter::FormatCohereParam( + const std::optional& name, + const std::optional& key_pattern_expr, + const SchemaSpecPtr& schema, + int32_t value_rule_id +) { + return Sequence( + {CohereParamPrefix(name, key_pattern_expr), FormatCohereParamSuffix(schema, value_rule_id)} + ); +} + +int32_t CohereXMLToolCallingConverter::FormatCohereParamSuffix( + const SchemaSpecPtr& schema, int32_t value_rule_id +) { + // Copy the name before generating: GenerateFromSpec may add rules and reallocate the + // builder's rule storage, invalidating references into it. + std::string value_rule_name = builder_.GetRule(value_rule_id).name; + if (const auto* ref = std::get_if(&schema->spec); ref != nullptr) { + auto cached = cohere_param_ref_rules_.find(ref->uri); + if (cached != cohere_param_ref_rules_.end()) { + return RuleRef(cached->second); + } + // Register the rule before resolving the reference. A schema that leads back to this URI, + // directly or through nested dict/list items, then reuses the rule instead of expanding + // again without bound, and shared acyclic references are built only once. + int32_t param_rule_id = builder_.AddEmptyRuleWithHint(value_rule_name + "_cohere_param"); + cohere_param_ref_rules_.emplace(ref->uri, param_rule_id); + SchemaSpecPtr resolved = ResolveRefSchema(*ref, value_rule_name); + builder_.UpdateRuleBody(param_rule_id, FormatCohereParamSuffix(resolved, value_rule_id)); + return RuleRef(param_rule_id); + } + + // CreateRule may return any aggregate rule (cached or freshly generated), but the schema is + // retained along this call path so Any can select correlated wrappers here. + if (std::holds_alternative(schema->spec)) { + return FormatAnyCohereSuffix(); + } + if (const auto* all_of = std::get_if(&schema->spec); + all_of != nullptr && all_of->schemas.size() != 1) { + // The base converter intentionally falls back to Any while multi-branch allOf support is + // incomplete. Keep that fallback canonical instead of wrapping its aggregate body as json. + return FormatAnyCohereSuffix(); + } + + auto options = GetCohereCompositeOptions(schema); + if (!options.has_value()) { + return FormatCohereSuffixWithType(GetCohereTypePattern(schema), value_rule_id); + } + + std::vector choices; + choices.reserve(options->size()); + for (size_t index = 0; index < options->size(); ++index) { + const SchemaSpecPtr& option = (*options)[index]; + int32_t option_rule_id = + CreateRule(option, value_rule_name + "_cohere_case_" + std::to_string(index)); + choices.push_back(FormatCohereParamSuffix(option, option_rule_id)); + } + return choices.size() == 1 ? choices[0] : Choice(choices); +} + +int32_t CohereXMLToolCallingConverter::GenerateString( + const StringSpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateString(spec, rule_name); + } + if (!spec.pattern.has_value() && !spec.format.has_value() && spec.min_length == 0 && + spec.max_length == -1) { + return RuleRef(kXMLString); + } + if (spec.format.has_value()) { + const std::string& format = *spec.format; + auto regex_pattern = JSONFormatToRegexPattern(format); + if (regex_pattern.has_value()) { + return RegexExpression(regex_pattern.value(), false, true); + } + } + if (spec.pattern.has_value()) { + return RegexExpression(*spec.pattern, false, /*force_cfg_expansion=*/true); + } + if (spec.min_length != 0 || spec.max_length != -1) { + return Repeat( + rule_name + "_characters", + builder_.AddCharacterClass({{0, 0x10ffff}}), + spec.min_length, + spec.max_length + ); + } + return JSONSchemaConverter::GenerateString(spec, rule_name); +} + +int32_t CohereXMLToolCallingConverter::GenerateAny( + const AnySpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateAny(spec, rule_name); + } + if (AtCohereRoot()) { + return RuleRef(kXMLObject); + } + return Choice( + {RuleRef(kXMLString), RuleRef(kCohereAnyScalar), RuleRef(kXMLObject), RuleRef(kCohereAnyList)} + ); +} + +int32_t CohereXMLToolCallingConverter::GenerateConst( + const ConstSpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateConst(spec, rule_name); + } + return ByteString(XMLValue(spec.json_value)); +} + +int32_t CohereXMLToolCallingConverter::GenerateEnum( + const EnumSpec& spec, const std::string& rule_name +) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateEnum(spec, rule_name); + } + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& value : spec.json_values) { + values.push_back(ByteString(XMLValue(value))); + } + return Choice(values); +} + +int32_t CohereXMLToolCallingConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces +) { + nested_object_level_++; + bool use_cohere_object = InCohereValueContext(); + + int32_t result; + if (use_cohere_object) { + SchemaSpecPtr additional_property; + if (spec.allow_additional_properties && spec.additional_properties_schema) { + additional_property = spec.additional_properties_schema; + } else if (spec.allow_unevaluated_properties && spec.unevaluated_properties_schema) { + additional_property = spec.unevaluated_properties_schema; + } else if (spec.allow_additional_properties || spec.allow_unevaluated_properties) { + additional_property = SchemaSpec::Make(AnySpec{}, "", "any"); + } + + object_stack_.push_back(&spec); + additional_property_stack_.push_back(additional_property); + result = JSONSchemaConverter::GenerateObject(spec, rule_name, false); + additional_property_stack_.pop_back(); + object_stack_.pop_back(); + } else { + result = JSONSchemaConverter::GenerateObject(spec, rule_name, nested_object_level_ > 1); + } + + nested_object_level_--; + return result; +} + +int32_t CohereXMLToolCallingConverter::GenerateArray( + const ArraySpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + nested_object_level_++; + auto result = JSONSchemaConverter::GenerateArray(spec, rule_name); + nested_object_level_--; + return result; + } + + cohere_array_level_++; + std::vector item_patterns; + for (size_t i = 0; i < spec.prefix_items.size(); ++i) { + int32_t item_rule_id = + CreateRule(spec.prefix_items[i], rule_name + "_item_" + std::to_string(i)); + item_patterns.push_back( + FormatCohereParam(std::nullopt, std::nullopt, spec.prefix_items[i], item_rule_id) + ); + } + + std::optional additional_item_pattern; + if (spec.allow_additional_items && spec.additional_items) { + int32_t additional_rule_id = CreateRule(spec.additional_items, rule_name + "_additional"); + additional_item_pattern = + FormatCohereParam(std::nullopt, std::nullopt, spec.additional_items, additional_rule_id); + } + cohere_array_level_--; + + if (item_patterns.empty()) { + if (!additional_item_pattern.has_value() || spec.max_items == 0) { + return Empty(); + } + return Repeat( + rule_name + "_items", + *additional_item_pattern, + static_cast(spec.min_items), + spec.max_items == -1 ? -1 : static_cast(spec.max_items) + ); + } + + int32_t prefix_part = Sequence(item_patterns); + if (!additional_item_pattern.has_value()) { + return prefix_part; + } + + int64_t min_additional = std::max( + static_cast(0), spec.min_items - static_cast(item_patterns.size()) + ); + int64_t max_additional = + spec.max_items == -1 ? -1 : spec.max_items - static_cast(item_patterns.size()); + return Sequence( + {prefix_part, + Repeat( + rule_name + "_additional_items", + *additional_item_pattern, + static_cast(min_additional), + max_additional == -1 ? -1 : static_cast(max_additional) + )} + ); +} + +int32_t CohereXMLToolCallingConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + if (!object_stack_.empty() && idx >= 0 && + idx < static_cast(object_stack_.back()->properties.size())) { + const auto& prop = object_stack_.back()->properties[idx]; + return FormatCohereParam(prop.name, std::nullopt, prop.schema, value_rule_id); + } + return XMLToolCallingConverter::FormatProperty(key, value_rule_id, rule_name, idx, schema); +} + +int32_t CohereXMLToolCallingConverter::FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema +) { + SchemaSpecPtr value_schema = schema; + if (!value_schema && !additional_property_stack_.empty()) { + value_schema = additional_property_stack_.back(); + } + if (!value_schema && InCohereValueContext()) { + value_schema = SchemaSpec::Make(AnySpec{}, "", "any"); + value_rule_id = CreateRule(value_schema, rule_name + "_" + rule_name_suffix + "_cohere_any"); + } + if (value_schema) { + return FormatCohereParam(std::nullopt, key_pattern_expr, value_schema, value_rule_id); + } + return XMLToolCallingConverter::FormatOtherProperty( + key_pattern_expr, value_rule_id, rule_name, rule_name_suffix, schema + ); +} + +std::string CohereXMLToolCallingConverter::GetKeyPattern() const { + if (InCohereValueContext()) { + return kCohereKey; + } + return JSONSchemaConverter::GetKeyPattern(); +} + +int32_t CohereXMLToolCallingConverter::BuildCohereKeyExcludingBody( + const CohereKeyTrieNode& node, int depth +) { + std::vector choices; + if (depth > 0 && !node.is_terminal) { + choices.push_back(Empty()); + } + + int32_t optional_key_suffix = Choice({Empty(), RuleRef(kCohereKey)}); + int32_t ordinary_key_unit = + builder_.AddCharacterClass(CohereOrdinaryKeyRangesExcluding(node.children)); + choices.push_back(Sequence({ordinary_key_unit, optional_key_suffix})); + for (const auto& [codepoint, entity] : kCohereKeyEntities) { + if (!node.children.count(codepoint)) { + int32_t entity_key_unit = ByteString(entity); + choices.push_back(Sequence({entity_key_unit, optional_key_suffix})); + } + } + + for (const auto& [codepoint, child] : node.children) { + choices.push_back(Sequence( + {ByteString(SerializeCohereKeyCodepoint(codepoint)), + BuildCohereKeyExcludingBody(child, depth + 1)} + )); + } + + return Choice(choices); +} + +int32_t CohereXMLToolCallingConverter::GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name +) { + if (InCohereValueContext()) { + if (properties.empty()) { + return RuleRef(GetKeyPattern()); + } + CohereKeyTrieNode root; + for (const auto& prop : properties) { + CohereKeyTrieNode* cur = &root; + auto codepoints = ParseCohereKeyCodepoints(prop.name); + for (TCodepoint codepoint : codepoints) { + cur = &cur->children[codepoint]; + } + if (!codepoints.empty()) { + cur->is_terminal = true; + } + } + int32_t key_rule_id = builder_.AddEmptyRuleWithHint(rule_name + "_cohere_addl_key"); + builder_.UpdateRuleBody(key_rule_id, BuildCohereKeyExcludingBody(root, 0)); + return RuleRef(key_rule_id); + } + return JSONSchemaConverter::GetKeyPatternExcluding(properties, rule_name); +} + +std::string CohereXMLToolCallingConverter::NextSeparator(bool is_end) { + if (InCohereValueContext()) { + return GetWhitespacePattern(); + } + return JSONSchemaConverter::NextSeparator(is_end); +} + +void CohereXMLToolCallingConverter::AddCache(const std::string& key, int32_t rule_id) { + if (key.empty()) { + return; + } + rule_cache_manager_.AddCache(key, nested_object_level_ > 1 && !InCohereValueContext(), rule_id); +} + +std::optional CohereXMLToolCallingConverter::GetCache(const std::string& key) const { + if (key.empty()) { + return std::nullopt; + } + // "true" and {} are equivalent schemas. At the tool-arguments root both are unrestricted + // dictionaries, while nested {} keeps using the aggregate Any body rule. + if (AtCohereRoot() && (key == "{}" || key == "true")) { + return builder_.GetRuleId(kXMLObject); + } + return rule_cache_manager_.GetCache(key, nested_object_level_ > 1 && !InCohereValueContext()); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/json_schema_converter_ext.h b/third_party/xgrammar/cpp/json_schema_converter_ext.h new file mode 100644 index 0000000000..1386225d18 --- /dev/null +++ b/third_party/xgrammar/cpp/json_schema_converter_ext.h @@ -0,0 +1,292 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter_ext.h + * \brief Extended format converters for JSON Schema, including XML Tool Calling format. + */ + +#ifndef XGRAMMAR_JSON_SCHEMA_CONVERTER_EXT_H_ +#define XGRAMMAR_JSON_SCHEMA_CONVERTER_EXT_H_ + +#include +#include +#include +#include +#include +#include + +#include "json_schema_converter.h" + +namespace xgrammar { + +/*! + * \brief Converter for MiniMax M3's recursive namespace-prefixed XML format. + * + * This initial implementation supports schemas whose object property names are + * known when the grammar is built. Schemas requiring runtime element names are + * rejected explicitly. + */ +class MiniMaxM3XMLToolCallingConverter : public JSONSchemaConverter { + public: + MiniMaxM3XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + bool any_order = false + ); + + protected: + int32_t GenerateString(const StringSpec& spec, const std::string& rule_name) override; + int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name) override; + int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces = false + ) override; + int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name) override; + int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name) override; + int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name) override; + + int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ) override; + std::string NextSeparator(bool is_end = false) override; + void AddBasicRules() override; + + private: + int32_t FormatElement(const std::string& name, int32_t value_rule_id); + int32_t GenerateLiteral(const picojson::value& value); + void ValidateObject(const ObjectSpec& spec) const; + static void ValidateElementName(const std::string& name); +}; + +/*! + * \brief Converter for XML Tool Calling format (e.g., Qwen style). + * + * This converter generates a grammar where: + * - The outermost object uses XML format: value + * - Inner values use standard JSON format + */ +class XMLToolCallingConverter : public JSONSchemaConverter { + public: + XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + JSONFormat json_format = JSONFormat::kQwenXML, + bool any_order = false + ); + + /*! \brief Convert SchemaSpec to grammar with XML format for root object. Note that this function + * is not thread-safe.*/ + Grammar Convert(const SchemaSpecPtr& spec); + + protected: + // Override methods for XML format + int32_t GenerateString(const StringSpec& spec, const std::string& rule_name) override; + int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces = false + ) override; + int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name) override; + int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name) override; + int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name) override; + int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name) override; + + // Override format hooks + int32_t FormatPropertyKey(const std::string& key, const SchemaSpecPtr& schema) override; + int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ) override; + int32_t FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema + ) override; + + std::string GetKeyPattern() const override; + std::string GetBasicAnyRuleName() const override; + int32_t GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name + ) override; + + std::string NextSeparator(bool is_end = false) override; + + void AddBasicRules() override; + + void AddCache(const std::string& key, int32_t rule_id) override; + std::optional GetCache(const std::string& key) const override; + + protected: + // Wrapper strings for XML parameter tags (key prefix/suffix, value prefix, closing suffix) + struct XMLWrapper { + std::string key_wrapper_prefix; + std::string key_wrapper_suffix; + std::string value_wrapper_prefix; + std::string parameter_suffix; + }; + + static const std::unordered_map kKeyWrapperMap; + static const std::string kXMLString; + static const std::string kXMLAny; + static const std::string kXMLObject; + static const std::string kXMLVariableName; + + std::string XMLValue(const std::string& json_value) const; + std::string EscapeAttrValue(const std::string& value) const; + + /*! + * \brief Return the Kimi-K3 `type` attribute a value of \p spec is rendered with, or + * std::nullopt if the schema does not pin down a single type (\p spec may be nullptr, which + * is how free-form keys end up unconstrained). + * + * The Kimi-K3 tool-call parser reads the attribute as a decoding switch: type="string" + * keeps the value as raw text, anything else JSON-decodes it. So the attribute must agree + * with the value grammar, otherwise the decoded argument changes type (e.g. a string + * property tagged type="number" with body 123 decodes to the integer 123). Mirrors the + * model's renderer (_xtml_type), which maps both ints and floats to "number". + */ + static std::optional KimiK3TypeAttr(const SchemaSpecPtr& spec); + + /*! + * \brief Build the expression between the property key and its value. + * \param pinned_type For kimi_k3_xml, the single type attribute this property must carry. + * std::nullopt keeps every type allowed, which is what free-form keys + * (additionalProperties / patternProperties) need. + */ + int32_t XMLKeySuffix(const std::optional& pinned_type = std::nullopt); + + /*! + * \brief Build a deepseek_v4_1_xml parameter's string attribute, value and closing tag. + * string="true" wraps raw strings, string="false" wraps JSON values. Unions and mixed enums + * produce one alternative per option; KimiK3TypeAttr supplies the type classification. + */ + int32_t FormatDeepSeekV41ParamSuffix(const SchemaSpecPtr& schema, int32_t value_rule_id); + + // Parameter suffix rules are independent of the key, so references can be shared across + // named and dynamic parameters. Allocate them before resolving refs to handle cycles. + std::unordered_map deepseek_v41_param_ref_rules_; + + JSONFormat json_format_; + // Track if we're at the root object level + int nested_object_level_ = 0; + const XMLWrapper xml_wrapper_; +}; + +/*! + * \brief Converter for Cohere XML Tool Calling format. + * + * This converter generates recursive Cohere value tags: + * value. + * Object properties use named value tags. Array items use unnamed value tags. + */ +class CohereXMLToolCallingConverter : public XMLToolCallingConverter { + public: + CohereXMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + bool any_order = false + ); + + protected: + int32_t GenerateString(const StringSpec& spec, const std::string& rule_name) override; + int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces = false + ) override; + int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name) override; + int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name) override; + int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name) override; + int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name) override; + + int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ) override; + int32_t FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema + ) override; + + std::string GetKeyPattern() const override; + int32_t GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name + ) override; + std::string NextSeparator(bool is_end = false) override; + + void AddBasicRules() override; + void AddCache(const std::string& key, int32_t rule_id) override; + std::optional GetCache(const std::string& key) const override; + + private: + struct CohereKeyTrieNode { + bool is_terminal = false; + std::map children; + }; + + static const std::string kCohereKey; + static const std::string kCohereAnyScalar; + static const std::string kCohereAnyList; + + /*! \brief `& name, const std::optional& key_pattern_expr + ); + /*! + * \brief The type attribute, value and closing tag of one parameter, correlated with + * \p schema. The suffix does not depend on the parameter name, so `$ref` schemas get one + * memoized rule per URI that is registered before the reference is resolved; recursive + * references (directly, or through nested dict/list items) then point back at that rule. + */ + int32_t FormatCohereParamSuffix(const SchemaSpecPtr& schema, int32_t value_rule_id); + int32_t FormatCohereSuffixWithType(int32_t type_expression, int32_t value_rule_id); + int32_t FormatAnyCohereSuffix(); + int32_t FormatCohereParam( + const std::optional& name, + const std::optional& key_pattern_expr, + const SchemaSpecPtr& schema, + int32_t value_rule_id + ); + int32_t FormatAnyCohereParam( + const std::optional& name, const std::optional& key_pattern_expr + ); + int32_t FormatCohereValue(int32_t value_rule_id); + int32_t GetCohereTypePattern(const SchemaSpecPtr& schema); + static std::string CohereTypeForJSONLiteral(const std::string& json_value); + static std::optional CommonCohereTypeForJSONLiterals( + const std::vector& json_values + ); + std::optional> GetCohereCompositeOptions(const SchemaSpecPtr& schema + ) const; + int32_t BuildCohereKeyExcludingBody(const CohereKeyTrieNode& node, int depth); + bool AtCohereRoot() const; + bool InCohereValueContext() const; + + std::vector object_stack_; + std::vector additional_property_stack_; + // Parameter suffix rules keyed by `$ref` URI; see FormatCohereParamSuffix. + std::unordered_map cohere_param_ref_rules_; + int cohere_array_level_ = 0; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_JSON_SCHEMA_CONVERTER_EXT_H_ diff --git a/third_party/xgrammar/cpp/lark_converter.cc b/third_party/xgrammar/cpp/lark_converter.cc new file mode 100644 index 0000000000..e4408d4cae --- /dev/null +++ b/third_party/xgrammar/cpp/lark_converter.cc @@ -0,0 +1,2714 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/lark_converter.cc + */ + +#include "lark_converter.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_functor.h" +#include "support/encoding.h" +#include "support/json_parse.h" +#include "support/logging.h" + +namespace xgrammar { +namespace { + +struct Location { + int line = 1; + int column = 1; +}; + +[[noreturn]] void RaiseLarkError( + const std::string& source, const Location& location, const std::string& message +) { + size_t line_start = 0; + int current_line = 1; + while (current_line < location.line && line_start < source.size()) { + size_t newline = source.find('\n', line_start); + if (newline == std::string::npos) { + line_start = source.size(); + break; + } + line_start = newline + 1; + ++current_line; + } + size_t line_end = source.find('\n', line_start); + if (line_end == std::string::npos) { + line_end = source.size(); + } + std::string line_text = source.substr(line_start, line_end - line_start); + std::ostringstream os; + os << "Lark error at line " << location.line << ", column " << location.column << ": " << message; + if (!line_text.empty()) { + os << "\n" << line_text << "\n" << std::string(std::max(0, location.column - 1), ' ') << "^"; + } + throw XGrammarError(os.str()); +} + +enum class TokenType { + kName, + kString, + kRegex, + kNumber, + kSpecialToken, + kGrammarRef, + kJson, + kRegexExt, + kGrammarOptions, + kImport, + kIgnore, + kLark, + kIf, + kUnsupportedDirective, + kColon, + kDoubleColon, + kComma, + kDot, + kDotDot, + kArrow, + kEquals, + kLParen, + kRParen, + kLBracket, + kRBracket, + kLBrace, + kRBrace, + kPipe, + kAnd, + kTilde, + kQuestion, + kStar, + kPlus, + kNewline, + kEnd, +}; + +struct Token { + TokenType type; + std::string text; + std::string flags; + Location location; +}; + +class LarkLexer { + public: + explicit LarkLexer(const std::string& source) : source_(source) {} + + std::vector Tokenize() { + std::vector result; + while (position_ < source_.size()) { + char c = source_[position_]; + if (c == ' ' || c == '\t' || c == '\f') { + Advance(); + continue; + } + if (c == '\r' || c == '\n') { + Location location = CurrentLocation(); + if (c == '\r') { + Advance(); + if (position_ < source_.size() && source_[position_] == '\n') { + Advance(); + } + } else { + Advance(); + } + result.push_back({TokenType::kNewline, "\n", "", location}); + continue; + } + if (c == '#') { + SkipComment(); + continue; + } + if (c == '/' && PeekChar(1) == '/') { + SkipComment(); + continue; + } + + Location location = CurrentLocation(); + switch (c) { + case ':': + if (PeekChar(1) == ':') { + result.push_back(SimpleToken(TokenType::kDoubleColon, 2)); + } else { + result.push_back(SimpleToken(TokenType::kColon, 1)); + } + break; + case ',': + result.push_back(SimpleToken(TokenType::kComma, 1)); + break; + case '.': + if (PeekChar(1) == '.') { + result.push_back(SimpleToken(TokenType::kDotDot, 2)); + } else { + result.push_back(SimpleToken(TokenType::kDot, 1)); + } + break; + case '-': + if (PeekChar(1) == '>') { + result.push_back(SimpleToken(TokenType::kArrow, 2)); + } else if (std::isdigit(static_cast(PeekChar(1)))) { + result.push_back(LexNumber()); + } else { + RaiseLarkError(source_, location, "unexpected '-' character"); + } + break; + case '+': + if (std::isdigit(static_cast(PeekChar(1)))) { + result.push_back(LexNumber()); + } else { + result.push_back(SimpleToken(TokenType::kPlus, 1)); + } + break; + case '=': + result.push_back(SimpleToken(TokenType::kEquals, 1)); + break; + case '(': + result.push_back(SimpleToken(TokenType::kLParen, 1)); + break; + case ')': + result.push_back(SimpleToken(TokenType::kRParen, 1)); + break; + case '[': + result.push_back(SimpleToken(TokenType::kLBracket, 1)); + break; + case ']': + result.push_back(SimpleToken(TokenType::kRBracket, 1)); + break; + case '{': + result.push_back(SimpleToken(TokenType::kLBrace, 1)); + break; + case '}': + result.push_back(SimpleToken(TokenType::kRBrace, 1)); + break; + case '|': + result.push_back(SimpleToken(TokenType::kPipe, 1)); + break; + case '&': + result.push_back(SimpleToken(TokenType::kAnd, 1)); + break; + case '~': + result.push_back(SimpleToken(TokenType::kTilde, 1)); + break; + case '?': + if (std::isalpha(static_cast(PeekChar(1))) || PeekChar(1) == '_') { + result.push_back(LexName()); + } else { + result.push_back(SimpleToken(TokenType::kQuestion, 1)); + } + break; + case '*': + result.push_back(SimpleToken(TokenType::kStar, 1)); + break; + case '!': + if (std::isalpha(static_cast(PeekChar(1))) || PeekChar(1) == '_') { + result.push_back(LexName()); + } else { + RaiseLarkError(source_, location, "unexpected '!' character"); + } + break; + case '"': + result.push_back(LexString()); + break; + case '/': + result.push_back(LexRegex()); + break; + case '<': + result.push_back(LexSpecialToken()); + break; + case '@': + result.push_back(LexGrammarRef()); + break; + case '%': + result.push_back(LexDirective()); + break; + default: + if (std::isalpha(static_cast(c)) || c == '_') { + result.push_back(LexName()); + } else if (std::isdigit(static_cast(c))) { + result.push_back(LexNumber()); + } else { + RaiseLarkError(source_, location, std::string("unexpected character '") + c + "'"); + } + } + } + result.push_back({TokenType::kEnd, "", "", CurrentLocation()}); + return result; + } + + private: + char PeekChar(size_t offset) const { + size_t index = position_ + offset; + return index < source_.size() ? source_[index] : '\0'; + } + + Location CurrentLocation() const { return {line_, column_}; } + + void Advance() { + if (position_ >= source_.size()) { + return; + } + char c = source_[position_++]; + if (c == '\n') { + ++line_; + column_ = 1; + } else { + ++column_; + } + } + + void AdvanceTo(size_t end_position) { + while (position_ < end_position) { + Advance(); + } + } + + Token SimpleToken(TokenType type, size_t length) { + Location location = CurrentLocation(); + std::string text = source_.substr(position_, length); + AdvanceTo(position_ + length); + return {type, std::move(text), "", location}; + } + + void SkipComment() { + while (position_ < source_.size() && source_[position_] != '\n' && source_[position_] != '\r') { + Advance(); + } + } + + Token LexName() { + Location location = CurrentLocation(); + size_t start = position_; + if (source_[position_] == '!' || source_[position_] == '?') { + Advance(); + } + while (position_ < source_.size()) { + char c = source_[position_]; + if (!std::isalnum(static_cast(c)) && c != '_' && c != '-') { + break; + } + Advance(); + } + return {TokenType::kName, source_.substr(start, position_ - start), "", location}; + } + + Token LexNumber() { + Location location = CurrentLocation(); + size_t start = position_; + if (source_[position_] == '+' || source_[position_] == '-') { + Advance(); + } + while (std::isdigit(static_cast(PeekChar(0)))) { + Advance(); + } + if (PeekChar(0) == '.' && PeekChar(1) != '.') { + Advance(); + while (std::isdigit(static_cast(PeekChar(0)))) { + Advance(); + } + } + if (PeekChar(0) == 'e' || PeekChar(0) == 'E') { + Advance(); + if (PeekChar(0) == '+' || PeekChar(0) == '-') { + Advance(); + } + while (std::isdigit(static_cast(PeekChar(0)))) { + Advance(); + } + } + return {TokenType::kNumber, source_.substr(start, position_ - start), "", location}; + } + + Token LexString() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + bool escaped = false; + while (position_ < source_.size()) { + char c = source_[position_]; + if (!escaped && c == '"') { + Advance(); + if (PeekChar(0) == 'i') { + Advance(); + } + return {TokenType::kString, source_.substr(start, position_ - start), "", location}; + } + if (c == '\n' || c == '\r') { + RaiseLarkError(source_, location, "unterminated string literal"); + } + if (!escaped && c == '\\') { + escaped = true; + } else { + escaped = false; + } + Advance(); + } + RaiseLarkError(source_, location, "unterminated string literal"); + } + + Token LexRegex() { + Location location = CurrentLocation(); + Advance(); + size_t pattern_start = position_; + bool escaped = false; + while (position_ < source_.size()) { + char c = source_[position_]; + if (!escaped && c == '/') { + std::string pattern = source_.substr(pattern_start, position_ - pattern_start); + Advance(); + size_t flags_start = position_; + while (std::isalpha(static_cast(PeekChar(0)))) { + Advance(); + } + return { + TokenType::kRegex, + std::move(pattern), + source_.substr(flags_start, position_ - flags_start), + location + }; + } + if (!escaped && c == '\\') { + escaped = true; + } else { + escaped = false; + } + Advance(); + } + RaiseLarkError(source_, location, "unterminated regular expression"); + } + + Token LexSpecialToken() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + while (position_ < source_.size() && source_[position_] != '>') { + char c = source_[position_]; + if (std::isspace(static_cast(c)) || c == '<') { + RaiseLarkError(source_, location, "invalid special token"); + } + Advance(); + } + if (position_ == source_.size()) { + RaiseLarkError(source_, location, "unterminated special token"); + } + Advance(); + return {TokenType::kSpecialToken, source_.substr(start, position_ - start), "", location}; + } + + Token LexGrammarRef() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + while (std::isalnum(static_cast(PeekChar(0))) || PeekChar(0) == '_' || + PeekChar(0) == '-') { + Advance(); + } + if (position_ == start + 1) { + RaiseLarkError(source_, location, "empty grammar reference"); + } + return {TokenType::kGrammarRef, source_.substr(start, position_ - start), "", location}; + } + + Token LexDirective() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + while (std::isalpha(static_cast(PeekChar(0))) || PeekChar(0) == '_') { + Advance(); + } + std::string directive = source_.substr(start, position_ - start); + if (directive == "%json") { + return LexJSONValue(TokenType::kJson, location, directive); + } + if (directive == "%regex") { + return LexJSONValue(TokenType::kRegexExt, location, directive); + } + if (directive == "%grammar_options") { + return LexJSONValue(TokenType::kGrammarOptions, location, directive); + } + if (directive == "%import") { + return {TokenType::kImport, directive, "", location}; + } + if (directive == "%ignore") { + return {TokenType::kIgnore, directive, "", location}; + } + if (directive == "%lark") { + return {TokenType::kLark, directive, "", location}; + } + if (directive == "%if") { + return {TokenType::kIf, directive, "", location}; + } + return {TokenType::kUnsupportedDirective, directive, "", location}; + } + + Token LexJSONValue(TokenType type, const Location& location, const std::string& directive) { + while (position_ < source_.size() && + std::isspace(static_cast(source_[position_]))) { + Advance(); + } + auto begin = source_.begin() + static_cast(position_); + auto end = source_.end(); + picojson::value value; + std::string error; + auto parsed_end = ParseJSON(value, begin, end, &error); + if (!error.empty() || parsed_end == begin) { + RaiseLarkError( + source_, location, "failed to parse JSON value after " + directive + ": " + error + ); + } + size_t new_position = static_cast(parsed_end - source_.begin()); + AdvanceTo(new_position); + return {type, value.serialize(), "", location}; + } + + const std::string& source_; + size_t position_ = 0; + int line_ = 1; + int column_ = 1; +}; + +struct Document; + +struct Node { + enum class Kind { + kSequence, + kChoice, + kRepeat, + kString, + kRegex, + kRange, + kName, + kJson, + kRegexExt, + kNestedLark, + kSpecialToken, + kGrammarRef, + kNot, + }; + + Kind kind = Kind::kSequence; + Location location; + std::string text; + std::string text2; + std::string flags; + int32_t min_repeat = 0; + int32_t max_repeat = 0; + std::vector children; + std::shared_ptr nested; +}; + +struct Definition { + std::string name; + bool is_terminal = false; + bool lazy = false; + std::optional suffix; + std::optional temperature; + Location suffix_location; + std::optional stop; + Location stop_location; + std::optional stop_capture_name; + Location stop_capture_location; + std::optional max_tokens; + Location max_tokens_location; + std::optional max_chars; + Location max_chars_location; + std::optional capture_name; + Location capture_location; + Node body; + Location location; +}; + +struct Import { + std::string path; + std::string local_name; + Location location; +}; + +struct Document { + std::vector definitions; + std::vector ignores; + std::vector imports; + std::vector> options; +}; + +class LarkParser { + public: + LarkParser(const std::string& source, std::vector tokens) + : source_(source), tokens_(std::move(tokens)) {} + + Document Parse() { return ParseDocument(false); } + + private: + const Token& Peek(size_t offset = 0) const { + size_t index = std::min(position_ + offset, tokens_.size() - 1); + return tokens_[index]; + } + + bool Match(TokenType type) { + if (Peek().type != type) { + return false; + } + ++position_; + return true; + } + + Token Consume(TokenType type, const std::string& message) { + if (Peek().type != type) { + RaiseLarkError(source_, Peek().location, message); + } + return tokens_[position_++]; + } + + void ConsumeNewlines() { + while (Match(TokenType::kNewline)) { + } + } + + Document ParseDocument(bool stop_at_rbrace) { + Document document; + ConsumeNewlines(); + while (Peek().type != TokenType::kEnd && !(stop_at_rbrace && Peek().type == TokenType::kRBrace) + ) { + switch (Peek().type) { + case TokenType::kImport: + ParseImport(&document); + break; + case TokenType::kIgnore: + ParseIgnore(&document); + break; + case TokenType::kGrammarOptions: + ParseOptions(&document); + break; + case TokenType::kUnsupportedDirective: + RaiseLarkError( + source_, Peek().location, "directive " + Peek().text + " is not supported" + ); + default: + document.definitions.push_back(ParseDefinition()); + break; + } + if (Peek().type != TokenType::kNewline && Peek().type != TokenType::kEnd && + !(stop_at_rbrace && Peek().type == TokenType::kRBrace)) { + RaiseLarkError(source_, Peek().location, "expected end of grammar item"); + } + ConsumeNewlines(); + } + return document; + } + + static bool IsTerminalName(const std::string& raw_name) { + size_t index = 0; + if (!raw_name.empty() && (raw_name[0] == '!' || raw_name[0] == '?')) { + index = 1; + } + if (index < raw_name.size() && raw_name[index] == '_') { + ++index; + } + return index < raw_name.size() && std::isupper(static_cast(raw_name[index])); + } + + static std::string NormalizeRuleName(std::string name) { + if (!name.empty() && (name[0] == '!' || name[0] == '?')) { + name.erase(name.begin()); + } + return name; + } + + void ValidateCaptureName(const std::string& capture_name, const Location& location) const { + if (capture_name.empty()) { + RaiseLarkError(source_, location, "capture name must not be empty"); + } + for (char c : capture_name) { + bool valid = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || + c == '_' || c == '-' || c == '.'; + if (!valid) { + RaiseLarkError( + source_, location, "capture name must only contain letters, digits, '_', '-' and '.'" + ); + } + } + } + + Node ParseStopLikeValue(const std::string& attribute_name) { + Token token = Peek(); + if (Match(TokenType::kString)) { + return ParseStringNode(token); + } + if (Match(TokenType::kRegex)) { + Node result; + result.kind = Node::Kind::kRegex; + result.location = token.location; + result.text = token.text; + result.flags = token.flags; + return result; + } + if (Match(TokenType::kName)) { + if (!IsTerminalName(token.text)) { + RaiseLarkError( + source_, token.location, attribute_name + " terminal name must be uppercase" + ); + } + Node result; + result.kind = Node::Kind::kName; + result.location = token.location; + result.text = NormalizeRuleName(token.text); + return result; + } + RaiseLarkError( + source_, + token.location, + "expected string literal, regular expression, or uppercase terminal name after " + + attribute_name + "=" + ); + } + + void ParseImport(Document* document) { + Location location = Consume(TokenType::kImport, "expected %import").location; + Token first = Consume(TokenType::kName, "expected import path"); + std::string path = first.text; + while (Match(TokenType::kDot)) { + path += "." + Consume(TokenType::kName, "expected name after '.'").text; + } + + if (Match(TokenType::kLParen)) { + do { + Token name = Consume(TokenType::kName, "expected imported terminal name"); + document->imports.push_back({path + "." + name.text, name.text, location}); + } while (Match(TokenType::kComma)); + Consume(TokenType::kRParen, "expected ')' after import list"); + return; + } + + std::string local_name = path.substr(path.find_last_of('.') + 1); + if (Match(TokenType::kArrow)) { + local_name = Consume(TokenType::kName, "expected import alias").text; + } + document->imports.push_back({path, local_name, location}); + } + + void ParseIgnore(Document* document) { + Consume(TokenType::kIgnore, "expected %ignore"); + document->ignores.push_back(ParseChoice()); + } + + void ParseOptions(Document* document) { + Token token = Consume(TokenType::kGrammarOptions, "expected %grammar_options"); + picojson::value value; + std::string error = ParseJSON(value, token.text); + if (!error.empty()) { + RaiseLarkError(source_, token.location, "invalid %grammar_options value: " + error); + } + document->options.push_back({std::move(value), token.location}); + } + + Definition ParseDefinition() { + Token name_token = Consume(TokenType::kName, "expected rule or terminal name"); + Definition result; + result.name = NormalizeRuleName(name_token.text); + result.is_terminal = IsTerminalName(name_token.text); + result.location = name_token.location; + + if (Peek().type == TokenType::kLBracket) { + if (result.is_terminal) { + RaiseLarkError(source_, Peek().location, "attributes are only supported on rules"); + } + ParseAttributes(&result); + } + if (Peek().type == TokenType::kDot) { + RaiseLarkError(source_, Peek().location, "rule and terminal priorities are not supported"); + } + if (Peek().type == TokenType::kDoubleColon) { + RaiseLarkError(source_, Peek().location, "parametric grammar is not supported"); + } + if (Peek().type == TokenType::kLBrace) { + RaiseLarkError(source_, Peek().location, "Lark templates are not supported"); + } + Consume(TokenType::kColon, "expected ':' after rule name"); + result.body = ParseChoice(); + return result; + } + + void ParseAttributes(Definition* definition) { + Consume(TokenType::kLBracket, "expected '['"); + while (Peek().type != TokenType::kRBracket) { + Token key = Consume(TokenType::kName, "expected rule attribute"); + if (key.text == "lazy" && Peek().type != TokenType::kEquals) { + definition->lazy = true; + } else if (key.text == "max_tokens") { + Consume(TokenType::kEquals, "expected '=' after max_tokens attribute"); + Location value_location = Peek().location; + int32_t value = ParseInteger(); + if (value <= 0) { + RaiseLarkError(source_, value_location, "max_tokens must be positive"); + } + if (value > 1'000'000) { + RaiseLarkError(source_, value_location, "max_tokens is too large"); + } + if (definition->max_tokens.has_value()) { + RaiseLarkError(source_, key.location, "max_tokens attribute is specified more than once"); + } + definition->max_tokens = value; + definition->max_tokens_location = key.location; + } else if (key.text == "max_chars") { + Consume(TokenType::kEquals, "expected '=' after max_chars attribute"); + int32_t value = ParseInteger(); + if (definition->max_chars.has_value()) { + RaiseLarkError(source_, key.location, "max_chars attribute is specified more than once"); + } + definition->max_chars = value; + definition->max_chars_location = key.location; + } else if (key.text == "capture") { + std::string capture_name; + Location capture_location = key.location; + if (Match(TokenType::kEquals)) { + Token name_token = Consume(TokenType::kString, "expected string literal after capture="); + Node name_node = ParseStringNode(name_token); + if (!name_node.flags.empty()) { + RaiseLarkError( + source_, name_node.location, "case-insensitive flags are not supported on capture" + ); + } + capture_name = std::move(name_node.text); + capture_location = name_node.location; + } else { + capture_name = definition->name; + } + ValidateCaptureName(capture_name, capture_location); + if (definition->capture_name.has_value()) { + RaiseLarkError(source_, key.location, "capture attribute is specified more than once"); + } + definition->capture_name = std::move(capture_name); + definition->capture_location = capture_location; + } else if (key.text == "suffix") { + Consume(TokenType::kEquals, "expected '=' after suffix attribute"); + Node suffix = ParseStopLikeValue("suffix"); + if (suffix.kind == Node::Kind::kString && suffix.text.empty()) { + RaiseLarkError(source_, suffix.location, "suffix must not be empty"); + } + if (definition->suffix.has_value()) { + RaiseLarkError(source_, key.location, "suffix attribute is specified more than once"); + } + if (definition->stop.has_value()) { + RaiseLarkError(source_, key.location, "suffix cannot be combined with stop"); + } + Location suffix_location = suffix.location; + definition->suffix = std::move(suffix); + definition->suffix_location = suffix_location; + } else if (key.text == "stop") { + Consume(TokenType::kEquals, "expected '=' after stop attribute"); + Node stop = ParseStopLikeValue("stop"); + if (stop.kind == Node::Kind::kString && stop.text.empty()) { + RaiseLarkError(source_, stop.location, "stop must not be empty"); + } + if (definition->stop.has_value()) { + RaiseLarkError(source_, key.location, "stop attribute is specified more than once"); + } + if (definition->suffix.has_value()) { + RaiseLarkError(source_, key.location, "stop cannot be combined with suffix"); + } + Location stop_location = stop.location; + definition->stop = std::move(stop); + definition->stop_location = stop_location; + } else if (key.text == "stop_capture") { + Consume(TokenType::kEquals, "expected '=' after stop_capture attribute"); + Token name_token = + Consume(TokenType::kString, "expected string literal after stop_capture="); + Node name_node = ParseStringNode(name_token); + if (!name_node.flags.empty()) { + RaiseLarkError( + source_, + name_node.location, + "case-insensitive flags are not supported on stop_capture" + ); + } + ValidateCaptureName(name_node.text, name_node.location); + if (definition->stop_capture_name.has_value()) { + RaiseLarkError( + source_, key.location, "stop_capture attribute is specified more than once" + ); + } + Location stop_capture_location = name_node.location; + definition->stop_capture_name = std::move(name_node.text); + definition->stop_capture_location = stop_capture_location; + } else if (key.text == "temperature") { + Consume(TokenType::kEquals, "expected '=' after temperature attribute"); + Token value = Consume(TokenType::kNumber, "expected number after temperature="); + if (definition->temperature.has_value()) { + RaiseLarkError( + source_, key.location, "temperature attribute is specified more than once" + ); + } + try { + size_t parsed_length = 0; + float temperature = std::stof(value.text, &parsed_length); + if (parsed_length != value.text.size() || !std::isfinite(temperature) || + temperature < 0) { + RaiseLarkError( + source_, value.location, "temperature must be a finite non-negative number" + ); + } + definition->temperature = temperature; + } catch (const std::exception&) { + RaiseLarkError( + source_, value.location, "temperature must be a finite non-negative number" + ); + } + } else { + RaiseLarkError( + source_, + key.location, + "rule attribute '" + key.text + "' is not supported by XGrammar Lark" + ); + } + if (!Match(TokenType::kComma)) { + break; + } + } + Consume(TokenType::kRBracket, "expected ']' after rule attributes"); + if (definition->stop_capture_name.has_value() && !definition->suffix.has_value() && + !definition->stop.has_value()) { + RaiseLarkError( + source_, definition->stop_capture_location, "stop_capture requires stop or suffix" + ); + } + } + + Node ParseChoice() { + Location location = Peek().location; + std::vector alternatives; + alternatives.push_back(ParseSequence()); + while (MatchAlternativeSeparator()) { + alternatives.push_back(ParseSequence()); + } + if (alternatives.size() == 1) { + return std::move(alternatives[0]); + } + Node result; + result.kind = Node::Kind::kChoice; + result.location = location; + result.children = std::move(alternatives); + return result; + } + + bool MatchAlternativeSeparator() { + if (Match(TokenType::kPipe)) { + return true; + } + size_t saved_position = position_; + ConsumeNewlines(); + if (Match(TokenType::kPipe)) { + return true; + } + if (Peek().type == TokenType::kRParen || Peek().type == TokenType::kRBracket || + Peek().type == TokenType::kRBrace) { + return false; + } + position_ = saved_position; + return false; + } + + Node ParseSequence() { + Location location = Peek().location; + std::vector elements; + while (!IsSequenceEnd(Peek().type)) { + elements.push_back(ParseExpr()); + } + if (Match(TokenType::kArrow)) { + Consume(TokenType::kName, "expected alias name after '->'"); + } + if (Peek().type == TokenType::kAnd) { + RaiseLarkError(source_, Peek().location, "terminal intersection '&' is not supported"); + } + if (Peek().type == TokenType::kIf) { + RaiseLarkError(source_, Peek().location, "parametric %if conditions are not supported"); + } + Node result; + result.kind = Node::Kind::kSequence; + result.location = location; + result.children = std::move(elements); + return result; + } + + static bool IsSequenceEnd(TokenType type) { + return type == TokenType::kNewline || type == TokenType::kPipe || type == TokenType::kRParen || + type == TokenType::kRBracket || type == TokenType::kRBrace || type == TokenType::kEnd || + type == TokenType::kArrow || type == TokenType::kAnd || type == TokenType::kIf; + } + + int32_t ParseInteger() { + Token token = Consume(TokenType::kNumber, "expected integer"); + try { + size_t parsed = 0; + long long value = std::stoll(token.text, &parsed); + if (parsed != token.text.size() || value < 0 || value > std::numeric_limits::max()) { + RaiseLarkError(source_, token.location, "invalid non-negative repetition count"); + } + return static_cast(value); + } catch (const std::exception&) { + RaiseLarkError(source_, token.location, "invalid repetition count"); + } + } + + Node ParseExpr() { + Location location = Peek().location; + bool negated = Match(TokenType::kTilde); + Node atom = ParseAtom(); + if (negated) { + Node not_node; + not_node.kind = Node::Kind::kNot; + not_node.location = location; + not_node.children.push_back(std::move(atom)); + atom = std::move(not_node); + } + + int32_t min_repeat = -1; + int32_t max_repeat = -1; + if (Match(TokenType::kQuestion)) { + min_repeat = 0; + max_repeat = 1; + } else if (Match(TokenType::kStar)) { + min_repeat = 0; + max_repeat = -1; + } else if (Match(TokenType::kPlus)) { + min_repeat = 1; + max_repeat = -1; + } else if (Match(TokenType::kTilde)) { + min_repeat = ParseInteger(); + max_repeat = min_repeat; + if (Match(TokenType::kDotDot)) { + max_repeat = ParseInteger(); + } + } else if (Match(TokenType::kLBrace)) { + min_repeat = Peek().type == TokenType::kComma ? 0 : ParseInteger(); + if (Match(TokenType::kComma)) { + max_repeat = Peek().type == TokenType::kRBrace ? -1 : ParseInteger(); + } else { + max_repeat = min_repeat; + } + Consume(TokenType::kRBrace, "expected '}' after repetition range"); + } + + if (min_repeat == -1) { + return atom; + } + if (max_repeat != -1 && max_repeat < min_repeat) { + RaiseLarkError(source_, location, "repetition end must be greater than or equal to start"); + } + Node repeat; + repeat.kind = Node::Kind::kRepeat; + repeat.location = location; + repeat.min_repeat = min_repeat; + repeat.max_repeat = max_repeat; + repeat.children.push_back(std::move(atom)); + return repeat; + } + + // Groups recurse ParseChoice -> ParseSequence -> ParseExpr -> ParseAtom once per level, so bound + // the depth instead of overflowing the stack. Each level uses a few KB of stack; the limit keeps + // a wide margin below the 1 MB main-thread stack of Windows. + Node ParseGroupBody(const Token& open) { + if (group_depth_ >= kMaxGroupDepth) { + RaiseLarkError( + source_, + open.location, + "groups are nested deeper than " + std::to_string(kMaxGroupDepth) + " levels" + ); + } + ++group_depth_; + Node result = ParseChoice(); + --group_depth_; + return result; + } + + Node ParseAtom() { + Token token = Peek(); + if (Match(TokenType::kLParen)) { + Node result = ParseGroupBody(token); + Consume(TokenType::kRParen, "expected ')' after group"); + return result; + } + if (Match(TokenType::kLBracket)) { + Node inner = ParseGroupBody(token); + Consume(TokenType::kRBracket, "expected ']' after optional group"); + Node result; + result.kind = Node::Kind::kRepeat; + result.location = token.location; + result.min_repeat = 0; + result.max_repeat = 1; + result.children.push_back(std::move(inner)); + return result; + } + if (Match(TokenType::kString)) { + Node result = ParseStringNode(token); + if (Match(TokenType::kDotDot)) { + Token end = Consume(TokenType::kString, "expected string after '..'"); + Node end_node = ParseStringNode(end); + if (!result.flags.empty()) { + RaiseLarkError(source_, token.location, "flags are not allowed on character ranges"); + } + if (!end_node.flags.empty()) { + RaiseLarkError(source_, end.location, "flags are not allowed on character ranges"); + } + Node range; + range.kind = Node::Kind::kRange; + range.location = token.location; + range.text = result.text; + range.text2 = end_node.text; + return range; + } + return result; + } + if (Match(TokenType::kRegex)) { + Node result; + result.kind = Node::Kind::kRegex; + result.location = token.location; + result.text = token.text; + result.flags = token.flags; + return result; + } + if (Match(TokenType::kName)) { + Node result; + result.kind = Node::Kind::kName; + result.location = token.location; + result.text = NormalizeRuleName(token.text); + if (Peek().type == TokenType::kDoubleColon) { + RaiseLarkError(source_, Peek().location, "parametric grammar is not supported"); + } + if (Peek().type == TokenType::kLBrace && Peek(1).type != TokenType::kComma && + Peek(1).type != TokenType::kNumber) { + RaiseLarkError(source_, Peek().location, "Lark templates are not supported"); + } + return result; + } + if (Match(TokenType::kJson)) { + Node result; + result.kind = Node::Kind::kJson; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kRegexExt)) { + Node result; + result.kind = Node::Kind::kRegexExt; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kSpecialToken)) { + Node result; + result.kind = Node::Kind::kSpecialToken; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kGrammarRef)) { + Node result; + result.kind = Node::Kind::kGrammarRef; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kLark)) { + Consume(TokenType::kLBrace, "expected '{' after %lark"); + Node result; + result.kind = Node::Kind::kNestedLark; + result.location = token.location; + result.nested = std::make_shared(ParseDocument(true)); + Consume(TokenType::kRBrace, "expected '}' after nested Lark grammar"); + return result; + } + if (token.type == TokenType::kUnsupportedDirective) { + RaiseLarkError(source_, token.location, "directive " + token.text + " is not supported"); + } + RaiseLarkError(source_, token.location, "expected grammar expression"); + } + + Node ParseStringNode(const Token& token) { + std::string json_string = token.text; + std::string flags; + if (!json_string.empty() && json_string.back() == 'i') { + flags = "i"; + json_string.pop_back(); + } + picojson::value value; + std::string error = ParseJSON(value, json_string); + if (!error.empty() || !value.is()) { + RaiseLarkError(source_, token.location, "invalid string literal: " + error); + } + Node result; + result.kind = Node::Kind::kString; + result.location = token.location; + result.text = value.get(); + result.flags = std::move(flags); + return result; + } + + const std::string& source_; + std::vector tokens_; + size_t position_ = 0; + int group_depth_ = 0; + static constexpr int kMaxGroupDepth = 200; +}; + +const std::unordered_map& CommonRegexes() { + static const std::unordered_map regexes = { + {"common.DIGIT", "[0-9]"}, + {"common.HEXDIGIT", "[a-fA-F0-9]"}, + {"common.INT", "[0-9]+"}, + {"common.SIGNED_INT", "(\\+|-)?[0-9]+"}, + {"common.DECIMAL", "([0-9]+\\.[0-9]*)|(\\.[0-9]+)"}, + {"common._EXP", "[eE](\\+|-)?[0-9]+"}, + {"common.FLOAT", "([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+"}, + {"common.SIGNED_FLOAT", + "(\\+|-)?(([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+)"}, + {"common.NUMBER", + "([0-9]+)|([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+"}, + {"common.SIGNED_NUMBER", + "(\\+|-)?(([0-9]+)|([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+" + ")"}, + {"common.ESCAPED_STRING", "\\\"([^\\\"\\\\]|\\\\.)*\\\""}, + {"common.LCASE_LETTER", "[a-z]"}, + {"common.UCASE_LETTER", "[A-Z]"}, + {"common.LETTER", "[A-Za-z]"}, + {"common.WORD", "[A-Za-z]+"}, + {"common.CNAME", "[_A-Za-z][_A-Za-z0-9]*"}, + {"common.WS_INLINE", "[ \\t]+"}, + {"common.WS", "[ \\t\\f\\r\\n]+"}, + {"common.CR", "\\r"}, + {"common.LF", "\\n"}, + {"common.NEWLINE", "(\\r?\\n)+"}, + {"common.SH_COMMENT", "#[^\\n]*"}, + {"common.CPP_COMMENT", "//[^\\n]*"}, + {"common.C_COMMENT", "\\/\\*[^*]*\\*+(?:[^/*][^*]*\\*+)*\\/"}, + {"common.SQL_COMMENT", "--[^\\n]*"}, + }; + return regexes; +} + +std::string Trim(std::string value) { + size_t begin = 0; + while (begin < value.size() && std::isspace(static_cast(value[begin]))) { + ++begin; + } + size_t end = value.size(); + while (end > begin && std::isspace(static_cast(value[end - 1]))) { + --end; + } + return value.substr(begin, end - begin); +} + +std::optional ParseFixedRegexLiteral(const std::string& pattern) { + std::string result; + for (size_t i = 0; i < pattern.size();) { + char c = pattern[i++]; + if (c != '\\') { + if (std::string(".^$*+?()[]{}|").find(c) != std::string::npos) { + return std::nullopt; + } + result.push_back(c); + continue; + } + if (i == pattern.size()) { + return std::nullopt; + } + char escaped = pattern[i++]; + switch (escaped) { + case 'n': + result.push_back('\n'); + break; + case 'r': + result.push_back('\r'); + break; + case 't': + result.push_back('\t'); + break; + case 'f': + result.push_back('\f'); + break; + case 'v': + result.push_back('\v'); + break; + case '0': + result.push_back('\0'); + break; + case '^': + case '$': + case '.': + case '*': + case '+': + case '?': + case '(': + case ')': + case '[': + case ']': + case '{': + case '}': + case '|': + case '\\': + case '/': + case '-': + result.push_back(escaped); + break; + case 'x': + case 'u': { + bool braced = escaped == 'u' && i < pattern.size() && pattern[i] == '{'; + TCodepoint codepoint = 0; + if (braced) { + ++i; + int digit_count = 0; + while (i < pattern.size() && HexCharToInt(pattern[i]) != -1 && digit_count < 6) { + codepoint = codepoint * 16 + HexCharToInt(pattern[i++]); + ++digit_count; + } + if (digit_count == 0 || i >= pattern.size() || pattern[i++] != '}') { + return std::nullopt; + } + } else { + int digit_count = escaped == 'x' ? 2 : 4; + if (i + static_cast(digit_count) > pattern.size()) { + return std::nullopt; + } + for (int digit = 0; digit < digit_count; ++digit) { + int value = HexCharToInt(pattern[i++]); + if (value == -1) { + return std::nullopt; + } + codepoint = codepoint * 16 + value; + } + } + if (codepoint > 0x10FFFF || (codepoint >= 0xD800 && codepoint <= 0xDFFF)) { + return std::nullopt; + } + result += CharToUTF8(codepoint); + break; + } + default: + return std::nullopt; + } + } + return result; +} + +struct NamedGrammarRegistry { + std::unordered_map> inputs; + std::unordered_map compiled; + std::vector active; +}; + +class LarkCompiler { + public: + LarkCompiler( + const std::string& source, + Document document, + const std::optional& tokenizer_info, + NamedGrammarRegistry& named_grammars + ) + : source_(source), + document_(std::move(document)), + tokenizer_info_(tokenizer_info), + named_grammars_(named_grammars) {} + + Grammar Compile() { + ExpandImports(); + ParseOptions(); + IndexDefinitions(); + ValidateTerminalCycles(); + + for (const auto& definition : document_.definitions) { + rule_ids_[definition.name] = builder_.AddEmptyRule(definition.name); + } + + for (const auto& definition : document_.definitions) { + if (definition.is_terminal) { + builder_.UpdateRuleBody( + rule_ids_.at(definition.name), CompileNode(definition.body, definition.name, true) + ); + } + } + + CompileIgnore(); + + const Definition& start_definition = *definition_by_name_.at("start"); + std::optional dynamic_start_body = start_definition.temperature.has_value() + ? std::nullopt + : CompileDynamicStart(start_definition); + + for (const auto& definition : document_.definitions) { + if (definition.is_terminal) { + continue; + } + if (dynamic_unused_rules_.count(definition.name)) { + if (definition.max_tokens.has_value()) { + RaiseLarkError( + source_, + definition.max_tokens_location, + "max_tokens is not supported on rules consumed by dynamic dispatch" + ); + } + if (definition.max_chars.has_value()) { + XGRAMMAR_LOG(WARNING) << "Ignoring max_chars on rule '" << definition.name + << "' because it is consumed by dynamic dispatch."; + } + if (definition.capture_name.has_value()) { + RaiseLarkError( + source_, + definition.capture_location, + "capture is not supported on rules consumed by dynamic dispatch" + ); + } + builder_.UpdateRuleBody(rule_ids_.at(definition.name), builder_.AddEmptyStr()); + continue; + } + int32_t body_expr_id; + if (definition.temperature.has_value()) { + body_expr_id = CompileTemperatureRule(definition); + if (definition.max_chars.has_value()) { + builder_.UpdateMaxChars(rule_ids_.at(definition.name), definition.max_chars.value()); + } + } else if (definition.name == "start") { + if (dynamic_start_body.has_value()) { + if (definition.max_tokens.has_value()) { + RaiseLarkError( + source_, + definition.max_tokens_location, + "max_tokens is not supported on a dynamic dispatch start rule" + ); + } + if (definition.max_chars.has_value()) { + XGRAMMAR_LOG(WARNING) << "Ignoring max_chars on dynamic dispatch start rule '" + << definition.name << "'."; + } + body_expr_id = dynamic_start_body.value(); + } else if (definition.max_tokens.has_value() || definition.max_chars.has_value()) { + body_expr_id = CompileBudgetRule(definition); + } else if (HasLazySemantics(definition)) { + body_expr_id = CompileLazyRule(definition); + } else { + body_expr_id = CompileNode(definition.body, definition.name, false); + } + if (allow_initial_skip_ && skip_rule_id_ != -1) { + body_expr_id = builder_.AddSequence({builder_.AddRuleRef(skip_rule_id_), body_expr_id}); + } + } else if (definition.max_tokens.has_value() || definition.max_chars.has_value()) { + body_expr_id = CompileBudgetRule(definition); + } else if (HasLazySemantics(definition)) { + body_expr_id = CompileLazyRule(definition); + } else { + body_expr_id = CompileNode(definition.body, definition.name, false); + } + int32_t rule_id = rule_ids_.at(definition.name); + builder_.UpdateRuleBody(rule_id, body_expr_id); + if (definition.capture_name.has_value()) { + builder_.UpdateCaptureName(rule_id, definition.capture_name.value()); + } + builder_.UpdateRuleTemperature(rule_id, definition.temperature); + } + + auto start_it = rule_ids_.find("start"); + if (start_it == rule_ids_.end()) { + RaiseLarkError(source_, {1, 1}, "no start rule found"); + } + int32_t root_rule_id = start_it->second; + if (start_definition.temperature.has_value() && skip_rule_id_ != -1) { + std::vector elements; + if (allow_initial_skip_) { + elements.push_back(builder_.AddRuleRef(skip_rule_id_)); + } + elements.push_back(builder_.AddRuleRef(root_rule_id)); + elements.push_back(builder_.AddRuleRef(skip_rule_id_)); + root_rule_id = + builder_.AddRuleWithHint("start_with_skip", builder_.AddSequence(std::move(elements))); + } + return DeadCodeEliminator::Apply(GrammarNormalizer().Apply(builder_.Get(root_rule_id))); + } + + private: + struct SpecialTokenSet { + bool excluded = false; + std::vector token_ids; + }; + + struct Trigger { + enum class Level { kString, kToken } level; + std::string string; + std::vector token_ids; + Location location; + }; + + struct DynamicAlternative { + Trigger trigger; + Node remainder; + int32_t marker_event_rule_id = -1; + }; + + static bool HasLazySemantics(const Definition& definition) { + return definition.lazy || definition.suffix.has_value() || definition.stop.has_value(); + } + + int32_t CompileTemperatureRule(const Definition& definition) { + const Node* body = UnwrapSingle(&definition.body); + if (body->kind == Node::Kind::kJson || body->kind == Node::Kind::kNestedLark || + body->kind == Node::Kind::kGrammarRef) { + return CompileNode(*body, definition.name, false, false); + } + try { + return CompileNode(definition.body, definition.name, true, false); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + definition.location, + std::string(error.what()) + "; temperature is only supported on terminals and subgrammars" + ); + } + } + + void ExpandImports() { + for (const auto& import : document_.imports) { + auto it = CommonRegexes().find(import.path); + if (it == CommonRegexes().end()) { + RaiseLarkError(source_, import.location, "unknown common import '" + import.path + "'"); + } + Node regex; + regex.kind = Node::Kind::kRegex; + regex.location = import.location; + regex.text = it->second; + Definition definition; + definition.name = import.local_name; + definition.is_terminal = true; + definition.body = std::move(regex); + definition.location = import.location; + document_.definitions.push_back(std::move(definition)); + } + } + + void ParseOptions() { + for (const auto& [value, location] : document_.options) { + if (!value.is()) { + RaiseLarkError(source_, location, "%grammar_options value must be an object"); + } + for (const auto& [key, option] : value.get()) { + if (key == "allow_initial_skip") { + if (!option.is()) { + RaiseLarkError(source_, location, "allow_initial_skip must be a boolean"); + } + allow_initial_skip_ = allow_initial_skip_ || option.get(); + } else if (key == "no_forcing" || key == "allow_invalid_utf8") { + if (!option.is()) { + RaiseLarkError(source_, location, key + " must be a boolean"); + } + if (option.get()) { + RaiseLarkError( + source_, location, "%grammar_options option '" + key + "' is not supported" + ); + } + } else { + RaiseLarkError(source_, location, "unknown %grammar_options option '" + key + "'"); + } + } + } + } + + void IndexDefinitions() { + for (auto& definition : document_.definitions) { + if (definition_by_name_.count(definition.name)) { + RaiseLarkError( + source_, definition.location, "duplicate rule or terminal '" + definition.name + "'" + ); + } + definition_by_name_[definition.name] = &definition; + } + if (!definition_by_name_.count("start")) { + RaiseLarkError(source_, {1, 1}, "no start rule found"); + } + if (definition_by_name_.at("start")->is_terminal) { + RaiseLarkError(source_, definition_by_name_.at("start")->location, "start must be a rule"); + } + } + + void CollectReferencedNames(const Node& node, std::vector* names) const { + if (node.kind == Node::Kind::kName) { + names->push_back(node.text); + } + for (const Node& child : node.children) { + CollectReferencedNames(child, names); + } + } + + void ValidateTerminalCycles() { + std::unordered_map states; + for (const auto& definition : document_.definitions) { + if (definition.is_terminal && states[definition.name] == 0) { + VisitTerminal(definition, &states); + } + } + } + + void VisitTerminal(const Definition& definition, std::unordered_map* states) { + (*states)[definition.name] = 1; + std::vector names; + CollectReferencedNames(definition.body, &names); + for (const std::string& name : names) { + auto it = definition_by_name_.find(name); + if (it == definition_by_name_.end()) { + RaiseLarkError(source_, definition.location, "unknown name '" + name + "'"); + } + if (!it->second->is_terminal) { + RaiseLarkError( + source_, + definition.location, + "terminal '" + definition.name + "' cannot reference rule '" + name + "'" + ); + } + if ((*states)[name] == 1) { + RaiseLarkError( + source_, definition.location, "circular reference in terminal '" + name + "'" + ); + } + if ((*states)[name] == 0) { + VisitTerminal(*it->second, states); + } + } + (*states)[definition.name] = 2; + } + + void CompileIgnore() { + if (document_.ignores.empty()) { + return; + } + std::vector ignore_choices; + for (const Node& ignore : document_.ignores) { + ignore_choices.push_back(CompileNode(ignore, "lark_ignore", true)); + } + int32_t ignore_body = + ignore_choices.size() == 1 ? ignore_choices[0] : builder_.AddChoices(ignore_choices); + int32_t ignore_item_rule = builder_.AddRuleWithHint("lark_ignore_item", ignore_body); + int32_t ignore_repeat = builder_.AddRepeat(ignore_item_rule, 0, -1); + skip_rule_id_ = builder_.AddRuleWithHint("lark_ignore", ignore_repeat); + } + + int32_t CompileStringLiteral(const Node& node) { + if (node.flags.empty()) { + return node.text.empty() ? builder_.AddEmptyStr() : builder_.AddByteString(node.text); + } + if (node.flags != "i") { + RaiseLarkError( + source_, node.location, "unsupported string literal flags '" + node.flags + "'" + ); + } + std::vector codepoints = ParseUTF8(node.text.c_str()); + if (!node.text.empty() && + (codepoints.empty() || codepoints[0] == CharHandlingError::kInvalidUTF8)) { + RaiseLarkError(source_, node.location, "case-insensitive string is not valid UTF-8"); + } + std::vector elements; + elements.reserve(codepoints.size()); + for (TCodepoint codepoint : codepoints) { + if (codepoint > 0x7F) { + RaiseLarkError( + source_, + node.location, + "case-insensitive string literals currently support ASCII characters only" + ); + } + if ((codepoint >= 'a' && codepoint <= 'z') || (codepoint >= 'A' && codepoint <= 'Z')) { + TCodepoint lowercase = + static_cast(std::tolower(static_cast(codepoint))); + TCodepoint uppercase = + static_cast(std::toupper(static_cast(codepoint))); + elements.push_back( + builder_.AddCharacterClass({{lowercase, lowercase}, {uppercase, uppercase}}) + ); + } else { + elements.push_back(builder_.AddByteString(CharToUTF8(codepoint))); + } + } + if (elements.empty()) { + return builder_.AddEmptyStr(); + } + return elements.size() == 1 ? elements[0] : builder_.AddSequence(elements); + } + + struct RegexFlags { + bool case_insensitive = false; + bool dot_all = false; + }; + + RegexFlags ParseRegexFlags(const Node& node) const { + RegexFlags result; + for (char flag : node.flags) { + if (flag == 'i') { + result.case_insensitive = true; + } else if (flag == 's') { + result.dot_all = true; + } else if (flag == 'u') { + // XGrammar regular expressions use Unicode codepoint semantics by default. + } else if (flag == 'l') { + RaiseLarkError(source_, node.location, "regular-expression flag 'l' is not supported"); + } else { + RaiseLarkError( + source_, + node.location, + "regular-expression flag '" + std::string(1, flag) + "' is not supported" + ); + } + } + return result; + } + + std::string PrepareRegexPattern(const Node& node) const { + return RewriteRegexDots(node.text, ParseRegexFlags(node).dot_all); + } + + static std::string EscapeRegexLiteral(const std::string& value) { + static const std::string kRegexMeta = R"(\.^$|()[]{}*+?)"; + static constexpr char kHex[] = "0123456789ABCDEF"; + std::string result; + for (unsigned char byte : value) { + if (byte == '\n') { + result += "\\n"; + } else if (byte == '\r') { + result += "\\r"; + } else if (byte == '\t') { + result += "\\t"; + } else if (byte < 0x20 || byte == 0x7F) { + result += "\\x"; + result += kHex[byte >> 4]; + result += kHex[byte & 0x0F]; + } else { + char character = static_cast(byte); + if (kRegexMeta.find(character) != std::string::npos) { + result += '\\'; + } + result += character; + } + } + return result; + } + + std::string StringLiteralToRegex(const Node& node) { + if (node.flags.empty()) { + return EscapeRegexLiteral(node.text); + } + if (node.flags != "i") { + RaiseLarkError( + source_, node.location, "unsupported string literal flags '" + node.flags + "'" + ); + } + std::vector codepoints = ParseUTF8(node.text.c_str()); + if (!node.text.empty() && + (codepoints.empty() || codepoints[0] == CharHandlingError::kInvalidUTF8)) { + RaiseLarkError(source_, node.location, "case-insensitive string is not valid UTF-8"); + } + std::string result; + for (TCodepoint codepoint : codepoints) { + if (codepoint > 0x7F) { + RaiseLarkError( + source_, + node.location, + "case-insensitive string literals currently support ASCII characters only" + ); + } + char character = static_cast(codepoint); + if ((character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z')) { + char lower = static_cast(std::tolower(static_cast(character))); + char upper = static_cast(std::toupper(static_cast(character))); + result += "["; + result += lower; + result += upper; + result += "]"; + } else { + result += EscapeRegexLiteral(std::string(1, character)); + } + } + return result; + } + + std::string TerminalNodeToRegex( + const Node& node, std::unordered_set* visiting = nullptr + ) { + auto wrap = [](const std::string& pattern) { + return pattern.empty() ? std::string() : "(?:" + pattern + ")"; + }; + switch (node.kind) { + case Node::Kind::kSequence: { + std::string result; + for (const Node& child : node.children) { + result += wrap(TerminalNodeToRegex(child, visiting)); + } + return result; + } + case Node::Kind::kChoice: { + std::string result = "(?:"; + for (size_t i = 0; i < node.children.size(); ++i) { + if (i != 0) { + result += "|"; + } + result += TerminalNodeToRegex(node.children[i], visiting); + } + return result + ")"; + } + case Node::Kind::kRepeat: { + std::string child = TerminalNodeToRegex(node.children[0], visiting); + if (child.empty()) { + return ""; + } + std::string result = "(?:" + child + ")"; + if (node.min_repeat == 0 && node.max_repeat == -1) { + return result + "*"; + } + if (node.min_repeat == 1 && node.max_repeat == -1) { + return result + "+"; + } + if (node.min_repeat == 0 && node.max_repeat == 1) { + return result + "?"; + } + result += "{" + std::to_string(node.min_repeat); + if (node.max_repeat != node.min_repeat) { + result += ","; + if (node.max_repeat != -1) { + result += std::to_string(node.max_repeat); + } + } + return result + "}"; + } + case Node::Kind::kString: + return StringLiteralToRegex(node); + case Node::Kind::kRegex: + if (ParseRegexFlags(node).case_insensitive) { + RaiseLarkError( + source_, + node.location, + "regular-expression flag 'i' is not supported with suffix or stop attributes" + ); + } + return "(?:" + PrepareRegexPattern(node) + ")"; + case Node::Kind::kRange: { + std::vector begin = ParseUTF8(node.text.c_str()); + std::vector end = ParseUTF8(node.text2.c_str()); + if (begin.size() != 1 || end.size() != 1 || begin[0] == CharHandlingError::kInvalidUTF8 || + end[0] == CharHandlingError::kInvalidUTF8) { + RaiseLarkError(source_, node.location, "character range endpoints must be one character"); + } + if (begin[0] > end[0]) { + RaiseLarkError(source_, node.location, "character range start must not exceed end"); + } + auto escape_class_character = [](TCodepoint codepoint) { + std::string value = CharToUTF8(codepoint); + if (value == "\\" || value == "]" || value == "-" || value == "^") { + return "\\" + value; + } + return value; + }; + return "[" + escape_class_character(begin[0]) + "-" + escape_class_character(end[0]) + "]"; + } + case Node::Kind::kName: { + auto definition_it = definition_by_name_.find(node.text); + if (definition_it == definition_by_name_.end()) { + RaiseLarkError(source_, node.location, "unknown name '" + node.text + "'"); + } + if (!definition_it->second->is_terminal) { + RaiseLarkError( + source_, node.location, "terminal cannot reference rule '" + node.text + "'" + ); + } + std::unordered_set local_visiting; + if (visiting == nullptr) { + visiting = &local_visiting; + } + if (!visiting->insert(node.text).second) { + RaiseLarkError( + source_, node.location, "recursive terminal '" + node.text + "' is not supported" + ); + } + std::string result = TerminalNodeToRegex(definition_it->second->body, visiting); + visiting->erase(node.text); + return result; + } + case Node::Kind::kSpecialToken: + RaiseLarkError(source_, node.location, "special tokens cannot be used in terminals"); + case Node::Kind::kJson: + RaiseLarkError(source_, node.location, "%json cannot be used in terminals"); + case Node::Kind::kNestedLark: + RaiseLarkError(source_, node.location, "nested %lark cannot be used in terminals"); + case Node::Kind::kRegexExt: + RaiseLarkError( + source_, node.location, "structured %regex cannot be used with suffix or stop" + ); + case Node::Kind::kGrammarRef: + RaiseLarkError(source_, node.location, "named grammars cannot be used in terminals"); + case Node::Kind::kNot: + RaiseLarkError( + source_, node.location, "regular-expression complement '~' is not supported" + ); + } + RaiseLarkError(source_, node.location, "unsupported terminal node"); + } + + std::vector ParseStructuredRegexChunks(const Node& node) const { + picojson::value value; + std::string error = ParseJSON(value, node.text); + if (!error.empty()) { + RaiseLarkError(source_, node.location, "failed to parse %regex: " + error); + } + if (!value.is()) { + RaiseLarkError(source_, node.location, "%regex value must be an object"); + } + + const auto& object = value.get(); + std::vector fields; + for (const auto& [key, field_value] : object) { + if (key != "substring_chunks" && key != "substring_chars" && key != "substring_words") { + RaiseLarkError(source_, node.location, "unknown field '" + key + "' in %regex"); + } + if (!field_value.is()) { + fields.push_back(key); + } + } + if (fields.empty()) { + RaiseLarkError(source_, node.location, "no fields set on %regex"); + } + if (fields.size() != 1) { + RaiseLarkError(source_, node.location, "only one field can be set on %regex"); + } + + const std::string& field = fields[0]; + const picojson::value& field_value = object.at(field); + if (field == "substring_words") { + if (!field_value.is()) { + RaiseLarkError(source_, node.location, "substring_words must be a string"); + } + RaiseLarkError(source_, node.location, "substring_words is not supported yet"); + } + if (field == "substring_chars") { + if (!field_value.is()) { + RaiseLarkError(source_, node.location, "substring_chars must be a string"); + } + const std::string& text = field_value.get(); + std::vector chunks; + for (size_t offset = 0; offset < text.size();) { + if (text[offset] == '\0') { + chunks.emplace_back(1, '\0'); + ++offset; + continue; + } + auto [codepoint, length] = ParseNextUTF8(text.c_str() + offset); + if (codepoint == CharHandlingError::kInvalidUTF8) { + RaiseLarkError(source_, node.location, "substring_chars must be valid UTF-8"); + } + chunks.push_back(text.substr(offset, length)); + offset += length; + } + return chunks; + } + + if (!field_value.is()) { + RaiseLarkError(source_, node.location, "substring_chunks must be an array of strings"); + } + std::vector chunks; + const auto& array = field_value.get(); + chunks.reserve(array.size()); + for (const picojson::value& chunk : array) { + if (!chunk.is()) { + RaiseLarkError(source_, node.location, "substring_chunks must be an array of strings"); + } + chunks.push_back(chunk.get()); + } + return chunks; + } + + int32_t CompileStructuredRegex(const Node& node, const std::string& rule_hint) { + std::vector chunks = ParseStructuredRegexChunks(node); + // A substring expr can only be the body of a rule, so wrap it and return a reference. + int32_t substring_expr_id = builder_.AddSubstring(chunks); + int32_t rule_id = builder_.AddRuleWithHint(rule_hint + "_substring", substring_expr_id); + return builder_.AddRuleRef(rule_id); + } + + const Grammar& ResolveNamedGrammar(const std::string& name, const Location& location) { + auto input_it = named_grammars_.inputs.find(name); + if (input_it == named_grammars_.inputs.end()) { + RaiseLarkError(source_, location, "unknown named grammar '@" + name + "'"); + } + if (std::holds_alternative(input_it->second)) { + return std::get(input_it->second); + } + auto compiled_it = named_grammars_.compiled.find(name); + if (compiled_it != named_grammars_.compiled.end()) { + return compiled_it->second; + } + + auto active_it = std::find(named_grammars_.active.begin(), named_grammars_.active.end(), name); + if (active_it != named_grammars_.active.end()) { + std::ostringstream cycle; + for (auto it = active_it; it != named_grammars_.active.end(); ++it) { + if (it != active_it) { + cycle << " -> "; + } + cycle << "@" << *it; + } + cycle << " -> @" << name; + RaiseLarkError(source_, location, "circular named grammar reference: " + cycle.str()); + } + + named_grammars_.active.push_back(name); + try { + const std::string& named_source = std::get(input_it->second); + auto tokens = LarkLexer(named_source).Tokenize(); + auto document = LarkParser(named_source, std::move(tokens)).Parse(); + Grammar compiled = + LarkCompiler(named_source, std::move(document), tokenizer_info_, named_grammars_) + .Compile(); + auto compiled_it = named_grammars_.compiled.emplace(name, std::move(compiled)).first; + named_grammars_.active.pop_back(); + return compiled_it->second; + } catch (const std::exception& error) { + named_grammars_.active.pop_back(); + RaiseLarkError( + source_, + location, + "failed to compile named grammar '@" + name + "': " + std::string(error.what()) + ); + } + } + + int32_t CompileNode( + const Node& node, const std::string& rule_hint, bool terminal_mode, bool append_skip = true + ) { + switch (node.kind) { + case Node::Kind::kSequence: { + if (node.children.empty()) { + return builder_.AddEmptyStr(); + } + std::vector elements; + elements.reserve(node.children.size()); + for (const Node& child : node.children) { + elements.push_back(CompileNode(child, rule_hint, terminal_mode, append_skip)); + } + return elements.size() == 1 ? elements[0] : builder_.AddSequence(elements); + } + case Node::Kind::kChoice: { + std::vector choices; + choices.reserve(node.children.size()); + for (const Node& child : node.children) { + choices.push_back(CompileNode(child, rule_hint, terminal_mode, append_skip)); + } + return choices.size() == 1 ? choices[0] : builder_.AddChoices(choices); + } + case Node::Kind::kRepeat: { + int32_t child = + CompileNode(node.children[0], rule_hint + "_repeat", terminal_mode, append_skip); + return builder_.AddRepeatFromExpr( + rule_hint + "_repeat", child, node.min_repeat, node.max_repeat + ); + } + case Node::Kind::kName: { + auto definition_it = definition_by_name_.find(node.text); + if (definition_it == definition_by_name_.end()) { + RaiseLarkError(source_, node.location, "unknown name '" + node.text + "'"); + } + if (terminal_mode && !definition_it->second->is_terminal) { + RaiseLarkError( + source_, node.location, "terminal cannot reference rule '" + node.text + "'" + ); + } + int32_t result = builder_.AddRuleRef(rule_ids_.at(node.text)); + // Lazy rules are compiled like terminals (lexemes), so they also take a trailing skip. + // Temperature rules are also compiled like terminals. + bool is_lexeme = definition_it->second->is_terminal || + HasLazySemantics(*definition_it->second) || + definition_it->second->temperature.has_value(); + return !terminal_mode && append_skip && is_lexeme ? AppendSkip(result) : result; + } + case Node::Kind::kString: { + int32_t result = CompileStringLiteral(node); + return !terminal_mode && append_skip && !node.text.empty() ? AppendSkip(result) : result; + } + case Node::Kind::kRange: { + auto begin = ParseUTF8(node.text.c_str()); + auto end = ParseUTF8(node.text2.c_str()); + if (begin.size() != 1 || end.size() != 1 || begin[0] == CharHandlingError::kInvalidUTF8 || + end[0] == CharHandlingError::kInvalidUTF8) { + RaiseLarkError(source_, node.location, "character range endpoints must be one character"); + } + if (begin[0] > end[0]) { + RaiseLarkError(source_, node.location, "character range start must not exceed end"); + } + int32_t result = builder_.AddCharacterClass({{begin[0], end[0]}}); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } + case Node::Kind::kRegex: { + RegexFlags flags = ParseRegexFlags(node); + std::string pattern = RewriteRegexDots(node.text, flags.dot_all); + if (flags.case_insensitive) { + // The FSM regex engine handles the (?i) prefix with ASCII case folding. Validate the + // pattern eagerly so that errors carry the source location. + std::string flagged_pattern = "(?i)" + pattern; + auto matches_empty = RegexFSMBuilder::MatchesEmpty(flagged_pattern); + if (matches_empty.IsErr()) { + RaiseLarkError( + source_, + node.location, + "failed to compile regular expression: " + + std::string(std::move(matches_empty).UnwrapErr().what()) + ); + } + int32_t regex_rule_id = + builder_.AddRuleWithHint(rule_hint + "_regex", builder_.AddRegex(flagged_pattern)); + int32_t result = builder_.AddRuleRef(regex_rule_id); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } + try { + int32_t root = SubGrammarAdder::Apply(&builder_, Grammar::FromRegex(pattern)); + int32_t result = builder_.AddRuleRef(root); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + node.location, + "failed to compile regular expression: " + std::string(error.what()) + ); + } + } + case Node::Kind::kJson: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "%json cannot be used in terminals"); + } + try { + int32_t root = SubGrammarAdder::Apply(&builder_, Grammar::FromJSONSchema(node.text)); + int32_t result = builder_.AddRuleRef(root); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + node.location, + "failed to compile inline JSON schema: " + std::string(error.what()) + ); + } + } + case Node::Kind::kNestedLark: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "nested %lark cannot be used in terminals"); + } + try { + LarkCompiler compiler(source_, *node.nested, tokenizer_info_, named_grammars_); + int32_t root = SubGrammarAdder::Apply(&builder_, compiler.Compile()); + int32_t result = builder_.AddRuleRef(root); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + node.location, + "failed to compile nested Lark grammar: " + std::string(error.what()) + ); + } + } + case Node::Kind::kSpecialToken: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "special tokens cannot be used in terminals"); + } + SpecialTokenSet token_set = ResolveSpecialToken(node.text, node.location); + int32_t result = token_set.excluded ? builder_.AddExcludeTokenSet(token_set.token_ids) + : builder_.AddTokenSet(token_set.token_ids); + return append_skip ? AppendSkip(result) : result; + } + case Node::Kind::kRegexExt: { + int32_t result = CompileStructuredRegex(node, rule_hint); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } + case Node::Kind::kGrammarRef: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "named grammars cannot be used in terminals"); + } + std::string name = node.text.substr(1); + auto root_it = named_grammar_roots_.find(name); + if (root_it == named_grammar_roots_.end()) { + int32_t root = + SubGrammarAdder::Apply(&builder_, ResolveNamedGrammar(name, node.location)); + root_it = named_grammar_roots_.emplace(name, root).first; + } + int32_t result = builder_.AddRuleRef(root_it->second); + return append_skip ? AppendSkip(result) : result; + } + case Node::Kind::kNot: + RaiseLarkError( + source_, node.location, "regular-expression complement '~' is not supported" + ); + } + RaiseLarkError(source_, node.location, "unsupported grammar node"); + } + + int32_t AppendSkip(int32_t expression) { + if (skip_rule_id_ == -1) { + return expression; + } + return builder_.AddSequence({expression, builder_.AddRuleRef(skip_rule_id_)}); + } + + /*! \brief Compile a rule with a token or character budget. */ + int32_t CompileBudgetRule(const Definition& definition) { + int32_t rule_id = rule_ids_.at(definition.name); + if (definition.max_tokens.has_value()) { + builder_.UpdateMaxTokens(rule_id, definition.max_tokens.value()); + } + if (definition.max_chars.has_value()) { + builder_.UpdateMaxChars(rule_id, definition.max_chars.value()); + } + if (HasLazySemantics(definition)) { + return CompileLazyRule(definition); + } + return CompileNode(definition.body, definition.name, false); + } + + SpecialTokenSet ResolveSpecialToken(const std::string& token, const Location& location) const { + if (token.size() >= 4 && token.substr(0, 2) == "<[" && token.substr(token.size() - 2) == "]>") { + std::string contents = token.substr(2, token.size() - 4); + SpecialTokenSet result; + if (!contents.empty() && contents[0] == '^') { + result.excluded = true; + contents.erase(contents.begin()); + } + if (contents == "*") { + if (result.excluded) { + RaiseLarkError(source_, location, "negated wildcard special token is not supported"); + } + if (!tokenizer_info_.has_value()) { + RaiseLarkError(source_, location, "wildcard special token requires tokenizer_info"); + } + result.token_ids.reserve(tokenizer_info_->GetVocabSize()); + for (int32_t token_id = 0; token_id < tokenizer_info_->GetVocabSize(); ++token_id) { + result.token_ids.push_back(token_id); + } + return result; + } + if (contents.find('*') != std::string::npos) { + RaiseLarkError(source_, location, "wildcard cannot be mixed with token ranges"); + } + size_t offset = 0; + while (offset <= contents.size()) { + size_t comma = contents.find(',', offset); + std::string range = Trim(contents.substr(offset, comma - offset)); + if (!range.empty()) { + size_t dash = range.find('-'); + if (dash != std::string::npos && range.find('-', dash + 1) != std::string::npos) { + RaiseLarkError( + source_, location, "invalid numeric special-token range '" + range + "'" + ); + } + int64_t first; + int64_t last; + try { + auto parse_token_id = [](const std::string& value) { + std::string trimmed = Trim(value); + size_t parsed = 0; + int64_t result = std::stoll(trimmed, &parsed); + if (parsed != trimmed.size()) { + throw std::invalid_argument("trailing characters"); + } + return result; + }; + first = parse_token_id(range.substr(0, dash)); + last = dash == std::string::npos ? first : parse_token_id(range.substr(dash + 1)); + } catch (const std::exception&) { + RaiseLarkError( + source_, location, "invalid numeric special-token range '" + range + "'" + ); + } + if (first < 0 || last < first || last > std::numeric_limits::max()) { + RaiseLarkError( + source_, location, "invalid numeric special-token range '" + range + "'" + ); + } + if (last - first > 1'000'000) { + RaiseLarkError(source_, location, "special-token range is too large"); + } + for (int64_t token_id = first; token_id <= last; ++token_id) { + result.token_ids.push_back(static_cast(token_id)); + } + } + if (comma == std::string::npos) { + break; + } + offset = comma + 1; + } + if (result.token_ids.empty()) { + RaiseLarkError(source_, location, "empty numeric special-token range"); + } + std::sort(result.token_ids.begin(), result.token_ids.end()); + result.token_ids.erase( + std::unique(result.token_ids.begin(), result.token_ids.end()), result.token_ids.end() + ); + return result; + } + + if (!tokenizer_info_.has_value()) { + RaiseLarkError( + source_, location, "named special token " + token + " requires tokenizer_info" + ); + } + SpecialTokenSet result; + const auto& decoded_vocab = tokenizer_info_->GetDecodedVocab(); + for (int32_t token_id = 0; token_id < static_cast(decoded_vocab.size()); ++token_id) { + if (decoded_vocab[token_id] == token) { + result.token_ids.push_back(token_id); + } + } + if (result.token_ids.empty()) { + RaiseLarkError(source_, location, "unknown special token " + token); + } + return result; + } + + static const Node* UnwrapSingle(const Node* node) { + while (node->kind == Node::Kind::kSequence && node->children.size() == 1) { + node = &node->children[0]; + } + return node; + } + + bool IsAnyText(const Node& node, std::unordered_set* visiting = nullptr) const { + if (node.kind == Node::Kind::kRegex) { + std::string pattern; + for (char c : node.text) { + if (c != ' ' && c != '\t' && c != '\r') { + pattern.push_back(c); + } + } + if (node.flags.find_first_not_of("isu") != std::string::npos) { + return false; + } + if (node.flags.find('s') != std::string::npos) { + return pattern == ".*"; + } + return pattern == "(.|\\n)*" || pattern == "(\\n|.)*" || pattern == "(?s:.*)" || + pattern == "(?:.|\\n)*" || pattern == "(?:\\n|.)*" || pattern == "[\\s\\S]*"; + } + if (node.kind == Node::Kind::kSequence && node.children.size() == 1) { + return IsAnyText(node.children[0], visiting); + } + if (node.kind == Node::Kind::kName) { + std::unordered_set local_visiting; + if (visiting == nullptr) { + visiting = &local_visiting; + } + if (visiting->count(node.text)) { + return false; + } + auto it = definition_by_name_.find(node.text); + if (it == definition_by_name_.end()) { + return false; + } + visiting->insert(node.text); + bool result = IsAnyText(it->second->body, visiting); + visiting->erase(node.text); + return result; + } + return false; + } + + std::optional ExtractLazyRegexTrigger(const Node& node) const { + if (node.kind != Node::Kind::kRegex) { + return std::nullopt; + } + if (node.flags.find('i') != std::string::npos || + node.flags.find_first_not_of("su") != std::string::npos) { + return std::nullopt; + } + std::vector prefixes; + if (node.flags.find('s') != std::string::npos) { + prefixes = {".*"}; + } else { + prefixes = {"(.|\\n)*", "(\\n|.)*", "(?:.|\\n)*", "(?:\\n|.)*", "[\\s\\S]*", "(?s:.*)"}; + } + for (const std::string& prefix : prefixes) { + if (node.text.size() <= prefix.size() || node.text.compare(0, prefix.size(), prefix) != 0) { + continue; + } + auto trigger = ParseFixedRegexLiteral(node.text.substr(prefix.size())); + if (trigger.has_value() && !trigger->empty()) { + return Trigger{Trigger::Level::kString, std::move(trigger.value()), {}, node.location}; + } + } + return std::nullopt; + } + + std::optional ExtractLazyTrigger(const Definition& definition) const { + if (definition.stop.has_value()) { + const Node& marker = definition.stop.value(); + if (!IsAnyText(definition.body) || marker.kind != Node::Kind::kString || + !marker.flags.empty()) { + return std::nullopt; + } + return Trigger{Trigger::Level::kString, marker.text, {}, definition.stop_location}; + } + if (definition.suffix.has_value()) { + const Node& marker = definition.suffix.value(); + if (!IsAnyText(definition.body) || marker.kind != Node::Kind::kString || + !marker.flags.empty()) { + return std::nullopt; + } + return Trigger{Trigger::Level::kString, marker.text, {}, definition.suffix_location}; + } + if (!definition.lazy) { + return std::nullopt; + } + const Node* body = UnwrapSingle(&definition.body); + if (body->kind == Node::Kind::kRegex) { + auto regex_trigger = ExtractLazyRegexTrigger(*body); + if (regex_trigger.has_value()) { + return regex_trigger; + } + } + if (definition.body.kind != Node::Kind::kSequence || definition.body.children.size() != 2 || + !IsAnyText(definition.body.children[0])) { + return std::nullopt; + } + const Node& trigger = definition.body.children[1]; + if (trigger.kind == Node::Kind::kString && !trigger.text.empty() && trigger.flags.empty()) { + return Trigger{Trigger::Level::kString, trigger.text, {}, trigger.location}; + } + if (trigger.kind == Node::Kind::kSpecialToken) { + SpecialTokenSet token_set = ResolveSpecialToken(trigger.text, trigger.location); + if (token_set.excluded) { + RaiseLarkError(source_, trigger.location, "lazy special-token trigger cannot be negated"); + } + return Trigger{Trigger::Level::kToken, "", token_set.token_ids, trigger.location}; + } + return std::nullopt; + } + + int32_t CompileLazyRule(const Definition& definition) { + int32_t rule_id = rule_ids_.at(definition.name); + const Node* marker = definition.suffix.has_value() + ? &definition.suffix.value() + : (definition.stop.has_value() ? &definition.stop.value() : nullptr); + bool marker_has_fixed_byte_length = marker != nullptr && marker->kind == Node::Kind::kString; + int32_t hidden_bytes = + marker_has_fixed_byte_length ? static_cast(marker->text.size()) : 1; + Grammar::Impl::SuffixStopInfo suffix_stop_info; + if (definition.suffix.has_value()) { + suffix_stop_info.hidden_suffix_bytes = hidden_bytes; + } else if (definition.stop.has_value()) { + suffix_stop_info.hidden_stop_bytes = hidden_bytes; + } + if (definition.stop_capture_name.has_value()) { + suffix_stop_info.stop_capture_name = definition.stop_capture_name.value(); + } + const Node* body = UnwrapSingle(&definition.body); + if (!definition.suffix.has_value() && !definition.stop.has_value() && + body->kind == Node::Kind::kRegex && ExtractLazyRegexTrigger(*body).has_value()) { + RaiseLarkError( + source_, + definition.location, + "lazy regex suffix is only supported on a head used by dynamic dispatch" + ); + } + std::optional body_pattern; + std::optional marker_pattern; + if (marker != nullptr && (!marker_has_fixed_byte_length || definition.max_tokens.has_value() || + definition.max_chars.has_value())) { + body_pattern = TerminalNodeToRegex(definition.body); + marker_pattern = TerminalNodeToRegex(*marker); + int32_t body_helper_expr = builder_.AddRegex(body_pattern.value()); + int32_t body_helper_rule = + builder_.AddRuleWithHint(definition.name + "_stop_body", body_helper_expr); + int32_t marker_helper_expr = builder_.AddRegex(marker_pattern.value()); + int32_t marker_helper_rule = + builder_.AddRuleWithHint(definition.name + "_stop_marker", marker_helper_expr); + suffix_stop_info.body_rule_id = body_helper_rule; + suffix_stop_info.marker_rule_id = marker_helper_rule; + } + builder_.UpdateSuffixStopInfo(rule_id, suffix_stop_info); + auto trigger = ExtractLazyTrigger(definition); + if (!trigger.has_value()) { + // General committed-shortest lazy rule: compiled like a terminal (no skip insertion); + // the terminal-like requirement is validated after grammar optimization. suffix="s" and + // stop="s" both desugar to the lazy rule over (body "s"); their only difference is capture + // scope, represented by the metadata set above. + builder_.UpdateLazy(rule_id, true); + if (marker != nullptr) { + if (!body_pattern.has_value()) { + body_pattern = TerminalNodeToRegex(definition.body); + marker_pattern = TerminalNodeToRegex(*marker); + } + return builder_.AddRegex( + "(?:" + body_pattern.value() + ")(?:" + marker_pattern.value() + ")" + ); + } + return CompileNode(definition.body, definition.name, true); + } + int32_t empty_rule = builder_.AddRuleWithHint("lark_lazy_end", builder_.AddEmptyStr()); + int32_t result; + if (trigger->level == Trigger::Level::kString) { + result = builder_.AddTagDispatch({{{trigger->string, empty_rule}}, false, {}}); + } else { + Grammar::Impl::TokenTagDispatch dispatch; + for (int32_t token_id : trigger->token_ids) { + dispatch.trigger_rule_pairs.push_back({token_id, empty_rule}); + } + dispatch.loop_after_dispatch = false; + result = builder_.AddTokenTagDispatch(dispatch); + } + return AppendSkip(result); + } + + static std::vector FlattenSequence(const Node& node) { + if (node.kind == Node::Kind::kSequence) { + return node.children; + } + return {node}; + } + + std::optional CompileDynamicStart(const Definition& start) { + std::unordered_set unused_rules; + std::vector start_elements = FlattenSequence(start.body); + if (start_elements.size() != 2) { + return std::nullopt; + } + const Node* loop = UnwrapSingle(&start_elements[0]); + if (loop->kind != Node::Kind::kRepeat || loop->min_repeat != 0 || loop->max_repeat != -1) { + return std::nullopt; + } + const Node* loop_body = UnwrapSingle(&loop->children[0]); + std::vector tool_names; + if (loop_body->kind == Node::Kind::kChoice) { + for (const Node& alternative : loop_body->children) { + const Node* name = UnwrapSingle(&alternative); + if (name->kind != Node::Kind::kName) { + return std::nullopt; + } + tool_names.push_back(name->text); + } + } else if (loop_body->kind == Node::Kind::kName) { + tool_names.push_back(loop_body->text); + } else { + return std::nullopt; + } + + const Node* tail_name = UnwrapSingle(&start_elements[1]); + if (tail_name->kind != Node::Kind::kName) { + return std::nullopt; + } + auto tail_it = definition_by_name_.find(tail_name->text); + if (tail_it == definition_by_name_.end() || !IsAnyText(tail_it->second->body)) { + return std::nullopt; + } + unused_rules.insert(tail_name->text); + + std::vector alternatives; + for (const std::string& tool_name : tool_names) { + auto tool_it = definition_by_name_.find(tool_name); + if (tool_it == definition_by_name_.end() || tool_it->second->is_terminal) { + return std::nullopt; + } + unused_rules.insert(tool_name); + std::vector tool_elements = FlattenSequence(tool_it->second->body); + if (tool_elements.empty()) { + return std::nullopt; + } + + std::optional trigger; + int32_t marker_event_rule_id = -1; + size_t remainder_begin = 0; + const Node* first = UnwrapSingle(&tool_elements[0]); + if (first->kind == Node::Kind::kName) { + auto head_it = definition_by_name_.find(first->text); + if (head_it != definition_by_name_.end()) { + trigger = ExtractLazyTrigger(*head_it->second); + if (trigger.has_value()) { + unused_rules.insert(first->text); + remainder_begin = 1; + const Definition& head = *head_it->second; + if (head.stop.has_value() || head.stop_capture_name.has_value()) { + // The dispatch FSM consumes the trigger before entering the remainder. Insert a + // zero-width rule there so capture materialization can recover that preceding + // marker without giving up the deterministic dispatch path. + int32_t event_expr = builder_.AddEmptyStr(); + marker_event_rule_id = + builder_.AddRuleWithHint(head.name + "_dynamic_marker", event_expr); + int32_t marker_expr = builder_.AddByteString(trigger->string); + int32_t marker_rule_id = + builder_.AddRuleWithHint(head.name + "_dynamic_marker_text", marker_expr); + Grammar::Impl::SuffixStopInfo suffix_stop_info; + suffix_stop_info.body_rule_id = marker_event_rule_id; + suffix_stop_info.marker_rule_id = marker_rule_id; + int32_t hidden_bytes = static_cast(trigger->string.size()); + if (head.stop.has_value()) { + suffix_stop_info.hidden_stop_bytes = hidden_bytes; + } else { + suffix_stop_info.hidden_suffix_bytes = hidden_bytes; + } + if (head.stop_capture_name.has_value()) { + suffix_stop_info.stop_capture_name = head.stop_capture_name.value(); + } + builder_.UpdateSuffixStopInfo(marker_event_rule_id, suffix_stop_info); + } + } + } + } + if (!trigger.has_value() && tool_elements.size() >= 2 && IsAnyText(tool_elements[0])) { + const Node* token_trigger = UnwrapSingle(&tool_elements[1]); + if (token_trigger->kind == Node::Kind::kSpecialToken) { + SpecialTokenSet token_set = + ResolveSpecialToken(token_trigger->text, token_trigger->location); + if (token_set.excluded) { + RaiseLarkError( + source_, token_trigger->location, "dynamic special-token trigger cannot be negated" + ); + } + trigger = + Trigger{Trigger::Level::kToken, "", token_set.token_ids, token_trigger->location}; + remainder_begin = 2; + } + } + if (!trigger.has_value()) { + return std::nullopt; + } + + Node remainder; + remainder.kind = Node::Kind::kSequence; + remainder.location = tool_it->second->location; + remainder.children.assign( + tool_elements.begin() + static_cast(remainder_begin), tool_elements.end() + ); + alternatives.push_back( + {std::move(trigger.value()), std::move(remainder), marker_event_rule_id} + ); + } + + if (alternatives.empty()) { + return std::nullopt; + } + Trigger::Level level = alternatives[0].trigger.level; + for (const auto& alternative : alternatives) { + if (alternative.trigger.level != level) { + RaiseLarkError( + source_, + start.location, + "a dynamic Lark start rule cannot mix string and token triggers" + ); + } + } + + if (level == Trigger::Level::kString) { + std::unordered_map> grouped; + std::vector trigger_order; + for (const auto& alternative : alternatives) { + if (!grouped.count(alternative.trigger.string)) { + trigger_order.push_back(alternative.trigger.string); + } + grouped[alternative.trigger.string].push_back(&alternative); + } + Grammar::Impl::TagDispatch dispatch; + dispatch.loop_after_dispatch = true; + for (const std::string& trigger : trigger_order) { + std::vector remainder_choices; + for (const DynamicAlternative* alternative : grouped.at(trigger)) { + int32_t remainder = CompileNode(alternative->remainder, "lark_dynamic_body", false); + if (alternative->marker_event_rule_id >= 0) { + remainder = builder_.AddSequence( + {builder_.AddRuleRef(alternative->marker_event_rule_id), remainder} + ); + } + remainder_choices.push_back(remainder); + } + int32_t body = remainder_choices.size() == 1 ? remainder_choices[0] + : builder_.AddChoices(remainder_choices); + int32_t body_rule = builder_.AddRuleWithHint("lark_dynamic_body", body); + dispatch.tag_rule_pairs.push_back({trigger, body_rule}); + } + dynamic_unused_rules_ = std::move(unused_rules); + return builder_.AddTagDispatch(dispatch); + } + + std::unordered_map> grouped; + std::vector token_order; + for (const auto& alternative : alternatives) { + for (int32_t token_id : alternative.trigger.token_ids) { + if (!grouped.count(token_id)) { + token_order.push_back(token_id); + } + grouped[token_id].push_back(&alternative); + } + } + Grammar::Impl::TokenTagDispatch dispatch; + dispatch.loop_after_dispatch = true; + for (int32_t token_id : token_order) { + std::vector remainder_choices; + for (const DynamicAlternative* alternative : grouped.at(token_id)) { + remainder_choices.push_back( + CompileNode(alternative->remainder, "lark_dynamic_token_body", false) + ); + } + int32_t body = remainder_choices.size() == 1 ? remainder_choices[0] + : builder_.AddChoices(remainder_choices); + int32_t body_rule = builder_.AddRuleWithHint("lark_dynamic_token_body", body); + dispatch.trigger_rule_pairs.push_back({token_id, body_rule}); + } + dynamic_unused_rules_ = std::move(unused_rules); + return builder_.AddTokenTagDispatch(dispatch); + } + + const std::string& source_; + Document document_; + const std::optional& tokenizer_info_; + NamedGrammarRegistry& named_grammars_; + GrammarBuilder builder_; + std::unordered_map definition_by_name_; + std::unordered_map rule_ids_; + std::unordered_map named_grammar_roots_; + int32_t skip_rule_id_ = -1; + bool allow_initial_skip_ = false; + std::unordered_set dynamic_unused_rules_; +}; + +} // namespace + +Grammar LarkToGrammar( + const std::string& lark_string, + const std::optional& tokenizer_info, + const std::vector& named_grammars +) { + NamedGrammarRegistry named_grammar_registry; + for (const auto& [name, grammar_or_source] : named_grammars) { + if (name.empty()) { + throw XGrammarError("Named grammar names must not be empty"); + } + if (!std::all_of(name.begin(), name.end(), [](unsigned char character) { + return std::isalnum(character) || character == '_' || character == '-'; + })) { + throw XGrammarError( + "Invalid named grammar name '" + name + + "': names may contain only letters, digits, underscores, and hyphens" + ); + } + if (!named_grammar_registry.inputs.emplace(name, grammar_or_source).second) { + throw XGrammarError("Duplicate named grammar '" + name + "'"); + } + } + auto tokens = LarkLexer(lark_string).Tokenize(); + auto document = LarkParser(lark_string, std::move(tokens)).Parse(); + return LarkCompiler(lark_string, std::move(document), tokenizer_info, named_grammar_registry) + .Compile(); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/lark_converter.h b/third_party/xgrammar/cpp/lark_converter.h new file mode 100644 index 0000000000..411652047f --- /dev/null +++ b/third_party/xgrammar/cpp/lark_converter.h @@ -0,0 +1,26 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/lark_converter.h + * \brief Convert Lark syntax to XGrammar Grammar IR. + */ + +#ifndef XGRAMMAR_LARK_CONVERTER_H_ +#define XGRAMMAR_LARK_CONVERTER_H_ + +#include + +#include +#include +#include + +namespace xgrammar { + +Grammar LarkToGrammar( + const std::string& lark_string, + const std::optional& tokenizer_info = std::nullopt, + const std::vector& named_grammars = {} +); + +} // namespace xgrammar + +#endif // XGRAMMAR_LARK_CONVERTER_H_ diff --git a/third_party/xgrammar/cpp/regex_converter.cc b/third_party/xgrammar/cpp/regex_converter.cc new file mode 100644 index 0000000000..b46bd76ada --- /dev/null +++ b/third_party/xgrammar/cpp/regex_converter.cc @@ -0,0 +1,413 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/regex_converter.cc + */ +#include "regex_converter.h" + +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/logging.h" +#include "support/utils.h" + +namespace xgrammar { + +/*! + * \brief Convert a regex to EBNF. + * \details The implementation refers to the regex described in + * https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Regular_expressions + */ +class RegexConverter { + public: + explicit RegexConverter(const std::string& regex) : regex_(regex) { + if (!regex.empty()) { + // ParseUTF8 takes a C string and stops at the first NUL, so a regex + // containing one would either yield an empty codepoint vector (leading + // NUL -- an out-of-bounds read below) or be silently truncated to the + // prefix. A NUL is never meaningful in a regex, so reject it here. + if (regex.find('\0') != std::string::npos) { + XGRAMMAR_LOG(FATAL) << "The regex must not contain null characters."; + XGRAMMAR_UNREACHABLE(); + } + regex_codepoints_ = ParseUTF8(regex_.c_str(), false); + if (regex_codepoints_[0] == kInvalidUTF8) { + XGRAMMAR_LOG(FATAL) << "The regex is not a valid UTF-8 string."; + XGRAMMAR_UNREACHABLE(); + } + } + regex_codepoints_.push_back(0); // Add a null terminator + } + std::string Convert(); + + private: + /** + * \brief Add a segment string to the result EBNF string. It especially adds a space if needed + * and add_space is true. + */ + void AddEBNFSegment(const std::string& element); + + [[noreturn]] void RaiseError(const std::string& message); + void RaiseWarning(const std::string& message); + + std::string HandleCharacterClass(); + std::string HandleRepetitionRange(); + std::string HandleCharEscape(); + std::string HandleEscape(); + std::string HandleEscapeInCharClass(); + /** + * \brief Handle group modifier. The general format is "(?" + modifier + content + ")". E.g. + * "(?:abc)" is a non-capturing group. + */ + void HandleGroupModifier(); + + std::string regex_; + std::vector regex_codepoints_; + TCodepoint* start_; + TCodepoint* current_; + TCodepoint* end_; + std::string result_ebnf_; + int parenthesis_level_ = 0; +}; + +void RegexConverter::AddEBNFSegment(const std::string& element) { + if (!result_ebnf_.empty()) { + result_ebnf_ += ' '; + } + result_ebnf_ += element; +} + +void RegexConverter::RaiseError(const std::string& message) { + XGRAMMAR_LOG(FATAL) << "Regex parsing error at position " << current_ - start_ + 1 << ": " + << message; + XGRAMMAR_UNREACHABLE(); +} + +void RegexConverter::RaiseWarning(const std::string& message) { + XGRAMMAR_LOG(WARNING) << "Regex parsing warning at position " << current_ - start_ + 1 << ": " + << message; +} + +std::string RegexConverter::HandleCharacterClass() { + std::string char_class = "["; + ++current_; + if (*current_ == ']') { + RaiseError("Empty character class is not allowed in regex."); + } + while (*current_ != ']' && current_ != end_) { + if (*current_ == '\\') { + char_class += HandleEscapeInCharClass(); + } else { + char_class += CharToUTF8(*current_); + ++current_; + } + } + if (current_ == end_) { + RaiseError("Unclosed '['"); + } + char_class += ']'; + ++current_; + return char_class; +} + +// {x}: Match exactly x occurrences of the preceding regular expression. +// {x,} +// {x,y} +std::string RegexConverter::HandleRepetitionRange() { + std::string result = "{"; + ++current_; + if (!isdigit(*current_)) { + RaiseError("Invalid repetition count."); + } + while (isdigit(*current_)) { + result += static_cast(*current_); + ++current_; + } + if (*current_ != ',' && *current_ != '}') { + RaiseError("Invalid repetition count."); + } + result += static_cast(*current_); + ++current_; + if (current_[-1] == '}') { + // Matches {x} + return result; + } + if (!isdigit(*current_) && *current_ != '}') { + RaiseError("Invalid repetition count."); + } + while (isdigit(*current_)) { + result += static_cast(*current_); + ++current_; + } + if (*current_ != '}') { + RaiseError("Invalid repetition count."); + } + result += '}'; + ++current_; + return result; +} + +std::string RegexConverter::HandleCharEscape() { + // clang-format off + static const std::unordered_map CUSTOM_ESCAPE_MAP = { + {'^', '^'}, {'$', '$'}, {'.', '.'}, {'*', '*'}, {'+', '+'}, {'?', '?'}, {'\\', '\\'}, + {'(', '('}, {')', ')'}, {'[', '['}, {']', ']'}, {'{', '{'}, {'}', '}'}, {'|', '|'}, + {'/', '/'}, {'-', '-'} + }; + // clang-format on + if (end_ - current_ < 2 || (current_[1] == 'u' && end_ - current_ < 5) || + (current_[1] == 'x' && end_ - current_ < 4) || (current_[1] == 'c' && end_ - current_ < 3)) { + RaiseError("Escape sequence is not finished."); + } + auto [codepoint, len] = ParseNextEscaped(current_, CUSTOM_ESCAPE_MAP); + if (codepoint != CharHandlingError::kInvalidEscape) { + current_ += len; + return EscapeString(codepoint); + } else if (current_[1] == 'u' && current_[2] == '{') { + current_ += 3; + int len = 0; + TCodepoint value = 0; + while (HexCharToInt(current_[len]) != -1 && len <= 6) { + value = value * 16 + HexCharToInt(current_[len]); + ++len; + } + if (len == 0 || len > 6 || current_[len] != '}') { + RaiseError("Invalid Unicode escape sequence."); + } + current_ += len + 1; + return EscapeString(value); + } else if (current_[1] == 'c') { + current_ += 2; + if (!std::isalpha(*current_)) { + RaiseError("Invalid control character escape sequence."); + } + ++current_; + return EscapeString((*(current_ - 1)) % 32); + } else { + RaiseWarning( + "Escape sequence '\\" + EscapeString(current_[1]) + + "' is not recognized. The character itself will be matched" + ); + current_ += 2; + return EscapeString(current_[-1]); + } +} + +std::string RegexConverter::HandleEscapeInCharClass() { + if (end_ - current_ < 2) { + RaiseError("Escape sequence is not finished."); + } + if (current_[1] == 'd') { + current_ += 2; + return "0-9"; + } else if (current_[1] == 'D') { + current_ += 2; + return R"(\x00-\x2F\x3A-\U0010FFFF)"; + } else if (current_[1] == 'w') { + current_ += 2; + return "a-zA-Z0-9_"; + } else if (current_[1] == 'W') { + current_ += 2; + return R"(\x00-\x2F\x3A-\x40\x5B-\x5E\x60\x7B-\U0010FFFF)"; + } else if (current_[1] == 's') { + current_ += 2; + return R"(\f\n\r\t\v\u0020\u00a0)"; + } else if (current_[1] == 'S') { + current_ += 2; + return R"(\x00-\x08\x0E-\x1F\x21-\x9F\xA1-\U0010FFFF)"; + } else { + auto res = HandleCharEscape(); + if (res == "]" || res == "-") { + return "\\" + res; + } else { + return res; + } + } +} + +std::string RegexConverter::HandleEscape() { + // clang-format off + static const std::unordered_map CUSTOM_ESCAPE_MAP = { + {'^', '^'}, {'$', '$'}, {'.', '.'}, {'*', '*'}, {'+', '+'}, {'?', '?'}, {'\\', '\\'}, + {'(', '('}, {')', ')'}, {'[', '['}, {']', ']'}, {'{', '{'}, {'}', '}'}, {'|', '|'}, + {'/', '/'} + }; + // clang-format on + if (end_ - current_ < 2) { + RaiseError("Escape sequence is not finished."); + } + if (current_[1] == 'd') { + current_ += 2; + return "[0-9]"; + } else if (current_[1] == 'D') { + current_ += 2; + return "[^0-9]"; + } else if (current_[1] == 'w') { + current_ += 2; + return "[a-zA-Z0-9_]"; + } else if (current_[1] == 'W') { + current_ += 2; + return "[^a-zA-Z0-9_]"; + } else if (current_[1] == 's') { + current_ += 2; + return R"([\f\n\r\t\v\u0020\u00a0])"; + } else if (current_[1] == 'S') { + current_ += 2; + return R"([^\f\n\r\t\v\u0020\u00a0])"; + } else if ((current_[1] >= '1' && current_[1] <= '9') || current_[1] == 'k') { + RaiseError("Backreference is not supported yet."); + } else if (current_[1] == 'p' || current_[1] == 'P') { + RaiseError("Unicode character class escape sequence is not supported yet."); + } else if (current_[1] == 'b' || current_[1] == 'B') { + RaiseError("Word boundary is not supported yet."); + } else { + return "\"" + HandleCharEscape() + "\""; + } +} + +void RegexConverter::HandleGroupModifier() { + if (current_ == end_) { + RaiseError("Group modifier is not finished."); + } + if (*current_ == ':') { + // Non-capturing group. + ++current_; + } else if (*current_ == '=' || *current_ == '!') { + // Positive or negative lookahead. + RaiseError("Lookahead is not supported yet."); + } else if (*current_ == '<' && current_ + 1 != end_ && + (current_[1] == '=' || current_[1] == '!')) { + // Positive or negative lookbehind. + RaiseError("Lookbehind is not supported yet."); + } else if (*current_ == '<') { + ++current_; + while (current_ != end_ && isalpha(*current_)) { + ++current_; + } + if (current_ == end_ || *current_ != '>') { + RaiseError("Invalid named capturing group."); + } + // Just ignore the named of the group. + ++current_; + } else { + // Group modifier flag. + RaiseError("Group modifier flag is not supported yet."); + } +} + +std::string RegexConverter::Convert() { + start_ = regex_codepoints_.data(); + current_ = start_; + end_ = start_ + regex_codepoints_.size() - 1; + bool is_empty = true; + while (current_ != end_) { + if (*current_ == '^') { + if (current_ != start_) { + RaiseWarning( + "'^' should be at the start of the regex, but found in the middle. It is ignored." + ); + } + ++current_; + } else if (*current_ == '$') { + if (current_ != end_ - 1) { + RaiseWarning( + "'$' should be at the end of the regex, but found in the middle. It is ignored." + ); + } + ++current_; + } else if (*current_ == '[') { + is_empty = false; + AddEBNFSegment(HandleCharacterClass()); + } else if (*current_ == '(') { + is_empty = false; + ++current_; + ++parenthesis_level_; + AddEBNFSegment("("); + if (current_ != end_ && *current_ == '?') { + ++current_; + HandleGroupModifier(); + } + } else if (*current_ == ')') { + is_empty = false; + if (parenthesis_level_ == 0) { + RaiseError("Unmatched ')'"); + } + // Empty alternative before ')' (e.g. "(a|)" or "(a|$)"): emit "" so it isn't a bare '|'. + if (!result_ebnf_.empty() && result_ebnf_.back() == '|') { + AddEBNFSegment("\"\""); + } + --parenthesis_level_; + AddEBNFSegment(")"); + ++current_; + } else if (*current_ == '*' || *current_ == '+' || *current_ == '?') { + is_empty = false; + result_ebnf_ += static_cast(*current_); + ++current_; + if (current_ != end_ && *current_ == '?') { + // Ignore the non-greedy modifier because our grammar handles all repetition numbers + // non-deterministically. + ++current_; + } + if (current_ != end_ && + (*current_ == '{' || *current_ == '*' || *current_ == '+' || *current_ == '?')) { + RaiseError("Two consecutive repetition modifiers are not allowed."); + } + } else if (*current_ == '{') { + is_empty = false; + result_ebnf_ += HandleRepetitionRange(); + if (current_ != end_ && *current_ == '?') { + // Still ignore the non-greedy modifier. + ++current_; + } + if (current_ != end_ && + (*current_ == '{' || *current_ == '*' || *current_ == '+' || *current_ == '?')) { + RaiseError("Two consecutive repetition modifiers are not allowed."); + } + } else if (*current_ == '|') { + is_empty = false; + // Empty alternative before '|': emit "" so there's no bare '|' on the left. + // Covers leading ("^$|abc"), consecutive ("a||b") and group-start ("(|a)") cases. + if (result_ebnf_.empty() || result_ebnf_.back() == '|' || result_ebnf_.back() == '(') { + AddEBNFSegment("\"\""); + } + AddEBNFSegment("|"); + ++current_; + } else if (*current_ == '\\') { + is_empty = false; + AddEBNFSegment(HandleEscape()); + } else if (*current_ == '.') { + is_empty = false; + AddEBNFSegment(R"([\u0000-\U0010FFFF])"); + ++current_; + } else { + is_empty = false; + // Non-special characters are matched literally. + AddEBNFSegment("\"" + EscapeString(*current_) + "\""); + ++current_; + } + } + if (parenthesis_level_ != 0) { + RaiseError("The parenthesis is not closed."); + } + // Trailing empty alternative, e.g. "abc|": emit "" so it doesn't end with a bare '|'. + if (!result_ebnf_.empty() && result_ebnf_.back() == '|') { + AddEBNFSegment("\"\""); + } + if (is_empty) { + AddEBNFSegment("\"\""); + } + return result_ebnf_; +} + +std::string RegexToEBNF(const std::string& regex, bool with_rule_name) { + RegexConverter converter(regex); + if (with_rule_name) { + return "root ::= " + converter.Convert() + "\n"; + } else { + return converter.Convert(); + } +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/regex_converter.h b/third_party/xgrammar/cpp/regex_converter.h new file mode 100644 index 0000000000..fcb3442ac1 --- /dev/null +++ b/third_party/xgrammar/cpp/regex_converter.h @@ -0,0 +1,21 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/regex_converter.h + * \brief Convert a regex string to EBNF grammar string. + */ + +#ifndef XGRAMMAR_REGEX_CONVERTER_H_ +#define XGRAMMAR_REGEX_CONVERTER_H_ + +#include + +namespace xgrammar { + +/*! + * \brief Convert a regex string to EBNF grammar string. + */ +std::string RegexToEBNF(const std::string& regex, bool with_rule_name = true); + +} // namespace xgrammar + +#endif // XGRAMMAR_REGEX_CONVERTER_H_ diff --git a/third_party/xgrammar/cpp/structural_tag.cc b/third_party/xgrammar/cpp/structural_tag.cc new file mode 100644 index 0000000000..3b9c616093 --- /dev/null +++ b/third_party/xgrammar/cpp/structural_tag.cc @@ -0,0 +1,2520 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/structural_tag.cc + */ +#include "structural_tag.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_builder.h" +#include "grammar_functor.h" +#include "grammar_impl.h" +#include "json_schema_converter.h" +#include "support/json_parse.h" +#include "support/logging.h" +#include "support/recursion_guard.h" +#include "support/utils.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +// Short alias for the error type. +using ISTError = InvalidStructuralTagError; + +// Forward declaration for helpers that convert Format to picojson::value. +picojson::value FormatToJSONValue(const Format& format); + +picojson::value StringVectorToJSONArray(const std::vector& vector) { + picojson::array array; + array.reserve(vector.size()); + for (const auto& string : vector) { + array.push_back(picojson::value(string)); + } + return picojson::value(std::move(array)); +} + +picojson::value FormatVectorToJSONArray(const std::vector& vector) { + picojson::array array; + array.reserve(vector.size()); + for (const auto& format : vector) { + array.push_back(xgrammar::FormatToJSONValue(format)); + } + return picojson::value(std::move(array)); +} + +picojson::value TagVectorToJSONArray(const std::vector& vector) { + picojson::array array; + array.reserve(vector.size()); + for (const auto& tag : vector) { + array.push_back(tag.ToJSON()); + } + return picojson::value(std::move(array)); +} + +picojson::value IntOrStringVectorToJSONArray( + const std::vector>& vec +) { + picojson::array array; + array.reserve(vec.size()); + for (const auto& item : vec) { + if (std::holds_alternative(item)) { + array.push_back(picojson::value(static_cast(std::get(item)))); + } else { + array.push_back(picojson::value(std::get(item))); + } + } + return picojson::value(std::move(array)); +} + +/******************** Format To JSON ********************/ +std::string FormatToJSON(const Format& format) { return FormatToJSONValue(format).serialize(); } + +picojson::value FormatToJSONValue(const Format& format) { + return std::visit([&](auto&& arg) { return arg.ToJSON(); }, format); +} + +picojson::value ConstStringFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["value"] = picojson::value(value); + return picojson::value(std::move(obj)); +} + +picojson::value JSONSchemaFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + picojson::value schema_val; + if (ParseJSON(schema_val, json_schema).empty()) { + obj["json_schema"] = schema_val; + } else { + obj["json_schema"] = picojson::value(json_schema); + } + obj["style"] = picojson::value(style); + obj["any_order"] = picojson::value(any_order); + if (max_whitespace_cnt.has_value()) { + obj["max_whitespace_cnt"] = picojson::value(static_cast(*max_whitespace_cnt)); + } else { + obj["max_whitespace_cnt"] = picojson::value(); + } + return picojson::value(std::move(obj)); +} + +picojson::value GrammarFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["grammar"] = picojson::value(grammar); + return picojson::value(std::move(obj)); +} + +picojson::value RegexFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["pattern"] = picojson::value(pattern); + return picojson::value(std::move(obj)); +} + +picojson::value AnyTextFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["excludes"] = StringVectorToJSONArray(excludes); + obj["detected_end_strs"] = StringVectorToJSONArray(detected_end_strs_); + if (max_tokens >= 0) { + obj["max_tokens"] = picojson::value(static_cast(max_tokens)); + } + if (max_chars >= 0) { + obj["max_chars"] = picojson::value(static_cast(max_chars)); + } + return picojson::value(std::move(obj)); +} + +// These two constructors are defined here rather than inline because instantiating +// vector against the still-incomplete Format variant is ill-formed under C++20. +SequenceFormat::SequenceFormat(std::vector elements) : elements(std::move(elements)) {} + +OrFormat::OrFormat(std::vector elements) : elements(std::move(elements)) {} + +picojson::value SequenceFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["elements"] = FormatVectorToJSONArray(elements); + return picojson::value(std::move(obj)); +} + +picojson::value OrFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["elements"] = FormatVectorToJSONArray(elements); + return picojson::value(std::move(obj)); +} + +picojson::value TagFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + if (std::holds_alternative(begin)) { + obj["begin"] = picojson::value(std::get(begin)); + } else { + obj["begin"] = std::get(begin).ToJSON(); + } + if (content) { + obj["content"] = FormatToJSONValue(*content); + } else { + obj["content"] = picojson::value(); + } + if (std::holds_alternative(end)) { + obj["end"] = std::get(end).ToJSON(); + } else { + const auto& end_strs = std::get>(end); + if (end_strs.size() == 1) { + obj["end"] = picojson::value(end_strs[0]); + } else { + obj["end"] = StringVectorToJSONArray(end_strs); + } + } + return picojson::value(std::move(obj)); +} + +picojson::value TokenFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + if (std::holds_alternative(token)) { + obj["token"] = picojson::value(static_cast(std::get(token))); + } else { + obj["token"] = picojson::value(std::get(token)); + } + return picojson::value(std::move(obj)); +} + +picojson::value ExcludeTokenFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + return picojson::value(std::move(obj)); +} + +picojson::value AnyTokensFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + if (max_tokens >= 0) { + obj["max_tokens"] = picojson::value(static_cast(max_tokens)); + } + return picojson::value(std::move(obj)); +} + +picojson::value TokenTriggeredTagsFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["trigger_tokens"] = IntOrStringVectorToJSONArray(trigger_tokens); + obj["tags"] = TagVectorToJSONArray(tags); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + obj["at_least_one"] = picojson::value(at_least_one); + obj["stop_after_first"] = picojson::value(stop_after_first); + return picojson::value(std::move(obj)); +} + +picojson::value TriggeredTagsFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["triggers"] = StringVectorToJSONArray(triggers); + obj["tags"] = TagVectorToJSONArray(tags); + obj["excludes"] = StringVectorToJSONArray(excludes); + obj["at_least_one"] = picojson::value(at_least_one); + obj["stop_after_first"] = picojson::value(stop_after_first); + obj["detected_end_strs"] = StringVectorToJSONArray(detected_end_strs_); + return picojson::value(std::move(obj)); +} + +picojson::value TagsWithSeparatorFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["tags"] = TagVectorToJSONArray(tags); + obj["separator"] = picojson::value(separator); + obj["at_least_one"] = picojson::value(at_least_one); + obj["stop_after_first"] = picojson::value(stop_after_first); + return picojson::value(std::move(obj)); +} + +picojson::value OptionalFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value PlusFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value StarFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value RepeatFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["min"] = picojson::value(static_cast(min)); + obj["max"] = picojson::value(static_cast(max)); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value DispatchFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + picojson::array rules_arr; + rules_arr.reserve(rules.size()); + for (const auto& pair : rules) { + picojson::array pair_arr; + pair_arr.push_back(picojson::value(pair.first)); + if (pair.second) { + pair_arr.push_back(FormatToJSONValue(*pair.second)); + } else { + pair_arr.push_back(picojson::value()); + } + rules_arr.push_back(picojson::value(std::move(pair_arr))); + } + obj["rules"] = picojson::value(std::move(rules_arr)); + obj["loop"] = picojson::value(loop); + obj["excludes"] = StringVectorToJSONArray(excludes); + return picojson::value(std::move(obj)); +} + +picojson::value TokenDispatchFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + picojson::array rules_arr; + rules_arr.reserve(rules.size()); + for (const auto& pair : rules) { + picojson::array pair_arr; + if (std::holds_alternative(pair.first)) { + pair_arr.push_back(picojson::value(static_cast(std::get(pair.first)))); + } else { + pair_arr.push_back(picojson::value(std::get(pair.first))); + } + if (pair.second) { + pair_arr.push_back(FormatToJSONValue(*pair.second)); + } else { + pair_arr.push_back(picojson::value()); + } + rules_arr.push_back(picojson::value(std::move(pair_arr))); + } + obj["rules"] = picojson::value(std::move(rules_arr)); + obj["loop"] = picojson::value(loop); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + return picojson::value(std::move(obj)); +} + +/************** StructuralTag Parser **************/ + +class StructuralTagParser { + public: + static Result FromJSON(const std::string& json); + + private: + Result ParseStructuralTag(const picojson::value& value); + + /*! + * \brief Parse a Format object from a JSON value. + * \param value The JSON value to parse. + * \return A Format object if the JSON is valid, otherwise an error message in std::runtime_error. + * \note The "type" field is checked in this function, and not checked in the Parse*Format + * functions. + */ + Result ParseFormat(const picojson::value& value); + Result ParseConstStringFormat(const picojson::object& value); + Result ParseJSONSchemaFormat( + const picojson::object& value, std::optional style_override = std::nullopt + ); + Result ParseAnyTextFormat(const picojson::object& value); + Result ParseGrammarFormat(const picojson::object& value); + Result ParseRegexFormat(const picojson::object& value); + Result ParseSequenceFormat(const picojson::object& value); + Result ParseOrFormat(const picojson::object& value); + /*! \brief ParseTagFormat with extra check for object and the type field. */ + Result ParseTagFormat(const picojson::value& value); + Result ParseTagFormat(const picojson::object& value); + Result ParseTriggeredTagsFormat(const picojson::object& value); + Result ParseTagsWithSeparatorFormat( + const picojson::object& value + ); + Result ParseOptionalFormat(const picojson::object& value); + Result ParsePlusFormat(const picojson::object& value); + Result ParseStarFormat(const picojson::object& value); + Result ParseRepeatFormat(const picojson::object& value); + Result ParseTokenFormat(const picojson::object& value); + Result ParseExcludeTokenFormat(const picojson::object& value); + Result ParseAnyTokensFormat(const picojson::object& value); + Result ParseTokenTriggeredTagsFormat( + const picojson::object& value + ); + Result ParseDispatchFormat(const picojson::object& value); + Result ParseTokenDispatchFormat(const picojson::object& value); + + int parse_format_recursion_depth_ = 0; +}; + +Result StructuralTagParser::FromJSON(const std::string& json) { + picojson::value value; + std::string err = ParseJSON(value, json); + if (!err.empty()) { + return ResultErr("Failed to parse JSON: " + err); + } + return Result::Convert( + StructuralTagParser().ParseStructuralTag(value) + ); +} + +Result StructuralTagParser::ParseStructuralTag(const picojson::value& value +) { + if (!value.is()) { + return ResultErr("Structural tag must be an object"); + } + const auto& obj = value.get(); + // The type field is optional but must be "structural_tag" if present. + if (obj.find("type") != obj.end()) { + if (!obj["type"].is() || obj["type"].get() != "structural_tag") { + return ResultErr("Structural tag's type must be a string \"structural_tag\""); + } + } + // The format field is required. + if (obj.find("format") == obj.end()) { + return ResultErr("Structural tag must have a format field"); + } + auto format = ParseFormat(obj["format"]); + if (format.IsErr()) { + return ResultErr(std::move(format).UnwrapErr()); + } + return ResultOk(std::move(format).Unwrap()); +} + +Result StructuralTagParser::ParseFormat(const picojson::value& value) { + RecursionGuard guard(&parse_format_recursion_depth_); + // The global recursion limit is far deeper than the native stack allows for the recursive passes + // over the format tree (a few KB per level, and Windows has a 1 MB main-thread stack), so cap + // the nesting of formats explicitly. + static constexpr int kMaxFormatDepth = 100; + if (parse_format_recursion_depth_ > kMaxFormatDepth) { + return ResultErr( + "Formats are nested deeper than " + std::to_string(kMaxFormatDepth) + " levels" + ); + } + if (!value.is()) { + return ResultErr("Format must be an object"); + } + const auto& obj = value.get(); + // If type is present, use it to determine the format. + if (obj.find("type") != obj.end()) { + if (!obj["type"].is()) { + return ResultErr("Format's type must be a string"); + } + auto type = obj["type"].get(); + if (type == "const_string") { + return Result::Convert(ParseConstStringFormat(obj)); + } else if (type == "json_schema") { + return Result::Convert(ParseJSONSchemaFormat(obj)); + } else if (type == "any_text") { + return Result::Convert(ParseAnyTextFormat(obj)); + } else if (type == "sequence") { + return Result::Convert(ParseSequenceFormat(obj)); + } else if (type == "or") { + return Result::Convert(ParseOrFormat(obj)); + } else if (type == "tag") { + return Result::Convert(ParseTagFormat(obj)); + } else if (type == "triggered_tags") { + return Result::Convert(ParseTriggeredTagsFormat(obj)); + } else if (type == "tags_with_separator") { + return Result::Convert(ParseTagsWithSeparatorFormat(obj)); + } else if (type == "optional") { + return Result::Convert(ParseOptionalFormat(obj)); + } else if (type == "plus") { + return Result::Convert(ParsePlusFormat(obj)); + } else if (type == "star") { + return Result::Convert(ParseStarFormat(obj)); + } else if (type == "repeat") { + return Result::Convert(ParseRepeatFormat(obj)); + } else if (type == "qwen_xml_parameter") { + return Result::Convert(ParseJSONSchemaFormat(obj, "qwen_xml")); + } else if (type == "grammar") { + return Result::Convert(ParseGrammarFormat(obj)); + } else if (type == "regex") { + return Result::Convert(ParseRegexFormat(obj)); + } else if (type == "token") { + return Result::Convert(ParseTokenFormat(obj)); + } else if (type == "exclude_token") { + return Result::Convert(ParseExcludeTokenFormat(obj)); + } else if (type == "any_tokens") { + return Result::Convert(ParseAnyTokensFormat(obj)); + } else if (type == "token_triggered_tags") { + return Result::Convert(ParseTokenTriggeredTagsFormat(obj)); + } else if (type == "dispatch") { + return Result::Convert(ParseDispatchFormat(obj)); + } else if (type == "token_dispatch") { + return Result::Convert(ParseTokenDispatchFormat(obj)); + } else { + return ResultErr("Format type not recognized: " + type); + } + } + + // If type is not present, try every format type one by one. Tag is prioritized. + auto tag_format = ParseTagFormat(obj); + if (!tag_format.IsErr()) { + return ResultOk(std::move(tag_format).Unwrap()); + } + auto const_string_format = ParseConstStringFormat(obj); + if (!const_string_format.IsErr()) { + return ResultOk(std::move(const_string_format).Unwrap()); + } + auto json_schema_format = ParseJSONSchemaFormat(obj); + if (!json_schema_format.IsErr()) { + return ResultOk(std::move(json_schema_format).Unwrap()); + } + auto any_text_format = ParseAnyTextFormat(obj); + if (!any_text_format.IsErr()) { + return ResultOk(std::move(any_text_format).Unwrap()); + } + auto sequence_format = ParseSequenceFormat(obj); + if (!sequence_format.IsErr()) { + return ResultOk(std::move(sequence_format).Unwrap()); + } + auto or_format = ParseOrFormat(obj); + if (!or_format.IsErr()) { + return ResultOk(std::move(or_format).Unwrap()); + } + auto triggered_tags_format = ParseTriggeredTagsFormat(obj); + if (!triggered_tags_format.IsErr()) { + return ResultOk(std::move(triggered_tags_format).Unwrap()); + } + auto tags_with_separator_format = ParseTagsWithSeparatorFormat(obj); + if (!tags_with_separator_format.IsErr()) { + return ResultOk(std::move(tags_with_separator_format).Unwrap()); + } + auto optional_format = ParseOptionalFormat(obj); + if (!optional_format.IsErr()) { + return ResultOk(std::move(optional_format).Unwrap()); + } + auto plus_format = ParsePlusFormat(obj); + if (!plus_format.IsErr()) { + return ResultOk(std::move(plus_format).Unwrap()); + } + auto star_format = ParseStarFormat(obj); + if (!star_format.IsErr()) { + return ResultOk(std::move(star_format).Unwrap()); + } + auto repeat_format = ParseRepeatFormat(obj); + if (!repeat_format.IsErr()) { + return ResultOk(std::move(repeat_format).Unwrap()); + } + auto tag_dispatch_format = ParseDispatchFormat(obj); + if (!tag_dispatch_format.IsErr()) { + return ResultOk(std::move(tag_dispatch_format).Unwrap()); + } + auto token_tag_dispatch_format = ParseTokenDispatchFormat(obj); + if (!token_tag_dispatch_format.IsErr()) { + return ResultOk(std::move(token_tag_dispatch_format).Unwrap()); + } + return ResultErr("Invalid format: " + value.serialize(false)); +} + +Result StructuralTagParser::ParseConstStringFormat( + const picojson::object& obj +) { + // value is required. + auto value_it = obj.find("value"); + if (value_it == obj.end() || !value_it->second.is()) { + return ResultErr("ConstString format must have a value field with a string"); + } + return ResultOk(value_it->second.get()); +} + +Result StructuralTagParser::ParseJSONSchemaFormat( + const picojson::object& obj, std::optional style_override +) { + // json_schema is required. + auto json_schema_it = obj.find("json_schema"); + if (json_schema_it == obj.end() || + !(json_schema_it->second.is() || json_schema_it->second.is())) { + return ResultErr( + "JSON schema format must have a json_schema field with a object or boolean value" + ); + } + std::string style = "json"; + if (style_override.has_value()) { + style = *style_override; + } else { + auto it = obj.find("style"); + if (it != obj.end() && it->second.is()) { + style = it->second.get(); + if (style != "json" && style != "qwen_xml" && style != "minimax_xml" && + style != "minimax_m3_xml" && style != "deepseek_xml" && style != "glm_xml" && + style != "cohere_xml" && style != "kimi_k3_xml" && style != "deepseek_v4_1_xml") { + return ResultErr( + "style must be \"json\", \"qwen_xml\", \"minimax_xml\", \"minimax_m3_xml\", " + "\"deepseek_xml\", \"glm_xml\", \"cohere_xml\", \"kimi_k3_xml\", or " + "\"deepseek_v4_1_xml\"" + ); + } + } + } + bool any_order = false; + auto any_order_it = obj.find("any_order"); + if (any_order_it != obj.end()) { + if (!any_order_it->second.is()) { + return ResultErr("any_order must be a boolean"); + } + any_order = any_order_it->second.get(); + } + std::optional max_whitespace_cnt = std::nullopt; + auto max_whitespace_cnt_it = obj.find("max_whitespace_cnt"); + if (max_whitespace_cnt_it != obj.end() && !max_whitespace_cnt_it->second.is()) { + if (!max_whitespace_cnt_it->second.is()) { + return ResultErr("max_whitespace_cnt must be an integer or null"); + } + max_whitespace_cnt = static_cast(max_whitespace_cnt_it->second.get()); + } + // here introduces a serialization/deserialization overhead; try to avoid it in the future. + return ResultOk( + json_schema_it->second.serialize(false), style, any_order, max_whitespace_cnt + ); +} + +Result ParseOptionalBudget( + const picojson::object& obj, const std::string& field_name, const std::string& format_name +) { + auto it = obj.find(field_name); + if (it == obj.end() || it->second.is()) { + return ResultOk(-1); + } + if (!it->second.is()) { + return ResultErr( + field_name + " in " + format_name + " must be a non-negative 32-bit integer or null" + ); + } + double value = it->second.get(); + if (value < 0 || value > std::numeric_limits::max() || value != std::floor(value)) { + return ResultErr( + field_name + " in " + format_name + " must be a non-negative 32-bit integer or null" + ); + } + return ResultOk(static_cast(value)); +} + +Result StructuralTagParser::ParseAnyTextFormat(const picojson::object& obj +) { + auto max_tokens_result = ParseOptionalBudget(obj, "max_tokens", "any_text"); + if (max_tokens_result.IsErr()) { + return ResultErr(std::move(max_tokens_result).UnwrapErr()); + } + auto max_chars_result = ParseOptionalBudget(obj, "max_chars", "any_text"); + if (max_chars_result.IsErr()) { + return ResultErr(std::move(max_chars_result).UnwrapErr()); + } + int32_t max_tokens = std::move(max_tokens_result).Unwrap(); + int32_t max_chars = std::move(max_chars_result).Unwrap(); + + auto excluded_strs_it = obj.find("excludes"); + if (excluded_strs_it == obj.end()) { + if ((obj.find("type") == obj.end())) { + return ResultErr("Any text format should not have any fields other than type"); + } + return ResultOk(std::vector{}, max_tokens, max_chars); + } + if (!excluded_strs_it->second.is()) { + return ResultErr("AnyText format's excluded_strs field must be an array"); + } + const auto& excluded_strs_array = excluded_strs_it->second.get(); + std::vector excluded_strs; + excluded_strs.reserve(excluded_strs_array.size()); + for (const auto& excluded_str : excluded_strs_array) { + if (!excluded_str.is()) { + return ResultErr("AnyText format's excluded_strs array must contain strings"); + } + excluded_strs.push_back(excluded_str.get()); + } + return ResultOk(std::move(excluded_strs), max_tokens, max_chars); +} + +Result StructuralTagParser::ParseGrammarFormat(const picojson::object& obj +) { + // grammar is required. + auto grammar_it = obj.find("grammar"); + if (grammar_it == obj.end() || !grammar_it->second.is() || + grammar_it->second.get().empty()) { + return ResultErr("Grammar format must have a grammar field with a non-empty string"); + } + return ResultOk(grammar_it->second.get()); +} + +Result StructuralTagParser::ParseRegexFormat(const picojson::object& obj) { + // pattern is required. + auto pattern_it = obj.find("pattern"); + if (pattern_it == obj.end() || !pattern_it->second.is() || + pattern_it->second.get().empty()) { + return ResultErr("Regex format must have a pattern field with a non-empty string"); + } + return ResultOk(pattern_it->second.get()); +} + +Result StructuralTagParser::ParseSequenceFormat( + const picojson::object& obj +) { + // elements is required. + auto elements_it = obj.find("elements"); + if (elements_it == obj.end() || !elements_it->second.is()) { + return ResultErr("Sequence format must have an elements field with an array"); + } + const auto& elements_array = elements_it->second.get(); + std::vector elements; + elements.reserve(elements_array.size()); + for (const auto& element : elements_array) { + auto format = ParseFormat(element); + if (format.IsErr()) { + return ResultErr(std::move(format).UnwrapErr()); + } + elements.push_back(std::move(format).Unwrap()); + } + if (elements.size() == 0) { + return ResultErr("Sequence format must have at least one element"); + } + return ResultOk(std::move(elements)); +} + +Result StructuralTagParser::ParseOrFormat(const picojson::object& obj) { + // elements is required. + auto elements_it = obj.find("elements"); + if (elements_it == obj.end() || !elements_it->second.is()) { + return ResultErr("Or format must have an elements field with an array"); + } + const auto& elements_array = elements_it->second.get(); + std::vector elements; + elements.reserve(elements_array.size()); + for (const auto& element : elements_array) { + auto format = ParseFormat(element); + if (format.IsErr()) { + return ResultErr(std::move(format).UnwrapErr()); + } + elements.push_back(std::move(format).Unwrap()); + } + if (elements.size() == 0) { + return ResultErr("Or format must have at least one element"); + } + return ResultOk(std::move(elements)); +} + +Result StructuralTagParser::ParseTagFormat(const picojson::value& value) { + if (!value.is()) { + return ResultErr("Tag format must be an object"); + } + const auto& obj = value.get(); + if (obj.find("type") != obj.end() && + (!obj["type"].is() || obj["type"].get() != "tag")) { + return ResultErr("Tag format's type must be a string \"tag\""); + } + return ParseTagFormat(obj); +} + +Result StructuralTagParser::ParseTagFormat(const picojson::object& obj) { + // begin is required: string or TokenFormat object + auto begin_it = obj.find("begin"); + if (begin_it == obj.end()) { + return ResultErr("Tag format's begin field must be a string"); + } + std::variant begin; + if (begin_it->second.is()) { + begin = begin_it->second.get(); + } else if (begin_it->second.is()) { + auto tf = ParseTokenFormat(begin_it->second.get()); + if (tf.IsErr()) { + return ResultErr(std::move(tf).UnwrapErr()); + } + begin = std::move(tf).Unwrap(); + } else { + return ResultErr("Tag format's begin field must be a string"); + } + + // content is required. + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Tag format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + + // end is required: string, array of strings, or TokenFormat object + auto end_it = obj.find("end"); + if (end_it == obj.end()) { + return ResultErr("Tag format must have an end field"); + } + + std::variant, TokenFormat> end; + if (end_it->second.is()) { + end = std::vector{end_it->second.get()}; + } else if (end_it->second.is()) { + const auto& end_array = end_it->second.get(); + if (end_array.empty()) { + return ResultErr("Tag format's end array cannot be empty"); + } + std::vector end_strings; + for (const auto& item : end_array) { + if (!item.is()) { + return ResultErr("Tag format's end array must contain only strings"); + } + end_strings.push_back(item.get()); + } + end = std::move(end_strings); + } else if (end_it->second.is()) { + auto tf = ParseTokenFormat(end_it->second.get()); + if (tf.IsErr()) { + return ResultErr(std::move(tf).UnwrapErr()); + } + end = std::move(tf).Unwrap(); + } else { + return ResultErr("Tag format's end field must be a string or array of strings"); + } + + return ResultOk( + std::move(begin), std::make_shared(std::move(content).Unwrap()), std::move(end) + ); +} + +Result StructuralTagParser::ParseTriggeredTagsFormat( + const picojson::object& obj +) { + // triggers is required. + auto triggers_it = obj.find("triggers"); + if (triggers_it == obj.end() || !triggers_it->second.is()) { + return ResultErr("Triggered tags format must have a triggers field with an array"); + } + const auto& triggers_array = triggers_it->second.get(); + std::vector excluded_strs; + std::vector triggers; + triggers.reserve(triggers_array.size()); + for (const auto& trigger : triggers_array) { + if (!trigger.is() || trigger.get().empty()) { + return ResultErr("Triggered tags format's triggers must be non-empty strings"); + } + triggers.push_back(trigger.get()); + } + if (triggers.size() == 0) { + return ResultErr("Triggered tags format's triggers must be non-empty"); + } + // tags is required. + auto tags_it = obj.find("tags"); + if (tags_it == obj.end() || !tags_it->second.is()) { + return ResultErr("Triggered tags format must have a tags field with an array"); + } + const auto& tags_array = tags_it->second.get(); + std::vector tags; + tags.reserve(tags_array.size()); + for (const auto& tag : tags_array) { + auto tag_format = ParseTagFormat(tag); + if (tag_format.IsErr()) { + return ResultErr(std::move(tag_format).UnwrapErr()); + } + tags.push_back(std::move(tag_format).Unwrap()); + } + if (tags.size() == 0) { + return ResultErr("Triggered tags format's tags must be non-empty"); + } + // excludes is optional. + auto excludes_it = obj.find("excludes"); + if (excludes_it != obj.end()) { + if (!excludes_it->second.is()) { + return ResultErr("Triggered tags format should have a excludes field with an array" + ); + } + const auto& excludes_array = excludes_it->second.get(); + excluded_strs.reserve(excludes_array.size()); + for (const auto& excluded_str : excludes_array) { + if (!excluded_str.is() || excluded_str.get().empty()) { + return ResultErr("Triggered tags format's excluded_strs must be non-empty strings" + ); + } + excluded_strs.push_back(excluded_str.get()); + } + } + + // at_least_one is optional. + bool at_least_one = false; + auto at_least_one_it = obj.find("at_least_one"); + if (at_least_one_it != obj.end()) { + if (!at_least_one_it->second.is()) { + return ResultErr("at_least_one must be a boolean"); + } + at_least_one = at_least_one_it->second.get(); + } + // stop_after_first is optional. + bool stop_after_first = false; + auto stop_after_first_it = obj.find("stop_after_first"); + if (stop_after_first_it != obj.end()) { + if (!stop_after_first_it->second.is()) { + return ResultErr("stop_after_first must be a boolean"); + } + stop_after_first = stop_after_first_it->second.get(); + } + return ResultOk( + std::move(triggers), std::move(tags), std::move(excluded_strs), at_least_one, stop_after_first + ); +} + +Result StructuralTagParser::ParseTagsWithSeparatorFormat( + const picojson::object& obj +) { + // tags is required. + auto tags_it = obj.find("tags"); + if (tags_it == obj.end() || !tags_it->second.is()) { + return ResultErr("Tags with separator format must have a tags field with an array"); + } + const auto& tags_array = tags_it->second.get(); + std::vector tags; + tags.reserve(tags_array.size()); + for (const auto& tag : tags_array) { + auto tag_format = ParseTagFormat(tag); + if (tag_format.IsErr()) { + return ResultErr(std::move(tag_format).UnwrapErr()); + } + tags.push_back(std::move(tag_format).Unwrap()); + } + if (tags.size() == 0) { + return ResultErr("Tags with separator format's tags must be non-empty"); + } + // separator is required (can be empty string). + auto separator_it = obj.find("separator"); + if (separator_it == obj.end() || !separator_it->second.is()) { + return ResultErr("Tags with separator format's separator field must be a string"); + } + // at_least_one is optional. + bool at_least_one = false; + auto at_least_one_it = obj.find("at_least_one"); + if (at_least_one_it != obj.end()) { + if (!at_least_one_it->second.is()) { + return ResultErr("at_least_one must be a boolean"); + } + at_least_one = at_least_one_it->second.get(); + } + // stop_after_first is optional. + bool stop_after_first = false; + auto stop_after_first_it = obj.find("stop_after_first"); + if (stop_after_first_it != obj.end()) { + if (!stop_after_first_it->second.is()) { + return ResultErr("stop_after_first must be a boolean"); + } + stop_after_first = stop_after_first_it->second.get(); + } + return ResultOk( + std::move(tags), separator_it->second.get(), at_least_one, stop_after_first + ); +} + +Result StructuralTagParser::ParseOptionalFormat( + const picojson::object& obj +) { + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Optional format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk(std::make_shared(std::move(content).Unwrap())); +} + +Result StructuralTagParser::ParsePlusFormat(const picojson::object& obj) { + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Plus format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk(std::make_shared(std::move(content).Unwrap())); +} + +Result StructuralTagParser::ParseStarFormat(const picojson::object& obj) { + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Star format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk(std::make_shared(std::move(content).Unwrap())); +} + +Result StructuralTagParser::ParseTokenFormat(const picojson::object& obj) { + auto token_it = obj.find("token"); + if (token_it == obj.end()) { + return ResultErr("TokenFormat must have a token field"); + } + if (token_it->second.is()) { + double d = token_it->second.get(); + if (d != static_cast(static_cast(d))) { + return ResultErr("Token ID must be an integer"); + } + int32_t id = static_cast(d); + if (id < 0) { + return ResultErr("Token ID must be non-negative"); + } + return ResultOk(std::variant(id)); + } else if (token_it->second.is()) { + auto s = token_it->second.get(); + if (s.empty()) { + return ResultErr("Token string must be non-empty"); + } + return ResultOk(std::variant(std::move(s))); + } + return ResultErr("TokenFormat's token must be an integer or string"); +} + +Result>, ISTError> ParseIntOrStringArray( + const picojson::value& val, const std::string& field_name +) { + std::vector> result; + if (!val.is()) { + return ResultErr(field_name + " must be an array"); + } + for (const auto& v : val.get()) { + if (v.is()) { + double d = v.get(); + if (d != static_cast(static_cast(d))) { + return ResultErr(field_name + " elements must be integers, not floats"); + } + int32_t id = static_cast(d); + if (id < 0) { + return ResultErr( + field_name + " elements must be non-negative integers or strings" + ); + } + result.emplace_back(std::in_place_type, id); + } else if (v.is()) { + auto s = v.get(); + if (s.empty()) { + return ResultErr(field_name + " string elements must be non-empty"); + } + result.push_back(std::move(s)); + } else { + return ResultErr(field_name + " elements must be integers or strings"); + } + } + return ResultOk(std::move(result)); +} + +Result StructuralTagParser::ParseExcludeTokenFormat( + const picojson::object& obj +) { + std::vector> exclude_tokens; + auto it = obj.find("exclude_tokens"); + if (it != obj.end()) { + auto parsed = ParseIntOrStringArray(it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + return ResultOk(std::move(exclude_tokens)); +} + +Result StructuralTagParser::ParseAnyTokensFormat( + const picojson::object& obj +) { + std::vector> exclude_tokens; + auto it = obj.find("exclude_tokens"); + if (it != obj.end()) { + auto parsed = ParseIntOrStringArray(it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + auto max_tokens_result = ParseOptionalBudget(obj, "max_tokens", "any_tokens"); + if (max_tokens_result.IsErr()) { + return ResultErr(std::move(max_tokens_result).UnwrapErr()); + } + return ResultOk( + std::move(exclude_tokens), std::move(max_tokens_result).Unwrap() + ); +} + +Result StructuralTagParser::ParseTokenTriggeredTagsFormat( + const picojson::object& obj +) { + // trigger_tokens is required + auto triggers_it = obj.find("trigger_tokens"); + if (triggers_it == obj.end()) { + return ResultErr("TokenTriggeredTagsFormat must have a trigger_tokens field"); + } + auto triggers = ParseIntOrStringArray(triggers_it->second, "trigger_tokens"); + if (triggers.IsErr()) { + return ResultErr(std::move(triggers).UnwrapErr()); + } + auto trigger_tokens = std::move(triggers).Unwrap(); + if (trigger_tokens.empty()) { + return ResultErr("trigger_tokens must be non-empty"); + } + + // tags is required + auto tags_it = obj.find("tags"); + if (tags_it == obj.end() || !tags_it->second.is()) { + return ResultErr("TokenTriggeredTagsFormat must have a tags field with an array"); + } + std::vector tags; + for (const auto& tag : tags_it->second.get()) { + auto tag_format = ParseTagFormat(tag); + if (tag_format.IsErr()) { + return ResultErr(std::move(tag_format).UnwrapErr()); + } + tags.push_back(std::move(tag_format).Unwrap()); + } + if (tags.empty()) { + return ResultErr("TokenTriggeredTagsFormat tags must be non-empty"); + } + + // exclude_tokens is optional + std::vector> exclude_tokens; + auto excludes_it = obj.find("exclude_tokens"); + if (excludes_it != obj.end()) { + auto parsed = ParseIntOrStringArray(excludes_it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + + bool at_least_one = false; + auto alo_it = obj.find("at_least_one"); + if (alo_it != obj.end()) { + if (!alo_it->second.is()) { + return ResultErr("at_least_one must be a boolean"); + } + at_least_one = alo_it->second.get(); + } + + bool stop_after_first = false; + auto saf_it = obj.find("stop_after_first"); + if (saf_it != obj.end()) { + if (!saf_it->second.is()) { + return ResultErr("stop_after_first must be a boolean"); + } + stop_after_first = saf_it->second.get(); + } + + return ResultOk( + std::move(trigger_tokens), + std::move(tags), + std::move(exclude_tokens), + at_least_one, + stop_after_first + ); +} + +Result StructuralTagParser::ParseDispatchFormat( + const picojson::object& obj +) { + auto rules_it = obj.find("rules"); + if (rules_it == obj.end() || !rules_it->second.is()) { + return ResultErr("TagDispatch format must have a rules field with an array"); + } + const auto& rules_array = rules_it->second.get(); + if (rules_array.empty()) { + return ResultErr("TagDispatch format rules must be non-empty"); + } + std::vector>> rules; + rules.reserve(rules_array.size()); + for (const auto& item : rules_array) { + if (!item.is()) { + return ResultErr("TagDispatch pair must be a 2-element array"); + } + const auto& pair_arr = item.get(); + if (pair_arr.size() != 2) { + return ResultErr("TagDispatch pair must be a 2-element array"); + } + if (!pair_arr[0].is()) { + return ResultErr("TagDispatch pair first element must be a string"); + } + std::string trigger = pair_arr[0].get(); + auto content = ParseFormat(pair_arr[1]); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + rules.push_back({std::move(trigger), std::make_shared(std::move(content).Unwrap())}); + } + + bool loop = true; + auto loop_it = obj.find("loop"); + if (loop_it != obj.end()) { + if (!loop_it->second.is()) { + return ResultErr("loop must be a boolean"); + } + loop = loop_it->second.get(); + } + + std::vector excludes; + auto excludes_it = obj.find("excludes"); + if (excludes_it != obj.end()) { + if (!excludes_it->second.is()) { + return ResultErr("excludes must be an array"); + } + for (const auto& e : excludes_it->second.get()) { + if (!e.is() || e.get().empty()) { + return ResultErr("excludes must contain non-empty strings"); + } + excludes.push_back(e.get()); + } + } + + return ResultOk(std::move(rules), loop, std::move(excludes)); +} + +Result StructuralTagParser::ParseTokenDispatchFormat( + const picojson::object& obj +) { + auto rules_it = obj.find("rules"); + if (rules_it == obj.end() || !rules_it->second.is()) { + return ResultErr("TokenTagDispatch format must have a rules field with an array"); + } + const auto& rules_array = rules_it->second.get(); + if (rules_array.empty()) { + return ResultErr("TokenTagDispatch format rules must be non-empty"); + } + std::vector, std::shared_ptr>> rules; + rules.reserve(rules_array.size()); + for (const auto& item : rules_array) { + if (!item.is()) { + return ResultErr("TokenTagDispatch pair must be a 2-element array"); + } + const auto& pair_arr = item.get(); + if (pair_arr.size() != 2) { + return ResultErr("TokenTagDispatch pair must be a 2-element array"); + } + std::variant trigger; + if (pair_arr[0].is()) { + double d = pair_arr[0].get(); + if (d != static_cast(static_cast(d))) { + return ResultErr("Token ID must be an integer"); + } + trigger = static_cast(d); + } else if (pair_arr[0].is()) { + trigger = pair_arr[0].get(); + } else { + return ResultErr("TokenTagDispatch pair first element must be an integer or string" + ); + } + auto content = ParseFormat(pair_arr[1]); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + rules.push_back({std::move(trigger), std::make_shared(std::move(content).Unwrap())}); + } + + bool loop = true; + auto loop_it = obj.find("loop"); + if (loop_it != obj.end()) { + if (!loop_it->second.is()) { + return ResultErr("loop must be a boolean"); + } + loop = loop_it->second.get(); + } + + std::vector> exclude_tokens; + auto excludes_it = obj.find("exclude_tokens"); + if (excludes_it != obj.end()) { + auto parsed = ParseIntOrStringArray(excludes_it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + + return ResultOk(std::move(rules), loop, std::move(exclude_tokens)); +} + +/************** StructuralTagTokenResolver **************/ + +class StructuralTagTokenResolver { + public: + static std::optional Resolve( + StructuralTag* structural_tag, const std::optional& tokenizer_info + ); + + private: + explicit StructuralTagTokenResolver(const std::optional& tokenizer_info) + : tokenizer_info_(tokenizer_info) {} + + std::optional ResolveFormat(Format* format); + std::optional ResolveTagFormat(TagFormat* tag); + std::optional ResolveTokenFormat(TokenFormat* tf); + std::optional ResolveIntOrStringVec( + const std::vector>& input, std::vector* output + ); + + const std::optional& tokenizer_info_; +}; + +std::optional StructuralTagTokenResolver::Resolve( + StructuralTag* structural_tag, const std::optional& tokenizer_info +) { + return StructuralTagTokenResolver(tokenizer_info).ResolveFormat(&structural_tag->format); +} + +std::optional StructuralTagTokenResolver::ResolveTokenFormat(TokenFormat* tf) { + if (tf->resolved_token_id_ >= 0) return std::nullopt; + if (!std::holds_alternative(tf->token)) return std::nullopt; + if (!tokenizer_info_) { + return ISTError("Token string resolution requires tokenizer_info"); + } + const auto& token_str = std::get(tf->token); + const auto& vocab = tokenizer_info_->GetDecodedVocab(); + for (int32_t i = 0; i < static_cast(vocab.size()); ++i) { + if (vocab[i] == token_str) { + tf->resolved_token_id_ = i; + return std::nullopt; + } + } + return ISTError("Token string \"" + token_str + "\" not found in vocabulary"); +} + +std::optional StructuralTagTokenResolver::ResolveIntOrStringVec( + const std::vector>& input, std::vector* output +) { + output->clear(); + output->reserve(input.size()); + for (const auto& item : input) { + if (std::holds_alternative(item)) { + output->push_back(std::get(item)); + } else { + if (!tokenizer_info_) { + return ISTError("Token string resolution requires tokenizer_info"); + } + const auto& s = std::get(item); + const auto& vocab = tokenizer_info_->GetDecodedVocab(); + bool found = false; + for (int32_t i = 0; i < static_cast(vocab.size()); ++i) { + if (vocab[i] == s) { + output->push_back(i); + found = true; + break; + } + } + if (!found) { + return ISTError("Token string \"" + s + "\" not found in vocabulary"); + } + } + } + return std::nullopt; +} + +std::optional StructuralTagTokenResolver::ResolveTagFormat(TagFormat* tag) { + if (std::holds_alternative(tag->begin)) { + auto err = ResolveTokenFormat(&std::get(tag->begin)); + if (err) return err; + } + if (std::holds_alternative(tag->end)) { + auto err = ResolveTokenFormat(&std::get(tag->end)); + if (err) return err; + } + return ResolveFormat(tag->content.get()); +} + +std::optional StructuralTagTokenResolver::ResolveFormat(Format* format) { + return std::visit( + [&](auto&& arg) -> std::optional { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return ResolveTokenFormat(&arg); + } else if constexpr (std::is_same_v) { + return ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_token_ids_); + } else if constexpr (std::is_same_v) { + return ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_exclude_token_ids_); + } else if constexpr (std::is_same_v) { + auto err = ResolveIntOrStringVec(arg.trigger_tokens, &arg.resolved_trigger_token_ids_); + if (err) return err; + err = ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_exclude_token_ids_); + if (err) return err; + for (auto& tag : arg.tags) { + err = ResolveTagFormat(&tag); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + std::vector> trigger_tokens; + trigger_tokens.reserve(arg.rules.size()); + for (const auto& p : arg.rules) { + trigger_tokens.push_back(p.first); + } + auto err = ResolveIntOrStringVec(trigger_tokens, &arg.resolved_trigger_token_ids_); + if (err) return err; + err = ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_exclude_token_ids_); + if (err) return err; + for (auto& p : arg.rules) { + if (p.second) { + auto e = ResolveFormat(p.second.get()); + if (e) return e; + } + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& p : arg.rules) { + if (p.second) { + auto err = ResolveFormat(p.second.get()); + if (err) return err; + } + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + return ResolveTagFormat(&arg); + } else if constexpr (std::is_same_v) { + for (auto& elem : arg.elements) { + auto err = ResolveFormat(&elem); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& elem : arg.elements) { + auto err = ResolveFormat(&elem); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& tag : arg.tags) { + auto err = ResolveTagFormat(&tag); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& tag : arg.tags) { + auto err = ResolveTagFormat(&tag); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v || std::is_same_v || + std::is_same_v) { + return ResolveFormat(arg.content.get()); + } else { + return std::nullopt; + } + }, + *format + ); +} + +Result StructuralTagParser::ParseRepeatFormat(const picojson::object& obj) { + auto min_it = obj.find("min"); + if (min_it == obj.end() || !min_it->second.is()) { + return ResultErr("Repeat format must have a min field (number)"); + } + auto max_it = obj.find("max"); + if (max_it == obj.end() || !max_it->second.is()) { + return ResultErr("Repeat format must have a max field (number)"); + } + int64_t min = min_it->second.get(); + int64_t max = max_it->second.get(); + int32_t max_value_int32 = std::numeric_limits::max(); + if (max >= 0 && min > max) { + return ResultErr("Repeat min must be <= max"); + } + if (min < 0) { + return ResultErr("Repeat min must be >= 0"); + } + if (max < -1) { + return ResultErr("Repeat max must be -1 (unbounded) or >= 0"); + } + if (max > static_cast(max_value_int32)) { + XGRAMMAR_LOG(WARNING) << "Repeat max is too large, will be set as not limited"; + max = -1; // -1 means unlimited + } + if (min > static_cast(max_value_int32)) { + return ResultErr( + "Repeat min is too large, must be <= " + std::to_string(max_value_int32) + ); + } + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Repeat format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk( + static_cast(min), + static_cast(max), + std::make_shared(std::move(content).Unwrap()) + ); +} + +/************** StructuralTag Analyzer **************/ + +/*! + * \brief Analyze a StructuralTag and extract useful information for conversion to Grammar. + */ +class StructuralTagAnalyzer { + public: + static std::optional Analyze(StructuralTag* structural_tag); + + private: + /*! \brief A variant that can hold the pointer of any Format types. */ + using FormatPtrVariant = std::variant< + ConstStringFormat*, + JSONSchemaFormat*, + AnyTextFormat*, + GrammarFormat*, + RegexFormat*, + SequenceFormat*, + OrFormat*, + TagFormat*, + TriggeredTagsFormat*, + TagsWithSeparatorFormat*, + OptionalFormat*, + PlusFormat*, + StarFormat*, + RepeatFormat*, + TokenFormat*, + ExcludeTokenFormat*, + AnyTokensFormat*, + TokenTriggeredTagsFormat*, + DispatchFormat*, + TokenDispatchFormat*>; + + // Call this if we have a pointer to a Format. + std::optional Visit(Format* format); + // Call this if we have a pointer to a variant of Format. + std::optional Visit(FormatPtrVariant format); + + // The following is dispatched from Visit. Don't call them directly because they don't handle + // stack logics. + std::optional VisitSub(ConstStringFormat* format); + std::optional VisitSub(JSONSchemaFormat* format); + std::optional VisitSub(AnyTextFormat* format); + std::optional VisitSub(GrammarFormat* format); + std::optional VisitSub(RegexFormat* format); + std::optional VisitSub(SequenceFormat* format); + std::optional VisitSub(OrFormat* format); + std::optional VisitSub(TagFormat* format); + std::optional VisitSub(TriggeredTagsFormat* format); + std::optional VisitSub(TagsWithSeparatorFormat* format); + std::optional VisitSub(OptionalFormat* format); + std::optional VisitSub(PlusFormat* format); + std::optional VisitSub(StarFormat* format); + std::optional VisitSub(TokenFormat* format); + std::optional VisitSub(ExcludeTokenFormat* format); + std::optional VisitSub(AnyTokensFormat* format); + std::optional VisitSub(TokenTriggeredTagsFormat* format); + std::optional VisitSub(RepeatFormat* format); + std::optional VisitSub(DispatchFormat* format); + std::optional VisitSub(TokenDispatchFormat* format); + + std::vector DetectEndStrings(); + std::vector DetectEndTokenIds(); + bool IsUnlimited(const Format& format); + bool IsExcluded(const Format& format); + + int visit_format_recursion_depth_ = 0; + std::vector stack_; +}; + +std::optional StructuralTagAnalyzer::Analyze(StructuralTag* structural_tag) { + return StructuralTagAnalyzer().Visit(&structural_tag->format); +} + +std::vector StructuralTagAnalyzer::DetectEndStrings() { + for (int i = static_cast(stack_.size()) - 1; i >= 0; --i) { + auto& format = stack_[i]; + if (std::holds_alternative(format)) { + auto* tag = std::get(format); + if (std::holds_alternative>(tag->end)) { + return std::get>(tag->end); + } + return {}; // TokenFormat end — propagated via DetectEndTokenIds + } + } + return {}; +} + +std::vector StructuralTagAnalyzer::DetectEndTokenIds() { + for (int i = static_cast(stack_.size()) - 1; i >= 0; --i) { + auto& format = stack_[i]; + if (std::holds_alternative(format)) { + auto* tag = std::get(format); + if (std::holds_alternative(tag->end)) { + auto& tf = std::get(tag->end); + return {tf.resolved_token_id_}; + } + return {}; + } + } + return {}; +} + +bool StructuralTagAnalyzer::IsUnlimited(const Format& format) { + return std::visit( + [&](auto&& arg) -> bool { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return arg.is_unlimited_; + } else if constexpr (std::is_same_v) { + return arg.is_unlimited_; + } else if constexpr (std::is_same_v) { + return IsUnlimited(*arg.content); + } else if constexpr (std::is_same_v || std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return arg.max == -1 || (arg.max != 0 && IsUnlimited(*arg.content)); + } else { + return false; + } + }, + format + ); +} + +bool StructuralTagAnalyzer::IsExcluded(const Format& format) { + return std::visit( + [&](auto&& arg) -> bool { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return !arg.excludes.empty(); + } else if constexpr (std::is_same_v) { + return !arg.excludes.empty(); + } else if constexpr (std::is_same_v) { + return !arg.exclude_tokens.empty(); + } else if constexpr (std::is_same_v) { + return !arg.excludes.empty(); + } else if constexpr (std::is_same_v) { + return !arg.exclude_tokens.empty(); + } else if constexpr (std::is_same_v) { + return !arg.exclude_tokens.empty(); + } else { + return false; + } + }, + format + ); +} + +std::optional StructuralTagAnalyzer::Visit(Format* format) { + FormatPtrVariant format_ptr_variant = + std::visit([&](auto&& arg) -> FormatPtrVariant { return &arg; }, *format); + return Visit(format_ptr_variant); +} + +std::optional StructuralTagAnalyzer::Visit(FormatPtrVariant format) { + RecursionGuard guard(&visit_format_recursion_depth_); + + // Push format to stack + stack_.push_back(format); + + // Dispatch to the corresponding visit function + auto result = + std::visit([&](auto&& arg) -> std::optional { return VisitSub(arg); }, format); + + // Pop format from stack + stack_.pop_back(); + + return result; +} + +std::optional StructuralTagAnalyzer::VisitSub(ConstStringFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(JSONSchemaFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(AnyTextFormat* format) { + format->detected_end_strs_ = DetectEndStrings(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(GrammarFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(RegexFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(SequenceFormat* format) { + bool is_any_unlimited = false; + for (auto& element : format->elements) { + auto err = Visit(&element); + if (err.has_value()) { + return err; + } + is_any_unlimited |= IsUnlimited(element) && !IsExcluded(element); + } + format->is_unlimited_ = is_any_unlimited; + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(OrFormat* format) { + bool is_any_unlimited = false; + for (auto& element : format->elements) { + auto err = Visit(&element); + if (err.has_value()) { + return err; + } + is_any_unlimited |= IsUnlimited(element) && !IsExcluded(element); + } + format->is_unlimited_ = is_any_unlimited; + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TagFormat* format) { + auto err = Visit(format->content.get()); + if (err.has_value()) { + return err; + } + auto is_content_unlimited = IsUnlimited(*(format->content)); + if (is_content_unlimited) { + if (std::holds_alternative>(format->end)) { + const auto& ends = std::get>(format->end); + bool has_non_empty_end = false; + for (const auto& end_str : ends) { + if (!end_str.empty()) { + has_non_empty_end = true; + break; + } + } + if (!has_non_empty_end && !IsExcluded(*format->content)) { + return ISTError("When the content is unlimited, at least one end string must be non-empty"); + } + } + // TokenFormat end is always non-empty → no error needed + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TriggeredTagsFormat* format) { + for (auto& tag : format->tags) { + auto err = Visit(&tag); + if (err.has_value()) { + return err; + } + } + format->detected_end_strs_ = DetectEndStrings(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TagsWithSeparatorFormat* format) { + for (auto& tag : format->tags) { + auto err = Visit(&tag); + if (err.has_value()) { + return err; + } + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(OptionalFormat* format) { + return Visit(format->content.get()); +} + +std::optional StructuralTagAnalyzer::VisitSub(PlusFormat* format) { + return Visit(format->content.get()); +} + +std::optional StructuralTagAnalyzer::VisitSub(StarFormat* format) { + return Visit(format->content.get()); +} + +std::optional StructuralTagAnalyzer::VisitSub(TokenFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(ExcludeTokenFormat* format) { + format->detected_end_token_ids_ = DetectEndTokenIds(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(AnyTokensFormat* format) { + format->detected_end_token_ids_ = DetectEndTokenIds(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TokenTriggeredTagsFormat* format) { + for (auto& tag : format->tags) { + auto err = Visit(&tag); + if (err.has_value()) { + return err; + } + } + format->detected_end_token_ids_ = DetectEndTokenIds(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(DispatchFormat* format) { + for (auto& pair : format->rules) { + if (pair.second) { + auto err = Visit(pair.second.get()); + if (err.has_value()) return err; + } + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TokenDispatchFormat* format) { + for (auto& pair : format->rules) { + if (pair.second) { + auto err = Visit(pair.second.get()); + if (err.has_value()) return err; + } + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(RepeatFormat* format) { + return Visit(format->content.get()); +} + +/************** StructuralTag to Grammar Converter **************/ + +class StructuralTagGrammarConverter { + public: + static Result Convert(const StructuralTag& structural_tag); + + private: + StructuralTagGrammarConverter() = default; + + /*! + * \brief Visit a Format and return the rule id of the added rule. + * \param format The Format to visit. + * \return The rule id of the added rule. If the visit fails, the error is returned. + * \note This method uses serialization to deduplicate identical formats. + */ + Result Visit(const Format& format); + Result VisitSub(const ConstStringFormat& format); + Result VisitSub(const JSONSchemaFormat& format); + Result VisitSub(const AnyTextFormat& format); + Result VisitSub(const GrammarFormat& format); + Result VisitSub(const RegexFormat& format); + Result VisitSub(const SequenceFormat& format); + Result VisitSub(const OrFormat& format); + Result VisitSub(const TagFormat& format); + Result VisitSub(const TriggeredTagsFormat& format); + Result VisitSub(const TagsWithSeparatorFormat& format); + Result VisitSub(const OptionalFormat& format); + Result VisitSub(const PlusFormat& format); + Result VisitSub(const StarFormat& format); + Result VisitSub(const TokenFormat& format); + Result VisitSub(const ExcludeTokenFormat& format); + Result VisitSub(const AnyTokensFormat& format); + Result VisitSub(const TokenTriggeredTagsFormat& format); + Result VisitSub(const RepeatFormat& format); + Result VisitSub(const DispatchFormat& format); + Result VisitSub(const TokenDispatchFormat& format); + Grammar AddRootRuleAndGetGrammar(int ref_rule_id); + void SetRuleBodyAndBudgets( + int32_t rule_id, int body_expr_id, int32_t max_tokens, int32_t max_chars = -1 + ); + + bool IsPrefix(const std::string& prefix, const std::string& full_str); + int BuildBeginExpr(const TagFormat& tag); + int BuildEndExpr(const TagFormat& tag); + + GrammarBuilder grammar_builder_; + + /*! + * \brief Cache from format serialization to rule id. + * This enables deduplication of identical formats to reduce grammar size. + */ + std::unordered_map serialization_to_rule_id_; +}; + +bool StructuralTagGrammarConverter::IsPrefix( + const std::string& prefix, const std::string& full_str +) { + return prefix.size() <= full_str.size() && + std::string_view(full_str).substr(0, prefix.size()) == prefix; +} + +Result StructuralTagGrammarConverter::Convert(const StructuralTag& structural_tag +) { + StructuralTagGrammarConverter converter; + auto result = converter.Visit(structural_tag.format); + if (result.IsErr()) { + return ResultErr(std::move(result).UnwrapErr()); + } + // Add a root rule + auto root_rule_id = std::move(result).Unwrap(); + return ResultOk(converter.AddRootRuleAndGetGrammar(root_rule_id)); +} + +Grammar StructuralTagGrammarConverter::AddRootRuleAndGetGrammar(int ref_rule_id) { + auto expr = grammar_builder_.AddRuleRef(ref_rule_id); + auto sequence_expr = grammar_builder_.AddSequence({expr}); + auto choices_expr = grammar_builder_.AddChoices({sequence_expr}); + auto root_rule_id = grammar_builder_.AddRuleWithHint("root", choices_expr); + return grammar_builder_.Get(root_rule_id); +} + +void StructuralTagGrammarConverter::SetRuleBodyAndBudgets( + int32_t rule_id, int body_expr_id, int32_t max_tokens, int32_t max_chars +) { + // A zero-token region is empty by definition. Materialize that in the grammar because -1 is + // also the parser's pre-first-token deadline sentinel. + if (max_tokens == 0) { + body_expr_id = grammar_builder_.AddChoices({grammar_builder_.AddEmptyStr()}); + } + grammar_builder_.UpdateRuleBody(rule_id, body_expr_id); + if (max_tokens >= 0) { + grammar_builder_.UpdateMaxTokens(rule_id, max_tokens); + } + if (max_chars >= 0) { + grammar_builder_.UpdateMaxChars(rule_id, max_chars); + } +} + +Result StructuralTagGrammarConverter::Visit(const Format& format) { + std::string fingerprint = FormatToJSONValue(format).serialize(); + + // Check if we've already processed an identical format + auto it = serialization_to_rule_id_.find(fingerprint); + if (it != serialization_to_rule_id_.end()) { + return ResultOk(it->second); + } + + // Process the format and cache the result + auto result = + std::visit([&](auto&& arg) -> Result { return VisitSub(arg); }, format); + if (result.IsOk()) { + int rule_id = std::move(result).Unwrap(); + serialization_to_rule_id_[fingerprint] = rule_id; + return ResultOk(rule_id); + } + return result; +} + +Result StructuralTagGrammarConverter::VisitSub(const ConstStringFormat& format) { + auto expr = format.value.empty() ? grammar_builder_.AddEmptyStr() + : grammar_builder_.AddByteString(format.value); + auto sequence_expr = grammar_builder_.AddSequence({expr}); + auto choices_expr = grammar_builder_.AddChoices({sequence_expr}); + return ResultOk(grammar_builder_.AddRuleWithHint("const_string", choices_expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const JSONSchemaFormat& format) { + auto json_format = JSONFormatFromString(format.style); + if (!json_format.has_value()) { + return ResultErr("Unsupported parsing type: " + format.style); + } + // The whitespace cap comes from the JSONSchemaFormat node (per-tag). + auto sub_grammar = GrammarNormalizer::Apply(JSONSchemaToGrammar( + format.json_schema, + /*any_whitespace=*/true, + /*indent=*/std::nullopt, + /*separators=*/std::nullopt, + /*strict_mode=*/true, + /*max_whitespace_cnt=*/format.max_whitespace_cnt, + /*any_order=*/format.any_order, + /*json_format=*/*json_format + )); + auto added_root_rule_id = SubGrammarAdder().Apply(&grammar_builder_, sub_grammar); + return ResultOk(added_root_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const GrammarFormat& format) { + auto sub_grammar = Grammar::FromEBNF(format.grammar); + auto added_root_rule_id = SubGrammarAdder().Apply(&grammar_builder_, sub_grammar); + return ResultOk(added_root_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const RegexFormat& format) { + auto sub_grammar = Grammar::FromRegex(format.pattern); + auto added_root_rule_id = SubGrammarAdder().Apply(&grammar_builder_, sub_grammar); + return ResultOk(added_root_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const AnyTextFormat& format) { + std::vector all_excludes = format.excludes; + for (const auto& s : format.detected_end_strs_) { + if (!s.empty()) { + all_excludes.push_back(s); + } + } + int body_expr_id; + if (!all_excludes.empty()) { + body_expr_id = + grammar_builder_.AddTagDispatch(Grammar::Impl::TagDispatch{{}, false, all_excludes}); + } else { + auto any_text_expr = grammar_builder_.AddCharacterClassStar({{0, 0x10FFFF}}, false); + auto sequence_expr = grammar_builder_.AddSequence({any_text_expr}); + body_expr_id = grammar_builder_.AddChoices({sequence_expr}); + } + int rule_id = grammar_builder_.AddEmptyRuleWithHint("any_text"); + SetRuleBodyAndBudgets(rule_id, body_expr_id, format.max_tokens, format.max_chars); + return ResultOk(rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const SequenceFormat& format) { + std::vector rule_ref_ids; + rule_ref_ids.reserve(format.elements.size()); + for (const auto& element : format.elements) { + auto result = Visit(element); + if (result.IsErr()) { + return result; + } + int sub_rule_id = std::move(result).Unwrap(); + rule_ref_ids.push_back(grammar_builder_.AddRuleRef(sub_rule_id)); + } + auto expr = grammar_builder_.AddChoices({grammar_builder_.AddSequence(rule_ref_ids)}); + return ResultOk(grammar_builder_.AddRuleWithHint("sequence", expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const OrFormat& format) { + std::vector sequence_ids; + sequence_ids.reserve(format.elements.size()); + for (const auto& element : format.elements) { + auto result = Visit(element); + if (result.IsErr()) { + return result; + } + int sub_rule_id = std::move(result).Unwrap(); + auto rule_ref_expr = grammar_builder_.AddRuleRef(sub_rule_id); + sequence_ids.push_back(grammar_builder_.AddSequence({rule_ref_expr})); + } + auto expr = grammar_builder_.AddChoices(sequence_ids); + return ResultOk(grammar_builder_.AddRuleWithHint("or", expr)); +} + +int StructuralTagGrammarConverter::BuildBeginExpr(const TagFormat& tag) { + if (std::holds_alternative(tag.begin)) { + return grammar_builder_.AddByteString(std::get(tag.begin)); + } + return grammar_builder_.AddTokenSet({std::get(tag.begin).resolved_token_id_}); +} + +int StructuralTagGrammarConverter::BuildEndExpr(const TagFormat& tag) { + if (std::holds_alternative(tag.end)) { + return grammar_builder_.AddTokenSet({std::get(tag.end).resolved_token_id_}); + } + const auto& ends = std::get>(tag.end); + if (ends.size() == 1) { + return ends[0].empty() ? grammar_builder_.AddEmptyStr() + : grammar_builder_.AddByteString(ends[0]); + } + std::vector end_seq_ids; + for (const auto& s : ends) { + auto e = s.empty() ? grammar_builder_.AddEmptyStr() : grammar_builder_.AddByteString(s); + end_seq_ids.push_back(grammar_builder_.AddSequence({e})); + } + auto choice = grammar_builder_.AddChoices(end_seq_ids); + auto rule = grammar_builder_.AddRuleWithHint("tag_end", choice); + return grammar_builder_.AddRuleRef(rule); +} + +Result StructuralTagGrammarConverter::VisitSub(const TagFormat& format) { + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + auto sub_rule_id = std::move(result).Unwrap(); + auto begin_expr = BuildBeginExpr(format); + auto rule_ref_expr = grammar_builder_.AddRuleRef(sub_rule_id); + auto end_expr = BuildEndExpr(format); + + auto sequence_expr_id = grammar_builder_.AddSequence({begin_expr, rule_ref_expr, end_expr}); + auto choices_expr = grammar_builder_.AddChoices({sequence_expr_id}); + return ResultOk(grammar_builder_.AddRuleWithHint("tag", choices_expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const TriggeredTagsFormat& format) { + // Step 1. Visit all tags and add to grammar + std::vector> trigger_to_tag_ids(format.triggers.size()); + std::vector tag_content_rule_ids; + tag_content_rule_ids.reserve(format.tags.size()); + + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + if (!std::holds_alternative(tag.begin)) { + return ResultErr( + "Tags in triggered_tags must have a string begin, not a token format" + ); + } + const auto& tag_begin = std::get(tag.begin); + int matched_trigger_id = -1; + for (int it_trigger = 0; it_trigger < static_cast(format.triggers.size()); ++it_trigger) { + const auto& trigger = format.triggers[it_trigger]; + if (IsPrefix(trigger, tag_begin)) { + if (matched_trigger_id != -1) { + return ResultErr("One tag matches multiple triggers in a triggered tags format" + ); + } + matched_trigger_id = it_trigger; + } + } + if (matched_trigger_id == -1) { + return ResultErr("One tag does not match any trigger in a triggered tags format"); + } + trigger_to_tag_ids[matched_trigger_id].push_back(it_tag); + + auto result = Visit(*tag.content); + if (result.IsErr()) { + return result; + } + tag_content_rule_ids.push_back(std::move(result).Unwrap()); + } + + // Step 2. Special Case: at_least_one && stop_after_first. + if (format.at_least_one && format.stop_after_first) { + std::vector choice_elements; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr_id = BuildBeginExpr(tag); + auto rule_ref_expr_id = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr_id = BuildEndExpr(tag); + choice_elements.push_back( + grammar_builder_.AddSequence({begin_expr_id, rule_ref_expr_id, end_expr_id}) + ); + } + auto choice_expr_id = grammar_builder_.AddChoices(choice_elements); + return ResultOk(grammar_builder_.AddRuleWithHint("triggered_tags", choice_expr_id)); + } + + // Step 3. Normal Case. + // Step 3.1 Get tag_rule_pairs. + std::vector> tag_rule_pairs; + for (int it_trigger = 0; it_trigger < static_cast(format.triggers.size()); ++it_trigger) { + const auto& trigger = format.triggers[it_trigger]; + std::vector choice_elements; + for (const auto& tag_id : trigger_to_tag_ids[it_trigger]) { + const auto& tag = format.tags[tag_id]; + const auto& tag_begin = std::get(tag.begin); + int begin_expr_id = grammar_builder_.AddByteString(tag_begin.substr(trigger.size())); + int rule_ref_expr_id = grammar_builder_.AddRuleRef(tag_content_rule_ids[tag_id]); + int end_expr_id = BuildEndExpr(tag); + choice_elements.push_back( + grammar_builder_.AddSequence({begin_expr_id, rule_ref_expr_id, end_expr_id}) + ); + } + auto choice_expr_id = grammar_builder_.AddChoices(choice_elements); + auto sub_rule_id = grammar_builder_.AddRuleWithHint("triggered_tags_group", choice_expr_id); + tag_rule_pairs.push_back(std::make_pair(trigger, sub_rule_id)); + } + + // Step 3.2 Add TagDispatch. + int32_t rule_expr_id; + bool loop_after_dispatch = !format.stop_after_first; + std::vector all_excludes = format.excludes; + for (const auto& s : format.detected_end_strs_) { + if (!s.empty()) { + all_excludes.push_back(s); + } + } + rule_expr_id = grammar_builder_.AddTagDispatch( + Grammar::Impl::TagDispatch{tag_rule_pairs, loop_after_dispatch, all_excludes} + ); + + // Step 3.3 Consider at_least_one + if (format.at_least_one) { + std::vector first_choice_elements; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr_id = BuildBeginExpr(tag); + auto rule_ref_expr_id = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr_id = BuildEndExpr(tag); + first_choice_elements.push_back( + grammar_builder_.AddSequence({begin_expr_id, rule_ref_expr_id, end_expr_id}) + ); + } + auto first_choice_expr_id = grammar_builder_.AddChoices(first_choice_elements); + auto first_rule_id = + grammar_builder_.AddRuleWithHint("triggered_tags_first", first_choice_expr_id); + + auto tag_dispatch_rule_id = + grammar_builder_.AddRuleWithHint("triggered_tags_sub", rule_expr_id); + auto ref_first_rule_expr_id = grammar_builder_.AddRuleRef(first_rule_id); + auto ref_tag_dispatch_rule_expr_id = grammar_builder_.AddRuleRef(tag_dispatch_rule_id); + auto sequence_expr_id = + grammar_builder_.AddSequence({ref_first_rule_expr_id, ref_tag_dispatch_rule_expr_id}); + rule_expr_id = grammar_builder_.AddChoices({sequence_expr_id}); + } + + auto rule_id = grammar_builder_.AddRuleWithHint("triggered_tags", rule_expr_id); + return ResultOk(rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const TagsWithSeparatorFormat& format +) { + // The grammar: + // Step 1. tags_rule: call tags + // tags_rule ::= tag1 | tag2 | ... | tagN + // Step 2. Special handling (stop_after_first is true): + // if at_least_one is false: + // root ::= tags_rule | "" + // if at_least_one is true: + // root ::= tags_rule + // Step 3. Normal handling (stop_after_first is false): + // if at_least_one is false: + // root ::= tags_rule tags_rule_sub | "" + // if at_least_one is true: + // root ::= tags_rule tags_rule_sub + // tags_rule_sub ::= sep tags_rule tags_rule_sub | "" + + // Step 1. Construct a rule representing any tag + std::vector choice_ids; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + auto tag_rule_id = Visit(format.tags[it_tag]); + if (tag_rule_id.IsErr()) { + return tag_rule_id; + } + auto tag_rule_ref_id = grammar_builder_.AddRuleRef(std::move(tag_rule_id).Unwrap()); + auto sequence_expr_id = grammar_builder_.AddSequence({tag_rule_ref_id}); + choice_ids.push_back(sequence_expr_id); + } + auto choice_expr_id = grammar_builder_.AddChoices(choice_ids); + auto all_tags_rule_id = + grammar_builder_.AddRuleWithHint("tags_with_separator_tags", choice_expr_id); + + auto all_tags_rule_ref_id = grammar_builder_.AddRuleRef(all_tags_rule_id); + + // Step 2. Special case (stop_after_first is true): + if (format.stop_after_first) { + int32_t rule_body_expr_id; + if (format.at_least_one) { + // root ::= tags_rule + rule_body_expr_id = + grammar_builder_.AddChoices({grammar_builder_.AddSequence({all_tags_rule_ref_id})}); + } else { + // root ::= tags_rule | "" + rule_body_expr_id = grammar_builder_.AddChoices( + {grammar_builder_.AddSequence({all_tags_rule_ref_id}), grammar_builder_.AddEmptyStr()} + ); + } + + auto rule_id = grammar_builder_.AddRuleWithHint("tags_with_separator", rule_body_expr_id); + return ResultOk(rule_id); + } + + // Step 3. Normal handling (stop_after_first is false): + // Step 3.1 Construct sub rule: sub ::= sep tags sub | "" + auto sub_rule_id = grammar_builder_.AddEmptyRuleWithHint("tags_with_separator_sub"); + + auto end_str_sequence_id = grammar_builder_.AddEmptyStr(); + + std::vector sub_sequence_elements; + if (!format.separator.empty()) { + sub_sequence_elements.push_back(grammar_builder_.AddByteString(format.separator)); + } + sub_sequence_elements.push_back(all_tags_rule_ref_id); + sub_sequence_elements.push_back(grammar_builder_.AddRuleRef(sub_rule_id)); + + auto sub_rule_body_id = grammar_builder_.AddChoices( + {grammar_builder_.AddSequence(sub_sequence_elements), end_str_sequence_id} + ); + grammar_builder_.UpdateRuleBody(sub_rule_id, sub_rule_body_id); + + // Step 3.2 Construct root rule + std::vector choices = { + grammar_builder_.AddSequence({all_tags_rule_ref_id, grammar_builder_.AddRuleRef(sub_rule_id)} + ), + }; + if (!format.at_least_one) { + choices.push_back(end_str_sequence_id); + } + auto rule_body_expr_id = grammar_builder_.AddChoices(choices); + auto rule_id = grammar_builder_.AddRuleWithHint("tags_with_separator", rule_body_expr_id); + return ResultOk(rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const OptionalFormat& format) { + // optional: 0 or 1 occurrence -> Choice(content, "") + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + auto content_ref = grammar_builder_.AddRuleRef(content_rule_id); + auto expr = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({content_ref})} + ); + return ResultOk(grammar_builder_.AddRuleWithHint("optional", expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const PlusFormat& format) { + // plus: 1 or more occurrences -> content content_star, where content_star = content content_star + // | "" + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + auto content_ref = grammar_builder_.AddRuleRef(content_rule_id); + auto star_rule_id = grammar_builder_.AddEmptyRuleWithHint("plus_star"); + auto star_ref = grammar_builder_.AddRuleRef(star_rule_id); + auto star_body = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({content_ref, star_ref})} + ); + grammar_builder_.UpdateRuleBody(star_rule_id, star_body); + auto plus_expr = grammar_builder_.AddSequence({content_ref, star_ref}); + return ResultOk(grammar_builder_.AddRuleWithHint("plus", plus_expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const StarFormat& format) { + // star: 0 or more occurrences -> content_star, where content_star = content content_star | "" + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + auto content_ref = grammar_builder_.AddRuleRef(content_rule_id); + auto star_rule_id = grammar_builder_.AddEmptyRuleWithHint("star"); + auto star_ref = grammar_builder_.AddRuleRef(star_rule_id); + auto star_body = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({content_ref, star_ref})} + ); + grammar_builder_.UpdateRuleBody(star_rule_id, star_body); + return ResultOk(grammar_builder_.AddRuleWithHint("star", star_ref)); +} + +Result StructuralTagGrammarConverter::VisitSub(const TokenFormat& format) { + XGRAMMAR_DCHECK(format.resolved_token_id_ >= 0) + << "TokenFormat must be resolved before conversion"; + auto token_set_expr = grammar_builder_.AddTokenSet({format.resolved_token_id_}); + auto seq = grammar_builder_.AddSequence({token_set_expr}); + auto choices = grammar_builder_.AddChoices({seq}); + return ResultOk(grammar_builder_.AddRuleWithHint("token", choices)); +} + +Result StructuralTagGrammarConverter::VisitSub(const ExcludeTokenFormat& format) { + std::vector all_excludes = format.resolved_token_ids_; + for (auto tid : format.detected_end_token_ids_) { + all_excludes.push_back(tid); + } + int expr = grammar_builder_.AddExcludeTokenSet(all_excludes); + auto seq = grammar_builder_.AddSequence({expr}); + auto choices = grammar_builder_.AddChoices({seq}); + return ResultOk(grammar_builder_.AddRuleWithHint("exclude_token", choices)); +} + +Result StructuralTagGrammarConverter::VisitSub(const AnyTokensFormat& format) { + std::vector all_excludes = format.resolved_exclude_token_ids_; + for (auto tid : format.detected_end_token_ids_) { + all_excludes.push_back(tid); + } + int exclude_expr = grammar_builder_.AddExcludeTokenSet(all_excludes); + int exclude_seq = grammar_builder_.AddSequence({exclude_expr}); + int exclude_choices = grammar_builder_.AddChoices({exclude_seq}); + int inner_rule = grammar_builder_.AddRuleWithHint("any_tokens_inner", exclude_choices); + auto inner_ref = grammar_builder_.AddRuleRef(inner_rule); + auto star_rule_id = grammar_builder_.AddEmptyRuleWithHint("any_tokens"); + auto star_ref = grammar_builder_.AddRuleRef(star_rule_id); + auto star_body = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({inner_ref, star_ref})} + ); + SetRuleBodyAndBudgets(star_rule_id, star_body, format.max_tokens); + return ResultOk(star_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const TokenTriggeredTagsFormat& format +) { + // Step 1. Visit all tags, map trigger → tag IDs + std::vector> trigger_to_tag_ids(format.trigger_tokens.size()); + std::vector tag_content_rule_ids; + tag_content_rule_ids.reserve(format.tags.size()); + + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + if (!std::holds_alternative(tag.begin)) { + return ResultErr( + "Tags in token_triggered_tags must have a token format begin, not a string" + ); + } + auto begin_token_id = std::get(tag.begin).resolved_token_id_; + + int matched = -1; + for (int it_t = 0; it_t < static_cast(format.resolved_trigger_token_ids_.size()); ++it_t) { + if (format.resolved_trigger_token_ids_[it_t] == begin_token_id) { + if (matched != -1) { + return ResultErr("Tag matches multiple triggers"); + } + matched = it_t; + } + } + if (matched == -1) { + return ResultErr("Tag does not match any trigger"); + } + trigger_to_tag_ids[matched].push_back(it_tag); + + auto result = Visit(*tag.content); + if (result.IsErr()) return result; + tag_content_rule_ids.push_back(std::move(result).Unwrap()); + } + + // Step 2. Special case: at_least_one && stop_after_first + if (format.at_least_one && format.stop_after_first) { + std::vector choice_elements; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr = BuildBeginExpr(tag); + auto ref = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr = BuildEndExpr(tag); + choice_elements.push_back(grammar_builder_.AddSequence({begin_expr, ref, end_expr})); + } + auto choice = grammar_builder_.AddChoices(choice_elements); + return ResultOk(grammar_builder_.AddRuleWithHint("token_triggered_tags", choice)); + } + + // Step 3. Normal case — TokenTagDispatch + std::vector> trigger_rule_pairs; + for (int it_t = 0; it_t < static_cast(format.trigger_tokens.size()); ++it_t) { + std::vector choice_elements; + for (auto tag_id : trigger_to_tag_ids[it_t]) { + const auto& tag = format.tags[tag_id]; + auto ref = grammar_builder_.AddRuleRef(tag_content_rule_ids[tag_id]); + auto end_expr = BuildEndExpr(tag); + choice_elements.push_back(grammar_builder_.AddSequence({ref, end_expr})); + } + auto choice = grammar_builder_.AddChoices(choice_elements); + auto sub_rule = grammar_builder_.AddRuleWithHint("token_triggered_tags_group", choice); + trigger_rule_pairs.push_back({format.resolved_trigger_token_ids_[it_t], sub_rule}); + } + + bool loop = !format.stop_after_first; + std::vector all_excludes = format.resolved_exclude_token_ids_; + for (auto tid : format.detected_end_token_ids_) { + all_excludes.push_back(tid); + } + auto ttd_expr = grammar_builder_.AddTokenTagDispatch( + Grammar::Impl::TokenTagDispatch{trigger_rule_pairs, loop, all_excludes} + ); + int32_t rule_expr_id = ttd_expr; + + if (format.at_least_one) { + std::vector first_choices; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr = BuildBeginExpr(tag); + auto ref = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr = BuildEndExpr(tag); + first_choices.push_back(grammar_builder_.AddSequence({begin_expr, ref, end_expr})); + } + auto first_choice = grammar_builder_.AddChoices(first_choices); + auto first_rule = grammar_builder_.AddRuleWithHint("token_triggered_tags_first", first_choice); + auto dispatch_rule = grammar_builder_.AddRuleWithHint("token_triggered_tags_sub", rule_expr_id); + auto seq = grammar_builder_.AddSequence( + {grammar_builder_.AddRuleRef(first_rule), grammar_builder_.AddRuleRef(dispatch_rule)} + ); + rule_expr_id = grammar_builder_.AddChoices({seq}); + } + + return ResultOk(grammar_builder_.AddRuleWithHint("token_triggered_tags", rule_expr_id)); +} + +Result StructuralTagGrammarConverter::VisitSub(const RepeatFormat& format) { + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + int repeat_expr_id = grammar_builder_.AddRepeat(content_rule_id, format.min, format.max); + return ResultOk(grammar_builder_.AddRuleWithHint("repeat", repeat_expr_id)); +} + +Result StructuralTagGrammarConverter::VisitSub(const DispatchFormat& format) { + std::vector> tag_rule_pairs; + tag_rule_pairs.reserve(format.rules.size()); + for (const auto& pair : format.rules) { + if (!pair.second) { + return ResultErr("TagDispatch pair must have content"); + } + auto result = Visit(*pair.second); + if (result.IsErr()) { + return result; + } + tag_rule_pairs.push_back({pair.first, std::move(result).Unwrap()}); + } + auto rule_expr_id = grammar_builder_.AddTagDispatch( + Grammar::Impl::TagDispatch{std::move(tag_rule_pairs), format.loop, format.excludes} + ); + return ResultOk(grammar_builder_.AddRuleWithHint("tag_dispatch", rule_expr_id)); +} + +Result StructuralTagGrammarConverter::VisitSub(const TokenDispatchFormat& format) { + XGRAMMAR_DCHECK(format.resolved_trigger_token_ids_.size() == format.rules.size()) + << "TokenDispatchFormat must be resolved before conversion"; + std::vector> trigger_rule_pairs; + trigger_rule_pairs.reserve(format.rules.size()); + for (size_t i = 0; i < format.rules.size(); ++i) { + const auto& pair = format.rules[i]; + if (!pair.second) { + return ResultErr("TokenTagDispatch pair must have content"); + } + auto result = Visit(*pair.second); + if (result.IsErr()) { + return result; + } + trigger_rule_pairs.push_back({format.resolved_trigger_token_ids_[i], std::move(result).Unwrap()} + ); + } + std::vector all_excludes = format.resolved_exclude_token_ids_; + auto rule_expr_id = grammar_builder_.AddTokenTagDispatch( + Grammar::Impl::TokenTagDispatch{trigger_rule_pairs, format.loop, all_excludes} + ); + return ResultOk(grammar_builder_.AddRuleWithHint("token_tag_dispatch", rule_expr_id)); +} + +/************** StructuralTag Conversion Public API **************/ + +Result StructuralTagToGrammar( + const std::string& structural_tag_json, const std::optional& tokenizer_info +) { + auto structural_tag_result = StructuralTagParser::FromJSON(structural_tag_json); + if (structural_tag_result.IsErr()) { + return ResultErr(std::move(structural_tag_result).UnwrapErr()); + } + auto structural_tag = std::move(structural_tag_result).Unwrap(); + + auto resolve_err = StructuralTagTokenResolver::Resolve(&structural_tag, tokenizer_info); + if (resolve_err.has_value()) { + return ResultErr(std::move(resolve_err).value()); + } + + auto err = StructuralTagAnalyzer().Analyze(&structural_tag); + if (err.has_value()) { + return ResultErr(std::move(err).value()); + } + + auto result = StructuralTagGrammarConverter::Convert(structural_tag); + if (result.IsErr()) { + return ResultErr(std::move(result).UnwrapErr()); + } + return ResultOk(GrammarNormalizer::Apply(std::move(result).Unwrap())); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/structural_tag.h b/third_party/xgrammar/cpp/structural_tag.h new file mode 100644 index 0000000000..bf22e34948 --- /dev/null +++ b/third_party/xgrammar/cpp/structural_tag.h @@ -0,0 +1,411 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/structural_tag_impl.h + * \brief The implementation header for the structural tag. + */ + +#ifndef XGRAMMAR_STRUCTURAL_TAG_H_ +#define XGRAMMAR_STRUCTURAL_TAG_H_ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "support/utils.h" +#include "xgrammar/tokenizer_info.h" + +namespace xgrammar { + +/******************** Structural Tag Definition ********************/ + +// TODO(yixin): Consider moving the definition to Public API. + +struct ConstStringFormat; +struct JSONSchemaFormat; +struct AnyTextFormat; +struct GrammarFormat; +struct RegexFormat; +struct SequenceFormat; +struct OrFormat; +struct TagFormat; +struct TriggeredTagsFormat; +struct TagsWithSeparatorFormat; +struct OptionalFormat; +struct PlusFormat; +struct StarFormat; +struct TokenFormat; +struct ExcludeTokenFormat; +struct AnyTokensFormat; +struct TokenTriggeredTagsFormat; +struct RepeatFormat; +struct DispatchFormat; +struct TokenDispatchFormat; + +using Format = std::variant< + ConstStringFormat, + JSONSchemaFormat, + AnyTextFormat, + GrammarFormat, + RegexFormat, + SequenceFormat, + OrFormat, + TagFormat, + TriggeredTagsFormat, + TagsWithSeparatorFormat, + OptionalFormat, + PlusFormat, + StarFormat, + TokenFormat, + ExcludeTokenFormat, + AnyTokensFormat, + TokenTriggeredTagsFormat, + RepeatFormat, + DispatchFormat, + TokenDispatchFormat>; + +/******************** Basic Formats ********************/ + +struct ConstStringFormat { + static constexpr const char* type = "const_string"; + std::string value; + ConstStringFormat(std::string value) : value(std::move(value)) {} + picojson::value ToJSON() const; +}; + +struct JSONSchemaFormat { + static constexpr const char* type = "json_schema"; + std::string json_schema; + // "json", "qwen_xml", "minimax_xml", "minimax_m3_xml", "deepseek_xml", "glm_xml", + // "cohere_xml", "kimi_k3_xml", "deepseek_v4_1_xml" + std::string style = "json"; + // Whether to allow object properties to appear in any order. See + // Grammar::FromJSONSchema / JSONSchemaToEBNF for the semantics. + bool any_order = false; + // Per-tag cap on consecutive whitespace characters in the JSON-schema content. + std::optional max_whitespace_cnt = std::nullopt; + JSONSchemaFormat( + std::string json_schema, + std::string style = "json", + bool any_order = false, + std::optional max_whitespace_cnt = std::nullopt + ) + : json_schema(std::move(json_schema)), + style(std::move(style)), + any_order(any_order), + max_whitespace_cnt(max_whitespace_cnt) {} + picojson::value ToJSON() const; +}; + +struct GrammarFormat { + static constexpr const char* type = "grammar"; + std::string grammar; + GrammarFormat(std::string grammar) : grammar(std::move(grammar)) {} + picojson::value ToJSON() const; +}; + +struct RegexFormat { + static constexpr const char* type = "regex"; + std::string pattern; + RegexFormat(std::string pattern) : pattern(std::move(pattern)) {} + picojson::value ToJSON() const; +}; + +struct AnyTextFormat { + static constexpr const char* type = "any_text"; + std::vector excludes; + int32_t max_tokens = -1; + int32_t max_chars = -1; + AnyTextFormat( + std::vector excluded_strs, int32_t max_tokens = -1, int32_t max_chars = -1 + ) + : excludes(std::move(excluded_strs)), max_tokens(max_tokens), max_chars(max_chars) {} + picojson::value ToJSON() const; + + private: + std::vector detected_end_strs_; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TokenFormat { + static constexpr const char* type = "token"; + std::variant token; + TokenFormat(std::variant token) : token(std::move(token)) { + if (std::holds_alternative(this->token)) { + resolved_token_id_ = std::get(this->token); + } + } + picojson::value ToJSON() const; + + private: + int32_t resolved_token_id_ = -1; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct ExcludeTokenFormat { + static constexpr const char* type = "exclude_token"; + std::vector> exclude_tokens; + ExcludeTokenFormat(std::vector> exclude_tokens) + : exclude_tokens(std::move(exclude_tokens)) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_token_ids_; + std::vector detected_end_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct AnyTokensFormat { + static constexpr const char* type = "any_tokens"; + std::vector> exclude_tokens; + int32_t max_tokens = -1; + AnyTokensFormat( + std::vector> exclude_tokens, int32_t max_tokens = -1 + ) + : exclude_tokens(std::move(exclude_tokens)), max_tokens(max_tokens) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_exclude_token_ids_; + std::vector detected_end_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +/******************** Combinatorial Formats ********************/ + +struct SequenceFormat { + static constexpr const char* type = "sequence"; + std::vector elements; + SequenceFormat(std::vector elements); + picojson::value ToJSON() const; + + private: + // Detected in StructuralTagAnalyzer + bool is_unlimited_ = false; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct OrFormat { + static constexpr const char* type = "or"; + std::vector elements; + OrFormat(std::vector elements); + picojson::value ToJSON() const; + + private: + // Detected in StructuralTagAnalyzer + bool is_unlimited_ = false; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TagFormat { + static constexpr const char* type = "tag"; + std::variant begin; + std::shared_ptr content; + std::variant, TokenFormat> end; + + TagFormat( + std::variant begin, + std::shared_ptr content, + std::variant, TokenFormat> end + ) + : begin(std::move(begin)), content(std::move(content)), end(std::move(end)) {} + picojson::value ToJSON() const; +}; + +struct TriggeredTagsFormat { + static constexpr const char* type = "triggered_tags"; + std::vector triggers; + std::vector tags; + std::vector excludes; + bool at_least_one = false; + bool stop_after_first = false; + + TriggeredTagsFormat( + std::vector triggers, + std::vector tags, + std::vector excludes, + bool at_least_one, + bool stop_after_first + ) + : triggers(std::move(triggers)), + tags(std::move(tags)), + excludes(std::move(excludes)), + at_least_one(at_least_one), + stop_after_first(stop_after_first) {} + picojson::value ToJSON() const; + + private: + std::vector detected_end_strs_; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TagsWithSeparatorFormat { + static constexpr const char* type = "tags_with_separator"; + std::vector tags; + std::string separator; + bool at_least_one = false; + bool stop_after_first = false; + + TagsWithSeparatorFormat( + std::vector tags, std::string separator, bool at_least_one, bool stop_after_first + ) + : tags(std::move(tags)), + separator(std::move(separator)), + at_least_one(at_least_one), + stop_after_first(stop_after_first) {} + picojson::value ToJSON() const; + + private: + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TokenTriggeredTagsFormat { + static constexpr const char* type = "token_triggered_tags"; + std::vector> trigger_tokens; + std::vector tags; + std::vector> exclude_tokens; + bool at_least_one = false; + bool stop_after_first = false; + + TokenTriggeredTagsFormat( + std::vector> trigger_tokens, + std::vector tags, + std::vector> exclude_tokens, + bool at_least_one, + bool stop_after_first + ) + : trigger_tokens(std::move(trigger_tokens)), + tags(std::move(tags)), + exclude_tokens(std::move(exclude_tokens)), + at_least_one(at_least_one), + stop_after_first(stop_after_first) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_trigger_token_ids_; + std::vector resolved_exclude_token_ids_; + std::vector detected_end_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct OptionalFormat { + static constexpr const char* type = "optional"; + std::shared_ptr content; + OptionalFormat(std::shared_ptr content) : content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +struct PlusFormat { + static constexpr const char* type = "plus"; + std::shared_ptr content; + PlusFormat(std::shared_ptr content) : content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +struct StarFormat { + static constexpr const char* type = "star"; + std::shared_ptr content; + StarFormat(std::shared_ptr content) : content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +struct RepeatFormat { + static constexpr const char* type = "repeat"; + int32_t min; + int32_t max; + std::shared_ptr content; + RepeatFormat(int32_t min, int32_t max, std::shared_ptr content) + : min(min), max(max), content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +/*! + * \brief A format that maps directly to a TagDispatch grammar. + * Accepts ``[trigger string, content format]`` pairs in JSON; each content is converted to a rule + * and the result is a single TagDispatch(loop, excludes). + */ +struct DispatchFormat { + static constexpr const char* type = "dispatch"; + std::vector>> rules; + bool loop = true; + std::vector excludes; + + DispatchFormat( + std::vector>> rules, + bool loop = true, + std::vector excludes = {} + ) + : rules(std::move(rules)), loop(loop), excludes(std::move(excludes)) {} + picojson::value ToJSON() const; +}; + +/*! + * \brief A format that maps directly to a TokenTagDispatch grammar. + * Accepts ``[trigger token, content format]`` pairs in JSON; trigger can be token ID or token + * string (resolved via tokenizer_info). Each content is converted to a rule. + */ +struct TokenDispatchFormat { + static constexpr const char* type = "token_dispatch"; + std::vector, std::shared_ptr>> rules; + bool loop = true; + std::vector> exclude_tokens; + + TokenDispatchFormat( + std::vector, std::shared_ptr>> rules, + bool loop = true, + std::vector> exclude_tokens = {} + ) + : rules(std::move(rules)), loop(loop), exclude_tokens(std::move(exclude_tokens)) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_trigger_token_ids_; + std::vector resolved_exclude_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagGrammarConverter; +}; + +/******************** Top Level ********************/ + +struct StructuralTag { + static constexpr const char* type = "structural_tag"; + Format format; + + StructuralTag(Format format) : format(std::move(format)) {} +}; + +/******************** Conversion API ********************/ + +/*! + * \brief Convert a structural tag JSON string to a grammar. + * \param structural_tag_json The JSON string of the structural tag. + * \return A grammar if the JSON is valid, otherwise an error message in std::string. + */ +Result StructuralTagToGrammar( + const std::string& structural_tag_json, + const std::optional& tokenizer_info = std::nullopt +); + +} // namespace xgrammar + +#endif // XGRAMMAR_STRUCTURAL_TAG_H_ diff --git a/third_party/xgrammar/cpp/suffix_automata.cc b/third_party/xgrammar/cpp/suffix_automata.cc new file mode 100644 index 0000000000..a3b4c85bac --- /dev/null +++ b/third_party/xgrammar/cpp/suffix_automata.cc @@ -0,0 +1,88 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/suffix_automata.cc + */ + +#include "suffix_automata.h" + +#include +#include +#include +#include + +namespace xgrammar { + +FSMWithStartEnd SuffixAutomata::Build(const std::vector& chunks) { + // Step 1. Build the suffix automaton over the chunk sequence with the standard online + // construction. Each chunk is treated as one symbol of the alphabet. + struct State { + int32_t length = 0; + int32_t suffix_link = -1; + std::map transitions; + }; + + std::vector states(1); + int32_t last = 0; + for (const std::string& chunk : chunks) { + int32_t current = static_cast(states.size()); + states.push_back({states[last].length + 1, -1, {}}); + + int32_t parent = last; + while (parent != -1 && !states[parent].transitions.count(chunk)) { + states[parent].transitions[chunk] = current; + parent = states[parent].suffix_link; + } + if (parent == -1) { + states[current].suffix_link = 0; + } else { + int32_t target = states[parent].transitions.at(chunk); + if (states[parent].length + 1 == states[target].length) { + states[current].suffix_link = target; + } else { + int32_t clone = static_cast(states.size()); + states.push_back(states[target]); + states[clone].length = states[parent].length + 1; + while (parent != -1) { + auto transition = states[parent].transitions.find(chunk); + if (transition == states[parent].transitions.end() || transition->second != target) { + break; + } + transition->second = clone; + parent = states[parent].suffix_link; + } + states[target].suffix_link = clone; + states[current].suffix_link = clone; + } + } + last = current; + } + + // Step 2. Expand the chunk-level automaton into a byte-level FSM. Automaton state i maps to + // FSM state i; every automaton state is accepting. A chunk-labeled transition becomes a chain + // of byte transitions through fresh intermediate states; an empty chunk becomes an epsilon + // transition. + FSM fsm(static_cast(states.size())); + std::vector end_states; + end_states.reserve(states.size()); + for (int32_t index = 0; index < static_cast(states.size()); ++index) { + end_states.push_back(index); + } + for (int32_t index = 0; index < static_cast(states.size()); ++index) { + for (const auto& [chunk, target] : states[index].transitions) { + if (chunk.empty()) { + fsm.AddEpsilonEdge(index, target); + continue; + } + int current_state = index; + for (size_t offset = 0; offset < chunk.size(); ++offset) { + int next_state = offset + 1 == chunk.size() ? static_cast(target) : fsm.AddState(); + uint8_t byte = static_cast(chunk[offset]); + fsm.AddEdge(current_state, next_state, byte, byte); + current_state = next_state; + } + } + } + return FSMWithStartEnd(fsm, 0, std::move(end_states)); +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/suffix_automata.h b/third_party/xgrammar/cpp/suffix_automata.h new file mode 100644 index 0000000000..d4d9dac536 --- /dev/null +++ b/third_party/xgrammar/cpp/suffix_automata.h @@ -0,0 +1,36 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/suffix_automata.h + * \brief Suffix automaton construction for substring expressions. + */ +#ifndef XGRAMMAR_SUFFIX_AUTOMATA_H_ +#define XGRAMMAR_SUFFIX_AUTOMATA_H_ + +#include +#include + +#include "fsm.h" + +namespace xgrammar { + +/*! + * \brief Builds the automaton of a substring expression via a chunk-level suffix automaton. + */ +class SuffixAutomata { + public: + /*! + * \brief Build an FSM that accepts exactly the contiguous subsequences of the chunk list, + * including the empty one. + * \details A suffix automaton is built over the chunk sequence (each chunk is one symbol), so + * the number of automaton states grows linearly with the number of chunks. Every automaton + * state is accepting. Each chunk-labeled transition is then expanded into a chain of byte + * transitions; an empty chunk becomes an epsilon transition. + * \param chunks The list of byte string chunks. Chunks may be empty or repeated. + * \return The FSM with start and end states. + */ + static FSMWithStartEnd Build(const std::vector& chunks); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUFFIX_AUTOMATA_H_ diff --git a/third_party/xgrammar/cpp/support/compact_2d_array.h b/third_party/xgrammar/cpp/support/compact_2d_array.h new file mode 100644 index 0000000000..f95c20e9c1 --- /dev/null +++ b/third_party/xgrammar/cpp/support/compact_2d_array.h @@ -0,0 +1,397 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/compact_2d_array.h + */ +#ifndef XGRAMMAR_SUPPORT_COMPACT_2D_ARRAY_H_ +#define XGRAMMAR_SUPPORT_COMPACT_2D_ARRAY_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" +#include "memory_size.h" +#include "reflection.h" + +namespace xgrammar { + +/*! + * \brief This class implements a Compressed Sparse Row (CSR) array data structure. It stores + * a 2D array in a compressed format, where each row can have a variable number of elements, and + * all rows are stored contiguously in memory. The inserted row is immutable. + * + * \note Inserting new rows into the Compact2DArray will invalidate the existing Row objects. + * + * \tparam DataType The type of elements stored in the Compact2DArray. + * + * \details + * The Compact2DArray stores elements of type DataType in a compressed format, + * where each row can have a variable number of elements. It uses two vectors: + * - data_: stores all elements contiguously + * - indptr_: stores the starting index of each row in data_. Its last element is the size of data_ + * representing the ending index. + * + * This structure allows efficient storage and access for sparse data. + */ +template +class Compact2DArray { + public: + /*! + * \brief The struct representing a row in the Compact2DArray. + */ + struct Row { + /*! \brief The value type is DataType. */ + using value_type = DataType; + + /*! \brief Pointer to the data of the row. */ + const DataType* data; + /*! \brief Length of the row data. */ + int32_t data_len; + + /*! + * \brief Access an element in the row. + * \param i Index of the element to access. + * \return Reference to the element at index i. + */ + const DataType& operator[](int32_t i) const { + XGRAMMAR_DCHECK(i >= 0 && i < data_len) + << "Index " << i << " of the Compact2DArray Row is out of bound"; + return data[i]; + } + + /*! \brief Get the beginning iterator of the row. */ + const DataType* begin() const { return data; } + /*! \brief Get the end iterator of the row. */ + const DataType* end() const { return data + data_len; } + /*! \brief Get the size of the row. */ + int32_t size() const { return data_len; } + + /*! \brief Get a sub-row in [begin, end). */ + Row Slice(int32_t begin, int32_t end) const { + XGRAMMAR_DCHECK(begin >= 0 && begin <= end && end <= data_len) + << "Compact2DArray Row slice is out of bound"; + return {data + begin, end - begin}; + } + + friend std::ostream& operator<<(std::ostream& os, const Row& row) { + os << "["; + for (auto i = 0; i < row.data_len; ++i) { + if (i > 0) { + os << ", "; + } + os << row[i]; + } + os << "]"; + return os; + } + }; + + /*! + * \brief The mutable struct representing a row in the Compact2DArray. + */ + struct MutableRow { + /*! \brief The value type is DataType. */ + using value_type = DataType; + + /*! \brief Pointer to the data of the row. */ + DataType* data; + /*! \brief Length of the row data. */ + int32_t data_len; + + /*! + * \brief Access an element in the row. + * \param i Index of the element to access. + * \return Reference to the element at index i. + */ + DataType& operator[](int32_t i) const { + XGRAMMAR_DCHECK(i >= 0 && i < data_len) + << "Index " << i << " of the Compact2DArray MutableRow is out of bound"; + return data[i]; + } + + /*! \brief Get the beginning iterator of the row. */ + DataType* begin() const { return data; } + /*! \brief Get the end iterator of the row. */ + DataType* end() const { return data + data_len; } + /*! \brief Get the size of the row. */ + int32_t size() const { return data_len; } + }; + + /*! \brief The value type is Row. */ + using value_type = Row; + + /*! \brief Default constructor. */ + Compact2DArray() = default; + + /*! + * \brief Construct a Compact2DArray from an existing CSR representation. + * \param data All row elements stored contiguously. + * \param indptr Row start offsets. Must start with 0, be non-decreasing, and end with + * data.size(). + * \return The constructed Compact2DArray. + */ + static Compact2DArray FromDataAndIndptr(std::vector data, std::vector indptr); + + /*! + * \brief Construct a Compact2DArray from row sizes with default-constructed data. + * \param row_sizes The size of each row. + * \return The constructed Compact2DArray. + */ + static Compact2DArray FromRowSizes(const std::vector& row_sizes); + + /*! + * \brief Reset the Compact2DArray from row sizes with default-constructed data. + * \param row_sizes The size of each row. + */ + void ResetWithRowSizes(const std::vector& row_sizes); + + /****************** Accessors ******************/ + + /*! \brief Get the number of rows in the Compact2DArray. */ + int32_t size() const { return static_cast(indptr_.size()) - 1; } + + /*! + * \brief Check the CSR invariants after deserialization: indptr starts with 0, is non-decreasing + * and ends with the data size. Otherwise row accesses would read out of bounds. + * \return An error message if the invariants are violated. + */ + std::optional Validate() const { + if (indptr_.empty() || indptr_.front() != 0 || + indptr_.back() != static_cast(data_.size()) || + !std::is_sorted(indptr_.begin(), indptr_.end())) { + return "Invalid indptr: it must start with 0, be non-decreasing and end with the data size"; + } + return std::nullopt; + } + + friend std::size_t MemorySize(const Compact2DArray& arr) { + return MemorySize(arr.data_) + MemorySize(arr.indptr_); + } + + /*! + * \brief Access a row in the Compact2DArray. + * \param i Index of the row to access. + * \return Row struct representing the i-th row. + */ + Row operator[](int32_t i) const; + + /*! + * \brief Access a mutable row in the Compact2DArray. + * \param i Index of the row to access. + * \return MutableRow struct representing the i-th row. + */ + MutableRow MutableRowAt(int32_t i); + + /****************** Modifiers ******************/ + + /*! + * \brief Insert a new row of data into the Compact2DArray. + * \param data Pointer to the data to be inserted. + * \param data_len Length of the data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBack(const DataType* new_data, int32_t new_data_len); + + /*! + * \brief Insert a new row of data into the Compact2DArray from a vector. + * \param data Vector containing the data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBack(const std::vector& new_data); + + /*! + * \brief Insert a new row of data into the Compact2DArray from a Row struct. + * \param row The Row struct containing the data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBack(const Row& row) { return PushBack(row.data, row.data_len); } + + /*! + * \brief Push back a new element in the latest row. + * \param new_data the element to be pushed. + */ + void PushBackInLatestRow(const DataType& new_data) { + XGRAMMAR_DCHECK(!indptr_.empty()) << "Cannot push back in an empty Compact2DArray"; + CheckCanAppendData(1); + data_.push_back(new_data); + indptr_.back()++; + } + + Row Back() { return (*this)[size() - 1]; } + + /*! + * \brief Insert a new row of non-contiguous data into the Compact2DArray. This method inserts a + * single element followed by a sequence of elements. This is useful in the GrammarExpr data + * structure. + * \param data_1 The first element to be inserted. + * \param data_2 Pointer to the remaining data to be inserted. + * \param data_2_len Length of the remaining data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBackNonContiguous(DataType data_1, const DataType* data_2, int32_t data_2_len); + + /*! + * \brief Pop back the last one or multiple rows of the Compact2DArray. + * \param cnt The number of rows to be popped. + */ + void PopBack(const int32_t& cnt) { + indptr_.erase(indptr_.end() - cnt, indptr_.end()); + data_.erase(data_.begin() + indptr_.back(), data_.end()); + return; + } + + /****************** Internal Accessors ******************/ + + /*! \brief Get a pointer to the underlying data array. */ + const DataType* data() const { return data_.data(); } + /*! \brief Get a pointer to the underlying index pointer array. */ + const int32_t* indptr() const { return indptr_.data(); } + + /****************** Printing ******************/ + + friend std::ostream& operator<<(std::ostream& os, const Compact2DArray& compact_2d_array) { + os << "Compact2DArray(["; + for (auto i = 0; i < compact_2d_array.size(); ++i) { + if (i > 0) { + os << ", "; + } + os << compact_2d_array[i]; + } + os << "])"; + return os; + } + + private: + static constexpr size_t kMaxRepresentableSize = static_cast(INT32_MAX); + + inline void CheckCanAppendData(size_t additional_size) const { + XGRAMMAR_ICHECK(additional_size <= kMaxRepresentableSize - data_.size()) + << "Compact2DArray data size would exceed the int32_t limit, likely due to an unbounded " + "grammar pattern causing state explosion"; + } + + /*! \brief Vector storing all elements contiguously. */ + std::vector data_; + /*! \brief Vector storing the starting index of each row in data_. */ + std::vector indptr_{0}; + friend struct member_trait>; +}; + +template +inline typename Compact2DArray::Row Compact2DArray::operator[](int32_t i +) const { + XGRAMMAR_DCHECK(i >= 0 && i < size()) << "Compact2DArray index " << i << " is out of bound"; + int32_t start = indptr_[i]; + int32_t end = indptr_[i + 1]; + return {data_.data() + start, end - start}; +} + +template +inline typename Compact2DArray::MutableRow Compact2DArray::MutableRowAt( + int32_t i +) { + XGRAMMAR_DCHECK(i >= 0 && i < size()) << "Compact2DArray index " << i << " is out of bound"; + int32_t start = indptr_[i]; + int32_t end = indptr_[i + 1]; + return {data_.data() + start, end - start}; +} + +template +inline Compact2DArray Compact2DArray::FromDataAndIndptr( + std::vector data, std::vector indptr +) { + XGRAMMAR_CHECK(!indptr.empty()) << "Compact2DArray indptr cannot be empty"; + XGRAMMAR_CHECK(indptr.front() == 0) << "Compact2DArray indptr must start with 0"; + for (int32_t i = 1; i < static_cast(indptr.size()); ++i) { + XGRAMMAR_CHECK(indptr[i - 1] <= indptr[i]) << "Compact2DArray indptr must be non-decreasing"; + } + XGRAMMAR_CHECK(indptr.back() == static_cast(data.size())) + << "Compact2DArray indptr must end with data.size()"; + + Compact2DArray result; + result.data_ = std::move(data); + result.indptr_ = std::move(indptr); + return result; +} + +template +inline Compact2DArray Compact2DArray::FromRowSizes( + const std::vector& row_sizes +) { + Compact2DArray result; + result.ResetWithRowSizes(row_sizes); + return result; +} + +template +inline void Compact2DArray::ResetWithRowSizes(const std::vector& row_sizes) { + indptr_.resize(row_sizes.size() + 1); + indptr_[0] = 0; + for (int32_t i = 0; i < static_cast(row_sizes.size()); ++i) { + XGRAMMAR_CHECK(row_sizes[i] >= 0) << "Compact2DArray row size cannot be negative"; + XGRAMMAR_CHECK(static_cast(row_sizes[i]) <= kMaxRepresentableSize - indptr_[i]) + << "Compact2DArray data size exceeds the int32_t limit"; + indptr_[i + 1] = indptr_[i] + row_sizes[i]; + } + data_.resize(indptr_.back()); +} + +template +inline int32_t Compact2DArray::PushBack(const DataType* new_data, int32_t new_data_len) { + CheckCanAppendData(static_cast(new_data_len)); + // If the new data is already in the Compact2DArray, we need to copy it to the new memory + // location. + if (new_data >= data_.data() && new_data < data_.data() + data_.size()) { + std::vector new_data_copied(new_data, new_data + new_data_len); + data_.insert(data_.end(), new_data_copied.begin(), new_data_copied.end()); + } else { + data_.insert(data_.end(), new_data, new_data + new_data_len); + } + indptr_.push_back(static_cast(data_.size())); + return static_cast(indptr_.size()) - 2; +} + +template +inline int32_t Compact2DArray::PushBack(const std::vector& new_data) { + CheckCanAppendData(new_data.size()); + data_.insert(data_.end(), new_data.begin(), new_data.end()); + indptr_.push_back(static_cast(data_.size())); + return static_cast(indptr_.size()) - 2; +} + +template +inline int32_t Compact2DArray::PushBackNonContiguous( + DataType data_1, const DataType* data_2, int32_t data_2_len +) { + CheckCanAppendData(static_cast(data_2_len) + 1); + if (data_2 >= data_.data() && data_2 < data_.data() + data_.size()) { + std::vector new_data_copied(data_2, data_2 + data_2_len); + data_.push_back(data_1); + data_.insert(data_.end(), new_data_copied.begin(), new_data_copied.end()); + } else { + data_.push_back(data_1); + data_.insert(data_.end(), data_2, data_2 + data_2_len); + } + indptr_.push_back(static_cast(data_.size())); + return static_cast(indptr_.size()) - 2; +} + +template +XGRAMMAR_MEMBER_TABLE_TEMPLATE( + Compact2DArray, + "data_", + &Compact2DArray::data_, + "indptr_", + &Compact2DArray::indptr_ +); + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_COMPACT_2D_ARRAY_H_ diff --git a/third_party/xgrammar/cpp/support/container.h b/third_party/xgrammar/cpp/support/container.h new file mode 100644 index 0000000000..2ca660e90b --- /dev/null +++ b/third_party/xgrammar/cpp/support/container.h @@ -0,0 +1,166 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/container.h + * \brief The header for container. + */ +#ifndef XGRAMMAR_SUPPORT_CONTAINER_H_ +#define XGRAMMAR_SUPPORT_CONTAINER_H_ +#include + +#include "logging.h" + +namespace xgrammar { + +namespace details { + +template +class NodePool { + public: + NodePool() = default; + + void Reserve(int n) { node_pool_.reserve(n); } + + [[nodiscard]] + int Allocate() { + if (free_list_.empty()) { + int node = Size(); + node_pool_.emplace_back(); + return node; + } else { + int node = free_list_.back(); + free_list_.pop_back(); + return node; + } + } + + void Deallocate(int node) { free_list_.push_back(node); } + + void Clear() { + node_pool_.clear(); + free_list_.clear(); + } + + Node& operator[](int node) { + XGRAMMAR_DCHECK(0 <= node && node < Size()); + return node_pool_[node]; + } + + int Size() const { return static_cast(node_pool_.size()); } + + private: + std::vector node_pool_; + std::vector free_list_; +}; + +} // namespace details + +template +class List { + private: + struct Node { + int prev; + int next; + Value value; + }; + + public: + struct iterator { + public: + iterator(int n, List& c) : node_(n), list_(&c) { + XGRAMMAR_DCHECK(0 <= node_ && node_ < list_->node_pool_.Size()); + } + iterator& operator++() { + node_ = GetNode().next; + return *this; + } + iterator operator++(int) { + iterator tmp = *this; + ++*this; + return tmp; + } + Value& operator*() const { return GetNode().value; } + Value* operator->() const { return &GetNode().value; } + bool operator==(const iterator& rhs) const { + XGRAMMAR_DCHECK(list_ == rhs.list_) << "compare different container is UB"; + return node_ == rhs.node_; // compare different container is UB + } + bool operator!=(const iterator& rhs) const { + XGRAMMAR_DCHECK(list_ == rhs.list_) << "compare different container is UB"; + return node_ != rhs.node_; // compare different container is UB + } + + int Index() const { return node_; } + + private: + friend class List; + Node& GetNode() const { return list_->node_pool_[node_]; } + + int node_; + List* list_; + }; + + List(int reserved = 0) { + node_pool_.Reserve(reserved); + InitGuard(); + } + + iterator PushBack(const Value& value) { + int node = node_pool_.Allocate(); + XGRAMMAR_DCHECK(0 < node && node < node_pool_.Size()); + node_pool_[node].value = value; + LinkBefore(node, 0); + return iterator(node, *this); + } + + void MoveBack(int node) { + XGRAMMAR_DCHECK(0 < node && node < node_pool_.Size()); + Unlink(node); + LinkBefore(node, 0); + } + + iterator Erase(iterator it) { + int node = it.Index(); + XGRAMMAR_DCHECK(0 < node && node < node_pool_.Size()); + int next = node_pool_[node].next; + Unlink(node); + node_pool_.Deallocate(node); + return iterator(next, *this); + } + + void Clear() { + node_pool_.Clear(); + InitGuard(); + } + + iterator begin() { return iterator(node_pool_[0].next, *this); } + iterator end() { return iterator(0, *this); } + + private: + void InitGuard() { + int node_id = node_pool_.Allocate(); + XGRAMMAR_DCHECK(node_id == 0) << "node 0 should be reserved as guard node"; + node_pool_[0].prev = 0; + node_pool_[0].next = 0; + } + + void LinkBefore(int node, int next) { + int prev = node_pool_[next].prev; + node_pool_[node].prev = prev; + node_pool_[node].next = next; + node_pool_[prev].next = node; + node_pool_[next].prev = node; + } + + void Unlink(int node) { + int prev = node_pool_[node].prev; + int next = node_pool_[node].next; + node_pool_[prev].next = next; + node_pool_[next].prev = prev; + } + + details::NodePool node_pool_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_CONTAINER_H_ diff --git a/third_party/xgrammar/cpp/support/cpptrace.h b/third_party/xgrammar/cpp/support/cpptrace.h new file mode 100644 index 0000000000..c1d7b7a415 --- /dev/null +++ b/third_party/xgrammar/cpp/support/cpptrace.h @@ -0,0 +1,39 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/cpptrace.h + * \details This file is an encapsulation of the cpptrace library. It helps debugging. This file + * takes effect only when XGRAMMAR_ENABLE_CPPTRACE is set to 1, and only support Linux and + * RelWithDebugInfo or Debug build. + */ +#ifndef XGRAMMAR_SUPPORT_CPPTRACE_H_ +#define XGRAMMAR_SUPPORT_CPPTRACE_H_ + +#if XGRAMMAR_ENABLE_CPPTRACE == 1 +#include +#endif + +#include + +namespace xgrammar { + +#if XGRAMMAR_ENABLE_CPPTRACE == 1 + +// Flag to check if cpptrace feature is enabled +static constexpr bool CPPTRACE_ENABLED = true; + +inline void PrintTrace() { cpptrace::generate_trace().print(); } +inline std::string GetTraceString() { return cpptrace::generate_trace().to_string(true); } + +#else + +static constexpr bool CPPTRACE_ENABLED = false; + +// Provide empty implementation when cpptrace is disabled +inline void PrintTrace() {} +inline std::string GetTraceString() { return ""; } + +#endif + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_CPPTRACE_H_ diff --git a/third_party/xgrammar/cpp/support/dynamic_bitset.h b/third_party/xgrammar/cpp/support/dynamic_bitset.h new file mode 100644 index 0000000000..592a612398 --- /dev/null +++ b/third_party/xgrammar/cpp/support/dynamic_bitset.h @@ -0,0 +1,363 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/dynamic_bitset.h + * \brief The header for utilities used in grammar-guided generation. + */ +#ifndef XGRAMMAR_SUPPORT_DYNAMIC_BITSET_H_ +#define XGRAMMAR_SUPPORT_DYNAMIC_BITSET_H_ + +#include + +#include +#include +#include +#include +#include + +// For __popcnt +#ifdef _MSC_VER +#include +#endif + +#include "json_serializer.h" +#include "logging.h" + +namespace xgrammar { + +/*! + * \brief A bitset whose length is specified at runtime. Note the size cannot be changed after + * construction. + * \details The buffer of the bitset is a uint32_t array. There are two uses for this class: + * - When passing nullptr to data, it maintains an internal buffer for the bitset. + * - When passing a pointer to a buffer with enough size, it uses the external buffer for the + * bitset. + * \details Part of the implementation is adopted from Boost::dynamic_bitset. + */ +class DynamicBitset { + public: + /*! + * \brief Calculate the minimal size of the uint32_t buffer for the bitset with the given size. + * \param element_size The size of the bitset. + * \return The minimal buffer size. + */ + static int GetBufferSize(int element_size) { return (element_size + 31) / 32; } + + /*! + * \brief Construct a empty bitset. This object should be assigned to a valid bitset before using. + */ + DynamicBitset() : size_(0), buffer_size_(0), data_(nullptr), is_internal_(true) {} + + /*! + * \brief Construct a bitset with the given size. + * \param size The size of the bitset. + * \param data The buffer for the bitset. If nullptr, the bitset will maintain an internal buffer. + */ + DynamicBitset(int size, uint32_t* data = nullptr) + : size_(size), buffer_size_(GetBufferSize(size)) { + if (data == nullptr) { + internal_buffer_.resize(buffer_size_, 0); + data_ = internal_buffer_.data(); + is_internal_ = true; + } else { + data_ = data; + is_internal_ = false; + } + } + + /*! \brief Copy constructor. Copy the buffer and manage the memory internally. */ + DynamicBitset(const DynamicBitset& other) + : size_(other.size_), + buffer_size_(other.buffer_size_), + data_(), + internal_buffer_(), + is_internal_(other.is_internal_) { + if (other.is_internal_) { + // copy the internal buffer + internal_buffer_ = other.internal_buffer_; + data_ = internal_buffer_.data(); + } else { + // simply point to the same external buffer + data_ = other.data_; + } + } + + /*! \brief Move constructor. Reset other and take ownership of its buffer. */ + DynamicBitset(DynamicBitset&& other) noexcept + : size_(std::exchange(other.size_, 0)), + buffer_size_(std::exchange(other.buffer_size_, 0)), + data_(std::exchange(other.data_, nullptr)), + internal_buffer_(std::move(other.internal_buffer_)), + is_internal_(std::exchange(other.is_internal_, true)) {} + + /*! \brief Copy assignment. */ + DynamicBitset& operator=(const DynamicBitset& other) { + XGRAMMAR_DCHECK(is_internal_ || size_ >= other.size_) + << "Expanding bitset size is not allowed when the " + "memory of the bitset is externally managed"; + size_ = other.size_; + buffer_size_ = other.buffer_size_; + if (is_internal_) { + internal_buffer_.resize(buffer_size_); + data_ = internal_buffer_.data(); + } + if (data_ != other.data_) { + std::memcpy(data_, other.data_, buffer_size_ * sizeof(uint32_t)); + } + return *this; + } + + /*! \brief Move assignment. */ + DynamicBitset& operator=(DynamicBitset&& other) noexcept { + size_ = other.size_; + buffer_size_ = other.buffer_size_; + is_internal_ = other.is_internal_; + if (is_internal_) { + internal_buffer_ = std::move(other.internal_buffer_); + data_ = internal_buffer_.data(); + } else { + data_ = other.data_; + } + return *this; + } + + /*! \brief Get the value of the bit at the given index. */ + bool operator[](int index) const { + XGRAMMAR_DCHECK(data_ && index >= 0 && index < size_); + return (data_[index / 32] >> (index % 32)) & 1; + } + + /*! \brief Get the size of the bitset. */ + int Size() const { return size_; } + + /*! \brief Set the whole bitset to true. */ + void Set() { + XGRAMMAR_DCHECK(data_); + std::memset(data_, 0xFF, buffer_size_ * sizeof(uint32_t)); + } + + /*! \brief Set the bit at the given index to the given value. */ + void Set(int index, bool value = true) { + XGRAMMAR_DCHECK(data_ && index >= 0 && index < size_); + if (value) { + data_[index / 32] |= 1 << (index % 32); + } else { + data_[index / 32] &= ~(1 << (index % 32)); + } + } + + /*! \brief Set the whole bitset to false. */ + void Reset() { + XGRAMMAR_DCHECK(data_); + std::memset(data_, 0, buffer_size_ * sizeof(uint32_t)); + } + + /*! \brief Set the bit at the given index to false. */ + void Reset(int index) { Set(index, false); } + + /*! \brief Perform a bitwise OR operation between the current bitset and another bitset. */ + DynamicBitset& operator|=(const DynamicBitset& other) { + XGRAMMAR_DCHECK(buffer_size_ <= other.buffer_size_); + for (int i = 0; i < buffer_size_; ++i) { + data_[i] |= other.data_[i]; + } + return *this; + } + + int FindFirstOne() const { return DoFindOneFrom(0); } + + int FindNextOne(int pos) const { + if (pos >= size_ - 1 || size_ == 0) return -1; + ++pos; + int blk = pos / BITS_PER_BLOCK; + int ind = pos % BITS_PER_BLOCK; + uint32_t fore = data_[blk] >> ind; + int result = fore ? pos + LowestBit(fore) : DoFindOneFrom(blk + 1); + return result < size_ ? result : -1; + } + + int FindFirstZero() const { return DoFindZeroFrom(0); } + + int FindNextZero(int pos) const { + if (pos >= size_ - 1 || size_ == 0) return -1; + ++pos; + int blk = pos / BITS_PER_BLOCK; + int ind = pos % BITS_PER_BLOCK; + uint32_t fore = (~data_[blk]) >> ind; + int result = fore ? pos + LowestBit(fore) : DoFindZeroFrom(blk + 1); + return result < size_ ? result : -1; + } + + int Count() const { + int count = 0; + for (int i = 0; i < buffer_size_; ++i) { + count += PopCount(data_[i]); + } + return count; + } + + bool All() const { + if (size_ == 0) return true; + // Check all complete blocks except the last one + for (int i = 0; i < buffer_size_ - 1; ++i) { + if (data_[i] != ~static_cast(0)) { + return false; + } + } + // For the last block, create a mask for valid bits only + int remaining_bits = size_ % BITS_PER_BLOCK; + uint32_t last_block_mask = remaining_bits ? (static_cast(1) << remaining_bits) - 1 + : ~static_cast(0); + return (data_[buffer_size_ - 1] & last_block_mask) == last_block_mask; + } + + bool Any() const { + if (size_ == 0) return false; + // Check all complete blocks except the last one + for (int i = 0; i < buffer_size_ - 1; ++i) { + if (data_[i] != 0) { + return true; + } + } + // For the last block, only consider the valid bits + int remaining_bits = size_ % BITS_PER_BLOCK; + uint32_t last_block_mask = remaining_bits ? (static_cast(1) << remaining_bits) - 1 + : ~static_cast(0); + return (data_[buffer_size_ - 1] & last_block_mask) != 0; + } + + static constexpr int BITS_PER_BLOCK = 32; + + friend std::size_t MemorySize(const DynamicBitset& bitset) { + return bitset.buffer_size_ * sizeof(bitset.data_[0]); + } + + friend picojson::value SerializeJSONValue(const DynamicBitset& bitset) { + XGRAMMAR_DCHECK(bitset.buffer_size_ == GetBufferSize(bitset.size_)); + picojson::array result; + result.reserve(2 + bitset.buffer_size_); + result.emplace_back(picojson::value(static_cast(bitset.size_))); + result.emplace_back(picojson::value(static_cast(bitset.buffer_size_))); + for (int i = 0; i < bitset.buffer_size_; ++i) { + result.emplace_back(picojson::value(static_cast(bitset.data_[i]))); + } + return picojson::value(std::move(result)); + } + + friend std::optional DeserializeJSONValue( + DynamicBitset* bitset, const picojson::value& value, const std::string& type_name + ) { + if (!value.is()) { + return ConstructDeserializeError("Expect an array", type_name); + } + const auto& arr = value.get(); + if (arr.size() < 2) { + return ConstructDeserializeError("Except at least 2 elements in the array", type_name); + } + if (!arr[0].is()) { + return ConstructDeserializeError("Expect an integer for size", type_name); + } + int size = static_cast(arr[0].get()); + if (!arr[1].is()) { + return ConstructDeserializeError("Expect an integer for buffer_size", type_name); + } + int buffer_size = static_cast(arr[1].get()); + if (size < 0 || buffer_size != GetBufferSize(size)) { + return ConstructDeserializeError( + "Invalid buffer_size. Buffer size should be ceil(size / 32)", type_name + ); + } + if (static_cast(arr.size()) != static_cast(buffer_size) + 2) { + return ConstructDeserializeError( + "Expect exactly buffer_size + 2 elements in the array", type_name + ); + } + + DynamicBitset result(size); + for (int i = 0; i < buffer_size; ++i) { + if (!arr[i + 2].is()) { + return ConstructDeserializeError("Expect an integer in the array", type_name); + } + int64_t value = arr[i + 2].get(); + if (value < 0 || value > std::numeric_limits::max()) { + return ConstructDeserializeError( + "Integer in the array is " + std::to_string(value) + " and out of the uint32_t range", + type_name + ); + } + result.data_[i] = static_cast(value); + } + *bitset = std::move(result); + return std::nullopt; + } + + bool operator==(const DynamicBitset& other) const { + if (size_ != other.size_) return false; + if (buffer_size_ != other.buffer_size_) return false; + for (int i = 0; i < buffer_size_; ++i) { + if (data_[i] != other.data_[i]) return false; + } + return true; + } + + private: + static int LowestBit(uint32_t value) { +#ifdef __GNUC__ + return __builtin_ctz(value); +#else // __GNUC__ + // From https://stackoverflow.com/a/757266 + static const int MultiplyDeBruijnBitPosition[32] = {0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, + 15, 25, 17, 4, 8, 31, 27, 13, 23, 21, 19, + 16, 7, 26, 12, 18, 6, 11, 5, 10, 9}; + return MultiplyDeBruijnBitPosition[((uint32_t)((value & -value) * 0x077CB531U)) >> 27]; +#endif // __GNUC__ + } + + static int PopCount(uint32_t value) { +#ifdef __GNUC__ + return __builtin_popcount(value); +#elif defined(_MSC_VER) + return __popcnt(value); +#else + XGRAMMAR_LOG(FATAL) << "PopCount is not supported on this platform"; +#endif + } + + int DoFindZeroFrom(int first_block) const { + int position = -1; + for (int i = first_block; i < buffer_size_; ++i) { + if (data_[i] != ~static_cast(0)) { + position = i; + break; + } + } + if (position == -1) return -1; + return position * BITS_PER_BLOCK + LowestBit(~data_[position]); + } + + int DoFindOneFrom(int first_block) const { + int position = -1; + for (int i = first_block; i < buffer_size_; ++i) { + if (data_[i] != 0) { + position = i; + break; + } + } + if (position == -1) return -1; + return position * BITS_PER_BLOCK + LowestBit(data_[position]); + } + + // The size of the bitset. + int size_; + // The size of the buffer. + int buffer_size_; + // The buffer for the bitset. + uint32_t* data_; + // The internal buffer. It is empty if not needed. + std::vector internal_buffer_; + // Whether the buffer is internally managed. + bool is_internal_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_DYNAMIC_BITSET_H_ diff --git a/third_party/xgrammar/cpp/support/encoding.h b/third_party/xgrammar/cpp/support/encoding.h new file mode 100644 index 0000000000..195fde5228 --- /dev/null +++ b/third_party/xgrammar/cpp/support/encoding.h @@ -0,0 +1,459 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/encoding.h + * \brief Encoding and decoding from/to UTF-8 and escape sequence to/from codepoints. + */ +#ifndef XGRAMMAR_SUPPORT_ENCODING_H_ +#define XGRAMMAR_SUPPORT_ENCODING_H_ +// TODO(yixin): enhance performance + +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +/*! \brief Represents a unicode codepoint. */ +using TCodepoint = int32_t; + +/*! + * \brief Represents an error when handling characters. Will be returned as a special TCodepoint + * value. + */ +enum CharHandlingError : TCodepoint { + /*! \brief The UTF-8 string is invalid. */ + kInvalidUTF8 = -10, + /*! \brief The escape sequence is invalid. */ + kInvalidEscape = -11, + /*! \brief The Latin-1 string is invalid. */ + kInvalidLatin1 = -12, +}; + +/******************** UTF-8 Handling ********************/ + +/*! + * \brief Print a codepoint to a UTF-8 string. + * \param codepoint The codepoint. + * \return The UTF-8 string. + */ +std::string CharToUTF8(TCodepoint codepoint); + +/*! + * \brief Handle the utf-8 first byte. + * \returns (is_valid, total_number_of_bytes, initial_codepoint). + */ +std::tuple HandleUTF8FirstByte(uint8_t byte); + +/*! + * \brief Parse all codepoints in a UTF-8 string. + * \param utf8 The UTF-8 string. + * \param perserve_invalid_bytes If the invalid UTF8 bytes will be preserved in the result. + * \return All codepoints. If the UTF-8 string is invalid, when perserve_invalid_bytes is false, + * the invalid bytes will be added to the result as a TCodepoint. Otherwise, the function will + * return {CharHandlingError::kInvalidUTF8}. + */ +std::vector ParseUTF8(const char* utf8, bool perserve_invalid_bytes = false); + +/*! + * \brief Parse the first codepoint in a UTF-8 string. + * \param utf8 The UTF-8 string. + * \return The codepoint and the number of bytes consumed. If the UTF-8 string is invalid, return + * {CharHandlingError::kInvalidUTF8, 0}. + */ +std::pair ParseNextUTF8(const char* utf8); + +/*! + * \brief Convert a Latin-1 string to a byte sequence. + * \param latin1 The Latin-1 string. + * \return The byte sequence. + */ +std::optional Latin1ToBytes(const std::string& latin1, std::string* result); + +/******************** Escape Handling ********************/ + +/*! + * \brief Convert a codepoint to a escaped string. If the codepoint is not printable, it will be + * escaped. By default the function support escape sequences in C ("\n", "\t", "\u0123"). User + * can specify more escape sequences using additional_escape_map. + * \param codepoint The codepoint. + * \param additional_escape_map A map from codepoint to escape sequence. If the codepoint is in + * the map, it will be escaped using the corresponding escape sequence. e.g. {{'-', "\\-"}}. + * \return The printable string. + */ +std::string EscapeString( + TCodepoint codepoint, + const std::unordered_map& additional_escape_map = {} +); + +/*! + * \brief Convert the given char to a escaped string that can be printed. + * \return The escaped string. + */ +std::string EscapeString(uint8_t raw_char); + +/*! + * \brief Convert the given string to a escaped string that can be printed. + * \return The escaped string. + */ +std::string EscapeString(std::string raw_str); + +/*! + * \brief Convert a hex character to an integer. + * \param c The hex character: 0-9, a-f, A-F. + * \return The integer value of the hex character. If the character is not a valid hex character, + * return -1. + */ +int HexCharToInt(char c); + +/*! + * \brief Parse the first escaped codepoint from a escaped string. data must start with a '\' + * character. + * \param data The escaped string. Can be TCodepoint* (e.g. string decoded from UTF-8) or char*. + * \param additional_escape_map A map from escape sequence to codepoint. If the escape sequence is + * in the map, it will be converted to the corresponding codepoint. e.g. {{"\\-", '-'}}. + * \return The codepoint and the number of bytes consumed. + */ +template +std::pair ParseNextEscaped( + const CharType* data, const std::unordered_map& additional_escape_map = {} +); + +/*! + * \brief Parse the first codepoint from a UTF-8 string. Also checks escape sequences and converts + * the escaped char to its original value. + * \param utf8 The UTF-8 string or the escape sequence. + * \param additional_escape_map A map from escape sequence to codepoint. If the escape sequence is + * in the map, it will be converted to the corresponding codepoint. e.g. {{"\\-", '-'}}. + * \return The codepoint and the number of bytes consumed. If the UTF-8 string is invalid, the + * function returns (CharHandlingError::kInvalidUTF8, 0). If the escape sequence is invalid, the + * function returns (CharHandlingError::kInvalidEscape, 0). + */ +std::pair ParseNextUTF8OrEscaped( + const char* utf8, const std::unordered_map& additional_escape_map = {} +); + +/******************** Implementation ********************/ + +inline std::string CharToUTF8(TCodepoint codepoint) { + XGRAMMAR_DCHECK(codepoint <= 0x10FFFF) << "Invalid codepoint: " << codepoint; + std::string utf8; + if (codepoint <= 0x7F) { + // 1-byte sequence + utf8 += static_cast(codepoint); + } else if (codepoint <= 0x7FF) { + // 2-byte sequence + utf8 += static_cast(0xC0 | ((codepoint >> 6) & 0x1F)); + utf8 += static_cast(0x80 | (codepoint & 0x3F)); + } else if (codepoint <= 0xFFFF) { + // 3-byte sequence + utf8 += static_cast(0xE0 | ((codepoint >> 12) & 0x0F)); + utf8 += static_cast(0x80 | ((codepoint >> 6) & 0x3F)); + utf8 += static_cast(0x80 | (codepoint & 0x3F)); + } else { + // 4-byte sequence + utf8 += static_cast(0xF0 | ((codepoint >> 18) & 0x07)); + utf8 += static_cast(0x80 | ((codepoint >> 12) & 0x3F)); + utf8 += static_cast(0x80 | ((codepoint >> 6) & 0x3F)); + utf8 += static_cast(0x80 | (codepoint & 0x3F)); + } + return utf8; +} + +inline std::tuple HandleUTF8FirstByte(uint8_t byte) { + static const std::array kFirstByteMask = {0x00, 0x7F, 0x1F, 0x0F, 0x07}; + // clang-format off + static const std::array kUtf8Bytes = { + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, + 4, 4, 4, 4, 4, 4, 4, 4, -1, -1, -1, -1, -1, -1, -1, -1, + }; + // clang-format on + auto num_bytes = kUtf8Bytes[static_cast(byte)]; + if (num_bytes == -1) { + return {false, 0, 0}; + } + return {true, num_bytes, byte & kFirstByteMask[num_bytes]}; +} + +inline std::pair ParseNextUTF8(const char* utf8) { + auto [accepted, num_bytes, res] = HandleUTF8FirstByte(utf8[0]); + if (accepted) { + for (int i = 1; i < num_bytes; ++i) { + if (utf8[i] == 0 || (static_cast(utf8[i]) & 0xC0) != 0x80) { + // invalid utf8 + accepted = false; + break; + } + res = (res << 6) | (static_cast(utf8[i]) & 0x3F); + } + } + + if (!accepted) { + // invalid utf8 + return {CharHandlingError::kInvalidUTF8, 0}; + } + + return {res, num_bytes}; +} + +inline std::vector ParseUTF8(const char* utf8, bool perserve_invalid_bytes) { + std::vector codepoints; + while (*utf8 != 0) { + auto [codepoint, num_bytes] = ParseNextUTF8(utf8); + if (codepoint == CharHandlingError::kInvalidUTF8) { + if (perserve_invalid_bytes) { + codepoints.push_back(static_cast(static_cast(utf8[0]))); + utf8 += 1; + continue; + } else { + return {CharHandlingError::kInvalidUTF8}; + } + } + codepoints.push_back(codepoint); + utf8 += num_bytes; + } + return codepoints; +} + +/*! + \brief Convert a Latin-1 string to a byte sequence. + \param latin1 The Latin-1 string. + \param result The output byte sequence. + The function will convert each Latin-1 character to its corresponding byte(s). + For characters in the range [0x00, 0x7F], the corresponding byte is the same as the character. + Otherwise, the character should be encoded in two bytes in UTF-8: + - First byte: 110xxxxx (0xC0 | (char >> 6)) + - Second byte: 10xxxxxx (0x80 | (char & 0x3F)) + Example: + 0xC3 0xBF -> 0xFF + 'A' -> 'A' + \return std::nullopt if the conversion is successful. Otherwise, return + CharHandlingError::kInvalidLatin1 if the Latin-1 string is invalid. +*/ +inline std::optional Latin1ToBytes( + const std::string& latin1, std::string* result +) { + result->clear(); + result->reserve(latin1.size()); + + const size_t len = latin1.size(); + for (size_t i = 0; i < len; ++i) { + unsigned char c1 = static_cast(latin1[i]); + if (c1 < 0x80) { + result->push_back(static_cast(c1)); + } else { + if (i + 1 >= len) { + return CharHandlingError::kInvalidLatin1; + } + + unsigned char c2 = static_cast(latin1[i + 1]); + if ((c2 & 0xC0) != 0x80) { + return CharHandlingError::kInvalidLatin1; + } + + int code = ((c1 & 0x1F) << 6) | (c2 & 0x3F); + if (code < 0x80 || code > 0xFF) { + return CharHandlingError::kInvalidLatin1; + } + + result->push_back(static_cast(code)); + ++i; + } + } + + return std::nullopt; +} + +/*! + \brief Convert a byte sequence to a Latin-1 string. + \param Bytes The input byte sequence. + \param result The output Latin-1 string. + The function will convert each byte in the input to a Latin-1 character. + For bytes in the range [0x00, 0x7F], the corresponding Latin-1 character is the same as the byte. + For bytes in the range [0x80, 0xFF], the corresponding Latin-1 character is represented by two + bytes in UTF-8: + - First byte: 110xxxxx (0xC0 | (byte >> 6)) + - Second byte: 10xxxxxx (0x80 | (byte & 0x3F)) + Example: + 0xFF -> 0xC3 0xBF + 'A' -> 'A' +*/ +inline void ByteToLatin1(const std::string& bytes, std::string* result) { + result->clear(); + result->reserve(bytes.size()); + + for (unsigned char current_char : bytes) { + // Ascii character, directly add to result. + if (current_char <= 0x7F) { + result->push_back(static_cast(current_char)); + continue; + } + + // not Ascii character, convert to Latin-1. + unsigned char latin1_first_byte = 0; + unsigned char latin1_second_byte = 0; + + latin1_first_byte = 0xC0 | (current_char >> 6); + latin1_second_byte = 0x80 | (current_char & 0x3F); + result->push_back(static_cast(latin1_first_byte)); + result->push_back(static_cast(latin1_second_byte)); + } +} + +inline int HexCharToInt(char c) { + if (c >= '0' && c <= '9') { + return c - '0'; + } else if (c >= 'a' && c <= 'f') { + return c - 'a' + 10; + } else if (c >= 'A' && c <= 'F') { + return c - 'A' + 10; + } else { + return -1; + } +} + +inline std::string EscapeString( + TCodepoint codepoint, const std::unordered_map& additional_escape_map +) { + static const std::unordered_map kCodepointToEscape = { + {'\'', "\\\'"}, + {'\"', "\\\""}, + {'\?', "\\?"}, + {'\\', "\\\\"}, + {'\a', "\\a"}, + {'\b', "\\b"}, + {'\f', "\\f"}, + {'\n', "\\n"}, + {'\r', "\\r"}, + {'\t', "\\t"}, + {'\v', "\\v"}, + {'\0', "\\0"}, + {'\x1B', "\\e"} + }; + + if (auto it = additional_escape_map.find(codepoint); it != additional_escape_map.end()) { + return it->second; + } + + if (auto it = kCodepointToEscape.find(codepoint); it != kCodepointToEscape.end()) { + return it->second; + } + + if (codepoint >= 0x20 && codepoint <= 0x7E) { + return std::string({static_cast(codepoint)}); + } + + // convert codepoint to hex + char prefix = codepoint <= 0xFF ? 'x' : codepoint <= 0xFFFF ? 'u' : 'U'; + int width = codepoint <= 0xFF ? 2 : codepoint <= 0xFFFF ? 4 : 8; + std::stringstream ss; + ss << std::setfill('0') << std::setw(width) << std::hex << codepoint; + auto hex = ss.str(); + return std::string("\\") + prefix + hex; +} + +inline std::string EscapeString(uint8_t raw_char) { + return EscapeString(static_cast(raw_char)); +} + +inline std::string EscapeString(std::string raw_str) { + std::string res; + auto codepoints = ParseUTF8(raw_str.c_str(), true); + for (auto c : codepoints) { + res += EscapeString(c); + } + return res; +} + +template +std::pair ParseNextEscaped( + const CharType* data, const std::unordered_map& additional_escape_map +) { + // C escape characters + static const std::unordered_map kEscapeToCodepoint = { + // clang-format off + {'\'', '\''}, {'\"', '\"'}, {'?', '\?'}, {'\\', '\\'}, {'/', '/'}, {'a', '\a'}, + {'b', '\b'}, {'f', '\f'}, {'n', '\n'}, {'r', '\r'}, {'t', '\t'}, {'v', '\v'}, + {'0', '\0'}, {'e', '\x1B'} // clang-format on + }; + if (data[0] != '\\') { + return {CharHandlingError::kInvalidEscape, 0}; + } + + bool escape_char_in_escape_range = + static_cast(static_cast(data[1])) <= 128; + if (!escape_char_in_escape_range) { + return {CharHandlingError::kInvalidEscape, 0}; + } + + if (auto it = additional_escape_map.find(static_cast(data[1])); + it != additional_escape_map.end()) { + return {it->second, 2}; + } + if (auto it = kEscapeToCodepoint.find(static_cast(data[1])); + it != kEscapeToCodepoint.end()) { + return {it->second, 2}; + } + + if (data[1] == 'x') { + // arbitrary length hex + int len = 0; + TCodepoint codepoint = 0; + int32_t digit; + while ((digit = HexCharToInt(data[2 + len])) != -1) { + codepoint = codepoint * 16 + digit; + ++len; + } + if (len == 0) { + return {CharHandlingError::kInvalidEscape, 0}; + } + return {codepoint, len + 2}; + } else if (data[1] == 'u' || data[1] == 'U') { + // 4- or 8-digit hex + int len = data[1] == 'u' ? 4 : 8; + TCodepoint codepoint = 0; + + for (int i = 0; i < len; ++i) { + auto digit = HexCharToInt(data[i + 2]); + if (digit == -1) { + return {CharHandlingError::kInvalidEscape, 0}; + } + codepoint = codepoint * 16 + digit; + } + return {codepoint, len + 2}; + } else { + return {CharHandlingError::kInvalidEscape, 0}; + } +} + +inline std::pair ParseNextUTF8OrEscaped( + const char* utf8, const std::unordered_map& additional_escape_map +) { + if (utf8[0] != '\\') { + return ParseNextUTF8(utf8); + } + return ParseNextEscaped(utf8, additional_escape_map); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_ENCODING_H_ diff --git a/third_party/xgrammar/cpp/support/int_set.h b/third_party/xgrammar/cpp/support/int_set.h new file mode 100644 index 0000000000..f37f9a83ab --- /dev/null +++ b/third_party/xgrammar/cpp/support/int_set.h @@ -0,0 +1,132 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/int_set.h + * \brief The header for utilities used in grammar-guided generation. + */ +#ifndef XGRAMMAR_SUPPORT_INT_SET_H_ +#define XGRAMMAR_SUPPORT_INT_SET_H_ + +#include +#include +#include +#include + +namespace xgrammar { + +/*! + * \brief Let lhs be the union of lhs and rhs. Suppose that both sets are sorted. + * \note No additional vectors are allocated, and the time complexity is O(n) + */ +inline void IntsetUnion(std::vector* lhs, const std::vector& rhs) { + int original_lhs_size = lhs->size(); + int rhs_size = rhs.size(); + + lhs->resize(original_lhs_size + rhs_size); + + auto it_lhs = lhs->rbegin() + rhs_size; + auto it_rhs = rhs.rbegin(); + auto it_result = lhs->rbegin(); + + while (it_lhs != lhs->rend() && it_rhs != rhs.rend()) { + if (*it_lhs > *it_rhs) { + *it_result = *it_lhs; + ++it_lhs; + } else if (*it_lhs < *it_rhs) { + *it_result = *it_rhs; + ++it_rhs; + } else { + *it_result = *it_lhs; + ++it_lhs; + ++it_rhs; + } + ++it_result; + } + + while (it_rhs != rhs.rend()) { + *it_result = *it_rhs; + ++it_result; + ++it_rhs; + } + + auto last = std::unique(lhs->begin(), lhs->end()); + lhs->erase(last, lhs->end()); +} + +/*! + * \brief Let lhs be the intersection of lhs and rhs. Suppose that both sets are sorted. + * \note No additional vector is allocated, and the time complexity is O(n). + * \note Support the case where lhs is the universal set by setting lhs to {-1}. The result will be + * rhs then. + */ +inline void IntsetIntersection(std::vector* lhs, const std::vector& rhs) { + if (lhs->size() == 1 && (*lhs)[0] == -1) { + *lhs = rhs; + return; + } + + auto it_lhs = lhs->begin(); + auto it_rhs = rhs.begin(); + auto it_result = lhs->begin(); + + while (it_lhs != lhs->end() && it_rhs != rhs.end()) { + if (*it_lhs < *it_rhs) { + ++it_lhs; + } else if (*it_lhs > *it_rhs) { + ++it_rhs; + } else { + *it_result = *it_lhs; + ++it_lhs; + ++it_rhs; + ++it_result; + } + } + lhs->erase(it_result, lhs->end()); +} + +/*! + * \brief Let lhs = lhs - rhs. Both sets must be sorted. + * \note In-place, no additional vector allocated, O(n) time. + */ +inline void IntsetDifference(std::vector* lhs, const std::vector& rhs) { + auto it_lhs = lhs->begin(); + auto it_rhs = rhs.begin(); + auto it_result = lhs->begin(); + + while (it_lhs != lhs->end() && it_rhs != rhs.end()) { + if (*it_lhs < *it_rhs) { + *it_result++ = *it_lhs++; + } else if (*it_lhs > *it_rhs) { + ++it_rhs; + } else { + ++it_lhs; + ++it_rhs; + } + } + while (it_lhs != lhs->end()) { + *it_result++ = *it_lhs++; + } + lhs->erase(it_result, lhs->end()); +} + +/*! + * \brief Compute result = [0, n) - excluded. excluded must be sorted with values in [0, n). + * \note O(n) time. + */ +inline void IntsetComplement( + std::vector* result, int32_t n, const std::vector& excluded +) { + result->clear(); + result->reserve(n - static_cast(excluded.size())); + auto it = excluded.begin(); + for (int32_t i = 0; i < n; ++i) { + if (it != excluded.end() && *it == i) { + ++it; + } else { + result->push_back(i); + } + } +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_INT_SET_H_ diff --git a/third_party/xgrammar/cpp/support/json_parse.h b/third_party/xgrammar/cpp/support/json_parse.h new file mode 100644 index 0000000000..3deba6916d --- /dev/null +++ b/third_party/xgrammar/cpp/support/json_parse.h @@ -0,0 +1,118 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/json_parse.h + * \brief picojson::parse wrappers whose nesting depth is bounded by RecursionGuard. + * picojson recurses once per nesting level and has no depth limit of its own, so a deeply nested + * input would overflow the stack before any of the recursion guards in xgrammar run. + */ +#ifndef XGRAMMAR_SUPPORT_JSON_PARSE_H_ +#define XGRAMMAR_SUPPORT_JSON_PARSE_H_ + +#include + +#include +#include + +#include "logging.h" +#include "recursion_guard.h" + +namespace xgrammar { + +namespace detail { + +/*! + * \brief A picojson parse context that behaves like picojson::default_parse_context, but counts + * every nested array or object in a RecursionGuard. The maximum recursion depth of xgrammar thus + * applies while parsing, and exceeding it throws before the stack overflows. + */ +class DepthGuardedParseContext { + public: + DepthGuardedParseContext(picojson::value* out, int* depth) : out_(out), depth_(depth) {} + DepthGuardedParseContext(const DepthGuardedParseContext&) = delete; + DepthGuardedParseContext& operator=(const DepthGuardedParseContext&) = delete; + + bool set_null() { + *out_ = picojson::value(); + return true; + } + bool set_bool(bool b) { + *out_ = picojson::value(b); + return true; + } + bool set_int64(int64_t i) { + *out_ = picojson::value(i); + return true; + } + bool set_number(double f) { + *out_ = picojson::value(f); + return true; + } + template + bool parse_string(picojson::input& in) { + *out_ = picojson::value(picojson::string_type, false); + return picojson::_parse_string(out_->get(), in); + } + bool parse_array_start() { + *out_ = picojson::value(picojson::array_type, false); + return true; + } + template + bool parse_array_item(picojson::input& in, size_t) { + picojson::array& a = out_->get(); + a.push_back(picojson::value()); + RecursionGuard guard(depth_); + DepthGuardedParseContext child(&a.back(), depth_); + return picojson::_parse(child, in); + } + bool parse_array_stop(size_t) { return true; } + bool parse_object_start() { + *out_ = picojson::value(picojson::object_type, false); + return true; + } + template + bool parse_object_item(picojson::input& in, const std::string& key) { + picojson::object& o = out_->get(); + RecursionGuard guard(depth_); + DepthGuardedParseContext child(&o[key], depth_); + return picojson::_parse(child, in); + } + + private: + picojson::value* out_; + int* depth_; +}; + +} // namespace detail + +/*! + * \brief Parse the JSON value at the start of [begin, end). Same as + * picojson::parse(out, begin, end, err), but the nesting depth is bounded by the maximum + * recursion depth (see RecursionGuard). + * \return The iterator past the parsed value; begin if the depth limit was exceeded. + */ +template +inline Iter ParseJSON(picojson::value& out, Iter begin, Iter end, std::string* err) { + int depth = 0; + detail::DepthGuardedParseContext ctx(&out, &depth); + try { + return picojson::_parse(ctx, begin, end, err); + } catch (const LogFatalError& error) { + *err = error.what(); + return begin; + } +} + +/*! + * \brief Parse a JSON string. Same as picojson::parse(out, json), but the nesting depth is bounded + * by the maximum recursion depth (see RecursionGuard). + * \return The error message, empty on success. + */ +inline std::string ParseJSON(picojson::value& out, const std::string& json) { + std::string err; + ParseJSON(out, json.begin(), json.end(), &err); + return err; +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_JSON_PARSE_H_ diff --git a/third_party/xgrammar/cpp/support/json_serializer.h b/third_party/xgrammar/cpp/support/json_serializer.h new file mode 100644 index 0000000000..aaff03cccf --- /dev/null +++ b/third_party/xgrammar/cpp/support/json_serializer.h @@ -0,0 +1,640 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/json_serializer.h + * \brief A JSON-based serializer. Automatically generates serialization and deserialization logic + * from reflection. + */ +#ifndef XGRAMMAR_SUPPORT_JSON_SERIALIZER_H_ +#define XGRAMMAR_SUPPORT_JSON_SERIALIZER_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "encoding.h" +#include "json_parse.h" +#include "logging.h" +#include "reflection.h" +#include "utils.h" +#include "xgrammar/exception.h" +#include "xgrammar/object.h" + +namespace xgrammar { + +/******************** Interfaces ********************/ + +/*! + * \brief Manages the version of the serialized object. The version will be added to the serialized + * object, and during deserialization, the object's version must match the current serialization + * version in xgrammar. + */ +class SerializeVersion { + public: + /*! + * \brief Returns the current serialization version. + */ + static std::string_view GetVersion() { return kXGrammarSerializeVersion; } + + /*! + * \brief Adds the version info to the serialized object. + */ + static void Apply(picojson::object* object); + + /*! + * \brief Checks if the serialized object's version matches the current serialization version. + * \return An error if the version does not exist or does not match. + */ + static std::optional Check(const picojson::object& object); + + private: + /*! + * \brief The key of the version info in the serialized object. + */ + static constexpr const char kXGrammarSerializeVersionKey[] = "__VERSION__"; + + /*! + * \brief The current serialization version. When the serialization result of any object in + * XGrammar is changed, this version should be bumped. + */ + static constexpr const char kXGrammarSerializeVersion[] = "v16"; +}; + +/*! + * \brief Serializes a value to a JSON value. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * SerializeJSONValue function. For reflection-based types, the serialization logic is automatically + * generated from the defined members. + * \param value The value to be serialized. + * \return The serialized JSON value. + */ +template +picojson::value AutoSerializeJSONValue(const T& value); + +/*! + * \brief Serializes a value to a JSON string. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * SerializeJSONValue function. For reflection-based types, the serialization logic is automatically + * generated from the defined members. + * \param value The value to be serialized. + * \param add_version Whether to add the version info to the serialized object. The addition is + * valid only when the serialized result is an object. + * \return The serialized JSON string. + */ +template +std::string AutoSerializeJSON(const T& value, bool add_version = false); + +/*! + * \brief Deserializes a value from a JSON value. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * DeserializeJSONValue function. For reflection-based types, the deserialization logic is + * automatically generated from the defined members. If a reflection-based type also defines + * `std::optional Validate() const`, it is called after all members are deserialized, + * and a returned message is reported as a deserialization error. This is where a type checks the + * invariants (index ranges, CSR layout, ...) that the constructors normally guarantee. + * \param result The pointer to the result to be deserialized. + * \param value The JSON value to be deserialized. + * \param type_name The name of the type to be deserialized. Used for error message. + * \return The deserialization error if any. + */ +template +std::optional AutoDeserializeJSONValue( + T* result, const picojson::value& value, const std::string& type_name = "" +); + +/*! + * \brief Deserializes a value from a JSON string. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * DeserializeJSONValue function. For reflection-based types, the deserialization logic is + * automatically generated from the defined members. + * \param result The pointer to the result to be deserialized. + * \param json_string The JSON string to be deserialized. + * \param check_version Whether to check the version info in the serialized object. The check is + * valid only when the serialized object is an object. + * \param type_name The name of the type to be deserialized. Used for error message. + * \return The deserialization error if any. + */ +template +std::optional AutoDeserializeJSON( + T* result, + const std::string& json_string, + bool check_version = false, + const std::string& type_name = "" +); + +/*! + * \brief Constructs a deserialize error with the given error message and type name. + * \param error_message The error message. + * \param type_name The name of the type. + * \return The constructed runtime error. + */ +inline SerializationError ConstructDeserializeError( + const std::string& error_message, const std::string& type_name +); + +/******************** Implementations ********************/ + +inline void SerializeVersion::Apply(picojson::object* object) { + XGRAMMAR_DCHECK(object != nullptr); + XGRAMMAR_DCHECK(object->find(kXGrammarSerializeVersionKey) == object->end()); + (*object)[kXGrammarSerializeVersionKey] = picojson::value(std::string(GetVersion())); +} + +inline std::optional SerializeVersion::Check(const picojson::object& object) { + if (object.find(kXGrammarSerializeVersionKey) == object.end()) { + return DeserializeVersionError( + std::string("Missing version in serialized object: ") + kXGrammarSerializeVersionKey + ); + } + if (object.at(kXGrammarSerializeVersionKey).get() != GetVersion()) { + return DeserializeVersionError( + std::string("Wrong version in serialized object: Got ") + + object.at(kXGrammarSerializeVersionKey).get() + ", expected " + + std::string(GetVersion()) + ); + } + return std::nullopt; +} + +/******************** Template Implementations ********************/ + +namespace detail::json_serializer { + +template +struct has_serialize_json_global : std::false_type {}; + +template +struct has_serialize_json_global< + T, + std::void_t()))>> : std::true_type { + static_assert( + std::is_same_v())), picojson::value>, + "SerializeJSONValue must be a global function returning picojson::value" + ); +}; + +template +struct has_deserialize_json_global : std::false_type {}; + +template +struct has_deserialize_json_global< + T, + std::void_t(), picojson::value{}, std::string{}) + )>> : std::true_type { + static_assert( + std::is_same_v< + decltype(DeserializeJSONValue(std::declval(), picojson::value{}, std::string{})), + std::optional>, + "DeserializeJSONValue must be a global function returning std::optional" + ); + static_assert( + std::is_default_constructible_v, + "global deserializer can only apply to a default constructible type" + ); +}; + +template +struct has_validate : std::false_type {}; + +template +struct has_validate().Validate())>> + : std::true_type { + static_assert( + std::is_same_v().Validate()), std::optional>, + "Validate must be a const member function returning std::optional" + ); +}; + +/*! \brief Runs T::Validate() on a deserialized value if the type defines one. */ +template +inline std::optional ValidateDeserialized( + const T& value, const std::string& type_name +) { + if constexpr (has_validate::value) { + if (auto error = value.Validate()) { + return ConstructDeserializeError(*error, type_name); + } + } + return std::nullopt; +} + +template +inline constexpr bool false_v = false; + +template +inline picojson::value TraitSerializeJSONValue(const T& value) { + using Functor = member_functor; + if constexpr (Functor::value == member_type::kConfig) { + if constexpr (Functor::has_names) { + // normal named struct + picojson::object obj; + obj.reserve(Functor::member_count); + visit_config([&](auto ptr, const char* name, std::size_t) { + XGRAMMAR_DCHECK(obj.find(name) == obj.end()); + obj[name] = AutoSerializeJSONValue(value.*ptr); + }); + return picojson::value(std::move(obj)); + } else if constexpr (Functor::member_count == 1) { + // optimize for single member unnamed structs + constexpr auto member_ptr = std::get<0>(Functor::members); + return AutoSerializeJSONValue(value.*member_ptr); + } else { + // normal unnamed struct + picojson::array arr; + arr.resize(Functor::member_count); + visit_config([&](auto ptr, const char*, std::size_t idx) { + arr[idx] = AutoSerializeJSONValue(value.*ptr); + }); + return picojson::value(std::move(arr)); + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Invalid trait type"); + return picojson::value{}; + } +} + +template +inline std::optional TraitDeserializeJSONValue( + T* result, const picojson::value& value, const std::string& type_name +) { + using Functor = member_functor; + if constexpr (Functor::value == member_type::kConfig) { + if constexpr (Functor::has_names) { + // normal named struct + if (!value.is()) { + return ConstructDeserializeError("Expect an object", type_name); + } + const auto& obj = value.get(); + std::optional err = std::nullopt; + visit_config([&](auto ptr, const char* name, std::size_t idx) { + if (err) { + return; + } else if (obj.find(name) == obj.end()) { + err = ConstructDeserializeError("Missing member " + std::string(name), type_name); + } else if (auto e = AutoDeserializeJSONValue(&(result->*ptr), obj.at(name), type_name)) { + err = e; + } + }); + return err; + } else if constexpr (Functor::member_count == 1) { + // optimize for single member unnamed structs + constexpr auto member_ptr = std::get<0>(Functor::members); + return AutoDeserializeJSONValue(&(result->*member_ptr), value, type_name); + } else { + // normal unnamed struct + if (!value.is()) { + return ConstructDeserializeError("Expect an array", type_name); + } + const auto& arr = value.get(); + if (arr.size() != Functor::member_count) { + return ConstructDeserializeError( + "Wrong number of elements in array: Expected " + std::to_string(Functor::member_count) + + ", but got " + std::to_string(arr.size()), + type_name + ); + } + std::optional err = std::nullopt; + visit_config([&](auto ptr, const char*, std::size_t idx) { + if (err) { + return; + } else if (auto e = AutoDeserializeJSONValue(&(result->*ptr), arr[idx], type_name)) { + err = e; + } + }); + return err; + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Invalid trait type"); + XGRAMMAR_UNREACHABLE(); + } +} + +/******************** Customized Serialization ********************/ + +template > +inline picojson::value AutoSerializeJSONValuePImpl(const T& value) { + if (value.IsNull()) return picojson::value{}; + return AutoSerializeJSONValue(*value.ImplPtr()); +} + +template > +inline std::optional AutoDeserializeJSONValuePImpl( + T* result, const picojson::value& value, const std::string& type_name +) { + XGRAMMAR_DCHECK(result->IsNull()); + if (value.is()) { + *result = T{NullObj{}}; + return std::nullopt; + } + auto ptr = std::make_shared(); + if (auto error = AutoDeserializeJSONValue(ptr.get(), value, type_name)) { + return error; + } + *result = T(std::move(ptr)); + return std::nullopt; +} + +} // namespace detail::json_serializer + +inline SerializationError ConstructDeserializeError( + const std::string& error_message, const std::string& type_name +) { + if (type_name.empty()) { + return DeserializeFormatError("Deserialize error: " + error_message); + } else { + return DeserializeFormatError("Deserialize error for type " + type_name + ": " + error_message); + } +} + +template +inline picojson::value AutoSerializeJSONValue(const T& value) { + if constexpr (detail::json_serializer::has_serialize_json_global::value) { + // User-defined SerializeJSONValue (highest priority) + return SerializeJSONValue(value); + } else if constexpr (is_pimpl_class::value) { + // Library-customized serialization methods + return detail::json_serializer::AutoSerializeJSONValuePImpl(value); + } else if constexpr (member_trait::value != member_type::kNone) { + // Trait serialization methods + return detail::json_serializer::TraitSerializeJSONValue(value); + } else if constexpr (std::is_same_v) { + // Below is primitive types + return picojson::value(value); + } else if constexpr (std::is_integral_v || std::is_enum_v) { + return picojson::value(static_cast(value)); + } else if constexpr (std::is_floating_point_v) { + return picojson::value(static_cast(value)); + } else if constexpr (std::is_same_v) { + std::string result; + ByteToLatin1(value, &result); + return picojson::value(result); + } else if constexpr (is_std_optional::value) { + if (value.has_value()) { + return AutoSerializeJSONValue(*value); + } else { + return picojson::value{}; + } + } else if constexpr (is_std_pair::value) { + // std::pair: serialize as an array of size 2 + picojson::array arr; + arr.resize(2); + arr[0] = AutoSerializeJSONValue(value.first); + arr[1] = AutoSerializeJSONValue(value.second); + return picojson::value(std::move(arr)); + } else if constexpr (is_std_vector::value) { + picojson::array arr; + arr.reserve(value.size()); + for (const auto& item : value) { + arr.push_back(AutoSerializeJSONValue(item)); + } + return picojson::value(std::move(arr)); + } else if constexpr (is_std_unordered_set::value) { + std::vector ptr_vec; + ptr_vec.reserve(value.size()); + for (const auto& item : value) { + ptr_vec.push_back(&item); + } + std::sort(ptr_vec.begin(), ptr_vec.end(), [](const auto* a, const auto* b) { return *a < *b; }); + picojson::array arr; + arr.reserve(value.size()); + for (const auto* ptr : ptr_vec) { + arr.push_back(AutoSerializeJSONValue(*ptr)); + } + return picojson::value(std::move(arr)); + } else if constexpr (is_std_unordered_map::value) { + if constexpr (std::is_same_v) { + // unordered_map: map to json object + picojson::object obj; + obj.reserve(value.size()); + for (const auto& item : value) { + obj[item.first] = AutoSerializeJSONValue(item.second); + } + return picojson::value(std::move(obj)); + } else { + // unordered_map (T1 is not string): map to json array of array of size 2 + std::vector ptr_vec; + ptr_vec.reserve(value.size()); + for (const auto& item : value) { + ptr_vec.push_back(&item); + } + std::sort(ptr_vec.begin(), ptr_vec.end(), [](const auto* a, const auto* b) { + return a->first < b->first; + }); + picojson::array arr; + arr.reserve(value.size()); + for (const auto* ptr : ptr_vec) { + const auto& [key, item] = *ptr; + picojson::array sub_arr{AutoSerializeJSONValue(key), AutoSerializeJSONValue(item)}; + arr.push_back(picojson::value(std::move(sub_arr))); + } + return picojson::value(std::move(arr)); + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Cannot serialize this type"); + XGRAMMAR_UNREACHABLE(); + } +} + +template +inline std::string AutoSerializeJSON(const T& value, bool add_version) { + picojson::value json_value = AutoSerializeJSONValue(value); + if (add_version) { + XGRAMMAR_DCHECK(json_value.is()); + SerializeVersion::Apply(&json_value.get()); + } + return picojson::value(json_value).serialize(); +} + +template +inline std::optional AutoDeserializeJSONValue( + T* result, const picojson::value& value, const std::string& type_name +) { + static_assert(!std::is_const_v, "Cannot deserialize into a const type"); + if constexpr (detail::json_serializer::has_deserialize_json_global::value) { + return DeserializeJSONValue(result, value, type_name); + } else if constexpr (is_pimpl_class::value) { + return detail::json_serializer::AutoDeserializeJSONValuePImpl(result, value, type_name); + } else if constexpr (member_trait::value != member_type::kNone) { + if (auto error = detail::json_serializer::TraitDeserializeJSONValue(result, value, type_name)) { + return error; + } + return detail::json_serializer::ValidateDeserialized(*result, type_name); + } else if constexpr (std::is_same_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect a boolean", type_name); + } + *result = value.get(); + return std::nullopt; + } else if constexpr (std::is_integral_v || std::is_enum_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect an integer", type_name); + } + *result = static_cast(value.get()); + return std::nullopt; + } else if constexpr (std::is_floating_point_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect a floating point number", type_name); + } + *result = static_cast(value.get()); + return std::nullopt; + } else if constexpr (std::is_same_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect a string", type_name); + } + // Now PicoJSON will convert byte sequence to latin-1 string. Convert it back to byte sequence. + auto error = Latin1ToBytes(value.get(), result); + if (error) { + return ConstructDeserializeError( + "XGramamr serializer will serialize byte sequence as latin-1 string, but got invalid " + "latin-1 string", + type_name + ); + } + return std::nullopt; + } else if constexpr (is_std_optional::value) { + // for the following container, T must be default constructible + if (value.is()) { + result->reset(); + return std::nullopt; + } else { + return AutoDeserializeJSONValue(&(result->emplace()), value, type_name); + } + } else if constexpr (is_std_pair::value) { + // std::pair: deserialize from an array of size 2 + if (!value.is()) { + return ConstructDeserializeError("Expect an array for deserializing pair", type_name); + } + const auto& arr = value.get(); + if (arr.size() != 2) { + return ConstructDeserializeError( + "Expect an array of size 2 for deserializing pair", type_name + ); + } + if (auto error = AutoDeserializeJSONValue(&(result->first), arr[0], type_name)) { + return error; + } + if (auto error = AutoDeserializeJSONValue(&(result->second), arr[1], type_name)) { + return error; + } + return std::nullopt; + } else if constexpr (is_std_vector::value) { + if (!value.is()) { + return ConstructDeserializeError("Expect an array", type_name); + } + const auto& arr = value.get(); + result->clear(); + result->reserve(arr.size()); + for (const auto& item : arr) { + if (auto error = AutoDeserializeJSONValue(&(result->emplace_back()), item, type_name)) { + return error; + } + } + return std::nullopt; + } else if constexpr (is_std_unordered_set::value) { + if (!value.is()) { + return ConstructDeserializeError( + "Expect an array for deserializing unordered set", type_name + ); + } + const auto& arr = value.get(); + result->clear(); + result->reserve(arr.size()); + for (const auto& item : arr) { + typename T::value_type item_value{}; + if (auto error = AutoDeserializeJSONValue(&item_value, item, type_name)) { + return error; + } + result->emplace(std::move(item_value)); + } + return std::nullopt; + } else if constexpr (is_std_unordered_map::value) { + if constexpr (std::is_same_v) { + // unordered_map: convert from json object + if (!value.is()) { + return ConstructDeserializeError("Expect an object", type_name); + } + const auto& obj = value.get(); + result->clear(); + result->reserve(obj.size()); + for (const auto& [key, item] : obj) { + typename T::mapped_type item_value{}; + if (auto error = AutoDeserializeJSONValue(&item_value, item, type_name)) { + return error; + } + result->try_emplace(key, std::move(item_value)); + } + return std::nullopt; + } else { + // unordered_map (T1 is not string): convert from json array of array of size 2 + if (!value.is()) { + return ConstructDeserializeError( + "Expect an array for deserializing unordered map", type_name + ); + } + const auto& arr = value.get(); + result->clear(); + result->reserve(arr.size()); + for (const auto& item : arr) { + if (!item.is()) { + return ConstructDeserializeError( + "Expect an array of array of size 2 for deserializing unordered map", type_name + ); + } + const auto& sub_arr = item.get(); + if (sub_arr.size() != 2) { + return ConstructDeserializeError( + "Expect an array of array of size 2 for deserializing unordered map", type_name + ); + } + typename T::key_type key_value{}; + if (auto error = AutoDeserializeJSONValue(&key_value, sub_arr[0], type_name)) { + return error; + } + typename T::mapped_type item_value{}; + if (auto error = AutoDeserializeJSONValue(&item_value, sub_arr[1], type_name)) { + return error; + } + result->emplace(std::move(key_value), std::move(item_value)); + } + return std::nullopt; + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Cannot deserialize this type"); + XGRAMMAR_UNREACHABLE(); + } +} + +template +inline std::optional AutoDeserializeJSON( + T* result, const std::string& json_string, bool check_version, const std::string& type_name +) { + picojson::value json_value; + if (auto error = ParseJSON(json_value, json_string); !error.empty()) { + return InvalidJSONError(error); + } + if (check_version) { + XGRAMMAR_DCHECK(json_value.is()); + if (auto error = SerializeVersion::Check(json_value.get())) { + return error; + } + } + return AutoDeserializeJSONValue(result, json_value, type_name); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_JSON_SERIALIZER_H_ diff --git a/third_party/xgrammar/cpp/support/logging.cc b/third_party/xgrammar/cpp/support/logging.cc new file mode 100644 index 0000000000..f14cf8a2c1 --- /dev/null +++ b/third_party/xgrammar/cpp/support/logging.cc @@ -0,0 +1,24 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/logging.cc + */ +#include "logging.h" + +namespace xgrammar { + +#if XGRAMMAR_LOG_CUSTOMIZE == 0 + +LogFatal::Entry& LogFatal::GetEntry() { + static thread_local LogFatal::Entry result; + return result; +} + +const char* LogMessage::level_strings_[] = { + ": ", // XGRAMMAR_LOG_LEVEL_INFO + ": Debug: ", // XGRAMMAR_LOG_LEVEL_DEBUG + ": Warning: ", // XGRAMMAR_LOG_LEVEL_WARNING +}; + +#endif // XGRAMMAR_LOG_CUSTOMIZE + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/support/logging.h b/third_party/xgrammar/cpp/support/logging.h new file mode 100644 index 0000000000..27bd30018d --- /dev/null +++ b/third_party/xgrammar/cpp/support/logging.h @@ -0,0 +1,236 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/logging.h + * \brief A logging library that supports logging at different levels. + */ +#ifndef XGRAMMAR_SUPPORT_LOGGING_H_ +#define XGRAMMAR_SUPPORT_LOGGING_H_ + +#include +#include +#include +#include +#include + +#include "cpptrace.h" // IWYU pragma: keep + +/*! + * \brief Whether or not customize the logging output. + * If log customize is enabled, the user must implement + * xgrammar::LogFatalImpl and xgrammar::LogMessageImpl. + */ +#ifndef XGRAMMAR_LOG_CUSTOMIZE +#define XGRAMMAR_LOG_CUSTOMIZE 0 +#endif + +namespace xgrammar { + +/*! + * \brief Error type for errors from XGRAMMAR_CHECK, XGRAMMAR_ICHECK, and XGRAMMAR_LOG(FATAL). This + * error contains a backtrace of where it occurred. + */ +class LogFatalError : public std::runtime_error { + public: + /*! \brief Construct an error. Not recommended to use directly. Instead use XGRAMMAR_LOG(FATAL). + * + * \param file The file where the error occurred. + * \param lineno The line number where the error occurred. + * \param message The error message to display. + * \param time The time at which the error occurred. This should be in local time. + */ + LogFatalError( + const std::string& file, + int lineno, + const std::string& message, + std::time_t time = std::time(nullptr) + ) + : std::runtime_error(message), file_(file), lineno_(lineno), time_(time) { + std::ostringstream s; + s << "[" << std::put_time(std::localtime(&time), "%H:%M:%S") << "] " << file << ":" << lineno + << ": " << message << "\n"; + full_message_ = s.str(); + } + + /*! \return The file in which the error occurred. */ + const std::string& file() const { return file_; } + /*! \return The time at which this error occurred. */ + const std::time_t& time() const { return time_; } + /*! \return The line number at which this error occurred. */ + int lineno() const { return lineno_; } + /*! \return The error message. */ + const char* what() const noexcept override { return full_message_.c_str(); } + + private: + std::string file_; + int lineno_; + std::time_t time_; + std::string full_message_; +}; + +// Provide support for customized logging. +#if XGRAMMAR_LOG_CUSTOMIZE +/*! + * \brief Custom implementations of LogFatal. + * + * \sa XGRAMMAR_LOG_CUSTOMIZE + */ +[[noreturn]] void LogFatalImpl(const std::string& file, int lineno, const std::string& message); + +/*! + * \brief Custom implementations of LogMessage. + * + * \sa XGRAMMAR_LOG_CUSTOMIZE + */ +void LogMessageImpl(const std::string& file, int lineno, int level, const std::string& message); + +/*! + * \brief Class to accumulate an error message and throw it. Do not use + * directly, instead use LOG(FATAL). + */ +class LogFatal { + public: + LogFatal(const std::string& file, int lineno) : file_(file), lineno_(lineno) {} +#ifdef _MSC_VER +#pragma disagnostic push +#pragma warning(disable : 4722) +#endif + [[noreturn]] ~LogFatal() noexcept(false) { LogFatalImpl(file_, lineno_, stream_.str()); } +#ifdef _MSC_VER +#pragma disagnostic pop +#endif + std::ostringstream& stream() { return stream_; } + + private: + std::ostringstream stream_; + std::string file_; + int lineno_; +}; + +/*! + * \brief Class to accumulate an log message. Do not use directly, instead use + * LOG(INFO), LOG(WARNING), LOG(ERROR). + */ +class LogMessage { + public: + LogMessage(const std::string& file, int lineno, int level) + : file_(file), lineno_(lineno), level_(level) {} + ~LogMessage() { LogMessageImpl(file_, lineno_, level_, stream_.str()); } + std::ostringstream& stream() { return stream_; } + + private: + std::string file_; + int lineno_; + int level_; + std::ostringstream stream_; +}; + +#else // if XGRAMMAR_LOG_CUSTOMIZE + +/*! + * \brief Class to accumulate an error message and throw it. Do not use + * directly, instead use XGRAMMAR_LOG(FATAL). + * \note The `LogFatal` class is designed to be an empty class to reduce stack size usage. + * To play this trick, we use the thread-local storage to store its internal data. + */ +class LogFatal { + public: + LogFatal(const std::string& file, int lineno) { GetEntry().Init(file, lineno); } +#ifdef _MSC_VER +#pragma disagnostic push +#pragma warning(disable : 4722) +#endif + [[noreturn]] ~LogFatal() noexcept(false) { + GetEntry().Finalize(); + throw; + } +#ifdef _MSC_VER +#pragma disagnostic pop +#endif + std::ostringstream& stream() { return GetEntry().stream_; } + + private: + struct Entry { + void Init(const std::string& file, int lineno) { + this->stream_.str(""); + this->file_ = file; + this->lineno_ = lineno; + } + [[noreturn]] LogFatalError Finalize() noexcept(false) { + LogFatalError error(file_, lineno_, stream_.str()); + throw error; + } + std::ostringstream stream_; + std::string file_; + int lineno_; + }; + + static Entry& GetEntry(); +}; + +/*! + * \brief Class to accumulate an log message. Do not use directly, instead use + * XGRAMMAR_LOG(INFO), XGRAMMAR_LOG(WARNING), XGRAMMAR_LOG(ERROR). + */ +class LogMessage { + public: + LogMessage(const std::string& file, int lineno, int level) { + std::time_t t = std::time(nullptr); + stream_ << "[" << std::put_time(std::localtime(&t), "%H:%M:%S") << "] " << file << ":" << lineno + << level_strings_[level]; + } + ~LogMessage() { std::cerr << (stream_.str() + "\n"); } + std::ostringstream& stream() { return stream_; } + + private: + std::ostringstream stream_; + static const char* level_strings_[]; +}; + +#endif // XGRAMMAR_LOG_CUSTOMIZE + +#define XGRAMMAR_LOG_LEVEL_INFO 0 +#define XGRAMMAR_LOG_LEVEL_DEBUG 1 +#define XGRAMMAR_LOG_LEVEL_WARNING 2 +#define XGRAMMAR_LOG_LEVEL_FATAL 3 + +#define XGRAMMAR_LOG_INFO LogMessage(__FILE__, __LINE__, XGRAMMAR_LOG_LEVEL_INFO).stream() +#define XGRAMMAR_LOG_DEBUG LogMessage(__FILE__, __LINE__, XGRAMMAR_LOG_LEVEL_DEBUG).stream() +#define XGRAMMAR_LOG_WARNING LogMessage(__FILE__, __LINE__, XGRAMMAR_LOG_LEVEL_WARNING).stream() +#define XGRAMMAR_LOG_FATAL LogFatal(__FILE__, __LINE__).stream() + +/*! + * \brief Log a message at the given level. + * \param level The level of the message. Can be INFO, DEBUG, WARNING, FATAL. + */ +#define XGRAMMAR_LOG(level) XGRAMMAR_LOG_##level + +/*! + * \brief Check if the condition is true. Used for checking the correctness of user inputs. + * \param x The condition to check. + */ +#define XGRAMMAR_CHECK(x) \ + if (!(x)) LogFatal(__FILE__, __LINE__).stream() << "Check failed: (" #x << ") is false: " + +/*! + * \brief Check if the condition is true. Used to guarantee some internal conditions in the code. + * \param x The condition to check. + */ +#define XGRAMMAR_ICHECK(x) \ + if (!(x)) LogFatal(__FILE__, __LINE__).stream() << "Internal check failed: (" #x << ") is false: " + +/*! + * \brief Check if the condition is true. Used to guarantee some internal conditions in the code. + * \note This check is only enabled in debug mode. In release mode, it will be disabled for + * efficiency. This should be used in preference to XGRAMMAR_ICHECK. + * \param x The condition to check. + */ +#if XGRAMMAR_ENABLE_INTERNAL_CHECK +#define XGRAMMAR_DCHECK(x) XGRAMMAR_ICHECK(x) +#else +#define XGRAMMAR_DCHECK(x) \ + while (false) XGRAMMAR_ICHECK(x) +#endif // XGRAMMAR_ENABLE_DCHECK + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_LOGGING_H_ diff --git a/third_party/xgrammar/cpp/support/memory_size.h b/third_party/xgrammar/cpp/support/memory_size.h new file mode 100644 index 0000000000..d5ad99529f --- /dev/null +++ b/third_party/xgrammar/cpp/support/memory_size.h @@ -0,0 +1,129 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/memory_size.h + * \brief Compute the memory consumption of a container in heap memory. + */ + +#ifndef XGRAMMAR_SUPPORT_MEMORY_SIZE_H_ +#define XGRAMMAR_SUPPORT_MEMORY_SIZE_H_ + +#include +#include +#include +#include +#include +#include + +#include "reflection.h" + +namespace xgrammar { + +/******************* MemorySize Procotol *******************/ + +template +inline constexpr std::size_t MemorySize(const T& value); + +template +inline constexpr std::size_t MemorySize(const std::pair& pair); + +template +inline constexpr std::size_t MemorySize(const std::tuple& tpl); + +template +inline constexpr std::size_t MemorySize(const std::optional& optional_value); + +inline std::size_t MemorySize(const std::vector& value); + +/******************* MemorySize Implementations *******************/ + +namespace detail::memory_size { + +/*! + * \brief Get the element type of a container. + */ +template +using ElementType = std::decay_t; + +/*! + * \brief A false value for static_assert. + */ +template +inline constexpr bool false_v = false; + +} // namespace detail::memory_size + +/*! + * \brief Compute the memory consumption of a value. + * \tparam T The type of the value. + * \param value The value. + * \return The memory consumption in heap memory of the value in bytes. + */ +template +inline constexpr std::size_t MemorySize(const T& value) { + if constexpr (is_pimpl_class::value) { + // Customized MemorySize + return MemorySize(*value.ImplPtr()); + } else if constexpr (std::is_trivially_copyable_v) { + // Primitive type + return 0; + } else if constexpr (std::is_trivially_copyable_v>) { + // Container of primitive type + return sizeof(detail::memory_size::ElementType) * std::size(value); + } else if constexpr (!std::is_trivially_copyable_v>) { + // Container of non-primitive type: sum up the memory size of all elements + std::size_t size = sizeof(detail::memory_size::ElementType) * std::size(value); + for (const auto& element : value) { + size += MemorySize(element); + } + return size; + } else { + static_assert(detail::memory_size::false_v, "MemorySize is not implemented for this type"); + } +} + +/*! + * \brief Compute the memory consumption of a pair. + * \tparam T1 The type of the first element. + * \tparam T2 The type of the second element. + * \param pair The pair. + * \return The memory consumption in heap memory of the pair. + */ +template +inline constexpr std::size_t MemorySize(const std::pair& pair) { + return MemorySize(pair.first) + MemorySize(pair.second); +} + +/*! + * \brief Compute the memory consumption of a tuple. + * \tparam Ts The types of the tuple. + * \param tpl The tuple. + * \return The memory consumption in heap memory of the tuple. + */ +template +inline constexpr std::size_t MemorySize(const std::tuple& tpl) { + return std::apply([](auto&&... elems) { return (MemorySize(elems) + ... + 0); }, tpl); +} + +/*! + * \brief Compute the memory consumption in heap memory. This function is specialized for + * std::optional. + * \tparam Tp The type of the optional. + * \param range The optional. + * \return The memory consumption in heap memory of the optional. + */ +template +inline constexpr std::size_t MemorySize(const std::optional& optional_value) { + return optional_value.has_value() ? MemorySize(*optional_value) : 0; +} + +/*! + * \brief Compute the memory consumption of a std::vector, which is bit-packed: it stores + * one bit (not one byte) per element. The generic container overload would over-count it by 8x. + * \param value The vector. + * \return The memory consumption in heap memory of the vector in bytes. + */ +inline std::size_t MemorySize(const std::vector& value) { return (value.size() + 7) / 8; } + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_MEMORY_SIZE_H_ diff --git a/third_party/xgrammar/cpp/support/recursion_guard.cc b/third_party/xgrammar/cpp/support/recursion_guard.cc new file mode 100644 index 0000000000..0a182863b4 --- /dev/null +++ b/third_party/xgrammar/cpp/support/recursion_guard.cc @@ -0,0 +1,50 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/recursion_guard.cc + */ + +#include "recursion_guard.h" + +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +int RecursionGuard::LoadMaxRecursionDepthFromEnv() { + const char* env_value = std::getenv(kMaxRecursionDepthEnvVar); + if (env_value == nullptr) { + return kDefaultMaxRecursionDepth; + } + + int value = 0; + std::string_view sv(env_value); + + // Convert the string to an integer + auto result = std::from_chars(sv.data(), sv.data() + sv.size(), value); + + // Check if the conversion is successful + if (result.ec == std::errc::invalid_argument || result.ec == std::errc::result_out_of_range || + result.ptr != sv.data() + sv.size() || value <= 0) { + XGRAMMAR_LOG(WARNING) << "Env variable XGRAMMAR_MAX_RECURSION_DEPTH is not a valid " + "integer or out of range: '" + << env_value << "', using default " << kDefaultMaxRecursionDepth; + return kDefaultMaxRecursionDepth; + } + + // Check if the value is too large + if (value > kMaxReasonableDepth) { + XGRAMMAR_LOG(WARNING) << "Env variable XGRAMMAR_MAX_RECURSION_DEPTH too large: " << value + << ", clamping to " << kMaxReasonableDepth; + return kMaxReasonableDepth; + } + + return value; +} + +std::atomic RecursionGuard::max_recursion_depth_{LoadMaxRecursionDepthFromEnv()}; + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/support/recursion_guard.h b/third_party/xgrammar/cpp/support/recursion_guard.h new file mode 100644 index 0000000000..3c1ac4ba05 --- /dev/null +++ b/third_party/xgrammar/cpp/support/recursion_guard.h @@ -0,0 +1,127 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/recursion_guard.h + * \brief The header for recursion depth guard. + */ + +#ifndef XGRAMMAR_SUPPORT_RECURSION_GUARD_H_ +#define XGRAMMAR_SUPPORT_RECURSION_GUARD_H_ + +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +/*! + * \brief Thread-safe recursion guard to prevent stack overflow + * + * This class provides a RAII-style guard that tracks recursion depth + * and prevents excessive recursion that could lead to stack overflow. + * It uses atomic operations for thread safety and supports configurable + * maximum recursion depth. + */ +class RecursionGuard { + public: + /*! + * \brief Constructor that increments recursion depth + * \param current_recursion_depth Pointer to the current recursion depth counter + * \throws Logs fatal error if max recursion depth is exceeded + */ + explicit RecursionGuard(int* current_recursion_depth) + : current_depth_ptr_(current_recursion_depth) { + auto error = AddRecursionDepth(current_depth_ptr_); + XGRAMMAR_CHECK(error == std::nullopt) << error.value().what(); + } + + /*! + * \brief Reset the recursion depth to 0 + * \param current_recursion_depth Pointer to the current recursion depth counter + */ + static void ResetRecursionDepth(int* current_recursion_depth) { + XGRAMMAR_DCHECK(current_recursion_depth != nullptr); + *current_recursion_depth = 0; + } + + /*! + * \brief Destructor that decrements recursion depth + */ + ~RecursionGuard() { SubtractRecursionDepth(current_depth_ptr_); } + + /*! + * \brief Get the maximum allowed recursion depth + * \return Current maximum recursion depth limit + */ + static int GetMaxRecursionDepth() { return max_recursion_depth_.load(std::memory_order_relaxed); } + + /*! + * \brief Set the maximum allowed recursion depth + * \param max_depth New maximum recursion depth limit (must be positive) + */ + static void SetMaxRecursionDepth(int max_depth) { + if (max_depth <= 0 || max_depth > kMaxReasonableDepth) { + XGRAMMAR_LOG(FATAL + ) << "RecursionGuard: Maximum recursion depth must be positive and less than " + << kMaxReasonableDepth << ", got: " << max_depth; + } + max_recursion_depth_.store(max_depth, std::memory_order_relaxed); + } + + static std::optional AddRecursionDepth(int* current_recursion_depth) { + XGRAMMAR_DCHECK(current_recursion_depth != nullptr); + int current_depth = ++(*current_recursion_depth); + int max_depth = max_recursion_depth_.load(std::memory_order_relaxed); + if (current_depth > max_depth) { + return std::runtime_error( + "RecursionGuard: Maximum recursion depth exceeded. " + "Current depth: " + + std::to_string(current_depth) + ", Max allowed: " + std::to_string(max_depth) + ); + } + return std::nullopt; + } + + static void SubtractRecursionDepth(int* current_recursion_depth) { + XGRAMMAR_DCHECK(current_recursion_depth != nullptr && *current_recursion_depth > 0); + --(*current_recursion_depth); + } + + private: + /*! + * \brief Get the maximum allowed recursion depth from the environment variable. Used to + * initialize max_recursion_depth_. + * \return Current maximum recursion depth limit + */ + static int LoadMaxRecursionDepthFromEnv(); + + /*! + * \brief Pointer to the recursion depth counter + */ + int* current_depth_ptr_; + + /*! + * \brief Thread-safe global configuration + */ + static std::atomic max_recursion_depth_; + + /*! + * \brief Environment variable name for the maximum recursion depth + */ + inline constexpr static char kMaxRecursionDepthEnvVar[] = "XGRAMMAR_MAX_RECURSION_DEPTH"; + + /*! + * \brief Default maximum recursion depth + */ + inline constexpr static int kDefaultMaxRecursionDepth = 10000; + + /*! + * \brief Maximum reasonable recursion depth + */ + inline constexpr static int kMaxReasonableDepth = 1000000; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_RECURSION_GUARD_H_ diff --git a/third_party/xgrammar/cpp/support/reflection.h b/third_party/xgrammar/cpp/support/reflection.h new file mode 100644 index 0000000000..4703d9a746 --- /dev/null +++ b/third_party/xgrammar/cpp/support/reflection.h @@ -0,0 +1,300 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/reflection.h + * \brief The header for compile-time reflection. + */ + +#ifndef XGRAMMAR_SUPPORT_REFLECTION_H_ +#define XGRAMMAR_SUPPORT_REFLECTION_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace xgrammar { + +/******************** Core Reflection Types ********************/ + +/*! + * \brief The type of the member trait. + */ +enum class member_type { + kNone = 0, // this is default, which has no member trait + kConfig = 1, // this is a config with member pointers +}; + +/** + * \brief Base trait for member traits. + * + * \tparam T the type whose members are being reflected + * \details Provides a default trait indicating no members. + */ +template +struct member_trait { + static constexpr auto value = member_type::kNone; +}; + +/******************** STL and Custom Type Traits ********************/ + +template +struct is_std_array : std::false_type {}; + +template +struct is_std_array> : std::true_type {}; + +template +struct is_std_pair : std::false_type {}; + +template +struct is_std_pair> : std::true_type {}; + +template +struct is_std_tuple : std::false_type {}; + +template +struct is_std_tuple> : std::true_type {}; + +template +struct is_std_optional : std::false_type {}; + +template +struct is_std_optional> : std::true_type {}; + +template +struct is_std_vector : std::false_type {}; + +template +struct is_std_vector> : std::true_type {}; + +template +struct is_std_unordered_map : std::false_type {}; + +template +struct is_std_unordered_map> : std::true_type {}; + +template +struct is_std_unordered_set : std::false_type {}; + +template +struct is_std_unordered_set> : std::true_type {}; + +/*! + * \brief XGrammar specific: Check if a class is a PImpl class. + */ +template +struct is_pimpl_class : std::false_type {}; + +/*! + * \brief XGrammar specific: Check if a class is a PImpl class. It's true iff the class has a + * member `Impl` and the class is not the same as the `Impl` type. + */ +template +struct is_pimpl_class< + T, + std::void_t, void>>> + : std::true_type {}; + +/*! + * \brief A helper class to print the value when the condition is false. + */ +template +struct DebugAssert { + static_assert(condition); +}; + +/******************** Implementation Details ********************/ + +namespace detail::reflection { + +// We cannot use `static_assert(false)` even in unreachable code in `if constexpr`. +// See https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2023/p2593r1.html +// for more details. +// TL;DR: We use the following `false_v` as a workaround. +template +inline constexpr bool false_v = false; + +// Note that we don't allow empty tables now (that's uncommon). +template +inline constexpr auto make_member_table(X, Y second, Args... args) { + static_assert(sizeof...(args) % 2 == 0, "member table must be even"); + static_assert(std::is_same_v, "first member must be a c-string"); + static_assert(std::is_member_pointer_v, "second member must be a member pointer"); + if constexpr (sizeof...(args) == 0) { + return std::make_tuple(second); + } else { + return std::tuple_cat(std::make_tuple(second), make_member_table(args...)); + } +} + +template +inline constexpr auto make_name_table_aux(std::index_sequence, Tuple tuple) { + return std::array{std::get(tuple)...}; +} + +template +inline constexpr auto make_name_table(Args... args) { + constexpr auto N = sizeof...(args); + static_assert(N % 2 == 0, "name table must be even"); + return make_name_table_aux(std::make_index_sequence{}, std::make_tuple(args...)); +} + +template +inline void visit_config_impl(Fn&& fn, std::index_sequence) { + // This is a helper function to visit each member of the config. + // It uses fold expression to apply the function to each member. + static_assert(Ftor::value == member_type::kConfig, "T must be a config type"); + static constexpr auto get_name = [](std::size_t idx) { + if constexpr (Ftor::has_names) { + return Ftor::names[idx]; + } else { + return ""; + } + }; + return (fn(std::get(Ftor::members), get_name(Idx), Idx), ...); +} + +} // namespace detail::reflection + +/******************** Member Functors and Visitors ********************/ + +/*! + * \brief A functor that provides access to the members of a config type. + * It extracts the members from the `member_trait` specialization for the type `T`. + * A valid `member_trait` specialization must meet the following requirements: + * - It must have a static member `value` of type `member_type`, + * which must be either `kNone` or `kConfig`. + */ +template ::value> +struct member_functor { + static_assert(detail::reflection::false_v, "This specialization should never be used"); +}; + +/*! + * \brief A specialization of `member_functor` for config types. + * A valid `member_trait` specialization for a config type must meet the following: + * - It must have a static member `value` of type `member_type::kConfig`. + * - It must have a static tuple `members` that contains the member pointers. + * - It must have a static array `names` that contains the names of the members. + * - The size of `names` must be either 0 or equal to the number of members in `members`. + * - In the first case, `names` will be empty. + * - In the second case, `names` represent the printed name of each member. + */ +template +struct member_functor { + private: + using _trait_t = member_trait; + using _members_t = std::decay_t; + using _names_t = std::decay_t; + + public: + static constexpr auto value = member_type::kConfig; + static constexpr auto members = _trait_t::members; + static constexpr auto names = _trait_t::names; + static constexpr auto member_count = std::tuple_size_v<_members_t>; + static constexpr auto has_names = names.size() == member_count; + + // some static_asserts to check the member list and name list + static_assert(is_std_tuple<_members_t>::value, "Member list must be a tuple"); + static_assert(is_std_array<_names_t>::value, "Name list must be an array"); + static_assert(member_count > 0, "Member list must not be empty"); + static_assert( + names.size() == member_count || names.size() == 0, + "Name list must be empty or have the same size as member list" + ); +}; + +/*! + * \brief Visit the members of a config type. + * \tparam T The type of the config. + * \tparam Fn The type of the function to visit the members. + * \param fn The function to visit the members. fn's signature should be: + * \code{.cpp} + * (auto ptr, const char* name, size_t idx) -> void + * \endcode + * where `ptr` is the pointer to the member, `name` is the name of the member, and `idx` is the + * index of the member. + */ +template +inline void visit_config(Fn&& fn) { + using Ftor = member_functor; + return detail::reflection::visit_config_impl( + fn, std::make_index_sequence{} + ); +} + +/******************** Registration Macros ********************/ + +/** + * \brief Macros to define member traits for types. + * \details These macros are used to define the structural information of types + * for serialization and reflection purposes. + * + * Macros: + * - \c XGRAMMAR_MEMBER_TABLE: Defines a type with a table of (name, member pointer) pairs. + * - \c XGRAMMAR_MEMBER_ARRAY: Defines a type with an array of member pointers. + * + * Use the `_TEMPLATE` variants for template types. + * + * \example + * \code{.cpp} + * // Example of using XGRAMMAR_MEMBER_TABLE to register (name, member pointer) pairs + * struct SimpleClass { + * int a; + * double b; + * }; + * XGRAMMAR_MEMBER_TABLE(SimpleClass, "name_a", &SimpleClass::a, "name_b", &SimpleClass::b); + * + * // Or register members as an array with XGRAMMAR_MEMBER_ARRAY + * XGRAMMAR_MEMBER_ARRAY(SimpleClass, &SimpleClass::a, &SimpleClass::b); + * + * // Example of using XGRAMMAR_MEMBER_ARRAY to register members from a derived class + * struct Derived : SimpleClass { + * std::string c; + * }; + * XGRAMMAR_MEMBER_TABLE(Derived, "name_a", &Derived::a, "name_b", &Derived::b, "name_c", + * &Derived::c); + * + * // Example of using XGRAMMAR_MEMBER_ARRAY_TEMPLATE for a template type + * // If the default constructor/member is private, you need to declare a friend for member_trait. + * template + * struct TemplateClass { + * private: + * T value; + * TemplateClass() = default; + * friend struct member_trait; + * }; + * template + * XGRAMMAR_MEMBER_ARRAY_TEMPLATE(TemplateClass, &TemplateClass::value); + * \endcode + */ +#define XGRAMMAR_MEMBER_TABLE_TEMPLATE(Type, ...) \ + struct member_trait { \ + static constexpr auto value = member_type::kConfig; \ + static constexpr auto members = detail::reflection::make_member_table(__VA_ARGS__); \ + static constexpr auto names = detail::reflection::make_name_table(__VA_ARGS__); \ + } + +#define XGRAMMAR_MEMBER_ARRAY_TEMPLATE(Type, ...) \ + struct member_trait { \ + static constexpr auto value = member_type::kConfig; \ + static constexpr auto members = std::make_tuple(__VA_ARGS__); \ + static constexpr auto names = std::array{}; \ + } + +#define XGRAMMAR_MEMBER_TABLE(Type, ...) \ + template <> \ + XGRAMMAR_MEMBER_TABLE_TEMPLATE(Type, __VA_ARGS__) + +#define XGRAMMAR_MEMBER_ARRAY(Type, ...) \ + template <> \ + XGRAMMAR_MEMBER_ARRAY_TEMPLATE(Type, __VA_ARGS__) + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_REFLECTION_H_ diff --git a/third_party/xgrammar/cpp/support/thread_pool.h b/third_party/xgrammar/cpp/support/thread_pool.h new file mode 100644 index 0000000000..6bbec3c2d1 --- /dev/null +++ b/third_party/xgrammar/cpp/support/thread_pool.h @@ -0,0 +1,239 @@ +/*! + * Copyright (c) 2023 by Contributors + * \file xgrammar/support/thread_pool.h + * \brief Thread pool. + */ +#ifndef XGRAMMAR_SUPPORT_THREAD_POOL_H_ +#define XGRAMMAR_SUPPORT_THREAD_POOL_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +/*! + * \brief A thread pool implementation for parallel task execution. + * + * ThreadPool manages a pool of worker threads that can execute tasks asynchronously. + * Tasks are submitted to a queue and executed by available threads from the pool. + * The pool automatically handles thread synchronization and task distribution. + */ +class ThreadPool { + public: + /*! + * \brief Construct a new thread pool with the specified number of threads. + * \param num_threads Number of worker threads to create. Defaults to hardware concurrency. + * \note The pool starts the worker threads immediately upon construction. + */ + ThreadPool(size_t num_threads = std::thread::hardware_concurrency()) { + // Initialize thread pool with num_threads threads + for (size_t i = 0; i < num_threads; ++i) { + workers_.emplace_back([this] { + while (true) { + std::function task; + { + // Lock queue while waiting for new task + std::unique_lock lock(queue_mutex_); + queue_condition_.wait(lock, [this] { return shutdown_ || !task_queue_.empty(); }); + + // Exit thread if shutdown and queue is empty + if (shutdown_ && task_queue_.empty()) return; + + // Get task from queue + task = std::move(task_queue_.front()); + task_queue_.pop(); + } + try { + task(); + } catch (...) { + std::unique_lock lock(queue_mutex_); + if (!first_exception_) { + first_exception_ = std::current_exception(); + } + } + TaskComplete(); + } + }); + } + } + + /*! + * \brief Add a new task to be executed by the thread pool. + * \tparam F Type of the function to execute + * \tparam Args Types of the arguments to pass to the function + * \param f Function to execute + * \param args Arguments to pass to the function + * \return std::shared_future containing the result of the function call + * \note Tasks are executed in FIFO order but may complete in any order. + */ + template + auto Submit(F&& f, Args&&... args) -> std::shared_future> { + using return_type = std::invoke_result_t; + + // Package the task with its arguments into a shared pointer + auto task = std::make_shared>( + std::bind(std::forward(f), std::forward(args)...) + ); + + std::shared_future res = task->get_future().share(); + + { + std::unique_lock lock(queue_mutex_); + XGRAMMAR_CHECK(!shutdown_) << "Cannot submit task to stopped ThreadPool"; + ++unfinished_task_count_; // Increment task count + + // Directly add the task without wrapping + task_queue_.emplace([task]() { (*task)(); }); + } + queue_condition_.notify_one(); + return res; + } + + /*! + * \brief Add a new task to be executed by the thread pool without returning a future. + * \tparam F Type of the function to execute + * \tparam Args Types of the arguments to pass to the function + * \param f Function to execute + * \param args Arguments to pass to the function + * \note Tasks are executed asynchronously by the worker threads. + */ + template + void Execute(F&& f, Args&&... args) { + { + std::unique_lock lock(queue_mutex_); + XGRAMMAR_CHECK(!shutdown_) << "Cannot execute task in stopped ThreadPool"; + ++unfinished_task_count_; // Increment task count + + // Directly add the task without wrapping + task_queue_.emplace(std::bind(std::forward(f), std::forward(args)...)); + } + queue_condition_.notify_one(); + } + + /*! + * \brief Wait until all submitted tasks have finished. + * \note If a task submitted with Execute threw, the first such exception is rethrown here on the + * calling thread. An exception escaping a worker thread would otherwise terminate the process. + */ + void Wait() { + { + std::unique_lock lock(queue_mutex_); + tasks_done_condition_.wait(lock, [this] { return unfinished_task_count_ == 0; }); + } + RethrowTaskException(); + } + + /*! + * \brief Join all threads in the pool. + * + * Sets shutdown flag and waits for all threads to complete their current tasks + * before destroying the pool. Any remaining tasks in the queue will be executed + * before shutdown completes. + * \note If a task submitted with Execute threw, the first such exception is rethrown here on the + * calling thread. An exception escaping a worker thread would otherwise terminate the process. + */ + void Join() { + Shutdown(); + RethrowTaskException(); + } + + /*! + * \brief Destructor that ensures graceful shutdown of the thread pool. + */ + ~ThreadPool() { Shutdown(); } + + // Prevent copying or moving of the thread pool + ThreadPool(const ThreadPool&) = delete; + ThreadPool(ThreadPool&&) = delete; + ThreadPool& operator=(const ThreadPool&) = delete; + ThreadPool& operator=(ThreadPool&&) = delete; + + private: + void Shutdown() { + { + std::unique_lock lock(queue_mutex_); + if (shutdown_) return; // Already shut down + shutdown_ = true; + } + + queue_condition_.notify_all(); // Wake up all threads so they can exit + for (std::thread& worker : workers_) { + if (worker.joinable()) worker.join(); // Wait for thread to finish + } + } + + void RethrowTaskException() { + std::exception_ptr exception; + { + std::unique_lock lock(queue_mutex_); + std::swap(exception, first_exception_); + } + if (exception) { + std::rethrow_exception(exception); + } + } + + void TaskComplete() { + std::unique_lock lock(queue_mutex_); + --unfinished_task_count_; + if (unfinished_task_count_ == 0) { + tasks_done_condition_.notify_all(); // Notify waiting threads + } + } + + /*! \brief Thread container */ + std::vector workers_; + /*! \brief Task queue */ + std::queue> task_queue_; + /*! \brief Mutex to protect task queue */ + std::mutex queue_mutex_; + /*! \brief Condition variable for thread synchronization */ + std::condition_variable queue_condition_; + /*! \brief Condition variable for task completion */ + std::condition_variable tasks_done_condition_; + /*! \brief Flag to indicate thread pool shutdown */ + bool shutdown_ = false; + /*! \brief Number of unfinished tasks */ + int unfinished_task_count_ = 0; + /*! \brief The first exception thrown by a task, rethrown by Wait() or Join() */ + std::exception_ptr first_exception_ = nullptr; +}; + +inline void ParallelFor(int low, int high, int num_threads, std::function f) { + if (high - low == 1) { + f(low); + return; + } + + ThreadPool pool(num_threads); + + int total = high - low; + int chunk_size = (total + num_threads - 1) / num_threads; + + for (int t = 0; t < num_threads; ++t) { + int start = low + t * chunk_size; + int end = std::min(start + chunk_size, high); + + if (start >= end) break; // No more iterations to process + + pool.Execute([f, start, end]() { + for (int i = start; i < end; ++i) { + f(i); + } + }); + } + pool.Join(); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_THREAD_POOL_H_ diff --git a/third_party/xgrammar/cpp/support/thread_safe_cache.h b/third_party/xgrammar/cpp/support/thread_safe_cache.h new file mode 100644 index 0000000000..28859db66c --- /dev/null +++ b/third_party/xgrammar/cpp/support/thread_safe_cache.h @@ -0,0 +1,404 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/thread_safe_cache.h + * \brief The header for thread-safe caching functionality. + */ +#ifndef XGRAMMAR_SUPPORT_THREAD_SAFE_CACHE_H_ +#define XGRAMMAR_SUPPORT_THREAD_SAFE_CACHE_H_ + +#include +#include // IWYU pragma: keep +#include +#include +#include +#include +#include +#include +#include +#include + +#include "container.h" + +namespace xgrammar { + +/*! + * \brief Primary template for ThreadSafeCache + * \details This class provides thread-safe caching functionality in two forms: + * 1. Single value cache when only Value template parameter is provided + * 2. Key-value cache when both Key and Value template parameters are provided + */ +template +class ThreadSafeCache; + +/*! + * \brief Thread-safe cache for a single computed value + * \tparam Value The type of value being cached + * \details Specialization that provides: + * - Thread-safe access to a single cached value + * - Lazy computation on first access + * - Reader-writer locking for concurrent reads + */ +template +class ThreadSafeCache { + public: + /*! + * \brief Constructs a new single-value cache + * \param compute The function that computes the cached value + */ + explicit ThreadSafeCache(std::function compute) : compute_(std::move(compute)) {} + + /*! + * \brief Gets or computes the cached value + * \return The cached or newly computed value + */ + Value Get() { + // First try reading from cache with shared lock + { + std::shared_lock cache_lock(cache_mutex_); + if (cache_.has_value()) { + return cache_.value(); // Cache hit + } + } + + // Acquire exclusive lock to compute value + std::unique_lock cache_lock(cache_mutex_); + + // Double-check to prevent redundant computation + if (cache_.has_value()) { + return cache_.value(); + } + + Value value = compute_(); + XGRAMMAR_DCHECK(!cache_.has_value()); + cache_ = value; + return value; + } + + /*! + * \brief Clears the cached value + * This function removes the cached value, so the next call to Get() will recompute it. + */ + void Clear() { + std::unique_lock cache_lock(cache_mutex_); + cache_.reset(); + } + + private: + /*! \brief Optional container holding the cached value */ + std::optional cache_; + /*! \brief Function used to compute the value when not cached */ + std::function compute_; + /*! \brief Reader-writer lock protecting access to cache_ */ + std::shared_mutex cache_mutex_; +}; + +/*! + * \brief A thread-safe key-value cache with on-demand computation + * \tparam Key The type of keys used to lookup values. Should be hashable. + * \tparam Value The type of values stored in the cache + * \details This cache provides thread-safe access to computed values with the following features: + * - Lazy computation: Values are only computed when first requested + * - Thread safety: Uses reader-writer locks for concurrent reads + * - Parallel computation: Different keys can be computed simultaneously + * - Double-checked locking: Prevents redundant computation + */ +template +class ThreadSafeCache { + public: + /*! + * \brief Constructs a new thread-safe cache + * \param compute The function that computes values for uncached keys + */ + explicit ThreadSafeCache(std::function compute) + : compute_(std::move(compute)) {} + + /*! + * \brief Gets or computes the value for a key + * \param key The key to lookup + * \return The cached or newly computed value of the key + */ + Value Get(const Key& key) { + // Why we need this: + // - When adding new elements to a unordered_map, the map may be rehashed, + // - which means all the iterators may be invalidated. + // - However, cppreference says: + // - "References and pointers to either key or data stored in the container are only invalidated + // - by erasing that element, even when the corresponding iterator is invalidated." + // - (See https://en.cppreference.com/w/cpp/container/unordered_map) + // - Therefore, we should maintain 2 locks. + // - When we add something to the cache, we should hold the cache_mutex_. + // - When we erase something from the cache, we should hold the clear_mutex_. + + auto erase_lock = std::shared_lock(erase_mutex_); + + // First attempt to read from cache_ + { + auto cache_lock = std::shared_lock(cache_mutex_); + auto it = cache_.find(key); + if (it != cache_.end()) { // Cache hit + auto& entry = it->second; // The iterator is invalidated after releasing the lock + cache_lock.unlock(); // Therefore, we should hold the entry by reference first + + // We should not hold lock here, since this function may be blocking. + return entry.get(compute_, key); + } + } + + // Acquire exclusive lock to compute value + { + auto cache_lock = std::unique_lock(cache_mutex_); + auto& entry = cache_[key]; // Create a new entry + cache_lock.unlock(); // Release the lock before blocking + + // We should not hold lock here, since this function may be blocking. + return entry.get(compute_, key); + } + } + + /*! + * \brief Clears all cached values and associated per-key mutexes + * This function removes all cached key-value pairs, so subsequent calls to Get() will recompute + * them. + */ + void Clear() { + auto erase_lock = std::unique_lock(erase_mutex_); + cache_.clear(); + } + + private: + struct Entry { + Value value; + std::once_flag flag; + const Value& get(const std::function& f, const Key& key) { + // block in this lambda until the value is computed + std::call_once(flag, [&] { value = f(key); }); + return value; + } + }; + + /*! \brief The cache mapping keys to computed values */ + std::unordered_map cache_; + /*! \brief The function used to compute values for uncached keys */ + std::function compute_; + /*! \brief Reader-writer lock protecting access to cache_ */ + std::shared_mutex cache_mutex_; + /*! \brief Mutex protecting removing elements */ + std::shared_mutex erase_mutex_; +}; + +namespace details { + +template +class LRUCacheImpl { + public: + struct Entry { + Value value; // value of the node + int index; // node index + }; + + /*! \brief Visits the node and moves it to the back of the LRU list. Return its value. */ + const Value& LRUVisit(const std::pair& pair) { + const auto& entry = pair.second; + lru_list_.MoveBack(entry.index); + return entry.value; + } + + /*! \brief Initializes the node with the given value and moves it to the back of the LRU list. */ + void LRUInit(std::pair& pair, const Value& init) { + auto& entry = pair.second; + entry.value = init; + entry.index = lru_list_.PushBack(&pair).Index(); + } + + /*! + * \brief Evicts the least recently used nodes until the predicate returns false. + * \param predicate The function that returns true if eviction should continue. + * \param evict The function takes a value and returns true if the value can be evicted. + * This will be only called when the predicate returns true. + * If this function returns true, it should update the size information before return. + * \details This function will evict the least recently used nodes until the predicate returns + * false. The evict function will be called for each node to determine if it should be evicted. + */ + template + void LRUEvict(const Predicate& predicate, const Evict& evict) { + if (!predicate()) return; + + auto iter = lru_list_.begin(); + if (iter == lru_list_.end()) return; + + do { + auto& [key, entry] = **iter; + if (evict(entry.value)) { + iter = lru_list_.Erase(iter); + map_.erase(key); + } else { + ++iter; // simply skip those waiting for computation + } + } while (predicate() && iter != lru_list_.end()); + } + + std::unordered_map& GetMap() { return map_; } + + private: + std::unordered_map map_; + List*> lru_list_; +}; + +} // namespace details + +/** + * \brief A thread-safe key-value cache with on-demand computation and LRU eviction + * \tparam Key The type of keys used to lookup values. Should be hashable. + * \tparam Value The type of values stored in the cache + * \tparam Computer The functor that computes values for uncached keys + * \tparam SizeEstimator The functor that estimates the size of a value in bytes + * \details This cache provides thread-safe access to computed values with the following features: + * - Lazy computation: Values are only computed when first requested + * - LRU eviction: When the cache is full, the least recently used value is evicted + * - Thread safety: Uses reader-writer locks for concurrent reads + * \attention User should guarantee the following: + * 1. The policy class should provide a compute method that takes a key and returns a value. + * 2. The value type should have a MemorySize method that returns the size of the value in bytes. + */ +template +class ThreadSafeLRUCache { + private: + struct SizedValue { + Value value; + std::size_t size; + }; + + public: + inline static constexpr std::size_t kUnlimitedSize = static_cast(-1); + + explicit ThreadSafeLRUCache( + std::size_t max_size = kUnlimitedSize, + const Computer& computer = Computer{}, + const SizeEstimator& size_estimator = SizeEstimator{} + ) + : max_size_(max_size), computer_(computer), size_estimator_(size_estimator), cache_() {} + + std::size_t MaxMemorySize() const { return max_size_; } + std::size_t MemorySize() const { return current_size_; } + + Value Get(const Key& key) { + auto future = GetFuture(key); + return future.get().value; + } + + void Clear() { + // Remove all the ready entries. + const auto lock_map = std::lock_guard{map_mutex_}; + if (this->max_size_ == kUnlimitedSize) + cache_.GetMap().clear(); + else + cache_.LRUEvict( + [] { return true; }, + [&](const std::shared_future& value) { + // always evict and block until the value is ready + try { + current_size_ -= value.get().size; + } catch (...) { + // fine, just ignore the exception, size is not updated + } + return true; + } + ); + } + + private: + std::shared_future GetFuture(const Key& key) { + if (this->max_size_ == kUnlimitedSize) return GetFutureUnlimited(key); + auto& map = cache_.GetMap(); + + { + auto lock_map = std::shared_lock{map_mutex_}; + auto it = map.find(key); + if (it != map.end()) { + // We only need to hold LRU lock when shared lock is held here. + // When unique lock of map_mutex_ is held, only 1 thread can access the + // LRU list at the same time, so we do not need to hold the LRU lock then. + const auto lock_lru = std::lock_guard{lru_mutex_}; + return cache_.LRUVisit(*it); + } + } + + auto task = std::packaged_task{[this, &key] { + auto value = computer_(key); + auto result = SizedValue{value, size_estimator_(value)}; + current_size_ += result.size; + return result; + }}; + + auto lock_map = std::unique_lock{map_mutex_}; + auto [it, success] = map.try_emplace(key); + if (!success) return cache_.LRUVisit(*it); + + // in this case, we insert the task, and we need to compute the value + auto future = task.get_future().share(); + + // perform eviction if the cache is full + cache_.LRUInit(*it, future); + cache_.LRUEvict( + [&] { return current_size_ > max_size_; }, + [&](const std::shared_future& value) { + using namespace std::chrono_literals; + // if not ready, then do not wait and block here + if (value.wait_for(0s) != std::future_status::ready) return false; + try { + current_size_ -= value.get().size; + } catch (...) { + // fine, just ignore the exception, size is not updated + } + return true; + } + ); + + // perform the costly computation outside all locks + lock_map.unlock(); + task(); + return future; + } + + std::shared_future GetFutureUnlimited(const Key& key) { + auto& map = cache_.GetMap(); + + { + auto lock_map = std::shared_lock{map_mutex_}; + auto it = map.find(key); + if (it != map.end()) return it->second.value; + } + + auto task = std::packaged_task{[this, &key] { + auto value = computer_(key); + auto result = SizedValue{value, size_estimator_(value)}; + current_size_ += result.size; + return result; + }}; + + auto lock_map = std::unique_lock{map_mutex_}; + auto [it, success] = map.try_emplace(key); + if (!success) return it->second.value; + + auto future = task.get_future().share(); + it->second.value = future; + + // perform the costly computation outside all locks + lock_map.unlock(); + task(); + return future; + } + + private: + const std::size_t max_size_; + const Computer computer_; + const SizeEstimator size_estimator_; + details::LRUCacheImpl> cache_; + std::atomic_size_t current_size_{0}; + std::shared_mutex map_mutex_; + std::mutex lru_mutex_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_THREAD_SAFE_CACHE_H_ diff --git a/third_party/xgrammar/cpp/support/union_find_set.h b/third_party/xgrammar/cpp/support/union_find_set.h new file mode 100644 index 0000000000..be82a51674 --- /dev/null +++ b/third_party/xgrammar/cpp/support/union_find_set.h @@ -0,0 +1,105 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/union_find_set.h + */ +#ifndef XGRAMMAR_SUPPORT_UNION_FIND_SET_H_ +#define XGRAMMAR_SUPPORT_UNION_FIND_SET_H_ + +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +template +class UnionFindSet { + private: + std::unordered_map> element_to_parent_and_size_; + + public: + UnionFindSet() = default; + + /*! + * \brief Add a new element to the union-find set. + * \param element The element to add. + * \return True if the element was added successfully, false if it already exists. + */ + bool Add(const T& element) { + if (element_to_parent_and_size_.find(element) != element_to_parent_and_size_.end()) { + return false; // Element already exists. + } + element_to_parent_and_size_[element] = {element, 1}; + return true; + } + + /*! \brief Clear the union find set.*/ + void Clear() { element_to_parent_and_size_.clear(); } + + /*! + * \brief Find the representative of the set containing the element. + * \param element The element to find. + * \return The representative of the set containing the element. + */ + T Find(const T& element) { + XGRAMMAR_CHECK(element_to_parent_and_size_.find(element) != element_to_parent_and_size_.end()) + << "Element not found in union-find set."; + if (element_to_parent_and_size_[element].first != element) { + // Path compression. + element_to_parent_and_size_[element].first = Find(element_to_parent_and_size_[element].first); + } + return element_to_parent_and_size_[element].first; + } + + /*! + * \brief Union two elements into the same set. + * \param a The first element. + * \param b The second element. + */ + void Union(const T& a, const T& b) { + XGRAMMAR_CHECK(element_to_parent_and_size_.find(a) != element_to_parent_and_size_.end()) + << "Element " << a << " not found in union-find set."; + XGRAMMAR_CHECK(element_to_parent_and_size_.find(b) != element_to_parent_and_size_.end()) + << "Element " << b << " not found in union-find set."; + T root_a = Find(a); + T root_b = Find(b); + if (root_a == root_b) { + return; + } + if (element_to_parent_and_size_[root_a].second < element_to_parent_and_size_[root_b].second) { + std::swap(root_a, root_b); + // Make sure root_a is the larger set. + } + element_to_parent_and_size_[root_b].first = root_a; + element_to_parent_and_size_[root_a].second += element_to_parent_and_size_[root_b].second; + } + + int Count(const T& element) const { return element_to_parent_and_size_.count(element); } + + std::vector> GetAllSets() { + std::vector> result; + std::unordered_map root_to_set; + for (const auto& [value, _] : element_to_parent_and_size_) { + auto root = Find(value); + if (root_to_set.find(root) == root_to_set.end()) { + result.emplace_back(); + root_to_set[root] = result.size() - 1; + } + result[root_to_set[root]].push_back(value); + } + // Sort result to make it deterministic + for (auto& vec : result) { + std::sort(vec.begin(), vec.end()); + } + std::sort(result.begin(), result.end(), [](const std::vector& v1, const std::vector& v2) { + return v1.front() < v2.front(); + }); + return result; + } +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_UNION_FIND_SET_H_ diff --git a/third_party/xgrammar/cpp/support/utils.h b/third_party/xgrammar/cpp/support/utils.h new file mode 100644 index 0000000000..d8c974adee --- /dev/null +++ b/third_party/xgrammar/cpp/support/utils.h @@ -0,0 +1,451 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/utils.h + * \brief Utility functions. + */ +#ifndef XGRAMMAR_SUPPORT_UTILS_H_ +#define XGRAMMAR_SUPPORT_UTILS_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" + +/****************** Hash Library ******************/ + +namespace xgrammar { + +/*! + * \brief Hash and combine value into seed. + * \ref https://www.boost.org/doc/libs/1_84_0/boost/intrusive/detail/hash_combine.hpp + */ +inline void HashCombineBinary(uint64_t& seed, uint64_t value) { + seed ^= value + 0x9e3779b97f4a7c15ull + (seed << 6) + (seed >> 2); +} + +/*! + * \brief Find the hash sum of several size_t args. + */ +template +inline uint64_t HashCombine(Args... args) { + uint64_t seed = 0; + (..., HashCombineBinary(seed, args)); + return seed; +} + +/*! + * \brief Helper class to define the hash function for a struct by its members. + */ +template +struct HashByMembers { + std::size_t operator()(T const& x) const noexcept { + return HashCombine(std::hash>{}(x.*Members)...); + } +}; + +} // namespace xgrammar + +/*! + * \brief Define a hash function for a struct by its members in namespace std. Should be used + * outside of namespace xgrammar. + * \param Type The type of the struct. + * \param ... The member pointers of the struct. + * \example + * \code + * // In the global namespace + * XGRAMMAR_HASH_BY_MEMBERS(Type, &Type::member1, &Type::member2, &Type::member3); + * \endcode + */ +#define XGRAMMAR_HASH_BY_MEMBERS(Type, ...) \ + namespace std { \ + template <> \ + struct hash : public xgrammar::HashByMembers {}; \ + } + +/*! + * \brief Empty specialization of XGRAMMAR_HASH_BY_MEMBERS. + */ +#define XGRAMMAR_HASH_BY_MEMBERS_EMPTY(Type) \ + namespace std { \ + template <> \ + struct hash : public xgrammar::HashByMembers {}; \ + } + +namespace std { + +/*! + * \brief Define the hash function for std::pair. + */ +template +struct hash> { + size_t operator()(const std::pair& pair) const noexcept { + return xgrammar::HashCombine(std::hash{}(pair.first), std::hash{}(pair.second)); + } +}; + +/*! + * \brief Define the hash function for std::tuple. + */ +template +struct hash> { + size_t operator()(const std::tuple& tuple) const noexcept { + return std::apply( + [](const Args&... args) { return xgrammar::HashCombine(std::hash{}(args)...); }, tuple + ); + } +}; + +/*! + * \brief Define the hash function for std::vector. + */ +template +struct hash> { + size_t operator()(const std::vector& vec) const { + uint32_t seed = 0; + for (const auto& item : vec) { + xgrammar::HashCombineBinary(seed, std::hash{}(item)); + } + return seed; + } +}; + +} // namespace std + +namespace xgrammar { + +/****************** Result Library ******************/ + +/*! + * \brief A partial result type that can be used to construct a Result. Holds a result value or an + * error value. + * \tparam T The type of the value + * \tparam IsOk Whether the result is ok + */ +template +struct PartialResult { + template + PartialResult(Args&&... args) : value(std::forward(args)...) {} + T value; +}; + +/*! + * \brief Construct a success result with the arguments to construct a T. + * \tparam T The type of the success value + * \tparam Args The types of the arguments to construct a T + * \param args The arguments to construct a T + * \return A PartialResult with the arguments to construct a T + * \example + * \code + * // Call the constructor of T with the arguments + * return ResultOk(1, 2, 3); + * \endcode + */ +template +inline PartialResult ResultOk(Args&&... args) { + return PartialResult{std::forward(args)...}; +} + +/*! + * \brief Construct a success result with a universal reference (both lvalue and rvalue) + * \tparam T The type of the success value + * \param value The universal reference to the success value + * \return A PartialResult with the universal reference to the success value + * \example + * \code + * T value = T(1, 2, 3); + * // Move the value to the PartialResult + * return ResultOk(std::move(value)); + * \endcode + */ +template +inline PartialResult ResultOk(T&& value) { + return PartialResult{std::forward(value)}; +} + +/*! + * \brief Construct a error result with the arguments to construct a E. + * \tparam E The type of the error value. Default to std::runtime_error. + * \tparam Args The types of the arguments to construct a E + * \param args The arguments to construct a E + * \return A PartialResult with the arguments to construct a E + * \example + * \code + * // Construct a std::runtime_error with a error + * std::runtime_error error("Message"); + * return ResultErr(std::move(error)); + * \endcode + * \code + * // Construct a std::runtime_error with its argument + * return ResultErr("Error"); + * \endcode + * \code + * // Construct an E error with its argument + * return ResultErr("Error"); + * \endcode + */ +template +inline PartialResult ResultErr(Args&&... args) { + return PartialResult{std::forward(args)...}; +} + +/*! + * \brief Construct a error result with a universal reference (both lvalue and rvalue) + * \tparam E The type of the error value + * \param err The universal reference to the error value + * \return A PartialResult with the universal reference to the error value + * \example + * \code + * E err = E("Error"); + * // Move the err to the PartialResult + * return ResultErr(std::move(err)); + * \endcode + */ +template +inline PartialResult ResultErr(E&& err) { + return PartialResult{std::forward(err)}; +} + +/*! + * \brief An always-move Result type similar to Rust's Result, representing either success (Ok) or + * failure (Err). It always uses move semantics for the success and error values. + * \tparam T The type of the success value + * \tparam E The type of the error value + * + * \note The Ok and Err constructor, and all methods of this class (except for ValueRef and ErrRef) + * accept only rvalue references as parameters for performance reasons. You should use std::move to + * convert a Result to an rvalue reference before invoking these methods. Examples for move + * semantics are shown below. + * + * \example Construct a success result with a rvalue reference + * \code + * T value; + * return Result::Ok(std::move(value)); + * \endcode + * \example Construct a error result with a rvalue reference of std::runtime_error + * \code + * std::runtime_error error_msg = std::runtime_error("Error"); + * return Result::Err(std::move(error_msg)); + * \endcode + * \example Construct a error result with a std::runtime_error object constructed with a string + * \code + * std::string error_msg = "Error"; + * return Result::Err(std::move(error_msg)); + * \endcode + * \example Unwrap the rvalue reference of the result + * \code + * Result result = func(); + * if (result.IsOk()) { + * T result_val = std::move(result).Unwrap(); + * } else { + * std::runtime_error error_msg = std::move(result).UnwrapErr(); + * } + * \endcode + */ +template +class Result { + private: + static_assert(!std::is_same_v, "T and E cannot be the same type"); + + public: + /*! \brief Default constructor is deleted to avoid accidental use */ + Result() = delete; + + /*! \brief Construct from Result::Ok */ + template >>> + Result(PartialResult&& partial_result) + : data_(std::in_place_type, std::forward(partial_result.value)) {} + + /*! \brief Construct from Result::Err */ + template >>> + Result(PartialResult&& partial_result) + : data_(std::in_place_type, std::forward(partial_result.value)) {} + + /*! \brief Check if Result contains success value */ + bool IsOk() const { return std::holds_alternative(data_); } + + /*! \brief Check if Result contains error */ + bool IsErr() const { return std::holds_alternative(data_); } + + /*! \brief Get the success value. It assumes (or checks if in debug mode) the result is ok. */ + T Unwrap() && { + XGRAMMAR_DCHECK(IsOk()) << "Called Unwrap() on an Err value"; + return std::get(std::move(data_)); + } + + /*! \brief Get the error value. It assumes (or checks if in debug mode) the result is an error. */ + E UnwrapErr() && { + XGRAMMAR_DCHECK(IsErr()) << "Called UnwrapErr() on an Ok value"; + return std::get(std::move(data_)); + } + + /*! \brief Get the success value if present, otherwise return the provided default */ + T UnwrapOr(T default_value) && { + return IsOk() ? std::get(std::move(data_)) : std::move(default_value); + } + + /*! \brief Map success value to new type using provided function */ + template >> + Result Map(F&& f) && { + if (IsOk()) { + return ResultOk(f(std::get(std::move(data_)))); + } + return ResultErr(std::get(std::move(data_))); + } + + /*! \brief Map error value to new type using provided function */ + template >> + Result MapErr(F&& f) && { + if (IsErr()) { + return ResultErr(f(std::get(std::move(data_)))); + } + return ResultOk(std::get(std::move(data_))); + } + + /*! + * \brief Convert a Result to a Result. U should be convertible to T, and V should be + * convertible to E. + */ + template + static Result Convert(Result&& result) { + if (result.IsOk()) { + return ResultOk(std::move(result).Unwrap()); + } + return ResultErr(std::move(result).UnwrapErr()); + } + + /*! \brief Get a std::variant from the result. */ + std::variant ToVariant() && { return std::move(data_); } + + /*! + * \brief Get a reference to the success value. It assumes (or checks if in debug mode) the + * result is ok. + */ + T& ValueRef() & { + XGRAMMAR_DCHECK(IsOk()) << "Called ValueRef() on an Err value"; + return std::get(data_); + } + + /*! + * \brief Get a reference to the error value. It assumes (or checks if in debug mode) the + * result is an error. + */ + E& ErrRef() & { + XGRAMMAR_DCHECK(IsErr()) << "Called ErrRef() on an Ok value"; + return std::get(data_); + } + + private: + // in-place construct T in variant + template + explicit Result(std::in_place_type_t, Args&&... args) + : data_(std::in_place_type, std::forward(args)...) {} + + // in-place construct E in variant + template + explicit Result(std::in_place_type_t, Args&&... args) + : data_(std::in_place_type, std::forward(args)...) {} + + std::variant data_; +}; + +/****************** Misc ******************/ + +// Sometimes GCC fails to detect some branches will not return, such as when we use LOG(FATAL) +// to raise an error. This macro manually mark them as unreachable to avoid warnings. +#ifdef __GNUC__ +#define XGRAMMAR_UNREACHABLE() __builtin_unreachable() +#else +#define XGRAMMAR_UNREACHABLE() +#endif + +/*! + * \brief An error class that contains a type. The type can be an enum. + */ +template +class TypedError : public std::runtime_error { + public: + explicit TypedError(T type, const std::string& msg) : std::runtime_error(msg), type_(type) {} + const T& Type() const noexcept { return type_; } + + private: + T type_; +}; + +/** + * \brief Helper function to compare two objects by their members. + */ +template +constexpr bool EqualByMembers(const T& lhs, const T& rhs) noexcept { + return std::tie(lhs.*Ms...) == std::tie(rhs.*Ms...); +} + +/** + * \brief Define == and != operator for a struct by its members. + * \param Type The type of the struct. Must be under namespace xgrammar. + * \param ... The member pointers of the struct. + * \example + * \code + * struct Type { + * int member1; + * std::string member2; + * double member3; + * + * XGRAMMAR_EQUAL_BY_MEMBERS(Type, &Type::member1, &Type::member2, &Type::member3); + * }; + * \endcode + */ +#define XGRAMMAR_EQUAL_BY_MEMBERS(Type, ...) \ + friend bool operator==(const Type& lhs, const Type& rhs) noexcept { \ + return EqualByMembers(lhs, rhs); \ + } \ + friend bool operator!=(const Type& lhs, const Type& rhs) noexcept { return !(lhs == rhs); } + +/*! + * \brief Empty specialization of XGRAMMAR_EQUAL_BY_MEMBERS. + */ +#define XGRAMMAR_EQUAL_BY_MEMBERS_EMPTY(Type) \ + friend bool operator==(const Type& lhs, const Type& rhs) noexcept { return true; } \ + friend bool operator!=(const Type& lhs, const Type& rhs) noexcept { return false; } + +/*! + * \brief Throw an error from a variant of multiple error types. + * \param error_variant The variant of multiple error types. + * \tparam Args The types of the error types. Each type should inherit from std::runtime_error. + */ +template +[[noreturn]] void ThrowVariantError(const std::variant& error_variant) { + std::visit([](const auto& e) { throw e; }, error_variant); + XGRAMMAR_UNREACHABLE(); +} + +/*! + * \brief Get the message from a variant of multiple error types. + * \param error_variant The variant of multiple error types. + * \return The message from the error variant. + * \tparam Args The types of the error types. Each type should inherit from std::runtime_error. + */ +template +std::string GetMessageFromVariantError(const std::variant& error_variant) { + return std::visit([](const auto& e) { return e.what(); }, error_variant); +} + +/*! + * \brief Get the type name from a variant of XGrammarError types (each has GetType()). + * \param error_variant The variant of multiple error types. + * \return The type string from the error variant (e.g. "DeserializeVersionError"). + * \tparam Args The types of the error types. Each type should have GetType() const. + */ +template +std::string GetTypeFromVariantError(const std::variant& error_variant) { + return std::visit([](const auto& e) { return e.GetType(); }, error_variant); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_UTILS_H_ diff --git a/third_party/xgrammar/cpp/testing.cc b/third_party/xgrammar/cpp/testing.cc new file mode 100644 index 0000000000..b7b5bdf156 --- /dev/null +++ b/third_party/xgrammar/cpp/testing.cc @@ -0,0 +1,63 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/testing.cc + */ +#include "testing.h" + +#include + +#include +#include +#include +#include +#include + +#include "grammar_impl.h" +#include "grammar_parser.h" +#include "support/encoding.h" + +namespace xgrammar { + +std::string PrintTokenByIds( + const std::vector& token_ids, const TokenizerInfo& tokenizer_info, int max_print_num +) { + std::stringstream ss; + const auto& sorted_decoded_vocab = tokenizer_info.GetDecodedVocab(); + ss << "["; + int print_num = std::min(static_cast(token_ids.size()), max_print_num); + for (int i = 0; i < print_num; ++i) { + ss << "#" << token_ids[i] << " <" << EscapeString(sorted_decoded_vocab[token_ids[i]]) << ">"; + if (i < print_num - 1) { + ss << ", "; + } + } + if (static_cast(token_ids.size()) > max_print_num) { + ss << ", ..."; + } + ss << "]"; + return ss.str(); +} + +Grammar _EBNFToGrammarNoNormalization( + const std::string& ebnf_string, const std::string& root_rule_name +) { + return ParseEBNF(ebnf_string, root_rule_name); +} + +std::string _PrintGrammarFSMs(const Grammar& grammar) { + XGRAMMAR_CHECK(static_cast(grammar->per_rule_fsms.size()) == grammar->NumRules()) + << "The grammar has no per-rule FSMs; build them first"; + std::string result; + for (int i = 0; i < grammar->NumRules(); i++) { + result += "Rule " + std::to_string(i) + ": " + grammar->GetRule(i).name + ", FSM: "; + if (grammar->per_rule_fsms[i].has_value()) { + result += grammar->per_rule_fsms[i]->GetFsm().ToString(); + } else { + result += "None"; + } + result += "\n"; + } + return result; +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/testing.h b/third_party/xgrammar/cpp/testing.h new file mode 100644 index 0000000000..71b3c6331a --- /dev/null +++ b/third_party/xgrammar/cpp/testing.h @@ -0,0 +1,30 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/testing.h + * \brief The header testing utilities. + */ +#ifndef XGRAMMAR_TESTING_H_ +#define XGRAMMAR_TESTING_H_ + +#include +#include + +#include +#include +#include + +namespace xgrammar { + +std::string PrintTokenByIds( + const std::vector& token_ids, const TokenizerInfo& tokenizer_info, int max_print_num +); + +Grammar _EBNFToGrammarNoNormalization( + const std::string& ebnf_string, const std::string& root_rule_name +); + +std::string _PrintGrammarFSMs(const Grammar& grammar); + +} // namespace xgrammar + +#endif // XGRAMMAR_TESTING_H_ diff --git a/third_party/xgrammar/cpp/tokenizer_info.cc b/third_party/xgrammar/cpp/tokenizer_info.cc new file mode 100644 index 0000000000..b580d3c30f --- /dev/null +++ b/third_party/xgrammar/cpp/tokenizer_info.cc @@ -0,0 +1,562 @@ +/*! + * Copyright (c) 2023 by Contributors + * \file xgrammar/tokenizer_info.cc + */ + +#include "xgrammar/tokenizer_info.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/json_parse.h" +#include "support/json_serializer.h" +#include "support/logging.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/************* Token decoders: ByteFallback and ByteLevel *************/ + +class TokenDecoder { + public: + /*! + * \brief Post-process a raw token to the actual token with the given post-processing method. + */ + static std::string DecodeToken(const std::string& token, VocabType vocab_type) { + // TODO(yixin): Avoid allocating new string in decoder calls + if (vocab_type == VocabType::BYTE_FALLBACK) { + return SpaceReplacerDecoder(ByteFallbackDecoder(token)); + } else if (vocab_type == VocabType::BYTE_LEVEL) { + return ByteLevelDecoder(token); + } else { + return token; + } + } + + private: + /*! \brief ByteFallback decoder: transform tokens like <0x1B> to hex char byte 1B */ + static std::string ByteFallbackDecoder(const std::string& token) { + if (token.length() == 6 && token.substr(0, 3) == "<0x" && token.back() == '>') { + int byte = 0; + for (int i = 0; i < 2; ++i) { + byte *= 16; + byte += token[3 + i] >= '0' && token[3 + i] <= '9' ? token[3 + i] - '0' + : token[3 + i] - 'A' + 10; + } + XGRAMMAR_CHECK(byte >= 0 && byte < 256); + return std::string(/*n=*/1, static_cast(byte)); + } + return token; + } + + /*! \brief SpaceReplacer decoder: transform "\u2581" back to space */ + static std::string SpaceReplacerDecoder(const std::string& token) { + // \u2581 is the unicode for "lower one eighth block" + // UTF8 encoding for \u2581 is 0xE2 0x96 0x81 + std::string result; + for (int i = 0; i < static_cast(token.size()); ++i) { + if (i + 2 < static_cast(token.size()) && token[i] == char(0xE2) && + token[i + 1] == char(0x96) && token[i + 2] == char(0x81)) { + result += ' '; + i += 2; + } else { + result += token[i]; + } + } + return result; + } + + /*! + * \brief ByteLevel decoder: inverses the bytes-to-unicode transformation in the encoding + * process as in + * https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59 + */ + static std::string ByteLevelDecoder(const std::string& token) { + // The inverse map of bytes_to_unicode. -1 means there is no mapping to this unicode. + static const std::array char_to_byte_map = { + // clang-format off + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, + 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, + 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, + 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, + 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, -1, + 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, + 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, + 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, + 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, + 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 127, 128, + 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, + 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 173 + // clang-format on + }; + + auto unicode_codepoints = ParseUTF8(token.c_str(), false); + if (unicode_codepoints.size() == 1 && unicode_codepoints[0] == kInvalidUTF8) { + return token; + } + + std::string decoded; + decoded.reserve(unicode_codepoints.size()); + + for (auto unicode_codepoint : unicode_codepoints) { + XGRAMMAR_CHECK(unicode_codepoint >= 0); + if (unicode_codepoint >= static_cast(char_to_byte_map.size()) || + char_to_byte_map[unicode_codepoint] == -1) { + // If there is no mapping, return the original token + return token; + } + decoded += static_cast(char_to_byte_map[unicode_codepoint]); + } + return decoded; + } +}; + +/************* Metadata detection from huggingface tokenizer.json *************/ + +class HFTokenizerAnalyzer { + public: + /*! + * \brief Detect the vocabulary type from tokenizer.json. + * \details Find {"type": "ByteFallback"} or {"type": "ByteLevel"} in "decoder" field of the + * tokenizer. + */ + static VocabType DetectVocabType(const picojson::object& hf_tokenizer_obj) { +#define CHECK_AND_WARNING(condition, message) \ + if (!(condition)) { \ + XGRAMMAR_LOG(WARNING) << "Vocab type detection failed: (" #condition \ + << ") is false: " << (message) << " Using RAW VocabType by default."; \ + return VocabType::RAW; \ + } + + CHECK_AND_WARNING( + hf_tokenizer_obj.count("decoder") && hf_tokenizer_obj.at("decoder").is(), + "Decoder field is not found in tokenizer.json." + ); + + auto decoder_obj = hf_tokenizer_obj.at("decoder").get(); + CHECK_AND_WARNING( + decoder_obj.count("type") && decoder_obj.at("type").is(), + "Type field is not found in decoder field" + ); + auto type = decoder_obj.at("type").get(); + + std::vector decoders; + if (type == "Sequence") { + CHECK_AND_WARNING( + decoder_obj.count("decoders") && decoder_obj.at("decoders").is(), + "Decoders field is not found in a Sequence decoder" + ); + decoders = decoder_obj.at("decoders").get(); + } else { + decoders.emplace_back(hf_tokenizer_obj.at("decoder")); + } + + for (const auto& decoder : decoders) { + CHECK_AND_WARNING(decoder.is(), "Decoder is not an object"); + auto decoder_obj = decoder.get(); + CHECK_AND_WARNING( + decoder_obj.count("type") && decoder_obj.at("type").is(), + "Type field is not found in decoder field" + ); + auto type = decoder_obj.at("type").get(); + if (type == "ByteLevel") { + return VocabType::BYTE_LEVEL; + } else if (type == "ByteFallback") { + return VocabType::BYTE_FALLBACK; + } + } + + // If neither byte_level nor byte_fallback decoder is detected, return RAW. + return VocabType::RAW; + +#undef CHECK_AND_WARNING + } + + static bool DetectPrependNormalizer(const picojson::object& hf_tokenizer_obj) { + if (!hf_tokenizer_obj.count("normalizer") || + !hf_tokenizer_obj.at("normalizer").is()) { + return false; + } + + const picojson::value& normalizer_value = hf_tokenizer_obj.at("normalizer"); + if (!normalizer_value.is()) { + return false; + } + const picojson::object& normalizer_obj = normalizer_value.get(); + if (!normalizer_obj.count("type") || !normalizer_obj.at("type").is()) { + return false; + } + auto type = normalizer_obj.at("type").get(); + + std::vector normalizers; + if (type == "Sequence") { + if (!normalizer_obj.count("normalizers") || + !normalizer_obj.at("normalizers").is()) { + return false; + } + normalizers = normalizer_obj.at("normalizers").get(); + } else { + normalizers.emplace_back(normalizer_value); + } + + for (const auto& normalizer : normalizers) { + if (!normalizer.is()) { + continue; + } + auto normalizer_obj = normalizer.get(); + if (!normalizer_obj.count("type") || !normalizer_obj.at("type").is()) { + continue; + } + auto type = normalizer_obj.at("type").get(); + if (type == "Prepend" && normalizer_obj.count("prepend") && + normalizer_obj.at("prepend").is() && + normalizer_obj.at("prepend").get() == "▁") { + return true; + } + } + return false; + } + + static bool DetectMetaspacePreTokenizer(const picojson::object& hf_tokenizer_obj) { + if (!hf_tokenizer_obj.count("pre_tokenizer") || + !hf_tokenizer_obj.at("pre_tokenizer").is()) { + return false; + } + auto pre_tokenizer_obj = hf_tokenizer_obj.at("pre_tokenizer").get(); + if (!pre_tokenizer_obj.count("type") || !pre_tokenizer_obj.at("type").is()) { + return false; + } + auto type = pre_tokenizer_obj.at("type").get(); + if (!pre_tokenizer_obj.count("prepend_scheme") || + !pre_tokenizer_obj.at("prepend_scheme").is()) { + return false; + } + auto prepend_scheme = pre_tokenizer_obj.at("prepend_scheme").get(); + return type == "Metaspace" && (prepend_scheme == "always" || prepend_scheme == "first"); + } + + /*! + * \brief Detect whether add prefix space from tokenizer.json. + * \details Find {"type": "Prepend", "prepend": "▁"} in "normalizer" field of the tokenizer, or + * "pre_tokenizer": {"type": "Metaspace", "prepend_scheme": "always" | "first"} in the tokenizer. + */ + static bool DetectAddPrefixSpace(const picojson::object& hf_tokenizer_obj) { + return DetectPrependNormalizer(hf_tokenizer_obj) || + DetectMetaspacePreTokenizer(hf_tokenizer_obj); + } +}; + +/************* TokenizerInfo::Impl *************/ + +bool TokenizerInfo::Impl::IsSpecialToken(const std::string& token) { return token == ""; } + +TokenizerInfo::Impl::Impl( + const std::vector& encoded_vocab, + VocabType vocab_type, + std::optional vocab_size, + std::optional> stop_token_ids, + bool add_prefix_space +) + : vocab_type_(vocab_type), + vocab_size_(vocab_size.value_or(encoded_vocab.size())), + add_prefix_space_(add_prefix_space) { + // vocab_size only ever pads the encoded vocab with special ids; a value below the number of + // real tokens would leave token ids without a slot and corrupt the id -> index table below. + XGRAMMAR_CHECK(vocab_size_ >= static_cast(encoded_vocab.size())) + << "vocab_size (" << vocab_size_ << ") must be at least the number of tokens in the vocab (" + << encoded_vocab.size() << ")."; + decoded_vocab_.reserve(encoded_vocab.size()); + sorted_decoded_vocab_.reserve(encoded_vocab.size()); + for (int i = 0; i < static_cast(encoded_vocab.size()); ++i) { + const std::string& token = TokenDecoder::DecodeToken(encoded_vocab[i], vocab_type_); + decoded_vocab_.push_back(token); + if ((!stop_token_ids && DETECTION_STOP_TOKENS.count(token)) || + (stop_token_ids && + std::find(stop_token_ids->begin(), stop_token_ids->end(), i) != stop_token_ids->end())) { + stop_token_ids_.push_back(i); + } else if (IsSpecialToken(token)) { + special_token_ids_.push_back(i); + } else { + sorted_decoded_vocab_.push_back({i, token}); + } + } + for (int i = encoded_vocab.size(); i < vocab_size_; ++i) { + special_token_ids_.push_back(i); + } + + auto f_compare_token = [](const std::pair& a, + const std::pair& b) { + return a.second < b.second; + }; + std::sort(sorted_decoded_vocab_.begin(), sorted_decoded_vocab_.end(), f_compare_token); + + BuildTokenIdToSortedVocabIndex(); + + // The value means: the subtree is [i, trie_subtree_nodes_range[i]). + trie_subtree_nodes_range_.resize(sorted_decoded_vocab_.size(), 0); + std::stack> prefix_stack; + for (size_t i = 0; i < sorted_decoded_vocab_.size(); ++i) { + const auto& token = sorted_decoded_vocab_[i].second; + while ((!prefix_stack.empty()) && (token.find(prefix_stack.top().first) == std::string::npos)) { + const auto& top_pair = prefix_stack.top(); + trie_subtree_nodes_range_[top_pair.second] = i; + prefix_stack.pop(); + } + prefix_stack.push({token, i}); + } + while (!prefix_stack.empty()) { + const auto& top_pair = prefix_stack.top(); + trie_subtree_nodes_range_[top_pair.second] = sorted_decoded_vocab_.size(); + prefix_stack.pop(); + } + BuildTokenCharData(); +} + +std::optional TokenizerInfo::Impl::Validate() const { + const int64_t num_tokens = decoded_vocab_.size(); + if (vocab_size_ < num_tokens) { + return "vocab_size " + std::to_string(vocab_size_) + " is smaller than the number of tokens " + + std::to_string(num_tokens); + } + for (const auto& [token_id, token] : sorted_decoded_vocab_) { + if (token_id < 0 || token_id >= num_tokens) { + return "sorted_decoded_vocab contains token id " + std::to_string(token_id) + " out of range"; + } + } + if (trie_subtree_nodes_range_.size() != sorted_decoded_vocab_.size()) { + return "trie_subtree_nodes_range must have one entry per sorted token"; + } + auto id_ok = [&](int32_t token_id) { return token_id >= 0 && token_id < vocab_size_; }; + if (!std::all_of(stop_token_ids_.begin(), stop_token_ids_.end(), id_ok) || + !std::all_of(special_token_ids_.begin(), special_token_ids_.end(), id_ok)) { + return "stop_token_ids or special_token_ids contains a token id out of range"; + } + return std::nullopt; +} + +void TokenizerInfo::Impl::BuildTokenIdToSortedVocabIndex() { + token_id_to_sorted_vocab_index_.assign(vocab_size_, -1); + for (int32_t i = 0; i < static_cast(sorted_decoded_vocab_.size()); ++i) { + token_id_to_sorted_vocab_index_[sorted_decoded_vocab_[i].first] = i; + } +} + +void TokenizerInfo::Impl::BuildTokenCharData() { + token_char_counts_.assign(sorted_decoded_vocab_.size(), 0); + int32_t max_chars = 0; + for (int32_t index = 0; index < static_cast(sorted_decoded_vocab_.size()); ++index) { + int32_t count = 0; + for (uint8_t byte : sorted_decoded_vocab_[index].second) { + count += (byte & 0xC0) != 0x80; + } + token_char_counts_[index] = count; + max_chars = std::max(max_chars, count); + } + max_token_chars_ = max_chars; +} + +const std::vector& TokenizerInfo::Impl::GetTokenCharCounts() const { + return token_char_counts_; +} + +int32_t TokenizerInfo::Impl::GetMaxTokenChars() const { return max_token_chars_; } + +std::string TokenizerInfo::Impl::DumpMetadata() const { + return DumpMetadataValue().serialize(false); +} + +picojson::value TokenizerInfo::Impl::DumpMetadataValue() const { + picojson::object obj; + obj["vocab_type"] = picojson::value(static_cast(vocab_type_)); + obj["vocab_size"] = picojson::value(static_cast(vocab_size_)); + obj["add_prefix_space"] = picojson::value(add_prefix_space_); + picojson::array stop_token_ids_array; + for (auto id : stop_token_ids_) { + stop_token_ids_array.push_back(picojson::value(static_cast(id))); + } + obj["stop_token_ids"] = picojson::value(std::move(stop_token_ids_array)); + + return picojson::value(std::move(obj)); +} + +std::optional TokenizerInfo::Impl::CheckMetadataMatch( + const picojson::value& metadata +) const { + if (!metadata.is()) { + return std::runtime_error("Expect an object"); + } + const auto& object = metadata.get(); + if (object.find("vocab_type") == object.end()) { + return std::runtime_error("Missing 'vocab_type' in metadata"); + } + auto vocab_type = object.at("vocab_type").get(); + if (vocab_type != static_cast(vocab_type_)) { + return std::runtime_error( + "Vocab type mismatch: " + std::to_string(vocab_type) + + " != " + std::to_string(static_cast(vocab_type_)) + ); + } + if (object.find("vocab_size") == object.end()) { + return std::runtime_error("Missing 'vocab_size' in metadata"); + } + auto vocab_size = object.at("vocab_size").get(); + if (vocab_size != vocab_size_) { + return std::runtime_error( + "Vocab size mismatch: " + std::to_string(vocab_size) + " != " + std::to_string(vocab_size_) + ); + } + if (object.find("add_prefix_space") == object.end()) { + return std::runtime_error("Missing 'add_prefix_space' in metadata"); + } + auto add_prefix_space = object.at("add_prefix_space").get(); + if (add_prefix_space != add_prefix_space_) { + return std::runtime_error( + "Add prefix space mismatch: " + std::to_string(add_prefix_space) + + " != " + std::to_string(add_prefix_space_) + ); + } + if (object.find("stop_token_ids") == object.end()) { + return std::runtime_error("Missing 'stop_token_ids' in metadata"); + } + auto stop_token_ids = object.at("stop_token_ids").get(); + std::vector stop_token_ids_vec; + stop_token_ids_vec.reserve(stop_token_ids.size()); + for (const auto& id : stop_token_ids) { + if (!id.is()) { + return std::runtime_error("Stop token id is not an integer"); + } + stop_token_ids_vec.push_back(static_cast(id.get())); + } + if (stop_token_ids_vec != stop_token_ids_) { + return std::runtime_error("Stop token ids mismatch"); + } + return std::nullopt; +} + +std::shared_ptr TokenizerInfo::Impl::FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata +) { + picojson::value v; + std::string err = ParseJSON(v, metadata); + XGRAMMAR_CHECK(err.empty()) << "Failed to parse metadata: " << err; + + const picojson::object& obj = v.get(); + + XGRAMMAR_CHECK(obj.count("vocab_type") && obj["vocab_type"].is()) + << "Missing or invalid 'vocab_type' in metadata"; + int vocab_type_int = static_cast(obj["vocab_type"].get()); + XGRAMMAR_CHECK(vocab_type_int == 0 || vocab_type_int == 1 || vocab_type_int == 2) + << "Invalid vocab_type in metadata: " << vocab_type_int; + VocabType vocab_type = static_cast(vocab_type_int); + + XGRAMMAR_CHECK(obj.count("vocab_size") && obj["vocab_size"].is()) + << "Missing or invalid 'vocab_size' in metadata"; + int vocab_size = static_cast(obj["vocab_size"].get()); + + XGRAMMAR_CHECK(obj.count("add_prefix_space") && obj["add_prefix_space"].is()) + << "Missing or invalid 'add_prefix_space' in metadata"; + bool add_prefix_space = obj["add_prefix_space"].get(); + + std::vector stop_token_ids; + XGRAMMAR_CHECK(obj.count("stop_token_ids") && obj["stop_token_ids"].is()) + << "Missing or invalid 'stop_token_ids' in metadata"; + for (const auto& id : obj["stop_token_ids"].get()) { + XGRAMMAR_CHECK(id.is()) << "Stop token id is not an integer"; + stop_token_ids.push_back(static_cast(id.get())); + } + return std::make_shared( + encoded_vocab, vocab_type, vocab_size, stop_token_ids, add_prefix_space + ); +} + +std::string TokenizerInfo::Impl::DetectMetadataFromHF(const std::string& backend_str) { + picojson::value v; + std::string err = ParseJSON(v, backend_str); + XGRAMMAR_CHECK(err.empty() && v.is()) << "Failed to parse JSON object: " << err; + const picojson::object& obj = v.get(); + VocabType vocab_type = HFTokenizerAnalyzer::DetectVocabType(obj); + bool add_prefix_space = HFTokenizerAnalyzer::DetectAddPrefixSpace(obj); + + // Serialize the metadata + picojson::object metadata_obj; + metadata_obj["vocab_type"] = picojson::value(static_cast(vocab_type)); + metadata_obj["add_prefix_space"] = picojson::value(add_prefix_space); + return picojson::value(metadata_obj).serialize(false); +} + +/************* TokenizerInfo *************/ + +TokenizerInfo::TokenizerInfo( + const std::vector& encoded_vocab, + VocabType vocab_type, + std::optional vocab_size, + std::optional> stop_token_ids, + bool add_prefix_space +) + : pimpl_(std::make_shared( + encoded_vocab, vocab_type, vocab_size, stop_token_ids, add_prefix_space + )) {} + +int TokenizerInfo::GetVocabSize() const { return pimpl_->GetVocabSize(); } +VocabType TokenizerInfo::GetVocabType() const { return pimpl_->GetVocabType(); } +bool TokenizerInfo::GetAddPrefixSpace() const { return pimpl_->GetAddPrefixSpace(); } +const std::vector& TokenizerInfo::GetDecodedVocab() const { + return pimpl_->GetDecodedVocab(); +} +const std::vector& TokenizerInfo::GetStopTokenIds() const { + return pimpl_->GetStopTokenIds(); +} +const std::vector& TokenizerInfo::GetSpecialTokenIds() const { + return pimpl_->GetSpecialTokenIds(); +} +const std::vector>& TokenizerInfo::GetSortedDecodedVocab() const { + return pimpl_->GetSortedDecodedVocab(); +} + +const std::vector& TokenizerInfo::GetTrieSubtreeNodesRange() const { + return pimpl_->GetTrieSubtreeNodesRange(); +} + +std::string TokenizerInfo::DumpMetadata() const { return pimpl_->DumpMetadata(); } + +TokenizerInfo TokenizerInfo::FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata +) { + return TokenizerInfo(Impl::FromVocabAndMetadata(encoded_vocab, metadata)); +} + +std::string TokenizerInfo::DetectMetadataFromHF(const std::string& backend_str) { + return Impl::DetectMetadataFromHF(backend_str); +} + +std::string TokenizerInfo::SerializeJSON() const { return AutoSerializeJSON(*this, true); } + +std::variant TokenizerInfo::DeserializeJSON( + const std::string& json_string +) { + TokenizerInfo tokenizer_info{NullObj()}; + if (auto err = AutoDeserializeJSON(&tokenizer_info, json_string, true, "TokenizerInfo")) { + return err.value(); + } + // Derived per-token arrays are not serialized; rebuild them like the constructor does. + tokenizer_info->BuildTokenIdToSortedVocabIndex(); + tokenizer_info->BuildTokenCharData(); + return tokenizer_info; +} + +} // namespace xgrammar diff --git a/third_party/xgrammar/cpp/tokenizer_info_impl.h b/third_party/xgrammar/cpp/tokenizer_info_impl.h new file mode 100644 index 0000000000..863254aff4 --- /dev/null +++ b/third_party/xgrammar/cpp/tokenizer_info_impl.h @@ -0,0 +1,143 @@ +#ifndef XGRAMMAR_TOKENIZER_INFO_IMPL_H_ +#define XGRAMMAR_TOKENIZER_INFO_IMPL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "support/reflection.h" +#include "xgrammar/tokenizer_info.h" + +namespace xgrammar { + +class TokenizerInfo::Impl { + public: + explicit Impl() = default; + + Impl( + const std::vector& encoded_vocab, + VocabType vocab_type, + std::optional vocab_size, + std::optional> stop_token_ids, + bool add_prefix_space + ); + + VocabType GetVocabType() const { return vocab_type_; } + bool GetAddPrefixSpace() const { return add_prefix_space_; } + int GetVocabSize() const { return vocab_size_; } + const std::vector& GetDecodedVocab() { return decoded_vocab_; } + const std::vector& GetStopTokenIds() const { return stop_token_ids_; } + const std::vector& GetSpecialTokenIds() const { return special_token_ids_; } + const std::vector>& GetSortedDecodedVocab() const { + return sorted_decoded_vocab_; + } + const std::vector& GetTrieSubtreeNodesRange() const { return trie_subtree_nodes_range_; } + const std::vector& GetTokenIdToSortedVocabIndex() const { + return token_id_to_sorted_vocab_index_; + } + const std::vector& GetTokenCharCounts() const; + int32_t GetMaxTokenChars() const; + void BuildTokenIdToSortedVocabIndex(); + void BuildTokenCharData(); + + /*! + * \brief Check that the token ids and per-token arrays are consistent with the vocabulary. Used + * after deserialization, where the fields are restored verbatim. + * \return An error message if the tokenizer info is malformed. + */ + std::optional Validate() const; + + std::string DumpMetadata() const; + picojson::value DumpMetadataValue() const; + + static std::shared_ptr FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata + ); + + std::optional CheckMetadataMatch(const picojson::value& metadata) const; + + static std::string DetectMetadataFromHF(const std::string& backend_str); + + bool operator==(const Impl& other) const; + + private: + static bool IsSpecialToken(const std::string& decoded_token); + + /*! \brief The vocabulary type. */ + VocabType vocab_type_; + /*! \brief The size of the vocabulary. */ + int vocab_size_; + /*! \brief Whether to add prefix space. */ + bool add_prefix_space_; + + /*! \brief The vocabulary. Special tokens are included. */ + std::vector decoded_vocab_; + /*! \brief All (id, token) pairs sorted in lexicographic order. This sorting is done to + * maximize prefix reuse during matching. Special tokens and stop tokens are not included. */ + std::vector> sorted_decoded_vocab_; + /*! \brief A pesudo-trie. trie_subtree_nodes_range[i] stores how many nodes there are in the + * subtree. */ + std::vector trie_subtree_nodes_range_; + /*! \brief The stop tokens. When the GrammarMatcher can reach the end of the grammar, + * stop tokens can be accepted. */ + std::vector stop_token_ids_; + /*! \brief The special tokens. These tokens are ignored (masked out) during the grammar-guided + * generation. */ + std::vector special_token_ids_; + /*! \brief Reverse mapping: token_id -> index in sorted_decoded_vocab_. -1 if not present. */ + std::vector token_id_to_sorted_vocab_index_; + /*! \brief Unicode codepoint counts for the sorted decoded vocabulary. */ + int32_t max_token_chars_ = 0; + std::vector token_char_counts_; + + /*! + * \brief The tokens used to detect stop tokens from the vocabulary. + * + * LLaMA2: + * LLaMA3: <|end_of_text|>, <|eot_id|> + * Phi-2: <|endoftext|> + * Gemma: , + * DeepSeek-V2: <｜end▁of▁sentence｜> + */ + inline static const std::unordered_set DETECTION_STOP_TOKENS = { + "", + "<|end_of_text|>", + "<|eot_id|>", + "<|endoftext|>", + "", + "<|eos|>", + "", + "<｜end▁of▁sentence｜>" + }; + + friend struct member_trait; +}; + +XGRAMMAR_MEMBER_TABLE( + TokenizerInfo::Impl, + "vocab_type", + &TokenizerInfo::Impl::vocab_type_, + "vocab_size", + &TokenizerInfo::Impl::vocab_size_, + "add_prefix_space", + &TokenizerInfo::Impl::add_prefix_space_, + "stop_token_ids", + &TokenizerInfo::Impl::stop_token_ids_, + "special_token_ids", + &TokenizerInfo::Impl::special_token_ids_, + "decoded_vocab", + &TokenizerInfo::Impl::decoded_vocab_, + "sorted_decoded_vocab", + &TokenizerInfo::Impl::sorted_decoded_vocab_, + "trie_subtree_nodes_range", + &TokenizerInfo::Impl::trie_subtree_nodes_range_ +); + +} // namespace xgrammar + +#endif // XGRAMMAR_TOKENIZER_INFO_IMPL_H_ diff --git a/third_party/xgrammar/include/module.modulemap b/third_party/xgrammar/include/module.modulemap new file mode 100644 index 0000000000..619525a81a --- /dev/null +++ b/third_party/xgrammar/include/module.modulemap @@ -0,0 +1,4 @@ +module XGrammar { + umbrella header "xgrammar/xgrammar.h" + export * +} diff --git a/third_party/xgrammar/include/xgrammar/compiler.h b/third_party/xgrammar/include/xgrammar/compiler.h new file mode 100644 index 0000000000..2609e4b4d0 --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/compiler.h @@ -0,0 +1,125 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/compiler.h + * \brief The header for the compiler. + */ + +#ifndef XGRAMMAR_COMPILER_H_ +#define XGRAMMAR_COMPILER_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "xgrammar/exception.h" + +namespace xgrammar { + +/*! + * \brief The compiled grammar of a GrammarMatcher. It contains the preprocessing results of the + * grammar and tokenizer. + */ +class CompiledGrammar { + public: + /*! \brief Get the associated grammar. */ + Grammar GetGrammar() const; + + /*! \brief Get the associated tokenizer info. */ + TokenizerInfo GetTokenizerInfo() const; + + /*! \brief Return the approximate memory usage of the grammar in bytes. */ + std::size_t MemorySizeBytes() const; + + /*! \brief Return the serialized JSON string of the compiled grammar. */ + std::string SerializeJSON() const; + + /*! \brief Deserialize a compiled grammar from a JSON string and tokenizer info. */ + static std::variant DeserializeJSON( + const std::string& json_string, const TokenizerInfo& tokenizer_info + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(CompiledGrammar); +}; + +/*! + * \brief A cache to get the compiled grammar for grammar or schema. This class avoids + * redundant preprocessing of the grammar or schema when constructing a CompiledGrammar. + * \note This class is associated with a vocabulary when constructed. The vocabulary is used to + * create every compiled grammar. If multiple toke tables are used to create init + * contexts, an instance of this class for each vocabulary should be created. + */ +class GrammarCompiler { + public: + /*! + * \brief Construct a GrammarCompiler with a vocabulary. This class will always + * create compiled grammars with this vocabulary. + * \param tokenizer_info The tokenizer info. + * \param max_threads The maximum number of threads to use for compiling grammars. + * \param cache_enabled Whether to enable the cache. + * \param max_memory_bytes The maximum memory usage in bytes. + */ + GrammarCompiler( + const TokenizerInfo& tokenizer_info, + int max_threads = 8, + bool cache_enabled = true, + int64_t max_memory_bytes = -1 // unlimited + ); + + /*! \brief Get the compiled grammar for a JSON schema string. */ + CompiledGrammar CompileJSONSchema( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + bool any_order = false + ); + + /*! \brief Get the compiled grammar for pure JSON. */ + CompiledGrammar CompileBuiltinJSONGrammar(); + + /*! + * \brief Get the compiled grammar for a Lark grammar string. + * \param lark_string The Lark grammar. The root rule must be named "start". + * \param named_grammars Grammar objects or Lark sources that can be referenced with `@name`. + */ + CompiledGrammar CompileLark( + const std::string& lark_string, const std::vector& named_grammars = {} + ); + + /*! \brief Get the compiled grammar for a grammar. */ + CompiledGrammar CompileGrammar(const Grammar& grammar); + + /*! \brief Get the compiled grammar for a grammar. */ + CompiledGrammar CompileGrammar( + const std::string& ebnf_str, const std::string& root_rule_name = "root" + ); + + /*! \brief Get the compiled grammar for a structural tag. */ + CompiledGrammar CompileStructuralTag(const std::string& structural_tag_json); + + /*! \brief Get the compiled grammar for a regex. */ + CompiledGrammar CompileRegex(const std::string& regex); + + /*! \brief Clear the internal cache of compiled grammars. */ + void ClearCache(); + + /*! \brief Return the approximate memory usage of the compiler in bytes. */ + int64_t GetCacheSizeBytes() const; + + /*! \brief Return the approximate memory usage of the compiler in bytes. -1 means unlimited. */ + int64_t CacheLimitBytes() const; + + XGRAMMAR_DEFINE_PIMPL_METHODS(GrammarCompiler); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_COMPILER_H_ diff --git a/third_party/xgrammar/include/xgrammar/config.h b/third_party/xgrammar/include/xgrammar/config.h new file mode 100644 index 0000000000..4e5de5e84b --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/config.h @@ -0,0 +1,35 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/config.h + * \brief Global configuration for XGrammar. + */ + +#ifndef XGRAMMAR_CONFIG_H_ +#define XGRAMMAR_CONFIG_H_ + +#include + +namespace xgrammar { + +/*! + * \brief Set the maximum recursion depth for the grammar. + * \param max_recursion_depth The maximum recursion depth. + */ +void SetMaxRecursionDepth(int max_recursion_depth); + +/*! + * \brief Get the maximum recursion depth for the grammar. + * \return The maximum recursion depth. + */ +int GetMaxRecursionDepth(); + +/*! + * \brief Get the serialization version for the grammar. + * \return The serialization version. + * \note This is used to check the compatibility of the serialized grammar. + */ +std::string GetSerializationVersion(); + +} // namespace xgrammar + +#endif // XGRAMMAR_CONFIG_H_ diff --git a/third_party/xgrammar/include/xgrammar/exception.h b/third_party/xgrammar/include/xgrammar/exception.h new file mode 100644 index 0000000000..7efb3b0007 --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/exception.h @@ -0,0 +1,80 @@ +#ifndef XGRAMMAR_EXCEPTION_H +#define XGRAMMAR_EXCEPTION_H + +#include +#include +#include + +namespace xgrammar { + +/************** Exception Definitions **************/ + +/*! + * \brief Exception thrown when the version in the serialized data does not follow the current + * serialization version. + */ + +struct XGrammarError : std::runtime_error { + XGrammarError(const std::string& message) : std::runtime_error(message) {} + virtual std::string GetType() const { return "XGrammarError"; } +}; + +struct DeserializeVersionError : XGrammarError { + DeserializeVersionError(const std::string& message) + : XGrammarError(std::string("Deserialize version error: ") + message) {} + std::string GetType() const override { return "DeserializeVersionError"; } +}; + +/*! + * \brief Exception thrown when the JSON is invalid. + */ +struct InvalidJSONError : XGrammarError { + InvalidJSONError(const std::string& message) + : XGrammarError(std::string("Invalid JSON error: ") + message) {} + std::string GetType() const override { return "InvalidJSONError"; } +}; + +/*! + * \brief Exception thrown when the serialized data does not follow the expected format. + */ +struct DeserializeFormatError : XGrammarError { + DeserializeFormatError(const std::string& message) + : XGrammarError(std::string("Deserialize format error: ") + message) {} + std::string GetType() const override { return "DeserializeFormatError"; } +}; + +/*! + * \brief Exception thrown when the JSON schema is invalid or not satisfiable. + */ +struct InvalidJSONSchemaError : XGrammarError { + InvalidJSONSchemaError(const std::string& message) + : XGrammarError(std::string("Invalid JSON schema error: ") + message) {} + std::string GetType() const override { return "InvalidJSONSchemaError"; } +}; + +/*! + * \brief Exception thrown when the structural tag is invalid. + */ +struct InvalidStructuralTagError : XGrammarError { + InvalidStructuralTagError(const std::string& message) + : XGrammarError(std::string("Invalid structural tag error: ") + message) {} + std::string GetType() const override { return "InvalidStructuralTagError"; } +}; + +/************** Union Exceptions **************/ + +/*! + * \brief Represents a serialization error. + */ +using SerializationError = + std::variant; + +/*! + * \brief Represents an error from the structural tag conversion. + */ +using StructuralTagError = + std::variant; + +} // namespace xgrammar + +#endif // XGRAMMAR_EXCEPTION_H diff --git a/third_party/xgrammar/include/xgrammar/grammar.h b/third_party/xgrammar/include/xgrammar/grammar.h new file mode 100644 index 0000000000..ffb54672a8 --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/grammar.h @@ -0,0 +1,228 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar.h + * \brief The header for the definition and construction of BNF grammar. + */ + +#ifndef XGRAMMAR_GRAMMAR_H_ +#define XGRAMMAR_GRAMMAR_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "xgrammar/exception.h" + +namespace xgrammar { + +struct StructuralTagItem { + std::string begin; + std::string schema; + std::string end; + + bool operator==(const StructuralTagItem& other) const { + return begin == other.begin && schema == other.schema && end == other.end; + } +}; + +struct NamedGrammar; + +/*! + * \brief This class stores the abstract syntax tree (AST) of the Backus-Naur Form (BNF) grammar. + * The BNF definition here is standard BNF, and the characters are represented using regex-style + * character classes (e.g. [a-z], [^a-z]). + * + * \details + * ### Rules + * The BNF grammar AST consists of a set of rules. Each rule contains a name and a definition, and + * corresponds to a production in the grammar. The definition of a rule is a GrammarExpr. Each rule + * has a rule_id for reference. + * + * ### GrammarExprs + * GrammarExpr is the definition of a rule or part of the definition of a rule. It can contain + * elements, empty string, reference to other GrammarExprs, or reference to other rules. Each + * GrammarExpr corresponds to a grammar_expr_id for reference. + * + * For example, in the following rule: rule ::= ("a" "b") | "c" + * ("a" "b"), "c", ("a" "b") | "c" are all GrammarExprs. + * + * #### Types of GrammarExprs + * Every GrammarExpr is represented by a type as well as a variable-length array containing its + * data. GrammarExpr has several types: + * - Byte string: a string of bytes (0~255). Supports UTF-8 strings. + * - Character class: a range of characters (each character is a unicode codepoint), e.g. [a-z], + * [ac-z]. Can be negated: [^a-z], [^ac-z]. Now only ascii chars is allowed in [], but this + * expression can accept/reject unicode chars. + * - Character class star: a star quantifier of a character class. e.g. [a-z]*, [^a-z]*. + * - EmptyStr: an empty string, i.e. "" + * - Rule reference: a reference to another rule + * - Sequence: a sequence of grammar_exprs, e.g. ("a" "b"). These grammar_exprs are concatenated + * together. + * - Choices: a choice of grammar_exprs, e.g. ("a" "b") | "c". Each grammar_expr can be matched. + * + * #### Storage of GrammarExprs + * Each type of GrammarExpr has a different data format. For the format of each type of GrammarExpr, + * see docs in Grammar::Impl::GrammarExprType. + * + * We store all GrammarExprs in csr_matrix style. That is, they are stored consecutively in one + * vector (data vector) and the starting position of each GrammarExpr is recorded in the indptr + * vector. + * + * \remark The character class star GrammarExpr is for the special support for elements like [a-z]* + * in the grammar. We add it to make the matching more efficient, as we can avoid recursion into + * rules when matching a sequence of characters. It should be used like: + * rule1 ::= ((element1 element2 rule2 ...) | ...) + * rule2 ::= character_class_star_grammar_expr(id_of_a_character_class_grammar_expr) + */ +class Grammar { + static const std::vector& EmptyNamedGrammars(); + + public: + /*! + * \brief Get the EBNF string of the grammar. + */ + std::string ToString() const; + + /*! + * \brief Construct a BNF grammar with a EBNF-formatted string. The grammar will be normalized + * (simplified) by default. + * \param ebnf_string The EBNF-formatted string. + * \param root_rule_name The name of the root rule. + */ + static Grammar FromEBNF( + const std::string& ebnf_string, const std::string& root_rule_name = "root" + ); + + /*! + * \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 set to std::nullopt, the output will be + * in one line. Default: 2. + * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"}, + * {", ", ": "}. If std::nullopt, the default separators will be used: {",", ": "} when the + * indent is not nullopt, 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 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 Grammar FromJSONSchema( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + bool print_converted_ebnf = false, + bool any_order = false + ); + + /*! + * \brief Construct a grammar from a regular expression string. + * \param regex The regular expression string. + * \param print_converted_ebnf This method will convert the regex to EBNF first. If this is true, + * the converted EBNF string will be printed. For debugging purpose. Default: false. + */ + static Grammar FromRegex(const std::string& regex, bool print_converted_ebnf = false); + + /*! + * \brief Construct a grammar from Lark syntax. + * \param lark_string The Lark grammar. The root rule must be named "start". + * \param tokenizer_info Optional tokenizer metadata used to resolve named special tokens. + * \param named_grammars Grammar objects or Lark sources that can be referenced with `@name`. + */ + static Grammar FromLark( + const std::string& lark_string, + const std::optional& tokenizer_info = std::nullopt, + const std::vector& named_grammars = EmptyNamedGrammars() + ); + + /*! + * \brief Construct a grammar from a structural tag string. + * \param structural_tag_json The structural tag string. + * \param tokenizer_info Optional tokenizer info for resolving string token references. + */ + static std::variant FromStructuralTag( + const std::string& structural_tag_json, + const std::optional& tokenizer_info = std::nullopt + ); + + /*! + * \brief Get the grammar of standard JSON format. We have built-in support for JSON. + * \return The grammar of standard JSON format. + */ + static Grammar BuiltinJSONGrammar(); + + /*! + * \brief Create a grammar that matches any of the grammars in the list. That is equivalent to + * using the `|` operator to concatenate the grammars in the list. + * \param grammars The grammars to create the union of. + * \returns The union of the grammars. + */ + static Grammar Union(const std::vector& grammars); + + /*! + * \brief Create a grammar that matches the concatenation of the grammars in the list. That is + * equivalent to using the `+` operator to concatenate the grammars in the list. + * \param grammars The grammars to create the concatenation of. + * \returns The concatenation of the grammars. + */ + static Grammar Concat(const std::vector& grammars); + + /*! + * \brief Print a BNF grammar. + * \param os The output stream. + * \param grammar The grammar to print. + * \return The output stream. + */ + friend std::ostream& operator<<(std::ostream& os, const Grammar& grammar); + + /*! + * \brief Return the serialized JSON string of the grammar. + * \return The serialized JSON string. + */ + std::string SerializeJSON() const; + + /*! + * \brief Deserialize a grammar from a JSON string. + * \param json_string The JSON string to deserialize. + * \return If the deserialization is successful, return the grammar. Otherwise, return a runtime + * error with the error message. + */ + static std::variant DeserializeJSON(const std::string& json_string); + + // Allow the builder to bind to this grammar's impl for in-place editing (FromMutableGrammar). + friend class GrammarBuilder; + + XGRAMMAR_DEFINE_PIMPL_METHODS(Grammar); +}; + +/*! + * \brief A grammar or Lark source that can be referenced by name from another Lark grammar. + */ +struct NamedGrammar { + /*! \brief The reference name, without the leading `@`. */ + std::string name; + + /*! \brief An existing grammar or a Lark source with its own `start` rule. */ + std::variant grammar; +}; + +inline const std::vector& Grammar::EmptyNamedGrammars() { + static const std::vector empty; + return empty; +} + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_H_ diff --git a/third_party/xgrammar/include/xgrammar/matcher.h b/third_party/xgrammar/include/xgrammar/matcher.h new file mode 100644 index 0000000000..d3a12b02f0 --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/matcher.h @@ -0,0 +1,300 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/matcher.h + * \brief The header for the matcher. + */ + +#ifndef XGRAMMAR_MATCHER_H_ +#define XGRAMMAR_MATCHER_H_ + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace xgrammar { + +int32_t GetBitmaskSize(int vocab_size); + +DLDataType GetBitmaskDLType(); + +void _DebugGetMaskedTokensFromBitmask( + std::vector* rejected_tokens, const DLTensor& token_bitmask, int vocab_size, int index = 0 +); + +std::pair _IsSingleTokenBitmask(const DLTensor& bitmask, int vocab_size, int index); + +void ApplyTokenBitmaskInplaceCPU( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size = -1, + std::optional> indices = std::nullopt +); + +/*! + * \brief A stateful matcher to match tokens to the specified BNF grammar. This class is the core + * logic of the grammar-guided generation. + * + * \details This class implements the non-deterministic pushdown automaton (NPDA) matching algorithm + * to match characters to a BNF grammar. It keep track of the current state of the matching process + * by maintaining several stacks internally as possible paths in the NPDA. It also supports + * backtracking. + * + * It is particularly capable of finding the set of tokens that are acceptable for the next step + * and storing them in a bitmask. This aids in grammar-guided generation. + * + * \example + * \code + * Tokenizer tokenizer = ...; + * auto compiled_grammar = GrammarMatcher::CreateCompiledGrammar(grammar, + * tokenizer->PostProcessedVocab()); + * GrammarMatcher matcher(compiled_grammar, 10); + * matcher->AcceptToken(67); + * + * // Construct a DLTensor with shape (tokenizer.GetVocabSize() + 31) / 32, and dtype int32. + * DLTensor next_token_bitmask = ...; + * matcher->FillNextTokenBitmask(&next_token_bitmask); + * + * // Rollback is supported + * matcher->Rollback(1); + * \endcode + */ +class GrammarMatcher { + public: + /*! + * \brief Construct a GrammarMatcher from the preprocessing result of type + * CompiledGrammar. + * \param compiled_grammar The compiled grammar. It is obtained through + * CreateCompiledGrammar as a result of preprocessing the grammar and tokenizer. + * \param override_stop_tokens Optional stop token ids that replace those from the tokenizer. + * \param terminate_without_stop_token Whether to terminate when the root rule is complete. + * \param max_rollback_tokens Deprecated and unused. + * \param default_temperature The temperature used when no active rule specifies one. + */ + GrammarMatcher( + const CompiledGrammar& compiled_grammar, + std::optional> override_stop_tokens = std::nullopt, + bool terminate_without_stop_token = false, + int max_rollback_tokens = -1, + std::optional default_temperature = std::nullopt + ); + + /*! + * \brief Accept one token and update the state of the matcher. + * \param token_id The id of the token to accept. + * \return Whether the token is accepted. + * \note Termination state. + * When the end of the root rule is reached, the matcher can only accept the stop token. + * The matcher is terminated after accepting the stop token, i.e. no AcceptToken or + * FindNextTokenMask operations can be performed. The termination state can be canceled + * using Rollback(). + */ + bool AcceptToken(int32_t token_id, bool debug_print = false); + + /*! + * \brief Accept a string and update the state of the matcher. The whole string is considered + * as one step in rollback. It is used to complement the functionality of AcceptToken, and + * AcceptToken should always be used to accept tokens. + * \param input_str The string to be accepted. + * \param debug_print Whether to print information about the internal state of the matcher. + * \return Whether the string is accepted. + */ + bool AcceptString(const std::string& input_str, bool debug_print = false); + + /*! + * \brief Get the set of tokens that are acceptable for the next step and store them in a + * bitmask. + * \param next_token_bitmask The bitmask to store the result. The bitmask must be pre-allocated + * and with shape (GetBitmaskSize(),) and dtype int32. + * \return Whether the bitmask need to be applied (not all-true). + */ + bool FillNextTokenBitmask(DLTensor* next_token_bitmask, int index = 0, bool debug_print = false); + + /*! + * \brief Traverse a draft token tree and fill the token bitmask for each node. + * + * This function performs a DFS traversal of the speculative decoding tree and fills + * the token bitmask for each node based on grammar constraints. + * + * \param retrieve_next_token DLTensor where retrieve_next_token[i] gives the index of + * the child node of node i, or -1 if no child exists. + * \param retrieve_next_sibling DLTensor where retrieve_next_sibling[i] gives the index of + * the sibling node of node i, or -1 if no sibling exists. + * \param draft_tokens DLTensor of draft token ids at each node. + * \param token_bitmask DLTensor to store the bitmask (2D: num_nodes x bitmask_size). + * \param time_threshold Maximum allowed time in seconds for the DFS traversal. + * If the traversal exceeds this threshold, it returns false. + * A value <= 0 disables the timeout (default: -1.0). + * \param temperatures Optional DLTensor to store the effective temperature for each node + * (1D float32 with num_nodes elements). -1 represents no effective temperature; it is + * also written for nodes that were not visited or were rejected. + * \return true if the traversal completed successfully, false if it timed out. + */ + bool TraverseDraftTree( + const DLTensor* retrieve_next_token, + const DLTensor* retrieve_next_sibling, + const DLTensor* draft_tokens, + DLTensor* token_bitmask, + double time_threshold = -1.0, + DLTensor* temperatures = nullptr + ); + + /*! + * \brief Find the jump-forward string for jump-forward decoding. This is the longest string that + will be valid according to the current syntax. + * \note This method does not change the grammar state. + */ + std::string FindJumpForwardString(); + + /*! + * \brief Rollback the matcher to a previous state. + * \param num_tokens The number of tokens to rollback. It cannot exceed the current number of + * steps, nor can it exceed the specified maximum number of rollback tokens. + */ + void Rollback(int num_tokens = 1); + + /*! + * \brief Get the capture groups recorded so far, ordered by completion position. + * \param deduplicate The Earley parser explores parse hypotheses in parallel, so one + * occurrence of a captured rule may complete at several candidate end positions (e.g. a + * /[0-9]+/ body completes after every digit). If true (default), only the last (longest) + * completion of each occurrence is kept. Distinct occurrences — repeated matches of the same + * rule at different positions — are always kept. If false, the raw completion events are + * returned. + * \return A list of (capture_name, matched_bytes) pairs. Rules gain a capture name via the + * Lark attribute rule[capture] / rule[capture="name"] or the EBNF form rule[capture="name"] + * ::= ... The bytes are the input span that the rule matched. + * \note Captures are recorded when a rule is completed during AcceptToken / AcceptString, and + * are rolled back together with Rollback. Mask computation never records captures. Completions + * on parse paths that are later abandoned may still be recorded; for unambiguous grammars + * whose captured rules have unambiguous boundaries, the deduplicated result is exact. + */ + std::vector> GetCaptures(bool deduplicate = true) const; + + /*! + * \brief Check if the matcher has accepted the stop token and terminated. + * \sa AcceptToken + */ + bool IsTerminated() const; + + /*! + * \brief Check if the grammar's root rule has been fully matched by the input accepted so far. + * Unlike IsTerminated(), this does not require the stop token to have been accepted. + * \sa IsTerminated, AcceptToken + */ + bool IsCompleted() const; + + /*! \brief Reset the matcher to the initial state. */ + void Reset(); + + /*! + * \brief Fork the matcher. Returns a new GrammarMatcher with a deep copy of all state except + * compiled_grammar and tokenizer_info, which are shared with this matcher. + */ + GrammarMatcher Fork() const; + + /*! \brief Get the maximum number of rollback tokens allowed. */ + int GetMaxRollbackTokens() const; + + /*! \brief Get the effective sampling temperature for the next token. */ + std::optional GetTemperature() const; + + const std::vector& GetStopTokenIds() const; + + /*! \brief Print the internal state of the matcher. This is only used for debugging. The + * representation of the internal state is subject to change. + */ + std::string _DebugPrintInternalState() const; + + XGRAMMAR_DEFINE_PIMPL_METHODS(GrammarMatcher); +}; + +/*! + * \brief A batched version of GrammarMatcher for better efficiency. It supports batch processing + * of multiple GrammarMatcher objects in parallel. + * + * \details This class provides batched versions of the core methods of GrammarMatcher, including + * FillNextTokenBitmask, AcceptString, and AcceptToken. It utilizes multi-threading to process + * multiple GrammarMatcher objects simultaneously, significantly improving efficiency when dealing + * with a large number of matchers. + */ +class BatchGrammarMatcher { + public: + BatchGrammarMatcher(std::variant max_threads = "auto"); + + /*! + \brief A batched version of FillNextTokenBitmask for better efficiency. + \param matchers The array of GrammarMatcher objects. + \param next_token_bitmask The pre-allocated DLTensor to store the result bitmasks. + \param indices The optional array of indices to specify which matcher corresponds to which slice + of the bitmask tensor. If not provided, all matchers will write to the corresponding + indices(matchers[i] to next_token_bitmask[i]). + \param debug_print Whether to print debug information. Default is false. + */ + void BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices = std::nullopt, + bool debug_print = false + ); + + /*! + * \brief Fill the effective sampling temperature of each matcher into a tensor. + * \param matchers The array of GrammarMatcher objects. + * \param temperatures The pre-allocated 1D float32 CPU tensor to store the result. The entry of + * a matcher without an effective temperature is set to -1. + * \param indices The optional array of indices to specify which element each matcher writes to. + * If not provided, matchers[i] writes to temperatures[i]. + */ + static void BatchFillTemperature( + const std::vector& matchers, + DLTensor* temperatures, + const std::optional>& indices = std::nullopt + ); + + /*! + * \brief A batched version of AcceptString for better efficiency. + * \param matchers The array of GrammarMatcher objects. + * \param input_strs The array of input strings to be accepted. + * \param debug_print Whether to print debug information. Default is false. + * \return A vector of bytes indicating whether each string is accepted. + */ + static std::vector BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print = false + ); + + /*! + * \brief A batched version of AcceptToken for better efficiency. + * \param matchers The array of GrammarMatcher objects. + * \param token_ids The array of token ids to be accepted. + * \param debug_print Whether to print debug information. Default is false. + * \return A vector of bytes indicating whether each token is accepted. + */ + static std::vector BatchAcceptToken( + std::vector* matchers, + const std::vector& token_ids, + bool debug_print = false + ); + + /*! + * \brief A batched version of Rollback for better efficiency. + * \param matchers The array of GrammarMatcher objects. + * \param num_tokens The array of the number of tokens to rollback for each matcher. + */ + static void BatchRollback( + std::vector* matchers, const std::vector& num_tokens + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(BatchGrammarMatcher); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_MATCHER_H_ diff --git a/third_party/xgrammar/include/xgrammar/object.h b/third_party/xgrammar/include/xgrammar/object.h new file mode 100644 index 0000000000..641b5bc526 --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/object.h @@ -0,0 +1,51 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/object.h + * \brief Utilities for creating objects. + */ + +#ifndef XGRAMMAR_OBJECT_H_ +#define XGRAMMAR_OBJECT_H_ + +#include // IWYU pragma: keep +#include // IWYU pragma: keep + +namespace xgrammar { + +/*! + * \brief A tag type for creating a null object. + */ +struct NullObj {}; + +/*! + * \brief This macro defines the methods for the PImpl classes. + * \details Many classes in xgrammar are PImpl classes. PImpl classes only stores a shared pointer + * to the implementation. This allows reference-counter-based memory management and efficient + * object copy and passing. We always expose PImpl classes to Python to control over object sharing + * and memory management. Note simple and critical classes should not be defined as PImpl classes, + * but as normal classes for better efficiency. + */ +#define XGRAMMAR_DEFINE_PIMPL_METHODS(TypeName) \ + public: \ + class Impl; \ + /* Construct a null object. Note operating on a null object will fail. */ \ + explicit TypeName(NullObj) : pimpl_(nullptr) {} \ + /* Construct object with a shared pointer to impl. */ \ + explicit TypeName(std::shared_ptr pimpl) : pimpl_(std::move(pimpl)) {} \ + TypeName(const TypeName& other) = default; \ + TypeName(TypeName&& other) noexcept = default; \ + TypeName& operator=(const TypeName& other) = default; \ + TypeName& operator=(TypeName&& other) noexcept = default; \ + bool IsNull() const { return pimpl_ == nullptr; } \ + /* Access the impl pointer. Useful in implementation. */ \ + Impl* ImplPtr() { return pimpl_.get(); } \ + const Impl* ImplPtr() const { return pimpl_.get(); } \ + Impl* operator->() { return pimpl_.get(); } \ + const Impl* operator->() const { return pimpl_.get(); } \ + \ + private: \ + std::shared_ptr pimpl_ + +} // namespace xgrammar + +#endif // XGRAMMAR_OBJECT_H_ diff --git a/third_party/xgrammar/include/xgrammar/tokenizer_info.h b/third_party/xgrammar/include/xgrammar/tokenizer_info.h new file mode 100644 index 0000000000..6a8e67444b --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/tokenizer_info.h @@ -0,0 +1,86 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/tokenizer_info.h + * \brief The header for the tokenizer info. + */ + +#ifndef XGRAMMAR_TOKENIZER_INFO_H_ +#define XGRAMMAR_TOKENIZER_INFO_H_ + +#include + +#include +#include +#include +#include +#include + +#include "xgrammar/exception.h" + +namespace xgrammar { + +enum class VocabType : int { + RAW = 0, + BYTE_FALLBACK = 1, + BYTE_LEVEL = 2, +}; + +class TokenizerInfo { + public: + TokenizerInfo( + const std::vector& encoded_vocab, + VocabType vocab_type = VocabType::RAW, + std::optional vocab_size = std::nullopt, + std::optional> stop_token_ids = std::nullopt, + bool add_prefix_space = false + ); + + VocabType GetVocabType() const; + bool GetAddPrefixSpace() const; + int GetVocabSize() const; + const std::vector& GetDecodedVocab() const; + const std::vector& GetStopTokenIds() const; + const std::vector& GetSpecialTokenIds() const; + const std::vector>& GetSortedDecodedVocab() const; + const std::vector& GetTrieSubtreeNodesRange() const; + std::string DumpMetadata() const; + + /*! + * \brief Create a tokenizer info from a vocabulary and metadata. + * \param encoded_vocab The encoded vocabulary. + * \param metadata The metadata. + * \return The tokenizer info. + */ + static TokenizerInfo FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata + ); + + /*! + * \brief Detect the metadata from a Hugging Face backend string. + * \param backend_str The Hugging Face backend string. + * \return The metadata. + */ + static std::string DetectMetadataFromHF(const std::string& backend_str); + + /*! + * \brief Return the serialized JSON string of the tokenizer info. + * \return The serialized JSON string. + */ + std::string SerializeJSON() const; + + /*! + * \brief Deserialize a tokenizer info from a JSON string. + * \param json_string The JSON string to deserialize. + * \return If the deserialization is successful, return the tokenizer info. Otherwise, return a + * runtime error with the error message. + */ + static std::variant DeserializeJSON( + const std::string& json_string + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(TokenizerInfo); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_TOKENIZER_INFO_H_ diff --git a/third_party/xgrammar/include/xgrammar/xgrammar.h b/third_party/xgrammar/include/xgrammar/xgrammar.h new file mode 100644 index 0000000000..8513b3a9c5 --- /dev/null +++ b/third_party/xgrammar/include/xgrammar/xgrammar.h @@ -0,0 +1,17 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/xgrammar.h + * \brief The header for the support of grammar-guided generation. + */ + +#ifndef XGRAMMAR_XGRAMMAR_H_ +#define XGRAMMAR_XGRAMMAR_H_ + +#include +#include +#include +#include +#include +#include + +#endif // XGRAMMAR_XGRAMMAR_H_ From 79d229656b297ff9932a58c3664dd56d87a73522 Mon Sep 17 00:00:00 2001 From: Andrey Shvartsman Date: Sat, 19 Sep 2026 18:31:31 -0400 Subject: [PATCH 2/2] fix: support bounded JSON after reasoning and tool calls --- apps/cli/main.cpp | 1 + apps/cli/options.cpp | 11 +- docs/cli.md | 8 +- docs/maintainer/engine-architecture.md | 14 +++ docs/serving.md | 43 +++++-- src/models/qwen3_5/frontend/frontend.cpp | 16 +-- .../qwen3_5/frontend/output_session.cpp | 6 +- .../qwen3_5/frontend/tool_call_parser.cpp | 46 ++++++- .../qwen3_5/frontend/tool_call_parser.h | 4 + src/product/prompt_input/prompt_input.cpp | 35 ++++++ src/product/prompt_input/prompt_input.h | 5 + src/serve/CMakeLists.txt | 2 +- src/serve/translate.cpp | 15 ++- src/text/structured_output.cpp | 73 ++++++++++- src/text/structured_output.h | 10 +- tests/README.md | 8 +- tests/models/qwen3_5/test_frontend.cpp | 34 +++++ tests/test_cli_options.cpp | 5 + tests/test_openai_schema.cpp | 20 ++- tests/test_prompt_input.cpp | 27 ++++ tests/test_structured_output_live.py | 117 +++++++++++++++++- tests/test_tool_call_parser.cpp | 55 ++++++++ tests/text/test_structured_output.cpp | 81 +++++++++++- third_party/xgrammar/NINFER.md | 4 + .../xgrammar/cpp/json_schema_converter.cc | 7 +- 25 files changed, 596 insertions(+), 51 deletions(-) diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index 1ccba542f0..e0c1852aa4 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -257,6 +257,7 @@ int main(int argc, char** argv) { : ninfer::product::prompt_from_messages(cli.messages_path, cli.enable_thinking, cli.enable_vision); input.options.reasoning_effort = cli.reasoning_effort; + ninfer::product::apply_structured_output_instruction(input, cli.structured_output); ninfer::RequestOptions request; request.execution.sampling = cli.sampling; diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index 61d3164cf5..dd09e6d52b 100644 --- a/apps/cli/options.cpp +++ b/apps/cli/options.cpp @@ -103,7 +103,7 @@ std::string usage_text(const char* argv0) { "media sources may be local paths, HTTP(S) URLs, or base64 data URIs.\n" "--vision enables image/video input and loads the fixed Vision GPU allocations.\n" "--json constrains output to a JSON object; --json-schema FILE enforces a supported " - "JSON schema. Both disable thinking.\n" + "JSON schema. Thinking defaults off; an explicit effort or budget enables it.\n" "--thinking-budget caps model-origin thinking tokens; inserted control tokens count " "toward --max-new.\n" "--kv-capacity auto leaves " + @@ -241,12 +241,13 @@ Options parse_options(int argc, char** argv) { product::validate_speculative_cli_options(options.speculative); if (options.structured_output.kind != StructuredOutputKind::None) { if (options.raw_output || !options.stop_strings.empty() || - !options.stop_token_ids.empty() || options.thinking_budget || - (options.reasoning_effort && options.reasoning_effort != ReasoningEffort::None)) { + !options.stop_token_ids.empty()) { throw std::invalid_argument( - "structured output requires decoded text, default stops, and thinking disabled"); + "structured output requires decoded text and default stops"); + } + if (!options.reasoning_effort && !options.thinking_budget) { + options.enable_thinking = false; } - options.enable_thinking = false; } if (options.enable_thinking == false && options.reasoning_effort && *options.reasoning_effort != ReasoningEffort::None) { diff --git a/docs/cli.md b/docs/cli.md index 10c2337186..613b2c3c00 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -284,9 +284,11 @@ All weight, sequence, workspace, and graph allocations are released when the Eng ## JSON and JSON Schema output `--json` constrains output to a JSON object. `--json-schema FILE` constrains it to the supported -JSON Schema subset described in [serving](serving.md#structured-output). Both disable thinking. -They are mutually exclusive and reject raw output, custom stops, thinking budgets, and enabled -reasoning effort. They work with ordinary, MTP, DFlash and DFlash2 execution. +JSON Schema subset described in [serving](serving.md#structured-output). Thinking defaults off; +an explicit `--reasoning-effort` or `--thinking-budget` retains reasoning before the constrained answer. +The two format flags are mutually exclusive and reject raw output and custom stops. The CLI adds +the requested format/schema to the model's instructions before tokenization. They work with +ordinary, MTP, DFlash and DFlash2 execution. ```bash ./build/apps/ninfer model.ninfer --prompt 'Return the city as JSON' --json --max-new 128 diff --git a/docs/maintainer/engine-architecture.md b/docs/maintainer/engine-architecture.md index ba7f82f11c..e1284fbd5d 100644 --- a/docs/maintainer/engine-architecture.md +++ b/docs/maintainer/engine-architecture.md @@ -611,6 +611,20 @@ control. Sequence/checkpoint/cache state never owns it. OutputSession previews o moves the fork into the committed state only after Program commit succeeds. Cancellation before output preview advances no grammar state. +CLI and serving use the shared product prompt adapter to expose an explicit response format and +schema as a leading instruction before preparation. It retains caller message contents and remaps +explicit cache boundaries if a system message is inserted. This aligns model behavior with the +constraint, especially when an automatic tool-call alternative remains available; enforcement +still belongs to the matcher. Requests without a format are unchanged. + +The model frontend supplies its reasoning-close delimiter and native tool envelope. The text layer +composes these with the final JSON grammar into one matcher. Thus a candidate token, or a speculative +block, may cross from reasoning into tool calls or final JSON without an out-of-band phase switch. +The first reasoning-close delimiter ends unrestricted reasoning. The model's tool grammar uses +declared names and optional non-strict parameters in prompt declaration order. OutputSession also previews and +commits forced thinking-control tokens through the matcher; the injected closure cannot leave the +grammar in the reasoning phase. Framing and grammar state have identical rollback/commit ownership. + Program reserves vocabulary bitsets in its planned persistent device arena and owns pinned host staging. Ordinary/prefill sampling uses the current mask. MTP supplies each current draft prefix to a forked matcher before launch. DFlash/DFlash2 must first generate device drafts: their graph diff --git a/docs/serving.md b/docs/serving.md index fb6085374d..7f1218c635 100644 --- a/docs/serving.md +++ b/docs/serving.md @@ -91,10 +91,14 @@ reports that format in its Response object. Anthropic Messages accepts Supported schema constraints are `type`, `properties`, `required`, `additionalProperties`, `items`, `prefixItems`, `minItems`, `maxItems`, `minLength`, `maxLength`, `enum`, `const`, -`anyOf`, `$defs`, `definitions`, and local fragment `$ref` (including recursive schemas). +`anyOf`, `minimum`, `maximum`, `exclusiveMinimum`, `exclusiveMaximum`, `$defs`, `definitions`, +and local fragment `$ref` (including recursive schemas). Annotations `$schema`, `title`, `description`, `default`, `examples`, and `$comment` do not impose generation constraints. Other keywords are rejected with HTTP 400: this includes -numeric bounds, `multipleOf`, `pattern`, `format`, `oneOf`, `allOf`, `uniqueItems`, and conditionals. +`multipleOf`, `pattern`, `format`, `oneOf`, `allOf`, `uniqueItems`, and conditionals. +Numeric bounds require an explicit `integer` or `number` type and finite values within +`+/- (2^53-1)`. Integer bounds must be whole numbers. Bounded numbers use ordinary decimal +notation with at most six fractional digits; a range with no representable value is rejected. `$id` and external references are rejected. An explicit `$schema` must be JSON Schema 2020-12 or draft-07. Local references use `#` or literal object paths such as `#/$defs/node`; escaped or empty path segments are rejected. Bounded strings use unescaped Unicode characters; escaped quotes, backslashes, @@ -108,14 +112,29 @@ eight characters per run to prevent whitespace-only generation loops. Additional where necessary to prevent escaped aliases from overwriting declared typed properties. `strict:false` does not disable enforcement. -Structured responses default to thinking disabled even if the server's ordinary default enables -thinking. Explicit enabled thinking, a thinking budget, active tool generation, custom stops, -and assistant-prefill continuation are rejected in combination with structured output. Strict -tool argument generation and arbitrary grammar/regex aliases remain unsupported. Public C++ -callers select `ExecutionOptions::structured_output`, prepare a prompt with thinking disabled, -and use default stops and decoded text output. - -Only normal completed responses guarantee a complete JSON document satisfying the supported +Structured responses default to thinking disabled when no reasoning mode or budget is requested. +Explicit thinking and thinking budgets are supported: reasoning is returned in its normal channel, +and the first reasoning-close delimiter transitions to constrained final content. Active tools are +also supported. After reasoning, a response can emit native Qwen tool calls or the final JSON value. +The schema applies to final content, not reasoning or tool arguments. Tool responses have their normal +tool-call finish reason and need not contain a JSON content value. Tool definitions and thinking +settings may remain enabled on later requests after tool results are supplied. + +Serving also exposes the requested format and schema to the model as a leading instruction, +asking for raw JSON without Markdown fences and retaining tools for information gathering. +This helps the model choose the intended final-answer branch when tools remain enabled; the +token grammar enforces validity independently. Tool selection and termination remain model +decisions under `tool_choice:"auto"`; use `tool_choice:"none"` when no further tools are needed. +Requests without an explicit response format receive no format instruction or grammar. + +The native tool envelope restricts calls to declared names and emits each declared parameter at most +once, in the declaration order shown in the prompt. Parameter values remain non-strict; delimiter spellings that would make +Qwen XML ambiguous are excluded. Strict tool argument generation and arbitrary grammar/regex aliases +remain unsupported. Custom stops and assistant-prefill continuation are rejected with structured +output. Public C++ callers select `ExecutionOptions::structured_output` and use default stops and +decoded text output. + +Only normal completed final-content responses guarantee a complete JSON document satisfying the supported schema. Token/context limits, cancellation, transport failure, or generation errors can leave a partial document; inspect the finish reason or Responses status before parsing it as complete. SSE content deltas are ordinary partial JSON bytes; concatenate them before parsing. Constraints @@ -711,7 +730,7 @@ curl http://127.0.0.1:8080/v1/responses/input_tokens \ ``` Unsupported Create fields include Conversations, prompt templates, context management, hosted -moderation, Structured Outputs/JSON mode, non-empty `include`, background execution, compaction, +moderation, non-empty `include`, background execution, compaction, files/audio, and OpenAI-hosted/MCP/custom tools. These are compatibility boundaries, not silently accepted placeholders. @@ -787,7 +806,7 @@ emits `message_start` after Engine admission commits the prefix selection and be transfer/prefill output, so its uncached/cache-read split is already exact; terminal cumulative usage matches the aggregate response. -Documents, Search Results, Files, Structured Outputs, server-tool results, container uploads, and +Documents, Search Results, Files, server-tool results, container uploads, and other execution-dependent blocks are rejected with the missing capability identified. Metadata, service tier, inference geography, protocol-version/beta headers, cache TTL, and unknown advisory fields do not block an otherwise executable request. The request `model` is any non-empty local diff --git a/src/models/qwen3_5/frontend/frontend.cpp b/src/models/qwen3_5/frontend/frontend.cpp index 52159ff9d8..02aa04baf3 100644 --- a/src/models/qwen3_5/frontend/frontend.cpp +++ b/src/models/qwen3_5/frontend/frontend.cpp @@ -914,32 +914,32 @@ OutputSession Frontend::make_output_session(const PreparedPrompt& prompt, std::shared_ptr grammar; text::validate_structured_output(structured); if (structured.kind != StructuredOutputKind::None) { - if (prompt.data_->starts_in_reasoning || thinking.budget) { - throw std::invalid_argument("structured output requires thinking disabled"); - } if (!caller_stop.token_ids.empty() || !caller_stop.strings.empty() || !caller_stop.include_model_defaults || output.raw || output.preserve_special_tokens || caller_stop.publish_stop_token) { throw std::invalid_argument( "structured output requires default stops and decoded text output"); } - if (prompt.data_->tool_call_output && !prompt.data_->tool_call_output->tools.empty()) { - throw std::invalid_argument( - "structured output cannot be combined with tool generation"); + text::StructuredOutputEnvelope envelope; + if (prompt.data_->starts_in_reasoning) { envelope.reasoning_close = ""; } + if (prompt.data_->tool_call_output) { + envelope.alternative_format = + fi::structured_tool_call_format(*prompt.data_->tool_call_output); } std::lock_guard lock(impl_->grammar_mutex); if (!impl_->grammar_compiler) { std::vector vocab(impl_->tokenizer->vocab_size()); for (std::size_t id = 0; id < vocab.size(); ++id) { if (impl_->tokenizer->is_valid_token(id) && - !impl_->tokenizer->is_special_token(id)) { + (!impl_->tokenizer->is_special_token(id) || + impl_->tokenizer->decode_token_bytes(id, false) == "")) { vocab[id] = impl_->tokenizer->decode_token_bytes(id, false); } } impl_->grammar_compiler = std::make_unique( std::move(vocab), impl_->tokenizer->default_stop_token_ids()); } - grammar = impl_->grammar_compiler->compile(structured); + grammar = impl_->grammar_compiler->compile(structured, envelope); } if (output.raw) { policy.publish_stop_token = true; } return OutputSession( diff --git a/src/models/qwen3_5/frontend/output_session.cpp b/src/models/qwen3_5/frontend/output_session.cpp index e355418e28..394db4d10b 100644 --- a/src/models/qwen3_5/frontend/output_session.cpp +++ b/src/models/qwen3_5/frontend/output_session.cpp @@ -591,7 +591,11 @@ runtime::OutputDecision OutputSession::preview_control(std::span impl_->preview_semantic.control_pending = false; impl_->preview_semantic.applied = true; impl_->preview_semantic.injected_tokens = static_cast(tokens.size()); - impl_->preview_ready = true; + if (impl_->grammar) { + impl_->preview_grammar = impl_->grammar->fork(); + impl_->preview_grammar->accept(tokens); + } + impl_->preview_ready = true; return runtime::OutputDecision{ .accepted_tokens = static_cast(tokens.size()), .prefix_execution_split_after = impl_->preview_execution_split_after, diff --git a/src/models/qwen3_5/frontend/tool_call_parser.cpp b/src/models/qwen3_5/frontend/tool_call_parser.cpp index 124281d417..8af1c7df89 100644 --- a/src/models/qwen3_5/frontend/tool_call_parser.cpp +++ b/src/models/qwen3_5/frontend/tool_call_parser.cpp @@ -10,7 +10,7 @@ namespace ninfer::models::qwen3_5::frontend { namespace { -using Json = nlohmann::json; +using Json = nlohmann::ordered_json; using Contract = ToolCallOutputContract; using FallbackReason = ToolCallParseFallbackReason; using NormalizationPolicy = Contract::NormalizationPolicy; @@ -595,6 +595,50 @@ build_tool_call_output_contract(std::span tool_jsons, bool en return contract; } +std::string structured_tool_call_format(const ToolCallOutputContract& contract) { + const Json whitespace{{"type", "regex"}, {"pattern", "[ \\t\\r\\n]{0,8}"}}; + const auto literal = [](const std::string& text) { + return Json{{"type", "const_string"}, {"value", text}}; + }; + Json alternatives = Json::array(); + for (const auto& tool : contract.tools) { + if (!tool.unambiguous) { + throw std::invalid_argument("structured output requires unambiguous tool definitions"); + } + Json elements = + Json::array({literal(std::string(kToolOpen)), whitespace, + literal(std::string(kFunctionOpen) + tool.name + ">"), whitespace}); + for (const auto& parameter : tool.parameters) { + if (parameter.name.empty() || parameter.name.find('>') != std::string::npos) { + throw std::invalid_argument("tool parameter cannot be represented in Qwen XML"); + } + Json value{{"type", "any_text"}, + {"excludes", Json::array({kParamClose, kParamOpen, kFunctionClose, kToolOpen, + kToolClose})}}; + Json tagged{{"type", "tag"}, + {"begin", std::string(kParamOpen) + parameter.name + ">"}, + {"content", value}, + {"end", kParamClose}}; + elements.push_back( + Json{{"type", "optional"}, + {"content", Json{{"type", "sequence"}, + {"elements", Json::array({tagged, whitespace})}}}}); + } + elements.push_back(literal(std::string(kFunctionClose))); + elements.push_back(whitespace); + elements.push_back(literal(std::string(kToolClose))); + elements.push_back(whitespace); + alternatives.push_back(Json{{"type", "sequence"}, {"elements", elements}}); + } + if (alternatives.empty()) { return {}; } + return Json{{"type", "sequence"}, + {"elements", + Json::array({whitespace, Json{{"type", "plus"}, + {"content", Json{{"type", "or"}, + {"elements", alternatives}}}}})}} + .dump(); +} + ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, std::size_t max_tool_name_length, const ToolCallOutputContract& contract) { diff --git a/src/models/qwen3_5/frontend/tool_call_parser.h b/src/models/qwen3_5/frontend/tool_call_parser.h index 6624669aac..8469986350 100644 --- a/src/models/qwen3_5/frontend/tool_call_parser.h +++ b/src/models/qwen3_5/frontend/tool_call_parser.h @@ -62,6 +62,10 @@ struct ParsedToolCallOutput { [[nodiscard]] std::shared_ptr build_tool_call_output_contract(std::span tool_jsons, bool enabled); +// Native Qwen tool-call envelope for a structured final-response request. Parameters remain +// non-strict values, with each declared parameter emitted at most once in declaration order. +[[nodiscard]] std::string structured_tool_call_format(const ToolCallOutputContract& contract); + [[nodiscard]] ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, std::size_t max_tool_name_length, const ToolCallOutputContract& contract); diff --git a/src/product/prompt_input/prompt_input.cpp b/src/product/prompt_input/prompt_input.cpp index a4b62c76e2..6164003f43 100644 --- a/src/product/prompt_input/prompt_input.cpp +++ b/src/product/prompt_input/prompt_input.cpp @@ -273,4 +273,39 @@ PromptInput prompt_from_messages(const std::filesystem::path& path, return input; } +void apply_structured_output_instruction(PromptInput& input, + const StructuredOutputOptions& options) { + if (options.kind == StructuredOutputKind::None) { return; } + MessagePart instruction; + instruction.text = + "\n\nThe response format is JSON. Return the final answer as raw JSON without Markdown " + "fences or surrounding text."; + if (!input.options.tool_jsons.empty()) { + instruction.text += + " Tool calls remain available when needed to obtain information; do not call tools " + "merely to format JSON."; + } + if (options.kind == StructuredOutputKind::JsonSchema) { + instruction.text += "\nThe final answer must match this JSON Schema:\n" + options.schema; + } else { + instruction.text += " The final answer must be a JSON object."; + } + if (!input.messages.empty() && (input.messages.front().role == ChatRole::System || + input.messages.front().role == ChatRole::Developer)) { + // Append a part so explicit boundaries inside the caller's instruction stay unchanged. + input.messages.front().parts.push_back(std::move(instruction)); + } else { + ChatMessage system; + system.role = ChatRole::System; + system.parts.push_back(std::move(instruction)); + input.messages.insert(input.messages.begin(), std::move(system)); + for (auto& marker : input.context_cache.markers) { + if (marker.location == PromptCacheMarkerLocation::MessageBoundary || + marker.location == PromptCacheMarkerLocation::MessagePartBoundary) { + ++marker.after_message_count; + } + } + } +} + } // namespace ninfer::product diff --git a/src/product/prompt_input/prompt_input.h b/src/product/prompt_input/prompt_input.h index 9dc799ed6c..b7c987089e 100644 --- a/src/product/prompt_input/prompt_input.h +++ b/src/product/prompt_input/prompt_input.h @@ -12,4 +12,9 @@ namespace ninfer::product { std::optional enable_thinking, bool vision_enabled); +// Make the requested final-response contract visible to the model before preparation. +// Grammar enforcement remains an execution responsibility. +void apply_structured_output_instruction(PromptInput& input, + const StructuredOutputOptions& options); + } // namespace ninfer::product diff --git a/src/serve/CMakeLists.txt b/src/serve/CMakeLists.txt index 1bcdea46e8..abf3a75908 100644 --- a/src/serve/CMakeLists.txt +++ b/src/serve/CMakeLists.txt @@ -25,4 +25,4 @@ add_library(ninfer_serve STATIC ninfer_internal_includes(ninfer_serve) target_link_libraries(ninfer_serve PUBLIC ninfer_engine Threads::Threads ninfer::json ninfer::httplib - PRIVATE ninfer_media_acquire ninfer_product_logging CUDA::cudart) + PRIVATE ninfer_media_acquire ninfer_product_logging ninfer_product_prompt_input CUDA::cudart) diff --git a/src/serve/translate.cpp b/src/serve/translate.cpp index f1bb594648..0317c24380 100644 --- a/src/serve/translate.cpp +++ b/src/serve/translate.cpp @@ -1,5 +1,6 @@ #include "serve/translate.h" #include "serve/request_json.h" +#include "product/prompt_input/prompt_input.h" #include @@ -151,14 +152,15 @@ ResolvedPromptSemantics resolve_prompt_semantics(const GenerationRequest& reques } kwargs.erase("reasoning_effort"); if (request.structured_output.kind != StructuredOutputKind::None) { - if (thinking == true || (effort && *effort != RequestedReasoningEffort::None) || - request.thinking_budget || request.uses_tools() || !request.stop_strings.empty() || + if (!request.stop_strings.empty() || request.continuation != PromptContinuationMode::NewAssistantTurn) { - invalid_prompt_option("structured output requires thinking disabled, default stops, a " - "new assistant turn, and no active tools", - "response_format", "incompatible_structured_output"); + invalid_prompt_option( + "structured output requires default stops and a new assistant turn", + "response_format", "incompatible_structured_output"); } - thinking = false; + // Preserve explicit reasoning requests; retain the economical final-only default when + // the caller has not selected a reasoning mode or budget. + if (!thinking && !effort && !request.thinking_budget) { thinking = false; } } ResolvedPromptSemantics result{ .enable_thinking = thinking ? thinking : server.enable_thinking, @@ -315,6 +317,7 @@ ninfer::PromptInput to_prompt_input(const GenerationRequest& request, } input.context_cache.allow_engine_automatic_shared_prefixes = request.allow_engine_automatic_shared_prefixes; + product::apply_structured_output_instruction(input, request.structured_output); return input; } diff --git a/src/text/structured_output.cpp b/src/text/structured_output.cpp index 5251d1f089..1c9f2020ff 100644 --- a/src/text/structured_output.cpp +++ b/src/text/structured_output.cpp @@ -2,7 +2,9 @@ #include #include #include +#include #include +#include #include #include @@ -16,11 +18,24 @@ void schema_check(const Json& s) { static const std::unordered_set annotations = { "$schema", "title", "description", "default", "examples", "$comment", "$defs", "definitions"}; - static const std::unordered_set supported = { - "type", "properties", "required", "additionalProperties", - "items", "prefixItems", "minItems", "maxItems", - "minLength", "maxLength", "enum", "const", - "anyOf", "$ref"}; + static const std::unordered_set supported = {"type", + "properties", + "required", + "additionalProperties", + "items", + "prefixItems", + "minItems", + "maxItems", + "minLength", + "maxLength", + "enum", + "const", + "anyOf", + "$ref", + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum"}; for (const auto& [key, value] : s.items()) { if (key == "$schema" && value != "https://json-schema.org/draft/2020-12/schema" && value != "http://json-schema.org/draft-07/schema#") { @@ -30,6 +45,19 @@ void schema_check(const Json& s) { (!value.is_number_integer() || value < 0 || value > 2147483647)) { throw std::invalid_argument(key + " must be a nonnegative 32-bit integer"); } + if (key == "minimum" || key == "maximum" || key == "exclusiveMinimum" || + key == "exclusiveMaximum") { + // The pinned compiler represents bounds as doubles. Keep integer bounds exact + // and reject non-finite values before they reach its range arithmetic. + if (!value.is_number() || !std::isfinite(value.get()) || + std::abs(value.get()) > 9007199254740991.0L) { + throw std::invalid_argument(key + " requires a finite bound within +/- (2^53-1)"); + } + if (!s.contains("type") || (s.at("type") != "integer" && s.at("type") != "number")) { + throw std::invalid_argument( + "numeric bounds require explicit integer or number type"); + } + } if (!annotations.contains(key) && !supported.contains(key)) { throw std::invalid_argument("unsupported JSON schema keyword: " + key); } @@ -174,7 +202,9 @@ StructuredCompiler::StructuredCompiler(std::vector vocab, std::vect StructuredCompiler::~StructuredCompiler() = default; -std::shared_ptr StructuredCompiler::compile(const StructuredOutputOptions& options) { +std::shared_ptr +StructuredCompiler::compile(const StructuredOutputOptions& options, + const StructuredOutputEnvelope& envelope) { validate_structured_output(options); if (options.kind == StructuredOutputKind::None) { return {}; } std::lock_guard lock(impl_->mutex); @@ -184,6 +214,37 @@ std::shared_ptr StructuredCompiler::compile(const StructuredOutput options.kind == StructuredOutputKind::JsonObject ? "{\"type\":\"object\"}" : options.schema, true, std::nullopt, std::nullopt, false, 8); + if (!envelope.reasoning_close.empty() || !envelope.alternative_format.empty()) { + std::ostringstream ebnf; + ebnf << grammar.GetGrammar(); + Json content{{"type", "grammar"}, {"grammar", ebnf.str()}}; + // The schema root starts at the JSON value; Qwen's reasoning close is followed by + // whitespace (including the canonical budget-control suffix's two newlines). + content = Json{ + {"type", "sequence"}, + {"elements", Json::array({Json{{"type", "regex"}, {"pattern", "[ \\t\\r\\n]{0,8}"}}, + content})}}; + if (!envelope.alternative_format.empty()) { + content = Json{ + {"type", "or"}, + {"elements", Json::array({content, Json::parse(envelope.alternative_format)})}}; + } + if (!envelope.reasoning_close.empty()) { + // any_text excludes the first closing delimiter, including split-token and + // overlapping prefixes. A wildcard repetition would allow reasoning to consume + // the delimiter and bypass the final-content constraint. + content = + Json{{"type", "sequence"}, + {"elements", + Json::array( + {Json{{"type", "any_text"}, + {"excludes", Json::array({envelope.reasoning_close})}}, + Json{{"type", "const_string"}, {"value", envelope.reasoning_close}}, + content})}}; + } + grammar = impl_->compiler.CompileStructuralTag( + Json{{"type", "structural_tag"}, {"format", content}}.dump()); + } return std::shared_ptr(new GrammarState(std::make_unique( xgrammar::GrammarMatcher(grammar), impl_->tokenizer.GetVocabSize()))); } catch (const std::exception& e) { diff --git a/src/text/structured_output.h b/src/text/structured_output.h index 02eb09f050..f5c78a1775 100644 --- a/src/text/structured_output.h +++ b/src/text/structured_output.h @@ -10,6 +10,13 @@ namespace ninfer::text { // Reject unsupported constraints instead of letting a compiler silently weaken a schema. void validate_structured_output(const StructuredOutputOptions& options); +// Model-owned framing around the final response. The alternate format is an XGrammar +// structural-tag descriptor (e.g. a model's native tool-call serialization), not a user API. +struct StructuredOutputEnvelope { + std::string reasoning_close; + std::string alternative_format; +}; + class GrammarState { public: ~GrammarState(); @@ -30,7 +37,8 @@ class StructuredCompiler { public: StructuredCompiler(std::vector decoded_vocab, std::vector stop_tokens); ~StructuredCompiler(); - std::shared_ptr compile(const StructuredOutputOptions& options); + std::shared_ptr compile(const StructuredOutputOptions& options, + const StructuredOutputEnvelope& envelope = {}); private: struct Impl; std::unique_ptr impl_; diff --git a/tests/README.md b/tests/README.md index 0598ce0356..28c0c3b21a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -232,7 +232,7 @@ the fixed Engine fixture does not define bit parity across arbitrary floating-po The CPU grammar and protocol tests and the GPU sampling/speculative tests qualify this path: ```bash -ctest --test-dir build --output-on-failure -R 'ninfer_(structured_output|sampling|speculative_round|openai_schema|openai_responses|anthropic_schema|cli_options)_test$' +ctest --test-dir build --output-on-failure -R 'ninfer_(structured_output|sampling|speculative_round|openai_schema|openai_responses|anthropic_schema|cli_options|prompt_input|tool_call_parser|qwen3_5_frontend|qwen3_5_structured_round)_test$' ``` `tests/test_structured_output_live.py` runs a temporary loopback server and validates completed @@ -248,6 +248,10 @@ python tests/test_structured_output_live.py \ It checks conflicting prompts, greedy and stochastic generation, schema versus JSON object mode, concurrent and mixed traffic, prefix reuse, SSE, token limits, disconnect cleanup, compile errors, -Responses, and Anthropic Messages. `--concurrency 8 --draft-tokens 15 --modes dflash2` exercises +Responses, and Anthropic Messages. It also runs a tool-result-to-bounded-JSON exchange while +reasoning and multiple tools remain enabled, including a Markdown-formatted example in the prompt. +CPU tests cover numeric endpoints, schema enforcement with tool alternatives, split and mixed-token +reasoning transitions, speculative mask previews, injected thinking-budget closure, and prompt cache +boundary preservation. `--concurrency 8 --draft-tokens 15 --modes dflash2` exercises the largest draft and batch dimensions. A separate DFlash-capable artifact can use `--modes dflash`. The server is terminated on success or failure. An occupied test port causes the test to stop. diff --git a/tests/models/qwen3_5/test_frontend.cpp b/tests/models/qwen3_5/test_frontend.cpp index 1590797df3..7fc793d47c 100644 --- a/tests/models/qwen3_5/test_frontend.cpp +++ b/tests/models/qwen3_5/test_frontend.cpp @@ -8,6 +8,7 @@ #include "models/qwen3_5/frontend/test_access.h" #include "models/qwen3_5/frontend/tokenizer.h" #include "text/unicode.h" +#include "text/structured_output.h" #include @@ -1688,6 +1689,38 @@ ninfer::models::qwen3_5::PreparedPrompt thinking_prompt(const Frontend& frontend return frontend.prepare(std::move(input)); } +int test_structured_thinking_control(const Frontend& frontend) { + auto prompt = thinking_prompt(frontend); + auto session = + frontend.make_output_session(prompt, {}, {}, ninfer::ThinkingControlOptions{.budget = 2}, + {ninfer::StructuredOutputKind::JsonObject, {}}); + const auto grammar = session.grammar_state(); + const auto words = (fixture_tokenizer().vocab_size() + 31) / 32; + std::vector before(words), after(words); + grammar->fill_masks(before, {}); + const auto decision = session.preview_model(std::array{0, 0}, 1024, + ninfer::FinishReason::OutputLimit); + (void)session.commit_preview(); + int failures = + check(decision.continuation == ninfer::runtime::ContinuationAction::ApplyTargetControl, + "structured reasoning did not reach budget control"); + const auto control = session.pending_control_tokens(); + (void)session.preview_control(control, 1022); + grammar->fill_masks(after, {}); + const auto x = fixture_byte_token('x'); + failures += check(after[x / 32] & (1U << (x % 32)), "control preview advanced grammar"); + (void)session.commit_preview(); + grammar->fill_masks(after, {}); + failures += + check(!(after[x / 32] & (1U << (x % 32))), "forced closure left reasoning unconstrained"); + const auto json_tokens = fixture_tokenizer().encode("{\"ok\":true}"); + (void)session.preview_model(json_tokens, 900, ninfer::FinishReason::OutputLimit); + const auto output = session.commit_preview(); + failures += check(channel_text(output, ninfer::OutputChannel::Content) == "{\"ok\":true}", + "structured final content lost after forced reasoning close"); + return failures; +} + int test_thinking_budget_control(const Frontend& frontend) { auto prompt = thinking_prompt(frontend); ninfer::StopPolicy stop; @@ -2193,6 +2226,7 @@ int main() { failures += test_structured_tool_output(); failures += test_reasoning_split(frontend); failures += test_thinking_budget_control(frontend); + failures += test_structured_thinking_control(frontend); failures += test_utf8_and_hidden_eos(frontend); failures += test_media_cache_reuses_immutable_payload(); failures += test_media_payload_outlives_frontend_cache(); diff --git a/tests/test_cli_options.cpp b/tests/test_cli_options.cpp index d878b380fc..40d53e3923 100644 --- a/tests/test_cli_options.cpp +++ b/tests/test_cli_options.cpp @@ -109,6 +109,11 @@ int main() { check(structured.structured_output.kind == ninfer::StructuredOutputKind::JsonObject && structured.enable_thinking == false, "CLI JSON mode"); + const auto structured_reasoning = parse( + {"ninfer", "model.ninfer", "--prompt", "hello", "--json", "--reasoning-effort", "medium"}); + failures += check(structured_reasoning.enable_thinking != false && + structured_reasoning.reasoning_effort == ninfer::ReasoningEffort::Medium, + "CLI structured output preserves explicit reasoning"); failures += check( rejects([] { (void)parse({"ninfer", "model.ninfer", "--prompt", "x", "--json", "--raw-output"}); diff --git a/tests/test_openai_schema.cpp b/tests/test_openai_schema.cpp index 855c4059e7..b61befb691 100644 --- a/tests/test_openai_schema.cpp +++ b/tests/test_openai_schema.cpp @@ -142,6 +142,11 @@ int test_structured_output() { "JSON mode reaches Engine"); failures += check(semantics(generation).enable_thinking == false, "JSON mode disables thinking default"); + const auto json_prompt = prompt(generation); + failures += check(json_prompt.messages.front().role == ninfer::ChatRole::System && + json_prompt.messages.front().parts.back().text.find("raw JSON") != + std::string::npos, + "JSON response contract is visible to the model"); body["response_format"] = Json{ {"type", "json_schema"}, {"json_schema", Json{{"name", "answer"}, @@ -151,11 +156,22 @@ int test_structured_output() { {"required", Json::array({"x"})}, {"additionalProperties", false}}}}}}; generation = parse(body).generation; + failures += check(prompt(generation) + .messages.front() + .parts.back() + .text.find(generation.structured_output.schema) != std::string::npos, + "API-only response schema reaches the prompt"); failures += check(options(generation).execution.structured_output.kind == ninfer::StructuredOutputKind::JsonSchema, "JSON schema reaches Engine"); - for (const auto& extra : {Json{{"enable_thinking", true}}, Json{{"stop", "}"}}, - Json{{"reasoning_effort", "high"}}}) { + for (const auto& extra : + {Json{{"enable_thinking", true}}, Json{{"reasoning_effort", "high"}}}) { + auto enabled = body; + enabled.update(extra); + failures += check(semantics(parse(enabled).generation).enable_thinking == true, + "structured output preserves explicit reasoning"); + } + for (const auto& extra : {Json{{"stop", "}"}}}) { auto invalid = body; invalid.update(extra); failures += diff --git a/tests/test_prompt_input.cpp b/tests/test_prompt_input.cpp index 9c544d4b49..2d720140e6 100644 --- a/tests/test_prompt_input.cpp +++ b/tests/test_prompt_input.cpp @@ -60,6 +60,33 @@ int main() { std::cerr << "local messages JSON changed prompt-bearing object member order\n"; return 1; } + auto constrained = prompt; + constrained.context_cache.markers.push_back( + {.after_message_count = 1, + .location = ninfer::PromptCacheMarkerLocation::MessagePartBoundary, + .after_message_part_count = 1}); + ninfer::product::apply_structured_output_instruction( + constrained, {ninfer::StructuredOutputKind::JsonObject, {}}); + if (constrained.messages.size() != 3 || + constrained.messages.front().role != ninfer::ChatRole::System || + constrained.messages[1].parts[0].text != prompt.messages[0].parts[0].text || + constrained.context_cache.markers[0].after_message_count != 2 || + constrained.context_cache.markers[0].after_message_part_count != 1) { + std::cerr << "structured instruction changed caller content or cache boundary\n"; + return 1; + } + auto leading = ninfer::product::prompt_from_text("answer", false); + leading.messages.front().role = ninfer::ChatRole::System; + leading.context_cache.markers.push_back( + {.location = ninfer::PromptCacheMarkerLocation::LeadingInstructionBoundary, + .leading_instruction_bytes = 3}); + ninfer::product::apply_structured_output_instruction( + leading, {ninfer::StructuredOutputKind::JsonObject, {}}); + if (leading.messages.size() != 1 || leading.messages[0].parts[0].text != "answer" || + leading.context_cache.markers[0].leading_instruction_bytes != 3) { + std::cerr << "structured instruction moved caller instruction boundary\n"; + return 1; + } std::cout << "ok\n"; return 0; } diff --git a/tests/test_structured_output_live.py b/tests/test_structured_output_live.py index 047dc7332a..9e65e8fe8e 100644 --- a/tests/test_structured_output_live.py +++ b/tests/test_structured_output_live.py @@ -170,6 +170,122 @@ def check_json(body, schema=SCHEMA): check_json(chat()) check_json(chat(temperature=0)) check_json(chat({"type": "json_object"}), {"type": "object"}) + # Reasoning and tools remain enabled throughout a real tool -> constrained-content + # exchange, matching clients that keep their global settings on every request. + weather_schema = { + "type": "object", + "properties": { + "city": {"type": "string"}, + "temperature": {"type": "number", "minimum": -100, "maximum": 100}, + }, + "required": ["city", "temperature"], + "additionalProperties": False, + } + weather_format = { + "type": "json_schema", + "json_schema": { + "name": "weather", + "schema": weather_schema, + "strict": True, + }, + } + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + }, + }, + } + ] + tools.append( + { + "type": "function", + "function": { + "name": "run_code", + "description": "Execute code when computation is needed.", + "parameters": { + "type": "object", + "properties": { + "language": {"type": "string"}, + "code": {"type": "string"}, + }, + "required": ["language", "code"], + }, + }, + } + ) + tool_messages = [ + { + "role": "user", + "content": "Use get_weather to obtain the current temperature in Tokyo, then return city and " + "temperature as JSON. You must call the tool; do not invent live weather.\n" + 'Example format:\n```json\n{"city":"Paris","temperature":20}\n```', + } + ] + reasoning_kwargs = { + "enable_thinking": True, + "preserve_thinking": True, + "reasoning_effort": "medium", + } + first, _ = request( + chat( + weather_format, + messages=tool_messages, + tools=tools, + tool_choice="auto", + chat_template_kwargs=reasoning_kwargs, + temperature=0, + max_tokens=1024, + ) + ) + first_message = first["choices"][0]["message"] + calls = first_message.get("tool_calls", []) + assert calls and all( + c["function"]["name"] == "get_weather" for c in calls + ), first + tool_messages.append( + { + k: first_message[k] + for k in ("role", "content", "reasoning_content", "tool_calls") + if k in first_message + } + ) + for call in calls: + assert json.loads(call["function"]["arguments"])["city"] == "Tokyo", call + tool_messages.append( + { + "role": "tool", + "tool_call_id": call["id"], + "content": '{"city":"Tokyo","temperature":28}', + } + ) + final = check_json( + chat( + weather_format, + messages=tool_messages, + tools=tools, + tool_choice="auto", + chat_template_kwargs=reasoning_kwargs, + temperature=0, + max_tokens=1024, + ), + weather_schema, + ) + assert final == {"city": "Tokyo", "temperature": 28}, final + check_json(chat(enable_thinking=True, max_tokens=1024)) + print( + "PASS", + mode, + "reasoning, native tool envelope, bounded final JSON", + flush=True, + ) # Same prompt, fresh grammar, prefix reuse and mixed compact batch membership. with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: jobs = [pool.submit(check_json, chat(seed=25 + i)) for i in range(2)] @@ -207,7 +323,6 @@ def check_json(body, schema=SCHEMA): short, _ = request(chat(max_tokens=2)) assert short["choices"][0]["finish_reason"] == "length", short for bad in [ - chat(enable_thinking=True), chat(stop=["}"]), chat( { diff --git a/tests/test_tool_call_parser.cpp b/tests/test_tool_call_parser.cpp index 06c47f9b40..970c178155 100644 --- a/tests/test_tool_call_parser.cpp +++ b/tests/test_tool_call_parser.cpp @@ -1,4 +1,5 @@ #include "models/qwen3_5/frontend/tool_call_parser.h" +#include "text/structured_output.h" #include @@ -735,8 +736,62 @@ int test_incremental_embedded_parameter_markup() { } // namespace +int test_constrained_tool_envelope() { + const std::vector tools{ + tool_definition("weather", Json{{"city", Json{{"type", "string"}}}, + {"units", Json{{"type", "string"}}}}), + tool_definition("clock", Json::object()), + R"({"type":"function","function":{"name":"run_code","parameters":{"type":"object","properties":{"language":{"type":"string"},"code":{"type":"string"}}}}})"}; + auto contract = fi::build_tool_call_output_contract(tools, true); + std::vector vocab(257); + for (int i = 0; i < 256; ++i) { vocab[i] = std::string(1, static_cast(i)); } + ninfer::text::StructuredCompiler compiler(vocab, {256}); + const auto format = fi::structured_tool_call_format(*contract); + auto grammar = compiler.compile( + {ninfer::StructuredOutputKind::JsonSchema, + R"({"type":"object","properties":{"location":{"type":"string"},"temperature":{"type":"number","minimum":-100,"maximum":100}},"required":["location","temperature"],"additionalProperties":false})"}, + {.reasoning_close = "", .alternative_format = format}); + const auto accepts = [&](const std::string& value) { + auto trial = grammar->fork(); + std::vector tokens; + for (unsigned char ch : value) { tokens.push_back(ch); } + tokens.push_back(256); + try { + trial->accept(tokens); + return true; + } catch (const std::logic_error&) { return false; } + }; + int failures = check(accepts("plan{\"location\":\"Tokyo\",\"temperature\":28}"), + "reasoning + final schema rejected with multiple tools enabled"); + failures += check(!accepts("plan{\"location\":\"Tokyo\",\"temperature\":101}"), + "tool alternative weakened final schema"); + failures += check(accepts("" + "python" + "print(1)"), + "tool grammar reordered prompt-declared parameters"); + const std::string call = "\n\nTokyo\n" + "celsius\n\n"; + const std::string parallel = call + "\n"; + failures += check(accepts("plan\n" + parallel), "parallel native tools rejected"); + const auto parsed = fi::parse_qwen_tool_call_output(parallel, 128, *contract); + failures += check(parsed.is_tool_call_response && parsed.tool_calls.size() == 2 && + parsed.content.empty(), + "constrained native calls did not round-trip"); + failures += check(!accepts("" + call + " prose"), "tool suffix allowed prose"); + failures += check(!accepts(""), + "undeclared tool accepted"); + failures += check(!accepts("x" + "y"), + "duplicate parameter accepted"); + failures += + check(!accepts(""), "incomplete tool can stop"); + failures += check(!accepts("not JSON"), "reasoning alternative escaped constraints"); + return failures; +} + int main() { int failures = 0; + failures += test_constrained_tool_envelope(); failures += test_basic_legacy_parsing(); failures += test_multiple_calls(); failures += test_declared_strings_preserve_text(); diff --git a/tests/text/test_structured_output.cpp b/tests/text/test_structured_output.cpp index c2eb024ea5..fc5d061325 100644 --- a/tests/text/test_structured_output.cpp +++ b/tests/text/test_structured_output.cpp @@ -94,12 +94,91 @@ int main() { require(masks[drafts.size() * words + 8] & 1U, "bonus mask missing EOS"); object->fill_masks(before, {}); require(before[0] == masks[0], "draft traversal advanced committed grammar"); + const auto tokens_for = [](std::string_view value) { + std::vector tokens; + for (unsigned char ch : value) { tokens.push_back(ch); } + return tokens; + }; + const auto accepts_value = [&](const std::shared_ptr& grammar, + const std::string& value) { + auto trial = grammar->fork(); + auto tokens = tokens_for(value); + tokens.push_back(256); + try { + trial->accept(tokens); + return true; + } catch (const std::logic_error&) { return false; } + }; + auto rating = compiler.compile( + {StructuredOutputKind::JsonSchema, R"({"type":"number","minimum":0,"maximum":10})"}); + require(accepts_value(rating, "0") && accepts_value(rating, "10") && + accepts_value(rating, "9.5"), + "inclusive number bounds lost valid values"); + for (int i = -80; i <= 240; ++i) { + const double value = i / 16.0; + const auto spelling = nlohmann::json(value).dump(); + if (accepts_value(rating, spelling)) { + const double decoded = nlohmann::json::parse(spelling).get(); + require(decoded >= 0 && decoded <= 10, "number escaped independent range oracle"); + } + } + auto exclusive = compiler.compile( + {StructuredOutputKind::JsonSchema, + R"({"type":"number","exclusiveMinimum":-0.25,"exclusiveMaximum":0.25})"}); + require(accepts_value(exclusive, "0") && accepts_value(exclusive, "0.249999") && + !accepts_value(exclusive, "0.25") && !accepts_value(exclusive, "-0.25") && + !accepts_value(exclusive, "1e2"), + "exclusive number bounds weakened"); + auto integer = + compiler.compile({StructuredOutputKind::JsonSchema, + R"({"type":"integer","minimum":-2,"exclusiveMaximum":3})"}); + for (int i = -10; i <= 10; ++i) { + require(accepts_value(integer, std::to_string(i)) == (i >= -2 && i < 3), + "integer range disagrees with independent oracle"); + } + require(!accepts_value(integer, "1.5"), "integer accepted a fraction"); + auto tiny = + compiler.compile({StructuredOutputKind::JsonSchema, + R"({"type":"number","minimum":0.000001,"maximum":0.000002})"}); + require(accepts_value(tiny, "0.000001") && accepts_value(tiny, "0.000002") && + !accepts_value(tiny, "0") && !accepts_value(tiny, "0.000003"), + "number precision boundary escaped bounds"); + + auto reasoning = compiler.compile({StructuredOutputKind::JsonObject, {}}, + {.reasoning_close = ""}); + reasoning->accept(tokens_for("reasoning with << overlap {\"ok\":true}"); + std::vector crossing_masks(words * (crossing.size() + 1)); + reasoning->fill_masks(crossing_masks, crossing); + for (std::size_t i = 0; i < crossing.size(); ++i) { + require(crossing_masks[i * words + crossing[i] / 32] & (1U << (crossing[i] % 32)), + "reasoning-to-JSON speculative transition masked a valid token"); + } + require(!(crossing_masks[3 * words + 'p' / 32] & (1U << ('p' % 32))), + "reasoning grammar escaped into final prose"); + require(crossing_masks[crossing.size() * words + 8] & 1U, "final bonus EOS missing"); + require(accepts_value(reasoning, "nk>{}"), "draft masks advanced committed phase"); + require(accepts_value(reasoning, "nk>\n\n{}"), "canonical close whitespace rejected"); + require(!accepts_value(reasoning, "nk>plain text{}"), + "a second reasoning close bypassed the first boundary"); + auto mixed_vocab = vocab; + mixed_vocab.push_back("{"); + mixed_vocab.push_back("prose"); + StructuredCompiler mixed_compiler(mixed_vocab, {256}); + auto mixed = mixed_compiler.compile({StructuredOutputKind::JsonObject, {}}, + {.reasoning_close = ""}); + mixed->fill_masks(before, {}); + require((before[8] & 2U) && !(before[8] & 4U), + "single token spanning reasoning and content escaped grammar"); + mixed->accept(std::vector{257, '}', 256}); for (const char* bad : {R"({"type":"array","uniqueItems":true})", R"({"$ref":"#/$defs/a~1b","$defs":{"a/b":{"const":1},"a~1b":{"const":2}}})", R"({"oneOf":[{},{}]})", R"({"$ref":"https://example.org/schema"})", R"({"const":1,"type":"string"})", R"({"anyOf":[{}],"type":"object"})", - R"({"type":"integer","minimum":0})"}) { + R"({"type":"integer","minimum":3,"maximum":2})", R"({"type":"number","minimum":"0"})", + R"({"type":"number","minimum":0.0000001,"maximum":0.0000002})", R"({"minimum":0})", + R"({"type":"number","minimum":1e30})"}) { bool failed = false; try { compiler.compile({StructuredOutputKind::JsonSchema, bad}); diff --git a/third_party/xgrammar/NINFER.md b/third_party/xgrammar/NINFER.md index 53971a0c12..fc6144c9d4 100644 --- a/third_party/xgrammar/NINFER.md +++ b/third_party/xgrammar/NINFER.md @@ -9,3 +9,7 @@ Regression: tests/text/test_structured_output.cpp. The additional-property exclu declared properties before divergence, preventing duplicate-key overwrites from bypassing a property schema. This restricts some otherwise valid key spellings. Other upstream sources are unmodified. + +Numeric-range correctness patch: GenerateNumber rejects the empty-range regex sentinel instead of +interpreting it as an empty JSON number. This matters when the six-decimal generation grid contains +no value inside an otherwise nonempty real interval. Covered by the structured-output grammar tests. diff --git a/third_party/xgrammar/cpp/json_schema_converter.cc b/third_party/xgrammar/cpp/json_schema_converter.cc index c65529ee0a..ffd7df99f3 100644 --- a/third_party/xgrammar/cpp/json_schema_converter.cc +++ b/third_party/xgrammar/cpp/json_schema_converter.cc @@ -2361,8 +2361,13 @@ int32_t JSONSchemaConverter::GenerateNumber(const NumberSpec& spec, const std::s exclusive_end = true; } if (start.has_value() || end.has_value()) { + const auto regex = + GenerateFloatRangeRegex(start, end, /*precision=*/6, exclusive_start, exclusive_end); + if (regex == "^()$") { + throw std::invalid_argument("numeric range has no value at six decimal places"); + } return RegexExpression( - GenerateFloatRangeRegex(start, end, /*precision=*/6, exclusive_start, exclusive_end), + regex, false, /*force_cfg_expansion=*/true );