From daaf8f1ccbcae1064823edf22c3c403fa7aa11f4 Mon Sep 17 00:00:00 2001
From: Jake Wharton <jw@squareup.com>
Date: Fri, 19 Jul 2024 09:37:57 -0400
Subject: [PATCH] Deduplicate schema dependency loading

We need to support this in both JSON-based and FIR-based loading, so extract a common function for the logic.
---
 .../redwood/tooling/schema/schemaParser.kt    | 43 +++++++------------
 .../redwood/tooling/schema/schemaParserFir.kt | 19 +-------
 .../tooling/schema/SchemaParserTest.kt        |  2 +-
 3 files changed, 17 insertions(+), 47 deletions(-)

diff --git a/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParser.kt b/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParser.kt
index 039ede5b20..8e52246823 100644
--- a/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParser.kt
+++ b/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParser.kt
@@ -53,13 +53,19 @@ internal fun loadProtocolSchemaSet(
   classLoader: ClassLoader,
 ): ProtocolSchemaSet {
   val schema = loadProtocolSchema(type, classLoader)
-  val dependencies = schema.taggedDependencies.map { (tag, dependency) ->
-    loadProtocolSchema(dependency, classLoader, tag)
-  }
-  return ParsedProtocolSchemaSet(
-    schema = schema,
-    dependencies = dependencies.associateBy { it.type },
-  )
+  return loadProtocolSchemaDependencies(schema, classLoader)
+}
+
+internal fun loadProtocolSchemaDependencies(
+  schema: ProtocolSchema,
+  classLoader: ClassLoader,
+): ParsedProtocolSchemaSet {
+  val dependencies = schema.taggedDependencies.entries
+    .associate { (tag, type) ->
+      require(tag != 0) { "Dependency $type tag must be non-zero" }
+      type to loadProtocolSchema(type, classLoader, tag)
+    }
+  return ParsedProtocolSchemaSet(schema, dependencies)
 }
 
 internal fun loadProtocolSchema(
@@ -228,33 +234,14 @@ internal fun parseProtocolSchemaSet(schemaType: KClass<*>): ProtocolSchemaSet {
     )
   }
 
-  val dependencies = schemaAnnotation.dependencies
-    .associate {
-      val dependencyTag = it.tag
-      val dependencyType = it.schema.toFqType()
-      require(dependencyTag != 0) {
-        "Dependency $dependencyType tag must not be non-zero"
-      }
-
-      val schema = loadProtocolSchema(
-        type = dependencyType,
-        classLoader = schemaType.java.classLoader,
-        tag = dependencyTag,
-      )
-      dependencyTag to schema
-    }
-
   val schema = ParsedProtocolSchema(
     type = schemaType.toFqType(),
     scopes = scopes.toList(),
     widgets = widgets,
     modifiers = modifiers,
-    taggedDependencies = dependencies.mapValues { (_, schema) -> schema.type },
-  )
-  val schemaSet = ParsedProtocolSchemaSet(
-    schema,
-    dependencies.values.associateBy { it.type },
+    taggedDependencies = schemaAnnotation.dependencies.associate { it.tag to it.schema.toFqType() },
   )
+  val schemaSet = loadProtocolSchemaDependencies(schema, schemaType.java.classLoader)
 
   val duplicatedWidgets = schemaSet.all
     .flatMap { it.widgets.map { widget -> widget to it } }
diff --git a/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParserFir.kt b/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParserFir.kt
index 6d253120bd..079a176892 100644
--- a/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParserFir.kt
+++ b/redwood-tooling-schema/src/main/kotlin/app/cash/redwood/tooling/schema/schemaParserFir.kt
@@ -190,24 +190,7 @@ public fun parseProtocolSchema(
   disposable.dispose()
 
   val dependencyClassLoader = URLClassLoader(dependencies.map { it.toURI().toURL() }.toTypedArray())
-  val dependencySchemas = schema.taggedDependencies.entries
-    .associate { (dependencyTag, dependencyType) ->
-      require(dependencyTag != 0) {
-        "Dependency $dependencyType tag must not be non-zero"
-      }
-
-      val dependency = loadProtocolSchema(
-        type = dependencyType,
-        classLoader = dependencyClassLoader,
-        tag = dependencyTag,
-      )
-      dependencyTag to dependency
-    }
-
-  val schemaSet = ParsedProtocolSchemaSet(
-    schema,
-    dependencySchemas.values.associateBy { it.type },
-  )
+  val schemaSet = loadProtocolSchemaDependencies(schema, dependencyClassLoader)
 
   val duplicatedWidgets = schemaSet.all
     .flatMap { it.widgets.map { widget -> widget to it } }
diff --git a/redwood-tooling-schema/src/test/kotlin/app/cash/redwood/tooling/schema/SchemaParserTest.kt b/redwood-tooling-schema/src/test/kotlin/app/cash/redwood/tooling/schema/SchemaParserTest.kt
index ad6185c329..b3effc0047 100644
--- a/redwood-tooling-schema/src/test/kotlin/app/cash/redwood/tooling/schema/SchemaParserTest.kt
+++ b/redwood-tooling-schema/src/test/kotlin/app/cash/redwood/tooling/schema/SchemaParserTest.kt
@@ -947,7 +947,7 @@ class SchemaParserTest(
     assertFailure { parser.parse(SchemaDependencyTagZero::class) }
       .isInstanceOf<IllegalArgumentException>()
       .hasMessage(
-        "Dependency app.cash.redwood.layout.RedwoodLayout tag must not be non-zero",
+        "Dependency app.cash.redwood.layout.RedwoodLayout tag must be non-zero",
       )
   }