From 3b3ad949f5b7cf02b5fe418d0f30cadd76a90e39 Mon Sep 17 00:00:00 2001
From: ochafik <ochafik@google.com>
Date: Fri, 15 Mar 2024 00:52:36 +0000
Subject: [PATCH] json: fix top-level $refs

---
 examples/json-schema-to-grammar.py            |  2 +-
 examples/server/json-schema-to-grammar.cpp    |  2 +-
 .../server/public/json-schema-to-grammar.mjs  |  2 +-
 tests/test-json-schema-to-grammar.cpp         | 66 ++++++++++++++++++-
 4 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/examples/json-schema-to-grammar.py b/examples/json-schema-to-grammar.py
index 0c9b54754..cdbee7d39 100755
--- a/examples/json-schema-to-grammar.py
+++ b/examples/json-schema-to-grammar.py
@@ -303,7 +303,7 @@ class SchemaConverter:
         rule_name = name or 'root'
 
         if (ref := schema.get('$ref')) is not None:
-            return self._resolve_ref(ref)
+            return self._add_rule(rule_name, self._resolve_ref(ref))
 
         elif 'oneOf' in schema or 'anyOf' in schema:
             return self._add_rule(rule_name, self._generate_union_rule(name, schema.get('oneOf') or schema['anyOf']))
diff --git a/examples/server/json-schema-to-grammar.cpp b/examples/server/json-schema-to-grammar.cpp
index 919e26c35..054f49a5b 100644
--- a/examples/server/json-schema-to-grammar.cpp
+++ b/examples/server/json-schema-to-grammar.cpp
@@ -526,7 +526,7 @@ public:
         string rule_name = name.empty() ? "root" : name;
 
         if (schema.contains("$ref")) {
-            return _resolve_ref(schema["$ref"]);
+            return _add_rule(rule_name, _resolve_ref(schema["$ref"]));
         } else if (schema.contains("oneOf") || schema.contains("anyOf")) {
             vector<json> alt_schemas = schema.contains("oneOf") ? schema["oneOf"].get<vector<json>>() : schema["anyOf"].get<vector<json>>();
             return _add_rule(rule_name, _generate_union_rule(name, alt_schemas));
diff --git a/examples/server/public/json-schema-to-grammar.mjs b/examples/server/public/json-schema-to-grammar.mjs
index 4f921b688..862a9079f 100644
--- a/examples/server/public/json-schema-to-grammar.mjs
+++ b/examples/server/public/json-schema-to-grammar.mjs
@@ -330,7 +330,7 @@ export class SchemaConverter {
 
     const ref = schema.$ref;
     if (ref !== undefined) {
-      return this._resolveRef(ref);
+      return this._addRule(ruleName, this._resolveRef(ref));
     } else if (schema.oneOf || schema.anyOf) {
       return this._addRule(ruleName, this._generateUnionRule(name, schema.oneOf || schema.anyOf));
     } else if (Array.isArray(schemaType)) {
diff --git a/tests/test-json-schema-to-grammar.cpp b/tests/test-json-schema-to-grammar.cpp
index 557234c6c..e3819c112 100755
--- a/tests/test-json-schema-to-grammar.cpp
+++ b/tests/test-json-schema-to-grammar.cpp
@@ -84,7 +84,7 @@ int main() {
   run_all({
     .name = "exotic formats",
     .schema = R"""({
-      "prefixItems": [
+      "items": [
         { "format": "date" },
         { "format": "uuid" },
         { "format": "time" },
@@ -140,6 +140,37 @@ int main() {
     )"""
   });
 
+  run_all({
+    .name = "tuple1",
+    .schema = R"""({
+      "prefixItems": [{ "type": "string" }]
+    })""",
+    .expected = R"""(
+      root ::= "[" space string "]" space
+      space ::= " "?
+      string ::=  "\"" (
+              [^"\\] |
+              "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
+            )* "\"" space
+    )"""
+  });
+
+  run_all({
+    .name = "tuple2",
+    .schema = R"""({
+      "prefixItems": [{ "type": "string" }, { "type": "number" }]
+    })""",
+    .expected = R"""(
+      number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? space
+      root ::= "[" space string "," space number "]" space
+      space ::= " "?
+      string ::=  "\"" (
+              [^"\\] |
+              "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
+            )* "\"" space
+    )"""
+  });
+
   run_all({
     .name = "number",
     .schema = R"""({
@@ -298,4 +329,37 @@ int main() {
             )* "\"" space
     )"""
   });
+  
+  run_all({
+    .name = "top-level $ref",
+    .schema = R"""({
+      "$schema": "http://json-schema.org/draft-07/schema#",
+      "$ref": "#/definitions/MyType",
+      "definitions": {
+        "MyType": {
+          "type": "object",
+          "properties": {
+            "a": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "a"
+          ],
+          "additionalProperties": false
+        }
+      }
+    })""",
+    .expected = R"""(
+      MyType ::= "{" space MyType-a-kv "}" space
+      MyType-a-kv ::= "\"a\"" space ":" space string
+      root ::= MyType
+      space ::= " "?
+      string ::=  "\"" (
+              [^"\\] |
+              "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
+            )* "\"" space
+    )"""
+  });
+
 }