From 284ba87aa05399848ee24dcbf625c84c1918834a Mon Sep 17 00:00:00 2001
From: David Morgan <davidmorgan@google.com>
Date: Wed, 7 Aug 2024 10:42:57 +0200
Subject: [PATCH] Fix unions schema to match JSON.

---
 schemas/dart_model.schema.json                | 20 +++--
 schemas/macro_service.schema.json             | 86 +++++++++++++++----
 .../lib/generate_dart_model.dart              | 12 ++-
 3 files changed, 92 insertions(+), 26 deletions(-)

diff --git a/schemas/dart_model.schema.json b/schemas/dart_model.schema.json
index e96229fb..dc18e5f9 100644
--- a/schemas/dart_model.schema.json
+++ b/schemas/dart_model.schema.json
@@ -1,7 +1,9 @@
 {
   "$schema": "https://json-schema.org/draft/2020-12/schema",
   "oneOf": [
-    {"$ref": "#/$defs/Model"}
+    {
+      "$ref": "#/$defs/Model"
+    }
   ],
   "$defs": {
     "Augmentation": {
@@ -31,12 +33,16 @@
         "metadataAnnotations": {
           "description": "The metadata annotations attached to this iterface.",
           "type": "array",
-          "items": {"$ref": "#/$defs/MetadataAnnotation"}
+          "items": {
+            "$ref": "#/$defs/MetadataAnnotation"
+          }
         },
         "members": {
           "type": "object",
           "description": "Map of members by name.",
-          "additionalProperties": {"$ref": "#/$defs/Member"}
+          "additionalProperties": {
+            "$ref": "#/$defs/Member"
+          }
         },
         "properties": {
           "$comment": "The properties of this interface.",
@@ -51,7 +57,9 @@
         "scopes": {
           "type": "object",
           "description": "Scopes by name.",
-          "additionalProperties": {"$ref": "#/$defs/Interface"}
+          "additionalProperties": {
+            "$ref": "#/$defs/Interface"
+          }
         }
       }
     },
@@ -72,7 +80,9 @@
         "uris": {
           "type": "object",
           "description": "Libraries by URI.",
-          "additionalProperties": {"$ref": "#/$defs/Library"}
+          "additionalProperties": {
+            "$ref": "#/$defs/Library"
+          }
         }
       }
     },
diff --git a/schemas/macro_service.schema.json b/schemas/macro_service.schema.json
index 4283b967..f7188bd2 100644
--- a/schemas/macro_service.schema.json
+++ b/schemas/macro_service.schema.json
@@ -1,8 +1,15 @@
 {
   "$schema": "https://json-schema.org/draft/2020-12/schema",
-  "oneOf": [
-    {"$ref": "#/$defs/AugmentRequest"},
-    {"$ref": "#/$defs/AugmentResponse"}
+  "anyOf": [
+    {
+      "$ref": "#/$defs/HostRequest"
+    },
+    {
+      "$ref": "#/$defs/MacroRequest"
+    },
+    {
+      "$ref": "#/$defs/Response"
+    }
   ],
   "$defs": {
     "AugmentRequest": {
@@ -22,7 +29,9 @@
         "augmentations": {
           "description": "The augmentations.",
           "type": "array",
-          "items": {"$ref": "file:dart_model.schema.json#/$defs/Augmentation"}
+          "items": {
+            "$ref": "file:dart_model.schema.json#/$defs/Augmentation"
+          }
         }
       }
     },
@@ -38,9 +47,18 @@
     },
     "HostRequest": {
       "$description": "A request sent from host to macro.",
-      "oneOf": [
-        {"$ref": "#/$defs/AugmentRequest"}
-      ]
+      "properties": {
+        "type": {
+          "type": "string"
+        },
+        "value": {
+          "oneOf": [
+            {
+              "$ref": "#/$defs/AugmentRequest"
+            }
+          ]
+        }
+      }
     },
     "MacroDescription": {
       "type": "object",
@@ -59,7 +77,9 @@
       "type": "object",
       "description": "Informs the host that a macro has started.",
       "properties": {
-        "macroDescription": {"$ref": "#/$defs/MacroDescription"}
+        "macroDescription": {
+          "$ref": "#/$defs/MacroDescription"
+        }
       }
     },
     "MacroStartedResponse": {
@@ -68,32 +88,60 @@
     },
     "MacroRequest": {
       "$description": "A request sent from macro to host.",
-      "oneOf": [
-        {"$ref": "#/$defs/MacroStartedRequest"},
-        {"$ref": "#/$defs/QueryRequest"}
-      ]
+      "properties": {
+        "type": {
+          "type": "string"
+        },
+        "value": {
+          "oneOf": [
+            {
+              "$ref": "#/$defs/MacroStartedRequest"
+            },
+            {
+              "$ref": "#/$defs/QueryRequest"
+            }
+          ]
+        }
+      }
     },
     "QueryRequest": {
       "type": "object",
       "description": "Macro's query about the code it should augment.",
       "properties": {
-        "query": {"$ref": "file:dart_model.schema.json#/$defs/Query"}
+        "query": {
+          "$ref": "file:dart_model.schema.json#/$defs/Query"
+        }
       }
     },
     "QueryResponse": {
       "type": "object",
       "description": "Host's response to a [QueryRequest].",
       "properties": {
-        "model": {"$ref": "file:dart_model.schema.json#/$defs/Model"}
+        "model": {
+          "$ref": "file:dart_model.schema.json#/$defs/Model"
+        }
       }
     },
     "Response": {
       "$description": "A response to a [MacroRequest] or [HostRequest].",
-      "oneOf": [
-        {"$ref": "#/$defs/AugmentResponse"},
-        {"$ref": "#/$defs/MacroStartedResponse"},
-        {"$ref": "#/$defs/QueryResponse"}
-      ]
+      "properties": {
+        "type": {
+          "type": "string"
+        },
+        "value": {
+          "oneOf": [
+            {
+              "$ref": "#/$defs/AugmentResponse"
+            },
+            {
+              "$ref": "#/$defs/MacroStartedResponse"
+            },
+            {
+              "$ref": "#/$defs/QueryResponse"
+            }
+          ]
+        }
+      }
     }
   }
 }
diff --git a/tool/dart_model_generator/lib/generate_dart_model.dart b/tool/dart_model_generator/lib/generate_dart_model.dart
index 70cd116f..9ff69250 100644
--- a/tool/dart_model_generator/lib/generate_dart_model.dart
+++ b/tool/dart_model_generator/lib/generate_dart_model.dart
@@ -38,8 +38,8 @@ String generate(String schemaJson,
       refProvider: LocalRefProvider(dartModelJson ??
           File('schemas/dart_model.schema.json').readAsStringSync()));
   for (final def in schema.defs.entries) {
-    if (def.value.oneOf.isNotEmpty) {
-      result.add(_generateUnion(def.key, def.value.oneOf));
+    if (_isUnion(def.value)) {
+      result.add(_generateUnion(def.key, def.value.properties['value']!.oneOf));
     } else {
       result.add(_generateExtensionType(def.key, def.value));
     }
@@ -132,6 +132,14 @@ String _generateExtensionType(String name, JsonSchema definition) {
   return result.toString();
 }
 
+/// Whether [schema] represents a union type.
+///
+/// To be a union type it must have exactly two properties, "type" and "value",
+/// where "type" is a "string" and "value" is a "oneOf".
+bool _isUnion(JsonSchema schema) =>
+    schema.properties['type']?.schemaMap!['type'] == 'string' &&
+    schema.properties['value']?.oneOf != null;
+
 /// Generates a type called [name] that is a union of the specified [oneOf]
 /// types, which must all be `$ref`s to class definitions.
 ///