ViewRegistry.Key allows multiple factory types per rendering type

Paves the way for parallel Classic and Compose implementations of wrappers like `NamedScreen`, to fix #546.

workflow-ui/compose/src/main/java/com/squareup/workflow1/ui/compose/ComposeScreenViewFactory.kt

@@ -11,6 +11,7 @@ import com.squareup.workflow1.ui.Screen
 import com.squareup.workflow1.ui.ScreenViewFactory
 import com.squareup.workflow1.ui.ScreenViewHolder
 import com.squareup.workflow1.ui.ViewEnvironment
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
 import kotlin.reflect.KClass
 
@@ -113,6 +114,10 @@ internal fun <RenderingT : Screen> composeScreenViewFactory(
 @WorkflowUiExperimentalApi
 public abstract class ComposeScreenViewFactory<RenderingT : Screen> :
   ScreenViewFactory<RenderingT> {
+
+  final override val key: Key<RenderingT, ScreenViewFactory<*>>
+    get() = Key(type, ComposeScreenViewFactory::class)
+
   /**
    * The composable content of this [ScreenViewFactory]. This method will be called
    * any time [rendering] or [viewEnvironment] change. It is the Compose-based analogue of

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/ScreenViewFactory.kt

@@ -9,6 +9,8 @@ import androidx.viewbinding.ViewBinding
 import com.squareup.workflow1.ui.ScreenViewFactory.Companion.fromCode
 import com.squareup.workflow1.ui.ScreenViewFactory.Companion.fromLayout
 import com.squareup.workflow1.ui.ScreenViewFactory.Companion.fromViewBinding
+import com.squareup.workflow1.ui.ViewRegistry.Key
+import kotlin.reflect.KClass
 
 @WorkflowUiExperimentalApi
 public typealias ViewBindingInflater<BindingT> = (LayoutInflater, ViewGroup?, Boolean) -> BindingT
@@ -48,6 +50,10 @@ public typealias ViewBindingInflater<BindingT> = (LayoutInflater, ViewGroup?, Bo
  */
 @WorkflowUiExperimentalApi
 public interface ScreenViewFactory<in ScreenT : Screen> : ViewRegistry.Entry<ScreenT> {
+  public val type: KClass<in ScreenT>
+
+  override val key: Key<ScreenT, ScreenViewFactory<*>> get() = Key(type, ScreenViewFactory::class)
+
   /**
    * It is rare to call this method directly. Instead the most common path is to pass [Screen]
    * instances to [WorkflowViewStub.show], which will apply the [ScreenViewFactory] machinery for

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/ScreenViewFactoryFinder.kt

@@ -1,6 +1,7 @@
 package com.squareup.workflow1.ui
 
 import com.squareup.workflow1.ui.ScreenViewFactory.Companion.forWrapper
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import com.squareup.workflow1.ui.container.BackStackScreen
 import com.squareup.workflow1.ui.container.BackStackScreenViewFactory
 import com.squareup.workflow1.ui.container.BodyAndOverlaysContainer
@@ -54,7 +55,8 @@ public interface ScreenViewFactoryFinder {
     environment: ViewEnvironment,
     rendering: ScreenT
   ): ScreenViewFactory<ScreenT> {
-    val entry = environment[ViewRegistry].getEntryFor(rendering::class)
+    val entry = environment[ViewRegistry]
+      .getEntryFor(Key(rendering::class, ScreenViewFactory::class))
 
     @Suppress("UNCHECKED_CAST")
     return (entry as? ScreenViewFactory<ScreenT>)

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/container/OverlayDialogFactory.kt

@@ -4,7 +4,9 @@ import android.app.Dialog
 import android.content.Context
 import com.squareup.workflow1.ui.ViewEnvironment
 import com.squareup.workflow1.ui.ViewRegistry
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
+import kotlin.reflect.KClass
 
 /**
  * Factory for [Dialog] instances that can show renderings of type [OverlayT].
@@ -64,6 +66,10 @@ import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
  */
 @WorkflowUiExperimentalApi
 public interface OverlayDialogFactory<OverlayT : Overlay> : ViewRegistry.Entry<OverlayT> {
+  public val type: KClass<in OverlayT>
+
+  override val key: Key<OverlayT, *> get() = Key(type, OverlayDialogFactory::class)
+
   /** Builds a [Dialog], but does not show it. */
   public fun buildDialog(
     initialRendering: OverlayT,

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/container/OverlayDialogFactoryFinder.kt

@@ -4,6 +4,7 @@ import com.squareup.workflow1.ui.Screen
 import com.squareup.workflow1.ui.ViewEnvironment
 import com.squareup.workflow1.ui.ViewEnvironmentKey
 import com.squareup.workflow1.ui.ViewRegistry
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
 
 /**
@@ -17,7 +18,8 @@ public interface OverlayDialogFactoryFinder {
     environment: ViewEnvironment,
     rendering: OverlayT
   ): OverlayDialogFactory<OverlayT> {
-    val entry = environment[ViewRegistry].getEntryFor(rendering::class)
+    val entry = environment[ViewRegistry]
+      .getEntryFor(Key(rendering::class, OverlayDialogFactory::class))
 
     @Suppress("UNCHECKED_CAST")
     return entry as? OverlayDialogFactory<OverlayT>

workflow-ui/core-common/src/main/java/com/squareup/workflow1/ui/CompositeViewRegistry.kt

@@ -1,7 +1,7 @@
 package com.squareup.workflow1.ui
 
 import com.squareup.workflow1.ui.ViewRegistry.Entry
-import kotlin.reflect.KClass
+import com.squareup.workflow1.ui.ViewRegistry.Key
 
 /**
  * A [ViewRegistry] that contains only other registries and delegates to their [getEntryFor]
@@ -20,26 +20,26 @@ import kotlin.reflect.KClass
  */
 @WorkflowUiExperimentalApi
 internal class CompositeViewRegistry private constructor(
-  private val registriesByKey: Map<KClass<*>, ViewRegistry>
+  private val registriesByKey: Map<Key<*, *>, ViewRegistry>
 ) : ViewRegistry {
 
   constructor (vararg registries: ViewRegistry) : this(mergeRegistries(*registries))
 
-  override val keys: Set<KClass<*>> get() = registriesByKey.keys
+  override val keys: Set<Key<*, *>> get() = registriesByKey.keys
 
-  override fun <RenderingT : Any> getEntryFor(
-    renderingType: KClass<out RenderingT>
-  ): Entry<RenderingT>? = registriesByKey[renderingType]?.getEntryFor(renderingType)
+  override fun <RenderingT : Any, FactoryT : Any> getEntryFor(
+    key: Key<RenderingT, FactoryT>
+  ): Entry<RenderingT>? = registriesByKey[key]?.getEntryFor(key)
 
   override fun toString(): String {
     return "CompositeViewRegistry(${registriesByKey.values.toSet().map { it.toString() }})"
   }
 
   companion object {
-    private fun mergeRegistries(vararg registries: ViewRegistry): Map<KClass<*>, ViewRegistry> {
-      val registriesByKey = mutableMapOf<KClass<*>, ViewRegistry>()
+    private fun mergeRegistries(vararg registries: ViewRegistry): Map<Key<*, *>, ViewRegistry> {
+      val registriesByKey = mutableMapOf<Key<*, *>, ViewRegistry>()
 
-      fun putAllUnique(other: Map<KClass<*>, ViewRegistry>) {
+      fun putAllUnique(other: Map<Key<*, *>, ViewRegistry>) {
         val duplicateKeys = registriesByKey.keys.intersect(other.keys)
         require(duplicateKeys.isEmpty()) {
           "Must not have duplicate entries: $duplicateKeys. Use merge to replace existing entries."

workflow-ui/core-common/src/main/java/com/squareup/workflow1/ui/TypedViewRegistry.kt

@@ -1,6 +1,7 @@
 package com.squareup.workflow1.ui
 
 import com.squareup.workflow1.ui.ViewRegistry.Entry
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import kotlin.reflect.KClass
 
 /**
@@ -9,29 +10,29 @@ import kotlin.reflect.KClass
  */
 @WorkflowUiExperimentalApi
 internal class TypedViewRegistry private constructor(
-  private val bindings: Map<KClass<*>, Entry<*>>
+  private val bindings: Map<Key<*, *>, Entry<*>>
 ) : ViewRegistry {
 
   constructor(vararg bindings: Entry<*>) : this(
-    bindings.associateBy { it.type }
+    bindings.associateBy { it.key }
       .apply {
         check(keys.size == bindings.size) {
-          "${bindings.map { it.type }} must not have duplicate entries."
+          "${bindings.map { it.key }} must not have duplicate entries."
         }
-      } as Map<KClass<*>, Entry<*>>
+      } as Map<Key<*, *>, Entry<*>>
   )
 
-  override val keys: Set<KClass<*>> get() = bindings.keys
+  override val keys: Set<Key<*, *>> get() = bindings.keys
 
-  override fun <RenderingT : Any> getEntryFor(
-    renderingType: KClass<out RenderingT>
+  override fun <RenderingT : Any, FactoryT : Any> getEntryFor(
+    key: Key<RenderingT, FactoryT>
   ): Entry<RenderingT>? {
     @Suppress("UNCHECKED_CAST")
-    return bindings[renderingType] as? Entry<RenderingT>
+    return bindings[key] as? Entry<RenderingT>
   }
 
   override fun toString(): String {
-    val map = bindings.map { "${it.key.simpleName}=${it.value::class.qualifiedName}" }
+    val map = bindings.map { "${it.key}=${it.value::class.qualifiedName}" }
     return "TypedViewRegistry(bindings=$map)"
   }
 }

workflow-ui/core-common/src/main/java/com/squareup/workflow1/ui/ViewRegistry.kt

@@ -1,6 +1,7 @@
 package com.squareup.workflow1.ui
 
 import com.squareup.workflow1.ui.ViewRegistry.Entry
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import kotlin.reflect.KClass
 
 /**
@@ -61,8 +62,33 @@ import kotlin.reflect.KClass
  */
 @WorkflowUiExperimentalApi
 public interface ViewRegistry {
+  public class Key<in RenderingT : Any, out FactoryT : Any>(
+    public val renderingType: KClass<in RenderingT>,
+    public val factoryType: KClass<out FactoryT>
+  ) {
+    override fun equals(other: Any?): Boolean {
+      if (this === other) return true
+      if (javaClass != other?.javaClass) return false
+
+      other as Key<*, *>
+
+      if (renderingType != other.renderingType) return false
+      return factoryType == other.factoryType
+    }
+
+    override fun hashCode(): Int {
+      var result = renderingType.hashCode()
+      result = 31 * result + factoryType.hashCode()
+      return result
+    }
+
+    override fun toString(): String {
+      return "Key(renderingType=$renderingType, factoryType=$factoryType)"
+    }
+  }
+
   public interface Entry<in RenderingT : Any> {
-    public val type: KClass<in RenderingT>
+    public val key: Key<RenderingT, *>
   }
 
   /**
@@ -71,14 +97,14 @@ public interface ViewRegistry {
    *
    * Used to ensure that duplicate bindings are never registered.
    */
-  public val keys: Set<KClass<*>>
+  public val keys: Set<Key<*, *>>
 
   /**
-   * Returns the [Entry] that was registered for the given [renderingType], or null
+   * Returns the [Entry] that was registered for the given [key], or null
    * if none was found.
    */
-  public fun <RenderingT : Any> getEntryFor(
-    renderingType: KClass<out RenderingT>
+  public fun <RenderingT : Any, FactoryT : Any> getEntryFor(
+    key: Key<RenderingT, FactoryT>
   ): Entry<RenderingT>?
 
   public companion object : ViewEnvironmentKey<ViewRegistry>() {
@@ -90,9 +116,10 @@ public interface ViewRegistry {
   }
 }
 
-@WorkflowUiExperimentalApi public inline operator fun <reified RenderingT : Any> ViewRegistry.get(
-  renderingType: KClass<out RenderingT>
-): Entry<RenderingT>? = getEntryFor(renderingType)
+@WorkflowUiExperimentalApi
+public inline operator fun <reified RenderingT : Any, reified FactoryT : Any> ViewRegistry.get(
+  key: Key<RenderingT, FactoryT>
+): Entry<RenderingT>? = getEntryFor(key)
 
 @WorkflowUiExperimentalApi
 public fun ViewRegistry(vararg bindings: Entry<*>): ViewRegistry =
@@ -115,8 +142,10 @@ public operator fun ViewRegistry.plus(entry: Entry<*>): ViewRegistry =
   this + ViewRegistry(entry)
 
 /**
- * Transforms the receiver to add all entries from [other], throwing [IllegalArgumentException]
- * if the receiver already has any matching [entry]. Use [merge] to replace existing entries.
+ * Transforms the receiver to add all entries from [other].
+ *
+ * @throws [IllegalArgumentException] if the receiver already has an matching [Entry].
+ * Use [merge] to replace existing entries instead.
  */
 @WorkflowUiExperimentalApi
 public operator fun ViewRegistry.plus(other: ViewRegistry): ViewRegistry {
@@ -125,6 +154,11 @@ public operator fun ViewRegistry.plus(other: ViewRegistry): ViewRegistry {
   return CompositeViewRegistry(this, other)
 }
 
+/**
+ * Returns a new [ViewEnvironment] that adds [registry] to the receiver.
+ * If the receiver already has a [ViewRegistry], [ViewEnvironmentKey.combine]
+ * is applied as usual to [merge] its entries.
+ */
 @WorkflowUiExperimentalApi
 public operator fun ViewEnvironment.plus(registry: ViewRegistry): ViewEnvironment {
   if (this[ViewRegistry] === registry) return this

workflow-ui/core-common/src/test/java/com/squareup/workflow1/ui/CompositeViewRegistryTest.kt

@@ -2,6 +2,7 @@ package com.squareup.workflow1.ui
 
 import com.google.common.truth.Truth.assertThat
 import com.squareup.workflow1.ui.ViewRegistry.Entry
+import com.squareup.workflow1.ui.ViewRegistry.Key
 import org.junit.Test
 import kotlin.reflect.KClass
 import kotlin.test.assertFailsWith
@@ -28,26 +29,26 @@ internal class CompositeViewRegistryTest {
     val bazFactory = TestEntry(BazRendering::class)
     val fooBarRegistry = TestRegistry(
       mapOf(
-        FooRendering::class to fooFactory,
-        BarRendering::class to barFactory
+        fooFactory.key to fooFactory,
+        barFactory.key to barFactory
       )
     )
-    val bazRegistry = TestRegistry(factories = mapOf(BazRendering::class to bazFactory))
+    val bazRegistry = TestRegistry(factories = mapOf(bazFactory.key to bazFactory))
     val registry = fooBarRegistry + bazRegistry
 
-    assertThat(registry.getEntryFor(FooRendering::class))
+    assertThat(registry.getEntryFor(Key(FooRendering::class, TestEntry::class)))
       .isSameInstanceAs(fooFactory)
-    assertThat(registry.getEntryFor(BarRendering::class))
+    assertThat(registry.getEntryFor(Key(BarRendering::class, TestEntry::class)))
       .isSameInstanceAs(barFactory)
-    assertThat(registry.getEntryFor(BazRendering::class))
+    assertThat(registry.getEntryFor(Key(BazRendering::class, TestEntry::class)))
       .isSameInstanceAs(bazFactory)
   }
 
   @Test fun `getFactoryFor returns null on missing registry`() {
     val fooRegistry = TestRegistry(setOf(FooRendering::class))
     val registry = CompositeViewRegistry(ViewRegistry(), fooRegistry)
 
-    assertThat(registry.getEntryFor(BarRendering::class)).isNull()
+    assertThat(registry.getEntryFor(Key(BarRendering::class, TestEntry::class))).isNull()
   }
 
   @Test fun `keys includes all composite registries' keys`() {
@@ -56,28 +57,33 @@ internal class CompositeViewRegistryTest {
     val registry = CompositeViewRegistry(fooBarRegistry, bazRegistry)
 
     assertThat(registry.keys).containsExactly(
-      FooRendering::class,
-      BarRendering::class,
-      BazRendering::class
+      Key(FooRendering::class, TestEntry::class),
+      Key(BarRendering::class, TestEntry::class),
+      Key(BazRendering::class, TestEntry::class)
     )
   }
 
-  private class TestEntry<T : Any>(
-    override val type: KClass<in T>
-  ) : Entry<T>
+  private class TestEntry<T : Any>(type: KClass<in T>) : Entry<T> {
+    override val key = Key(type, TestEntry::class)
+  }
 
   private object FooRendering
   private object BarRendering
   private object BazRendering
 
-  private class TestRegistry(private val factories: Map<KClass<*>, Entry<*>>) : ViewRegistry {
-    constructor(keys: Set<KClass<*>>) : this(keys.associateWith { TestEntry(it) })
+  private class TestRegistry(private val factories: Map<Key<*, *>, Entry<*>>) : ViewRegistry {
+    constructor(keys: Set<KClass<*>>) : this(
+      keys.associate {
+        val entry = TestEntry(it)
+        entry.key to entry
+      }
+    )
 
-    override val keys: Set<KClass<*>> get() = factories.keys
+    override val keys: Set<Key<*, *>> get() = factories.keys
 
     @Suppress("UNCHECKED_CAST")
-    override fun <RenderingT : Any> getEntryFor(
-      renderingType: KClass<out RenderingT>
-    ): Entry<RenderingT> = factories.getValue(renderingType) as Entry<RenderingT>
+    override fun <RenderingT : Any, FactoryT : Any> getEntryFor(
+      key: Key<RenderingT, FactoryT>
+    ): Entry<RenderingT> = factories.getValue(key) as Entry<RenderingT>
   }
 }